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:

teacher_model — an additional VLM loaded in eval mode
kd_loss_fn — KL-divergence between temperature-scaled distributions
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`	Fine-tune a student VLM via knowledge distillation from a teacher VLM.

Functions

Name	Description
`_build_kd_loss_fn`	-
`_build_teacher_model`	Build and initialize the teacher VLM for knowledge distillation.
`_verify_tokenizer_compatibility`	-
`main`	Run the VLM KD recipe from CLI or directly.

Data

logger

API

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

Bases: FinetuneRecipeForVLM

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

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.

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

Configure the teacher after construction and check its mask contract.

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

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

boolDefaults to True

Whether to run backward.

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

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

Serve teacher forwards until the student mesh broadcasts stop.

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.

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

Run one validation pass with KD loss computation.

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

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

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

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

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

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

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

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

Build student & teacher, dataloaders, optimizers, etc.

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

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 | NoneDefaults to None

Resolved distributed topology and policy object.

device

Defaults to None

Device to place the teacher model on.

Returns: torch.nn.Module

The frozen teacher model ready for inference.

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

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

Run the VLM KD recipe from CLI or directly.

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