nemo_automodel.components.loss.mtp

Module Contents

Classes

Name	Description
`MTPLossConfig`	Typed config for the Multi-Token-Prediction auxiliary loss.
`MTPLossOutput`	Aggregate MTP loss and its per-depth components.
`PipelineCausalLMLoss`	Pipeline schedule loss that can add MTP auxiliary CE on the last stage.

Functions

Name	Description
`calculate_mtp_loss`	Compute the DeepSeek-V3 Multi-Token Prediction auxiliary loss.

API

class nemo_automodel.components.loss.mtp.MTPLossConfig(
    scaling_factor: float | None = None,
    ignore_index: int | None = None
)

Dataclass

Typed config for the Multi-Token-Prediction auxiliary loss.

MTP is gated on the model emitting per-depth outputs; this config only carries its hyperparameters. scaling_factor=None keeps the model-provided value (out.mtp_loss_scaling_factor / get_mtp_loss_scaling_factor); set it to override. ignore_index=None inherits the configured loss’s sentinel; an explicit value must match it.

ignore_index

int | None = None

scaling_factor

float | None = None

nemo_automodel.components.loss.mtp.MTPLossConfig.build(
    loss_fn: torch.nn.Module,
    model: torch.nn.Module,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> nemo_automodel.components.loss.mtp.PipelineCausalLMLoss

Build the pipeline-schedule, MTP-aware loss for loss_fn/model.

class nemo_automodel.components.loss.mtp.MTPLossOutput(
    loss: torch.Tensor,
    per_depth_losses: list[torch.Tensor]
)

Dataclass

Aggregate MTP loss and its per-depth components.

loss

Tensor

per_depth_losses

list[Tensor]

class nemo_automodel.components.loss.mtp.PipelineCausalLMLoss(
    loss_fn: torch.nn.Module,
    model: torch.nn.Module,
    scaling_factor: float | None = None,
    ignore_index: int | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
)

Bases: Module

Pipeline schedule loss that can add MTP auxiliary CE on the last stage.

Per-microbatch seq_idx is read from a trailing element of the last-stage output tuple — the model appends an [B, S] int32 tail when MTP is enabled. This binds each microbatch’s seq_idx to its loss call via the PP runtime’s output→loss contract, so the wiring is schedule-agnostic. Legacy cu_seqlens (THD path) is a fallback for models that don’t emit a seq_idx tail.

cu_seqlens

Tensor | None = None

nemo_automodel.components.loss.mtp.PipelineCausalLMLoss._extract_seq_idx_tail(
    output
) -> tuple[torch.Tensor | None, object]

staticmethod

Detect and strip a trailing per-microbatch seq_idx from output.

Convention: with MTP enabled the last-stage output is (logits, *mtp_per_depth_h, seq_idx) with an [B, S] int32 tail — dtype alone discriminates.

nemo_automodel.components.loss.mtp.PipelineCausalLMLoss.forward(
    output,
    labels: torch.Tensor
) -> torch.Tensor

Last-stage pipeline loss (main CE plus optional MTP aux CE).

B=microbatch, S=seq, H=hidden, V=vocab.

Parameters:

output

bare hidden states [B, S, H] (FusedLinearCrossEntropy path), a HF output with logits [B, S, V], or an MTP tuple (logits, *mtp_per_depth_h[, seq_idx]) with seq_idx [B, S] int32. A tuple with FusedLinearCrossEntropy raises.

labels

torch.Tensor

target token ids [B, S] int64.

Returns: torch.Tensor

Scalar loss tensor.

nemo_automodel.components.loss.mtp.calculate_mtp_loss(
    loss_fn: torch.nn.Module,
    mtp_per_depth_h: list[torch.Tensor] | None = None,
    mtp_per_depth_logits: list[torch.Tensor] | None = None,
    mtp_per_depth_targets: collections.abc.Sequence[torch.Tensor] | None = None,
    labels: torch.Tensor,
    model: torch.nn.Module,
    scaling_factor: float = 0.1,
    num_label_tokens: int | None = None,
    ignore_index: int | None = None,
    cu_seqlens: torch.Tensor | None = None,
    seq_idx: torch.Tensor | None = None,
    lm_weight: torch.Tensor | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None,
    loss_weights: torch.Tensor | None = None,
    return_per_depth: bool = False
) -> torch.Tensor | nemo_automodel.components.loss.mtp.MTPLossOutput

Compute the DeepSeek-V3 Multi-Token Prediction auxiliary loss.

Each depth’s CE is dispatched through :func:calculate_loss with the same loss class as the main path, so MTP inherits FusedLinearCrossEntropy / MaskedCrossEntropy memory and numerical characteristics.

Parameters:

loss_fn

nn.Module

Configured per-token loss class (same instance the main path uses).

mtp_per_depth_h

list[torch.Tensor] | NoneDefaults to None

Per-depth hidden-state tensors of shape [batch, sequence, hidden], or [1, tokens, hidden] for a flattened THD-packed stream.

mtp_per_depth_logits

list[torch.Tensor] | NoneDefaults to None

Per-depth logit tensors of shape [batch, sequence, vocab], or [1, tokens, vocab] for a flattened THD-packed stream.

mtp_per_depth_targets

Sequence[torch.Tensor] | NoneDefaults to None

Optional precomputed target tensors, one per MTP depth, with the same local shape and token layout as labels. Context-parallel callers must shift and boundary-mask them in global sequence order before applying the model input shard. These targets are authoritative: this function cannot infer or repair global packed-sequence boundaries from a rank-local CP shard.

labels

torch.Tensor

Original unshifted label tensor of shape [batch, sequence] or [tokens] for a flattened THD-packed stream.

model

nn.Module

The wrapped model; used to fetch the shared LM head when the loss class needs materialized logits (non-FusedLinearCE path).

scaling_factor

floatDefaults to 0.1

Coefficient applied to the summed per-depth CE.

num_label_tokens

int | NoneDefaults to None

Total non-ignore label-token count used for sum-reduction normalization.

ignore_index

int | NoneDefaults to None

Label value masked out of the CE loss for the trailing k+1 rolled positions at depth k. Defaults to the configured loss’s ignore index and must match it when provided.

cu_seqlens

torch.Tensor | NoneDefaults to None

Optional cumulative sequence lengths [num_seqs+1] (THD-pack layout). When supplied and seq_idx is not, builds a per-token sub-sequence index via searchsorted. Without packing this can be omitted.

seq_idx

torch.Tensor | NoneDefaults to None

Optional per-token sub-sequence index [B, S] (or [S]). Equality classes are what matter; absolute values can be any ints. Takes precedence over cu_seqlens. Used to mask label rolls whose source position lies in a different sub-sequence.

lm_weight

torch.Tensor | NoneDefaults to None

Optional LM-head weight tensor of shape [vocab, hidden]. Supplying it lets the main loss and all MTP depths share one DTensor full_tensor() gather on the FusedLinearCrossEntropy path.

grad_reduce_group

dist.ProcessGroup | NoneDefaults to None

Group that contributes independent loss shards when the shared LM-head weight is a DTensor.

loss_weights

torch.Tensor | NoneDefaults to None

Optional per-token objective multipliers matching labels.shape.

return_per_depth

boolDefaults to False

Return the aggregate loss together with the unscaled loss for each MTP depth. Defaults to False to preserve the scalar return expected by existing callers.

Returns: torch.Tensor | MTPLossOutput

Scalar MTP loss tensor. When return_per_depth=True, returns an