> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.loss.dllm_loss

Loss functions for diffusion LLM (dLLM) training.

All loss classes return :class:`DLLMLossOutput` so the recipe can handle them
uniformly without branching on model type.

## Module Contents

### Classes

| Name                                                                                                         | Description                                                               |
| ------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------- |
| [`BlockDiffusionCrossEntropyLoss`](#nemo_automodel-components-loss-dllm_loss-BlockDiffusionCrossEntropyLoss) | Flat cross-entropy loss for block-diffusion (`diffusion_gemma`) training. |
| [`DFlashDecayLoss`](#nemo_automodel-components-loss-dllm_loss-DFlashDecayLoss)                               | Position-decay cross-entropy loss for DFlash draft model training.        |
| [`DLLMLossOutput`](#nemo_automodel-components-loss-dllm_loss-DLLMLossOutput)                                 | Unified return type for all dLLM loss functions.                          |
| [`HybridDiffusionLLMLoss`](#nemo_automodel-components-loss-dllm_loss-HybridDiffusionLLMLoss)                 | Combined diffusion + optional AR loss for hybrid diffusion LLM models.    |
| [`IDLMLoss`](#nemo_automodel-components-loss-dllm_loss-IDLMLoss)                                             | Introspective DLM all-masked loss (Yu et al., 2026; arXiv:2604.11035).    |
| [`MDLMCrossEntropyLoss`](#nemo_automodel-components-loss-dllm_loss-MDLMCrossEntropyLoss)                     | Cross-entropy loss for MDLM training.                                     |
| [`SCDDLoss`](#nemo_automodel-components-loss-dllm_loss-SCDDLoss)                                             | Discrete-time NELBO for SCDD (openreview\.net/forum?id=zQKlzKB6I9).       |
| [`SCDDSchedule`](#nemo_automodel-components-loss-dllm_loss-SCDDSchedule)                                     | Marginal of the SCDD forward process at a diffusion time.                 |

### Functions

| Name                                                                                         | Description                                                       |
| -------------------------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`_compute_per_token_nll`](#nemo_automodel-components-loss-dllm_loss-_compute_per_token_nll) | Compute per-token negative log-likelihood, shape `[B, L]`.        |
| [`encoder_ar_loss`](#nemo_automodel-components-loss-dllm_loss-encoder_ar_loss)               | Autoregressive next-token CE on the encoder's causal logits.      |
| [`scdd_schedule`](#nemo_automodel-components-loss-dllm_loss-scdd_schedule)                   | Evaluate the SCDD forward-process marginal at diffusion time *t*. |

### Data

[`_SCDD_TINY`](#nemo_automodel-components-loss-dllm_loss-_SCDD_TINY)

### API

```python
class nemo_automodel.components.loss.dllm_loss.BlockDiffusionCrossEntropyLoss(
    fp32_upcast: bool = True
)
```

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

```python
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
```

Compute the flat block-diffusion cross-entropy loss.

**Parameters:**

**`logits`** `torch.Tensor`

Model output logits over the canvas, shape `[B, L, V]`.

---

**`target_ids`** `torch.Tensor`

Clean (uncorrupted) canvas token IDs, shape `[B, L]`.

---

**`noise_mask`** `torch.Tensor`

Boolean mask of corrupted positions, shape `[B, L]`.

---

**`p_mask`** `torch.Tensor`

Ignored (flat loss has no per-token weight).

---

**`loss_mask`** `torch.Tensor`

Supervised positions mask, shape `[B, L]`.

---

**`num_diffusion_tokens`** `int | None` — default: None

If provided, the global corrupted-token count
used as the normalization denominator (summed across grad-acc
microbatches). If `None`, normalizes by the local corrupted
count in this microbatch.

---

**`noisy_input_ids`** `torch.Tensor | None` — default: None

Ignored (the flat loss scores the clean targets),
shape `[B, L]` when supplied.

---

**Returns:** `DLLMLossOutput`

class:`DLLMLossOutput` where `total_loss == dllm_loss` (no AR).

```python
class nemo_automodel.components.loss.dllm_loss.DFlashDecayLoss(
    loss_gamma: float | None = 7.0,
    use_fused_linear_ce: bool = False,
    chunk_size: int = 1024,
    normalize: str = 'tokens'
)
```

**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:**

**`loss_gamma`** `float | None` — default: 7.0

Decay parameter γ.

---

**`use_fused_linear_ce`** `bool` — default: False

When True, compute the per-token NLL with the
chunked linear-CE path (:meth:`forward_fused`) — projects the
LM head and runs cross-entropy in position chunks, each wrapped in
:func:`torch.utils.checkpoint` so the full `[B, T, vocab]` logits
tensor is never materialised (peak is one chunk). Keeps large
`num_blocks_per_sample` (e.g. paper-default 512) within memory on
full-vocab targets.

We deliberately do NOT use `liger_kernel`'s
`LigerFusedLinearCrossEntropyLoss` here: its custom autograd
Function computes `grad_input` eagerly in forward and only
integrates with FSDP via the model-patching redirection
(`apply_liger_kernel_to_*`). Used standalone under FSDP2 the
gradient does not reach the sharded model params (grad\_norm 0).
The chunked path is plain autograd, so FSDP2 handles it correctly.

---

**`chunk_size`** `int` — default: 1024

Number of predicted positions per chunk in the chunked
linear-CE path. Smaller = lower peak memory, more recompute.

---

**`normalize`** `str` — default: 'tokens'

Loss denominator. `"tokens"` (default) divides the
decay-weighted sum by `num_tokens`, a global all-reduced count
that keeps the loss consistent across DP replicas and grad-accum.
`"mean"` divides by the effective weight sum
`(w_k * block_mask).sum()` for a per-call decay-weighted mean.

---

**`loss_gamma`** `float | None` — default: 7.0

Decay parameter γ. `None` disables decay (all predicted
positions weighted equally).

---

**`chunk_size`** `= int(chunk_size)`

---

**`loss_gamma`** `= None if loss_gamma is None else float(loss_gamma)`

---

**`use_fused_linear_ce`** `= bool(use_fused_linear_ce)`

---

```python
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]
```

staticmethod

Project one position chunk; return its per-token NLL and argmax-matches.

Wrapped in :func:`torch.utils.checkpoint` by the caller, so the
`[chunk, vocab]` logits are recomputed in backward rather than held.
The argmax is non-differentiable, so it adds no backward cost.

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

Eq. 4 weights for `T` predicted positions, resetting per block.

Returns all-ones (uniform) when `loss_gamma is None` (decay disabled).

```python
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]
```

staticmethod

Per-rank (correct, count) sums per block offset k=1..block\_size-1.

`correct` is a `[B, T]` bool/float tensor of argmax matches and
`block_mask` excludes padding (T = N \* (block\_size - 1) when
`block_size` is provided). Reshape to `[B, N, block_size-1]` and
sum over `(B, N)` to get per-offset counts of shape
`[block_size-1]`. Returns `(None, None)` when `block_size` is
unknown (single-block / legacy path).

```python
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
```

Apply decay weights + block mask, sum, and normalise.

```python
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
```

Compute the DFlash decay-weighted loss from pre-computed logits.

**Parameters:**

**`logits`** `torch.Tensor`

Draft model logits for the predicted block positions,
shape `[B, T, V]` where `T = N * (block_size - 1)`.

---

**`target_ids`** `torch.Tensor`

Ground-truth token IDs, shape `[B, T]`.

---

**`block_mask`** `torch.Tensor`

Float/bool valid-position mask, shape `[B, T]`.
Zero entries (padding) are excluded from the loss.

---

**`num_tokens`** `int | None` — default: None

Optional global token count for loss normalisation.

---

**`block_size`** `int | None` — default: None

When provided, the decay weights reset at each block
boundary so that every block's first predicted position has
weight 1.  Required for multi-block training (N > 1).

---

**Returns:** `DLLMLossOutput`

class:`DLLMLossOutput`.

```python
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
```

Chunked linear-CE: never materialises the full logits tensor.

Projects the LM head + cross-entropy in chunks of `chunk_size`
predicted positions, each wrapped in :func:`torch.utils.checkpoint` so
the `[chunk, vocab]` logits are recomputed in backward instead of
held — peak logit memory is one chunk, not `[B*T, vocab]`. Pure
autograd, so the gradient flows correctly through FSDP2 (unlike a
standalone liger fused-CE Function).

**Parameters:**

**`hidden`** `torch.Tensor`

Draft hidden states for the predicted positions,
shape `[B, T, D]` (`D` = model dim, NOT vocab).

---

**`lm_head_weight`** `torch.Tensor`

LM-head projection weight, shape `[V, D]`.

---

**`target_ids`** `torch.Tensor`

Ground-truth token IDs, shape `[B, T]`.

---

**`block_mask`** `torch.Tensor`

Valid-position mask, shape `[B, T]`.

---

**`num_tokens / block_size`**

as in :meth:`forward`.

---

**`lm_head_bias`** `torch.Tensor | None` — default: None

Optional LM-head bias, shape `[V]`.

---

**Returns:** `DLLMLossOutput`

class:`DLLMLossOutput`.

```python
class nemo_automodel.components.loss.dllm_loss.DLLMLossOutput()
```

**Bases:** `NamedTuple`

Unified return type for all dLLM loss functions.

**`dllm_loss`** `Tensor`

---

**`draft_correct_per_pos`** `Tensor | None = None`

---

**`draft_count_per_pos`** `Tensor | None = None`

---

**`total_loss`** `Tensor`

---

```python
class nemo_automodel.components.loss.dllm_loss.HybridDiffusionLLMLoss(
    alpha: float = 1.0,
    fp32_upcast: bool = True
)
```

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

```python
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
```

Compute the hybrid diffusion + AR loss.

**Parameters:**

**`logits`** `torch.Tensor`

Model output logits, shape `[B, L, V]` or
`[B, L+L_ar, V]` if the model produces both diffusion and AR
logits in a single concatenated tensor (legacy path).

---

**`target_ids`** `torch.Tensor`

Clean token IDs, shape `[B, L]`.

---

**`noise_mask`** `torch.Tensor`

Boolean mask of corrupted positions, shape `[B, L]`.

---

**`p_mask`** `torch.Tensor`

Per-position masking probability, shape `[B, L]`.

---

**`loss_mask`** `torch.Tensor`

Diffusion loss mask (supervised positions), shape `[B, L]`.

---

**`loss_mask_ar`** `torch.Tensor | None` — default: None

AR loss mask, shape `[B, L]`. If None, no AR loss.

---

**`num_diffusion_tokens`** `int | None` — default: None

Total diffusion label tokens for normalization.

---

**`num_ar_tokens`** `int | None` — default: None

Total AR label tokens for normalization.

---

**`causal_logits`** `torch.Tensor | None` — default: None

Optional separate AR logits, shape `[B, L, V]`.
When provided, avoids the concat/split of the legacy layout.

---

**`noisy_input_ids`** `torch.Tensor | None` — default: None

Ignored (the model applies masking internally),
shape `[B, L]` when supplied.

---

**Returns:** `DLLMLossOutput`

class:`DLLMLossOutput` with combined `total_loss` and the pure

```python
class nemo_automodel.components.loss.dllm_loss.IDLMLoss(
    clean_loss_weight: float = 0.2,
    auto_balance: bool = False
)
```

**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:**

**`clean_loss_weight`** `float` — default: 0.2

Fixed `alpha` for the clean-copy CE.

---

**`auto_balance`** `bool` — default: False

Replace `alpha` with `(CE_noisy / CE_clean).detach()`.

---

**`auto_balance`** `= bool(auto_balance)`

---

**`clean_loss_weight`** `= float(clean_loss_weight)`

---

```python
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
```

Compute the I-DLM block-diffusion loss.

**Parameters:**

**`logits`** `torch.Tensor`

Concatenated forward logits, shape `[B, 2L, V]` ordered
`[x_t | x_0]`.

---

**`target_ids`** `torch.Tensor`

Clean token IDs for one copy, shape `[B, L]`.

---

**`answer_mask`** `torch.Tensor`

Bool mask of supervised (response) positions, `[B, L]`.

---

**`valid_mask`** `torch.Tensor`

Bool/long padding-validity mask, shape `[B, L]`.

---

**`seq_len`** `int`

Length `L` of one copy.

---

**`num_diffusion_tokens`** `int | None` — default: None

Global, DP-all-reduced supervised-token count
used as the loss denominator (summed across grad-accum
microbatches and data-parallel ranks). Pass this so the loss is a
proper global token-mean — required for the recipe's
`(loss * dp_group_size).backward()` scaling to give
DP/grad-accum-invariant gradients. Falls back to the local
supervised count when `None` (single-process use / unit tests).

---

**Returns:** `DLLMLossOutput`

class:`DLLMLossOutput` with the combined `total_loss` and

```python
class nemo_automodel.components.loss.dllm_loss.MDLMCrossEntropyLoss(
    fp32_upcast: bool = True
)
```

**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).

```python
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
```

Compute the MDLM cross-entropy loss.

**Parameters:**

**`logits`** `torch.Tensor`

Model output logits, shape `[B, L, V]`.

---

**`target_ids`** `torch.Tensor`

Clean (uncorrupted) token IDs, shape `[B, L]`.

---

**`noise_mask`** `torch.Tensor`

Boolean mask of corrupted positions, shape `[B, L]`.

---

**`p_mask`** `torch.Tensor`

Per-position masking probability, shape `[B, L]`.

---

**`loss_mask`** `torch.Tensor`

Supervised positions mask, shape `[B, L]`.

---

**`num_diffusion_tokens`** `int | None` — default: None

If provided, used for global normalization
(total supervised tokens across all grad-acc microbatches).

---

**`noisy_input_ids`** `torch.Tensor | None` — default: None

Ignored (the absorbing kernel needs only
`noise_mask`), shape `[B, L]` when supplied.

---

**Returns:** `DLLMLossOutput`

class:`DLLMLossOutput` where `total_loss == dllm_loss`.

```python
class nemo_automodel.components.loss.dllm_loss.SCDDLoss(
    mask_token_id: int,
    num_timesteps: int = 1000,
    max_ratio: float = 0.1,
    gamma_shape: float = 1.0,
    t_peak: float = 0.5,
    chunk_size: int | None = 1024
)
```

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

**`chunk_size`**

---

**`gamma_shape`** `= float(gamma_shape)`

---

**`mask_token_id`** `= int(mask_token_id)`

---

**`max_ratio`** `= float(max_ratio)`

---

**`num_timesteps`** `= int(num_timesteps)`

---

**`t_peak`** `= float(t_peak)`

---

```python
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]
```

staticmethod

Reduce one position chunk over the vocabulary axis.

This is the only part of the ELBO whose working set scales with the
vocabulary, so it is the part the caller wraps in
:func:`torch.utils.checkpoint`: the two `[chunk, vocab]` fp32
intermediates are then recomputed in backward instead of held. Every
position is independent, so chunking is exact.

**Parameters:**

**`logits_chunk`** `torch.Tensor`

Model logits, shape `[chunk, vocab]`.

---

**`x_0_chunk`** `torch.Tensor`

Clean token IDs, shape `[chunk]`.

---

**`z_t_chunk`** `torch.Tensor`

Corrupted token IDs seen by the model, shape `[chunk]`.

---

**`log_base_s_chunk`** `torch.Tensor`

`log` of the uniform base mass at `s`, shape
`[chunk]`.

---

**`log_rho_s_chunk`** `torch.Tensor`

`log` of the retained-mass ratio at `s`, shape
`[chunk]`.

---

**`mask_token_id`** `int`

Token ID of the absorbing `[MASK]` state.

---

**Returns:** `torch.Tensor`

Tuple of `(sum_log, log_at_x0, log_at_zt, log_p_zt)`, each of

```python
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
```

Compute the SCDD discrete-time NELBO.

**Parameters:**

**`logits`** `torch.Tensor`

Model output logits, shape `[batch, sequence, vocab]`.

---

**`target_ids`** `torch.Tensor`

Clean token IDs `x_0`, shape `[batch, sequence]`.

---

**`noise_mask`** `torch.Tensor`

Boolean mask of corrupted positions, shape
`[batch, sequence]`. Ignored — the SCDD ELBO is supported on
every supervised position, corrupted or not.

---

**`p_mask`** `torch.Tensor`

Per-position diffusion time `t`, shape
`[batch, sequence]`, constant along the sequence axis (the
SCDD forward process draws one `t` per sequence). This is the
contract with
:meth:`~nemo_automodel.recipes.dllm.strategy.SCDDStrategy.apply_corruption`,
which samples `t` on the `1/T` grid and broadcasts it here;
unlike the absorbing kernels this slot carries the time itself,
because the ELBO weights need the full schedule at `t` and at
the previous grid point.

---

**`loss_mask`** `torch.Tensor`

Supervised positions mask, shape `[batch, sequence]`.

---

**`loss_mask_ar`** `torch.Tensor | None` — default: None

Ignored (SCDD has no autoregressive term).

---

**`num_diffusion_tokens`** `int | None` — default: None

If provided, the global supervised-token count
used as the normalisation denominator (summed across grad-acc
microbatches). If `None`, normalises by the local supervised
count.

---

**`num_ar_tokens`** `int | None` — default: None

Ignored (SCDD has no autoregressive term).

---

**`causal_logits`** `torch.Tensor | None` — default: None

Ignored (SCDD has no autoregressive term).

---

**`noisy_input_ids`** `torch.Tensor | None` — default: None

Corrupted token IDs `z_t` the model was fed, shape
`[batch, sequence]`. Required: the correction term is a
function of the visible token, which cannot be recovered from
`noise_mask` alone.

---

**Returns:** `DLLMLossOutput`

class:`DLLMLossOutput` where `total_loss == dllm_loss`.

```python
class nemo_automodel.components.loss.dllm_loss.SCDDSchedule(
    clean_mass: torch.Tensor,
    uniform_mass: torch.Tensor,
    absorbed_mass: torch.Tensor,
    gamma: torch.Tensor,
    rho: torch.Tensor
)
```

Dataclass

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.

**`absorbed_mass`** `Tensor`

---

**`clean_mass`** `Tensor`

---

**`gamma`** `Tensor`

---

**`rho`** `Tensor`

---

**`uniform_mass`** `Tensor`

---

```python
nemo_automodel.components.loss.dllm_loss._compute_per_token_nll(
    logits: torch.Tensor,
    target_ids: torch.Tensor
) -> torch.Tensor
```

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

```python
nemo_automodel.components.loss.dllm_loss.encoder_ar_loss(
    encoder_logits: torch.Tensor,
    input_ids: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
    num_tokens: int | None = None
) -> torch.Tensor
```

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:**

**`encoder_logits`** `torch.Tensor`

Encoder logits over the clean sequence, `[B, S, V]`.

---

**`input_ids`** `torch.Tensor`

Clean token IDs, `[B, S]`.

---

**`valid_mask`** `torch.Tensor | None` — default: None

Boolean non-pad mask `[B, S]`. If `None`, all positions count.

---

**`num_tokens`** `int | None` — default: None

Optional global denominator (summed across grad-acc microbatches);
defaults to the local valid next-token count.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.loss.dllm_loss.scdd_schedule(
    t: torch.Tensor,
    max_ratio: float,
    gamma_shape: float,
    t_peak: float
) -> nemo_automodel.components.loss.dllm_loss.SCDDSchedule
```

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:**

**`t`** `torch.Tensor`

Diffusion times in `[0, 1]`, shape `[batch]`. Values outside the
unit interval are clamped (fractional powers of a negative base are
undefined).

---

**`max_ratio`** `float`

Peak uniform-to-retained mass ratio, in `[0, 1)`. `0`
degenerates the process to pure absorbing masked diffusion (MDLM).

---

**`gamma_shape`** `float`

Total shape mass of the bump; larger values concentrate the
uniform noise around *t\_peak*.

---

**`t_peak`** `float`

Time in `(0, 1)` at which the uniform-noise ratio peaks.

---

**Returns:** `SCDDSchedule`

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

```python
nemo_automodel.components.loss.dllm_loss._SCDD_TINY = 1e-30
```