Files
Model-Optimizer/modelopt
jingyu-mlandClaude Opus 4.8 aec72ffa68 Add DMD2 distillation for Qwen-Image (fastgen) (#1326)
### What does this PR do?

**Type of change:** New example + new `modelopt.torch.fastgen` library
module.

Adds **DMD2 (Distribution Matching Distillation) for Qwen-Image** —
distilling the base model into a few-step (1–4) generator. Includes the
framework-agnostic `modelopt.torch.fastgen` loss library (DMD pipeline,
EMA, optional GAN discriminator) and a NeMo AutoModel–based training
example with a mock-data smoke config, a real-data config, and inference
/ export scripts.

**Noted**: the example script will be migrated to AutoModel repo

### Usage

```bash
# Mock-data wiring smoke — runs end-to-end with no dataset to prepare
torchrun --nproc-per-node=8 \
    examples/diffusers/fastgen/dmd2_finetune.py \
    --config examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml
```

See `examples/diffusers/fastgen/README.md` for real-data training and
inference.

### Testing

Unit tests under `tests/unit/torch/fastgen/`; `pre-commit` /
code-quality clean.

### Before your PR is "*Ready for review*"

- Backward compatible?: ✅ (new, additive module)
- Followed `CONTRIBUTING.md` for any copied code / new deps: ✅
- New tests added?: ✅
- Updated Changelog?: N/A


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Adds a FastGen-based distillation framework (DMD2) with
student/fake-score training, EMA support, GAN discriminator branch,
inference pipeline, and export utilities.
* Qwen-Image integration with latent packing and feature-capture for
plugin-enabled pipelines.

* **Documentation**
* New README, example configs, and runnable example scripts for
Qwen-Image distillation and inference.

* **Tests**
* Comprehensive unit tests covering math parity, gradient routing,
plugins, hooks, EMA, and recipe setup.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 16:03:35 -07:00
..
…