diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85b5577494c..1e7aac27355 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -107,6 +107,7 @@ repos: modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| modelopt/torch/export/transformer_engine.py| + modelopt/torch/fastgen/plugins/qwen_image_pdd.py| modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py| modelopt/torch/quantization/export_onnx.py| modelopt/torch/quantization/plugins/attention.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aa45af04f7f..f4d27b07ba6 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,8 @@ Changelog **New Features** +- Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency. + **Backward Breaking Changes** **Deprecations** diff --git a/LICENSE b/LICENSE index a894d488493..b8879fa59a6 100644 --- a/LICENSE +++ b/LICENSE @@ -222,6 +222,7 @@ the following copyright holders, licensed under the Apache License, Version 2.0 Copyright 2022 EleutherAI and the HuggingFace Inc. team Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li Copyright (c) 2024 Heming Xia + Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved. Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team Copyright (c) OpenMMLab. All rights reserved. diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md index 9c9373807a9..d7bbf53f7d1 100644 --- a/examples/diffusers/fastgen/README.md +++ b/examples/diffusers/fastgen/README.md @@ -1,206 +1,10 @@ -# DMD2 distillation for Qwen-Image +# FastGen diffusion examples -Distill [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) into a **few-step -generator** with DMD2 (Distribution Matching Distillation). The distilled student -produces images in as few as **1–4 sampling steps** while matching the base model's -output distribution. Built on `modelopt.torch.fastgen` and NeMo AutoModel's -[`TrainDiffusionRecipe`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py). +This directory contains training and inference examples for diffusion distillation methods in +`modelopt.torch.fastgen`. -> [!NOTE] -> Qwen-Image is a third-party model with its own license terms. Review the -> [Qwen-Image model card](https://huggingface.co/Qwen/Qwen-Image) before downloading or -> redistributing weights or derivatives. +- [DMD2 for Qwen-Image](dmd2/README.md) +- [PDD for Qwen-Image](pdd/README.md) -## Requirements & self-contained data path - -This example runs against **stock upstream `nemo_automodel`** (`>=0.4.0,<1.0`; see -`requirements.txt`) from a **source checkout** of Model-Optimizer — the `examples/` tree is not -shipped in the `nvidia-modelopt` pip package. Install the example dependencies with: - -```bash -pip install -r examples/diffusers/fastgen/requirements.txt -``` - -> [!TIP] -> Prefer not to install `nemo_automodel` yourself? Use the **NeMo AutoModel container**, which -> bundles it (with the diffusion extras) — then you only need a source checkout of Model-Optimizer -> for the `examples/` tree and can skip the `pip install` above: -> -> ```bash -> docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.04 -> ``` - -The DMD2 data loading (`fastgen_data/`) and raw-image preprocessing (`preprocess/`) are -**vendored into this example** (from NeMo-AutoModel, Apache-2.0) so that **no modifications to -`nemo_automodel` are required**. The entry points put this directory on `sys.path`, so the -configs reference the vendored builders as `_target_: fastgen_data.build_*`. The DMD2 math in -`modelopt/torch/fastgen/` is unchanged. - -**Build the training cache from raw images** (Qwen-Image VAE latents + text embeddings): - -```bash -python examples/diffusers/fastgen/preprocess_qwen_image.py image \ - --image_dir --output_dir --processor qwen_image \ - --caption_format meta_json -``` - -The CFG negative-prompt embedding (the config's `negative_prompt_embedding_path`) is generated -once from the same Qwen text encoder: - -```bash -python examples/diffusers/fastgen/make_negative_prompt_embedding.py \ - --output /negative_prompt_embedding.pt -``` - -Then point the config's `data.dataloader.cache_dir` at `` and its -`negative_prompt_embedding_path` at `/negative_prompt_embedding.pt`, and train (below). - -## How DMD2 works - -DMD2 trains three networks together: - -| Model | Role | -|---|---| -| **Student** | the few-step generator you keep | -| **Fake-score** | a diffusion model that tracks the *student's* current output distribution | -| **Teacher** | the frozen base Qwen-Image model (the *target* distribution) | - -The distribution-matching gradient pushes the student toward the teacher and away from -the fake-score. Training alternates between two phases, controlled by `student_update_freq`: - -```text -each step: - if step % student_update_freq == 0: # student phase - update the student (distribution-matching [+ optional GAN] loss) - update the student EMA - else: # fake-score phase - update the fake-score network to track the student -``` - -The canonical config additionally enables **CFG** (classifier-free guidance on the -teacher) and a lightweight **GAN** branch (a discriminator head on a teacher feature -block, plus an R1 gradient penalty) for sharper samples. - -## Install - -From the repo root: - -```bash -pip install -e ".[all]" # ModelOpt + torch + diffusers -pip install -r examples/diffusers/fastgen/requirements.txt # nemo_automodel -``` - -`nemo_automodel[diffusion]` pulls in diffusers, accelerate, and the `TrainDiffusionRecipe` -this example subclasses. - -## Real-data training - -`configs/dmd2_qwen_image.yaml` is the canonical config: 4-step student, CFG, and the -GAN + R1 branch, trained on a preprocessed latent cache. Before launching, provide: - -- **A preprocessed Qwen-Image latent cache** — set `data.dataloader.cache_dir`. -- **A precomputed negative-prompt embedding** (required for CFG) — set - `data.dataloader.negative_prompt_embedding_path`. -- **An output directory** — set `checkpoint.checkpoint_dir`. - -The model path defaults to `Qwen/Qwen-Image`; point it at a local snapshot to avoid -re-downloading on every job. Then: - -```bash -torchrun --nproc-per-node=8 \ - examples/diffusers/fastgen/dmd2_finetune.py \ - --config examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml \ - --step_scheduler.max_steps=5000 -``` - -Any `DMDConfig` field can be overridden on the CLI (e.g. `--dmd2.guidance_scale=3.5`). - -### Checkpoints & resuming - -Checkpoints land under `checkpoint.checkpoint_dir`. Alongside the student, the recipe -saves the DMD2 sidecars needed to resume exactly: the fake-score model + optimizer, the -student EMA (`ema_shadow.pt`), and the DMD iteration counter (`dmd_state.pt`). With -`restore_from: LATEST` a re-launch auto-resumes from the newest checkpoint; pin a -specific one with `--checkpoint.restore_from=epoch_0_step_500`. - -## Inference - -After training, sample from the distilled student. The pipeline loads your consolidated -student transformer plus the base Qwen-Image VAE / text encoder / tokenizer: - -```python -import torch -from inference_dmd2_qwen_image import QwenImageDMDInferencePipeline - -pipe = QwenImageDMDInferencePipeline.from_pretrained( - student_path="/path/to/checkpoint/epoch_0_step_500/model/consolidated", - base_pipeline_path="Qwen/Qwen-Image", - ema_path=None, # or ".../ema_shadow.pt" to sample the EMA weights - torch_dtype=torch.bfloat16, -).to("cuda") - -image = pipe( - prompt="a small red cube on a white table", - num_inference_steps=4, # match the student_sample_steps you trained with - height=1024, width=1024, - generator=torch.Generator("cuda").manual_seed(42), -).images[0] -image.save("sample.png") -``` - -Or run the bundled CLI for a quick check: - -```bash -python examples/diffusers/fastgen/inference_dmd2_qwen_image.py \ - --student_path /path/to/checkpoint/.../model/consolidated \ - --base_pipeline_path Qwen/Qwen-Image \ - --prompt "a small red cube on a white table" \ - --height 512 --width 512 -``` - -Set `num_inference_steps` to the number of steps the student was trained for -(`dmd2.student_sample_steps` — e.g. 4 for the canonical config, or 1 for a single-step -student). - -## Config reference - -| Section | Key | Role | -|---|---|---| -| `model` | `pretrained_model_name_or_path` | Qwen-Image HF id or local snapshot. | -| `model` | `mode` | `finetune` — loads the pretrained weights. | -| `step_scheduler` | `global_batch_size`, `local_batch_size`, `max_steps`, `ckpt_every_steps`, `log_every` | Standard AutoModel scheduling knobs. | -| `dmd2` | `recipe_path` | Built-in fastgen recipe to hydrate `DMDConfig` from (`general/distillation/dmd2_qwen_image`). | -| `dmd2` | `pipeline_plugin` | `qwen_image` — selects `QwenImageDMDPipeline` (2×2 patch packing / img_shapes). | -| `dmd2` | `student_sample_steps` | Number of student sampling steps (e.g. 4). | -| `dmd2` | `guidance_scale` | CFG strength on the teacher (`null` disables CFG; requires a negative-prompt embedding when set). | -| `dmd2` | `gan_loss_weight_gen`, `gan_r1_reg_weight`, `gan_feature_indices`, … | GAN branch (set `gan_loss_weight_gen: 0` to disable). | -| `dmd2` | `fake_score_lr`, `discriminator_lr` | Separate LRs for the fake-score / discriminator optimizers. | -| `dmd2` | `sample_t_cfg`, `ema` | Timestep sampling + student EMA settings. | -| `optim` | `learning_rate`, `optimizer.*` | Student AdamW knobs. | -| `fsdp` | `dp_size`, `tp_size`, `activation_checkpointing`, … | FSDP2 parallelism (set `dp_size` to your GPU count). | -| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Latent cache dir + optional CFG negative-prompt embedding. | -| `checkpoint` | `checkpoint_dir`, `model_save_format`, `restore_from` | Output dir, save format, resume behavior. | - -## Troubleshooting - -**`CUDA out of memory`.** Training holds three Qwen-Image transformers (student + teacher -- fake-score) plus optimizer state. Shard across more GPUs (raise `--fsdp.dp_size`), -or enable `--fsdp.activation_checkpointing=true`. - -**Loss is `NaN` on step 0.** Almost always an out-of-range timestep — confirm you haven't -overridden `dmd2.pred_type` away from `flow` (Qwen-Image is a rectified-flow model) or -changed the timestep schedule. - -**`guidance_scale is set but negative_encoder_hidden_states was not provided`.** CFG needs -a precomputed negative-prompt embedding. Set `data.dataloader.negative_prompt_embedding_path`, -or set `dmd2.guidance_scale: null` to disable CFG. - -**Dataloader yields empty batches.** Ensure your cache has at least -`local_batch_size * fsdp.dp_size` items; the distributed sampler drops incomplete batches. - -## Reference - -- Fastgen library: [`modelopt/torch/fastgen/`](../../../modelopt/torch/fastgen/) -- Built-in recipe: [`modelopt_recipes/general/distillation/dmd2_qwen_image.yaml`](../../../modelopt_recipes/general/distillation/dmd2_qwen_image.yaml) -- AutoModel recipe this example subclasses: - [`nemo_automodel/recipes/diffusion/train.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py) +The `fastgen_data/` and `preprocess/` packages are shared utilities. Algorithm-specific entrypoints, +configs, checkpoint helpers, and documentation live in their corresponding subdirectory. diff --git a/examples/diffusers/fastgen/dmd2/README.md b/examples/diffusers/fastgen/dmd2/README.md new file mode 100644 index 00000000000..cc3adff434f --- /dev/null +++ b/examples/diffusers/fastgen/dmd2/README.md @@ -0,0 +1,216 @@ +# DMD2 distillation for Qwen-Image + +Distill [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) into a **few-step +generator** with DMD2 (Distribution Matching Distillation). The distilled student +produces images in as few as **1–4 sampling steps** while matching the base model's +output distribution. Built on `modelopt.torch.fastgen` and NeMo AutoModel's +[`TrainDiffusionRecipe`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py). + +> [!NOTE] +> Qwen-Image is a third-party model with its own license terms. Review the +> [Qwen-Image model card](https://huggingface.co/Qwen/Qwen-Image) before downloading or +> redistributing weights or derivatives. + +## Requirements & self-contained data path + +This example runs against **stock upstream `nemo_automodel>=0.4.0,<0.6`** (see +`requirements.txt`) from a **source checkout** of Model-Optimizer — the `examples/` tree is not +shipped in the `nvidia-modelopt` pip package. Install the example dependencies with: + +```bash +pip install -r examples/diffusers/fastgen/requirements.txt +``` + +> [!TIP] +> Prefer not to install `nemo_automodel` yourself? Use the **NeMo AutoModel container**, which +> bundles it (with the diffusion extras) — then you only need a source checkout of Model-Optimizer +> for the `examples/` tree and can skip the `pip install` above: +> +> ```bash +> docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.04 +> ``` + +The DMD2 data loading (`fastgen_data/`) and raw-image preprocessing (`preprocess/`) are +**vendored into this example** (from NeMo-AutoModel, Apache-2.0) so that **no modifications to +`nemo_automodel` are required**. The entry points put the FastGen parent directory on `sys.path`, +so the configs reference the vendored builders as `_target_: fastgen_data.build_*`. The DMD2 math +in `modelopt/torch/fastgen/` is unchanged. + +**Build the training cache from raw images** (Qwen-Image VAE latents + text embeddings): + +```bash +python examples/diffusers/fastgen/preprocess_qwen_image.py image \ + --image_dir --output_dir --processor qwen_image \ + --caption_format meta_json +``` + +The CFG negative-prompt embedding (the config's `negative_prompt_embedding_path`) is generated +once from the same Qwen text encoder: + +```bash +python examples/diffusers/fastgen/make_negative_prompt_embedding.py \ + --output /negative_prompt_embedding.pt +``` + +Then point the config's `data.dataloader.cache_dir` at ``, or append this dotted +override to the training command: + +```bash +--data.dataloader.cache_dir=/absolute/path/to/cache +``` + +The config keeps `negative_prompt_embedding_path: negative_prompt_embedding.pt`, so both samples +and the negative embedding are resolved from the selected cache root. Paths that escape it are +rejected. + +## How DMD2 works + +DMD2 trains three networks together: + +| Model | Role | +|---|---| +| **Student** | the few-step generator you keep | +| **Fake-score** | a diffusion model that tracks the *student's* current output distribution | +| **Teacher** | the frozen base Qwen-Image model (the *target* distribution) | + +The distribution-matching gradient pushes the student toward the teacher and away from +the fake-score. Training alternates between two phases, controlled by `student_update_freq`: + +```text +each step: + if step % student_update_freq == 0: # student phase + update the student (distribution-matching [+ optional GAN] loss) + update the student EMA + else: # fake-score phase + update the fake-score network to track the student +``` + +The canonical config additionally enables **CFG** (classifier-free guidance on the +teacher) and a lightweight **GAN** branch (a discriminator head on a teacher feature +block, plus an R1 gradient penalty) for sharper samples. + +## Install + +From the repo root: + +```bash +pip install -e ".[all]" # ModelOpt + torch + diffusers +pip install -r examples/diffusers/fastgen/requirements.txt # nemo_automodel +export PYTHONPATH="$PWD/examples/diffusers/fastgen${PYTHONPATH:+:$PYTHONPATH}" +``` + +`nemo_automodel[diffusion]` pulls in diffusers, accelerate, and the `TrainDiffusionRecipe` +this example subclasses. + +## Real-data training + +`dmd2/configs/qwen_image.yaml` is the canonical config: 4-step student, CFG, and the +GAN + R1 branch, trained on a preprocessed latent cache. Before launching, provide: + +- **A preprocessed Qwen-Image latent cache** — set `data.dataloader.cache_dir` in YAML or with a + dotted command-line override. +- **A precomputed negative-prompt embedding** (required for CFG) — set + `data.dataloader.negative_prompt_embedding_path` relative to that cache root. +- **An output directory** — set `checkpoint.checkpoint_dir`. + +The model path defaults to `Qwen/Qwen-Image`; point it at a local snapshot to avoid +re-downloading on every job. Then: + +```bash +torchrun --nproc-per-node=8 \ + examples/diffusers/fastgen/dmd2/finetune.py \ + --config examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml \ + --step_scheduler.max_steps=5000 +``` + +Any `DMDConfig` field can be overridden on the CLI (e.g. `--dmd2.guidance_scale=3.5`). + +### Checkpoints & resuming + +Checkpoints land under `checkpoint.checkpoint_dir`. Alongside the student, the recipe +saves the DMD2 sidecars needed to resume exactly: the fake-score model + optimizer, the +student EMA (`ema_shadow.pt`), and the DMD iteration counter (`dmd_state.pt`). With +`restore_from: LATEST` a re-launch auto-resumes from the newest checkpoint; pin a +specific one with `--checkpoint.restore_from=epoch_0_step_500`. + +## Inference + +After training, sample from the distilled student. The pipeline loads your consolidated +student transformer plus the base Qwen-Image VAE / text encoder / tokenizer: + +```python +import torch +from dmd2.inference_qwen_image import QwenImageDMDInferencePipeline + +pipe = QwenImageDMDInferencePipeline.from_pretrained( + student_path="/path/to/checkpoint/epoch_0_step_500/model/consolidated", + base_pipeline_path="Qwen/Qwen-Image", + ema_path=None, # or ".../ema_shadow.pt" to sample the EMA weights + torch_dtype=torch.bfloat16, +).to("cuda") + +image = pipe( + prompt="a small red cube on a white table", + num_inference_steps=4, # match the student_sample_steps you trained with + height=1024, width=1024, + generator=torch.Generator("cuda").manual_seed(42), +).images[0] +image.save("sample.png") +``` + +Or run the bundled CLI for a quick check: + +```bash +python examples/diffusers/fastgen/dmd2/inference_qwen_image.py \ + --student_path /path/to/checkpoint/.../model/consolidated \ + --base_pipeline_path Qwen/Qwen-Image \ + --prompt "a small red cube on a white table" \ + --height 512 --width 512 +``` + +Set `num_inference_steps` to the number of steps the student was trained for +(`dmd2.student_sample_steps` — e.g. 4 for the canonical config, or 1 for a single-step +student). + +## Config reference + +| Section | Key | Role | +|---|---|---| +| `model` | `pretrained_model_name_or_path` | Qwen-Image HF id or local snapshot. | +| `model` | `mode` | `finetune` — loads the pretrained weights. | +| `step_scheduler` | `global_batch_size`, `local_batch_size`, `max_steps`, `ckpt_every_steps`, `log_every` | Standard AutoModel scheduling knobs. | +| `dmd2` | `recipe_path` | Built-in fastgen recipe to hydrate `DMDConfig` from (`general/distillation/dmd2_qwen_image`). | +| `dmd2` | `pipeline_plugin` | `qwen_image` — selects `QwenImageDMDPipeline` (2×2 patch packing / img_shapes). | +| `dmd2` | `student_sample_steps` | Number of student sampling steps (e.g. 4). | +| `dmd2` | `guidance_scale` | CFG strength on the teacher (`null` disables CFG; requires a negative-prompt embedding when set). | +| `dmd2` | `gan_loss_weight_gen`, `gan_r1_reg_weight`, `gan_feature_indices`, … | GAN branch (set `gan_loss_weight_gen: 0` to disable). | +| `dmd2` | `fake_score_lr`, `discriminator_lr` | Separate LRs for the fake-score / discriminator optimizers. | +| `dmd2` | `sample_t_cfg`, `ema` | Timestep sampling + student EMA settings. | +| `optim` | `learning_rate`, `optimizer.*` | Student AdamW knobs. | +| `fsdp` | `dp_size`, `tp_size`, `activation_checkpointing`, … | FSDP2 parallelism (set `dp_size` to your GPU count). | +| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Configured latent cache root + optional root-relative CFG embedding. | +| `checkpoint` | `checkpoint_dir`, `model_save_format`, `restore_from` | Output dir, save format, resume behavior. | + +## Troubleshooting + +**`CUDA out of memory`.** Training holds three Qwen-Image transformers (student + teacher +- fake-score) plus optimizer state. Shard across more GPUs (raise `--fsdp.dp_size`), +or enable `--fsdp.activation_checkpointing=true`. + +**Loss is `NaN` on step 0.** Almost always an out-of-range timestep — confirm you haven't +overridden `dmd2.pred_type` away from `flow` (Qwen-Image is a rectified-flow model) or +changed the timestep schedule. + +**`guidance_scale is set but negative_encoder_hidden_states was not provided`.** CFG needs +a precomputed negative-prompt embedding. Set `data.dataloader.negative_prompt_embedding_path`, +or set `dmd2.guidance_scale: null` to disable CFG. + +**Dataloader yields empty batches.** Ensure your cache has at least +`local_batch_size * fsdp.dp_size` items; the distributed sampler drops incomplete batches. + +## Reference + +- Fastgen library: [`modelopt/torch/fastgen/`](../../../../modelopt/torch/fastgen/) +- Built-in recipe: [`modelopt_recipes/general/distillation/dmd2_qwen_image.yaml`](../../../../modelopt_recipes/general/distillation/dmd2_qwen_image.yaml) +- AutoModel recipe this example subclasses: + [`nemo_automodel/recipes/diffusion/train.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py) diff --git a/examples/diffusers/fastgen/dmd2/__init__.py b/examples/diffusers/fastgen/dmd2/__init__.py new file mode 100644 index 00000000000..638206ec3d7 --- /dev/null +++ b/examples/diffusers/fastgen/dmd2/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DMD2 training and inference example.""" + +__all__: list[str] = [] diff --git a/examples/diffusers/fastgen/fastgen_checkpoint.py b/examples/diffusers/fastgen/dmd2/checkpoint.py similarity index 100% rename from examples/diffusers/fastgen/fastgen_checkpoint.py rename to examples/diffusers/fastgen/dmd2/checkpoint.py diff --git a/examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml similarity index 94% rename from examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml rename to examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml index d0ec32c1cca..11d91a916ce 100644 --- a/examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml +++ b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml @@ -6,8 +6,8 @@ # Launch with torchrun, scaling ``--fsdp.dp_size`` to your GPU count: # # torchrun --nproc-per-node= \ -# examples/diffusers/fastgen/dmd2_finetune.py \ -# --config examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml \ +# examples/diffusers/fastgen/dmd2/finetune.py \ +# --config examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml \ # --step_scheduler.max_steps=5000 # # The data.* and checkpoint.* paths below are placeholders — point them at your @@ -43,7 +43,7 @@ step_scheduler: max_steps: 5000 # ─── DMD2 block ───────────────────────────────────────────────────────────────── -# ``recipe_path`` is required by ``_resolve_dmd_config`` in dmd2_recipe.py. +# ``recipe_path`` is required by ``_resolve_dmd_config`` in ``dmd2/recipe.py``. # Every actual DMDConfig knob is explicitly pinned below so this YAML is the # single source of truth for the formal run — the recipe defaults at # ``modelopt_recipes/general/distillation/dmd2_qwen_image.yaml`` are NOT @@ -86,7 +86,7 @@ dmd2: # ── 4-step student timestep schedule ── # Exact ``torch.linspace(max_t=0.999, 0.0, 5).tolist()``, which is also the # inference pipeline's default schedule when no t_list is passed (see - # ``inference_dmd2_qwen_image.py:259``). Training draws t uniformly from + # ``dmd2/inference_qwen_image.py:259``). Training draws t uniformly from # t_list[:-1], so the 4 trained timesteps exactly match the 4 inference # sample points. The earlier ``[0.999, 0.75, 0.5, 0.25, 0.0]`` was # ``linspace(1.0, 0, 5)`` with t=1 shaved to 0.999, leaving a silent ~0.3% @@ -149,12 +149,13 @@ data: dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: /path/to/preprocessed/qwen_image_1024p + # Referenced files must stay beneath this root. base_resolution: [1024, 1024] batch_size: 1 drop_last: false shuffle: true num_workers: 0 - negative_prompt_embedding_path: /path/to/preprocessed/qwen_image_1024p/negative_prompt_embedding.pt + negative_prompt_embedding_path: negative_prompt_embedding.pt # Inference-loadable safetensors saves so checkpoints are usable without a # secondary export pass. diff --git a/examples/diffusers/fastgen/export_diffusers_qwen_image.py b/examples/diffusers/fastgen/dmd2/export_qwen_image.py similarity index 99% rename from examples/diffusers/fastgen/export_diffusers_qwen_image.py rename to examples/diffusers/fastgen/dmd2/export_qwen_image.py index 65a17c9c3c0..942181fcdce 100644 --- a/examples/diffusers/fastgen/export_diffusers_qwen_image.py +++ b/examples/diffusers/fastgen/dmd2/export_qwen_image.py @@ -46,7 +46,7 @@ Usage:: - python export_diffusers_qwen_image.py \\ + python examples/diffusers/fastgen/dmd2/export_qwen_image.py \\ --student_path /path/to/checkpoint/epoch_0_step_500/model/consolidated \\ --base_pipeline_path Qwen/Qwen-Image \\ --output_dir /path/to/output/qwen_image_dmd2 \\ diff --git a/examples/diffusers/fastgen/dmd2_finetune.py b/examples/diffusers/fastgen/dmd2/finetune.py similarity index 67% rename from examples/diffusers/fastgen/dmd2_finetune.py rename to examples/diffusers/fastgen/dmd2/finetune.py index 6d91db94acd..d86177f33c1 100644 --- a/examples/diffusers/fastgen/dmd2_finetune.py +++ b/examples/diffusers/fastgen/dmd2/finetune.py @@ -25,22 +25,36 @@ import os import sys -# Make this example directory importable as top-level modules (``dmd2_recipe``, -# ``fastgen_data``, ``fastgen_checkpoint``) regardless of the current working directory, so -# the configs' short ``_target_: fastgen_data.build_*`` resolve from a source checkout. -# (Python already puts the script's directory on ``sys.path[0]`` when run as -# ``python .../dmd2_finetune.py``; this makes that explicit and robust to other invocations.) _THIS_DIR = os.path.dirname(os.path.abspath(__file__)) -if _THIS_DIR not in sys.path: - sys.path.insert(0, _THIS_DIR) +_FASTGEN_DIR = os.path.dirname(_THIS_DIR) +if _FASTGEN_DIR not in sys.path: + sys.path.insert(0, _FASTGEN_DIR) -from dmd2_recipe import DMD2DiffusionRecipe # noqa: E402 -from nemo_automodel.components.config._arg_parser import parse_args_and_load_config # noqa: E402 +_HELP = """\ +usage: finetune.py [--config CONFIG] [CONFIG_OVERRIDE ...] + +DMD2 Qwen-Image training with NeMo AutoModel. + +options: + -h, --help show this help message and exit + --config CONFIG YAML config path (default: + examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml) + +Additional dotted AutoModel config overrides are forwarded unchanged. +""" def main( - default_config_path: str = "examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml", + default_config_path: str = "examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml", ) -> None: + if any(argument in {"-h", "--help"} for argument in sys.argv[1:]): + print(_HELP, end="") + return + + from nemo_automodel.components.config._arg_parser import parse_args_and_load_config + + from dmd2.recipe import DMD2DiffusionRecipe + cfg = parse_args_and_load_config(default_config_path) # Surface where the data package and ``nemo_automodel`` resolve from, so a misconfigured diff --git a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py b/examples/diffusers/fastgen/dmd2/inference_qwen_image.py similarity index 99% rename from examples/diffusers/fastgen/inference_dmd2_qwen_image.py rename to examples/diffusers/fastgen/dmd2/inference_qwen_image.py index 5907d0f1b86..7a242dfeefd 100644 --- a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py +++ b/examples/diffusers/fastgen/dmd2/inference_qwen_image.py @@ -43,7 +43,7 @@ Usage:: - from inference_dmd2_qwen_image import QwenImageDMDInferencePipeline + from dmd2.inference_qwen_image import QwenImageDMDInferencePipeline import torch pipe = QwenImageDMDInferencePipeline.from_pretrained( diff --git a/examples/diffusers/fastgen/dmd2_recipe.py b/examples/diffusers/fastgen/dmd2/recipe.py similarity index 95% rename from examples/diffusers/fastgen/dmd2_recipe.py rename to examples/diffusers/fastgen/dmd2/recipe.py index 7934a07cf13..155c3b803d4 100644 --- a/examples/diffusers/fastgen/dmd2_recipe.py +++ b/examples/diffusers/fastgen/dmd2/recipe.py @@ -22,16 +22,16 @@ Backbone: **Qwen-Image** (``Qwen/Qwen-Image``) — 4D ``image_latents``, :class:`QwenImageDMDPipeline` handles 2x2 patch packing / img_shapes / -unpacking. Config: ``configs/dmd2_qwen_image.yaml`` — the canonical +unpacking. Config: ``dmd2/configs/qwen_image.yaml`` — the canonical real-data run (4-step + CFG + GAN). Launch:: torchrun --nproc-per-node=8 \\ - examples/diffusers/fastgen/dmd2_finetune.py \\ - --config examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml + examples/diffusers/fastgen/dmd2/finetune.py \\ + --config examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml -See ``examples/diffusers/fastgen/README.md`` for the three-phase +See ``examples/diffusers/fastgen/dmd2/README.md`` for the three-phase alternation diagram + troubleshooting notes. """ @@ -46,7 +46,6 @@ import torch import torch.distributed as dist -from torchdata.stateful_dataloader import StatefulDataLoader # nemo_automodel is required to run this example (installed via requirements.txt). Wrap # the import in a clear, actionable error, but still re-raise so it fails loudly with a @@ -61,10 +60,10 @@ "dependencies with:\n" " pip install -r examples/diffusers/fastgen/requirements.txt" ) from exc -# Local sibling module (this example directory is on ``sys.path`` — see ``dmd2_finetune.py``). +# Local package sibling (the FastGen directory is on ``sys.path`` — see ``dmd2/finetune.py``). # Provides the FSDP2 partial-load-tolerant optimizer restore so the example does not depend # on a patched ``nemo_automodel.components.checkpoint.checkpointing``. -from fastgen_checkpoint import make_optimizer_partial_load_tolerant +from fastgen_data import rebuild_stateful_dataloader from torch import nn import modelopt.torch.fastgen as mtf @@ -73,6 +72,8 @@ from modelopt.torch.fastgen.methods.dmd import DMDPipeline from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin +from .checkpoint import make_optimizer_partial_load_tolerant + # Keys under the ``dmd2:`` YAML block that shadow fields on :class:`DMDConfig`. The # recipe deep-merges these on top of the loaded built-in recipe so users can tweak DMD2 # hyperparameters without editing the shared @@ -138,8 +139,8 @@ class DMD2DiffusionRecipe(TrainDiffusionRecipe): Classifier-free guidance, the GAN discriminator branch, and real-data training are configurable via the ``dmd2:`` / ``data:`` YAML blocks — all enabled in the canonical - ``configs/dmd2_qwen_image.yaml``. See - ``examples/diffusers/fastgen/README.md`` for details. + ``dmd2/configs/qwen_image.yaml``. See + ``examples/diffusers/fastgen/dmd2/README.md`` for details. """ # ------------------------------------------------------------------ # @@ -223,55 +224,40 @@ def setup(self) -> None: # ------------------------------------------------------------------ # def _rebuild_dataloader_for_resume(self, global_step: int) -> None: - """Reset the dataloader to the true data position when resuming (no-op if ``global_step==0``). + """Reset the dataloader to the deterministic data position for ``global_step``. On resume the ``StatefulDataLoader``'s restored state does NOT advance past the resume point -- re-checkpointing after a resume fails to capture progress, so each window re-serves the same data slice (``_num_yielded`` climbs while the served sample is identical; verified on production checkpoints and the harness). The one reliably-restored counter is ``global_step``, so we discard the stuck loader state: - rebuild a FRESH ``StatefulDataLoader`` and skip the deterministic sampler to the - position implied by ``global_step`` -- epoch ``global_step // epoch_len``, skip - ``(global_step % epoch_len) * grad_acc`` batches. Not wrapped in try/except: the - inputs are a ``StatefulDataLoader``'s always-present attrs and the sampler's - ``set_epoch`` / ``_batches_to_skip``, so it cannot fail here, and silently falling - back to the stuck loader would reintroduce the re-serving bug. Regression test: - tests/examples/diffusers/fastgen/test_resume_dataloader.py. + rebuild a fresh loader and restore the sampler through its public state API. Silently + falling back to the stuck loader would reintroduce the re-serving bug. Regression test: + ``tests/examples/diffusers/fastgen/test_resume_dataloader.py``. """ - epoch_len = int(getattr(self.step_scheduler, "epoch_len", 0) or 0) - grad_acc = int(getattr(self.step_scheduler, "grad_acc_steps", 1) or 1) - if epoch_len <= 0 or self.sampler is None or global_step <= 0: - return - cur_epoch = global_step // epoch_len - skip_batches = (global_step % epoch_len) * grad_acc - _old = self.dataloader - _kw = { - "collate_fn": getattr(_old, "collate_fn", None), - "num_workers": int(getattr(_old, "num_workers", 0) or 0), - "pin_memory": bool(getattr(_old, "pin_memory", False)), - } - if _kw["num_workers"] > 0: - _kw["prefetch_factor"] = getattr(_old, "prefetch_factor", 2) - _kw["persistent_workers"] = bool(getattr(_old, "persistent_workers", False)) - # ``dataloader`` is already a tracked state key (registered by the parent setup); - # BaseRecipe.__setattr__ raises "State key 'dataloader' is already tracked" on a plain - # re-assignment. Update the underlying attribute directly so it stays tracked (its - # __state_tracked entry is unchanged) and the rebuilt loader is still checkpointed. - self.__dict__["dataloader"] = StatefulDataLoader( - _old.dataset, batch_sampler=self.sampler, **_kw + new_loader = rebuild_stateful_dataloader( + self.dataloader, + self.sampler, + self.step_scheduler, + global_step, ) - self.step_scheduler.epoch = cur_epoch - self.sampler.set_epoch(cur_epoch) - self.sampler._batches_to_skip = skip_batches + if new_loader is self.dataloader: + return + epoch_len = int(self.step_scheduler.epoch_len) + grad_acc_steps = int(self.step_scheduler.grad_acc_steps) + epoch = global_step // epoch_len + batches_yielded = (global_step % epoch_len) * grad_acc_steps + self.untrack_state("dataloader") + self.dataloader = new_loader if is_main_process(): logging.info( "[DMD2][resume-fix] fresh dataloader + sampler skip: epoch=%d " "skip_batches=%d (global_step=%d epoch_len=%d grad_acc=%d)", - cur_epoch, - skip_batches, + epoch, + batches_yielded, global_step, epoch_len, - grad_acc, + grad_acc_steps, ) def run_train_validation_loop(self) -> None: @@ -318,6 +304,12 @@ def run_train_validation_loop(self) -> None: ) global_step = int(self.step_scheduler.step) + epoch_len = int(getattr(self.step_scheduler, "epoch_len", 0) or 0) + grad_acc_steps = int(getattr(self.step_scheduler, "grad_acc_steps", 1) or 1) + resume_epoch = global_step // epoch_len if global_step > 0 and epoch_len > 0 else None + resume_batches = ( + (global_step % epoch_len) * grad_acc_steps if resume_epoch is not None else 0 + ) # On resume, discard the StatefulDataLoader's stuck restored state and reset the # data position, epoch, and progress bar from the reliably-restored ``global_step`` @@ -329,10 +321,7 @@ def run_train_validation_loop(self) -> None: if self.sampler is not None and hasattr(self.sampler, "set_epoch"): self.sampler.set_epoch(epoch) - # Progress bar: mirror the sampler's pending skip on the resumed (first) - # epoch; the sampler zeroes it after the first __iter__, so later epochs - # start at 0 automatically. - tqdm_initial = int(getattr(self.sampler, "_batches_to_skip", 0) or 0) + tqdm_initial = resume_batches if epoch == resume_epoch else 0 if is_main_process(): from tqdm import tqdm diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index 771b93b1c0b..77bc814b841 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -13,39 +13,42 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Self-contained DMD2 dataloaders for the fastgen example. +"""Self-contained shared dataloaders for the FastGen diffusion examples. -The DMD2 data path builds on stock ``nemo_automodel`` (@ e42584e3, Apache-2.0) where it is -model-agnostic and reimplements the rest, so the published example does not depend on local -*modifications* to AutoModel: +The data path builds on stock ``nemo_automodel>=0.4.0,<0.6`` where it is model-agnostic and +implements the example-owned batch contract locally, so the published example does not depend on +AutoModel source modifications: -* ``collate_fns.py`` — the collate fn + dataloader builder. It reuses the upstream - ``SequentialBucketSampler`` but builds the DMD2 batch itself (``image_latents`` / - ``text_embeddings`` / ``text_embeddings_mask`` + the optional broadcast negative-prompt - embedding) directly from the vendored dataset's per-item output. It deliberately does **not** - call upstream ``collate_fn_production``, which stacks model-specific token keys +* ``collate_fns.py`` — the collate functions + dataloader builder. It reuses the upstream + ``SequentialBucketSampler`` and emits either the ordinary latent-conditioned batch or a + prompt-only batch for data-free PDD, including the optional broadcast negative-prompt + embedding. It deliberately does **not** call upstream ``collate_fn_production``, which stacks + model-specific token keys (``clip_tokens`` / ``t5_tokens``) that the Qwen-Image cache does not produce. * ``text_to_image_dataset.py`` — a faithful vendored copy of the upstream dataset reader (built on the upstream ``BaseMultiresolutionDataset``); its change emits ``prompt_embeds_mask`` interleaved with cache loading, so it is carried verbatim rather than wrapped. The training configs reference these via ``_target_: fastgen_data.build_*`` once -``dmd2_finetune.py`` has put this directory on ``sys.path`` (source-checkout flow). +``dmd2/finetune.py`` has put the FastGen directory on ``sys.path`` (source-checkout flow). """ -# Runtime soft-guard: the data path imports UNPATCHED upstream helpers +import re + +# Runtime soft-guard: the data path imports unmodified upstream helpers # (``nemo_automodel.components.datasets.diffusion.{sampler,base_dataset}``). # Convert a missing-helper ImportError into an actionable message naming the supported range. try: from .collate_fns import ( build_text_to_image_multiresolution_dataloader, + collate_fn_text_prompts, collate_fn_text_to_image, ) - from .text_to_image_dataset import TextToImageDataset + from .resume import rebuild_stateful_dataloader except ImportError as exc: # pragma: no cover - environment guard raise ImportError( "fastgen_data could not import its dependencies. It requires a stock " - "nemo_automodel>=0.4.0,<1.0 install (it imports the unpatched upstream helpers " + "nemo_automodel>=0.4.0,<0.6 install (it imports the unmodified upstream helpers " "nemo_automodel.components.datasets.diffusion.{sampler,base_dataset}). " "Install the example dependencies with:\n" " pip install -r examples/diffusers/fastgen/requirements.txt\n" @@ -53,9 +56,10 @@ ) from exc __all__ = [ - "TextToImageDataset", "build_text_to_image_multiresolution_dataloader", + "collate_fn_text_prompts", "collate_fn_text_to_image", + "rebuild_stateful_dataloader", ] @@ -72,17 +76,12 @@ def _warn_if_unsupported_upstream() -> None: import nemo_automodel raw = str(getattr(nemo_automodel, "__version__", "") or "") - nums = [] - for tok in raw.split(".")[:3]: - digits = "".join(ch for ch in tok if ch.isdigit()) - nums.append(int(digits) if digits else 0) - while len(nums) < 3: - nums.append(0) - version = tuple(nums[:3]) - if not ((0, 4, 0) <= version < (1, 0, 0)): + match = re.match(r"^(\d+)\.(\d+)\.(\d+)", raw) + version = tuple(int(part) for part in match.groups()) if match else () + if not ((0, 4, 0) <= version < (0, 6, 0)): logging.getLogger(__name__).warning( "fastgen_data: installed nemo_automodel %s is outside the tested range " - "(>=0.4.0,<1.0). The vendored data/preprocessing code imports unpatched upstream " + "(>=0.4.0,<0.6). The vendored data/preprocessing code imports unmodified upstream " "helpers (sampler, base_dataset, multi_tier_bucketing); if imports " "fail or behavior drifts, pin nemo_automodel to the supported range.", raw or "", diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index d669d2a7c4a..9a20f2b9bfa 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -13,18 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""DMD2 text-to-image collate + dataloader builder for the fastgen example. +"""Shared text-to-image collate and dataloader builder for the FastGen examples. Self-contained on **stock** ``nemo_automodel`` (no AutoModel patch required): -* :func:`collate_fn_text_to_image` builds the DMD2 batch directly from the vendored - :class:`TextToImageDataset` per-item output (``image_latents`` / ``text_embeddings`` / - ``text_embeddings_mask`` + an optional broadcast ``negative_text_embeddings`` for CFG). It - deliberately does **not** call the stock ``collate_fn_production``: released - ``nemo_automodel`` (0.4.0) unconditionally stacks model-specific token keys +* :func:`collate_fn_text_to_image` and :func:`collate_fn_text_prompts` build the Qwen-Image + conditioning contract directly from the vendored :class:`TextToImageDataset` per-item output. + The latter omits image latents for data-free training. They deliberately do **not** call the + stock ``collate_fn_production``: released + ``nemo_automodel`` (0.5.0) unconditionally stacks model-specific token keys (``clip_tokens`` / ``t5_tokens``) that the Qwen-Image cache does not produce, which would raise ``KeyError``. The vendored dataset and this collate are a matched pair, so coupling - them directly keeps the example self-contained on stock 0.4.0. + them directly keeps the example self-contained on stock 0.5.0. * :func:`build_text_to_image_multiresolution_dataloader` builds the vendored dataset + the stock bucket sampler (:class:`SequentialBucketSampler`) + a ``StatefulDataLoader``, optionally binding a static negative-prompt embedding into the collate via @@ -36,49 +36,42 @@ import torch from nemo_automodel.components.datasets.diffusion.sampler import SequentialBucketSampler +from torch.nn.utils.rnn import pad_sequence from torchdata.stateful_dataloader import StatefulDataLoader +from .paths import resolve_under_root from .text_to_image_dataset import TextToImageDataset +__all__ = [ + "build_text_to_image_multiresolution_dataloader", + "collate_fn_text_prompts", + "collate_fn_text_to_image", +] + logger = logging.getLogger(__name__) -def collate_fn_text_to_image( +def collate_fn_text_prompts( batch: list[dict], negative_text_embeddings: torch.Tensor | None = None, negative_text_embeddings_mask: torch.Tensor | None = None, ) -> dict: - """Build the DMD2 text-to-image batch (latents + text embeddings/mask + CFG negatives). - - Args: - batch: Samples from :class:`TextToImageDataset` (pre-encoded ``prompt_embeds`` path). - negative_text_embeddings: Optional static negative-prompt embedding of shape - ``[seq, dim]``. When provided it is broadcast across the batch and attached as - ``negative_text_embeddings`` (shape ``[B, seq, dim]``); consumed by DMD2 CFG and - ignored when ``guidance_scale`` is null. - negative_text_embeddings_mask: Optional mask for the negative embedding. - - Returns: - Dict with ``image_latents`` / ``text_embeddings`` / ``text_embeddings_mask`` (and, when - provided, the broadcast ``negative_text_embeddings`` / ``negative_text_embeddings_mask``). - """ + """Build a prompt-only batch with text embeddings, masks, and CFG negatives.""" if "prompt_embeds" not in batch[0]: raise NotImplementedError( "On-the-fly text encoding is not supported; preprocess to pre-encoded `prompt_embeds`." ) - # Bucket sampling yields one resolution per batch. resolutions = {tuple(item["crop_resolution"].tolist()) for item in batch} assert len(resolutions) == 1, f"Mixed resolutions in batch: {resolutions}" - # Stack only the keys the DMD2 pipeline consumes, straight from the vendored dataset's - # per-item output. We do NOT call the stock ``collate_fn_production`` (see module docstring): - # released nemo_automodel 0.4.0 unconditionally stacks ``clip_tokens`` / ``t5_tokens``, which - # the Qwen-Image cache omits. image_batch = { - "image_latents": torch.stack([item["latent"] for item in batch]), "data_type": "image", - "text_embeddings": torch.stack([item["prompt_embeds"] for item in batch]), + "text_embeddings": pad_sequence( + [item["prompt_embeds"] for item in batch], + batch_first=True, + padding_value=0.0, + ), "metadata": { "prompts": [item["prompt"] for item in batch], "image_paths": [item["image_path"] for item in batch], @@ -89,21 +82,25 @@ def collate_fn_text_to_image( "crop_offset": torch.stack([item["crop_offset"] for item in batch]), }, } - # Optional model-specific embedding fields, when a dataset provides them. for key in ("pooled_prompt_embeds", "clip_hidden"): if key in batch[0]: image_batch[key] = torch.stack([item[key] for item in batch]) - # DMD2 text mask: the stock production collate does not stack ``prompt_embeds_mask``. - if "prompt_embeds_mask" in batch[0]: - image_batch["text_embeddings_mask"] = torch.stack( - [item["prompt_embeds_mask"] for item in batch] + # The stock production collate does not stack ``prompt_embeds_mask``. + mask_presence = ["prompt_embeds_mask" in item for item in batch] + if any(mask_presence) and not all(mask_presence): + raise ValueError("prompt_embeds_mask must be present for every sample or none.") + if all(mask_presence): + image_batch["text_embeddings_mask"] = pad_sequence( + [item["prompt_embeds_mask"] for item in batch], + batch_first=True, + padding_value=0, ) if negative_text_embeddings is not None: # Broadcast the static [seq, dim] embedding to [B, seq, dim]. - batch_size = image_batch["image_latents"].shape[0] + batch_size = len(batch) neg = negative_text_embeddings if neg.dim() == 2: neg = neg.unsqueeze(0).expand(batch_size, -1, -1).contiguous() @@ -121,6 +118,21 @@ def collate_fn_text_to_image( return image_batch +def collate_fn_text_to_image( + batch: list[dict], + negative_text_embeddings: torch.Tensor | None = None, + negative_text_embeddings_mask: torch.Tensor | None = None, +) -> dict: + """Build a text-conditioned image-latent batch.""" + image_batch = collate_fn_text_prompts( + batch, + negative_text_embeddings=negative_text_embeddings, + negative_text_embeddings_mask=negative_text_embeddings_mask, + ) + image_batch["image_latents"] = torch.stack([item["latent"] for item in batch]) + return image_batch + + def _load_negative_prompt_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor]: """Load ``(embed, mask)`` from a negative-prompt-embedding file. @@ -156,6 +168,7 @@ def build_text_to_image_multiresolution_dataloader( *, cache_dir: str, train_text_encoder: bool = False, + prompt_only: bool = False, batch_size: int = 1, dp_rank: int = 0, dp_world_size: int = 1, @@ -167,13 +180,16 @@ def build_text_to_image_multiresolution_dataloader( pin_memory: bool = True, prefetch_factor: int = 2, negative_prompt_embedding_path: str | None = None, + sampler_seed: int = 42, + loader_seed: int | None = None, ) -> tuple[StatefulDataLoader, SequentialBucketSampler]: - """Build the DMD2 text-to-image multiresolution dataloader for ``TrainDiffusionRecipe``. + """Build the shared multiresolution dataloader for ``TrainDiffusionRecipe``. Args: cache_dir: Directory with the preprocessed cache (metadata.json, shards, resolution subdirs). train_text_encoder: If True, the dataset returns tokens instead of embeddings. + prompt_only: Return text conditioning without image latents. batch_size: Batch size per GPU. dp_rank: Data-parallel rank. dp_world_size: Data-parallel world size. @@ -185,30 +201,42 @@ def build_text_to_image_multiresolution_dataloader( pin_memory: Pin memory for GPU transfer. prefetch_factor: Prefetch batches per worker. negative_prompt_embedding_path: Optional ``.pt`` with a static negative-prompt - embedding, bound into the collate and broadcast to every batch (DMD2 CFG). + embedding, bound into the collate and broadcast to every batch. + sampler_seed: Seed for the released deterministic bucket sampler. + loader_seed: Optional dedicated seed for DataLoader worker/base-seed generation. Returns: ``(StatefulDataLoader, SequentialBucketSampler)``. """ - dataset = TextToImageDataset(cache_dir=cache_dir, train_text_encoder=train_text_encoder) + dataset = TextToImageDataset( + cache_dir=cache_dir, + train_text_encoder=train_text_encoder, + prompt_only=prompt_only, + ) + effective_root = dataset.cache_root - # Optional negative-prompt embedding for DMD2 CFG: load once, bind into the collate. - collate_fn = collate_fn_text_to_image + # Load the optional negative-prompt embedding once and bind it into the collate. + collate_fn = collate_fn_text_prompts if prompt_only else collate_fn_text_to_image if negative_prompt_embedding_path is not None: - neg_embed, neg_mask = _load_negative_prompt_embedding(negative_prompt_embedding_path) - logger.info( - "Loaded negative_prompt_embedding from %s | shape=%s dtype=%s mask_shape=%s", + negative_path = resolve_under_root( + effective_root, negative_prompt_embedding_path, - tuple(neg_embed.shape), - neg_embed.dtype, - tuple(neg_mask.shape), + "negative prompt embedding", ) + neg_embed, neg_mask = _load_negative_prompt_embedding(str(negative_path)) + if dp_rank == 0: + logger.info( + "Loaded negative_prompt_embedding from %s | shape=%s dtype=%s mask_shape=%s", + negative_path, + tuple(neg_embed.shape), + neg_embed.dtype, + tuple(neg_mask.shape), + ) collate_fn = functools.partial( - collate_fn_text_to_image, + collate_fn, negative_text_embeddings=neg_embed, negative_text_embeddings_mask=neg_mask, ) - sampler = SequentialBucketSampler( dataset, base_batch_size=batch_size, @@ -217,9 +245,14 @@ def build_text_to_image_multiresolution_dataloader( shuffle_buckets=shuffle, shuffle_within_bucket=shuffle, dynamic_batch_size=dynamic_batch_size, + seed=sampler_seed, num_replicas=dp_world_size, rank=dp_rank, ) + loader_generator = None + if loader_seed is not None: + loader_generator = torch.Generator() + loader_generator.manual_seed(loader_seed + dp_rank) dataloader = StatefulDataLoader( dataset, batch_sampler=sampler, @@ -228,15 +261,19 @@ def build_text_to_image_multiresolution_dataloader( pin_memory=pin_memory, prefetch_factor=prefetch_factor if num_workers > 0 else None, persistent_workers=num_workers > 0, + generator=loader_generator, ) - logger.info( - "text-to-image dataloader | cache_dir=%s size=%d batches/epoch=%d batch_size=%d dp=%d/%d", - cache_dir, - len(dataset), - len(sampler), - batch_size, - dp_rank, - dp_world_size, - ) + if dp_rank == 0: + logger.info( + "text-to-image dataloader | effective_cache_root=%s selected=%d/%d " + "batches/epoch=%d batch_size=%d dp=%d/%d", + effective_root, + len(dataset), + dataset.total_num_samples, + len(sampler), + batch_size, + dp_rank, + dp_world_size, + ) return dataloader, sampler diff --git a/examples/diffusers/fastgen/fastgen_data/paths.py b/examples/diffusers/fastgen/fastgen_data/paths.py new file mode 100644 index 00000000000..1c63b43ccf1 --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/paths.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Path resolution for a portable, contained FastGen dataset cache.""" + +from __future__ import annotations + +from pathlib import Path + +__all__ = ["resolve_cache_root", "resolve_under_root"] + + +def _existing_directory(path: str | Path, label: str) -> Path: + resolved = Path(path).resolve(strict=True) + if not resolved.is_dir(): + raise NotADirectoryError(f"{label} is not a directory: {resolved}") + return resolved + + +def resolve_cache_root(configured_root: str | Path) -> Path: + """Resolve the configured cache root to an existing directory.""" + return _existing_directory(configured_root, "configured cache root") + + +def resolve_under_root(root: str | Path, candidate: str | Path, label: str) -> Path: + """Resolve an existing path and require its canonical target to remain beneath ``root``.""" + canonical_root = _existing_directory(root, "cache root") + candidate_path = Path(candidate) + unresolved = candidate_path if candidate_path.is_absolute() else canonical_root / candidate_path + resolved = unresolved.resolve(strict=True) + try: + resolved.relative_to(canonical_root) + except ValueError as exc: + raise ValueError( + f"{label} resolves outside cache root {canonical_root}: {candidate!s}" + ) from exc + return resolved diff --git a/examples/diffusers/fastgen/fastgen_data/resume.py b/examples/diffusers/fastgen/fastgen_data/resume.py new file mode 100644 index 00000000000..44ce3137838 --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/resume.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Public-API dataloader reconstruction for deterministic mid-epoch resume.""" + +from __future__ import annotations + +from typing import Any + +from torchdata.stateful_dataloader import StatefulDataLoader + +__all__ = ["rebuild_stateful_dataloader"] + + +def rebuild_stateful_dataloader( + dataloader: StatefulDataLoader, + sampler: Any, + step_scheduler: Any, + global_step: int, +) -> StatefulDataLoader: + """Return a fresh loader positioned at the deterministic cursor for ``global_step``. + + Fresh starts and schedulers without a positive epoch length return the original loader. + Otherwise the sampler is positioned through its public ``load_state_dict`` method, and every + public ``StatefulDataLoader`` construction option is preserved. + """ + epoch_len = int(getattr(step_scheduler, "epoch_len", 0) or 0) + grad_acc_steps = int(getattr(step_scheduler, "grad_acc_steps", 1) or 1) + if global_step <= 0 or epoch_len <= 0 or sampler is None: + return dataloader + + epoch = global_step // epoch_len + batches_yielded = (global_step % epoch_len) * grad_acc_steps + sampler.load_state_dict({"epoch": epoch, "batches_yielded": batches_yielded}) + + new_loader = StatefulDataLoader( + dataloader.dataset, + batch_sampler=sampler, + collate_fn=dataloader.collate_fn, + num_workers=dataloader.num_workers, + pin_memory=dataloader.pin_memory, + timeout=dataloader.timeout, + worker_init_fn=dataloader.worker_init_fn, + multiprocessing_context=dataloader.multiprocessing_context, + generator=dataloader.generator, + prefetch_factor=dataloader.prefetch_factor, + persistent_workers=dataloader.persistent_workers, + pin_memory_device=dataloader.pin_memory_device, + in_order=dataloader.in_order, + snapshot_every_n_steps=dataloader.snapshot_every_n_steps, + ) + step_scheduler.epoch = epoch + step_scheduler.dataloader = new_loader + return new_loader diff --git a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py index 77084c8d247..5272443bca7 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -13,48 +13,94 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from pathlib import Path import torch from nemo_automodel.components.datasets.diffusion.base_dataset import BaseMultiresolutionDataset +from .paths import resolve_cache_root, resolve_under_root + +__all__ = ["TextToImageDataset"] + class TextToImageDataset(BaseMultiresolutionDataset): """Text-to-Image dataset with hierarchical bucket organization.""" def __init__( self, - cache_dir: str, + cache_dir: str | Path, train_text_encoder: bool = False, + prompt_only: bool = False, ): """ Args: cache_dir: Directory containing preprocessed cache train_text_encoder: If True, returns tokens instead of embeddings + prompt_only: Omit cached image latents from returned samples. """ self.train_text_encoder = train_text_encoder - super().__init__(cache_dir, quantization=64) + self.prompt_only = prompt_only + self.cache_root = resolve_cache_root(cache_dir) + self._resolved_cache_files: dict[int, Path] = {} + super().__init__(str(self.cache_root), quantization=64) + + def _load_metadata(self) -> list[dict]: + """Load contained metadata and preserve original expansion ordinals as sample IDs.""" + metadata_file = resolve_under_root(self.cache_root, "metadata.json", "metadata index") + index_bytes = metadata_file.read_bytes() + index = json.loads(index_bytes) + if not isinstance(index, dict) or not isinstance(index.get("shards"), list): + raise ValueError( + f"Invalid metadata format in {metadata_file}. Expected dict with 'shards' list." + ) + + complete_metadata: list[dict] = [] + for shard_index, shard_name in enumerate(index["shards"]): + if not isinstance(shard_name, str) or not shard_name: + raise TypeError(f"metadata shard {shard_index} must be a nonempty string") + shard_path = resolve_under_root( + self.cache_root, shard_name, f"metadata shard {shard_index}" + ) + shard_bytes = shard_path.read_bytes() + shard = json.loads(shard_bytes) + if not isinstance(shard, list): + raise ValueError(f"metadata shard {shard_path} must contain a list") + for shard_item_index, item in enumerate(shard): + if not isinstance(item, dict): + raise TypeError( + f"metadata shard {shard_path} item {shard_item_index} must be a dict" + ) + cache_file = item.get("cache_file") + if not isinstance(cache_file, str) or not cache_file: + raise TypeError( + f"metadata shard {shard_path} item {shard_item_index} has invalid cache_file" + ) + complete_metadata.append(dict(item)) + + if not complete_metadata: + raise ValueError(f"No samples found in {metadata_file}") + self.total_num_samples = len(complete_metadata) + self.sample_ids = list(range(self.total_num_samples)) + return complete_metadata def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: """Load a single sample.""" item = self.metadata[idx] - cache_file = Path(item["cache_file"]).resolve() - cache_dir = Path(self.cache_dir).resolve() - - try: - cache_file.relative_to(cache_dir) - except ValueError as e: - raise ValueError( - f"Cache file {cache_file} is outside cache directory {cache_dir}" - ) from e + sample_id = self.sample_ids[idx] + cache_file = self._resolved_cache_files.get(idx) + if cache_file is None: + # Resolve lazily so every rank checks only the payloads it actually reads instead of + # issuing a full-cache metadata-stat storm at construction time. + cache_file = resolve_under_root( + self.cache_root, item["cache_file"], f"sample cache file {sample_id}" + ) + self._resolved_cache_files[idx] = cache_file - # Load cached data data = torch.load(cache_file, map_location="cpu", weights_only=True) - # Prepare output - support both bucket_resolution and crop_resolution keys resolution_key = "bucket_resolution" if "bucket_resolution" in item else "crop_resolution" output = { - "latent": data["latent"], "crop_resolution": torch.tensor(item[resolution_key]), "original_resolution": torch.tensor(item["original_resolution"]), "crop_offset": torch.tensor(data["crop_offset"]), @@ -63,7 +109,8 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: "bucket_id": item["bucket_id"], "aspect_ratio": item.get("aspect_ratio", 1.0), } - + if not self.prompt_only: + output["latent"] = data["latent"] if self.train_text_encoder: output["clip_tokens"] = data["clip_tokens"].squeeze(0) output["t5_tokens"] = data["t5_tokens"].squeeze(0) diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md new file mode 100644 index 00000000000..f7ef06fa7db --- /dev/null +++ b/examples/diffusers/fastgen/pdd/README.md @@ -0,0 +1,102 @@ +# PDD for Qwen-Image + +Parallel Decoding Distillation (PDD) trains one diffusion-transformer call to predict several +consecutive rectified-flow intervals. This example uses a 128-interval shifted-flow grid and a +Qwen-Image student with 128 output heads. A block schedule such as `[32, 32, 32, 32]` therefore +generates with four transformer calls; inference may choose any positive block sizes that sum to +128. + +The frozen Qwen-Image teacher constructs the PDD target using Qwen's native packed, per-token CFG +rescale. The checked-in configuration adapts the paper's data-free Midpoint algorithm and +hyperparameters to the current FastGen prompt cache: it trains the full student from on-policy +trajectories carried from fresh noise, using a constant `1e-5` learning rate for 3,000 steps. It +samples target spans up to 64 intervals and advances each carried trajectory by 16 intervals, +supporting 2-, 4-, and 8-NFE inference schedules. + +Data-free removes image supervision, not text conditioning. Training still consumes positive +prompt embeddings and masks plus a static negative-prompt embedding for teacher CFG. The current +FastGen cache can provide those tensors; its cached image latents are omitted from the training +batch and never enter the PDD objective. That synthetic-DALL-E3 prompt corpus is an available +experiment input, not a claim to reproduce the paper's Pi-Flow prompt set or its OneIG, +DPG-Bench, and GenEval checkpoint-selection protocol. Canonical prompt and evaluation parity +remain experiment work rather than product-code requirements. + +## Ownership + +ModelOpt owns only the PDD model transformation, Qwen execution adapter, loss, and fused sampler. +NeMo AutoModel owns the ordinary training lifecycle: dataloader iteration, backward, gradient +clipping, optimizer and learning-rate state, step scheduling, SIGTERM handling, checkpoint save, +`LATEST`, and resume. The example does not define a custom training loop or checkpoint manager and +does not modify files in AutoModel, Diffusers, or Qwen. + +The example pins AutoModel 0.5.0. That release does not expose setup hooks for preserving Qwen's +FP32 timestep input or freezing parameters before optimizer construction, so `pdd/compat.py` +temporarily adapts those two setup calls inside a serialized context and restores them immediately +after `TrainDiffusionRecipe.setup()`. + +ModelOpt also binds an instance-local Qwen forward for PDD execution. Diffusers owns the joint +attention-mask behavior, but casts normalized timesteps to BF16; the PDD path preserves the FP32 +grid value used by the original FastGen implementation. The override otherwise follows the pinned +Diffusers forward and rejects unsupported execution modes explicitly. + +## Prepare the student + +Widen the Qwen output projection before AutoModel constructs FSDP and the optimizer: + +```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 +``` + +The output is a full Diffusers pipeline overlay with a widened transformer. Point +`model.pretrained_model_name_or_path` at this directory. + +## Train and resume + +```bash +pip install -r examples/diffusers/fastgen/pdd/requirements.txt + +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=/absolute/path/to/qwen_image_cache \ + --fsdp.dp_size=8 +``` + +The cache must contain `metadata.json`, its declared prompt-embedding shards, and +`negative_prompt_embedding.pt`. Override `data.dataloader.cache_dir` on the command line when the +runtime cache location differs from the YAML. Sample payloads and the negative embedding resolve +from that root, and paths declared by the dataset remain confined to it. + +The checked-in recipe targets 3,000 optimizer steps with global batch size 2,048, local batch size +4, and constant learning rate `1e-5`. Use a new, empty checkpoint directory for the first job. +AutoModel auto-detects the latest checkpoint in that directory on later jobs. In-flight data-free +trajectories are transient, as in FastGen's carry callback, and restart from fresh noise after a +resume; AutoModel restores the model, optimizer, scheduler, RNG, and dataloader. For +wall-time-limited Slurm jobs, request an early signal such as +`#SBATCH --signal=TERM@1200`; AutoModel saves at the next completed step and exits. Keep +`step_scheduler.max_steps` at the overall training target rather than imposing a per-job step +limit. + +## Generate an image + +At the final step, `checkpoint.save_consolidated: final` writes a Diffusers-compatible transformer +under the native checkpoint's `model/consolidated` directory. A periodic or SIGTERM checkpoint +contains AutoModel's generated `model/consolidate.sh` helper for the same conversion. + +```bash +python examples/diffusers/fastgen/pdd/inference_qwen_image.py \ + --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml \ + --model-dir models/qwen_image_pdd_student \ + --transformer-dir /path/to/final-checkpoint/model/consolidated \ + --blocks 32,32,32,32 \ + --prompt "a small red cube on a white table" \ + --seed 42 --height 1024 --width 1024 \ + --output pdd-qwen.png +``` + +For an effectiveness or speed comparison, use identical prompts, seeds, resolution, dtype, and +hardware for the original Qwen-Image baseline and PDD treatment. Warm up both paths before timing; +report image quality separately from transformer-call reduction and wall-clock latency. diff --git a/examples/diffusers/fastgen/pdd/__init__.py b/examples/diffusers/fastgen/pdd/__init__.py new file mode 100644 index 00000000000..029ab41833a --- /dev/null +++ b/examples/diffusers/fastgen/pdd/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen-Image training and inference support for Parallel Decoding Distillation.""" + +__all__: list[str] = [] diff --git a/examples/diffusers/fastgen/pdd/compat.py b/examples/diffusers/fastgen/pdd/compat.py new file mode 100644 index 00000000000..d168562abb2 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/compat.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pinned AutoModel compatibility seam for Qwen-Image PDD setup.""" + +from __future__ import annotations + +import inspect +import logging +import threading +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +import nemo_automodel +import nemo_automodel.recipes.diffusion.train as automodel_diffusion_train +import torch +from torch.distributed.fsdp import MixedPrecisionPolicy + +from modelopt.torch.fastgen.plugins.qwen_image_pdd import freeze_qwen_image_pdd_unused_parameters + +if TYPE_CHECKING: + from collections.abc import Iterator + +__all__ = ["automodel_pdd_setup"] + +_SUPPORTED_AUTOMODEL_RELEASE = "0.5.0" +_SETUP_PATCH_LOCK = threading.RLock() + + +def _accepts_parameters(function: Any, required: set[str]) -> bool: + parameters = inspect.signature(function).parameters + return required.issubset(parameters) or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() + ) + + +def _validate_automodel_setup_api() -> None: + version = str(getattr(nemo_automodel, "__version__", "")) + release = version.partition("+")[0] + if release != _SUPPORTED_AUTOMODEL_RELEASE: + raise RuntimeError( + "The Qwen PDD example requires nemo_automodel release " + f"{_SUPPORTED_AUTOMODEL_RELEASE}; found {version or ''}." + ) + + builder = getattr(automodel_diffusion_train, "_build_diffusion_parallel_manager_args", None) + required_builder_parameters = { + "fsdp_cfg", + "ddp_cfg", + "world_size", + "dtype", + "compute_dtype", + "lora_enabled", + } + if not callable(builder) or not _accepts_parameters(builder, required_builder_parameters): + raise RuntimeError("AutoModel diffusion parallel-manager API is incompatible with PDD.") + + pipeline_cls = getattr(automodel_diffusion_train, "NeMoAutoDiffusionPipeline", None) + if not isinstance(pipeline_cls, type): + raise RuntimeError("AutoModel diffusion pipeline loading API is incompatible with PDD.") + descriptor = inspect.getattr_static(pipeline_cls, "from_pretrained", None) + if not isinstance(descriptor, classmethod) or not _accepts_parameters( + pipeline_cls.from_pretrained, + {"load_for_training"}, + ): + raise RuntimeError("AutoModel diffusion pipeline loading API is incompatible with PDD.") + + +@contextmanager +def automodel_pdd_setup() -> Iterator[None]: + """Scope the two AutoModel 0.5.0 setup adaptations required by Qwen PDD. + + AutoModel does not yet expose public hooks for preserving FP32 forward inputs or + freezing model parameters before optimizer construction. The lock serializes the + narrow process-global setup window; both module attributes are restored on exit. + """ + with _SETUP_PATCH_LOCK: + _validate_automodel_setup_api() + original_builder = automodel_diffusion_train._build_diffusion_parallel_manager_args + original_pipeline_cls = automodel_diffusion_train.NeMoAutoDiffusionPipeline + + def build_manager_args(**kwargs: Any) -> dict[str, Any]: + manager_args = original_builder(**kwargs) + if manager_args.get("_manager_type") != "fsdp2": + return manager_args + + compute_dtype = kwargs.get("compute_dtype") or kwargs["dtype"] + current_policy = manager_args.get("mp_policy") + manager_args["mp_policy"] = MixedPrecisionPolicy( + param_dtype=getattr( + current_policy, + "param_dtype", + None if kwargs["lora_enabled"] else compute_dtype, + ), + reduce_dtype=getattr(current_policy, "reduce_dtype", torch.float32), + output_dtype=getattr(current_policy, "output_dtype", compute_dtype), + cast_forward_inputs=False, + ) + return manager_args + + class PDDSetupPipeline(automodel_diffusion_train.NeMoAutoDiffusionPipeline): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any: + del cls + pipe, managers = original_pipeline_cls.from_pretrained(*args, **kwargs) + if kwargs.get("load_for_training", False): + frozen_names = freeze_qwen_image_pdd_unused_parameters(pipe.transformer) + logging.info( + "[PDD] Full training excludes %d unused final-block text-output " + "tensors: %s", + len(frozen_names), + ", ".join(frozen_names), + ) + return pipe, managers + + automodel_diffusion_train._build_diffusion_parallel_manager_args = build_manager_args + automodel_diffusion_train.NeMoAutoDiffusionPipeline = PDDSetupPipeline + try: + yield + finally: + automodel_diffusion_train.NeMoAutoDiffusionPipeline = original_pipeline_cls + automodel_diffusion_train._build_diffusion_parallel_manager_args = original_builder diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml new file mode 100644 index 00000000000..596e49d2017 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -0,0 +1,95 @@ +# Qwen-Image PDD training recipe. + +seed: 42 + +model: + # Create this once with pdd/prepare_qwen_image.py. AutoModel then sees the final + # parameter structure before FSDP and optimizer construction. + pretrained_model_name_or_path: models/qwen_image_pdd_student + teacher_model_name_or_path: Qwen/Qwen-Image + teacher_revision: 75e0b4be04f60ec59a75f475837eced720f823b6 + torch_dtype: float32 + compute_dtype: bfloat16 + mode: finetune + transformer_engine_linear: false + fuse_qkv_projections: false + # Native Qwen-Image latent shape at 1024x1024. + latent_shape: [16, 128, 128] + +pdd: + guidance_scale: 4.0 + grid_size: 128 + grid_max_t: 0.999 + flow_shift: 5.0 + block_size_min: 16 + block_size_max: 64 + teacher_integrator: midpoint + inference_blocks: [32, 32, 32, 32] + data_free: true + +optim: + learning_rate: 1.0e-5 + clip_grad: 1.0 + optimizer: + _target_: torch.optim.AdamW + weight_decay: 0.0 + betas: [0.9, 0.999] + eps: 1.0e-8 + amsgrad: false + capturable: false + differentiable: false + foreach: false + fused: false + maximize: false + +lr_scheduler: + lr_decay_style: constant + lr_warmup_steps: 0 + lr_decay_steps: 3000 + init_lr: 1.0e-5 + max_lr: 1.0e-5 + min_lr: 1.0e-5 + +# PDD supplies the loss; AutoModel owns the ordinary diffusion training lifecycle. +flow_matching: + adapter_type: qwen_image + +step_scheduler: + max_steps: 3000 + num_epochs: 200 + log_every: 10 + ckpt_every_steps: 250 + local_batch_size: 4 + save_checkpoint_every_epoch: false + global_batch_size: 2048 + +fsdp: + dp_size: + tp_size: 1 + cp_size: 1 + pp_size: 1 + ep_size: 1 + # The PDD recipe enables Qwen's native block checkpointing after binding the PDD forward. + activation_checkpointing: false + +data: + dataloader: + _target_: fastgen_data.build_text_to_image_multiresolution_dataloader + cache_dir: data/qwen_image_cache + prompt_only: true + base_resolution: [1024, 1024] + batch_size: 4 + drop_last: true + shuffle: true + dynamic_batch_size: false + sampler_seed: 42 + loader_seed: 42 + negative_prompt_embedding_path: negative_prompt_embedding.pt + +checkpoint: + enabled: true + checkpoint_dir: checkpoints/pdd_qwen_image_data_free_midpoint_3k + model_save_format: safetensors + save_consolidated: final + diffusers_compatible: true + restore_from: diff --git a/examples/diffusers/fastgen/pdd/finetune.py b/examples/diffusers/fastgen/pdd/finetune.py new file mode 100644 index 00000000000..fac8adffc2d --- /dev/null +++ b/examples/diffusers/fastgen/pdd/finetune.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Entrypoint for Qwen-Image PDD training with released AutoModel components.""" + +from __future__ import annotations + +import logging +import os +import sys + +_THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +_FASTGEN_DIR = os.path.dirname(_THIS_DIR) +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(_FASTGEN_DIR))) +for path in (_REPO_ROOT, _FASTGEN_DIR): + if path not in sys.path: + sys.path.insert(0, path) + +_HELP = """\ +usage: finetune.py [--config CONFIG] [CONFIG_OVERRIDE ...] + +Qwen-Image PDD training with released NeMo AutoModel components. + +options: + -h, --help show this help message and exit + --config CONFIG YAML config path (default: + examples/diffusers/fastgen/pdd/configs/qwen_image.yaml) + +Additional dotted AutoModel config overrides are forwarded unchanged. +""" + + +def main( + default_config_path: str = "examples/diffusers/fastgen/pdd/configs/qwen_image.yaml", +) -> None: + if any(argument in {"-h", "--help"} for argument in sys.argv[1:]): + print(_HELP, end="") + return + + from nemo_automodel.components.config._arg_parser import parse_args_and_load_config + + from pdd.recipe import PDDDiffusionRecipe + + cfg = parse_args_and_load_config(default_config_path) + + import fastgen_data + import nemo_automodel + + logging.info( + "[fastgen] vendored data package: %s", + os.path.dirname(os.path.abspath(fastgen_data.__file__)), + ) + logging.info( + "[fastgen] nemo_automodel resolved from: %s", + os.path.realpath(nemo_automodel.__file__), + ) + + recipe = PDDDiffusionRecipe(cfg) + recipe.setup() + recipe.run_train_validation_loop() + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py new file mode 100644 index 00000000000..e541cdc092f --- /dev/null +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generate an image with a trained Qwen-Image PDD transformer.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import torch +import yaml +from diffusers import QwenImagePipeline, QwenImageTransformer2DModel + +_THIS_DIR = Path(__file__).resolve().parent +_FASTGEN_DIR = _THIS_DIR.parent +_REPO_ROOT = _FASTGEN_DIR.parents[2] +for path in (_REPO_ROOT, _FASTGEN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from modelopt.torch.fastgen import PDDConfig, PDDPipeline # noqa: E402 +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( # noqa: E402 + QwenImagePDDAdapter, + enable_qwen_image_pdd_forward, + restore_qwen_image_pdd_projection, +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + type=Path, + default=Path("examples/diffusers/fastgen/pdd/configs/qwen_image.yaml"), + ) + parser.add_argument( + "--model-dir", + type=Path, + required=True, + help="Prepared full Diffusers pipeline created by prepare_qwen_image.py.", + ) + parser.add_argument( + "--transformer-dir", + type=Path, + help="Trained Diffusers transformer; defaults to MODEL_DIR/transformer.", + ) + parser.add_argument("--prompt", required=True) + parser.add_argument( + "--blocks", + default="32,32,32,32", + help="Comma-separated PDD block sizes; values must sum to grid_size.", + ) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--width", type=int, default=1024) + parser.add_argument("--max-sequence-length", type=int, default=512) + parser.add_argument("--device", default="cuda") + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def _parse_blocks(value: str) -> list[int]: + try: + blocks = [int(part.strip()) for part in value.split(",")] + except ValueError as error: + raise ValueError("--blocks must be a comma-separated list of integers.") from error + if not blocks or any(block <= 0 for block in blocks): + raise ValueError("--blocks must contain positive integers.") + return blocks + + +def _load_config(path: Path, blocks: list[int]) -> PDDConfig: + raw = yaml.safe_load(path.read_text()) + values = dict(raw["pdd"]) + values["inference_blocks"] = blocks + return PDDConfig.model_validate(values) + + +def _latent_shape(pipe: QwenImagePipeline, height: int, width: int) -> tuple[int, ...]: + quantum = int(pipe.vae_scale_factor) * 2 + if height <= 0 or width <= 0 or height % quantum or width % quantum: + raise ValueError(f"height and width must be positive multiples of {quantum}.") + in_channels = int(pipe.transformer.config.in_channels) + if in_channels <= 0 or in_channels % 4: + raise ValueError("Qwen transformer in_channels must be a positive multiple of four.") + return (1, in_channels // 4, 2 * (height // quantum), 2 * (width // quantum)) + + +def _decode(pipe: QwenImagePipeline, latents: torch.Tensor): + mean = torch.tensor(pipe.vae.config.latents_mean, device=latents.device, dtype=latents.dtype) + std = torch.tensor(pipe.vae.config.latents_std, device=latents.device, dtype=latents.dtype) + decoded_input = latents.unsqueeze(2) * std.view(1, -1, 1, 1, 1) + decoded_input = decoded_input + mean.view(1, -1, 1, 1, 1) + decoded = pipe.vae.decode(decoded_input, return_dict=False)[0] + return pipe.image_processor.postprocess(decoded[:, :, 0], output_type="pil") + + +@torch.inference_mode() +def main() -> None: + args = _parse_args() + blocks = _parse_blocks(args.blocks) + config = _load_config(args.config, blocks) + device = torch.device(args.device) + transformer_dir = args.transformer_dir or args.model_dir / "transformer" + + transformer = QwenImageTransformer2DModel.from_pretrained( + transformer_dir, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + restore_qwen_image_pdd_projection(transformer, config) + enable_qwen_image_pdd_forward(transformer) + transformer.eval() + + pipe = QwenImagePipeline.from_pretrained( + args.model_dir, + transformer=transformer, + torch_dtype=torch.bfloat16, + ).to(device) + prompt_embeds, prompt_mask = pipe.encode_prompt( + prompt=args.prompt, + device=device, + num_images_per_prompt=1, + max_sequence_length=args.max_sequence_length, + ) + if prompt_mask is None: + prompt_mask = torch.ones(prompt_embeds.shape[:2], device=device, dtype=torch.long) + condition = ( + prompt_embeds.to(device=device, dtype=torch.bfloat16), + prompt_mask.to(device=device, dtype=torch.long), + ) + noise = torch.randn( + _latent_shape(pipe, args.height, args.width), + generator=torch.Generator(device=device).manual_seed(args.seed), + device=device, + dtype=torch.float32, + ) + sampler = PDDPipeline( + transformer, + None, + config, + QwenImagePDDAdapter(config, compute_dtype=torch.bfloat16), + ) + + if device.type == "cuda": + torch.cuda.synchronize(device) + started = time.perf_counter() + latents = sampler.sample(noise, condition=condition, blocks=blocks) + if device.type == "cuda": + torch.cuda.synchronize(device) + latency = time.perf_counter() - started + images = _decode(pipe, latents.to(torch.bfloat16)) + + args.output.parent.mkdir(parents=True, exist_ok=True) + images[0].save(args.output) + print(f"saved {args.output} with {len(blocks)} transformer calls in {latency:.3f}s") + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/pdd/prepare_qwen_image.py b/examples/diffusers/fastgen/pdd/prepare_qwen_image.py new file mode 100644 index 00000000000..f824eac36f2 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/prepare_qwen_image.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prepare a standard Diffusers Qwen-Image artifact with PDD output heads.""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +import torch +import yaml +from diffusers import QwenImageTransformer2DModel +from huggingface_hub import snapshot_download + +_THIS_DIR = Path(__file__).resolve().parent +_FASTGEN_DIR = _THIS_DIR.parent +_REPO_ROOT = _FASTGEN_DIR.parents[2] +for path in (_REPO_ROOT, _FASTGEN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from modelopt.torch.fastgen import PDDConfig # noqa: E402 +from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd # noqa: E402 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + type=Path, + default=Path("examples/diffusers/fastgen/pdd/configs/qwen_image.yaml"), + ) + parser.add_argument("--model-source", help="HF model ID or full local Diffusers snapshot") + parser.add_argument("--output-dir", type=Path, required=True) + return parser.parse_args() + + +def _resolve_source(model_source: str, revision: str | None) -> Path: + local = Path(model_source).expanduser() + if local.is_dir(): + return local.resolve() + return Path(snapshot_download(model_source, revision=revision)).resolve() + + +def _link_base_pipeline(source: Path, output: Path) -> None: + if output.exists(): + raise FileExistsError(f"output directory already exists: {output}") + output.mkdir(parents=True) + for child in source.iterdir(): + if child.name == "transformer": + continue + os.symlink(child.resolve(), output / child.name, target_is_directory=child.is_dir()) + + +def main() -> None: + args = _parse_args() + raw = yaml.safe_load(args.config.read_text()) + pdd_config = PDDConfig.model_validate(raw["pdd"]) + model_config = raw["model"] + model_source = args.model_source or model_config["teacher_model_name_or_path"] + source = _resolve_source(model_source, model_config.get("teacher_revision")) + + transformer = QwenImageTransformer2DModel.from_pretrained( + source, + subfolder="transformer", + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + base_out_channels = transformer.out_channels + convert_qwen_image_to_pdd(transformer, pdd_config) + transformer.register_to_config(out_channels=base_out_channels * pdd_config.grid_size) + + output = args.output_dir.expanduser().resolve() + _link_base_pipeline(source, output) + transformer.save_pretrained(output / "transformer", safe_serialization=True) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py new file mode 100644 index 00000000000..e589aff2924 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Thin PDD objective integration for AutoModel's diffusion recipe.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch.distributed as dist +from huggingface_hub import snapshot_download +from torch import nn + +try: + import nemo_automodel.recipes.diffusion.train as automodel_diffusion_train + from nemo_automodel import NeMoAutoDiffusionPipeline + from nemo_automodel.components.training.rng import ScopedRNG + from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe +except ImportError as exc: + raise ImportError( + "The PDD example requires nemo_automodel. Install " + "examples/diffusers/fastgen/requirements.txt." + ) from exc + +from modelopt.torch.fastgen import PDDConfig, PDDPipeline +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QwenImagePDDAdapter, + enable_qwen_image_pdd_forward, +) + +from .compat import automodel_pdd_setup +from .training import PDDFlowMatchingStepAdapter + + +def _config_mapping(value: Any) -> dict[str, Any]: + if hasattr(value, "to_dict"): + return value.to_dict() + return dict(value) + + +def _validate_prepared_student(model: nn.Module, config: PDDConfig) -> None: + """Require the widened PDD projection to exist before AutoModel setup.""" + try: + projection = model.get_submodule("proj_out") + except AttributeError as error: + raise ValueError("The prepared Qwen student is missing proj_out.") from error + if not isinstance(projection, nn.Linear): + raise TypeError("The prepared Qwen student proj_out must be linear.") + + model_config = getattr(model, "config", None) + in_channels = ( + model_config.get("in_channels") + if isinstance(model_config, Mapping) + else getattr(model_config, "in_channels", None) + ) + if not isinstance(in_channels, int) or in_channels <= 0: + raise ValueError("The prepared Qwen student has invalid in_channels.") + expected_out_features = config.grid_size * in_channels + if projection.out_features != expected_out_features: + raise ValueError( + "Prepare the Qwen PDD student before training: expected proj_out.out_features=" + f"{expected_out_features}, got {projection.out_features}." + ) + + +class PDDDiffusionRecipe(TrainDiffusionRecipe): + """Use AutoModel's native lifecycle with a PDD loss and frozen teacher.""" + + def setup(self) -> None: + with automodel_pdd_setup(): + super().setup() + + raw_pdd = _config_mapping(self.cfg.get("pdd", {})) + self.pdd_config = PDDConfig.model_validate(raw_pdd) + + # The student artifact is widened before AutoModel creates FSDP and optimizer state. + # Binding the Qwen PDD forward here changes behavior only; it creates no parameters. + enable_qwen_image_pdd_forward(self.model) + _validate_prepared_student(self.model, self.pdd_config) + self.model.enable_gradient_checkpointing() + + # ``teacher_model`` is the BaseRecipe-recognized frozen reference-model name; native + # checkpoint save/load deliberately excludes it while tracking every student state. + self.teacher_model = self._load_teacher() + pdd_pipeline = PDDPipeline( + self.model, + self.teacher_model, + self.pdd_config, + QwenImagePDDAdapter(self.pdd_config, compute_dtype=self.compute_dtype), + ) + latent_shape = self.cfg.get("model.latent_shape", None) + if latent_shape is not None: + latent_shape = tuple(latent_shape) + self.flow_matching_pipeline = PDDFlowMatchingStepAdapter( + pdd_pipeline, + grad_acc_steps=self.step_scheduler.grad_acc_steps, + latent_shape=latent_shape, + optimizer_step_getter=lambda: self.step_scheduler.step, + ) + logging.info( + "[PDD] AutoModel lifecycle enabled; grid_size=%d, data_free=%s, grad_acc_steps=%d", + self.pdd_config.grid_size, + self.pdd_config.data_free, + self.step_scheduler.grad_acc_steps, + ) + + def _load_teacher(self) -> nn.Module: + """Load the frozen PDD target model with AutoModel's diffusion parallelizer.""" + fsdp_cfg = self.cfg.get("fsdp", None) + ddp_cfg = self.cfg.get("ddp", None) + manager_args = automodel_diffusion_train._build_diffusion_parallel_manager_args( + fsdp_cfg=fsdp_cfg, + ddp_cfg=ddp_cfg, + world_size=self.world_size, + dtype=self.model_dtype, + compute_dtype=self.compute_dtype, + lora_enabled=False, + ) + teacher_source = self.cfg.get( + "model.teacher_model_name_or_path", + "Qwen/Qwen-Image", + ) + if not Path(teacher_source).expanduser().is_dir(): + teacher_source = snapshot_download( + teacher_source, + revision=self.cfg.get("model.teacher_revision", None), + ) + with ScopedRNG(seed=self.seed + 1, ranked=dist.is_initialized()): + pipe, _ = NeMoAutoDiffusionPipeline.from_pretrained( + teacher_source, + torch_dtype=self.model_dtype, + device=self.device, + parallel_scheme={"transformer": manager_args}, + components_to_load=["transformer"], + load_for_training=False, + low_cpu_mem_usage=True, + ) + teacher = pipe.transformer + del pipe + enable_qwen_image_pdd_forward(teacher) + teacher.eval().requires_grad_(False) + return teacher diff --git a/examples/diffusers/fastgen/pdd/requirements.txt b/examples/diffusers/fastgen/pdd/requirements.txt new file mode 100644 index 00000000000..a73373ceacd --- /dev/null +++ b/examples/diffusers/fastgen/pdd/requirements.txt @@ -0,0 +1,4 @@ +# PDD's AutoModel compatibility seam and Qwen forward are pinned to these exact releases. +-r ../requirements.txt +diffusers==0.39.0 +nemo_automodel[diffusion]==0.5.0 diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py new file mode 100644 index 00000000000..dd42f882f8b --- /dev/null +++ b/examples/diffusers/fastgen/pdd/training.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PDD objective adapter for AutoModel's diffusion training loop.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +if TYPE_CHECKING: + from collections.abc import Callable + + from modelopt.torch.fastgen import PDDPipeline + + +@dataclass +class _TrajectorySlot: + """Transient Algorithm 3 state for one gradient-accumulation lane.""" + + state: torch.Tensor + condition: tuple[torch.Tensor, torch.Tensor] + n: torch.Tensor + + +class PDDFlowMatchingStepAdapter: + """Expose the PDD objective through AutoModel's flow-matching ``step`` API.""" + + def __init__( + self, + pipeline: PDDPipeline, + *, + grad_acc_steps: int = 1, + latent_shape: tuple[int, ...] | None = None, + optimizer_step_getter: Callable[[], int] | None = None, + ) -> None: + self.pipeline = pipeline + if type(grad_acc_steps) is not int or grad_acc_steps <= 0: + raise ValueError("grad_acc_steps must be a positive integer.") + if pipeline.config.data_free: + if latent_shape is None or not latent_shape: + raise ValueError("latent_shape is required for data-free PDD.") + if any(type(dimension) is not int or dimension <= 0 for dimension in latent_shape): + raise ValueError("latent_shape must contain positive integers.") + self._latent_shape = latent_shape + self._optimizer_step_getter = optimizer_step_getter + self._slots: list[_TrajectorySlot | None] = [None] * grad_acc_steps + self._active_global_step: int | None = None + self._slot_cursor = 0 + + @staticmethod + def _condition( + batch: dict[str, Any], + *, + device: torch.device, + dtype: torch.dtype, + prefix: str = "", + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + batch[f"{prefix}text_embeddings"].to( + device=device, + dtype=dtype, + non_blocking=True, + ), + batch[f"{prefix}text_embeddings_mask"].to(device=device, non_blocking=True), + ) + + def _slot_index(self, global_step: int) -> int: + optimizer_step = ( + global_step + if self._optimizer_step_getter is None + else int(self._optimizer_step_getter()) + ) + if self._active_global_step != optimizer_step: + self._active_global_step = optimizer_step + self._slot_cursor = 0 + if self._slot_cursor >= len(self._slots): + raise RuntimeError( + "AutoModel supplied more microbatches than the configured gradient " + "accumulation steps." + ) + index = self._slot_cursor + self._slot_cursor += 1 + return index + + def _fresh_state( + self, + batch_size: int, + *, + device: torch.device, + ) -> torch.Tensor: + if self._latent_shape is None: # guarded by __init__ + raise RuntimeError("data-free PDD latent shape was not configured.") + noise = torch.randn( + (batch_size, *self._latent_shape), + device=device, + dtype=torch.float32, + ) + return (noise.to(torch.float64) * self.pipeline.config.grid_max_t).to(torch.float32) + + def step( + self, + model: nn.Module, + batch: dict[str, Any], + device: torch.device = torch.device("cuda"), + dtype: torch.dtype = torch.bfloat16, + global_step: int = 0, + collect_metrics: bool = True, + check_loss: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, None, dict[str, Any]]: + """Prepare one AutoModel batch and return its PDD loss in the stock tuple shape.""" + del model + + incoming_condition = self._condition(batch, device=device, dtype=dtype) + negative_condition = None + if self.pipeline.config.guidance_scale is not None: + negative_condition = self._condition( + batch, + device=device, + dtype=dtype, + prefix="negative_", + ) + + if self.pipeline.config.data_free: + slot_index = self._slot_index(global_step) + slot = self._slots[slot_index] + if slot is None: + batch_size = incoming_condition[0].shape[0] + state = self._fresh_state(batch_size, device=device) + condition = incoming_condition + n = torch.zeros(batch_size, device=device, dtype=torch.long) + else: + state = slot.state + condition = slot.condition + n = slot.n + + loss, metrics, next_state, next_n = self.pipeline.compute_data_free_loss( + state, + n=n, + condition=condition, + negative_condition=negative_condition, + collect_metrics=collect_metrics, + ) + if bool(torch.all(next_n == self.pipeline.config.grid_size)): + self._slots[slot_index] = None + else: + self._slots[slot_index] = _TrajectorySlot( + state=next_state, + condition=condition, + n=next_n, + ) + else: + data = batch["image_latents"].to(device=device, dtype=dtype, non_blocking=True) + loss, metrics = self.pipeline.compute_loss( + data, + condition=incoming_condition, + negative_condition=negative_condition, + collect_metrics=collect_metrics, + ) + if check_loss and not bool(torch.isfinite(loss)): + raise FloatingPointError("PDD loss is non-finite.") + + per_sample_loss = metrics["student_target_mse"] + return per_sample_loss, loss, None, metrics if collect_metrics else {} diff --git a/examples/diffusers/fastgen/preprocess/__init__.py b/examples/diffusers/fastgen/preprocess/__init__.py index 60a4a1abf85..d727176bc06 100644 --- a/examples/diffusers/fastgen/preprocess/__init__.py +++ b/examples/diffusers/fastgen/preprocess/__init__.py @@ -34,7 +34,7 @@ ) except ImportError as exc: # pragma: no cover - environment guard raise ImportError( - "fastgen preprocessing requires a stock nemo_automodel>=0.4.0,<1.0 install providing " + "fastgen preprocessing requires a stock nemo_automodel==0.5.0 install providing " "nemo_automodel.components.datasets.diffusion.multi_tier_bucketing. Install the example " "dependencies with:\n" " pip install -r examples/diffusers/fastgen/requirements.txt\n" diff --git a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py index d11efe30f9e..fdd2e9fe20a 100644 --- a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py +++ b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py @@ -113,24 +113,42 @@ def _save_metadata_shards( each other. Merge the per-rank index files afterwards with a separate script to produce a single unified metadata.json. """ + output_root = output_dir.resolve(strict=True) + normalized_metadata = [] + for item_index, item in enumerate(all_metadata): + cache_file = Path(item["cache_file"]).resolve(strict=True) + try: + relative_cache_file = cache_file.relative_to(output_root) + except ValueError as exc: + raise ValueError( + f"cache_file for metadata item {item_index} is outside output root " + f"{output_root}: {cache_file}" + ) from exc + normalized_metadata.append( + { + **item, + "cache_file": relative_cache_file.as_posix(), + } + ) + sharded = shard_world > 1 shard_prefix = f"r{shard_rank:02d}_" if sharded else "" index_filename = f"metadata_r{shard_rank:02d}.json" if sharded else "metadata.json" shard_files = [] - for chunk_start in range(0, len(all_metadata), shard_size): - chunk_data = all_metadata[chunk_start : chunk_start + shard_size] + for chunk_start in range(0, len(normalized_metadata), shard_size): + chunk_data = normalized_metadata[chunk_start : chunk_start + shard_size] chunk_idx = chunk_start // shard_size shard_file = output_dir / f"metadata_shard_{shard_prefix}s{chunk_idx:04d}.json" with open(shard_file, "w") as f: - json.dump(chunk_data, f, indent=2) + json.dump(chunk_data, f, indent=2, allow_nan=False) shard_files.append(shard_file.name) metadata = { "processor": processor_name, "model_name": model_name, "model_type": model_type, - "total_items": len(all_metadata), + "total_items": len(normalized_metadata), "num_shards": len(shard_files), "shard_size": shard_size, "shards": shard_files, @@ -141,7 +159,7 @@ def _save_metadata_shards( metadata["shard_world"] = shard_world with open(output_dir / index_filename, "w") as f: - json.dump(metadata, f, indent=2) + json.dump(metadata, f, indent=2, allow_nan=False) def _print_bucket_distribution(all_metadata: list[dict]) -> None: diff --git a/examples/diffusers/fastgen/preprocess_qwen_image.py b/examples/diffusers/fastgen/preprocess_qwen_image.py index 73f2d3fb9d0..46c76398260 100644 --- a/examples/diffusers/fastgen/preprocess_qwen_image.py +++ b/examples/diffusers/fastgen/preprocess_qwen_image.py @@ -18,7 +18,7 @@ Builds the VAE + text-embed ``.pt`` cache that the DMD2 training dataloader reads, using only stock ``nemo_automodel`` (no dependency on the un-packaged AutoModel ``tools/`` tree). -Mirrors ``dmd2_finetune.py``: it puts this example directory on ``sys.path`` so the +Mirrors ``dmd2/finetune.py``: it puts this example directory on ``sys.path`` so the ``preprocess`` package imports cleanly from a source checkout, then dispatches to the vendored driver's ``main`` (argparse with ``image`` / ``video`` subcommands; this example uses ``image`` with ``--processor qwen_image``). @@ -36,7 +36,7 @@ import sys # Make the ``preprocess`` package importable as a top-level package regardless of the current -# working directory (same seam as dmd2_finetune.py). +# working directory (same seam as ``dmd2/finetune.py``). _THIS_DIR = os.path.dirname(os.path.abspath(__file__)) if _THIS_DIR not in sys.path: sys.path.insert(0, _THIS_DIR) diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index 036d85fb3e0..59c6c3c9fec 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -1,15 +1,10 @@ -# Runtime requirements for the DMD2 Qwen-Image AutoModel example. +# Shared runtime requirements for the Qwen-Image AutoModel examples. # Torch + diffusers are already pulled in via Model-Optimizer's ``[all]`` extras. # The one thing that's NOT shipped with Model-Optimizer is nemo_automodel. -# NeMo AutoModel (parent recipe, FSDP2 wrapping, and the UNPATCHED upstream helpers that the -# vendored data/preprocessing code imports: components.datasets.diffusion.{sampler,base_dataset, -# multi_tier_bucketing,text_to_video_dataset}). The diffusion extras install diffusers + -# accelerate with matching pins. Bounded to a tested range (0.4.0 == the validated commit); -# fastgen_data/__init__.py adds a runtime guard with an actionable message if the helpers move. -# Capped below 0.6.0, which dropped ``recipes.diffusion.train.is_main_process`` (imported by -# dmd2_recipe.py) without a replacement. +# NeMo AutoModel supplies the parent recipe, FSDP2 wrapping, checkpointer, and upstream diffusion +# data helpers. Keep the existing DMD2 compatibility range; PDD has a narrower requirements file. nemo_automodel[diffusion]>=0.4.0,<0.6 -# Optional but recommended for the smoke logs. +# Optional but recommended for training logs. wandb diff --git a/modelopt/torch/fastgen/__init__.py b/modelopt/torch/fastgen/__init__.py index 507acfddd72..961c39c04e6 100644 --- a/modelopt/torch/fastgen/__init__.py +++ b/modelopt/torch/fastgen/__init__.py @@ -57,12 +57,17 @@ from .config import * from .ema import * from .factory import * +from .flow_matching import ( + fusion_coefficients, + integrate_interval_velocities, + make_shifted_flow_grid, +) from .loader import * from .methods.dmd import * +from .methods.pdd import * from .pipeline import * # isort: off -# Plugins must be imported after the core exports so the plugin hooks can reference -# DMDPipeline if needed in the future; also matches the ordering used by -# modelopt.torch.distill. +# Plugins must be imported after the core exports so plugin hooks can reference +# the public pipeline types; this matches the ordering used by modelopt.torch.distill. from . import plugins diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py index 30d78b8720f..83a55b81fac 100644 --- a/modelopt/torch/fastgen/config.py +++ b/modelopt/torch/fastgen/config.py @@ -15,10 +15,9 @@ """Pydantic configuration classes for the fastgen distillation pipelines. -Configurations are layered so a method-specific config (e.g. :class:`DMDConfig`) inherits -shared diffusion-distillation hyperparameters from :class:`DistillationConfig`. All classes -inherit :class:`modelopt.torch.opt.config.ModeloptBaseConfig`, which provides torch-safe -serialization and dict-like iteration. +Configurations inherit :class:`modelopt.torch.opt.config.ModeloptBaseConfig`, which provides +torch-safe serialization and dict-like iteration. DMD builds on shared diffusion-distillation +settings, while PDD exposes only its fixed-grid method settings. The default values in :class:`DMDConfig` mirror the FastGen Wan 2.2 5B experiment at ``FastGen/fastgen/configs/experiments/WanT2V/config_dmd2_wan22_5b.py``. @@ -26,9 +25,10 @@ from __future__ import annotations +import math from typing import TYPE_CHECKING, Literal -from pydantic import Field, model_validator +from pydantic import Field, field_validator, model_validator from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField @@ -39,6 +39,7 @@ "DMDConfig", "DistillationConfig", "EMAConfig", + "PDDConfig", "SampleTimestepConfig", ] @@ -170,8 +171,7 @@ class EMAConfig(ModeloptBaseConfig): class DistillationConfig(ModeloptBaseConfig): """Shared hyperparameters for diffusion step-distillation methods. - Concrete methods subclass this config to add method-specific fields - (see :class:`DMDConfig`). + DMD subclasses this config to add method-specific fields. """ pred_type: PredType = ModeloptField( @@ -218,6 +218,110 @@ class DistillationConfig(ModeloptBaseConfig): ) +class PDDConfig(ModeloptBaseConfig): + """Hyperparameters for Parallel Decoding Distillation (PDD). + + PDD trains one velocity head per interval on a fixed shifted rectified-flow + grid. The explicit inference block schedule partitions that same grid; it does + not define a second timestep schedule. + """ + + guidance_scale: float | None = ModeloptField( + default=None, + title="CFG scale", + description="Teacher classifier-free guidance scale. If ``None``, CFG is disabled.", + ) + grid_size: int = ModeloptField( + default=128, + title="PDD grid size", + description="Number of rectified-flow intervals and student output heads.", + ) + grid_max_t: float = ModeloptField( + default=0.999, + title="PDD grid maximum time", + description="Maximum rectified-flow time used to construct the fixed PDD grid.", + ) + flow_shift: float = ModeloptField( + default=5.0, + title="Rectified-flow grid shift", + description="Shift applied to the fixed decreasing rectified-flow grid.", + ) + block_size_min: int = ModeloptField( + default=4, + title="Minimum training block alignment", + description="Alignment of sampled training start indices.", + ) + block_size_max: int = ModeloptField( + default=64, + title="Maximum trained block size", + description="Largest target span sampled during training.", + ) + teacher_integrator: Literal["euler", "midpoint"] = ModeloptField( + default="euler", + title="Teacher target integrator", + description="Integrator used to estimate the teacher mean velocity for an interval.", + ) + inference_blocks: tuple[int, ...] = ModeloptField( + default=(32, 32, 32, 32), + title="Fused inference block schedule", + description="Contiguous interval counts that partition the complete PDD grid.", + ) + data_free: bool = ModeloptField( + default=False, + title="Data-free training", + description=( + "Carry student-generated trajectories from fresh noise instead of noising real latents." + ), + ) + + @field_validator("grid_max_t", mode="before") + @classmethod + def _check_grid_max_t_type(cls, value: object) -> object: + if type(value) is not float: + raise ValueError(f"grid_max_t must be a float, got {value!r}.") + return value + + @model_validator(mode="after") + def _check_pdd(self) -> PDDConfig: + if self.grid_size <= 0: + raise ValueError(f"grid_size must be > 0, got {self.grid_size}.") + if not math.isfinite(self.grid_max_t) or not 0.0 < self.grid_max_t <= 1.0: + raise ValueError( + "grid_max_t must be finite and satisfy 0 < grid_max_t <= 1, got " + f"{self.grid_max_t!r}." + ) + if not math.isfinite(self.flow_shift) or self.flow_shift < 1.0: + raise ValueError(f"flow_shift must be finite and >= 1, got {self.flow_shift}.") + if not 0 < self.block_size_min <= self.block_size_max <= self.grid_size: + raise ValueError( + "require 0 < block_size_min <= block_size_max <= grid_size, got " + f"{self.block_size_min}, {self.block_size_max}, {self.grid_size}." + ) + if self.grid_size % self.block_size_min != 0: + raise ValueError( + f"grid_size={self.grid_size} must be divisible by " + f"block_size_min={self.block_size_min}." + ) + if not self.inference_blocks: + raise ValueError("inference_blocks must contain at least one block.") + for index, block in enumerate(self.inference_blocks): + if block <= 0: + raise ValueError(f"inference_blocks[{index}] must be > 0, got {block}.") + if sum(self.inference_blocks) != self.grid_size: + raise ValueError( + f"inference_blocks must sum to grid_size={self.grid_size}, got " + f"{sum(self.inference_blocks)}." + ) + return self + + @classmethod + def from_yaml(cls, config_file: str | Path) -> PDDConfig: + """Construct a :class:`PDDConfig` from a filesystem or built-in YAML file.""" + from .loader import load_pdd_config + + return load_pdd_config(config_file) + + class DMDConfig(DistillationConfig): """Hyperparameters for DMD / DMD2 distribution-matching distillation. @@ -305,8 +409,8 @@ def from_yaml(cls, config_file: str | Path) -> DMDConfig: """Construct a :class:`DMDConfig` from a YAML file. Thin wrapper around :func:`modelopt.torch.fastgen.loader.load_dmd_config`. - The resolver searches the built-in ``modelopt_recipes/`` package first, then - the filesystem. Suffixes (``.yml`` / ``.yaml``) may be omitted. + The resolver searches the filesystem first, then the built-in + ``modelopt_recipes/`` package. Suffixes (``.yml`` / ``.yaml``) may be omitted. """ # Imported lazily to avoid a circular import between this module and # ``modelopt.torch.fastgen.loader`` (which imports :class:`DMDConfig`). diff --git a/modelopt/torch/fastgen/flow_matching.py b/modelopt/torch/fastgen/flow_matching.py index 66925dd00d1..735dd3f9398 100644 --- a/modelopt/torch/fastgen/flow_matching.py +++ b/modelopt/torch/fastgen/flow_matching.py @@ -20,8 +20,9 @@ fastgen into any training stack without adopting a new scheduler object. RF convention used throughout: ``alpha_t = 1 - t`` and ``sigma_t = t``, so -``x_t = (1 - t) * x_0 + t * eps`` with ``t in [0, 1]``. Internally all arithmetic is in -``float64`` for numerical stability, and the result is cast back to the input dtype. +``x_t = (1 - t) * x_0 + t * eps`` with ``t in [0, 1]``. Existing RF conversions use +``float64`` intermediates and cast back to the input dtype. PDD grid and interval helpers +instead preserve float32-or-higher outputs because their result remains in the decoding path. """ from __future__ import annotations @@ -39,6 +40,9 @@ __all__ = [ "add_noise", + "fusion_coefficients", + "integrate_interval_velocities", + "make_shifted_flow_grid", "pred_noise_to_pred_x0", "pred_x0_from_flow", "rf_alpha", @@ -50,6 +54,178 @@ ] +def _pdd_math_dtype(dtype: torch.dtype) -> torch.dtype: + """Promote low-precision floating-point dtypes for PDD interval math.""" + if dtype == torch.float64: + return dtype + if dtype in (torch.float16, torch.bfloat16, torch.float32): + return torch.float32 + raise TypeError(f"PDD interval math requires a real floating-point dtype, got {dtype}.") + + +def make_shifted_flow_grid( + grid_size: int, + shift: float, + *, + max_t: float, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Construct the fixed decreasing shifted rectified-flow grid. + + Schedule construction and upper clamps use float64, matching FastGen's RF + scheduler. The completed grid is cast to the requested PDD math dtype. Low + precision requests are promoted to float32 so distinct early intervals do not + collapse for the canonical 128-node, shift-5 schedule. + """ + if isinstance(grid_size, bool) or not isinstance(grid_size, int) or grid_size <= 0: + raise ValueError(f"grid_size must be a positive integer, got {grid_size!r}.") + if type(max_t) is not float: + raise TypeError(f"max_t must be a float, got {max_t!r}.") + if not math.isfinite(max_t) or not 0.0 < max_t <= 1.0: + raise ValueError(f"max_t must be finite and satisfy 0 < max_t <= 1, got {max_t!r}.") + if not math.isfinite(shift) or shift < 1.0: + raise ValueError(f"shift must be finite and >= 1, got {shift!r}.") + + math_dtype = _pdd_math_dtype(dtype) + schedule_dtype = torch.float64 + upper = torch.as_tensor(max_t, device=device, dtype=schedule_dtype) + unshifted = torch.linspace( + max_t, + 0.0, + grid_size + 1, + device=device, + dtype=schedule_dtype, + ) + unshifted = torch.minimum(unshifted, upper) + grid = shift * unshifted / (1.0 + (shift - 1.0) * unshifted) + grid = torch.minimum(grid, upper).to(math_dtype) + torch._assert_async( + torch.all(torch.diff(grid) < 0), + "shifted grid is not strictly decreasing for " + f"grid_size={grid_size}, max_t={max_t}, shift={shift}.", + ) + return grid + + +def _batch_indices( + value: int | torch.Tensor, + *, + name: str, + batch_size: int, + device: torch.device, +) -> torch.Tensor: + """Normalize a scalar or per-sample interval index to a batch tensor.""" + if isinstance(value, bool): + raise TypeError(f"{name} must be an integer or integer tensor, got bool.") + if isinstance(value, int): + return torch.full((batch_size,), value, device=device, dtype=torch.long) + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be an integer or tensor, got {type(value).__name__}.") + if value.dtype == torch.bool or value.dtype.is_floating_point or value.dtype.is_complex: + raise TypeError(f"{name} must use an integer dtype, got {value.dtype}.") + if value.ndim == 0: + return value.to(device=device, dtype=torch.long).expand(batch_size) + if value.shape != (batch_size,): + raise ValueError( + f"{name} must be scalar or shape ({batch_size},), got {tuple(value.shape)}." + ) + return value.to(device=device, dtype=torch.long) + + +def integrate_interval_velocities( + state: torch.Tensor, + velocities: torch.Tensor, + grid: torch.Tensor, + start: int | torch.Tensor, + end: int | torch.Tensor, +) -> torch.Tensor: + """Integrate per-sample velocity heads over half-open blocks ``[start, end)``. + + ``state`` is batch-first with shape ``[B, ...]`` and ``velocities`` has shape + ``[B, N, ...]``. ``start`` and ``end`` may be scalar integers or integer + tensors of shape ``[B]``. The empty block ``start == end`` returns ``state`` + promoted to the PDD math dtype. + """ + if state.ndim < 1: + raise ValueError(f"state must be batch-first, got shape {tuple(state.shape)}.") + if velocities.ndim != state.ndim + 1: + raise ValueError( + f"velocities must have one interval axis after batch; got state {tuple(state.shape)} " + f"and velocities {tuple(velocities.shape)}." + ) + if velocities.shape[0] != state.shape[0] or velocities.shape[2:] != state.shape[1:]: + raise ValueError( + f"velocities must have shape [B, N, *state.shape[1:]], got " + f"state {tuple(state.shape)} and velocities {tuple(velocities.shape)}." + ) + if grid.ndim != 1 or grid.shape[0] != velocities.shape[1] + 1: + raise ValueError( + f"grid must be one-dimensional with N + 1={velocities.shape[1] + 1} nodes, " + f"got shape {tuple(grid.shape)}." + ) + if not grid.dtype.is_floating_point: + raise ValueError("grid must be a strictly decreasing real floating-point tensor.") + torch._assert_async( + torch.all(torch.diff(grid) < 0), + "grid must be a strictly decreasing real floating-point tensor.", + ) + + batch_size, grid_size = state.shape[0], velocities.shape[1] + if isinstance(start, int) and isinstance(end, int) and not 0 <= start <= end <= grid_size: + raise ValueError(f"require 0 <= start <= end <= {grid_size}, got {start}, {end}.") + start_tensor = _batch_indices(start, name="start", batch_size=batch_size, device=state.device) + end_tensor = _batch_indices(end, name="end", batch_size=batch_size, device=state.device) + torch._assert_async( + torch.all((start_tensor >= 0) & (start_tensor <= end_tensor) & (end_tensor <= grid_size)), + f"require 0 <= start <= end <= {grid_size} for every batch element.", + ) + + result_dtype = _pdd_math_dtype(torch.promote_types(state.dtype, velocities.dtype)) + result_dtype = torch.promote_types(result_dtype, _pdd_math_dtype(grid.dtype)) + interval_ids = torch.arange(grid_size, device=state.device) + mask = (interval_ids[None] >= start_tensor[:, None]) & ( + interval_ids[None] < end_tensor[:, None] + ) + widths = torch.diff(grid.to(device=state.device, dtype=result_dtype)) + weights = mask.to(result_dtype) * widths[None] + update = torch.einsum( + "bn,bn...->b...", + weights, + velocities.to(device=state.device, dtype=result_dtype), + ) + return state.to(dtype=result_dtype) + update + + +def fusion_coefficients(grid: torch.Tensor, start: int, end: int) -> torch.Tensor: + """Return step-width-normalized coefficients for a contiguous block ``[start, end)``.""" + if grid.ndim != 1 or grid.shape[0] < 2: + raise ValueError(f"grid must be one-dimensional with at least two nodes, got {grid.shape}.") + if not grid.dtype.is_floating_point: + raise ValueError("grid must be a strictly decreasing real floating-point tensor.") + torch._assert_async( + torch.all(torch.diff(grid) < 0), + "grid must be a strictly decreasing real floating-point tensor.", + ) + if isinstance(start, bool) or isinstance(end, bool): + raise TypeError("start and end must be integers, not bool.") + if not isinstance(start, int) or not isinstance(end, int): + raise TypeError("start and end must be Python integers.") + grid_size = grid.shape[0] - 1 + if not 0 <= start < end <= grid_size: + raise ValueError(f"require 0 <= start < end <= {grid_size}, got {start}, {end}.") + + result_dtype = _pdd_math_dtype(grid.dtype) + math_grid = grid.to(dtype=result_dtype) + widths = torch.diff(math_grid)[start:end] + coefficients = widths / (math_grid[end] - math_grid[start]) + valid_coefficients = torch.all(coefficients > 0) & torch.isclose( + coefficients.sum(), torch.ones((), device=grid.device, dtype=result_dtype) + ) + torch._assert_async(valid_coefficients, "fusion coefficients must be positive and sum to one.") + return coefficients + + def rf_alpha(t: torch.Tensor) -> torch.Tensor: """Rectified-flow data coefficient ``alpha_t = 1 - t``.""" return 1.0 - t diff --git a/modelopt/torch/fastgen/loader.py b/modelopt/torch/fastgen/loader.py index 175ffa23661..845e59da7a1 100644 --- a/modelopt/torch/fastgen/loader.py +++ b/modelopt/torch/fastgen/loader.py @@ -15,10 +15,10 @@ """YAML-driven configuration loading for fastgen distillation pipelines. -YAML is the first-class entry point for DMD configurations — the fastgen library +YAML is the first-class entry point for fastgen configurations — the library does not expect callers to hand-build Python dicts. Typical usage:: - from modelopt.torch.fastgen import DMDConfig, load_dmd_config + from modelopt.torch.fastgen import DMDConfig, PDDConfig, load_dmd_config, load_pdd_config # (a) Load a built-in recipe by relative path cfg = load_dmd_config("general/distillation/dmd2_qwen_image") @@ -29,11 +29,14 @@ # (c) Equivalent classmethod cfg = DMDConfig.from_yaml("/path/to/my_dmd.yaml") + # (d) PDD uses the same resolver with its own schema + pdd_cfg = PDDConfig.from_yaml("general/distillation/pdd_qwen_image") + The loader resolves paths in two places, in order: -1. ``modelopt_recipes/`` (the built-in recipes package shipped with ModelOpt) — resolved - via :func:`importlib.resources.files`. Suffixes ``.yml`` / ``.yaml`` may be omitted. -2. The filesystem (absolute or working-directory-relative). +1. The filesystem (absolute or working-directory-relative). +2. ``modelopt_recipes/`` (the built-in recipes package shipped with ModelOpt) — resolved + via :func:`importlib.resources.files` for relative paths. Suffixes ``.yml`` and ``.yaml`` are both accepted. """ @@ -52,12 +55,12 @@ import yaml -from .config import DMDConfig +from .config import DMDConfig, PDDConfig if TYPE_CHECKING: from importlib.abc import Traversable -__all__ = ["load_config", "load_dmd_config"] +__all__ = ["load_config", "load_dmd_config", "load_pdd_config"] # Root to all built-in recipes shipped with modelopt. @@ -105,8 +108,8 @@ def load_config(config_file: str | Path) -> dict[str, Any]: the ExMy-num-bits post-processing that is specific to quantization recipes. Args: - config_file: YAML path. Suffix is optional; resolution searches the built-in - ``modelopt_recipes/`` package first, then the filesystem. + config_file: YAML path. Suffix is optional; resolution searches the filesystem + first, then the built-in ``modelopt_recipes/`` package. Returns: The parsed dictionary. An empty file yields ``{}``. @@ -122,8 +125,8 @@ def load_config(config_file: str | Path) -> dict[str, Any]: ) return data raise FileNotFoundError( - f"Cannot locate config file {config_file!r}; searched both the built-in " - f"recipe library and the filesystem." + f"Cannot locate config file {config_file!r}; searched the filesystem and then " + f"the built-in recipe library." ) @@ -147,3 +150,9 @@ def load_dmd_config(config_file: str | Path) -> DMDConfig: """ data = load_config(config_file) return DMDConfig(**data) + + +def load_pdd_config(config_file: str | Path) -> PDDConfig: + """Load a YAML file and construct a validated :class:`PDDConfig`.""" + data = load_config(config_file) + return PDDConfig(**data) diff --git a/modelopt/torch/fastgen/methods/__init__.py b/modelopt/torch/fastgen/methods/__init__.py index c999ca87f6f..87c23f46b69 100644 --- a/modelopt/torch/fastgen/methods/__init__.py +++ b/modelopt/torch/fastgen/methods/__init__.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Concrete distillation method implementations (DMD, future: Self-Forcing, CausVid, ...).""" +"""Concrete diffusion-distillation method implementations.""" from .dmd import * +from .pdd import * diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py new file mode 100644 index 00000000000..bd173c785ae --- /dev/null +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -0,0 +1,799 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Framework-neutral projection, training, and sampling primitives for PDD. + +This module owns projection layout and fusion plus the data-dependent and +data-free objectives and block sampler. Model calls and architecture-specific +packing remain behind the adapter protocol and belong in +``modelopt.torch.fastgen.plugins``. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Literal, Protocol + +import torch +import torch.nn.functional as F +from torch import nn + +from ..config import PDDConfig +from ..flow_matching import ( + add_noise, + fusion_coefficients, + integrate_interval_velocities, + make_shifted_flow_grid, +) + +__all__ = [ + "PDDLayerSpec", + "PDDModelAdapter", + "PDDOutputProjection", + "PDDPipeline", + "convert_to_pdd_output_projection", +] + +PDDHeadLayout = Literal["channel_major", "patch_major"] +_HEAD_LAYOUTS = ("channel_major", "patch_major") + + +def _require_int(value: Any, *, name: str, minimum: int = 1) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}, got {value!r}.") + return value + + +@dataclass(frozen=True) +class PDDLayerSpec: + """Immutable description of an architecture's final PDD projection. + + ``channel_major`` stores widened outputs as ``[head, base_output]``. + ``patch_major`` stores them as ``[patch, head, output_channel]`` and + therefore requires the unpatched ``output_channels`` count. + """ + + projection_path: str + head_layout: PDDHeadLayout + output_channels: int | None = None + + def __post_init__(self) -> None: + if not isinstance(self.projection_path, str) or not self.projection_path: + raise ValueError("projection_path must be a non-empty dotted module path.") + if any(not part for part in self.projection_path.split(".")): + raise ValueError( + f"projection_path contains an empty component: {self.projection_path!r}." + ) + if self.head_layout not in _HEAD_LAYOUTS: + raise ValueError( + f"head_layout must be one of {_HEAD_LAYOUTS}, got {self.head_layout!r}." + ) + if self.head_layout == "channel_major": + if self.output_channels is not None: + raise ValueError("channel_major layout does not use output_channels.") + else: + _require_int(self.output_channels, name="output_channels") + + +class PDDOutputProjection(nn.Linear): + """A widened linear projection with one output head per PDD interval. + + ``forward`` returns the full widened output unless an explicit fusion tuple + ``(start, end, grid)`` is supplied. Fused parameters are computed in float32 + and applied without mutating the module or replacing its registered weights. + """ + + def __init__( + self, + in_features: int, + base_out_features: int, + grid_size: int, + layer_spec: PDDLayerSpec, + *, + bias: bool = True, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + """Initialize an unpopulated widened projection with validated layout metadata.""" + _require_int(in_features, name="in_features") + _require_int(base_out_features, name="base_out_features") + _require_int(grid_size, name="grid_size") + if not isinstance(layer_spec, PDDLayerSpec): + raise TypeError(f"layer_spec must be PDDLayerSpec, got {type(layer_spec).__name__}.") + if type(bias) is not bool: + raise ValueError(f"bias must be bool, got {bias!r}.") + if layer_spec.head_layout == "patch_major": + output_channels = layer_spec.output_channels + if output_channels is None or base_out_features % output_channels != 0: + raise ValueError( + f"base_out_features={base_out_features} must be divisible by " + f"output_channels={output_channels} for patch_major layout." + ) + + super().__init__( + in_features, + base_out_features * grid_size, + bias=bias, + device=device, + dtype=dtype, + ) + self.base_out_features = base_out_features + self.grid_size = grid_size + self.layer_spec = layer_spec + + @property + def patch_factor(self) -> int: + """Number of output patches represented by the base projection.""" + if self.layer_spec.head_layout == "channel_major": + return 1 + output_channels = self.layer_spec.output_channels + if output_channels is None: # guarded by PDDLayerSpec validation + raise RuntimeError("patch_major projection is missing output_channels.") + return self.base_out_features // output_channels + + @classmethod + def from_linear( + cls, + base_linear: nn.Linear, + grid_size: int, + layer_spec: PDDLayerSpec, + ) -> PDDOutputProjection: + """Convert a loaded base linear without modifying it. + + Repeating a compatible conversion is idempotent. A PDD projection with + conflicting grid, layout, path, or channel metadata is rejected. + """ + if isinstance(base_linear, cls): + if base_linear.grid_size != grid_size or base_linear.layer_spec != layer_spec: + raise ValueError( + "existing PDDOutputProjection is incompatible with the requested " + f"grid/spec: existing=({base_linear.grid_size}, {base_linear.layer_spec}), " + f"requested=({grid_size}, {layer_spec})." + ) + return base_linear + if not isinstance(base_linear, nn.Linear): + raise TypeError(f"base_linear must be nn.Linear, got {type(base_linear).__name__}.") + + projection = cls( + base_linear.in_features, + base_linear.out_features, + grid_size, + layer_spec, + bias=base_linear.bias is not None, + device=base_linear.weight.device, + dtype=base_linear.weight.dtype, + ) + with torch.no_grad(): + projection.weight.copy_(projection._repeat_base_tensor(base_linear.weight)) + if projection.bias is not None and base_linear.bias is not None: + projection.bias.copy_(projection._repeat_base_tensor(base_linear.bias)) + projection.weight.requires_grad_(base_linear.weight.requires_grad) + if projection.bias is not None and base_linear.bias is not None: + projection.bias.requires_grad_(base_linear.bias.requires_grad) + projection.train(base_linear.training) + return projection + + def _repeat_base_tensor(self, tensor: torch.Tensor) -> torch.Tensor: + """Repeat a base weight or bias according to the configured head layout.""" + if tensor.shape[0] != self.base_out_features: + raise ValueError( + f"base tensor first dimension must be {self.base_out_features}, " + f"got {tensor.shape[0]}." + ) + trailing_shape = tensor.shape[1:] + if self.layer_spec.head_layout == "channel_major": + return ( + tensor.reshape(1, self.base_out_features, *trailing_shape) + .expand(self.grid_size, self.base_out_features, *trailing_shape) + .reshape(self.out_features, *trailing_shape) + .clone() + ) + + output_channels = self.layer_spec.output_channels + if output_channels is None: # guarded by PDDLayerSpec validation + raise RuntimeError("patch_major projection is missing output_channels.") + return ( + tensor.reshape(self.patch_factor, output_channels, *trailing_shape) + .unsqueeze(1) + .expand(self.patch_factor, self.grid_size, output_channels, *trailing_shape) + .reshape(self.out_features, *trailing_shape) + .clone() + ) + + def _tensor_by_head(self, tensor: torch.Tensor) -> torch.Tensor: + """View widened weight or bias as ``[head, base_output, ...]``.""" + trailing_shape = tensor.shape[1:] + if self.layer_spec.head_layout == "channel_major": + return tensor.reshape(self.grid_size, self.base_out_features, *trailing_shape) + + output_channels = self.layer_spec.output_channels + if output_channels is None: # guarded by PDDLayerSpec validation + raise RuntimeError("patch_major projection is missing output_channels.") + return ( + tensor.reshape( + self.patch_factor, + self.grid_size, + output_channels, + *trailing_shape, + ) + .movedim(1, 0) + .reshape(self.grid_size, self.base_out_features, *trailing_shape) + ) + + def _fused_parameters( + self, + start: int, + end: int, + grid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Compute block-fused parameters for the half-open interval ``[start, end)``.""" + if isinstance(start, bool) or isinstance(end, bool): + raise TypeError("start and end must be integers, not bool.") + if not isinstance(start, int) or not isinstance(end, int): + raise TypeError("start and end must be Python integers.") + if grid.ndim != 1 or grid.shape[0] != self.grid_size + 1: + raise ValueError( + f"grid must contain {self.grid_size + 1} nodes, got shape {tuple(grid.shape)}." + ) + coefficients = fusion_coefficients(grid, start, end).to( + device=self.weight.device, + dtype=torch.float32, + ) + head_weights = self._tensor_by_head(self.weight)[start:end] + fused_weight = torch.einsum("n,n...->...", coefficients, head_weights.float()).to( + self.weight.dtype + ) + if self.bias is None: + return fused_weight, None + head_bias = self._tensor_by_head(self.bias)[start:end] + fused_bias = torch.einsum("n,n...->...", coefficients, head_bias.float()).to( + self.bias.dtype + ) + return fused_weight, fused_bias + + def forward( + self, + input: torch.Tensor, + *, + fusion: tuple[int, int, torch.Tensor] | None = None, + ) -> torch.Tensor: + """Apply the widened projection or an explicitly selected fused block.""" + if fusion is None: + return F.linear(input, self.weight, self.bias) + start, end, grid = fusion + fused_weight, fused_bias = self._fused_parameters(start, end, grid) + return F.linear(input, fused_weight, fused_bias) + + +def _get_module_by_path(model: nn.Module, path: str) -> nn.Module: + """Return an already registered nested module at ``path``.""" + if not isinstance(model, nn.Module): + raise TypeError(f"model must be nn.Module, got {type(model).__name__}.") + if not isinstance(path, str) or not path or any(not part for part in path.split(".")): + raise ValueError(f"path must be a non-empty dotted module path, got {path!r}.") + try: + return model.get_submodule(path) + except AttributeError as error: + raise ValueError( + f"module path {path!r} does not resolve to a registered module." + ) from error + + +def _replace_module_by_path(model: nn.Module, path: str, replacement: nn.Module) -> nn.Module: + """Replace an existing nested module and return the previous module.""" + if not isinstance(replacement, nn.Module): + raise TypeError(f"replacement must be nn.Module, got {type(replacement).__name__}.") + previous = _get_module_by_path(model, path) + parent_path, _, name = path.rpartition(".") + parent = _get_module_by_path(model, parent_path) if parent_path else model + setattr(parent, name, replacement) + return previous + + +def convert_to_pdd_output_projection( + model: nn.Module, + layer_spec: PDDLayerSpec, + grid_size: int, +) -> PDDOutputProjection: + """Explicitly replace ``layer_spec.projection_path`` with a PDD projection.""" + current = _get_module_by_path(model, layer_spec.projection_path) + if not isinstance(current, nn.Linear): + raise TypeError( + f"PDD projection at {layer_spec.projection_path!r} must be nn.Linear, " + f"got {type(current).__name__}." + ) + projection = PDDOutputProjection.from_linear(current, grid_size, layer_spec) + if projection is not current: + replaced = _replace_module_by_path(model, layer_spec.projection_path, projection) + if replaced is not current: + raise RuntimeError("projection changed during synchronous PDD conversion.") + return projection + + +class PDDModelAdapter(Protocol): + """Architecture adapter used by the framework-neutral PDD pipeline.""" + + def student_all_heads( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Return canonical ``[batch, head, *latent_shape]`` student velocities.""" + ... + + def student_fused_block( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + start: int, + end: int, + grid: torch.Tensor, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Return one base-shaped velocity from the fused projection block.""" + ... + + def teacher_velocity( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Return the adapter-specific guided teacher velocity.""" + ... + + +class PDDPipeline: + """PDD losses and fused sampler over a single core-owned grid.""" + + def __init__( + self, + student: nn.Module, + teacher: nn.Module | None, + config: PDDConfig, + adapter: PDDModelAdapter, + ) -> None: + """Store the models/config/adapter and freeze the optional training teacher.""" + if not isinstance(config, PDDConfig): + raise TypeError(f"config must be PDDConfig, got {type(config).__name__}.") + self.student = student + self.teacher = None if teacher is None else teacher.eval().requires_grad_(False) + self.config = config + self.adapter = adapter + + def time_grid(self, device: torch.device | str | None = None) -> torch.Tensor: + """Construct this pipeline's sole shifted rectified-flow grid.""" + return make_shifted_flow_grid( + self.config.grid_size, + self.config.flow_shift, + max_t=self.config.grid_max_t, + device=device, + dtype=torch.float32, + ) + + @staticmethod + def _validate_state(state: torch.Tensor, *, name: str) -> None: + if not isinstance(state, torch.Tensor): + raise TypeError(f"{name} must be a tensor, got {type(state).__name__}.") + if state.ndim < 2 or state.shape[0] <= 0: + raise ValueError(f"{name} must have shape [batch, *latent_shape], got {state.shape}.") + if not state.dtype.is_floating_point: + raise TypeError(f"{name} must use a real floating-point dtype, got {state.dtype}.") + + @staticmethod + def _model_kwargs(model_kwargs: Mapping[str, Any] | None) -> dict[str, Any]: + if model_kwargs is None: + return {} + if not isinstance(model_kwargs, Mapping): + raise TypeError( + f"model_kwargs must be a mapping or None, got {type(model_kwargs).__name__}." + ) + return dict(model_kwargs) + + @staticmethod + def _normalize_velocity( + velocity: torch.Tensor, + *, + expected_shape: torch.Size, + device: torch.device, + name: str, + ) -> torch.Tensor: + if not isinstance(velocity, torch.Tensor): + raise TypeError(f"{name} must return a tensor, got {type(velocity).__name__}.") + if velocity.shape != expected_shape: + raise ValueError( + f"{name} must return shape {tuple(expected_shape)}, got {tuple(velocity.shape)}." + ) + if velocity.device != device: + raise ValueError(f"{name} returned device {velocity.device}, expected {device}.") + if not velocity.dtype.is_floating_point: + raise TypeError( + f"{name} must return a real floating-point tensor, got {velocity.dtype}." + ) + return velocity.to(torch.float32) + + @staticmethod + def _validate_explicit_index( + index: torch.Tensor, + *, + name: str, + batch_size: int, + device: torch.device, + ) -> torch.Tensor: + if not isinstance(index, torch.Tensor): + raise TypeError(f"{name} must be an integer tensor, got {type(index).__name__}.") + if index.shape != (batch_size,): + raise ValueError(f"{name} must have shape ({batch_size},), got {tuple(index.shape)}.") + if index.device != device: + raise ValueError(f"{name} must be on {device}, got {index.device}.") + if index.dtype == torch.bool or index.dtype.is_floating_point or index.dtype.is_complex: + raise TypeError(f"{name} must use an integer dtype, got {index.dtype}.") + return index.to(torch.long) + + def _resolve_indices( + self, + *, + batch_size: int, + device: torch.device, + n: torch.Tensor | None, + k: torch.Tensor | None, + generator: torch.Generator | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + grid_size = self.config.grid_size + block_min = self.config.block_size_min + block_max = self.config.block_size_max + if n is None and k is not None: + raise ValueError( + "explicit k requires explicit n so their joint support is deterministic." + ) + if n is None: + n = block_min * torch.randint( + 0, + grid_size // block_min, + (batch_size,), + device=device, + generator=generator, + ) + else: + n = self._validate_explicit_index( + n, + name="n", + batch_size=batch_size, + device=device, + ) + torch._assert_async( + torch.all((n >= 0) & (n < grid_size) & (n.remainder(block_min) == 0)), + f"n must be aligned to {block_min} and satisfy 0 <= n < {grid_size}.", + ) + + upper = torch.minimum(n + block_max, torch.full_like(n, grid_size)) + if k is None: + interval_ids = torch.arange(grid_size, device=device) + support = (interval_ids[None] >= n[:, None]) & (interval_ids[None] < upper[:, None]) + k = torch.multinomial(support.to(torch.float32), 1, generator=generator).squeeze(1) + else: + k = self._validate_explicit_index( + k, + name="k", + batch_size=batch_size, + device=device, + ) + torch._assert_async( + torch.all((k >= n) & (k < upper)), + "k must satisfy n <= k < min(n + block_size_max, grid_size).", + ) + return n, k + + @staticmethod + def _rms_per_sample(value: torch.Tensor) -> torch.Tensor: + dims = tuple(range(1, value.ndim)) + return value.square().mean(dim=dims).sqrt() + + def _compute_loss_from_state( + self, + state: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + model_kwargs: Mapping[str, Any] | None = None, + n: torch.Tensor, + k: torch.Tensor, + collect_metrics: bool = True, + advance_intervals: int | None = None, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor], torch.Tensor | None]: + """Compute the shared PDD target from an already constructed state ``x_n``.""" + if self.teacher is None: + raise RuntimeError("PDD training requires a teacher model.") + kwargs = self._model_kwargs(model_kwargs) + state_fp32 = state.to(torch.float32) + batch_size = state.shape[0] + grid = self.time_grid(state.device) + time_n = grid[n] + broadcast_shape = (batch_size,) + (1,) * (state.ndim - 1) + + student_heads = self.adapter.student_all_heads( + self.student, + state_fp32, + time_n, + condition=condition, + **kwargs, + ) + expected_head_shape = torch.Size((batch_size, self.config.grid_size, *state.shape[1:])) + student_heads = self._normalize_velocity( + student_heads, + expected_shape=expected_head_shape, + device=state.device, + name="student_all_heads", + ) + with torch.no_grad(): + x_bar_k = integrate_interval_velocities(state_fp32, student_heads, grid, n, k) + + batch_ids = torch.arange(batch_size, device=state.device) + student_target = student_heads[batch_ids, k] + time_k = grid[k] + with torch.no_grad(): + teacher_query = x_bar_k.detach() + teacher_first = self.adapter.teacher_velocity( + self.teacher, + teacher_query, + time_k, + condition=condition, + negative_condition=negative_condition, + **kwargs, + ) + teacher_first = self._normalize_velocity( + teacher_first, + expected_shape=state.shape, + device=state.device, + name="teacher_velocity", + ) + if self.config.teacher_integrator == "euler": + teacher_target = teacher_first + else: + delta_k = grid[k + 1] - time_k + midpoint_state = ( + teacher_query + 0.5 * delta_k.reshape(broadcast_shape) * teacher_first + ) + midpoint_time = time_k + 0.5 * delta_k + teacher_target = self.adapter.teacher_velocity( + self.teacher, + midpoint_state, + midpoint_time, + condition=condition, + negative_condition=negative_condition, + **kwargs, + ) + teacher_target = self._normalize_velocity( + teacher_target, + expected_shape=state.shape, + device=state.device, + name="teacher_velocity", + ) + + next_state = None + if advance_intervals is not None: + next_n = n + advance_intervals + torch._assert_async( + torch.all(next_n <= self.config.grid_size), + "data-free state advance must not pass the final PDD interval.", + ) + next_state = integrate_interval_velocities( + state_fp32, + student_heads, + grid, + n, + next_n, + ).detach() + + squared_error = (student_target - teacher_target).square() + loss = squared_error.mean() + metric_dims = tuple(range(1, squared_error.ndim)) + with torch.no_grad(): + metrics = { + "student_target_mse": squared_error.mean(dim=metric_dims).detach(), + } + if collect_metrics: + all_head_dims = tuple(range(1, student_heads.ndim)) + metrics.update( + n=n.detach(), + k=k.detach(), + target_span=(k - n + 1).detach(), + student_velocity_rms=self._rms_per_sample(student_target).detach(), + teacher_velocity_rms=self._rms_per_sample(teacher_target).detach(), + reconstructed_state_rms=self._rms_per_sample(x_bar_k).detach(), + all_student_heads_finite=torch.isfinite(student_heads) + .all(dim=all_head_dims) + .detach(), + student_target_finite=torch.isfinite(student_target) + .all(dim=metric_dims) + .detach(), + teacher_target_finite=torch.isfinite(teacher_target) + .all(dim=metric_dims) + .detach(), + reconstructed_state_finite=torch.isfinite(x_bar_k) + .all(dim=metric_dims) + .detach(), + loss_finite=torch.isfinite(loss).detach(), + ) + return loss, metrics, next_state + + def compute_loss( + self, + data: torch.Tensor, + *, + noise: torch.Tensor | None = None, + condition: Any = None, + negative_condition: Any = None, + model_kwargs: Mapping[str, Any] | None = None, + n: torch.Tensor | None = None, + k: torch.Tensor | None = None, + generator: torch.Generator | None = None, + collect_metrics: bool = True, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute the exact data-dependent PDD objective for one batch.""" + if type(collect_metrics) is not bool: + raise TypeError("collect_metrics must be a bool.") + self._validate_state(data, name="data") + if noise is None: + noise_fp32 = torch.randn( + data.shape, + device=data.device, + dtype=torch.float32, + generator=generator, + ) + else: + self._validate_state(noise, name="noise") + if noise.shape != data.shape: + raise ValueError( + f"noise must match data shape {tuple(data.shape)}, got {tuple(noise.shape)}." + ) + if noise.device != data.device: + raise ValueError(f"noise must be on {data.device}, got {noise.device}.") + noise_fp32 = noise.to(torch.float32) + + batch_size = data.shape[0] + grid = self.time_grid(data.device) + n, k = self._resolve_indices( + batch_size=batch_size, + device=data.device, + n=n, + k=k, + generator=generator, + ) + state = add_noise(data.to(torch.float32), noise_fp32, grid[n]) + loss, metrics, _ = self._compute_loss_from_state( + state, + condition=condition, + negative_condition=negative_condition, + model_kwargs=model_kwargs, + n=n, + k=k, + collect_metrics=collect_metrics, + ) + return loss, metrics + + def compute_data_free_loss( + self, + state: torch.Tensor, + *, + n: torch.Tensor, + condition: Any = None, + negative_condition: Any = None, + model_kwargs: Mapping[str, Any] | None = None, + k: torch.Tensor | None = None, + generator: torch.Generator | None = None, + collect_metrics: bool = True, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor], torch.Tensor, torch.Tensor]: + """Compute Algorithm 3 and return the detached state carried to ``n + L_min``.""" + if type(collect_metrics) is not bool: + raise TypeError("collect_metrics must be a bool.") + self._validate_state(state, name="state") + n, k = self._resolve_indices( + batch_size=state.shape[0], + device=state.device, + n=n, + k=k, + generator=generator, + ) + loss, metrics, next_state = self._compute_loss_from_state( + state, + condition=condition, + negative_condition=negative_condition, + model_kwargs=model_kwargs, + n=n, + k=k, + collect_metrics=collect_metrics, + advance_intervals=self.config.block_size_min, + ) + if next_state is None: # guarded by advance_intervals above + raise RuntimeError("data-free PDD did not produce a carried state.") + return loss, metrics, next_state, (n + self.config.block_size_min).detach() + + def _validate_blocks(self, blocks: Sequence[int] | None) -> tuple[int, ...]: + if blocks is None: + resolved = tuple(self.config.inference_blocks) + else: + if isinstance(blocks, str | bytes) or not isinstance(blocks, Sequence): + raise TypeError("blocks must be a sequence of integer interval counts.") + resolved = tuple(blocks) + if not resolved: + raise ValueError("blocks must contain at least one interval count.") + for index, block in enumerate(resolved): + if type(block) is not int or block <= 0: + raise ValueError(f"blocks[{index}] must be a positive integer, got {block!r}.") + total = sum(resolved) + if total != self.config.grid_size: + raise ValueError(f"blocks must sum to grid_size={self.config.grid_size}, got {total}.") + return resolved + + @torch.no_grad() + def sample( + self, + noise: torch.Tensor, + *, + condition: Any = None, + blocks: Sequence[int] | None = None, + model_kwargs: Mapping[str, Any] | None = None, + ) -> torch.Tensor: + """Sample from raw RF noise with one fused call per contiguous block.""" + self._validate_state(noise, name="noise") + kwargs = self._model_kwargs(model_kwargs) + resolved_blocks = self._validate_blocks(blocks) + grid = self.time_grid(noise.device) + # Derive fusion coefficients from the high-precision schedule, then cast + # them to the FP32 decoding dtype used for state/time integration. + fusion_grid = make_shifted_flow_grid( + self.config.grid_size, + self.config.flow_shift, + max_t=self.config.grid_max_t, + device=noise.device, + dtype=torch.float64, + ) + current = (noise.to(torch.float64) * self.config.grid_max_t).to(torch.float32) + start = 0 + for block in resolved_blocks: + end = start + block + time = grid[start].expand(noise.shape[0]) + velocity = self.adapter.student_fused_block( + self.student, + current, + time, + start=start, + end=end, + grid=fusion_grid, + condition=condition, + **kwargs, + ) + velocity = self._normalize_velocity( + velocity, + expected_shape=current.shape, + device=noise.device, + name="student_fused_block", + ) + current = current + (grid[end] - grid[start]) * velocity + start = end + return current diff --git a/modelopt/torch/fastgen/plugins/__init__.py b/modelopt/torch/fastgen/plugins/__init__.py index 8810470b26f..9c139ee8281 100644 --- a/modelopt/torch/fastgen/plugins/__init__.py +++ b/modelopt/torch/fastgen/plugins/__init__.py @@ -15,13 +15,14 @@ """Optional plugins for the fastgen subpackage (gated via ``import_plugin``). -``qwen_image`` holds the Qwen-Image pipeline plus the forward-hook helpers that expose -intermediate teacher activations to the DMD2 GAN discriminator. The import is gated so -environments that choose not to install the optional fastgen dependencies still see a -clean package import. +``qwen_image`` holds the shared Qwen latent helpers and DMD pipeline, while +``qwen_image_pdd`` holds the explicit PDD conversion and adapter. Imports are +gated so environments without optional fastgen dependencies retain a clean +package import. """ from modelopt.torch.utils import import_plugin with import_plugin("qwen_image"): from .qwen_image import * + from .qwen_image_pdd import * diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py new file mode 100644 index 00000000000..056f6e46943 --- /dev/null +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -0,0 +1,683 @@ +# Adapted from https://github.com/huggingface/diffusers/blob/275869dcae4ebcfee6a80253fdabc56033335020/src/diffusers/models/transformers/transformer_qwenimage.py + +# Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen-Image adapter and explicit output-projection conversion for PDD. + +The masked joint-attention execution follows NVIDIA FastGen merge request 210. +""" + +from __future__ import annotations + +import math +import types +from collections.abc import Mapping +from typing import Any + +import torch +from torch import nn + +from ..config import PDDConfig +from ..methods.pdd import PDDLayerSpec, PDDOutputProjection +from .qwen_image import build_img_shapes, pack_latents, unpack_latents + +__all__ = [ + "QWEN_IMAGE_PDD_EXECUTION", + "QWEN_IMAGE_PDD_LAYER_SPEC", + "QwenImagePDDAdapter", + "convert_qwen_image_to_pdd", + "enable_qwen_image_pdd_forward", + "freeze_qwen_image_pdd_unused_parameters", + "require_qwen_image_pdd_forward", + "restore_qwen_image_pdd_projection", +] + +QWEN_IMAGE_PDD_EXECUTION = "qwen_image_pdd_masked_joint_attention_v1" + +QWEN_IMAGE_PDD_LAYER_SPEC = PDDLayerSpec( + projection_path="transformer.proj_out", + head_layout="channel_major", +) + +_CONTROLLED_MODEL_KWARGS = { + "additional_t_cond", + "attention_kwargs", + "controlnet_block_samples", + "encoder_hidden_states", + "encoder_hidden_states_mask", + "guidance", + "hidden_states", + "img_shapes", + "max_txt_seq_len", + "return_dict", + "timestep", + "txt_seq_lens", + "_pdd_fusion", +} + +_QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE = "_modelopt_qwen_image_pdd_execution" + + +def _config_guidance_embeds(transformer: nn.Module) -> bool: + config = getattr(transformer, "config", None) + if isinstance(config, Mapping): + return bool(config.get("guidance_embeds", False)) + return bool(getattr(config, "guidance_embeds", False)) + + +def _config_value(transformer: nn.Module, name: str, default: Any = None) -> Any: + config = getattr(transformer, "config", None) + if isinstance(config, Mapping): + return config.get(name, default) + return getattr(config, name, default) + + +def _qwen_image_pdd_forward( + self: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + encoder_hidden_states_mask: torch.Tensor | None = None, + timestep: torch.Tensor | None = None, + img_shapes: list[Any] | None = None, + guidance: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + controlnet_block_samples: Any = None, + additional_t_cond: torch.Tensor | None = None, + return_dict: bool = True, + _pdd_fusion: tuple[int, int, torch.Tensor] | None = None, +) -> Any: + """Run Qwen-Image with the masked joint-attention contract required by PDD.""" + if not isinstance(timestep, torch.Tensor) or timestep.dtype != torch.float32: + raise TypeError("Qwen PDD timestep must remain FP32 at transformer entry.") + if not isinstance(encoder_hidden_states, torch.Tensor) or not isinstance( + encoder_hidden_states_mask, torch.Tensor + ): + raise TypeError("Qwen PDD requires text embeddings and their attention mask.") + if attention_kwargs: + raise ValueError("Qwen PDD does not support nonempty attention_kwargs.") + if guidance is not None: + raise ValueError("Qwen PDD does not support transformer guidance embeddings.") + if controlnet_block_samples is not None: + raise ValueError("Qwen PDD does not support ControlNet block samples.") + if additional_t_cond is not None: + raise ValueError("Qwen PDD does not support additional timestep conditioning.") + max_txt_seq_len = encoder_hidden_states.shape[1] + encoder_hidden_states_mask = encoder_hidden_states_mask.to(torch.bool) + + hidden_states = self.img_in(hidden_states) + encoder_hidden_states = self.txt_in(self.txt_norm(encoder_hidden_states)) + temb = self.time_text_embed(timestep, hidden_states) + image_rotary_emb = self.pos_embed( + img_shapes, + max_txt_seq_len=max_txt_seq_len, + device=hidden_states.device, + ) + for block in self.transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + encoder_hidden_states_mask, + temb, + image_rotary_emb, + ) + else: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=encoder_hidden_states_mask, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = self.norm_out(hidden_states, temb) + if _pdd_fusion is None: + output = self.proj_out(hidden_states) + else: + if not isinstance(self.proj_out, PDDOutputProjection): + raise TypeError("Qwen PDD fusion requires a PDDOutputProjection.") + output = self.proj_out(hidden_states, fusion=_pdd_fusion) + if not return_dict: + return (output,) + + from diffusers.models.modeling_outputs import Transformer2DModelOutput + + return Transformer2DModelOutput(sample=output) + + +def _is_qwen_image_pdd_forward(model: nn.Module) -> bool: + forward = model.__dict__.get("forward") + return ( + isinstance(forward, types.MethodType) + and forward.__func__ is _qwen_image_pdd_forward + and forward.__self__ is model + and getattr(model, _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE, None) == QWEN_IMAGE_PDD_EXECUTION + ) + + +def require_qwen_image_pdd_forward(model: nn.Module) -> str: + """Require and return the semantic label for the bound Qwen PDD forward.""" + if not isinstance(model, nn.Module) or not _is_qwen_image_pdd_forward(model): + raise RuntimeError("Qwen-Image PDD requires its masked joint-attention forward.") + return QWEN_IMAGE_PDD_EXECUTION + + +def freeze_qwen_image_pdd_unused_parameters(transformer: nn.Module) -> tuple[str, ...]: + """Freeze final-block text outputs that the Qwen PDD forward does not consume.""" + blocks = getattr(transformer, "transformer_blocks", None) + if not isinstance(blocks, nn.ModuleList) or not blocks: + raise RuntimeError("Qwen PDD requires a nonempty transformer_blocks ModuleList.") + + final_block = blocks[-1] + local_names = ( + "attn.to_add_out.weight", + "attn.to_add_out.bias", + "txt_mlp.net.0.proj.weight", + "txt_mlp.net.0.proj.bias", + "txt_mlp.net.2.weight", + "txt_mlp.net.2.bias", + ) + frozen_names = [] + for local_name in local_names: + try: + parameter = final_block.get_parameter(local_name) + except AttributeError as error: + raise RuntimeError( + f"Qwen PDD final block is missing required parameter {local_name!r}." + ) from error + parameter.requires_grad_(False) + frozen_names.append(f"transformer_blocks.{len(blocks) - 1}.{local_name}") + return tuple(frozen_names) + + +def enable_qwen_image_pdd_forward(transformer: nn.Module) -> nn.Module: + """Bind the masked joint-attention PDD forward to a loaded Qwen transformer.""" + if not isinstance(transformer, nn.Module): + raise TypeError(f"transformer must be nn.Module, got {type(transformer).__name__}.") + if _is_qwen_image_pdd_forward(transformer): + return transformer + + # Diffusers is an optional dependency used only by the Qwen example. + from diffusers import QwenImageTransformer2DModel + + if not isinstance(transformer, QwenImageTransformer2DModel): + raise TypeError( + "Qwen PDD forward binding requires QwenImageTransformer2DModel, " + f"got {type(transformer).__name__}." + ) + existing_forward = transformer.__dict__.get("forward") + if existing_forward is not None: + raise RuntimeError("Qwen root already has a different instance-level forward override.") + if _config_guidance_embeds(transformer): + raise ValueError("Qwen PDD does not support transformer guidance embeddings.") + if getattr(transformer, "peft_config", None): + raise ValueError("Qwen PDD does not support active PEFT adapters.") + if any(getattr(module, "fused_projections", False) for module in transformer.modules()): + raise ValueError("Qwen PDD does not support fused QKV projections.") + for name in ("zero_cond_t", "use_additional_t_cond", "use_layer3d_rope"): + if bool(_config_value(transformer, name, False)): + raise ValueError(f"Qwen PDD requires {name}=False.") + + transformer.forward = types.MethodType(_qwen_image_pdd_forward, transformer) + setattr(transformer, _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE, QWEN_IMAGE_PDD_EXECUTION) + require_qwen_image_pdd_forward(transformer) + return transformer + + +def _validate_qwen_pdd_config(config: PDDConfig) -> None: + if not isinstance(config, PDDConfig): + raise TypeError(f"config must be PDDConfig, got {type(config).__name__}.") + + +def _output_projection(transformer: nn.Module) -> nn.Linear: + projection = getattr(transformer, "proj_out", None) + if not isinstance(projection, nn.Linear): + raise TypeError( + "Qwen-Image transformer must register an nn.Linear at 'proj_out', got " + f"{type(projection).__name__}." + ) + return projection + + +def _require_pdd_projection( + model: nn.Module, + grid_size: int, + base_out_features: int, + *, + fused: bool = False, +) -> nn.Linear: + projection = _output_projection(model) + if projection.out_features != grid_size * base_out_features: + raise ValueError( + f"Qwen PDD proj_out has {projection.out_features} outputs; expected " + f"{grid_size * base_out_features} ({grid_size} heads x {base_out_features})." + ) + if isinstance(projection, PDDOutputProjection): + if projection.grid_size != grid_size or projection.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: + raise ValueError("Qwen PDD projection metadata does not match its configuration.") + elif fused: + raise TypeError("Qwen fused PDD inference requires a PDDOutputProjection.") + return projection + + +def convert_qwen_image_to_pdd( + transformer: nn.Module, + config: PDDConfig, +) -> PDDOutputProjection: + """Replace a loaded, unwrapped Qwen transformer's ``proj_out`` for PDD. + + The full metadata path remains ``transformer.proj_out`` even though this + helper receives the active transformer component directly. Callers must run + conversion before device/distributed wrappers and optimizer construction. + """ + _validate_qwen_pdd_config(config) + if not isinstance(transformer, nn.Module): + raise TypeError(f"transformer must be nn.Module, got {type(transformer).__name__}.") + if _config_guidance_embeds(transformer): + raise ValueError("Qwen-Image PDD does not support transformer guidance embeddings.") + current = _output_projection(transformer) + + projection = PDDOutputProjection.from_linear( + current, + config.grid_size, + QWEN_IMAGE_PDD_LAYER_SPEC, + ) + if projection is not current: + transformer.proj_out = projection + return projection + + +def restore_qwen_image_pdd_projection( + transformer: nn.Module, + config: PDDConfig, +) -> PDDOutputProjection: + """Restore PDD fusion behavior on a serialized, already-widened Qwen projection.""" + _validate_qwen_pdd_config(config) + if not isinstance(transformer, nn.Module): + raise TypeError(f"transformer must be nn.Module, got {type(transformer).__name__}.") + if _config_guidance_embeds(transformer): + raise ValueError("Qwen-Image PDD does not support transformer guidance embeddings.") + current = _output_projection(transformer) + if isinstance(current, PDDOutputProjection): + return PDDOutputProjection.from_linear( + current, + config.grid_size, + QWEN_IMAGE_PDD_LAYER_SPEC, + ) + + base_out_features = _config_value(transformer, "in_channels") + if type(base_out_features) is not int or base_out_features <= 0: + raise ValueError("Qwen-Image transformer config must define positive in_channels.") + expected_out_features = config.grid_size * base_out_features + if current.out_features != expected_out_features: + raise ValueError( + "Serialized Qwen PDD proj_out has the wrong width: expected " + f"{expected_out_features}, got {current.out_features}." + ) + + projection = PDDOutputProjection( + current.in_features, + base_out_features, + config.grid_size, + QWEN_IMAGE_PDD_LAYER_SPEC, + bias=current.bias is not None, + device="meta", + dtype=current.weight.dtype, + ) + projection.weight = current.weight + projection.bias = current.bias + projection.train(current.training) + transformer.proj_out = projection + return projection + + +class QwenImagePDDAdapter: + """Adapt raw Qwen packed-token calls to the framework-neutral PDD protocol.""" + + def __init__( + self, + config: PDDConfig, + *, + compute_dtype: torch.dtype | None = None, + ) -> None: + """Validate the fixed Qwen continuous-time and packed-CFG contract.""" + _validate_qwen_pdd_config(config) + if config.guidance_scale is not None and not math.isfinite(config.guidance_scale): + raise ValueError("guidance_scale must be finite when Qwen teacher CFG is enabled.") + if compute_dtype is not None and ( + not isinstance(compute_dtype, torch.dtype) or not compute_dtype.is_floating_point + ): + raise TypeError("compute_dtype must be a real floating-point torch dtype or None.") + + self.config = config + self.guidance_scale = ( + None if config.guidance_scale is None else float(config.guidance_scale) + ) + self.compute_dtype = compute_dtype + + @staticmethod + def _parse_condition( + condition: Any, + *, + state: torch.Tensor, + name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance(condition, tuple) or len(condition) != 2: + raise TypeError(f"{name} must be a tuple of (encoder_hidden_states, attention_mask).") + encoder_hidden_states, attention_mask = condition + if not isinstance(encoder_hidden_states, torch.Tensor) or not isinstance( + attention_mask, torch.Tensor + ): + raise TypeError(f"{name} entries must be tensors.") + if encoder_hidden_states.ndim != 3 or attention_mask.ndim != 2: + raise ValueError( + f"{name} requires batched embeddings and a 2D mask, got " + f"{tuple(encoder_hidden_states.shape)} and {tuple(attention_mask.shape)}." + ) + if not encoder_hidden_states.dtype.is_floating_point: + raise TypeError(f"{name} embeddings must use a real floating-point dtype.") + if attention_mask.dtype.is_floating_point or attention_mask.dtype.is_complex: + raise TypeError(f"{name} mask must use an integer or bool dtype.") + if attention_mask.shape != encoder_hidden_states.shape[:2]: + raise ValueError(f"{name} mask shape must match the embedding batch and sequence axes.") + if encoder_hidden_states.shape[0] != state.shape[0]: + raise ValueError(f"{name} batch size must match state batch size {state.shape[0]}.") + if encoder_hidden_states.device != state.device or attention_mask.device != state.device: + raise ValueError(f"{name} tensors must be on {state.device}.") + return encoder_hidden_states, attention_mask + + def _model_dtype(self, model: nn.Module, fallback: torch.dtype) -> torch.dtype: + if self.compute_dtype is not None: + return self.compute_dtype + for parameter in model.parameters(): + if parameter.dtype.is_floating_point: + return parameter.dtype + return fallback + + @staticmethod + def _extract_packed_output(output: Any) -> torch.Tensor: + if ( + not isinstance(output, tuple) + or len(output) != 1 + or not isinstance(output[0], torch.Tensor) + ): + raise TypeError("Qwen-Image PDD requires a one-tensor tuple from return_dict=False.") + packed = output[0] + if packed.ndim != 3: + raise ValueError( + f"Qwen-Image model output must be packed [B, P, F], got {tuple(packed.shape)}." + ) + return packed + + def _prepare_model_call( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + model_kwargs: Mapping[str, Any], + ) -> tuple[torch.Tensor, torch.dtype, list[list[tuple[int, int, int]]]]: + require_qwen_image_pdd_forward(model) + if _config_guidance_embeds(model): + raise ValueError("Qwen-Image PDD does not support transformer guidance embeddings.") + conflicts = sorted(_CONTROLLED_MODEL_KWARGS.intersection(model_kwargs)) + if conflicts: + raise ValueError(f"Qwen-Image PDD model_kwargs contains controlled keys: {conflicts}.") + model_dtype = self._model_dtype(model, state.dtype) + if model_dtype != torch.bfloat16: + raise TypeError("Qwen PDD execution requires BF16 compute.") + if time.dtype != torch.float32: + raise TypeError("Qwen PDD execution requires FP32 time.") + packed_state = pack_latents(state).to(model_dtype) + if time.shape != (packed_state.shape[0],) or time.device != packed_state.device: + raise ValueError("Qwen PDD time must have shape [batch] on the state device.") + return ( + packed_state, + model_dtype, + build_img_shapes(state.shape[0], state.shape[2], state.shape[3]), + ) + + def _call_packed( + self, + model: nn.Module, + packed_state: torch.Tensor, + time: torch.Tensor, + model_kwargs: Mapping[str, Any], + *, + prepared_condition: tuple[torch.Tensor, torch.Tensor], + model_dtype: torch.dtype, + img_shapes: list[list[tuple[int, int, int]]], + fusion: tuple[int, int, torch.Tensor] | None = None, + ) -> torch.Tensor: + encoder_hidden_states, attention_mask = prepared_condition + + encoder_hidden_states = encoder_hidden_states.to(model_dtype) + call_kwargs = dict(model_kwargs) + if fusion is not None: + call_kwargs["_pdd_fusion"] = fusion + output = model( + hidden_states=packed_state, + timestep=time, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=attention_mask, + img_shapes=img_shapes, + return_dict=False, + **call_kwargs, + ) + return self._extract_packed_output(output) + + @staticmethod + def _expected_packed_shape(state: torch.Tensor, *, output_features: int) -> torch.Size: + return torch.Size( + ( + state.shape[0], + (state.shape[2] // 2) * (state.shape[3] // 2), + output_features, + ) + ) + + def _unpack_all_heads(self, packed: torch.Tensor, state: torch.Tensor) -> torch.Tensor: + batch_size, channels, height, width = state.shape + base_packed_features = channels * 4 + expected = self._expected_packed_shape( + state, + output_features=self.config.grid_size * base_packed_features, + ) + if packed.shape != expected: + raise ValueError( + f"unfused Qwen PDD output must have shape {tuple(expected)}, " + f"got {tuple(packed.shape)}." + ) + by_head = packed.reshape( + batch_size, + packed.shape[1], + self.config.grid_size, + base_packed_features, + ).permute(0, 2, 1, 3) + flat = by_head.reshape( + batch_size * self.config.grid_size, + packed.shape[1], + base_packed_features, + ) + unpacked = unpack_latents(flat, height, width) + return unpacked.reshape(batch_size, self.config.grid_size, channels, height, width) + + def _unpack_single(self, packed: torch.Tensor, state: torch.Tensor) -> torch.Tensor: + expected = self._expected_packed_shape(state, output_features=state.shape[1] * 4) + if packed.shape != expected: + raise ValueError( + f"fused/base Qwen output must have shape {tuple(expected)}, " + f"got {tuple(packed.shape)}." + ) + return unpack_latents(packed, state.shape[2], state.shape[3]) + + def student_all_heads( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Return unpacked PDD interval velocities from one Qwen call.""" + _require_pdd_projection( + model, + self.config.grid_size, + base_out_features=state.shape[1] * 4, + ) + packed_state, model_dtype, img_shapes = self._prepare_model_call( + model, state, time, model_kwargs + ) + prepared_condition = self._parse_condition(condition, state=state, name="condition") + packed = self._call_packed( + model, + packed_state, + time, + model_kwargs, + prepared_condition=prepared_condition, + model_dtype=model_dtype, + img_shapes=img_shapes, + ) + return self._unpack_all_heads(packed, state) + + def student_fused_block( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + start: int, + end: int, + grid: torch.Tensor, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Run one conditional Qwen call with its final projection fused for a block.""" + projection = _output_projection(model) + base_out_features = ( + projection.base_out_features + if isinstance(projection, PDDOutputProjection) + else state.shape[1] * 4 + ) + _require_pdd_projection( + model, + self.config.grid_size, + base_out_features, + fused=True, + ) + packed_state, model_dtype, img_shapes = self._prepare_model_call( + model, state, time, model_kwargs + ) + prepared_condition = self._parse_condition(condition, state=state, name="condition") + packed = self._call_packed( + model, + packed_state, + time, + model_kwargs, + prepared_condition=prepared_condition, + model_dtype=model_dtype, + img_shapes=img_shapes, + fusion=(start, end, grid), + ) + return self._unpack_single(packed, state) + + @torch.no_grad() + def teacher_velocity( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Return conditional or fixed two-pass packed-CFG Qwen teacher velocity.""" + packed_state, model_dtype, img_shapes = self._prepare_model_call( + model, state, time, model_kwargs + ) + prepared_condition = self._parse_condition(condition, state=state, name="condition") + guidance_scale = self.guidance_scale + prepared_negative_condition = prepared_condition + if guidance_scale is not None: + prepared_negative_condition = self._parse_condition( + negative_condition, + state=state, + name="negative_condition", + ) + + conditional = self._call_packed( + model, + packed_state, + time, + model_kwargs, + prepared_condition=prepared_condition, + model_dtype=model_dtype, + img_shapes=img_shapes, + ) + if guidance_scale is None: + return self._unpack_single(conditional, state) + + unconditional = self._call_packed( + model, + packed_state, + time, + model_kwargs, + prepared_condition=prepared_negative_condition, + model_dtype=model_dtype, + img_shapes=img_shapes, + ) + expected = self._expected_packed_shape(state, output_features=state.shape[1] * 4) + if conditional.shape != expected or unconditional.shape != expected: + raise ValueError( + f"Qwen teacher outputs must both have shape {tuple(expected)}, got " + f"{tuple(conditional.shape)} and {tuple(unconditional.shape)}." + ) + + # Qwen-Image applies its native CFG rescale independently to every + # packed image token before unpacking the velocity. + guided_low_precision = conditional + (float(guidance_scale) - 1.0) * ( + conditional - unconditional + ) + conditional_fp32 = conditional.to(torch.float32) + guided_fp32 = guided_low_precision.to(torch.float32) + conditional_norm = torch.linalg.vector_norm( + conditional_fp32, + dim=-1, + keepdim=True, + ) + guided_norm = torch.linalg.vector_norm( + guided_fp32, + dim=-1, + keepdim=True, + ).clamp_min(1e-5) + guided = (guided_fp32 * (conditional_norm / guided_norm)).to(conditional.dtype) + return self._unpack_single(guided, state) diff --git a/modelopt_recipes/general/distillation/dmd2_qwen_image.yaml b/modelopt_recipes/general/distillation/dmd2_qwen_image.yaml index 79f8124e6fe..acd50f6ecb3 100644 --- a/modelopt_recipes/general/distillation/dmd2_qwen_image.yaml +++ b/modelopt_recipes/general/distillation/dmd2_qwen_image.yaml @@ -63,7 +63,8 @@ sample_t_cfg: p_std: 1.0 # Exact ``torch.linspace(max_t=0.999, 0.0, 5).tolist()`` — the inference # pipeline's default schedule when no t_list is passed (see - # ``inference_dmd2_qwen_image.py:259``). Stride = 0.999 / 4 = 0.249975, so + # ``examples/diffusers/fastgen/dmd2/inference_qwen_image.py:259``). Stride = 0.999 / 4 = + # 0.249975, so # the four training timesteps drawn from ``t_list[:-1]`` exactly match the # four inference sample points. The previous values # ``[0.999, 0.75, 0.5, 0.25, 0.0]`` looked like a uniform 0.25 stride but diff --git a/modelopt_recipes/general/distillation/pdd_qwen_image.yaml b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml new file mode 100644 index 00000000000..ecc4c2b9e4e --- /dev/null +++ b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml @@ -0,0 +1,18 @@ +# Parallel Decoding Distillation defaults for the Qwen-Image example. +# +# Maps to modelopt.torch.fastgen.PDDConfig. This file contains algorithm and +# model-call defaults only; data roots, training topology, checkpoint paths, and +# cluster settings belong to the framework example and run manifest. + +# Canonical Qwen teacher guidance. Packed-space norm rescaling remains owned by +# the Qwen plugin/example rather than the framework-neutral PDD config. +guidance_scale: 4.0 + +grid_size: 128 +grid_max_t: 0.999 +flow_shift: 5.0 +block_size_min: 4 +block_size_max: 64 +teacher_integrator: euler +inference_blocks: [32, 32, 32, 32] +data_free: false diff --git a/tests/examples/diffusers/fastgen/conftest.py b/tests/examples/diffusers/fastgen/conftest.py new file mode 100644 index 00000000000..e4e7d7d04c2 --- /dev/null +++ b/tests/examples/diffusers/fastgen/conftest.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + + +@pytest.fixture +def make_fastgen_cache(): + """Create a tiny, fully local FastGen latent cache.""" + torch = pytest.importorskip("torch") + + def _make( + root: Path, + *, + count: int = 6, + marker: float = 0.0, + absolute_payloads: bool = False, + ) -> Path: + root.mkdir(parents=True, exist_ok=True) + payload_dir = root / "payloads" + payload_dir.mkdir() + + metadata = [] + for sample_id in range(count): + payload_path = payload_dir / f"sample_{sample_id}.pt" + torch.save( + { + "latent": torch.full((4, 2, 2), marker + sample_id), + "crop_offset": (0, 0), + "prompt": f"prompt-{marker}-{sample_id}", + "image_path": f"/source/image_{sample_id}.png", + "prompt_embeds": torch.full((1, 2, 3), marker + sample_id), + "prompt_embeds_mask": torch.ones((1, 2), dtype=torch.long), + }, + payload_path, + ) + cache_file = payload_path if absolute_payloads else payload_path.relative_to(root) + resolution = [64, 64] if sample_id % 2 == 0 else [64, 128] + metadata.append( + { + "cache_file": str(cache_file), + "bucket_resolution": resolution, + "original_resolution": resolution, + "bucket_id": sample_id % 2, + "aspect_ratio": resolution[0] / resolution[1], + } + ) + + midpoint = max(1, count // 2) + shard_names = ["metadata_shard_0.json", "metadata_shard_1.json"] + for shard_name, items in zip( + shard_names, (metadata[:midpoint], metadata[midpoint:]), strict=True + ): + (root / shard_name).write_text(json.dumps(items)) + (root / "metadata.json").write_text(json.dumps({"shards": shard_names})) + torch.save( + { + "embed": torch.full((2, 3), marker), + "mask": torch.ones(2, dtype=torch.long), + }, + root / "negative_prompt_embedding.pt", + ) + return root + + return _make diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py new file mode 100644 index 00000000000..86cada06dff --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Portable and contained path contract for the shared FastGen cache.""" + +from __future__ import annotations + +import json +import logging +import pathlib +import sys +import types + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("nemo_automodel") +pytest.importorskip("torchdata") + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +from fastgen_data import build_text_to_image_multiresolution_dataloader +from fastgen_data.paths import resolve_cache_root, resolve_under_root +from fastgen_data.text_to_image_dataset import TextToImageDataset + + +def test_cache_root_resolves_configured_directory(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + + assert resolve_cache_root(cache) == cache.resolve() + + +def test_cache_root_rejects_missing_or_non_directory(tmp_path): + with pytest.raises(FileNotFoundError): + resolve_cache_root(tmp_path / "missing") + + regular_file = tmp_path / "file" + regular_file.write_text("not a directory") + with pytest.raises(NotADirectoryError): + resolve_cache_root(regular_file) + + +def test_resolve_under_root_rejects_traversal_absolute_and_symlink_escape(tmp_path): + root = tmp_path / "cache" + root.mkdir() + inside = root / "inside.pt" + inside.write_bytes(b"inside") + outside = tmp_path / "outside.pt" + outside.write_bytes(b"outside") + (root / "escape.pt").symlink_to(outside) + + assert resolve_under_root(root, "inside.pt", "sample") == inside.resolve() + for candidate in ("../outside.pt", outside, "escape.pt"): + with pytest.raises(ValueError, match="sample"): + resolve_under_root(root, candidate, "sample") + + +def test_dataset_rejects_metadata_shard_escape(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + safe_item = json.loads((cache / "metadata_shard_0.json").read_text())[0] + outside_shard = tmp_path / "outside.json" + outside_shard.write_text(json.dumps([safe_item])) + (cache / "metadata.json").write_text(json.dumps({"shards": [str(outside_shard)]})) + + with pytest.raises(ValueError, match="metadata shard"): + TextToImageDataset(cache) + + +def test_dataset_rejects_payload_symlink_escape(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + outside = tmp_path / "outside.pt" + torch.save({"latent": torch.zeros(1)}, outside) + (cache / "escape.pt").symlink_to(outside) + shard = json.loads((cache / "metadata_shard_0.json").read_text()) + shard[0]["cache_file"] = "escape.pt" + (cache / "metadata_shard_0.json").write_text(json.dumps(shard)) + + dataset = TextToImageDataset(cache) + with pytest.raises(ValueError, match="sample cache file"): + dataset[0] + + +def test_dataset_accepts_absolute_payload_beneath_root(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache", absolute_payloads=True) + dataset = TextToImageDataset(cache) + + assert "latent" in dataset[0] + + +def test_prompt_only_dataset_and_loader_do_not_emit_image_latents(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + dataset = TextToImageDataset(cache, prompt_only=True) + assert "latent" not in dataset[0] + + loader, _ = build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), + prompt_only=True, + batch_size=1, + num_workers=0, + shuffle=False, + negative_prompt_embedding_path="negative_prompt_embedding.pt", + ) + batch = next(iter(loader)) + + assert "image_latents" not in batch + assert {"text_embeddings", "text_embeddings_mask"}.issubset(batch) + assert "negative_text_embeddings" in batch + + +def test_builder_rejects_negative_embedding_escape(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + outside = tmp_path / "negative.pt" + torch.save(torch.zeros(2, 3), outside) + + with pytest.raises(ValueError, match="negative prompt embedding"): + build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), + num_workers=0, + negative_prompt_embedding_path=str(outside), + ) + + +def test_builder_logs_effective_root_once_on_rank_zero(make_fastgen_cache, caplog, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + caplog.set_level(logging.INFO, logger="fastgen_data.collate_fns") + + build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), dp_rank=1, dp_world_size=2, num_workers=0 + ) + assert not [record for record in caplog.records if "effective_cache_root=" in record.message] + + caplog.clear() + build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), dp_rank=0, dp_world_size=2, num_workers=0 + ) + messages = [ + record.message for record in caplog.records if "effective_cache_root=" in record.message + ] + assert len(messages) == 1 + assert str(cache.resolve()) in messages[0] + assert "selected=6/6" in messages[0] + + +def test_preprocessing_publishes_relocatable_paths(monkeypatch, tmp_path): + # The metadata publisher does not use OpenCV. Stub that optional video dependency so this + # CPU-only test executes the real publisher in the lean AutoModel test environment. + monkeypatch.setitem(sys.modules, "cv2", types.ModuleType("cv2")) + from preprocess.preprocessing_multiprocess import _save_metadata_shards + + monkeypatch.chdir(tmp_path) + output = pathlib.Path("relative-cache") + output.mkdir() + payload = output / "sample.pt" + torch.save({"latent": torch.zeros(1)}, payload) + + _save_metadata_shards( + [{"cache_file": str(payload)}], + output, + "qwen_image", + "Qwen/Qwen-Image", + "qwen_image", + 10, + {}, + ) + + shard = json.loads((output / "metadata_shard_s0000.json").read_text()) + published = pathlib.Path(shard[0]["cache_file"]) + assert published == pathlib.Path("sample.pt") + + relocated = tmp_path / "relocated-cache" + output.rename(relocated) + resolved = resolve_under_root(relocated, published, "sample cache file") + assert torch.equal(torch.load(resolved, weights_only=True)["latent"], torch.zeros(1)) diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference.py b/tests/examples/diffusers/fastgen/test_pdd_inference.py new file mode 100644 index 00000000000..cab6b8142c9 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_inference.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Diffusers serialization seam for a trained Qwen-Image PDD projection.""" + +from __future__ import annotations + +import pathlib +import sys + +import torch +from _test_utils.torch.diffusers_models import get_tiny_qwen_image_transformer +from diffusers import QwenImageTransformer2DModel + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +from pdd.inference_qwen_image import _load_config, _parse_blocks + +from modelopt.torch.fastgen import PDDConfig, PDDOutputProjection +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + convert_qwen_image_to_pdd, + restore_qwen_image_pdd_projection, +) + + +def test_widened_diffusers_projection_restores_pdd_fusion_metadata(tmp_path) -> None: + config = PDDConfig( + grid_size=4, + block_size_min=1, + block_size_max=4, + inference_blocks=[2, 2], + ) + transformer = get_tiny_qwen_image_transformer(num_layers=1) + base_out_channels = transformer.out_channels + projection = convert_qwen_image_to_pdd(transformer, config) + base_projection_features = projection.base_out_features + transformer.register_to_config(out_channels=base_out_channels * config.grid_size) + expected = {name: value.detach().clone() for name, value in projection.state_dict().items()} + transformer.save_pretrained(tmp_path) + + restored = QwenImageTransformer2DModel.from_pretrained(tmp_path) + assert not isinstance(restored.proj_out, PDDOutputProjection) + restored_weight = restored.proj_out.weight + restored_bias = restored.proj_out.bias + adopted = restore_qwen_image_pdd_projection(restored, config) + + assert restored.proj_out is adopted + assert adopted.weight is restored_weight + assert adopted.bias is restored_bias + assert adopted.grid_size == config.grid_size + assert adopted.base_out_features == base_projection_features + for name, value in adopted.state_dict().items(): + torch.testing.assert_close(value, expected[name]) + + +def test_inference_block_override_is_validated(tmp_path) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text( + "pdd:\n" + " grid_size: 8\n" + " block_size_min: 1\n" + " block_size_max: 8\n" + " inference_blocks: [4, 4]\n" + ) + + blocks = _parse_blocks("2, 2,4") + assert blocks == [2, 2, 4] + assert _load_config(config_path, blocks).inference_blocks == tuple(blocks) diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py new file mode 100644 index 00000000000..ebc6ae3e1bd --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused tests for the PDD loss seam used by AutoModel's diffusion recipe.""" + +from __future__ import annotations + +import pathlib +import sys +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +from pdd import compat as pdd_compat +from pdd.recipe import _validate_prepared_student +from pdd.training import PDDFlowMatchingStepAdapter + +from modelopt.torch.fastgen import PDDConfig + + +class _PreparedStudent(nn.Module): + def __init__(self, out_features: int) -> None: + super().__init__() + self.config = SimpleNamespace(in_channels=4) + self.proj_out = nn.Linear(3, out_features) + + +class _PreparedQwenStudent(_PreparedStudent): + def __init__(self, out_features: int) -> None: + super().__init__(out_features) + block = nn.Module() + block.attn = nn.Module() + block.attn.to_add_out = nn.Linear(3, 3) + block.txt_mlp = nn.Module() + block.txt_mlp.net = nn.ModuleList( + [ + nn.ModuleDict({"proj": nn.Linear(3, 4)}), + nn.Identity(), + nn.Linear(4, 3), + ] + ) + self.transformer_blocks = nn.ModuleList([block]) + self.backbone = nn.Linear(3, 3) + + +class _LossPipeline: + def __init__(self, *, guidance_scale: float | None = None) -> None: + self.config = SimpleNamespace(guidance_scale=guidance_scale, data_free=False) + self.scale = nn.Parameter(torch.tensor(2.0)) + self.last_call = None + + def compute_loss(self, data, *, condition, negative_condition, collect_metrics): + self.last_call = (data, condition, negative_condition, collect_metrics) + per_sample = (data * self.scale).square().flatten(1).mean(1) + return per_sample.mean(), {"student_target_mse": per_sample} + + +class _DataFreeLossPipeline: + def __init__(self) -> None: + self.config = SimpleNamespace( + guidance_scale=None, + data_free=True, + grid_max_t=0.999, + grid_size=4, + ) + self.scale = nn.Parameter(torch.tensor(2.0)) + self.calls = [] + + def compute_data_free_loss( + self, + state, + *, + n, + condition, + negative_condition, + collect_metrics, + ): + self.calls.append( + { + "state": state.detach().clone(), + "n": n.detach().clone(), + "condition": tuple(value.detach().clone() for value in condition), + "negative_condition": negative_condition, + "collect_metrics": collect_metrics, + } + ) + per_sample = (state * self.scale).square().flatten(1).mean(1) + return ( + per_sample.mean(), + {"student_target_mse": per_sample}, + state.detach() + 1, + n + 2, + ) + + +def _batch() -> dict[str, torch.Tensor]: + return { + "image_latents": torch.ones(2, 1, 2, 2), + "text_embeddings": torch.ones(2, 3, 4), + "text_embeddings_mask": torch.ones(2, 3, dtype=torch.long), + "negative_text_embeddings": torch.zeros(2, 3, 4), + "negative_text_embeddings_mask": torch.ones(2, 3, dtype=torch.long), + } + + +def test_prepared_student_width_is_validated_before_training() -> None: + config = PDDConfig( + grid_size=8, + block_size_min=1, + block_size_max=8, + inference_blocks=[4, 4], + ) + _validate_prepared_student(_PreparedStudent(out_features=32), config) + + with pytest.raises(ValueError, match="Prepare the Qwen PDD student"): + _validate_prepared_student(_PreparedStudent(out_features=4), config) + + +def test_pdd_fsdp_preserves_fp32_timestep_inputs(monkeypatch) -> None: + def build_manager_args(**kwargs): + del kwargs + return {"_manager_type": "fsdp2"} + + monkeypatch.setattr( + pdd_compat.automodel_diffusion_train, + "_build_diffusion_parallel_manager_args", + build_manager_args, + ) + with pdd_compat.automodel_pdd_setup(): + manager_args = pdd_compat.automodel_diffusion_train._build_diffusion_parallel_manager_args( + dtype=torch.float32, + compute_dtype=torch.bfloat16, + lora_enabled=False, + ) + + policy = manager_args["mp_policy"] + assert policy.param_dtype == torch.bfloat16 + assert policy.reduce_dtype == torch.float32 + assert policy.output_dtype == torch.bfloat16 + assert policy.cast_forward_inputs is False + assert ( + pdd_compat.automodel_diffusion_train._build_diffusion_parallel_manager_args + is build_manager_args + ) + + +def test_unused_final_text_outputs_are_frozen_before_optimizer_collection(monkeypatch) -> None: + class _Pipeline: + @classmethod + def from_pretrained(cls, *args, **kwargs): + del cls, args, kwargs + return SimpleNamespace(transformer=_PreparedQwenStudent(out_features=32)), {} + + monkeypatch.setattr( + pdd_compat.automodel_diffusion_train, + "NeMoAutoDiffusionPipeline", + _Pipeline, + ) + + original_descriptor = _Pipeline.__dict__["from_pretrained"] + with pdd_compat.automodel_pdd_setup(): + setup_pipeline = pdd_compat.automodel_diffusion_train.NeMoAutoDiffusionPipeline + assert issubclass(setup_pipeline, _Pipeline) + student, _ = setup_pipeline.from_pretrained("student", load_for_training=True) + teacher, _ = setup_pipeline.from_pretrained("teacher", load_for_training=False) + + frozen_names = { + name + for name, parameter in student.transformer.named_parameters() + if not parameter.requires_grad + } + assert frozen_names == { + "transformer_blocks.0.attn.to_add_out.weight", + "transformer_blocks.0.attn.to_add_out.bias", + "transformer_blocks.0.txt_mlp.net.0.proj.weight", + "transformer_blocks.0.txt_mlp.net.0.proj.bias", + "transformer_blocks.0.txt_mlp.net.2.weight", + "transformer_blocks.0.txt_mlp.net.2.bias", + } + assert all(parameter.requires_grad for parameter in teacher.transformer.parameters()) + assert _Pipeline.__dict__["from_pretrained"] is original_descriptor + assert pdd_compat.automodel_diffusion_train.NeMoAutoDiffusionPipeline is _Pipeline + + +def test_automodel_setup_rejects_an_unvalidated_release(monkeypatch) -> None: + monkeypatch.setattr(pdd_compat.nemo_automodel, "__version__", "0.6.0") + + with ( + pytest.raises(RuntimeError, match=r"requires nemo_automodel release 0\.5\.0"), + pdd_compat.automodel_pdd_setup(), + ): + pass + + +@pytest.mark.parametrize("guidance_scale", [None, 4.0]) +def test_step_adapter_returns_native_tuple_and_preserves_pdd_gradient(guidance_scale) -> None: + pipeline = _LossPipeline(guidance_scale=guidance_scale) + adapter = PDDFlowMatchingStepAdapter(pipeline) + + per_sample, loss, prediction, metrics = adapter.step( + model=nn.Identity(), + batch=_batch(), + device=torch.device("cpu"), + dtype=torch.float32, + ) + loss.backward() + + assert prediction is None + torch.testing.assert_close(per_sample, torch.full((2,), 4.0)) + assert metrics["student_target_mse"] is per_sample + torch.testing.assert_close(pipeline.scale.grad, torch.tensor(4.0)) + _, _, negative_condition, collect_metrics = pipeline.last_call + assert (negative_condition is not None) is (guidance_scale is not None) + assert collect_metrics is True + + +def test_data_free_step_uses_independent_accumulation_slots_and_reuses_prompts() -> None: + pipeline = _DataFreeLossPipeline() + optimizer_step = [0] + adapter = PDDFlowMatchingStepAdapter( + pipeline, + grad_acc_steps=2, + latent_shape=(1, 2, 2), + optimizer_step_getter=lambda: optimizer_step[0], + ) + + def prompt_batch(value: float) -> dict[str, torch.Tensor]: + return { + "text_embeddings": torch.full((1, 3, 4), value), + "text_embeddings_mask": torch.ones(1, 3, dtype=torch.long), + } + + adapter.step( + nn.Identity(), + prompt_batch(1.0), + device=torch.device("cpu"), + dtype=torch.float32, + global_step=0, + ) + adapter.step( + nn.Identity(), + prompt_batch(2.0), + device=torch.device("cpu"), + dtype=torch.float32, + global_step=0, + ) + first_states = [call["state"] for call in pipeline.calls] + optimizer_step[0] = 1 + + adapter.step( + nn.Identity(), + prompt_batch(101.0), + device=torch.device("cpu"), + dtype=torch.float32, + global_step=0, + ) + adapter.step( + nn.Identity(), + prompt_batch(102.0), + device=torch.device("cpu"), + dtype=torch.float32, + global_step=0, + ) + + assert torch.equal(pipeline.calls[0]["n"], torch.tensor([0])) + assert torch.equal(pipeline.calls[1]["n"], torch.tensor([0])) + assert torch.equal(pipeline.calls[2]["n"], torch.tensor([2])) + assert torch.equal(pipeline.calls[3]["n"], torch.tensor([2])) + torch.testing.assert_close(pipeline.calls[2]["state"], first_states[0] + 1) + torch.testing.assert_close(pipeline.calls[3]["state"], first_states[1] + 1) + assert torch.all(pipeline.calls[2]["condition"][0] == 1.0) + assert torch.all(pipeline.calls[3]["condition"][0] == 2.0) + + optimizer_step[0] = 2 + adapter.step( + nn.Identity(), + prompt_batch(3.0), + device=torch.device("cpu"), + dtype=torch.float32, + global_step=0, + ) + assert torch.equal(pipeline.calls[4]["n"], torch.tensor([0])) + assert torch.all(pipeline.calls[4]["condition"][0] == 3.0) diff --git a/tests/examples/diffusers/fastgen/test_resume_dataloader.py b/tests/examples/diffusers/fastgen/test_resume_dataloader.py index 3eccb103985..2610affdfaf 100644 --- a/tests/examples/diffusers/fastgen/test_resume_dataloader.py +++ b/tests/examples/diffusers/fastgen/test_resume_dataloader.py @@ -40,8 +40,8 @@ import pytest -# Put the example dir on sys.path so ``dmd2_recipe`` imports as a top-level module, -# exactly as dmd2_finetune.py does (mirrors test_vendored_migration.py). +# Put the FastGen directory on sys.path so ``dmd2.recipe`` and shared ``fastgen_data`` imports +# resolve exactly as ``dmd2/finetune.py`` configures them. _REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] _FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" if str(_FASTGEN_DIR) not in sys.path: @@ -52,7 +52,7 @@ pytest.importorskip("torch") _sampler_mod = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") _stateful_dataloader_mod = pytest.importorskip("torchdata.stateful_dataloader") -dmd2_recipe = pytest.importorskip("dmd2_recipe") +dmd2_recipe = pytest.importorskip("dmd2.recipe") SequentialBucketSampler = _sampler_mod.SequentialBucketSampler StatefulDataLoader = _stateful_dataloader_mod.StatefulDataLoader @@ -76,6 +76,12 @@ def __getitem__(self, i): return int(i) # identity: the served value IS the global sample index +class _RecordingSampler(SequentialBucketSampler): + def load_state_dict(self, state_dict): + self.loaded_state = dict(state_dict) + super().load_state_dict(state_dict) + + def _build(n, sampler_cls, loader_cls): """A real sampler + StatefulDataLoader over one shared synthetic dataset.""" ds = _Dataset(n) @@ -91,7 +97,12 @@ def _build(n, sampler_cls, loader_cls): num_replicas=1, rank=0, ) - loader = loader_cls(ds, batch_sampler=sampler, collate_fn=lambda b: b, num_workers=0) + loader = loader_cls( + ds, + batch_sampler=sampler, + collate_fn=lambda b: b, + num_workers=0, + ) return sampler, loader @@ -129,18 +140,15 @@ def test_resume_rebuild_serves_clean_run_position(monkeypatch, epoch_len, grad_a # Mid-epoch, epoch-boundary, and cross-epoch resume points (in optimizer steps). for global_step in resume_points: - sampler, loader = _build(n, SequentialBucketSampler, StatefulDataLoader) - # Use a REAL recipe instance (object.__new__ skips __init__) so BaseRecipe.__setattr__ - # state-tracking is exercised: ``dataloader`` is registered as a tracked key here and the - # reset re-assigns it. A plain stub (no __setattr__) misses the "State key 'dataloader' - # is already tracked" guard that crashed the real run on resume. + sampler, loader = _build(n, _RecordingSampler, StatefulDataLoader) + # Use a real recipe instance (object.__new__ skips __init__) so public BaseRecipe state + # untracking and normal reassignment are exercised. recipe = object.__new__(dmd2_recipe.DMD2DiffusionRecipe) recipe.sampler = sampler recipe.step_scheduler = SimpleNamespace( epoch_len=epoch_len, grad_acc_steps=grad_acc, epoch=0 ) - recipe.dataloader = loader # registers "dataloader" in __state_tracked - assert "dataloader" in recipe.__dict__["__state_tracked"] + recipe.dataloader = loader recipe._rebuild_dataloader_for_resume(global_step) # must not raise "already tracked" @@ -149,11 +157,14 @@ def test_resume_rebuild_serves_clean_run_position(monkeypatch, epoch_len, grad_a cur_epoch = global_step // epoch_len skip_batches = (global_step % epoch_len) * grad_acc assert recipe.step_scheduler.epoch == cur_epoch - assert recipe.sampler._batches_to_skip == skip_batches - assert "dataloader" in recipe.__dict__["__state_tracked"] # still tracked after rebuild + assert recipe.sampler.loaded_state == { + "epoch": cur_epoch, + "batches_yielded": skip_batches, + } + assert recipe.step_scheduler.dataloader is recipe.dataloader # The real training loop calls ``set_epoch(cur_epoch)`` AFTER the rebuild and BEFORE the - # first ``__iter__`` (dmd2_recipe.py). The fix relies on ``set_epoch`` NOT clearing + # first ``__iter__`` (``dmd2/recipe.py``). The fix relies on ``set_epoch`` NOT clearing # ``_batches_to_skip``; replicate that call here so a future sampler that reset the skip # in ``set_epoch`` (silently re-serving from the epoch start) would fail this test. recipe.sampler.set_epoch(cur_epoch) @@ -178,5 +189,5 @@ def test_resume_reset_is_noop_on_fresh_start(monkeypatch): recipe._rebuild_dataloader_for_resume(0) assert recipe.dataloader is loader, "fresh start must not rebuild the dataloader" - assert getattr(recipe.sampler, "_batches_to_skip", 0) == 0 + assert recipe.sampler.state_dict() == {"epoch": 0, "batches_yielded": 0} assert recipe.step_scheduler.epoch == 0 diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index d881a6e49a6..a869552154a 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -37,9 +37,8 @@ import pytest # Resolve the example dir (examples/diffusers/fastgen) from this test's location -# (tests/examples/diffusers/fastgen/) and put it on sys.path so ``fastgen_data`` / -# ``fastgen_checkpoint`` / ``preprocess`` import as top-level modules, exactly as -# dmd2_finetune.py / preprocess_qwen_image.py do. +# (tests/examples/diffusers/fastgen/) and put it on sys.path so ``fastgen_data`` / ``dmd2`` / +# ``preprocess`` imports resolve exactly as ``dmd2/finetune.py`` / ``preprocess_qwen_image.py`` do. _REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] _FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" if str(_FASTGEN_DIR) not in sys.path: @@ -51,7 +50,7 @@ # relative to this example dir) or ``None`` when the patch is intentionally excluded (unused on # the DMD2 path, or not needed for real training). STAGED_AUTOMODEL_DISPOSITION = { - "components/checkpoint/checkpointing.py": "fastgen_checkpoint.py", # subclass override + "components/checkpoint/checkpointing.py": "dmd2/checkpoint.py", # subclass override "components/datasets/diffusion/__init__.py": "fastgen_data/__init__.py", "components/datasets/diffusion/collate_fns.py": "fastgen_data/collate_fns.py", # thin wrapper "components/datasets/diffusion/mock_dataloader.py": None, # excluded: mock smoke (not real training) @@ -74,8 +73,8 @@ def test_all_configs_target_vendored_builders(): Enumerates all YAMLs so a newly added config cannot silently reintroduce the upstream dependence (which would break on stock nemo_automodel). """ - configs = sorted((_FASTGEN_DIR / "configs").glob("*.yaml")) - assert configs, "no configs found under configs/" + configs = sorted(_FASTGEN_DIR.glob("*/configs/*.yaml")) + assert configs, "no algorithm configs found" for cfg in configs: text = cfg.read_text() assert "nemo_automodel.components.datasets.diffusion.build_" not in text, ( @@ -149,8 +148,8 @@ def test_formerly_vendored_files_use_standard_nvidia_header(): # --------------------------------------------------------------------------------------------- # -def test_data_builders_importable_and_accept_negative_prompt_path(): - """The real-data builder exists and accepts ``negative_prompt_embedding_path``.""" +def test_data_builders_importable_and_accept_shared_cache_options(): + """The real-data builder exposes the shared negative-embedding seam.""" pytest.importorskip("nemo_automodel") pytest.importorskip("torch") @@ -172,7 +171,7 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): seq, dim, c, h, w = 5, 16, 4, 8, 8 # A per-item sample matching what TextToImageDataset emits — collate_fn_production requires - # crop_resolution / original_resolution / crop_offset / prompt / image_path / bucket_id / + # crop_resolution / original_resolution / crop_offset / prompt / sample_id / bucket_id / # aspect_ratio in addition to the latent + text embeds. sample = { "latent": torch.randn(c, h, w), @@ -180,7 +179,7 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): "original_resolution": torch.tensor([h, w]), "crop_offset": torch.tensor([0, 0]), "prompt": "a test prompt", - "image_path": "img.png", + "image_path": "/source/image.png", "bucket_id": 0, "aspect_ratio": 1.0, "prompt_embeds": torch.randn(seq, dim), @@ -199,10 +198,40 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): ) # broadcast [seq,dim]->[B,seq,dim] +def test_collate_zero_pads_variable_length_qwen_embeddings_and_masks(): + pytest.importorskip("nemo_automodel") + torch = pytest.importorskip("torch") + + from fastgen_data import collate_fn_text_to_image + + def sample(sample_id, sequence_length): + return { + "latent": torch.full((4, 8, 8), float(sample_id)), + "crop_resolution": torch.tensor([8, 8]), + "original_resolution": torch.tensor([8, 8]), + "crop_offset": torch.tensor([0, 0]), + "prompt": f"prompt-{sample_id}", + "image_path": f"/source/{sample_id}.png", + "bucket_id": 0, + "aspect_ratio": 1.0, + "prompt_embeds": torch.full((sequence_length, 3), float(sample_id)), + "prompt_embeds_mask": torch.ones(sequence_length, dtype=torch.long), + } + + result = collate_fn_text_to_image([sample(1, 5), sample(2, 3)]) + + assert result["text_embeddings"].shape == (2, 5, 3) + assert result["text_embeddings_mask"].tolist() == [ + [1, 1, 1, 1, 1], + [1, 1, 1, 0, 0], + ] + torch.testing.assert_close(result["text_embeddings"][1, 3:], torch.zeros(2, 3)) + + def test_partial_load_checkpointer_overrides_only_load_optimizer(): """The subclass relaxes only optimizer load; model-state load stays strict (inherited).""" pytest.importorskip("nemo_automodel") - from fastgen_checkpoint import PartialLoadCheckpointer, make_optimizer_partial_load_tolerant + from dmd2.checkpoint import PartialLoadCheckpointer, make_optimizer_partial_load_tolerant from nemo_automodel.components.checkpoint.checkpointing import Checkpointer assert issubclass(PartialLoadCheckpointer, Checkpointer) @@ -220,8 +249,8 @@ def test_partial_load_checkpointer_overrides_only_load_optimizer(): def test_recipe_injects_partial_load_checkpointer_in_load_checkpoint(): """The recipe upgrades self.checkpointer in load_checkpoint (before the parent restore).""" - src = (_FASTGEN_DIR / "dmd2_recipe.py").read_text() - assert "from fastgen_checkpoint import make_optimizer_partial_load_tolerant" in src + src = (_FASTGEN_DIR / "dmd2" / "recipe.py").read_text() + assert "from .checkpoint import make_optimizer_partial_load_tolerant" in src assert "make_optimizer_partial_load_tolerant(self.checkpointer)" in src diff --git a/tests/gpu/torch/fastgen/test_pdd_toy.py b/tests/gpu/torch/fastgen/test_pdd_toy.py new file mode 100644 index 00000000000..bf19a1f28e3 --- /dev/null +++ b/tests/gpu/torch/fastgen/test_pdd_toy.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-GPU BF16 proof for the framework-neutral PDD core.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from modelopt.torch.fastgen import ( + PDDConfig, + PDDLayerSpec, + PDDOutputProjection, + PDDPipeline, + convert_to_pdd_output_projection, +) + +_WIDTH = 8 +_GRID_SIZE = 4 + + +class _Student(nn.Module): + def __init__(self) -> None: + super().__init__() + self.backbone = nn.Linear(_WIDTH, _WIDTH) + self.projection = nn.Linear(_WIDTH, _WIDTH) + + def forward( + self, + state: torch.Tensor, + *, + fusion: tuple[int, int, torch.Tensor] | None = None, + ) -> torch.Tensor: + return self.projection(torch.tanh(self.backbone(state)), fusion=fusion) + + +class _Teacher(nn.Module): + def __init__(self) -> None: + super().__init__() + self.projection = nn.Linear(_WIDTH, _WIDTH) + + def forward(self, state: torch.Tensor, time: torch.Tensor) -> torch.Tensor: + return self.projection(state) + 0.125 * time[:, None] + + +class _Adapter: + def __init__(self) -> None: + self.fused_calls = 0 + + @staticmethod + def _model_dtype(model: nn.Module) -> torch.dtype: + return next(model.parameters()).dtype + + def student_all_heads( + self, + model: _Student, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del time, condition, model_kwargs + output = model(state.to(self._model_dtype(model))) + return output.reshape(state.shape[0], _GRID_SIZE, _WIDTH) + + def student_fused_block( + self, + model: _Student, + state: torch.Tensor, + time: torch.Tensor, + *, + start: int, + end: int, + grid: torch.Tensor, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del time, condition, model_kwargs + projection = model.projection + assert isinstance(projection, PDDOutputProjection) + self.fused_calls += 1 + return model(state.to(self._model_dtype(model)), fusion=(start, end, grid)) + + def teacher_velocity( + self, + model: _Teacher, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del condition, negative_condition, model_kwargs + dtype = self._model_dtype(model) + return model(state.to(dtype), time.to(dtype)) + + +def _fill_parameters(model: nn.Module, *, offset: float) -> None: + with torch.no_grad(): + for index, parameter in enumerate(model.parameters()): + values = torch.linspace( + -0.2 + offset + 0.01 * index, + 0.2 + offset + 0.01 * index, + parameter.numel(), + dtype=torch.float32, + device=parameter.device, + ) + parameter.copy_(values.reshape_as(parameter).to(parameter.dtype)) + + +def _build( + device: torch.device, +) -> tuple[_Student, _Teacher, PDDOutputProjection, PDDPipeline, _Adapter]: + config = PDDConfig( + grid_size=_GRID_SIZE, + grid_max_t=0.999, + flow_shift=5.0, + block_size_min=1, + block_size_max=_GRID_SIZE, + inference_blocks=[2, 2], + guidance_scale=None, + ) + student = _Student().to(device=device, dtype=torch.bfloat16) + teacher = _Teacher().to(device=device, dtype=torch.bfloat16) + _fill_parameters(student, offset=0.0) + _fill_parameters(teacher, offset=0.05) + projection = convert_to_pdd_output_projection( + student, + PDDLayerSpec("projection", "channel_major"), + config.grid_size, + ) + adapter = _Adapter() + pipeline = PDDPipeline(student, teacher, config, adapter) + return student, teacher, projection, pipeline, adapter + + +def test_bf16_loss_backward_and_fused_sample() -> None: + assert torch.cuda.is_available(), "BF16 test requires a real CUDA device" + device = torch.device("cuda", 0) + assert torch.cuda.get_device_capability(device)[0] >= 8, "BF16 requires Ampere or newer" + + student, teacher, projection, pipeline, adapter = _build(device) + assert {parameter.dtype for parameter in student.parameters()} == {torch.bfloat16} + assert pipeline.time_grid(device).dtype == torch.float32 + + data = torch.linspace(-0.75, 0.75, _WIDTH, device=device, dtype=torch.bfloat16).reshape(1, -1) + noise = torch.linspace(0.5, -0.5, _WIDTH, device=device, dtype=torch.float32).reshape(1, -1) + n = torch.tensor([0], device=device, dtype=torch.int64) + k = torch.tensor([2], device=device, dtype=torch.int64) + loss, metrics = pipeline.compute_loss(data, noise=noise, n=n, k=k) + assert loss.dtype == torch.float32 + assert torch.isfinite(loss) + for name in ( + "all_student_heads_finite", + "student_target_finite", + "teacher_target_finite", + "reconstructed_state_finite", + "loss_finite", + ): + assert bool(metrics[name].all()), name + loss.backward() + + assert all(parameter.grad is None for parameter in teacher.parameters()) + assert projection.weight.grad is not None + assert student.backbone.weight.grad is not None + assert torch.isfinite(projection.weight.grad).all() + assert torch.isfinite(student.backbone.weight.grad).all() + + sampled = pipeline.sample(noise, blocks=[2, 2]) + assert sampled.dtype == torch.float32 + assert torch.isfinite(sampled).all() + assert adapter.fused_calls == 2 + torch.cuda.synchronize(device) diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 5945187dd02..7c4c8574b3b 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -35,6 +35,7 @@ RecipeType, ) from modelopt.recipe.loader import _apply_dotlist, load_config, load_recipe +from modelopt.torch.fastgen import PDDConfig, load_pdd_config from modelopt.torch.opt.config_loader import _load_raw_config, _schema_type from modelopt.torch.quantization.config import QuantizerAttributeConfig, normalize_quant_cfg_list from modelopt.torch.quantization.mode import CalibrateModeRegistry, get_modelike_from_algo_cfg @@ -73,6 +74,18 @@ quantize: {} """ + +def test_load_pdd_config_builtin_recipe(): + """The public PDD loader resolves and validates its built-in recipe.""" + config = load_pdd_config("general/distillation/pdd_qwen_image") + + assert isinstance(config, PDDConfig) + assert config.guidance_scale == 4.0 + assert config.grid_max_t == 0.999 + assert "grid_max_t" in config.model_fields_set + assert config.inference_blocks == (32, 32, 32, 32) + + QUANTIZER_ATTRIBUTE_SCHEMA = ( "# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig\n" ) diff --git a/tests/unit/torch/fastgen/test_pdd_config.py b/tests/unit/torch/fastgen/test_pdd_config.py new file mode 100644 index 00000000000..ebbfaae9fd6 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Validation and loading tests for framework-neutral PDD configuration.""" + +from __future__ import annotations + +import pytest + +from modelopt.torch.fastgen import PDDConfig, load_pdd_config + + +def test_default_pdd_config_is_canonical(): + config = PDDConfig() + + assert config.guidance_scale is None + assert config.grid_size == 128 + assert config.grid_max_t == 0.999 + assert config.flow_shift == 5.0 + assert config.block_size_min == 4 + assert config.block_size_max == 64 + assert config.teacher_integrator == "euler" + assert config.inference_blocks == (32, 32, 32, 32) + assert config.data_free is False + + +def test_pdd_config_accepts_supported_schedule(): + config = PDDConfig( + inference_blocks=[64, 64], + teacher_integrator="midpoint", + data_free=True, + ) + + assert config.inference_blocks == (64, 64) + assert config.teacher_integrator == "midpoint" + assert config.data_free is True + + +@pytest.mark.parametrize("value", [True, 1]) +def test_pdd_config_rejects_non_float_grid_max_t_before_coercion(value): + with pytest.raises(ValueError, match="grid_max_t must be a float"): + PDDConfig(grid_max_t=value) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"grid_size": 0}, "grid_size must be > 0"), + ({"grid_max_t": 0.0}, "0 < grid_max_t <= 1"), + ({"grid_max_t": 1.0001}, "0 < grid_max_t <= 1"), + ({"grid_max_t": float("nan")}, "0 < grid_max_t <= 1"), + ({"flow_shift": 0.5}, "flow_shift must be finite and >= 1"), + ({"block_size_min": 0}, "0 < block_size_min"), + ({"block_size_max": 129}, "block_size_max <= grid_size"), + ({"grid_size": 130}, "must be divisible"), + ({"inference_blocks": []}, "at least one block"), + ({"inference_blocks": [32, 32, 32]}, "must sum to grid_size"), + ], +) +def test_pdd_config_rejects_invalid_grid_and_block_boundaries(overrides, message): + with pytest.raises(ValueError, match=message): + PDDConfig(**overrides) + + +def test_pdd_config_accepts_inference_partition_outside_training_block_support(): + config = PDDConfig(inference_blocks=[1, 127]) + + assert config.inference_blocks == (1, 127) + + +@pytest.mark.parametrize( + "overrides", [{"teacher_integrator": "heun"}, {"teacher_integrator": "rk4"}] +) +def test_pdd_config_locks_algorithm_modes(overrides): + with pytest.raises(ValueError): + PDDConfig(**overrides) + + +def test_pdd_config_loads_filesystem_yaml_with_optional_suffix(tmp_path): + config_path = tmp_path / "pdd.yaml" + config_path.write_text( + "inference_blocks: [64, 64]\nteacher_integrator: midpoint\n", + encoding="utf-8", + ) + + loaded = load_pdd_config(config_path.with_suffix("")) + from_class = PDDConfig.from_yaml(config_path) + + assert loaded == from_class + assert loaded.inference_blocks == (64, 64) + assert loaded.teacher_integrator == "midpoint" diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py new file mode 100644 index 00000000000..fb5c98d2c72 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -0,0 +1,502 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Analytic tests for the framework-neutral PDD objective and fused sampler.""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch +from torch import nn + +from modelopt.torch.fastgen import PDDConfig, PDDPipeline +from modelopt.torch.fastgen.flow_matching import fusion_coefficients, make_shifted_flow_grid + + +class _HeadModel(nn.Module): + """Small state-dependent multi-head velocity model with explicit parameters.""" + + def __init__(self, grid_size: int, width: int) -> None: + super().__init__() + self.state_scale = nn.Parameter(torch.tensor(0.25)) + self.head_bias = nn.Parameter( + torch.arange(grid_size * width, dtype=torch.float32).reshape(grid_size, width) / 7 + ) + + def all_heads(self, state: torch.Tensor) -> torch.Tensor: + return self.state_scale * state[:, None] + self.head_bias[None] + + +class _Teacher(nn.Module): + """Analytic teacher velocity ``scale * state + time + bias``.""" + + def __init__(self, width: int) -> None: + super().__init__() + self.scale = nn.Parameter(torch.tensor(-0.375)) + self.register_buffer("bias", torch.arange(width, dtype=torch.float32) / 13) + + def forward(self, state: torch.Tensor, time: torch.Tensor) -> torch.Tensor: + return self.scale * state + time[:, None] + self.bias + + +class _RecordingAdapter: + """Canonical toy adapter that records every architecture boundary call.""" + + def __init__(self) -> None: + self.student_calls: list[dict[str, Any]] = [] + self.fused_calls: list[dict[str, Any]] = [] + self.teacher_calls: list[dict[str, Any]] = [] + self.bad_student_shape = False + self.bad_fused_dtype = False + self.low_precision_outputs = False + + def student_all_heads( + self, + model: _HeadModel, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + self.student_calls.append( + { + "state": state.detach().clone(), + "time": time.detach().clone(), + "condition": condition, + "kwargs": model_kwargs, + } + ) + output = model.all_heads(state) + if self.bad_student_shape: + output = output[:, :-1] + return output.to(torch.bfloat16) if self.low_precision_outputs else output + + def student_fused_block( + self, + model: _HeadModel, + state: torch.Tensor, + time: torch.Tensor, + *, + start: int, + end: int, + grid: torch.Tensor, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + self.fused_calls.append( + { + "state": state.detach().clone(), + "time": time.detach().clone(), + "start": start, + "end": end, + "grid": grid.detach().clone(), + "condition": condition, + "kwargs": model_kwargs, + } + ) + # The real PDD projection derives coefficients from the float64 grid, + # then casts them to FP32 before fusing FP32 master parameters. + coefficients = fusion_coefficients(grid, start, end).to(torch.float32) + heads = model.all_heads(state)[:, start:end] + output = torch.einsum("n,bnd->bd", coefficients, heads) + return output.to(torch.int64) if self.bad_fused_dtype else output + + def teacher_velocity( + self, + model: _Teacher, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + self.teacher_calls.append( + { + "state": state.detach().clone(), + "time": time.detach().clone(), + "requires_grad": state.requires_grad, + "condition": condition, + "negative_condition": negative_condition, + "kwargs": model_kwargs, + } + ) + output = model(state, time) + return output.to(torch.bfloat16) if self.low_precision_outputs else output + + +def _config(*, teacher_integrator: str = "euler") -> PDDConfig: + return PDDConfig( + grid_size=8, + grid_max_t=0.999, + flow_shift=5.0, + block_size_min=2, + block_size_max=4, + inference_blocks=[4, 4], + teacher_integrator=teacher_integrator, + ) + + +def _pipeline(*, teacher_integrator: str = "euler") -> tuple[PDDPipeline, _RecordingAdapter]: + adapter = _RecordingAdapter() + pipeline = PDDPipeline( + _HeadModel(grid_size=8, width=3), + _Teacher(width=3), + _config(teacher_integrator=teacher_integrator), + adapter, + ) + return pipeline, adapter + + +def _explicit_integrate_per_sample( + state: torch.Tensor, + heads: torch.Tensor, + grid: torch.Tensor, + n: torch.Tensor, + k: torch.Tensor, +) -> torch.Tensor: + result = state.clone() + for batch_index in range(state.shape[0]): + for head_index in range(int(n[batch_index]), int(k[batch_index])): + result[batch_index] += (grid[head_index + 1] - grid[head_index]) * heads[ + batch_index, head_index + ] + return result + + +def _reference_rf_forward_process( + data: torch.Tensor, + noise: torch.Tensor, + time: torch.Tensor, +) -> torch.Tensor: + """Reproduce the PDD RF input staging without production helpers.""" + data_64 = data.to(torch.float32).to(torch.float64) + noise_64 = noise.to(torch.float32).to(torch.float64) + time_64 = time.to(torch.float32).to(torch.float64) + while time_64.ndim < data_64.ndim: + time_64 = time_64.unsqueeze(-1) + return (data_64 * (1.0 - time_64) + noise_64 * time_64).to(torch.float32) + + +def test_euler_loss_matches_analytic_empty_and_tail_reconstruction() -> None: + pipeline, adapter = _pipeline() + data = torch.tensor([[1.0, -2.0, 0.5], [-1.5, 0.25, 2.0]]) + noise = torch.tensor([[0.25, 1.5, -0.5], [2.0, -1.0, 0.75]]) + n = torch.tensor([0, 6]) + k = torch.tensor([0, 7]) + + loss, metrics = pipeline.compute_loss( + data, + noise=noise, + condition="positive", + negative_condition="negative", + model_kwargs={"tag": 17}, + n=n, + k=k, + ) + + grid = pipeline.time_grid() + x_n = _reference_rf_forward_process(data, noise, grid[n]) + heads = pipeline.student.all_heads(x_n) + x_bar_k = _explicit_integrate_per_sample(x_n, heads, grid, n, k) + teacher_target = pipeline.teacher(x_bar_k, grid[k]) + student_target = heads[torch.arange(data.shape[0]), k] + expected_loss = (student_target - teacher_target).square().mean() + + torch.testing.assert_close(loss, expected_loss) + torch.testing.assert_close(adapter.student_calls[0]["state"], x_n) + torch.testing.assert_close(adapter.student_calls[0]["time"], grid[n]) + torch.testing.assert_close(adapter.teacher_calls[0]["state"], x_bar_k) + torch.testing.assert_close(adapter.teacher_calls[0]["time"], grid[k]) + assert adapter.teacher_calls[0]["requires_grad"] is False + assert adapter.student_calls[0]["condition"] == "positive" + assert adapter.student_calls[0]["kwargs"] == {"tag": 17} + assert adapter.teacher_calls[0]["negative_condition"] == "negative" + assert len(adapter.student_calls) == len(adapter.teacher_calls) == 1 + assert torch.equal(metrics["n"], n) + assert torch.equal(metrics["k"], k) + assert torch.equal(metrics["target_span"], torch.tensor([1, 2])) + assert metrics["student_target_mse"].shape == (2,) + assert metrics["all_student_heads_finite"].shape == (2,) + assert bool(metrics["loss_finite"]) + assert all(not value.requires_grad for value in metrics.values()) + + +def test_midpoint_target_uses_exact_final_interval_midpoint() -> None: + pipeline, adapter = _pipeline(teacher_integrator="midpoint") + data = torch.tensor([[1.0, -2.0, 0.5], [-1.5, 0.25, 2.0]]) + noise = torch.tensor([[0.25, 1.5, -0.5], [2.0, -1.0, 0.75]]) + n = torch.tensor([4, 6]) + k = torch.tensor([6, 7]) + + loss, _ = pipeline.compute_loss(data, noise=noise, n=n, k=k) + + grid = pipeline.time_grid() + x_n = _reference_rf_forward_process(data, noise, grid[n]) + heads = pipeline.student.all_heads(x_n) + x_bar_k = _explicit_integrate_per_sample(x_n, heads, grid, n, k) + first_velocity = pipeline.teacher(x_bar_k, grid[k]) + delta = grid[k + 1] - grid[k] + midpoint_state = x_bar_k + 0.5 * delta.reshape(2, 1) * first_velocity + midpoint_time = grid[k] + 0.5 * delta + midpoint_target = pipeline.teacher(midpoint_state, midpoint_time) + expected_loss = (heads[torch.arange(data.shape[0]), k] - midpoint_target).square().mean() + + torch.testing.assert_close(loss, expected_loss) + assert len(adapter.student_calls) == 1 + assert len(adapter.teacher_calls) == 2 + torch.testing.assert_close(adapter.teacher_calls[0]["state"], x_bar_k) + torch.testing.assert_close(adapter.teacher_calls[0]["time"], grid[k]) + torch.testing.assert_close(adapter.teacher_calls[1]["state"], midpoint_state) + torch.testing.assert_close(adapter.teacher_calls[1]["time"], midpoint_time) + + +def test_data_free_loss_matches_algorithm_3_and_advances_exact_minimum_block() -> None: + pipeline, adapter = _pipeline(teacher_integrator="midpoint") + state = torch.tensor([[0.25, 1.5, -0.5], [2.0, -1.0, 0.75]]) + n = torch.tensor([0, 6]) + k = torch.tensor([3, 7]) + + loss, metrics, next_state, next_n = pipeline.compute_data_free_loss( + state, + n=n, + k=k, + condition="positive", + negative_condition="negative", + ) + + grid = pipeline.time_grid() + heads = pipeline.student.all_heads(state) + x_k = _explicit_integrate_per_sample(state, heads, grid, n, k) + first_velocity = pipeline.teacher(x_k, grid[k]) + delta = grid[k + 1] - grid[k] + midpoint_state = x_k + 0.5 * delta[:, None] * first_velocity + midpoint_target = pipeline.teacher(midpoint_state, grid[k] + 0.5 * delta) + selected = heads[torch.arange(state.shape[0]), k] + expected_next = _explicit_integrate_per_sample( + state, + heads, + grid, + n, + n + pipeline.config.block_size_min, + ) + + torch.testing.assert_close(loss, (selected - midpoint_target).square().mean()) + torch.testing.assert_close(next_state, expected_next) + assert next_state.requires_grad is False + assert torch.equal(next_n, torch.tensor([2, 8])) + assert torch.equal(metrics["n"], n) + assert torch.equal(metrics["k"], k) + assert adapter.student_calls[0]["condition"] == "positive" + assert adapter.teacher_calls[0]["negative_condition"] == "negative" + + +def test_selected_head_low_precision_outputs_use_float32_mse() -> None: + pipeline, adapter = _pipeline() + adapter.low_precision_outputs = True + data = torch.tensor([[1.0, -2.0, 0.5]]) + noise = torch.tensor([[0.25, 1.5, -0.5]]) + n = torch.tensor([0]) + k = torch.tensor([0]) + + loss, _ = pipeline.compute_loss(data, noise=noise, n=n, k=k) + + grid = pipeline.time_grid() + x_n = _reference_rf_forward_process(data, noise, grid[n]) + selected = pipeline.student.all_heads(x_n)[:, 0].to(torch.bfloat16).float() + teacher = pipeline.teacher(x_n, grid[k]).to(torch.bfloat16).float() + expected = (selected - teacher).square().mean() + + assert loss.dtype == torch.float32 + torch.testing.assert_close(loss, expected) + + +def test_sampled_indices_stay_on_exact_uniform_support() -> None: + pipeline, _ = _pipeline() + generator = torch.Generator().manual_seed(1234) + + _, metrics = pipeline.compute_loss( + torch.ones(4096, 3), + noise=torch.zeros(4096, 3), + generator=generator, + ) + n, k = metrics["n"], metrics["k"] + + assert set(n.tolist()) == {0, 2, 4, 6} + assert torch.all(n.remainder(2) == 0) + assert torch.all(k >= n) + assert torch.all(k < torch.minimum(n + 4, torch.full_like(n, 8))) + for n_value in (0, 2, 4, 6): + observed = set(k[n == n_value].tolist()) + assert observed == set(range(n_value, min(n_value + 4, 8))) + + +@pytest.mark.parametrize("blocks", [None, [2, 2, 2, 2], [1, 7]]) +def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: + pipeline, adapter = _pipeline() + noise = torch.tensor([[1.0, -2.0, 0.5], [-0.25, 0.75, 1.5]], dtype=torch.bfloat16) + + actual = pipeline.sample( + noise, + condition="prompt", + blocks=blocks, + model_kwargs={"tag": 23}, + ) + + resolved = [4, 4] if blocks is None else blocks + grid = pipeline.time_grid() + fusion_grid = make_shifted_flow_grid( + pipeline.config.grid_size, + pipeline.config.flow_shift, + max_t=pipeline.config.grid_max_t, + dtype=torch.float64, + ) + expected = (noise.to(torch.float64) * pipeline.config.grid_max_t).to(torch.float32) + start = 0 + for block in resolved: + end = start + block + heads = pipeline.student.all_heads(expected) + for index in range(start, end): + expected = expected + (grid[index + 1] - grid[index]) * heads[:, index] + start = end + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=2e-6, atol=2e-6) + assert len(adapter.fused_calls) == len(resolved) + start = 0 + for call, block in zip(adapter.fused_calls, resolved): + end = start + block + assert (call["start"], call["end"]) == (start, end) + torch.testing.assert_close(call["time"], grid[start].expand(noise.shape[0])) + torch.testing.assert_close(call["grid"], fusion_grid) + assert call["grid"].dtype == torch.float64 + assert call["condition"] == "prompt" + assert call["kwargs"] == {"tag": 23} + start = end + + +def test_sampling_does_not_require_a_teacher() -> None: + pipeline, adapter = _pipeline() + sampler = PDDPipeline(pipeline.student, None, pipeline.config, adapter) + + result = sampler.sample(torch.ones(1, 3)) + + assert result.shape == (1, 3) + with pytest.raises(RuntimeError, match="training requires a teacher"): + sampler.compute_loss(torch.ones(1, 3)) + + +def test_fused_sampler_uses_precast_max_time_once_for_raw_noise() -> None: + pipeline, adapter = _pipeline() + noise = torch.tensor( + [[-0.21963761746883392, -1.409722924232483, 1.8951480388641357]], + dtype=torch.float32, + ) + + pipeline.sample(noise) + + expected = (noise.to(torch.float64) * 0.999).to(torch.float32) + from_cast_grid = (noise.to(torch.float64) * pipeline.time_grid()[0].to(torch.float64)).to( + torch.float32 + ) + double_scaled = (expected.to(torch.float64) * 0.999).to(torch.float32) + assert not torch.equal(expected, from_cast_grid) + assert not torch.equal(expected, double_scaled) + assert torch.equal(adapter.fused_calls[0]["state"], expected) + + +@pytest.mark.parametrize( + ("n", "k", "message"), + [ + (torch.tensor([1]), torch.tensor([1]), "n must be aligned"), + (torch.tensor([8]), torch.tensor([8]), "n must be aligned"), + (torch.tensor([2]), torch.tensor([1]), "k must satisfy"), + (torch.tensor([2]), torch.tensor([6]), "k must satisfy"), + (torch.tensor([6]), torch.tensor([8]), "k must satisfy"), + ], +) +def test_loss_rejects_indices_outside_trained_support(n, k, message) -> None: + pipeline, _ = _pipeline() + with pytest.raises(RuntimeError, match=message): + pipeline.compute_loss(torch.ones(1, 3), noise=torch.zeros(1, 3), n=n, k=k) + + +def test_pipeline_rejects_invalid_shapes_dtypes_and_blocks() -> None: + pipeline, adapter = _pipeline() + with pytest.raises(TypeError, match="data must be a tensor"): + pipeline.compute_loss({"state": torch.ones(1, 3)}) # type: ignore[arg-type] + with pytest.raises(TypeError, match="noise must be a tensor"): + pipeline.sample({"state": torch.ones(1, 3)}) # type: ignore[arg-type] + with pytest.raises(TypeError, match="real floating-point"): + pipeline.compute_loss(torch.ones(1, 3, dtype=torch.int64)) + with pytest.raises(ValueError, match="noise must match"): + pipeline.compute_loss(torch.ones(1, 3), noise=torch.ones(1, 2)) + with pytest.raises(TypeError, match="n must use an integer dtype"): + pipeline.compute_loss( + torch.ones(1, 3), + noise=torch.zeros(1, 3), + n=torch.tensor([0.0]), + k=torch.tensor([0]), + ) + with pytest.raises(TypeError, match="model_kwargs must be a mapping"): + pipeline.sample(torch.ones(1, 3), model_kwargs=[]) # type: ignore[arg-type] + for blocks in ([2, 2], [], [0, 8], [4, 4.0], [2, 2, 2, 4]): + with pytest.raises(ValueError): + pipeline.sample(torch.ones(1, 3), blocks=blocks) + + adapter.bad_student_shape = True + with pytest.raises(ValueError, match="student_all_heads must return shape"): + pipeline.compute_loss( + torch.ones(1, 3), + noise=torch.zeros(1, 3), + n=torch.tensor([0]), + k=torch.tensor([0]), + ) + adapter.bad_fused_dtype = True + with pytest.raises(TypeError, match="student_fused_block must return a real floating-point"): + pipeline.sample(torch.ones(1, 3)) + + +def test_pipeline_freezes_teacher_but_not_student() -> None: + pipeline, _ = _pipeline() + + assert pipeline.teacher.training is False + assert all(not parameter.requires_grad for parameter in pipeline.teacher.parameters()) + assert all(parameter.requires_grad for parameter in pipeline.student.parameters()) + + +def test_only_selected_head_and_shared_backbone_receive_gradients() -> None: + pipeline, _ = _pipeline() + + loss, metrics = pipeline.compute_loss( + torch.tensor([[1.0, -0.5, 0.25]]), + noise=torch.tensor([[-0.25, 2.0, 0.5]]), + n=torch.tensor([0]), + k=torch.tensor([3]), + ) + loss.backward() + + student = pipeline.student + assert student.state_scale.grad is not None + assert torch.isfinite(student.state_scale.grad) + assert student.head_bias.grad is not None + assert torch.count_nonzero(student.head_bias.grad[3]) > 0 + assert torch.count_nonzero(student.head_bias.grad[:3]) == 0 + assert torch.count_nonzero(student.head_bias.grad[4:]) == 0 + assert all(parameter.grad is None for parameter in pipeline.teacher.parameters()) + assert all(not value.requires_grad for value in metrics.values()) diff --git a/tests/unit/torch/fastgen/test_pdd_projection.py b/tests/unit/torch/fastgen/test_pdd_projection.py new file mode 100644 index 00000000000..d6c9d233005 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_projection.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Independent shape, conversion, fusion, and metadata tests for PDD projections.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from modelopt.torch.fastgen import ( + PDDLayerSpec, + PDDOutputProjection, + convert_to_pdd_output_projection, +) + + +class _NestedModel(nn.Module): + def __init__(self, *, bias: bool = True): + super().__init__() + self.transformer = nn.Module() + self.transformer.blocks = nn.Sequential(nn.Identity(), nn.Linear(2, 6, bias=bias)) + + def forward(self, inputs): + return self.transformer.blocks(inputs) + + +def _base_linear(*, bias: bool = True) -> nn.Linear: + linear = nn.Linear(2, 6, bias=bias) + with torch.no_grad(): + linear.weight.copy_(torch.arange(12, dtype=torch.float32).reshape(6, 2) / 7) + if linear.bias is not None: + linear.bias.copy_(torch.arange(6, dtype=torch.float32) / 11) + return linear + + +def _spec(layout: str) -> PDDLayerSpec: + return PDDLayerSpec( + projection_path="transformer.blocks.1", + head_layout=layout, + output_channels=2 if layout == "patch_major" else None, + ) + + +def _decode_heads( + raw: torch.Tensor, + *, + layout: str, + grid_size: int, + base_out_features: int = 6, + output_channels: int = 2, +) -> torch.Tensor: + """Independent raw-layout decoder returning ``[batch, head, base_output]``.""" + if layout == "channel_major": + return raw.reshape(raw.shape[0], grid_size, base_out_features) + patch_factor = base_out_features // output_channels + return ( + raw.reshape(raw.shape[0], patch_factor, grid_size, output_channels) + .permute(0, 2, 1, 3) + .reshape(raw.shape[0], grid_size, base_out_features) + ) + + +def _encode_head_parameters( + head_values: torch.Tensor, + *, + layout: str, + output_channels: int = 2, +) -> torch.Tensor: + """Independent ``[head, base_output, ...]`` encoder for widened storage.""" + if layout == "channel_major": + return head_values.reshape(-1, *head_values.shape[2:]) + grid_size, base_out_features = head_values.shape[:2] + patch_factor = base_out_features // output_channels + trailing_shape = head_values.shape[2:] + return ( + head_values.reshape(grid_size, patch_factor, output_channels, *trailing_shape) + .permute(1, 0, 2, *range(3, head_values.ndim + 1)) + .reshape(-1, *trailing_shape) + ) + + +@pytest.mark.parametrize("layout", ["channel_major", "patch_major"]) +@pytest.mark.parametrize("bias", [False, True]) +def test_from_linear_repeats_every_head_without_mutating_base(layout, bias): + base = _base_linear(bias=bias) + base.eval() + base.weight.requires_grad_(False) + if base.bias is not None: + base.bias.requires_grad_(False) + original = {name: tensor.clone() for name, tensor in base.state_dict().items()} + inputs = torch.tensor([[0.5, -1.0], [2.0, 0.25]]) + expected = base(inputs) + + projection = PDDOutputProjection.from_linear(base, 3, _spec(layout)) + actual_heads = _decode_heads(projection(inputs), layout=layout, grid_size=3) + + torch.testing.assert_close(actual_heads, expected[:, None].expand_as(actual_heads)) + assert projection.out_features == 18 + assert projection.base_out_features == 6 + assert projection.training is False + assert projection.weight.requires_grad is False + assert (projection.bias is None) is (base.bias is None) + for name, tensor in base.state_dict().items(): + assert torch.equal(tensor, original[name]) + + +def test_patch_major_requires_divisible_output_channels(): + spec = PDDLayerSpec("projection", "patch_major", output_channels=4) + with pytest.raises(ValueError, match="must be divisible"): + PDDOutputProjection.from_linear(_base_linear(), 3, spec) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"projection_path": "", "head_layout": "channel_major"}, + {"projection_path": "a..b", "head_layout": "channel_major"}, + {"projection_path": "a", "head_layout": "unknown"}, + {"projection_path": "a", "head_layout": "channel_major", "output_channels": 2}, + {"projection_path": "a", "head_layout": "patch_major"}, + ], +) +def test_layer_spec_rejects_unsupported_layout_metadata(kwargs): + with pytest.raises(ValueError): + PDDLayerSpec(**kwargs) + + +def test_nested_conversion_is_explicit_idempotent_and_conflict_safe(): + model = _NestedModel() + spec = _spec("channel_major") + original = model.get_submodule(spec.projection_path) + + projection = convert_to_pdd_output_projection(model, spec, grid_size=3) + repeated = convert_to_pdd_output_projection(model, spec, grid_size=3) + + assert projection is repeated + assert model.get_submodule(spec.projection_path) is projection + assert original is not projection + with pytest.raises(ValueError, match="incompatible"): + convert_to_pdd_output_projection(model, spec, grid_size=4) + with pytest.raises(ValueError, match="incompatible"): + convert_to_pdd_output_projection(model, _spec("patch_major"), grid_size=3) + assert model.get_submodule(spec.projection_path) is projection + + +@pytest.mark.parametrize("layout", ["channel_major", "patch_major"]) +@pytest.mark.parametrize("bias", [False, True]) +def test_fused_forward_matches_independent_weighted_head_sum(layout, bias): + projection = PDDOutputProjection.from_linear(_base_linear(bias=bias), 3, _spec(layout)) + head_weights = torch.arange(36, dtype=torch.float32).reshape(3, 6, 2) / 13 + head_bias = torch.arange(18, dtype=torch.float32).reshape(3, 6) / 17 + with torch.no_grad(): + projection.weight.copy_(_encode_head_parameters(head_weights, layout=layout)) + if projection.bias is not None: + projection.bias.copy_(_encode_head_parameters(head_bias, layout=layout)) + + inputs = torch.tensor([[0.25, -1.0], [2.0, 0.5]]) + original_inputs = inputs.clone() + grid = torch.tensor([1.0, 0.8, 0.3, 0.0], dtype=torch.float64) + original_grid = grid.clone() + weight_id = id(projection.weight) + weight_pointer = projection.weight.data_ptr() + state_keys = tuple(projection.state_dict()) + explicit_heads = torch.stack( + [ + F.linear(inputs, head_weights[index], head_bias[index] if bias else None) + for index in range(3) + ], + dim=1, + ) + coefficients = torch.tensor([0.5 / 0.8, 0.3 / 0.8]) + expected = torch.einsum("n,bno->bo", coefficients, explicit_heads[:, 1:3]) + + actual = projection(inputs, fusion=(1, 3, grid)) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + _decode_heads(projection(inputs), layout=layout, grid_size=3), explicit_heads + ) + assert id(projection.weight) == weight_id + assert projection.weight.data_ptr() == weight_pointer + assert tuple(projection.state_dict()) == state_keys + assert torch.equal(inputs, original_inputs) + assert torch.equal(grid, original_grid) + + +def test_fused_forward_does_not_change_subsequent_projection_calls(): + projection = PDDOutputProjection.from_linear(_base_linear(), 3, _spec("channel_major")) + inputs = torch.tensor([[1.0, -0.5]]) + grid = torch.tensor([1.0, 0.7, 0.2, 0.0]) + + assert projection(inputs, fusion=(1, 3, grid)).shape == (1, 6) + assert projection(inputs).shape == (1, 18) + + +@pytest.mark.parametrize( + ("start", "end", "message"), + [(-1, 2, "0 <= start"), (1, 1, "0 <= start"), (1, 4, "0 <= start")], +) +def test_fused_forward_rejects_invalid_blocks(start, end, message): + projection = PDDOutputProjection.from_linear(_base_linear(), 3, _spec("channel_major")) + grid = torch.tensor([1.0, 0.7, 0.2, 0.0]) + with pytest.raises(ValueError, match=message): + projection(torch.zeros(1, 2), fusion=(start, end, grid)) + + +def _state_clone(state): + return {name: tensor.clone() for name, tensor in state.items()} + + +def _assert_state_unchanged(state, original): + assert state.keys() == original.keys() + for name, tensor in state.items(): + assert torch.equal(tensor, original[name]) + + +def test_base_and_widened_checkpoint_load_order_is_strict_and_nonmutating(): + spec = _spec("channel_major") + base_model = _NestedModel() + base_state = _state_clone(base_model.state_dict()) + base_original = _state_clone(base_state) + + student = _NestedModel() + student.load_state_dict(base_state, strict=True) + projection = convert_to_pdd_output_projection(student, spec, grid_size=3) + _assert_state_unchanged(base_state, base_original) + + widened_state = _state_clone(student.state_dict()) + widened_original = _state_clone(widened_state) + restored = _NestedModel() + restored_projection = convert_to_pdd_output_projection(restored, spec, grid_size=3) + restored.load_state_dict(widened_state, strict=True) + _assert_state_unchanged(widened_state, widened_original) + torch.testing.assert_close(restored_projection.weight, projection.weight) + + with pytest.raises(RuntimeError, match="size mismatch"): + student.load_state_dict(base_state, strict=True) + _assert_state_unchanged(base_state, base_original) + with pytest.raises(RuntimeError, match="size mismatch"): + _NestedModel().load_state_dict(widened_state, strict=True) + _assert_state_unchanged(widened_state, widened_original) diff --git a/tests/unit/torch/fastgen/test_pdd_reference_math.py b/tests/unit/torch/fastgen/test_pdd_reference_math.py new file mode 100644 index 00000000000..e467fe35555 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_reference_math.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""First-principles reference math for Parallel Decoding Distillation (PDD). + +These deliberately simple CPU oracles are independent of the future PDD implementation. +Production grid, integration, and projection-fusion code should be tested against them. +""" + +from __future__ import annotations + +import pytest +import torch + +from modelopt.torch.fastgen.flow_matching import ( + add_noise, + fusion_coefficients, + integrate_interval_velocities, + make_shifted_flow_grid, +) + + +def _reference_shifted_grid(grid_size: int, shift: float, max_t: float) -> torch.Tensor: + """Reproduce the frozen FastGen float64 schedule without production helpers.""" + unshifted = torch.linspace(max_t, 0.0, grid_size + 1, dtype=torch.float64) + unshifted = unshifted.clamp(max=max_t) + shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted) + return shifted.clamp(max=max_t) + + +def _reference_rf_forward_process( + data: torch.Tensor, + noise: torch.Tensor, + time: torch.Tensor, +) -> torch.Tensor: + """Reproduce FastGen's float64 RF forward process without production helpers.""" + original_dtype = data.dtype + data_64 = data.to(torch.float64) + noise_64 = noise.to(torch.float64) + time_64 = time.to(torch.float64) + while time_64.ndim < data_64.ndim: + time_64 = time_64.unsqueeze(-1) + state_64 = data_64 * (1.0 - time_64) + noise_64 * time_64 + return state_64.to(original_dtype) + + +def _reference_integrate( + state: torch.Tensor, + velocities: torch.Tensor, + grid: torch.Tensor, + start: int, + end: int, +) -> torch.Tensor: + """Integrate exactly the half-open interval-head block ``[start, end)``.""" + result = state.to(torch.float64).clone() + for index in range(start, end): + result += (grid[index + 1] - grid[index]) * velocities[index].to(torch.float64) + return result + + +def test_production_shifted_grid_matches_independent_oracle(): + grid = make_shifted_flow_grid(grid_size=128, shift=5.0, max_t=0.999) + oracle = _reference_shifted_grid(grid_size=128, shift=5.0, max_t=0.999).to(torch.float32) + + assert grid.dtype == torch.float32 + assert grid.shape == (129,) + assert grid[0].item() == 0.9990000128746033 + assert grid[-1] == 0.0 + assert torch.all(torch.diff(grid) < 0) + torch.testing.assert_close(grid, oracle, rtol=0, atol=0) + assert make_shifted_flow_grid(128, 5.0, max_t=0.999, dtype=torch.bfloat16).dtype == ( + torch.float32 + ) + + direct_fp32 = torch.linspace(0.999, 0.0, 129, dtype=torch.float32) + upper = torch.tensor(0.999, dtype=torch.float32) + if upper.item() > 0.999: + upper = torch.nextafter(upper, torch.tensor(float("-inf"))) + direct_fp32 = direct_fp32.clamp(max=upper) + direct_fp32 = (5.0 * direct_fp32 / (1.0 + 4.0 * direct_fp32)).clamp(max=upper) + assert torch.count_nonzero(grid != direct_fp32).item() == 52 + + +def test_production_rf_forward_process_matches_float64_intermediate_oracle(): + data = torch.tensor( + [[-0.2654421329498291, 0.5161616802215576, -0.7285917401313782]], + dtype=torch.float32, + ) + noise = torch.tensor( + [[0.3856363296508789, -0.34849217534065247, -0.11881951987743378]], + dtype=torch.float32, + ) + grid = make_shifted_flow_grid(128, 5.0, max_t=0.999) + time = grid[:1] + original = (data.clone(), noise.clone(), time.clone(), grid.clone()) + + expected = _reference_rf_forward_process(data, noise, time) + actual = add_noise(data, noise, time) + stale_direct = (1.0 - time[:, None]) * data + time[:, None] * noise + + assert time.item() == 0.9990000128746033 + assert torch.equal( + stale_direct, + torch.tensor([[0.3849852383136749, -0.34762755036354065, -0.11942929029464722]]), + ) + assert torch.equal( + expected, + torch.tensor([[0.3849852681159973, -0.34762752056121826, -0.11942928284406662]]), + ) + assert not torch.equal(stale_direct, expected) + assert torch.equal(actual, expected) + assert actual.dtype == torch.float32 + for value, unchanged in zip((data, noise, time, grid), original): + assert torch.equal(value, unchanged) + + +@pytest.mark.parametrize( + ("args", "kwargs", "error"), + [ + ((0, 5.0), {"max_t": 0.999}, ValueError), + ((128, 0.5), {"max_t": 0.999}, ValueError), + ((4, 5.0), {"max_t": 1}, TypeError), + ((4, 5.0), {"max_t": 0.0}, ValueError), + ((4, 5.0), {"max_t": 0.999, "dtype": torch.int64}, TypeError), + ], +) +def test_production_grid_rejects_invalid_inputs(args, kwargs, error): + with pytest.raises(error): + make_shifted_flow_grid(*args, **kwargs) + + +def test_production_half_open_integration_matches_independent_oracle_per_sample(): + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999, dtype=torch.float64) + state = torch.tensor([[3.0, -2.0], [1.0, 4.0]], dtype=torch.bfloat16) + velocities = torch.tensor( + [ + [[1000.0, 1000.0], [2.0, -1.0], [-3.0, 4.0], [-1000.0, -1000.0]], + [[1.0, 2.0], [-2.0, 3.0], [4.0, -5.0], [6.0, 7.0]], + ], + dtype=torch.bfloat16, + ) + starts = torch.tensor([1, 2]) + ends = torch.tensor([3, 2]) + + actual = integrate_interval_velocities(state, velocities, grid, starts, ends) + expected = torch.stack( + [ + _reference_integrate(state[0], velocities[0], grid, start=1, end=3), + _reference_integrate(state[1], velocities[1], grid, start=2, end=2), + ] + ) + + assert actual.dtype == torch.float64 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +def test_production_integration_promotes_bfloat16_math_to_float32(): + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999) + state = torch.zeros(1, 2, dtype=torch.bfloat16) + velocities = torch.ones(1, 4, 2, dtype=torch.bfloat16) + + result = integrate_interval_velocities(state, velocities, grid, start=0, end=4) + + assert result.dtype == torch.float32 + torch.testing.assert_close(result, torch.full((1, 2), -0.999)) + + +def test_production_helpers_reject_invalid_blocks(): + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999) + state = torch.zeros(1, 2) + velocities = torch.ones(1, 4, 2) + + with pytest.raises(ValueError, match="0 <= start <= end <= 4"): + integrate_interval_velocities(state, velocities, grid, start=3, end=5) + with pytest.raises(ValueError, match="0 <= start < end <= 4"): + fusion_coefficients(grid, start=2, end=2) + + +def test_production_fusion_coefficients_match_independent_oracle(): + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999, dtype=torch.float64) + actual = fusion_coefficients(grid, start=1, end=4) + oracle_grid = _reference_shifted_grid(4, 5.0, 0.999) + expected = torch.stack( + [ + (oracle_grid[index + 1] - oracle_grid[index]) / (oracle_grid[4] - oracle_grid[1]) + for index in range(1, 4) + ] + ) + + assert torch.all(actual > 0) + torch.testing.assert_close(actual.sum(), torch.tensor(1.0, dtype=torch.float64)) + torch.testing.assert_close(actual, expected, rtol=1e-14, atol=1e-14) diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py new file mode 100644 index 00000000000..625816cac53 --- /dev/null +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -0,0 +1,743 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hermetic Qwen-like tests for the ModelOpt PDD adapter and conversion.""" + +from __future__ import annotations + +import copy +from types import MethodType, SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +import modelopt.torch.fastgen.plugins.qwen_image_pdd as qwen_image_pdd_plugin +from modelopt.torch.fastgen import PDDConfig, PDDOutputProjection, PDDPipeline +from modelopt.torch.fastgen.flow_matching import fusion_coefficients +from modelopt.torch.fastgen.plugins import QwenImagePDDAdapter +from modelopt.torch.fastgen.plugins.qwen_image import build_img_shapes, pack_latents, unpack_latents +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QWEN_IMAGE_PDD_EXECUTION, + convert_qwen_image_to_pdd, + enable_qwen_image_pdd_forward, + freeze_qwen_image_pdd_unused_parameters, + require_qwen_image_pdd_forward, +) + + +class _QwenImageTestDouble(nn.Module): + """Explicit unit-test scope for adapter protocol doubles.""" + + +@pytest.fixture(autouse=True) +def _allow_qwen_image_test_doubles(monkeypatch): + require_production_forward = qwen_image_pdd_plugin.require_qwen_image_pdd_forward + + def require_forward(model: nn.Module) -> str: + if isinstance(model, _QwenImageTestDouble): + if ( + getattr(model, "_modelopt_qwen_image_pdd_execution", None) + != QWEN_IMAGE_PDD_EXECUTION + ): + raise RuntimeError("Qwen-Image PDD requires its masked joint-attention forward.") + return QWEN_IMAGE_PDD_EXECUTION + return require_production_forward(model) + + monkeypatch.setattr(qwen_image_pdd_plugin, "require_qwen_image_pdd_forward", require_forward) + + +class _TinyQwenTransformer(_QwenImageTestDouble): + """Qwen-shaped packed transformer with a real registered final linear.""" + + def __init__(self, *, packed_channels: int = 4, hidden_width: int = 5) -> None: + super().__init__() + self.config = SimpleNamespace(guidance_embeds=False) + self._modelopt_qwen_image_pdd_execution = QWEN_IMAGE_PDD_EXECUTION + self.backbone = nn.Linear(packed_channels, hidden_width, dtype=torch.bfloat16) + self.proj_out = nn.Linear(hidden_width, packed_channels, dtype=torch.bfloat16) + self.calls: list[dict[str, object]] = [] + + def forward( + self, + *, + hidden_states, + timestep, + encoder_hidden_states, + encoder_hidden_states_mask, + img_shapes, + max_txt_seq_len=None, + **kwargs, + ): + if max_txt_seq_len is None: + max_txt_seq_len = encoder_hidden_states.shape[1] + fusion = kwargs.pop("_pdd_fusion", None) + condition_value = encoder_hidden_states.mean(dim=(1, 2), keepdim=True) + condition_value = condition_value + 0.01 * encoder_hidden_states_mask.sum( + dim=1, keepdim=True + ).unsqueeze(-1) + hidden = torch.tanh(self.backbone(hidden_states)) + hidden = hidden + condition_value.to(hidden.dtype) + hidden = hidden + (0.1 * timestep[:, None, None]).to(hidden.dtype) + if fusion is None: + output = self.proj_out(hidden) + else: + assert isinstance(self.proj_out, PDDOutputProjection) + output = self.proj_out(hidden, fusion=fusion) + self.calls.append( + { + "hidden_states": hidden_states.detach().clone(), + "timestep": timestep.detach().clone(), + "encoder_hidden_states": encoder_hidden_states.detach().clone(), + "encoder_hidden_states_mask": encoder_hidden_states_mask.detach().clone(), + "img_shapes": img_shapes, + "max_txt_seq_len": max_txt_seq_len, + "projection_input": hidden.detach().clone(), + "kwargs": kwargs, + "output": output.detach().clone(), + } + ) + return (output,) + + +def _config(*, guidance_scale: float | None = 4.0, grid_size: int = 4) -> PDDConfig: + return PDDConfig( + grid_size=grid_size, + grid_max_t=0.999, + flow_shift=5.0, + block_size_min=1, + block_size_max=grid_size, + inference_blocks=[2, 2] if grid_size == 4 else [grid_size], + guidance_scale=guidance_scale, + ) + + +def _inputs(batch_size: int = 2): + torch.manual_seed(7) + state = torch.randn(batch_size, 1, 4, 4) + time = torch.tensor([0.875, 0.25])[:batch_size] + embeddings = torch.randn(batch_size, 3, 2, dtype=torch.bfloat16) + mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long)[:batch_size] + embeddings = embeddings * mask.unsqueeze(-1) + negative_embeddings = torch.randn(batch_size, 3, 2, dtype=torch.bfloat16) + negative_mask = torch.tensor([[1, 0, 0], [1, 1, 1]], dtype=torch.long)[:batch_size] + negative_embeddings = negative_embeddings * negative_mask.unsqueeze(-1) + return state, time, (embeddings, mask), (negative_embeddings, negative_mask) + + +def _pack_oracle(latents: torch.Tensor) -> torch.Tensor: + batch, channels, height, width = latents.shape + return ( + latents.reshape(batch, channels, height // 2, 2, width // 2, 2) + .permute(0, 2, 4, 1, 3, 5) + .reshape(batch, (height // 2) * (width // 2), channels * 4) + ) + + +def _unpack_oracle(packed: torch.Tensor, height: int, width: int) -> torch.Tensor: + batch, _patches, packed_channels = packed.shape + channels = packed_channels // 4 + return ( + packed.reshape(batch, height // 2, width // 2, channels, 2, 2) + .permute(0, 3, 1, 4, 2, 5) + .reshape(batch, channels, height, width) + ) + + +def _mr210_rollout_oracle( + state: torch.Tensor, + heads: torch.Tensor, + grid: torch.Tensor, + n: torch.Tensor, + k: torch.Tensor, +) -> torch.Tensor: + interval_ids = torch.arange(grid.numel() - 1, device=state.device) + velocity_mask = (interval_ids[None] >= n[:, None]) & (interval_ids[None] < k[:, None]) + weighted_intervals = velocity_mask.to(torch.float32) * torch.diff(grid.float())[None] + return state.float() + torch.einsum("bn,bn...->b...", weighted_intervals, heads.float()) + + +def test_all_head_training_accepts_a_serialized_widened_linear() -> None: + student = _TinyQwenTransformer() + config = _config() + converted = convert_qwen_image_to_pdd(student, config) + serialized = nn.Linear( + converted.in_features, + converted.out_features, + bias=converted.bias is not None, + dtype=converted.weight.dtype, + ) + serialized.load_state_dict(converted.state_dict()) + student.proj_out = serialized + state, time, condition, _ = _inputs(batch_size=1) + + heads = QwenImagePDDAdapter(config).student_all_heads( + student, + state, + time, + condition=condition, + ) + heads.float().square().mean().backward() + + assert heads.shape == (1, 4, 1, 4, 4) + assert serialized.weight.grad is not None + + +def test_fused_student_matches_explicit_packed_weight_fusion() -> None: + student = _TinyQwenTransformer() + config = _config() + projection = convert_qwen_image_to_pdd(student, config) + generator = torch.Generator().manual_seed(91) + with torch.no_grad(): + projection.weight.copy_(torch.randn(projection.weight.shape, generator=generator) / 3) + projection.bias.copy_(torch.randn(projection.bias.shape, generator=generator) / 7) + adapter = QwenImagePDDAdapter(config) + state, time, condition, _ = _inputs() + grid = torch.tensor([1.0, 0.85, 0.55, 0.2, 0.0]) + + actual = adapter.student_fused_block( + student, + state, + time, + start=1, + end=4, + grid=grid, + condition=condition, + ) + coefficients = fusion_coefficients(grid, 1, 4).float() + head_weights = projection.weight.reshape(4, 4, 5) + head_bias = projection.bias.reshape(4, 4) + fused_weight = torch.einsum("n,n...->...", coefficients, head_weights[1:4].float()).to( + torch.bfloat16 + ) + fused_bias = torch.einsum("n,n...->...", coefficients, head_bias[1:4].float()).to( + torch.bfloat16 + ) + expected_packed = F.linear(student.calls[0]["projection_input"], fused_weight, fused_bias) + expected = unpack_latents(expected_packed, 4, 4) + + torch.testing.assert_close(actual, expected, rtol=2e-6, atol=2e-6) + assert len(student.calls) == 1 + assert student.proj_out is projection + assert student.proj_out(projection.weight.new_zeros(1, 5)).shape[-1] == 16 + + +def test_teacher_cfg_zero_guided_norm_uses_qwen_clamp() -> None: + class ZeroGuidedTeacher(_QwenImageTestDouble): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(guidance_embeds=False) + self._modelopt_qwen_image_pdd_execution = QWEN_IMAGE_PDD_EXECUTION + self.anchor = nn.Parameter(torch.zeros((), dtype=torch.bfloat16), requires_grad=False) + self.calls = 0 + + def forward(self, *, hidden_states, **_kwargs): + self.calls += 1 + value = 3.0 if self.calls == 1 else 4.0 + return (torch.full_like(hidden_states, value, dtype=torch.bfloat16),) + + teacher = ZeroGuidedTeacher() + state, time, condition, negative_condition = _inputs(batch_size=1) + actual = QwenImagePDDAdapter(_config(guidance_scale=4.0)).teacher_velocity( + teacher, + state, + time, + condition=condition, + negative_condition=negative_condition, + ) + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, torch.zeros_like(actual), rtol=0, atol=0) + + +def test_guidance_disabled_teacher_is_one_conditional_call_without_negative_condition() -> None: + teacher = _TinyQwenTransformer() + config = _config(guidance_scale=None) + adapter = QwenImagePDDAdapter(config) + state, time, condition, _ = _inputs() + + actual = adapter.teacher_velocity( + teacher, + state, + time, + condition=condition, + ) + + assert len(teacher.calls) == 1 + expected = unpack_latents(teacher.calls[0]["output"], 4, 4) + torch.testing.assert_close(actual, expected) + + +def _tiny_diffusers_qwen(): + diffusers = pytest.importorskip("diffusers") + return diffusers.QwenImageTransformer2DModel( + patch_size=2, + in_channels=8, + out_channels=2, + num_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=12, + guidance_embeds=False, + axes_dims_rope=(2, 2, 4), + ) + + +def test_mr210_freezes_exactly_the_structurally_unused_parameters() -> None: + student = _tiny_diffusers_qwen() + + frozen_names = freeze_qwen_image_pdd_unused_parameters(student) + + assert set(frozen_names) == { + "transformer_blocks.0.attn.to_add_out.weight", + "transformer_blocks.0.attn.to_add_out.bias", + "transformer_blocks.0.txt_mlp.net.0.proj.weight", + "transformer_blocks.0.txt_mlp.net.0.proj.bias", + "transformer_blocks.0.txt_mlp.net.2.weight", + "transformer_blocks.0.txt_mlp.net.2.bias", + } + assert { + name for name, parameter in student.named_parameters() if not parameter.requires_grad + } == set(frozen_names) + optimizer = torch.optim.AdamW( + [parameter for parameter in student.parameters() if parameter.requires_grad] + ) + optimized_parameter_ids = { + id(parameter) for group in optimizer.param_groups for parameter in group["params"] + } + assert all( + id(student.get_parameter(name)) not in optimized_parameter_ids for name in frozen_names + ) + + +def _mr210_qwen_forward_oracle( + model: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor, + timestep: torch.Tensor, + img_shapes: list, + max_txt_seq_len: int, +) -> torch.Tensor: + """Test-local MR210 operation order; intentionally independent of production binding.""" + encoder_hidden_states_mask = encoder_hidden_states_mask.to(torch.bool) + hidden_states = model.img_in(hidden_states) + encoder_hidden_states = model.txt_in(model.txt_norm(encoder_hidden_states)) + temb = model.time_text_embed(timestep, hidden_states) + image_rotary_emb = model.pos_embed( + img_shapes, + max_txt_seq_len=max_txt_seq_len, + device=hidden_states.device, + ) + for block in model.transformer_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=encoder_hidden_states_mask, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + return model.proj_out(model.norm_out(hidden_states, temb)) + + +def test_mr210_real_qwen_loss_and_backward_match_independent_graph() -> None: + class CapturingAdapter(QwenImagePDDAdapter): + def _call_packed(self, *args, **kwargs): + with torch.autocast(device_type="cpu", dtype=torch.bfloat16): + return super()._call_packed(*args, **kwargs) + + def student_all_heads(self, *args, **kwargs): + value = super().student_all_heads(*args, **kwargs) + self.captured_heads = value.detach().clone() + return value + + def teacher_velocity(self, model, state, time, **kwargs): + self.captured_teacher_state = state.detach().clone() + value = super().teacher_velocity(model, state, time, **kwargs) + self.captured_teacher = value.detach().clone() + return value + + torch.manual_seed(20260716) + base = _tiny_diffusers_qwen().eval() + actual_student = enable_qwen_image_pdd_forward(copy.deepcopy(base)) + actual_student.enable_gradient_checkpointing() + actual_teacher = copy.deepcopy(actual_student).eval().requires_grad_(False) + oracle_student = copy.deepcopy(base) + oracle_teacher = copy.deepcopy(base).eval().requires_grad_(False) + config = _config(guidance_scale=4.0) + convert_qwen_image_to_pdd(actual_student, config) + convert_qwen_image_to_pdd(oracle_student, config) + adapter = CapturingAdapter(config, compute_dtype=torch.bfloat16) + pipeline = PDDPipeline(actual_student, actual_teacher, config, adapter) + + generator = torch.Generator().manual_seed(20260716) + data = torch.randn(1, 2, 4, 4, generator=generator) + noise = torch.randn(1, 2, 4, 4, generator=generator) + condition = ( + torch.randn(1, 3, 12, generator=generator).to(torch.bfloat16), + torch.tensor([[1, 1, 1]], dtype=torch.long), + ) + negative_condition = ( + torch.randn(1, 2, 12, generator=generator).to(torch.bfloat16), + torch.tensor([[1, 1]], dtype=torch.long), + ) + n = torch.tensor([1], dtype=torch.long) + k = torch.tensor([3], dtype=torch.long) + + actual_loss, _ = pipeline.compute_loss( + data, + noise=noise, + condition=condition, + negative_condition=negative_condition, + n=n, + k=k, + ) + actual_loss.backward() + + unshifted = torch.linspace(0.999, 0.0, 5, dtype=torch.float64) + grid = (5.0 * unshifted / (1.0 + 4.0 * unshifted)).clamp_max(0.999).float() + time_n = grid[n] + broadcast_time = time_n.to(torch.float64).reshape(1, 1, 1, 1) + x_n = ( + data.float().to(torch.float64) * (1.0 - broadcast_time) + + noise.float().to(torch.float64) * broadcast_time + ).float() + + def oracle_forward(model, state, time, current_condition): + embeddings, mask = current_condition + with torch.autocast(device_type="cpu", dtype=torch.bfloat16): + return _mr210_qwen_forward_oracle( + model, + hidden_states=_pack_oracle(state).to(torch.bfloat16), + encoder_hidden_states=embeddings, + encoder_hidden_states_mask=mask, + timestep=time, + img_shapes=build_img_shapes(state.shape[0], state.shape[2], state.shape[3]), + max_txt_seq_len=int(mask.sum(dim=1).max().to(torch.int32).item()), + ) + + packed_heads = oracle_forward(oracle_student, x_n, time_n, condition) + batch, patches, _features = packed_heads.shape + packed_heads = packed_heads.reshape(batch, patches, 4, 8).permute(0, 2, 1, 3) + oracle_heads = _unpack_oracle( + packed_heads.reshape(4, patches, 8), + 4, + 4, + ).reshape(1, 4, 2, 4, 4) + oracle_heads_fp32 = oracle_heads.float() + with torch.no_grad(): + x_bar_k = _mr210_rollout_oracle(x_n, oracle_heads_fp32, grid, n, k) + student_target = oracle_heads_fp32[:, int(k.item())] + time_k = grid[k] + conditional = oracle_forward(oracle_teacher, x_bar_k, time_k, condition) + unconditional = oracle_forward(oracle_teacher, x_bar_k, time_k, negative_condition) + guided_low_precision = conditional + 3.0 * (conditional - unconditional) + conditional_fp32 = conditional.float() + guided_fp32 = guided_low_precision.float() + teacher_target_packed = ( + guided_fp32 + * ( + torch.linalg.vector_norm(conditional_fp32, dim=-1, keepdim=True) + / torch.linalg.vector_norm(guided_fp32, dim=-1, keepdim=True).clamp_min(1e-5) + ) + ).to(torch.bfloat16) + teacher_target_low_precision = _unpack_oracle(teacher_target_packed, 4, 4) + teacher_target = teacher_target_low_precision.float().detach() + oracle_loss = (student_target - teacher_target).square().mean() + oracle_loss.backward() + + torch.testing.assert_close( + adapter.captured_heads[:, int(k.item())], + oracle_heads[:, int(k.item())], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + adapter.captured_teacher, teacher_target_low_precision, rtol=0, atol=0 + ) + torch.testing.assert_close(adapter.captured_teacher_state, x_bar_k, rtol=1e-6, atol=1e-7) + torch.testing.assert_close(actual_loss, oracle_loss, rtol=1e-6, atol=1e-7) + for actual_parameter, oracle_parameter in ( + (actual_student.proj_out.weight, oracle_student.proj_out.weight), + (actual_student.img_in.weight, oracle_student.img_in.weight), + ): + assert actual_parameter.grad is not None and oracle_parameter.grad is not None + assert actual_parameter.grad.dtype == torch.float32 + assert oracle_parameter.grad.dtype == torch.float32 + torch.testing.assert_close( + actual_parameter.grad, + oracle_parameter.grad, + rtol=1e-6, + atol=1e-7, + ) + + +def test_qwen_pdd_forward_binding_preserves_root_state_and_deepcopy() -> None: + source = _tiny_diffusers_qwen().eval() + source_type = type(source) + source_state = {name: value.detach().clone() for name, value in source.state_dict().items()} + + adopted = enable_qwen_image_pdd_forward(source) + + assert adopted is source + assert type(adopted) is source_type + assert enable_qwen_image_pdd_forward(adopted) is adopted + assert require_qwen_image_pdd_forward(adopted) == QWEN_IMAGE_PDD_EXECUTION + for name, value in adopted.state_dict().items(): + torch.testing.assert_close(value, source_state[name], rtol=0, atol=0) + + round_trip = copy.deepcopy(adopted) + assert round_trip.forward.__self__ is round_trip + require_qwen_image_pdd_forward(round_trip) + + conflicting = _tiny_diffusers_qwen() + conflicting.forward = MethodType(lambda self, **_kwargs: self, conflicting) + with pytest.raises(RuntimeError, match="instance-level forward override"): + enable_qwen_image_pdd_forward(conflicting) + + +def test_mr210_qwen_conversion_preserves_every_initialized_head() -> None: + base = _tiny_diffusers_qwen().eval().to(torch.bfloat16) + student = copy.deepcopy(base) + student = enable_qwen_image_pdd_forward(student) + config = _config() + generator = torch.Generator().manual_seed(20260715) + state = torch.randn(2, 2, 4, 4, generator=generator) + time = torch.tensor([0.875, 0.25], dtype=torch.float32) + embeddings = torch.randn(2, 3, 12, generator=generator).to(torch.bfloat16) + mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long) + model_kwargs = { + "hidden_states": pack_latents(state).to(torch.bfloat16), + "timestep": time, + "encoder_hidden_states": embeddings, + "encoder_hidden_states_mask": mask, + "img_shapes": build_img_shapes(2, 4, 4), + "max_txt_seq_len": 3, + } + + with torch.no_grad(): + expected = unpack_latents( + _mr210_qwen_forward_oracle(base, **model_kwargs), + 4, + 4, + ) + convert_qwen_image_to_pdd(student, config) + actual = QwenImagePDDAdapter(config).student_all_heads( + student, + state, + time, + condition=(embeddings, mask), + ) + + torch.testing.assert_close(actual, expected[:, None].expand_as(actual), rtol=0, atol=0) + + +def test_mr210_joint_mask_ignores_padded_token_values() -> None: + canonical = _tiny_diffusers_qwen().eval().to(torch.bfloat16) + student = copy.deepcopy(canonical) + student = enable_qwen_image_pdd_forward(student) + config = _config() + convert_qwen_image_to_pdd(student, config) + adapter = QwenImagePDDAdapter(config) + generator = torch.Generator().manual_seed(20260715) + state = torch.randn(2, 2, 4, 4, generator=generator) + time = torch.tensor([0.875, 0.25], dtype=torch.float32) + encoder_hidden_states = torch.randn(2, 4, 12, generator=generator).to(torch.bfloat16) + mask = torch.tensor([[1, 1, 0, 0], [1, 0, 1, 0]], dtype=torch.long) + poisoned = encoder_hidden_states.clone() + poisoned[~mask.bool()] = ( + torch.randn( + poisoned[~mask.bool()].shape, + generator=generator, + ) + * 100 + ).to(torch.bfloat16) + canonical_kwargs = { + "hidden_states": pack_latents(state).to(torch.bfloat16), + "timestep": time, + "encoder_hidden_states_mask": mask, + "img_shapes": build_img_shapes(2, 4, 4), + "guidance": None, + "return_dict": False, + } + captured_masks: list[torch.Tensor] = [] + + def capture_block_mask(_module, _args, kwargs): + assert kwargs.get("joint_attention_kwargs") is None + captured_masks.append(kwargs["encoder_hidden_states_mask"].detach().clone()) + + hook = student.transformer_blocks[0].register_forward_pre_hook( + capture_block_mask, + with_kwargs=True, + ) + with torch.no_grad(): + canonical_baseline = canonical( + encoder_hidden_states=encoder_hidden_states, + **canonical_kwargs, + )[0] + canonical_poisoned = canonical( + encoder_hidden_states=poisoned, + **canonical_kwargs, + )[0] + strict_baseline = adapter.student_all_heads( + student, + state, + time, + condition=(encoder_hidden_states, mask), + ) + strict_poisoned = adapter.student_all_heads( + student, + state, + time, + condition=(poisoned, mask), + ) + hook.remove() + + torch.testing.assert_close(canonical_poisoned, canonical_baseline, rtol=0, atol=0) + torch.testing.assert_close(strict_poisoned, strict_baseline, rtol=0, atol=0) + assert len(captured_masks) == 2 + assert all(torch.equal(captured, mask.bool()) for captured in captured_masks) + + +def test_mr210_preserves_diffusers_output_contract() -> None: + student = enable_qwen_image_pdd_forward(_tiny_diffusers_qwen().eval().to(torch.bfloat16)) + generator = torch.Generator().manual_seed(20260716) + kwargs = { + "hidden_states": pack_latents(torch.randn(2, 2, 4, 4, generator=generator)).to( + torch.bfloat16 + ), + "encoder_hidden_states": torch.randn(2, 3, 12, generator=generator).to(torch.bfloat16), + "encoder_hidden_states_mask": torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long), + "timestep": torch.tensor([0.875, 0.25], dtype=torch.float32), + "img_shapes": build_img_shapes(2, 4, 4), + "guidance": None, + } + + with torch.no_grad(): + tuple_output = student(**kwargs, return_dict=False) + model_output = student(**kwargs, return_dict=True) + + assert isinstance(tuple_output, tuple) and len(tuple_output) == 1 + assert hasattr(model_output, "sample") + torch.testing.assert_close(model_output.sample, tuple_output[0], rtol=0, atol=0) + + +def test_mr210_time_embed_receives_fp32_grid_value() -> None: + student = enable_qwen_image_pdd_forward(_tiny_diffusers_qwen().eval().to(torch.bfloat16)) + config = _config() + convert_qwen_image_to_pdd(student, config) + captured: list[torch.Tensor] = [] + + def capture_time(_module, args): + captured.append(args[0].detach().clone()) + + hook = student.time_text_embed.register_forward_pre_hook(capture_time) + generator = torch.Generator().manual_seed(20260715) + state = torch.randn(1, 2, 4, 4, generator=generator) + time = torch.tensor([0.999], dtype=torch.float32) + embeddings = torch.randn(1, 3, 12, generator=generator).to(torch.bfloat16) + mask = torch.ones(1, 3, dtype=torch.long) + with torch.no_grad(): + QwenImagePDDAdapter(config).student_all_heads( + student, + state, + time, + condition=(embeddings, mask), + ) + hook.remove() + + assert len(captured) == 1 + assert captured[0].dtype == torch.float32 + torch.testing.assert_close(captured[0], time, rtol=0, atol=0) + assert captured[0].item() != time.to(torch.bfloat16).float().item() + + +def test_mr210_qwen_teacher_cfg_matches_per_token_reference() -> None: + teacher = enable_qwen_image_pdd_forward(_tiny_diffusers_qwen().eval().to(torch.bfloat16)) + config = _config(guidance_scale=4.0) + adapter = QwenImagePDDAdapter(config) + generator = torch.Generator().manual_seed(20260716) + state = torch.randn(2, 2, 4, 4, generator=generator) + time = torch.tensor([0.75, 0.125], dtype=torch.float32) + condition = ( + torch.randn(2, 3, 12, generator=generator).to(torch.bfloat16), + torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long), + ) + negative_condition = ( + torch.randn(2, 2, 12, generator=generator).to(torch.bfloat16), + torch.tensor([[1, 0], [1, 1]], dtype=torch.long), + ) + + def direct_packed(current_condition): + embeddings, mask = current_condition + return teacher( + hidden_states=pack_latents(state).to(torch.bfloat16), + timestep=time, + encoder_hidden_states=embeddings, + encoder_hidden_states_mask=mask, + img_shapes=build_img_shapes(2, 4, 4), + return_dict=False, + )[0] + + with torch.no_grad(): + conditional = direct_packed(condition) + unconditional = direct_packed(negative_condition) + guided_low_precision = conditional + 3.0 * (conditional - unconditional) + conditional_fp32 = conditional.float() + guided_fp32 = guided_low_precision.float() + factor = torch.linalg.vector_norm( + conditional_fp32, dim=-1, keepdim=True + ) / torch.linalg.vector_norm(guided_fp32, dim=-1, keepdim=True).clamp_min(1e-5) + expected = unpack_latents((guided_fp32 * factor).to(torch.bfloat16), 4, 4) + global_factor = torch.linalg.vector_norm( + conditional_fp32, dim=(1, 2), keepdim=True + ) / torch.linalg.vector_norm(guided_fp32, dim=(1, 2), keepdim=True).clamp_min(1e-5) + global_expected = unpack_latents((guided_fp32 * global_factor).to(torch.bfloat16), 4, 4) + actual = adapter.teacher_velocity( + teacher, + state, + time, + condition=condition, + negative_condition=negative_condition, + ) + + assert not torch.equal(expected, global_expected) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_qwen_pdd_rejects_unsupported_configuration_and_inputs() -> None: + transformer = _TinyQwenTransformer() + config = _config() + adapter = QwenImagePDDAdapter(config) + state, time, condition, _ = _inputs() + with pytest.raises(TypeError, match="negative_condition must be a tuple"): + adapter.teacher_velocity(transformer, state, time, condition=condition) + assert transformer.calls == [] + + convert_qwen_image_to_pdd(transformer, config) + with pytest.raises(TypeError, match="FP32 time"): + adapter.student_all_heads( + transformer, + state, + time.to(torch.bfloat16), + condition=condition, + ) + + unmarked = _TinyQwenTransformer() + delattr(unmarked, "_modelopt_qwen_image_pdd_execution") + convert_qwen_image_to_pdd(unmarked, config) + with pytest.raises(RuntimeError, match="masked joint-attention forward"): + adapter.student_all_heads(unmarked, state, time, condition=condition)