> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.recipes.retrieval.distill_bi_encoder

## Module Contents

### Classes

| Name                                                                                                    | Description                                                        |
| ------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------ |
| [`EmbeddingDistillRecipe`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-EmbeddingDistillRecipe) | Recipe for Stage-1 embedding distillation on bi-encoder backbones. |

### Functions

| Name                                                                                                                  | Description                                                                      |
| --------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- |
| [`_build_or_none`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_build_or_none)                               | -                                                                                |
| [`_cfg_get_path`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_cfg_get_path)                                 | Read either OmegaConf-style dotted keys or nested dict/config values.            |
| [`_clean_path`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_clean_path)                                     | -                                                                                |
| [`_copy_checkpoint_metadata`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_copy_checkpoint_metadata)         | Copy HF metadata/tokenizer files needed by AutoModel.from\_pretrained.           |
| [`_dp_group_src_rank`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_dp_group_src_rank)                       | -                                                                                |
| [`_export_hf_student_checkpoint`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_export_hf_student_checkpoint) | Materialize an evaluator-facing HF checkpoint from Automodel wrapper weights.    |
| [`_mirror_hf_metadata`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_mirror_hf_metadata)                     | Mirror non-weight HF artifacts from the original student checkpoint.             |
| [`_move_to_device`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_move_to_device)                             | -                                                                                |
| [`_strip_student_prefix`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_strip_student_prefix)                 | Return the HF backbone tensors from a RetrieverStudentWithProjection checkpoint. |
| [`_unpack_qpn`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-_unpack_qpn)                                     | -                                                                                |
| [`main`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-main)                                                   | Entry point: load config, build the distillation recipe, and run training.       |

### Data

[`logger`](#nemo_automodel-recipes-retrieval-distill_bi_encoder-logger)

### API

```python
class nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe()
```

**Bases:** [TrainBiEncoderRecipe](/nemo-automodel/nemo_automodel/recipes/retrieval/train_bi_encoder#nemo_automodel-recipes-retrieval-train_bi_encoder-TrainBiEncoderRecipe)

Recipe for Stage-1 embedding distillation on bi-encoder backbones.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._build_optimizer_param_groups() -> list[dict[str, typing.Any]]
```

Build optimizer groups with projection params isolated before checkpoint restore.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._extract_scoring_reps(
    model_output
)
```

Select the pooled student embedding for validation scoring.

`RetrieverStudentWithProjection.forward` returns
`(pooled, projected, intermediate_outputs)`. The pooled embedding is the student's
native retrieval representation (and what training's InfoNCE terms score with), so it
is what the inherited validation loop should compare against.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._forward_backward_step(
    idx,
    batch,
    loss_buffer,
    num_batches,
    is_train: bool = True
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._projection_parameters() -> list[torch.nn.Parameter]
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._run_train_optim_step(
    batches,
    max_grad_norm = None
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._sync_projection_gradients() -> None
```

Average projection gradients that are outside FSDP/DDP wrapping.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe._sync_projection_parameters() -> None
```

Keep the rank-local projection head replicated across DP ranks.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe.save_checkpoint(
    epoch: int,
    step: int,
    train_loss: float,
    val_loss: dict[str, float] | None = None,
    best_metric_key: str = 'default'
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.EmbeddingDistillRecipe.setup()
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._build_or_none(
    cfg_section
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._cfg_get_path(
    cfg,
    path: str,
    default = None
)
```

Read either OmegaConf-style dotted keys or nested dict/config values.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._clean_path(
    path: pathlib.Path
) -> None
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._copy_checkpoint_metadata(
    src_dir: pathlib.Path,
    dst_dir: pathlib.Path
) -> None
```

Copy HF metadata/tokenizer files needed by AutoModel.from\_pretrained.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._dp_group_src_rank(
    group
) -> int
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._export_hf_student_checkpoint(
    src_dir: pathlib.Path,
    dst_dir: pathlib.Path
) -> None
```

Materialize an evaluator-facing HF checkpoint from Automodel wrapper weights.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._mirror_hf_metadata(
    src_model_dir: pathlib.Path,
    dst_dir: pathlib.Path,
    overwrite: bool = False
) -> None
```

Mirror non-weight HF artifacts from the original student checkpoint.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._move_to_device(
    batch: dict,
    device: torch.device
) -> dict
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._strip_student_prefix(
    state: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]
```

Return the HF backbone tensors from a RetrieverStudentWithProjection checkpoint.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder._unpack_qpn(
    batch: dict[str, torch.Tensor]
)
```

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.main(
    default_config_path = 'examples/retrieval/distill...
)
```

Entry point: load config, build the distillation recipe, and run training.

```python
nemo_automodel.recipes.retrieval.distill_bi_encoder.logger = logging.getLogger(__name__)
```