diff --git a/docs/guides/dllm/finetune.mdx b/docs/guides/dllm/finetune.mdx index 3f238a6820..0d0f8b362a 100644 --- a/docs/guides/dllm/finetune.mdx +++ b/docs/guides/dllm/finetune.mdx @@ -31,7 +31,7 @@ The following table outlines the key steps in the fine-tuning workflow: | Step | Section | What You Do | |------|---------|-------------| -| **1. Install** | [Install NeMo AutoModel](#install-nemo-automodel) | Install the package with uv or Docker | +| **1. Install** | [Install NeMo AutoModel](#install-nemo-automodel) | Install the package with `uv` or Docker | | **2. Configure** | [Configure Your Training Recipe](#configure-your-training-recipe) | Write a YAML config specifying model, data, dLLM mode, and training settings | | **3. Train** | [Fine-Tune the Model](#fine-tune-the-model) | Launch training with `torchrun` | | **4. Generate** | [Run Inference](#run-inference) | Generate text from a fine-tuned checkpoint | @@ -44,6 +44,7 @@ The following table lists the supported models, their training modes, loss funct |---|---|---|---|---| | LLaDA | `mdlm` | MDLM cross-entropy | Standalone full-forward denoising without a key-value (KV) cache | [llada_sft.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/llada_sft.yaml) | | LLaDA2 | `mdlm` | MDLM cross-entropy | Built-in block-refinement generation | [llada2_sft.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/llada2_sft.yaml) | +| SCDD (LLaDA backbone) | `scdd` | Self-correcting discrete-diffusion NELBO (denoise + correction terms) | Ancestral sampling over the whole canvas, with per-step self-correction | [llada_scdd.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/llada_scdd.yaml) | | Nemotron-Labs-Diffusion | `hybrid` | Diffusion and AR (alpha-weighted) | Block diffusion with KV cache | [nemotron_labs_diffusion_sft.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/nemotron_labs_diffusion_sft.yaml) | | DiffusionGemma | `block_diffusion` | Flat block-diffusion cross-entropy and encoder AR | Built-in Hugging Face diffusion sampler (entropy-bounded denoising) | [diffusion_gemma_sft.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/diffusion_gemma_sft.yaml) | | DFlash | `dflash` | Decay-weighted cross-entropy (Equation 4) | Training only (decoding occurs in the speculative-decoding stack) | [dflash_sft.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/dflash_sft.yaml) | @@ -61,7 +62,7 @@ source .venv/bin/activate uv pip install "nemo-automodel" ``` -Alternatively, use the prebuilt Docker container: +Alternatively, use the prebuilt Docker container. ```bash docker pull nvcr.io/nvidia/nemo-automodel:26.06.00 @@ -77,11 +78,12 @@ The following components drive dLLM fine-tuning: 1. A recipe script ([`train_ft.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/dllm/train_ft.py)) orchestrates the training loop with dLLM-specific corruption, loss, and batch handling. 2. A YAML configuration file specifies the model, data, optimizer, dLLM-specific settings, and distributed training strategy. -The recipe uses a strategy pattern to handle differences between model families. The `dllm.mode` field in the YAML configuration selects the strategy: +The recipe uses a strategy pattern to handle differences between model families. The `dllm.mode` field in the YAML configuration selects the strategy. | Mode | Strategy | Description | |------|----------|-------------| | `mdlm` | `MDLMStrategy` | LLaDA-style: model receives corrupted tokens, MDLM cross-entropy loss | +| `scdd` | `SCDDStrategy` | SCDD: absorbing `[MASK]` noise mixed with uniform token transitions, trained with the self-correcting NELBO | | `hybrid` | `HybridStrategy` | Nemotron-Labs-Diffusion-style: model receives clean tokens and `masked_indices`, with combined diffusion and AR loss | | `block_diffusion` | `BlockDiffusionStrategy` | DiffusionGemma-style: uniform random-token corruption over a response canvas, with flat cross-entropy and co-trained encoder AR loss | | `dflash` | `DFlashStrategy` | DFlash: frozen target LM provides hidden states, and the draft model trains with decay-weighted loss | @@ -106,6 +108,55 @@ dataset: unshifted: true # Required for dLLM training ``` +### Configure SCDD + +SCDD (see the [SCDD paper](https://openreview.net/forum?id=zQKlzKB6I9)) trains self-correction into the model +instead of bolting it on at inference. Its forward process mixes the usual absorbing `[MASK]` noise +with **uniform token transitions**, so the model is trained on contexts that contain +wrong-but-plausible tokens and learns to overwrite them. The objective adds a correction term at +every visible position on top of the familiar denoising term at `[MASK]` positions, which is what +lets the sampler decode many tokens per step without the quality collapse a pure absorbing model +shows under parallel decoding. + +Because the objective covers every supervised position (corrupted or not), the loss denominator is +the supervised-token count, not the corrupted-token count. + +Refer to [`llada_scdd.yaml`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/llada_scdd.yaml) +for the full working configuration. The following example shows the key dLLM-specific sections. + +```yaml +model: + pretrained_model_name_or_path: GSAI-ML/LLaDA-8B-Base + torch_dtype: float32 + trust_remote_code: true + +dllm: + mode: scdd + mask_token_id: 126336 # LLaDA mask token + vocab_size: 126464 # Required: the uniform channel samples over the vocabulary + eps: 0.001 # Minimum diffusion time + num_timesteps: 1000 # Discrete diffusion steps T + uniform_ratio: 0.1 # Peak uniform-noise share (0 degenerates to MDLM) + schedule_shape: 1.0 + schedule_peak: 0.5 + +dataset: + unshifted: true +``` + +The schedule values above match the config shipped with the authors' released checkpoint. Their +only other released setting differs solely in `uniform_ratio: 0.2`, so that is the first knob to +sweep. + +Decode SCDD checkpoints with `--sampler scdd`, passing the same `uniform_ratio`, `schedule_shape`, +and `schedule_peak` used during training. The sampler rebuilds the reverse posterior from that +schedule, so a mismatch degrades generation. + +Unlike the absorbing losses, the SCDD objective needs the model's probability for every +non-`[MASK]` token, so it cannot use a fused cross-entropy kernel. That vocabulary reduction runs in +checkpointed chunks of `dllm.chunk_size` positions, which keeps its fp32 intermediates off the +backward tape. Lower `chunk_size` before reducing sequence length if the loss runs out of memory. + ### Configure Nemotron-Labs-Diffusion Refer to [`nemotron_labs_diffusion_sft.yaml`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/nemotron_labs_diffusion_sft.yaml) for the full working configuration. The following example shows the key dLLM-specific sections: @@ -160,9 +211,8 @@ The following table describes the key configuration fields for dLLM fine-tuning: | Field | Description | |-------|-------------| -| `dllm.mode` | Training strategy (`mdlm`, `hybrid`, `block_diffusion`, `dflash`, or `idlm`) | +| `dllm.mode` | Training strategy (`mdlm`, `scdd`, `hybrid`, `block_diffusion`, `dflash`, or `idlm`) | | `dllm.mask_token_id` | Token ID used for masking (`126336` for LLaDA, `156895` for LLaDA2.1, `100` for Nemotron-Labs-Diffusion). Unused by `block_diffusion`. | -| `dllm.vocab_size` | Vocabulary size for uniform-random token corruption. Required for `block_diffusion` (`262144` for DiffusionGemma). | | `dllm.eps` | Minimum corruption ratio to avoid zero-corruption samples | | `dllm.block_size` | Hybrid: when set, use blockwise corruption (otherwise uniform). Block-diffusion: response-block size for one-canvas-per-step selection (default `256`). | | `dllm.encoder_loss_weight` | Weight on the co-trained encoder AR loss. Block-diffusion only (default `1.0`). | @@ -172,6 +222,12 @@ The following table describes the key configuration fields for dLLM fine-tuning: | `dllm.block_length` | Diffusion block size for the I-DLM block-diffusion mask (paper curriculum 1→2→3). I-DLM mode only. | | `dllm.clean_loss_weight` | Fixed `α` on the clean-copy verification CE (paper `0.2`). I-DLM mode only. | | `dllm.auto_balance_clean_loss` | Replaces `α` with `(CE_noisy/CE_clean).detach()` (paper Equation 2, b3 stage). I-DLM mode only. | +| `dllm.vocab_size` | Vocabulary size, including `[MASK]`. Required by `scdd` and `block_diffusion`, whose corruption draws replacement tokens over the vocabulary (`262144` for DiffusionGemma). | +| `dllm.num_timesteps` | Number of discrete diffusion steps `T` in the SCDD NELBO (default `1000`). SCDD mode only. | +| `dllm.uniform_ratio` | Peak share of uniform (correctable) noise. `0` degenerates SCDD to plain MDLM. SCDD mode only. | +| `dllm.schedule_shape` | Shape mass of the uniform-noise bump; larger values concentrate the noise near the peak. SCDD mode only. | +| `dllm.schedule_peak` | Time in `(0, 1)` at which the uniform-noise ratio peaks. SCDD mode only. | +| `dllm.chunk_size` | Positions per checkpointed chunk of the SCDD loss's vocabulary reduction. Lower it first if the loss runs out of memory; `null` disables chunking. SCDD mode only. | | `dataset.unshifted` | Must be `true` for dLLM. Disables the autoregressive input and target shift. | ### Configure DFlash @@ -229,7 +285,7 @@ Both metrics are computed from the logits that the chunked linear cross-entropy The DFlash paper recommends training on responses regenerated by the target model, as described in Section 5.1. Rather than directly using the original dataset, you can construct the training set with responses generated by the target model to achieve better target alignment. Skipping this step trains the draft model on a different output distribution than the target model produces at inference, which directly reduces the acceptance length. -The existing `nemo_automodel.components.speculative.regenerate` script handles this process. Start an SGLang server that hosts the target model, and then regenerate the assistant turns: +The existing `nemo_automodel.components.speculative.regenerate` script handles this process. Start an SGLang server that hosts the target model, and then regenerate the assistant turns. ```bash # 1. Serve the target model on the local node (default port 30000) @@ -257,8 +313,7 @@ You can then point the recipe configuration `dataset.path_or_dataset_id` to the I-DLM ([Introspective Diffusion LM](https://arxiv.org/abs/2604.11035), Yu et al., 2026) converts a pretrained autoregressive LM into a diffusion LM by all-masked fine-tuning. Each step concatenates a fully masked copy `x_t` and the clean copy `x_0` into a length-`2L` sequence run under a block-diffusion attention mask, with a Dream-style next-token logit shift. Two cross-entropy terms, both over the response tokens, are combined: `CE_noisy` (decode `q`, the masked copy conditioned on the clean ground-truth prefix) and `CE_clean` (verify `p`, the clean copy under strict causal attention). -See [`qwen3_8b_idlm.yaml`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/qwen3_8b_idlm.yaml) for the full -working config. The key I-DLM-specific sections are: +See [`qwen3_8b_idlm.yaml`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_sft/qwen3_8b_idlm.yaml) for the full working configuration. The following example shows the key I-DLM-specific sections. ```yaml model: @@ -285,11 +340,19 @@ The mask is built for `sdpa` or `eager` (dense additive) or `flex_attention` (sp ### Fine-Tune LLaDA2 ```bash -torchrun --nproc-per-node=8 \ +uv run torchrun --nproc-per-node=8 \ examples/dllm_sft/finetune.py \ -c examples/dllm_sft/llada2_sft.yaml ``` +### Fine-Tune with SCDD + +```bash +uv run torchrun --nproc-per-node=8 \ + examples/dllm_sft/finetune.py \ + -c examples/dllm_sft/llada_scdd.yaml +``` + ### Fine-Tune DiffusionGemma Prepare the GSM8K chat JSONL once, then launch full SFT or LoRA. See the [DiffusionGemma Fine-Tuning Guide](/recipes-e2e-examples/diffusiongemma) for the training objective and LoRA target modules. @@ -305,7 +368,7 @@ torchrun --standalone --nproc-per-node=8 \ ### Fine-Tune with DFlash ```bash -torchrun --nproc-per-node=8 \ +uv run torchrun --nproc-per-node=8 \ examples/dllm_sft/finetune.py \ -c examples/dllm_sft/dflash_sft.yaml ``` @@ -313,7 +376,7 @@ torchrun --nproc-per-node=8 \ ### Fine-Tune with I-DLM ```bash -torchrun --nproc-per-node=8 \ +uv run torchrun --nproc-per-node=8 \ nemo_automodel/recipes/dllm/train_ft.py \ -c examples/dllm_sft/qwen3_8b_idlm.yaml ``` @@ -321,32 +384,49 @@ torchrun --nproc-per-node=8 \ ### Fine-Tune Nemotron-Labs-Diffusion ```bash -torchrun --nproc-per-node=8 \ +uv run torchrun --nproc-per-node=8 \ nemo_automodel/recipes/dllm/train_ft.py \ -c examples/dllm_sft/nemotron_labs_diffusion_sft.yaml ``` ## Run Inference -The generation script ([`generate.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_generate/generate.py)) supports chat and raw generation. Select the sampler that matches the trained family by using the `--sampler {llada,llada2,nemotron,gemma,idlm}` argument. Infilling (`--infill`) is available with the `llada` sampler only. +The generation script ([`generate.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/dllm_generate/generate.py)) supports chat and raw generation. Select the sampler that matches the trained family by using the `--sampler {llada,scdd,llada2,nemotron,gemma,idlm}` argument. Infilling (`--infill`) is available with the `llada` sampler only. The `--checkpoint` argument accepts several path types, including a path to a `consolidated/` directory, a step directory such as `.../epoch_0_step_499`, or the top-level checkpoint directory. The script automatically resolves the path to `LATEST/model/consolidated/`. Training with the provided example configs automatically writes this consolidated Hugging Face-format directory at the final checkpoint (`checkpoint.save_consolidated: final`). You can pass the directory printed at the end of training directly to `--checkpoint`. ### Generate with LLaDA ```bash -python examples/dllm_generate/generate.py \ +uv run python examples/dllm_generate/generate.py \ --checkpoint \ --prompt "Explain what a neural network is." \ --sampler llada ``` +### Generate with SCDD + +The SCDD sampler resamples every generated position at each step from the exact reverse posterior, +so a token it now believes is wrong can be replaced. `--uniform_ratio`, `--schedule_shape`, and +`--schedule_peak` must match the training configuration. `--block_size`, `--remasking`, +`--threshold`, and KV caching do not apply. + +```bash +uv run python examples/dllm_generate/generate.py \ + --checkpoint dllm_checkpoints/llada_scdd//model/consolidated \ + --prompt "Explain what a neural network is." \ + --sampler scdd \ + --steps 128 \ + --max_new_tokens 128 \ + --uniform_ratio 0.1 +``` + ### Generate with LLaDA2 LLaDA2 generation calls the model's built-in block-refinement `generate()` method. ```bash -python examples/dllm_generate/generate.py \ +uv run python examples/dllm_generate/generate.py \ --checkpoint \ --prompt "Explain what a neural network is." \ --sampler llada2 \ @@ -359,7 +439,7 @@ python examples/dllm_generate/generate.py \ ### Generate with Nemotron-Labs-Diffusion ```bash -python examples/dllm_generate/generate.py \ +uv run python examples/dllm_generate/generate.py \ --checkpoint \ --prompt "Explain what a neural network is." \ --sampler nemotron @@ -371,7 +451,7 @@ DiffusionGemma generation calls the diffusion sampler that ships with `transform (entropy-bounded denoising with adaptive stopping). ```bash -python examples/dllm_generate/generate.py \ +uv run python examples/dllm_generate/generate.py \ --checkpoint \ --prompt "Explain what a neural network is." \ --sampler gemma diff --git a/examples/dllm_generate/generate.py b/examples/dllm_generate/generate.py index e5c4b1a6e0..5366a3b5ba 100644 --- a/examples/dllm_generate/generate.py +++ b/examples/dllm_generate/generate.py @@ -17,6 +17,7 @@ Provides ``DLLMSampler`` (core logic) with preset subclasses: - ``LLaDASampler``: no-cache, full-forward defaults. +- ``SCDDSampler``: self-correcting ancestral sampling over the whole canvas. - ``LLaDA2Sampler``: built-in block-refinement generation defaults. - ``NemotronLabsDLLMSampler``: KV-cache block-diffusion defaults. - ``IDLMSampler``: I-DLM introspective strided decoding (Dream logit shift). @@ -32,6 +33,13 @@ --prompt "Explain what a neural network is." \ --sampler llada +SCDD generation (the schedule flags must match the training config):: + + python examples/dllm_generate/generate.py \ + --checkpoint \ + --prompt "Explain what a neural network is." \ + --sampler scdd --steps 128 --uniform_ratio 0.1 + I-DLM generation (``--mask_id`` is the reserved token used at training, e.g. 151669 for the Qwen3-based I-DLM checkpoint):: @@ -109,6 +117,13 @@ trim_response, ) +from nemo_automodel.components.loss.dllm_loss import scdd_schedule + +# Probability floor for the SCDD reverse posterior: the schedule hits exact +# zeros at t = 1 (no retained mass) and t = 0 (no absorbed mass), and those +# divisions are only ever taken on the branch the result is discarded from. +_SCDD_EPS = 1e-30 + # --------------------------------------------------------------------------- # Sampler config # --------------------------------------------------------------------------- @@ -127,6 +142,14 @@ class SamplerConfig: threshold: float | None = None causal_context: bool = False eos_token_id: int | None = None + # SCDD forward-process hyperparameters; must match the values the checkpoint + # was trained with (``dllm.uniform_ratio`` / ``schedule_shape`` / + # ``schedule_peak`` / ``eps`` in the training config). Ignored by the + # absorbing samplers. + uniform_ratio: float = 0.1 + schedule_shape: float = 1.0 + schedule_peak: float = 0.5 + eps: float = 1e-3 # --------------------------------------------------------------------------- @@ -531,6 +554,182 @@ class IDLMSampler(DLLMSampler): ) +class SCDDSampler(DLLMSampler): + """SCDD ancestral sampler (openreview.net/forum?id=zQKlzKB6I9). + + Where the absorbing samplers commit tokens irreversibly — each step unmasks + a few positions and never revisits them — SCDD resamples *every* generated + position from the exact reverse posterior at each step. A token the model + now believes is wrong can be replaced by a better one; that self-correction + is what keeps quality up when many positions are decoded per step. + + Two cases make up the posterior, both derived from the mixed forward + schedule (:func:`~nemo_automodel.components.loss.dllm_loss.scdd_schedule`): + + * a visible token can stay or be rewritten to another visible token. It is + never sent back to ``[MASK]``: ``[MASK]`` is absorbing in the forward + process, so the reverse posterior puts no mass on it, + * a ``[MASK]`` position un-absorbs into the denoiser's distribution or stays + masked. + + A final argmax denoise at the last time point (SCDD's ``noise_removal``) + rewrites every generated position, clearing residual ``[MASK]`` and any + wrong-but-visible tokens the uniform channel left behind. + + Unsupported knobs: ``use_kv_cache`` (every position is rewritten each step, + so nothing can be cached), ``block_size``, ``remasking`` and ``threshold`` + (there is no top-k transfer schedule to gate). ``temperature`` sharpens the + denoiser distribution before the posterior is formed; ``0`` makes it + one-hot. + """ + + default_config = SamplerConfig( + steps=128, + max_new_tokens=128, + block_size=128, + temperature=1.0, + remasking="low_confidence", + use_kv_cache=False, + threshold=None, + causal_context=False, + eos_token_id=None, + ) + + def _denoiser_log_probs(self, x: torch.Tensor, attention_mask: torch.Tensor, temperature: float) -> torch.Tensor: + """Run the model and return log ``p_theta(. | x)`` over non-``[MASK]`` tokens. + + Args: + x: Current token sequence, shape ``[batch, sequence]``. + attention_mask: Binary attention mask, shape ``[batch, sequence]``. + temperature: Softmax temperature. ``0`` yields a one-hot argmax + distribution (in log space, ``0`` and ``-inf``). + + Returns: + Log-probabilities of shape ``[batch, sequence, vocab]`` with + ``-inf`` at the ``[MASK]`` column. + """ + logits = self.model(x, attention_mask=attention_mask).logits.float() + mask_col = torch.tensor([self.mask_id], device=logits.device) + logits = logits.index_fill(-1, mask_col, float("-inf")) + if temperature == 0.0: + return torch.log_softmax( + torch.full_like(logits, float("-inf")).scatter(-1, logits.argmax(-1, keepdim=True), 0.0), + dim=-1, + ) + return torch.log_softmax(logits / temperature, dim=-1) + + @torch.no_grad() + def sample( + self, + inputs, + config: SamplerConfig | None = None, + **overrides, + ) -> torch.Tensor: + """Generate by iterating the SCDD reverse posterior over the full canvas. + + Args: + inputs: List of prompt token tensors (each shape ``[prompt]``) or + lists of token IDs. + config: Full config. If ``None``, uses :attr:`default_config`. + **overrides: Override individual fields on the config. + + Returns: + Token tensor of shape ``[batch, max_prompt + max_new_tokens]``. + """ + cfg = config or self.default_config + if overrides: + cfg = replace(cfg, **overrides) + if cfg.use_kv_cache: + raise ValueError("SCDD resamples every position each step; use_kv_cache is not supported.") + + if isinstance(inputs[0], list): + inputs = [torch.as_tensor(p, dtype=torch.long, device=self.device) for p in inputs] + prompt_lens = [p.shape[0] for p in inputs] + max_prompt_len = max(prompt_lens) + B = len(inputs) + gen_length = cfg.max_new_tokens + T = max_prompt_len + gen_length + + x = torch.full((B, T), self.eos_id, dtype=torch.long, device=self.device) + attention_mask = torch.zeros((B, T), dtype=torch.long, device=self.device) + # Positions the sampler owns: the generation window of each prompt. Left + # of it is the prompt; right of it is padding for shorter prompts. + canvas = torch.zeros((B, T), dtype=torch.bool, device=self.device) + for i, p in enumerate(inputs): + x[i, : prompt_lens[i]] = p + end = min(prompt_lens[i] + gen_length, T) + x[i, prompt_lens[i] : end] = self.mask_id + canvas[i, prompt_lens[i] : end] = True + attention_mask[i, :end] = 1 + + timesteps = torch.linspace(1.0, cfg.eps, cfg.steps + 1, device=self.device) + schedule_kwargs = { + "max_ratio": cfg.uniform_ratio, + "gamma_shape": cfg.schedule_shape, + "t_peak": cfg.schedule_peak, + } + + # Every sequence shares the same time grid, so the schedule is scalar and + # only the canvas positions are ever rewritten. Working on the flattened + # canvas rows keeps the posterior at [canvas_tokens, vocab] instead of + # [batch, sequence, vocab], and keeps the prompt out of it: at t = 1 the + # visible-token posterior there is identically zero (all mass sits on the + # absorbing state), which multinomial rejects. + canvas_flat = canvas.reshape(-1) + + for i in range(cfg.steps): + sched_t = scdd_schedule(timesteps[i].reshape(1), **schedule_kwargs) + sched_s = scdd_schedule(timesteps[i + 1].reshape(1), **schedule_kwargs) + + clean_transition = sched_t.clean_mass / sched_s.clean_mass.clamp(min=_SCDD_EPS) + uniform_transition = sched_t.gamma / sched_s.gamma.clamp(min=_SCDD_EPS) - clean_transition + absorbing_transition = 1.0 - clean_transition - uniform_transition + + log_p = self._denoiser_log_probs(x, attention_mask, cfg.temperature) + vocab = log_p.size(-1) + # Domain of the denoiser: the vocabulary minus the absorbing state. + num_states = vocab - 1 + p_theta = log_p.reshape(-1, vocab)[canvas_flat].exp() # [canvas_tokens, vocab] + del log_p + current = x[canvas] # [canvas_tokens] + p_at_current = p_theta.gather(-1, current[:, None]) # [canvas_tokens, 1] + + # --- visible token: stay, or be rewritten (the correction channel) --- + denom = sched_t.uniform_mass / num_states + sched_t.clean_mass * p_at_current + numer = (sched_s.clean_mass * uniform_transition / num_states) * p_theta + ( + sched_s.uniform_mass * uniform_transition + ) / num_states**2 + stay = (sched_s.clean_mass * clean_transition) * p_at_current + ( + clean_transition * sched_s.uniform_mass + ) / num_states + numer = numer.scatter_add(-1, current[:, None], stay) + numer[:, self.mask_id] = 0.0 + visible_probs = numer / denom.clamp(min=_SCDD_EPS) + + # --- [MASK]: un-absorb into the denoiser, or stay masked --- + release = absorbing_transition / sched_t.absorbed_mass.clamp(min=_SCDD_EPS) + masked_probs = release * sched_s.uniform_mass / num_states + (release * sched_s.clean_mass) * p_theta + masked_probs[:, self.mask_id] = sched_s.absorbed_mass / sched_t.absorbed_mass.clamp(min=_SCDD_EPS) + + rows = torch.where((current == self.mask_id)[:, None], masked_probs, visible_probs).clamp(min=0) + # A posterior with no reachable state (possible only for a degenerate + # schedule, e.g. uniform_ratio=0) keeps the current token. + keep = torch.zeros_like(rows).scatter(-1, current[:, None], 1.0) + rows = torch.where(rows.sum(-1, keepdim=True) > 0, rows, keep) + x = x.clone() + x[canvas] = torch.multinomial(rows, 1).squeeze(-1) + + # Final noise-removal denoise: argmax every canvas position, not just the + # residual [MASK]. At t = eps the uniform channel still leaves wrong-but- + # visible tokens that only this pass corrects; the mask column is -inf, so + # leftover [MASK] resolves too. The prompt is off-canvas and untouched. + if canvas.any(): + final = self._denoiser_log_probs(x, attention_mask, 0.0).argmax(-1) + x = torch.where(canvas, final, x) + + return x + + class DiffusionGemmaSampler(DLLMSampler): """Config-preset holder for DiffusionGemma generation. @@ -551,6 +750,7 @@ class DiffusionGemmaSampler(DLLMSampler): SAMPLERS = { "llada": LLaDASampler, + "scdd": SCDDSampler, "llada2": LLaDA2Sampler, "nemotron": NemotronLabsDLLMSampler, "idlm": IDLMSampler, @@ -661,6 +861,24 @@ def main(): parser.add_argument("--temperature", type=float, default=None) parser.add_argument("--remasking", default=None, choices=["low_confidence", "random"]) parser.add_argument("--threshold", type=float, default=None) + parser.add_argument( + "--uniform_ratio", + type=float, + default=None, + help="SCDD peak uniform-noise ratio; must match dllm.uniform_ratio used in training.", + ) + parser.add_argument( + "--schedule_shape", + type=float, + default=None, + help="SCDD uniform-noise bump shape; must match dllm.schedule_shape used in training.", + ) + parser.add_argument( + "--schedule_peak", + type=float, + default=None, + help="SCDD uniform-noise peak time; must match dllm.schedule_peak used in training.", + ) parser.add_argument( "--mask_id", type=int, @@ -689,6 +907,8 @@ def main(): parser.error("--infill is not supported by the Nemotron generation path (the tokenizer has no mask token)") if args.infill and args.sampler == "gemma": parser.error("--infill is not supported by the DiffusionGemma generation path") + if args.infill and args.sampler == "scdd": + parser.error("--infill is not supported by the SCDD sampler (it has no top-k transfer schedule)") try: checkpoint_path = resolve_checkpoint(args.checkpoint) @@ -724,6 +944,9 @@ def main(): "temperature", "remasking", "threshold", + "uniform_ratio", + "schedule_shape", + "schedule_peak", ]: val = getattr(args, key) if val is not None: diff --git a/examples/dllm_generate/utils.py b/examples/dllm_generate/utils.py index 5ae791f7ef..208cc031e9 100644 --- a/examples/dllm_generate/utils.py +++ b/examples/dllm_generate/utils.py @@ -141,9 +141,9 @@ def load_model_and_tokenizer(checkpoint_path: str, sampler_name: str = "llada", Args: checkpoint_path: Path to the HF-format checkpoint directory. - sampler_name: ``"llada"``, ``"llada2"``, ``"nemotron"``, ``"idlm"``, or - ``"gemma"``. Adjusts tokenizer setup and model construction kwargs - for the chosen family. + sampler_name: ``"llada"``, ``"scdd"``, ``"llada2"``, ``"nemotron"``, + ``"idlm"``, or ``"gemma"``. Adjusts tokenizer setup and model + construction kwargs for the chosen family. mask_id_override: Explicit mask token id. Takes precedence over the tokenizer/config lookup. Required for I-DLM, whose base Qwen3 tokenizer has no mask token (training reuses a reserved id). @@ -199,7 +199,9 @@ def load_model_and_tokenizer(checkpoint_path: str, sampler_name: str = "llada", return model.eval(), tokenizer, None, tokenizer.eos_token_id - if sampler_name == "llada": + if sampler_name in ("llada", "scdd"): + # SCDD fine-tunes LLaDA-family checkpoints, whose tokenizer ships the + # mask token only in the model config. if tokenizer.mask_token is None: tokenizer.add_special_tokens({"mask_token": "<|mdm_mask|>"}) diff --git a/examples/dllm_sft/llada_scdd.yaml b/examples/dllm_sft/llada_scdd.yaml new file mode 100644 index 0000000000..8e8520bdd0 --- /dev/null +++ b/examples/dllm_sft/llada_scdd.yaml @@ -0,0 +1,175 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SCDD (Self-Correcting Discrete Diffusion) fine-tuning of LLaDA-8B. +# Paper: https://openreview.net/forum?id=zQKlzKB6I9 +# +# The forward process mixes absorbing [MASK] noise with uniform token +# transitions, so the model trains on contexts holding wrong-but-plausible +# tokens and learns to overwrite them. That self-correction is what sustains +# quality when many positions are decoded per step. +# +# Scope note: the paper *pretrains* GPT-2-scale DiT backbones from scratch with +# this objective. This recipe instead applies it as SFT on LLaDA-8B, which was +# pretrained with the pure absorbing kernel, so fine-tuning has to teach +# self-correction on top of an absorbing-only prior. Treat the schedule settings +# below as a starting point, not as reproduced paper numbers. +# +# The forward-process block below matches the config shipped with the authors' +# released checkpoint: T=1000, ratio 0.1, gamma 1, t_peak 0.5, sampling_eps 1e-3. +# Their only other released setting differs solely in ratio: 0.2 -- every other +# field is identical, so ratio is the one schedule knob worth sweeping first. +# +# Their optimizer settings are NOT copied here: lr 5e-4, weight_decay 0.02, +# eps 1e-9, global batch 256, sequence length 512 all belong to a from-scratch +# pretraining run of a 90M-parameter DiT. The values below mirror llada_sft.yaml +# instead, so an SCDD run is directly comparable against the MDLM baseline on the +# same 8B model and data. +# +# Known gap: the authors' runs keep an EMA of the weights (decay 0.9999) and +# evaluate the EMA copy. AutoModel ships ``components/training/ema.py``, but the +# dLLM recipe does not wire it, so there is no ``ema:`` key to set here; expect +# somewhat noisier samples than the released checkpoints until it is threaded in. +# +# Usage (8-GPU single node): +# python -m torch.distributed.run --nproc-per-node=8 \ +# nemo_automodel/recipes/dllm/train_ft.py \ +# -c examples/dllm_sft/llada_scdd.yaml +# +# Decode with the matching sampler, passing the same schedule: +# python examples/dllm_generate/generate.py \ +# --checkpoint --sampler scdd \ +# --uniform_ratio 0.1 --schedule_shape 1.0 --schedule_peak 0.5 + +recipe: DiffusionLMSFTRecipe + +step_scheduler: + global_batch_size: 32 # per-device 1 x grad-accum 4 x 8 GPUs + local_batch_size: 1 + ckpt_every_steps: 500 + val_every_steps: 500 + max_steps: null + num_epochs: 2 + +dist_env: + backend: nccl + timeout_minutes: 30 + +seed: 42 + +wandb: + enable: false + project: dllm-scdd + name: am-llada-scdd_run + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: GSAI-ML/LLaDA-8B-Base + torch_dtype: float32 # fp32 master weights; compute stays bf16 via mp_policy + trust_remote_code: true + +checkpoint: + enabled: true + checkpoint_dir: dllm_checkpoints/llada_scdd/ + model_save_format: safetensors + save_consolidated: final # final step exports the consolidated HF dir generate.py loads + +distributed: + strategy: fsdp2 + dp_size: none + tp_size: 1 + cp_size: 1 # SCDD does not support context parallelism (SCDDStrategy raises) + pp_size: 1 + sequence_parallel: false + activation_checkpointing: true # 8B full FT; trade compute for memory + mp_policy: + param_dtype: bfloat16 + reduce_dtype: float32 + output_dtype: float32 + autocast_dtype: bfloat16 + offload_policy: null + # Per-layer FSDP2 prefetch overlaps all-gather / reduce-scatter with compute. + enable_fsdp2_prefetch: true + fsdp2_forward_prefetch_depth: 1 + fsdp2_backward_prefetch_depth: 2 + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +# mode: scdd selects SCDDStrategy -- mixed absorbing + uniform corruption scored +# by the discrete-time SCDD NELBO. The ELBO is supported at every supervised +# position (corrupted or not), so the loss denominator is the supervised-token +# count rather than the corrupted-token count. +dllm: + mode: scdd + mask_token_id: 126336 # LLaDA <|mdm_mask|> + # Required: the uniform channel draws replacements over the vocabulary minus + # [MASK] and the clean token, and the ELBO normalises over that same domain. + # Must equal the model's vocab_size -- SCDDStrategy.setup_extra rejects a mismatch. + vocab_size: 126464 + # --- schedule: matches the authors' released checkpoint config --- + eps: 0.001 # their training.sampling_eps + num_timesteps: 1000 # their T; the ELBO compares t against t - 1/T + # Peak share of uniform (correctable) noise -- their forward.ratio. 0 + # degenerates the objective to plain MDLM. The authors release 0.1 and 0.2; + # start at 0.1 when fine-tuning an absorbing-pretrained checkpoint, it being + # the smaller shift away from LLaDA's prior. + uniform_ratio: 0.1 + schedule_shape: 1.0 # their forward.gamma; larger concentrates noise at the peak + schedule_peak: 0.5 # their forward.t_peak + # Positions per checkpointed chunk of the loss's vocabulary reduction. The + # SCDD ELBO needs p_theta over the whole vocabulary, so it cannot use a fused + # cross-entropy; chunking keeps the two [chunk, vocab] fp32 intermediates off + # the backward tape. Lower this first if the loss runs out of memory; null + # disables chunking. + chunk_size: 1024 + supervise_padding: true + +optimizer: + _target_: torch.optim.AdamW + betas: [0.9, 0.999] + eps: 1.0e-8 + lr: 2.0e-5 # matches llada_sft.yaml so SCDD vs MDLM stays comparable + weight_decay: 0.0 + +clip_grad_norm: + max_norm: 1.0 + +lr_scheduler: + lr_decay_style: cosine + init_lr: 0.0 + min_lr: 0.0 + lr_warmup_steps: 200 + +dataset: + _target_: nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset + path_or_dataset_id: allenai/tulu-3-sft-mixture + split: train + shuffle_seed: 42 + seq_length: 1024 + truncation: true + unshifted: true + # LLaDA-8B-Base ships a generation-only chat template that unconditionally + # appends an assistant prompt and lacks {% generation %} tags, breaking answer-mask + # extraction on multi-turn samples. Use a Llama-3 template with generation tags so + # the tokenizer returns the assistant mask directly (supervises response + EOS). + chat_template: examples/dllm_sft/llada_chat_template.jinja + tokenizer: + pretrained_model_name_or_path: GSAI-ML/LLaDA-8B-Base + trust_remote_code: true + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: nemo_automodel.components.datasets.utils.default_collater + group_by_length: true diff --git a/nemo_automodel/components/datasets/dllm/corruption.py b/nemo_automodel/components/datasets/dllm/corruption.py index 5b371fa559..45b132b99c 100644 --- a/nemo_automodel/components/datasets/dllm/corruption.py +++ b/nemo_automodel/components/datasets/dllm/corruption.py @@ -19,6 +19,7 @@ - ``corrupt_blockwise``: per-block weighted corruption with exponential position bias - ``corrupt_uniform_random``: per-block random-token (D3PM-uniform) corruption - ``corrupt_all_masked``: deterministic all-masked corruption (I-DLM) +- ``corrupt_mix``: two-channel absorbing + uniform-transition corruption (SCDD) """ from __future__ import annotations @@ -151,6 +152,123 @@ def corrupt_all_masked( return noisy_input_ids, noise_mask, p_mask +def corrupt_mix( + input_ids: torch.Tensor, + loss_mask: torch.Tensor, + mask_token_id: int, + vocab_size: int, + *, + mask_prob: torch.Tensor, + uniform_prob: torch.Tensor, + generator: torch.Generator | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Two-channel absorbing + uniform-transition corruption. + + Each supervised position independently lands in exactly one of three states, + drawn from a single uniform variate so the two channels are mutually + exclusive: + + * ``[MASK]`` with probability ``mask_prob`` (the absorbing channel), + * a **different** non-``[MASK]`` token with probability ``uniform_prob`` + (the uniform-transition channel), drawn uniformly over the + ``vocab_size - 2`` tokens that are neither ``[MASK]`` nor the clean token, + * unchanged otherwise. + + This is the forward kernel that gives a diffusion LM corrupted-but-plausible + context to correct, as opposed to the pure absorbing kernel of + :func:`corrupt_uniform` where every corrupted position is visibly ``[MASK]``. + Callers own the schedule that turns a diffusion time into the two + probabilities; this function applies whatever it is given. + + The replacement token is drawn by sampling an index over ``vocab_size - 2`` + and shifting it past the two excluded ids, so peak memory stays + ``O(batch * sequence)`` rather than materialising a ``[batch, sequence, + vocab]`` probability table. + + Args: + input_ids: Clean token IDs, shape ``[batch, sequence]``. + loss_mask: Binary mask of supervised positions, shape + ``[batch, sequence]``. Only supervised positions are ever corrupted. + mask_token_id: Token ID of the absorbing ``[MASK]`` state. + vocab_size: Vocabulary size, including ``[MASK]``. Must be at least 3. + mask_prob: Per-sequence absorbing probability, shape ``[batch]`` or + ``[batch, 1]`` (broadcast over positions). + uniform_prob: Per-sequence uniform-transition probability, same shape as + *mask_prob*. ``mask_prob + uniform_prob`` must not exceed 1. + generator: Optional ``torch.Generator`` (on ``input_ids.device``) used for + ALL random draws (the channel variate and the replacement tokens). + Pass a step-seeded generator so the corruption is a deterministic + function of the training step and reproduces exactly on checkpoint + resume; ``None`` falls back to the global RNG (not resume-safe). + + Returns: + Tuple of ``(noisy_input_ids, noise_mask)``, each of shape + ``[batch, sequence]``. + + * ``noisy_input_ids`` — ``input_ids`` with absorbed positions replaced by + ``mask_token_id`` and transitioned positions replaced by a different + non-``[MASK]`` token. + * ``noise_mask`` — bool mask of positions changed by either channel. + + No ``p_mask`` is returned: the two probabilities alone do not determine + the loss weight for a mixed kernel, so the caller supplies whatever + per-position quantity its loss needs. + + Raises: + ValueError: If ``vocab_size < 3``, or if any ``mask_prob + uniform_prob`` + exceeds 1. + + Note: + Supervised positions whose clean token already **is** ``mask_token_id`` + are never routed to the uniform channel (there would be no well-defined + "different non-``[MASK]``" replacement); they may still be absorbed, + which leaves them unchanged and hence out of ``noise_mask``. + """ + if vocab_size < 3: + raise ValueError(f"corrupt_mix requires vocab_size >= 3 (got {vocab_size})") + + B, L = input_ids.shape + device = input_ids.device + + mask_prob = mask_prob.reshape(B, 1).to(torch.float32) + uniform_prob = uniform_prob.reshape(B, 1).to(torch.float32) + # The two channels are carved out of a single [0, 1) variate, so anything + # past 1 is unrepresentable: the uniform channel would be silently truncated + # and the realised noise would no longer match the schedule the loss weights + # assume. Costs one host sync per call, which is worth an explicit error. + channel_split = mask_prob + uniform_prob + if bool((channel_split > 1.0 + 1e-5).any()): + raise ValueError( + f"corrupt_mix requires mask_prob + uniform_prob <= 1 (got up to {channel_split.max().item():.6f})" + ) + + # One variate per position splits the two channels without double-drawing. + u = torch.rand((B, L), device=device, generator=generator) + supervised = loss_mask.bool() + absorbed = (u < mask_prob) & supervised + transitioned = (u >= mask_prob) & (u < channel_split) & supervised + transitioned &= input_ids != mask_token_id + + # Uniform over the vocabulary minus {mask_token_id, input_ids}: draw over + # vocab_size - 2 slots, then shift past each excluded id in sorted order. + draw = torch.randint(0, vocab_size - 2, (B, L), device=device, dtype=input_ids.dtype, generator=generator) + mask_id_t = torch.full_like(input_ids, mask_token_id) + lo = torch.minimum(input_ids, mask_id_t) + hi = torch.maximum(input_ids, mask_id_t) + draw = draw + (draw >= lo).to(draw.dtype) + draw = draw + (draw >= hi).to(draw.dtype) + # A clean token that already is [MASK] gives lo == hi, so both shifts fire on + # the same draw and the result can reach vocab_size. Those positions are + # excluded from `transitioned` above and are never selected below, but clamp + # so the tensor cannot carry an out-of-range id regardless. + draw = draw.clamp_(max=vocab_size - 1) + + noisy_input_ids = torch.where(absorbed, mask_id_t, input_ids) + noisy_input_ids = torch.where(transitioned, draw, noisy_input_ids) + + return noisy_input_ids, absorbed | transitioned + + def corrupt_blockwise( input_ids: torch.Tensor, loss_mask: torch.Tensor, diff --git a/nemo_automodel/components/loss/dllm_loss.py b/nemo_automodel/components/loss/dllm_loss.py index b4003f6bb1..504aaabd94 100644 --- a/nemo_automodel/components/loss/dllm_loss.py +++ b/nemo_automodel/components/loss/dllm_loss.py @@ -20,6 +20,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import NamedTuple, Tuple import torch @@ -27,6 +28,15 @@ import torch.nn.functional as F from torch.distributed.tensor import DTensor +from nemo_automodel.components.loss.chunked_ce import _validate_chunk_len + +# Probability floor used throughout the SCDD schedule/ELBO. Quantities that are +# exactly zero at the schedule boundaries (rho -> 1 gives a zero uniform base) +# are clamped to this before a log, which keeps every term finite; the resulting +# bias is far below fp32 resolution and the affected terms carry a zero +# coefficient anyway. +_SCDD_TINY = 1e-30 + def _compute_per_token_nll( logits: torch.Tensor, @@ -127,6 +137,7 @@ def forward( 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, ) -> DLLMLossOutput: """Compute the MDLM cross-entropy loss. @@ -138,6 +149,8 @@ def forward( loss_mask: Supervised positions mask, shape ``[B, L]``. num_diffusion_tokens: If provided, used for global normalization (total supervised tokens across all grad-acc microbatches). + noisy_input_ids: Ignored (the absorbing kernel needs only + ``noise_mask``), shape ``[B, L]`` when supplied. Returns: :class:`DLLMLossOutput` where ``total_loss == dllm_loss``. @@ -161,6 +174,379 @@ def forward( return DLLMLossOutput(total_loss=loss, dllm_loss=loss.detach().clone()) +@dataclass(frozen=True) +class SCDDSchedule: + """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. + + Attributes: + clean_mass: ``gamma_t * rho_t`` — probability the token is *retained*, + shape ``[batch]``. + uniform_mass: ``gamma_t * (1 - rho_t)`` — probability the token was + redrawn from the uniform distribution, shape ``[batch]``. + absorbed_mass: ``1 - gamma_t`` — probability the token is ``[MASK]``, + shape ``[batch]``. + gamma: Probability the token is not ``[MASK]``, shape ``[batch]``. + rho: Probability the token is retained given that it is not ``[MASK]``, + shape ``[batch]``. + """ + + clean_mass: torch.Tensor + uniform_mass: torch.Tensor + absorbed_mass: torch.Tensor + gamma: torch.Tensor + rho: torch.Tensor + + +def scdd_schedule( + t: torch.Tensor, + *, + max_ratio: float, + gamma_shape: float, + t_peak: float, +) -> 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). + + Args: + t: Diffusion times in ``[0, 1]``, shape ``[batch]``. Values outside the + unit interval are clamped (fractional powers of a negative base are + undefined). + max_ratio: Peak uniform-to-retained mass ratio, in ``[0, 1)``. ``0`` + degenerates the process to pure absorbing masked diffusion (MDLM). + gamma_shape: Total shape mass of the bump; larger values concentrate the + uniform noise around *t_peak*. + t_peak: Time in ``(0, 1)`` at which the uniform-noise ratio peaks. + + Returns: + The :class:`SCDDSchedule` at *t*; every field has shape ``[batch]``. + """ + if not 0.0 <= max_ratio < 1.0: + raise ValueError(f"scdd_schedule requires 0 <= max_ratio < 1 (got {max_ratio})") + if not 0.0 < t_peak < 1.0: + raise ValueError(f"scdd_schedule requires 0 < t_peak < 1 (got {t_peak})") + + t = t.clamp(0.0, 1.0) + a = gamma_shape * t_peak + b = gamma_shape * (1.0 - t_peak) + peak = (t_peak**a) * ((1.0 - t_peak) ** b) + scale = (max_ratio / (1.0 - max_ratio)) / peak + + c = scale * torch.pow(t, a) * torch.pow(1.0 - t, b) + clean_mass = (1.0 - t) / (1.0 + c) + uniform_mass = c / (1.0 + c) + absorbed_mass = 1.0 - clean_mass - uniform_mass + gamma = clean_mass + uniform_mass + rho = clean_mass / gamma.clamp(min=_SCDD_TINY) + + return SCDDSchedule( + clean_mass=clean_mass, + uniform_mass=uniform_mass, + absorbed_mass=absorbed_mass, + gamma=gamma, + rho=rho, + ) + + +class SCDDLoss(nn.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. + """ + + def __init__( + self, + 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, + ): + """Initialise the SCDD loss. + + Args: + mask_token_id: Token ID of the absorbing ``[MASK]`` state. + num_timesteps: Number of discrete diffusion steps ``T``; the loss is + the ``T``-step NELBO and the reverse step is ``1/T``. At least 2, + so the grid holds a usable point below the fully absorbed ``t = 1``. + max_ratio: Peak uniform-to-retained mass ratio of the forward + process (``0`` degenerates to MDLM). + gamma_shape: Shape mass of the uniform-noise bump. + t_peak: Time at which the uniform-noise ratio peaks. + chunk_size: Number of positions whose vocabulary-sized terms are + computed at once, each chunk wrapped in + :func:`torch.utils.checkpoint`. Smaller means lower peak memory + and more recompute. ``None`` computes every position in one + shot with no checkpointing — numerically identical, but it holds + two fp32 ``[batch * sequence, vocab]`` tensors at once. + """ + super().__init__() + if num_timesteps < 2: + raise ValueError(f"SCDDLoss requires num_timesteps >= 2 (got {num_timesteps})") + self.mask_token_id = int(mask_token_id) + self.num_timesteps = int(num_timesteps) + self.max_ratio = float(max_ratio) + self.gamma_shape = float(gamma_shape) + self.t_peak = float(t_peak) + # Same positive-int contract the chunked cross-entropy kernel uses. + self.chunk_size = None if chunk_size is None else _validate_chunk_len(chunk_size) + + @staticmethod + def _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]: + """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. + + Args: + logits_chunk: Model logits, shape ``[chunk, vocab]``. + x_0_chunk: Clean token IDs, shape ``[chunk]``. + z_t_chunk: Corrupted token IDs seen by the model, shape ``[chunk]``. + log_base_s_chunk: ``log`` of the uniform base mass at ``s``, shape + ``[chunk]``. + log_rho_s_chunk: ``log`` of the retained-mass ratio at ``s``, shape + ``[chunk]``. + mask_token_id: Token ID of the absorbing ``[MASK]`` state. + + Returns: + Tuple of ``(sum_log, log_at_x0, log_at_zt, log_p_zt)``, each of + shape ``[chunk]``: + + * ``sum_log`` — the posterior numerator summed over the non-``[MASK]`` + domain. + * ``log_at_x0`` / ``log_at_zt`` — that numerator at the clean and at + the corrupted token. + * ``log_p_zt`` — the denoiser's log-probability of the corrupted token. + """ + # Driving the [MASK] logit to -inf removes the absorbing state from the + # denoiser's domain. The fill is done in the logits' own dtype so only + # the float() cast below pays vocabulary-sized fp32. + mask_col = torch.tensor([mask_token_id], device=logits_chunk.device) + logits_chunk = logits_chunk.index_fill(-1, mask_col, float("-inf")).float() # [chunk, vocab] + log_denom = torch.logsumexp(logits_chunk, dim=-1) # [chunk] + + # log p_theta(v) = logits(v) - logsumexp(logits), so the log-softmax + # never has to be materialised: its per-position normaliser folds into a + # scalar shift, and it is only needed pointwise at z_t. + shift = log_rho_s_chunk - log_denom # [chunk] + log_term = torch.logaddexp( + log_base_s_chunk[:, None].expand_as(logits_chunk), + shift[:, None] + logits_chunk, + ) # [chunk, vocab] — log( base_s + rho_s * p_theta(v) ) + + # At the [MASK] column the logit is -inf, so the term collapses to + # log(base_s): excluding that column from the vocabulary sum is a scalar + # subtraction, not a gather. + sum_log = log_term.sum(dim=-1) - log_base_s_chunk + log_at_x0 = log_term.gather(-1, x_0_chunk[:, None]).squeeze(-1) + log_at_zt = log_term.gather(-1, z_t_chunk[:, None]).squeeze(-1) + log_p_zt = logits_chunk.gather(-1, z_t_chunk[:, None]).squeeze(-1) - log_denom + return sum_log, log_at_x0, log_at_zt, log_p_zt + + def forward( + self, + 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, + ) -> DLLMLossOutput: + """Compute the SCDD discrete-time NELBO. + + Args: + logits: Model output logits, shape ``[batch, sequence, vocab]``. + target_ids: Clean token IDs ``x_0``, shape ``[batch, sequence]``. + noise_mask: Boolean mask of corrupted positions, shape + ``[batch, sequence]``. Ignored — the SCDD ELBO is supported on + every supervised position, corrupted or not. + p_mask: 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: Supervised positions mask, shape ``[batch, sequence]``. + loss_mask_ar: Ignored (SCDD has no autoregressive term). + num_diffusion_tokens: 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: Ignored (SCDD has no autoregressive term). + causal_logits: Ignored (SCDD has no autoregressive term). + noisy_input_ids: 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: + :class:`DLLMLossOutput` where ``total_loss == dllm_loss``. + """ + del noise_mask, loss_mask_ar, num_ar_tokens, causal_logits + + if noisy_input_ids is None: + raise ValueError("SCDDLoss requires noisy_input_ids (the corrupted tokens z_t seen by the model).") + + if isinstance(logits, DTensor): + logits = logits.full_tensor() + z_t = noisy_input_ids.to(logits.device) + x_0 = target_ids.to(logits.device) + + vocab = logits.size(-1) + # Domain of the denoiser: every token except the absorbing state. + num_states = vocab - 1 + + # --- schedule at t and at the previous grid point s = t - 1/T --- + step = 1.0 / self.num_timesteps + t = p_mask[:, 0].to(torch.float32).clamp(0.0, 1.0) # [batch] + s = (t - step).clamp(min=0.0) + sched_t = scdd_schedule(t, max_ratio=self.max_ratio, gamma_shape=self.gamma_shape, t_peak=self.t_peak) + sched_s = scdd_schedule(s, max_ratio=self.max_ratio, gamma_shape=self.gamma_shape, t_peak=self.t_peak) + + rho_t, rho_s = sched_t.rho, sched_s.rho + rho_s_safe = rho_s.clamp(min=_SCDD_TINY) + # Backward transition of the retained/uniform split between s and t. + clean_transition = rho_t / rho_s_safe + uniform_transition = (rho_s - rho_t) / rho_s_safe + # Fraction of the absorbed mass released over one reverse step. + unmask_coeff = (sched_s.gamma - sched_t.gamma) / (1.0 - sched_t.gamma).clamp(min=_SCDD_TINY) + base_s = (1.0 - rho_s) / num_states + base_t = (1.0 - rho_t) / num_states + + # --- vocabulary-sized work, one position chunk at a time --- + batch, seq_len = x_0.shape + # Broadcast the per-sequence schedule onto positions so a chunk can span + # the batch boundary. + log_base_s = base_s.clamp(min=_SCDD_TINY).log().repeat_interleave(seq_len) # [batch * sequence] + log_rho_s = rho_s.clamp(min=_SCDD_TINY).log().repeat_interleave(seq_len) # [batch * sequence] + flat_logits = logits.reshape(-1, vocab) + flat_x0 = x_0.reshape(-1) + flat_zt = z_t.reshape(-1) + del logits + + num_positions = flat_logits.size(0) + chunk = num_positions if self.chunk_size is None else self.chunk_size + parts = [] + for start in range(0, num_positions, chunk): + end = start + chunk + args = ( + flat_logits[start:end], + flat_x0[start:end], + flat_zt[start:end], + log_base_s[start:end], + log_rho_s[start:end], + self.mask_token_id, + ) + if self.chunk_size is None: + parts.append(self._vocab_terms(*args)) + else: + parts.append(torch.utils.checkpoint.checkpoint(self._vocab_terms, *args, use_reentrant=False)) + sum_log, log_at_x0, log_at_zt, log_p_zt = (torch.cat(term).reshape(batch, seq_len) for term in zip(*parts)) + + # --- z_t == [MASK]: standard denoising term --- + absorbed_loss = -unmask_coeff[:, None] * (base_s[:, None] * sum_log + rho_s[:, None] * log_at_x0) + + # --- z_t != [MASK]: correction term --- + log_denom = torch.logaddexp( + base_t.clamp(min=_SCDD_TINY).log()[:, None].expand_as(log_p_zt), + rho_t.clamp(min=_SCDD_TINY).log()[:, None] + log_p_zt, + ) # [batch, sequence] + + # Expectation of log(q/p) over z_s ~ q(. | z_t, x_0), expanded into the + # four (uniform|retained) x (x_0|z_t) coefficient blocks. + retained = (z_t == x_0).to(log_denom.dtype) # [batch, sequence] + coeff_uniform = (uniform_transition / num_states)[:, None] + coeff_clean = clean_transition[:, None] + total = ( + (base_s[:, None] * coeff_uniform) * (sum_log - num_states * log_denom) + + (rho_s[:, None] * coeff_uniform) * (log_at_x0 - log_denom) + + (base_s[:, None] * coeff_clean) * (log_at_zt - log_denom) + + (rho_s[:, None] * coeff_clean * retained) * (log_at_x0 - log_denom) + ) + correction_loss = -total / (base_t[:, None] + rho_t[:, None] * retained).clamp(min=_SCDD_TINY) + + per_token = torch.where(z_t == self.mask_token_id, absorbed_loss, correction_loss) * self.num_timesteps + + mask = loss_mask.bool().to(per_token.dtype) + loss = (per_token * mask).sum() + denom = num_diffusion_tokens if num_diffusion_tokens is not None else int(mask.sum().item()) + loss = loss / max(denom, 1) + + return DLLMLossOutput(total_loss=loss, dllm_loss=loss.detach().clone()) + + class BlockDiffusionCrossEntropyLoss(nn.Module): """Flat cross-entropy loss for block-diffusion (``diffusion_gemma``) training. @@ -200,6 +586,7 @@ def forward( 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, ) -> DLLMLossOutput: """Compute the flat block-diffusion cross-entropy loss. @@ -213,6 +600,8 @@ def forward( used as the normalization denominator (summed across grad-acc microbatches). If ``None``, normalizes by the local corrupted count in this microbatch. + noisy_input_ids: Ignored (the flat loss scores the clean targets), + shape ``[B, L]`` when supplied. Returns: :class:`DLLMLossOutput` where ``total_loss == dllm_loss`` (no AR). @@ -266,6 +655,7 @@ def forward( 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, ) -> DLLMLossOutput: """Compute the hybrid diffusion + AR loss. @@ -282,6 +672,8 @@ def forward( num_ar_tokens: Total AR label tokens for normalization. causal_logits: Optional separate AR logits, shape ``[B, L, V]``. When provided, avoids the concat/split of the legacy layout. + noisy_input_ids: Ignored (the model applies masking internally), + shape ``[B, L]`` when supplied. Returns: :class:`DLLMLossOutput` with combined ``total_loss`` and the pure diff --git a/nemo_automodel/recipes/dllm/strategy.py b/nemo_automodel/recipes/dllm/strategy.py index e9bb781983..676c9ab83b 100644 --- a/nemo_automodel/recipes/dllm/strategy.py +++ b/nemo_automodel/recipes/dllm/strategy.py @@ -44,6 +44,7 @@ from nemo_automodel.components.datasets.dllm.corruption import ( corrupt_all_masked, corrupt_blockwise, + corrupt_mix, corrupt_uniform, corrupt_uniform_random, ) @@ -54,6 +55,8 @@ HybridDiffusionLLMLoss, IDLMLoss, MDLMCrossEntropyLoss, + SCDDLoss, + scdd_schedule, ) logger = logging.getLogger(__name__) @@ -200,6 +203,147 @@ def prepare_batch(self, batch, noisy_input_ids, noise_mask, clean_input_ids): return batch +class SCDDStrategy(DLLMStrategy): + """Strategy for SCDD (Self-Correcting Discrete Diffusion). + + Paper: https://openreview.net/forum?id=zQKlzKB6I9 + + SCDD generalises MDLM by adding a uniform-transition channel to the + absorbing forward process, so the denoiser is trained on contexts that + contain wrong-but-plausible tokens and learns to overwrite them. That + self-correction is what lets it decode many tokens per step without the + quality collapse a pure absorbing model shows under parallel decoding. + + - Loss: :class:`SCDDLoss` — the discrete-time NELBO with a denoising term at + ``[MASK]`` positions and a correction term everywhere else. + - Corruption: :func:`corrupt_mix` driven by :func:`scdd_schedule` at a + diffusion time drawn on the ``1/T`` grid. + - Normalization: ``"supervised"`` — the ELBO is supported on every + supervised position, not only the corrupted ones. + - Batch: like MDLM, the model receives the corrupted tokens as ``input_ids`` + and attends bidirectionally. + + Time conditioning: the SCDD reference backbone takes the noise level as an + input. Pretrained masked-dLLM checkpoints in Automodel (LLaDA and friends) + are time-free — they read the corruption level off the number of visible + ``[MASK]`` tokens — so no time embedding is threaded into the forward pass + here, matching :class:`MDLMStrategy`. The schedule still enters the + objective through the ELBO weights. + + Requires ``dllm.vocab_size`` and ``dllm.mask_token_id``: the uniform channel + samples replacements over the vocabulary minus ``[MASK]``, and the loss + re-parameterises the model output over that same domain. Context parallelism + is unsupported — the ELBO scores the corrupted tokens against the clean + targets, which the recipe keeps unsharded. + """ + + def __init__(self) -> None: + # vocab_size and the schedule hyperparameters are not part of the + # apply_corruption ABC signature, so they are captured from the dllm + # config in create_loss_fn, which the recipe always calls during setup + # before any corruption runs. + self._vocab_size: int | None = None + self._num_timesteps: int = 1000 + self._max_ratio: float = 0.1 + self._gamma_shape: float = 1.0 + self._t_peak: float = 0.5 + + def create_loss_fn(self, dllm_cfg: dict) -> nn.Module: + vocab_size = dllm_cfg.get("vocab_size", None) + if vocab_size is None: + raise ValueError( + "SCDDStrategy requires dllm.vocab_size to be set in the config " + "(the uniform-transition channel draws replacements over the " + "vocabulary excluding [MASK] and the clean token)." + ) + self._vocab_size = int(vocab_size) + self._num_timesteps = int(dllm_cfg.get("num_timesteps", 1000)) + self._max_ratio = float(dllm_cfg.get("uniform_ratio", 0.1)) + self._gamma_shape = float(dllm_cfg.get("schedule_shape", 1.0)) + self._t_peak = float(dllm_cfg.get("schedule_peak", 0.5)) + # Positions per checkpointed chunk of the loss's vocabulary reduction; + # ``null`` in YAML disables chunking. This is the memory knob for long + # sequences on a large vocabulary. + chunk_size = dllm_cfg.get("chunk_size", 1024) + # mask_token_id may still be unresolved here (the recipe falls back to + # the tokenizer); setup_extra below installs the resolved value. + return SCDDLoss( + mask_token_id=int(dllm_cfg.get("mask_token_id", 0)), + num_timesteps=self._num_timesteps, + max_ratio=self._max_ratio, + gamma_shape=self._gamma_shape, + t_peak=self._t_peak, + chunk_size=None if chunk_size is None else int(chunk_size), + ) + + def setup_extra(self, recipe) -> None: + if getattr(recipe.distributed_config, "cp_size", 1) > 1: + raise ValueError("SCDD does not support context parallelism (cp_size must be 1).") + if recipe.mask_token_id is None: + raise ValueError("SCDD requires dllm.mask_token_id, or a tokenizer that resolves a mask token.") + model_config = getattr(recipe.model_parts[0], "config", None) + vocab_size = getattr(model_config, "vocab_size", None) + if vocab_size is not None: + # A wrong id silently corrupts with a real token and trains garbage. + if not 0 <= int(recipe.mask_token_id) < int(vocab_size): + raise ValueError( + f"dllm.mask_token_id={recipe.mask_token_id} is outside the model vocab (size {vocab_size})." + ) + # The uniform channel and the ELBO's non-[MASK] domain must both be + # the model's own output domain, or the objective is inconsistent. + if int(vocab_size) != self._vocab_size: + raise ValueError(f"dllm.vocab_size={self._vocab_size} does not match the model vocab ({vocab_size}).") + # The recipe may only resolve the mask id from the tokenizer, after + # create_loss_fn has already built the loss module. + recipe.dllm_loss_fn.mask_token_id = int(recipe.mask_token_id) + + def apply_corruption( + self, input_ids, loss_mask, mask_token_id, *, eps, block_size, half_life_ratio, generator=None + ): + del block_size, half_life_ratio # SCDD corrupts the whole sequence at one time + if self._vocab_size is None: + raise ValueError("SCDDStrategy.create_loss_fn must run before corruption (it captures dllm.vocab_size).") + + batch = input_ids.shape[0] + # t ~ U(eps, 1) snapped onto the discrete grid {1/T, ..., 1}: SCDD is + # derived in discrete time, and the ELBO weights compare t against the + # previous grid point s = t - 1/T. + u = torch.rand((batch,), device=input_ids.device, generator=generator) + u = (1.0 - eps) * u + eps + t = ((u * self._num_timesteps).to(torch.int64).float() + 1.0) / self._num_timesteps + # Drop the top point t = 1, where the schedule is fully absorbed and rho + # is degenerate. Clamping to the previous grid point rather than to + # 1 - 1e-4 keeps both t and s = t - 1/T on the grid. + t = t.clamp(max=1.0 - 1.0 / self._num_timesteps) + + sched = scdd_schedule(t, max_ratio=self._max_ratio, gamma_shape=self._gamma_shape, t_peak=self._t_peak) + # A uniform draw that lands back on the clean token leaves the position + # unchanged, so only the (K-1)/K share of the uniform mass is routed + # through the "replace with a different token" channel. + num_states = self._vocab_size - 1 + uniform_prob = sched.uniform_mass * (num_states - 1) / num_states + + noisy_input_ids, noise_mask = corrupt_mix( + input_ids, + loss_mask, + mask_token_id, + self._vocab_size, + mask_prob=sched.absorbed_mass, + uniform_prob=uniform_prob, + generator=generator, + ) + # p_mask carries the diffusion time itself (see SCDDLoss): the mixed + # kernel's ELBO weights need the full schedule at t, not a single + # per-position corruption probability. + p_mask = t[:, None].expand_as(input_ids).float() + return noisy_input_ids, noise_mask, p_mask + + def prepare_batch(self, batch, noisy_input_ids, noise_mask, clean_input_ids): + batch["input_ids"] = noisy_input_ids + batch.pop("attention_mask", None) # SCDD models are bidirectional + return batch + + class HybridStrategy(DLLMStrategy): """Strategy for hybrid diffusion + AR models (e.g., Nemotron-Labs-Diffusion). @@ -1008,6 +1152,7 @@ def prepare_batch(self, batch, noisy_input_ids, noise_mask, clean_input_ids): DLLM_STRATEGIES: Dict[str, type] = { "mdlm": MDLMStrategy, + "scdd": SCDDStrategy, "hybrid": HybridStrategy, "idlm": IDLMStrategy, "dflash": DFlashStrategy, diff --git a/nemo_automodel/recipes/dllm/train_ft.py b/nemo_automodel/recipes/dllm/train_ft.py index c9ecb141f2..e55390266b 100644 --- a/nemo_automodel/recipes/dllm/train_ft.py +++ b/nemo_automodel/recipes/dllm/train_ft.py @@ -430,6 +430,9 @@ def _forward_backward_step( num_diffusion_tokens=num_diffusion_tokens, num_ar_tokens=num_ar_tokens if has_causal else None, causal_logits=causal_logits, + # Mixed forward kernels (scdd) score the corrupted token itself, + # which noise_mask alone cannot recover; absorbing losses ignore it. + noisy_input_ids=noisy_input_ids, ) microbatch_loss = loss_result.total_loss dllm_loss = loss_result.dllm_loss.detach().clone() diff --git a/tests/functional_tests/dllm/L2_DLLM_SCDD_Smoke.sh b/tests/functional_tests/dllm/L2_DLLM_SCDD_Smoke.sh new file mode 100644 index 0000000000..ef66b93652 --- /dev/null +++ b/tests/functional_tests/dllm/L2_DLLM_SCDD_Smoke.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -xeuo pipefail # Exit immediately if a command exits with a non-zero status + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="0" + +# Propagate -s flag if PYTEST_PROPAGATE_S is set +PYTEST_S_FLAG="" +if [ "${PYTEST_PROPAGATE_S:-}" = "1" ]; then + PYTEST_S_FLAG="-s" +fi + +# Tiny public checkpoint + committed chat fixture: the point is to prove the +# scdd wiring runs on one GPU, not to converge. Override the model with +# SCDD_SMOKE_MODEL to smoke a real dLLM checkpoint instead (its vocab size must +# then be passed via --dllm.vocab_size and its mask id via --dllm.mask_token_id). +SCDD_SMOKE_MODEL=${SCDD_SMOKE_MODEL:-hf-internal-testing/tiny-random-LlamaForCausalLM} + +python \ +-m coverage run \ +-m pytest $PYTEST_S_FLAG tests/functional_tests/training/test_scdd_smoke.py \ + --config tests/functional_tests/dllm/scdd_smoke.yaml \ + --model.pretrained_model_name_or_path "$SCDD_SMOKE_MODEL" \ + --dataset.tokenizer.pretrained_model_name_or_path "$SCDD_SMOKE_MODEL" \ + --step_scheduler.max_steps 3 diff --git a/tests/functional_tests/dllm/__init__.py b/tests/functional_tests/dllm/__init__.py new file mode 100644 index 0000000000..341a77c5bc --- /dev/null +++ b/tests/functional_tests/dllm/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/functional_tests/dllm/scdd_smoke.yaml b/tests/functional_tests/dllm/scdd_smoke.yaml new file mode 100644 index 0000000000..1ce3303ef3 --- /dev/null +++ b/tests/functional_tests/dllm/scdd_smoke.yaml @@ -0,0 +1,90 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Single-GPU SCDD smoke: a few optimizer steps of the `scdd` dLLM strategy on a +# tiny public checkpoint and a committed 64-row chat fixture, so the whole path +# (corruption -> bidirectional forward -> SCDD NELBO -> backward -> step) runs in +# seconds without a large download. Real training configs live in +# examples/dllm_sft/ -- this one exists to prove the wiring, not to converge. + +recipe: DiffusionLMSFTRecipe + +step_scheduler: + global_batch_size: 4 + local_batch_size: 2 + max_steps: 3 + num_epochs: 1 + ckpt_every_steps: 1000 + val_every_steps: 1000 + +dist_env: + backend: nccl + timeout_minutes: 5 + +seed: 42 + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: hf-internal-testing/tiny-random-LlamaForCausalLM + torch_dtype: float32 + # This smoke runs in fp32 with no mp_policy, so compute stays fp32 too. The + # default backend is flash_attention_2 whenever flash-attn is installed (as in + # the CI container), and FA2 accepts only fp16/bf16 -- pin the dtype-agnostic + # sdpa backend so the run does not depend on whether flash-attn is present. + attn_implementation: sdpa + +checkpoint: + enabled: false + +distributed: + strategy: fsdp2 + dp_size: none + tp_size: 1 + cp_size: 1 + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +dllm: + mode: scdd + # The tiny Llama tokenizer has no mask token; reserve the last vocab slot, + # exactly as a real SCDD run would reserve the checkpoint's [MASK] id. + mask_token_id: 31999 + vocab_size: 32000 + eps: 0.001 + num_timesteps: 1000 + uniform_ratio: 0.1 + schedule_shape: 1.0 + schedule_peak: 0.5 + +optimizer: + _target_: torch.optim.AdamW + lr: 1.0e-4 + +dataset: + _target_: nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset + path_or_dataset_id: tests/functional_tests/dllm/scdd_smoke_data.jsonl + seq_length: 64 + padding: max_length + truncation: true + unshifted: true + # The tiny test tokenizer ships no chat template; reuse the LLaDA one so the + # assistant turn is tagged and the answer-only loss mask is populated. + chat_template: examples/dllm_sft/llada_chat_template.jinja + tokenizer: + pretrained_model_name_or_path: hf-internal-testing/tiny-random-LlamaForCausalLM + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: nemo_automodel.components.datasets.utils.default_collater diff --git a/tests/functional_tests/dllm/scdd_smoke_data.jsonl b/tests/functional_tests/dllm/scdd_smoke_data.jsonl new file mode 100644 index 0000000000..b13ebc7750 --- /dev/null +++ b/tests/functional_tests/dllm/scdd_smoke_data.jsonl @@ -0,0 +1,64 @@ +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} +{"messages": [{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "The capital of France is Paris."}]} +{"messages": [{"role": "user", "content": "Explain gravity in one sentence."}, {"role": "assistant", "content": "Gravity is the mutual attraction between objects with mass."}]} +{"messages": [{"role": "user", "content": "Name three primary colors."}, {"role": "assistant", "content": "Red, green, and blue are the additive primary colors."}]} +{"messages": [{"role": "user", "content": "What does a compiler do?"}, {"role": "assistant", "content": "A compiler translates source code into machine code."}]} +{"messages": [{"role": "user", "content": "Why does the sky look blue?"}, {"role": "assistant", "content": "Shorter blue wavelengths scatter most off air molecules."}]} +{"messages": [{"role": "user", "content": "Define entropy briefly."}, {"role": "assistant", "content": "Entropy measures the disorder of a physical system."}]} +{"messages": [{"role": "user", "content": "What is a prime number?"}, {"role": "assistant", "content": "A prime number is divisible only by one and itself."}]} +{"messages": [{"role": "user", "content": "How do plants make food?"}, {"role": "assistant", "content": "Plants make food from sunlight through photosynthesis."}]} diff --git a/tests/functional_tests/dllm/test_dllm_scdd.py b/tests/functional_tests/dllm/test_dllm_scdd.py new file mode 100644 index 0000000000..ba768a54bd --- /dev/null +++ b/tests/functional_tests/dllm/test_dllm_scdd.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tests.utils.test_utils import run_test_script + +TEST_FOLDER = "dllm" +SCDD_SMOKE_FILENAME = "L2_DLLM_SCDD_Smoke.sh" + + +class TestDLLMSCDD: + def test_scdd_smoke(self): + run_test_script(TEST_FOLDER, SCDD_SMOKE_FILENAME) diff --git a/tests/functional_tests/training/test_scdd_smoke.py b/tests/functional_tests/training/test_scdd_smoke.py new file mode 100644 index 0000000000..8430412892 --- /dev/null +++ b/tests/functional_tests/training/test_scdd_smoke.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-GPU SCDD (self-correcting discrete diffusion) training smoke test.""" + +from __future__ import annotations + +import sys + +import datasets +import pytest +import torch + +from nemo_automodel.components.config._arg_parser import parse_args_and_load_config +from nemo_automodel.components.loss.dllm_loss import SCDDLoss +from nemo_automodel.recipes.dllm.strategy import SCDDStrategy +from nemo_automodel.recipes.dllm.train_ft import DiffusionLMSFTRecipe + +datasets.disable_caching() + + +def _get_cfg_path() -> str: + argv = sys.argv[1:] + for i, tok in enumerate(argv): + if tok in ("--config", "-c"): + if i + 1 >= len(argv): + raise ValueError("Expected a path after --config") + return argv[i + 1] + raise ValueError("Expected --config/-c to be provided by the functional-test launcher") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="SCDD functional test requires CUDA") +def test_scdd_smoke(): + """End-to-end smoke test of ``dllm.mode: scdd``: + + - build the dLLM recipe from a config that selects the SCDD strategy + - assert: the SCDD strategy and loss are the ones actually wired up, with the + mask token id resolved and the schedule taken from the config + - run a couple of training steps + """ + cfg = parse_args_and_load_config(_get_cfg_path()) + recipe = DiffusionLMSFTRecipe(cfg) + recipe.setup() + + assert isinstance(recipe.dllm_strategy, SCDDStrategy) + loss_fn = recipe.dllm_loss_fn + assert isinstance(loss_fn, SCDDLoss) + # setup_extra installs the resolved id; a stale 0 here would corrupt with a + # real token and train garbage without ever failing loudly. + assert loss_fn.mask_token_id == recipe.mask_token_id + assert loss_fn.max_ratio > 0, "uniform_ratio must be > 0 or SCDD degenerates to MDLM" + # The ELBO scores uncorrupted positions too, so the denominator is the + # supervised-token count. + assert recipe.dllm_strategy.normalization_mode == "supervised" + + # Run a very short training loop (max_steps is controlled by the config/CLI overrides). + # Per-step loss/grad_norm finiteness is asserted by the CPU unit tests in + # tests/unit_tests/loss/test_dllm_loss.py; here the loop itself is the check. + recipe.run_train_validation_loop() diff --git a/tests/unit_tests/datasets/dllm/test_corruption.py b/tests/unit_tests/datasets/dllm/test_corruption.py index 30118454d2..1264e877ce 100644 --- a/tests/unit_tests/datasets/dllm/test_corruption.py +++ b/tests/unit_tests/datasets/dllm/test_corruption.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for dLLM corruption functions (corrupt_uniform, corrupt_blockwise, corrupt_uniform_random).""" +"""Tests for dLLM corruption functions (corrupt_uniform, corrupt_blockwise, corrupt_uniform_random, corrupt_mix).""" import pytest import torch @@ -20,6 +20,7 @@ from nemo_automodel.components.datasets.dllm.corruption import ( corrupt_all_masked, corrupt_blockwise, + corrupt_mix, corrupt_uniform, corrupt_uniform_random, gumbel_topk, @@ -328,3 +329,175 @@ def test_corrupt_uniform_generator_does_not_consume_global_rng(self, inputs): torch.manual_seed(0) corrupt_uniform(input_ids, loss_mask, MASK_TOKEN_ID, generator=torch.Generator().manual_seed(7)) assert torch.equal(torch.rand(4), expected) + + +class TestCorruptMix: + """Two-channel absorbing + uniform-transition corruption (SCDD forward kernel).""" + + def test_channels_are_mutually_exclusive_and_gated_by_loss_mask(self, inputs): + input_ids, loss_mask = inputs + noisy, noise_mask = corrupt_mix( + input_ids, + loss_mask, + MASK_TOKEN_ID, + vocab_size=1000, + mask_prob=torch.full((B,), 0.4), + uniform_prob=torch.full((B,), 0.4), + generator=torch.Generator().manual_seed(0), + ) + assert noisy.shape == input_ids.shape + assert noise_mask.dtype == torch.bool + # Nothing outside the supervised span may change. + assert torch.equal(noisy[loss_mask == 0], input_ids[loss_mask == 0]) + assert not noise_mask[loss_mask == 0].any() + # Every flagged position actually changed, and vice versa. + assert torch.equal(noise_mask, noisy != input_ids) + + def test_uniform_replacements_exclude_mask_and_clean_token(self): + # Every supervised position goes through the uniform channel. + input_ids = torch.randint(0, 50, (8, 64)) + loss_mask = torch.ones_like(input_ids) + noisy, noise_mask = corrupt_mix( + input_ids, + loss_mask, + MASK_TOKEN_ID, + vocab_size=1000, + mask_prob=torch.zeros(8), + uniform_prob=torch.ones(8), + generator=torch.Generator().manual_seed(1), + ) + assert noise_mask.all() + assert (noisy != MASK_TOKEN_ID).all() + assert (noisy != input_ids).all() + assert (noisy >= 0).all() and (noisy < 1000).all() + + def test_replacements_cover_the_allowed_support_uniformly(self): + # vocab 8, mask id 7, clean token 0 -> replacements must be {1..6}. + input_ids = torch.zeros(1, 60000, dtype=torch.long) + noisy, _ = corrupt_mix( + input_ids, + torch.ones_like(input_ids), + mask_token_id=7, + vocab_size=8, + mask_prob=torch.zeros(1), + uniform_prob=torch.ones(1), + generator=torch.Generator().manual_seed(2), + ) + counts = torch.bincount(noisy.flatten(), minlength=8).float() + assert counts[0] == 0 and counts[7] == 0 + expected = input_ids.numel() / 6 + assert torch.allclose(counts[1:7], torch.full((6,), expected), rtol=0.05) + + def test_channel_marginals_match_requested_probabilities(self): + input_ids = torch.randint(0, 50, (2, 40000)) + loss_mask = torch.ones_like(input_ids) + noisy, noise_mask = corrupt_mix( + input_ids, + loss_mask, + MASK_TOKEN_ID, + vocab_size=1000, + mask_prob=torch.tensor([0.3, 0.1]), + uniform_prob=torch.tensor([0.5, 0.2]), + generator=torch.Generator().manual_seed(3), + ) + absorbed = (noisy == MASK_TOKEN_ID).float().mean(dim=1) + transitioned = noise_mask.float().mean(dim=1) - absorbed + assert torch.allclose(absorbed, torch.tensor([0.3, 0.1]), atol=0.01) + assert torch.allclose(transitioned, torch.tensor([0.5, 0.2]), atol=0.01) + + def test_clean_mask_tokens_are_never_routed_to_the_uniform_channel(self): + # A clean token that already IS [MASK] has no well-defined "different + # non-[MASK]" replacement, so it must be left alone by that channel. + input_ids = torch.full((1, 32), MASK_TOKEN_ID) + noisy, noise_mask = corrupt_mix( + input_ids, + torch.ones_like(input_ids), + MASK_TOKEN_ID, + vocab_size=1000, + mask_prob=torch.zeros(1), + uniform_prob=torch.ones(1), + generator=torch.Generator().manual_seed(4), + ) + assert torch.equal(noisy, input_ids) + assert not noise_mask.any() + + def test_same_seed_is_identical_and_global_rng_untouched(self, inputs): + input_ids, loss_mask = inputs + kwargs = dict( + vocab_size=1000, + mask_prob=torch.full((B,), 0.3), + uniform_prob=torch.full((B,), 0.3), + ) + torch.manual_seed(0) + expected = torch.rand(4) + torch.manual_seed(0) + a = corrupt_mix(input_ids, loss_mask, MASK_TOKEN_ID, generator=torch.Generator().manual_seed(9), **kwargs) + assert torch.equal(torch.rand(4), expected) + b = corrupt_mix(input_ids, loss_mask, MASK_TOKEN_ID, generator=torch.Generator().manual_seed(9), **kwargs) + for x, y in zip(a, b): + assert torch.equal(x, y) + + def test_rejects_degenerate_vocab(self, inputs): + input_ids, loss_mask = inputs + with pytest.raises(ValueError, match="vocab_size >= 3"): + corrupt_mix( + input_ids, + loss_mask, + MASK_TOKEN_ID, + vocab_size=2, + mask_prob=torch.zeros(B), + uniform_prob=torch.zeros(B), + ) + + def test_rejects_channel_probabilities_beyond_one(self, inputs): + """The two channels split one [0, 1) variate, so their sum cannot exceed 1. + + Without the check the uniform channel is silently truncated and the + realised noise stops matching the schedule the ELBO weights assume. + """ + input_ids, loss_mask = inputs + with pytest.raises(ValueError, match="mask_prob \\+ uniform_prob <= 1"): + corrupt_mix( + input_ids, + loss_mask, + MASK_TOKEN_ID, + vocab_size=1000, + mask_prob=torch.full((B,), 0.7), + uniform_prob=torch.full((B,), 0.4), + ) + + def test_probabilities_summing_to_exactly_one_are_accepted(self, inputs): + """A schedule whose masses sum to 1 must not trip the tolerance.""" + input_ids, loss_mask = inputs + noisy, _ = corrupt_mix( + input_ids, + loss_mask, + MASK_TOKEN_ID, + vocab_size=1000, + mask_prob=torch.full((B,), 0.3), + uniform_prob=torch.full((B,), 0.7), + ) + assert noisy.shape == input_ids.shape + + def test_ids_stay_in_range_when_every_clean_token_is_mask(self): + """Clean ``[MASK]`` tokens make the two exclusion shifts collide. + + ``lo == hi`` there, so an unclamped draw could reach ``vocab_size``. + Those positions are never routed to the uniform channel, so they must + come back unchanged and every id must stay addressable. + """ + vocab_size, mask_id = 16, 5 + input_ids = torch.full((B, L), mask_id, dtype=torch.long) + loss_mask = torch.ones(B, L, dtype=torch.long) + noisy, noise_mask = corrupt_mix( + input_ids, + loss_mask, + mask_id, + vocab_size=vocab_size, + mask_prob=torch.zeros(B), + uniform_prob=torch.ones(B), + generator=torch.Generator().manual_seed(0), + ) + assert torch.equal(noisy, input_ids) + assert not noise_mask.any() + assert int(noisy.max()) < vocab_size diff --git a/tests/unit_tests/loss/test_dllm_loss.py b/tests/unit_tests/loss/test_dllm_loss.py index 3cbf616d60..c596edcfd3 100644 --- a/tests/unit_tests/loss/test_dllm_loss.py +++ b/tests/unit_tests/loss/test_dllm_loss.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for dLLM loss functions (MDLMCrossEntropyLoss, DFlashDecayLoss).""" +"""Tests for dLLM loss functions (MDLMCrossEntropyLoss, DFlashDecayLoss, SCDDLoss).""" import pytest import torch @@ -24,7 +24,9 @@ DLLMLossOutput, HybridDiffusionLLMLoss, MDLMCrossEntropyLoss, + SCDDLoss, _compute_per_token_nll, + scdd_schedule, ) # --------------------------------------------------------------------------- @@ -786,3 +788,421 @@ def test_global_normalization_denominator(self): # global denominator = 10 r_global = loss_fn(logits, target_ids, noise_mask, p_mask, loss_mask, num_diffusion_tokens=10) assert torch.allclose(r_global.total_loss, summed / 10, atol=1e-6) + + +# --------------------------------------------------------------------------- +# SCDD schedule + loss +# --------------------------------------------------------------------------- + + +class TestSCDDSchedule: + """Closed-form properties the SCDD derivation depends on.""" + + def test_boundaries_are_clean_and_fully_absorbed(self): + t = torch.tensor([0.0, 1.0]) + s = scdd_schedule(t, max_ratio=0.3, gamma_shape=1.0, t_peak=0.5) + assert torch.allclose(s.clean_mass, torch.tensor([1.0, 0.0])) + assert torch.allclose(s.uniform_mass, torch.tensor([0.0, 0.0])) + assert torch.allclose(s.absorbed_mass, torch.tensor([0.0, 1.0])) + assert torch.allclose(s.gamma, torch.tensor([1.0, 0.0])) + + def test_masses_form_a_distribution(self): + s = scdd_schedule(torch.linspace(0, 1, 51), max_ratio=0.25, gamma_shape=2.0, t_peak=0.4) + total = s.clean_mass + s.uniform_mass + s.absorbed_mass + assert torch.allclose(total, torch.ones_like(total), atol=1e-6) + assert (s.clean_mass >= 0).all() and (s.uniform_mass >= 0).all() and (s.absorbed_mass >= 0).all() + + def test_uniform_mass_peaks_at_t_peak_with_the_configured_ratio(self): + # The normalisation is chosen so uniform_mass(t_peak) == max_ratio exactly. + for t_peak, ratio in [(0.5, 0.1), (0.3, 0.25)]: + s = scdd_schedule(torch.tensor([t_peak]), max_ratio=ratio, gamma_shape=2.0, t_peak=t_peak) + assert torch.allclose(s.uniform_mass, torch.tensor([ratio]), atol=1e-6) + + def test_rho_and_gamma_are_monotonically_decreasing(self): + # Monotonicity is what makes [MASK] absorbing (no remasking at sampling). + s = scdd_schedule(torch.linspace(0, 1, 201), max_ratio=0.2, gamma_shape=1.0, t_peak=0.5) + assert (s.gamma.diff() <= 1e-6).all() + assert (s.rho.diff() <= 1e-6).all() + + def test_zero_ratio_degenerates_to_absorbing_masking(self): + t = torch.linspace(0, 0.9, 10) + s = scdd_schedule(t, max_ratio=0.0, gamma_shape=1.0, t_peak=0.5) + assert torch.allclose(s.uniform_mass, torch.zeros_like(t)) + assert torch.allclose(s.absorbed_mass, t) + # rho is a conditional given "not [MASK]", so it is 1 everywhere the + # conditioning event has mass (i.e. everywhere except t = 1). + assert torch.allclose(s.rho, torch.ones_like(t)) + + @pytest.mark.parametrize( + "kwargs, match", + [ + (dict(max_ratio=1.0, gamma_shape=1.0, t_peak=0.5), "max_ratio"), + (dict(max_ratio=0.1, gamma_shape=1.0, t_peak=0.0), "t_peak"), + ], + ) + def test_rejects_out_of_range_hyperparameters(self, kwargs, match): + with pytest.raises(ValueError, match=match): + scdd_schedule(torch.tensor([0.5]), **kwargs) + + +SCDD_MASK_ID = V - 1 + + +def _scdd_batch(seed=0, t=0.5): + """Build (logits, x0, z_t, loss_mask, p_mask) for the SCDD loss. + + Returns tensors of shape ``[B, L, V]``, ``[B, L]``, ``[B, L]``, ``[B, L]`` + and ``[B, L]`` respectively; ``p_mask`` carries the diffusion time. + """ + torch.manual_seed(seed) + logits = torch.randn(B, L, V) + x0 = torch.randint(0, V - 1, (B, L)) # never the mask id + z_t = x0.clone() + loss_mask = torch.ones(B, L, dtype=torch.long) + p_mask = torch.full((B, L), t) + return logits, x0, z_t, loss_mask, p_mask + + +class TestSCDDLoss: + @pytest.mark.parametrize("num_timesteps", [0, 1]) + def test_rejects_grids_without_a_usable_point(self, num_timesteps): + """T = 1 leaves only t = 1, where the forward process is fully absorbed.""" + with pytest.raises(ValueError, match="num_timesteps >= 2"): + SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=num_timesteps) + + def test_reduces_to_mdlm_when_the_uniform_channel_is_off(self): + """With max_ratio=0 the NELBO collapses to -log p(x0)/t at masked + positions and exactly zero elsewhere — the MDLM objective.""" + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=1, t=0.4) + z_t[:, ::2] = SCDD_MASK_ID # every other position absorbed + + loss = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.0)( + logits=logits, + target_ids=x0, + noise_mask=z_t == SCDD_MASK_ID, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + + ref_logits = logits.clone() + ref_logits[:, :, SCDD_MASK_ID] = float("-inf") + log_p = torch.log_softmax(ref_logits, dim=-1) + nll = -log_p.gather(-1, x0.unsqueeze(-1)).squeeze(-1) + expected = (nll * (z_t == SCDD_MASK_ID) / 0.4).sum() / loss_mask.sum() + torch.testing.assert_close(loss.total_loss, expected, rtol=1e-4, atol=1e-5) + + def test_visible_wrong_token_drives_the_correction_term(self): + """At a visible position that disagrees with x0, putting mass on x0 + must cost less than doubling down on the corrupted token.""" + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=2, t=0.5) + z_t = (x0 + 1) % (V - 1) # every visible token is wrong + + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.2) + confident_x0 = torch.full_like(logits, -10.0).scatter(-1, x0.unsqueeze(-1), 10.0) + confident_zt = torch.full_like(logits, -10.0).scatter(-1, z_t.unsqueeze(-1), 10.0) + + common = dict( + target_ids=x0, + noise_mask=torch.ones_like(x0, dtype=torch.bool), + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + good = loss_fn(logits=confident_x0, **common).total_loss + bad = loss_fn(logits=confident_zt, **common).total_loss + assert good < bad + + def test_correction_term_decreases_with_confidence_in_x0(self): + """At an uncorrupted visible position the correction term is still live + (it is a cross-entropy against the true posterior, so it is negative at + the optimum rather than zero) and must fall monotonically as the model + concentrates on x0.""" + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=3, t=0.5) + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.2) + common = dict( + target_ids=x0, + noise_mask=torch.zeros_like(x0, dtype=torch.bool), + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + losses = [ + loss_fn( + logits=torch.zeros_like(logits).scatter(-1, x0.unsqueeze(-1), peak), + **common, + ).total_loss.item() + for peak in (0.0, 2.0, 10.0) + ] + assert losses[0] > losses[1] > losses[2] + + @pytest.mark.parametrize("t", [0.001, 0.5, 0.9999]) + def test_finite_forward_and_backward_across_the_schedule(self, t): + """t=1/T puts rho_s at 1 (zero uniform base) and t~1 empties the + retained mass; both boundaries must stay clear of inf/NaN.""" + logits, x0, z_t, loss_mask, _ = _scdd_batch(seed=4, t=t) + logits = logits.requires_grad_(True) + z_t[:, ::3] = SCDD_MASK_ID + z_t[:, 1::3] = (x0[:, 1::3] + 5) % (V - 1) + + loss = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.1)( + logits=logits, + target_ids=x0, + noise_mask=z_t != x0, + p_mask=torch.full((B, L), t), + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + assert torch.isfinite(loss.total_loss) + loss.total_loss.backward() + assert torch.isfinite(logits.grad).all() + + def test_unsupervised_positions_are_excluded(self): + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=5, t=0.6) + z_t[:, ::2] = SCDD_MASK_ID + loss_mask[:, L // 2 :] = 0 + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.1) + common = dict( + target_ids=x0, + noise_mask=z_t == SCDD_MASK_ID, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + num_diffusion_tokens=int(loss_mask.sum()), + ) + base = loss_fn(logits=logits, **common).total_loss + perturbed = logits.clone() + perturbed[:, L // 2 :] += 5.0 + assert torch.equal(loss_fn(logits=perturbed, **common).total_loss, base) + + def test_global_denominator_is_honoured(self): + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=6, t=0.5) + z_t[:, ::2] = SCDD_MASK_ID + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.1) + common = dict( + target_ids=x0, + noise_mask=z_t == SCDD_MASK_ID, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + local = loss_fn(logits=logits, **common).total_loss + halved = loss_fn(logits=logits, num_diffusion_tokens=2 * int(loss_mask.sum()), **common).total_loss + torch.testing.assert_close(halved, local / 2, rtol=1e-5, atol=1e-7) + + def test_never_predicts_the_mask_token(self): + """The SCDD parameterisation removes [MASK] from the denoiser domain, so + the mask logit must not influence the loss at all.""" + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=7, t=0.5) + z_t[:, ::2] = SCDD_MASK_ID + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.1) + common = dict( + target_ids=x0, + noise_mask=z_t == SCDD_MASK_ID, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + base = loss_fn(logits=logits, **common).total_loss + shifted = logits.clone() + shifted[:, :, SCDD_MASK_ID] += 50.0 + torch.testing.assert_close(loss_fn(logits=shifted, **common).total_loss, base, rtol=1e-5, atol=1e-6) + + def test_caller_logits_are_not_mutated(self): + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=8, t=0.5) + snapshot = logits.clone() + SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.1)( + logits=logits, + target_ids=x0, + noise_mask=torch.zeros_like(x0, dtype=torch.bool), + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + assert torch.equal(logits, snapshot) + + def test_requires_noisy_input_ids(self): + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=9, t=0.5) + with pytest.raises(ValueError, match="noisy_input_ids"): + SCDDLoss(mask_token_id=SCDD_MASK_ID)( + logits=logits, + target_ids=x0, + noise_mask=torch.zeros_like(x0, dtype=torch.bool), + p_mask=p_mask, + loss_mask=loss_mask, + ) + + def test_rejects_zero_timesteps(self): + with pytest.raises(ValueError, match="num_timesteps"): + SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=0) + + def test_returns_dllm_loss_output(self): + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=10, t=0.5) + z_t[:, ::2] = SCDD_MASK_ID + out = SCDDLoss(mask_token_id=SCDD_MASK_ID, max_ratio=0.1)( + logits=logits, + target_ids=x0, + noise_mask=z_t == SCDD_MASK_ID, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + assert isinstance(out, DLLMLossOutput) + torch.testing.assert_close(out.dllm_loss, out.total_loss.detach()) + assert out.draft_correct_per_pos is None and out.draft_count_per_pos is None + + +class TestSCDDLossDistributed: + """Properties SCDD needs to hold under FSDP2/TP + gradient accumulation.""" + + def test_grad_accumulation_matches_the_full_batch(self): + """DP + grad-accum splits the global batch and sums per-microbatch + losses. With the global supervised-token denominator, that sum must + equal the single-batch loss, otherwise the gradient is mis-scaled by + the number of microbatches.""" + torch.manual_seed(11) + big = 6 + logits = torch.randn(big, L, V) + x0 = torch.randint(0, V - 1, (big, L)) + z_t = x0.clone() + z_t[:, ::2] = SCDD_MASK_ID + z_t[:, 1::4] = (x0[:, 1::4] + 3) % (V - 1) + loss_mask = torch.ones(big, L, dtype=torch.long) + p_mask = torch.rand(big, 1).clamp(0.05, 0.95).expand(big, L).contiguous() + global_denom = int(loss_mask.sum()) + + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.15) + + def _call(sl): + return loss_fn( + logits=logits[sl], + target_ids=x0[sl], + noise_mask=(z_t != x0)[sl], + p_mask=p_mask[sl], + loss_mask=loss_mask[sl], + noisy_input_ids=z_t[sl], + num_diffusion_tokens=global_denom, + ).total_loss + + whole = _call(slice(None)) + accumulated = sum(_call(slice(i, i + 2)) for i in range(0, big, 2)) + torch.testing.assert_close(accumulated, whole, rtol=1e-5, atol=1e-6) + + def test_per_sequence_times_are_read_independently(self): + """Each sequence carries its own diffusion time; batching sequences at + different times must not couple them (a DP rank sees a mixed batch).""" + torch.manual_seed(12) + logits = torch.randn(2, L, V) + x0 = torch.randint(0, V - 1, (2, L)) + z_t = x0.clone() + z_t[:, ::2] = SCDD_MASK_ID + loss_mask = torch.ones(2, L, dtype=torch.long) + times = torch.tensor([0.2, 0.8]) + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.15) + + def _call(sl, t): + return loss_fn( + logits=logits[sl], + target_ids=x0[sl], + noise_mask=(z_t == SCDD_MASK_ID)[sl], + p_mask=t[:, None].expand(-1, L), + loss_mask=loss_mask[sl], + noisy_input_ids=z_t[sl], + num_diffusion_tokens=int(loss_mask.sum()), + ).total_loss + + together = _call(slice(None), times) + apart = _call(slice(0, 1), times[:1]) + _call(slice(1, 2), times[1:]) + torch.testing.assert_close(together, apart, rtol=1e-5, atol=1e-6) + + def test_vocab_sharded_dtensor_matches_plain(self, trivial_pg): + """TP shards logits over the vocab axis; the SCDD parameterisation and + the sum over the non-[MASK] domain must survive the materialisation.""" + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import Shard, distribute_tensor + + logits, x0, z_t, loss_mask, p_mask = _scdd_batch(seed=13, t=0.5) + z_t[:, ::2] = SCDD_MASK_ID + loss_fn = SCDDLoss(mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.15) + common = dict( + target_ids=x0, + noise_mask=z_t == SCDD_MASK_ID, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + ) + + plain = loss_fn(logits=logits, **common).total_loss + mesh = init_device_mesh("cpu", (1,)) + sharded = loss_fn(logits=distribute_tensor(logits, mesh, [Shard(-1)]), **common).total_loss + torch.testing.assert_close(sharded, plain, rtol=1e-5, atol=1e-6) + + +class TestSCDDLossChunking: + """The chunked vocabulary reduction must be numerically transparent. + + The ELBO needs the model's probability of every non-[MASK] token, so it + cannot use a fused cross-entropy kernel; the vocab-sized work is instead + checkpointed in position chunks. Chunking changes peak memory, never the + result — each position is independent. + """ + + @staticmethod + def _inputs(seed=20): + """Build a batch big enough that a small chunk size spans the batch axis. + + Returns tensors of shape ``[4, 16, V]``, ``[4, 16]`` x 3 and ``[4, 16]``. + """ + torch.manual_seed(seed) + big_b, big_l = 4, 16 + logits = torch.randn(big_b, big_l, V) + x0 = torch.randint(0, V - 1, (big_b, big_l)) + z_t = x0.clone() + z_t[:, ::3] = SCDD_MASK_ID + z_t[:, 1::3] = (x0[:, 1::3] + 7) % (V - 1) + loss_mask = torch.ones(big_b, big_l, dtype=torch.long) + p_mask = torch.rand(big_b, 1).clamp(0.05, 0.95).expand(big_b, big_l).contiguous() + return logits, x0, z_t, loss_mask, p_mask + + def _run(self, chunk_size, logits, x0, z_t, loss_mask, p_mask): + logits = logits.clone().requires_grad_(True) + loss_fn = SCDDLoss( + mask_token_id=SCDD_MASK_ID, num_timesteps=1000, max_ratio=0.15, chunk_size=chunk_size + ) + out = loss_fn( + logits=logits, + target_ids=x0, + noise_mask=z_t != x0, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=z_t, + num_diffusion_tokens=int(loss_mask.sum()), + ) + out.total_loss.backward() + return out.total_loss.detach(), logits.grad + + @pytest.mark.parametrize("chunk_size", [1, 7, 32, 1024]) + def test_matches_the_unchunked_result(self, chunk_size): + """Chunk boundaries that split a sequence, align with it, and exceed the + whole batch must all give the same loss and the same gradient.""" + args = self._inputs() + ref_loss, ref_grad = self._run(None, *args) + loss, grad = self._run(chunk_size, *args) + torch.testing.assert_close(loss, ref_loss, rtol=1e-6, atol=1e-7) + torch.testing.assert_close(grad, ref_grad, rtol=1e-5, atol=1e-7) + + def test_gradient_reaches_every_supervised_position(self): + """Checkpoint recompute must not drop gradient for chunks after the + first — a stale-boundary bug would leave later positions at zero.""" + logits, x0, z_t, loss_mask, p_mask = self._inputs(seed=21) + _, grad = self._run(7, logits, x0, z_t, loss_mask, p_mask) + per_position = grad.abs().sum(dim=-1) + assert (per_position > 0).all(), f"zero-gradient positions: {(per_position == 0).nonzero().tolist()}" + + def test_rejects_nonpositive_chunk_size(self): + with pytest.raises(ValueError, match="chunk_len"): + SCDDLoss(mask_token_id=SCDD_MASK_ID, chunk_size=0) + + def test_defaults_to_chunking(self): + assert SCDDLoss(mask_token_id=SCDD_MASK_ID).chunk_size == 1024 + assert SCDDLoss(mask_token_id=SCDD_MASK_ID, chunk_size=None).chunk_size is None diff --git a/tests/unit_tests/recipes/dllm/test_generation.py b/tests/unit_tests/recipes/dllm/test_generation.py index f916363988..609061313d 100644 --- a/tests/unit_tests/recipes/dllm/test_generation.py +++ b/tests/unit_tests/recipes/dllm/test_generation.py @@ -31,6 +31,7 @@ IDLMSampler, LLaDA2Sampler, LLaDASampler, + SCDDSampler, encode_generation_prompts, generate_gemma, generate_llada2, @@ -458,3 +459,173 @@ def test_translate_adapter_reparents_gemma_module_paths(tmp_path): cfg = json.loads((out / "adapter_config.json").read_text()) assert cfg["target_modules"] == ["model.decoder.layers.3.self_attn.q_proj"] assert cfg["r"] == 2 # non-key fields untouched + + +class _FakeSCDDDenoiser(torch.nn.Module): + """Denoiser that puts (almost) all mass on one token, position-independent. + + ``switch_after`` flips the prediction from ``target`` to ``revised`` once + that many forward passes have run, which simulates a model changing its mind + about tokens it already committed. + """ + + def __init__(self, target: int, vocab_size: int = 16, revised: int | None = None, switch_after: int = 0): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(1)) + self.target = target + self.revised = revised + self.switch_after = switch_after + self.vocab_size = vocab_size + self.calls = 0 + + def forward(self, x, attention_mask=None): + """Run the fake denoiser. + + Args: + x: Token IDs of shape ``[batch, sequence]``. + attention_mask: Binary mask of shape ``[batch, sequence]``; unused. + + Returns: + Namespace whose ``logits`` have shape ``[batch, sequence, vocab]``. + """ + B, L = x.shape + token = self.target + if self.revised is not None and self.calls >= self.switch_after: + token = self.revised + self.calls += 1 + logits = torch.zeros(B, L, self.vocab_size) + logits[:, :, token] = 30.0 + return types.SimpleNamespace(logits=logits) + + +def test_scdd_sampler_is_registered(): + assert SAMPLERS["scdd"] is SCDDSampler + cfg = SCDDSampler.default_config + # Ancestral sampling must be stochastic, and every position is rewritten each + # step so there is nothing to cache. + assert cfg.temperature == 1.0 + assert cfg.use_kv_cache is False + + +def test_scdd_rejects_kv_cache(): + sampler = SCDDSampler(_FakeSCDDDenoiser(target=7), mask_id=9, eos_id=0) + with pytest.raises(ValueError, match="use_kv_cache"): + sampler.sample([[1, 2]], use_kv_cache=True) + + +def test_scdd_leaves_no_mask_tokens_and_preserves_the_prompt(): + mask_id = 9 + sampler = SCDDSampler( + _FakeSCDDDenoiser(target=7), + mask_id=mask_id, + eos_id=0, + steps=8, + max_new_tokens=6, + temperature=1.0, + uniform_ratio=0.1, + ) + + out = sampler.sample([[1, 2]]) + + assert out.shape == (1, 8) + assert torch.equal(out[0, :2], torch.tensor([1, 2])), "the prompt must never be resampled" + assert (out == mask_id).sum().item() == 0, "residual [MASK] survived the final denoise" + assert (out[0, 2:] == 7).all(), "a near-deterministic denoiser should drive every canvas position" + + +def test_scdd_corrects_tokens_it_already_committed(): + """The defining SCDD behaviour: a non-[MASK] token the denoiser no longer + agrees with is overwritten. An absorbing sampler, which only ever fills + [MASK], would leave the stale token in place.""" + mask_id = 9 + model = _FakeSCDDDenoiser(target=3, revised=7, switch_after=2) + sampler = SCDDSampler(model, mask_id=mask_id, eos_id=0, steps=16, max_new_tokens=4, uniform_ratio=0.2) + + torch.manual_seed(0) + out = sampler.sample([[1]]) + + assert model.calls > 2, "the switch must happen mid-decode for this to test anything" + assert (out[0, 1:] == 7).all(), "positions committed as 3 were never revised to 7" + + +def test_scdd_infill_is_rejected_before_loading(monkeypatch, capsys): + monkeypatch.setattr( + sys, + "argv", + ["generate.py", "--checkpoint", "unused", "--prompt", "hello", "--sampler", "scdd", "--infill"], + ) + with pytest.raises(SystemExit, match="2"): + main() + assert "--infill is not supported by the SCDD sampler" in capsys.readouterr().err + + +def test_scdd_schedule_flags_reach_the_sampler_config(monkeypatch): + captured = {} + + def _fake_load(checkpoint_path, sampler_name="llada", mask_id_override=None): + return _FakeSCDDDenoiser(target=7), _FakeChatTokenizer(), 9, 0 + + import generate as generate_mod + + monkeypatch.setattr(generate_mod, "load_model_and_tokenizer", _fake_load) + monkeypatch.setattr(generate_mod, "resolve_checkpoint", lambda p: p) + monkeypatch.setattr( + generate_mod, + "encode_generation_prompts", + lambda tokenizer, prompts, raw: (_ for _ in ()).throw(_Captured(captured)), + ) + + class _Captured(Exception): + def __init__(self, store): + super().__init__("stop") + + monkeypatch.setattr( + sys, + "argv", + [ + "generate.py", + "--checkpoint", + "unused", + "--prompt", + "hello", + "--sampler", + "scdd", + "--uniform_ratio", + "0.3", + "--schedule_shape", + "2.0", + "--schedule_peak", + "0.25", + ], + ) + with pytest.raises(_Captured): + main() + + +def test_scdd_zero_temperature_collapses_the_denoiser(): + """``temperature=0`` must give a one-hot distribution over the argmax and + never put mass on [MASK] — the path the final residual-mask denoise uses.""" + mask_id = 9 + sampler = SCDDSampler(_FakeSCDDDenoiser(target=7), mask_id=mask_id, eos_id=0) + x = torch.tensor([[1, mask_id, mask_id]]) + attention_mask = torch.ones_like(x) + + log_p = sampler._denoiser_log_probs(x, attention_mask, 0.0) + probs = log_p.exp() + + assert torch.allclose(probs.sum(-1), torch.ones(1, 3), atol=1e-5) + assert (probs.argmax(-1) == 7).all() + assert torch.allclose(probs[..., mask_id], torch.zeros(1, 3)), "[MASK] must be outside the domain" + + +def test_scdd_degenerate_schedule_still_terminates(): + """uniform_ratio=0 removes the correction channel, which makes the + visible-token posterior unreachable; the sampler must fall back to keeping + the current token rather than handing multinomial an all-zero row.""" + mask_id = 9 + sampler = SCDDSampler( + _FakeSCDDDenoiser(target=7), mask_id=mask_id, eos_id=0, steps=4, max_new_tokens=4, uniform_ratio=0.0 + ) + out = sampler.sample([[1]]) + assert out.shape == (1, 5) + assert (out == mask_id).sum().item() == 0 diff --git a/tests/unit_tests/recipes/dllm/test_scdd_recipe_smoke.py b/tests/unit_tests/recipes/dllm/test_scdd_recipe_smoke.py new file mode 100644 index 0000000000..846a3fd9c4 --- /dev/null +++ b/tests/unit_tests/recipes/dllm/test_scdd_recipe_smoke.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU smoke test for the SCDD dLLM training path. + +The full ``DiffusionLMSFTRecipe`` loop is CUDA-only, so this exercises the same +sequence the recipe drives per micro-batch — strategy corruption, batch +preparation, a real transformer forward, the SCDD loss, and backward — against a +hermetically constructed tiny model. It answers "would this actually train?" +without a GPU or a checkpoint download; the GPU counterpart is +``tests/functional_tests/dllm/L2_DLLM_SCDD_Smoke.sh``. +""" + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from nemo_automodel.recipes.dllm.strategy import get_dllm_strategy + +VOCAB = 128 +MASK_TOKEN_ID = VOCAB - 1 +SEQ_LEN = 24 +BATCH = 4 + +DLLM_CFG = { + "mode": "scdd", + "mask_token_id": MASK_TOKEN_ID, + "vocab_size": VOCAB, + "eps": 1e-3, + "num_timesteps": 1000, + "uniform_ratio": 0.1, + "schedule_shape": 1.0, + "schedule_peak": 0.5, +} + + +def _tiny_model() -> LlamaForCausalLM: + """Build a 2-layer randomly-initialised causal LM with no network access.""" + config = LlamaConfig( + vocab_size=VOCAB, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=SEQ_LEN, + attn_implementation="eager", + ) + torch.manual_seed(0) + return LlamaForCausalLM(config) + + +def _batch() -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(input_ids, loss_mask)``, both of shape ``[batch, sequence]``. + + The first third of each row stands in for a prompt (unsupervised); the rest + is the supervised response. + """ + torch.manual_seed(1) + input_ids = torch.randint(0, VOCAB - 1, (BATCH, SEQ_LEN)) + loss_mask = torch.zeros(BATCH, SEQ_LEN, dtype=torch.long) + loss_mask[:, SEQ_LEN // 3 :] = 1 + return input_ids, loss_mask + + +def _step(strategy, loss_fn, model, input_ids, loss_mask, seed): + """Run one micro-batch exactly as ``_forward_backward_step`` does. + + Args: + strategy: The ``SCDDStrategy`` under test. + loss_fn: The ``SCDDLoss`` built by the strategy. + model: The tiny causal LM. + input_ids: Clean token IDs of shape ``[batch, sequence]``. + loss_mask: Supervised-position mask of shape ``[batch, sequence]``. + seed: Seed for the corruption generator; reuse it to hold the noise fixed. + + Returns: + The scalar loss tensor for this micro-batch. + """ + noisy_input_ids, noise_mask, p_mask = strategy.apply_corruption( + input_ids, + loss_mask, + MASK_TOKEN_ID, + eps=DLLM_CFG["eps"], + block_size=None, + half_life_ratio=None, + generator=torch.Generator().manual_seed(seed), + ) + batch = strategy.prepare_batch( + {"input_ids": input_ids.clone(), "attention_mask": torch.ones_like(input_ids)}, + noisy_input_ids, + noise_mask, + input_ids, + ) + logits = model(**batch).logits + return loss_fn( + logits=logits, + target_ids=input_ids, + noise_mask=noise_mask, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=noisy_input_ids, + num_diffusion_tokens=int(loss_mask.sum()), + ).total_loss + + +def test_scdd_micro_batch_runs_end_to_end(): + """Corruption -> forward -> loss -> backward must produce finite gradients on + the parameters the optimizer would update.""" + strategy = get_dllm_strategy("scdd") + loss_fn = strategy.create_loss_fn(DLLM_CFG) + model = _tiny_model() + input_ids, loss_mask = _batch() + + loss = _step(strategy, loss_fn, model, input_ids, loss_mask, seed=0) + assert torch.isfinite(loss) + loss.backward() + + grads = [p.grad for p in model.parameters() if p.grad is not None] + assert grads, "no parameter received a gradient" + assert all(torch.isfinite(g).all() for g in grads) + assert sum(g.abs().sum() for g in grads) > 0 + + +def test_scdd_loss_decreases_under_optimization(): + """With the corruption held fixed, a few Adam steps must drive the SCDD + objective down — the signal that its gradients point the right way rather + than merely being finite.""" + strategy = get_dllm_strategy("scdd") + loss_fn = strategy.create_loss_fn(DLLM_CFG) + model = _tiny_model() + input_ids, loss_mask = _batch() + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2) + + losses = [] + for _ in range(12): + optimizer.zero_grad() + loss = _step(strategy, loss_fn, model, input_ids, loss_mask, seed=0) + loss.backward() + optimizer.step() + losses.append(loss.item()) + + assert all(torch.isfinite(torch.tensor(losses))) + assert losses[-1] < losses[0], f"SCDD loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" + + +def test_scdd_batch_feeds_corrupted_tokens_to_the_model(): + """The model must see the corrupted sequence (and no attention mask, since + it attends bidirectionally), while the loss scores the clean targets.""" + strategy = get_dllm_strategy("scdd") + strategy.create_loss_fn(DLLM_CFG) + input_ids, loss_mask = _batch() + + noisy_input_ids, noise_mask, _ = strategy.apply_corruption( + input_ids, + loss_mask, + MASK_TOKEN_ID, + eps=DLLM_CFG["eps"], + block_size=None, + half_life_ratio=None, + generator=torch.Generator().manual_seed(3), + ) + batch = strategy.prepare_batch( + {"input_ids": input_ids.clone(), "attention_mask": torch.ones_like(input_ids)}, + noisy_input_ids, + noise_mask, + input_ids, + ) + assert torch.equal(batch["input_ids"], noisy_input_ids) + assert "attention_mask" not in batch + assert noise_mask.any(), "corruption produced nothing to learn from" + + +@pytest.mark.parametrize("uniform_ratio", [0.0, 0.05, 0.4]) +def test_scdd_runs_across_the_uniform_ratio_range(uniform_ratio): + """Both ends of the schedule are reachable from config: 0 degenerates to + MDLM and a large ratio makes most corrupted tokens visible-but-wrong.""" + strategy = get_dllm_strategy("scdd") + loss_fn = strategy.create_loss_fn({**DLLM_CFG, "uniform_ratio": uniform_ratio}) + model = _tiny_model() + input_ids, loss_mask = _batch() + + loss = _step(strategy, loss_fn, model, input_ids, loss_mask, seed=5) + assert torch.isfinite(loss) + loss.backward() + assert all(torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None) diff --git a/tests/unit_tests/recipes/dllm/test_strategy.py b/tests/unit_tests/recipes/dllm/test_strategy.py index 997654d59b..43cbd5846a 100644 --- a/tests/unit_tests/recipes/dllm/test_strategy.py +++ b/tests/unit_tests/recipes/dllm/test_strategy.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for dLLM strategies (MDLMStrategy, HybridStrategy, DFlashStrategy) and get_dllm_strategy.""" +"""Tests for dLLM strategies (MDLMStrategy, SCDDStrategy, HybridStrategy, DFlashStrategy) and get_dllm_strategy.""" import types +from pathlib import Path import pytest import torch @@ -24,6 +25,7 @@ DFlashDecayLoss, IDLMLoss, MDLMCrossEntropyLoss, + SCDDLoss, ) from nemo_automodel.recipes.dllm.strategy import ( DLLM_STRATEGIES, @@ -32,11 +34,15 @@ HybridStrategy, IDLMStrategy, MDLMStrategy, + SCDDStrategy, _build_target_layer_ids, get_dllm_strategy, ) +REPO_ROOT = Path(__file__).resolve().parents[4] + + def test_get_dllm_strategy_rejects_unknown_mode(): """Unknown mode must raise a clear ValueError (recipe entry point relies on this).""" with pytest.raises(ValueError, match="Unknown dllm.mode"): @@ -1002,3 +1008,299 @@ def test_prepare_batch_sets_canvas_and_clean_inputs(self, strategy): # Bidirectional model -> attention_mask / use_cache dropped. assert "attention_mask" not in result assert "use_cache" not in result + + +# --------------------------------------------------------------------------- +# SCDDStrategy +# --------------------------------------------------------------------------- + + +SCDD_VOCAB = 64 +SCDD_MASK_ID = 63 + + +def _scdd_strategy(**cfg): + strategy = SCDDStrategy() + strategy.create_loss_fn({"vocab_size": SCDD_VOCAB, "mask_token_id": SCDD_MASK_ID, **cfg}) + return strategy + + +def test_get_dllm_strategy_resolves_scdd(): + assert isinstance(get_dllm_strategy("scdd"), SCDDStrategy) + assert DLLM_STRATEGIES["scdd"] is SCDDStrategy + + +def test_scdd_create_loss_fn_builds_scdd_loss_from_config(): + strategy = SCDDStrategy() + loss_fn = strategy.create_loss_fn( + { + "vocab_size": SCDD_VOCAB, + "mask_token_id": SCDD_MASK_ID, + "num_timesteps": 256, + "uniform_ratio": 0.25, + "schedule_shape": 2.0, + "schedule_peak": 0.4, + } + ) + assert isinstance(loss_fn, SCDDLoss) + assert (loss_fn.num_timesteps, loss_fn.max_ratio) == (256, 0.25) + assert (loss_fn.gamma_shape, loss_fn.t_peak) == (2.0, 0.4) + + +def test_scdd_apply_corruption_keeps_t_on_the_discrete_grid(): + """t must land on {1/T, ..., (T-1)/T}. + + SCDDLoss reads t back out of p_mask and forms s = t - 1/T, so an off-grid t + (the old 1 - 1e-4 clamp) puts s off-grid too. The top point t = 1 stays + excluded because the schedule is fully absorbed there. + """ + num_timesteps = 8 + strategy = _scdd_strategy(num_timesteps=num_timesteps) + input_ids = torch.randint(0, SCDD_VOCAB - 1, (256, 4)) + loss_mask = torch.ones(256, 4, dtype=torch.long) + + _, _, p_mask = strategy.apply_corruption( + input_ids, + loss_mask, + SCDD_MASK_ID, + eps=1e-3, + block_size=None, + half_life_ratio=None, + generator=torch.Generator().manual_seed(0), + ) + + t = p_mask[:, 0] + steps = t * num_timesteps + assert torch.allclose(steps, steps.round()) + assert int(steps.min()) >= 1 + assert int(steps.max()) == num_timesteps - 1 + + +def test_scdd_create_loss_fn_requires_vocab_size(): + with pytest.raises(ValueError, match="dllm.vocab_size"): + SCDDStrategy().create_loss_fn({"mask_token_id": SCDD_MASK_ID}) + + +def test_scdd_apply_corruption_before_create_loss_fn_raises(): + with pytest.raises(ValueError, match="create_loss_fn"): + SCDDStrategy().apply_corruption( + torch.zeros(1, 4, dtype=torch.long), + torch.ones(1, 4, dtype=torch.long), + SCDD_MASK_ID, + eps=1e-3, + block_size=None, + half_life_ratio=None, + ) + + +def _scdd_recipe(mask_token_id=SCDD_MASK_ID, vocab_size=SCDD_VOCAB, cp_size=1, loss_fn=None): + """Minimal recipe stand-in exposing what SCDDStrategy.setup_extra reads.""" + model = types.SimpleNamespace(config=types.SimpleNamespace(vocab_size=vocab_size)) + return types.SimpleNamespace( + mask_token_id=mask_token_id, + dllm_loss_fn=loss_fn if loss_fn is not None else SCDDLoss(mask_token_id=0), + model_parts=[model], + distributed_config=types.SimpleNamespace(cp_size=cp_size), + ) + + +def test_scdd_setup_extra_installs_the_resolved_mask_token_id(): + """mask_token_id may only be known after the tokenizer is built, so the + loss module must pick up the recipe's resolved value.""" + loss_fn = SCDDLoss(mask_token_id=0) + _scdd_strategy().setup_extra(_scdd_recipe(loss_fn=loss_fn)) + assert loss_fn.mask_token_id == SCDD_MASK_ID + + +def test_scdd_setup_extra_requires_a_mask_token_id(): + with pytest.raises(ValueError, match="mask_token_id"): + _scdd_strategy().setup_extra(_scdd_recipe(mask_token_id=None)) + + +def test_scdd_setup_extra_rejects_a_mask_id_outside_the_model_vocab(): + with pytest.raises(ValueError, match="outside the model vocab"): + _scdd_strategy().setup_extra(_scdd_recipe(mask_token_id=SCDD_VOCAB + 10)) + + +def test_scdd_setup_extra_rejects_a_vocab_size_mismatch(): + """The corruption domain and the ELBO's non-[MASK] domain must both be the + model's output domain; a stale dllm.vocab_size silently changes the loss.""" + with pytest.raises(ValueError, match="does not match the model vocab"): + _scdd_strategy().setup_extra(_scdd_recipe(vocab_size=SCDD_VOCAB + 128)) + + +def test_scdd_setup_extra_rejects_context_parallelism(): + """The loss scores unsharded clean targets against the model logits, so a + sequence-sharded forward would silently mis-align them.""" + with pytest.raises(ValueError, match="context parallelism"): + _scdd_strategy().setup_extra(_scdd_recipe(cp_size=2)) + + +def test_scdd_apply_corruption_contract(): + strategy = _scdd_strategy(num_timesteps=100, uniform_ratio=0.2) + input_ids = torch.randint(0, SCDD_VOCAB - 1, (3, 64)) + loss_mask = torch.zeros(3, 64, dtype=torch.long) + loss_mask[:, 16:] = 1 + + noisy, noise_mask, p_mask = strategy.apply_corruption( + input_ids, + loss_mask, + SCDD_MASK_ID, + eps=1e-3, + block_size=None, + half_life_ratio=None, + generator=torch.Generator().manual_seed(0), + ) + + assert noisy.shape == input_ids.shape and noise_mask.shape == input_ids.shape + assert p_mask.shape == input_ids.shape and p_mask.dtype == torch.float32 + # Only supervised positions may be corrupted. + assert torch.equal(noisy[loss_mask == 0], input_ids[loss_mask == 0]) + assert not noise_mask[loss_mask == 0].any() + # p_mask carries one diffusion time per sequence, snapped to the 1/T grid. + t = p_mask[:, 0] + assert torch.equal(p_mask, t[:, None].expand_as(p_mask)) + assert ((t > 0) & (t <= 1.0)).all() + torch.testing.assert_close(t * 100, (t * 100).round(), rtol=0, atol=1e-3) + + +def test_scdd_apply_corruption_produces_both_channels(): + """The point of SCDD is that some corrupted positions are wrong-but-visible + tokens, not just [MASK].""" + strategy = _scdd_strategy(num_timesteps=1000, uniform_ratio=0.4) + input_ids = torch.randint(0, SCDD_VOCAB - 1, (16, 256)) + loss_mask = torch.ones_like(input_ids) + noisy, noise_mask, _ = strategy.apply_corruption( + input_ids, + loss_mask, + SCDD_MASK_ID, + eps=1e-3, + block_size=None, + half_life_ratio=None, + generator=torch.Generator().manual_seed(1), + ) + absorbed = noisy == SCDD_MASK_ID + transitioned = noise_mask & ~absorbed + assert absorbed.any() and transitioned.any() + assert (noisy[transitioned] != input_ids[transitioned]).all() + + +def test_scdd_apply_corruption_is_seed_reproducible(): + strategy = _scdd_strategy() + args = (torch.randint(0, SCDD_VOCAB - 1, (2, 32)), torch.ones(2, 32, dtype=torch.long), SCDD_MASK_ID) + kwargs = dict(eps=1e-3, block_size=None, half_life_ratio=None) + a = strategy.apply_corruption(*args, generator=torch.Generator().manual_seed(5), **kwargs) + b = strategy.apply_corruption(*args, generator=torch.Generator().manual_seed(5), **kwargs) + for x, y in zip(a, b): + assert torch.equal(x, y) + + +def test_scdd_prepare_batch_feeds_the_corrupted_tokens(): + strategy = _scdd_strategy() + clean = torch.randint(0, SCDD_VOCAB - 1, (2, 8)) + noisy = clean.clone() + noisy[:, 0] = SCDD_MASK_ID + batch = {"input_ids": clean, "attention_mask": torch.ones(2, 8, dtype=torch.long)} + out = strategy.prepare_batch(batch, noisy, noisy != clean, clean) + assert torch.equal(out["input_ids"], noisy) + assert "attention_mask" not in out + + +def test_scdd_normalizes_over_all_supervised_tokens(): + """The SCDD ELBO has a term at every supervised position, corrupted or not, + so the denominator must not be the corrupted-only count.""" + assert SCDDStrategy().normalization_mode == "supervised" + + +def test_scdd_corruption_does_not_touch_the_global_rng(): + """TP/CP peers share a data shard but not their global RNG state, so the + corruption must come exclusively from the step-seeded generator — otherwise + peers feed the sharded forward different inputs.""" + strategy = _scdd_strategy() + input_ids = torch.randint(0, SCDD_VOCAB - 1, (2, 32)) + loss_mask = torch.ones(2, 32, dtype=torch.long) + kwargs = dict(eps=1e-3, block_size=None, half_life_ratio=None) + + torch.manual_seed(0) + expected = torch.rand(4) + + torch.manual_seed(0) + peer_a = strategy.apply_corruption( + input_ids, loss_mask, SCDD_MASK_ID, generator=torch.Generator().manual_seed(1234), **kwargs + ) + assert torch.equal(torch.rand(4), expected), "global RNG was consumed" + + torch.manual_seed(999) # a peer with a diverged global RNG state + peer_b = strategy.apply_corruption( + input_ids, loss_mask, SCDD_MASK_ID, generator=torch.Generator().manual_seed(1234), **kwargs + ) + for a, b in zip(peer_a, peer_b): + assert torch.equal(a, b) + + +def test_scdd_corruption_and_loss_agree_on_the_p_mask_contract(): + """End-to-end guard on the strategy<->loss handshake: the strategy writes + the diffusion time into ``p_mask`` and the loss reads the schedule back out + of it. A drift in either direction silently trains the wrong objective.""" + strategy = _scdd_strategy(num_timesteps=1000, uniform_ratio=0.2) + loss_fn = strategy.create_loss_fn( + {"vocab_size": SCDD_VOCAB, "mask_token_id": SCDD_MASK_ID, "num_timesteps": 1000, "uniform_ratio": 0.2} + ) + input_ids = torch.randint(0, SCDD_VOCAB - 1, (4, 32)) + loss_mask = torch.ones(4, 32, dtype=torch.long) + + noisy, noise_mask, p_mask = strategy.apply_corruption( + input_ids, + loss_mask, + SCDD_MASK_ID, + eps=1e-3, + block_size=None, + half_life_ratio=None, + generator=torch.Generator().manual_seed(7), + ) + + logits = torch.randn(4, 32, SCDD_VOCAB, requires_grad=True) + out = loss_fn( + logits=logits, + target_ids=input_ids, + noise_mask=noise_mask, + p_mask=p_mask, + loss_mask=loss_mask, + noisy_input_ids=noisy, + num_diffusion_tokens=int(loss_mask.sum()), + ) + assert torch.isfinite(out.total_loss) + out.total_loss.backward() + assert torch.isfinite(logits.grad).all() and logits.grad.abs().sum() > 0 + + +def test_scdd_create_loss_fn_threads_chunk_size(): + """The loss's memory knob must be reachable from the recipe YAML, including + the explicit ``null`` that turns chunking off.""" + strategy = SCDDStrategy() + base = {"vocab_size": SCDD_VOCAB, "mask_token_id": SCDD_MASK_ID} + assert strategy.create_loss_fn(base).chunk_size == 1024 + assert strategy.create_loss_fn({**base, "chunk_size": 256}).chunk_size == 256 + assert strategy.create_loss_fn({**base, "chunk_size": None}).chunk_size is None + + +def test_shipped_scdd_recipe_builds_its_strategy_and_loss(): + """The go-to recipe must stay loadable: a typo in dllm.mode or a schedule key + would otherwise only surface on an 8-GPU run.""" + import yaml + + config_path = REPO_ROOT / "examples" / "dllm_sft" / "llada_scdd.yaml" + cfg = yaml.safe_load(config_path.read_text()) + dllm_cfg = cfg["dllm"] + + strategy = get_dllm_strategy(dllm_cfg["mode"]) + assert isinstance(strategy, SCDDStrategy) + loss_fn = strategy.create_loss_fn(dllm_cfg) + assert isinstance(loss_fn, SCDDLoss) + # Schedule values match the authors' released checkpoint config. + assert (loss_fn.num_timesteps, loss_fn.max_ratio) == (1000, 0.1) + assert (loss_fn.gamma_shape, loss_fn.t_peak) == (1.0, 0.5) + assert loss_fn.max_ratio > 0, "uniform_ratio 0 would silently degenerate SCDD to MDLM" + # The mask id must be inside the vocabulary the uniform channel samples over. + assert 0 <= dllm_cfg["mask_token_id"] < dllm_cfg["vocab_size"] + assert cfg["distributed"]["cp_size"] == 1, "SCDD rejects context parallelism"