Add FastGen quantization-aware distillation

[jingyu-ml]

Summary

• Add a standalone FastGen quantization-aware distillation (QAD) recipe built on the existing ModelOpt kd_loss API, without adding ModelOpt core APIs.
• Load a frozen BF16 teacher through the standard Diffusers checkpoint interface, 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.
• Support additive final-output MSE, optional layerwise MSE, and an independently weighted flow-matching task loss.
• Restore ModelOpt/SVDQuant topology before FSDP, bind the optimizer to post-FSDP live student parameters, and support both all and SVDQuant lora_only train scopes.
• Save and resume complete sharded student DCP checkpoints while keeping the teacher and transient distillation controller out of checkpoint and optimizer state.

Why

FastGen previously exposed the Qwen-Image DMD2 training path but did not provide a conventional quantization-aware distillation recipe. QAD is intentionally separate from DMD2: teacher and student receive the same noisy latent, timestep, and conditioning, and the objective comes from teacher/student MSE plus an optional task loss. Attention Grill is intentionally out of scope.

Stack

This draft is stacked on #2069 and targets its head branch, jingyux/qwen-image-dmd2-svdquant-peft, because it reuses the FastGen Qwen-Image, checkpoint-resume, and generic SVDQuant + PEFT support introduced there.

Validation

• All repository pre-commit hooks passed for the nine added files, including Ruff, mypy, YAML formatting, license checks, Bandit, and Markdown lint.
• Real ModelOpt CPU smoke passed for controller construction, output + layerwise loss aggregation, backward ownership, controller export/capture cleanup, and strict NVFP4 artifact rejection.
• Focused CPU smoke passed for SVDQuant pre_quant_scale registered-buffer validation.
• Both YAML recipes and dotted CLI overrides parsed with the actual AutoModel e42584e3 config parser.
• Independent read-only subagent audit and follow-up re-audit found no blocker, high, or medium severity issue.

No Slurm job was launched. A full Qwen-Image/GPU/FSDP multi-rank smoke was not run from the login node because the available training container is ARM64 while the login node is x86_64.

[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: 8404f08e-288a-4104-a269-88b721adfb93

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.}

[jingyu-ml]

Superseded by #2085, which uses the unprotected jingyux/qwen-image-fastgen-qad branch and preserves the same commit history.