Megatron-LM Diffusion Language Model Training Guide

Industrial-grade distributed training for Diffusion Language Models (dLLMs) built on top of Megatron-LM's tensor/pipeline/data parallelism infrastructure.

This extension enables training of masked diffusion LMs (MDLM), block diffusion LMs (BD3LM), and autoregressive-to-diffusion (A2D) converted models with the same efficiency guarantees as Megatron-LM's GPT training.

Overview
Supported Models
Architecture
Quick Start
A2D: Converting AR Models to Diffusion
MDLM Training (Pretrain & SFT)
BD3LM Training (Block Diffusion)
Distributed Training
Configuration Reference
Precision Alignment with dllm
Testing
Roadmap

Overview

This framework brings three diffusion language modeling paradigms to Megatron-LM:

Paradigm	Script	Description	Reference
MDLM	`pretrain_diffusion.py`	Masked Diffusion LM — stochastic token masking with timestep scheduling	arxiv:2406.07524
BD3LM	`pretrain_bd3lm.py`	Block Diffusion — block-wise attention with x_t/x_0 concatenation	arxiv:2503.09573
A2D	`tools/convert_a2d_to_megatron.py`	AutoRegressive-to-Diffusion — convert any AR model to diffusion training	dllm framework

All implementations are precision-aligned with the dllm reference framework (< 1e-5 relative error) and support Megatron-LM's full parallelism stack:

Tensor Parallelism (TP) — split attention heads and MLP across GPUs
Pipeline Parallelism (PP) — split layers across GPUs with 1F1B scheduling
Data Parallelism (DP) — replicate model across GPU groups
Context Parallelism (CP) — split sequence dimension

Supported Models

Model	Type	Sizes	Training Script	Converter
LLaDA	Native Diffusion (MDLM)	8B	`pretrain_diffusion.py`	`convert_a2d_to_megatron.py` (LLaDA mode)
SDAR	Block Diffusion (BD3LM)	1.7B, 4B, 8B, 30B-MoE	`pretrain_bd3lm.py`	`convert_a2d_to_megatron.py` (Qwen-style)
Qwen3 (A2D)	AR→Diffusion	0.6B, 1.7B, 4B, 8B, 32B	`pretrain_diffusion.py`	`convert_a2d_to_megatron.py`
LLaMA (A2D)	AR→Diffusion	7B, 13B, 70B	`pretrain_diffusion.py`	`convert_a2d_to_megatron.py`
Dream	Native Diffusion (MDLM)	7B+	`pretrain_diffusion.py`	`convert_a2d_to_megatron.py` (Qwen-style)

Architecture

File Structure

Megatron-LM/
├── pretrain_diffusion.py                     # MDLM training entry point (Pretrain + SFT)
├── pretrain_bd3lm.py                         # BD3LM training entry point (Pretrain + SFT)
├── diffusion_builders.py                     # Model builder (like gpt_builders.py)
├── tools/
│   └── convert_a2d_to_megatron.py            # HF → Megatron weight converter
├── megatron/core/models/diffusion/
│   ├── __init__.py
│   ├── alpha_scheduler.py                    # Linear & Cosine diffusion schedulers
│   ├── diffusion_model.py                    # DiffusionLanguageModel (extends LanguageModule)
│   └── diffusion_layer_specs.py              # Bidirectional attention layer specs
├── scripts/
│   └── test_diffusion_multigpu.sh            # Multi-GPU validation script
└── tests/unit_tests/models/
    ├── test_diffusion_alpha_scheduler.py      # 23 scheduler precision tests
    ├── test_diffusion_precision.py            # 23 loss computation precision tests
    └── test_diffusion_e2e_precision.py        # 18 end-to-end model precision tests

How Diffusion Training Differs from GPT

In standard GPT training, the model predicts the next token autoregressively with causal attention. In diffusion training:

