Model CoverageDiffusion Language Models

DiffusionGemma

View as Markdown

DiffusionGemma is a block-diffusion language model from Google. Instead of generating tokens left-to-right, it denoises a fixed-length canvas of tokens in parallel: a causal encoder reads the prompt and a bidirectional decoder iteratively refines the response canvas. The released checkpoint is a Mixture-of-Experts model with 26B total parameters and ~4B active per token.

TaskText Generation (Block Diffusion, MoE)
ArchitectureDiffusionGemmaForBlockDiffusion
Parameters26B total / ~4B active
HF Orggoogle

Available Models

  • DiffusionGemma 26B-A4B-it (DiffusionGemmaForBlockDiffusion): instruction-tuned block-diffusion MoE.

Architectures

  • DiffusionGemmaForBlockDiffusion — block-diffusion MoE (causal prompt encoder + bidirectional canvas decoder).

Example HF Models

ModelHF ID
DiffusionGemma 26B-A4B-itgoogle/diffusiongemma-26B-A4B-it

Example Recipes

RecipeDescription
diffusion_gemma_sft.yamlFull SFT — DiffusionGemma 26B-A4B with FSDP2 + Expert Parallelism
diffusion_gemma_lora.yamlLoRA SFT — DiffusionGemma 26B-A4B with FSDP2 + Expert Parallelism

Try with NeMo AutoModel

1. Clone the source checkout that contains the example recipes:

$git clone https://github.com/NVIDIA-NeMo/Automodel.git
$cd Automodel

2. Install that checkout’s locked environment (full instructions):

$uv sync --locked --all-groups --extra all

3. Generate the GSM8K chat dataset required by both DiffusionGemma recipes. Run this from the repository root; it writes gsm8k_chat_train.jsonl, the path used by both YAML files:

$uv run python examples/dllm_sft/prep_gsm8k.py

This recipe was validated with Expert Parallelism (EP=8) on a single 8×H100 node. See the Launcher Guide for multi-node setup.

4. Run either recipe from inside the repository:

$# Full SFT
$uv run automodel examples/dllm_sft/diffusion_gemma_sft.yaml --nproc-per-node 8
$
$# LoRA SFT
$uv run automodel examples/dllm_sft/diffusion_gemma_lora.yaml --nproc-per-node 8

Fine-Tuning

See the DiffusionGemma Fine-Tuning Guide for the block-diffusion training objective, self-conditioning, and the full list of supported features (SFT, LoRA, Expert Parallelism, activation checkpointing).

Hugging Face Model Cards