Add FastGen quantization-aware distillation

[jingyu-ml]

Summary

• Add a standalone FastGen quantization-aware distillation (QAD) recipe built on the existing ModelOpt distillation API, without adding ModelOpt core APIs.
• Keep QAD separate from DMD2: a frozen BF16 teacher and quantized student receive the same noisy latent, timestep, and conditioning.
• Load teachers through standard Diffusers checkpoints, including Qwen/Qwen-Image and nvidia/Qwen-Image-Flash.
• Support both regular block-16 NVFP4 W4A4 students and user-prepared NVFP4 SVDQuant + Hugging Face PEFT students.
• Default to final-output MSE with weight 1.0; optionally enable configurable layerwise MSE and flow-matching task loss.
• Restore the complete student bundle before FSDP and support both full-student and SVDQuant lora_only training.
• Add an independent quantize.py --output-bundle path that saves the calibrated pipeline with pipe.save_pretrained(...) for QAD.
• Add Qwen-Image calibration handling for latent output, base-model true CFG, and Flash four-step inference.

Scope

The new bundle path is additive. Existing weight-free checkpoint save/restore, ONNX/TensorRT, DeviceModel, generic Diffusers tests, and Hugging Face deployment export behavior are unchanged.

Attention Grill is intentionally out of scope.

Stack

This draft is stacked on #2069 and targets jingyux/qwen-image-dmd2-svdquant-peft. Commit history is preserved:

1. 1678e3629 — Add FastGen quantization-aware distillation
2. bc47080df — Save quantized Diffusers training bundles
3. 8e21cd43a — Limit QAD changes to training flow

Validation

• Repository pre-commit hooks passed, including Ruff, mypy, license checks, Bandit, and Markdown lint.
• Python syntax and git diff --check passed.
• The final diff was independently audited and reduced from 18 to 12 files.
• The resulting PR has zero diff in diffusion_trt.py, ONNX/TRT scripts and documentation, modelopt/torch/export, generic Diffusers integration tests, and compress handling.
• Real ModelOpt CPU smoke covered QAD controller construction, loss aggregation, backward ownership, artifact validation, and SVDQuant registered-buffer validation.
• Four-GPU Calib32 validation is in progress. Both Qwen-Image-Flash regular NVFP4 and SVDQuant bundles have completed successfully; the two original Qwen-Image calibrations are still running.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8d25a530-419f-46b6-8a0e-f1635c61bc4c

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}