Bidirectional attention — all tokens attend to all other tokens (no causal mask)
Stochastic masking — random tokens are replaced with [MASK] based on diffusion timestep t
Weighted loss — loss is weighted by the diffusion scheduler: w(t) = -α'(t) / (1 - α(t))
Labels are clean tokens — the model predicts original (unmasked) tokens from noised input

The diffusion-specific logic lives entirely in pretrain_diffusion.py::diffusion_forward_step(), not in the model. The DiffusionLanguageModel is a standard bidirectional transformer.

Quick Start

Prerequisites

Megatron-LM with dependencies installed
CUDA-capable GPU(s)
Python 3.10+, PyTorch 2.6+

1. Train with Mock Data (No Checkpoint Needed)

# Single GPU — MDLM with random weights
CUDA_VISIBLE_DEVICES=0 torchrun --nproc-per-node=1 pretrain_diffusion.py \
    --num-layers 4 --hidden-size 256 --num-attention-heads 4 \
    --seq-length 128 --max-position-embeddings 128 \
    --micro-batch-size 4 --global-batch-size 4 \
    --train-iters 100 --lr 1e-4 --min-lr 1e-5 \
    --lr-decay-style cosine --lr-warmup-iters 10 \
    --bf16 --mock-data --vocab-size 32000 \
    --position-embedding-type rope --transformer-impl local \
    --no-persist-layer-norm --no-masked-softmax-fusion \
    --no-bias-gelu-fusion --attention-softmax-in-fp32 \
    --diffusion-mask-token-id 0 \
    --tokenizer-type NullTokenizer --split 100,0,0 \
    --log-interval 1

2. Train with a Real Model (Qwen3-1.7B A2D)

# Step 1: Convert HF checkpoint to Megatron format
python tools/convert_a2d_to_megatron.py \
    --hf-model-path /path/to/Qwen3-1.7B-Base \
    --output-path /path/to/megatron_ckpt_qwen3_a2d \
    --mask-token-id 151936 --add-mask-at-end

# Step 2: Train
CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun --nproc-per-node=1 pretrain_diffusion.py \
    --load /path/to/megatron_ckpt_qwen3_a2d \
    --no-load-optim --no-load-rng \
    --num-layers 28 --hidden-size 2048 \
    --num-attention-heads 16 --num-query-groups 8 \
    --group-query-attention \
    --ffn-hidden-size 6144 --swiglu \
    --seq-length 2048 --max-position-embeddings 32768 \
    --micro-batch-size 2 --global-batch-size 16 \
    --train-iters 1000 --lr 1e-4 --min-lr 1e-6 \
    --lr-decay-style cosine --lr-warmup-iters 50 \
    --bf16 --normalization RMSNorm --disable-bias-linear \
    --position-embedding-type rope --rotary-base 1000000 \
    --qk-layernorm \
    --transformer-impl local --no-persist-layer-norm \
    --no-masked-softmax-fusion --no-bias-gelu-fusion \
    --attention-softmax-in-fp32 \
    --diffusion-mask-token-id 151936 \
    --diffusion-alpha-scheduler LinearAlphaScheduler \
    --vocab-size 151936 --padded-vocab-size 151937 \
    --tokenizer-type NullTokenizer --split 100,0,0 \
    --data-path /path/to/megatron-format-data \
    --log-interval 10 --save-interval 500 \
    --save /path/to/checkpoints

A2D: Converting AR Models to Diffusion

The convert_a2d_to_megatron.py tool converts HuggingFace autoregressive models into Megatron's DiffusionLanguageModel checkpoint format.

Supported Architectures

Source Model	HF Config `model_type`	Weight Naming	GQA	QK Norm
Qwen3	`qwen3`	`model.layers.{i}.self_attn.*`	Yes	Yes
Qwen2	`qwen2`	`model.layers.{i}.self_attn.*`	Yes	No
LLaMA	`llama`	`model.layers.{i}.self_attn.*`	Yes	No
SDAR	`sdar`	`model.layers.{i}.self_attn.*`	Yes	Yes
LLaDA	`llada`	`model.transformer.blocks.{i}.*`	No (MHA)	No

Usage

# Qwen3 (adds mask token at end of vocab)
python tools/convert_a2d_to_megatron.py \
    --hf-model-path /path/to/Qwen3-1.7B-Base \
    --output-path /path/to/output \
    --mask-token-id 151936 --add-mask-at-end

# SDAR (mask token already in vocab)
python tools/convert_a2d_to_megatron.py \
    --hf-model-path /path/to/SDAR-1.7B-Chat \
    --output-path /path/to/output \
    --mask-token-id 151667

# LLaDA (different weight naming convention)
python tools/convert_a2d_to_megatron.py \
    --hf-model-path /path/to/LLaDA-8B-Base \
    --output-path /path/to/output \
    --mask-token-id 126336

What the Converter Does

Loads HF model weights (safetensors or pytorch_model.bin, sharded or single file)
Maps weight names to Megatron's DiffusionLanguageModel layout
Handles GQA (Grouped-Query Attention) by interleaving Q/K/V into Megatron's per-group layout
Handles QK LayerNorm (Qwen3/SDAR per-head RMS normalization)
Extends embedding table for [MASK] token (randomly initialized)
Saves in Megatron's release checkpoint format

Output Structure

output/
├── release/
│   └── mp_rank_00/
│       └── model_optim_rng.pt     # Megatron checkpoint
├── latest_checkpointed_iteration.txt  # "release"
└── a2d_conversion_summary.json        # Conversion metadata

MDLM Training (Pretrain & SFT)

Pretraining

torchrun --nproc-per-node=8 pretrain_diffusion.py \
    --load /path/to/megatron-ckpt \
    --no-load-optim --no-load-rng \
    [MODEL ARGS] \
    --diffusion-mask-token-id <MASK_ID> \
    --diffusion-alpha-scheduler LinearAlphaScheduler \
    --diffusion-loss-weight-type scheduler \
    --diffusion-loss-norm-type token \
    --diffusion-time-epsilon 0.001 \
    [DATA ARGS] \
    [TRAINING ARGS]

Supervised Fine-Tuning (SFT)

SFT mode masks the loss on prompt tokens, training only on response tokens.

torchrun --nproc-per-node=8 pretrain_diffusion.py \
    --sft \
    --load /path/to/megatron-ckpt \
    [same args as pretrain] \
    --data-path /path/to/sft-data

In SFT mode, the dataset produces loss_mask=0 for prompt tokens. The diffusion masking logic uses maskable_mask = (loss_mask > 0), so prompt tokens are never masked or included in loss. This matches dllm's mask_prompt_loss=True behavior.

Training Pipeline (What Happens Inside)

1. get_batch() → tokens, labels, loss_mask, attention_mask, position_ids
2. diffusion_forward_step():
   a. Sample timestep:     t = ε + (1-ε) × rand(b)
   b. Compute mask rate:   p_mask = 1 - α(t)
   c. Stochastic masking:  masked_mask = (rand < p_mask) & maskable_mask
   d. Create noised input: noised_ids = where(mask, [MASK], tokens)
   e. Compute loss weight: w(t) = -α'(t) / (1 - α(t) + 1e-6)
   f. Build loss_mask:     diffusion_loss_mask = w(t) × masked_mask / norm_factor
   g. Forward pass:        model(noised_ids, ..., labels=clean_tokens)
3. loss_func() → sum(CE × diffusion_loss_mask)

BD3LM Training (Block Diffusion)

BD3LM extends MDLM with block-wise attention where the model attends to both noised blocks (x_t) and clean context blocks (x_0).

Usage

torchrun --nproc-per-node=8 pretrain_bd3lm.py \
    --load /path/to/sdar-megatron-ckpt \
    --no-load-optim --no-load-rng \
    [MODEL ARGS] \
    --diffusion-mask-token-id <MASK_ID> \
    --diffusion-block-size 32 \
    [DATA ARGS] \
    [TRAINING ARGS]

Key Differences from MDLM

Aspect	MDLM	BD3LM
Input	`[x_t]` — shape `[b, L]`	`[x_t, x_0]` — shape `[b, 2L]`
Position IDs	`[0..L-1]`	`[0..L-1, 0..L-1]` (repeated)
Attention	Full bidirectional	3-part block mask
Loss	Over all masked tokens	Over first L tokens only (x_t half)
Block size	N/A	Configurable via `--diffusion-block-size`

Block Attention Mask

The BD3LM attention mask combines three sub-masks:

M_BD  (Block Diagonal):      self-attention within same block, same segment
M_OBC (Offset Block Causal): x_t attends to earlier x_0 blocks
M_BC  (Block Causal):        x_0 attends to current and earlier x_0 blocks

Final mask = M_BD | M_OBC | M_BC

Distributed Training

Tensor Parallelism (TP)

Split model across GPUs along the head/hidden dimension.

CUDA_DEVICE_MAX_CONNECTIONS=1 \
torchrun --nproc-per-node=2 pretrain_diffusion.py \
    --tensor-model-parallel-size 2 \
    [other args]

Important: Diffusion masking uses a deterministic CPU RNG seeded by (iteration, dp_rank, cp_rank) to ensure all TP ranks generate identical masks. The noised_input_ids and diffusion_loss_mask are computed on TP rank 0 and broadcast to other TP ranks.

Pipeline Parallelism (PP)

Split layers across GPUs vertically.

torchrun --nproc-per-node=4 pretrain_diffusion.py \
    --tensor-model-parallel-size 2 \
    --pipeline-model-parallel-size 2 \
    [other args]

The first PP stage applies diffusion masking and sends noised tokens through the pipeline. The last PP stage recomputes the diffusion_loss_mask using the same deterministic RNG seed for loss computation.

Example: 8 GPU Configurations

# TP=8 (maximum intra-layer parallelism)
CUDA_DEVICE_MAX_CONNECTIONS=1 \
torchrun --nproc-per-node=8 pretrain_diffusion.py \
    --tensor-model-parallel-size 8 ...

# TP=4, DP=2 (balanced)
CUDA_DEVICE_MAX_CONNECTIONS=1 \
torchrun --nproc-per-node=8 pretrain_diffusion.py \
    --tensor-model-parallel-size 4 ...

# TP=2, PP=2, DP=2 (all three)
CUDA_DEVICE_MAX_CONNECTIONS=1 \
torchrun --nproc-per-node=8 pretrain_diffusion.py \
    --tensor-model-parallel-size 2 \
    --pipeline-model-parallel-size 2 ...

Configuration Reference

Diffusion-Specific Arguments

Argument	Default	Description
`--diffusion-mask-token-id`	None (required)	Token ID for `[MASK]`. Model-specific.
`--diffusion-alpha-scheduler`	`LinearAlphaScheduler`	Scheduler for masking rate α(t). Options: `LinearAlphaScheduler`, `CosineAlphaScheduler`.
`--diffusion-time-epsilon`	`1e-3`	Lower bound for timestep sampling: t ~ U[ε, 1).
`--diffusion-loss-weight-type`	`scheduler`	Loss weighting: `scheduler` (w(t) from α) or `uniform` (all ones).
`--diffusion-loss-norm-type`	`token`	Loss normalization: `token` (by total valid tokens), `sequence` (per-seq), `batch` (by batch size).
`--diffusion-block-size`	`32`	Block size for BD3LM block diffusion.
`--diffusion-a2d-source-model`	None	Path to A2D-converted checkpoint.

Model Architecture Arguments

For Qwen3/SDAR-style models:

--num-layers 28
--hidden-size 2048
--num-attention-heads 16
--num-query-groups 8          # KV heads for GQA
--group-query-attention       # Enable GQA
--ffn-hidden-size 6144        # Intermediate size
--swiglu                      # SiLU-gated MLP
--normalization RMSNorm       # RMS normalization
--disable-bias-linear         # No bias in linear layers
--qk-layernorm                # Per-head Q/K RMS norms (Qwen3/SDAR)
--position-embedding-type rope
--rotary-base 1000000

For LLaDA-style models:

--num-layers 32
--hidden-size 4096
--num-attention-heads 32
--ffn-hidden-size 12288
--swiglu
--normalization RMSNorm
--disable-bias-linear
--position-embedding-type rope
--rotary-base 500000

Common Training Arguments

--bf16                             # BFloat16 mixed precision
--transformer-impl local           # Use MCore local backend (no TE required)
--no-persist-layer-norm
--no-masked-softmax-fusion
--no-bias-gelu-fusion
--attention-softmax-in-fp32
--tokenizer-type NullTokenizer     # For mock data / pre-tokenized data
--split 100,0,0                    # Train-only split

Precision Alignment with dllm

All loss computations are verified to match the dllm framework within < 1e-5 relative error.

What is Verified

Component	Test	Tolerance
Alpha schedulers (Linear, Cosine)	`alpha(t)`, `weight(t)`, `derivative(t)`	< 1e-6
Timestep sampling	`t = ε + (1-ε) × rand(b)`	Exact
Masking pattern	`(rand < p_mask) & maskable_mask`	Bitwise identical
Loss weights	`w(t) × masked_mask / norm_factor`	< 1e-5
CE loss	`sum(CE × diffusion_loss_mask)`	< 1e-5
All scheduler × norm type combos	6 combinations (2 schedulers × 3 norms)	< 1e-5
BFloat16 logits	Forward pass in bf16	< 1e-3
GPU consistency	CPU vs GPU results	< 1e-5
BD3LM block mask	3-part mask structure	Bitwise identical to dllm

dllm Reference Baselines

Model	Type	Steps	Avg Loss
Qwen3-1.7B A2D	MDLM PT	10	8.2060
LLaDA-8B	MDLM PT	5	13.4685

Testing

Run All Diffusion Tests

# Alpha scheduler precision tests (23 tests)
python -m pytest tests/unit_tests/models/test_diffusion_alpha_scheduler.py -v

# Loss computation precision tests (23 tests)
python -m pytest tests/unit_tests/models/test_diffusion_precision.py -v

# End-to-end model precision tests (18 tests)
python -m pytest tests/unit_tests/models/test_diffusion_e2e_precision.py -v

Multi-GPU Validation

bash scripts/test_diffusion_multigpu.sh

Tests the following configurations with mock data:

Single GPU (TP=1, PP=1)
TP=2
PP=2
TP=2, PP=2 (4 GPUs)
TP=2, DP=2 (4 GPUs)

Roadmap

Completed (Phase 1 & 2)

MDLM training (pretrain + SFT)
BD3LM block diffusion training (pretrain + SFT)
A2D (AR→Diffusion) model conversion
HF→Megatron weight converter (Qwen3, SDAR, LLaDA)
Alpha schedulers (Linear, Cosine) with precision verification
Full TP/PP/DP/CP distributed training support
64 unit tests with dllm precision alignment
Bidirectional attention layer specs (local + TE backends)
SFT with prompt loss masking

Phase 3: MoE Support (Next)

Expert Parallelism (EP) for diffusion MoE — integrate Megatron's existing MoE infrastructure with diffusion training
LLaDA-MoE (7B-A1B) — sparse expert routing with MDLM loss
LLaDA 2.0 (100B+) — large-scale MoE diffusion training
SDAR-30B-A3B-Chat — MoE block diffusion with expert parallelism
MoE-specific load balancing with diffusion masking (sparse token distributions)

Phase 4: Inference & Sampling

MDLM Sampler — multi-step iterative demasking with confidence-based remasking
BD3LM Sampler — block-by-block sequential generation with inner diffusion loop
Classifier-Free Guidance (CFG) — unconditional + conditional logit mixing
vLLM/TRT-LLM integration — efficient inference serving for diffusion LMs
Infill/editing mode — fill masked positions in existing sequences

Phase 5: Advanced Training

FP8 training via Transformer Engine — quantized diffusion training for throughput
Sequence Parallelism — further scale sequence lengths for diffusion models
Activation checkpointing — memory optimization for long sequences
Multi-Token Prediction (MTP) for diffusion — predict multiple tokens per denoising step
Right-shift logits (AR-style MDLM variant, right_shift_logits=True)
Megatron→HF checkpoint converter — export trained Megatron models back to HF format
Real data pipeline — Megatron binary format preprocessor for diffusion training data
Curriculum learning — progressive sequence length / masking rate scheduling
Evaluation harness integration — lm-eval-harness with diffusion loglikelihood estimation

Phase 6: Ecosystem

Comprehensive benchmarks — throughput comparison (tokens/sec, MFU) vs dllm+FSDP
Pre-built training recipes — ready-to-use configs for common model sizes (1.7B, 8B, 70B)
Wandb/TensorBoard logging — diffusion-specific metrics (NLL, PPL, mask ratio, timestep distribution)
RLHF/DPO for diffusion models — reward model training and alignment
Continuous pretraining — resume diffusion training from AR checkpoints with progressive masking

Citation

If you use this framework, please cite:

@article{mdlm2024,
    title={Simple and Effective Masked Diffusion Language Models},
    author={Sahoo, Subham and Arriola, Marianne and Schiff, Yair and Gokaslan, Aaron and Marroquin, Edgar and Chiu, Justin T and Rush, Alexander and Kuleshov, Volodymyr},
    journal={arXiv preprint arXiv:2406.07524},
    year={2024}
}

@article{llada2025,
    title={Large Language Diffusion Models},
    author={Nie, Shen and Zhu, Fengqi and You, Chao and Zhang, Xiaojie and Meng, Chenyang and others},
    journal={arXiv preprint arXiv:2502.09992},
    year={2025}
}

@article{bd3lm2025,
    title={Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models},
    author={Arriola, Marianne and Sahoo, Subham and Schiff, Yair and Rush, Alexander and Kuleshov, Volodymyr},
    journal={arXiv preprint arXiv:2503.09573},
    year={2025}
}

FilesExpand file tree

diffusion_training.md

Latest commit

History

diffusion_training.md

File metadata and controls

Megatron-LM Diffusion Language Model Training Guide

Table of Contents

Overview

Supported Models

Architecture

File Structure

How Diffusion Training Differs from GPT

Quick Start

Prerequisites

1. Train with Mock Data (No Checkpoint Needed)

2. Train with a Real Model (Qwen3-1.7B A2D)

A2D: Converting AR Models to Diffusion

Supported Architectures

Usage

What the Converter Does

Output Structure

MDLM Training (Pretrain & SFT)

Pretraining

Supervised Fine-Tuning (SFT)

Training Pipeline (What Happens Inside)

BD3LM Training (Block Diffusion)

Usage

Key Differences from MDLM

Block Attention Mask

Distributed Training

Tensor Parallelism (TP)

Pipeline Parallelism (PP)

Example: 8 GPU Configurations

Configuration Reference

Diffusion-Specific Arguments

Model Architecture Arguments

Common Training Arguments

Precision Alignment with dllm

What is Verified

dllm Reference Baselines

Testing

Run All Diffusion Tests

Multi-GPU Validation

Roadmap

Completed (Phase 1 & 2)

Phase 3: MoE Support (Next)

Phase 4: Inference & Sampling

Phase 5: Advanced Training

Phase 6: Ecosystem

Citation