dLLM Fine-Tuning

Introduction

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

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

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

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 config. The key I-DLM-specific sections are:

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 (it ignores arbitrary masks), as is context parallelism. The paper trains a block_length 1→2→3 curriculum (one epoch each). Run the stages as successive fine-tunes, enabling auto_balance_clean_loss at the b3 stage.

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 with I-DLM

$
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

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

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

Generate with LLaDA2

LLaDA2 generation calls the model’s built-in block-refinement generate() method.

$
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

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

$
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 and Nemotron call the built-in generate() methods of their models. The --remasking and --no_kv_cache arguments do not affect LLaDA2 generation. The following table lists the default and preset values for the generation parameters.