mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-19 05:49:19 +08:00
mxinO
9ceca982a6
### 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 ```bash 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: ```bash 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. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added Parallel Decoding Distillation (PDD) support for Qwen-Image, including student preparation, training, configurable inference blocks, and reusable recipes. * Added prompt-only data loading and support for dataset caches that can be relocated while keeping paths within the cache. * **Documentation** * Added PDD and updated DMD2 Qwen-Image guides. * **Bug Fixes** * Improved checkpoint resumption, dataset-path validation, configuration checks, and Qwen-Image inference compatibility. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Meng Xin <mxin@nvidia.com>