distillation_provider

Noise-embedding distillation provider for Megatron-Bridge.

NoiseDistillationProvider adapts Megatron-Bridge's ModelOpt distillation flow for Stained Glass noise training: the student is wrapped in MegatronNoisyModel, the teacher is the clean model, the KD criterion set gains a MegatronNoiseEmbeddingLoss term (Slot 1), and the losses are combined by FixedWeightedKDLossBalancer (Slot 2 = ModelOpt's logits KL loss).

The provider is grafted onto the student provider's real class at runtime by convert_to_noise_distillation_provider (it must not be instantiated directly), mirroring Megatron-Bridge's own convert_to_distillation_provider.

Classes:

Name	Description
NoiseDistillationProvider	Noise-embedding distillation provider built on top of ModelOpt KD.

Functions:

Name	Description
convert_to_noise_distillation_provider	Graft NoiseDistillationProvider onto a student provider in place.

NoiseDistillationProvider dataclass

Bases: DistillationProvider

Noise-embedding distillation provider built on top of ModelOpt KD.

Subclasses DistillationProvider so that distill()'s isinstance(config.model, DistillationProvider) check passes and the to_cfg_dict/__post_init__/__setattr__ behavior is inherited. Like its parent, the class is grafted onto the student provider's real class at runtime via convert_to_noise_distillation_provider; it must not be instantiated directly.

Methods:

Name	Description
__init__	Reject direct instantiation.
provide	Build the distillation model: noisy student + clean teacher + criteria.

__init__

__init__(*args: Any, **kwargs: Any) -> None

Reject direct instantiation.

Parameters:

Name	Type	Description	Default
*args	Any	Ignored.	()
**kwargs	Any	Ignored.	{}

Raises:

Type	Description
NotImplementedError	Always; use convert_to_noise_distillation_provider.

provide

provide(
    pre_process: Any = None,
    post_process: Any = None,
    vp_stage: Any = None,
) -> Any

Build the distillation model: noisy student + clean teacher + criteria.

Parameters:

Name	Type	Description	Default
pre_process	Any	Whether this stage runs the pre-process (embedding) step.	None
post_process	Any	Whether this stage runs the post-process (output) step.	None
vp_stage	Any	The virtual pipeline stage; unsupported (must be None).	None

Returns:

Type	Description
Any	The ModelOpt distillation model.

Raises:

Type	Description
ValueError	If vp_stage is not None (virtual pipeline parallelism is unsupported).

convert_to_noise_distillation_provider

convert_to_noise_distillation_provider(
    student_provider: GPTModelProvider,
    teacher_provider: GPTModelProvider,
    *,
    std_scale: tuple[float, float],
    std_init_shift: float,
    std_shallow: float = 1.0,
    mean_bound: float | None = None,
    mean_init_shift: float = 0.0,
    kd_config: Any | None = None,
    noise_teacher_layer: str = "embedding",
    loss_config: (
        MegatronDistillationLossConfig | None
    ) = None,
    kd_logits_weight: float = 1.0,
    kd_intermediate_weight: float = 1.0,
    kd_original_loss_weight: float = 1.0
) -> NoiseDistillationProvider

Graft NoiseDistillationProvider onto a student provider in place.

Mirrors Megatron-Bridge's convert_to_distillation_provider but keeps DistillationProvider in the base tuple so isinstance(model, DistillationProvider) holds, while __bases__[0] remains the student's real provider class (so _super_class.provide builds the clean student).

Parameters:

Name	Type	Description	Default
student_provider	GPTModelProvider	The student model provider (mutated in place).	required
teacher_provider	GPTModelProvider	The teacher (clean) model provider.	required
std_scale	tuple[float, float]	Asymptotic (min, max) bounds of the noise std field (required; governs the privacy floor/ceiling, so it is always chosen explicitly).	required
std_init_shift	float	Shift added to the std estimator's raw output before the tanh (required); a negative value starts the std field near min.	required
std_shallow	float	Temperature of the std parameterization's tanh.	1.0
mean_bound	float | None	Magnitude bound of the mean field, or None (default) for an unbounded mean (matching core's TransformerCloak).	None
mean_init_shift	float	Shift added to the mean estimator's raw output (only applied when mean_bound is set).	0.0
kd_config	Any | None	The ModelOpt distillation config.	None
noise_teacher_layer	str	The teacher-side layer the noise criterion pairs with.	'embedding'
loss_config	MegatronDistillationLossConfig | None	The noise-embedding loss weights.	None
kd_logits_weight	float	Weight on the logits KL loss.	1.0
kd_intermediate_weight	float	Weight on the intermediate (noise) loss.	1.0
kd_original_loss_weight	float	Weight on the student's original task loss.	1.0

Returns:

Type	Description
NoiseDistillationProvider	The mutated student_provider, now a NoiseDistillationProvider.