bridge.training.distill

Module Contents

Functions

_check_step_handles_context_parallelism	A VLM that shards the sequence in its own forward needs its own step under CP.
distill	Main function to run knowledge distillation (KD).

API

bridge.training.distill._check_step_handles_context_parallelism(

config: megatron.bridge.training.config.ConfigContainer,

forward_step_func: megatron.bridge.training.forward_step_func_types.ForwardStepCallable,

) → None

A VLM that shards the sequence in its own forward needs its own step under CP.

Distilling only the language submodule keeps that forward out of the loop, so the GPT step is correct there – it hands mrope the full-length position ids and lets the rotary embedding shard.

bridge.training.distill.distill(

config: megatron.bridge.training.config.ConfigContainer,

forward_step_func: megatron.bridge.training.forward_step_func_types.ForwardStepCallable = forward_step_modelopt,

) → None

Main function to run knowledge distillation (KD).

Parameters:

• config – The main configuration container holding all necessary parameters.

• forward_step_func – Step function to distill with, defaulting to the GPT step. Pass qwen3_vl_step.forward_step_modelopt when distilling a whole Qwen3-VL model under context parallelism; for distill_submodule="language_model" the default is correct, since that forward never re-shards the sequence – though packed data is not supported there under context parallelism. A custom step must attach the KD loss, or training silently runs without distillation.

.. warning::

This is an experimental API and is subject to change in backwards incompatible ways without notice.