nemo_automodel.components.loss.kd_loss

Module Contents

Classes

Name	Description
`KDLoss`	Forward KL divergence loss for knowledge distillation.

Functions

Name	Description
`_forward_kl_from_log_probs`	Compute forward KL from full-vocabulary log probabilities.
`_infer_tp_group_from_dtensor`	If logits is a DTensor sharded on the vocab (last) dimension, return its TP process group.
`_kl_forward_chunked`	Compute per-token forward KL in chunks to reduce peak memory.
`_kl_forward_tp`	Compute per-token forward KL with tensor parallelism.

Data

_HAVE_DTENSOR

API

class nemo_automodel.components.loss.kd_loss.KDLoss(
    ignore_index: int = -100,
    temperature: float = 1.0,
    fp32_upcast: bool = True,
    tp_group: typing.Optional[torch.distributed.ProcessGroup] = None,
    chunk_size: int = 0
)

Bases: Module

Forward KL divergence loss for knowledge distillation.

Computes KL(P_teacher ‖ P_student) averaged over valid (non-padding) tokens.

Supports tensor-parallel (TP) training: when logits are vocab-sharded DTensors, the TP group is inferred automatically and a distributed softmax is used to avoid gathering the full vocabulary on each rank. A tp_group can also be supplied explicitly.

Parameters:

ignore_index

intDefaults to -100

Label value marking padding tokens (default -100).

temperature

floatDefaults to 1.0

Softmax temperature T. Both teacher and student logits are divided by T before computing probabilities. The loss is then multiplied by T² so that gradient magnitudes remain independent of the chosen temperature (Hinton et al., 2015).

fp32_upcast

boolDefaults to True

Cast logits to float32 before computing softmax / log-softmax for numerical stability (default True).

tp_group

Optional[torch.distributed.ProcessGroup]Defaults to None

Explicit TP process group. When None (default) the group is inferred from the DTensor placement of student_logits, or the non-TP path is used for plain tensors.

chunk_size

intDefaults to 0

When positive, valid tokens are processed in chunks of this size to avoid materializing the full [num_valid_tokens, vocab_size] probability matrix in fp32. Reduces peak memory at the cost of slightly more kernel launches. 0 (default) disables chunking. Ignored when using the TP path.

nemo_automodel.components.loss.kd_loss.KDLoss.forward(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    labels: torch.Tensor,
    num_batch_labels: int | None = None
) -> torch.Tensor

Compute the KD loss.

Parameters:

student_logits

torch.Tensor

Tensor of shape [..., vocab] with arbitrary leading dimensions. Under TP, the local tensor has shape [..., local_vocab] with Shard(-1) placement.

teacher_logits

torch.Tensor

Tensor with the same global and local shape contract as student_logits.

labels

torch.Tensor

Tensor of shape [...] matching the logits’ leading dimensions. Positions equal to ignore_index are excluded.

num_batch_labels

int | NoneDefaults to None

Total number of valid tokens across all gradient-accumulation steps. When provided the loss is sum(kl_per_token) / num_batch_labels; otherwise it is mean(kl_per_token) over the valid tokens in this micro-batch.

Returns: torch.Tensor

Scalar tensor containing zero-based forward KL.

nemo_automodel.components.loss.kd_loss._forward_kl_from_log_probs(
    teacher_log_prob: torch.Tensor,
    student_log_prob: torch.Tensor
) -> torch.Tensor

Compute forward KL from full-vocabulary log probabilities.

Parameters:

teacher_log_prob

torch.Tensor

Tensor of shape [..., vocab] containing teacher log probabilities with arbitrary leading dimensions.

student_log_prob

torch.Tensor

Tensor of shape [..., vocab] containing student log probabilities on the same dtype and device.

Returns: torch.Tensor

Tensor of shape [...] containing per-token

nemo_automodel.components.loss.kd_loss._infer_tp_group_from_dtensor(
    logits: torch.Tensor
) -> typing.Optional[torch.distributed.ProcessGroup]

If logits is a DTensor sharded on the vocab (last) dimension, return its TP process group.

Iterates over the DTensor placements to find the mesh dimension that holds a vocab-dim Shard and returns the corresponding process group. Returns None for plain tensors or DTensors that are not vocab-sharded.

nemo_automodel.components.loss.kd_loss._kl_forward_chunked(
    t_logits: torch.Tensor,
    s_logits: torch.Tensor,
    chunk_size: int
) -> torch.Tensor

Compute per-token forward KL in chunks to reduce peak memory.

Processes chunk_size tokens at a time so that only one chunk’s worth of the [chunk_size, vocab_size] fp32 probability matrix is live at any moment.

Parameters:

t_logits

torch.Tensor

Tensor of shape [tokens, vocab] containing teacher logits.

s_logits

torch.Tensor

Tensor of shape [tokens, vocab] containing student logits.

chunk_size

int

Number of tokens per chunk.

Returns: torch.Tensor

Tensor of shape [tokens] containing per-token forward KL.

nemo_automodel.components.loss.kd_loss._kl_forward_tp(
    t_logits: torch.Tensor,
    s_logits: torch.Tensor,
    tp_group: torch.distributed.ProcessGroup
) -> torch.Tensor

Compute per-token forward KL with tensor parallelism.

Parameters:

t_logits

torch.Tensor

Tensor of shape [tokens, local_vocab] containing the local teacher vocabulary shard.

s_logits

torch.Tensor

Tensor of shape [tokens, local_vocab] containing the local student vocabulary shard.

tp_group

torch.distributed.ProcessGroup

Process group spanning the tensor-parallel ranks.

Returns: torch.Tensor

Tensor of shape [tokens] containing per-token forward KL. The local

nemo_automodel.components.loss.kd_loss._HAVE_DTENSOR = True