dLLM Fine-Tuning

Introduction

Diffusion language models (dLLMs) generate text by iteratively denoising a corrupted sequence rather than generating one token at a time from left to right, as autoregressive (AR) models do. Masked-diffusion families such as LLaDA start from [MASK] tokens and progressively unmask the most confident positions. DiffusionGemma instead corrupts the response canvas with uniform-random vocabulary tokens (no [MASK]) and denoises that canvas in blocks.

This approach enables parallel token generation and bidirectional attention, giving the model more context for each prediction than AR models provide.

NeMo AutoModel currently supports the following dLLM model families:

• LLaDA or LLaDA2 (Masked Diffusion Language Model or MDLM): This bidirectional masked diffusion model receives corrupted tokens and predicts the clean token at each masked position. For details, refer to the LLaDA2 paper.
• Nemotron-Labs-Diffusion (hybrid): This model combines diffusion with an autoregressive loss. During training, the model processes clean tokens and a masked_indices sidecar, learning both a diffusion objective and an autoregressive objective simultaneously.
• DiffusionGemma (block diffusion): A causal encoder reads the clean prompt and response; a bidirectional decoder denoises a response canvas corrupted with uniform-random vocabulary tokens. Training adds a co-trained encoder AR loss. See the DiffusionGemma Fine-Tuning Guide.
• DFlash: This speculative block diffusion model uses a small draft model that proposes tokens for a block conditioned on hidden states from a frozen target language model (LM). A decay-weighted loss trains the draft model to predict target tokens. See the DFlash paper.

Workflow Overview

┌──────────────┐    ┌──────────────┐    ┌──────────────┐    ┌──────────────┐
│  1. Install  │--->│ 2. Configure │--->│   3. Train   │--->│ 4. Generate  │
│              │    │    YAML      │    │              │    │              │
│ uv sync      │    │  Recipe +    │    │  torchrun    │    │  Run dLLM    │
│ or Docker    │    │  dLLM config │    │              │    │  inference   │
└──────────────┘    └──────────────┘    └──────────────┘    └──────────────┘

The following table outlines the key steps in the fine-tuning workflow:

Supported Models

The following table lists the supported models, their training modes, loss functions, inference methods, and example configurations:

For a dedicated walkthrough of DiffusionGemma fine-tuning, including full fine-tuning and LoRA with expert parallelism for its 26B-A4B MoE, refer to the DiffusionGemma Fine-Tuning Guide.

Install NeMo AutoModel

Install NeMo AutoModel with uv:

$
uv venv
$
source .venv/bin/activate
$
uv pip install "nemo-automodel"

Alternatively, use the prebuilt Docker container.

$
docker pull nvcr.io/nvidia/nemo-automodel:26.06.00
$
docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.06.00

For the full set of installation methods, see the Installation Guide.

Configure Your Training Recipe

The following components drive dLLM fine-tuning:

1. A recipe script (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.

Configure LLaDA

Refer to llada_sft.yaml for the full working configuration. The following example shows the key dLLM-specific sections:

1
model:
2
  pretrained_model_name_or_path: GSAI-ML/LLaDA-8B-Base
3
  torch_dtype: float32
4
  trust_remote_code: true
5
6
dllm:
7
  mode: mdlm
8
  mask_token_id: 126336       # LLaDA mask token
9
  eps: 0.001                  # Minimum corruption ratio
10
11
dataset:
12
  unshifted: true             # Required for dLLM training

Configure SCDD

SCDD (see the SCDD paper) 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 for the full working configuration. The following example shows the key dLLM-specific sections.

1
model:
2
  pretrained_model_name_or_path: GSAI-ML/LLaDA-8B-Base
3
  torch_dtype: float32
4
  trust_remote_code: true
5
6
dllm:
7
  mode: scdd
8
  mask_token_id: 126336       # LLaDA mask token
9
  vocab_size: 126464          # Required: the uniform channel samples over the vocabulary
10
  eps: 0.001                  # Minimum diffusion time
11
  num_timesteps: 1000         # Discrete diffusion steps T
12
  uniform_ratio: 0.1          # Peak uniform-noise share (0 degenerates to MDLM)
13
  schedule_shape: 1.0
14
  schedule_peak: 0.5
15
16
dataset:
17
  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 for the full working configuration. The following example shows the key dLLM-specific sections:

1
model:
2
  pretrained_model_name_or_path: nvidia/Nemotron-Labs-Diffusion-8B-Base
3
  torch_dtype: float32          # Master-weight dtype. Use `float32` for an fp32 master copy or `bfloat16` for BF16.
4
  trust_remote_code: true
5
  dlm_paradigm: block_diff       # Required for SFT: HF default "bidirectional" is the inference mode
6
  block_size: 32
7
8
dllm:
9
  mode: hybrid
10
  mask_token_id: 100              # Nemotron-Labs-Diffusion mask token
11
  eps: 0.001
12
  ar_loss_alpha: 0.3              # Weight on the diffusion branch (AR branch is unweighted)
13
  pad_seq_len_divisible: 1024
14
15
dataset:
16
  unshifted: true

Configure DiffusionGemma

Refer to diffusion_gemma_sft.yaml for the full working configuration. The following example shows the key block-diffusion sections:

1
model:
2
  pretrained_model_name_or_path: google/diffusiongemma-26B-A4B-it
3
  torch_dtype: float32          # fp32 master weights; compute is bf16 via mp_policy
4
  canvas_length: 256
5
  self_conditioning: true
6
  freeze_router: true
7
8
dllm:
9
  mode: block_diffusion
10
  block_size: 256
11
  vocab_size: 262144
12
  eps: 0.001
13
  pad_block_size: 256
14
  pad_seq_len_divisible: 256
15
16
dataset:
17
  unshifted: true
18
  mask_history: true            # Supervise only the final turn (single-turn SFT)

Key dLLM Configuration Fields

The following table describes the key configuration fields for dLLM fine-tuning:

Configure DFlash

DFlash trains a small draft model to predict tokens conditioned on a frozen causal target language model. Only the draft model weights are updated. The target language model loads one time and remains frozen.

Refer to dflash_sft.yaml for the full working configuration. The following example shows the key DFlash-specific sections:

1
model:                                          # Draft model
2
  _target_: transformers.AutoModel.from_pretrained
3
  pretrained_model_name_or_path: z-lab/Qwen3-4B-DFlash-b16
4
  trust_remote_code: true
5
  torch_dtype: bfloat16
6
7
dllm:
8
  mode: dflash
9
  mask_token_id: null                           # Resolved automatically from target tokenizer
10
  eps: 0.001
11
12
dflash:
13
  target_model_id: Qwen/Qwen3-4B               # Frozen causal LM
14
  target_torch_dtype: bfloat16
15
  block_size: 0                                 # 0 reads from draft model config
16
  loss_decay_gamma: 0.0                         # 0 uses paper defaults (γ=7 for block_size=16)
17
  num_blocks_per_sample: 512                    # Paper default (Appendix A.1)
18
  attention_backend: flex_attention             # required for N > ~64; sdpa OOMs
19
  overlap_anchors: true                         # paper samples anchors independently

The following table describes the key configuration fields for DFlash fine-tuning:

DFlash Training Metrics

In addition to the shared metrics (such as loss, grad_norm, lr, mem, tps, and mfu), DFlash training runs log a draft top-1 accuracy proxy for the acceptance length, as described in the following table:

Both metrics are computed from the logits that the chunked linear cross-entropy path already produces, without an additional model forward pass. They are reduced across data-parallel and context-parallel ranks using per-rank raw (correct, count) sums. An all-reduce operation sums these values before division. This process ensures that the values are correct across arbitrary per-rank token distributions in any AutoModel distributed mode.

Prepare DFlash Training Data

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.

$
# 1. Serve the target model on the local node (default port 30000)
$
python -m sglang.launch_server \
>
    --model-path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
>
    --served-model-name nemotron-30b \
>
    --trust-remote-code
$
$
# 2. Regenerate the dataset's assistant turns through the target (separate shell)
>
python -m nemo_automodel.components.speculative.regenerate \
>
    --input-data nvidia/Nemotron-Post-Training-Dataset-v2 \
>
    --output-dir /data/dflash-train-regen \
>
    --model nemotron-30b \
>
    --temperature 0.8 \
>
    --shard-size 1000 \
>
    --concurrency 64 \
>
    --resume

The use of --temperature 0.8 (compared to the EAGLE-oriented default value of 0.0 for the script) follows the DFlash paper. Sampling diversity in the supervised tokens teaches the draft model to handle a wider target distribution, which improves the acceptance length. The --concurrency 64 setting better saturates a vLLM or SGLang server.

You can then point the recipe configuration dataset.path_or_dataset_id to the regenerated Parquet shards in the /data/dflash-train-regen directory instead of using the raw Hugging Face dataset.

Configure I-DLM

I-DLM (Introspective Diffusion LM, 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 for the full working configuration. The following example shows the key I-DLM-specific sections.

1
model:
2
  pretrained_model_name_or_path: Qwen/Qwen3-8B
3
  trust_remote_code: true
4
  dtype: bfloat16
5
  attn_implementation: sdpa      # Honors the 4D block-diffusion mask; use flex_attention at scale
6
7
dllm:
8
  mode: idlm
9
  mask_token_id: 151669          # Reserved Qwen3 special token used as the mask
10
  block_length: 1                # Diffusion block size (paper curriculum 1 -> 2 -> 3)
11
  clean_loss_weight: 0.2         # Fixed alpha on CE_clean
12
  auto_balance_clean_loss: false # Set true for the auto-balanced alpha (Eq. 2, b3 stage)
13
14
dataset:
15
  unshifted: true

The mask is built for sdpa or eager (dense additive) or flex_attention (sparse BlockMask, preferred at scale). FlashAttention-2 is unsupported because it ignores arbitrary masks, and context parallelism is also unsupported. The paper trains a block_length 1→2→3 curriculum (one epoch each). Run the stages as successive fine-tunes, and enable auto_balance_clean_loss at the b3 stage.

Fine-Tune the Model

Fine-Tune LLaDA2

$
uv run torchrun --nproc-per-node=8 \
>
    examples/dllm_sft/finetune.py \
>
    -c examples/dllm_sft/llada2_sft.yaml

Fine-Tune with SCDD

$
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 for the training objective and LoRA target modules.

$
python examples/dllm_sft/prep_gsm8k.py
$
$
torchrun --standalone --nproc-per-node=8 \
>
    examples/dllm_sft/finetune.py \
>
    -c examples/dllm_sft/diffusion_gemma_sft.yaml

Fine-Tune with DFlash

$
uv run torchrun --nproc-per-node=8 \
>
    examples/dllm_sft/finetune.py \
>
    -c examples/dllm_sft/dflash_sft.yaml

Fine-Tune with I-DLM

$
uv run torchrun --nproc-per-node=8 \
>
    nemo_automodel/recipes/dllm/train_ft.py \
>
    -c examples/dllm_sft/qwen3_8b_idlm.yaml

Fine-Tune Nemotron-Labs-Diffusion

$
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) 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

$
uv run python examples/dllm_generate/generate.py \
>
    --checkpoint <path> \
>
    --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.

$
uv run python examples/dllm_generate/generate.py \
>
    --checkpoint dllm_checkpoints/llada_scdd/<step>/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.

$
uv run python examples/dllm_generate/generate.py \
>
    --checkpoint <path> \
>
    --prompt "Explain what a neural network is." \
>
    --sampler llada2 \
>
    --steps 32 \
>
    --max_new_tokens 128 \
>
    --block_size 32 \
>
    --threshold 0.5

Generate with Nemotron-Labs-Diffusion

$
uv run python examples/dllm_generate/generate.py \
>
    --checkpoint <path> \
>
    --prompt "Explain what a neural network is." \
>
    --sampler nemotron

Generate with DiffusionGemma

DiffusionGemma generation calls the diffusion sampler that ships with transformers (entropy-bounded denoising with adaptive stopping).

$
uv run python examples/dllm_generate/generate.py \
>
    --checkpoint <path> \
>
    --prompt "Explain what a neural network is." \
>
    --sampler gemma

Generation Parameters

The sampler selected with the --sampler argument supplies these preset values. LLaDA uses the standalone AutoModel sampler. LLaDA2, Nemotron, and DiffusionGemma call the built-in generate() methods of their models. DiffusionGemma only forwards --steps (as max_denoising_steps, default 48) and --max_new_tokens (default 256). The remaining flags keep their upstream Hugging Face defaults. The --remasking and --no_kv_cache arguments do not affect LLaDA2 or DiffusionGemma generation.

The following table lists the default and preset values for the generation parameters: