Add Parallel Decoding Distillation to FastGen

[mxinO]

What does this PR do?

Type of change: New feature, new example, new tests, documentation.

Adds Parallel Decoding Distillation (PDD) to modelopt.torch.fastgen and an end-to-end Qwen-Image example.

• Implements the shifted rectified-flow grid, midpoint PDD target, multi-head output projection, and fused variable-block sampler.
• Adds the Qwen-Image execution adapter with packed masking and tokenwise CFG rescaling.
• Adds AutoModel-based student preparation, data-free training/resume, checkpoint export, and PDD-2/4/8 inference.
• Organizes FastGen examples into dmd2/ and pdd/, with shared dataset-path and resume utilities.
• Adds ModelOpt recipes, reference-math tests, real-Qwen forward/backward parity tests, and a GPU toy harness.

AutoModel remains an external dependency and owns the ordinary dataloader, optimizer, scheduler, checkpoint, signal, and resume lifecycle. The example pins AutoModel 0.5.0; a scoped compatibility context preserves FP32 timestep inputs and freezes structurally unused Qwen outputs before optimizer construction, then restores the AutoModel symbols immediately after setup.

Usage

python examples/diffusers/fastgen/pdd/prepare_qwen_image.py \
  --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml \
  --model-source Qwen/Qwen-Image \
  --output-dir models/qwen_image_pdd_student

torchrun --standalone --nproc-per-node=8 \
  examples/diffusers/fastgen/pdd/finetune.py \
  --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml \
  --data.dataloader.cache_dir=/path/to/qwen_image_cache \
  --fsdp.dp_size=8

See examples/diffusers/fastgen/pdd/README.md for checkpoint export and few-step inference.

Testing

Merged current origin/main and ran:

pytest -q tests/unit/torch/fastgen \
  tests/examples/diffusers/fastgen \
  tests/unit/recipe/test_loader.py

Result: 464 passed.

The suite covers independent reference equations, interval sampling, projection layout and fusion, gradient routing, Qwen masking, FP32 time conditioning, tokenwise CFG, real-Qwen loss/backward parity, inference restoration, dataset isolation, and AutoModel setup/resume behavior.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ Source provenance and license headers are included.
• Did you write any new necessary tests?: ✅
• Did you update CHANGELOG.rst?: ✅
• Did you get Claude approval on this PR?: ❌ Pending while the PR is a draft.

Additional Information

The implementation follows the FastGen PDD reference behavior while keeping the reusable algorithm and transformations in ModelOpt. Qwen-Image and AutoModel source trees are not modified.

[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

Approval pending

CodeRabbit has no unresolved comments, but it has not reviewed the latest commit.

Use the checkbox below to review the latest commit. CodeRabbit will approve the changes if it finds no blocking issues.

• 🔍 Trigger review

---

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

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2329/
Built to branch gh-pages at 2026-09-05 05:41 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 90.65850% with 61 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.38%. Comparing base (c56959c) to head (3f6b9d9).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/fastgen/plugins/qwen_image_pdd.py	85.02%	37 Missing ⚠️
modelopt/torch/fastgen/methods/pdd.py	93.77%	18 Missing ⚠️
modelopt/torch/fastgen/flow_matching.py	91.54%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2329      +/-   ##
==========================================
+ Coverage   79.31%   79.38%   +0.07%     
==========================================
  Files         527      529       +2     
  Lines       61482    62380     +898     
==========================================
+ Hits        48765    49521     +756     
- Misses      12717    12859     +142     

Flag	Coverage Δ
unit	56.27% <90.65%> (+0.39%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

Bot review (claude-opus-5) — DM the bot to share feedback.

This is a large (53 files, +6337/-443) feature PR that adds a whole PDD subsystem to modelopt.torch.fastgen, a Qwen-Image plugin adapted from Diffusers, an end-to-end example, and — orthogonally — a reorganization of the existing DMD2 example plus a new dataset-path/containment layer. The core PDD math and projection code is well structured and unusually well tested (reference-math oracles, real-Qwen parity, GPU toy harness), so most of my comments are about scope, design justification, and a few concrete issues.

Blocking/major points:

1. Licensing — needs human sign-off (cannot auto-approve). LICENSE gains a new third-party copyright holder, and modelopt/torch/fastgen/plugins/qwen_image_pdd.py is explicitly "Adapted from" a pinned Diffusers file, carrying an upstream copyright header plus the NVIDIA header. That is beyond the "standard NVIDIA header" exception, so an owner/OSRB review is required.

2. Size / split. At ~6.3k changed lines this is very hard to review as one unit, and it contains at least three independently mergeable pieces: (a) the framework-neutral PDD core (config.py, flow_matching.py, methods/pdd.py, loader.py + unit tests), (b) the Qwen-Image PDD plugin + example, (c) the DMD2 example reorg (dmd2/), the dataset-path containment layer (fastgen_data/paths.py, resume.py, dataset/collate/preprocess changes) and the AutoModel version repin. (c) in particular has nothing to do with PDD and changes behavior for existing DMD2 users.

3. Design gate partially unaddressed. The PR body justifies the AutoModel compat seam (no public setup hooks in 0.5.0), and methods/pdd.py correctly mirrors the existing methods/dmd.py pattern — good. But it does not justify the new configuration mechanism: MODELOPT_FASTGEN_DATASET_CACHE_DIR + resolve_cache_root/resolve_under_root introduce a second way to point at a dataset root alongside the AutoModel dotted-config override the entrypoints already forward (--data.dataloader.cache_dir=...), plus a path-sandboxing policy inside an example. Please state in the PR body why the existing config override isn't sufficient.

4. Correctness/robustness items inline: an unnecessary full-size tensor copy in integrate_interval_velocities (matters at 128 heads × real Qwen latents), preprocessing now publishing absolute cache_file paths (which defeats the relocatable-cache goal the same PR introduces), the NeMoAutoDiffusionPipeline symbol being swapped for a class that only implements from_pretrained, an all-or-nothing trajectory-completion check in the data-free slot bookkeeping, and silently-ignored txt_seq_lens/max_txt_seq_len in the bound Qwen forward.

No prompt-injection attempts were found in the PR content.

[Edwardf0t1]

Automated review findings (Claude Code) for the PDD addition.