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 MDLMCrossEntropyLoss / 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).

The normalization policy is selected by normalize. Pass num_tokens as a global all-reduced denominator for consistency 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(
    k: int,
    device,
    dtype
) -> torch.Tensor

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._dpace_weight(
    token_nll: torch.Tensor,
    binary_mask: torch.Tensor,
    binary_mask_bool: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._draft_acc_per_pos(
    correct: torch.Tensor,
    block_mask: torch.Tensor,
    report_per_pos: bool = True
) -> typing.Tuple[torch.Tensor | None, torch.Tensor | None]

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._mean_denominator(
    weights: torch.Tensor,
    bsz: int,
    n_blocks: int,
    total_blocks: int | None = None
) -> torch.Tensor

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

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss._reshape_block_inputs(
    values: torch.Tensor,
    target_ids: torch.Tensor,
    block_mask: torch.Tensor,
    block_size: int | None,
    value_name: str
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]

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,
    total_blocks: 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,
    total_blocks: int | None = None
) -> nemo_automodel.components.loss.dllm_loss.DLLMLossOutput

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

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss.position_weights(
    token_nll: torch.Tensor,
    block_mask: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss.weighted_mean(
    values: torch.Tensor,
    token_nll: torch.Tensor,
    block_mask: torch.Tensor,
    value_mask: torch.Tensor | None = None,
    total_blocks: int | None = None
) -> tuple[torch.Tensor, torch.Tensor]

DFlash loss internals reused by the online trainer wrappers.

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 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 torch.utils.checkpoint (the same treatment 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.

Bases: Module

Uno diffusion-adapter loss (Sahoo et al., 2026; arXiv:2609.04010).

Operates on the concatenated [x_t (L) | x_0 (L)] forward output produced under the block-diffusion attention mask, where the LoRA adapter is gated on for the noisy x_t half only. The clean half therefore gives the frozen AR distribution, which serves as the detached teacher; the noisy half is the diffusion student. Both are compared at the same index over the response positions:

.. math:: L = \gamma \cdot \text{TV} + \beta \cdot \text{KL}(p_\text{student} | p_\text{teacher})

• TV — the blockwise total-variation objective of the paper (Sec. 3, after Leviathan et al. 2023, Cor. 3.6): the full L1 distance sum_v |p_student - p_teacher| per position. Computed in position chunks wrapped in torch.utils.checkpoint so only one [chunk, vocab] pair of fp32 probabilities is live at a time.
• KL — reverse KL, computed by KDLoss with the arguments swapped (KDLoss computes KL(P_teacher_arg || P_student_arg)).

The released Uno-Qwen3-8B recipe uses TV only (tv_weight=1, kl_weight=0).

Parameters:

nemo_automodel.components.loss.dllm_loss.UnoDistillLoss.forward(
    logits: 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

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

Per-row L1 distance between the student and teacher softmax distributions.

Parameters:

Returns: torch.Tensor

Tensor of shape [rows] with sum_v |softmax(student) - softmax(teacher)| in float32.

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].