> 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.vlm.kd

Knowledge Distillation recipe for Vision-Language Models with NeMo AutoModel.

This recipe fine-tunes a *student* VLM using the logits of a frozen *teacher* VLM. It
extends `FinetuneRecipeForVLM` adding:

1. teacher\_model — an additional VLM loaded in `eval` mode
2. kd\_loss\_fn    — KL-divergence between temperature-scaled distributions
3. kd\_ratio      — linear mix between CE loss and KD loss

The training loop preserves all VLM-specific input handling (pixel\_values, image\_grid\_thw,
etc.) and passes multimodal inputs to both teacher and student models.

The loss becomes:
loss = (1-kd\_ratio) \* ce\_loss + kd\_ratio \* kd\_loss

Pipeline parallelism is not supported in this recipe.

The file exposes `KnowledgeDistillationRecipeForVLM` and a `main` entry-point
so it can be launched exactly the same way as other recipes:

python -m torch.distributed.run --nproc-per-node=8 \
nemo\_automodel/recipes/vlm/kd.py \
-c examples/vlm\_kd/qwen3\_5/qwen3\_5\_vl\_4b\_kd.yaml

## Module Contents

### Classes

| Name                                                                                                    | Description                                                            |
| ------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`KnowledgeDistillationRecipeForVLM`](#nemo_automodel-recipes-vlm-kd-KnowledgeDistillationRecipeForVLM) | Fine-tune a student VLM via knowledge distillation from a teacher VLM. |

### Functions

| Name                                                                                                | Description                                                      |
| --------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- |
| [`_build_kd_loss_fn`](#nemo_automodel-recipes-vlm-kd-_build_kd_loss_fn)                             | -                                                                |
| [`_build_teacher_model`](#nemo_automodel-recipes-vlm-kd-_build_teacher_model)                       | Build and initialize the teacher VLM for knowledge distillation. |
| [`_verify_tokenizer_compatibility`](#nemo_automodel-recipes-vlm-kd-_verify_tokenizer_compatibility) | -                                                                |
| [`main`](#nemo_automodel-recipes-vlm-kd-main)                                                       | Run the VLM KD recipe from CLI or directly.                      |

### Data

[`logger`](#nemo_automodel-recipes-vlm-kd-logger)

### API

```python
class nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM()
```

**Bases:** [FinetuneRecipeForVLM](/nemo-automodel/nemo_automodel/recipes/vlm/finetune#nemo_automodel-recipes-vlm-finetune-FinetuneRecipeForVLM)

Fine-tune a student VLM via knowledge distillation from a teacher VLM.

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._configure_packing() -> nemo_automodel.components.models.common.packing.PackingCapabilities
```

Emit metadata for the teacher as well as the student.

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._configure_teacher_packing() -> None
```

Configure the teacher after construction and check its mask contract.

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._create_distributed_setup() -> nemo_automodel.components.distributed.config.DistributedSetup
```

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._forward_backward_step(
    idx,
    batch,
    loss_buffer,
    num_label_tokens,
    num_batches,
    is_train: bool = True
)
```

Run one student VLM microbatch with KD.

Teacher and student logits have global shape
`[batch, sequence, vocab]` and may use local TP/CP layouts inside the
step.

**Parameters:**

**`idx`**

Zero-based accumulation microbatch index.

---

**`batch`**

Mapping containing text tensors of shape
`[batch, sequence]`. Multimodal tensor leaves may have
arbitrary rank and axis order.

---

**`loss_buffer`**

Output list receiving one detached scalar tensor.

---

**`num_label_tokens`**

Valid-label count across the optimizer step.

---

**`num_batches`**

Number of accumulation microbatches in the step.

---

**`is_train`** `bool` — default: True

Whether to run backward.

---

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._get_separate_teacher_logits(
    batch: dict[str, typing.Any]
) -> torch.Tensor
```

Request teacher logits for one student VLM batch.

**Parameters:**

**`batch`** `dict[str, Any]`

Mapping containing `input_ids` and `labels` as tensors of
shape `[batch, sequence]`. Multimodal tensor leaves may have
arbitrary rank and axis order and are transported unchanged.

---

**Returns:** `torch.Tensor`

Replicated tensor of shape `[batch, sequence, vocab]` containing

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._run_teacher_worker() -> None
```

Serve teacher forwards until the student mesh broadcasts stop.

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._run_train_optim_step(
    batches,
    max_grad_norm: float | None = None
)
```

Execute a single training step with KD loss tracking.

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._run_validation_epoch(
    val_dataloader
)
```

Run one validation pass with KD loss computation.

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._setup_kd_state() -> None
```

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._should_setup_training_components() -> bool
```

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM._teacher_forward_separate(
    batch: dict[str, typing.Any]
) -> torch.Tensor
```

Run one teacher VLM batch and materialize full logits.

**Parameters:**

**`batch`** `dict[str, Any]`

Mapping containing `input_ids` and `labels` as tensors of
shape `[batch, sequence]`. Multimodal tensor leaves may have
arbitrary rank and axis order; their model processor owns those
layouts.

---

**Returns:** `torch.Tensor`

Detached tensor of shape `[batch, sequence, vocab]` containing

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM.log_train_metrics(
    log_data
) -> float
```

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM.log_val_metrics(
    log_data
)
```

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM.run_train_validation_loop()
```

Run the student loop or serve teacher forwards on a separate mesh.

```python
nemo_automodel.recipes.vlm.kd.KnowledgeDistillationRecipeForVLM.setup()
```

Build student & teacher, dataloaders, optimizers, etc.

```python
nemo_automodel.recipes.vlm.kd._build_kd_loss_fn(
    cfg_kd
)
```

```python
nemo_automodel.recipes.vlm.kd._build_teacher_model(
    cfg_teacher,
    cfg_freeze,
    seed: int,
    distributed_setup: nemo_automodel.components.distributed.config.DistributedSetup | None = None,
    device = None
) -> torch.nn.Module
```

Build and initialize the teacher VLM for knowledge distillation.

Uses the same `build_model` as the student but without PEFT, FP8, or QAT
since the teacher should be frozen in full precision.

**Parameters:**

**`cfg_teacher`**

Configuration for teacher model instantiation.

---

**`cfg_freeze`**

Freeze configuration for the teacher model.

---

**`seed`** `int`

Random seed for reproducibility.

---

**`distributed_setup`** `DistributedSetup | None` — default: None

Resolved distributed topology and policy object.

---

**`device`** — default: None

Device to place the teacher model on.

---

**Returns:** `torch.nn.Module`

The frozen teacher model ready for inference.

```python
nemo_automodel.recipes.vlm.kd._verify_tokenizer_compatibility(
    student_cfg,
    teacher_cfg,
    trust_remote_code = True
)
```

```python
nemo_automodel.recipes.vlm.kd.main(
    config_path = None
)
```

Run the VLM KD recipe from CLI or directly.

```python
nemo_automodel.recipes.vlm.kd.logger = logging.getLogger(__name__)
```