Skip to content

Add Parallel Decoding Distillation to FastGen #2329

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mxinO wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add Parallel Decoding Distillation to FastGen #2329

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6454898
Add Parallel Decoding Distillation to FastGen
mxinO Sep 9, 2026
c437ef9
docs: focus PDD guide on usage
mxinO Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ repos:
modelopt/onnx/quantization/ort_patching.py|
modelopt/torch/_deploy/utils/onnx_utils.py|
modelopt/torch/export/transformer_engine.py|
modelopt/torch/fastgen/plugins/qwen_image_pdd.py|
modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py|
modelopt/torch/quantization/export_onnx.py|
modelopt/torch/quantization/plugins/attention.py|
Expand Down
4 changes: 4 additions & 0 deletions CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@ Changelog

- Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts.

*Misc*

- Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency.
Comment thread
mxinO marked this conversation as resolved.
Comment thread
mxinO marked this conversation as resolved.

**Backward Breaking Changes**

**Deprecations**
Expand Down
1 change: 1 addition & 0 deletions LICENSE
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ the following copyright holders, licensed under the Apache License, Version 2.0
Copyright 2022 EleutherAI and the HuggingFace Inc. team
Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
Copyright (c) 2024 Heming Xia
Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved.
Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team
Copyright (c) OpenMMLab. All rights reserved.
Copyright (c) 2022 megvii-model. All Rights Reserved.
Expand Down
210 changes: 7 additions & 203 deletions examples/diffusers/fastgen/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,206 +1,10 @@
# DMD2 distillation for Qwen-Image
# FastGen diffusion examples

Distill [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) into a **few-step
generator** with DMD2 (Distribution Matching Distillation). The distilled student
produces images in as few as **1–4 sampling steps** while matching the base model's
output distribution. Built on `modelopt.torch.fastgen` and NeMo AutoModel's
[`TrainDiffusionRecipe`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py).
This directory contains training and inference examples for diffusion distillation methods in
`modelopt.torch.fastgen`.

> [!NOTE]
> Qwen-Image is a third-party model with its own license terms. Review the
> [Qwen-Image model card](https://huggingface.co/Qwen/Qwen-Image) before downloading or
> redistributing weights or derivatives.
- [DMD2 for Qwen-Image](dmd2/README.md)
- [PDD for Qwen-Image](pdd/README.md)

## Requirements & self-contained data path

This example runs against **stock upstream `nemo_automodel`** (`>=0.4.0,<1.0`; see
`requirements.txt`) from a **source checkout** of Model-Optimizer — the `examples/` tree is not
shipped in the `nvidia-modelopt` pip package. Install the example dependencies with:

```bash
pip install -r examples/diffusers/fastgen/requirements.txt
```

> [!TIP]
> Prefer not to install `nemo_automodel` yourself? Use the **NeMo AutoModel container**, which
> bundles it (with the diffusion extras) — then you only need a source checkout of Model-Optimizer
> for the `examples/` tree and can skip the `pip install` above:
>
> ```bash
> docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.04
> ```

The DMD2 data loading (`fastgen_data/`) and raw-image preprocessing (`preprocess/`) are
**vendored into this example** (from NeMo-AutoModel, Apache-2.0) so that **no modifications to
`nemo_automodel` are required**. The entry points put this directory on `sys.path`, so the
configs reference the vendored builders as `_target_: fastgen_data.build_*`. The DMD2 math in
`modelopt/torch/fastgen/` is unchanged.

**Build the training cache from raw images** (Qwen-Image VAE latents + text embeddings):

```bash
python examples/diffusers/fastgen/preprocess_qwen_image.py image \
--image_dir <raw images> --output_dir <cache dir> --processor qwen_image \
--caption_format meta_json
```

The CFG negative-prompt embedding (the config's `negative_prompt_embedding_path`) is generated
once from the same Qwen text encoder:

```bash
python examples/diffusers/fastgen/make_negative_prompt_embedding.py \
--output <cache dir>/negative_prompt_embedding.pt
```

Then point the config's `data.dataloader.cache_dir` at `<cache dir>` and its
`negative_prompt_embedding_path` at `<cache dir>/negative_prompt_embedding.pt`, and train (below).

## How DMD2 works

DMD2 trains three networks together:

| Model | Role |
|---|---|
| **Student** | the few-step generator you keep |
| **Fake-score** | a diffusion model that tracks the *student's* current output distribution |
| **Teacher** | the frozen base Qwen-Image model (the *target* distribution) |

The distribution-matching gradient pushes the student toward the teacher and away from
the fake-score. Training alternates between two phases, controlled by `student_update_freq`:

```text
each step:
if step % student_update_freq == 0: # student phase
update the student (distribution-matching [+ optional GAN] loss)
update the student EMA
else: # fake-score phase
update the fake-score network to track the student
```

The canonical config additionally enables **CFG** (classifier-free guidance on the
teacher) and a lightweight **GAN** branch (a discriminator head on a teacher feature
block, plus an R1 gradient penalty) for sharper samples.

## Install

From the repo root:

```bash
pip install -e ".[all]" # ModelOpt + torch + diffusers
pip install -r examples/diffusers/fastgen/requirements.txt # nemo_automodel
```

`nemo_automodel[diffusion]` pulls in diffusers, accelerate, and the `TrainDiffusionRecipe`
this example subclasses.

## Real-data training

`configs/dmd2_qwen_image.yaml` is the canonical config: 4-step student, CFG, and the
GAN + R1 branch, trained on a preprocessed latent cache. Before launching, provide:

- **A preprocessed Qwen-Image latent cache** — set `data.dataloader.cache_dir`.
- **A precomputed negative-prompt embedding** (required for CFG) — set
`data.dataloader.negative_prompt_embedding_path`.
- **An output directory** — set `checkpoint.checkpoint_dir`.

The model path defaults to `Qwen/Qwen-Image`; point it at a local snapshot to avoid
re-downloading on every job. Then:

```bash
torchrun --nproc-per-node=8 \
examples/diffusers/fastgen/dmd2_finetune.py \
--config examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml \
--step_scheduler.max_steps=5000
```

Any `DMDConfig` field can be overridden on the CLI (e.g. `--dmd2.guidance_scale=3.5`).

### Checkpoints & resuming

Checkpoints land under `checkpoint.checkpoint_dir`. Alongside the student, the recipe
saves the DMD2 sidecars needed to resume exactly: the fake-score model + optimizer, the
student EMA (`ema_shadow.pt`), and the DMD iteration counter (`dmd_state.pt`). With
`restore_from: LATEST` a re-launch auto-resumes from the newest checkpoint; pin a
specific one with `--checkpoint.restore_from=epoch_0_step_500`.

## Inference

After training, sample from the distilled student. The pipeline loads your consolidated
student transformer plus the base Qwen-Image VAE / text encoder / tokenizer:

```python
import torch
from inference_dmd2_qwen_image import QwenImageDMDInferencePipeline

pipe = QwenImageDMDInferencePipeline.from_pretrained(
student_path="/path/to/checkpoint/epoch_0_step_500/model/consolidated",
base_pipeline_path="Qwen/Qwen-Image",
ema_path=None, # or ".../ema_shadow.pt" to sample the EMA weights
torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
prompt="a small red cube on a white table",
num_inference_steps=4, # match the student_sample_steps you trained with
height=1024, width=1024,
generator=torch.Generator("cuda").manual_seed(42),
).images[0]
image.save("sample.png")
```

Or run the bundled CLI for a quick check:

```bash
python examples/diffusers/fastgen/inference_dmd2_qwen_image.py \
--student_path /path/to/checkpoint/.../model/consolidated \
--base_pipeline_path Qwen/Qwen-Image \
--prompt "a small red cube on a white table" \
--height 512 --width 512
```

Set `num_inference_steps` to the number of steps the student was trained for
(`dmd2.student_sample_steps` — e.g. 4 for the canonical config, or 1 for a single-step
student).

## Config reference

| Section | Key | Role |
|---|---|---|
| `model` | `pretrained_model_name_or_path` | Qwen-Image HF id or local snapshot. |
| `model` | `mode` | `finetune` — loads the pretrained weights. |
| `step_scheduler` | `global_batch_size`, `local_batch_size`, `max_steps`, `ckpt_every_steps`, `log_every` | Standard AutoModel scheduling knobs. |
| `dmd2` | `recipe_path` | Built-in fastgen recipe to hydrate `DMDConfig` from (`general/distillation/dmd2_qwen_image`). |
| `dmd2` | `pipeline_plugin` | `qwen_image` — selects `QwenImageDMDPipeline` (2×2 patch packing / img_shapes). |
| `dmd2` | `student_sample_steps` | Number of student sampling steps (e.g. 4). |
| `dmd2` | `guidance_scale` | CFG strength on the teacher (`null` disables CFG; requires a negative-prompt embedding when set). |
| `dmd2` | `gan_loss_weight_gen`, `gan_r1_reg_weight`, `gan_feature_indices`, … | GAN branch (set `gan_loss_weight_gen: 0` to disable). |
| `dmd2` | `fake_score_lr`, `discriminator_lr` | Separate LRs for the fake-score / discriminator optimizers. |
| `dmd2` | `sample_t_cfg`, `ema` | Timestep sampling + student EMA settings. |
| `optim` | `learning_rate`, `optimizer.*` | Student AdamW knobs. |
| `fsdp` | `dp_size`, `tp_size`, `activation_checkpointing`, … | FSDP2 parallelism (set `dp_size` to your GPU count). |
| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Latent cache dir + optional CFG negative-prompt embedding. |
| `checkpoint` | `checkpoint_dir`, `model_save_format`, `restore_from` | Output dir, save format, resume behavior. |

## Troubleshooting

**`CUDA out of memory`.** Training holds three Qwen-Image transformers (student + teacher
- fake-score) plus optimizer state. Shard across more GPUs (raise `--fsdp.dp_size`),
or enable `--fsdp.activation_checkpointing=true`.

**Loss is `NaN` on step 0.** Almost always an out-of-range timestep — confirm you haven't
overridden `dmd2.pred_type` away from `flow` (Qwen-Image is a rectified-flow model) or
changed the timestep schedule.

**`guidance_scale is set but negative_encoder_hidden_states was not provided`.** CFG needs
a precomputed negative-prompt embedding. Set `data.dataloader.negative_prompt_embedding_path`,
or set `dmd2.guidance_scale: null` to disable CFG.

**Dataloader yields empty batches.** Ensure your cache has at least
`local_batch_size * fsdp.dp_size` items; the distributed sampler drops incomplete batches.

## Reference

- Fastgen library: [`modelopt/torch/fastgen/`](../../../modelopt/torch/fastgen/)
- Built-in recipe: [`modelopt_recipes/general/distillation/dmd2_qwen_image.yaml`](../../../modelopt_recipes/general/distillation/dmd2_qwen_image.yaml)
- AutoModel recipe this example subclasses:
[`nemo_automodel/recipes/diffusion/train.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py)
The `fastgen_data/` and `preprocess/` packages are shared utilities. Algorithm-specific entrypoints,
configs, checkpoint helpers, and documentation live in their corresponding subdirectory.
Loading
Loading