Skip to content

feat: add Ministral3 embedding distillation recipe#3058

Merged
rnyak merged 15 commits into
NVIDIA-NeMo:mainfrom
vinay-raman:vinayraman/distillation-clean
Jul 15, 2026
Merged

feat: add Ministral3 embedding distillation recipe#3058
rnyak merged 15 commits into
NVIDIA-NeMo:mainfrom
vinay-raman:vinayraman/distillation-clean

Conversation

@vinay-raman

@vinay-raman vinay-raman commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Support for knowledge distillation (KD) of embedding/retriever models

Changelog

  • Add EmbeddingDistillRecipe with RetrieverStudentWithProjection / RetrieverTeacherEmbeddingEncoder, distillation losses (cosine / MSE / InfoNCE-distill), intermediate-layer distillation
  • cross-tokenizer cached-teacher support
  • bi-encoder distill collator, example config, and unit tests.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

  • Related to # (issue)

Add EmbeddingDistillRecipe with RetrieverStudentWithProjection /
RetrieverTeacherEmbeddingEncoder, distillation losses (cosine / MSE /
InfoNCE-distill), intermediate-layer distillation, cross-tokenizer
cached-teacher support, bi-encoder distill collator, example config,
and unit tests.

Signed-off-by: Vinay Raman <viraman@nvidia.com>
@vinay-raman vinay-raman requested review from a team as code owners July 13, 2026 14:20
@copy-pr-bot

copy-pr-bot Bot commented Jul 13, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@rnyak

rnyak commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

/claude review

@rnyak rnyak added the enhancement New feature or request label Jul 13, 2026
Comment thread tests/unit_tests/models/biencoder/test_nemotron_embedding.py Outdated
Comment thread nemo_automodel/recipes/retrieval/distill_bi_encoder.py Outdated
Comment thread tests/unit_tests/models/biencoder/test_nemotron_embedding.py Outdated
Comment thread nemo_automodel/recipes/retrieval/distill_bi_encoder.py Outdated
Comment thread examples/retrieval/distillation/nemotron3_embed_1b_distill.yaml Outdated
Comment thread nemo_automodel/_transformers/retrieval.py
Comment thread nemo_automodel/recipes/retrieval/distill_bi_encoder.py
Comment thread nemo_automodel/recipes/retrieval/train_bi_encoder.py
Comment thread examples/retrieval/distillation/nemotron3_embed_1b_distill.yaml
Signed-off-by: Vinay Raman <viraman@nvidia.com>
Signed-off-by: Vinay Raman <viraman@nvidia.com>
Signed-off-by: Vinay Raman <viraman@nvidia.com>
@rnyak

rnyak commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

/ok to test fbcddc8

@rnyak

rnyak commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

/ok to test 25e219e

@thomasdhc

Copy link
Copy Markdown
Contributor

/ok to test 500b983

@thomasdhc

Copy link
Copy Markdown
Contributor

/ok to test 1db9e76

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request Feature

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants