dLLM Fine-Tuning

Introduction

Diffusion language models (dLLMs) generate text by iteratively denoising masked tokens, rather than generating one token at a time left-to-right like autoregressive (AR) models. Starting from a sequence of [MASK] tokens, the model progressively unmasks the most confident positions over multiple denoising steps until the full response is revealed.

This approach enables parallel token generation and bidirectional attention, which gives the model more context for each prediction compared to AR models.

NeMo AutoModel currently supports the following dLLM model families:

• LLaDA / LLaDA2 (MDLM) — Bidirectional masked diffusion. The model receives corrupted tokens and predicts the clean token at each masked position (see LLaDA2 paper).
• Nemotron-Labs-Diffusion (Hybrid) — Combines diffusion with an autoregressive loss. During training, the model processes clean tokens plus a masked_indices sidecar and learns both a diffusion objective and an AR objective simultaneously.
• DFlash — Speculative block diffusion. A small draft model proposes tokens for a block conditioned on frozen target LM hidden states; a decay-weighted loss trains it to predict the target’s distribution (see DFlash paper).

Workflow Overview

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

Supported Models

Install NeMo AutoModel

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

Alternatively, use the pre-built 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

dLLM fine-tuning is driven by:

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 selects the strategy:

LLaDA Configuration

See llada_sft.yaml for the full working config. The key dLLM-specific sections are:

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

Nemotron-Labs-Diffusion Configuration

See nemotron_labs_diffusion_sft.yaml for the full working config. The key dLLM-specific sections are:

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

Key dLLM Config Fields

DFlash Configuration

DFlash trains a small draft model to predict tokens conditioned on a frozen causal target LM. Only the draft model’s weights are updated; the target LM is loaded once and kept frozen.

See dflash_sft.yaml for the full working config. The key DFlash-specific sections are:

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

DFlash Training Metrics

In addition to the shared metrics (loss, grad_norm, lr, mem, tps, mfu), DFlash runs log a draft top-1 accuracy proxy for acceptance length:

Both are computed for free inside the chunked linear-CE path (same logits used for the loss) and DP/CP-reduced via per-rank raw (correct, count) sums that are SUM-allreduced and then divided post-reduction, so the values are correct across arbitrary per-rank token distributions under any of AutoModel’s distributed modes.

Prepare DFlash Training Data

The paper trains on responses regenerated by the target model (§5.1): “Instead of directly using the original dataset, we construct our training set with the responses generated by the target model for better target alignment.” Skipping this step trains the draft on a different output distribution than the target produces at inference, which directly reduces acceptance length.

The existing nemo_automodel.components.speculative.regenerate script handles this. Stand up an SGLang server hosting the target, then re-roll 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

--temperature 0.8 (vs the script’s EAGLE-oriented 0.0 default) follows the DFlash paper: sampling diversity in the supervised tokens teaches the draft to handle a wider target distribution, improving acceptance length. --concurrency 64 better saturates one vLLM/SGLang server.

Then point the recipe’s dataset.path_or_dataset_id at the regenerated parquet shards (/data/dflash-train-regen) instead of the raw HF dataset.

Fine-Tune the Model

Fine-Tune LLaDA2

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

Fine-Tune with DFlash

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

Fine-Tune Nemotron-Labs-Diffusion

$
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, raw, and infilling modes. Pick the sampler that matches the trained family with --sampler {llada,nemotron}.

--checkpoint accepts any of: a path to a consolidated/ directory, a step directory (.../epoch_0_step_499), or the top-level checkpoint dir (the script will follow LATEST/model/consolidated/).

Generate with LLaDA

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

Generate with Nemotron-Labs-Diffusion

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

Generation Parameters

The sampler selected with --sampler supplies these preset values. For LLaDA chat and raw generation, the standalone sampler consumes every field shown below and honors its command-line override. LLaDA infilling uses --steps, --block_size, --temperature, and --remasking while filling masks already present in the input, so it does not use --max_new_tokens. For Nemotron chat and raw generation, the script calls the model’s built-in generate(...), which consumes --max_new_tokens, --block_size, and --temperature from this table (as well as --threshold, which is not shown). Although --steps and --remasking are accepted and stored in the Nemotron preset config, the built-in generation path does not use them, so overriding those two flags does not affect that path.