loss_balancer

Fixed-weight KD loss balancer for Megatron noise distillation.

FixedWeightedKDLossBalancer combines ModelOpt's per-criterion losses with fixed weights: the mean of all logits (KL) losses, the mean of all intermediate losses (which includes the noise-embedding loss), and — optionally — the student's original task loss. This is Slot 2's combiner in the two-slot loss design.

Classes:

Name	Description
FixedWeightedKDLossBalancer	Combine logits and intermediate KD losses with fixed weights.

FixedWeightedKDLossBalancer

Bases: DistillationLossBalancer

Combine logits and intermediate KD losses with fixed weights.

Methods:

Name	Description
__init__	Initialize the balancer.
forward	Combine the per-criterion losses into the total KD loss.

__init__

__init__(
    logits_weight: float = 1.0,
    intermediate_weight: float = 1.0,
    original_loss_weight: float = 1.0,
    skip_original_loss: bool = True,
) -> None

Initialize the balancer.

Parameters:

Name	Type	Description	Default
logits_weight	float	Weight on the mean logits-distillation (KL) loss.	1.0
intermediate_weight	float	Weight on the mean intermediate loss (includes the noise-embedding loss).	1.0
original_loss_weight	float	Weight on the student's original task loss.	1.0
skip_original_loss	bool	If True, the student's original task loss is excluded from the total (pure distillation).	True

forward

forward(loss: dict[str, Tensor]) -> dict[str, Tensor]

Combine the per-criterion losses into the total KD loss.

Parameters:

Name	Type	Description	Default
loss	dict[str, Tensor]	Mapping from criterion key to its scalar loss. Logits losses are keyed by the LogitsKLLoss class name prefix; the student's original loss (if present) is keyed by STUDENT_LOSS_KEY; all other keys are treated as intermediate losses.	required

Returns:

Type	Description
dict[str, Tensor]	A mapping with the combined kd_loss plus the component
dict[str, Tensor]	logits_loss and intermediate_loss for reporting.

Raises:

Type	Description
ValueError	If no logits distillation loss is present.