nemo_automodel.components.loss.dllm_loss

Bases: Module

Flat cross-entropy loss for block-diffusion (diffusion_gemma) training.

The diffusion_gemma checkpoint uses uniform random-token (D3PM-uniform) corruption, not absorbing [MASK]. Its loss is plain mean cross-entropy over all supervised canvas positions (corrupted AND uncorrupted): the loss support is the full selected canvas (target_mask = canvas_mask), which is NOT noise-gated. noise_mask is accepted (for diagnostics) but does NOT gate the loss support:

.. math:: \text{loss} = \frac{\sum_{i \in \text{supervised (canvas)}} \text{CE}_i}{N}

where N is the supervised canvas-token count. There is no 1/p (1/t) reweighting (that is the absorbing-kernel ELBO weight, which does not apply to the uniform kernel) and no autoregressive term. Flatness is a property of this class, not of a caller passing p_mask = 1.

The signature matches :class:MDLMCrossEntropyLoss / :class:HybridDiffusionLLMLoss so the recipe can call it uniformly; the p_mask / causal_logits / loss_mask_ar / num_ar_tokens arguments are accepted but ignored.

nemo_automodel.components.loss.dllm_loss.BlockDiffusionCrossEntropyLoss.forward(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    noise_mask: torch.Tensor,
    p_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    loss_mask_ar: torch.Tensor | None = None,
    num_diffusion_tokens: int | None = None,
    num_ar_tokens: int | None = None,
    causal_logits: torch.Tensor | None = None,
    noisy_input_ids: torch.Tensor | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

Bases: Module

Position-decay cross-entropy loss for DFlash draft model training.

Implements Eq. 4 of the DFlash paper:

.. math:: w_k = \exp!\left(-\frac{k-1}{\gamma}\right), \quad k = 1, \dots, T

where k indexes the predicted positions within a block (k=0 is the clean anchor and is not predicted; k=1 is the first masked position).

Loss is normalised by the sum of effective weights (w_k * block_mask). Pass num_tokens (a global all-reduced count) for normalisation consistent across DP replicas and gradient-accumulation steps.

Paper default γ values (Appendix A.1):

• block size 16 → γ = 7
• block size 10 → γ = 5
• block size 8 → γ = 4

Parameters:

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._chunk_nll(
    hidden_chunk: torch.Tensor,
    lm_head_weight: torch.Tensor,
    lm_head_bias: torch.Tensor | None,
    target_chunk: torch.Tensor
) -> typing.Tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._decay_weights(
    T: int,
    block_size: int | None,
    device,
    dtype
) -> torch.Tensor

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._draft_acc_per_pos(
    correct: torch.Tensor,
    block_mask: torch.Tensor,
    block_size: int | None
) -> typing.Tuple[torch.Tensor | None, torch.Tensor | None]

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._reduce(
    token_nll: torch.Tensor,
    block_mask: torch.Tensor,
    num_tokens: int | None,
    block_size: int | None,
    draft_correct_per_pos: torch.Tensor | None = None,
    draft_count_per_pos: torch.Tensor | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss.forward(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    block_mask: torch.Tensor,
    num_tokens: int | None = None,
    block_size: int | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss.forward_fused(
    hidden: torch.Tensor,
    lm_head_weight: torch.Tensor,
    target_ids: torch.Tensor,
    block_mask: torch.Tensor,
    num_tokens: int | None = None,
    block_size: int | None = None,
    lm_head_bias: torch.Tensor | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

Bases: NamedTuple

Unified return type for all dLLM loss functions.

Bases: Module

Combined diffusion + optional AR loss for hybrid diffusion LLM models.

Used by Nemotron-Labs-Diffusion. The diffusion component computes MDLM-style loss at noise-masked positions, weighted by 1/p_mask. An optional autoregressive (AR) component adds standard cross-entropy at AR positions (the causal branch of model output).

Total loss = alpha * diffusion_loss + ar_loss.

nemo_automodel.components.loss.dllm_loss.HybridDiffusionLLMLoss.forward(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    noise_mask: torch.Tensor,
    p_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    loss_mask_ar: torch.Tensor | None = None,
    num_diffusion_tokens: int | None = None,
    num_ar_tokens: int | None = None,
    causal_logits: torch.Tensor | None = None,
    noisy_input_ids: torch.Tensor | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

Bases: Module

Introspective DLM all-masked loss (Yu et al., 2026; arXiv:2604.11035).

Operates on the concatenated [x_t (L) | x_0 (L)] forward output produced under the block-diffusion attention mask, where x_t is the noisy (masked) copy and x_0 the clean copy. With a next-token “logit shift” (the hidden state at position i predicts token i+1) the objective combines two cross-entropy terms, both supervised on the response (answer) tokens:

.. math:: L = \text{CE}\text{noisy} + \alpha \cdot \text{CE}\text{clean}

• CE_noisy — decode CE on the x_t half (distribution q): each masked token is conditioned on the clean ground-truth prefix.
• CE_clean — verify CE on the x_0 half (distribution p): the clean copy of the response under strict causal attention.

With auto_balance=True the fixed weight is replaced by the detached ratio CE_noisy / CE_clean each step so the two terms stay comparable in magnitude (paper Eq. 2, used for the later stride expansions). Otherwise the fixed clean_loss_weight is used (the paper’s 0.2 for early training).

Parameters:

nemo_automodel.components.loss.dllm_loss.IDLMLoss.forward(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    answer_mask: torch.Tensor,
    valid_mask: torch.Tensor,
    seq_len: int,
    num_diffusion_tokens: int | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

Bases: Module

Cross-entropy loss for MDLM training.

Matches the reference dllm framework (dllm/core/trainers/mdlm.py):

.. math:: \text{loss} = \frac{\sum_{i \in \text{masked}} \text{CE}_i \cdot w(t)}{\sum \text{maskable}}

where :math:w(t) = 1/t for the scheduler weight type (linear schedule).

nemo_automodel.components.loss.dllm_loss.MDLMCrossEntropyLoss.forward(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    noise_mask: torch.Tensor,
    p_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    loss_mask_ar: torch.Tensor | None = None,
    num_diffusion_tokens: int | None = None,
    num_ar_tokens: int | None = None,
    causal_logits: torch.Tensor | None = None,
    noisy_input_ids: torch.Tensor | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

Bases: Module

Discrete-time NELBO for SCDD (openreview.net/forum?id=zQKlzKB6I9).

The forward process mixes an absorbing [MASK] channel with uniform transitions (see :func:scdd_schedule), so a position at time t is either [MASK] or a possibly-wrong non-[MASK] token. The two cases contribute different terms to the ELBO:

• z_t = [MASK] — the familiar denoising term, the reverse-KL mass that the model must place on the clean token when it un-absorbs.
• z_t != [MASK] — the correction term, the reverse KL of the true posterior against the model posterior at an already-visible token. This is what trains the model to overwrite its own earlier mistakes, and it is scored at every non-[MASK] supervised position, including uncorrupted ones (where it vanishes only in the degenerate max_ratio = 0 limit).

Both terms are scaled by num_timesteps so the loss is the discrete-time NELBO per token rather than a per-step increment.

Setting max_ratio = 0 removes the uniform channel entirely and the loss reduces exactly to the MDLM objective -log p(x_0) / t at masked positions with zero correction term — the invariant the unit tests pin.

The model output is re-parameterised as a distribution over non-[MASK] tokens (the [MASK] logit is driven to -inf before the log-softmax), matching the SCDD backbone parameterisation: the denoiser never predicts the absorbing state.

Unlike the absorbing losses, the ELBO needs the model’s probability of every non-[MASK] token, so it cannot be reduced by a fused cross-entropy kernel. The vocabulary-sized work is instead done in position chunks wrapped in :func:torch.utils.checkpoint (the same treatment :meth:DFlashDecayLoss.forward_fused gives its LM-head projection), so the two [positions, vocab] fp32 intermediates are recomputed in backward and peak activation is one chunk rather than the whole batch.

nemo_automodel.components.loss.dllm_loss.SCDDLoss._vocab_terms(
    logits_chunk: torch.Tensor,
    x_0_chunk: torch.Tensor,
    z_t_chunk: torch.Tensor,
    log_base_s_chunk: torch.Tensor,
    log_rho_s_chunk: torch.Tensor,
    mask_token_id: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.loss.dllm_loss.SCDDLoss.forward(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    noise_mask: torch.Tensor,
    p_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    loss_mask_ar: torch.Tensor | None = None,
    num_diffusion_tokens: int | None = None,
    num_ar_tokens: int | None = None,
    causal_logits: torch.Tensor | None = None,
    noisy_input_ids: torch.Tensor | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

Marginal of the SCDD forward process at a diffusion time.

SCDD (Self-Correcting Discrete Diffusion, openreview.net/forum?id=zQKlzKB6I9) generalises the absorbing masked-diffusion forward process by mixing in uniform transitions, so the denoiser sees corrupted-but-plausible tokens during training and learns to correct them rather than only to fill [MASK]. The marginal of a clean token x at time t is

.. math:: q(z_t \mid x) = \gamma_t\bigl(\rho_t x + (1-\rho_t) u\bigr)

• (1-\gamma_t),m

where u is uniform over the non-[MASK] vocabulary and m is the absorbing [MASK] state.

Compute per-token negative log-likelihood, shape [B, L].

Autoregressive next-token CE on the encoder’s causal logits.

The co-trained encoder loss for diffusion_gemma SFT: a standard causal LM cross-entropy over the clean full sequence, scored where both the current and next position are valid (non-pad).

Parameters:

Returns: torch.Tensor

Scalar AR loss (mean CE over valid next-token positions).

Evaluate the SCDD forward-process marginal at diffusion time t.

The uniform-noise mass follows a Beta-shaped bump c(t) = B t^a (1-t)^b with a = gamma_shape * t_peak and b = gamma_shape * (1 - t_peak), normalised so that its ratio against the retained mass peaks at max_ratio at t = t_peak. The retained mass decays linearly, giving the closed form

clean = (1-t)/(1+c), uniform = c/(1+c), absorbed = t/(1+c).

Both rho and gamma are monotonically decreasing in t, which is what makes [MASK] an absorbing state of the induced Markov chain (no remasking during sampling).

Parameters:

Returns: SCDDSchedule

class:SCDDSchedule at t; every field has shape [batch].