loss

Megatron noise-embedding distillation loss (Slot 1 of the two-slot loss).

MegatronNoiseEmbeddingLoss is a ModelOpt MCore criterion attached to the noise layer. It consumes the noise layer's (clean, noisy, std, mean) output and computes the per-token std_log term, delegating the math to core's gaussian_entropy_log_loss (same constants, same formula) rather than reimplementing it.

The loss is deliberately pluggable: MegatronDistillationLossConfig carries weights for the std_log term (active), an absolute-cosine-similarity penalty (available, off by default), and the mutual-information term (Monte-Carlo Gaussian-mixture entropy — planned follow-up, still gated). The KD logits divergence is Slot 2, handled separately by the distillation provider / loss balancer.

Classes:

Name	Description
MegatronDistillationLossConfig	Weights for the noise-embedding (Slot 1) loss terms.
MegatronNoiseEmbeddingLoss	Noise-embedding distillation loss over the noise layer's 4-tuple output.

MegatronDistillationLossConfig dataclass

Weights for the noise-embedding (Slot 1) loss terms.

Attributes:

Name	Type	Description
std_log_weight	float	Weight on the std_log differential-entropy term (the core teaching signal, active by default).
absolute_cosine_similarity_weight	float	Weight on the absolute-cosine-similarity penalty |cos(clean, noisy)| over the hidden dimension, which pushes noisy embeddings toward orthogonality with clean ones. Off by default.
mutual_information_weight	float	Weight on the mutual-information term (std_log + Gaussian-mixture entropy). Off by default; gated until a later story (requires a Monte-Carlo micro-batch split).

MegatronNoiseEmbeddingLoss

Bases: BaseLoss

Noise-embedding distillation loss over the noise layer's 4-tuple output.

The student-side hook output is expected to be (clean_embeddings, noisy_embeddings, stds, means). Padding-token masking is applied downstream by Bridge's loss_reduction_fn using a 2D mask, so this returns a per-token loss.

Methods:

Name	Description
__init__	Initialize the loss.
forward	Compute the per-token noise-embedding loss.

__init__

__init__(
    model_config: Any,
    config: MegatronDistillationLossConfig | None = None,
) -> None

Initialize the loss.

Parameters:

Name	Type	Description	Default
model_config	Any	The MCore transformer config (passed to BaseLoss).	required
config	MegatronDistillationLossConfig | None	Weights for the noise-embedding loss terms.	None

Raises:

Type	Description
NotImplementedError	If the mutual-information term (still gated) is given a nonzero weight.

forward

forward(predictions: object, targets: object) -> Tensor

Compute the per-token noise-embedding loss.

Each weighted term is only computed (and attached to the autograd graph) when its weight is nonzero, so a disabled term never contributes a degenerate zero-gradient node.

Parameters:

Name	Type	Description	Default
predictions	object	The noise layer's (clean, noisy, std, mean) output (possibly wrapped by the distillation hook).	required
targets	object	The teacher-side embedding output (unused by this term).	required

Returns:

Type	Description
Tensor	The reduced loss as produced by BaseLoss.post_forward.