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

• LLaDA (MDLM) — Bidirectional masked diffusion. The model receives corrupted tokens and predicts the clean token at each masked position.

Workflow Overview

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

Supported Models

Install NeMo AutoModel

$
pip3 install nemo-automodel

Alternatively, use the pre-built Docker container:

$
docker pull nvcr.io/nvidia/nemo-automodel:26.04.00
$
docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.04.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

Key dLLM Config Fields

Fine-Tune the Model

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

Generation / Inference

The generation script (generate.py) supports chat, raw, and infilling modes for LLaDA checkpoints.

LLaDA Generation

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

Generation Parameters