bridge.models.distillation_provider#

Module Contents#

Classes#

DistillationProvider

Provider for Bridge language models in distillation mode.

Functions#

convert_to_distillation_provider

Convert a given model provider to a DistillationProvider.

Data#

logger

API#

bridge.models.distillation_provider.logger#

‘getLogger(…)’

class bridge.models.distillation_provider.DistillationProvider(*args, **kwargs)#

Bases: megatron.bridge.models.transformer_config.TransformerConfig

Provider for Bridge language models in distillation mode.

Please use convert_to_distillation_provider() to create an instance of this class.

Parameters:

  • teacher – The teacher model provider.

  • kd_config – Knowledge-distillation configuration.

  • distill_submodule – If set, distill only this submodule of the built model (e.g. "language_model" for VLMs); the rest of the model is exported unchanged (only the submodule is trained). None distills the whole model.

Initialization

teacher: Optional[megatron.bridge.models.gpt_provider.GPTModelProvider | megatron.bridge.models.hybrid.hybrid_provider.HybridModelProvider]#

None

kd_config: Optional[megatron.bridge.training.post_training.distillation.ModelOptDistillConfig]#

None

distill_submodule: Optional[str]#

None

__post_init__()#
provide(
pre_process=None,
post_process=None,
vp_stage=None,
) → megatron.core.models.common.language_module.language_module.LanguageModule#

Build the (un-converted) student model.

The knowledge-distillation conversion is deferred to _convert_hook (a pre-wrap hook registered by convert_to_distillation_provider) so it runs after the student’s weights are loaded. This lets a caller restore a quantized student (QAD) via an earlier pre-wrap hook, and lets a weight-loaded submodule (VLMs) be extracted before it is wrapped with the teacher.

Parameters:

  • pre_process – Whether to include pre-processing in the model, defaults to first pipeline stage

  • post_process – Whether to include post-processing in the model, defaults to last pipeline stage

  • vp_stage – Virtual pipeline stage

Returns:

The un-converted student model (converted later by _convert_hook).

_convert_hook(model_chunks: list) → list#

Pre-wrap hook that applies the KD conversion after the student is weight-loaded.

With distill_submodule set (e.g. VLMs), only that submodule is distilled and returned as the model; the full model is retained on full_model so the distilled submodule can be exported back within it. Registered after the bridge’s weight-load hook, so weights are present.

to_cfg_dict() → dict[str, Any]#

Custom method to save equivalent to the original provider class.

Used by _ConfigContainerBase to serialize the main ConfigContainer to YAML. There is no need to restore a DistillationProvider from the run config file, as it can always be re-converted using the original student provider.

Returns:

Dictionary representation of this provider class

__setattr__(name, value)#
bridge.models.distillation_provider.convert_to_distillation_provider(
student_provider: megatron.bridge.models.gpt_provider.GPTModelProvider | megatron.bridge.models.hybrid.hybrid_provider.HybridModelProvider,
teacher_provider: megatron.bridge.models.gpt_provider.GPTModelProvider | megatron.bridge.models.hybrid.hybrid_provider.HybridModelProvider,
kd_config: Optional[megatron.bridge.training.post_training.distillation.ModelOptDistillConfig] = None,
*,
distill_submodule: Optional[str] = None,
) → bridge.models.distillation_provider.DistillationProvider#

Convert a given model provider to a DistillationProvider.

The KD conversion runs in a pre-wrap hook (after the student’s weights are loaded), not in provide(). To initialize the student from a checkpoint before conversion (e.g. QAD), register your own pre-wrap hook with student_provider.register_pre_wrap_hook(fn, prepend=True) so it runs before the KD-conversion hook.

Parameters:

  • student_provider – The student model provider (also the base class of the returned provider).

  • teacher_provider – The teacher model provider.

  • kd_config – Knowledge-distillation configuration.

  • distill_submodule – If set, distill only this submodule of the built model (e.g. "language_model" for VLMs); the rest of the model is exported unchanged (only the submodule is trained).