feat: add Ministral3 embedding distillation recipe

examples/retrieval/distillation/README.md

@@ -0,0 +1,14 @@
+# Retrieval Distillation (Automodel)
+
+Run the baseline Stage-1 style recipe:
+
+```bash
+automodel --nproc-per-node 8 examples/retrieval/distillation/ministral3_2b_distill.yaml
+```
+
+The recipe writes both:
+
+- `epoch_<E>_step_<S>/...` standard Automodel checkpoints
+- `step_<S>/student` + `step_<S>/projection.pt` legacy-compatible HF checkpoint/sidecar
+
+

examples/retrieval/distillation/nemotron3_embed_1b_distill.yaml

@@ -0,0 +1,125 @@
+# Usage (8 GPUs, single node):
+#   automodel --nproc-per-node 8 examples/retrieval/distillation/ministral3_2b_distill.yaml
+
+recipe: EmbeddingDistillRecipe
+
+seed: 42
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 5
+
+step_scheduler:
+  global_batch_size: 64
+  local_batch_size: 4
+  ckpt_every_steps: 500
+  num_epochs: 1
+
+model:
+  _target_: nemo_automodel._transformers.retrieval.RetrieverStudentWithProjection.build
+  pretrained_model_name_or_path: <your_path>
+  teacher_hidden_size: 4096
+  pooling: avg
+  l2_normalize: false
+  capture_layers: []
+  trust_remote_code: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+
+teacher_model:
+  _target_: nemo_automodel._transformers.retrieval.RetrieverTeacherEmbeddingEncoder.build
+  pretrained_model_name_or_path: <your_path>
+  pooling: avg
+  l2_normalize: false
+  capture_layers: []
+  trust_remote_code: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+
+tokenizer:
+  _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
+  pretrained_model_name_or_path: <your_path>
+  add_eos_token: false
+  padding_side: left
+  force_hf: true
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  dataset:
+    _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
+    model_type: bi_encoder
+    n_passages: 5
+    seed: 42
+    do_shuffle: true
+    data_type: train
+    data_dir_list:
+      - hf://nvidia/embed-nemotron-dataset-v1/FEVER # 50k
+      - hf://nvidia/embed-nemotron-dataset-v1/SyntheticClassificationData # 100k
+  collate_fn:
+    _target_: nemo_automodel.components.datasets.llm.BiEncoderDistillCollator
+    q_max_len: 512
+    p_max_len: 512
+    query_prefix: "query:"
+    passage_prefix: "passage:"
+    pad_to_multiple_of: 8
+  shuffle: true
+  num_workers: 2
+
+# ------------------------------------------------------------------------
+# Loss building blocks (instances are built whether or not their weight is
+# active, so the recipe-level weights below drive the actual loss mix).
+# ------------------------------------------------------------------------
+distill_loss:
+  _target_: nemo_automodel.components.loss.embedding_distill.EmbeddingDistillLoss
+  reduction: mean
+
+mse_loss:
+  _target_: nemo_automodel.components.loss.embedding_distill.EmbeddingMSELoss
+  normalize: false
+  reduction: mean
+
+infonce_distill_loss:
+  _target_: nemo_automodel.components.loss.infonce.InfoNCEDistillLoss
+  temperature: 0.05
+  direction: q2d
+  use_in_batch_negatives: true
+  normalize: true
+  divergence: kl
+  cross_device_negatives: true
+
+# ------------------------------------------------------------------------
+# Active loss mix example:
+#   total = 1.0 * cosine + 1.0 * mse + 0 * nce_distill
+# (i.e. Cosine similarity + MSE alignment + InfoNCE-distill listwise KD)
+# ------------------------------------------------------------------------
+distill_loss_weight: 1.0
+mse_loss_weight: 1.0
+nce_distill_loss_weight: 0.0
+
+optimizer:
+  _target_: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
+  lr: 1.0e-5
+  weight_decay: 0.01
+  adam_w_mode: true
+  bias_correction: true
+  master_weights: true
+
+lr_scheduler:
+  lr_warmup_steps: 200
+  lr_decay_style: cosine
+
+clip_grad_norm:
+  max_norm: 1.0
+
+distributed:
+  strategy: fsdp2
+  dp_size: none
+  tp_size: 1
+  cp_size: 1
+  sequence_parallel: false
+
+checkpoint:
+  enabled: true
+  checkpoint_dir: <your_path>
+  model_save_format: safetensors
+  save_consolidated: true

nemo_automodel/_transformers/__init__.py

@@ -34,6 +34,14 @@
     "NeMoAutoModelCrossEncoder": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelCrossEncoder"),
     "NeMoAutoTokenizer": ("nemo_automodel._transformers.auto_tokenizer", "NeMoAutoTokenizer"),
     "AutoMFU": ("nemo_automodel._transformers.mfu", "AutoMFU"),
+    "RetrieverStudentWithProjection": (
+        "nemo_automodel._transformers.retrieval",
+        "RetrieverStudentWithProjection",
+    ),
+    "RetrieverTeacherEmbeddingEncoder": (
+        "nemo_automodel._transformers.retrieval",
+        "RetrieverTeacherEmbeddingEncoder",
+    ),
 }
 
 __all__ = [
@@ -47,6 +55,8 @@
     "NeMoAutoModelCrossEncoder",
     "NeMoAutoTokenizer",
     "AutoMFU",
+    "RetrieverStudentWithProjection",
+    "RetrieverTeacherEmbeddingEncoder",
 ]
 
 

nemo_automodel/_transformers/retrieval.py

@@ -16,6 +16,7 @@
 
 import inspect
 import os
+from collections.abc import Iterable, Sequence
 from typing import Optional
 
 import torch
@@ -26,6 +27,7 @@
 from transformers.utils import logging
 
 from nemo_automodel._transformers.registry import ModelRegistry
+from nemo_automodel.components.loss.intermediate_distill import LayerCapture
 from nemo_automodel.components.models.common.bidirectional import EncoderStateDictAdapter
 
 logger = logging.get_logger(__name__)
@@ -152,6 +154,8 @@ def pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor, pool_ty
         Pooled embeddings [batch_size, hidden_size]
     """
     last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    if pool_type == "mean":
+        pool_type = "avg"
 
     if pool_type == "avg":
         emb = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
@@ -489,3 +493,160 @@ def forward(self, input_dict: dict = None, **kwargs) -> Optional[torch.Tensor]:
         inputs = input_dict if input_dict is not None else kwargs
         inputs.setdefault("return_dict", True)
         return self.model(**inputs)
+
+
+def _get_layers(model: nn.Module) -> nn.ModuleList:
+    """Locate the transformer block list on a HuggingFace backbone."""
+    if hasattr(model, "layers") and isinstance(model.layers, nn.ModuleList):
+        return model.layers
+    if hasattr(model, "model") and hasattr(model.model, "layers") and isinstance(model.model.layers, nn.ModuleList):
+        return model.model.layers
+    if (
+        hasattr(model, "transformer")
+        and hasattr(model.transformer, "h")
+        and isinstance(model.transformer.h, nn.ModuleList)
+    ):
+        return model.transformer.h
+    raise AttributeError("Could not locate transformer layers on model (tried .layers, .model.layers, .transformer.h)")
+
+
+class RetrieverStudentWithProjection(nn.Module):
+    """Bi-encoder student with a trainable linear projection into teacher space."""
+
+    def __init__(
+        self,
+        student: BiEncoderModel,
+        teacher_hidden_size: int,
+        capture_layers: Sequence[int] | None = None,
+    ) -> None:
+        super().__init__()
+        self.student = student
+        student_hidden = int(self.student.model.config.hidden_size)
+        self.student_hidden = student_hidden
+        self.teacher_hidden = int(teacher_hidden_size)
+
+        self.projection = nn.Linear(student_hidden, self.teacher_hidden, bias=True)
+        nn.init.xavier_uniform_(self.projection.weight)
+        nn.init.zeros_(self.projection.bias)
+        self.projection = self.projection.float()
+        if torch.cuda.is_available():
+            # Keep projection colocated with rank-local activations under distributed training.
+            self.projection = self.projection.to(device=torch.cuda.current_device())
+
+        self._capture = LayerCapture(detach=False)
+        if capture_layers:
+            self.attach_intermediate_capture(capture_layers)
+
+    @classmethod
+    def build(
+        cls,
+        pretrained_model_name_or_path: str,
+        teacher_hidden_size: int,
+        pooling: str = "avg",
+        l2_normalize: bool = False,
+        capture_layers: Sequence[int] | None = None,
+        trust_remote_code: bool = True,
+        **kwargs,
+    ) -> "RetrieverStudentWithProjection":
+        from nemo_automodel import NeMoAutoModelBiEncoder
+
+        student = NeMoAutoModelBiEncoder.from_pretrained(
+            pretrained_model_name_or_path=pretrained_model_name_or_path,
+            pooling=pooling,
+            l2_normalize=l2_normalize,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        return cls(student=student, teacher_hidden_size=teacher_hidden_size, capture_layers=capture_layers)
+
+    def attach_intermediate_capture(self, layer_indices: Iterable[int]) -> "RetrieverStudentWithProjection":
+        self._capture.attach(_get_layers(self.student.model), layer_indices)
+        return self
+
+    def detach_intermediate_capture(self) -> None:
+        self._capture.detach_hooks()
+
+    def save_pretrained(self, save_directory: str, **kwargs) -> None:
+        # Keep the save format HF-compatible for evaluator tooling: this stores
+        # only the student backbone; the projection is saved by the recipe.
+        self.student.save_pretrained(save_directory, **kwargs)
+
+    def _encode(self, input_dict: dict) -> torch.Tensor:
+        if not input_dict:
+            return None
+        embeds = self.student(input_dict)
+        return embeds.contiguous()
+
+    def forward(
+        self,
+        input_dict: dict,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[int, torch.Tensor]]:
+        self._capture.reset()
+        pooled = self._encode(input_dict)
+        pooled_fp32 = pooled.float()
+
+        with torch.amp.autocast(device_type=pooled.device.type, enabled=False):
+            projected = self.projection(pooled_fp32)
+
+        return pooled_fp32, projected, dict(self._capture.outputs)
+
+
+class RetrieverTeacherEmbeddingEncoder(nn.Module):
+    """Frozen bi-encoder teacher with optional intermediate-layer capture."""
+
+    def __init__(
+        self,
+        teacher: BiEncoderModel,
+        capture_layers: Sequence[int] | None = None,
+    ) -> None:
+        super().__init__()
+        self.teacher = teacher
+        self.teacher.eval()
+        for param in self.teacher.parameters():
+            param.requires_grad_(False)
+
+        self.hidden_size = int(self.teacher.model.config.hidden_size)
+        self._capture = LayerCapture(detach=True)
+        if capture_layers:
+            self.attach_intermediate_capture(capture_layers)
+
+    @classmethod
+    def build(
+        cls,
+        pretrained_model_name_or_path: str,
+        pooling: str = "avg",
+        l2_normalize: bool = False,
+        capture_layers: Sequence[int] | None = None,
+        trust_remote_code: bool = True,
+        **kwargs,
+    ) -> "RetrieverTeacherEmbeddingEncoder":
+        from nemo_automodel import NeMoAutoModelBiEncoder
+
+        teacher = NeMoAutoModelBiEncoder.from_pretrained(
+            pretrained_model_name_or_path=pretrained_model_name_or_path,
+            pooling=pooling,
+            l2_normalize=l2_normalize,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        return cls(teacher=teacher, capture_layers=capture_layers)
+
+    def attach_intermediate_capture(self, layer_indices: Iterable[int]) -> "RetrieverTeacherEmbeddingEncoder":
+        self._capture.attach(_get_layers(self.teacher.model), layer_indices)
+        return self
+
+    def detach_intermediate_capture(self) -> None:
+        self._capture.detach_hooks()
+
+    @torch.no_grad()
+    def _encode(self, input_dict: dict) -> torch.Tensor:
+        if not input_dict:
+            return None
+        embeds = self.teacher(input_dict)
+        return embeds.contiguous()
+
+    @torch.no_grad()
+    def forward(self, input_dict: dict) -> tuple[torch.Tensor, dict[int, torch.Tensor]]:
+        self._capture.reset()
+        pooled = self._encode(input_dict)
+        return pooled.float(), dict(self._capture.outputs)

nemo_automodel/components/checkpoint/checkpointing.py

@@ -1298,6 +1298,8 @@ def _maybe_build_consolidated_index(
             pre_shard_hf_state_dict_keys = (
                 getattr(model, "_pre_shard_hf_state_dict_keys", None) or self.config.model_state_dict_keys
             )
+            if pre_shard_hf_state_dict_keys is None:
+                pre_shard_hf_state_dict_keys = list(state_dict.keys())
             if model_type and requires_tensor_merging(model_type) and not hasattr(model_part, "state_dict_adapter"):
                 # in this case, Transformers performed weight conversion so we will save the converted format in the checkpoint
                 num_shards = max(fqn_to_file_index_mapping.values()) if fqn_to_file_index_mapping else 1

nemo_automodel/components/datasets/llm/__init__.py

@@ -28,6 +28,7 @@
     make_vision_collator_from_processor_method,
 )
 from .retrieval_dataset import make_retrieval_dataset  # noqa: F401
+from .retrieval_distill_collator import BiEncoderDistillCollator  # noqa: F401
 from .squad import make_squad_dataset  # noqa: F401
 from .xlam import make_xlam_dataset  # noqa: F401
 
@@ -38,6 +39,7 @@
     "make_xlam_dataset",
     "make_agent_chat_dataset",
     "BiEncoderCollator",
+    "BiEncoderDistillCollator",
     "CrossEncoderCollator",
     "make_vision_collator_from_processor_method",
     "ColumnMappedTextInstructionDataset",