From 348cebf51b21f493f77648aeea48763c1f922f7d Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 03:23:27 -0700 Subject: [PATCH 01/45] Reorganize FastGen DMD2 example layout Signed-off-by: Meng Xin --- examples/diffusers/fastgen/README.md | 209 +----------------- examples/diffusers/fastgen/dmd2/README.md | 207 +++++++++++++++++ examples/diffusers/fastgen/dmd2/__init__.py | 6 + .../checkpoint.py} | 0 .../configs/qwen_image.yaml} | 8 +- .../export_qwen_image.py} | 2 +- .../{dmd2_finetune.py => dmd2/finetune.py} | 34 ++- .../inference_qwen_image.py} | 2 +- .../{dmd2_recipe.py => dmd2/recipe.py} | 17 +- .../fastgen/fastgen_data/__init__.py | 2 +- .../fastgen/preprocess_qwen_image.py | 4 +- .../general/distillation/dmd2_qwen_image.yaml | 3 +- .../examples/diffusers/fastgen/test_layout.py | 172 ++++++++++++++ .../fastgen/test_resume_dataloader.py | 8 +- .../fastgen/test_vendored_migration.py | 17 +- 15 files changed, 447 insertions(+), 244 deletions(-) create mode 100644 examples/diffusers/fastgen/dmd2/README.md create mode 100644 examples/diffusers/fastgen/dmd2/__init__.py rename examples/diffusers/fastgen/{fastgen_checkpoint.py => dmd2/checkpoint.py} (100%) rename examples/diffusers/fastgen/{configs/dmd2_qwen_image.yaml => dmd2/configs/qwen_image.yaml} (96%) rename examples/diffusers/fastgen/{export_diffusers_qwen_image.py => dmd2/export_qwen_image.py} (99%) rename examples/diffusers/fastgen/{dmd2_finetune.py => dmd2/finetune.py} (67%) rename examples/diffusers/fastgen/{inference_dmd2_qwen_image.py => dmd2/inference_qwen_image.py} (99%) rename examples/diffusers/fastgen/{dmd2_recipe.py => dmd2/recipe.py} (99%) create mode 100644 tests/examples/diffusers/fastgen/test_layout.py diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md index 9c9373807a9..1441b034f54 100644 --- a/examples/diffusers/fastgen/README.md +++ b/examples/diffusers/fastgen/README.md @@ -1,206 +1,9 @@ -# 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) -## 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..53f3f9440ab --- /dev/null +++ b/examples/diffusers/fastgen/dmd2/README.md @@ -0,0 +1,207 @@ +# 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,<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 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 `` 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 +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`. +- **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/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` | 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) diff --git a/examples/diffusers/fastgen/dmd2/__init__.py b/examples/diffusers/fastgen/dmd2/__init__.py new file mode 100644 index 00000000000..1184907e87c --- /dev/null +++ b/examples/diffusers/fastgen/dmd2/__init__.py @@ -0,0 +1,6 @@ +# 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 96% rename from examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml rename to examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml index d0ec32c1cca..26aa09bba9a 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% 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 99% rename from examples/diffusers/fastgen/dmd2_recipe.py rename to examples/diffusers/fastgen/dmd2/recipe.py index 7934a07cf13..261ff0a42c7 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. """ @@ -61,10 +61,9 @@ "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 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. """ # ------------------------------------------------------------------ # diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index 771b93b1c0b..c1baaf2b44a 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -30,7 +30,7 @@ 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 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/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/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py new file mode 100644 index 00000000000..6abcf36ea16 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Closed layout contract for the FastGen Diffusers example.""" + +from __future__ import annotations + +import hashlib +import json +import os +import pathlib +import re +import subprocess +import sys + +import yaml + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_ROOT = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +_EXPECTED_ROOT_ENTRIES = { + "README.md", + "dmd2", + "fastgen_data", + "make_negative_prompt_embedding.py", + "preprocess", + "preprocess_qwen_image.py", + "requirements.txt", +} +_EXPECTED_DMD2_FILES = { + "README.md", + "__init__.py", + "checkpoint.py", + "configs", + "export_qwen_image.py", + "finetune.py", + "inference_qwen_image.py", + "recipe.py", +} +_DMD2_CONFIG_DIGEST = "633799d328a710fa30e3c78ea230a29cabf8c3c01d6b86f265e2146c5c46a493" +_TEXT_SUFFIXES = {".json", ".md", ".py", ".rst", ".sh", ".toml", ".txt", ".yaml", ".yml"} + + +def _old_name(prefix: str, suffix: str) -> str: + return f"{prefix}_{suffix}" + + +def _old_modules() -> tuple[str, ...]: + return ( + _old_name("dmd2", "finetune"), + _old_name("dmd2", "recipe"), + _old_name("fastgen", "checkpoint"), + _old_name("export", "diffusers_qwen_image"), + _old_name("inference", "dmd2_qwen_image"), + ) + + +def _source_text_files() -> list[pathlib.Path]: + completed = subprocess.run( + ["git", "ls-files", "-co", "--exclude-standard", "-z"], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + ) + return [ + _REPO_ROOT / relative + for relative in completed.stdout.decode().split("\0") + if relative + and pathlib.Path(relative).suffix in _TEXT_SUFFIXES + and (_REPO_ROOT / relative).is_file() + ] + + +def test_fastgen_root_has_closed_shared_and_dmd2_ownership() -> None: + root_entries = {path.name for path in _FASTGEN_ROOT.iterdir() if path.name != "__pycache__"} + assert root_entries == _EXPECTED_ROOT_ENTRIES + assert not (_FASTGEN_ROOT / "configs").exists() + + dmd2_entries = { + path.name for path in (_FASTGEN_ROOT / "dmd2").iterdir() if path.name != "__pycache__" + } + assert dmd2_entries == _EXPECTED_DMD2_FILES + assert {path.name for path in (_FASTGEN_ROOT / "dmd2" / "configs").iterdir()} == { + "qwen_image.yaml" + } + + +def test_dmd2_config_retains_accepted_semantics() -> None: + config_path = _FASTGEN_ROOT / "dmd2" / "configs" / "qwen_image.yaml" + value = yaml.safe_load(config_path.read_text()) + digest = hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + assert digest == _DMD2_CONFIG_DIGEST + assert ( + value["data"]["dataloader"]["_target_"] + == "fastgen_data.build_text_to_image_multiresolution_dataloader" + ) + + +def test_repository_sources_have_no_flat_dmd2_paths() -> None: + old_modules = _old_modules() + stale = ( + *(f"{module}.py" for module in old_modules), + "configs/" + _old_name("dmd2", "qwen_image") + ".yaml", + ) + failures: list[str] = [] + for path in _source_text_files(): + if path == pathlib.Path(__file__): + continue + text = path.read_text(errors="strict") + relative = path.relative_to(_REPO_ROOT).as_posix() + failures.extend(f"{relative}: {token}" for token in stale if token in text) + if path.suffix == ".py": + failures.extend( + f"{relative}: stale import {module}" + for module in old_modules + if re.search(rf"(?m)^\s*(?:from|import)\s+{re.escape(module)}(?:\s|\.|$)", text) + or re.search(rf"['\"]{re.escape(module)}['\"]", text) + ) + assert not failures, "\n".join(failures) + + +def test_dmd2_package_is_import_light() -> None: + probe = f""" +import importlib, json, sys +sys.path.insert(0, {str(_FASTGEN_ROOT)!r}) +before = set(sys.modules) +module = importlib.import_module('dmd2') +forbidden = ('torch', 'modelopt', 'nemo_automodel', 'diffusers', 'transformers') +loaded = sorted(name for name in set(sys.modules) - before if name.split('.')[0] in forbidden) +print(json.dumps({{'loaded': loaded, 'all': getattr(module, '__all__', None)}})) +""" + completed = subprocess.run( + [sys.executable, "-c", probe], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + env=dict(os.environ, PYTHONDONTWRITEBYTECODE="1"), + ) + assert json.loads(completed.stdout) == {"loaded": [], "all": []} + + +def test_dmd2_help_works_from_repo_root_without_training_imports() -> None: + script = _FASTGEN_ROOT / "dmd2" / "finetune.py" + probe = f""" +import json, runpy, sys +sys.argv = [{str(script)!r}, '--help'] +before = set(sys.modules) +try: + runpy.run_path({str(script)!r}, run_name='__main__') +except SystemExit as error: + if error.code != 0: + raise +forbidden = ('torch', 'modelopt', 'nemo_automodel', 'diffusers', 'transformers') +loaded = sorted(name for name in set(sys.modules) - before if name.split('.')[0] in forbidden) +print('__FASTGEN_LOADED__=' + json.dumps(loaded)) +""" + completed = subprocess.run( + [sys.executable, "-c", probe], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + env=dict(os.environ, PYTHONDONTWRITEBYTECODE="1"), + ) + assert "DMD2 Qwen-Image training" in completed.stdout + assert "--config" in completed.stdout + loaded_line = next( + line for line in completed.stdout.splitlines() if line.startswith("__FASTGEN_LOADED__=") + ) + assert json.loads(loaded_line.partition("=")[2]) == [] diff --git a/tests/examples/diffusers/fastgen/test_resume_dataloader.py b/tests/examples/diffusers/fastgen/test_resume_dataloader.py index 3eccb103985..4e3e99981f0 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 @@ -153,7 +153,7 @@ def test_resume_rebuild_serves_clean_run_position(monkeypatch, epoch_len, grad_a assert "dataloader" in recipe.__dict__["__state_tracked"] # still tracked after rebuild # 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) diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index d881a6e49a6..a93636580ce 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, ( @@ -202,7 +201,7 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): 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 +219,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 From d5fac96927f8b102ba852884ade858a4c78c0332 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 04:05:06 -0700 Subject: [PATCH 02/45] Refactor shared FastGen data lifecycle Signed-off-by: Meng Xin --- examples/diffusers/fastgen/dmd2/README.md | 19 +- .../fastgen/dmd2/configs/qwen_image.yaml | 2 +- examples/diffusers/fastgen/dmd2/recipe.py | 68 +++--- .../fastgen/fastgen_data/__init__.py | 56 ++--- .../fastgen/fastgen_data/collate_fns.py | 56 +++-- .../diffusers/fastgen/fastgen_data/paths.py | 66 ++++++ .../diffusers/fastgen/fastgen_data/resume.py | 69 ++++++ .../diffusers/fastgen/fastgen_data/splits.py | 55 +++++ .../fastgen_data/text_to_image_dataset.py | 92 +++++++- .../preprocess/preprocessing_multiprocess.py | 19 +- examples/diffusers/fastgen/requirements.txt | 7 +- tests/examples/diffusers/fastgen/conftest.py | 88 ++++++++ .../diffusers/fastgen/test_dataset_paths.py | 211 ++++++++++++++++++ .../diffusers/fastgen/test_dataset_splits.py | 146 ++++++++++++ .../examples/diffusers/fastgen/test_layout.py | 20 +- .../fastgen/test_resume_dataloader.py | 83 ++++++- .../fastgen/test_vendored_migration.py | 7 +- 17 files changed, 942 insertions(+), 122 deletions(-) create mode 100644 examples/diffusers/fastgen/fastgen_data/paths.py create mode 100644 examples/diffusers/fastgen/fastgen_data/resume.py create mode 100644 examples/diffusers/fastgen/fastgen_data/splits.py create mode 100644 tests/examples/diffusers/fastgen/conftest.py create mode 100644 tests/examples/diffusers/fastgen/test_dataset_paths.py create mode 100644 tests/examples/diffusers/fastgen/test_dataset_splits.py diff --git a/examples/diffusers/fastgen/dmd2/README.md b/examples/diffusers/fastgen/dmd2/README.md index 53f3f9440ab..4569f5d4c07 100644 --- a/examples/diffusers/fastgen/dmd2/README.md +++ b/examples/diffusers/fastgen/dmd2/README.md @@ -52,8 +52,16 @@ 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). +Then point the config's `data.dataloader.cache_dir` at ``, or override it without +editing YAML: + +```bash +export MODELOPT_FASTGEN_DATASET_CACHE_DIR=/absolute/path/to/cache +``` + +The environment override must name an absolute existing directory. The config keeps +`negative_prompt_embedding_path: negative_prompt_embedding.pt`, so both samples and the negative +embedding are resolved from that same selected root. Paths that escape it are rejected. ## How DMD2 works @@ -99,9 +107,10 @@ this example subclasses. `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`. +- **A preprocessed Qwen-Image latent cache** — set `data.dataloader.cache_dir`, or export + `MODELOPT_FASTGEN_DATASET_CACHE_DIR` to an absolute existing cache directory. - **A precomputed negative-prompt embedding** (required for CFG) — set - `data.dataloader.negative_prompt_embedding_path`. + `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 @@ -179,7 +188,7 @@ student). | `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. | +| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Environment-overridable latent cache root + optional root-relative CFG embedding. | | `checkpoint` | `checkpoint_dir`, `model_save_format`, `restore_from` | Output dir, save format, resume behavior. | ## Troubleshooting diff --git a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml index 26aa09bba9a..847441f838a 100644 --- a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml @@ -154,7 +154,7 @@ data: 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/dmd2/recipe.py b/examples/diffusers/fastgen/dmd2/recipe.py index 261ff0a42c7..155c3b803d4 100644 --- a/examples/diffusers/fastgen/dmd2/recipe.py +++ b/examples/diffusers/fastgen/dmd2/recipe.py @@ -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 @@ -64,6 +63,7 @@ # 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_data import rebuild_stateful_dataloader from torch import nn import modelopt.torch.fastgen as mtf @@ -224,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: @@ -319,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`` @@ -330,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 c1baaf2b44a..934dc16ea28 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -13,11 +13,11 @@ # 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.5.0`` 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`` / @@ -33,30 +33,35 @@ ``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. +# Convert a missing-helper ImportError into an actionable message naming the supported release. try: - from .collate_fns import ( - build_text_to_image_multiresolution_dataloader, - collate_fn_text_to_image, - ) - from .text_to_image_dataset import TextToImageDataset + from . import collate_fns as _collate_fns + from . import paths as _paths + from . import resume as _resume + from . import splits as _splits + from . import text_to_image_dataset as _text_to_image_dataset + from .collate_fns import * + from .paths import * + from .resume import * + from .splits import * + from .text_to_image_dataset import * 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.5.0 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" f"Underlying import error: {exc!r}" ) from exc -__all__ = [ - "TextToImageDataset", - "build_text_to_image_multiresolution_dataloader", - "collate_fn_text_to_image", -] +__all__: list[str] = [] +for _module in (_collate_fns, _paths, _resume, _splits, _text_to_image_dataset): + __all__.extend(_module.__all__) def _warn_if_unsupported_upstream() -> None: @@ -72,19 +77,14 @@ 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 version != (0, 5, 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 " + "fastgen_data: installed nemo_automodel %s does not match the tested release " + "(==0.5.0). 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.", + "fail or behavior drifts, pin nemo_automodel to the supported release.", raw or "", ) except Exception: # pragma: no cover - never block import on a version probe diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index d669d2a7c4a..e7d91702efb 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -33,13 +33,20 @@ import functools import logging +from collections.abc import Sequence import torch from nemo_automodel.components.datasets.diffusion.sampler import SequentialBucketSampler 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_to_image", +] + logger = logging.getLogger(__name__) @@ -87,6 +94,7 @@ def collate_fn_text_to_image( "crop_resolution": torch.stack([item["crop_resolution"] for item in batch]), "original_resolution": torch.stack([item["original_resolution"] for item in batch]), "crop_offset": torch.stack([item["crop_offset"] for item in batch]), + "sample_ids": torch.tensor([item["sample_id"] for item in batch], dtype=torch.long), }, } @@ -167,6 +175,7 @@ def build_text_to_image_multiresolution_dataloader( pin_memory: bool = True, prefetch_factor: int = 2, negative_prompt_embedding_path: str | None = None, + selected_indices: Sequence[int] | None = None, ) -> tuple[StatefulDataLoader, SequentialBucketSampler]: """Build the DMD2 text-to-image multiresolution dataloader for ``TrainDiffusionRecipe``. @@ -186,23 +195,35 @@ def build_text_to_image_multiresolution_dataloader( 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). + selected_indices: Optional ordered original metadata ordinals to expose. 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, + selected_indices=selected_indices, + ) + 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 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, negative_text_embeddings=neg_embed, @@ -230,13 +251,16 @@ def build_text_to_image_multiresolution_dataloader( persistent_workers=num_workers > 0, ) - 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..aed0ce3911a --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/paths.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. + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Path resolution for a portable, contained FastGen dataset cache.""" + +from __future__ import annotations + +import os +from pathlib import Path + +__all__ = ["resolve_cache_root", "resolve_under_root"] + +_CACHE_ROOT_ENV = "MODELOPT_FASTGEN_DATASET_CACHE_DIR" + + +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: + """Return the effective cache root selected by config and environment. + + An unset or exactly empty ``MODELOPT_FASTGEN_DATASET_CACHE_DIR`` falls back to + ``configured_root``. A nonempty override must already be an absolute path to an existing + directory. + """ + override = os.environ.get(_CACHE_ROOT_ENV) + if override: + override_path = Path(override) + if not override_path.is_absolute(): + raise ValueError(f"{_CACHE_ROOT_ENV} must be an absolute path; got {override!r}") + return _existing_directory(override_path, _CACHE_ROOT_ENV) + 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..3756639e245 --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/resume.py @@ -0,0 +1,69 @@ +# 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 + +"""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/splits.py b/examples/diffusers/fastgen/fastgen_data/splits.py new file mode 100644 index 00000000000..e06790507c1 --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/splits.py @@ -0,0 +1,55 @@ +# 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 + +"""Deterministic train/validation membership for FastGen cache ordinals.""" + +from __future__ import annotations + +import torch + +__all__ = ["make_train_validation_indices"] + + +def _require_integer(name: str, value: int) -> int: + if type(value) is not int: + raise TypeError(f"{name} must be an integer; got {type(value).__name__}") + return value + + +def make_train_validation_indices( + num_samples: int, + validation_count: int, + seed: int, +) -> tuple[list[int], list[int]]: + """Return disjoint ordered metadata ordinals using a local CPU generator.""" + num_samples = _require_integer("num_samples", num_samples) + validation_count = _require_integer("validation_count", validation_count) + seed = _require_integer("seed", seed) + if num_samples <= 0: + raise ValueError("num_samples must be positive") + if not 1 <= validation_count < num_samples: + raise ValueError("validation_count must be in [1, num_samples)") + if seed < 0: + raise ValueError("seed must be nonnegative") + + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + permutation = torch.randperm(num_samples, generator=generator, device="cpu").tolist() + validation = sorted(permutation[:validation_count]) + train = sorted(permutation[validation_count:]) + return train, validation 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..6c0041a974e 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -13,40 +13,109 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +from collections.abc import Sequence 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, + selected_indices: Sequence[int] | None = None, ): """ Args: cache_dir: Directory containing preprocessed cache train_text_encoder: If True, returns tokens instead of embeddings + selected_indices: Optional ordered original metadata ordinals to expose. """ self.train_text_encoder = train_text_encoder - super().__init__(cache_dir, quantization=64) + self.cache_root = resolve_cache_root(cache_dir) + self._selected_indices = selected_indices + 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") + with metadata_file.open(encoding="utf-8") as file: + index = json.load(file) + 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}" + ) + with shard_path.open(encoding="utf-8") as file: + shard = json.load(file) + 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 = self._validate_selected_indices(self.total_num_samples) + return [complete_metadata[index] for index in self.sample_ids] + + def _validate_selected_indices(self, num_samples: int) -> list[int]: + if self._selected_indices is None: + return list(range(num_samples)) + if isinstance(self._selected_indices, str | bytes) or not isinstance( + self._selected_indices, Sequence + ): + raise TypeError("selected_indices must be a sequence of integers") + selected = list(self._selected_indices) + if not selected: + raise ValueError("selected_indices must not be empty") + for index in selected: + if type(index) is not int: + raise TypeError("selected_indices must contain only non-bool integers") + if not 0 <= index < num_samples: + raise ValueError(f"selected index {index} is outside [0, {num_samples})") + if len(set(selected)) != len(selected): + raise ValueError("selected_indices must be unique") + return selected 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) @@ -62,6 +131,7 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: "image_path": data["image_path"], "bucket_id": item["bucket_id"], "aspect_ratio": item.get("aspect_ratio", 1.0), + "sample_id": sample_id, } if self.train_text_encoder: diff --git a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py index d11efe30f9e..8bd0fd1b711 100644 --- a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py +++ b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py @@ -113,13 +113,26 @@ 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: + 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": str(cache_file)}) + 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: @@ -130,7 +143,7 @@ def _save_metadata_shards( "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, diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index 5e5e79f0d12..e8f4255634e 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -1,13 +1,14 @@ -# Runtime requirements for the DMD2 Qwen-Image AutoModel example. +# 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); +# accelerate with matching pins. Pinned to the public API tested by these examples; # fastgen_data/__init__.py adds a runtime guard with an actionable message if the helpers move. -nemo_automodel[diffusion]>=0.4.0,<1.0 +# Version 0.5.0 provides the public BaseRecipe.untrack_state API used by deterministic resume. +nemo_automodel[diffusion]==0.5.0 # Optional but recommended for the smoke logs. wandb diff --git a/tests/examples/diffusers/fastgen/conftest.py b/tests/examples/diffusers/fastgen/conftest.py new file mode 100644 index 00000000000..5090aa10b91 --- /dev/null +++ b/tests/examples/diffusers/fastgen/conftest.py @@ -0,0 +1,88 @@ +# 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 + +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..715d4ced4e6 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -0,0 +1,211 @@ +# 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 + +"""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 ( + TextToImageDataset, + build_text_to_image_multiresolution_dataloader, + resolve_cache_root, + resolve_under_root, +) + + +def test_cache_root_uses_unset_or_empty_fallback(make_fastgen_cache, monkeypatch, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + + monkeypatch.delenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", raising=False) + assert resolve_cache_root(cache) == cache.resolve() + + monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", "") + assert resolve_cache_root(cache) == cache.resolve() + + +@pytest.mark.parametrize("override", ["relative/cache", "~/cache", " "]) +def test_cache_root_rejects_nonempty_relative_override(monkeypatch, override, tmp_path): + fallback = tmp_path / "fallback" + fallback.mkdir() + monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", override) + + with pytest.raises(ValueError, match="absolute"): + resolve_cache_root(fallback) + + +def test_cache_root_rejects_missing_or_non_directory_override(monkeypatch, tmp_path): + fallback = tmp_path / "fallback" + fallback.mkdir() + + monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", str(tmp_path / "missing")) + with pytest.raises(FileNotFoundError): + resolve_cache_root(fallback) + + regular_file = tmp_path / "file" + regular_file.write_text("not a directory") + monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", str(regular_file)) + with pytest.raises(NotADirectoryError): + resolve_cache_root(fallback) + + +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 dataset[0]["sample_id"] == 0 + + +def test_environment_redirects_samples_and_relative_negative_embedding( + make_fastgen_cache, monkeypatch, tmp_path +): + fallback = make_fastgen_cache(tmp_path / "fallback", marker=1.0) + override = make_fastgen_cache(tmp_path / "override", marker=9.0) + monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", str(override.resolve())) + + loader, _ = build_text_to_image_multiresolution_dataloader( + cache_dir=str(fallback), + batch_size=1, + num_workers=0, + shuffle=False, + negative_prompt_embedding_path="negative_prompt_embedding.pt", + ) + batch = next(iter(loader)) + + assert loader.dataset.cache_root == override.resolve() + assert batch["metadata"]["prompts"][0].startswith("prompt-9.0-") + assert torch.equal(batch["negative_text_embeddings"], torch.full((1, 2, 3), 9.0)) + + +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_absolute_paths_for_relative_output(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.is_absolute() + assert published == payload.resolve() + published.relative_to(output.resolve()) diff --git a/tests/examples/diffusers/fastgen/test_dataset_splits.py b/tests/examples/diffusers/fastgen/test_dataset_splits.py new file mode 100644 index 00000000000..9123566d558 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_dataset_splits.py @@ -0,0 +1,146 @@ +# 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 + +"""Deterministic stable-ID split contract for the shared FastGen cache.""" + +from __future__ import annotations + +import pathlib +import sys + +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 ( + TextToImageDataset, + build_text_to_image_multiresolution_dataloader, + make_train_validation_indices, +) + + +def _snapshot(root: pathlib.Path) -> dict[str, bytes]: + return { + path.relative_to(root).as_posix(): path.read_bytes() + for path in root.rglob("*") + if path.is_file() + } + + +def test_split_has_frozen_membership_and_does_not_change_global_rng(): + torch.manual_seed(1234) + rng_before = torch.random.get_rng_state().clone() + + train, validation = make_train_validation_indices(10, validation_count=3, seed=17) + + assert validation == [0, 7, 9] + assert train == [1, 2, 3, 4, 5, 6, 8] + assert set(train).isdisjoint(validation) + assert sorted(train + validation) == list(range(10)) + assert torch.equal(torch.random.get_rng_state(), rng_before) + assert make_train_validation_indices(10, 3, 17) == (train, validation) + + +@pytest.mark.parametrize( + ("num_samples", "validation_count", "seed", "error"), + [ + (True, 1, 0, TypeError), + (1, 1, 0, ValueError), + (4, False, 0, TypeError), + (4, 0, 0, ValueError), + (4, 4, 0, ValueError), + (4, 1, True, TypeError), + (4, 1, -1, ValueError), + ], +) +def test_split_rejects_invalid_inputs(num_samples, validation_count, seed, error): + with pytest.raises(error): + make_train_validation_indices(num_samples, validation_count, seed) + + +@pytest.mark.parametrize( + "selected_indices", + [[], [0, 0], [-1], [6], [True], [1.5], ["1"]], +) +def test_dataset_rejects_invalid_selected_indices(make_fastgen_cache, selected_indices, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + with pytest.raises((TypeError, ValueError)): + TextToImageDataset(cache, selected_indices=selected_indices) + + +def test_dataset_preserves_selected_original_ordinals(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + dataset = TextToImageDataset(cache, selected_indices=[5, 1, 3]) + + assert dataset.total_num_samples == 6 + assert dataset.sample_ids == [5, 1, 3] + assert [dataset[index]["sample_id"] for index in range(len(dataset))] == [5, 1, 3] + + +def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( + make_fastgen_cache, tmp_path +): + cache = make_fastgen_cache(tmp_path / "cache") + before = _snapshot(cache) + train_ids, validation_ids = make_train_validation_indices(6, validation_count=2, seed=17) + + train_loader, train_sampler = build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), + selected_indices=train_ids, + batch_size=1, + num_workers=0, + shuffle=True, + drop_last=True, + ) + validation_loader, validation_sampler = build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), + selected_indices=validation_ids, + batch_size=1, + num_workers=0, + shuffle=False, + drop_last=False, + ) + + train_seen = torch.cat([batch["metadata"]["sample_ids"] for batch in train_loader]).tolist() + validation_seen = torch.cat( + [batch["metadata"]["sample_ids"] for batch in validation_loader] + ).tolist() + + assert sorted(train_seen) == train_ids + assert sorted(validation_seen) == validation_ids + assert set(train_seen).isdisjoint(validation_seen) + assert ( + train_loader.dataset.cache_root == validation_loader.dataset.cache_root == cache.resolve() + ) + assert train_sampler.shuffle_buckets and train_sampler.shuffle_within_bucket + assert not validation_sampler.shuffle_buckets + assert not validation_sampler.shuffle_within_bucket + assert validation_sampler.drop_last is False + assert all( + batch["metadata"]["sample_ids"].dtype == torch.long + and batch["metadata"]["sample_ids"].device.type == "cpu" + for batch in validation_loader + ) + assert _snapshot(cache) == before diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index 6abcf36ea16..0397e963c6e 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -1,3 +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. + # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 @@ -36,7 +51,7 @@ "inference_qwen_image.py", "recipe.py", } -_DMD2_CONFIG_DIGEST = "633799d328a710fa30e3c78ea230a29cabf8c3c01d6b86f265e2146c5c46a493" +_DMD2_CONFIG_DIGEST = "c802301bcaf1d861b3be4bf384ac437b4bea6cefb00865c3cf879355bcc42019" _TEXT_SUFFIXES = {".json", ".md", ".py", ".rst", ".sh", ".toml", ".txt", ".yaml", ".yml"} @@ -95,6 +110,9 @@ def test_dmd2_config_retains_accepted_semantics() -> None: value["data"]["dataloader"]["_target_"] == "fastgen_data.build_text_to_image_multiresolution_dataloader" ) + assert value["data"]["dataloader"]["negative_prompt_embedding_path"] == ( + "negative_prompt_embedding.pt" + ) def test_repository_sources_have_no_flat_dmd2_paths() -> None: diff --git a/tests/examples/diffusers/fastgen/test_resume_dataloader.py b/tests/examples/diffusers/fastgen/test_resume_dataloader.py index 4e3e99981f0..18fd9514635 100644 --- a/tests/examples/diffusers/fastgen/test_resume_dataloader.py +++ b/tests/examples/diffusers/fastgen/test_resume_dataloader.py @@ -34,6 +34,7 @@ from __future__ import annotations +import inspect import pathlib import sys from types import SimpleNamespace @@ -53,6 +54,7 @@ _sampler_mod = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") _stateful_dataloader_mod = pytest.importorskip("torchdata.stateful_dataloader") dmd2_recipe = pytest.importorskip("dmd2.recipe") +from fastgen_data import rebuild_stateful_dataloader SequentialBucketSampler = _sampler_mod.SequentialBucketSampler StatefulDataLoader = _stateful_dataloader_mod.StatefulDataLoader @@ -76,7 +78,13 @@ def __getitem__(self, i): return int(i) # identity: the served value IS the global sample index -def _build(n, sampler_cls, loader_cls): +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, **loader_kwargs): """A real sampler + StatefulDataLoader over one shared synthetic dataset.""" ds = _Dataset(n) sampler = sampler_cls( @@ -91,7 +99,13 @@ 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, + **loader_kwargs, + ) return sampler, loader @@ -129,18 +143,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,8 +160,11 @@ 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 @@ -178,5 +192,50 @@ 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 + + +def test_resume_helper_preserves_public_loader_options(): + generator = pytest.importorskip("torch").Generator().manual_seed(11) + sampler, loader = _build( + _N, + _RecordingSampler, + StatefulDataLoader, + timeout=0, + worker_init_fn=None, + generator=generator, + pin_memory=False, + in_order=True, + snapshot_every_n_steps=7, + ) + scheduler = SimpleNamespace(epoch_len=_N, grad_acc_steps=1, epoch=0, dataloader=loader) + + rebuilt = rebuild_stateful_dataloader(loader, sampler, scheduler, global_step=3) + + assert rebuilt is scheduler.dataloader + for name in ( + "collate_fn", + "num_workers", + "pin_memory", + "timeout", + "worker_init_fn", + "multiprocessing_context", + "generator", + "prefetch_factor", + "persistent_workers", + "pin_memory_device", + "in_order", + "snapshot_every_n_steps", + ): + assert getattr(rebuilt, name) is getattr(loader, name) or getattr(rebuilt, name) == getattr( + loader, name + ) + + +def test_recipe_resume_path_has_no_private_loader_or_sampler_access(): + source = inspect.getsource(dmd2_recipe.DMD2DiffusionRecipe._rebuild_dataloader_for_resume) + loop_source = inspect.getsource(dmd2_recipe.DMD2DiffusionRecipe.run_train_validation_loop) + + assert "_batches_to_skip" not in source + loop_source + assert '__dict__["dataloader"]' not in source diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index a93636580ce..6977fd2c1e0 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -148,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 negative embedding and stable-ID selection seams.""" pytest.importorskip("nemo_automodel") pytest.importorskip("torch") @@ -158,6 +158,7 @@ def test_data_builders_importable_and_accept_negative_prompt_path(): assert callable(fastgen_data.build_text_to_image_multiresolution_dataloader) sig = inspect.signature(fastgen_data.build_text_to_image_multiresolution_dataloader) assert "negative_prompt_embedding_path" in sig.parameters + assert "selected_indices" in sig.parameters # Default None => CFG-less construction works without the negative embedding (it is optional). assert sig.parameters["negative_prompt_embedding_path"].default is None @@ -184,6 +185,7 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): "aspect_ratio": 1.0, "prompt_embeds": torch.randn(seq, dim), "prompt_embeds_mask": torch.ones(seq, dtype=torch.long), + "sample_id": 7, } batch = [dict(sample), dict(sample)] neg = torch.randn(seq, dim) @@ -196,6 +198,7 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): assert out["negative_text_embeddings"].shape[0] == len( batch ) # broadcast [seq,dim]->[B,seq,dim] + assert out["metadata"]["sample_ids"].tolist() == [7, 7] def test_partial_load_checkpointer_overrides_only_load_optimizer(): From 81e32f4ccccd7eb6c205c365f3fbb5905213534b Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 03:02:40 -0700 Subject: [PATCH 03/45] test(fastgen): add PDD reference math oracle Signed-off-by: Meng Xin --- .../torch/fastgen/test_pdd_reference_math.py | 166 ++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 tests/unit/torch/fastgen/test_pdd_reference_math.py 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..6b9adc71825 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_reference_math.py @@ -0,0 +1,166 @@ +# 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 torch +import torch.nn.functional as F + + +def _reference_shifted_grid(grid_size: int, shift: float) -> torch.Tensor: + """Build the decreasing shifted rectified-flow grid with scalar arithmetic.""" + values = [] + for index in range(grid_size + 1): + unshifted = 1.0 - index / grid_size + values.append(shift * unshifted / (1.0 + (shift - 1.0) * unshifted)) + return torch.tensor(values, dtype=torch.float64) + + +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 _reference_fused_parameters( + weight: torch.Tensor, + bias: torch.Tensor | None, + grid: torch.Tensor, + start: int, + end: int, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Fuse per-interval linear parameters for the block ``[start, end)``.""" + denominator = grid[end] - grid[start] + fused_weight = torch.zeros_like(weight[0], dtype=torch.float64) + fused_bias = None if bias is None else torch.zeros_like(bias[0], dtype=torch.float64) + + for index in range(start, end): + coefficient = (grid[index + 1] - grid[index]) / denominator + fused_weight += coefficient * weight[index].to(torch.float64) + if fused_bias is not None: + fused_bias += coefficient * bias[index].to(torch.float64) + + return fused_weight, fused_bias + + +def test_shifted_grid_matches_hand_calculated_values_and_preserves_float32_intervals(): + grid = _reference_shifted_grid(grid_size=4, shift=5.0) + + expected = torch.tensor([1.0, 0.9375, 5.0 / 6.0, 0.625, 0.0], dtype=torch.float64) + torch.testing.assert_close(grid, expected, rtol=0.0, atol=0.0) + + canonical_grid = _reference_shifted_grid(grid_size=128, shift=5.0) + assert torch.all(torch.diff(canonical_grid.to(torch.float32)) < 0) + assert canonical_grid.to(torch.bfloat16)[1] == 1.0 + + for start in range(0, 128, 32): + assert canonical_grid[start + 32] - canonical_grid[start] != 0 + + +def test_half_open_integration_uses_only_selected_interval_heads(): + grid = _reference_shifted_grid(grid_size=4, shift=5.0) + state = torch.tensor([3.0, -2.0]) + velocities = torch.tensor( + [ + [1000.0, 1000.0], + [2.0, -1.0], + [-3.0, 4.0], + [-1000.0, -1000.0], + ] + ) + + result = _reference_integrate(state, velocities, grid, start=1, end=3) + expected = ( + state.to(torch.float64) + + (grid[2] - grid[1]) * velocities[1].to(torch.float64) + + (grid[3] - grid[2]) * velocities[2].to(torch.float64) + ) + + torch.testing.assert_close(result, expected, rtol=0.0, atol=1e-15) + torch.testing.assert_close( + _reference_integrate(state, velocities, grid, start=2, end=2), + state.to(torch.float64), + rtol=0.0, + atol=0.0, + ) + + +def test_final_half_open_block_advances_exactly_four_intervals(): + grid = _reference_shifted_grid(grid_size=8, shift=5.0) + state = torch.tensor([1.25]) + velocities = torch.ones(8, 1) + + result = _reference_integrate(state, velocities, grid, start=4, end=8) + expected = state.to(torch.float64) + grid[8] - grid[4] + + torch.testing.assert_close(result, expected, rtol=0.0, atol=1e-15) + + +def test_fused_projection_matches_weighted_sum_and_explicit_block_update(): + grid = _reference_shifted_grid(grid_size=4, shift=5.0) + inputs = torch.tensor([[2.0, -1.0], [-0.5, 3.0]], dtype=torch.float64) + weight = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.0, 2.0], [1.0, -1.0]], + [[-1.0, 1.0], [2.0, 0.5]], + [[3.0, -2.0], [-0.5, 1.5]], + ], + dtype=torch.float64, + ) + bias = torch.tensor( + [[0.0, 0.5], [1.0, -1.0], [2.0, 0.25], [-1.0, 3.0]], + dtype=torch.float64, + ) + original = (inputs.clone(), weight.clone(), bias.clone(), grid.clone()) + start, end = 1, 4 + + fused_weight, fused_bias = _reference_fused_parameters(weight, bias, grid, start, end) + fused_output = F.linear(inputs, fused_weight, fused_bias) + + head_outputs = torch.stack( + [F.linear(inputs, weight[index], bias[index]) for index in range(weight.shape[0])] + ) + coefficients = torch.tensor([1.0 / 9.0, 2.0 / 9.0, 2.0 / 3.0], dtype=torch.float64) + explicit_output = torch.einsum("i,ibo->bo", coefficients, head_outputs[start:end]) + explicit_update = ( + (-5.0 / 48.0) * head_outputs[1] + + (-5.0 / 24.0) * head_outputs[2] + + (-5.0 / 8.0) * head_outputs[3] + ) + + torch.testing.assert_close(fused_output, explicit_output, rtol=1e-14, atol=1e-14) + torch.testing.assert_close( + (grid[end] - grid[start]) * fused_output, + explicit_update, + rtol=1e-14, + atol=1e-14, + ) + for value, unchanged in zip((inputs, weight, bias, grid), original): + assert torch.equal(value, unchanged) From e65b2944f6de4d0783dad6fbb9cfb2f1626b71a4 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 04:17:13 -0700 Subject: [PATCH 04/45] feat(fastgen): add PDD config and flow grid Signed-off-by: Meng Xin --- modelopt/torch/fastgen/__init__.py | 5 + modelopt/torch/fastgen/config.py | 145 ++++++++++++++- modelopt/torch/fastgen/flow_matching.py | 168 +++++++++++++++++- modelopt/torch/fastgen/loader.py | 31 ++-- .../general/distillation/pdd_qwen_image.yaml | 24 +++ tests/unit/recipe/test_loader.py | 24 +++ tests/unit/torch/fastgen/test_pdd_config.py | 152 ++++++++++++++++ .../torch/fastgen/test_pdd_reference_math.py | 146 +++++++++++++++ 8 files changed, 680 insertions(+), 15 deletions(-) create mode 100644 modelopt_recipes/general/distillation/pdd_qwen_image.yaml create mode 100644 tests/unit/torch/fastgen/test_pdd_config.py diff --git a/modelopt/torch/fastgen/__init__.py b/modelopt/torch/fastgen/__init__.py index 507acfddd72..b8c9fe79fe7 100644 --- a/modelopt/torch/fastgen/__init__.py +++ b/modelopt/torch/fastgen/__init__.py @@ -57,6 +57,11 @@ 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 .pipeline import * diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py index 30d78b8720f..a0788cbddc5 100644 --- a/modelopt/torch/fastgen/config.py +++ b/modelopt/torch/fastgen/config.py @@ -26,6 +26,7 @@ from __future__ import annotations +import math from typing import TYPE_CHECKING, Literal from pydantic import Field, model_validator @@ -39,6 +40,7 @@ "DMDConfig", "DistillationConfig", "EMAConfig", + "PDDConfig", "SampleTimestepConfig", ] @@ -218,6 +220,145 @@ class DistillationConfig(ModeloptBaseConfig): ) +class PDDConfig(DistillationConfig): + """Hyperparameters for data-dependent 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. + """ + + pred_type: Literal["flow"] = ModeloptField( + default="flow", + title="Network prediction parameterization", + description="PDD is defined for rectified-flow velocity prediction.", + ) + student_sample_type: Literal["ode"] = ModeloptField( + default="ode", + title="Student sampling mode", + description="PDD fused inference follows the fixed rectified-flow ODE grid.", + ) + student_sample_steps: int = ModeloptField( + default=4, + title="Student inference steps", + description="Number of contiguous blocks in ``inference_blocks``.", + ) + grid_size: int = ModeloptField( + default=128, + title="PDD grid size", + description="Number of rectified-flow intervals and student output heads.", + ) + 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 and inference blocks.", + ) + block_size_max: int = ModeloptField( + default=64, + title="Maximum trained block size", + description="Largest target span and inference block supported by this training run.", + ) + 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: list[int] = Field( + default_factory=lambda: [32, 32, 32, 32], + title="Fused inference block schedule", + description="Contiguous interval counts that partition the complete PDD grid.", + ) + data_free: Literal[False] = ModeloptField( + default=False, + title="Data-free training", + description="Data-free PDD is unsupported; training uses noised real latents.", + ) + + def __setattr__(self, name: str, value: object) -> None: + """Validate a complete candidate config before changing an initialized field. + + Pydantic's after-model validators otherwise run after assignment and leave the + rejected value stored. PDD has cross-field schedule invariants, so attribute + and mutable-mapping updates must be transactional. + """ + if name in type(self).model_fields and name in self.__dict__: + candidate = self.model_dump() + candidate[name] = value + type(self).model_validate(candidate) + super().__setattr__(name, 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.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 block % self.block_size_min != 0: + raise ValueError( + f"inference_blocks[{index}]={block} must be aligned to " + f"block_size_min={self.block_size_min}." + ) + if block > self.block_size_max: + raise ValueError( + f"inference_blocks[{index}]={block} exceeds " + f"block_size_max={self.block_size_max}." + ) + 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)}." + ) + if self.student_sample_steps != len(self.inference_blocks): + raise ValueError( + "student_sample_steps must equal len(inference_blocks), got " + f"{self.student_sample_steps} and {len(self.inference_blocks)}." + ) + + start = 0 + for index, block in enumerate(self.inference_blocks): + if start % self.block_size_min != 0 or start > self.grid_size - self.block_size_min: + raise ValueError( + f"inference block {index} starts at {start}, which is not a valid " + f"training start aligned to block_size_min={self.block_size_min}." + ) + start += block + + default_sample_t_cfg = SampleTimestepConfig() + if self.sample_t_cfg.model_dump() != default_sample_t_cfg.model_dump(): + raise ValueError( + "sample_t_cfg is unused by PDD and cannot be overridden; PDD samples " + "discrete interval indices from its fixed shifted grid." + ) + 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 +446,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..adf6f25fe3c 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,166 @@ ] +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, + *, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Construct the fixed decreasing shifted rectified-flow grid. + + The returned tensor has ``grid_size + 1`` nodes from exactly 1 to 0. 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 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) + unshifted = torch.linspace(1.0, 0.0, grid_size + 1, device=device, dtype=math_dtype) + grid = shift * unshifted / (1.0 + (shift - 1.0) * unshifted) + torch._assert_async( + torch.all(torch.diff(grid) < 0), + f"shifted grid is not strictly decreasing for grid_size={grid_size}, 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)) + velocity_mask = mask.reshape(mask.shape + (1,) * (velocities.ndim - 2)) + selected_velocities = torch.where( + velocity_mask, + velocities.to(device=state.device, dtype=result_dtype), + torch.zeros((), device=state.device, dtype=result_dtype), + ) + update = torch.einsum( + "n,bn...->b...", + widths, + selected_velocities, + ) + 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_recipes/general/distillation/pdd_qwen_image.yaml b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml new file mode 100644 index 00000000000..673a8e5cb24 --- /dev/null +++ b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml @@ -0,0 +1,24 @@ +# 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. + +pred_type: flow +student_sample_type: ode +student_sample_steps: 4 + +# Qwen consumes the normalized continuous time supplied by the PDD grid. +num_train_timesteps: + +# 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 +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/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 3aaacaa3e0e..e66748acae4 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -34,6 +34,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 @@ -71,6 +72,29 @@ 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.inference_blocks == [32, 32, 32, 32] + + +def test_load_pdd_config_filesystem_precedes_same_named_builtin(tmp_path, monkeypatch): + """An explicit same-named filesystem recipe takes precedence over the built-in.""" + relative_path = tmp_path / "general" / "distillation" / "pdd_qwen_image.yaml" + relative_path.parent.mkdir(parents=True) + relative_path.write_text("guidance_scale: 7.0\n", encoding="utf-8") + monkeypatch.chdir(tmp_path) + + config = load_pdd_config("general/distillation/pdd_qwen_image") + + assert config.guidance_scale == 7.0 + 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..e396a27b6a1 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -0,0 +1,152 @@ +# 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, + SampleTimestepConfig, + fusion_coefficients, + integrate_interval_velocities, + load_pdd_config, + make_shifted_flow_grid, +) + + +def test_default_pdd_config_is_canonical_and_lists_are_independent(): + first = PDDConfig() + second = PDDConfig() + + assert first.pred_type == "flow" + assert first.student_sample_type == "ode" + assert first.student_sample_steps == 4 + assert first.grid_size == 128 + assert first.flow_shift == 5.0 + assert first.block_size_min == 4 + assert first.block_size_max == 64 + assert first.teacher_integrator == "euler" + assert first.inference_blocks == [32, 32, 32, 32] + assert first.data_free is False + assert first.inference_blocks is not second.inference_blocks + + first.inference_blocks[0] = 16 + assert second.inference_blocks == [32, 32, 32, 32] + + +def test_pdd_public_surface_exports_stateless_math_helpers(): + assert callable(make_shifted_flow_grid) + assert callable(integrate_interval_velocities) + assert callable(fusion_coefficients) + + +def test_pdd_config_accepts_supported_schedule_and_adapter_time_scale(): + config = PDDConfig( + inference_blocks=[64, 64], + student_sample_steps=2, + teacher_integrator="midpoint", + num_train_timesteps=1000, + ) + + assert config.inference_blocks == [64, 64] + assert config.teacher_integrator == "midpoint" + assert config.num_train_timesteps == 1000 + + +def test_rejected_attribute_assignment_leaves_pdd_config_unchanged(): + config = PDDConfig() + + with pytest.raises(ValueError, match="student_sample_steps must equal"): + config.student_sample_steps = 3 + + assert config.student_sample_steps == 4 + assert config.inference_blocks == [32, 32, 32, 32] + + +def test_rejected_mapping_assignment_leaves_pdd_config_unchanged(): + config = PDDConfig() + + with pytest.raises(ValueError, match="student_sample_steps must equal"): + config["inference_blocks"] = [64, 64] + + assert config.student_sample_steps == 4 + assert config.inference_blocks == [32, 32, 32, 32] + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"grid_size": 0}, "grid_size must be > 0"), + ({"flow_shift": 0.5}, "flow_shift must be finite and >= 1"), + ({"flow_shift": float("inf")}, "flow_shift must be finite and >= 1"), + ({"block_size_min": 0}, "0 < block_size_min"), + ({"block_size_min": 65}, "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": [30, 34, 32, 32]}, "must be aligned"), + ({"inference_blocks": [128], "student_sample_steps": 1}, "exceeds block_size_max"), + ({"inference_blocks": [32, 32, 32]}, "must sum to grid_size"), + ( + {"inference_blocks": [64, 64], "student_sample_steps": 4}, + "student_sample_steps must equal", + ), + ], +) +def test_pdd_config_rejects_invalid_grid_and_block_boundaries(overrides, message): + with pytest.raises(ValueError, match=message): + PDDConfig(**overrides) + + +@pytest.mark.parametrize( + "overrides", + [ + {"pred_type": "x0"}, + {"student_sample_type": "sde"}, + {"teacher_integrator": "heun"}, + {"data_free": True}, + ], +) +def test_pdd_config_locks_algorithm_modes(overrides): + with pytest.raises(ValueError): + PDDConfig(**overrides) + + +def test_pdd_config_rejects_nondefault_sample_timestep_config(): + with pytest.raises(ValueError, match="sample_t_cfg is unused by PDD"): + PDDConfig(sample_t_cfg=SampleTimestepConfig(shift=6.0)) + + +def test_pdd_config_accepts_explicit_default_sample_timestep_config(): + config = PDDConfig(sample_t_cfg=SampleTimestepConfig()) + assert config.sample_t_cfg == SampleTimestepConfig() + + +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]\nstudent_sample_steps: 2\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_reference_math.py b/tests/unit/torch/fastgen/test_pdd_reference_math.py index 6b9adc71825..d27b38a534e 100644 --- a/tests/unit/torch/fastgen/test_pdd_reference_math.py +++ b/tests/unit/torch/fastgen/test_pdd_reference_math.py @@ -21,9 +21,16 @@ from __future__ import annotations +import pytest import torch import torch.nn.functional as F +from modelopt.torch.fastgen.flow_matching import ( + fusion_coefficients, + integrate_interval_velocities, + make_shifted_flow_grid, +) + def _reference_shifted_grid(grid_size: int, shift: float) -> torch.Tensor: """Build the decreasing shifted rectified-flow grid with scalar arithmetic.""" @@ -164,3 +171,142 @@ def test_fused_projection_matches_weighted_sum_and_explicit_block_update(): ) for value, unchanged in zip((inputs, weight, bias, grid), original): assert torch.equal(value, unchanged) + + +def test_production_shifted_grid_matches_independent_oracle(): + grid = make_shifted_flow_grid(grid_size=128, shift=5.0) + oracle = _reference_shifted_grid(grid_size=128, shift=5.0).to(torch.float32) + + assert grid.dtype == torch.float32 + assert grid.shape == (129,) + assert grid[0] == 1.0 + assert grid[-1] == 0.0 + assert torch.all(torch.diff(grid) < 0) + torch.testing.assert_close(grid, oracle, rtol=2e-7, atol=1e-7) + + +def test_production_grid_promotes_low_precision_requests(): + grid = make_shifted_flow_grid(128, 5.0, dtype=torch.bfloat16) + + assert grid.dtype == torch.float32 + assert torch.all(torch.diff(grid) < 0) + + +@pytest.mark.parametrize( + ("grid_size", "shift", "message"), + [ + (0, 5.0, "positive integer"), + (128, 0.5, "finite and >= 1"), + (128, float("nan"), "finite and >= 1"), + ], +) +def test_production_grid_rejects_invalid_boundaries(grid_size, shift, message): + with pytest.raises(ValueError, match=message): + make_shifted_flow_grid(grid_size, shift) + + +def test_production_grid_rejects_non_floating_dtype(): + with pytest.raises(TypeError, match="floating-point dtype"): + make_shifted_flow_grid(128, 5.0, dtype=torch.int64) + + +def test_production_half_open_integration_matches_independent_oracle_per_sample(): + grid = make_shifted_flow_grid(4, 5.0, 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_does_not_consume_excluded_nonfinite_heads(): + grid = make_shifted_flow_grid(4, 5.0, dtype=torch.float64) + state = torch.tensor([[3.0, -2.0]], dtype=torch.float64) + velocities = torch.tensor( + [[[torch.nan, torch.nan], [2.0, -1.0], [-3.0, 4.0], [torch.inf, -torch.inf]]], + dtype=torch.float64, + ) + + actual = integrate_interval_velocities(state, velocities, grid, start=1, end=3) + expected = _reference_integrate(state[0], velocities[0], grid, start=1, end=3)[None] + + assert torch.all(torch.isfinite(actual)) + 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) + 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), -1.0)) + + +def test_production_integration_rejects_out_of_range_half_open_block(): + grid = make_shifted_flow_grid(4, 5.0) + 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) + + +def test_production_fusion_coefficients_match_independent_oracle(): + grid = make_shifted_flow_grid(4, 5.0, dtype=torch.float64) + actual = fusion_coefficients(grid, start=1, end=4) + oracle_grid = _reference_shifted_grid(4, 5.0) + 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) + + +def test_production_fusion_coefficients_reject_empty_block(): + grid = make_shifted_flow_grid(4, 5.0) + with pytest.raises(ValueError, match="0 <= start < end <= 4"): + fusion_coefficients(grid, start=2, end=2) + + +def test_production_helpers_do_not_extract_meta_tensor_scalars(): + grid = make_shifted_flow_grid(4, 5.0, device="meta") + state = torch.empty(2, 3, device="meta") + velocities = torch.empty(2, 4, 3, device="meta") + + integrated = integrate_interval_velocities( + state, + velocities, + grid, + start=torch.tensor([0, 2], device="meta"), + end=torch.tensor([2, 4], device="meta"), + ) + coefficients = fusion_coefficients(grid, start=1, end=4) + + assert integrated.device.type == "meta" + assert integrated.shape == state.shape + assert coefficients.device.type == "meta" + assert coefficients.shape == (3,) From db44ece17995eca5f82d55bf916e6e215830043b Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 04:36:31 -0700 Subject: [PATCH 05/45] feat(fastgen): add PDD output projection Signed-off-by: Meng Xin --- modelopt/torch/fastgen/__init__.py | 1 + modelopt/torch/fastgen/methods/__init__.py | 3 +- modelopt/torch/fastgen/methods/pdd.py | 584 ++++++++++++++++++ .../unit/torch/fastgen/test_pdd_projection.py | 421 +++++++++++++ 4 files changed, 1008 insertions(+), 1 deletion(-) create mode 100644 modelopt/torch/fastgen/methods/pdd.py create mode 100644 tests/unit/torch/fastgen/test_pdd_projection.py diff --git a/modelopt/torch/fastgen/__init__.py b/modelopt/torch/fastgen/__init__.py index b8c9fe79fe7..59d8165d062 100644 --- a/modelopt/torch/fastgen/__init__.py +++ b/modelopt/torch/fastgen/__init__.py @@ -64,6 +64,7 @@ ) from .loader import * from .methods.dmd import * +from .methods.pdd import * from .pipeline import * # isort: off 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..b4fdfa476bd --- /dev/null +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -0,0 +1,584 @@ +# 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 output projection primitives for PDD. + +This module owns only projection layout, conversion, reconstruction metadata, +and in-forward fusion. Model calls and architecture-specific packing belong to +adapters in ``modelopt.torch.fastgen.plugins``. +""" + +from __future__ import annotations + +import contextlib +import threading +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Any, Literal + +import torch +import torch.nn.functional as F +from torch import nn + +from ..config import PDDConfig +from ..flow_matching import fusion_coefficients + +__all__ = [ + "PDDLayerSpec", + "PDDMetadata", + "PDDOutputProjection", + "convert_to_pdd_output_projection", + "get_module_by_path", + "replace_module_by_path", +] + +PDDHeadLayout = Literal["channel_major", "patch_major"] +_HEAD_LAYOUTS = ("channel_major", "patch_major") +_METADATA_SCHEMA_VERSION = 1 + + +def _require_exact_keys(mapping: Mapping[str, Any], expected: set[str], *, name: str) -> None: + if any(not isinstance(key, str) for key in mapping): + raise ValueError(f"{name} keys must all be strings.") + actual = set(mapping) + if actual != expected: + missing = sorted(expected - actual) + extra = sorted(actual - expected) + raise ValueError(f"{name} keys mismatch: missing={missing}, extra={extra}.") + + +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") + + def to_dict(self) -> dict[str, Any]: + """Serialize to a strict primitive mapping.""" + return { + "projection_path": self.projection_path, + "head_layout": self.head_layout, + "output_channels": self.output_channels, + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PDDLayerSpec: + """Deserialize a strict primitive mapping.""" + if not isinstance(data, Mapping): + raise TypeError(f"layer_spec must be a mapping, got {type(data).__name__}.") + _require_exact_keys( + data, + {"projection_path", "head_layout", "output_channels"}, + name="layer_spec", + ) + if not isinstance(data["projection_path"], str): + raise ValueError("layer_spec.projection_path must be a string.") + if not isinstance(data["head_layout"], str): + raise ValueError("layer_spec.head_layout must be a string.") + output_channels = data["output_channels"] + if output_channels is not None and type(output_channels) is not int: + raise ValueError("layer_spec.output_channels must be an integer or null.") + return cls( + projection_path=data["projection_path"], + head_layout=data["head_layout"], + output_channels=output_channels, + ) + + +@dataclass(frozen=True) +class PDDMetadata: + """Versioned minimum metadata required to reconstruct a PDD projection.""" + + grid_size: int + flow_shift: float + block_size_min: int + block_size_max: int + inference_blocks: tuple[int, ...] + teacher_integrator: Literal["euler", "midpoint"] + layer_spec: PDDLayerSpec + projection_in_features: int + projection_out_features: int + projection_bias: bool + schema_version: int = _METADATA_SCHEMA_VERSION + + def __post_init__(self) -> None: + _require_int(self.schema_version, name="schema_version") + if self.schema_version != _METADATA_SCHEMA_VERSION: + raise ValueError( + f"unsupported PDD metadata schema_version={self.schema_version}; " + f"expected {_METADATA_SCHEMA_VERSION}." + ) + _require_int(self.grid_size, name="grid_size") + if type(self.flow_shift) is not float: + raise ValueError(f"flow_shift must be a float, got {self.flow_shift!r}.") + _require_int(self.block_size_min, name="block_size_min") + _require_int(self.block_size_max, name="block_size_max") + if not isinstance(self.inference_blocks, tuple) or any( + type(block) is not int for block in self.inference_blocks + ): + raise ValueError("inference_blocks must be a tuple of integers.") + if self.teacher_integrator not in ("euler", "midpoint"): + raise ValueError( + "teacher_integrator must be either 'euler' or 'midpoint', got " + f"{self.teacher_integrator!r}." + ) + _require_int(self.projection_in_features, name="projection_in_features") + _require_int(self.projection_out_features, name="projection_out_features") + if type(self.projection_bias) is not bool: + raise ValueError(f"projection_bias must be bool, got {self.projection_bias!r}.") + if not isinstance(self.layer_spec, PDDLayerSpec): + raise TypeError( + f"layer_spec must be PDDLayerSpec, got {type(self.layer_spec).__name__}." + ) + if self.layer_spec.head_layout == "patch_major": + output_channels = self.layer_spec.output_channels + if output_channels is None or self.projection_out_features % output_channels != 0: + raise ValueError( + f"projection_out_features={self.projection_out_features} must be divisible by " + f"output_channels={output_channels} for patch_major layout." + ) + + PDDConfig( + grid_size=self.grid_size, + flow_shift=self.flow_shift, + block_size_min=self.block_size_min, + block_size_max=self.block_size_max, + inference_blocks=list(self.inference_blocks), + student_sample_steps=len(self.inference_blocks), + teacher_integrator=self.teacher_integrator, + ) + + @classmethod + def from_config(cls, config: PDDConfig, projection: PDDOutputProjection) -> PDDMetadata: + """Build reconstruction metadata from a validated config and projection.""" + if not isinstance(config, PDDConfig): + raise TypeError(f"config must be PDDConfig, got {type(config).__name__}.") + if not isinstance(projection, PDDOutputProjection): + raise TypeError( + f"projection must be PDDOutputProjection, got {type(projection).__name__}." + ) + if config.grid_size != projection.grid_size: + raise ValueError( + f"config grid_size={config.grid_size} does not match projection " + f"grid_size={projection.grid_size}." + ) + return cls( + grid_size=config.grid_size, + flow_shift=config.flow_shift, + block_size_min=config.block_size_min, + block_size_max=config.block_size_max, + inference_blocks=tuple(config.inference_blocks), + teacher_integrator=config.teacher_integrator, + layer_spec=projection.layer_spec, + projection_in_features=projection.in_features, + projection_out_features=projection.base_out_features, + projection_bias=projection.bias is not None, + ) + + def to_dict(self) -> dict[str, Any]: + """Serialize to a strict, JSON/YAML-safe mapping.""" + return { + "schema_version": self.schema_version, + "grid_size": self.grid_size, + "flow_shift": self.flow_shift, + "block_size_min": self.block_size_min, + "block_size_max": self.block_size_max, + "inference_blocks": list(self.inference_blocks), + "teacher_integrator": self.teacher_integrator, + "layer_spec": self.layer_spec.to_dict(), + "base_projection": { + "in_features": self.projection_in_features, + "out_features": self.projection_out_features, + "bias": self.projection_bias, + }, + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PDDMetadata: + """Deserialize a strict, versioned metadata mapping.""" + if not isinstance(data, Mapping): + raise TypeError(f"PDD metadata must be a mapping, got {type(data).__name__}.") + _require_exact_keys( + data, + { + "schema_version", + "grid_size", + "flow_shift", + "block_size_min", + "block_size_max", + "inference_blocks", + "teacher_integrator", + "layer_spec", + "base_projection", + }, + name="PDD metadata", + ) + schema_version = _require_int(data["schema_version"], name="schema_version") + if type(data["flow_shift"]) is not float: + raise ValueError(f"flow_shift must be a float, got {data['flow_shift']!r}.") + if not isinstance(data["inference_blocks"], list) or any( + type(block) is not int for block in data["inference_blocks"] + ): + raise ValueError("inference_blocks must be a list of integers.") + if data["teacher_integrator"] not in ("euler", "midpoint"): + raise ValueError( + "teacher_integrator must be either 'euler' or 'midpoint', got " + f"{data['teacher_integrator']!r}." + ) + base_projection = data["base_projection"] + if not isinstance(base_projection, Mapping): + raise TypeError("base_projection must be a mapping.") + _require_exact_keys( + base_projection, + {"in_features", "out_features", "bias"}, + name="base_projection", + ) + if type(base_projection["bias"]) is not bool: + raise ValueError("base_projection.bias must be bool.") + + return cls( + schema_version=schema_version, + grid_size=_require_int(data["grid_size"], name="grid_size"), + flow_shift=data["flow_shift"], + block_size_min=_require_int(data["block_size_min"], name="block_size_min"), + block_size_max=_require_int(data["block_size_max"], name="block_size_max"), + inference_blocks=tuple(data["inference_blocks"]), + teacher_integrator=data["teacher_integrator"], + layer_spec=PDDLayerSpec.from_dict(data["layer_spec"]), + projection_in_features=_require_int( + base_projection["in_features"], name="base_projection.in_features" + ), + projection_out_features=_require_int( + base_projection["out_features"], name="base_projection.out_features" + ), + projection_bias=base_projection["bias"], + ) + + +@dataclass(frozen=True) +class _FusionRequest: + start: int + end: int + grid: torch.Tensor + + +class PDDOutputProjection(nn.Linear): + """A widened linear projection with one output head per PDD interval. + + Outside :meth:`fuse_block`, ``forward`` returns the full widened output. + Inside the context, ``forward`` computes the selected block's weighted + projection in float32 and returns the base-sized output. Fusion state is + synchronous and thread-owned; nested contexts in one thread are supported, + while concurrent access from another thread is rejected. + """ + + 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 + self._fusion_stack: list[_FusionRequest] = [] + self._fusion_owner_thread: int | None = None + self._fusion_lock = threading.Lock() + + def __getstate__(self) -> dict[str, Any]: + """Exclude the process-local lock while preserving ordinary module deepcopy.""" + with self._fusion_lock: + if self._fusion_stack: + raise RuntimeError("cannot copy or serialize an active PDD fusion context.") + state = super().__getstate__() + state.pop("_fusion_lock", None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore module state with a fresh process-local fusion lock.""" + super().__setstate__(state) + self._fusion_lock = threading.Lock() + + @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) + ) + + @contextlib.contextmanager + def fuse_block(self, start: int, end: int, grid: torch.Tensor) -> Iterator[PDDOutputProjection]: + """Temporarily return the fused base-sized projection for ``[start, end)``. + + Contexts may nest synchronously in one thread. Because fusion selection is + stored on the module, a second thread may neither enter a context nor call + ``forward`` until the owning context exits. + """ + 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)}." + ) + fusion_coefficients(grid, start, end) + + thread_id = threading.get_ident() + request = _FusionRequest(start=start, end=end, grid=grid) + with self._fusion_lock: + if self._fusion_stack and self._fusion_owner_thread != thread_id: + raise RuntimeError("PDD fusion context is already active in another thread.") + if not self._fusion_stack: + self._fusion_owner_thread = thread_id + self._fusion_stack.append(request) + try: + yield self + finally: + with self._fusion_lock: + if not self._fusion_stack or self._fusion_stack[-1] is not request: + raise RuntimeError("PDD fusion contexts exited out of order.") + self._fusion_stack.pop() + if not self._fusion_stack: + self._fusion_owner_thread = None + + def _fused_parameters( + self, request: _FusionRequest + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Compute block-fused parameters from the registered widened parameter.""" + coefficients = fusion_coefficients(request.grid, request.start, request.end).to( + device=self.weight.device, + dtype=torch.float32, + ) + head_weights = self._tensor_by_head(self.weight)[request.start : request.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)[request.start : request.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) -> torch.Tensor: + """Apply the widened or currently scoped fused projection.""" + with self._fusion_lock: + if not self._fusion_stack: + request = None + else: + if self._fusion_owner_thread != threading.get_ident(): + raise RuntimeError("PDD fused forward was called from a non-owning thread.") + request = self._fusion_stack[-1] + if request is None: + return F.linear(input, self.weight, self.bias) + fused_weight, fused_bias = self._fused_parameters(request) + 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 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..2e0f792bd86 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_projection.py @@ -0,0 +1,421 @@ +# 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 copy +import threading +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from dataclasses import replace + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from modelopt.torch.fastgen import ( + PDDConfig, + PDDLayerSpec, + PDDMetadata, + PDDOutputProjection, + convert_to_pdd_output_projection, + get_module_by_path, + replace_module_by_path, +) + + +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 = get_module_by_path(model, 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 get_module_by_path(model, 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 get_module_by_path(model, spec.projection_path) is projection + + +def test_nested_module_helpers_require_existing_registered_modules(): + model = _NestedModel() + replacement = nn.Linear(2, 6) + previous = replace_module_by_path(model, "transformer.blocks.1", replacement) + + assert isinstance(previous, nn.Linear) + assert get_module_by_path(model, "transformer.blocks.1") is replacement + with pytest.raises(ValueError, match="does not resolve"): + get_module_by_path(model, "transformer.missing") + with pytest.raises(ValueError, match="non-empty dotted"): + get_module_by_path(model, "") + + +@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]) + + with projection.fuse_block(1, 3, grid): + actual = projection(inputs) + + 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_fusion_contexts_nest_and_restore_in_order(): + 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]) + normal = projection(inputs) + + with projection.fuse_block(0, 2, grid): + outer_before = projection(inputs) + with projection.fuse_block(1, 3, grid): + inner = projection(inputs) + outer_after = projection(inputs) + + torch.testing.assert_close(outer_before, outer_after) + assert outer_before.shape == inner.shape == (1, 6) + assert projection(inputs).shape == normal.shape == (1, 18) + torch.testing.assert_close(projection(inputs), normal) + + +def test_active_fusion_rejects_forward_from_another_thread(): + 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]) + + with projection.fuse_block(0, 2, grid), ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(projection, inputs) + with pytest.raises(RuntimeError, match="non-owning thread"): + future.result() + + +def test_simultaneous_fusion_entry_admits_exactly_one_thread(): + projection = PDDOutputProjection.from_linear(_base_linear(), 3, _spec("channel_major")) + grid = torch.tensor([1.0, 0.7, 0.2, 0.0]) + start_barrier = threading.Barrier(3) + release_owner = threading.Event() + + def _enter(start, end): + start_barrier.wait() + try: + with projection.fuse_block(start, end, grid): + release_owner.wait(timeout=5) + return "admitted" + except RuntimeError as error: + return f"rejected: {error}" + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(_enter, 0, 2), executor.submit(_enter, 1, 3)] + start_barrier.wait() + done, _ = wait(futures, timeout=5, return_when=FIRST_COMPLETED) + release_owner.set() + results = [future.result(timeout=5) for future in futures] + + assert len(done) == 1 + assert results.count("admitted") == 1 + rejected = [result for result in results if result != "admitted"] + assert len(rejected) == 1 + assert "another thread" in rejected[0] + assert projection._fusion_stack == [] + assert projection._fusion_owner_thread is None + + +def test_fusion_context_exception_cleans_up_and_allows_reuse(): + 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]) + + with pytest.raises(RuntimeError, match="body failed"), projection.fuse_block(0, 2, grid): + raise RuntimeError("body failed") + + with projection.fuse_block(1, 3, grid): + assert projection(inputs).shape == (1, 6) + assert projection(inputs).shape == (1, 18) + + +def test_projection_deepcopy_recreates_inactive_fusion_lock(): + projection = PDDOutputProjection.from_linear(_base_linear(), 3, _spec("channel_major")) + copied = copy.deepcopy(projection) + inputs = torch.tensor([[1.0, -0.5]]) + grid = torch.tensor([1.0, 0.7, 0.2, 0.0]) + + with copied.fuse_block(0, 2, grid): + assert copied(inputs).shape == (1, 6) + assert copied._fusion_lock is not projection._fusion_lock + + +@pytest.mark.parametrize( + ("start", "end", "message"), + [(-1, 2, "0 <= start"), (1, 1, "0 <= start"), (1, 4, "0 <= start")], +) +def test_fusion_context_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.fuse_block(start, end, grid): + pass + + +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) + + +@pytest.mark.parametrize("layout", ["channel_major", "patch_major"]) +def test_metadata_round_trips_exactly_without_mutating_mapping(layout): + projection = PDDOutputProjection.from_linear(_base_linear(), 128, _spec(layout)) + metadata = PDDMetadata.from_config(PDDConfig(), projection) + payload = metadata.to_dict() + original = copy.deepcopy(payload) + + restored = PDDMetadata.from_dict(payload) + + assert restored == metadata + assert restored.to_dict() == payload + assert payload == original + + +@pytest.mark.parametrize( + "changes", + [ + {"schema_version": True}, + {"flow_shift": 5}, + {"inference_blocks": [32, 32, 32, 32]}, + {"projection_bias": 1}, + ], +) +def test_direct_metadata_construction_is_as_strict_as_deserialization(changes): + projection = PDDOutputProjection.from_linear(_base_linear(), 128, _spec("channel_major")) + metadata = PDDMetadata.from_config(PDDConfig(), projection) + with pytest.raises(ValueError): + replace(metadata, **changes) + + +def _valid_patch_metadata_payload(): + projection = PDDOutputProjection.from_linear(_base_linear(), 128, _spec("patch_major")) + return PDDMetadata.from_config(PDDConfig(), projection).to_dict() + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + (lambda data: data.update(schema_version=2), "unsupported.*schema_version"), + (lambda data: data.update(extra=True), "keys mismatch"), + (lambda data: data.update({1: "bad"}), "keys must all be strings"), + (lambda data: data.update(flow_shift=5), "flow_shift must be a float"), + (lambda data: data.update(grid_size=124), "inference_blocks must sum"), + (lambda data: data.update(inference_blocks=(32, 32, 32, 32)), "list of integers"), + ( + lambda data: data["layer_spec"].update(head_layout="unsupported"), + "head_layout must be one of", + ), + ( + lambda data: data["base_projection"].update(out_features=5), + "must be divisible", + ), + ], +) +def test_metadata_rejects_unsupported_schema_schedule_shape_and_layout(mutate, message): + payload = _valid_patch_metadata_payload() + mutate(payload) + with pytest.raises(ValueError, match=message): + PDDMetadata.from_dict(payload) + + +def test_metadata_rejects_projection_config_grid_mismatch(): + projection = PDDOutputProjection.from_linear(_base_linear(), 64, _spec("channel_major")) + with pytest.raises(ValueError, match="does not match projection"): + PDDMetadata.from_config(PDDConfig(), projection) From f4b202a3aaebd985e40f5bcf7abd5919f84f838a Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 04:52:48 -0700 Subject: [PATCH 06/45] feat(fastgen): add PDD training pipeline Signed-off-by: Meng Xin --- modelopt/torch/fastgen/methods/pdd.py | 410 ++++++++++++++++- .../fastgen/test_pdd_gradient_routing.py | 128 ++++++ tests/unit/torch/fastgen/test_pdd_pipeline.py | 423 ++++++++++++++++++ 3 files changed, 954 insertions(+), 7 deletions(-) create mode 100644 tests/unit/torch/fastgen/test_pdd_gradient_routing.py create mode 100644 tests/unit/torch/fastgen/test_pdd_pipeline.py diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index b4fdfa476bd..e4ea02b7b13 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -13,32 +13,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Framework-neutral output projection primitives for PDD. +"""Framework-neutral projection, training, and sampling primitives for PDD. -This module owns only projection layout, conversion, reconstruction metadata, -and in-forward fusion. Model calls and architecture-specific packing belong to -adapters in ``modelopt.torch.fastgen.plugins``. +This module owns projection layout and fusion plus the data-dependent objective +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 import contextlib import threading -from collections.abc import Iterator, Mapping +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass -from typing import Any, Literal +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 fusion_coefficients +from ..flow_matching import ( + fusion_coefficients, + integrate_interval_velocities, + make_shifted_flow_grid, +) +from ..pipeline import DistillationPipeline __all__ = [ "PDDLayerSpec", "PDDMetadata", + "PDDModelAdapter", "PDDOutputProjection", + "PDDPipeline", "convert_to_pdd_output_projection", "get_module_by_path", "replace_module_by_path", @@ -582,3 +589,392 @@ def convert_to_pdd_output_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(DistillationPipeline): + """Data-dependent PDD loss and fused sampler over a single core-owned grid.""" + + def __init__( + self, + student: nn.Module, + teacher: nn.Module, + config: PDDConfig, + adapter: PDDModelAdapter, + ) -> None: + """Store the models/config/adapter and freeze the teacher.""" + if not isinstance(config, PDDConfig): + raise TypeError(f"config must be PDDConfig, got {type(config).__name__}.") + super().__init__(student, teacher, 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, + 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( + 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, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute the exact data-dependent PDD objective for one batch.""" + 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) + + kwargs = self._model_kwargs(model_kwargs) + data_fp32 = data.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, + ) + time_n = grid[n] + broadcast_shape = (batch_size,) + (1,) * (data.ndim - 1) + time_n_expanded = time_n.reshape(broadcast_shape) + x_n = (1.0 - time_n_expanded) * data_fp32 + time_n_expanded * noise_fp32 + + student_heads = self.adapter.student_all_heads( + self.student, + x_n, + time_n, + condition=condition, + **kwargs, + ) + expected_head_shape = torch.Size((batch_size, self.config.grid_size, *data.shape[1:])) + student_heads = self._normalize_velocity( + student_heads, + expected_shape=expected_head_shape, + device=data.device, + name="student_all_heads", + ) + with torch.no_grad(): + x_bar_k = integrate_interval_velocities(x_n, student_heads, grid, n, k) + + batch_ids = torch.arange(batch_size, device=data.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=data.shape, + device=data.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=data.shape, + device=data.device, + name="teacher_velocity", + ) + + squared_error = (student_target - teacher_target).square() + loss = squared_error.mean() + metric_dims = tuple(range(1, squared_error.ndim)) + all_head_dims = tuple(range(1, student_heads.ndim)) + with torch.no_grad(): + metrics = { + "n": n.detach(), + "k": k.detach(), + "target_span": (k - n).detach(), + "student_target_mse": squared_error.mean(dim=metric_dims).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 + + 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.") + start = 0 + 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}.") + if block % self.config.block_size_min != 0: + raise ValueError( + f"blocks[{index}]={block} must be aligned to " + f"block_size_min={self.config.block_size_min}." + ) + if block > self.config.block_size_max: + raise ValueError( + f"blocks[{index}]={block} exceeds block_size_max={self.config.block_size_max}." + ) + if start > self.config.grid_size - self.config.block_size_min: + raise ValueError(f"block {index} starts outside the trained support at {start}.") + start += block + if start != self.config.grid_size: + raise ValueError(f"blocks must sum to grid_size={self.config.grid_size}, got {start}.") + return resolved + + @torch.no_grad() + def sample( + self, + state: torch.Tensor, + *, + condition: Any = None, + blocks: Sequence[int] | None = None, + model_kwargs: Mapping[str, Any] | None = None, + ) -> torch.Tensor: + """Sample with one fused student call per validated contiguous block.""" + self._validate_state(state, name="state") + kwargs = self._model_kwargs(model_kwargs) + resolved_blocks = self._validate_blocks(blocks) + grid = self.time_grid(state.device) + current = state.to(torch.float32) + start = 0 + for block in resolved_blocks: + end = start + block + time = grid[start].expand(state.shape[0]) + velocity = self.adapter.student_fused_block( + self.student, + current, + time, + start=start, + end=end, + grid=grid, + condition=condition, + **kwargs, + ) + velocity = self._normalize_velocity( + velocity, + expected_shape=current.shape, + device=state.device, + name="student_fused_block", + ) + current = current + (grid[end] - grid[start]) * velocity + start = end + return current diff --git a/tests/unit/torch/fastgen/test_pdd_gradient_routing.py b/tests/unit/torch/fastgen/test_pdd_gradient_routing.py new file mode 100644 index 00000000000..be679a5560e --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_gradient_routing.py @@ -0,0 +1,128 @@ +# 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. + +"""Gradient-routing contract for data-dependent PDD targets.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from modelopt.torch.fastgen import PDDConfig, PDDPipeline + + +class _GradientStudent(nn.Module): + def __init__(self) -> None: + super().__init__() + self.shared = nn.Parameter(torch.tensor(0.5)) + self.heads = nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(8, 2) / 11) + + +class _GradientTeacher(nn.Module): + def __init__(self) -> None: + super().__init__() + self.scale = nn.Parameter(torch.tensor(-0.25)) + + +class _GradientAdapter: + def __init__(self) -> None: + self.raw_heads: torch.Tensor | None = None + self.teacher_query: torch.Tensor | None = None + + def student_all_heads( + self, + model: _GradientStudent, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del time, condition, model_kwargs + output = model.shared * state[:, None] + model.heads[None] + output.retain_grad() + self.raw_heads = output + return output + + def student_fused_block(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise AssertionError("sampling is not part of this gradient test") + + def teacher_velocity( + self, + model: _GradientTeacher, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del condition, negative_condition, model_kwargs + self.teacher_query = state + return model.scale * state + time[:, None] + + +def test_only_selected_head_and_shared_backbone_receive_gradients() -> None: + student = _GradientStudent() + teacher = _GradientTeacher() + adapter = _GradientAdapter() + config = PDDConfig( + grid_size=8, + flow_shift=5.0, + block_size_min=2, + block_size_max=4, + inference_blocks=[4, 4], + student_sample_steps=2, + ) + pipeline = PDDPipeline(student, teacher, config, adapter) + data = torch.tensor([[1.0, -0.5]]) + noise = torch.tensor([[-0.25, 2.0]]) + original_heads = student.heads.detach().clone() + optimizer = torch.optim.SGD(student.parameters(), lr=0.1) + + loss, metrics = pipeline.compute_loss( + data, + noise=noise, + n=torch.tensor([0]), + k=torch.tensor([3]), + ) + loss.backward() + + assert student.shared.grad is not None + assert not torch.equal(student.shared.grad, torch.zeros_like(student.shared.grad)) + assert torch.all(torch.isfinite(student.shared.grad)) + assert student.heads.grad is not None + assert torch.count_nonzero(student.heads.grad[3]) > 0 + assert torch.count_nonzero(student.heads.grad[:3]) == 0 + assert torch.count_nonzero(student.heads.grad[4:]) == 0 + assert torch.all(torch.isfinite(student.heads.grad)) + assert adapter.raw_heads is not None + assert adapter.raw_heads.grad is not None + assert torch.count_nonzero(adapter.raw_heads.grad[:, 3]) > 0 + assert torch.count_nonzero(adapter.raw_heads.grad[:, :3]) == 0 + assert torch.count_nonzero(adapter.raw_heads.grad[:, 4:]) == 0 + assert adapter.teacher_query is not None + assert adapter.teacher_query.requires_grad is False + assert teacher.scale.requires_grad is False + assert teacher.scale.grad is None + assert loss.requires_grad is True + assert all(not value.requires_grad for value in metrics.values()) + + optimizer.step() + assert torch.equal(student.heads[:3], original_heads[:3]) + assert not torch.equal(student.heads[3], original_heads[3]) + assert torch.equal(student.heads[4:], original_heads[4:]) 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..8ddfc98a076 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -0,0 +1,423 @@ +# 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 + + +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 + + 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) + return output[:, :-1] if self.bad_student_shape 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, + } + ) + coefficients = fusion_coefficients(grid, start, end) + 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, + } + ) + return model(state, time) + + +def _config(*, teacher_integrator: str = "euler") -> PDDConfig: + return PDDConfig( + grid_size=8, + flow_shift=5.0, + block_size_min=2, + block_size_max=4, + inference_blocks=[4, 4], + student_sample_steps=2, + 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 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 = (1 - grid[n, None]) * data + grid[n, None] * noise + 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([0, 1])) + 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]]) + noise = torch.tensor([[0.25, 1.5, -0.5]]) + n = torch.tensor([6]) + k = torch.tensor([7]) + + loss, _ = pipeline.compute_loss(data, noise=noise, n=n, k=k) + + grid = pipeline.time_grid() + x_n = (1 - grid[n, None]) * data + grid[n, None] * noise + heads = pipeline.student.all_heads(x_n) + x_bar_k = x_n + (grid[7] - grid[6]) * heads[:, 6] + first_velocity = pipeline.teacher(x_bar_k, grid[k]) + delta = grid[8] - grid[7] + midpoint_state = x_bar_k + 0.5 * delta * first_velocity + midpoint_time = grid[k] + 0.5 * delta + midpoint_target = pipeline.teacher(midpoint_state, midpoint_time) + expected_loss = (heads[:, 7] - 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_small_grid_accepts_exactly_the_trained_index_support() -> None: + pipeline, _ = _pipeline() + data = torch.ones(1, 3) + noise = torch.zeros_like(data) + expected = { + (0, 0), + (0, 1), + (0, 2), + (0, 3), + (2, 2), + (2, 3), + (2, 4), + (2, 5), + (4, 4), + (4, 5), + (4, 6), + (4, 7), + (6, 6), + (6, 7), + } + accepted = set() + + for n_value in range(-1, 9): + for k_value in range(-1, 9): + pair = (n_value, k_value) + if pair not in expected: + with pytest.raises(RuntimeError): + pipeline.compute_loss( + data, + noise=noise, + n=torch.tensor([n_value]), + k=torch.tensor([k_value]), + ) + continue + pipeline.compute_loss( + data, + noise=noise, + n=torch.tensor([n_value]), + k=torch.tensor([k_value]), + ) + accepted.add(pair) + + assert accepted == expected + + +def test_explicit_k_requires_explicit_n_but_explicit_n_can_sample_k() -> None: + pipeline, _ = _pipeline() + data = torch.ones(16, 3) + noise = torch.zeros_like(data) + + with pytest.raises(ValueError, match="explicit k requires explicit n"): + pipeline.compute_loss(data, noise=noise, k=torch.zeros(16, dtype=torch.long)) + + _, metrics = pipeline.compute_loss( + data, + noise=noise, + n=torch.full((16,), 6, dtype=torch.long), + generator=torch.Generator().manual_seed(7), + ) + assert torch.equal(metrics["n"], torch.full((16,), 6, dtype=torch.long)) + assert set(metrics["k"].tolist()) == {6, 7} + + +def test_sampled_indices_stay_on_exact_uniform_support() -> None: + pipeline, _ = _pipeline() + generator = torch.Generator().manual_seed(1234) + + n, k = pipeline._resolve_indices( + batch_size=4096, + device=torch.device("cpu"), + n=None, + k=None, + generator=generator, + ) + + 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]]) +def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: + pipeline, adapter = _pipeline() + initial = torch.tensor([[1.0, -2.0, 0.5], [-0.25, 0.75, 1.5]], dtype=torch.bfloat16) + + actual = pipeline.sample( + initial, + condition="prompt", + blocks=blocks, + model_kwargs={"tag": 23}, + ) + + resolved = [4, 4] if blocks is None else blocks + grid = pipeline.time_grid() + expected = initial.float() + 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(initial.shape[0])) + torch.testing.assert_close(call["grid"], grid) + assert call["condition"] == "prompt" + assert call["kwargs"] == {"tag": 23} + start = end + + +@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="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 ([3, 5], [2, 2], [6, 2], [], [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()) From cabe88576da58e7e33d0ead4200f4ecb7ad51609 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 05:22:48 -0700 Subject: [PATCH 07/45] feat(fastgen): add Qwen-Image PDD adapter Signed-off-by: Meng Xin --- modelopt/torch/fastgen/plugins/__init__.py | 9 +- .../torch/fastgen/plugins/qwen_image_pdd.py | 465 ++++++++++++++++++ .../fastgen/test_qwen_image_pdd_plugin.py | 359 ++++++++++++++ 3 files changed, 829 insertions(+), 4 deletions(-) create mode 100644 modelopt/torch/fastgen/plugins/qwen_image_pdd.py create mode 100644 tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py 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..f13ce2de408 --- /dev/null +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -0,0 +1,465 @@ +# 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.""" + +from __future__ import annotations + +import math +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_LAYER_SPEC", + "QwenImagePDDAdapter", + "convert_qwen_image_to_pdd", +] + +QWEN_IMAGE_PDD_LAYER_SPEC = PDDLayerSpec( + projection_path="transformer.proj_out", + head_layout="channel_major", +) + +_CONTROLLED_MODEL_KWARGS = { + "encoder_hidden_states", + "encoder_hidden_states_mask", + "guidance", + "hidden_states", + "img_shapes", + "return_dict", + "timestep", + "txt_seq_lens", +} + + +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 _validate_qwen_pdd_config(config: PDDConfig) -> None: + if not isinstance(config, PDDConfig): + raise TypeError(f"config must be PDDConfig, got {type(config).__name__}.") + if config.num_train_timesteps is not None: + raise ValueError( + "Qwen-Image PDD requires num_train_timesteps=None because the adapter " + "forwards normalized continuous grid time." + ) + + +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.") + try: + current = transformer.get_submodule("proj_out") + except AttributeError as error: + raise ValueError( + "Qwen-Image transformer must register an nn.Linear at 'proj_out'." + ) from error + if not isinstance(current, nn.Linear): + raise TypeError(f"Qwen-Image proj_out must be nn.Linear, got {type(current).__name__}.") + + projection = PDDOutputProjection.from_linear( + current, + config.grid_size, + QWEN_IMAGE_PDD_LAYER_SPEC, + ) + if projection is not current: + transformer.proj_out = projection + if transformer.get_submodule("proj_out") is not projection: + raise RuntimeError("Qwen-Image proj_out replacement did not remain registered.") + return projection + + +class QwenImagePDDAdapter: + """Adapt raw Qwen packed-token calls to the framework-neutral PDD protocol.""" + + def __init__( + self, + config: PDDConfig, + *, + guidance_rescale: float = 1.0, + guidance_eps: float = 1e-5, + ) -> None: + """Validate the fixed Qwen continuous-time and packed-CFG contract.""" + _validate_qwen_pdd_config(config) + if isinstance(guidance_rescale, bool) or not isinstance(guidance_rescale, int | float): + raise TypeError("guidance_rescale must be a real number.") + if not math.isfinite(guidance_rescale) or not 0.0 <= guidance_rescale <= 1.0: + raise ValueError("guidance_rescale must be finite and in [0, 1].") + if isinstance(guidance_eps, bool) or not isinstance(guidance_eps, int | float): + raise TypeError("guidance_eps must be a real number.") + if not math.isfinite(guidance_eps) or guidance_eps <= 0.0: + raise ValueError("guidance_eps must be finite and > 0.") + 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.") + + self.config = config + self.guidance_scale = ( + None if config.guidance_scale is None else float(config.guidance_scale) + ) + self.guidance_rescale = float(guidance_rescale) + self.guidance_eps = float(guidance_eps) + + @staticmethod + def _validate_state_and_time(state: torch.Tensor, time: torch.Tensor) -> None: + if state.ndim != 4: + raise ValueError( + f"Qwen-Image PDD state must have shape [B, C, H, W], got {tuple(state.shape)}." + ) + if state.shape[2] % 2 or state.shape[3] % 2: + raise ValueError("Qwen-Image PDD requires even latent height and width.") + if time.shape != (state.shape[0],): + raise ValueError( + f"Qwen-Image PDD time must have shape ({state.shape[0]},), got {tuple(time.shape)}." + ) + if time.device != state.device: + raise ValueError(f"time must be on {state.device}, got {time.device}.") + if not time.dtype.is_floating_point: + raise TypeError(f"time must use a real floating-point dtype, got {time.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 < 2 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 + or attention_mask.shape[1] != encoder_hidden_states.shape[1] + ): + raise ValueError( + f"{name} mask must be an integer/bool tensor matching the embedding " + "sequence length." + ) + batch_size = state.shape[0] + if encoder_hidden_states.shape[0] != batch_size or attention_mask.shape[0] != batch_size: + raise ValueError(f"{name} batch size must match state batch size {batch_size}.") + 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 + + @staticmethod + def _model_dtype(model: nn.Module, fallback: torch.dtype) -> torch.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 isinstance(output, tuple): + if not output: + raise TypeError("Qwen-Image model returned an empty tuple.") + packed = output[0] + elif isinstance(output, torch.Tensor): + packed = output + elif hasattr(output, "sample"): + packed = output.sample + else: + raise TypeError( + "Qwen-Image PDD could not extract a tensor from model output of type " + f"{type(output).__name__}." + ) + if not isinstance(packed, torch.Tensor): + raise TypeError("Qwen-Image model output payload must be a tensor.") + 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_call( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + condition: Any, + model_kwargs: Mapping[str, Any], + *, + condition_name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_state_and_time(state, time) + if _config_guidance_embeds(model): + raise ValueError("Qwen-Image PDD does not support transformer guidance embeddings.") + encoder_hidden_states, attention_mask = self._parse_condition( + condition, + state=state, + name=condition_name, + ) + conflicts = sorted(_CONTROLLED_MODEL_KWARGS.intersection(model_kwargs)) + if conflicts: + raise ValueError(f"Qwen-Image PDD model_kwargs contains controlled keys: {conflicts}.") + return encoder_hidden_states, attention_mask + + def _call_packed( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + condition: Any, + model_kwargs: Mapping[str, Any], + *, + condition_name: str, + ) -> torch.Tensor: + encoder_hidden_states, attention_mask = self._prepare_call( + model, + state, + time, + condition, + model_kwargs, + condition_name=condition_name, + ) + + batch_size, _, height, width = state.shape + packed_state = pack_latents(state).to(self._model_dtype(model, state.dtype)) + txt_seq_lens = attention_mask.sum(dim=1).to(torch.int32).tolist() + output = model( + hidden_states=packed_state, + timestep=time, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=attention_mask, + img_shapes=build_img_shapes(batch_size, height, width), + txt_seq_lens=txt_seq_lens, + guidance=None, + return_dict=False, + **model_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]) + + @staticmethod + def _projection(model: nn.Module, grid_size: int) -> PDDOutputProjection: + try: + projection = model.get_submodule("proj_out") + except AttributeError as error: + raise ValueError( + "Qwen student must register a PDD projection at 'proj_out'." + ) from error + if not isinstance(projection, PDDOutputProjection): + raise TypeError( + "Qwen student proj_out must be converted to PDDOutputProjection before use." + ) + if projection.grid_size != grid_size: + raise ValueError( + f"Qwen PDD projection grid_size={projection.grid_size} does not match " + f"config grid_size={grid_size}." + ) + if projection.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: + raise ValueError("Qwen PDD projection carries an incompatible layer specification.") + return projection + + def student_all_heads( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Return unpacked canonical interval velocities from one Qwen call.""" + self._projection(model, self.config.grid_size) + packed = self._call_packed( + model, + state, + time, + condition, + model_kwargs, + condition_name="condition", + ) + 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 = self._projection(model, self.config.grid_size) + with projection.fuse_block(start, end, grid): + packed = self._call_packed( + model, + state, + time, + condition, + model_kwargs, + condition_name="condition", + ) + 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.""" + guidance_scale = self.guidance_scale + if guidance_scale is not None and negative_condition is None: + raise ValueError("negative_condition is required when Qwen teacher CFG is enabled.") + if guidance_scale is not None: + # Validate both collective-participating calls before either model + # call so malformed rank-local conditioning cannot split call counts. + self._prepare_call( + model, + state, + time, + condition, + model_kwargs, + condition_name="condition", + ) + self._prepare_call( + model, + state, + time, + negative_condition, + model_kwargs, + condition_name="negative_condition", + ) + + conditional = self._call_packed( + model, + state, + time, + condition, + model_kwargs, + condition_name="condition", + ) + if guidance_scale is None: + return self._unpack_single(conditional, state) + + unconditional = self._call_packed( + model, + state, + time, + negative_condition, + model_kwargs, + condition_name="negative_condition", + ) + 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)}." + ) + + conditional_fp32 = conditional.to(torch.float32) + guided = conditional_fp32 + (float(guidance_scale) - 1.0) * ( + conditional_fp32 - unconditional.to(torch.float32) + ) + conditional_norm = torch.linalg.vector_norm( + conditional_fp32, + dim=-1, + keepdim=True, + ) + guided_norm = torch.linalg.vector_norm(guided, dim=-1, keepdim=True) + factor = self.guidance_rescale * conditional_norm / guided_norm.clamp_min( + self.guidance_eps + ) + (1.0 - self.guidance_rescale) + return self._unpack_single(guided * factor, state) 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..6dc0ef3ba65 --- /dev/null +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -0,0 +1,359 @@ +# 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 SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from modelopt.torch.fastgen import PDDConfig, PDDOutputProjection +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_LAYER_SPEC, + convert_qwen_image_to_pdd, +) + + +class _TinyQwenTransformer(nn.Module): + """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.backbone = nn.Linear(packed_channels, hidden_width) + self.proj_out = nn.Linear(hidden_width, packed_channels) + self.calls: list[dict[str, object]] = [] + + def forward( + self, + *, + hidden_states, + timestep, + encoder_hidden_states, + encoder_hidden_states_mask, + img_shapes, + txt_seq_lens, + guidance, + return_dict, + **kwargs, + ): + 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 + 0.1 * timestep[:, None, None] + output = self.proj_out(hidden) + 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, + "txt_seq_lens": txt_seq_lens, + "guidance": guidance, + "return_dict": return_dict, + "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, + flow_shift=5.0, + block_size_min=1, + block_size_max=grid_size, + inference_blocks=[2, 2] if grid_size == 4 else [grid_size], + student_sample_steps=2 if grid_size == 4 else 1, + guidance_scale=guidance_scale, + num_train_timesteps=None, + ) + + +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) + mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long)[:batch_size] + negative_embeddings = torch.randn(batch_size, 3, 2) + negative_mask = torch.tensor([[1, 0, 0], [1, 1, 1]], dtype=torch.long)[:batch_size] + return state, time, (embeddings, mask), (negative_embeddings, negative_mask) + + +def _call_base_packed( + model: _TinyQwenTransformer, + state: torch.Tensor, + time: torch.Tensor, + condition: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + embeddings, mask = condition + return model( + hidden_states=pack_latents(state), + timestep=time, + encoder_hidden_states=embeddings, + encoder_hidden_states_mask=mask, + img_shapes=build_img_shapes(state.shape[0], state.shape[2], state.shape[3]), + txt_seq_lens=mask.sum(dim=1).to(torch.int32).tolist(), + guidance=None, + return_dict=False, + )[0] + + +def test_conversion_is_idempotent_and_every_initialized_head_matches_base() -> None: + base = _TinyQwenTransformer() + student = copy.deepcopy(base) + state, time, condition, _ = _inputs() + base_packed = _call_base_packed(base, state, time, condition) + base_velocity = unpack_latents(base_packed, 4, 4) + config = _config() + + projection = convert_qwen_image_to_pdd(student, config) + repeated = convert_qwen_image_to_pdd(student, config) + adapter = QwenImagePDDAdapter(config) + actual = adapter.student_all_heads(student, state, time, condition=condition) + + assert projection is repeated + assert student.proj_out is projection + assert isinstance(projection, PDDOutputProjection) + assert projection.layer_spec == QWEN_IMAGE_PDD_LAYER_SPEC + assert projection.layer_spec.projection_path == "transformer.proj_out" + assert actual.shape == (2, 4, 1, 4, 4) + torch.testing.assert_close(actual, base_velocity[:, None].expand_as(actual)) + assert len(student.calls) == 1 + torch.testing.assert_close(student.calls[0]["timestep"], time) + assert student.calls[0]["img_shapes"] == [[(1, 2, 2)], [(1, 2, 2)]] + assert student.calls[0]["txt_seq_lens"] == [2, 1] + assert student.calls[0]["guidance"] is None + + +def test_unfused_channel_major_output_maps_each_packed_head_in_order() -> None: + student = _TinyQwenTransformer() + config = _config() + projection = convert_qwen_image_to_pdd(student, config) + with torch.no_grad(): + projection.weight.zero_() + head_bias = torch.arange(16, dtype=torch.float32).reshape(4, 4) / 5 + projection.bias.copy_(head_bias.reshape(-1)) + state, time, condition, _ = _inputs(batch_size=1) + + actual = QwenImagePDDAdapter(config).student_all_heads( + student, + state, + time, + condition=condition, + ) + expected = torch.stack( + [ + unpack_latents( + head_bias[index].reshape(1, 1, 4).expand(1, 4, 4), + 4, + 4, + ) + for index in range(4) + ], + dim=1, + ) + + torch.testing.assert_close(actual, expected) + + +def test_fused_student_matches_explicit_packed_head_weighting() -> 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]) + all_heads = adapter.student_all_heads(student, state, time, condition=condition) + student.calls.clear() + + actual = adapter.student_fused_block( + student, + state, + time, + start=1, + end=4, + grid=grid, + condition=condition, + ) + coefficients = fusion_coefficients(grid, 1, 4) + expected = torch.einsum("n,bnchw->bchw", coefficients, all_heads[:, 1: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(state.new_zeros(1, 5)).shape[-1] == 16 + + +def test_teacher_cfg_and_packed_token_norm_rescale_match_direct_reference() -> None: + teacher = _TinyQwenTransformer() + config = _config(guidance_scale=4.0) + adapter = QwenImagePDDAdapter(config, guidance_rescale=1.0, guidance_eps=1e-5) + state, time, condition, negative_condition = _inputs() + + actual = adapter.teacher_velocity( + teacher, + state, + time, + condition=condition, + negative_condition=negative_condition, + ) + + assert len(teacher.calls) == 2 + conditional = teacher.calls[0]["output"].float() + unconditional = teacher.calls[1]["output"].float() + guided = conditional + 3.0 * (conditional - unconditional) + factor = torch.linalg.vector_norm( + conditional, + dim=-1, + keepdim=True, + ) / torch.linalg.vector_norm(guided, dim=-1, keepdim=True).clamp_min(1e-5) + expected = unpack_latents(guided * factor, 4, 4) + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(teacher.calls[0]["encoder_hidden_states"], condition[0]) + torch.testing.assert_close(teacher.calls[1]["encoder_hidden_states"], negative_condition[0]) + assert teacher.calls[0]["txt_seq_lens"] == [2, 1] + assert teacher.calls[1]["txt_seq_lens"] == [1, 3] + assert all(call["guidance"] is None for call in teacher.calls) + + +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 test_conversion_preserves_requires_grad_mode_and_rejects_conflicts() -> None: + transformer = _TinyQwenTransformer() + transformer.eval() + transformer.proj_out.weight.requires_grad_(False) + transformer.proj_out.bias.requires_grad_(False) + projection = convert_qwen_image_to_pdd(transformer, _config()) + + assert transformer.training is False + assert projection.training is False + assert projection.weight.requires_grad is False + assert projection.bias.requires_grad is False + with pytest.raises(ValueError, match="incompatible"): + convert_qwen_image_to_pdd(transformer, _config(grid_size=2)) + + +def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> None: + with pytest.raises(ValueError, match="num_train_timesteps=None"): + QwenImagePDDAdapter(_config().model_copy(update={"num_train_timesteps": 1000})) + with pytest.raises(ValueError, match="guidance_rescale"): + QwenImagePDDAdapter(_config(), guidance_rescale=1.1) + with pytest.raises(ValueError, match="guidance_eps"): + QwenImagePDDAdapter(_config(), guidance_eps=0.0) + + transformer = _TinyQwenTransformer() + transformer.config.guidance_embeds = True + with pytest.raises(ValueError, match="guidance embeddings"): + convert_qwen_image_to_pdd(transformer, _config()) + + transformer.config.guidance_embeds = False + config = _config() + adapter = QwenImagePDDAdapter(config) + state, time, condition, _ = _inputs() + with pytest.raises(ValueError, match="negative_condition is required"): + adapter.teacher_velocity(transformer, state, time, condition=condition) + with pytest.raises(TypeError, match="negative_condition must be a tuple"): + adapter.teacher_velocity( + transformer, + state, + time, + condition=condition, + negative_condition=condition[0], + ) + assert transformer.calls == [] + with pytest.raises(TypeError, match="converted to PDDOutputProjection"): + adapter.student_all_heads(transformer, state, time, condition=condition) + + convert_qwen_image_to_pdd(transformer, config) + with pytest.raises(TypeError, match="tuple"): + adapter.student_all_heads(transformer, state, time, condition=condition[0]) + with pytest.raises(ValueError, match="controlled keys"): + adapter.student_all_heads( + transformer, + state, + time, + condition=condition, + guidance=torch.ones(state.shape[0]), + ) + + +def test_raw_head_reference_uses_independent_linear_outputs() -> None: + """Pin the widened storage order without calling adapter reshape helpers.""" + student = _TinyQwenTransformer() + config = _config() + projection = convert_qwen_image_to_pdd(student, config) + state, time, condition, _ = _inputs(batch_size=1) + embeddings, mask = condition + packed = pack_latents(state) + hidden = torch.tanh(student.backbone(packed)) + hidden = hidden + embeddings.mean(dim=(1, 2), keepdim=True) + hidden = hidden + 0.01 * mask.sum(dim=1, keepdim=True).unsqueeze(-1) + hidden = hidden + 0.1 * time[:, None, None] + head_weights = projection.weight.reshape(4, 4, 5) + head_bias = projection.bias.reshape(4, 4) + expected_packed = torch.stack( + [F.linear(hidden, head_weights[index], head_bias[index]) for index in range(4)], + dim=1, + ) + expected = torch.stack( + [unpack_latents(expected_packed[:, index], 4, 4) for index in range(4)], + dim=1, + ) + + actual = QwenImagePDDAdapter(config).student_all_heads( + student, + state, + time, + condition=condition, + ) + + torch.testing.assert_close(actual, expected) From 35a608c5d2b6d749daed67dc10810385ae16ff1c Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 05:54:08 -0700 Subject: [PATCH 08/45] test(fastgen): complete PDD core audit Signed-off-by: Meng Xin --- modelopt/torch/fastgen/__init__.py | 16 +- modelopt/torch/fastgen/methods/pdd.py | 2 +- tests/unit/torch/fastgen/test_pdd_config.py | 4 + tests/unit/torch/fastgen/test_pdd_metadata.py | 193 ++++++++++++++++++ tests/unit/torch/fastgen/test_pdd_pipeline.py | 32 ++- .../unit/torch/fastgen/test_pdd_public_api.py | 167 +++++++++++++++ 6 files changed, 406 insertions(+), 8 deletions(-) create mode 100644 tests/unit/torch/fastgen/test_pdd_metadata.py create mode 100644 tests/unit/torch/fastgen/test_pdd_public_api.py diff --git a/modelopt/torch/fastgen/__init__.py b/modelopt/torch/fastgen/__init__.py index 59d8165d062..da576540145 100644 --- a/modelopt/torch/fastgen/__init__.py +++ b/modelopt/torch/fastgen/__init__.py @@ -67,8 +67,14 @@ 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. -from . import plugins + +def __getattr__(name: str): + """Load optional model plugins only when the public namespace is requested.""" + if name != "plugins": + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + import importlib + + loaded = importlib.import_module(f"{__name__}.plugins") + globals()[name] = loaded + return loaded diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index e4ea02b7b13..ad5712da7c6 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -915,7 +915,7 @@ 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): + 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: diff --git a/tests/unit/torch/fastgen/test_pdd_config.py b/tests/unit/torch/fastgen/test_pdd_config.py index e396a27b6a1..3824cf959a2 100644 --- a/tests/unit/torch/fastgen/test_pdd_config.py +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -119,6 +119,7 @@ def test_pdd_config_rejects_invalid_grid_and_block_boundaries(overrides, message {"pred_type": "x0"}, {"student_sample_type": "sde"}, {"teacher_integrator": "heun"}, + {"teacher_integrator": "rk4"}, {"data_free": True}, ], ) @@ -131,6 +132,9 @@ def test_pdd_config_rejects_nondefault_sample_timestep_config(): with pytest.raises(ValueError, match="sample_t_cfg is unused by PDD"): PDDConfig(sample_t_cfg=SampleTimestepConfig(shift=6.0)) + with pytest.raises(ValueError, match="sample_t_cfg is unused by PDD"): + PDDConfig(sample_t_cfg=SampleTimestepConfig(t_list=[1.0, 0.0])) + def test_pdd_config_accepts_explicit_default_sample_timestep_config(): config = PDDConfig(sample_t_cfg=SampleTimestepConfig()) diff --git a/tests/unit/torch/fastgen/test_pdd_metadata.py b/tests/unit/torch/fastgen/test_pdd_metadata.py new file mode 100644 index 00000000000..d2297d3d536 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_metadata.py @@ -0,0 +1,193 @@ +# 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. + +"""Plain-torch PDD lifecycle and serialized-metadata reconstruction evidence.""" + +from __future__ import annotations + +import copy +import json +from typing import Any + +import pytest +import torch +from torch import nn + +from modelopt.torch.fastgen import ( + PDDConfig, + PDDLayerSpec, + PDDMetadata, + PDDOutputProjection, + PDDPipeline, + convert_to_pdd_output_projection, +) + + +class _ToyStudent(nn.Module): + def __init__(self, width: int = 3) -> None: + super().__init__() + self.backbone = nn.Linear(width, width) + self.projection = nn.Linear(width, width) + + def forward(self, state: torch.Tensor) -> torch.Tensor: + return self.projection(torch.tanh(self.backbone(state))) + + +class _ToyTeacher(nn.Module): + def __init__(self) -> None: + super().__init__() + self.scale = nn.Parameter(torch.tensor(-0.25)) + + def forward(self, state: torch.Tensor, time: torch.Tensor) -> torch.Tensor: + return self.scale * state + 0.1 * time[:, None] + + +class _ToyAdapter: + def __init__(self, grid_size: int) -> None: + self.grid_size = grid_size + + def _projection(self, model: _ToyStudent) -> PDDOutputProjection: + projection = model.projection + assert isinstance(projection, PDDOutputProjection) + return projection + + def student_all_heads( + self, + model: _ToyStudent, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del time, condition, model_kwargs + raw = model(state) + return raw.reshape(state.shape[0], self.grid_size, state.shape[1]) + + def student_fused_block( + self, + model: _ToyStudent, + 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 + with self._projection(model).fuse_block(start, end, grid): + return model(state) + + def teacher_velocity( + self, + model: _ToyTeacher, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del condition, negative_condition, model_kwargs + return model(state, time) + + +def _config() -> PDDConfig: + return PDDConfig( + grid_size=4, + flow_shift=5.0, + block_size_min=1, + block_size_max=4, + inference_blocks=[2, 2], + student_sample_steps=2, + ) + + +def test_plain_torch_training_sampling_and_strict_metadata_reconstruction(tmp_path) -> None: + torch.manual_seed(19) + config = _config() + layer_spec = PDDLayerSpec("projection", "channel_major") + student = _ToyStudent() + projection = convert_to_pdd_output_projection(student, layer_spec, config.grid_size) + pipeline = PDDPipeline(student, _ToyTeacher(), config, _ToyAdapter(config.grid_size)) + optimizer = torch.optim.SGD(student.parameters(), lr=0.05) + data = torch.tensor([[0.5, -1.0, 0.25], [-0.75, 0.5, 1.25]]) + noise = torch.tensor([[-0.25, 0.75, 1.0], [0.5, -1.25, 0.0]]) + projection_before = projection.weight.detach().clone() + + loss, _ = pipeline.compute_loss( + data, + noise=noise, + n=torch.tensor([0, 1]), + k=torch.tensor([2, 3]), + ) + loss.backward() + optimizer.step() + + assert torch.isfinite(loss) + assert not torch.equal(projection.weight, projection_before) + initial = torch.tensor([[1.0, -0.5, 0.25], [-0.25, 0.75, 1.5]]) + expected_sample = pipeline.sample(initial) + + metadata = PDDMetadata.from_config(config, projection) + metadata_path = tmp_path / "pdd_metadata.json" + metadata_path.write_text(json.dumps(metadata.to_dict(), sort_keys=True), encoding="utf-8") + restored_metadata = PDDMetadata.from_dict(json.loads(metadata_path.read_text(encoding="utf-8"))) + assert restored_metadata == metadata + + restored_config = PDDConfig( + grid_size=restored_metadata.grid_size, + flow_shift=restored_metadata.flow_shift, + block_size_min=restored_metadata.block_size_min, + block_size_max=restored_metadata.block_size_max, + inference_blocks=list(restored_metadata.inference_blocks), + student_sample_steps=len(restored_metadata.inference_blocks), + teacher_integrator=restored_metadata.teacher_integrator, + ) + + restored_student = _ToyStudent(width=restored_metadata.projection_in_features) + restored_projection = convert_to_pdd_output_projection( + restored_student, + restored_metadata.layer_spec, + restored_metadata.grid_size, + ) + load_result = restored_student.load_state_dict(copy.deepcopy(student.state_dict()), strict=True) + assert load_result.missing_keys == [] + assert load_result.unexpected_keys == [] + assert restored_projection.base_out_features == restored_metadata.projection_out_features + assert (restored_projection.bias is not None) is restored_metadata.projection_bias + + restored_pipeline = PDDPipeline( + restored_student, + copy.deepcopy(pipeline.teacher), + restored_config, + _ToyAdapter(restored_metadata.grid_size), + ) + torch.testing.assert_close(restored_pipeline.sample(initial), expected_sample) + + +def test_strict_restore_rejects_checkpoint_with_different_projection_grid() -> None: + config = _config() + layer_spec = PDDLayerSpec("projection", "channel_major") + student = _ToyStudent() + convert_to_pdd_output_projection(student, layer_spec, config.grid_size) + checkpoint = copy.deepcopy(student.state_dict()) + incompatible = _ToyStudent() + convert_to_pdd_output_projection(incompatible, layer_spec, grid_size=2) + + with pytest.raises(RuntimeError, match="size mismatch"): + incompatible.load_state_dict(checkpoint, strict=True) diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index 8ddfc98a076..8fbf719a429 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -62,6 +62,7 @@ def __init__(self) -> None: 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, @@ -81,7 +82,9 @@ def student_all_heads( } ) output = model.all_heads(state) - return output[:, :-1] if self.bad_student_shape else output + 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, @@ -131,7 +134,8 @@ def teacher_velocity( "kwargs": model_kwargs, } ) - return model(state, time) + output = model(state, time) + return output.to(torch.bfloat16) if self.low_precision_outputs else output def _config(*, teacher_integrator: str = "euler") -> PDDConfig: @@ -246,6 +250,26 @@ def test_midpoint_target_uses_exact_final_interval_midpoint() -> None: torch.testing.assert_close(adapter.teacher_calls[1]["time"], midpoint_time) +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 = (1 - grid[n, None]) * data + grid[n, None] * noise + 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_small_grid_accepts_exactly_the_trained_index_support() -> None: pipeline, _ = _pipeline() data = torch.ones(1, 3) @@ -385,6 +409,10 @@ def test_loss_rejects_indices_outside_trained_support(n, k, message) -> None: 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="state 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"): diff --git a/tests/unit/torch/fastgen/test_pdd_public_api.py b/tests/unit/torch/fastgen/test_pdd_public_api.py new file mode 100644 index 00000000000..617b6a9b46c --- /dev/null +++ b/tests/unit/torch/fastgen/test_pdd_public_api.py @@ -0,0 +1,167 @@ +# 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-surface and optional-dependency isolation checks for core PDD.""" + +from __future__ import annotations + +import ast +import subprocess +import sys +from pathlib import Path + +import modelopt.torch.fastgen as fastgen +from modelopt.torch.fastgen.config import PDDConfig +from modelopt.torch.fastgen.flow_matching import ( + fusion_coefficients, + integrate_interval_velocities, + make_shifted_flow_grid, +) +from modelopt.torch.fastgen.loader import load_pdd_config +from modelopt.torch.fastgen.methods.pdd import ( + PDDLayerSpec, + PDDMetadata, + PDDModelAdapter, + PDDOutputProjection, + PDDPipeline, + convert_to_pdd_output_projection, + get_module_by_path, + replace_module_by_path, +) + +_CORE_SOURCES = ( + "modelopt/torch/fastgen/config.py", + "modelopt/torch/fastgen/flow_matching.py", + "modelopt/torch/fastgen/loader.py", + "modelopt/torch/fastgen/methods/pdd.py", +) +_FORBIDDEN_IMPORT_ROOTS = {"diffusers", "fastgen", "nemo_automodel", "transformers"} +_FORBIDDEN_RELATIVE_COMPONENTS = {"plugins", "qwen_image", "qwen_image_pdd"} + + +def test_core_pdd_sources_do_not_import_model_or_framework_packages() -> None: + repository = Path(__file__).resolve().parents[4] + violations = [] + + for relative_path in _CORE_SOURCES: + source_path = repository / relative_path + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported = [alias.name for alias in node.names] + forbidden = [ + name for name in imported if name.split(".", 1)[0] in _FORBIDDEN_IMPORT_ROOTS + ] + elif isinstance(node, ast.ImportFrom) and node.module is not None: + imported = [node.module] + if node.level == 0: + forbidden = [ + name + for name in imported + if name.split(".", 1)[0] in _FORBIDDEN_IMPORT_ROOTS + ] + else: + forbidden = [ + name + for name in imported + if set(name.split(".")).intersection(_FORBIDDEN_RELATIVE_COMPONENTS) + ] + else: + continue + violations.extend((relative_path, name) for name in forbidden) + + assert violations == [] + + +def test_core_pdd_symbols_are_exported_from_the_public_package() -> None: + expected = { + "PDDConfig": PDDConfig, + "PDDLayerSpec": PDDLayerSpec, + "PDDMetadata": PDDMetadata, + "PDDModelAdapter": PDDModelAdapter, + "PDDOutputProjection": PDDOutputProjection, + "PDDPipeline": PDDPipeline, + "convert_to_pdd_output_projection": convert_to_pdd_output_projection, + "fusion_coefficients": fusion_coefficients, + "get_module_by_path": get_module_by_path, + "integrate_interval_velocities": integrate_interval_velocities, + "load_pdd_config": load_pdd_config, + "make_shifted_flow_grid": make_shifted_flow_grid, + "replace_module_by_path": replace_module_by_path, + } + + assert {name: getattr(fastgen, name) for name in expected} == expected + + +def test_fresh_core_import_does_not_load_model_plugins_or_frameworks() -> None: + repository = Path(__file__).resolve().parents[4] + script = r""" +import importlib.abc +import sys + + +class _UnavailableOptionalFrameworks(importlib.abc.MetaPathFinder): + prefixes = ("diffusers", "fastgen", "nemo_automodel", "transformers") + + def __init__(self): + self.attempts = [] + + def find_spec(self, fullname, path=None, target=None): + del path, target + if any(fullname == prefix or fullname.startswith(prefix + ".") for prefix in self.prefixes): + self.attempts.append(fullname) + raise ModuleNotFoundError(f"unavailable optional framework {fullname}", name=fullname) + return None + + +blocker = _UnavailableOptionalFrameworks() +sys.meta_path.insert(0, blocker) +import modelopt.torch + +baseline = set(sys.modules) +blocker.attempts.clear() +from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline + +assert all(symbol is not None for symbol in (PDDConfig, PDDMetadata, PDDPipeline)) +assert blocker.attempts == [] +loaded = set(sys.modules) - baseline +for prefix in ( + "diffusers", + "fastgen", + "nemo_automodel", + "modelopt.torch.fastgen.plugins.qwen_image", + "transformers", +): + assert not any(name == prefix or name.startswith(prefix + ".") for name in loaded), ( + prefix, + sorted(name for name in loaded if name == prefix or name.startswith(prefix + ".")), + ) +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=repository, + check=False, + capture_output=True, + close_fds=True, + start_new_session=True, + stdin=subprocess.DEVNULL, + text=True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + + +def test_optional_plugins_remain_available_through_explicit_public_access() -> None: + assert fastgen.plugins.__name__ == "modelopt.torch.fastgen.plugins" From ea4765d052862e15228e2834f894611710b8764e Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 06:35:35 -0700 Subject: [PATCH 09/45] feat(fastgen): add portable Qwen cache contract Signed-off-by: Meng Xin --- .../fastgen/dmd2/configs/qwen_image.yaml | 4 + .../fastgen/fastgen_data/collate_fns.py | 1 - .../fastgen_data/text_to_image_dataset.py | 2 - .../fastgen/migrate_cache_manifest.py | 454 ++++++++++++++++++ examples/diffusers/fastgen/portable_cache.py | 292 +++++++++++ .../fastgen/preprocess/processors/base.py | 3 +- .../preprocess/processors/qwen_image.py | 4 +- .../fastgen/validate_cache_snapshot.py | 192 ++++++++ .../fastgen/test_migrate_cache_manifest.py | 352 ++++++++++++++ .../diffusers/fastgen/test_portable_cache.py | 418 ++++++++++++++++ .../fastgen/test_vendored_migration.py | 5 +- 11 files changed, 1720 insertions(+), 7 deletions(-) create mode 100644 examples/diffusers/fastgen/migrate_cache_manifest.py create mode 100644 examples/diffusers/fastgen/portable_cache.py create mode 100644 examples/diffusers/fastgen/validate_cache_snapshot.py create mode 100644 tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py create mode 100644 tests/examples/diffusers/fastgen/test_portable_cache.py diff --git a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml index 847441f838a..258c6d0f15d 100644 --- a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml @@ -149,6 +149,10 @@ data: dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: /path/to/preprocessed/qwen_image_1024p + # cache_dir must be a finalized migrate_cache_manifest.py snapshot. The split index and + # negative embedding are portable references beneath the effective root. + # MODELOPT_FASTGEN_DATASET_CACHE_DIR overrides cache_dir when set to a non-empty value. + metadata_index: metadata_train.json base_resolution: [1024, 1024] batch_size: 1 drop_last: false diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index e7d91702efb..812ee6c4e29 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -97,7 +97,6 @@ def collate_fn_text_to_image( "sample_ids": torch.tensor([item["sample_id"] for item in batch], dtype=torch.long), }, } - # Optional model-specific embedding fields, when a dataset provides them. for key in ("pooled_prompt_embeds", "clip_hidden"): if key in batch[0]: 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 6c0041a974e..98c4052b4a3 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -119,7 +119,6 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: # 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 = { @@ -133,7 +132,6 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: "aspect_ratio": item.get("aspect_ratio", 1.0), "sample_id": sample_id, } - 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/migrate_cache_manifest.py b/examples/diffusers/fastgen/migrate_cache_manifest.py new file mode 100644 index 00000000000..7651f59304b --- /dev/null +++ b/examples/diffusers/fastgen/migrate_cache_manifest.py @@ -0,0 +1,454 @@ +# 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. + +"""Migrate a legacy absolute-path FastGen cache into an immutable portable snapshot.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from portable_cache import ( + CACHE_SCHEMA_VERSION, + audit_no_absolute_paths, + resolve_cache_asset, + sha256_file, + stable_sample_id, + validate_relative_reference, +) +from validate_cache_snapshot import validate_snapshot + +if TYPE_CHECKING: + from collections.abc import Sequence + +_REMOVED_PATH_KEYS = { + "cache_dir", + "cache_file", + "image_path", + "output_dir", + "source_dir", + "source_path", + "video_path", +} +_SPLIT_DOMAIN = "modelopt-fastgen-split-v1" + + +@dataclass(frozen=True) +class MigrationRecord: + """Frozen mapping produced by read-only pass 1 and consumed by pass 2.""" + + sample_id: str + source_ref: str + source_payload: Path + source_sha256: str + destination_ref: str + manifest_fields: dict[str, Any] + + +def _load_json(path: Path, expected_type: type, label: str): + try: + with path.open(encoding="utf-8") as stream: + value = json.load(stream) + except json.JSONDecodeError as error: + raise ValueError(f"{label} is not valid JSON: {path}") from error + if not isinstance(value, expected_type): + raise ValueError(f"{label} must contain {expected_type.__name__}") + return value + + +def _relative_to_legacy_prefix(raw: str, prefix: Path, *, label: str) -> Path: + path = Path(raw).expanduser() + if not path.is_absolute(): + return validate_relative_reference(raw, label=label) + try: + relative = path.relative_to(prefix) + except ValueError as error: + raise ValueError(f"{label} is outside the declared legacy prefix: {raw}") from error + return validate_relative_reference(relative.as_posix(), label=label) + + +def _resolve_legacy_payload( + source_root: Path, + cache_file: Any, + legacy_cache_root: Path, + *, + label: str, +) -> Path: + if not isinstance(cache_file, str): + raise TypeError(f"{label} must be a string") + relative = _relative_to_legacy_prefix(cache_file, legacy_cache_root, label=label) + return resolve_cache_asset(source_root, relative.as_posix(), label=label) + + +def _legacy_source_ref( + entry: dict[str, Any], legacy_source_root: Path | None, *, label: str +) -> str: + if "source_ref" in entry: + return validate_relative_reference( + entry["source_ref"], label=f"{label}.source_ref" + ).as_posix() + image_path = entry.get("image_path") + if not isinstance(image_path, str): + raise ValueError(f"{label} needs source_ref or legacy image_path") + if Path(image_path).is_absolute() and legacy_source_root is None: + raise ValueError("--legacy-source-root is required for absolute legacy image_path values") + prefix = legacy_source_root or Path(".") + return _relative_to_legacy_prefix(image_path, prefix, label=f"{label}.image_path").as_posix() + + +def _sanitize_payload(value: Any) -> Any: + if isinstance(value, dict): + sanitized = {} + for key, nested in value.items(): + if not isinstance(key, str): + raise ValueError(f"payload contains non-string key {key!r}") + if key.lower() not in _REMOVED_PATH_KEYS: + sanitized[key] = _sanitize_payload(nested) + return sanitized + if isinstance(value, list): + return [_sanitize_payload(item) for item in value] + if isinstance(value, tuple): + return tuple(_sanitize_payload(item) for item in value) + if isinstance(value, set): + return {_sanitize_payload(item) for item in value} + if isinstance(value, frozenset): + return frozenset(_sanitize_payload(item) for item in value) + return value + + +def _portable_manifest_fields(entry: dict[str, Any], *, label: str) -> dict[str, Any]: + required = ("bucket_resolution", "original_resolution", "prompt", "bucket_id", "aspect_ratio") + missing = [name for name in required if name not in entry] + if missing: + raise ValueError(f"{label} is missing required fields: {missing}") + fields = {name: entry[name] for name in required} + for name in ("pixels", "model_type", "crop_resolution"): + if name in entry: + fields[name] = entry[name] + audit_no_absolute_paths(fields, context=label) + return fields + + +def plan_migration( + source_root: str | Path, + *, + source_index: str | Sequence[str] = "metadata.json", + legacy_cache_root: str | Path | None = None, + legacy_source_root: str | Path | None = None, +) -> tuple[MigrationRecord, ...]: + """Perform read-only pass 1 and return a deterministic frozen mapping.""" + root = Path(source_root).expanduser().resolve(strict=True) + if not root.is_dir(): + raise NotADirectoryError(f"source_root is not a directory: {root}") + cache_prefix = Path(legacy_cache_root).expanduser() if legacy_cache_root else root + source_prefix = Path(legacy_source_root).expanduser() if legacy_source_root else None + if not cache_prefix.is_absolute(): + raise ValueError("legacy_cache_root must be absolute") + if source_prefix is not None and not source_prefix.is_absolute(): + raise ValueError("legacy_source_root must be absolute") + + source_indexes = (source_index,) if isinstance(source_index, str) else tuple(source_index) + if not source_indexes or any(not isinstance(name, str) or not name for name in source_indexes): + raise ValueError("source_index must contain one or more index paths") + if len(source_indexes) != len(set(source_indexes)): + raise ValueError("source_index contains duplicates") + + parsed_indexes = [] + all_shards: set[str] = set() + for index_number, index_ref in enumerate(source_indexes): + label = f"source_index[{index_number}]" + index_path = resolve_cache_asset(root, index_ref, label=label) + index = _load_json(index_path, dict, label) + if "schema_version" in index and index["schema_version"] != CACHE_SCHEMA_VERSION: + raise ValueError(f"{label}.schema_version is unsupported") + shards = index.get("shards") + if ( + not isinstance(shards, list) + or not shards + or any(not isinstance(item, str) for item in shards) + ): + raise ValueError(f"{label}.shards must be a non-empty list of relative paths") + if len(shards) != len(set(shards)): + raise ValueError(f"{label}.shards contains duplicates") + if "num_shards" in index and index["num_shards"] != len(shards): + raise ValueError(f"{label}.num_shards does not match len(shards)") + overlap = all_shards.intersection(shards) + if overlap: + raise ValueError(f"source indices share metadata shards: {sorted(overlap)}") + all_shards.update(shards) + parsed_indexes.append((label, index, shards)) + + ranked_indexes = [ + index for _, index, _ in parsed_indexes if "shard_world" in index or "shard_rank" in index + ] + if ranked_indexes: + if len(ranked_indexes) != len(parsed_indexes): + raise ValueError("all source indices must declare shard_rank and shard_world") + worlds = {index.get("shard_world") for index in ranked_indexes} + ranks = {index.get("shard_rank") for index in ranked_indexes} + if worlds != {len(parsed_indexes)} or ranks != set(range(len(parsed_indexes))): + raise ValueError("source rank indices are incomplete or inconsistent") + + records: list[MigrationRecord] = [] + seen_ids: set[str] = set() + for index_label, index, shards in parsed_indexes: + index_record_count = 0 + source_sample_ids = [] + for shard_number, shard_ref in enumerate(shards): + shard_label = f"{index_label}.shards[{shard_number}]" + shard_path = resolve_cache_asset(root, shard_ref, label=shard_label) + entries = _load_json(shard_path, list, shard_label) + index_record_count += len(entries) + for entry_number, entry in enumerate(entries): + label = f"{shard_label}[{entry_number}]" + if not isinstance(entry, dict): + raise ValueError(f"{label} must be an object") + source_sample_ids.append(entry.get("sample_id")) + source_ref = _legacy_source_ref(entry, source_prefix, label=label) + source_payload = _resolve_legacy_payload( + root, + entry.get("cache_file"), + cache_prefix, + label=f"{label}.cache_file", + ) + resolution = entry.get("bucket_resolution", entry.get("crop_resolution")) + model_type = entry.get("model_type") + if not isinstance(model_type, str) or not model_type: + raise ValueError(f"{label}.model_type must be a non-empty string") + sample_id = stable_sample_id( + source_ref=source_ref, + resolution=resolution, + model_type=model_type, + ) + if sample_id in seen_ids: + raise ValueError(f"duplicate migrated sample_id: {sample_id}") + seen_ids.add(sample_id) + + source_digest = sha256_file(source_payload) + payload = torch.load(source_payload, map_location="cpu", weights_only=True) + if not isinstance(payload, dict): + raise TypeError(f"{label} payload must be a dict") + sanitized = _sanitize_payload(payload) + sanitized["sample_id"] = sample_id + sanitized["source_ref"] = source_ref + audit_no_absolute_paths(sanitized, context=f"payload[{sample_id}]") + + records.append( + MigrationRecord( + sample_id=sample_id, + source_ref=source_ref, + source_payload=source_payload, + source_sha256=source_digest, + destination_ref=f"payloads/{sample_id}.pt", + manifest_fields=_portable_manifest_fields(entry, label=label), + ) + ) + if "total_items" in index and index["total_items"] != index_record_count: + raise ValueError(f"{index_label}.total_items does not match its loaded entry count") + if "sample_ids" in index and index["sample_ids"] != source_sample_ids: + raise ValueError(f"{index_label}.sample_ids does not match its loaded entries") + + if not records: + raise ValueError("legacy source index contains no entries") + return tuple(sorted(records, key=lambda record: record.sample_id)) + + +def _write_json(path: Path, value: Any) -> None: + with path.open("w", encoding="utf-8") as stream: + json.dump(value, stream, indent=2, sort_keys=True) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + + +def _split_ids(records: tuple[MigrationRecord, ...], heldout_count: int, split_seed: str): + if heldout_count <= 0 or heldout_count >= len(records): + raise ValueError("heldout_count must be positive and smaller than the sample count") + ranked = sorted( + records, + key=lambda record: hashlib.sha256( + f"{_SPLIT_DOMAIN}\0{split_seed}\0{record.sample_id}".encode() + ).hexdigest(), + ) + heldout = {record.sample_id for record in ranked[:heldout_count]} + ordered = [record.sample_id for record in records] + return [sample_id for sample_id in ordered if sample_id not in heldout], [ + sample_id for sample_id in ordered if sample_id in heldout + ] + + +def migrate_cache( + source_root: str | Path, + output_root: str | Path, + *, + heldout_count: int, + split_seed: str = "0", + source_index: str | Sequence[str] = "metadata.json", + legacy_cache_root: str | Path | None = None, + legacy_source_root: str | Path | None = None, + negative_embedding: str | None = None, + shard_size: int = 10000, +) -> dict[str, Any]: + """Run two-pass migration and atomically publish the validated destination.""" + destination = Path(output_root).expanduser() + if destination.exists(): + raise FileExistsError(f"output_root already exists: {destination}") + if not destination.parent.exists(): + raise FileNotFoundError(f"output_root parent does not exist: {destination.parent}") + if shard_size <= 0: + raise ValueError("shard_size must be positive") + + # Pass 1 is intentionally complete before any destination or staging path is created. + records = plan_migration( + source_root, + source_index=source_index, + legacy_cache_root=legacy_cache_root, + legacy_source_root=legacy_source_root, + ) + train_ids, heldout_ids = _split_ids(records, heldout_count, split_seed) + + source = Path(source_root).expanduser().resolve(strict=True) + negative_source = None + negative_source_sha256 = None + if negative_embedding is not None: + negative_source = _resolve_legacy_payload( + source, + negative_embedding, + Path(legacy_cache_root).expanduser() if legacy_cache_root else source, + label="negative_embedding", + ) + negative_source_sha256 = sha256_file(negative_source) + negative_payload = torch.load(negative_source, map_location="cpu", weights_only=True) + audit_no_absolute_paths(negative_payload, context="negative_prompt_embedding") + + staging = Path(tempfile.mkdtemp(prefix=f".{destination.name}.staging-", dir=destination.parent)) + try: + (staging / "payloads").mkdir() + portable_entries = [] + for record in records: + if sha256_file(record.source_payload) != record.source_sha256: + raise RuntimeError(f"source payload changed after pass 1: {record.source_payload}") + payload = torch.load(record.source_payload, map_location="cpu", weights_only=True) + sanitized = _sanitize_payload(payload) + sanitized["sample_id"] = record.sample_id + sanitized["source_ref"] = record.source_ref + audit_no_absolute_paths(sanitized, context=f"payload[{record.sample_id}]") + + destination_path = staging / record.destination_ref + torch.save(sanitized, destination_path) + portable_entries.append( + { + "sample_id": record.sample_id, + "source_ref": record.source_ref, + "cache_file": record.destination_ref, + "payload_sha256": sha256_file(destination_path), + **record.manifest_fields, + } + ) + + shard_names = [] + for offset in range(0, len(portable_entries), shard_size): + name = f"metadata_shard_s{offset // shard_size:04d}.json" + _write_json(staging / name, portable_entries[offset : offset + shard_size]) + shard_names.append(name) + + negative_declaration = None + if negative_source is not None: + if sha256_file(negative_source) != negative_source_sha256: + raise RuntimeError("negative prompt embedding changed after pass 1") + negative_name = "negative_prompt_embedding.pt" + shutil.copyfile(negative_source, staging / negative_name) + negative_declaration = { + "path": negative_name, + "sha256": sha256_file(staging / negative_name), + } + + common = { + "schema_version": CACHE_SCHEMA_VERSION, + "shards": shard_names, + "num_shards": len(shard_names), + } + if negative_declaration is not None: + common["negative_prompt_embedding"] = negative_declaration + split_specs = { + "metadata.json": ("all", [record.sample_id for record in records]), + "metadata_train.json": ("train", train_ids), + "metadata_heldout.json": ("heldout", heldout_ids), + } + for name, (split, sample_ids) in split_specs.items(): + index = { + **common, + "split": split, + "sample_ids": sample_ids, + "total_items": len(sample_ids), + } + audit_no_absolute_paths(index, context=name) + _write_json(staging / name, index) + + validate_snapshot(staging) + os.replace(staging, destination) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + return { + "output_root": str(destination.resolve()), + "total_items": len(records), + "train_items": len(train_ids), + "heldout_items": len(heldout_ids), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-root", required=True) + parser.add_argument("--output-root", required=True) + parser.add_argument( + "--source-index", + action="append", + dest="source_indexes", + help="Relative source index; repeat to finalize all rank-local preprocessing indices.", + ) + parser.add_argument("--legacy-cache-root") + parser.add_argument("--legacy-source-root") + parser.add_argument("--negative-embedding") + parser.add_argument("--heldout-count", required=True, type=int) + parser.add_argument("--split-seed", default="0") + parser.add_argument("--shard-size", default=10000, type=int) + args = parser.parse_args() + report = migrate_cache( + args.source_root, + args.output_root, + heldout_count=args.heldout_count, + split_seed=args.split_seed, + source_index=tuple(args.source_indexes) if args.source_indexes else "metadata.json", + legacy_cache_root=args.legacy_cache_root, + legacy_source_root=args.legacy_source_root, + negative_embedding=args.negative_embedding, + shard_size=args.shard_size, + ) + print(report) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/portable_cache.py b/examples/diffusers/fastgen/portable_cache.py new file mode 100644 index 00000000000..f837fe424a3 --- /dev/null +++ b/examples/diffusers/fastgen/portable_cache.py @@ -0,0 +1,292 @@ +# 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 cache paths, identities, and metadata loading for FastGen examples.""" + +from __future__ import annotations + +import hashlib +import json +import os +from collections.abc import Mapping, Sequence +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import Any + +CACHE_SCHEMA_VERSION = 1 +DATASET_CACHE_ENV = "MODELOPT_FASTGEN_DATASET_CACHE_DIR" +SAMPLE_ID_DOMAIN = "modelopt-fastgen-sample-v1" +_BANNED_PATH_KEYS = { + "cache_dir", + "image_path", + "output_dir", + "source_dir", + "source_path", + "video_path", +} +_PATH_FIELD_NAMES = { + "cache_file", + "directory", + "directories", + "file", + "files", + "path", + "paths", + "source_ref", +} + + +def resolve_cache_root(configured_root: str | os.PathLike[str]) -> Path: + """Resolve the YAML root unless a valid absolute environment override is set.""" + override = os.environ.get(DATASET_CACHE_ENV) + selected = Path(override) if override else Path(configured_root) + if override and not selected.is_absolute(): + raise ValueError(f"{DATASET_CACHE_ENV} must be an absolute path, got {override!r}.") + try: + resolved = selected.expanduser().resolve(strict=True) + except FileNotFoundError as error: + source = DATASET_CACHE_ENV if override else "configured cache_dir" + raise FileNotFoundError(f"{source} does not exist: {selected}") from error + if not resolved.is_dir(): + raise NotADirectoryError(f"dataset cache root is not a directory: {resolved}") + return resolved + + +def validate_relative_reference(value: str | os.PathLike[str], *, label: str) -> Path: + """Validate one portable, normalized relative reference without resolving it.""" + if not isinstance(value, str | os.PathLike): + raise TypeError(f"{label} must be a relative path string, got {type(value).__name__}.") + raw = os.fspath(value) + if not raw or raw == ".": + raise ValueError(f"{label} must be a non-empty relative path.") + if "\0" in raw: + raise ValueError(f"{label} must not contain NUL bytes.") + if "\\" in raw: + raise ValueError(f"{label} must use portable '/' separators, got {raw!r}.") + windows_path = PureWindowsPath(raw) + if PurePosixPath(raw).is_absolute() or windows_path.is_absolute() or windows_path.drive: + raise ValueError(f"{label} must be relative, got absolute path {raw!r}.") + path = Path(raw) + if any(part in ("", ".", "..") for part in PurePosixPath(raw).parts): + raise ValueError(f"{label} contains traversal or non-normalized components: {raw!r}.") + return path + + +def resolve_cache_asset( + root: Path, + reference: str | os.PathLike[str], + *, + label: str, + kind: str = "file", +) -> Path: + """Resolve an existing relative asset and reject traversal or symlink escape.""" + relative = validate_relative_reference(reference, label=label) + root = root.resolve(strict=True) + try: + resolved = (root / relative).resolve(strict=True) + except FileNotFoundError as error: + raise FileNotFoundError(f"{label} does not exist beneath cache root: {relative}") from error + try: + resolved.relative_to(root) + except ValueError as error: + raise ValueError(f"{label} resolves outside cache root: {relative}") from error + if kind == "file" and not resolved.is_file(): + raise ValueError(f"{label} is not a file: {relative}") + if kind == "directory" and not resolved.is_dir(): + raise ValueError(f"{label} is not a directory: {relative}") + if kind not in ("file", "directory", "any"): + raise ValueError(f"unsupported asset kind {kind!r}.") + return resolved + + +def resolve_negative_embedding(root: Path, reference: str | os.PathLike[str]) -> Path: + """Resolve a negative embedding and require it to stay beneath the cache root. + + Portable configs use a relative reference. An absolute reference is accepted only for + compatibility with existing launch overrides and only when it resolves beneath the same + effective root. + """ + raw = os.fspath(reference) + path = Path(raw).expanduser() + windows_path = PureWindowsPath(raw) + if windows_path.drive and not PurePosixPath(raw).is_absolute(): + raise ValueError("negative_prompt_embedding_path must use the host path syntax") + if not (PurePosixPath(raw).is_absolute() or windows_path.is_absolute()): + return resolve_cache_asset(root, raw, label="negative_prompt_embedding_path") + + root = root.resolve(strict=True) + try: + resolved = path.resolve(strict=True) + except FileNotFoundError as error: + raise FileNotFoundError(f"negative prompt embedding does not exist: {path}") from error + try: + resolved.relative_to(root) + except ValueError as error: + raise ValueError( + "negative prompt embedding resolves outside the effective cache root" + ) from error + if not resolved.is_file(): + raise ValueError(f"negative prompt embedding is not a file: {resolved}") + return resolved + + +def sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str: + """Return the hexadecimal SHA-256 digest of a file.""" + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def stable_sample_id(*, source_ref: str, resolution: Sequence[int], model_type: str) -> str: + """Derive a root-independent sample ID from a logical source and processing identity.""" + logical = validate_relative_reference(source_ref, label="source_ref").as_posix() + if len(resolution) != 2 or any(type(value) is not int or value <= 0 for value in resolution): + raise ValueError(f"resolution must contain two positive integers, got {resolution!r}.") + if not isinstance(model_type, str) or not model_type: + raise ValueError("model_type must be a non-empty string.") + identity = ( + f"{SAMPLE_ID_DOMAIN}\0{model_type}\0{logical}\0{resolution[0]}x{resolution[1]}" + ).encode() + return hashlib.sha256(identity).hexdigest() + + +def _is_absolute_path(value: str) -> bool: + windows_path = PureWindowsPath(value) + return ( + PurePosixPath(value).is_absolute() + or windows_path.is_absolute() + or bool(windows_path.drive) + or value.lower().startswith("file://") + ) + + +def _is_path_field(key: str) -> bool: + lowered = key.lower() + return lowered in _PATH_FIELD_NAMES or lowered.endswith( + ("_path", "_paths", "_dir", "_dirs", "_file", "_files") + ) + + +def audit_no_absolute_paths( + value: Any, + *, + context: str = "value", + _path_context: bool = False, +) -> None: + """Reject absolute locations in path fields while leaving ordinary text untouched.""" + if isinstance(value, Mapping): + for key, nested in value.items(): + if not isinstance(key, str): + raise ValueError(f"{context} contains non-string key {key!r}.") + if _is_absolute_path(key): + raise ValueError(f"{context} contains absolute path key {key!r}.") + if key.lower() in _BANNED_PATH_KEYS: + raise ValueError(f"{context} contains forbidden personal-path key {key!r}.") + audit_no_absolute_paths( + nested, + context=f"{context}.{key}", + _path_context=_path_context or _is_path_field(key), + ) + return + if isinstance(value, list | tuple | set | frozenset): + for index, nested in enumerate(value): + audit_no_absolute_paths( + nested, + context=f"{context}[{index}]", + _path_context=_path_context, + ) + return + if _path_context and isinstance(value, str) and _is_absolute_path(value): + raise ValueError(f"{context} contains absolute path {value!r}.") + + +def _load_json(path: Path, *, expected_type: type, label: str): + try: + with path.open(encoding="utf-8") as stream: + value = json.load(stream) + except json.JSONDecodeError as error: + raise ValueError(f"{label} is not valid JSON: {path}") from error + if not isinstance(value, expected_type): + raise ValueError( + f"{label} must contain {expected_type.__name__}, got {type(value).__name__}." + ) + return value + + +def load_portable_metadata( + root: Path, + metadata_index: str = "metadata.json", +) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Load and validate one split index, filtering IDs before bucket grouping.""" + index_path = resolve_cache_asset(root, metadata_index, label="metadata_index") + index = _load_json(index_path, expected_type=dict, label="metadata_index") + audit_no_absolute_paths(index, context="metadata_index") + if index.get("schema_version") != CACHE_SCHEMA_VERSION: + raise ValueError( + f"metadata_index schema_version must be {CACHE_SCHEMA_VERSION}, " + f"got {index.get('schema_version')!r}." + ) + shards = index.get("shards") + sample_ids = index.get("sample_ids") + if not isinstance(shards, list) or not shards or any(not isinstance(v, str) for v in shards): + raise ValueError("metadata_index.shards must be a non-empty list of relative paths.") + if ( + not isinstance(sample_ids, list) + or not sample_ids + or any(not isinstance(v, str) or not v for v in sample_ids) + ): + raise ValueError("metadata_index.sample_ids must be a non-empty list of strings.") + if len(sample_ids) != len(set(sample_ids)): + raise ValueError("metadata_index.sample_ids contains duplicates.") + if len(shards) != len(set(shards)): + raise ValueError("metadata_index.shards contains duplicates.") + if "num_shards" in index and index["num_shards"] != len(shards): + raise ValueError("metadata_index.num_shards must equal len(shards).") + if index.get("total_items") != len(sample_ids): + raise ValueError("metadata_index.total_items must equal len(sample_ids).") + + entries_by_id: dict[str, dict[str, Any]] = {} + for shard_number, shard_ref in enumerate(shards): + shard_path = resolve_cache_asset(root, shard_ref, label=f"shards[{shard_number}]") + entries = _load_json(shard_path, expected_type=list, label=f"shards[{shard_number}]") + for entry_number, entry in enumerate(entries): + label = f"shards[{shard_number}][{entry_number}]" + if not isinstance(entry, dict): + raise ValueError(f"{label} must be an object.") + audit_no_absolute_paths(entry, context=label) + sample_id = entry.get("sample_id") + if not isinstance(sample_id, str) or not sample_id: + raise ValueError(f"{label}.sample_id must be a non-empty string.") + if sample_id in entries_by_id: + raise ValueError(f"duplicate sample_id in shards: {sample_id}") + resolve_cache_asset(root, entry.get("cache_file"), label=f"{label}.cache_file") + payload_sha256 = entry.get("payload_sha256") + if not isinstance(payload_sha256, str) or len(payload_sha256) != 64: + raise ValueError(f"{label}.payload_sha256 must be a hexadecimal SHA-256 digest.") + try: + int(payload_sha256, 16) + except ValueError as error: + raise ValueError( + f"{label}.payload_sha256 must be a hexadecimal SHA-256 digest." + ) from error + if "source_ref" in entry: + validate_relative_reference(entry["source_ref"], label=f"{label}.source_ref") + entries_by_id[sample_id] = entry + + missing = [sample_id for sample_id in sample_ids if sample_id not in entries_by_id] + if missing: + raise ValueError(f"metadata_index references missing sample_ids: {missing[:5]}") + return index, [entries_by_id[sample_id] for sample_id in sample_ids] diff --git a/examples/diffusers/fastgen/preprocess/processors/base.py b/examples/diffusers/fastgen/preprocess/processors/base.py index b0a1fafbd52..4ba806553f9 100644 --- a/examples/diffusers/fastgen/preprocess/processors/base.py +++ b/examples/diffusers/fastgen/preprocess/processors/base.py @@ -143,7 +143,8 @@ def get_cache_data( - bucket_resolution: Tuple[int, int] - crop_offset: Tuple[int, int] - prompt: str - - image_path: str + - sample_id: str + - source_ref: Optional[str] - bucket_id: str - tier: str - aspect_ratio: float diff --git a/examples/diffusers/fastgen/preprocess/processors/qwen_image.py b/examples/diffusers/fastgen/preprocess/processors/qwen_image.py index 61749d4284b..a8f37c5c8ed 100644 --- a/examples/diffusers/fastgen/preprocess/processors/qwen_image.py +++ b/examples/diffusers/fastgen/preprocess/processors/qwen_image.py @@ -252,12 +252,14 @@ def get_cache_data( "bucket_resolution": metadata["bucket_resolution"], "crop_offset": metadata["crop_offset"], "prompt": metadata["prompt"], - "image_path": metadata["image_path"], + "sample_id": metadata["sample_id"], "bucket_id": metadata["bucket_id"], "aspect_ratio": metadata["aspect_ratio"], # Model info "model_type": self.model_type, } + if "source_ref" in metadata: + cache["source_ref"] = metadata["source_ref"] # Carry the positive-prompt attention mask through to the cache when present, so the # dataset uses the real mask instead of synthesizing an all-ones one. if "prompt_embeds_mask" in text_encodings: diff --git a/examples/diffusers/fastgen/validate_cache_snapshot.py b/examples/diffusers/fastgen/validate_cache_snapshot.py new file mode 100644 index 00000000000..38ef7031860 --- /dev/null +++ b/examples/diffusers/fastgen/validate_cache_snapshot.py @@ -0,0 +1,192 @@ +# 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. + +"""Read-only integrity validation for a portable FastGen cache snapshot.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any + +import torch +from portable_cache import ( + audit_no_absolute_paths, + load_portable_metadata, + resolve_cache_asset, + sha256_file, + validate_relative_reference, +) + + +def _validate_payload(root: Path, entry: dict[str, Any]) -> Path: + payload_path = resolve_cache_asset(root, entry["cache_file"], label="cache_file") + actual_digest = sha256_file(payload_path) + if actual_digest != entry["payload_sha256"]: + raise ValueError( + f"payload SHA-256 mismatch for {entry['sample_id']}: " + f"expected {entry['payload_sha256']}, got {actual_digest}" + ) + payload = torch.load(payload_path, map_location="cpu", weights_only=True) + if not isinstance(payload, dict): + raise TypeError(f"cache payload must be a dict: {entry['cache_file']}") + audit_no_absolute_paths(payload, context=f"payload[{entry['sample_id']}]") + if payload.get("sample_id") != entry["sample_id"]: + raise ValueError(f"payload/manifest sample_id mismatch for {entry['cache_file']}") + if payload.get("source_ref") != entry.get("source_ref"): + raise ValueError(f"payload/manifest source_ref mismatch for {entry['cache_file']}") + return payload_path + + +def validate_snapshot( + cache_root: str | Path, + *, + all_index: str = "metadata.json", + train_index: str = "metadata_train.json", + heldout_index: str = "metadata_heldout.json", + reject_orphans: bool = True, +) -> dict[str, Any]: + """Validate manifests, payload hashes, splits, and declared snapshot inventory. + + The function performs no writes. All, train, and held-out indices are required so inventory, + split-disjointness, and split-union checks cannot be skipped accidentally. + """ + root = Path(cache_root).expanduser().resolve(strict=True) + if not root.is_dir(): + raise NotADirectoryError(f"cache_root is not a directory: {root}") + + expected_indexes = { + "all": validate_relative_reference(all_index, label="all_index").as_posix(), + "train": validate_relative_reference(train_index, label="train_index").as_posix(), + "heldout": validate_relative_reference(heldout_index, label="heldout_index").as_posix(), + } + missing_indexes = [name for name in expected_indexes.values() if not (root / name).is_file()] + if missing_indexes: + raise FileNotFoundError(f"required metadata indices do not exist: {missing_indexes}") + + split_ids: dict[str, set[str]] = {} + entries_by_id: dict[str, dict[str, Any]] = {} + declared_files = { + resolve_cache_asset(root, name, label="metadata_index") + for name in expected_indexes.values() + } + negative_declarations: set[tuple[str, str]] = set() + + for expected_split, index_name in expected_indexes.items(): + index, entries = load_portable_metadata(root, index_name) + split_name = index.get("split") + if split_name != expected_split: + raise ValueError(f"{index_name}.split must be {expected_split!r}, got {split_name!r}") + if split_name in split_ids: + raise ValueError(f"duplicate split declaration: {split_name}") + split_ids[split_name] = {entry["sample_id"] for entry in entries} + + for shard_ref in index["shards"]: + declared_files.add(resolve_cache_asset(root, shard_ref, label="metadata shard")) + negative = index.get("negative_prompt_embedding") + if negative is not None: + if not isinstance(negative, dict) or set(negative) != {"path", "sha256"}: + raise ValueError( + f"{index_name}.negative_prompt_embedding must contain path and sha256" + ) + negative_path = resolve_cache_asset( + root, + negative["path"], + label=f"{index_name}.negative_prompt_embedding.path", + ) + if sha256_file(negative_path) != negative["sha256"]: + raise ValueError(f"negative prompt embedding SHA-256 mismatch in {index_name}") + negative_payload = torch.load(negative_path, map_location="cpu", weights_only=True) + audit_no_absolute_paths(negative_payload, context="negative_prompt_embedding") + negative_declarations.add((negative["path"], negative["sha256"])) + declared_files.add(negative_path) + + for entry in entries: + sample_id = entry["sample_id"] + previous = entries_by_id.get(sample_id) + if previous is not None and previous != entry: + raise ValueError(f"inconsistent manifest entry for sample_id {sample_id}") + entries_by_id[sample_id] = entry + + if len(negative_declarations) > 1: + raise ValueError("metadata indices disagree on the negative prompt embedding") + + train_ids = split_ids.get("train") + heldout_ids = split_ids.get("heldout") + if train_ids is not None and heldout_ids is not None: + overlap = train_ids & heldout_ids + if overlap: + raise ValueError(f"train and heldout splits overlap: {sorted(overlap)[:5]}") + all_ids = split_ids.get("all") + if all_ids is not None and train_ids | heldout_ids != all_ids: + raise ValueError("train and heldout split union does not equal the all split") + + payload_files = {_validate_payload(root, entry) for entry in entries_by_id.values()} + declared_files.update(payload_files) + + if reject_orphans: + actual_files = set() + for path in root.rglob("*"): + if path.is_symlink(): + resolved = path.resolve(strict=True) + try: + resolved.relative_to(root) + except ValueError as error: + raise ValueError( + f"snapshot symlink resolves outside cache root: {path}" + ) from error + raise ValueError(f"snapshot contains unsupported symlink: {path}") + if not path.is_file(): + continue + resolved = path.resolve(strict=True) + try: + resolved.relative_to(root) + except ValueError as error: + raise ValueError(f"snapshot file resolves outside cache root: {path}") from error + actual_files.add(resolved) + undeclared = sorted( + path.relative_to(root).as_posix() for path in actual_files - declared_files + ) + if undeclared: + raise ValueError(f"snapshot contains undeclared files: {undeclared[:5]}") + + return { + "root": str(root), + "indexes": list(expected_indexes.values()), + "splits": {name: len(ids) for name, ids in split_ids.items()}, + "unique_payloads": len(payload_files), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cache-root", required=True) + parser.add_argument("--all-index", default="metadata.json") + parser.add_argument("--train-index", default="metadata_train.json") + parser.add_argument("--heldout-index", default="metadata_heldout.json") + parser.add_argument("--allow-orphans", action="store_true") + args = parser.parse_args() + report = validate_snapshot( + args.cache_root, + all_index=args.all_index, + train_index=args.train_index, + heldout_index=args.heldout_index, + reject_orphans=not args.allow_orphans, + ) + print(report) + + +if __name__ == "__main__": + main() diff --git a/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py b/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py new file mode 100644 index 00000000000..699f806f7e2 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py @@ -0,0 +1,352 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import pathlib +import shutil +import subprocess +import sys + +import pytest + +torch = pytest.importorskip("torch") + +_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)) + +import migrate_cache_manifest as migration +from portable_cache import load_portable_metadata +from validate_cache_snapshot import validate_snapshot + + +def _write_json(path: pathlib.Path, value) -> None: + path.write_text(json.dumps(value, indent=2) + "\n") + + +def _make_legacy_cache( + root: pathlib.Path, + *, + stored_cache_root: pathlib.Path, + stored_source_root: pathlib.Path, + reverse_shards: bool = False, +) -> None: + root.mkdir() + (root / "legacy_payloads").mkdir() + entries = [] + for index in range(4): + relative_payload = pathlib.Path("legacy_payloads") / f"item-{index}.pt" + source_path = stored_source_root / "class" / f"item-{index}.png" + torch.save( + { + "latent": torch.full((2, 2), index, dtype=torch.float32), + "prompt_embeds": torch.full((3, 4), index, dtype=torch.float32), + "crop_offset": (0, 0), + "prompt": f"prompt {index}", + "image_path": str(source_path), + "nested": {"source_path": str(source_path)}, + }, + root / relative_payload, + ) + entries.append( + { + "cache_file": str(stored_cache_root / relative_payload), + "image_path": str(source_path), + "bucket_resolution": [64, 64], + "original_resolution": [64, 64], + "prompt": f"prompt {index}", + "bucket_id": "square-64", + "aspect_ratio": 1.0, + "pixels": 4096, + "model_type": "qwen_image", + } + ) + + shards = [entries[:2], entries[2:]] + if reverse_shards: + shards.reverse() + shard_names = [] + for index, shard in enumerate(shards): + name = f"legacy-shard-{index}.json" + _write_json(root / name, shard) + shard_names.append(name) + _write_json(root / "metadata.json", {"shards": shard_names, "total_items": 4}) + torch.save({"embed": torch.arange(12).reshape(3, 4)}, root / "legacy-negative.pt") + + +def _manifest_signature(root: pathlib.Path): + index, entries = load_portable_metadata(root, "metadata.json") + return index["sample_ids"], [ + (entry["sample_id"], entry["source_ref"], entry["cache_file"]) for entry in entries + ] + + +def test_migration_is_path_independent_and_relocatable(tmp_path): + alice = tmp_path / "alice-legacy" + bob = tmp_path / "bob-legacy" + alice_cache_prefix = pathlib.Path("/legacy/alice/cache") + bob_cache_prefix = pathlib.Path("/different/bob/cache") + alice_source_prefix = pathlib.Path("/datasets/alice/images") + bob_source_prefix = pathlib.Path("/mnt/bob/source") + _make_legacy_cache( + alice, + stored_cache_root=alice_cache_prefix, + stored_source_root=alice_source_prefix, + ) + _make_legacy_cache( + bob, + stored_cache_root=bob_cache_prefix, + stored_source_root=bob_source_prefix, + reverse_shards=True, + ) + + alice_output = tmp_path / "alice-portable" + bob_output = tmp_path / "bob-portable" + migration.migrate_cache( + alice, + alice_output, + heldout_count=1, + split_seed="fixed", + legacy_cache_root=alice_cache_prefix, + legacy_source_root=alice_source_prefix, + negative_embedding="legacy-negative.pt", + shard_size=2, + ) + migration.migrate_cache( + bob, + bob_output, + heldout_count=1, + split_seed="fixed", + legacy_cache_root=bob_cache_prefix, + legacy_source_root=bob_source_prefix, + negative_embedding="legacy-negative.pt", + shard_size=2, + ) + + assert _manifest_signature(alice_output) == _manifest_signature(bob_output) + alice_train = load_portable_metadata(alice_output, "metadata_train.json")[0]["sample_ids"] + bob_train = load_portable_metadata(bob_output, "metadata_train.json")[0]["sample_ids"] + assert alice_train == bob_train + assert validate_snapshot(alice_output)["splits"] == {"all": 4, "train": 3, "heldout": 1} + assert validate_snapshot(bob_output)["splits"] == {"all": 4, "train": 3, "heldout": 1} + + cli = subprocess.run( + [ + sys.executable, + str(_FASTGEN_DIR / "validate_cache_snapshot.py"), + "--cache-root", + str(alice_output), + "--train-index", + "metadata_train.json", + "--heldout-index", + "metadata_heldout.json", + ], + check=False, + capture_output=True, + text=True, + ) + assert cli.returncode == 0, cli.stderr + + portable_text = "".join(path.read_text() for path in alice_output.glob("*.json")) + assert str(alice_cache_prefix) not in portable_text + assert str(alice_source_prefix) not in portable_text + for payload_path in alice_output.glob("payloads/*.pt"): + payload = torch.load(payload_path, map_location="cpu", weights_only=True) + assert "image_path" not in payload + assert "source_path" not in payload["nested"] + + relocated = tmp_path / "relocated" / "cache" + relocated.parent.mkdir() + shutil.copytree(alice_output, relocated) + assert _manifest_signature(relocated) == _manifest_signature(alice_output) + validate_snapshot(relocated) + + +def test_incomplete_pass_one_publishes_nothing(monkeypatch, tmp_path): + legacy = tmp_path / "legacy" + cache_prefix = pathlib.Path("/legacy/cache") + source_prefix = pathlib.Path("/legacy/images") + _make_legacy_cache( + legacy, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + (legacy / "legacy_payloads" / "item-3.pt").unlink() + output = tmp_path / "portable" + save_calls = [] + monkeypatch.setattr(migration.torch, "save", lambda *args, **kwargs: save_calls.append(args)) + + with pytest.raises(FileNotFoundError): + migration.migrate_cache( + legacy, + output, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + assert not output.exists() + assert not list(tmp_path.glob(".portable.staging-*")) + assert save_calls == [] + + +def test_changed_source_after_frozen_plan_cleans_staging(monkeypatch, tmp_path): + legacy = tmp_path / "legacy" + cache_prefix = pathlib.Path("/legacy/cache") + source_prefix = pathlib.Path("/legacy/images") + _make_legacy_cache( + legacy, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + frozen = migration.plan_migration( + legacy, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + changed = frozen[-1].source_payload + original_plan = migration.plan_migration + + def _return_frozen(*args, **kwargs): + changed.write_bytes(b"changed after pass one") + return frozen + + monkeypatch.setattr(migration, "plan_migration", _return_frozen) + output = tmp_path / "portable" + with pytest.raises(RuntimeError, match="changed after pass 1"): + migration.migrate_cache( + legacy, + output, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + monkeypatch.setattr(migration, "plan_migration", original_plan) + assert not output.exists() + assert not list(tmp_path.glob(".portable.staging-*")) + + +def test_invalid_legacy_reference_fails_before_publish(tmp_path): + legacy = tmp_path / "legacy" + cache_prefix = pathlib.Path("/legacy/cache") + source_prefix = pathlib.Path("/legacy/images") + _make_legacy_cache( + legacy, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + shard_path = legacy / "legacy-shard-0.json" + shard = json.loads(shard_path.read_text()) + shard[0]["cache_file"] = "/other/private/cache.pt" + _write_json(shard_path, shard) + + output = tmp_path / "portable" + with pytest.raises(ValueError, match="outside the declared legacy prefix"): + migration.migrate_cache( + legacy, + output, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + assert not output.exists() + + +def test_incomplete_source_counts_and_rank_indices_publish_nothing(tmp_path): + legacy = tmp_path / "legacy" + cache_prefix = pathlib.Path("/legacy/cache") + source_prefix = pathlib.Path("/legacy/images") + _make_legacy_cache( + legacy, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + index_path = legacy / "metadata.json" + index = json.loads(index_path.read_text()) + index["num_shards"] = 3 + _write_json(index_path, index) + output = tmp_path / "invalid-count-output" + with pytest.raises(ValueError, match="num_shards"): + migration.migrate_cache( + legacy, + output, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + assert not output.exists() + assert not list(tmp_path.glob(".invalid-count-output.staging-*")) + + shards = index["shards"] + for rank, shard in enumerate(shards): + _write_json( + legacy / f"metadata_r{rank:02d}.json", + { + "shards": [shard], + "num_shards": 1, + "total_items": 2, + "shard_rank": rank, + "shard_world": 2, + }, + ) + incomplete_output = tmp_path / "incomplete-ranks-output" + with pytest.raises(ValueError, match="incomplete or inconsistent"): + migration.migrate_cache( + legacy, + incomplete_output, + source_index="metadata_r00.json", + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + assert not incomplete_output.exists() + + complete_output = tmp_path / "complete-ranks-output" + migration.migrate_cache( + legacy, + complete_output, + source_index=("metadata_r00.json", "metadata_r01.json"), + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + assert validate_snapshot(complete_output)["splits"] == {"all": 4, "train": 3, "heldout": 1} + + +def test_migration_path_audit_allows_prompt_commands_but_rejects_set_paths(tmp_path): + legacy = tmp_path / "legacy" + cache_prefix = pathlib.Path("/legacy/cache") + source_prefix = pathlib.Path("/legacy/images") + _make_legacy_cache( + legacy, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + shard_path = legacy / "legacy-shard-0.json" + shard = json.loads(shard_path.read_text()) + shard[0]["prompt"] = "/imagine a cat" + _write_json(shard_path, shard) + payload_path = legacy / "legacy_payloads" / "item-0.pt" + payload = torch.load(payload_path, map_location="cpu", weights_only=True) + payload["prompt"] = "/imagine a cat" + torch.save(payload, payload_path) + migration.plan_migration( + legacy, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + + payload["paths"] = {"/home/alice/private.png"} + torch.save(payload, payload_path) + output = tmp_path / "portable" + with pytest.raises(ValueError, match="absolute path"): + migration.migrate_cache( + legacy, + output, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + assert not output.exists() diff --git a/tests/examples/diffusers/fastgen/test_portable_cache.py b/tests/examples/diffusers/fastgen/test_portable_cache.py new file mode 100644 index 00000000000..487cd55ac76 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_portable_cache.py @@ -0,0 +1,418 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import pathlib +import shutil +import sys +import types + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("nemo_automodel") + +_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 TextToImageDataset, collate_fn_text_to_image +from portable_cache import ( + DATASET_CACHE_ENV, + audit_no_absolute_paths, + load_portable_metadata, + resolve_cache_root, + resolve_negative_embedding, + sha256_file, +) +from validate_cache_snapshot import validate_snapshot + + +def _sample_id(source_ref: str, resolution: tuple[int, int]) -> str: + identity = ( + f"modelopt-fastgen-sample-v1\0qwen_image\0{source_ref}\0{resolution[0]}x{resolution[1]}" + ) + return hashlib.sha256(identity.encode()).hexdigest() + + +def _write_json(path: pathlib.Path, value) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n") + + +def _make_snapshot(root: pathlib.Path) -> dict[str, list[str]]: + root.mkdir() + entries = [] + for index, resolution in enumerate(((64, 64), (64, 64), (128, 64), (128, 64))): + source_ref = f"class/{chr(ord('a') + index)}.png" + sample_id = _sample_id(source_ref, resolution) + cache_ref = f"payloads/{sample_id}.pt" + payload_path = root / cache_ref + payload_path.parent.mkdir(exist_ok=True) + torch.save( + { + "latent": torch.full((2, 2, 2), index, dtype=torch.float32), + "prompt_embeds": torch.full((3, 4), index, dtype=torch.float32), + "prompt_embeds_mask": torch.ones(3, dtype=torch.long), + "crop_offset": (0, 0), + "prompt": f"prompt {index}", + "sample_id": sample_id, + "source_ref": source_ref, + }, + payload_path, + ) + entries.append( + { + "sample_id": sample_id, + "source_ref": source_ref, + "cache_file": cache_ref, + "payload_sha256": sha256_file(payload_path), + "bucket_resolution": list(resolution), + "original_resolution": list(resolution), + "bucket_id": f"bucket-{resolution[0]}", + "aspect_ratio": resolution[0] / resolution[1], + } + ) + + _write_json(root / "metadata_shard_s0000.json", entries) + all_ids = [entry["sample_id"] for entry in entries] + splits = { + "all": all_ids, + "train": [all_ids[2], all_ids[0], all_ids[1]], + "heldout": [all_ids[3]], + } + for split, ids in splits.items(): + name = "metadata.json" if split == "all" else f"metadata_{split}.json" + _write_json( + root / name, + { + "schema_version": 1, + "split": split, + "total_items": len(ids), + "num_shards": 1, + "shards": ["metadata_shard_s0000.json"], + "sample_ids": ids, + }, + ) + negative_path = root / "negative_prompt_embedding.pt" + torch.save({"embed": torch.arange(12).reshape(3, 4)}, negative_path) + negative_declaration = { + "path": negative_path.name, + "sha256": sha256_file(negative_path), + } + for name in ("metadata.json", "metadata_train.json", "metadata_heldout.json"): + index = json.loads((root / name).read_text()) + index["negative_prompt_embedding"] = negative_declaration + _write_json(root / name, index) + return splits + + +def _batch_signature(dataset: TextToImageDataset) -> list[tuple[str, float]]: + return [ + (dataset[index]["sample_id"], dataset[index]["latent"].sum().item()) + for index in range(len(dataset)) + ] + + +def test_cache_root_environment_precedence(monkeypatch, tmp_path): + configured = tmp_path / "configured" + override = tmp_path / "override" + configured.mkdir() + override.mkdir() + + monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) + assert resolve_cache_root(configured) == configured.resolve() + monkeypatch.setenv(DATASET_CACHE_ENV, "") + assert resolve_cache_root(configured) == configured.resolve() + monkeypatch.setenv(DATASET_CACHE_ENV, str(override)) + assert resolve_cache_root(configured) == override.resolve() + monkeypatch.setenv(DATASET_CACHE_ENV, "relative/cache") + with pytest.raises(ValueError, match="absolute"): + resolve_cache_root(configured) + monkeypatch.setenv(DATASET_CACHE_ENV, str(tmp_path / "missing")) + with pytest.raises(FileNotFoundError): + resolve_cache_root(configured) + not_a_directory = tmp_path / "cache-file" + not_a_directory.write_text("not a directory") + monkeypatch.setenv(DATASET_CACHE_ENV, str(not_a_directory)) + with pytest.raises(NotADirectoryError): + resolve_cache_root(configured) + + +def test_relocation_preserves_order_payloads_and_buckets(monkeypatch, tmp_path): + first = tmp_path / "alice" / "cache" + second = tmp_path / "bob" / "cache" + first.parent.mkdir() + splits = _make_snapshot(first) + second.parent.mkdir() + shutil.copytree(first, second) + + monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) + first_dataset = TextToImageDataset(str(first), metadata_index="metadata_train.json") + second_dataset = TextToImageDataset(str(second), metadata_index="metadata_train.json") + assert _batch_signature(first_dataset) == _batch_signature(second_dataset) + assert [entry["sample_id"] for entry in second_dataset.metadata] == splits["train"] + assert first_dataset.bucket_groups == second_dataset.bucket_groups + + monkeypatch.setenv(DATASET_CACHE_ENV, str(second)) + overridden = TextToImageDataset(str(first), metadata_index="metadata_train.json") + assert overridden.cache_dir == second.resolve() + assert _batch_signature(overridden) == _batch_signature(first_dataset) + + +def test_split_filters_before_inherited_bucket_grouping(monkeypatch, tmp_path): + root = tmp_path / "cache" + splits = _make_snapshot(root) + shard_path = root / "metadata_shard_s0000.json" + entries = json.loads(shard_path.read_text()) + heldout_entry = next(entry for entry in entries if entry["sample_id"] in splits["heldout"]) + del heldout_entry["bucket_resolution"] + _write_json(shard_path, entries) + + monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) + dataset = TextToImageDataset(str(root), metadata_index="metadata_train.json") + assert [entry["sample_id"] for entry in dataset.metadata] == splits["train"] + grouped = [index for bucket in dataset.bucket_groups.values() for index in bucket["indices"]] + assert sorted(grouped) == list(range(len(splits["train"]))) + + +@pytest.mark.parametrize( + "reference", + ["/etc/passwd", "C:/secret.pt", "C:secret.pt", "../escape.pt", "a/../b.pt"], +) +def test_manifest_rejects_nonportable_payload_references(monkeypatch, tmp_path, reference): + root = tmp_path / "cache" + _make_snapshot(root) + shard_path = root / "metadata_shard_s0000.json" + entries = json.loads(shard_path.read_text()) + entries[0]["cache_file"] = reference + _write_json(shard_path, entries) + + monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) + with pytest.raises((ValueError, FileNotFoundError)): + load_portable_metadata(root, "metadata.json") + + +def test_manifest_rejects_missing_and_symlink_escape(monkeypatch, tmp_path): + root = tmp_path / "cache" + _make_snapshot(root) + outside = tmp_path / "outside.pt" + torch.save({}, outside) + shard_path = root / "metadata_shard_s0000.json" + original = json.loads(shard_path.read_text()) + + monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) + for reference in ("payloads/missing.pt", "payloads/escape.pt"): + entries = json.loads(json.dumps(original)) + entries[0]["cache_file"] = reference + if reference.endswith("escape.pt"): + (root / reference).symlink_to(outside) + _write_json(shard_path, entries) + with pytest.raises((ValueError, FileNotFoundError)): + load_portable_metadata(root, "metadata.json") + + +def test_negative_embedding_is_resolved_under_effective_root(tmp_path): + root = tmp_path / "cache" + _make_snapshot(root) + assert ( + resolve_negative_embedding(root, "negative_prompt_embedding.pt") + == (root / "negative_prompt_embedding.pt").resolve() + ) + assert resolve_negative_embedding(root, root / "negative_prompt_embedding.pt").is_file() + + outside = tmp_path / "outside.pt" + torch.save({}, outside) + with pytest.raises(ValueError, match="outside"): + resolve_negative_embedding(root, outside) + (root / "negative_escape.pt").symlink_to(outside) + with pytest.raises(ValueError, match="outside"): + resolve_negative_embedding(root, "negative_escape.pt") + + +def test_collate_emits_logical_identity_without_source_paths(tmp_path): + root = tmp_path / "cache" + _make_snapshot(root) + dataset = TextToImageDataset(str(root), metadata_index="metadata_train.json") + samples = [dataset[1], dataset[2]] + output = collate_fn_text_to_image(samples) + assert output["metadata"]["sample_ids"] == [item["sample_id"] for item in samples] + assert output["metadata"]["source_refs"] == [item["source_ref"] for item in samples] + assert "image_paths" not in output["metadata"] + assert str(root) not in repr(output) + + +def test_validator_detects_hash_split_and_orphan_failures(tmp_path): + root = tmp_path / "cache" + splits = _make_snapshot(root) + assert validate_snapshot(root)["unique_payloads"] == 4 + + orphan = root / "payloads" / "orphan.pt" + torch.save({"image_path": "/home/alice/private.png"}, orphan) + with pytest.raises(ValueError, match="undeclared"): + validate_snapshot(root) + orphan.unlink() + + heldout_path = root / "metadata_heldout.json" + heldout = json.loads(heldout_path.read_text()) + heldout["sample_ids"] = [splits["train"][0]] + _write_json(heldout_path, heldout) + with pytest.raises(ValueError, match="overlap"): + validate_snapshot(root) + + +def test_validator_requires_complete_splits_and_rejects_directory_symlink(tmp_path): + root = tmp_path / "cache" + _make_snapshot(root) + (root / "metadata_heldout.json").unlink() + with pytest.raises(FileNotFoundError, match="required metadata indices"): + validate_snapshot(root) + + _make_snapshot(tmp_path / "complete") + complete = tmp_path / "complete" + outside = tmp_path / "outside-directory" + outside.mkdir() + (complete / "escaped-directory").symlink_to(outside, target_is_directory=True) + with pytest.raises(ValueError, match="symlink resolves outside"): + validate_snapshot(complete) + + +def test_recursive_path_audit_is_schema_aware_and_covers_containers(): + audit_no_absolute_paths({"prompt": "/imagine a cat"}) + with pytest.raises(ValueError, match="absolute path"): + audit_no_absolute_paths({"paths": {"/home/alice/private.png"}}) + with pytest.raises(ValueError, match="absolute path"): + audit_no_absolute_paths({"paths": {"primary": "/home/alice/private.png"}}) + with pytest.raises(ValueError, match="absolute path"): + audit_no_absolute_paths({"source_file": {"value": "/lustre/private.pt"}}) + with pytest.raises(ValueError, match="absolute path key"): + audit_no_absolute_paths({"/home/alice/private.png": "value"}) + + +def test_validator_detects_payload_hash_mismatch_and_is_read_only(tmp_path): + root = tmp_path / "cache" + _make_snapshot(root) + before = { + path.relative_to(root).as_posix(): (sha256_file(path), path.stat().st_mtime_ns) + for path in root.rglob("*") + if path.is_file() + } + validate_snapshot(root) + after = { + path.relative_to(root).as_posix(): (sha256_file(path), path.stat().st_mtime_ns) + for path in root.rglob("*") + if path.is_file() + } + assert after == before + + first_payload = next((root / "payloads").glob("*.pt")) + first_payload.write_bytes(first_payload.read_bytes() + b"tampered") + with pytest.raises(ValueError, match="SHA-256 mismatch"): + validate_snapshot(root) + + +def test_empty_split_is_rejected_before_bucket_grouping(tmp_path): + root = tmp_path / "cache" + _make_snapshot(root) + _write_json( + root / "metadata_empty.json", + { + "schema_version": 1, + "split": "empty", + "total_items": 0, + "shards": ["metadata_shard_s0000.json"], + "sample_ids": [], + }, + ) + with pytest.raises(ValueError, match="non-empty"): + TextToImageDataset(str(root), metadata_index="metadata_empty.json") + + +def test_qwen_preprocessor_payload_is_sanitized(): + from preprocess.processors.qwen_image import QwenImageProcessor + + processor = QwenImageProcessor() + metadata = { + "original_resolution": (64, 64), + "bucket_resolution": (64, 64), + "crop_offset": (0, 0), + "prompt": "portable", + "sample_id": "b" * 64, + "source_ref": "class/b.png", + "bucket_id": "square-64", + "aspect_ratio": 1.0, + } + payload = processor.get_cache_data( + torch.zeros(2, 2, 2), + {"prompt_embeds": torch.zeros(1, 3, 4)}, + metadata, + ) + assert payload["sample_id"] == metadata["sample_id"] + assert payload["source_ref"] == metadata["source_ref"] + assert "image_path" not in payload + + +def test_portable_index_writer_is_deterministic(monkeypatch, tmp_path): + try: + import cv2 # noqa: F401 + except ImportError: + monkeypatch.setitem(sys.modules, "cv2", types.ModuleType("cv2")) + from migrate_cache_manifest import migrate_cache + from preprocess.preprocessing_multiprocess import _save_metadata_shards + + staging = tmp_path / "staging" + (staging / "payloads").mkdir(parents=True) + entries = [] + for character in ("b", "a"): + source_ref = f"class/{character}.png" + sample_id = _sample_id(source_ref, (64, 64)) + payload_ref = f"payloads/{sample_id}.pt" + payload_path = staging / payload_ref + torch.save( + { + "latent": torch.zeros(2, 2), + "crop_offset": (0, 0), + "prompt": character, + "sample_id": sample_id, + "source_ref": source_ref, + }, + payload_path, + ) + entries.append( + { + "sample_id": sample_id, + "source_ref": source_ref, + "cache_file": payload_ref, + "payload_sha256": sha256_file(payload_path), + "bucket_resolution": [64, 64], + "original_resolution": [64, 64], + "prompt": character, + "bucket_id": "square-64", + "aspect_ratio": 1.0, + "model_type": "qwen_image", + } + ) + _save_metadata_shards( + entries, + staging, + "qwen_image", + "Qwen/Qwen-Image", + "qwen_image", + 10, + {}, + portable=True, + ) + index = json.loads((staging / "metadata.json").read_text()) + shard = json.loads((staging / index["shards"][0]).read_text()) + expected_ids = sorted(entry["sample_id"] for entry in entries) + assert index["sample_ids"] == expected_ids + assert [entry["sample_id"] for entry in shard] == expected_ids + assert not (staging / "metadata_train.json").exists() + assert str(staging) not in json.dumps([index, shard]) + + finalized = tmp_path / "finalized" + migrate_cache(staging, finalized, heldout_count=1) + assert validate_snapshot(finalized)["splits"] == {"all": 2, "train": 1, "heldout": 1} diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index 6977fd2c1e0..633cbd7b668 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -172,7 +172,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 +180,8 @@ 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", + "sample_id": "sample-0", + "source_ref": "images/img.png", "bucket_id": 0, "aspect_ratio": 1.0, "prompt_embeds": torch.randn(seq, dim), From 4d797186d3266d58c65ddd975039232244c9dfa1 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 07:18:42 -0700 Subject: [PATCH 10/45] feat(fastgen): add PDD AutoModel recipe setup Signed-off-by: Meng Xin --- examples/diffusers/fastgen/README.md | 1 + .../fastgen/automodel_dependency.json | 14 + .../fastgen/pdd/configs/qwen_image.yaml | 54 +++ examples/diffusers/fastgen/pdd_finetune.py | 54 +++ examples/diffusers/fastgen/pdd_recipe.py | 422 ++++++++++++++++++ examples/diffusers/fastgen/requirements.txt | 12 +- .../fastgen/verify_readonly_automodel.py | 197 ++++++++ .../fastgen/test_pdd_recipe_setup.py | 244 ++++++++++ 8 files changed, 991 insertions(+), 7 deletions(-) create mode 100644 examples/diffusers/fastgen/automodel_dependency.json create mode 100644 examples/diffusers/fastgen/pdd/configs/qwen_image.yaml create mode 100644 examples/diffusers/fastgen/pdd_finetune.py create mode 100644 examples/diffusers/fastgen/pdd_recipe.py create mode 100644 examples/diffusers/fastgen/verify_readonly_automodel.py create mode 100644 tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md index 1441b034f54..ccbe3b9119e 100644 --- a/examples/diffusers/fastgen/README.md +++ b/examples/diffusers/fastgen/README.md @@ -4,6 +4,7 @@ This directory contains training and inference examples for diffusion distillati `modelopt.torch.fastgen`. - [DMD2 for Qwen-Image](dmd2/README.md) +- PDD for Qwen-Image (integration in progress under `pdd/`) 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/automodel_dependency.json b/examples/diffusers/fastgen/automodel_dependency.json new file mode 100644 index 00000000000..4cb0535dbf9 --- /dev/null +++ b/examples/diffusers/fastgen/automodel_dependency.json @@ -0,0 +1,14 @@ +{ + "distribution": "nemo_automodel", + "import_name": "nemo_automodel", + "package_file_count": 490, + "package_tree_sha256": "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528", + "release_commit": "d02f49cb314554715aabb97e8dba6599c9f6e9e0", + "release_tag": "v0.5.0", + "runtime_versions": { + "diffusers": "0.38.0" + }, + "version": "0.5.0", + "wheel": "nemo_automodel-0.5.0-py3-none-any.whl", + "wheel_sha256": "881aebafc5145752842afbbfe0a42e1c33d06847c3e418ad3d6f154ddc8e0f45" +} 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..abd8a208e40 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -0,0 +1,54 @@ +# Qwen-Image PDD setup skeleton. Task 8 adds the training lifecycle. + +model: + pretrained_model_name_or_path: Qwen/Qwen-Image + revision: 75e0b4be04f60ec59a75f475837eced720f823b6 + torch_dtype: bfloat16 + device: cuda + transformer_engine_linear: false + peft: + guidance_embeds: false + # Accepted by the composition path, but disabled until the pinned Diffusers Qwen API + # performs an effective fusion rather than its current no-op. + fuse_qkv_projections: false + +pdd: + pred_type: flow + num_train_timesteps: + guidance_scale: 4.0 + student_sample_steps: 4 + student_sample_type: ode + grid_size: 128 + flow_shift: 5.0 + block_size_min: 4 + block_size_max: 64 + teacher_integrator: euler + inference_blocks: [32, 32, 32, 32] + data_free: false + +optim: + learning_rate: 2.0e-5 + weight_decay: 0.01 + +fsdp: + dp_size: + tp_size: 1 + cp_size: 1 + pp_size: 1 + ep_size: 1 + activation_checkpointing: true + +data: + dataloader: + _target_: fastgen_data.build_text_to_image_multiresolution_dataloader + cache_dir: data/qwen_image_cache + metadata_index: metadata_train.json + base_resolution: [1024, 1024] + batch_size: 1 + negative_prompt_embedding_path: negative_prompt_embedding.pt + +checkpoint: + enabled: true + checkpoint_dir: checkpoints/pdd_qwen_image + model_save_format: torch_save + save_consolidated: false diff --git a/examples/diffusers/fastgen/pdd_finetune.py b/examples/diffusers/fastgen/pdd_finetune.py new file mode 100644 index 00000000000..96e8f60bd72 --- /dev/null +++ b/examples/diffusers/fastgen/pdd_finetune.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build the released-AutoModel Qwen-Image PDD setup owned by ModelOpt.""" + +from __future__ import annotations + +import argparse +import logging +import sys +from pathlib import Path + +import yaml + +sys.dont_write_bytecode = True + +_THIS_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _THIS_DIR.parents[2] +for path in (_REPO_ROOT, _THIS_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + type=Path, + default=_THIS_DIR / "configs" / "pdd_qwen_image.yaml", + ) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + from pdd_recipe import build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config + + raw = yaml.safe_load(args.config.read_text()) + config = resolve_pdd_recipe_config(raw) + initialize_pdd_distributed( + backend="nccl" if config.device.type == "cuda" else "gloo", + timeout_minutes=60, + ) + setup = build_pdd_setup(config) + logging.info( + "PDD setup complete: lifecycle=%s student_keys=%d AutoModel=%s", + setup.lifecycle, + len(setup.checkpoint_keys), + setup.automodel_snapshot["version"], + ) + + +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..d4f4e7b8f16 --- /dev/null +++ b/examples/diffusers/fastgen/pdd_recipe.py @@ -0,0 +1,422 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ModelOpt-owned construction of a Qwen-Image PDD student and frozen teacher.""" + +from __future__ import annotations + +import copy +import logging +import math +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from torch import nn +from verify_readonly_automodel import snapshot_installed_distribution + +from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection +from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd + + +@dataclass(frozen=True) +class PDDParallelConfig: + """Pure-data-parallel FSDP2 settings for the first Qwen-Image example.""" + + dp_size: int | None = None + activation_checkpointing: bool = False + + +@dataclass(frozen=True) +class PDDCheckpointConfig: + """AutoModel Checkpointer settings needed by the PDD lifecycle.""" + + checkpoint_dir: str = "checkpoints/pdd_qwen_image" + enabled: bool = True + model_save_format: str = "torch_save" + save_consolidated: bool = False + + +@dataclass(frozen=True) +class PDDRecipeConfig: + """Resolved setup inputs; incompatible mutation modes have already been rejected.""" + + model_id: str + model_revision: str | None + pdd: PDDConfig + parallel: PDDParallelConfig + checkpoint: PDDCheckpointConfig + learning_rate: float + weight_decay: float + device: torch.device + dtype: torch.dtype + fuse_qkv_projections: bool + + +@dataclass(frozen=True) +class PDDSetupArtifacts: + """Objects produced in the required load-to-checkpoint construction order.""" + + pipe: Any + student: nn.Module + teacher: nn.Module + projection: PDDOutputProjection + optimizer: torch.optim.Optimizer + distributed_setup: Any + fsdp_manager: Any + checkpointer: Any + metadata: PDDMetadata + checkpoint_keys: tuple[str, ...] + lifecycle: tuple[str, ...] + automodel_snapshot: Mapping[str, Any] + + +def _as_mapping(value: Any, *, name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{name} must be a mapping, got {type(value).__name__}.") + return value + + +def _reject_enabled(value: Any, *, name: str) -> None: + if value is None or value is False or value == {}: + return + raise ValueError(f"PDD does not support {name}; disable it before model loading.") + + +def _require_bool(value: Any, *, name: str) -> bool: + if type(value) is not bool: + raise TypeError(f"{name} must be bool.") + return value + + +def _resolve_dtype(value: Any) -> torch.dtype: + if isinstance(value, torch.dtype): + return value + if not isinstance(value, str): + raise TypeError("model.torch_dtype must be a torch dtype name.") + dtypes = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + } + try: + return dtypes[value] + except KeyError as error: + raise ValueError( + f"Unsupported model.torch_dtype={value!r}; expected {sorted(dtypes)}." + ) from error + + +def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: + """Resolve PDD and reject TE-linear, PEFT, and guidance embeddings before loading.""" + raw = _as_mapping(raw, name="config") + model = _as_mapping(raw.get("model"), name="model") + pdd_raw = _as_mapping(raw.get("pdd"), name="pdd") + fsdp = _as_mapping(raw.get("fsdp", {}), name="fsdp") + optim = _as_mapping(raw.get("optim", {}), name="optim") + checkpoint = _as_mapping(raw.get("checkpoint", {}), name="checkpoint") + + _reject_enabled(model.get("transformer_engine_linear"), name="global TE-linear conversion") + _reject_enabled(model.get("peft"), name="PEFT/LoRA") + _reject_enabled(model.get("peft_cfg"), name="PEFT/LoRA") + _reject_enabled(raw.get("peft"), name="PEFT/LoRA") + _reject_enabled(raw.get("peft_cfg"), name="PEFT/LoRA") + _reject_enabled(model.get("guidance_embeds"), name="Qwen guidance embeddings") + _reject_enabled(model.get("guidance_embeddings"), name="Qwen guidance embeddings") + for option in ( + "device_map", + "load_in_4bit", + "load_in_8bit", + "offload_folder", + "offload_state_dict", + "quantization_config", + ): + _reject_enabled(model.get(option), name=f"model loader option {option!r}") + + model_id = model.get("pretrained_model_name_or_path") + if not isinstance(model_id, str) or not model_id: + raise ValueError("model.pretrained_model_name_or_path must be a non-empty string.") + model_revision = model.get("revision") + if model_revision is not None and ( + not isinstance(model_revision, str) + or len(model_revision) != 40 + or any(character not in "0123456789abcdefABCDEF" for character in model_revision) + ): + raise ValueError("model.revision must be null or a full 40-character commit hash.") + if not Path(model_id).is_dir(): + if model_revision is None: + raise ValueError("Remote PDD models require an exact model.revision commit hash.") + learning_rate = optim.get("learning_rate", 2.0e-5) + weight_decay = optim.get("weight_decay", 0.0) + if isinstance(learning_rate, bool) or not isinstance(learning_rate, int | float): + raise TypeError("optim.learning_rate must be a real number.") + if isinstance(weight_decay, bool) or not isinstance(weight_decay, int | float): + raise TypeError("optim.weight_decay must be a real number.") + if not math.isfinite(learning_rate) or not math.isfinite(weight_decay): + raise ValueError("optim.learning_rate and weight_decay must be finite.") + if learning_rate <= 0 or weight_decay < 0: + raise ValueError("optim.learning_rate must be > 0 and weight_decay must be >= 0.") + + dp_size = fsdp.get("dp_size") + if dp_size is not None and (type(dp_size) is not int or dp_size < 1): + raise ValueError("fsdp.dp_size must be null or an integer >= 1.") + activation_checkpointing = fsdp.get("activation_checkpointing", False) + if type(activation_checkpointing) is not bool: + raise TypeError("fsdp.activation_checkpointing must be bool.") + for dimension in ("tp_size", "cp_size", "pp_size", "ep_size"): + value = fsdp.get(dimension, 1) + if type(value) is not int or value != 1: + raise ValueError(f"PDD v1 supports pure data parallelism; fsdp.{dimension} must be 1.") + + pdd = PDDConfig(**dict(pdd_raw)) + if pdd.num_train_timesteps is not None: + raise ValueError("Qwen-Image PDD requires pdd.num_train_timesteps=null.") + + checkpoint_enabled = _require_bool(checkpoint.get("enabled", True), name="checkpoint.enabled") + save_consolidated = _require_bool( + checkpoint.get("save_consolidated", False), name="checkpoint.save_consolidated" + ) + if save_consolidated: + raise ValueError("PDD training checkpoints require checkpoint.save_consolidated=false.") + checkpoint_dir = checkpoint.get("checkpoint_dir", "checkpoints/pdd_qwen_image") + if not isinstance(checkpoint_dir, str) or not checkpoint_dir: + raise ValueError("checkpoint.checkpoint_dir must be a non-empty string.") + model_save_format = checkpoint.get("model_save_format", "torch_save") + if model_save_format != "torch_save": + raise ValueError("PDD training checkpoints require model_save_format='torch_save'.") + fuse_qkv_projections = _require_bool( + model.get("fuse_qkv_projections", False), name="model.fuse_qkv_projections" + ) + + return PDDRecipeConfig( + model_id=model_id, + model_revision=model_revision, + pdd=pdd, + parallel=PDDParallelConfig( + dp_size=dp_size, + activation_checkpointing=activation_checkpointing, + ), + checkpoint=PDDCheckpointConfig( + checkpoint_dir=checkpoint_dir, + enabled=checkpoint_enabled, + model_save_format=model_save_format, + save_consolidated=save_consolidated, + ), + learning_rate=float(learning_rate), + weight_decay=float(weight_decay), + device=torch.device(model.get("device", "cuda" if torch.cuda.is_available() else "cpu")), + dtype=_resolve_dtype(model.get("torch_dtype", "bfloat16")), + fuse_qkv_projections=fuse_qkv_projections, + ) + + +def _projection_identity(projection: PDDOutputProjection) -> tuple[int, int, int | None]: + return ( + id(projection), + id(projection.weight), + None if projection.bias is None else id(projection.bias), + ) + + +def _require_projection_identity( + student: nn.Module, + projection: PDDOutputProjection, + expected: tuple[int, int, int | None], + *, + stage: str, +) -> None: + if student.get_submodule("proj_out") is not projection: + raise RuntimeError(f"PDD projection was replaced during {stage}.") + if _projection_identity(projection) != expected: + raise RuntimeError(f"PDD projection parameter identity changed during {stage}.") + + +def _require_projection_module( + student: nn.Module, + projection: PDDOutputProjection, + *, + stage: str, +) -> None: + if student.get_submodule("proj_out") is not projection: + raise RuntimeError(f"PDD projection module was replaced during {stage}.") + + +def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: + """Compose released AutoModel APIs without editing or patching external packages.""" + if not isinstance(config, PDDRecipeConfig): + raise TypeError(f"config must be PDDRecipeConfig, got {type(config).__name__}.") + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("Initialize torch.distributed before building the PDD FSDP2 setup.") + + automodel_snapshot = snapshot_installed_distribution() + lifecycle: list[str] = [] + + # Imports are intentionally delayed until the exact installed wheel has passed verification. + from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.checkpoint.config import CheckpointingConfig + from nemo_automodel.components.distributed import ( + DistributedSetup, + FSDP2Config, + ParallelismSizes, + ) + from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager + + if Path(config.model_id).is_dir(): + model_source = config.model_id + else: + from huggingface_hub import snapshot_download + + if config.model_revision is None: + raise ValueError("Remote PDD models require a pinned model revision.") + model_source = snapshot_download(config.model_id, revision=config.model_revision) + if Path(model_source).resolve().name != config.model_revision: + raise RuntimeError( + "Hugging Face resolved a model snapshot that does not match the pinned revision." + ) + + pipe, loader_managers = NeMoAutoDiffusionPipeline.from_pretrained( + model_source, + parallel_scheme=None, + device=None, + torch_dtype=config.dtype, + move_to_device=False, + load_for_training=True, + components_to_load=["transformer"], + peft_cfg=None, + active_transformer="transformer", + transformer_engine_linear=False, + fuse_qkv_projections=False, + compact_fused_qkv_projections=False, + low_cpu_mem_usage=True, + text_encoder=None, + tokenizer=None, + vae=None, + ) + if loader_managers: + raise RuntimeError("Unwrapped AutoModel load unexpectedly created parallel managers.") + student = pipe.transformer + if not isinstance(student, nn.Module): + raise TypeError("AutoModel pipeline did not return an nn.Module transformer.") + teacher = copy.deepcopy(student).eval().requires_grad_(False) + lifecycle.append("load/select") + + projection = convert_qwen_image_to_pdd(student, config.pdd) + identity = _projection_identity(projection) + metadata = PDDMetadata.from_config(config.pdd, projection) + lifecycle.append("pdd_conversion") + + student.to(device=config.device, dtype=config.dtype) + teacher.to(device=config.device, dtype=config.dtype) + _require_projection_identity(student, projection, identity, stage="device placement") + lifecycle.append("device") + + if config.fuse_qkv_projections: + if not hasattr(student, "fuse_qkv_projections") or not hasattr( + teacher, "fuse_qkv_projections" + ): + raise AttributeError( + "QKV fusion requires both Qwen transformers to expose the object API." + ) + student.fuse_qkv_projections() + teacher.fuse_qkv_projections() + if not any(getattr(module, "fused_projections", False) for module in student.modules()): + logging.warning( + "Qwen fuse_qkv_projections() was accepted but produced no fused attention " + "modules in the pinned Diffusers release." + ) + _require_projection_identity(student, projection, identity, stage="QKV fusion") + lifecycle.append("qkv") + + world_size = dist.get_world_size() + dp_size = config.parallel.dp_size or world_size + if dp_size != world_size: + raise ValueError( + f"Pure-DP PDD requires fsdp.dp_size ({dp_size}) to equal world size ({world_size})." + ) + strategy = FSDP2Config(activation_checkpointing=config.parallel.activation_checkpointing) + distributed_setup = DistributedSetup.build( + strategy=strategy, + parallelism_sizes=ParallelismSizes(dp_size=dp_size), + activation_checkpointing=config.parallel.activation_checkpointing, + world_size=world_size, + ) + mesh_context = distributed_setup.mesh_context + manager = FSDP2Manager( + distributed_setup.strategy_config, + device_mesh=mesh_context.device_mesh, + moe_mesh=mesh_context.moe_mesh, + ) + student = manager.parallelize(student) + teacher = manager.parallelize(teacher) + pipe.transformer = student + # FSDP2 shards Parameters in place and may replace the Parameter objects. The registered + # projection module and FQN must survive; optimizer identity is checked against the new, + # live post-FSDP Parameters below. + _require_projection_module(student, projection, stage="FSDP2 parallelization") + lifecycle.append("parallelize") + + trainable = [parameter for parameter in student.parameters() if parameter.requires_grad] + if not trainable: + raise RuntimeError("PDD student has no trainable parameters after FSDP2 setup.") + if any(parameter.requires_grad for parameter in teacher.parameters()): + raise RuntimeError("PDD teacher became trainable during setup.") + optimizer = torch.optim.AdamW( + trainable, + lr=config.learning_rate, + weight_decay=config.weight_decay, + ) + optimizer_parameters = [ + parameter for group in optimizer.param_groups for parameter in group["params"] + ] + if not any(parameter is projection.weight for parameter in optimizer_parameters): + raise RuntimeError("PDD projection parameters are missing from the optimizer.") + lifecycle.append("optimizer") + + checkpoint_keys = tuple(student.state_dict().keys()) + projection_key = "proj_out.weight" + if projection_key not in checkpoint_keys: + raise RuntimeError( + f"PDD projection key {projection_key!r} is missing from checkpoint state." + ) + checkpoint_config = CheckpointingConfig( + enabled=config.checkpoint.enabled, + checkpoint_dir=config.checkpoint.checkpoint_dir, + model_save_format=config.checkpoint.model_save_format, + model_repo_id=config.model_id, + save_consolidated=config.checkpoint.save_consolidated, + is_peft=False, + model_state_dict_keys=list(checkpoint_keys), + ) + checkpointer = checkpoint_config.build( + dp_rank=dist.get_rank(), + tp_rank=0, + pp_rank=0, + moe_mesh=None, + ) + lifecycle.append("checkpoint") + + return PDDSetupArtifacts( + pipe=pipe, + student=student, + teacher=teacher, + projection=projection, + optimizer=optimizer, + distributed_setup=distributed_setup, + fsdp_manager=manager, + checkpointer=checkpointer, + metadata=metadata, + checkpoint_keys=checkpoint_keys, + lifecycle=tuple(lifecycle), + automodel_snapshot=automodel_snapshot, + ) + + +def initialize_pdd_distributed(*, backend: str, timeout_minutes: int = 60) -> Any: + """Verify the wheel, then initialize through AutoModel's released public API.""" + snapshot_installed_distribution() + from nemo_automodel.components.distributed import initialize_distributed + + return initialize_distributed(backend=backend, timeout_minutes=timeout_minutes) diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index e8f4255634e..676a8942a52 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -2,13 +2,11 @@ # 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. Pinned to the public API tested by these examples; -# fastgen_data/__init__.py adds a runtime guard with an actionable message if the helpers move. -# Version 0.5.0 provides the public BaseRecipe.untrack_state API used by deterministic resume. +# NeMo AutoModel supplies the parent recipe, FSDP2 wrapping, checkpointer, and the unmodified +# upstream diffusion helpers used by the shared data/preprocessing code. These versions are the +# public APIs tested by both DMD2 and PDD; PDD also verifies the exact AutoModel wheel at runtime. nemo_automodel[diffusion]==0.5.0 +diffusers==0.38.0 -# Optional but recommended for the smoke logs. +# Optional but recommended for training logs. wandb diff --git a/examples/diffusers/fastgen/verify_readonly_automodel.py b/examples/diffusers/fastgen/verify_readonly_automodel.py new file mode 100644 index 00000000000..841f9bc6e48 --- /dev/null +++ b/examples/diffusers/fastgen/verify_readonly_automodel.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Verify and snapshot the exact released AutoModel distribution used by PDD.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import importlib.util +import json +import os +from pathlib import Path +from typing import Any + +_MANIFEST_PATH = Path(__file__).with_name("automodel_dependency.json") +_GENERATED_NAMES = {"INSTALLER", "RECORD", "REQUESTED"} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def load_dependency_manifest(path: Path = _MANIFEST_PATH) -> dict[str, Any]: + """Load the immutable AutoModel dependency declaration.""" + data = json.loads(path.read_text()) + required = { + "distribution", + "import_name", + "package_file_count", + "package_tree_sha256", + "release_commit", + "release_tag", + "runtime_versions", + "version", + "wheel", + "wheel_sha256", + } + if set(data) != required: + raise RuntimeError( + f"AutoModel dependency manifest keys mismatch: expected={sorted(required)}, " + f"actual={sorted(data)}." + ) + return data + + +def _distribution_root( + distribution: importlib.metadata.Distribution, manifest: dict[str, Any] +) -> Path: + root = Path(distribution.locate_file("")).resolve() + dist_info = root / f"{manifest['distribution']}-{manifest['version']}.dist-info" + if not dist_info.is_dir() or not dist_info.name.endswith(".dist-info"): + raise RuntimeError(f"AutoModel has no regular wheel dist-info directory: {dist_info}.") + return root + + +def _package_files(root: Path, manifest: dict[str, Any]) -> list[Path]: + package_root = root / manifest["import_name"] + dist_info_root = root / f"{manifest['distribution']}-{manifest['version']}.dist-info" + if not package_root.is_dir() or not dist_info_root.is_dir(): + raise RuntimeError( + "AutoModel package or exact-version dist-info directory is missing from the " + f"installed distribution root {root}." + ) + + files: list[Path] = [] + for base in (package_root, dist_info_root): + for candidate in base.rglob("*"): + if candidate.is_symlink(): + raise RuntimeError(f"AutoModel distribution contains a symlink: {candidate}.") + if not candidate.is_file(): + continue + if "__pycache__" in candidate.parts or candidate.suffix == ".pyc": + continue + if candidate.name in _GENERATED_NAMES: + continue + files.append(candidate) + return sorted(files, key=lambda path: path.relative_to(root).as_posix()) + + +def snapshot_installed_distribution() -> dict[str, Any]: + """Return a deterministic content snapshot after enforcing the frozen wheel tree.""" + manifest = load_dependency_manifest() + distribution = importlib.metadata.distribution(manifest["distribution"]) + if distribution.version != manifest["version"]: + raise RuntimeError( + f"PDD requires {manifest['distribution']}=={manifest['version']}, " + f"found {distribution.version}." + ) + + runtime_versions = { + name: importlib.metadata.version(name) for name in manifest["runtime_versions"] + } + if runtime_versions != manifest["runtime_versions"]: + raise RuntimeError( + "PDD runtime dependency versions mismatch: " + f"expected {manifest['runtime_versions']}, found {runtime_versions}." + ) + + root = _distribution_root(distribution, manifest) + direct_url_text = distribution.read_text("direct_url.json") + if direct_url_text is not None: + direct_url = json.loads(direct_url_text) + if direct_url.get("dir_info", {}).get("editable", False): + raise RuntimeError("PDD rejects editable AutoModel installations.") + + spec = importlib.util.find_spec(manifest["import_name"]) + if spec is None or spec.origin is None: + raise RuntimeError(f"Cannot resolve import {manifest['import_name']!r}.") + import_origin = Path(spec.origin).resolve() + try: + import_origin.relative_to(root) + except ValueError as error: + raise RuntimeError( + f"AutoModel import {import_origin} is shadowing distribution root {root}." + ) from error + files = _package_files(root, manifest) + file_records: list[dict[str, Any]] = [] + tree_digest = hashlib.sha256() + for path in files: + relative = path.relative_to(root).as_posix() + digest = _sha256(path) + size = path.stat().st_size + tree_digest.update(relative.encode()) + tree_digest.update(b"\0") + tree_digest.update(digest.encode()) + tree_digest.update(b"\0") + tree_digest.update(str(size).encode()) + tree_digest.update(b"\n") + file_records.append({"path": relative, "sha256": digest, "size": size}) + + actual_tree_digest = tree_digest.hexdigest() + if len(file_records) != manifest["package_file_count"]: + raise RuntimeError( + "AutoModel package file count does not match the frozen wheel: " + f"expected {manifest['package_file_count']}, found {len(file_records)}." + ) + if actual_tree_digest != manifest["package_tree_sha256"]: + raise RuntimeError( + "AutoModel package tree does not match the frozen official wheel: " + f"expected {manifest['package_tree_sha256']}, found {actual_tree_digest}." + ) + + return { + "distribution": manifest["distribution"], + "files": file_records, + "import_origin": str(import_origin), + "package_file_count": len(file_records), + "package_tree_sha256": actual_tree_digest, + "release_commit": manifest["release_commit"], + "release_tag": manifest["release_tag"], + "root": str(root), + "runtime_versions": runtime_versions, + "version": distribution.version, + "wheel": manifest["wheel"], + "wheel_sha256": manifest["wheel_sha256"], + } + + +def write_snapshot(output: Path) -> None: + """Atomically write a distribution snapshot outside the installation.""" + snapshot = snapshot_installed_distribution() + output = output.resolve() + root = Path(snapshot["root"]) + try: + output.relative_to(root) + except ValueError: + pass + else: + raise ValueError("Snapshot output must be outside the AutoModel distribution.") + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_name(f".{output.name}.{os.getpid()}.tmp") + temporary.write_text(json.dumps(snapshot, indent=2, sort_keys=True) + "\n") + os.replace(temporary, output) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + snapshot = subparsers.add_parser("snapshot", help="verify and write a content snapshot") + snapshot.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if args.command == "snapshot": + write_snapshot(args.output) + + +if __name__ == "__main__": + main() 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..d109a4f1eff --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Released-AutoModel seam tests for the ModelOpt-owned PDD setup.""" + +from __future__ import annotations + +import importlib.metadata +import json +import os +import pathlib +import shutil +import subprocess +import sys + +import pytest +import torch +from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir + +_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_recipe import build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config +from verify_readonly_automodel import snapshot_installed_distribution + + +def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: + return { + "model": { + "pretrained_model_name_or_path": str(model_dir), + "torch_dtype": "float32", + "device": "cpu", + "transformer_engine_linear": False, + "peft": None, + "guidance_embeds": False, + "fuse_qkv_projections": qkv, + }, + "pdd": { + "pred_type": "flow", + "num_train_timesteps": None, + "guidance_scale": 4.0, + "student_sample_steps": 2, + "student_sample_type": "ode", + "grid_size": 4, + "flow_shift": 5.0, + "block_size_min": 1, + "block_size_max": 4, + "teacher_integrator": "euler", + "inference_blocks": [2, 2], + "data_free": False, + }, + "optim": {"learning_rate": 2.0e-5, "weight_decay": 0.01}, + "fsdp": { + "dp_size": 1, + "tp_size": 1, + "cp_size": 1, + "pp_size": 1, + "ep_size": 1, + "activation_checkpointing": False, + }, + "checkpoint": { + "enabled": True, + "checkpoint_dir": "checkpoints/test", + "model_save_format": "torch_save", + "save_consolidated": False, + }, + } + + +@pytest.mark.parametrize( + ("scope", "name", "value", "message"), + [ + ("model", "transformer_engine_linear", True, "TE-linear"), + ("model", "peft", {"rank": 8}, "PEFT/LoRA"), + ("root", "peft_cfg", {"rank": 8}, "PEFT/LoRA"), + ("model", "guidance_embeds", True, "guidance embeddings"), + ("model", "device_map", "auto", "device_map"), + ("model", "quantization_config", {"bits": 8}, "quantization_config"), + ], +) +def test_incompatible_modes_fail_during_config_resolution( + tmp_path, scope, name, value, message +) -> None: + raw = _raw_config(tmp_path) + target = raw if scope == "root" else raw[scope] + target[name] = value + + with pytest.raises(ValueError, match=message): + resolve_pdd_recipe_config(raw) + + +def test_remote_model_requires_full_revision_and_non_dp_parallelism_is_rejected(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" + with pytest.raises(ValueError, match=r"exact model\.revision"): + resolve_pdd_recipe_config(raw) + + raw["model"]["revision"] = "a" * 40 + raw["fsdp"]["tp_size"] = 2 + with pytest.raises(ValueError, match="tp_size must be 1"): + resolve_pdd_recipe_config(raw) + + +def test_frozen_automodel_distribution_snapshot_is_stable() -> None: + try: + version = importlib.metadata.version("nemo_automodel") + except importlib.metadata.PackageNotFoundError: + pytest.skip("nemo_automodel is not installed") + assert version == "0.5.0" + + before = snapshot_installed_distribution() + after = snapshot_installed_distribution() + + assert before == after + assert before["version"] == "0.5.0" + assert before["release_commit"] == "d02f49cb314554715aabb97e8dba6599c9f6e9e0" + assert before["runtime_versions"] == {"diffusers": "0.38.0"} + assert before["package_file_count"] == 490 + assert before["package_tree_sha256"] == ( + "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528" + ) + + +def test_exact_wheel_install_below_git_checkout_is_accepted(tmp_path) -> None: + try: + distribution = importlib.metadata.distribution("nemo_automodel") + except importlib.metadata.PackageNotFoundError: + pytest.skip("nemo_automodel is not installed") + assert distribution.version == "0.5.0" + + checkout = tmp_path / "checkout" + (checkout / ".git").mkdir(parents=True) + site_packages = checkout / ".venv" / "lib" / "python" / "site-packages" + site_packages.mkdir(parents=True) + installed_root = pathlib.Path(distribution.locate_file("")).resolve() + shutil.copytree(installed_root / "nemo_automodel", site_packages / "nemo_automodel") + dist_info_name = "nemo_automodel-0.5.0.dist-info" + shutil.copytree(installed_root / dist_info_name, site_packages / dist_info_name) + + output = tmp_path / "snapshot.json" + environment = os.environ.copy() + environment["PYTHONPATH"] = os.pathsep.join( + filter(None, (str(site_packages), environment.get("PYTHONPATH"))) + ) + subprocess.run( + [ + sys.executable, + str(_FASTGEN_DIR / "verify_readonly_automodel.py"), + "snapshot", + "--output", + str(output), + ], + check=True, + env=environment, + ) + + snapshot = json.loads(output.read_text()) + assert pathlib.Path(snapshot["root"]) == site_packages.resolve() + assert pathlib.Path(snapshot["import_origin"]).is_relative_to(site_packages.resolve()) + assert snapshot["package_tree_sha256"] == ( + "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528" + ) + + +def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: + before = snapshot_installed_distribution() + model_dir = create_tiny_qwen_image_pipeline_dir(tmp_path) + initialize_pdd_distributed(backend="gloo", timeout_minutes=1) + config = resolve_pdd_recipe_config(_raw_config(model_dir, qkv=True)) + + source = build_pdd_setup(config) + + assert source.lifecycle == ( + "load/select", + "pdd_conversion", + "device", + "qkv", + "parallelize", + "optimizer", + "checkpoint", + ) + assert type(source.pipe).__name__ == "QwenImagePipeline" + assert source.pipe.text_encoder is None + assert source.pipe.tokenizer is None + assert source.pipe.vae is None + assert source.pipe.transformer is source.student + assert source.student.get_submodule("proj_out") is source.projection + assert source.projection.out_features == source.projection.base_out_features * 4 + assert "proj_out.weight" in source.checkpoint_keys + assert source.student.state_dict()["proj_out.weight"].shape[0] == source.projection.out_features + assert not any(parameter.requires_grad for parameter in source.teacher.parameters()) + optimizer_parameters = [ + parameter for group in source.optimizer.param_groups for parameter in group["params"] + ] + assert any(parameter is source.projection.weight for parameter in optimizer_parameters) + # Diffusers 0.38 accepts the Qwen object API but currently performs no effective fusion. + assert not any( + getattr(module, "fused_projections", False) for module in source.student.modules() + ) + + source.optimizer.zero_grad(set_to_none=True) + # A real PDD forward touches the backbone and projection. Exercise the strict stock + # optimizer restore with complete Adam state rather than an artificial partial update. + sum(parameter.float().square().mean() for parameter in optimizer_parameters).backward() + source.optimizer.step() + expected_weight = source.projection.weight.detach().clone() + expected_exp_avg = source.optimizer.state[source.projection.weight]["exp_avg"].clone() + checkpoint_root = tmp_path / "checkpoint" + source.checkpointer.save_model(source.student, str(checkpoint_root)) + source.checkpointer.save_optimizer(source.optimizer, source.student, str(checkpoint_root)) + + destination = build_pdd_setup(config) + assert destination.metadata == source.metadata + destination_projection = destination.projection + destination_weight_id = id(destination_projection.weight) + destination.checkpointer.load_model( + destination.student, + str(checkpoint_root / "model"), + ) + destination.checkpointer.load_optimizer( + destination.optimizer, + destination.student, + str(checkpoint_root), + ) + + assert destination.student.get_submodule("proj_out") is destination_projection + assert id(destination_projection.weight) == destination_weight_id + torch.testing.assert_close(destination_projection.weight, expected_weight) + torch.testing.assert_close( + destination.optimizer.state[destination_projection.weight]["exp_avg"], + expected_exp_avg, + ) + assert snapshot_installed_distribution() == before + source.checkpointer.close() + destination.checkpointer.close() + + +def test_qwen_pdd_adapter_has_no_automodel_import() -> None: + source = ( + _REPO_ROOT / "modelopt" / "torch" / "fastgen" / "plugins" / "qwen_image_pdd.py" + ).read_text() + assert "nemo_automodel" not in source From 8906224334f386817c3fb963f4648f2f99115ae1 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 08:54:43 -0700 Subject: [PATCH 11/45] feat(fastgen): add PDD training lifecycle Signed-off-by: Meng Xin --- .../fastgen/fastgen_data/__init__.py | 11 +- .../fastgen/fastgen_data/collate_fns.py | 24 +- .../fastgen_data/replayable_sampler.py | 207 ++++ .../fastgen/pdd/configs/qwen_image.yaml | 27 +- examples/diffusers/fastgen/pdd_checkpoint.py | 707 +++++++++++++ examples/diffusers/fastgen/pdd_finetune.py | 580 ++++++++++- examples/diffusers/fastgen/pdd_recipe.py | 270 ++++- examples/diffusers/fastgen/pdd_training.py | 943 ++++++++++++++++++ .../fastgen/validate_cache_snapshot.py | 24 +- .../pdd_checkpoint_failure_distributed.py | 183 ++++ .../diffusers/fastgen/pdd_test_utils.py | 218 ++++ .../pdd_training_preflight_distributed.py | 119 +++ .../pdd_validation_oracle_distributed.py | 134 +++ .../fastgen/test_migrate_cache_manifest.py | 9 +- .../fastgen/test_pdd_recipe_setup.py | 48 + .../fastgen/test_pdd_training_lifecycle.py | 469 +++++++++ .../fastgen/test_pdd_validation_oracle.py | 146 +++ .../diffusers/fastgen/test_portable_cache.py | 31 +- 18 files changed, 4128 insertions(+), 22 deletions(-) create mode 100644 examples/diffusers/fastgen/fastgen_data/replayable_sampler.py create mode 100644 examples/diffusers/fastgen/pdd_checkpoint.py create mode 100644 examples/diffusers/fastgen/pdd_training.py create mode 100644 tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py create mode 100644 tests/examples/diffusers/fastgen/pdd_test_utils.py create mode 100644 tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py create mode 100644 tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py create mode 100644 tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py create mode 100644 tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index 934dc16ea28..75a88b4e1eb 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -41,11 +41,13 @@ try: from . import collate_fns as _collate_fns from . import paths as _paths + from . import replayable_sampler as _replayable_sampler from . import resume as _resume from . import splits as _splits from . import text_to_image_dataset as _text_to_image_dataset from .collate_fns import * from .paths import * + from .replayable_sampler import * from .resume import * from .splits import * from .text_to_image_dataset import * @@ -60,7 +62,14 @@ ) from exc __all__: list[str] = [] -for _module in (_collate_fns, _paths, _resume, _splits, _text_to_image_dataset): +for _module in ( + _collate_fns, + _paths, + _replayable_sampler, + _resume, + _splits, + _text_to_image_dataset, +): __all__.extend(_module.__all__) diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index 812ee6c4e29..a943b32c32e 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -40,6 +40,7 @@ from torchdata.stateful_dataloader import StatefulDataLoader from .paths import resolve_under_root +from .replayable_sampler import ReplayableBatchSampler from .text_to_image_dataset import TextToImageDataset __all__ = [ @@ -175,7 +176,10 @@ def build_text_to_image_multiresolution_dataloader( prefetch_factor: int = 2, negative_prompt_embedding_path: str | None = None, selected_indices: Sequence[int] | None = None, -) -> tuple[StatefulDataLoader, SequentialBucketSampler]: + exact_resume: bool = False, + sampler_seed: int = 42, + loader_seed: int | None = None, +) -> tuple[StatefulDataLoader, SequentialBucketSampler | ReplayableBatchSampler]: """Build the DMD2 text-to-image multiresolution dataloader for ``TrainDiffusionRecipe``. Args: @@ -195,6 +199,11 @@ def build_text_to_image_multiresolution_dataloader( negative_prompt_embedding_path: Optional ``.pt`` with a static negative-prompt embedding, bound into the collate and broadcast to every batch (DMD2 CFG). selected_indices: Optional ordered original metadata ordinals to expose. + exact_resume: Wrap the deterministic sampler with a committed cursor that is + independent of worker prefetch. Required by the PDD lifecycle. + sampler_seed: Seed for the released deterministic bucket sampler. + loader_seed: Optional dedicated seed for DataLoader worker/base-seed generation. PDD + supplies this so recreating an iterator cannot consume its restored training RNG. Returns: ``(StatefulDataLoader, SequentialBucketSampler)``. @@ -237,17 +246,24 @@ 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, ) + batch_sampler = ReplayableBatchSampler(sampler) if exact_resume else sampler + 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, + batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=num_workers, pin_memory=pin_memory, prefetch_factor=prefetch_factor if num_workers > 0 else None, persistent_workers=num_workers > 0, + generator=loader_generator, ) if dp_rank == 0: @@ -257,9 +273,9 @@ def build_text_to_image_multiresolution_dataloader( effective_root, len(dataset), dataset.total_num_samples, - len(sampler), + len(batch_sampler), batch_size, dp_rank, dp_world_size, ) - return dataloader, sampler + return dataloader, batch_sampler diff --git a/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py b/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py new file mode 100644 index 00000000000..46b3fee9cdb --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Committed-cursor wrapper for deterministic, prefetched batch samplers.""" + +from __future__ import annotations + +import hashlib +import struct +from collections.abc import Iterator, Mapping, Sequence +from typing import Any + +import torch +from torch.utils.data import Sampler + +__all__ = ["ReplayableBatchSampler"] + +_STATE_VERSION = 1 + + +class ReplayableBatchSampler(Sampler[list[int]]): + """Separate actually consumed batches from batches yielded to worker prefetch. + + The wrapped sampler remains the authority for deterministic rank/epoch batch plans. + This wrapper materializes that plan compactly and advances its committed cursor only + after the training loop confirms the logical sample IDs it consumed. + """ + + def __init__(self, sampler: Sampler[list[int]]) -> None: + if not isinstance(sampler, Sampler): + raise TypeError(f"sampler must be a torch Sampler, got {type(sampler).__name__}.") + dataset = getattr(sampler, "dataset", None) + metadata = getattr(dataset, "metadata", None) + if not isinstance(metadata, Sequence): + raise TypeError("sampler.dataset.metadata must be a sequence.") + + self.sampler = sampler + self.dataset = dataset + self.epoch = int(getattr(sampler, "epoch", 0)) + self.committed_batches = 0 + self.sample_slots_consumed = 0 + self._yielded_batches = 0 + self._flat_indices = torch.empty(0, dtype=torch.int64) + self._offsets = torch.zeros(1, dtype=torch.int64) + self._plan_sha256 = "" + self._build_plan() + + def _build_plan(self) -> None: + self.sampler.set_epoch(self.epoch) + self.sampler.load_state_dict({"epoch": self.epoch, "batches_yielded": 0}) + batches = [tuple(int(index) for index in batch) for batch in self.sampler] + if not batches: + raise ValueError("replayable batch plan must contain at least one batch.") + if any(not batch for batch in batches): + raise ValueError("replayable batch plan cannot contain an empty batch.") + + flat = [index for batch in batches for index in batch] + offsets = [0] + for batch in batches: + offsets.append(offsets[-1] + len(batch)) + self._flat_indices = torch.tensor(flat, dtype=torch.int64) + self._offsets = torch.tensor(offsets, dtype=torch.int64) + + digest = hashlib.sha256() + digest.update(b"modelopt-pdd-batch-plan-v1\0") + digest.update(struct.pack(">q", self.epoch)) + digest.update(struct.pack(">q", int(getattr(self.sampler, "rank", 0)))) + digest.update(struct.pack(">q", int(getattr(self.sampler, "num_replicas", 1)))) + for batch in batches: + digest.update(struct.pack(">q", len(batch))) + for index in batch: + digest.update(struct.pack(">q", index)) + self._plan_sha256 = digest.hexdigest() + self._yielded_batches = self.committed_batches + + @property + def plan_sha256(self) -> str: + return self._plan_sha256 + + @property + def remaining_batches(self) -> int: + return len(self) - self.committed_batches + + def _batch_indices(self, batch_index: int) -> list[int]: + if not 0 <= batch_index < len(self): + raise IndexError(f"batch_index={batch_index} is outside [0, {len(self)}).") + start = int(self._offsets[batch_index]) + end = int(self._offsets[batch_index + 1]) + return self._flat_indices[start:end].tolist() + + def _sample_ids(self, batch_index: int) -> tuple[str, ...]: + sample_ids: list[str] = [] + for index in self._batch_indices(batch_index): + item = self.dataset.metadata[index] + if not isinstance(item, Mapping) or not isinstance(item.get("sample_id"), str): + raise ValueError(f"dataset.metadata[{index}] has no string sample_id.") + sample_ids.append(item["sample_id"]) + return tuple(sample_ids) + + def expected_next_sample_ids(self) -> tuple[str, ...]: + """Return the next committed batch's logical IDs without consuming it.""" + if self.committed_batches == len(self): + return () + return self._sample_ids(self.committed_batches) + + def commit(self, sample_ids: Sequence[str]) -> None: + """Advance the durable cursor after verifying the collated logical IDs.""" + if isinstance(sample_ids, str) or not isinstance(sample_ids, Sequence): + raise TypeError("sample_ids must be a sequence of strings.") + actual = tuple(sample_ids) + if any(not isinstance(sample_id, str) for sample_id in actual): + raise TypeError("sample_ids must contain only strings.") + expected = self.expected_next_sample_ids() + if not expected: + raise RuntimeError("cannot commit beyond the end of the batch plan.") + if actual != expected: + raise RuntimeError( + f"consumed sample IDs do not match the committed cursor: " + f"expected={expected}, actual={actual}." + ) + self.committed_batches += 1 + self.sample_slots_consumed += len(actual) + + def set_epoch(self, epoch: int) -> None: + if type(epoch) is not int or epoch < 0: + raise ValueError("epoch must be an integer >= 0.") + if epoch == self.epoch: + return + if self.committed_batches != len(self): + raise RuntimeError("cannot change epoch before every planned batch is committed.") + self.epoch = epoch + self.committed_batches = 0 + self._build_plan() + + def state_dict(self) -> dict[str, Any]: + return { + "schema_version": _STATE_VERSION, + "epoch": self.epoch, + "committed_batches": self.committed_batches, + "sample_slots_consumed": self.sample_slots_consumed, + "plan_sha256": self.plan_sha256, + "next_sample_ids": list(self.expected_next_sample_ids()), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + if not isinstance(state, Mapping): + raise TypeError("replayable sampler state must be a mapping.") + expected_keys = { + "schema_version", + "epoch", + "committed_batches", + "sample_slots_consumed", + "plan_sha256", + "next_sample_ids", + } + if set(state) != expected_keys: + raise ValueError( + f"replayable sampler state keys mismatch: expected={sorted(expected_keys)}, " + f"actual={sorted(state)}." + ) + if state["schema_version"] != _STATE_VERSION: + raise ValueError(f"unsupported replayable sampler schema {state['schema_version']!r}.") + epoch = state["epoch"] + committed = state["committed_batches"] + consumed = state["sample_slots_consumed"] + if type(epoch) is not int or epoch < 0: + raise ValueError("saved epoch must be an integer >= 0.") + if type(committed) is not int or committed < 0: + raise ValueError("saved committed_batches must be an integer >= 0.") + if type(consumed) is not int or consumed < 0: + raise ValueError("saved sample_slots_consumed must be an integer >= 0.") + if not isinstance(state["plan_sha256"], str): + raise TypeError("saved plan_sha256 must be a string.") + if not isinstance(state["next_sample_ids"], list) or any( + not isinstance(sample_id, str) for sample_id in state["next_sample_ids"] + ): + raise TypeError("saved next_sample_ids must be a list of strings.") + + self.epoch = epoch + self.committed_batches = 0 + self._build_plan() + if committed > len(self): + raise ValueError( + f"saved committed_batches={committed} exceeds plan length {len(self)}." + ) + if self.plan_sha256 != state["plan_sha256"]: + raise RuntimeError("reconstructed batch plan does not match the saved plan hash.") + self.committed_batches = committed + self.sample_slots_consumed = consumed + self._yielded_batches = committed + if list(self.expected_next_sample_ids()) != state["next_sample_ids"]: + raise RuntimeError("reconstructed next sample IDs do not match the checkpoint.") + + def __iter__(self) -> Iterator[list[int]]: + start = self.committed_batches + flat_indices = self._flat_indices + offsets = self._offsets + total_batches = int(offsets.numel() - 1) + self._yielded_batches = start + for batch_index in range(start, total_batches): + batch_start = int(offsets[batch_index]) + batch_end = int(offsets[batch_index + 1]) + self._yielded_batches = batch_index + 1 + yield flat_indices[batch_start:batch_end].tolist() + + def __len__(self) -> int: + return int(self._offsets.numel() - 1) diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index abd8a208e40..546361fb235 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -1,4 +1,4 @@ -# Qwen-Image PDD setup skeleton. Task 8 adds the training lifecycle. +# Qwen-Image PDD training recipe. Result-bearing paths/topology remain externally gated. model: pretrained_model_name_or_path: Qwen/Qwen-Image @@ -29,6 +29,25 @@ pdd: optim: learning_rate: 2.0e-5 weight_decay: 0.01 + betas: [0.9, 0.999] + eps: 1.0e-8 + +guidance: + rescale: 1.0 + eps: 1.0e-5 + +training: + seed: 42 + max_steps: 10000 + max_grad_norm: 1.0 + zero_grad_warmup_steps: 0 + log_every_steps: 10 + checkpoint_every_steps: 1000 + validation_every_steps: 1000 + grad_accumulation_steps: 1 + # Freeze to 256 only after the production-topology/data gate is approved. + global_batch_size: + validation_seed: 2026 fsdp: dp_size: @@ -39,12 +58,17 @@ fsdp: activation_checkpointing: true data: + all_metadata_index: metadata.json + validation_metadata_index: metadata_heldout.json dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: data/qwen_image_cache metadata_index: metadata_train.json base_resolution: [1024, 1024] batch_size: 1 + drop_last: true + shuffle: true + dynamic_batch_size: false negative_prompt_embedding_path: negative_prompt_embedding.pt checkpoint: @@ -52,3 +76,4 @@ checkpoint: checkpoint_dir: checkpoints/pdd_qwen_image model_save_format: torch_save save_consolidated: false + restore_from: LATEST diff --git a/examples/diffusers/fastgen/pdd_checkpoint.py b/examples/diffusers/fastgen/pdd_checkpoint.py new file mode 100644 index 00000000000..48f0d96ea3e --- /dev/null +++ b/examples/diffusers/fastgen/pdd_checkpoint.py @@ -0,0 +1,707 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Atomic, strict PDD checkpoint publication around the stock AutoModel Checkpointer.""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import shutil +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch.distributed as dist + +from modelopt.torch.fastgen import PDDMetadata + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + +_CHECKPOINT_SCHEMA_VERSION = 1 +_COMPLETE_SCHEMA_VERSION = 1 +_FORBIDDEN_ARTIFACT_TOKENS = ("fake_score", "discriminator", "ema", "r1", "gan") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _require_sha256(value: Any, *, name: str) -> str: + if not isinstance(value, str) or len(value) != 64: + raise ValueError(f"{name} must be a 64-character SHA-256 digest.") + try: + int(value, 16) + except ValueError as error: + raise ValueError(f"{name} must be hexadecimal.") from error + return value.lower() + + +def _rank() -> int: + return dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + + +def _world_size() -> int: + return dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 + + +def _barrier() -> None: + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + +def _broadcast_rank0_payload(value: Any) -> Any: + payload = [value] + if dist.is_available() and dist.is_initialized(): + dist.broadcast_object_list(payload, src=0) + return payload[0] + + +def _gather_objects(value: Any) -> list[Any]: + if not dist.is_available() or not dist.is_initialized(): + return [value] + values: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(values, value) + return values + + +def _fsync_file(path: Path) -> None: + with path.open("rb") as stream: + os.fsync(stream.fileno()) + + +def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _fsync_tree(root: Path) -> None: + for path in sorted(root.rglob("*"), key=lambda candidate: len(candidate.parts), reverse=True): + if path.is_symlink(): + raise RuntimeError(f"PDD checkpoint staging contains a symlink: {path}.") + if path.is_file(): + _fsync_file(path) + elif path.is_dir(): + _fsync_directory(path) + _fsync_directory(root) + + +def _dcp_payload_hashes(checkpoint: Path) -> dict[str, str]: + hashes: dict[str, str] = {} + for component in ("model", "optim"): + root = checkpoint / component + if not root.is_dir() or root.is_symlink(): + raise RuntimeError(f"PDD checkpoint is missing the {component} DCP directory.") + files: list[Path] = [] + for path in root.rglob("*"): + if path.is_symlink(): + raise RuntimeError(f"PDD {component} DCP tree contains a symlink: {path}.") + if path.is_file(): + files.append(path) + relative_files = {path.relative_to(checkpoint).as_posix() for path in files} + if f"{component}/.metadata" not in relative_files: + raise RuntimeError(f"PDD checkpoint is missing strict {component} DCP metadata.") + if not any(relative != f"{component}/.metadata" for relative in relative_files): + raise RuntimeError(f"PDD checkpoint is missing {component} DCP payload shards.") + for path in files: + hashes[path.relative_to(checkpoint).as_posix()] = _sha256(path) + return hashes + + +def _atomic_text(path: Path, text: str) -> None: + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + with temporary.open("w") as stream: + stream.write(text) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + _fsync_directory(path.parent) + + +def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None: + _atomic_text(path, json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def _read_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError(f"cannot read PDD checkpoint JSON {path}.") from error + if not isinstance(value, dict): + raise RuntimeError(f"PDD checkpoint JSON {path} must contain an object.") + return value + + +def build_pdd_checkpoint_identity( + *, + metadata: PDDMetadata, + model_id: str, + model_revision: str | None, + guidance_scale: float | None, + guidance_rescale: float, + guidance_eps: float, + automodel_snapshot: Mapping[str, Any], + ordered_train_id_sha256: str, + ordered_heldout_id_sha256: str, + dataset_snapshot_sha256: str, + local_batch_size: int, + grad_accumulation_steps: int, + training_seed: int, + validation_seed: int, + validation_every_steps: int, + max_grad_norm: float, + zero_grad_warmup_steps: int, + activation_checkpointing: bool, + dtype: str, + optimizer: Any, + scheduler: Any, +) -> dict[str, Any]: + """Build the strict, path-independent compatibility identity for PDD resume.""" + if not isinstance(metadata, PDDMetadata): + raise TypeError("metadata must be PDDMetadata.") + if not isinstance(model_id, str) or not model_id: + raise ValueError("model_id must be a non-empty string.") + if model_revision is not None and not isinstance(model_revision, str): + raise TypeError("model_revision must be a string or null.") + for name, value in (("guidance_rescale", guidance_rescale), ("guidance_eps", guidance_eps)): + if isinstance(value, bool) or not isinstance(value, int | float): + raise TypeError(f"{name} must be a real number.") + if guidance_scale is not None and ( + isinstance(guidance_scale, bool) or not isinstance(guidance_scale, int | float) + ): + raise TypeError("guidance_scale must be a real number or null.") + if type(local_batch_size) is not int or local_batch_size < 1: + raise ValueError("local_batch_size must be an integer >= 1.") + if grad_accumulation_steps != 1: + raise ValueError("PDD v1 exact resume requires grad_accumulation_steps=1.") + for name, value, minimum in ( + ("training_seed", training_seed, 0), + ("validation_seed", validation_seed, 0), + ("validation_every_steps", validation_every_steps, 1), + ("zero_grad_warmup_steps", zero_grad_warmup_steps, 0), + ): + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}.") + if isinstance(max_grad_norm, bool) or not isinstance(max_grad_norm, int | float): + raise TypeError("max_grad_norm must be a real number.") + if not math.isfinite(max_grad_norm) or max_grad_norm <= 0: + raise ValueError("max_grad_norm must be finite and > 0.") + if type(activation_checkpointing) is not bool: + raise TypeError("activation_checkpointing must be bool.") + if not isinstance(dtype, str) or not dtype: + raise ValueError("dtype must be a non-empty string.") + if type(optimizer).__module__ != "torch.optim.adamw" or type(optimizer).__name__ != "AdamW": + raise TypeError("PDD checkpoint identity requires the stock torch.optim.AdamW optimizer.") + required_snapshot = { + "distribution", + "package_tree_sha256", + "runtime_versions", + "version", + "wheel_sha256", + } + missing = sorted(required_snapshot.difference(automodel_snapshot)) + if missing: + raise ValueError(f"AutoModel snapshot is missing identity keys: {missing}.") + return { + "schema_version": _CHECKPOINT_SCHEMA_VERSION, + "model": {"id": model_id, "revision": model_revision, "dtype": dtype}, + "pdd_metadata": metadata.to_dict(), + "guidance": { + "scale": None if guidance_scale is None else float(guidance_scale), + "rescale": float(guidance_rescale), + "eps": float(guidance_eps), + }, + "automodel": { + "distribution": automodel_snapshot["distribution"], + "version": automodel_snapshot["version"], + "package_tree_sha256": _require_sha256( + automodel_snapshot["package_tree_sha256"], + name="automodel.package_tree_sha256", + ), + "wheel_sha256": _require_sha256( + automodel_snapshot["wheel_sha256"], + name="automodel.wheel_sha256", + ), + "runtime_versions": dict(automodel_snapshot["runtime_versions"]), + }, + "data": { + "ordered_train_id_sha256": _require_sha256( + ordered_train_id_sha256, + name="ordered_train_id_sha256", + ), + "ordered_heldout_id_sha256": _require_sha256( + ordered_heldout_id_sha256, + name="ordered_heldout_id_sha256", + ), + "dataset_snapshot_sha256": _require_sha256( + dataset_snapshot_sha256, + name="dataset_snapshot_sha256", + ), + "local_batch_size": local_batch_size, + "grad_accumulation_steps": grad_accumulation_steps, + }, + "topology": {"world_size": _world_size(), "pure_data_parallel": True}, + "training": { + "seed": training_seed, + "validation_seed": validation_seed, + "validation_every_steps": validation_every_steps, + "max_grad_norm": float(max_grad_norm), + "zero_grad_warmup_steps": zero_grad_warmup_steps, + "activation_checkpointing": activation_checkpointing, + }, + "optimizer": { + "class": "torch.optim.AdamW", + "param_groups": [ + { + "lr": float(group["lr"]), + "betas": [float(beta) for beta in group["betas"]], + "eps": float(group["eps"]), + "weight_decay": float(group["weight_decay"]), + "amsgrad": bool(group["amsgrad"]), + "maximize": bool(group["maximize"]), + "capturable": bool(group["capturable"]), + "differentiable": bool(group["differentiable"]), + "foreach": bool(group["foreach"]), + "fused": bool(group["fused"]), + } + for group in optimizer.param_groups + ], + }, + "scheduler": { + "class": f"{type(scheduler).__module__}.{type(scheduler).__qualname__}", + "base_lrs": [float(value) for value in scheduler.base_lrs], + "policy": "constant", + }, + } + + +@dataclass(frozen=True) +class PDDResumeState: + """Restored progress plus the first logical IDs that must be served next.""" + + checkpoint_path: Path + completed_steps: int + sample_slots_consumed: int + expected_next_sample_ids: tuple[str, ...] + parent_checkpoint: str | None + + def verify_first_batch(self, sample_ids: Sequence[str]) -> None: + if tuple(sample_ids) != self.expected_next_sample_ids: + raise RuntimeError( + "first resumed sample IDs do not match the checkpoint: " + f"expected={self.expected_next_sample_ids}, actual={tuple(sample_ids)}." + ) + + +class PDDCheckpointManager: + """Publish and restore complete, metadata-compatible PDD checkpoints only.""" + + def __init__( + self, + *, + root: str | Path, + checkpointer: Any, + model: Any, + optimizer: Any, + scheduler: Any, + trainer: Any, + sampler: Any, + rng: Any, + identity: Mapping[str, Any], + ) -> None: + self.root = Path(root).resolve() + self.checkpointer = checkpointer + self.model = model + self.optimizer = optimizer + self.scheduler = scheduler + self.trainer = trainer + self.sampler = sampler + self.rng = rng + self.identity = json.loads(json.dumps(identity, sort_keys=True)) + if self.identity.get("schema_version") != _CHECKPOINT_SCHEMA_VERSION: + raise ValueError("PDD checkpoint identity has an unsupported schema version.") + topology = self.identity.get("topology") + if not isinstance(topology, dict) or topology.get("world_size") != _world_size(): + raise ValueError("PDD checkpoint identity world size does not match the process group.") + if bool(getattr(checkpointer.config, "is_async", False)): + raise ValueError("PDD v1 atomic publication requires synchronous checkpoint saves.") + + def _rank_summary(self) -> dict[str, Any]: + sampler_state = self.sampler.state_dict() + return { + "rank": _rank(), + "epoch": sampler_state["epoch"], + "committed_batches": sampler_state["committed_batches"], + "sample_slots_consumed": sampler_state["sample_slots_consumed"], + "plan_sha256": sampler_state["plan_sha256"], + "next_sample_ids": list(sampler_state["next_sample_ids"]), + } + + def _sidecar_paths(self, checkpoint: Path) -> list[Path]: + paths: list[Path] = [] + for rank in range(_world_size()): + paths.extend( + ( + checkpoint / "rng" / f"rng_dp_rank_{rank}.pt", + checkpoint / "sampler" / f"sampler_dp_rank_{rank}.pt", + checkpoint / "trainer" / f"trainer_dp_rank_{rank}.pt", + ) + ) + return paths + + def _manifest(self, checkpoint: Path) -> dict[str, Any]: + manifest = _read_json(checkpoint / "manifest.json") + expected = { + "schema_version", + "identity", + "completed_steps", + "learning_rates", + "parent_checkpoint", + "rank_progress", + "dcp_sha256", + "sidecar_sha256", + } + if set(manifest) != expected: + raise RuntimeError("PDD checkpoint manifest has incompatible keys.") + if manifest["schema_version"] != _CHECKPOINT_SCHEMA_VERSION: + raise RuntimeError("PDD checkpoint manifest schema is unsupported.") + return manifest + + def _validate_checkpoint(self, checkpoint: Path, *, require_identity: bool) -> dict[str, Any]: + if not checkpoint.is_dir() or checkpoint.is_symlink(): + raise RuntimeError(f"PDD checkpoint is not a regular directory: {checkpoint}.") + marker_path = checkpoint / "COMPLETE" + manifest_path = checkpoint / "manifest.json" + if not marker_path.is_file() or not manifest_path.is_file(): + raise RuntimeError(f"PDD checkpoint is incomplete: {checkpoint}.") + marker = _read_json(marker_path) + if ( + set(marker) != {"schema_version", "manifest_sha256"} + or marker.get("schema_version") != _COMPLETE_SCHEMA_VERSION + ): + raise RuntimeError("PDD COMPLETE marker is incompatible.") + if marker["manifest_sha256"] != _sha256(manifest_path): + raise RuntimeError("PDD COMPLETE marker does not match manifest content.") + manifest = self._manifest(checkpoint) + if require_identity and manifest["identity"] != self.identity: + raise RuntimeError("PDD checkpoint identity does not match the current run.") + if _read_json(checkpoint / "pdd_config.json") != manifest["identity"]: + raise RuntimeError("PDD checkpoint config sidecar does not match the manifest.") + trainer_state = _read_json(checkpoint / "trainer_state.json") + if trainer_state != { + "completed_steps": manifest["completed_steps"], + "learning_rates": manifest["learning_rates"], + "parent_checkpoint": manifest["parent_checkpoint"], + }: + raise RuntimeError("PDD trainer-state sidecar does not match the manifest.") + expected_dcp = manifest["dcp_sha256"] + if not isinstance(expected_dcp, dict) or any( + not isinstance(path, str) or not isinstance(digest, str) + for path, digest in expected_dcp.items() + ): + raise RuntimeError("PDD checkpoint DCP hash inventory is malformed.") + actual_dcp = _dcp_payload_hashes(checkpoint) + if actual_dcp != expected_dcp: + raise RuntimeError("PDD checkpoint DCP payload inventory or hash does not match.") + expected_sidecars = self._sidecar_paths(checkpoint) + expected_relative = {path.relative_to(checkpoint).as_posix() for path in expected_sidecars} + if ( + not isinstance(manifest["sidecar_sha256"], dict) + or set(manifest["sidecar_sha256"]) != expected_relative + ): + raise RuntimeError("PDD checkpoint sidecar inventory does not match the topology.") + for path in expected_sidecars: + relative = path.relative_to(checkpoint).as_posix() + if not path.is_file() or path.is_symlink(): + raise RuntimeError(f"PDD checkpoint sidecar is missing: {relative}.") + if _sha256(path) != manifest["sidecar_sha256"][relative]: + raise RuntimeError(f"PDD checkpoint sidecar hash mismatch: {relative}.") + for candidate in checkpoint.rglob("*"): + lowered = candidate.name.lower() + if any(token in lowered for token in _FORBIDDEN_ARTIFACT_TOKENS): + raise RuntimeError( + f"PDD checkpoint contains a forbidden DMD artifact: {candidate}." + ) + return manifest + + def _compatible_candidates( + self, + *, + after_completed_steps: int | None = None, + ) -> list[tuple[int, Path]]: + candidates: list[tuple[int, Path]] = [] + if not self.root.is_dir(): + return candidates + for path in self.root.iterdir(): + if not path.is_dir() or path.name.startswith("."): + continue + if after_completed_steps is not None: + prefix = "step_" + suffix = path.name.removeprefix(prefix) + if not path.name.startswith(prefix) or not suffix.isdigit(): + continue + if int(suffix) <= after_completed_steps: + continue + try: + manifest = self._validate_checkpoint(path, require_identity=True) + except RuntimeError: + continue + completed = manifest["completed_steps"] + if type(completed) is int and completed >= 0: + candidates.append((completed, path.resolve())) + return sorted(candidates, key=lambda item: (item[0], item[1].name), reverse=True) + + def resolve(self, restore_from: str | Path | None) -> Path | None: + """Resolve LATEST by scanning only complete, identity-compatible checkpoints.""" + if restore_from is None: + return None + if str(restore_from).upper() == "LATEST": + pointer = self.root / "LATEST" + pointed: tuple[int, Path] | None = None + if pointer.is_file() and not pointer.is_symlink(): + name = pointer.read_text().strip() + candidate = (self.root / name).resolve() + try: + candidate.relative_to(self.root) + manifest = self._validate_checkpoint(candidate, require_identity=True) + completed = manifest["completed_steps"] + if type(completed) is not int or completed < 0: + raise RuntimeError("PDD checkpoint completed_steps is invalid.") + pointed = (completed, candidate) + except (ValueError, RuntimeError): + pass + candidates = self._compatible_candidates( + after_completed_steps=None if pointed is None else pointed[0] + ) + if pointed is not None: + candidates.append(pointed) + candidates.sort(key=lambda item: (item[0], item[1].name), reverse=True) + return candidates[0][1] if candidates else None + + candidate = Path(restore_from) + if not candidate.is_absolute(): + candidate = self.root / candidate + candidate = candidate.resolve() + try: + candidate.relative_to(self.root) + except ValueError as error: + raise ValueError("explicit PDD checkpoint must be beneath checkpoint_dir.") from error + self._validate_checkpoint(candidate, require_identity=True) + return candidate + + def _collective_resolve(self, restore_from: str | Path | None) -> Path | None: + if _world_size() == 1: + return self.resolve(restore_from) + status = None + if _rank() == 0: + try: + resolved = self.resolve(restore_from) + status = { + "ok": True, + "path": None if resolved is None else str(resolved), + } + except BaseException as error: + status = { + "ok": False, + "error": f"{type(error).__name__}: {error}", + } + status = _broadcast_rank0_payload(status) + if not isinstance(status, dict) or type(status.get("ok")) is not bool: + raise RuntimeError("rank 0 broadcast a malformed checkpoint resolution status.") + if not status["ok"]: + raise RuntimeError(f"rank-0 checkpoint resolution failed: {status.get('error')}.") + return None if status["path"] is None else Path(status["path"]) + + def _prepare_staging(self, final: Path) -> str: + self.root.mkdir(parents=True, exist_ok=True) + if final.exists(): + try: + self._validate_checkpoint(final, require_identity=False) + except RuntimeError: + shutil.rmtree(final) + else: + raise FileExistsError(f"complete PDD checkpoint already exists: {final}.") + staging_name = f".{final.name}.{uuid.uuid4().hex}.staging" + (self.root / staging_name).mkdir() + return staging_name + + def _publish_staging( + self, + *, + staging: Path, + final: Path, + completed_steps: int, + learning_rates: list[float], + parent: Path | None, + rank_summaries: list[dict[str, Any]], + ) -> None: + sidecars = self._sidecar_paths(staging) + sidecar_sha256 = {path.relative_to(staging).as_posix(): _sha256(path) for path in sidecars} + manifest = { + "schema_version": _CHECKPOINT_SCHEMA_VERSION, + "identity": self.identity, + "completed_steps": completed_steps, + "learning_rates": learning_rates, + "parent_checkpoint": None if parent is None else parent.name, + "rank_progress": sorted(rank_summaries, key=lambda summary: summary["rank"]), + "dcp_sha256": _dcp_payload_hashes(staging), + "sidecar_sha256": sidecar_sha256, + } + _atomic_json(staging / "pdd_config.json", self.identity) + _atomic_json( + staging / "trainer_state.json", + { + "completed_steps": completed_steps, + "learning_rates": learning_rates, + "parent_checkpoint": manifest["parent_checkpoint"], + }, + ) + _atomic_json(staging / "manifest.json", manifest) + _fsync_tree(staging) + staging.rename(final) + _fsync_directory(self.root) + _atomic_json( + final / "COMPLETE", + { + "schema_version": _COMPLETE_SCHEMA_VERSION, + "manifest_sha256": _sha256(final / "manifest.json"), + }, + ) + self._validate_checkpoint(final, require_identity=True) + _atomic_text(self.root / "LATEST", final.name + "\n") + + def save(self) -> Path: + """Synchronously save into staging, publish atomically, mark complete, then update LATEST.""" + completed_steps = self.trainer.completed_steps + if type(completed_steps) is not int or completed_steps <= 0: + raise ValueError("PDD checkpoint requires at least one completed optimizer step.") + rank_summaries = _gather_objects(self._rank_summary()) + if len({summary["sample_slots_consumed"] for summary in rank_summaries}) != 1: + raise RuntimeError("PDD ranks disagree on consumed sample slots.") + learning_rates = [float(group["lr"]) for group in self.optimizer.param_groups] + final = self.root / f"step_{completed_steps:08d}" + parent = self._collective_resolve("LATEST") + + prepare_status = None + if _rank() == 0: + try: + prepare_status = {"ok": True, "staging": self._prepare_staging(final)} + except BaseException as error: + prepare_status = { + "ok": False, + "error": f"{type(error).__name__}: {error}", + } + prepare_status = _broadcast_rank0_payload(prepare_status) + if not isinstance(prepare_status, dict) or type(prepare_status.get("ok")) is not bool: + raise RuntimeError("rank 0 broadcast a malformed checkpoint preparation status.") + if not prepare_status["ok"]: + raise RuntimeError( + f"rank-0 checkpoint preparation failed: {prepare_status.get('error')}." + ) + staging = self.root / prepare_status["staging"] + + self.checkpointer.save_model(self.model, str(staging)) + self.checkpointer.save_optimizer( + self.optimizer, + self.model, + str(staging), + self.scheduler, + ) + sidecar_error = None + try: + self.checkpointer.save_on_dp_ranks(self.rng, "rng", str(staging)) + self.checkpointer.save_on_dp_ranks(self.sampler, "sampler", str(staging)) + self.checkpointer.save_on_dp_ranks(self.trainer, "trainer", str(staging)) + except BaseException as error: + sidecar_error = f"{type(error).__name__}: {error}" + sidecar_errors = _gather_objects(sidecar_error) + sidecar_failures = [ + f"rank {rank}: {message}" + for rank, message in enumerate(sidecar_errors) + if message is not None + ] + if sidecar_failures: + raise RuntimeError("PDD checkpoint sidecar save failed; " + "; ".join(sidecar_failures)) + _barrier() + + publish_status = None + if _rank() == 0: + try: + self._publish_staging( + staging=staging, + final=final, + completed_steps=completed_steps, + learning_rates=learning_rates, + parent=parent, + rank_summaries=rank_summaries, + ) + publish_status = {"ok": True} + except BaseException as error: + publish_status = { + "ok": False, + "error": f"{type(error).__name__}: {error}", + } + publish_status = _broadcast_rank0_payload(publish_status) + if not isinstance(publish_status, dict) or type(publish_status.get("ok")) is not bool: + raise RuntimeError("rank 0 broadcast a malformed checkpoint publication status.") + if not publish_status["ok"]: + raise RuntimeError( + f"rank-0 checkpoint publication failed: {publish_status.get('error')}." + ) + return final + + def load(self, restore_from: str | Path | None) -> PDDResumeState | None: + """Strictly restore model, optimizer/scheduler, cursor/trainer, and RNG last.""" + checkpoint = self._collective_resolve(restore_from) + if checkpoint is None: + return None + manifest = self._manifest(checkpoint) + self.checkpointer.load_model(self.model, str(checkpoint / "model")) + self.checkpointer.load_optimizer( + self.optimizer, + self.model, + str(checkpoint), + self.scheduler, + ) + self.checkpointer.load_on_dp_ranks(self.trainer, "trainer", str(checkpoint)) + self.checkpointer.load_on_dp_ranks(self.sampler, "sampler", str(checkpoint)) + rank_progress = manifest["rank_progress"] + if not isinstance(rank_progress, list) or len(rank_progress) != _world_size(): + raise RuntimeError("PDD checkpoint rank progress does not match world size.") + progress = rank_progress[_rank()] + if progress.get("rank") != _rank(): + raise RuntimeError("PDD checkpoint rank progress is not ordered by rank.") + sampler_state = self.sampler.state_dict() + for key in ( + "epoch", + "committed_batches", + "sample_slots_consumed", + "plan_sha256", + "next_sample_ids", + ): + if sampler_state[key] != progress[key]: + raise RuntimeError(f"PDD restored sampler {key} does not match the manifest.") + if self.trainer.completed_steps != manifest["completed_steps"]: + raise RuntimeError("PDD restored trainer step does not match the manifest.") + current_lrs = [float(group["lr"]) for group in self.optimizer.param_groups] + if current_lrs != manifest["learning_rates"]: + raise RuntimeError("PDD restored learning rate does not match the manifest.") + self.checkpointer.load_on_dp_ranks(self.rng, "rng", str(checkpoint)) + return PDDResumeState( + checkpoint_path=checkpoint, + completed_steps=manifest["completed_steps"], + sample_slots_consumed=progress["sample_slots_consumed"], + expected_next_sample_ids=tuple(progress["next_sample_ids"]), + parent_checkpoint=manifest["parent_checkpoint"], + ) diff --git a/examples/diffusers/fastgen/pdd_finetune.py b/examples/diffusers/fastgen/pdd_finetune.py index 96e8f60bd72..11ac80320b7 100644 --- a/examples/diffusers/fastgen/pdd_finetune.py +++ b/examples/diffusers/fastgen/pdd_finetune.py @@ -1,14 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Build the released-AutoModel Qwen-Image PDD setup owned by ModelOpt.""" +"""Train Qwen-Image with the ModelOpt-owned PDD lifecycle and released AutoModel APIs.""" from __future__ import annotations import argparse +import dataclasses +import hashlib import logging import sys +import time +from collections.abc import Mapping from pathlib import Path +from typing import Any import yaml @@ -31,9 +36,352 @@ def _parse_args() -> argparse.Namespace: return parser.parse_args() +def _metadata_sample_ids(metadata: Any, *, split: str) -> tuple[str, ...]: + if not isinstance(metadata, list): + raise TypeError(f"PDD {split} dataset metadata must be a list.") + sample_ids: list[str] = [] + for index, item in enumerate(metadata): + if not isinstance(item, Mapping) or not isinstance(item.get("sample_id"), str): + raise ValueError(f"{split} metadata[{index}] has no string sample_id.") + sample_ids.append(item["sample_id"]) + if len(set(sample_ids)) != len(sample_ids): + raise ValueError(f"PDD {split} metadata contains duplicate logical sample IDs.") + return tuple(sample_ids) + + +def _ordered_id_sha256(metadata: Any, *, split: str) -> str: + sample_ids = _metadata_sample_ids(metadata, split=split) + digest = hashlib.sha256() + digest.update(f"modelopt-pdd-ordered-{split}-ids-v1\0".encode()) + for sample_id in sample_ids: + digest.update(sample_id.encode()) + digest.update(b"\n") + return digest.hexdigest() + + +def _dataloader_options(raw: Mapping[str, Any]) -> dict[str, Any]: + data = raw.get("data") + if not isinstance(data, Mapping) or not isinstance(data.get("dataloader"), Mapping): + raise TypeError("PDD config requires a data.dataloader mapping.") + options = dict(data["dataloader"]) + target = options.pop("_target_", None) + expected_target = "fastgen_data.build_text_to_image_multiresolution_dataloader" + if target != expected_target: + raise ValueError(f"PDD data.dataloader._target_ must be {expected_target!r}.") + if "base_resolution" in options: + options["base_resolution"] = tuple(options["base_resolution"]) + return options + + +def _build_training_dataloader( + raw: Mapping[str, Any], + config: Any, + *, + dp_rank: int, + dp_world_size: int, +) -> tuple[Any, Any]: + from fastgen_data import ReplayableBatchSampler + from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader + + options = _dataloader_options(raw) + if options.get("drop_last", True) is not True: + raise ValueError("PDD exact sample accounting requires data.dataloader.drop_last=true.") + if options.get("dynamic_batch_size", False) is not False: + raise ValueError("PDD v1 requires data.dataloader.dynamic_batch_size=false.") + options.update( + dp_rank=dp_rank, + dp_world_size=dp_world_size, + exact_resume=True, + sampler_seed=config.training.seed, + loader_seed=config.training.seed, + ) + dataloader, sampler = build_text_to_image_multiresolution_dataloader(**options) + if not isinstance(sampler, ReplayableBatchSampler): + raise RuntimeError("PDD training requires the committed replayable batch sampler.") + if options.get("batch_size", 1) != config.training.local_batch_size: + raise RuntimeError("resolved local batch size does not match the built dataloader.") + return dataloader, sampler + + +def _build_validation_dataloader( + raw: Mapping[str, Any], + config: Any, + *, + dp_rank: int, + dp_world_size: int, +) -> tuple[Any, Any]: + from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader + + options = _dataloader_options(raw) + options.update( + metadata_index=config.validation_metadata_index, + dp_rank=dp_rank, + dp_world_size=dp_world_size, + drop_last=False, + shuffle=False, + dynamic_batch_size=False, + exact_resume=False, + sampler_seed=config.training.validation_seed, + loader_seed=config.training.validation_seed, + ) + return build_text_to_image_multiresolution_dataloader(**options) + + +def _validate_dataset_snapshot(raw: Mapping[str, Any], config: Any) -> Mapping[str, Any]: + import torch.distributed as dist + from portable_cache import resolve_cache_root + from validate_cache_snapshot import validate_snapshot + + options = _dataloader_options(raw) + root = resolve_cache_root(options["cache_dir"]) + roots: list[str] = [""] * dist.get_world_size() + dist.all_gather_object(roots, str(root)) + if any(candidate != roots[0] for candidate in roots[1:]): + raise RuntimeError(f"PDD ranks resolved different dataset cache roots: {roots}.") + + status = None + if dist.get_rank() == 0: + try: + report = validate_snapshot( + root, + all_index=config.all_metadata_index, + train_index=config.train_metadata_index, + heldout_index=config.validation_metadata_index, + ) + status = {"ok": True, "report": report} + except BaseException as error: + status = {"ok": False, "error": f"{type(error).__name__}: {error}"} + payload = [status] + dist.broadcast_object_list(payload, src=0) + status = payload[0] + if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: + raise RuntimeError("rank 0 broadcast a malformed dataset validation status.") + if not status["ok"]: + raise RuntimeError(f"PDD dataset snapshot validation failed: {status.get('error')}.") + report = status.get("report") + if not isinstance(report, Mapping): + raise RuntimeError("PDD dataset snapshot report is malformed.") + return report + + +def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[bool, ...], ...]]: + import torch.distributed as dist + from pdd_training import build_pdd_validation_assignments + + heldout_ids = _metadata_sample_ids(sampler.dataset.metadata, split="heldout") + assignments = build_pdd_validation_assignments( + heldout_ids, + config.pdd, + validation_seed=config.training.validation_seed, + ) + sampler.set_epoch(0) + sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) + local_plan = [ + tuple(sampler.dataset.metadata[index]["sample_id"] for index in batch) for batch in sampler + ] + sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) + plans: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(plans, local_plan) + batch_counts = {len(plan) for plan in plans} + if len(batch_counts) != 1: + raise RuntimeError("PDD validation sampler produced different batch counts across ranks.") + + masks = [[([False] * len(batch)) for batch in plan] for plan in plans] + seen: set[str] = set() + for batch_index in range(len(local_plan)): + for rank, plan in enumerate(plans): + for position, sample_id in enumerate(plan[batch_index]): + if sample_id not in seen: + masks[rank][batch_index][position] = True + seen.add(sample_id) + if seen != set(heldout_ids): + missing = sorted(set(heldout_ids) - seen) + extra = sorted(seen - set(heldout_ids)) + raise RuntimeError( + f"PDD validation sampler does not cover the held-out split: " + f"missing={missing[:5]}, extra={extra[:5]}." + ) + local_masks = tuple(tuple(batch) for batch in masks[dist.get_rank()]) + return assignments, local_masks + + +def _iter_validation_batches( + dataloader: Any, + masks: tuple[tuple[bool, ...], ...], + config: Any, +): + from pdd_training import prepare_qwen_pdd_batch + + count = 0 + for count, (raw_batch, valid_mask) in enumerate(zip(dataloader, masks, strict=True), start=1): + prepared = prepare_qwen_pdd_batch( + raw_batch, + device=config.device, + dtype=config.dtype, + require_negative_condition=config.pdd.guidance_scale is not None, + ) + yield dataclasses.replace(prepared, valid_mask=valid_mask) + if count != len(masks): + raise RuntimeError( + f"PDD validation loader produced {count} batches for a {len(masks)}-batch plan." + ) + + +def _coverage_axis(counts: Any, loss_sums: Any) -> dict[int, dict[str, float | int]]: + return { + index: {"count": int(count), "mean_loss": float(loss_sums[index] / count)} + for index, count in enumerate(counts.tolist()) + if count + } + + +def _collective_training_iterator(dataloader: Any, sampler: Any) -> Any: + """Advance epochs and construct rank-local iterators under a collective error gate.""" + import torch.distributed as dist + + iterator = None + error_message = None + try: + if sampler.remaining_batches == 0: + sampler.set_epoch(sampler.epoch + 1) + iterator = iter(dataloader) + if iterator is None: + raise RuntimeError("PDD dataloader returned no iterator.") + except BaseException as error: + error_message = f"{type(error).__name__}: {error}" + errors: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(errors, error_message) + failures = [f"rank {rank}: {message}" for rank, message in enumerate(errors) if message] + if failures: + raise RuntimeError( + "distributed PDD training iterator construction failed; " + "; ".join(failures) + ) + if iterator is None: + raise RuntimeError("local PDD iterator construction succeeded without an iterator.") + return iterator + + +def _collective_training_batch( + iterator: Any, + *, + sampler: Any, + resume: Any, + resume_pending: bool, + device: Any, + dtype: Any, + require_negative_condition: bool, + expected_batch_size: int, +) -> tuple[Any, tuple[str, ...]] | None: + """Prepare one rank-local batch, then agree on success before any model call.""" + import torch.distributed as dist + from pdd_training import prepare_qwen_pdd_batch + + prepared = None + sample_ids: tuple[str, ...] = () + status: dict[str, Any] + try: + raw_batch = next(iterator) + except StopIteration: + if resume_pending: + status = { + "state": "error", + "error": "RuntimeError: resumed dataloader ended before its first batch", + "resume_pending": True, + } + else: + status = {"state": "end", "resume_pending": False} + except BaseException as error: + status = { + "state": "error", + "error": f"{type(error).__name__}: {error}", + "resume_pending": resume_pending, + } + else: + try: + metadata = raw_batch["metadata"] + sample_ids = tuple(metadata["sample_ids"]) + expected_ids = sampler.expected_next_sample_ids() + if sample_ids != expected_ids: + raise RuntimeError( + "prefetched PDD batch does not match committed cursor: " + f"expected={expected_ids}, actual={sample_ids}." + ) + if resume_pending: + if resume is None: + raise RuntimeError("resume_pending is true without a PDD resume state.") + resume.verify_first_batch(sample_ids) + prepared = prepare_qwen_pdd_batch( + raw_batch, + device=device, + dtype=dtype, + require_negative_condition=require_negative_condition, + ) + if prepared is None: + raise RuntimeError("PDD batch preparation returned no prepared batch.") + if len(sample_ids) != expected_batch_size: + raise RuntimeError( + f"PDD training batch has {len(sample_ids)} samples; " + f"expected {expected_batch_size}." + ) + status = { + "state": "batch", + "batch_size": len(sample_ids), + "resume_pending": resume_pending, + } + except BaseException as error: + status = { + "state": "error", + "error": f"{type(error).__name__}: {error}", + "resume_pending": resume_pending, + } + + statuses: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(statuses, status) + malformed = [rank for rank, item in enumerate(statuses) if not isinstance(item, Mapping)] + if malformed: + raise RuntimeError(f"PDD training ranks returned malformed statuses: {malformed}.") + failures = [ + f"rank {rank}: {item.get('error')}" + for rank, item in enumerate(statuses) + if item.get("state") == "error" + ] + if failures: + raise RuntimeError( + "distributed PDD training batch preflight failed; " + "; ".join(failures) + ) + states = {item.get("state") for item in statuses} + if states == {"end"}: + return None + if states != {"batch"}: + raise RuntimeError( + "distributed PDD training ranks produced different dataloader lengths: " + f"{[item.get('state') for item in statuses]}." + ) + pending = {item.get("resume_pending") for item in statuses} + if len(pending) != 1: + raise RuntimeError("distributed PDD training ranks disagree on resume verification state.") + batch_sizes = {item.get("batch_size") for item in statuses} + if batch_sizes != {expected_batch_size}: + raise RuntimeError( + f"distributed PDD training ranks disagree on batch size: {sorted(batch_sizes)}." + ) + if prepared is None: + raise RuntimeError("local PDD batch preparation succeeded without a prepared batch.") + return prepared, sample_ids + + def main() -> None: args = _parse_args() - from pdd_recipe import build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config + import torch + import torch.distributed as dist + from pdd_checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity + from pdd_recipe import ( + build_pdd_setup, + build_pdd_training_artifacts, + initialize_pdd_distributed, + resolve_pdd_recipe_config, + ) + from pdd_training import run_pdd_validation raw = yaml.safe_load(args.config.read_text()) config = resolve_pdd_recipe_config(raw) @@ -41,13 +389,231 @@ def main() -> None: backend="nccl" if config.device.type == "cuda" else "gloo", timeout_minutes=60, ) + rank = dist.get_rank() + world_size = dist.get_world_size() + logging.basicConfig( + level=logging.INFO if rank == 0 else logging.WARNING, + format="%(asctime)s %(levelname)s %(message)s", + force=True, + ) + snapshot_report = _validate_dataset_snapshot(raw, config) + dataloader, sampler = _build_training_dataloader( + raw, + config, + dp_rank=rank, + dp_world_size=world_size, + ) + validation_dataloader, validation_sampler = _build_validation_dataloader( + raw, + config, + dp_rank=rank, + dp_world_size=world_size, + ) + validation_assignments, validation_masks = _build_validation_plan( + validation_sampler, + config, + ) setup = build_pdd_setup(config) - logging.info( - "PDD setup complete: lifecycle=%s student_keys=%d AutoModel=%s", - setup.lifecycle, - len(setup.checkpoint_keys), - setup.automodel_snapshot["version"], + training = build_pdd_training_artifacts(setup, config) + identity = build_pdd_checkpoint_identity( + metadata=setup.metadata, + model_id=config.model_id, + model_revision=config.model_revision, + guidance_scale=config.pdd.guidance_scale, + guidance_rescale=config.guidance.rescale, + guidance_eps=config.guidance.eps, + automodel_snapshot=setup.automodel_snapshot, + ordered_train_id_sha256=_ordered_id_sha256(sampler.dataset.metadata, split="train"), + ordered_heldout_id_sha256=_ordered_id_sha256( + validation_sampler.dataset.metadata, + split="heldout", + ), + dataset_snapshot_sha256=snapshot_report["snapshot_sha256"], + local_batch_size=config.training.local_batch_size, + grad_accumulation_steps=config.training.grad_accumulation_steps, + training_seed=config.training.seed, + validation_seed=config.training.validation_seed, + validation_every_steps=config.training.validation_every_steps, + max_grad_norm=config.training.max_grad_norm, + zero_grad_warmup_steps=config.training.zero_grad_warmup_steps, + activation_checkpointing=config.parallel.activation_checkpointing, + dtype=str(config.dtype).removeprefix("torch."), + optimizer=setup.optimizer, + scheduler=training.scheduler, + ) + checkpoint_manager = PDDCheckpointManager( + root=config.checkpoint.checkpoint_dir, + checkpointer=setup.checkpointer, + model=setup.student, + optimizer=setup.optimizer, + scheduler=training.scheduler, + trainer=training.trainer, + sampler=sampler, + rng=training.rng, + identity=identity, ) + resume = checkpoint_manager.load(config.checkpoint.restore_from) + resume_pending = resume is not None + if resume is not None and rank == 0: + logging.info( + "PDD resume selected: checkpoint=%s parent=%s step=%d sample_slots=%d " + "expected_first_sample_ids=%s", + resume.checkpoint_path, + resume.parent_checkpoint, + resume.completed_steps, + resume.sample_slots_consumed, + resume.expected_next_sample_ids, + ) + if rank == 0: + logging.info( + "PDD dataset snapshot verified: sha256=%s splits=%s declared_files=%s", + snapshot_report["snapshot_sha256"], + snapshot_report["splits"], + snapshot_report["declared_files"], + ) + logging.info( + "PDD setup complete: lifecycle=%s student_keys=%d AutoModel=%s", + setup.lifecycle, + len(setup.checkpoint_keys), + setup.automodel_snapshot["version"], + ) + last_saved_step = 0 if resume is None else resume.completed_steps + try: + while training.trainer.completed_steps < config.training.max_steps: + iterator = _collective_training_iterator(dataloader, sampler) + data_wait_started = time.perf_counter() + while training.trainer.completed_steps < config.training.max_steps: + next_batch = _collective_training_batch( + iterator, + sampler=sampler, + resume=resume, + resume_pending=resume_pending, + device=config.device, + dtype=config.dtype, + require_negative_condition=config.pdd.guidance_scale is not None, + expected_batch_size=config.training.local_batch_size, + ) + if next_batch is None: + break + data_wait_seconds = time.perf_counter() - data_wait_started + step_started = time.perf_counter() + batch, sample_ids = next_batch + if resume_pending: + if rank == 0: + logging.info( + "PDD resume first batch verified: checkpoint=%s sample_ids=%s", + resume.checkpoint_path, + sample_ids, + ) + resume_pending = False + measure_update = ( + training.trainer.completed_steps + 1 + ) % config.training.log_every_steps == 0 + diagnostics = training.trainer.train_step( + batch, + measure_updates=measure_update, + ) + training.scheduler.step() + sampler.commit(sample_ids) + if sampler.remaining_batches == 0: + sampler.set_epoch(sampler.epoch + 1) + step_seconds = time.perf_counter() - step_started + + if diagnostics.completed_step % config.training.log_every_steps == 0: + timing = torch.tensor( + [data_wait_seconds, step_seconds], + dtype=torch.float64, + device=config.device, + ) + dist.all_reduce(timing, op=dist.ReduceOp.MAX) + peak_memory = ( + torch.cuda.max_memory_allocated(config.device) + if config.device.type == "cuda" + else 0 + ) + memory = torch.tensor(peak_memory, dtype=torch.int64, device=config.device) + dist.all_reduce(memory, op=dist.ReduceOp.MAX) + global_samples = config.training.local_batch_size * dist.get_world_size() + throughput = global_samples / max(float(timing[1].item()), 1e-12) + coverage = training.trainer.coverage + bin_loss = [ + None if count == 0 else float(loss_sum / count) + for loss_sum, count in zip( + coverage.bin_loss_sums.tolist(), + coverage.bin_counts.tolist(), + ) + ] + if rank == 0: + logging.info( + "PDD step=%d loss=%.6g grad_norm=%.6g nominal_update_ratio=%.6g " + "projection_update_ratio=%s lr=%.6g student_rms=%.6g " + "teacher_rms=%.6g student_teacher_rms_ratio=%.6g " + "reconstruction_rms=%.6g pairs=%d n_coverage=%s k_coverage=%s " + "bins=%s bin_loss=%s samples_per_second=%.3f " + "data_wait_seconds=%.4f peak_memory_bytes=%d", + diagnostics.completed_step, + diagnostics.loss, + diagnostics.grad_norm, + diagnostics.student_adamw_nominal_update_ratio, + diagnostics.pdd_projection_update_ratio, + diagnostics.learning_rate, + diagnostics.student_velocity_rms, + diagnostics.teacher_velocity_rms, + diagnostics.student_teacher_velocity_rms_ratio, + diagnostics.reconstructed_state_rms, + int((coverage.pair_counts > 0).sum()), + _coverage_axis(coverage.n_counts, coverage.n_loss_sums), + _coverage_axis(coverage.k_counts, coverage.k_loss_sums), + coverage.bin_counts.tolist(), + bin_loss, + throughput, + float(timing[0].item()), + int(memory.item()), + ) + if config.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(config.device) + if ( + diagnostics.completed_step % config.training.validation_every_steps == 0 + or diagnostics.completed_step >= config.training.max_steps + ): + validation_sampler.set_epoch(0) + validation_sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) + validation_result = run_pdd_validation( + training.pipeline, + _iter_validation_batches( + validation_dataloader, + validation_masks, + config, + ), + validation_assignments, + validation_seed=config.training.validation_seed, + ) + if rank == 0: + logging.info( + "PDD validation step=%d loss=%.12g pairs=%d starts=%d heads=%d " + "ordered_id_sha256=%s records=%d", + diagnostics.completed_step, + validation_result.mean_loss, + validation_result.pair_count, + validation_result.start_count, + validation_result.head_count, + validation_result.ordered_id_sha256, + len(validation_result.records), + ) + if ( + config.checkpoint.enabled + and diagnostics.completed_step % config.training.checkpoint_every_steps == 0 + ): + checkpoint_manager.save() + last_saved_step = diagnostics.completed_step + if diagnostics.completed_step >= config.training.max_steps: + break + data_wait_started = time.perf_counter() + + if config.checkpoint.enabled and last_saved_step != training.trainer.completed_steps: + checkpoint_manager.save() + finally: + setup.checkpointer.close() if __name__ == "__main__": diff --git a/examples/diffusers/fastgen/pdd_recipe.py b/examples/diffusers/fastgen/pdd_recipe.py index d4f4e7b8f16..367ab9cbb93 100644 --- a/examples/diffusers/fastgen/pdd_recipe.py +++ b/examples/diffusers/fastgen/pdd_recipe.py @@ -15,11 +15,15 @@ import torch import torch.distributed as dist +from portable_cache import validate_relative_reference from torch import nn from verify_readonly_automodel import snapshot_installed_distribution -from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection -from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd +from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection, PDDPipeline +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QwenImagePDDAdapter, + convert_qwen_image_to_pdd, +) @dataclass(frozen=True) @@ -37,9 +41,35 @@ class PDDCheckpointConfig: checkpoint_dir: str = "checkpoints/pdd_qwen_image" enabled: bool = True model_save_format: str = "torch_save" + restore_from: str | None = None save_consolidated: bool = False +@dataclass(frozen=True) +class PDDTrainingConfig: + """Direct-update and observability settings for the standalone PDD lifecycle.""" + + seed: int = 42 + max_steps: int = 10_000 + max_grad_norm: float = 1.0 + zero_grad_warmup_steps: int = 0 + log_every_steps: int = 10 + checkpoint_every_steps: int = 1_000 + validation_every_steps: int = 1_000 + local_batch_size: int = 1 + global_batch_size: int | None = None + grad_accumulation_steps: int = 1 + validation_seed: int = 2026 + + +@dataclass(frozen=True) +class PDDGuidanceConfig: + """Resolved Qwen packed-CFG norm-rescaling settings.""" + + rescale: float = 1.0 + eps: float = 1e-5 + + @dataclass(frozen=True) class PDDRecipeConfig: """Resolved setup inputs; incompatible mutation modes have already been rejected.""" @@ -49,8 +79,15 @@ class PDDRecipeConfig: pdd: PDDConfig parallel: PDDParallelConfig checkpoint: PDDCheckpointConfig + training: PDDTrainingConfig + guidance: PDDGuidanceConfig learning_rate: float weight_decay: float + adam_betas: tuple[float, float] + adam_eps: float + all_metadata_index: str + train_metadata_index: str + validation_metadata_index: str device: torch.device dtype: torch.dtype fuse_qkv_projections: bool @@ -74,6 +111,16 @@ class PDDSetupArtifacts: automodel_snapshot: Mapping[str, Any] +@dataclass(frozen=True) +class PDDTrainingArtifacts: + """Direct-update objects layered on the already-constructed Task-7 setup.""" + + pipeline: PDDPipeline + trainer: Any + scheduler: torch.optim.lr_scheduler.LRScheduler + rng: Any + + def _as_mapping(value: Any, *, name: str) -> Mapping[str, Any]: if not isinstance(value, Mapping): raise TypeError(f"{name} must be a mapping, got {type(value).__name__}.") @@ -92,6 +139,22 @@ def _require_bool(value: Any, *, name: str) -> bool: return value +def _require_int_at_least(value: Any, *, name: str, minimum: int) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}.") + return value + + +def _require_finite_real(value: Any, *, name: str, minimum: float | None = None) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise TypeError(f"{name} must be a real number.") + resolved = float(value) + if not math.isfinite(resolved) or (minimum is not None and resolved < minimum): + qualifier = "finite" if minimum is None else f"finite and >= {minimum}" + raise ValueError(f"{name} must be {qualifier}.") + return resolved + + def _resolve_dtype(value: Any) -> torch.dtype: if isinstance(value, torch.dtype): return value @@ -118,6 +181,40 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: fsdp = _as_mapping(raw.get("fsdp", {}), name="fsdp") optim = _as_mapping(raw.get("optim", {}), name="optim") checkpoint = _as_mapping(raw.get("checkpoint", {}), name="checkpoint") + training = _as_mapping(raw.get("training", {}), name="training") + guidance = _as_mapping(raw.get("guidance", {}), name="guidance") + data = _as_mapping(raw.get("data", {}), name="data") + dataloader = _as_mapping(data.get("dataloader", {}), name="data.dataloader") + + target = dataloader.get("_target_") + expected_target = "fastgen_data.build_text_to_image_multiresolution_dataloader" + if target is not None and target != expected_target: + raise ValueError(f"data.dataloader._target_ must be {expected_target!r}.") + if _require_bool(dataloader.get("drop_last", True), name="data.dataloader.drop_last") is False: + raise ValueError("PDD exact sample accounting requires data.dataloader.drop_last=true.") + if _require_bool( + dataloader.get("dynamic_batch_size", False), + name="data.dataloader.dynamic_batch_size", + ): + raise ValueError("PDD v1 requires data.dataloader.dynamic_batch_size=false.") + if _require_bool( + dataloader.get("train_text_encoder", False), + name="data.dataloader.train_text_encoder", + ): + raise ValueError("PDD requires cached text embeddings; train_text_encoder must be false.") + _require_bool(dataloader.get("shuffle", True), name="data.dataloader.shuffle") + all_metadata_index = validate_relative_reference( + data.get("all_metadata_index", "metadata.json"), + label="data.all_metadata_index", + ).as_posix() + train_metadata_index = validate_relative_reference( + dataloader.get("metadata_index", "metadata_train.json"), + label="data.dataloader.metadata_index", + ).as_posix() + validation_metadata_index = validate_relative_reference( + data.get("validation_metadata_index", "metadata_heldout.json"), + label="data.validation_metadata_index", + ).as_posix() _reject_enabled(model.get("transformer_engine_linear"), name="global TE-linear conversion") _reject_enabled(model.get("peft"), name="PEFT/LoRA") @@ -159,6 +256,25 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: raise ValueError("optim.learning_rate and weight_decay must be finite.") if learning_rate <= 0 or weight_decay < 0: raise ValueError("optim.learning_rate must be > 0 and weight_decay must be >= 0.") + adam_betas_raw = optim.get("betas", [0.9, 0.999]) + if ( + not isinstance(adam_betas_raw, list | tuple) + or len(adam_betas_raw) != 2 + or any( + isinstance(beta, bool) or not isinstance(beta, int | float) for beta in adam_betas_raw + ) + ): + raise TypeError("optim.betas must contain two real numbers.") + adam_betas = tuple(float(beta) for beta in adam_betas_raw) + if any(not math.isfinite(beta) or not 0.0 <= beta < 1.0 for beta in adam_betas): + raise ValueError("optim.betas values must be finite and in [0, 1).") + adam_eps = _require_finite_real( + optim.get("eps", 1e-8), + name="optim.eps", + minimum=0.0, + ) + if adam_eps == 0.0: + raise ValueError("optim.eps must be > 0.") dp_size = fsdp.get("dp_size") if dp_size is not None and (type(dp_size) is not int or dp_size < 1): @@ -190,6 +306,81 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: fuse_qkv_projections = _require_bool( model.get("fuse_qkv_projections", False), name="model.fuse_qkv_projections" ) + restore_from = checkpoint.get("restore_from") + if restore_from is not None and (not isinstance(restore_from, str) or not restore_from): + raise ValueError("checkpoint.restore_from must be null or a non-empty string.") + if not checkpoint_enabled and restore_from is not None: + raise ValueError("checkpoint.restore_from requires checkpoint.enabled=true.") + + seed = _require_int_at_least(training.get("seed", 42), name="training.seed", minimum=0) + max_steps = _require_int_at_least( + training.get("max_steps", 10_000), name="training.max_steps", minimum=1 + ) + zero_grad_warmup_steps = _require_int_at_least( + training.get("zero_grad_warmup_steps", 0), + name="training.zero_grad_warmup_steps", + minimum=0, + ) + log_every_steps = _require_int_at_least( + training.get("log_every_steps", 10), + name="training.log_every_steps", + minimum=1, + ) + checkpoint_every_steps = _require_int_at_least( + training.get("checkpoint_every_steps", 1_000), + name="training.checkpoint_every_steps", + minimum=1, + ) + validation_every_steps = _require_int_at_least( + training.get("validation_every_steps", 1_000), + name="training.validation_every_steps", + minimum=1, + ) + grad_accumulation_steps = _require_int_at_least( + training.get("grad_accumulation_steps", 1), + name="training.grad_accumulation_steps", + minimum=1, + ) + if grad_accumulation_steps != 1: + raise ValueError("PDD v1 exact resume requires training.grad_accumulation_steps=1.") + local_batch_size = _require_int_at_least( + dataloader.get("batch_size", training.get("local_batch_size", 1)), + name="data.dataloader.batch_size", + minimum=1, + ) + global_batch_size = training.get("global_batch_size") + if global_batch_size is not None: + global_batch_size = _require_int_at_least( + global_batch_size, + name="training.global_batch_size", + minimum=1, + ) + validation_seed = _require_int_at_least( + training.get("validation_seed", 2026), + name="training.validation_seed", + minimum=0, + ) + max_grad_norm = _require_finite_real( + training.get("max_grad_norm", 1.0), + name="training.max_grad_norm", + minimum=0.0, + ) + if max_grad_norm == 0.0: + raise ValueError("training.max_grad_norm must be > 0.") + guidance_rescale = _require_finite_real( + guidance.get("rescale", 1.0), + name="guidance.rescale", + minimum=0.0, + ) + if guidance_rescale > 1.0: + raise ValueError("guidance.rescale must be <= 1.") + guidance_eps = _require_finite_real( + guidance.get("eps", 1e-5), + name="guidance.eps", + minimum=0.0, + ) + if guidance_eps == 0.0: + raise ValueError("guidance.eps must be > 0.") return PDDRecipeConfig( model_id=model_id, @@ -203,10 +394,30 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: checkpoint_dir=checkpoint_dir, enabled=checkpoint_enabled, model_save_format=model_save_format, + restore_from=restore_from, save_consolidated=save_consolidated, ), + training=PDDTrainingConfig( + seed=seed, + max_steps=max_steps, + max_grad_norm=max_grad_norm, + zero_grad_warmup_steps=zero_grad_warmup_steps, + log_every_steps=log_every_steps, + checkpoint_every_steps=checkpoint_every_steps, + validation_every_steps=validation_every_steps, + local_batch_size=local_batch_size, + global_batch_size=global_batch_size, + grad_accumulation_steps=grad_accumulation_steps, + validation_seed=validation_seed, + ), + guidance=PDDGuidanceConfig(rescale=guidance_rescale, eps=guidance_eps), learning_rate=float(learning_rate), weight_decay=float(weight_decay), + adam_betas=adam_betas, + adam_eps=adam_eps, + all_metadata_index=all_metadata_index, + train_metadata_index=train_metadata_index, + validation_metadata_index=validation_metadata_index, device=torch.device(model.get("device", "cuda" if torch.cuda.is_available() else "cpu")), dtype=_resolve_dtype(model.get("torch_dtype", "bfloat16")), fuse_qkv_projections=fuse_qkv_projections, @@ -331,6 +542,18 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: lifecycle.append("qkv") world_size = dist.get_world_size() + if config.training.global_batch_size is not None: + effective_global_batch = ( + config.training.local_batch_size * world_size * config.training.grad_accumulation_steps + ) + if effective_global_batch != config.training.global_batch_size: + raise ValueError( + "PDD global batch mismatch: " + f"local_batch_size={config.training.local_batch_size} * world_size={world_size} " + f"* grad_accumulation_steps={config.training.grad_accumulation_steps} " + f"= {effective_global_batch}, configured " + f"training.global_batch_size={config.training.global_batch_size}." + ) dp_size = config.parallel.dp_size or world_size if dp_size != world_size: raise ValueError( @@ -367,6 +590,14 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: trainable, lr=config.learning_rate, weight_decay=config.weight_decay, + betas=config.adam_betas, + eps=config.adam_eps, + amsgrad=False, + capturable=False, + differentiable=False, + foreach=False, + fused=False, + maximize=False, ) optimizer_parameters = [ parameter for group in optimizer.param_groups for parameter in group["params"] @@ -414,6 +645,41 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: ) +def build_pdd_training_artifacts( + setup: PDDSetupArtifacts, + config: PDDRecipeConfig, +) -> PDDTrainingArtifacts: + """Layer the direct-update pipeline, constant-LR scheduler, and ranked RNG on setup.""" + if not isinstance(setup, PDDSetupArtifacts): + raise TypeError("setup must be PDDSetupArtifacts.") + if not isinstance(config, PDDRecipeConfig): + raise TypeError("config must be PDDRecipeConfig.") + from nemo_automodel.components.training.rng import StatefulRNG + from pdd_training import PDDTrainer + + adapter = QwenImagePDDAdapter( + config.pdd, + guidance_rescale=config.guidance.rescale, + guidance_eps=config.guidance.eps, + ) + pipeline = PDDPipeline(setup.student, setup.teacher, config.pdd, adapter) + scheduler = torch.optim.lr_scheduler.LambdaLR(setup.optimizer, lr_lambda=lambda _: 1.0) + rng = StatefulRNG(config.training.seed, ranked=True) + trainer = PDDTrainer( + pipeline, + setup.optimizer, + projection=setup.projection, + max_grad_norm=config.training.max_grad_norm, + warmup_steps=config.training.zero_grad_warmup_steps, + ) + return PDDTrainingArtifacts( + pipeline=pipeline, + trainer=trainer, + scheduler=scheduler, + rng=rng, + ) + + def initialize_pdd_distributed(*, backend: str, timeout_minutes: int = 60) -> Any: """Verify the wheel, then initialize through AutoModel's released public API.""" snapshot_installed_distribution() diff --git a/examples/diffusers/fastgen/pdd_training.py b/examples/diffusers/fastgen/pdd_training.py new file mode 100644 index 00000000000..23869cb15e6 --- /dev/null +++ b/examples/diffusers/fastgen/pdd_training.py @@ -0,0 +1,943 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct PDD updates and a logical-ID-stable held-out validation oracle.""" + +from __future__ import annotations + +import hashlib +import math +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + +from modelopt.torch.fastgen import PDDConfig, PDDOutputProjection, PDDPipeline + +_TRAINER_STATE_VERSION = 1 +_VALIDATION_SCHEMA_VERSION = 1 +_VALIDATION_ORDER_DOMAIN = b"modelopt-pdd-validation-order-v1\0" +_VALIDATION_PAIR_DOMAIN = b"modelopt-pdd-validation-pair-v1\0" +_VALIDATION_NOISE_DOMAIN = b"modelopt-pdd-validation-noise-v1\0" + + +@dataclass(frozen=True) +class PreparedPDDBatch: + """Canonical PDD inputs extracted from one Qwen cache batch.""" + + data: torch.Tensor + condition: tuple[torch.Tensor, torch.Tensor] + negative_condition: tuple[torch.Tensor, torch.Tensor] | None + sample_ids: tuple[str, ...] + valid_mask: tuple[bool, ...] | None = None + + +@dataclass(frozen=True) +class PDDStepDiagnostics: + """Host-side diagnostics for one completed direct student update.""" + + completed_step: int + loss: float + grad_norm: float + student_adamw_nominal_update_ratio: float | None + pdd_projection_update_ratio: float | None + learning_rate: float + n: tuple[int, ...] + k: tuple[int, ...] + student_velocity_rms: float + teacher_velocity_rms: float + student_teacher_velocity_rms_ratio: float + reconstructed_state_rms: float + + +@dataclass(frozen=True) +class PDDValidationAssignment: + """One canonical held-out logical ID and its explicit PDD target indices.""" + + ordinal: int + sample_id: str + n: int + k: int + + +@dataclass(frozen=True) +class PDDValidationRecord: + """Per-sample deterministic held-out result.""" + + ordinal: int + sample_id: str + n: int + k: int + loss: float + + +@dataclass(frozen=True) +class PDDValidationResult: + """Rank-invariant held-out records and canonical float64 aggregate.""" + + records: tuple[PDDValidationRecord, ...] + mean_loss: float + ordered_id_sha256: str + pair_count: int + start_count: int + head_count: int + schema_version: int = _VALIDATION_SCHEMA_VERSION + + +class PDDCoverage: + """Exact host-side n/k/pair coverage with deterministic coarse head bins.""" + + def __init__(self, config: PDDConfig, *, bins: int = 8) -> None: + if type(bins) is not int or bins <= 0: + raise ValueError("bins must be a positive integer.") + self.grid_size = config.grid_size + self.block_size_min = config.block_size_min + self.block_size_max = config.block_size_max + self.bins = min(bins, config.grid_size) + self.n_counts = torch.zeros(config.grid_size, dtype=torch.int64) + self.k_counts = torch.zeros(config.grid_size, dtype=torch.int64) + self.pair_counts = torch.zeros((config.grid_size, config.grid_size), dtype=torch.int64) + self.bin_counts = torch.zeros(self.bins, dtype=torch.int64) + self.n_loss_sums = torch.zeros(config.grid_size, dtype=torch.float64) + self.k_loss_sums = torch.zeros(config.grid_size, dtype=torch.float64) + self.pair_loss_sums = torch.zeros((config.grid_size, config.grid_size), dtype=torch.float64) + self.bin_loss_sums = torch.zeros(self.bins, dtype=torch.float64) + + def update(self, n: torch.Tensor, k: torch.Tensor, losses: torch.Tensor) -> None: + n_cpu = n.detach().to(device="cpu", dtype=torch.int64).reshape(-1) + k_cpu = k.detach().to(device="cpu", dtype=torch.int64).reshape(-1) + losses_cpu = losses.detach().to(device="cpu", dtype=torch.float64).reshape(-1) + if n_cpu.shape != k_cpu.shape or n_cpu.shape != losses_cpu.shape: + raise ValueError("n, k, and loss coverage tensors must have identical shapes.") + upper = torch.minimum( + n_cpu + self.block_size_max, + torch.full_like(n_cpu, self.grid_size), + ) + valid = ( + (n_cpu >= 0) + & (n_cpu < self.grid_size) + & (n_cpu.remainder(self.block_size_min) == 0) + & (k_cpu >= n_cpu) + & (k_cpu < upper) + ) + if not bool(valid.all()): + raise RuntimeError("observed n/k lies outside the exact configured support.") + device = n.device + n_device = n.detach().to(device=device, dtype=torch.int64).reshape(-1) + k_device = k.detach().to(device=device, dtype=torch.int64).reshape(-1) + losses_device = losses.detach().to(device=device, dtype=torch.float64).reshape(-1) + bins_device = torch.minimum( + k_device * self.bins // self.grid_size, + torch.full_like(k_device, self.bins - 1), + ) + pair_indices = n_device * self.grid_size + k_device + counts = ( + torch.bincount(n_device, minlength=self.grid_size), + torch.bincount(k_device, minlength=self.grid_size), + torch.bincount(pair_indices, minlength=self.grid_size * self.grid_size), + torch.bincount(bins_device, minlength=self.bins), + ) + loss_sums = [] + for indices, size in ( + (n_device, self.grid_size), + (k_device, self.grid_size), + (pair_indices, self.grid_size * self.grid_size), + (bins_device, self.bins), + ): + sums = torch.zeros(size, dtype=torch.float64, device=device) + sums.index_add_(0, indices, losses_device) + loss_sums.append(sums) + if dist.is_available() and dist.is_initialized(): + for tensor in (*counts, *loss_sums): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + self.n_counts += counts[0].cpu() + self.k_counts += counts[1].cpu() + self.pair_counts += counts[2].reshape(self.grid_size, self.grid_size).cpu() + self.bin_counts += counts[3].cpu() + self.n_loss_sums += loss_sums[0].cpu() + self.k_loss_sums += loss_sums[1].cpu() + self.pair_loss_sums += loss_sums[2].reshape(self.grid_size, self.grid_size).cpu() + self.bin_loss_sums += loss_sums[3].cpu() + + def require_pairs(self, expected: Sequence[tuple[int, int]]) -> None: + missing = [pair for pair in expected if int(self.pair_counts[pair]) == 0] + if missing: + raise RuntimeError(f"targeted PDD smoke did not cover required n/k pairs: {missing}.") + + def state_dict(self) -> dict[str, Any]: + return { + "grid_size": self.grid_size, + "block_size_min": self.block_size_min, + "block_size_max": self.block_size_max, + "bins": self.bins, + "n_counts": self.n_counts.clone(), + "k_counts": self.k_counts.clone(), + "pair_counts": self.pair_counts.clone(), + "bin_counts": self.bin_counts.clone(), + "n_loss_sums": self.n_loss_sums.clone(), + "k_loss_sums": self.k_loss_sums.clone(), + "pair_loss_sums": self.pair_loss_sums.clone(), + "bin_loss_sums": self.bin_loss_sums.clone(), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + expected = { + "grid_size", + "block_size_min", + "block_size_max", + "bins", + "n_counts", + "k_counts", + "pair_counts", + "bin_counts", + "n_loss_sums", + "k_loss_sums", + "pair_loss_sums", + "bin_loss_sums", + } + if not isinstance(state, Mapping) or set(state) != expected: + raise ValueError("PDD coverage state has incompatible keys.") + identity = ( + state["grid_size"], + state["block_size_min"], + state["block_size_max"], + state["bins"], + ) + if identity != ( + self.grid_size, + self.block_size_min, + self.block_size_max, + self.bins, + ): + raise ValueError("PDD coverage state does not match the current configuration.") + for name in ( + "n_counts", + "k_counts", + "pair_counts", + "bin_counts", + "n_loss_sums", + "k_loss_sums", + "pair_loss_sums", + "bin_loss_sums", + ): + saved = state[name] + current = getattr(self, name) + if not isinstance(saved, torch.Tensor) or saved.shape != current.shape: + raise ValueError(f"PDD coverage {name} has an incompatible tensor shape.") + current.copy_(saved.to(device="cpu", dtype=current.dtype)) + + +def prepare_qwen_pdd_batch( + batch: Mapping[str, Any], + *, + device: torch.device, + dtype: torch.dtype, + require_negative_condition: bool, +) -> PreparedPDDBatch: + """Move a portable Qwen cache batch into the PDD adapter contract.""" + if not isinstance(batch, Mapping): + raise TypeError(f"batch must be a mapping, got {type(batch).__name__}.") + required = {"image_latents", "text_embeddings", "text_embeddings_mask", "metadata"} + missing = sorted(required.difference(batch)) + if missing: + raise KeyError(f"Qwen PDD batch is missing required keys: {missing}.") + data = batch["image_latents"] + text = batch["text_embeddings"] + mask = batch["text_embeddings_mask"] + metadata = batch["metadata"] + if not all(isinstance(value, torch.Tensor) for value in (data, text, mask)): + raise TypeError("Qwen PDD latent, text embedding, and mask values must be tensors.") + if data.ndim != 4: + raise ValueError(f"Qwen PDD image_latents must be 4D, got {tuple(data.shape)}.") + if not isinstance(metadata, Mapping): + raise TypeError("Qwen PDD batch metadata must be a mapping.") + sample_ids = metadata.get("sample_ids") + if isinstance(sample_ids, str) or not isinstance(sample_ids, Sequence): + raise TypeError("Qwen PDD metadata.sample_ids must be a sequence of strings.") + sample_ids = tuple(sample_ids) + if len(sample_ids) != data.shape[0] or any( + not isinstance(sample_id, str) or not sample_id for sample_id in sample_ids + ): + raise ValueError("Qwen PDD sample_ids must be non-empty strings matching batch size.") + + data = data.to(device=device, dtype=dtype) + text = text.to(device=device, dtype=dtype) + mask = mask.to(device=device) + if text.ndim == 2: + text = text.unsqueeze(0).expand(data.shape[0], -1, -1).contiguous() + if mask.ndim == 1: + mask = mask.unsqueeze(0).expand(data.shape[0], -1).contiguous() + if text.shape[0] != data.shape[0] or mask.shape[0] != data.shape[0]: + raise ValueError("Qwen PDD conditioning batch size must match image_latents.") + condition = (text, mask) + + negative: tuple[torch.Tensor, torch.Tensor] | None = None + negative_text = batch.get("negative_text_embeddings") + negative_mask = batch.get("negative_text_embeddings_mask") + if negative_text is not None or negative_mask is not None: + if not isinstance(negative_text, torch.Tensor) or not isinstance( + negative_mask, torch.Tensor + ): + raise TypeError("negative Qwen conditioning requires embedding and mask tensors.") + negative_text = negative_text.to(device=device, dtype=dtype) + negative_mask = negative_mask.to(device=device) + if negative_text.ndim == 2: + negative_text = negative_text.unsqueeze(0).expand(data.shape[0], -1, -1).contiguous() + if negative_mask.ndim == 1: + negative_mask = negative_mask.unsqueeze(0).expand(data.shape[0], -1).contiguous() + if negative_text.shape[0] != data.shape[0] or negative_mask.shape[0] != data.shape[0]: + raise ValueError("negative Qwen conditioning batch size must match image_latents.") + negative = (negative_text, negative_mask) + if require_negative_condition and negative is None: + raise ValueError("guided Qwen PDD training requires negative prompt conditioning.") + return PreparedPDDBatch(data, condition, negative, sample_ids, (True,) * len(sample_ids)) + + +def _local_tensor(value: torch.Tensor) -> torch.Tensor: + to_local = getattr(value, "to_local", None) + return to_local() if callable(to_local) else value + + +def _replication_factor(value: torch.Tensor) -> int: + placements = getattr(value, "placements", ()) + mesh = getattr(value, "device_mesh", None) + factor = 1 + if mesh is not None: + for dimension, placement in enumerate(placements): + if type(placement).__name__ == "Replicate": + factor *= int(mesh.size(dimension)) + return factor + + +def _global_squared_sum(values: Sequence[torch.Tensor], *, device: torch.device) -> torch.Tensor: + total = torch.zeros((), dtype=torch.float64, device=device) + for value in values: + local = _local_tensor(value.detach()) + contribution = local.float().square().sum(dtype=torch.float64) + total += contribution / _replication_factor(value) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(total, op=dist.ReduceOp.SUM) + return total + + +def _global_any(flag: bool, *, device: torch.device) -> bool: + tensor = torch.tensor(int(flag), dtype=torch.int64, device=device) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return bool(tensor.item()) + + +def _all_parameters_finite(parameters: Sequence[torch.Tensor], *, device: torch.device) -> bool: + local_finite = torch.ones((), dtype=torch.bool, device=device) + for parameter in parameters: + local_finite.logical_and_(torch.isfinite(_local_tensor(parameter.detach())).all()) + return not _global_any(not bool(local_finite.item()), device=device) + + +def _global_sample_mean(value: torch.Tensor) -> float: + value = value.detach().reshape(-1) + totals = torch.stack( + ( + value.double().sum(), + torch.tensor(float(value.numel()), dtype=torch.float64, device=value.device), + ) + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(totals, op=dist.ReduceOp.SUM) + if totals[1] <= 0: + raise RuntimeError("cannot aggregate an empty PDD sample metric.") + return float((totals[0] / totals[1]).item()) + + +def _require_supported_adamw(optimizer: torch.optim.Optimizer) -> None: + if type(optimizer) is not torch.optim.AdamW: + raise TypeError("PDD v1 diagnostics require the stock torch.optim.AdamW optimizer.") + rejected_truthy = ("amsgrad", "maximize", "capturable", "differentiable", "foreach", "fused") + for group_index, group in enumerate(optimizer.param_groups): + for name in rejected_truthy: + if group.get(name) is not False: + raise ValueError( + f"PDD v1 requires AdamW param_groups[{group_index}][{name!r}]=False." + ) + lr = group.get("lr") + betas = group.get("betas") + eps = group.get("eps") + weight_decay = group.get("weight_decay") + if not isinstance(lr, float) or not isinstance(weight_decay, float): + raise TypeError("PDD v1 AdamW learning rate and weight decay must be scalar floats.") + if ( + not isinstance(betas, tuple) + or len(betas) != 2 + or any(not isinstance(beta, float) for beta in betas) + or not isinstance(eps, float) + ): + raise TypeError("PDD v1 AdamW betas and epsilon must be scalar floats.") + if not math.isfinite(lr) or lr <= 0 or not math.isfinite(weight_decay) or weight_decay < 0: + raise ValueError("PDD v1 AdamW learning rate/weight decay is invalid.") + if any(not math.isfinite(beta) or not 0.0 <= beta < 1.0 for beta in betas): + raise ValueError("PDD v1 AdamW betas must be finite and in [0, 1).") + if not math.isfinite(eps) or eps <= 0: + raise ValueError("PDD v1 AdamW epsilon must be finite and > 0.") + + +def _adamw_nominal_update_ratio( + optimizer: torch.optim.Optimizer, + *, + device: torch.device, +) -> float: + """Stream the public AdamW equation over local shards without cloning the model.""" + update_squared = torch.zeros((), dtype=torch.float64, device=device) + parameter_squared = torch.zeros((), dtype=torch.float64, device=device) + for group in optimizer.param_groups: + lr = group["lr"] + beta1, beta2 = group["betas"] + eps = group["eps"] + decay = 1.0 - lr * group["weight_decay"] + if decay <= 0.0: + raise RuntimeError("AdamW decoupled weight decay factor must remain positive.") + for parameter in group["params"]: + state = optimizer.state.get(parameter) + if not state or "exp_avg" not in state or "exp_avg_sq" not in state: + continue + step_value = state.get("step") + if isinstance(step_value, torch.Tensor): + step = float(step_value.item()) + else: + step = float(step_value) + if step <= 0: + raise RuntimeError("AdamW state has a non-positive step after optimizer.step().") + parameter_local = _local_tensor(parameter.detach()).float() + exp_avg = _local_tensor(state["exp_avg"].detach()).float() + exp_avg_sq = _local_tensor(state["exp_avg_sq"].detach()).float() + bias_correction1 = 1.0 - beta1**step + bias_correction2_sqrt = math.sqrt(1.0 - beta2**step) + denominator = exp_avg_sq.sqrt().div_(bias_correction2_sqrt).add_(eps) + direction = exp_avg.div(denominator).div_(bias_correction1) + parameter_before = (parameter_local + lr * direction) / decay + nominal_delta = parameter_local - parameter_before + factor = _replication_factor(parameter) + update_squared += nominal_delta.square().sum(dtype=torch.float64) / factor + parameter_squared += parameter_before.square().sum(dtype=torch.float64) / factor + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(update_squared, op=dist.ReduceOp.SUM) + dist.all_reduce(parameter_squared, op=dist.ReduceOp.SUM) + if not bool(torch.isfinite(update_squared) & torch.isfinite(parameter_squared)): + raise RuntimeError("PDD AdamW nominal update diagnostics became non-finite.") + return float((update_squared.sqrt() / parameter_squared.sqrt().clamp_min(1e-30)).item()) + + +class PDDTrainer: + """Own direct student updates while leaving algorithm and checkpoint state separate.""" + + def __init__( + self, + pipeline: PDDPipeline, + optimizer: torch.optim.Optimizer, + *, + projection: PDDOutputProjection, + max_grad_norm: float, + warmup_steps: int = 0, + ) -> None: + if not isinstance(pipeline, PDDPipeline): + raise TypeError(f"pipeline must be PDDPipeline, got {type(pipeline).__name__}.") + if not isinstance(projection, PDDOutputProjection): + raise TypeError("projection must be PDDOutputProjection.") + if isinstance(max_grad_norm, bool) or not isinstance(max_grad_norm, int | float): + raise TypeError("max_grad_norm must be a real number.") + if not math.isfinite(max_grad_norm) or max_grad_norm <= 0: + raise ValueError("max_grad_norm must be finite and > 0.") + if type(warmup_steps) is not int or warmup_steps < 0: + raise ValueError("warmup_steps must be an integer >= 0.") + _require_supported_adamw(optimizer) + self.pipeline = pipeline + self.optimizer = optimizer + self.projection = projection + self.max_grad_norm = float(max_grad_norm) + self.warmup_steps = warmup_steps + self.completed_steps = 0 + self.consecutive_zero_grad_steps = 0 + self.coverage = PDDCoverage(pipeline.config) + + def _projection_snapshot(self) -> list[torch.Tensor]: + parameters = [self.projection.weight] + if self.projection.bias is not None: + parameters.append(self.projection.bias) + return [_local_tensor(parameter.detach()).clone() for parameter in parameters] + + def _projection_update_ratio(self, before: Sequence[torch.Tensor]) -> float: + parameters = [self.projection.weight] + if self.projection.bias is not None: + parameters.append(self.projection.bias) + update_squared = torch.zeros((), dtype=torch.float64, device=self.pipeline.device) + parameter_squared = torch.zeros((), dtype=torch.float64, device=self.pipeline.device) + for parameter, saved in zip(parameters, before): + local = _local_tensor(parameter.detach()).float() + factor = _replication_factor(parameter) + update_squared += (local - saved.float()).square().sum(dtype=torch.float64) / factor + parameter_squared += local.square().sum(dtype=torch.float64) / factor + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(update_squared, op=dist.ReduceOp.SUM) + dist.all_reduce(parameter_squared, op=dist.ReduceOp.SUM) + if not bool(torch.isfinite(update_squared) & torch.isfinite(parameter_squared)): + raise RuntimeError("PDD projection update diagnostics became non-finite.") + return float((update_squared.sqrt() / parameter_squared.sqrt().clamp_min(1e-30)).item()) + + def train_step( + self, + batch: PreparedPDDBatch, + *, + noise: torch.Tensor | None = None, + n: torch.Tensor | None = None, + k: torch.Tensor | None = None, + generator: torch.Generator | None = None, + measure_updates: bool = True, + ) -> PDDStepDiagnostics: + """Run one direct PDD student update and enforce all immediate hard aborts.""" + if not isinstance(batch, PreparedPDDBatch): + raise TypeError("batch must be PreparedPDDBatch.") + self.pipeline.student.train() + self.pipeline.teacher.eval() + self.optimizer.zero_grad(set_to_none=True) + before_projection = self._projection_snapshot() if measure_updates else None + loss, metrics = self.pipeline.compute_loss( + batch.data, + noise=noise, + condition=batch.condition, + negative_condition=batch.negative_condition, + n=n, + k=k, + generator=generator, + ) + finite_metrics = ( + "all_student_heads_finite", + "student_target_finite", + "teacher_target_finite", + "reconstructed_state_finite", + "loss_finite", + ) + local_nonfinite = not bool(torch.isfinite(loss)) or any( + not bool(metrics[name].all()) for name in finite_metrics + ) + if _global_any(local_nonfinite, device=batch.data.device): + raise FloatingPointError( + "PDD loss, prediction, target, or reconstruction is non-finite." + ) + self.coverage.update(metrics["n"], metrics["k"], metrics["student_target_mse"]) + loss.backward() + + teacher_gradient = any( + parameter.grad is not None for parameter in self.pipeline.teacher.parameters() + ) + if _global_any(teacher_gradient, device=batch.data.device): + raise RuntimeError("PDD frozen teacher received a gradient.") + trainable = [ + parameter for parameter in self.pipeline.student.parameters() if parameter.requires_grad + ] + gradients = [parameter.grad for parameter in trainable if parameter.grad is not None] + if not gradients: + grad_norm = 0.0 + else: + grad_squared = _global_squared_sum(gradients, device=batch.data.device) + if not bool(torch.isfinite(grad_squared)): + raise FloatingPointError("PDD student gradient became non-finite.") + grad_norm = float(grad_squared.sqrt().item()) + if self.completed_steps >= self.warmup_steps and grad_norm == 0.0: + self.consecutive_zero_grad_steps += 1 + if self.consecutive_zero_grad_steps >= 2: + raise RuntimeError("PDD student gradient was zero for two consecutive updates.") + else: + self.consecutive_zero_grad_steps = 0 + clip_coefficient = min(1.0, self.max_grad_norm / (grad_norm + 1e-6)) + if clip_coefficient < 1.0: + for gradient in gradients: + gradient.mul_(clip_coefficient) + + self.optimizer.step() + if not _all_parameters_finite(trainable, device=batch.data.device): + raise FloatingPointError("PDD student parameter update became non-finite.") + nominal_ratio = ( + _adamw_nominal_update_ratio( + self.optimizer, + device=batch.data.device, + ) + if measure_updates + else None + ) + projection_ratio = ( + self._projection_update_ratio(before_projection) + if before_projection is not None + else None + ) + if nominal_ratio == 0.0 and grad_norm > 0.0: + raise RuntimeError( + "PDD optimizer produced a zero nominal update from a nonzero gradient." + ) + self.completed_steps += 1 + student_velocity_rms = _global_sample_mean(metrics["student_velocity_rms"]) + teacher_velocity_rms = _global_sample_mean(metrics["teacher_velocity_rms"]) + + return PDDStepDiagnostics( + completed_step=self.completed_steps, + loss=_global_sample_mean(metrics["student_target_mse"]), + grad_norm=grad_norm, + student_adamw_nominal_update_ratio=nominal_ratio, + pdd_projection_update_ratio=projection_ratio, + learning_rate=float(self.optimizer.param_groups[0]["lr"]), + n=tuple(int(value) for value in metrics["n"].detach().cpu().tolist()), + k=tuple(int(value) for value in metrics["k"].detach().cpu().tolist()), + student_velocity_rms=student_velocity_rms, + teacher_velocity_rms=teacher_velocity_rms, + student_teacher_velocity_rms_ratio=student_velocity_rms + / max(teacher_velocity_rms, 1e-30), + reconstructed_state_rms=_global_sample_mean(metrics["reconstructed_state_rms"]), + ) + + def state_dict(self) -> dict[str, Any]: + return { + "schema_version": _TRAINER_STATE_VERSION, + "completed_steps": self.completed_steps, + "consecutive_zero_grad_steps": self.consecutive_zero_grad_steps, + "coverage": self.coverage.state_dict(), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + expected = { + "schema_version", + "completed_steps", + "consecutive_zero_grad_steps", + "coverage", + } + if not isinstance(state, Mapping) or set(state) != expected: + raise ValueError("PDD trainer state has incompatible keys.") + if state["schema_version"] != _TRAINER_STATE_VERSION: + raise ValueError(f"unsupported PDD trainer schema {state['schema_version']!r}.") + completed = state["completed_steps"] + zero_steps = state["consecutive_zero_grad_steps"] + if type(completed) is not int or completed < 0: + raise ValueError("saved completed_steps must be an integer >= 0.") + if type(zero_steps) is not int or zero_steps < 0: + raise ValueError("saved consecutive_zero_grad_steps must be an integer >= 0.") + self.completed_steps = completed + self.consecutive_zero_grad_steps = zero_steps + self.coverage.load_state_dict(state["coverage"]) + + +def _stable_digest(domain: bytes, validation_seed: int, payload: str) -> bytes: + digest = hashlib.sha256() + digest.update(domain) + digest.update(str(validation_seed).encode()) + digest.update(b"\0") + digest.update(payload.encode()) + return digest.digest() + + +def pdd_validation_support(config: PDDConfig) -> tuple[tuple[int, int], ...]: + """Return the exact lexicographic support of explicit PDD validation pairs.""" + return tuple( + (n, k) + for n in range(0, config.grid_size, config.block_size_min) + for k in range(n, min(n + config.block_size_max, config.grid_size)) + ) + + +def build_pdd_validation_assignments( + sample_ids: Sequence[str], + config: PDDConfig, + *, + validation_seed: int, + require_full_coverage: bool = True, +) -> tuple[PDDValidationAssignment, ...]: + """Assign every logical ID a rank/batch-order-independent explicit n/k pair.""" + if isinstance(sample_ids, str) or not isinstance(sample_ids, Sequence): + raise TypeError("sample_ids must be a sequence of strings.") + if any(not isinstance(sample_id, str) or not sample_id for sample_id in sample_ids): + raise ValueError("sample_ids must contain non-empty strings.") + if not sample_ids: + raise ValueError("sample_ids must contain at least one logical ID.") + if len(set(sample_ids)) != len(sample_ids): + raise ValueError("held-out validation sample_ids must be unique.") + if type(validation_seed) is not int or validation_seed < 0: + raise ValueError("validation_seed must be an integer >= 0.") + support = pdd_validation_support(config) + if require_full_coverage and len(sample_ids) < len(support): + raise ValueError( + f"full PDD validation coverage requires at least {len(support)} logical IDs, " + f"found {len(sample_ids)}." + ) + ordered_ids = sorted( + sample_ids, + key=lambda sample_id: ( + _stable_digest(_VALIDATION_ORDER_DOMAIN, validation_seed, sample_id), + sample_id, + ), + ) + permuted_support = sorted( + support, + key=lambda pair: ( + _stable_digest(_VALIDATION_PAIR_DOMAIN, validation_seed, f"{pair[0]}:{pair[1]}"), + pair, + ), + ) + return tuple( + PDDValidationAssignment(ordinal, sample_id, *permuted_support[ordinal % len(support)]) + for ordinal, sample_id in enumerate(ordered_ids) + ) + + +def pdd_validation_noise( + sample_id: str, + shape: Sequence[int], + *, + validation_seed: int, + device: torch.device, +) -> torch.Tensor: + """Generate per-ID CPU float32 noise without touching the global RNG.""" + digest = _stable_digest(_VALIDATION_NOISE_DOMAIN, validation_seed, sample_id) + seed = int.from_bytes(digest[:8], "big") & ((1 << 63) - 1) + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + return torch.randn(tuple(shape), generator=generator, dtype=torch.float32).to(device) + + +def _ordered_id_digest(records: Sequence[PDDValidationRecord]) -> str: + digest = hashlib.sha256() + digest.update(b"modelopt-pdd-ordered-validation-ids-v1\0") + for record in records: + digest.update(record.sample_id.encode()) + digest.update(b"\n") + return digest.hexdigest() + + +def _raise_collective_validation_error(error: BaseException | None, *, context: str) -> None: + if not dist.is_available() or not dist.is_initialized(): + if error is not None: + raise error + return + local = None if error is None else f"{type(error).__name__}: {error}" + errors: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(errors, local) + failures = [f"rank {rank}: {message}" for rank, message in enumerate(errors) if message] + if failures: + raise RuntimeError(f"distributed PDD validation {context} failed; " + "; ".join(failures)) + + +def run_pdd_validation( + pipeline: PDDPipeline, + batches: Iterable[PreparedPDDBatch], + assignments: Sequence[PDDValidationAssignment], + *, + validation_seed: int, +) -> PDDValidationResult: + """Evaluate explicit per-ID targets and aggregate identically across rank partitions.""" + if not isinstance(pipeline, PDDPipeline): + raise TypeError("pipeline must be PDDPipeline.") + distributed = dist.is_available() and dist.is_initialized() + assignment_by_id = {assignment.sample_id: assignment for assignment in assignments} + assignment_error: BaseException | None = None + if len(assignment_by_id) != len(assignments): + assignment_error = ValueError("validation assignments contain duplicate logical IDs.") + elif not assignments: + assignment_error = ValueError("validation assignments cannot be empty.") + _raise_collective_validation_error(assignment_error, context="assignment preflight") + if distributed: + assignment_identity = tuple( + (item.ordinal, item.sample_id, item.n, item.k) for item in assignments + ) + assignment_identities: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(assignment_identities, assignment_identity) + if any(identity != assignment_identities[0] for identity in assignment_identities[1:]): + raise RuntimeError("distributed PDD validation assignments differ across ranks.") + student_was_training = pipeline.student.training + teacher_was_training = pipeline.teacher.training + pipeline.student.eval() + pipeline.teacher.eval() + local_records: list[PDDValidationRecord] = [] + try: + with torch.no_grad(): + iterator = iter(batches) + batch_index = 0 + while True: + batch = None + next_error: BaseException | None = None + exhausted = False + try: + batch = next(iterator) + except StopIteration: + exhausted = True + except BaseException as error: + next_error = error + if distributed: + status = "error" if next_error is not None else "end" if exhausted else "batch" + statuses: list[str] = [""] * dist.get_world_size() + dist.all_gather_object(statuses, status) + if "error" in statuses: + _raise_collective_validation_error(next_error, context="iteration") + if all(item == "end" for item in statuses): + break + if any(item != "batch" for item in statuses): + raise RuntimeError( + "distributed PDD validation ranks produced different batch counts." + ) + else: + if next_error is not None: + raise next_error + if exhausted: + break + + local_error: BaseException | None = None + selected: list[PDDValidationAssignment] = [] + valid_mask: tuple[bool, ...] = () + signature: Any = None + try: + if not isinstance(batch, PreparedPDDBatch): + raise TypeError("validation batches must contain PreparedPDDBatch values.") + valid_mask = ( + (True,) * len(batch.sample_ids) + if batch.valid_mask is None + else batch.valid_mask + ) + if len(valid_mask) != len(batch.sample_ids) or any( + type(valid) is not bool for valid in valid_mask + ): + raise ValueError( + "validation valid_mask must contain one bool per sample ID." + ) + for position, (sample_id, valid) in enumerate( + zip(batch.sample_ids, valid_mask) + ): + if valid and sample_id not in assignment_by_id: + raise ValueError( + f"validation batch contains unassigned sample ID {sample_id!r}." + ) + if valid: + selected.append(assignment_by_id[sample_id]) + else: + template = assignments[(batch_index + position) % len(assignments)] + selected.append( + PDDValidationAssignment( + template.ordinal, + f"__pdd_dummy__:{batch_index}:{position}", + template.n, + template.k, + ) + ) + signature = ( + tuple(batch.data.shape), + tuple(batch.condition[0].shape), + tuple(batch.condition[1].shape), + None + if batch.negative_condition is None + else ( + tuple(batch.negative_condition[0].shape), + tuple(batch.negative_condition[1].shape), + ), + ) + except BaseException as error: + local_error = error + _raise_collective_validation_error(local_error, context="batch preflight") + assert isinstance(batch, PreparedPDDBatch) + + if distributed: + signatures: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(signatures, signature) + if any(item != signatures[0] for item in signatures[1:]): + raise RuntimeError( + "distributed PDD validation requires the same padded batch shape " + "on every rank." + ) + noise = torch.stack( + [ + pdd_validation_noise( + assignment.sample_id, + batch.data.shape[1:], + validation_seed=validation_seed, + device=batch.data.device, + ) + for assignment in selected + ] + ) + n = torch.tensor( + [assignment.n for assignment in selected], + dtype=torch.long, + device=batch.data.device, + ) + k = torch.tensor( + [assignment.k for assignment in selected], + dtype=torch.long, + device=batch.data.device, + ) + _, metrics = pipeline.compute_loss( + batch.data, + noise=noise, + condition=batch.condition, + negative_condition=batch.negative_condition, + n=n, + k=k, + ) + finite_metrics = ( + "all_student_heads_finite", + "student_target_finite", + "teacher_target_finite", + "reconstructed_state_finite", + "loss_finite", + ) + local_nonfinite = any(not bool(metrics[name].all()) for name in finite_metrics) + losses = metrics["student_target_mse"].double().cpu().tolist() + local_nonfinite = local_nonfinite or any(not math.isfinite(loss) for loss in losses) + if _global_any(local_nonfinite, device=batch.data.device): + raise FloatingPointError( + "deterministic PDD validation produced a non-finite prediction, target, " + "or loss." + ) + local_records.extend( + PDDValidationRecord( + assignment.ordinal, + assignment.sample_id, + assignment.n, + assignment.k, + float(loss), + ) + for assignment, loss, valid in zip(selected, losses, valid_mask) + if valid + ) + batch_index += 1 + finally: + pipeline.student.train(student_was_training) + pipeline.teacher.train(teacher_was_training) + + gathered: list[list[PDDValidationRecord]] + if distributed: + gathered = [[] for _ in range(dist.get_world_size())] + dist.all_gather_object(gathered, local_records) + else: + gathered = [local_records] + records = sorted( + (record for rank_records in gathered for record in rank_records), + key=lambda record: record.ordinal, + ) + if len(records) != len(assignments): + raise RuntimeError( + f"validation contributed {len(records)} records for {len(assignments)} assignments." + ) + if [record.ordinal for record in records] != list(range(len(assignments))): + raise RuntimeError("validation ordinals are duplicated or missing across ranks.") + for record, assignment in zip(records, assignments): + if (record.sample_id, record.n, record.k) != ( + assignment.sample_id, + assignment.n, + assignment.k, + ): + raise RuntimeError("validation record does not match its canonical assignment.") + pairs = {(record.n, record.k) for record in records} + starts = {record.n for record in records} + heads = {record.k for record in records} + return PDDValidationResult( + records=tuple(records), + mean_loss=math.fsum(record.loss for record in records) / len(records), + ordered_id_sha256=_ordered_id_digest(records), + pair_count=len(pairs), + start_count=len(starts), + head_count=len(heads), + ) diff --git a/examples/diffusers/fastgen/validate_cache_snapshot.py b/examples/diffusers/fastgen/validate_cache_snapshot.py index 38ef7031860..b7c465a94ac 100644 --- a/examples/diffusers/fastgen/validate_cache_snapshot.py +++ b/examples/diffusers/fastgen/validate_cache_snapshot.py @@ -18,6 +18,7 @@ from __future__ import annotations import argparse +import hashlib from pathlib import Path from typing import Any @@ -31,7 +32,7 @@ ) -def _validate_payload(root: Path, entry: dict[str, Any]) -> Path: +def _validate_payload(root: Path, entry: dict[str, Any]) -> tuple[Path, str]: payload_path = resolve_cache_asset(root, entry["cache_file"], label="cache_file") actual_digest = sha256_file(payload_path) if actual_digest != entry["payload_sha256"]: @@ -47,7 +48,7 @@ def _validate_payload(root: Path, entry: dict[str, Any]) -> Path: raise ValueError(f"payload/manifest sample_id mismatch for {entry['cache_file']}") if payload.get("source_ref") != entry.get("source_ref"): raise ValueError(f"payload/manifest source_ref mismatch for {entry['cache_file']}") - return payload_path + return payload_path, actual_digest def validate_snapshot( @@ -133,7 +134,8 @@ def validate_snapshot( if all_ids is not None and train_ids | heldout_ids != all_ids: raise ValueError("train and heldout split union does not equal the all split") - payload_files = {_validate_payload(root, entry) for entry in entries_by_id.values()} + payload_hashes = dict(_validate_payload(root, entry) for entry in entries_by_id.values()) + payload_files = set(payload_hashes) declared_files.update(payload_files) if reject_orphans: @@ -162,11 +164,27 @@ def validate_snapshot( if undeclared: raise ValueError(f"snapshot contains undeclared files: {undeclared[:5]}") + declared_hashes = { + path.relative_to(root).as_posix(): ( + payload_hashes[path] if path in payload_hashes else sha256_file(path) + ) + for path in declared_files + } + snapshot_digest = hashlib.sha256() + snapshot_digest.update(b"modelopt-fastgen-cache-snapshot-v1\0") + for relative, file_digest in sorted(declared_hashes.items()): + snapshot_digest.update(relative.encode()) + snapshot_digest.update(b"\0") + snapshot_digest.update(file_digest.encode()) + snapshot_digest.update(b"\n") + return { "root": str(root), "indexes": list(expected_indexes.values()), "splits": {name: len(ids) for name, ids in split_ids.items()}, "unique_payloads": len(payload_files), + "declared_files": len(declared_hashes), + "snapshot_sha256": snapshot_digest.hexdigest(), } diff --git a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py new file mode 100644 index 00000000000..20b059b9e65 --- /dev/null +++ b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Two-rank proof that rank-0 checkpoint failures propagate instead of deadlocking.""" + +from __future__ import annotations + +import pathlib +import shutil +import sys +import tempfile +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +import pdd_checkpoint as pdd_checkpoint_module +from pdd_checkpoint import PDDCheckpointManager + + +class _State: + def state_dict(self): + return {"value": 1} + + +class _Sampler(_State): + def state_dict(self): + return { + "epoch": 0, + "committed_batches": 1, + "sample_slots_consumed": 1, + "plan_sha256": "0" * 64, + "next_sample_ids": ["next"], + } + + +class _Trainer(_State): + def __init__(self, completed_steps: int = 1) -> None: + self.completed_steps = completed_steps + + +class _Checkpointer: + def __init__(self, rank: int, *, fail_sidecar: bool = False) -> None: + self.config = SimpleNamespace(is_async=False) + self.rank = rank + self.fail_sidecar = fail_sidecar + + def save_model(self, model, path: str) -> None: + del model + if self.rank == 0: + root = pathlib.Path(path) / "model" + root.mkdir(parents=True) + (root / ".metadata").write_bytes(b"metadata") + (root / "__0_0.distcp").write_bytes(b"model") + + def save_optimizer(self, optimizer, model, path: str, scheduler) -> None: + del optimizer, model, scheduler + if self.rank == 0: + root = pathlib.Path(path) / "optim" + root.mkdir(parents=True) + (root / ".metadata").write_bytes(b"metadata") + (root / "__0_0.distcp").write_bytes(b"optim") + + def save_on_dp_ranks(self, state, state_name: str, path: str) -> None: + if self.fail_sidecar and self.rank == 1 and state_name == "sampler": + raise OSError("injected rank-1 sidecar failure") + root = pathlib.Path(path) / state_name + root.mkdir(parents=True, exist_ok=True) + torch.save(state.state_dict(), root / f"{state_name}_dp_rank_{self.rank}.pt") + + +class _FailingManager(PDDCheckpointManager): + def __init__(self, *, failure_stage: str, **kwargs) -> None: + super().__init__(**kwargs) + self.failure_stage = failure_stage + + def _prepare_staging(self, final: pathlib.Path) -> str: + if self.failure_stage == "prepare": + raise OSError("injected preparation failure") + return super()._prepare_staging(final) + + def _publish_staging(self, **kwargs) -> None: + if self.failure_stage == "publish": + raise OSError("injected publication failure") + if self.failure_stage != "latest": + super()._publish_staging(**kwargs) + return + original = pdd_checkpoint_module._atomic_text + + def fail_latest(path: pathlib.Path, text: str) -> None: + if path.name == "LATEST": + raise OSError("injected LATEST update failure") + original(path, text) + + pdd_checkpoint_module._atomic_text = fail_latest + try: + super()._publish_staging(**kwargs) + finally: + pdd_checkpoint_module._atomic_text = original + + +def _run_failure(root: pathlib.Path, stage: str) -> None: + rank = dist.get_rank() + trainer = _Trainer() + if stage == "latest": + initial = _FailingManager( + failure_stage="none", + root=root / stage, + checkpointer=_Checkpointer(rank), + model=object(), + optimizer=SimpleNamespace(param_groups=[{"lr": 2.0e-5}]), + scheduler=object(), + trainer=trainer, + sampler=_Sampler(), + rng=_State(), + identity={"schema_version": 1, "topology": {"world_size": 2}}, + ) + initial.save() + trainer.completed_steps = 2 + manager = _FailingManager( + failure_stage=stage, + root=root / stage, + checkpointer=_Checkpointer(rank, fail_sidecar=stage == "sidecar"), + model=object(), + optimizer=SimpleNamespace(param_groups=[{"lr": 2.0e-5}]), + scheduler=object(), + trainer=trainer, + sampler=_Sampler(), + rng=_State(), + identity={"schema_version": 1, "topology": {"world_size": 2}}, + ) + message = None + try: + manager.save() + except RuntimeError as error: + message = str(error) + messages: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(messages, message) + if stage == "sidecar": + assert all( + item is not None and "checkpoint sidecar save failed" in item for item in messages + ) + else: + expected = "preparation" if stage == "prepare" else "publication" + assert all( + item is not None and f"rank-0 checkpoint {expected}" in item for item in messages + ) + if stage == "latest": + assert (root / stage / "LATEST").read_text().strip() == "step_00000001" + assert manager.resolve("LATEST").name == "step_00000002" + + +def main() -> None: + dist.init_process_group("gloo") + root_payload = [ + tempfile.mkdtemp(prefix="modelopt-pdd-rank0-failure-") if dist.get_rank() == 0 else None + ] + dist.broadcast_object_list(root_payload, src=0) + root = pathlib.Path(root_payload[0]) + try: + _run_failure(root, "prepare") + dist.barrier() + _run_failure(root, "publish") + dist.barrier() + _run_failure(root, "sidecar") + dist.barrier() + _run_failure(root, "latest") + finally: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(root) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/examples/diffusers/fastgen/pdd_test_utils.py b/tests/examples/diffusers/fastgen/pdd_test_utils.py new file mode 100644 index 00000000000..06d0b11c20d --- /dev/null +++ b/tests/examples/diffusers/fastgen/pdd_test_utils.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small plain-torch objects shared by PDD example lifecycle tests.""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from typing import Any + +import torch +from pdd_training import PDDTrainer, PreparedPDDBatch +from torch import nn + +from modelopt.torch.fastgen import ( + PDDConfig, + PDDLayerSpec, + PDDMetadata, + PDDOutputProjection, + PDDPipeline, + convert_to_pdd_output_projection, +) + + +class ToyStudent(nn.Module): + def __init__(self, width: int = 3) -> None: + super().__init__() + self.backbone = nn.Linear(width, width) + self.projection = nn.Linear(width, width) + + def forward(self, state: torch.Tensor) -> torch.Tensor: + return self.projection(torch.tanh(self.backbone(state))) + + +class ToyTeacher(nn.Module): + def __init__(self) -> None: + super().__init__() + self.scale = nn.Parameter(torch.tensor(-0.25)) + self.bias = nn.Parameter(torch.tensor(0.125)) + + def forward(self, state: torch.Tensor, time: torch.Tensor) -> torch.Tensor: + return self.scale * state + self.bias + 0.1 * time[:, None] + + +class ToyAdapter: + def __init__(self, grid_size: int, *, zero_student_gradient: bool = False) -> None: + self.grid_size = grid_size + self.zero_student_gradient = zero_student_gradient + + def student_all_heads( + self, + model: ToyStudent, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del time, condition, model_kwargs + raw = model(state) + if self.zero_student_gradient: + raw = raw * 0.0 + return raw.reshape(state.shape[0], self.grid_size, state.shape[1]) + + def student_fused_block( + self, + model: ToyStudent, + 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) + with projection.fuse_block(start, end, grid): + return model(state) + + def teacher_velocity( + self, + model: ToyTeacher, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + negative_condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del condition, negative_condition, model_kwargs + return model(state, time) + + +@dataclass +class ToyLifecycle: + config: PDDConfig + student: ToyStudent + teacher: ToyTeacher + projection: PDDOutputProjection + pipeline: PDDPipeline + optimizer: torch.optim.AdamW + scheduler: torch.optim.lr_scheduler.LambdaLR + trainer: PDDTrainer + metadata: PDDMetadata + + +def build_toy_lifecycle( + *, + seed: int = 17, + zero_student_gradient: bool = False, + weight_decay: float = 0.01, +) -> ToyLifecycle: + torch.manual_seed(seed) + config = PDDConfig( + grid_size=4, + flow_shift=5.0, + block_size_min=1, + block_size_max=4, + inference_blocks=[2, 2], + student_sample_steps=2, + guidance_scale=None, + ) + student = ToyStudent() + projection = convert_to_pdd_output_projection( + student, + PDDLayerSpec("projection", "channel_major"), + config.grid_size, + ) + teacher = ToyTeacher() + pipeline = PDDPipeline( + student, + teacher, + config, + ToyAdapter(config.grid_size, zero_student_gradient=zero_student_gradient), + ) + optimizer = torch.optim.AdamW( + student.parameters(), + lr=2e-3, + weight_decay=weight_decay, + amsgrad=False, + capturable=False, + differentiable=False, + foreach=False, + fused=False, + maximize=False, + ) + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda _: 1.0) + trainer = PDDTrainer( + pipeline, + optimizer, + projection=projection, + max_grad_norm=0.5, + ) + return ToyLifecycle( + config, + student, + teacher, + projection, + pipeline, + optimizer, + scheduler, + trainer, + PDDMetadata.from_config(config, projection), + ) + + +def make_batch(sample_ids: tuple[str, ...], *, offset: float = 0.0) -> PreparedPDDBatch: + data = torch.stack( + [ + torch.tensor([0.5 + offset + index / 10, -1.0, 0.25], dtype=torch.float32) + for index in range(len(sample_ids)) + ] + ) + condition = ( + torch.zeros((len(sample_ids), 1, 1), dtype=torch.float32), + torch.ones((len(sample_ids), 1), dtype=torch.long), + ) + return PreparedPDDBatch(data, condition, None, sample_ids) + + +class SamplerDataset: + def __init__(self, sample_ids: tuple[str, ...]) -> None: + self.metadata = [ + { + "sample_id": sample_id, + "bucket_id": "64x64", + "bucket_resolution": [64, 64], + } + for sample_id in sample_ids + ] + self.bucket_groups = { + (64, 64): { + "indices": list(range(len(sample_ids))), + "resolution": (64, 64), + "aspect_name": "square", + } + } + self.sorted_bucket_keys = [(64, 64)] + self.calculator = None + + def __len__(self) -> int: + return len(self.metadata) + + def __getitem__(self, index: int) -> int: + return index + + +def ordered_id_sha256(sample_ids: tuple[str, ...]) -> str: + digest = hashlib.sha256() + digest.update(b"modelopt-pdd-ordered-train-ids-v1\0") + for sample_id in sample_ids: + digest.update(sample_id.encode()) + digest.update(b"\n") + return digest.hexdigest() diff --git a/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py b/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py new file mode 100644 index 00000000000..e5a9a20b006 --- /dev/null +++ b/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Two-rank proof that training-input failures propagate before the model call.""" + +from __future__ import annotations + +import pathlib +import sys + +import torch +import torch.distributed as dist + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +from pdd_finetune import _collective_training_batch, _collective_training_iterator + + +class _Sampler: + def __init__(self, sample_ids: tuple[str, ...]) -> None: + self.sample_ids = sample_ids + self.epoch = 0 + self.remaining_batches = 1 + + def expected_next_sample_ids(self) -> tuple[str, ...]: + return self.sample_ids + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + self.remaining_batches = 1 + + +class _Loader: + def __init__(self, batch: dict, *, fail: bool) -> None: + self.batch = batch + self.fail = fail + + def __iter__(self): + if self.fail: + raise OSError("injected iterator construction failure") + return iter([self.batch]) + + +def _batch(sample_id: str) -> dict: + return { + "image_latents": torch.ones(1, 3, 4, 4), + "text_embeddings": torch.ones(1, 5, 6), + "text_embeddings_mask": torch.ones(1, 5, dtype=torch.bool), + "metadata": {"sample_ids": [sample_id]}, + } + + +def _expect_collective_failure(iterator, sampler: _Sampler, expected: str) -> None: + message = None + try: + _collective_training_batch( + iterator, + sampler=sampler, + resume=None, + resume_pending=False, + device=torch.device("cpu"), + dtype=torch.float32, + require_negative_condition=False, + expected_batch_size=1, + ) + except RuntimeError as error: + message = str(error) + messages: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(messages, message) + assert all(item is not None and expected in item for item in messages) + + +def main() -> None: + dist.init_process_group("gloo") + try: + rank = dist.get_rank() + sample_id = f"sample-rank-{rank}" + expected_ids = ("wrong-rank-0",) if rank == 0 else (sample_id,) + _expect_collective_failure( + iter([_batch(sample_id)]), + _Sampler(expected_ids), + "committed cursor", + ) + dist.barrier() + + malformed = _batch(sample_id) + if rank == 1: + malformed.pop("text_embeddings") + _expect_collective_failure( + iter([malformed]), + _Sampler((sample_id,)), + "missing required keys", + ) + dist.barrier() + + iterator_message = None + try: + _collective_training_iterator( + _Loader(_batch(sample_id), fail=rank == 0), + _Sampler((sample_id,)), + ) + except RuntimeError as error: + iterator_message = str(error) + iterator_messages: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(iterator_messages, iterator_message) + assert all( + item is not None and "iterator construction" in item for item in iterator_messages + ) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py new file mode 100644 index 00000000000..d1fda52c11c --- /dev/null +++ b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Two-rank CPU/Gloo equivalence harness for the deterministic PDD validation oracle.""" + +from __future__ import annotations + +import dataclasses +import pathlib +import sys + +import torch +import torch.distributed as dist + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) +if str(pathlib.Path(__file__).parent) not in sys.path: + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + +from pdd_test_utils import build_toy_lifecycle, make_batch +from pdd_training import build_pdd_validation_assignments, run_pdd_validation + + +def _expect_failure(error_type, callback) -> None: + try: + callback() + except error_type: + return + raise AssertionError(f"expected {error_type.__name__}") + + +def main() -> None: + lifecycle = build_toy_lifecycle() + sample_ids = tuple(f"distributed-validation-{index:02d}" for index in range(13)) + assignments = build_pdd_validation_assignments( + sample_ids, + lifecycle.config, + validation_seed=91, + require_full_coverage=False, + ) + all_batches = [ + make_batch((assignment.sample_id,), offset=assignment.ordinal / 100) + for assignment in assignments + ] + baseline = run_pdd_validation( + lifecycle.pipeline, + all_batches, + assignments, + validation_seed=91, + ) + + dist.init_process_group(backend="gloo") + try: + rank = dist.get_rank() + world_size = dist.get_world_size() + assert world_size == 2 + local_batches = all_batches[rank::world_size] + padded_batch_count = max( + len(all_batches[candidate_rank::world_size]) for candidate_rank in range(world_size) + ) + while len(local_batches) < padded_batch_count: + local_batches.append( + dataclasses.replace( + make_batch((f"dummy-rank-{rank}",)), + valid_mask=(False,), + ) + ) + distributed = run_pdd_validation( + lifecycle.pipeline, + local_batches, + assignments, + validation_seed=91, + ) + assert distributed.records == baseline.records + assert abs(distributed.mean_loss - baseline.mean_loss) <= 1e-12 + assert distributed.ordered_id_sha256 == baseline.ordered_id_sha256 + + invalid_mask = list(local_batches) + if rank == 0: + invalid_mask[0] = dataclasses.replace(invalid_mask[0], valid_mask=()) + _expect_failure( + RuntimeError, + lambda: run_pdd_validation( + lifecycle.pipeline, + invalid_mask, + assignments, + validation_seed=91, + ), + ) + dist.barrier() + + unassigned = list(local_batches) + if rank == 0: + unassigned[0] = dataclasses.replace( + unassigned[0], + sample_ids=("not-in-heldout-assignments",), + valid_mask=(True,), + ) + _expect_failure( + RuntimeError, + lambda: run_pdd_validation( + lifecycle.pipeline, + unassigned, + assignments, + validation_seed=91, + ), + ) + dist.barrier() + + nonfinite = list(local_batches) + if rank == 0: + nonfinite[0] = dataclasses.replace( + nonfinite[0], + data=torch.full_like(nonfinite[0].data, float("nan")), + ) + _expect_failure( + FloatingPointError, + lambda: run_pdd_validation( + lifecycle.pipeline, + nonfinite, + assignments, + validation_seed=91, + ), + ) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py b/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py index 699f806f7e2..9804faafe1e 100644 --- a/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py +++ b/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py @@ -130,8 +130,11 @@ def test_migration_is_path_independent_and_relocatable(tmp_path): alice_train = load_portable_metadata(alice_output, "metadata_train.json")[0]["sample_ids"] bob_train = load_portable_metadata(bob_output, "metadata_train.json")[0]["sample_ids"] assert alice_train == bob_train - assert validate_snapshot(alice_output)["splits"] == {"all": 4, "train": 3, "heldout": 1} - assert validate_snapshot(bob_output)["splits"] == {"all": 4, "train": 3, "heldout": 1} + alice_report = validate_snapshot(alice_output) + bob_report = validate_snapshot(bob_output) + assert alice_report["splits"] == {"all": 4, "train": 3, "heldout": 1} + assert bob_report["splits"] == {"all": 4, "train": 3, "heldout": 1} + assert alice_report["snapshot_sha256"] == bob_report["snapshot_sha256"] cli = subprocess.run( [ @@ -162,7 +165,7 @@ def test_migration_is_path_independent_and_relocatable(tmp_path): relocated.parent.mkdir() shutil.copytree(alice_output, relocated) assert _manifest_signature(relocated) == _manifest_signature(alice_output) - validate_snapshot(relocated) + assert validate_snapshot(relocated)["snapshot_sha256"] == alice_report["snapshot_sha256"] def test_incomplete_pass_one_publishes_nothing(monkeypatch, tmp_path): diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index d109a4f1eff..e45da3804c8 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -103,6 +103,54 @@ def test_remote_model_requires_full_revision_and_non_dp_parallelism_is_rejected( resolve_pdd_recipe_config(raw) +@pytest.mark.parametrize( + ("section", "name", "value", "message"), + [ + ("training", "grad_accumulation_steps", 2, "grad_accumulation_steps=1"), + ("training", "max_grad_norm", 0.0, "max_grad_norm must be > 0"), + ("training", "validation_every_steps", 0, "validation_every_steps"), + ("guidance", "rescale", 1.1, "guidance.rescale must be <= 1"), + ("optim", "betas", [0.9, 1.0], "optim.betas values"), + ("optim", "eps", 0.0, "optim.eps must be > 0"), + ], +) +def test_training_config_gates_fail_during_resolution( + tmp_path, section, name, value, message +) -> None: + raw = _raw_config(tmp_path) + raw.setdefault(section, {})[name] = value + + with pytest.raises(ValueError, match=message): + resolve_pdd_recipe_config(raw) + + +def test_restore_requires_enabled_checkpointing(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["checkpoint"]["enabled"] = False + raw["checkpoint"]["restore_from"] = "LATEST" + + with pytest.raises(ValueError, match="restore_from requires checkpoint.enabled=true"): + resolve_pdd_recipe_config(raw) + + +@pytest.mark.parametrize( + ("name", "value", "message"), + [ + ("drop_last", False, "drop_last=true"), + ("dynamic_batch_size", True, "dynamic_batch_size=false"), + ("train_text_encoder", True, "cached text embeddings"), + ], +) +def test_training_dataloader_modes_are_gated_during_resolution( + tmp_path, name, value, message +) -> None: + raw = _raw_config(tmp_path) + raw["data"] = {"dataloader": {name: value}} + + with pytest.raises(ValueError, match=message): + resolve_pdd_recipe_config(raw) + + def test_frozen_automodel_distribution_snapshot_is_stable() -> None: try: version = importlib.metadata.version("nemo_automodel") diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py new file mode 100644 index 00000000000..d8f487d4599 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -0,0 +1,469 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hermetic direct-update, committed-cursor, and strict PDD resume evidence.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import math +import pathlib +import shutil +import sys + +import pytest +import torch + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) +if str(pathlib.Path(__file__).parent) not in sys.path: + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + +from fastgen_data.replayable_sampler import ReplayableBatchSampler +from pdd_checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity +from pdd_recipe import initialize_pdd_distributed +from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 +from pdd_training import prepare_qwen_pdd_batch +from verify_readonly_automodel import snapshot_installed_distribution + + +def _released_sampler(sample_ids: tuple[str, ...]) -> ReplayableBatchSampler: + sampler_module = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") + dataset = SamplerDataset(sample_ids) + sampler = sampler_module.SequentialBucketSampler( + dataset, + base_batch_size=1, + base_resolution=(64, 64), + drop_last=True, + shuffle_buckets=True, + shuffle_within_bucket=True, + dynamic_batch_size=False, + seed=31, + num_replicas=1, + rank=0, + ) + return ReplayableBatchSampler(sampler) + + +def _run_next(lifecycle, sampler): + sample_ids = sampler.expected_next_sample_ids() + assert sample_ids + diagnostics = lifecycle.trainer.train_step(make_batch(sample_ids)) + lifecycle.scheduler.step() + sampler.commit(sample_ids) + if sampler.remaining_batches == 0: + sampler.set_epoch(sampler.epoch + 1) + return sample_ids, diagnostics + + +def _checkpointer(lifecycle, checkpoint_dir): + checkpoint_module = pytest.importorskip("nemo_automodel.components.checkpoint.config") + config = checkpoint_module.CheckpointingConfig( + enabled=True, + checkpoint_dir=str(checkpoint_dir), + model_save_format="torch_save", + model_repo_id="synthetic-pdd-toy", + save_consolidated=False, + is_peft=False, + model_state_dict_keys=list(lifecycle.student.state_dict()), + ) + return config.build(dp_rank=0, tp_rank=0, pp_rank=0, moe_mesh=None) + + +def _identity(lifecycle, scheduler, sample_ids): + return build_pdd_checkpoint_identity( + metadata=lifecycle.metadata, + model_id="synthetic-pdd-toy", + model_revision=None, + guidance_scale=None, + guidance_rescale=1.0, + guidance_eps=1e-5, + automodel_snapshot=snapshot_installed_distribution(), + ordered_train_id_sha256=ordered_id_sha256(sample_ids), + ordered_heldout_id_sha256="1" * 64, + dataset_snapshot_sha256="2" * 64, + local_batch_size=1, + grad_accumulation_steps=1, + training_seed=1234, + validation_seed=2026, + validation_every_steps=100, + max_grad_norm=0.5, + zero_grad_warmup_steps=0, + activation_checkpointing=False, + dtype="float32", + optimizer=lifecycle.optimizer, + scheduler=scheduler, + ) + + +def _manager(root, lifecycle, sampler, rng): + checkpointer = _checkpointer(lifecycle, root) + manager = PDDCheckpointManager( + root=root, + checkpointer=checkpointer, + model=lifecycle.student, + optimizer=lifecycle.optimizer, + scheduler=lifecycle.scheduler, + trainer=lifecycle.trainer, + sampler=sampler, + rng=rng, + identity=_identity(lifecycle, lifecycle.scheduler, tuple(f"sample-{i}" for i in range(8))), + ) + return manager, checkpointer + + +def _optimizer_state_by_name(lifecycle): + names = {parameter: name for name, parameter in lifecycle.student.named_parameters()} + return { + names[parameter]: { + key: value.detach().clone() if isinstance(value, torch.Tensor) else value + for key, value in state.items() + } + for parameter, state in lifecycle.optimizer.state.items() + } + + +def _refresh_complete_marker(checkpoint: pathlib.Path) -> None: + manifest_path = checkpoint / "manifest.json" + marker = { + "schema_version": 1, + "manifest_sha256": hashlib.sha256(manifest_path.read_bytes()).hexdigest(), + } + (checkpoint / "COMPLETE").write_text(json.dumps(marker, indent=2, sort_keys=True) + "\n") + + +def test_replayable_sampler_commits_consumed_batches_not_prefetch() -> None: + sample_ids = tuple(f"sample-{index}" for index in range(8)) + sampler = _released_sampler(sample_ids) + first_expected = sampler.expected_next_sample_ids() + iterator = iter(sampler) + next(iterator) + next(iterator) + + assert sampler.committed_batches == 0 + assert sampler.expected_next_sample_ids() == first_expected + sampler.commit(first_expected) + state = sampler.state_dict() + + restored = _released_sampler(sample_ids) + restored.load_state_dict(state) + assert restored.state_dict() == state + next_indices = next(iter(restored)) + assert tuple(restored.dataset.metadata[index]["sample_id"] for index in next_indices) == tuple( + state["next_sample_ids"] + ) + bad_hash = dict(state, plan_sha256="0" * 64) + with pytest.raises(RuntimeError, match="plan hash"): + _released_sampler(sample_ids).load_state_dict(bad_hash) + bad_ids = dict(state, next_sample_ids=["wrong-id"]) + with pytest.raises(RuntimeError, match="next sample IDs"): + _released_sampler(sample_ids).load_state_dict(bad_ids) + + +def test_qwen_batch_preparation_preserves_ids_masks_and_negative_condition() -> None: + batch = { + "image_latents": torch.ones(2, 3, 4, 4), + "text_embeddings": torch.ones(2, 5, 6), + "text_embeddings_mask": torch.ones(2, 5, dtype=torch.bool), + "negative_text_embeddings": torch.zeros(5, 6), + "negative_text_embeddings_mask": torch.ones(5, dtype=torch.bool), + "metadata": {"sample_ids": ["qwen-a", "qwen-b"]}, + } + + prepared = prepare_qwen_pdd_batch( + batch, + device=torch.device("cpu"), + dtype=torch.float32, + require_negative_condition=True, + ) + + assert prepared.sample_ids == ("qwen-a", "qwen-b") + assert prepared.valid_mask == (True, True) + assert prepared.data.shape == (2, 3, 4, 4) + assert prepared.condition[0].shape == (2, 5, 6) + assert prepared.negative_condition is not None + assert prepared.negative_condition[0].shape == (2, 5, 6) + + without_negative = dict(batch) + without_negative.pop("negative_text_embeddings") + without_negative.pop("negative_text_embeddings_mask") + with pytest.raises(ValueError, match="requires negative prompt conditioning"): + prepare_qwen_pdd_batch( + without_negative, + device=torch.device("cpu"), + dtype=torch.float32, + require_negative_condition=True, + ) + + +def test_two_direct_updates_have_finite_gradients_updates_and_targeted_coverage() -> None: + lifecycle = build_toy_lifecycle(weight_decay=0.02) + first_batch = make_batch(("a", "b")) + before = [parameter.detach().clone() for parameter in lifecycle.student.parameters()] + first = lifecycle.trainer.train_step( + first_batch, + noise=torch.tensor([[0.25, -0.5, 1.0], [-0.25, 0.5, -1.0]]), + n=torch.tensor([0, 1]), + k=torch.tensor([1, 3]), + ) + lifecycle.scheduler.step() + actual_update = math.sqrt( + sum( + (parameter.detach() - saved).double().square().sum().item() + for parameter, saved in zip(lifecycle.student.parameters(), before) + ) + ) + before_norm = math.sqrt(sum(saved.double().square().sum().item() for saved in before)) + + assert math.isfinite(first.loss) and first.loss > 0 + assert math.isfinite(first.grad_norm) and first.grad_norm > 0 + assert first.pdd_projection_update_ratio is not None + assert first.pdd_projection_update_ratio > 0 + assert math.isfinite(first.student_teacher_velocity_rms_ratio) + assert first.student_teacher_velocity_rms_ratio >= 0 + assert first.student_adamw_nominal_update_ratio == pytest.approx( + actual_update / before_norm, + rel=2e-5, + abs=1e-8, + ) + assert all(parameter.grad is None for parameter in lifecycle.teacher.parameters()) + + second = lifecycle.trainer.train_step( + make_batch(("c", "d"), offset=0.25), + noise=torch.tensor([[0.75, 0.0, -0.5], [-0.75, 0.0, 0.5]]), + n=torch.tensor([2, 3]), + k=torch.tensor([2, 3]), + ) + lifecycle.scheduler.step() + assert second.completed_step == 2 + lifecycle.trainer.coverage.require_pairs([(0, 1), (1, 3), (2, 2), (3, 3)]) + + +def test_training_hard_aborts_for_teacher_gradient_zero_gradient_and_missing_coverage() -> None: + teacher_gradient = build_toy_lifecycle() + teacher_gradient.teacher.scale.grad = torch.ones_like(teacher_gradient.teacher.scale) + with pytest.raises(RuntimeError, match="teacher received a gradient"): + teacher_gradient.trainer.train_step(make_batch(("teacher-grad",))) + + zero_gradient = build_toy_lifecycle(zero_student_gradient=True, weight_decay=0.0) + first = zero_gradient.trainer.train_step(make_batch(("zero-1",))) + assert first.grad_norm == 0.0 + with pytest.raises(RuntimeError, match="zero for two consecutive"): + zero_gradient.trainer.train_step(make_batch(("zero-2",))) + + with pytest.raises(RuntimeError, match="did not cover"): + zero_gradient.trainer.coverage.require_pairs([(3, 3)]) + + nonfinite = build_toy_lifecycle() + bad_batch = make_batch(("nan",)) + bad_batch.data.fill_(float("nan")) + with pytest.raises(FloatingPointError, match="non-finite"): + nonfinite.trainer.train_step(bad_batch) + + nonfinite_gradient = build_toy_lifecycle() + handle = nonfinite_gradient.projection.weight.register_hook( + lambda gradient: torch.full_like(gradient, float("inf")) + ) + with pytest.raises(FloatingPointError, match="gradient became non-finite"): + nonfinite_gradient.trainer.train_step(make_batch(("inf-gradient",))) + handle.remove() + + invalid_support = build_toy_lifecycle() + with pytest.raises(RuntimeError, match="k must satisfy"): + invalid_support.trainer.train_step( + make_batch(("invalid-support",)), + n=torch.tensor([2]), + k=torch.tensor([1]), + ) + + nonfinite_update = build_toy_lifecycle() + nonfinite_update.optimizer.param_groups[0]["eps"] = float("nan") + with pytest.raises(FloatingPointError, match="parameter update became non-finite"): + nonfinite_update.trainer.train_step( + make_batch(("inf-update",)), + measure_updates=False, + ) + + +def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) -> None: + pytest.importorskip("nemo_automodel") + rng_module = pytest.importorskip("nemo_automodel.components.training.rng") + if not torch.distributed.is_initialized(): + initialize_pdd_distributed(backend="gloo", timeout_minutes=1) + sample_ids = tuple(f"sample-{index}" for index in range(8)) + + source = build_toy_lifecycle() + source_sampler = _released_sampler(sample_ids) + source_rng = rng_module.StatefulRNG(1234, ranked=True) + source_manager, source_checkpointer = _manager( + tmp_path / "checkpoints", source, source_sampler, source_rng + ) + _run_next(source, source_sampler) + _run_next(source, source_sampler) + checkpoint = source_manager.save() + assert checkpoint.name == "step_00000002" + assert source_manager.identity["data"]["dataset_snapshot_sha256"] == "2" * 64 + assert source_manager.identity["training"] == { + "seed": 1234, + "validation_seed": 2026, + "validation_every_steps": 100, + "max_grad_norm": 0.5, + "zero_grad_warmup_steps": 0, + "activation_checkpointing": False, + } + expected_next_ids, expected_next = _run_next(source, source_sampler) + expected_model = copy.deepcopy(source.student.state_dict()) + expected_optimizer = _optimizer_state_by_name(source) + + destination = build_toy_lifecycle() + destination_sampler = _released_sampler(sample_ids) + destination_rng = rng_module.StatefulRNG(9999, ranked=True) + destination_manager, destination_checkpointer = _manager( + tmp_path / "checkpoints", + destination, + destination_sampler, + destination_rng, + ) + resume = destination_manager.load("LATEST") + assert resume is not None + assert resume.completed_steps == 2 + assert resume.sample_slots_consumed == 2 + assert resume.expected_next_sample_ids == expected_next_ids + resume.verify_first_batch(destination_sampler.expected_next_sample_ids()) + with pytest.raises(RuntimeError, match="first resumed sample IDs"): + resume.verify_first_batch(("wrong-id",)) + actual_ids, actual_next = _run_next(destination, destination_sampler) + + assert actual_ids == expected_next_ids + assert actual_next.n == expected_next.n + assert actual_next.k == expected_next.k + assert actual_next.loss == expected_next.loss + assert actual_next.learning_rate == expected_next.learning_rate + for name, tensor in destination.student.state_dict().items(): + torch.testing.assert_close(tensor, expected_model[name], rtol=0, atol=0) + actual_optimizer = _optimizer_state_by_name(destination) + assert actual_optimizer.keys() == expected_optimizer.keys() + for name in actual_optimizer: + for key, actual in actual_optimizer[name].items(): + expected = expected_optimizer[name][key] + if isinstance(actual, torch.Tensor): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + assert actual == expected + + resumed_checkpoint = destination_manager.save() + assert resumed_checkpoint.name == "step_00000003" + third = build_toy_lifecycle() + third_sampler = _released_sampler(sample_ids) + third_rng = rng_module.StatefulRNG(7, ranked=True) + third_manager, third_checkpointer = _manager( + tmp_path / "checkpoints", third, third_sampler, third_rng + ) + second_resume = third_manager.load("LATEST") + assert second_resume is not None + assert second_resume.completed_steps == 3 + assert second_resume.expected_next_sample_ids == destination_sampler.expected_next_sample_ids() + + inventory = {path.name.lower() for path in resumed_checkpoint.rglob("*")} + assert not any( + token in name + for name in inventory + for token in ("fake_score", "discriminator", "ema", "r1", "gan") + ) + + incomplete = tmp_path / "checkpoints" / "step_99999998" + incomplete.mkdir() + (tmp_path / "checkpoints" / "LATEST").write_text(incomplete.name + "\n") + assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() + with pytest.raises(RuntimeError, match="incomplete"): + third_manager.resolve(incomplete.name) + + mismatched = tmp_path / "checkpoints" / "step_99999999" + shutil.copytree(resumed_checkpoint, mismatched) + manifest_path = mismatched / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + manifest["identity"]["model"]["id"] = "different-model" + manifest["completed_steps"] = 99999999 + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + _refresh_complete_marker(mismatched) + (tmp_path / "checkpoints" / "LATEST").write_text(mismatched.name + "\n") + assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() + with pytest.raises(RuntimeError, match="identity"): + third_manager.resolve(mismatched.name) + + missing_model = tmp_path / "checkpoints" / "step_99999996" + shutil.copytree(resumed_checkpoint, missing_model) + model_payload = next( + path for path in (missing_model / "model").iterdir() if path.name != ".metadata" + ) + model_payload.unlink() + (tmp_path / "checkpoints" / "LATEST").write_text(missing_model.name + "\n") + assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() + with pytest.raises(RuntimeError, match="DCP"): + third_manager.resolve(missing_model.name) + + corrupt_optimizer = tmp_path / "checkpoints" / "step_99999997" + shutil.copytree(resumed_checkpoint, corrupt_optimizer) + optimizer_payload = next( + path for path in (corrupt_optimizer / "optim").iterdir() if path.name != ".metadata" + ) + with optimizer_payload.open("ab") as stream: + stream.write(b"corrupt") + (tmp_path / "checkpoints" / "LATEST").write_text(corrupt_optimizer.name + "\n") + assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() + with pytest.raises(RuntimeError, match="DCP"): + third_manager.resolve(corrupt_optimizer.name) + + step_mismatch = tmp_path / "checkpoints" / "step_00000004" + shutil.copytree(resumed_checkpoint, step_mismatch) + step_manifest_path = step_mismatch / "manifest.json" + step_manifest = json.loads(step_manifest_path.read_text()) + step_manifest["completed_steps"] = 4 + step_manifest_path.write_text(json.dumps(step_manifest, indent=2, sort_keys=True) + "\n") + trainer_state_path = step_mismatch / "trainer_state.json" + trainer_state = json.loads(trainer_state_path.read_text()) + trainer_state["completed_steps"] = 4 + trainer_state_path.write_text(json.dumps(trainer_state, indent=2, sort_keys=True) + "\n") + _refresh_complete_marker(step_mismatch) + with pytest.raises(RuntimeError, match="trainer step"): + third_manager.load(step_mismatch.name) + + lr_mismatch = tmp_path / "checkpoints" / "step_00000005" + shutil.copytree(resumed_checkpoint, lr_mismatch) + lr_manifest_path = lr_mismatch / "manifest.json" + lr_manifest = json.loads(lr_manifest_path.read_text()) + lr_manifest["learning_rates"] = [0.123] + lr_manifest_path.write_text(json.dumps(lr_manifest, indent=2, sort_keys=True) + "\n") + lr_trainer_path = lr_mismatch / "trainer_state.json" + lr_trainer = json.loads(lr_trainer_path.read_text()) + lr_trainer["learning_rates"] = [0.123] + lr_trainer_path.write_text(json.dumps(lr_trainer, indent=2, sort_keys=True) + "\n") + _refresh_complete_marker(lr_mismatch) + with pytest.raises(RuntimeError, match="learning rate"): + third_manager.load(lr_mismatch.name) + + cursor_mismatch = tmp_path / "checkpoints" / "step_00000006" + shutil.copytree(resumed_checkpoint, cursor_mismatch) + sampler_path = cursor_mismatch / "sampler" / "sampler_dp_rank_0.pt" + sampler_state = torch.load(sampler_path, weights_only=False) + sampler_state["plan_sha256"] = "0" * 64 + torch.save(sampler_state, sampler_path) + cursor_manifest_path = cursor_mismatch / "manifest.json" + cursor_manifest = json.loads(cursor_manifest_path.read_text()) + cursor_manifest["sidecar_sha256"]["sampler/sampler_dp_rank_0.pt"] = hashlib.sha256( + sampler_path.read_bytes() + ).hexdigest() + cursor_manifest_path.write_text(json.dumps(cursor_manifest, indent=2, sort_keys=True) + "\n") + _refresh_complete_marker(cursor_mismatch) + with pytest.raises(RuntimeError, match="plan hash"): + third_manager.load(cursor_mismatch.name) + + source_checkpointer.close() + destination_checkpointer.close() + third_checkpointer.close() diff --git a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py new file mode 100644 index 00000000000..07fe2eefabe --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic logical-ID PDD held-out oracle tests.""" + +from __future__ import annotations + +import pathlib +import sys + +import pytest +import torch + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) +if str(pathlib.Path(__file__).parent) not in sys.path: + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + +from pdd_test_utils import build_toy_lifecycle, make_batch +from pdd_training import ( + build_pdd_validation_assignments, + pdd_validation_noise, + pdd_validation_support, + run_pdd_validation, +) + +from modelopt.torch.fastgen import PDDConfig + + +def test_canonical_2k_assignment_covers_all_1568_pairs_32_starts_and_128_heads() -> None: + config = PDDConfig( + grid_size=128, + flow_shift=5.0, + block_size_min=4, + block_size_max=64, + inference_blocks=[32, 32, 32, 32], + student_sample_steps=4, + ) + sample_ids = [f"heldout-{index:04d}" for index in range(2000)] + assignments = build_pdd_validation_assignments( + list(reversed(sample_ids)), + config, + validation_seed=2026, + ) + + assert len(pdd_validation_support(config)) == 1568 + assert len(assignments) == 2000 + assert len({(assignment.n, assignment.k) for assignment in assignments}) == 1568 + assert len({assignment.n for assignment in assignments}) == 32 + assert len({assignment.k for assignment in assignments}) == 128 + assert assignments == build_pdd_validation_assignments( + sample_ids, + config, + validation_seed=2026, + ) + + +def test_assignment_rejects_duplicate_missing_coverage_and_noise_is_per_id_stable() -> None: + lifecycle = build_toy_lifecycle() + with pytest.raises(ValueError, match="unique"): + build_pdd_validation_assignments( + ["duplicate", "duplicate"], + lifecycle.config, + validation_seed=1, + require_full_coverage=False, + ) + with pytest.raises(ValueError, match="at least"): + build_pdd_validation_assignments( + ["too-small"], + lifecycle.config, + validation_seed=1, + ) + + before = torch.get_rng_state().clone() + first = pdd_validation_noise( + "logical-id", + (3,), + validation_seed=7, + device=torch.device("cpu"), + ) + second = pdd_validation_noise( + "logical-id", + (3,), + validation_seed=7, + device=torch.device("cpu"), + ) + different = pdd_validation_noise( + "different-id", + (3,), + validation_seed=7, + device=torch.device("cpu"), + ) + assert torch.equal(torch.get_rng_state(), before) + assert torch.equal(first, second) + assert not torch.equal(first, different) + + +def test_repeated_validation_is_exact_and_does_not_change_training_state_or_rng() -> None: + lifecycle = build_toy_lifecycle() + sample_ids = tuple(f"validation-{index:02d}" for index in range(12)) + assignments = build_pdd_validation_assignments( + sample_ids, + lifecycle.config, + validation_seed=44, + require_full_coverage=False, + ) + batches = [ + make_batch((assignment.sample_id,), offset=assignment.ordinal / 100) + for assignment in assignments + ] + parameter_before = { + name: parameter.detach().clone() for name, parameter in lifecycle.student.named_parameters() + } + optimizer_before = lifecycle.optimizer.state_dict() + scheduler_before = lifecycle.scheduler.state_dict() + rng_before = torch.get_rng_state().clone() + student_mode_before = lifecycle.student.training + teacher_mode_before = lifecycle.teacher.training + + first = run_pdd_validation( + lifecycle.pipeline, + (batch for batch in batches), + assignments, + validation_seed=44, + ) + second = run_pdd_validation( + lifecycle.pipeline, + list(reversed(batches)), + assignments, + validation_seed=44, + ) + + assert first == second + assert first.mean_loss == second.mean_loss + assert first.pair_count == len({(item.n, item.k) for item in assignments}) + assert torch.equal(torch.get_rng_state(), rng_before) + assert lifecycle.student.training is student_mode_before + assert lifecycle.teacher.training is teacher_mode_before + assert lifecycle.optimizer.state_dict() == optimizer_before + assert lifecycle.scheduler.state_dict() == scheduler_before + for name, parameter in lifecycle.student.named_parameters(): + torch.testing.assert_close(parameter, parameter_before[name], rtol=0, atol=0) diff --git a/tests/examples/diffusers/fastgen/test_portable_cache.py b/tests/examples/diffusers/fastgen/test_portable_cache.py index 487cd55ac76..01552a71e1f 100644 --- a/tests/examples/diffusers/fastgen/test_portable_cache.py +++ b/tests/examples/diffusers/fastgen/test_portable_cache.py @@ -20,7 +20,11 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -from fastgen_data import TextToImageDataset, collate_fn_text_to_image +from fastgen_data import ( + TextToImageDataset, + build_text_to_image_multiresolution_dataloader, + collate_fn_text_to_image, +) from portable_cache import ( DATASET_CACHE_ENV, audit_no_absolute_paths, @@ -156,6 +160,10 @@ def test_relocation_preserves_order_payloads_and_buckets(monkeypatch, tmp_path): assert _batch_signature(first_dataset) == _batch_signature(second_dataset) assert [entry["sample_id"] for entry in second_dataset.metadata] == splits["train"] assert first_dataset.bucket_groups == second_dataset.bucket_groups + first_report = validate_snapshot(first) + second_report = validate_snapshot(second) + assert first_report["snapshot_sha256"] == second_report["snapshot_sha256"] + assert first_report["declared_files"] == second_report["declared_files"] monkeypatch.setenv(DATASET_CACHE_ENV, str(second)) overridden = TextToImageDataset(str(first), metadata_index="metadata_train.json") @@ -163,6 +171,27 @@ def test_relocation_preserves_order_payloads_and_buckets(monkeypatch, tmp_path): assert _batch_signature(overridden) == _batch_signature(first_dataset) +def test_pdd_loader_iterator_does_not_advance_training_rng(monkeypatch, tmp_path): + root = tmp_path / "cache" + _make_snapshot(root) + monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) + loader, _ = build_text_to_image_multiresolution_dataloader( + cache_dir=str(root), + metadata_index="metadata_train.json", + batch_size=1, + base_resolution=(64, 64), + num_workers=0, + exact_resume=True, + sampler_seed=17, + loader_seed=17, + ) + before = torch.get_rng_state().clone() + + next(iter(loader)) + + assert torch.equal(torch.get_rng_state(), before) + + def test_split_filters_before_inherited_bucket_grouping(monkeypatch, tmp_path): root = tmp_path / "cache" splits = _make_snapshot(root) From a05777f45af89fb58a046b7596eb5c43bb644c14 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 10:26:24 -0700 Subject: [PATCH 12/45] feat(fastgen): add PDD Qwen inference example Signed-off-by: Meng Xin --- CHANGELOG.rst | 1 + .../diffusers/fastgen/analyze_pdd_results.py | 35 + .../fastgen/export_pdd_qwen_image.py | 380 ++++++ .../fastgen/inference_pdd_qwen_image.py | 320 +++++ examples/diffusers/fastgen/pdd_artifacts.py | 153 +++ examples/diffusers/fastgen/pdd_checkpoint.py | 272 ++++- examples/diffusers/fastgen/pdd_evaluation.py | 1088 +++++++++++++++++ examples/diffusers/fastgen/pdd_export.py | 607 +++++++++ examples/diffusers/fastgen/pdd_recipe.py | 191 ++- .../fastgen/seal_pdd_run_manifest.py | 45 + .../fastgen/validate_pdd_run_manifest.py | 32 + .../fastgen/pdd_export_distributed.py | 181 +++ .../diffusers/fastgen/test_pdd_evaluation.py | 500 ++++++++ .../fastgen/test_pdd_inference_checkpoint.py | 298 +++++ .../fastgen/test_pdd_recipe_setup.py | 29 +- .../fastgen/test_pdd_training_lifecycle.py | 14 +- 16 files changed, 4046 insertions(+), 100 deletions(-) create mode 100644 examples/diffusers/fastgen/analyze_pdd_results.py create mode 100644 examples/diffusers/fastgen/export_pdd_qwen_image.py create mode 100644 examples/diffusers/fastgen/inference_pdd_qwen_image.py create mode 100644 examples/diffusers/fastgen/pdd_artifacts.py create mode 100644 examples/diffusers/fastgen/pdd_evaluation.py create mode 100644 examples/diffusers/fastgen/pdd_export.py create mode 100644 examples/diffusers/fastgen/seal_pdd_run_manifest.py create mode 100644 examples/diffusers/fastgen/validate_pdd_run_manifest.py create mode 100644 tests/examples/diffusers/fastgen/pdd_export_distributed.py create mode 100644 tests/examples/diffusers/fastgen/test_pdd_evaluation.py create mode 100644 tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..a4421b1fd5a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -102,6 +102,7 @@ Changelog - Add offline DFlash speculative decoding training. Train the draft module from pre-computed base-model hidden states dumped by ``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py``; base-model transformer layers are deleted after conversion to save memory. Controlled by the auto-derived ``dflash_offline`` flag on ``DFlashConfig`` (derived from ``data_args.offline_data_path``). The dump scripts now share ``collect_hidden_states/common.py`` for aux-layer selection (``--aux-layers eagle|dflash|``) and optional assistant-token ``loss_mask`` for answer-only-loss training. - Add ``mtsa.config.SKIP_SOFTMAX_TRITON_CALIB`` for skip-softmax attention-sparsity calibration through the fused Triton ``attention_calibrate`` kernel (HF ``modelopt_triton`` backend), measuring multi-threshold tile-skip statistics the way the Triton inference kernel actually skips tiles for both prefill and decode. Exposed as ``--sparse_attn_cfg skip_softmax_triton_calib`` in ``examples/llm_sparsity/attention_sparsity/hf_sa.py`` (with a new ``--calib_data_dir`` flag for RULER calibration data). - Add DMD2 distillation for few-step diffusion models in ``examples/diffusers/fastgen/``: distill Qwen-Image into a 4/8-step student via Distribution Matching Distillation. See `examples/diffusers/fastgen/README.md `_ for details. +- Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with a Qwen-Image training, safe distributed-checkpoint export, PDD-2/4/8 inference, and paired effectiveness-evidence example. AutoModel remains an unmodified pinned runtime dependency. - Make ``.agents/skills/`` the canonical location for agent skills; agent-specific directories (``.claude/skills/``, etc.) are now relative symlinks into ``.agents/``, so one skill suite serves multiple coding agents (Claude Code, Codex). See ``.agents/README.md``. - Extend Claude Code agent skills for PTQ, deployment, evaluation, monitoring, and baseline-vs-quantized result comparison. Adds evaluation task references for additional benchmarks, stronger PTQ checkpoint validation gates, and session-scoped workspace/job tracking. - Add SLURM Quality of Service (QoS) support to the ModelOpt launcher. Users can set QoS via ``slurm_config.qos`` or ``SLURM_QOS`` and the value is forwarded to ``nemo_run.SlurmExecutor``. diff --git a/examples/diffusers/fastgen/analyze_pdd_results.py b/examples/diffusers/fastgen/analyze_pdd_results.py new file mode 100644 index 00000000000..56825e27443 --- /dev/null +++ b/examples/diffusers/fastgen/analyze_pdd_results.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Write deterministic paired summaries from claim-bearing PDD effectiveness evidence.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +sys.dont_write_bytecode = True + +_THIS_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _THIS_DIR.parents[2] +for path in (_REPO_ROOT, _THIS_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("manifest", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + from pdd_artifacts import write_canonical_json + from pdd_evaluation import summarize_effectiveness_bundle, validate_effectiveness_bundle + + validated = validate_effectiveness_bundle(args.manifest) + write_canonical_json(args.output, summarize_effectiveness_bundle(validated)) + print(args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/export_pdd_qwen_image.py b/examples/diffusers/fastgen/export_pdd_qwen_image.py new file mode 100644 index 00000000000..fc28e428fb6 --- /dev/null +++ b/examples/diffusers/fastgen/export_pdd_qwen_image.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Collectively restore a PDD checkpoint and publish a safe Qwen-Image export.""" + +from __future__ import annotations + +import argparse +import json +import math +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +import yaml + +sys.dont_write_bytecode = True + +_THIS_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _THIS_DIR.parents[2] +for path in (_REPO_ROOT, _THIS_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + type=Path, + default=_THIS_DIR / "configs" / "pdd_qwen_image.yaml", + ) + parser.add_argument( + "--checkpoint", + help="Checkpoint basename/path beneath checkpoint_dir, or LATEST; defaults to config.", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--max-shard-size-gib", type=float, default=5.0) + parser.add_argument("--memory-headroom", type=float, default=1.25) + return parser.parse_args() + + +def _read_integer(path: Path) -> int | None: + try: + value = path.read_text().strip() + except OSError: + return None + if value == "max": + return None + try: + return int(value) + except ValueError: + return None + + +def host_available_bytes() -> int: + """Return the strictest visible host/cgroup memory availability estimate.""" + candidates: list[int] = [] + try: + for line in Path("/proc/meminfo").read_text().splitlines(): + if line.startswith("MemAvailable:"): + candidates.append(int(line.split()[1]) * 1024) + break + except (OSError, ValueError, IndexError): + pass + for limit_path, used_path in ( + (Path("/sys/fs/cgroup/memory.max"), Path("/sys/fs/cgroup/memory.current")), + ( + Path("/sys/fs/cgroup/memory/memory.limit_in_bytes"), + Path("/sys/fs/cgroup/memory/memory.usage_in_bytes"), + ), + ): + limit = _read_integer(limit_path) + used = _read_integer(used_path) + if limit is not None and used is not None and 0 < limit < (1 << 62): + candidates.append(max(0, limit - used)) + if not candidates: + raise RuntimeError("cannot determine host memory availability.") + return min(candidates) + + +def _state_sizes(model: torch.nn.Module) -> tuple[int, int]: + sizes = [value.numel() * value.element_size() for value in model.state_dict().values()] + if not sizes: + raise RuntimeError("PDD export model has an empty state dictionary.") + return sum(sizes), max(sizes) + + +def collective_export_memory_preflight( + model: torch.nn.Module, + *, + max_shard_bytes: int, + headroom: float, + device: torch.device, +) -> tuple[int, int]: + """Abort collectively before full-state gathering when host/GPU headroom is insufficient.""" + if not math.isfinite(headroom) or headroom < 1.0: + raise ValueError("memory_headroom must be finite and >= 1.") + full_state_bytes, largest_tensor_bytes = _state_sizes(model) + local_error = None + try: + required_gpu = math.ceil(largest_tensor_bytes * headroom) + if device.type == "cuda": + free_gpu, _total_gpu = torch.cuda.mem_get_info(device) + if free_gpu < required_gpu: + raise MemoryError( + f"rank {dist.get_rank()} has {free_gpu} free GPU bytes; " + f"full-state gather requires at least {required_gpu}." + ) + if dist.get_rank() == 0: + required_host = math.ceil((full_state_bytes + max_shard_bytes) * headroom) + available_host = host_available_bytes() + if available_host < required_host: + raise MemoryError( + f"rank 0 has {available_host} available host bytes; export requires at " + f"least {required_host}." + ) + except BaseException as error: + local_error = f"{type(error).__name__}: {error}" + errors: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(errors, local_error) + failures = [f"rank {rank}: {error}" for rank, error in enumerate(errors) if error] + if failures: + raise RuntimeError("PDD export memory preflight failed; " + "; ".join(failures)) + return full_state_bytes, largest_tensor_bytes + + +def _git_source_identity() -> dict[str, Any]: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + dirty = bool( + subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=normal"], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout + ) + if dirty: + raise RuntimeError("PDD export requires a clean ModelOpt source checkout.") + return {"commit": commit, "dirty": False} + + +def _collective_publication_preflight(output_dir: Path) -> Mapping[str, Any]: + status = None + if dist.get_rank() == 0: + try: + if output_dir.is_symlink() or output_dir.resolve().exists(): + raise FileExistsError(f"PDD export output already exists: {output_dir}.") + status = {"ok": True, "modelopt_source": _git_source_identity()} + except BaseException as error: + status = {"ok": False, "error": f"{type(error).__name__}: {error}"} + payload = [status] + dist.broadcast_object_list(payload, src=0) + status = payload[0] + if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: + raise RuntimeError("rank 0 broadcast malformed PDD publication preflight status.") + if not status["ok"]: + raise RuntimeError(f"PDD publication preflight failed: {status.get('error')}.") + return status["modelopt_source"] + + +def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: + from modelopt.torch.fastgen import PDDMetadata + + identity = manifest.get("identity") + if not isinstance(identity, Mapping): + raise RuntimeError("PDD checkpoint has no identity mapping.") + if PDDMetadata.from_dict(identity.get("pdd_metadata")) != setup.metadata: + raise RuntimeError("PDD checkpoint metadata does not match the configured student.") + if identity.get("model") != { + "id": config.model_id, + "revision": config.model_revision, + "dtype": str(config.dtype).removeprefix("torch."), + }: + raise RuntimeError("PDD checkpoint model identity does not match the export config.") + checkpoint_automodel = identity.get("automodel") + if not isinstance(checkpoint_automodel, Mapping): + raise RuntimeError("PDD checkpoint has no AutoModel identity.") + for key in ( + "distribution", + "version", + "package_tree_sha256", + "wheel_sha256", + "runtime_versions", + ): + if checkpoint_automodel.get(key) != setup.automodel_snapshot.get(key): + raise RuntimeError(f"PDD checkpoint AutoModel identity mismatch for {key}.") + topology = identity.get("topology") + if not isinstance(topology, Mapping) or topology.get("world_size") != dist.get_world_size(): + raise RuntimeError("PDD checkpoint topology does not match the export process group.") + + +def _collective_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: + local_error = None + try: + _require_checkpoint_identity(config, setup, manifest) + except BaseException as error: + local_error = f"{type(error).__name__}: {error}" + errors: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(errors, local_error) + failures = [f"rank {rank}: {error}" for rank, error in enumerate(errors) if error] + if failures: + raise RuntimeError("PDD checkpoint identity validation failed; " + "; ".join(failures)) + + +def _checkpoint_selector_identity(config: Any, setup: Any) -> dict[str, Any]: + return { + "model": { + "id": config.model_id, + "revision": config.model_revision, + "dtype": str(config.dtype).removeprefix("torch."), + }, + "pdd_metadata": setup.metadata.to_dict(), + "guidance": { + "scale": config.pdd.guidance_scale, + "rescale": config.guidance.rescale, + "eps": config.guidance.eps, + }, + "automodel": { + key: setup.automodel_snapshot[key] + for key in ( + "distribution", + "version", + "package_tree_sha256", + "wheel_sha256", + "runtime_versions", + ) + }, + "topology": {"world_size": dist.get_world_size(), "pure_data_parallel": True}, + } + + +def _collective_checkpoint_resolution( + config: Any, setup: Any, restore_from: str +) -> tuple[Path, Mapping[str, Any]]: + from pdd_checkpoint import resolve_pdd_training_checkpoint + + status = None + if dist.get_rank() == 0: + try: + checkpoint, manifest = resolve_pdd_training_checkpoint( + config.checkpoint.checkpoint_dir, + restore_from, + expected_world_size=dist.get_world_size(), + expected_identity=_checkpoint_selector_identity(config, setup), + ) + status = {"ok": True, "checkpoint": str(checkpoint), "manifest": manifest} + except BaseException as error: + status = {"ok": False, "error": f"{type(error).__name__}: {error}"} + payload = [status] + dist.broadcast_object_list(payload, src=0) + status = payload[0] + if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: + raise RuntimeError("rank 0 broadcast malformed PDD checkpoint resolution status.") + if not status["ok"]: + raise RuntimeError(f"PDD checkpoint resolution failed: {status.get('error')}.") + return Path(status["checkpoint"]), status["manifest"] + + +def main() -> None: + args = _parse_args() + from pdd_artifacts import sha256_file + from pdd_export import write_pdd_export + from pdd_recipe import ( + build_pdd_export_setup, + initialize_pdd_distributed, + resolve_pdd_recipe_config, + ) + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + + raw = yaml.safe_load(args.config.read_text()) + config = resolve_pdd_recipe_config(raw) + if Path(config.model_id).is_dir() or config.model_revision is None: + raise ValueError( + "PDD inference export requires a remote model ID and pinned 40-character revision; " + "mutable local model directories are training-only inputs." + ) + if not math.isfinite(args.max_shard_size_gib) or args.max_shard_size_gib <= 0: + raise ValueError("max_shard_size_gib must be finite and > 0.") + if not math.isfinite(args.memory_headroom) or args.memory_headroom < 1.0: + raise ValueError("memory_headroom must be finite and >= 1.") + max_shard_bytes = int(args.max_shard_size_gib * (1 << 30)) + initialize_pdd_distributed( + backend="nccl" if config.device.type == "cuda" else "gloo", + timeout_minutes=60, + ) + modelopt_source = _collective_publication_preflight(args.output_dir) + restore_from = args.checkpoint or config.checkpoint.restore_from + if not restore_from: + raise ValueError("PDD export requires --checkpoint or checkpoint.restore_from.") + setup = build_pdd_export_setup(config) + try: + checkpoint, checkpoint_manifest = _collective_checkpoint_resolution( + config, setup, restore_from + ) + _collective_checkpoint_identity(config, setup, checkpoint_manifest) + setup.checkpointer.load_model(setup.student, str(checkpoint / "model")) + full_state_bytes, largest_tensor_bytes = collective_export_memory_preflight( + setup.student, + max_shard_bytes=max_shard_bytes, + headroom=args.memory_headroom, + device=config.device, + ) + state_dict = get_model_state_dict( + setup.student, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + local_error = None + try: + if dist.get_rank() == 0: + if set(state_dict) != set(setup.checkpoint_keys): + raise RuntimeError("gathered PDD full-state keys do not match the model.") + if ( + sum(tensor.numel() * tensor.element_size() for tensor in state_dict.values()) + != full_state_bytes + ): + raise RuntimeError( + "gathered PDD full-state byte count changed after preflight." + ) + gathered_largest = max( + tensor.numel() * tensor.element_size() for tensor in state_dict.values() + ) + if gathered_largest != largest_tensor_bytes: + raise RuntimeError("gathered PDD largest tensor changed after preflight.") + elif state_dict: + raise RuntimeError("nonzero rank received a full CPU state dictionary.") + except BaseException as error: + local_error = f"{type(error).__name__}: {error}" + gather_errors: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(gather_errors, local_error) + failures = [f"rank {rank}: {error}" for rank, error in enumerate(gather_errors) if error] + if failures: + raise RuntimeError("PDD full-state gather validation failed; " + "; ".join(failures)) + + publication = None + if dist.get_rank() == 0: + try: + output = write_pdd_export( + args.output_dir, + state_dict, + metadata=setup.metadata, + transformer_config=setup.transformer_config, + identity=checkpoint_manifest["identity"], + source_checkpoint={ + "name": checkpoint.name, + "manifest_sha256": sha256_file(checkpoint / "manifest.json"), + "completed_steps": checkpoint_manifest["completed_steps"], + }, + modelopt_source=modelopt_source, + max_shard_bytes=max_shard_bytes, + ) + publication = {"ok": True, "output": str(output)} + except BaseException as error: + publication = {"ok": False, "error": f"{type(error).__name__}: {error}"} + payload = [publication] + dist.broadcast_object_list(payload, src=0) + publication = payload[0] + if not isinstance(publication, Mapping) or type(publication.get("ok")) is not bool: + raise RuntimeError("rank 0 broadcast malformed PDD publication status.") + if not publication["ok"]: + raise RuntimeError(f"PDD export publication failed: {publication.get('error')}.") + if dist.get_rank() == 0: + print(json.dumps(publication, indent=2, sort_keys=True)) + finally: + setup.checkpointer.close() + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/inference_pdd_qwen_image.py b/examples/diffusers/fastgen/inference_pdd_qwen_image.py new file mode 100644 index 00000000000..75a80053137 --- /dev/null +++ b/examples/diffusers/fastgen/inference_pdd_qwen_image.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run conditional-only PDD inference from an authenticated Qwen-Image export.""" + +from __future__ import annotations + +import argparse +import hashlib +import math +import os +import sys +import time +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch +from torch import nn + +sys.dont_write_bytecode = True + +_THIS_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _THIS_DIR.parents[2] +for path in (_REPO_ROOT, _THIS_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--export-dir", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--prompt-id", required=True) + parser.add_argument("--schedule", choices=("pdd-2", "pdd-4", "pdd-8"), default="pdd-4") + 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) + parser.add_argument("--result-json", type=Path, required=True) + return parser.parse_args() + + +def _dtype_from_name(name: Any) -> torch.dtype: + if not isinstance(name, str): + raise ValueError("PDD export model dtype must be a string.") + dtypes = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + } + try: + return dtypes[name] + except KeyError as error: + raise ValueError(f"PDD inference does not support model dtype {name!r}.") from error + + +def _model_identity(descriptor: Any) -> Mapping[str, Any]: + identity = descriptor.manifest.get("identity") + if not isinstance(identity, Mapping): + raise RuntimeError("PDD export has no identity mapping.") + model = identity.get("model") + if not isinstance(model, Mapping) or set(model) != {"id", "revision", "dtype"}: + raise RuntimeError("PDD export model identity is malformed.") + if not isinstance(model["id"], str) or not model["id"]: + raise RuntimeError("PDD export model ID is invalid.") + revision = model["revision"] + if not isinstance(revision, str) or len(revision) != 40: + raise RuntimeError("PDD export requires a pinned 40-character model revision.") + try: + int(revision, 16) + except ValueError as error: + raise RuntimeError("PDD export model revision must be hexadecimal.") from error + return model + + +def _validate_qwen_projection(student: nn.Module, metadata: Any) -> nn.Linear: + """Validate the ordinary Qwen projection before widening it for PDD.""" + try: + base_projection = student.get_submodule("proj_out") + except AttributeError as error: + raise RuntimeError("reconstructed Qwen student has no proj_out linear layer.") from error + in_channels = getattr(getattr(student, "config", None), "in_channels", None) + if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: + raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") + if ( + not isinstance(base_projection, nn.Linear) + or base_projection.in_features != metadata.projection_in_features + or base_projection.out_features != metadata.projection_out_features + or (base_projection.bias is not None) != metadata.projection_bias + ): + raise RuntimeError("reconstructed Qwen proj_out does not match authenticated metadata.") + if base_projection.out_features != in_channels: + raise RuntimeError( + "Qwen proj_out width must equal transformer in_channels for 2x2 latent packing." + ) + return base_projection + + +def build_pdd_student(export_dir: str | Path) -> tuple[nn.Module, Any, torch.dtype]: + """Reconstruct and strictly load the converted Qwen student on CPU.""" + from diffusers import QwenImageTransformer2DModel + from pdd_export import inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata + + from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd + + descriptor = inspect_pdd_export(export_dir) + model_identity = _model_identity(descriptor) + dtype = _dtype_from_name(model_identity["dtype"]) + student = QwenImageTransformer2DModel.from_config(dict(descriptor.transformer_config)) + metadata = descriptor.metadata + _validate_qwen_projection(student, metadata) + config = pdd_config_from_metadata(metadata, blocks=metadata.inference_blocks) + convert_qwen_image_to_pdd(student, config) + student.to(dtype=dtype) + descriptor = load_pdd_export_into_model(export_dir, student) + return student, descriptor, dtype + + +def _normalize_prompt_condition( + prompt_embeds: Any, + prompt_mask: Any, + *, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Normalize the pinned Diffusers Qwen prompt-encoding contract for PDD.""" + if not isinstance(prompt_embeds, torch.Tensor) or prompt_embeds.ndim != 3: + raise RuntimeError("Qwen prompt embeddings must have shape [B, S, D].") + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + expected_shape = prompt_embeds.shape[:2] + if prompt_mask is None: + prompt_mask = torch.ones(expected_shape, device=device, dtype=torch.long) + elif not isinstance(prompt_mask, torch.Tensor) or prompt_mask.ndim != 2: + raise RuntimeError("Qwen prompt mask must have shape [B, S] or be None.") + elif tuple(prompt_mask.shape) != tuple(expected_shape): + raise RuntimeError("Qwen prompt mask shape does not match prompt embeddings.") + elif prompt_mask.dtype.is_floating_point or prompt_mask.dtype.is_complex: + raise RuntimeError("Qwen prompt mask must use an integer or boolean dtype.") + else: + prompt_mask = prompt_mask.to(device=device, dtype=torch.long) + return prompt_embeds, prompt_mask + + +def _latent_shape(pipe: Any, *, height: int, width: int) -> tuple[int, int, int, int]: + if type(height) is not int or type(width) is not int or height <= 0 or width <= 0: + raise ValueError("height and width must be positive integers.") + quantum = int(pipe.vae_scale_factor) * 2 + if height % quantum or width % quantum: + raise ValueError(f"height and width must be divisible by {quantum}.") + in_channels = getattr(pipe.transformer.config, "in_channels", None) + if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: + raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") + latent_height = 2 * (height // quantum) + latent_width = 2 * (width // quantum) + return 1, in_channels // 4, latent_height, latent_width + + +def _decode_qwen_latents(pipe: Any, latents: torch.Tensor) -> list[Any]: + if latents.ndim != 4: + raise ValueError("PDD Qwen latents must have shape [B, C, H, W].") + vae = pipe.vae + mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype) + std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype) + if mean.numel() != latents.shape[1] or std.numel() != latents.shape[1]: + raise RuntimeError("Qwen VAE latent statistics do not match the student channels.") + decoded_input = latents.unsqueeze(2) * std.view(1, -1, 1, 1, 1) + decoded_input = decoded_input + mean.view(1, -1, 1, 1, 1) + decoded = vae.decode(decoded_input, return_dict=False)[0] + if decoded.ndim != 5 or decoded.shape[2] != 1: + raise RuntimeError("Qwen VAE must return one-frame 5D image tensors.") + return pipe.image_processor.postprocess(decoded[:, :, 0], output_type="pil") + + +def _save_png(path: Path, image: Any) -> None: + if path.is_symlink(): + raise ValueError("PDD inference output cannot be a symlink.") + path = path.resolve() + if path.suffix.lower() != ".png": + raise ValueError("PDD inference output must use a .png suffix.") + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists() or path.is_symlink(): + raise FileExistsError(f"PDD inference output already exists: {path}.") + staging = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + try: + image.save(staging, format="PNG") + with staging.open("rb") as stream: + os.fsync(stream.fileno()) + staging.rename(path) + descriptor = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + finally: + staging.unlink(missing_ok=True) + + +@torch.no_grad() +def main() -> None: + args = _parse_args() + from diffusers import QwenImagePipeline + from pdd_artifacts import sha256_file, write_canonical_json + from pdd_export import pdd_config_from_metadata + + from modelopt.torch.fastgen import PDDPipeline + from modelopt.torch.fastgen.plugins.qwen_image_pdd import QwenImagePDDAdapter + + if args.output.is_symlink() or args.result_json.is_symlink(): + raise ValueError("PDD output and result JSON cannot be symlinks.") + output = args.output.resolve() + result_json = args.result_json.resolve() + if output.exists() or output.is_symlink(): + raise FileExistsError(f"PDD inference output already exists: {output}.") + if result_json.exists() or result_json.is_symlink(): + raise FileExistsError(f"PDD result JSON already exists: {result_json}.") + try: + output_reference = output.relative_to(result_json.parent).as_posix() + except ValueError as error: + raise ValueError("PDD output must be beneath the result JSON directory.") from error + if not isinstance(args.prompt_id, str) or not args.prompt_id.strip(): + raise ValueError("prompt_id must be non-empty.") + if args.seed < 0 or args.seed >= 2**63: + raise ValueError("seed must be in [0, 2**63).") + if args.max_sequence_length < 1: + raise ValueError("max_sequence_length must be positive.") + student, descriptor, dtype = build_pdd_student(args.export_dir) + model_identity = _model_identity(descriptor) + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable.") + + student.to(device=device) + pipe = QwenImagePipeline.from_pretrained( + model_identity["id"], + revision=model_identity["revision"], + transformer=student, + torch_dtype=dtype, + use_safetensors=True, + ) + pipe.to(device) + config = pdd_config_from_metadata(descriptor.metadata, schedule=args.schedule) + adapter = QwenImagePDDAdapter(config) + sampler = PDDPipeline(student, nn.Identity(), config, adapter) + prompt_embeds, prompt_mask = pipe.encode_prompt( + prompt=args.prompt, + device=device, + num_images_per_prompt=1, + max_sequence_length=args.max_sequence_length, + ) + condition = _normalize_prompt_condition( + prompt_embeds, + prompt_mask, + device=device, + dtype=dtype, + ) + generator = torch.Generator(device=device).manual_seed(args.seed) + shape = _latent_shape(pipe, height=args.height, width=args.width) + state = torch.randn(shape, generator=generator, device=device, dtype=torch.float32) + + transformer_invocations = 0 + + def count_invocation( + _module: nn.Module, _args: tuple[Any, ...], _kwargs: Mapping[str, Any] + ) -> None: + nonlocal transformer_invocations + transformer_invocations += 1 + + hook = student.register_forward_pre_hook(count_invocation, with_kwargs=True) + if device.type == "cuda": + torch.cuda.synchronize(device) + started = time.perf_counter() + try: + sampled = sampler.sample(state, condition=condition) + finally: + hook.remove() + images = _decode_qwen_latents(pipe, sampled.to(dtype)) + if device.type == "cuda": + torch.cuda.synchronize(device) + latency = time.perf_counter() - started + expected_invocations = len(config.inference_blocks) + if transformer_invocations != expected_invocations: + raise RuntimeError( + f"PDD sampler made {transformer_invocations} transformer calls; " + f"expected {expected_invocations}." + ) + if len(images) != 1: + raise RuntimeError(f"PDD single-prompt inference returned {len(images)} images.") + if not math.isfinite(latency) or latency <= 0: + raise RuntimeError("PDD inference latency measurement is invalid.") + _save_png(output, images[0]) + + result_json.parent.mkdir(parents=True, exist_ok=True) + result = { + "schema_version": 1, + "record_type": "pdd_inference", + "condition": args.schedule.replace("-", "_"), + "prompt_id": args.prompt_id, + "prompt_sha256": hashlib.sha256(args.prompt.encode("utf-8")).hexdigest(), + "seed": args.seed, + "schedule": args.schedule, + "blocks": list(config.inference_blocks), + "height": args.height, + "width": args.width, + "export_manifest_sha256": sha256_file(descriptor.root / "manifest.json"), + "output": {"path": output_reference, "sha256": sha256_file(output)}, + "scheduler_steps": expected_invocations, + "actual_transformer_invocations": transformer_invocations, + "batch_normalized_transformer_evaluations": transformer_invocations, + "latency_seconds": latency, + } + write_canonical_json(result_json, result) + print(result_json) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/pdd_artifacts.py b/examples/diffusers/fastgen/pdd_artifacts.py new file mode 100644 index 00000000000..54408de8194 --- /dev/null +++ b/examples/diffusers/fastgen/pdd_artifacts.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict canonical-JSON and relative-artifact helpers for the PDD example.""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +from collections.abc import Mapping, Sequence +from pathlib import Path, PurePosixPath +from typing import Any + + +def sha256_file(path: Path) -> str: + """Hash one regular file without following a symlink.""" + if not path.is_file() or path.is_symlink(): + raise RuntimeError(f"PDD artifact is not a regular file: {path}.") + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def require_sha256(value: Any, *, name: str) -> str: + """Validate and normalize a hexadecimal SHA-256 value.""" + if not isinstance(value, str) or len(value) != 64: + raise ValueError(f"{name} must be a 64-character SHA-256 digest.") + try: + int(value, 16) + except ValueError as error: + raise ValueError(f"{name} must be hexadecimal.") from error + return value.lower() + + +def _validate_json_value(value: Any, *, name: str = "JSON") -> None: + if value is None or isinstance(value, str | bool | int): + return + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError(f"{name} contains a non-finite number.") + return + if isinstance(value, Mapping): + if any(not isinstance(key, str) for key in value): + raise TypeError(f"{name} object keys must be strings.") + for key, item in value.items(): + _validate_json_value(item, name=f"{name}.{key}") + return + if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): + for index, item in enumerate(value): + _validate_json_value(item, name=f"{name}[{index}]") + return + raise TypeError(f"{name} contains unsupported type {type(value).__name__}.") + + +def canonical_json_bytes(value: Any) -> bytes: + """Serialize finite JSON deterministically with a trailing newline.""" + _validate_json_value(value) + return ( + json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + + "\n" + ).encode("utf-8") + + +def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + value: dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise ValueError(f"canonical JSON contains duplicate key {key!r}.") + value[key] = item + return value + + +def _reject_json_constant(token: str) -> None: + raise ValueError(f"canonical JSON contains {token}.") + + +def load_canonical_json(path: Path) -> Any: + """Load canonical JSON, rejecting duplicates, NaN/Inf, and noncanonical bytes.""" + if not path.is_file() or path.is_symlink(): + raise RuntimeError(f"PDD JSON artifact is not a regular file: {path}.") + raw = path.read_bytes() + try: + value = json.loads( + raw, + object_pairs_hook=_unique_object, + parse_constant=_reject_json_constant, + ) + except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as error: + raise RuntimeError(f"cannot parse canonical PDD JSON {path}.") from error + try: + expected = canonical_json_bytes(value) + except (TypeError, ValueError) as error: + raise RuntimeError(f"invalid canonical PDD JSON {path}.") from error + if raw != expected: + raise RuntimeError(f"PDD JSON is not in canonical form: {path}.") + return value + + +def write_canonical_json(path: Path, value: Any) -> None: + """Create one canonical JSON file and fsync its contents.""" + data = canonical_json_bytes(value) + with path.open("xb") as stream: + stream.write(data) + stream.flush() + os.fsync(stream.fileno()) + + +def resolve_relative_artifact(root: Path, reference: str) -> Path: + """Resolve a normalized POSIX reference beneath root with no symlink component.""" + if not isinstance(reference, str) or not reference: + raise ValueError("artifact reference must be a non-empty string.") + if "\\" in reference: + raise ValueError(f"artifact reference must use POSIX separators: {reference!r}.") + pure = PurePosixPath(reference) + if pure.is_absolute() or any(part in ("", ".", "..") for part in pure.parts): + raise ValueError(f"artifact reference must be normalized and relative: {reference!r}.") + root = root.resolve() + if not root.is_dir() or root.is_symlink(): + raise RuntimeError(f"PDD artifact root is not a regular directory: {root}.") + candidate = root + for part in pure.parts: + candidate = candidate / part + if candidate.is_symlink(): + raise RuntimeError(f"PDD artifact reference traverses a symlink: {reference!r}.") + resolved = candidate.resolve() + try: + resolved.relative_to(root) + except ValueError as error: + raise ValueError(f"artifact reference escapes its root: {reference!r}.") from error + if not resolved.is_file(): + raise FileNotFoundError(f"PDD artifact is missing: {reference!r}.") + return resolved + + +def validate_artifact_reference(root: Path, value: Any, *, name: str) -> Path: + """Validate an exact path/hash reference and return the verified regular file.""" + if not isinstance(value, Mapping) or set(value) != {"path", "sha256"}: + raise ValueError(f"{name} must contain exactly path and sha256.") + path = resolve_relative_artifact(root, value["path"]) + expected = require_sha256(value["sha256"], name=f"{name}.sha256") + if sha256_file(path) != expected: + raise RuntimeError(f"{name} SHA-256 does not match {value['path']!r}.") + return path diff --git a/examples/diffusers/fastgen/pdd_checkpoint.py b/examples/diffusers/fastgen/pdd_checkpoint.py index 48f0d96ea3e..8c0ae9a6be6 100644 --- a/examples/diffusers/fastgen/pdd_checkpoint.py +++ b/examples/diffusers/fastgen/pdd_checkpoint.py @@ -11,6 +11,7 @@ import os import shutil import uuid +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -20,7 +21,7 @@ from modelopt.torch.fastgen import PDDMetadata if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Sequence _CHECKPOINT_SCHEMA_VERSION = 1 _COMPLETE_SCHEMA_VERSION = 1 @@ -304,6 +305,204 @@ def verify_first_batch(self, sample_ids: Sequence[str]) -> None: ) +def _checkpoint_sidecar_paths(checkpoint: Path, world_size: int) -> list[Path]: + paths: list[Path] = [] + for rank in range(world_size): + paths.extend( + ( + checkpoint / "rng" / f"rng_dp_rank_{rank}.pt", + checkpoint / "sampler" / f"sampler_dp_rank_{rank}.pt", + checkpoint / "trainer" / f"trainer_dp_rank_{rank}.pt", + ) + ) + return paths + + +def validate_pdd_training_checkpoint( + checkpoint: str | Path, + *, + expected_identity: Mapping[str, Any] | None = None, + expected_world_size: int | None = None, +) -> dict[str, Any]: + """Validate a complete training checkpoint without deserializing pickle sidecars.""" + unresolved_checkpoint = Path(checkpoint) + if unresolved_checkpoint.is_symlink(): + raise RuntimeError(f"PDD checkpoint cannot be a symlink: {unresolved_checkpoint}.") + checkpoint = unresolved_checkpoint.resolve() + if not checkpoint.is_dir(): + raise RuntimeError(f"PDD checkpoint is not a regular directory: {checkpoint}.") + symlinks = [path for path in checkpoint.rglob("*") if path.is_symlink()] + if symlinks: + raise RuntimeError(f"PDD checkpoint contains a symlink: {symlinks[0]}.") + marker_path = checkpoint / "COMPLETE" + manifest_path = checkpoint / "manifest.json" + if ( + not marker_path.is_file() + or marker_path.is_symlink() + or not manifest_path.is_file() + or manifest_path.is_symlink() + ): + raise RuntimeError(f"PDD checkpoint is incomplete: {checkpoint}.") + marker = _read_json(marker_path) + if ( + set(marker) != {"schema_version", "manifest_sha256"} + or marker.get("schema_version") != _COMPLETE_SCHEMA_VERSION + ): + raise RuntimeError("PDD COMPLETE marker is incompatible.") + if marker["manifest_sha256"] != _sha256(manifest_path): + raise RuntimeError("PDD COMPLETE marker does not match manifest content.") + manifest = _read_json(manifest_path) + manifest_keys = { + "schema_version", + "identity", + "completed_steps", + "learning_rates", + "parent_checkpoint", + "rank_progress", + "dcp_sha256", + "sidecar_sha256", + } + if set(manifest) != manifest_keys or manifest["schema_version"] != _CHECKPOINT_SCHEMA_VERSION: + raise RuntimeError("PDD checkpoint manifest schema is incompatible.") + if expected_identity is not None and manifest["identity"] != expected_identity: + raise RuntimeError("PDD checkpoint identity does not match the current run.") + identity = manifest["identity"] + topology = identity.get("topology") if isinstance(identity, Mapping) else None + world_size = topology.get("world_size") if isinstance(topology, Mapping) else None + if type(world_size) is not int or world_size < 1: + raise RuntimeError("PDD checkpoint identity has an invalid world size.") + if expected_world_size is not None and world_size != expected_world_size: + raise RuntimeError( + f"PDD checkpoint world size {world_size} does not match {expected_world_size}." + ) + if _read_json(checkpoint / "pdd_config.json") != identity: + raise RuntimeError("PDD checkpoint config sidecar does not match the manifest.") + trainer_state = _read_json(checkpoint / "trainer_state.json") + if trainer_state != { + "completed_steps": manifest["completed_steps"], + "learning_rates": manifest["learning_rates"], + "parent_checkpoint": manifest["parent_checkpoint"], + }: + raise RuntimeError("PDD trainer-state sidecar does not match the manifest.") + rank_progress = manifest["rank_progress"] + if not isinstance(rank_progress, list) or len(rank_progress) != world_size: + raise RuntimeError("PDD checkpoint rank progress does not match its topology.") + expected_dcp = manifest["dcp_sha256"] + if not isinstance(expected_dcp, dict) or any( + not isinstance(path, str) or not isinstance(digest, str) + for path, digest in expected_dcp.items() + ): + raise RuntimeError("PDD checkpoint DCP hash inventory is malformed.") + if _dcp_payload_hashes(checkpoint) != expected_dcp: + raise RuntimeError("PDD checkpoint DCP payload inventory or hash does not match.") + expected_sidecars = _checkpoint_sidecar_paths(checkpoint, world_size) + expected_relative = {path.relative_to(checkpoint).as_posix() for path in expected_sidecars} + if ( + not isinstance(manifest["sidecar_sha256"], dict) + or set(manifest["sidecar_sha256"]) != expected_relative + ): + raise RuntimeError("PDD checkpoint sidecar inventory does not match the topology.") + for path in expected_sidecars: + relative = path.relative_to(checkpoint).as_posix() + if not path.is_file() or path.is_symlink(): + raise RuntimeError(f"PDD checkpoint sidecar is missing: {relative}.") + if _sha256(path) != manifest["sidecar_sha256"][relative]: + raise RuntimeError(f"PDD checkpoint sidecar hash mismatch: {relative}.") + for candidate in checkpoint.rglob("*"): + lowered = candidate.name.lower() + if any(token in lowered for token in _FORBIDDEN_ARTIFACT_TOKENS): + raise RuntimeError(f"PDD checkpoint contains a forbidden DMD artifact: {candidate}.") + return manifest + + +def resolve_pdd_training_checkpoint( + root: str | Path, + restore_from: str | Path, + *, + expected_world_size: int, + expected_identity: Mapping[str, Any] | None = None, +) -> tuple[Path, dict[str, Any]]: + """Resolve an explicit checkpoint or the newest compatible complete LATEST candidate.""" + unresolved_root = Path(root) + if unresolved_root.is_symlink(): + raise ValueError("PDD checkpoint_dir cannot be a symlink.") + root = unresolved_root.resolve() + if str(restore_from).upper() != "LATEST": + candidate = Path(restore_from) + if not candidate.is_absolute(): + candidate = root / candidate + if candidate.is_symlink(): + raise ValueError("explicit PDD checkpoint cannot be a symlink.") + candidate = candidate.resolve() + try: + candidate.relative_to(root) + except ValueError as error: + raise ValueError("explicit PDD checkpoint must be beneath checkpoint_dir.") from error + manifest = validate_pdd_training_checkpoint( + candidate, + expected_world_size=expected_world_size, + ) + if expected_identity is not None and not _identity_contains( + manifest.get("identity"), expected_identity + ): + raise RuntimeError("explicit PDD checkpoint identity does not match the selector.") + return candidate, manifest + + if not isinstance(expected_identity, Mapping) or not expected_identity: + raise ValueError("LATEST resolution requires a non-empty expected_identity selector.") + + pointed: tuple[int, Path, dict[str, Any]] | None = None + pointer = root / "LATEST" + if pointer.is_file() and not pointer.is_symlink(): + candidate = (root / pointer.read_text().strip()).resolve() + try: + candidate.relative_to(root) + manifest = validate_pdd_training_checkpoint( + candidate, + expected_world_size=expected_world_size, + ) + completed = manifest["completed_steps"] + if type(completed) is not int or completed < 0: + raise RuntimeError("PDD checkpoint completed_steps is invalid.") + if not _identity_contains(manifest.get("identity"), expected_identity): + raise RuntimeError("pointed PDD checkpoint identity does not match the selector.") + pointed = (completed, candidate, manifest) + except (ValueError, RuntimeError): + pass + + candidates: list[tuple[int, Path, dict[str, Any]]] = [] + if root.is_dir(): + for path in root.iterdir(): + suffix = path.name.removeprefix("step_") + if not path.is_dir() or not path.name.startswith("step_") or not suffix.isdigit(): + continue + if pointed is not None and int(suffix) <= pointed[0]: + continue + try: + manifest = validate_pdd_training_checkpoint( + path, + expected_world_size=expected_world_size, + ) + except RuntimeError: + continue + if not _identity_contains(manifest.get("identity"), expected_identity): + continue + completed = manifest["completed_steps"] + if type(completed) is int and completed >= 0: + candidates.append((completed, path.resolve(), manifest)) + if pointed is not None: + candidates.append(pointed) + if not candidates: + raise FileNotFoundError(f"no complete compatible PDD checkpoint exists beneath {root}.") + return max(candidates, key=lambda item: (item[0], item[1].name))[1:] + + +def _identity_contains(actual: Any, expected: Mapping[str, Any]) -> bool: + if not isinstance(actual, Mapping): + return False + return all(key in actual and actual[key] == value for key, value in expected.items()) + + class PDDCheckpointManager: """Publish and restore complete, metadata-compatible PDD checkpoints only.""" @@ -349,16 +548,7 @@ def _rank_summary(self) -> dict[str, Any]: } def _sidecar_paths(self, checkpoint: Path) -> list[Path]: - paths: list[Path] = [] - for rank in range(_world_size()): - paths.extend( - ( - checkpoint / "rng" / f"rng_dp_rank_{rank}.pt", - checkpoint / "sampler" / f"sampler_dp_rank_{rank}.pt", - checkpoint / "trainer" / f"trainer_dp_rank_{rank}.pt", - ) - ) - return paths + return _checkpoint_sidecar_paths(checkpoint, _world_size()) def _manifest(self, checkpoint: Path) -> dict[str, Any]: manifest = _read_json(checkpoint / "manifest.json") @@ -379,61 +569,11 @@ def _manifest(self, checkpoint: Path) -> dict[str, Any]: return manifest def _validate_checkpoint(self, checkpoint: Path, *, require_identity: bool) -> dict[str, Any]: - if not checkpoint.is_dir() or checkpoint.is_symlink(): - raise RuntimeError(f"PDD checkpoint is not a regular directory: {checkpoint}.") - marker_path = checkpoint / "COMPLETE" - manifest_path = checkpoint / "manifest.json" - if not marker_path.is_file() or not manifest_path.is_file(): - raise RuntimeError(f"PDD checkpoint is incomplete: {checkpoint}.") - marker = _read_json(marker_path) - if ( - set(marker) != {"schema_version", "manifest_sha256"} - or marker.get("schema_version") != _COMPLETE_SCHEMA_VERSION - ): - raise RuntimeError("PDD COMPLETE marker is incompatible.") - if marker["manifest_sha256"] != _sha256(manifest_path): - raise RuntimeError("PDD COMPLETE marker does not match manifest content.") - manifest = self._manifest(checkpoint) - if require_identity and manifest["identity"] != self.identity: - raise RuntimeError("PDD checkpoint identity does not match the current run.") - if _read_json(checkpoint / "pdd_config.json") != manifest["identity"]: - raise RuntimeError("PDD checkpoint config sidecar does not match the manifest.") - trainer_state = _read_json(checkpoint / "trainer_state.json") - if trainer_state != { - "completed_steps": manifest["completed_steps"], - "learning_rates": manifest["learning_rates"], - "parent_checkpoint": manifest["parent_checkpoint"], - }: - raise RuntimeError("PDD trainer-state sidecar does not match the manifest.") - expected_dcp = manifest["dcp_sha256"] - if not isinstance(expected_dcp, dict) or any( - not isinstance(path, str) or not isinstance(digest, str) - for path, digest in expected_dcp.items() - ): - raise RuntimeError("PDD checkpoint DCP hash inventory is malformed.") - actual_dcp = _dcp_payload_hashes(checkpoint) - if actual_dcp != expected_dcp: - raise RuntimeError("PDD checkpoint DCP payload inventory or hash does not match.") - expected_sidecars = self._sidecar_paths(checkpoint) - expected_relative = {path.relative_to(checkpoint).as_posix() for path in expected_sidecars} - if ( - not isinstance(manifest["sidecar_sha256"], dict) - or set(manifest["sidecar_sha256"]) != expected_relative - ): - raise RuntimeError("PDD checkpoint sidecar inventory does not match the topology.") - for path in expected_sidecars: - relative = path.relative_to(checkpoint).as_posix() - if not path.is_file() or path.is_symlink(): - raise RuntimeError(f"PDD checkpoint sidecar is missing: {relative}.") - if _sha256(path) != manifest["sidecar_sha256"][relative]: - raise RuntimeError(f"PDD checkpoint sidecar hash mismatch: {relative}.") - for candidate in checkpoint.rglob("*"): - lowered = candidate.name.lower() - if any(token in lowered for token in _FORBIDDEN_ARTIFACT_TOKENS): - raise RuntimeError( - f"PDD checkpoint contains a forbidden DMD artifact: {candidate}." - ) - return manifest + return validate_pdd_training_checkpoint( + checkpoint, + expected_identity=self.identity if require_identity else None, + expected_world_size=_world_size(), + ) def _compatible_candidates( self, diff --git a/examples/diffusers/fastgen/pdd_evaluation.py b/examples/diffusers/fastgen/pdd_evaluation.py new file mode 100644 index 00000000000..c188c84a9d1 --- /dev/null +++ b/examples/diffusers/fastgen/pdd_evaluation.py @@ -0,0 +1,1088 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict evidence-bundle validation and deterministic PDD effectiveness summaries.""" + +from __future__ import annotations + +import hashlib +import math +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from pdd_artifacts import ( + canonical_json_bytes, + load_canonical_json, + require_sha256, + sha256_file, + validate_artifact_reference, +) +from pdd_export import inspect_pdd_export + +from modelopt.torch.fastgen import make_shifted_flow_grid + +EVALUATION_CONDITIONS = ( + "teacher_guided", + "undistilled_euler_4", + "undistilled_2step_4eval", + "pdd_2", + "pdd_4", + "pdd_8", +) + + +def _grid_protocol(grid_size: int) -> Mapping[str, Any]: + nodes = [float(value) for value in make_shifted_flow_grid(grid_size, 5.0).tolist()] + return { + "builder": "modelopt.torch.fastgen.make_shifted_flow_grid", + "formula": "s*u/(1+(s-1)*u), u=1-i/grid_size, i=0..grid_size", + "grid_size": grid_size, + "flow_shift": 5.0, + "nodes": nodes, + "nodes_sha256": hashlib.sha256(canonical_json_bytes(nodes)).hexdigest(), + } + + +GRID_PROTOCOLS: Mapping[str, Mapping[str, Any]] = { + "pdd_grid_128_shift5": _grid_protocol(128), + "teacher_grid_50_shift5": _grid_protocol(50), +} + +_GUIDED_CFG_PROTOCOL = { + "execution": "sequential_conditional_unconditional", + "guidance_scale": 4.0, + "rescale": 1.0, + "eps": 1e-5, + "negative_condition": "manifest_negative_condition", +} + +_DISABLED_CFG_PROTOCOL = { + "execution": "disabled", + "guidance_scale": None, + "rescale": None, + "eps": None, + "negative_condition": None, +} + +INTEGRATOR_PROTOCOLS: Mapping[str, Mapping[str, Any]] = { + "euler_explicit": { + "math_dtype": "float32", + "velocity_evaluations_per_interval": 1, + "equations": [ + "dt=t_next-t_current", + "v_current=velocity(x_current,t_current)", + "x_next=x_current+dt*v_current", + ], + "terminal_rule": "apply the same Euler update when t_next=0; no special fallback", + }, + "heun_explicit_trapezoid": { + "math_dtype": "float32", + "velocity_evaluations_per_interval": 2, + "equations": [ + "dt=t_next-t_current", + "v_current=velocity(x_current,t_current)", + "x_predict=x_current+dt*v_current", + "v_predict=velocity(x_predict,t_next)", + "x_next=x_current+0.5*dt*(v_current+v_predict)", + ], + "terminal_rule": ( + "always evaluate v_predict at t_next, including t_next=0; no Euler fallback" + ), + }, + "pdd_fused_euler": { + "math_dtype": "float32", + "velocity_evaluations_per_interval": 1, + "implementation": "modelopt.torch.fastgen.methods.pdd.PDDPipeline.sample", + "source_identity": "manifest.modelopt.commit", + "equations": [ + "v_fused=student_fused_block(x_start,t_start,start,end,authenticated_grid)", + "x_end=x_start+(t_end-t_start)*v_fused", + ], + "terminal_rule": "apply the same fused update when t_end=0; no special fallback", + }, +} + +CONDITION_PROTOCOLS: Mapping[str, Mapping[str, Any]] = { + "teacher_guided": { + "artifact": "pdd_export", + "model_role": "frozen_teacher", + "integrator": "euler_explicit", + "cfg": _GUIDED_CFG_PROTOCOL, + "grid": { + "protocol": "teacher_grid_50_shift5", + "node_indices": list(range(51)), + }, + "pdd_blocks": [], + "scheduler_steps": 50, + "actual_transformer_invocations": 100, + "batch_normalized_transformer_evaluations": 100, + }, + "undistilled_euler_4": { + "artifact": "pdd_export", + "model_role": "pinned_base_model", + "integrator": "euler_explicit", + "cfg": _GUIDED_CFG_PROTOCOL, + "grid": { + "protocol": "pdd_grid_128_shift5", + "node_indices": [0, 32, 64, 96, 128], + }, + "pdd_blocks": [], + "scheduler_steps": 4, + "actual_transformer_invocations": 8, + "batch_normalized_transformer_evaluations": 8, + }, + "undistilled_2step_4eval": { + "artifact": "pdd_export", + "model_role": "pinned_base_model", + "integrator": "heun_explicit_trapezoid", + "cfg": _GUIDED_CFG_PROTOCOL, + "grid": { + "protocol": "pdd_grid_128_shift5", + "node_indices": [0, 64, 128], + }, + "pdd_blocks": [], + "scheduler_steps": 2, + "actual_transformer_invocations": 8, + "batch_normalized_transformer_evaluations": 8, + }, + "pdd_2": { + "artifact": "pdd_export", + "model_role": "pdd_student", + "integrator": "pdd_fused_euler", + "cfg": _DISABLED_CFG_PROTOCOL, + "grid": { + "protocol": "pdd_grid_128_shift5", + "node_indices": [0, 64, 128], + }, + "pdd_blocks": [64, 64], + "scheduler_steps": 2, + "actual_transformer_invocations": 2, + "batch_normalized_transformer_evaluations": 2, + }, + "pdd_4": { + "artifact": "pdd_export", + "model_role": "pdd_student", + "integrator": "pdd_fused_euler", + "cfg": _DISABLED_CFG_PROTOCOL, + "grid": { + "protocol": "pdd_grid_128_shift5", + "node_indices": [0, 32, 64, 96, 128], + }, + "pdd_blocks": [32, 32, 32, 32], + "scheduler_steps": 4, + "actual_transformer_invocations": 4, + "batch_normalized_transformer_evaluations": 4, + }, + "pdd_8": { + "artifact": "pdd_export", + "model_role": "pdd_student", + "integrator": "pdd_fused_euler", + "cfg": _DISABLED_CFG_PROTOCOL, + "grid": { + "protocol": "pdd_grid_128_shift5", + "node_indices": [0, 16, 32, 48, 64, 80, 96, 112, 128], + }, + "pdd_blocks": [16, 16, 16, 16, 16, 16, 16, 16], + "scheduler_steps": 8, + "actual_transformer_invocations": 8, + "batch_normalized_transformer_evaluations": 8, + }, +} + +_PROTOCOL_FIELDS = ( + "conditions", + "condition_protocols", + "grid_protocols", + "integrator_protocols", + "image_protocol", + "metric_protocols", + "timing_protocol", + "decision_rule", + "negative_condition", + "data_snapshot", + "stage_run_ids", + "prompt_set", + "bootstrap", +) + + +def _exact_mapping(value: Any, keys: set[str], *, name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or set(value) != keys: + actual = sorted(value) if isinstance(value, Mapping) else type(value).__name__ + raise ValueError(f"{name} keys mismatch: expected={sorted(keys)}, actual={actual}.") + return value + + +def _commit(value: Any, *, name: str) -> str: + if not isinstance(value, str) or len(value) != 40: + raise ValueError(f"{name} must be a 40-character Git commit.") + try: + int(value, 16) + except ValueError as error: + raise ValueError(f"{name} must be hexadecimal.") from error + return value.lower() + + +def _positive_finite(value: Any, *, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise TypeError(f"{name} must be a real number.") + resolved = float(value) + if not math.isfinite(resolved) or resolved <= 0: + raise ValueError(f"{name} must be finite and > 0.") + return resolved + + +def _nonnegative_int(value: Any, *, name: str) -> int: + if type(value) is not int or value < 0: + raise ValueError(f"{name} must be a nonnegative integer.") + return value + + +def _load_prompt_set(path: Path) -> tuple[dict[tuple[str, str, int], str], Mapping[str, Any]]: + prompt_set = _exact_mapping( + load_canonical_json(path), + {"schema_version", "prompts"}, + name="prompt set", + ) + if prompt_set["schema_version"] != 1 or not isinstance(prompt_set["prompts"], list): + raise ValueError("prompt set schema is unsupported.") + expected: dict[tuple[str, str, int], str] = {} + sort_keys: list[str] = [] + for index, raw_prompt in enumerate(prompt_set["prompts"]): + prompt = _exact_mapping( + raw_prompt, + {"prompt_id", "prompt", "prompt_sha256", "seeds"}, + name=f"prompts[{index}]", + ) + prompt_id = prompt["prompt_id"] + text = prompt["prompt"] + if not isinstance(prompt_id, str) or not prompt_id or not isinstance(text, str): + raise ValueError("prompt_id must be non-empty and prompt must be a string.") + digest = require_sha256(prompt["prompt_sha256"], name=f"prompts[{index}].prompt_sha256") + if hashlib.sha256(text.encode("utf-8")).hexdigest() != digest: + raise RuntimeError(f"prompt SHA-256 does not match for {prompt_id!r}.") + seeds = prompt["seeds"] + if ( + not isinstance(seeds, list) + or not seeds + or any(type(seed) is not int or seed < 0 or seed >= 2**63 for seed in seeds) + or seeds != sorted(set(seeds)) + ): + raise ValueError(f"prompt seeds must be sorted unique integers for {prompt_id!r}.") + sort_keys.append(prompt_id) + for seed in seeds: + key = (prompt_id, digest, seed) + if key in expected: + raise ValueError(f"duplicate prompt/seed pair: {key}.") + expected[key] = text + if sort_keys != sorted(set(sort_keys)): + raise ValueError("prompts must have unique IDs sorted lexicographically.") + if not expected: + raise ValueError("prompt set must contain at least one prompt/seed pair.") + return expected, prompt_set + + +def _protocol_sha256(manifest: Mapping[str, Any]) -> str: + payload = {name: manifest[name] for name in _PROTOCOL_FIELDS} + return hashlib.sha256(canonical_json_bytes(payload)).hexdigest() + + +def _validate_data_snapshot(root: Path, reference: Any) -> tuple[Path, Mapping[str, Any]]: + path = validate_artifact_reference(root, reference, name="data_snapshot") + snapshot = _exact_mapping( + load_canonical_json(path), + { + "schema_version", + "record_type", + "dataset_snapshot_sha256", + "train_ids_sha256", + "heldout_ids_sha256", + }, + name="data snapshot", + ) + if snapshot["schema_version"] != 1 or snapshot["record_type"] != "pdd_dataset_snapshot": + raise ValueError("data snapshot is not a schema-v1 PDD dataset snapshot.") + for name in ("dataset_snapshot_sha256", "train_ids_sha256", "heldout_ids_sha256"): + require_sha256(snapshot[name], name=f"data snapshot {name}") + return path, snapshot + + +def _validate_stage_evidence( + root: Path, + reference: Any, + *, + stage: str, + expected_run_id: str, + protocol_sha256: str, + model: Mapping[str, Any], + modelopt: Mapping[str, Any], + data_snapshot_sha256: str, + export_source_checkpoint: Mapping[str, Any], +) -> Mapping[str, Any]: + path = validate_artifact_reference(root, reference, name=f"stage_evidence.{stage}") + evidence = _exact_mapping( + load_canonical_json(path), + { + "schema_version", + "record_type", + "stage", + "status", + "run_id", + "model", + "modelopt", + "data_snapshot_sha256", + "evaluation_protocol_sha256", + "checkpoint", + "results", + }, + name=f"{stage} evidence", + ) + if ( + evidence["schema_version"] != 1 + or evidence["record_type"] != "pdd_stage_evidence" + or evidence["stage"] != stage + or evidence["status"] != "passed" + ): + raise ValueError(f"{stage} evidence is not a passed schema-v1 record.") + if evidence["run_id"] != expected_run_id: + raise RuntimeError(f"{stage} evidence run_id does not match frozen stage_run_ids.") + if evidence["model"] != model or evidence["modelopt"] != modelopt: + raise RuntimeError(f"{stage} evidence model/code lineage does not match evaluation.") + if ( + require_sha256(evidence["data_snapshot_sha256"], name=f"{stage} data snapshot SHA-256") + != data_snapshot_sha256 + ): + raise RuntimeError(f"{stage} evidence data lineage does not match evaluation.") + if ( + require_sha256(evidence["evaluation_protocol_sha256"], name=f"{stage} protocol SHA-256") + != protocol_sha256 + ): + raise RuntimeError(f"{stage} evidence does not carry the frozen evaluation protocol.") + checkpoint = _exact_mapping( + evidence["checkpoint"], + {"name", "manifest_sha256", "completed_steps"}, + name=f"{stage} checkpoint", + ) + if ( + not isinstance(checkpoint["name"], str) + or not checkpoint["name"] + or Path(checkpoint["name"]).name != checkpoint["name"] + or type(checkpoint["completed_steps"]) is not int + or checkpoint["completed_steps"] < 1 + ): + raise ValueError(f"{stage} checkpoint lineage is malformed.") + require_sha256(checkpoint["manifest_sha256"], name=f"{stage} checkpoint manifest SHA-256") + if stage == "training" and checkpoint != export_source_checkpoint: + raise RuntimeError("training evidence checkpoint does not match the exported checkpoint.") + results_path = validate_artifact_reference( + root, evidence["results"], name=f"{stage} evidence results" + ) + results = _exact_mapping( + load_canonical_json(results_path), + { + "schema_version", + "record_type", + "stage", + "status", + "slurm_job_ids", + "completed_updates", + "finite_loss", + "finite_gradients", + "resume_verified", + }, + name=f"{stage} results", + ) + expected_updates = 1_500 if stage == "canary" else 10_000 + expected_job_count = 3 if stage == "canary" else 1 + job_ids = results["slurm_job_ids"] + if ( + results["schema_version"] != 1 + or results["record_type"] != "pdd_stage_results" + or results["stage"] != stage + or results["status"] != "passed" + or not isinstance(job_ids, list) + or len(job_ids) != expected_job_count + or any(type(job_id) is not int or job_id <= 0 for job_id in job_ids) + or len(set(job_ids)) != len(job_ids) + or results["completed_updates"] != expected_updates + or results["finite_loss"] is not True + or results["finite_gradients"] is not True + or results["resume_verified"] is not True + ): + raise ValueError(f"{stage} results do not satisfy the frozen passed-stage contract.") + return evidence + + +def _validate_observations( + root: Path, + path: Path, + *, + prompt_pairs: Mapping[tuple[str, str, int], str], + metric_protocols: Mapping[str, Mapping[str, Any]], + image_protocol: Mapping[str, Any], + timing_protocol: Mapping[str, Any], + export_manifest_sha256: str, + evaluation_protocol_sha256: str, +) -> tuple[Mapping[str, Any], ...]: + document = _exact_mapping( + load_canonical_json(path), + {"schema_version", "records"}, + name="observations", + ) + if document["schema_version"] != 1 or not isinstance(document["records"], list): + raise ValueError("observation schema is unsupported.") + records: list[Mapping[str, Any]] = [] + observed: set[tuple[str, str, str, int]] = set() + order = {condition: index for index, condition in enumerate(EVALUATION_CONDITIONS)} + sort_keys: list[tuple[str, int, int]] = [] + keys = { + "condition", + "prompt_id", + "prompt_sha256", + "seed", + "metrics", + "output", + "scheduler_steps", + "actual_transformer_invocations", + "batch_normalized_transformer_evaluations", + "latency_seconds", + "throughput_images_per_second", + "peak_device_memory_bytes", + "height", + "width", + "protocol_sha256", + "evaluation_protocol_sha256", + "model_artifact_sha256", + } + for index, raw_record in enumerate(document["records"]): + record = _exact_mapping(raw_record, keys, name=f"records[{index}]") + condition = record["condition"] + if condition not in order: + raise ValueError(f"unknown evaluation condition {condition!r}.") + prompt_id = record["prompt_id"] + prompt_sha = require_sha256(record["prompt_sha256"], name=f"records[{index}].prompt_sha256") + seed = _nonnegative_int(record["seed"], name=f"records[{index}].seed") + pair = (prompt_id, prompt_sha, seed) + if pair not in prompt_pairs: + raise RuntimeError(f"observation does not match the prompt set: {pair}.") + key = (condition, *pair) + if key in observed: + raise ValueError(f"duplicate evaluation observation: {key}.") + observed.add(key) + protocol = CONDITION_PROTOCOLS[condition] + if ( + require_sha256(record["protocol_sha256"], name=f"records[{index}].protocol_sha256") + != hashlib.sha256(canonical_json_bytes(protocol)).hexdigest() + ): + raise RuntimeError(f"records[{index}] is not bound to its condition protocol.") + if ( + require_sha256( + record["evaluation_protocol_sha256"], + name=f"records[{index}].evaluation_protocol_sha256", + ) + != evaluation_protocol_sha256 + ): + raise RuntimeError(f"records[{index}] is not bound to the evaluation protocol.") + if ( + require_sha256( + record["model_artifact_sha256"], + name=f"records[{index}].model_artifact_sha256", + ) + != export_manifest_sha256 + ): + raise RuntimeError(f"records[{index}] model artifact does not match the PDD export.") + if ( + record["height"] != image_protocol["height"] + or record["width"] != image_protocol["width"] + ): + raise RuntimeError(f"records[{index}] resolution does not match image_protocol.") + metrics = record["metrics"] + if not isinstance(metrics, Mapping) or set(metrics) != set(metric_protocols): + raise ValueError(f"records[{index}].metrics does not match metric_protocols.") + if any( + isinstance(value, bool) + or not isinstance(value, int | float) + or not math.isfinite(float(value)) + for value in metrics.values() + ): + raise ValueError(f"records[{index}].metrics contains a non-finite value.") + validate_artifact_reference(root, record["output"], name=f"records[{index}].output") + counts = tuple( + _nonnegative_int(record[name], name=f"records[{index}].{name}") + for name in ( + "scheduler_steps", + "actual_transformer_invocations", + "batch_normalized_transformer_evaluations", + ) + ) + if not all(counts) or counts[1] > counts[2]: + raise ValueError(f"records[{index}] has invalid compute counters.") + expected_counts = tuple( + protocol[name] + for name in ( + "scheduler_steps", + "actual_transformer_invocations", + "batch_normalized_transformer_evaluations", + ) + ) + if counts != expected_counts: + raise RuntimeError( + f"{condition} compute counters must be {expected_counts}, got {counts}." + ) + latency = _positive_finite( + record["latency_seconds"], name=f"records[{index}].latency_seconds" + ) + throughput = _positive_finite( + record["throughput_images_per_second"], + name=f"records[{index}].throughput_images_per_second", + ) + expected_throughput = timing_protocol["batch_size"] / latency + if not math.isclose(throughput, expected_throughput, rel_tol=1e-6, abs_tol=0.0): + raise RuntimeError(f"records[{index}] throughput does not match batch size / latency.") + if ( + type(record["peak_device_memory_bytes"]) is not int + or record["peak_device_memory_bytes"] <= 0 + ): + raise ValueError(f"records[{index}].peak_device_memory_bytes must be positive.") + sort_keys.append((prompt_id, seed, order[condition])) + records.append(record) + if sort_keys != sorted(sort_keys): + raise ValueError("observations must be sorted by prompt_id, seed, and condition order.") + expected = {(condition, *pair) for pair in prompt_pairs for condition in EVALUATION_CONDITIONS} + if observed != expected: + missing = sorted(expected - observed) + extra = sorted(observed - expected) + raise RuntimeError( + f"effectiveness observations are incomplete: missing={missing[:5]}, extra={extra[:5]}." + ) + for condition in EVALUATION_CONDITIONS: + condition_counts = { + ( + record["scheduler_steps"], + record["actual_transformer_invocations"], + record["batch_normalized_transformer_evaluations"], + ) + for record in records + if record["condition"] == condition + } + if len(condition_counts) != 1: + raise RuntimeError(f"{condition} compute counters vary across paired observations.") + return tuple(records) + + +def _validate_automodel_snapshot( + snapshot: Any, + *, + export_automodel: Mapping[str, Any], +) -> None: + snapshot = _exact_mapping( + snapshot, + { + "distribution", + "files", + "import_origin", + "package_file_count", + "package_tree_sha256", + "release_commit", + "release_tag", + "root", + "runtime_versions", + "version", + "wheel", + "wheel_sha256", + }, + name="AutoModel environment snapshot", + ) + for key in ("distribution", "version", "runtime_versions"): + if snapshot[key] != export_automodel.get(key): + raise RuntimeError(f"AutoModel environment/export identity mismatch for {key}.") + for key in ("package_tree_sha256", "wheel_sha256"): + digest = require_sha256(snapshot[key], name=f"AutoModel snapshot {key}") + if digest != export_automodel.get(key): + raise RuntimeError(f"AutoModel environment/export identity mismatch for {key}.") + root = Path(snapshot["root"]) + import_origin = Path(snapshot["import_origin"]) + if not root.is_absolute() or not import_origin.is_absolute(): + raise ValueError("AutoModel snapshot root and import_origin must be absolute.") + try: + import_origin.relative_to(root) + except ValueError as error: + raise RuntimeError( + "AutoModel import origin is outside its installed distribution." + ) from error + files = snapshot["files"] + if ( + not isinstance(files, list) + or type(snapshot["package_file_count"]) is not int + or snapshot["package_file_count"] != len(files) + or not files + ): + raise ValueError("AutoModel snapshot file inventory is malformed.") + tree = hashlib.sha256() + previous = None + for index, raw_record in enumerate(files): + record = _exact_mapping( + raw_record, + {"path", "sha256", "size"}, + name=f"AutoModel files[{index}]", + ) + path = record["path"] + if ( + not isinstance(path, str) + or not path + or Path(path).is_absolute() + or "\\" in path + or any(part in ("", ".", "..") for part in path.split("/")) + or (previous is not None and path <= previous) + ): + raise ValueError("AutoModel snapshot file paths must be sorted normalized references.") + digest = require_sha256(record["sha256"], name=f"AutoModel files[{index}].sha256") + if type(record["size"]) is not int or record["size"] < 0: + raise ValueError(f"AutoModel files[{index}].size is invalid.") + tree.update(path.encode()) + tree.update(b"\0") + tree.update(digest.encode()) + tree.update(b"\0") + tree.update(str(record["size"]).encode()) + tree.update(b"\n") + previous = path + if tree.hexdigest() != snapshot["package_tree_sha256"]: + raise RuntimeError("AutoModel snapshot file inventory does not match its tree SHA-256.") + + +def _validate_evaluation_protocol(manifest: Mapping[str, Any]) -> None: + if manifest["conditions"] != list(EVALUATION_CONDITIONS): + raise ValueError("effectiveness manifest must contain the six fixed conditions in order.") + if manifest["condition_protocols"] != CONDITION_PROTOCOLS: + raise ValueError("condition_protocols must match the frozen Qwen PDD protocol exactly.") + if manifest["grid_protocols"] != GRID_PROTOCOLS: + raise ValueError("grid_protocols must contain the exact authenticated shifted-flow nodes.") + if manifest["integrator_protocols"] != INTEGRATOR_PROTOCOLS: + raise ValueError("integrator_protocols must contain the exact frozen update equations.") + stage_run_ids = _exact_mapping( + manifest["stage_run_ids"], {"canary", "training"}, name="stage_run_ids" + ) + if any(not isinstance(run_id, str) or not run_id for run_id in stage_run_ids.values()): + raise ValueError("stage_run_ids must contain non-empty run IDs.") + image = _exact_mapping( + manifest["image_protocol"], + {"height", "width", "batch_size", "max_sequence_length"}, + name="image_protocol", + ) + if image != {"height": 1024, "width": 1024, "batch_size": 1, "max_sequence_length": 512}: + raise ValueError("image_protocol must use the frozen 1024px single-image Qwen protocol.") + metrics = manifest["metric_protocols"] + if not isinstance(metrics, Mapping) or not metrics or list(metrics) != sorted(metrics): + raise ValueError("metric_protocols must be a non-empty, sorted mapping.") + for name, raw_protocol in metrics.items(): + if not isinstance(name, str) or not name: + raise ValueError("metric protocol names must be non-empty strings.") + protocol = _exact_mapping( + raw_protocol, + {"direction", "implementation", "revision"}, + name=f"metric_protocols.{name}", + ) + if protocol["direction"] not in ("higher", "lower"): + raise ValueError(f"metric_protocols.{name}.direction must be higher or lower.") + if not isinstance(protocol["implementation"], str) or not protocol["implementation"]: + raise ValueError(f"metric_protocols.{name}.implementation must be non-empty.") + _commit(protocol["revision"], name=f"metric_protocols.{name}.revision") + timing = _exact_mapping( + manifest["timing_protocol"], + {"batch_size", "warmup_runs", "measured_runs", "scope", "synchronize_device"}, + name="timing_protocol", + ) + if ( + timing["batch_size"] != 1 + or type(timing["warmup_runs"]) is not int + or timing["warmup_runs"] < 1 + or type(timing["measured_runs"]) is not int + or timing["measured_runs"] < 3 + or timing["scope"] != "transformer_sampling_and_vae_decode" + or timing["synchronize_device"] is not True + ): + raise ValueError("timing_protocol does not satisfy the frozen measurement contract.") + rule = _exact_mapping( + manifest["decision_rule"], + { + "primary_condition", + "primary_metric", + "quality_margin", + "quality_ci_rule", + "efficiency_measure", + "efficiency_baseline", + "minimum_relative_reduction", + "minimum_paired_samples", + }, + name="decision_rule", + ) + if rule["primary_condition"] not in ("pdd_2", "pdd_4", "pdd_8"): + raise ValueError("decision_rule.primary_condition must name a supported PDD schedule.") + if rule["primary_metric"] not in metrics: + raise ValueError("decision_rule.primary_metric is not in metric_protocols.") + if ( + isinstance(rule["quality_margin"], bool) + or not isinstance(rule["quality_margin"], int | float) + or not math.isfinite(float(rule["quality_margin"])) + or rule["quality_margin"] < 0 + or rule["quality_ci_rule"] != "paired_bootstrap_95_noninferiority" + ): + raise ValueError("decision_rule quality noninferiority contract is malformed.") + if ( + rule["efficiency_measure"] != "batch_normalized_transformer_evaluations" + or rule["efficiency_baseline"] != "teacher_guided" + or isinstance(rule["minimum_relative_reduction"], bool) + or not isinstance(rule["minimum_relative_reduction"], int | float) + or not 0 < float(rule["minimum_relative_reduction"]) < 1 + or type(rule["minimum_paired_samples"]) is not int + or rule["minimum_paired_samples"] < 16 + ): + raise ValueError("decision_rule efficiency/sample contract is malformed.") + + +def validate_effectiveness_bundle(manifest_path: str | Path) -> dict[str, Any]: + """Authenticate a complete, paired, claim-bearing effectiveness bundle.""" + unresolved_manifest = Path(manifest_path) + if unresolved_manifest.is_symlink(): + raise RuntimeError(f"effectiveness manifest cannot be a symlink: {unresolved_manifest}.") + manifest_path = unresolved_manifest.resolve() + root = manifest_path.parent + detached = manifest_path.with_suffix(manifest_path.suffix + ".sha256") + if not detached.is_file() or detached.is_symlink(): + raise RuntimeError(f"detached manifest SHA-256 is missing: {detached}.") + detached_bytes = detached.read_bytes() + expected_bytes = (sha256_file(manifest_path) + "\n").encode() + if detached_bytes != expected_bytes: + raise RuntimeError("detached effectiveness manifest SHA-256 does not match.") + manifest = _exact_mapping( + load_canonical_json(manifest_path), + { + "schema_version", + "stage", + "run_id", + "model", + "modelopt", + "pdd_export", + "prompt_set", + "observations", + "environment", + "data_snapshot", + "stage_evidence", + "conditions", + "condition_protocols", + "grid_protocols", + "integrator_protocols", + "image_protocol", + "metric_protocols", + "timing_protocol", + "decision_rule", + "negative_condition", + "stage_run_ids", + "bootstrap", + }, + name="effectiveness manifest", + ) + if manifest["schema_version"] != 1 or manifest["stage"] != "effectiveness_evaluation": + raise ValueError("only schema-v1 effectiveness_evaluation manifests can support claims.") + if not isinstance(manifest["run_id"], str) or not manifest["run_id"]: + raise ValueError("effectiveness run_id must be non-empty.") + model = _exact_mapping(manifest["model"], {"id", "revision"}, name="model") + if not isinstance(model["id"], str) or not model["id"]: + raise ValueError("model.id must be non-empty.") + _commit(model["revision"], name="model.revision") + modelopt = _exact_mapping(manifest["modelopt"], {"commit", "dirty"}, name="modelopt") + _commit(modelopt["commit"], name="modelopt.commit") + if modelopt["dirty"] is not False: + raise RuntimeError("claim-bearing effectiveness runs require a clean ModelOpt commit.") + _validate_evaluation_protocol(manifest) + bootstrap = _exact_mapping(manifest["bootstrap"], {"replicates", "seed"}, name="bootstrap") + if type(bootstrap["replicates"]) is not int or bootstrap["replicates"] < 1_000: + raise ValueError("bootstrap.replicates must be an integer >= 1000.") + _nonnegative_int(bootstrap["seed"], name="bootstrap.seed") + export_path = validate_artifact_reference(root, manifest["pdd_export"], name="pdd_export") + if export_path.name != "manifest.json": + raise ValueError("pdd_export must reference the export directory's manifest.json.") + export_descriptor = inspect_pdd_export(export_path.parent) + if export_descriptor.root / "manifest.json" != export_path: + raise RuntimeError("pdd_export does not identify the authenticated export manifest.") + environment_path = validate_artifact_reference( + root, manifest["environment"], name="environment" + ) + export_document = export_descriptor.manifest + environment_document = load_canonical_json(environment_path) + if not isinstance(export_document, Mapping): + raise ValueError("pdd_export must reference a canonical JSON object.") + if not isinstance(environment_document, Mapping): + raise ValueError("environment must reference a canonical JSON object.") + export_identity = export_document.get("identity") + export_modelopt = export_document.get("modelopt_source") + if ( + not isinstance(export_identity, Mapping) + or export_document.get("format") != "modelopt-pdd-safetensors" + ): + raise ValueError("pdd_export does not reference a ModelOpt PDD export manifest.") + export_model = export_identity.get("model") + if not isinstance(export_model, Mapping) or { + "id": export_model.get("id"), + "revision": export_model.get("revision"), + } != dict(model): + raise RuntimeError("effectiveness model does not match the PDD export identity.") + if not isinstance(export_modelopt, Mapping) or export_modelopt != modelopt: + raise RuntimeError("effectiveness ModelOpt source does not match the PDD export.") + export_automodel = export_identity.get("automodel") + if not isinstance(export_automodel, Mapping): + raise ValueError("PDD export has no AutoModel identity.") + export_guidance = export_identity.get("guidance") + if export_guidance != {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}: + raise RuntimeError("PDD export guidance does not match the frozen evaluation protocol.") + _validate_automodel_snapshot(environment_document, export_automodel=export_automodel) + data_snapshot_path, data_snapshot = _validate_data_snapshot(root, manifest["data_snapshot"]) + data_snapshot_sha256 = sha256_file(data_snapshot_path) + export_data = _exact_mapping( + export_identity.get("data"), + { + "ordered_train_id_sha256", + "ordered_heldout_id_sha256", + "dataset_snapshot_sha256", + "local_batch_size", + "grad_accumulation_steps", + }, + name="PDD export data identity", + ) + if { + "dataset_snapshot_sha256": export_data["dataset_snapshot_sha256"], + "train_ids_sha256": export_data["ordered_train_id_sha256"], + "heldout_ids_sha256": export_data["ordered_heldout_id_sha256"], + } != { + name: data_snapshot[name] + for name in ("dataset_snapshot_sha256", "train_ids_sha256", "heldout_ids_sha256") + }: + raise RuntimeError("PDD export training-data identity does not match data_snapshot.") + negative_path = validate_artifact_reference( + root, manifest["negative_condition"], name="negative_condition" + ) + negative = _exact_mapping( + load_canonical_json(negative_path), + {"schema_version", "record_type", "prompt_sha256", "embedding"}, + name="negative condition", + ) + if negative["schema_version"] != 1 or negative["record_type"] != "pdd_negative_condition": + raise ValueError("negative_condition must be a schema-v1 PDD negative condition.") + require_sha256(negative["prompt_sha256"], name="negative condition prompt SHA-256") + validate_artifact_reference(root, negative["embedding"], name="negative condition embedding") + prompt_path = validate_artifact_reference(root, manifest["prompt_set"], name="prompt_set") + observation_path = validate_artifact_reference( + root, manifest["observations"], name="observations" + ) + evidence = _exact_mapping( + manifest["stage_evidence"], {"canary", "training"}, name="stage_evidence" + ) + protocol_sha256 = _protocol_sha256(manifest) + export_source_checkpoint = export_document["source_checkpoint"] + for stage in ("canary", "training"): + _validate_stage_evidence( + root, + evidence[stage], + stage=stage, + expected_run_id=manifest["stage_run_ids"][stage], + protocol_sha256=protocol_sha256, + model=model, + modelopt=modelopt, + data_snapshot_sha256=data_snapshot_sha256, + export_source_checkpoint=export_source_checkpoint, + ) + prompt_pairs, prompt_set = _load_prompt_set(prompt_path) + records = _validate_observations( + root, + observation_path, + prompt_pairs=prompt_pairs, + metric_protocols=manifest["metric_protocols"], + image_protocol=manifest["image_protocol"], + timing_protocol=manifest["timing_protocol"], + export_manifest_sha256=sha256_file(export_path), + evaluation_protocol_sha256=protocol_sha256, + ) + if len(prompt_pairs) < manifest["decision_rule"]["minimum_paired_samples"]: + raise RuntimeError("paired sample count is below decision_rule.minimum_paired_samples.") + return { + "root": root, + "manifest": manifest, + "manifest_sha256": sha256_file(manifest_path), + "prompt_set": prompt_set, + "records": records, + } + + +def deterministic_bootstrap_mean_ci( + values: Sequence[float], + *, + replicates: int, + seed: int, +) -> tuple[float, float]: + """Return a deterministic SHA-256-index percentile interval for a sample mean.""" + if not values: + raise ValueError("bootstrap values must be non-empty.") + if replicates < 1: + raise ValueError("bootstrap replicates must be positive.") + samples: list[float] = [] + for replicate in range(replicates): + total = 0.0 + for draw in range(len(values)): + payload = f"{seed}:{replicate}:{draw}".encode() + index = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big") % len(values) + total += float(values[index]) + samples.append(total / len(values)) + samples.sort() + lower = samples[math.floor(0.025 * (replicates - 1))] + upper = samples[math.ceil(0.975 * (replicates - 1))] + return lower, upper + + +def summarize_effectiveness_bundle(validated: Mapping[str, Any]) -> dict[str, Any]: + """Compute paired aggregate metrics without promoting smoke output to evidence.""" + manifest = validated["manifest"] + records = validated["records"] + replicates = manifest["bootstrap"]["replicates"] + seed = manifest["bootstrap"]["seed"] + by_condition = { + condition: { + (record["prompt_id"], record["prompt_sha256"], record["seed"]): record + for record in records + if record["condition"] == condition + } + for condition in EVALUATION_CONDITIONS + } + pair_keys = sorted(by_condition["teacher_guided"]) + aggregates: dict[str, Any] = {} + for condition_index, condition in enumerate(EVALUATION_CONDITIONS): + condition_records = by_condition[condition] + metrics: dict[str, Any] = {} + for metric_index, (metric, protocol) in enumerate(manifest["metric_protocols"].items()): + direction = protocol["direction"] + values = [float(condition_records[key]["metrics"][metric]) for key in pair_keys] + teacher = [ + float(by_condition["teacher_guided"][key]["metrics"][metric]) for key in pair_keys + ] + deltas = [value - baseline for value, baseline in zip(values, teacher)] + metric_seed = seed + condition_index * 10_000 + metric_index * 2 + metrics[metric] = { + "direction": direction, + "mean": sum(values) / len(values), + "mean_ci95": list( + deterministic_bootstrap_mean_ci( + values, + replicates=replicates, + seed=metric_seed, + ) + ), + "paired_delta_vs_teacher": sum(deltas) / len(deltas), + "paired_delta_ci95": list( + deterministic_bootstrap_mean_ci( + deltas, + replicates=replicates, + seed=metric_seed + 1, + ) + ), + } + aggregates[condition] = { + "metrics": metrics, + "latency_seconds": {}, + "throughput_images_per_second": {}, + "peak_device_memory_bytes": {}, + "mean_batch_normalized_transformer_evaluations": sum( + condition_records[key]["batch_normalized_transformer_evaluations"] + for key in pair_keys + ) + / len(pair_keys), + } + latencies = [float(condition_records[key]["latency_seconds"]) for key in pair_keys] + teacher_latencies = [ + float(by_condition["teacher_guided"][key]["latency_seconds"]) for key in pair_keys + ] + latency_deltas = [value - baseline for value, baseline in zip(latencies, teacher_latencies)] + latency_seed = seed + condition_index * 10_000 + len(metrics) * 2 + aggregates[condition]["latency_seconds"] = { + "mean": sum(latencies) / len(latencies), + "mean_ci95": list( + deterministic_bootstrap_mean_ci( + latencies, + replicates=replicates, + seed=latency_seed, + ) + ), + "paired_delta_vs_teacher": sum(latency_deltas) / len(latency_deltas), + "paired_delta_ci95": list( + deterministic_bootstrap_mean_ci( + latency_deltas, + replicates=replicates, + seed=latency_seed + 1, + ) + ), + } + for telemetry_index, name in enumerate( + ("throughput_images_per_second", "peak_device_memory_bytes"), start=1 + ): + values = [float(condition_records[key][name]) for key in pair_keys] + teacher = [float(by_condition["teacher_guided"][key][name]) for key in pair_keys] + deltas = [value - baseline for value, baseline in zip(values, teacher)] + telemetry_seed = latency_seed + telemetry_index * 2 + aggregates[condition][name] = { + "mean": sum(values) / len(values), + "mean_ci95": list( + deterministic_bootstrap_mean_ci( + values, + replicates=replicates, + seed=telemetry_seed, + ) + ), + "paired_delta_vs_teacher": sum(deltas) / len(deltas), + "paired_delta_ci95": list( + deterministic_bootstrap_mean_ci( + deltas, + replicates=replicates, + seed=telemetry_seed + 1, + ) + ), + } + + rule = manifest["decision_rule"] + primary = aggregates[rule["primary_condition"]] + primary_metric = primary["metrics"][rule["primary_metric"]] + lower, upper = primary_metric["paired_delta_ci95"] + margin = float(rule["quality_margin"]) + direction = primary_metric["direction"] + if direction == "higher": + quality_state = ( + "pass" if lower >= -margin else "fail" if upper < -margin else "inconclusive" + ) + else: + quality_state = "pass" if upper <= margin else "fail" if lower > margin else "inconclusive" + baseline = aggregates[rule["efficiency_baseline"]] + candidate_value = primary["mean_batch_normalized_transformer_evaluations"] + baseline_value = baseline["mean_batch_normalized_transformer_evaluations"] + relative_reduction = 1.0 - candidate_value / baseline_value + efficiency_state = ( + "pass" if relative_reduction >= float(rule["minimum_relative_reduction"]) else "fail" + ) + if "fail" in (quality_state, efficiency_state): + conclusion = "not_effective" + elif (quality_state, efficiency_state) == ("pass", "pass"): + conclusion = "effective" + else: + conclusion = "inconclusive" + return { + "schema_version": 1, + "record_type": "effectiveness_summary", + "source_manifest_sha256": validated["manifest_sha256"], + "paired_sample_count": len(pair_keys), + "bootstrap": dict(manifest["bootstrap"]), + "aggregates": aggregates, + "decision": { + "label": conclusion, + "quality_state": quality_state, + "efficiency_state": efficiency_state, + "relative_efficiency_reduction": relative_reduction, + "rule": dict(rule), + }, + } diff --git a/examples/diffusers/fastgen/pdd_export.py b/examples/diffusers/fastgen/pdd_export.py new file mode 100644 index 00000000000..b62863fb92c --- /dev/null +++ b/examples/diffusers/fastgen/pdd_export.py @@ -0,0 +1,607 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Authenticated, bounded safetensors export and strict PDD reconstruction helpers.""" + +from __future__ import annotations + +import os +import re +import shutil +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from pdd_artifacts import ( + load_canonical_json, + require_sha256, + resolve_relative_artifact, + sha256_file, + write_canonical_json, +) +from safetensors import safe_open +from safetensors.torch import save_file + +from modelopt.torch.fastgen import PDDConfig, PDDMetadata +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC + +_EXPORT_SCHEMA_VERSION = 1 +_COMPLETE_SCHEMA_VERSION = 1 +_EXPORT_FORMAT = "modelopt-pdd-safetensors" +_CONFIG_FILE = "config.json" +_METADATA_FILE = "pdd_metadata.json" +_INDEX_FILE = "diffusion_pytorch_model.safetensors.index.json" +_MANIFEST_FILE = "manifest.json" +_COMPLETE_FILE = "COMPLETE" +_SHARD_PATTERN = "diffusion_pytorch_model-{index:05d}-of-{count:05d}.safetensors" + +PDD_INFERENCE_SCHEDULES: Mapping[str, tuple[int, ...]] = { + "pdd-2": (64, 64), + "pdd-4": (32, 32, 32, 32), + "pdd-8": (16, 16, 16, 16, 16, 16, 16, 16), +} + + +@dataclass(frozen=True) +class PDDExportDescriptor: + """Validated non-tensor export metadata.""" + + root: Path + manifest: Mapping[str, Any] + metadata: PDDMetadata + transformer_config: Mapping[str, Any] + + +def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _fsync_tree(root: Path) -> None: + for path in sorted(root.rglob("*"), key=lambda item: len(item.parts), reverse=True): + if path.is_symlink(): + raise RuntimeError(f"PDD export staging contains a symlink: {path}.") + if path.is_file(): + with path.open("rb") as stream: + os.fsync(stream.fileno()) + elif path.is_dir(): + _fsync_directory(path) + _fsync_directory(root) + + +def _tensor_nbytes(tensor: torch.Tensor) -> int: + return tensor.numel() * tensor.element_size() + + +def _validate_state_dict(state_dict: Mapping[str, Any]) -> dict[str, torch.Tensor]: + if not isinstance(state_dict, Mapping) or not state_dict: + raise ValueError("PDD export state_dict must be a non-empty mapping.") + tensors: dict[str, torch.Tensor] = {} + for key in sorted(state_dict): + tensor = state_dict[key] + if not isinstance(key, str) or not key: + raise ValueError("PDD export tensor keys must be non-empty strings.") + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"PDD export value {key!r} is not a tensor.") + if tensor.device.type != "cpu" or tensor.is_meta: + raise ValueError(f"PDD export tensor {key!r} must be a materialized CPU tensor.") + if tensor.layout != torch.strided or tensor.is_quantized: + raise TypeError(f"PDD export tensor {key!r} must be a dense, unquantized tensor.") + if tensor.dtype.is_floating_point and not torch.isfinite(tensor).all().item(): + raise FloatingPointError(f"PDD export tensor {key!r} is non-finite.") + tensors[key] = tensor.detach().contiguous() + return tensors + + +def _save_probe(staging: Path, keys: Sequence[str], tensors: Mapping[str, torch.Tensor]) -> int: + probe = staging / f".probe-{uuid.uuid4().hex}.safetensors" + payload = {key: tensors[key].clone() for key in keys} + try: + save_file(payload, str(probe), metadata={"format": "pt"}) + return probe.stat().st_size + finally: + probe.unlink(missing_ok=True) + + +def _bounded_shard_groups( + staging: Path, + tensors: Mapping[str, torch.Tensor], + max_shard_bytes: int, +) -> list[tuple[str, ...]]: + if type(max_shard_bytes) is not int or max_shard_bytes <= 0: + raise ValueError("max_shard_bytes must be a positive integer.") + initial: list[tuple[str, ...]] = [] + current: list[str] = [] + current_bytes = 0 + for key, tensor in tensors.items(): + size = _tensor_nbytes(tensor) + if size >= max_shard_bytes: + raise ValueError( + f"tensor {key!r} has {size} bytes and cannot fit beneath the physical " + f"shard bound {max_shard_bytes}." + ) + if current and current_bytes + size >= max_shard_bytes: + initial.append(tuple(current)) + current = [] + current_bytes = 0 + current.append(key) + current_bytes += size + if current: + initial.append(tuple(current)) + + bounded: list[tuple[str, ...]] = [] + pending = list(initial) + while pending: + keys = pending.pop(0) + if _save_probe(staging, keys, tensors) <= max_shard_bytes: + bounded.append(keys) + continue + if len(keys) == 1: + raise ValueError( + f"tensor {keys[0]!r} plus safetensors metadata exceeds max_shard_bytes." + ) + midpoint = len(keys) // 2 + pending[0:0] = [keys[:midpoint], keys[midpoint:]] + return bounded + + +def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> dict[str, Any]: + if not isinstance(identity, Mapping): + raise TypeError("PDD export identity must be a mapping.") + required = {"automodel", "guidance", "model", "pdd_metadata", "topology"} + missing = sorted(required.difference(identity)) + if missing: + raise ValueError(f"PDD export identity is missing keys: {missing}.") + if identity["pdd_metadata"] != metadata.to_dict(): + raise ValueError("PDD export metadata does not match the checkpoint identity.") + model = _require_exact_mapping( + identity["model"], {"id", "revision", "dtype"}, name="identity.model" + ) + if not isinstance(model["id"], str) or not model["id"] or not isinstance(model["dtype"], str): + raise ValueError("PDD export checkpoint identity has an invalid model ID or dtype.") + revision = model["revision"] + if not isinstance(revision, str) or len(revision) != 40: + raise ValueError("PDD export requires a pinned 40-character model revision.") + try: + int(revision, 16) + except ValueError as error: + raise ValueError("PDD export model revision must be hexadecimal.") from error + automodel = _require_exact_mapping( + identity["automodel"], + {"distribution", "version", "package_tree_sha256", "wheel_sha256", "runtime_versions"}, + name="identity.automodel", + ) + if ( + not isinstance(automodel["distribution"], str) + or not automodel["distribution"] + or not isinstance(automodel["version"], str) + or not isinstance(automodel["runtime_versions"], Mapping) + ): + raise ValueError("PDD export AutoModel identity is malformed.") + require_sha256(automodel["package_tree_sha256"], name="AutoModel package tree SHA-256") + require_sha256(automodel["wheel_sha256"], name="AutoModel wheel SHA-256") + guidance = _require_exact_mapping( + identity["guidance"], {"scale", "rescale", "eps"}, name="identity.guidance" + ) + for name, value in guidance.items(): + if value is not None and ( + isinstance(value, bool) + or not isinstance(value, int | float) + or not torch.isfinite(torch.tensor(float(value))).item() + ): + raise ValueError(f"PDD export guidance {name} must be finite or null.") + topology = identity["topology"] + if ( + not isinstance(topology, Mapping) + or type(topology.get("world_size")) is not int + or topology["world_size"] < 1 + or topology.get("pure_data_parallel") is not True + ): + raise ValueError("PDD export checkpoint identity has invalid pure-DP topology.") + return dict(identity) + + +def _validate_modelopt_source(source: Mapping[str, Any]) -> dict[str, Any]: + source = _require_exact_mapping(source, {"commit", "dirty"}, name="modelopt_source") + commit = source["commit"] + if not isinstance(commit, str) or len(commit) != 40: + raise ValueError("modelopt_source.commit must be a 40-character Git commit.") + try: + int(commit, 16) + except ValueError as error: + raise ValueError("modelopt_source.commit must be hexadecimal.") from error + if source["dirty"] is not False: + raise ValueError("modelopt_source.dirty must be false.") + return dict(source) + + +def write_pdd_export( + output_dir: str | Path, + state_dict: Mapping[str, Any], + *, + metadata: PDDMetadata, + transformer_config: Mapping[str, Any], + identity: Mapping[str, Any], + source_checkpoint: Mapping[str, Any], + modelopt_source: Mapping[str, Any], + max_shard_bytes: int, +) -> Path: + """Publish a complete PDD export into a previously absent directory.""" + if not isinstance(metadata, PDDMetadata): + raise TypeError("metadata must be PDDMetadata.") + if metadata.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: + raise ValueError("PDD export supports only the fixed Qwen-Image layer specification.") + if not isinstance(transformer_config, Mapping): + raise TypeError("transformer_config must be a mapping.") + checkpoint_keys = {"name", "manifest_sha256", "completed_steps"} + if not isinstance(source_checkpoint, Mapping) or set(source_checkpoint) != checkpoint_keys: + raise ValueError(f"source_checkpoint must contain exactly {sorted(checkpoint_keys)}.") + if ( + not isinstance(source_checkpoint["name"], str) + or not source_checkpoint["name"] + or Path(source_checkpoint["name"]).name != source_checkpoint["name"] + ): + raise ValueError("source_checkpoint.name must be a basename.") + require_sha256(source_checkpoint["manifest_sha256"], name="source manifest SHA-256") + if ( + type(source_checkpoint["completed_steps"]) is not int + or source_checkpoint["completed_steps"] < 1 + ): + raise ValueError("source_checkpoint.completed_steps must be an integer >= 1.") + resolved_modelopt_source = _validate_modelopt_source(modelopt_source) + resolved_identity = _validate_identity(identity, metadata) + tensors = _validate_state_dict(state_dict) + + unresolved_output = Path(output_dir) + if unresolved_output.is_symlink(): + raise ValueError("PDD export output cannot be a symlink.") + output = unresolved_output.resolve() + output.parent.mkdir(parents=True, exist_ok=True) + if output.exists() or output.is_symlink(): + raise FileExistsError(f"PDD export output already exists: {output}.") + staging = output.with_name(f".{output.name}.{uuid.uuid4().hex}.staging") + staging.mkdir() + published = False + try: + groups = _bounded_shard_groups(staging, tensors, max_shard_bytes) + shard_names = [ + _SHARD_PATTERN.format(index=index, count=len(groups)) + for index in range(1, len(groups) + 1) + ] + weight_map: dict[str, str] = {} + for name, keys in zip(shard_names, groups): + payload = {key: tensors[key].clone() for key in keys} + save_file(payload, str(staging / name), metadata={"format": "pt"}) + if (staging / name).stat().st_size > max_shard_bytes: + raise RuntimeError(f"PDD safetensors shard exceeds its physical bound: {name}.") + weight_map.update(dict.fromkeys(keys, name)) + + total_tensor_bytes = sum(_tensor_nbytes(tensor) for tensor in tensors.values()) + write_canonical_json(staging / _CONFIG_FILE, dict(transformer_config)) + write_canonical_json(staging / _METADATA_FILE, metadata.to_dict()) + write_canonical_json( + staging / _INDEX_FILE, + {"metadata": {"total_size": total_tensor_bytes}, "weight_map": weight_map}, + ) + file_names = [_CONFIG_FILE, _METADATA_FILE, _INDEX_FILE, *shard_names] + files = { + name: { + "sha256": sha256_file(staging / name), + "size": (staging / name).stat().st_size, + } + for name in file_names + } + tensor_specs = { + key: { + "dtype": str(tensor.dtype).removeprefix("torch."), + "nbytes": _tensor_nbytes(tensor), + "shape": list(tensor.shape), + "shard": weight_map[key], + } + for key, tensor in tensors.items() + } + manifest = { + "schema_version": _EXPORT_SCHEMA_VERSION, + "format": _EXPORT_FORMAT, + "identity": resolved_identity, + "source_checkpoint": dict(source_checkpoint), + "modelopt_source": resolved_modelopt_source, + "max_shard_bytes": max_shard_bytes, + "total_tensor_bytes": total_tensor_bytes, + "tensors": tensor_specs, + "files": files, + } + write_canonical_json(staging / _MANIFEST_FILE, manifest) + write_canonical_json( + staging / _COMPLETE_FILE, + { + "schema_version": _COMPLETE_SCHEMA_VERSION, + "manifest_sha256": sha256_file(staging / _MANIFEST_FILE), + }, + ) + _fsync_tree(staging) + inspect_pdd_export(staging) + staging.rename(output) + published = True + _fsync_directory(output.parent) + return output + except BaseException: + target = output if published else staging + if target.exists() and not target.is_symlink(): + shutil.rmtree(target) + raise + + +def _require_exact_mapping(value: Any, keys: set[str], *, name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or set(value) != keys: + actual = sorted(value) if isinstance(value, Mapping) else type(value).__name__ + raise ValueError(f"{name} keys mismatch: expected={sorted(keys)}, actual={actual}.") + return value + + +def inspect_pdd_export(export_dir: str | Path) -> PDDExportDescriptor: + """Authenticate a PDD export without loading its tensor payloads.""" + unresolved_root = Path(export_dir) + if unresolved_root.is_symlink(): + raise RuntimeError(f"PDD export cannot be a symlink: {unresolved_root}.") + root = unresolved_root.resolve() + if not root.is_dir(): + raise RuntimeError(f"PDD export is not a regular directory: {root}.") + for path in root.rglob("*"): + if path.is_symlink(): + raise RuntimeError(f"PDD export contains a symlink: {path}.") + if path.is_dir(): + raise RuntimeError(f"PDD export must be flat, found directory: {path}.") + + complete = _require_exact_mapping( + load_canonical_json(root / _COMPLETE_FILE), + {"schema_version", "manifest_sha256"}, + name="PDD COMPLETE", + ) + if complete["schema_version"] != _COMPLETE_SCHEMA_VERSION: + raise ValueError("PDD COMPLETE schema version is unsupported.") + if require_sha256( + complete["manifest_sha256"], name="PDD COMPLETE manifest SHA-256" + ) != sha256_file(root / _MANIFEST_FILE): + raise RuntimeError("PDD COMPLETE does not match the export manifest.") + manifest = _require_exact_mapping( + load_canonical_json(root / _MANIFEST_FILE), + { + "schema_version", + "format", + "identity", + "source_checkpoint", + "modelopt_source", + "max_shard_bytes", + "total_tensor_bytes", + "tensors", + "files", + }, + name="PDD export manifest", + ) + if manifest["schema_version"] != _EXPORT_SCHEMA_VERSION or manifest["format"] != _EXPORT_FORMAT: + raise ValueError("PDD export manifest schema or format is unsupported.") + _validate_modelopt_source(manifest["modelopt_source"]) + if type(manifest["max_shard_bytes"]) is not int or manifest["max_shard_bytes"] <= 0: + raise ValueError("PDD export max_shard_bytes is invalid.") + if type(manifest["total_tensor_bytes"]) is not int or manifest["total_tensor_bytes"] <= 0: + raise ValueError("PDD export total_tensor_bytes is invalid.") + source = _require_exact_mapping( + manifest["source_checkpoint"], + {"name", "manifest_sha256", "completed_steps"}, + name="source_checkpoint", + ) + if ( + not isinstance(source["name"], str) + or not source["name"] + or Path(source["name"]).name != source["name"] + ): + raise ValueError("source_checkpoint.name must be a basename.") + require_sha256(source["manifest_sha256"], name="source checkpoint manifest SHA-256") + if type(source["completed_steps"]) is not int or source["completed_steps"] < 1: + raise ValueError("source_checkpoint.completed_steps is invalid.") + + files = manifest["files"] + if not isinstance(files, Mapping) or not files: + raise ValueError("PDD export file inventory must be a non-empty mapping.") + expected_names = set(files) | {_MANIFEST_FILE, _COMPLETE_FILE} + actual_names = {path.name for path in root.iterdir() if path.is_file()} + if actual_names != expected_names: + raise RuntimeError( + f"PDD export file inventory mismatch: expected={sorted(expected_names)}, " + f"actual={sorted(actual_names)}." + ) + for name, record in files.items(): + if Path(name).name != name: + raise ValueError(f"PDD export file name must be a basename: {name!r}.") + record = _require_exact_mapping(record, {"sha256", "size"}, name=f"files[{name!r}]") + path = resolve_relative_artifact(root, name) + if type(record["size"]) is not int or record["size"] < 0: + raise ValueError(f"PDD export file size is invalid for {name!r}.") + if path.stat().st_size != record["size"]: + raise RuntimeError(f"PDD export file size mismatch for {name!r}.") + if sha256_file(path) != require_sha256(record["sha256"], name=f"files[{name!r}].sha256"): + raise RuntimeError(f"PDD export file SHA-256 mismatch for {name!r}.") + if name.endswith(".safetensors") and record["size"] > manifest["max_shard_bytes"]: + raise RuntimeError(f"PDD export shard exceeds max_shard_bytes: {name!r}.") + mandatory = {_CONFIG_FILE, _METADATA_FILE, _INDEX_FILE} + if not mandatory.issubset(files): + raise RuntimeError( + f"PDD export is missing mandatory files: {sorted(mandatory - set(files))}." + ) + + metadata_data = load_canonical_json(root / _METADATA_FILE) + metadata = PDDMetadata.from_dict(metadata_data) + if metadata.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: + raise ValueError("PDD export carries a non-Qwen layer specification.") + identity = manifest["identity"] + _validate_identity(identity, metadata) + transformer_config = load_canonical_json(root / _CONFIG_FILE) + if not isinstance(transformer_config, Mapping): + raise ValueError("PDD transformer config must contain an object.") + + tensors = manifest["tensors"] + if not isinstance(tensors, Mapping) or not tensors: + raise ValueError("PDD export tensor inventory must be a non-empty mapping.") + index = _require_exact_mapping( + load_canonical_json(root / _INDEX_FILE), + {"metadata", "weight_map"}, + name="safetensors index", + ) + index_metadata = _require_exact_mapping( + index["metadata"], {"total_size"}, name="index metadata" + ) + if index_metadata["total_size"] != manifest["total_tensor_bytes"]: + raise RuntimeError("PDD safetensors index total size does not match the manifest.") + weight_map = index["weight_map"] + if not isinstance(weight_map, Mapping) or set(weight_map) != set(tensors): + raise RuntimeError("PDD safetensors index keys do not match the tensor inventory.") + shard_names = {name for name in files if name.endswith(".safetensors")} + shard_pattern = re.compile(r"diffusion_pytorch_model-(\d{5})-of-(\d{5})\.safetensors") + shard_numbers = [] + for name in shard_names: + match = shard_pattern.fullmatch(name) + if match is None: + raise ValueError(f"PDD export has a noncanonical shard name: {name!r}.") + shard_numbers.append((int(match.group(1)), int(match.group(2)))) + if not shard_numbers: + raise RuntimeError("PDD export has no safetensors shards.") + shard_count = len(shard_numbers) + if sorted(shard_numbers) != [(index, shard_count) for index in range(1, shard_count + 1)]: + raise RuntimeError("PDD export safetensors shard numbering is inconsistent.") + if set(weight_map.values()) != shard_names: + raise RuntimeError("PDD safetensors index shard inventory does not match export files.") + + total = 0 + for key, spec in tensors.items(): + if not isinstance(key, str) or not key: + raise ValueError("PDD tensor inventory keys must be non-empty strings.") + spec = _require_exact_mapping( + spec, + {"dtype", "nbytes", "shape", "shard"}, + name=f"tensors[{key!r}]", + ) + if ( + not isinstance(spec["dtype"], str) + or type(spec["nbytes"]) is not int + or spec["nbytes"] <= 0 + or not isinstance(spec["shape"], list) + or any(type(size) is not int or size < 0 for size in spec["shape"]) + or spec["shard"] != weight_map[key] + ): + raise ValueError(f"PDD tensor specification is malformed for {key!r}.") + total += spec["nbytes"] + if total != manifest["total_tensor_bytes"]: + raise RuntimeError("PDD tensor byte inventory does not match total_tensor_bytes.") + return PDDExportDescriptor(root, manifest, metadata, transformer_config) + + +def pdd_config_from_metadata( + metadata: PDDMetadata, + *, + blocks: Sequence[int] | None = None, + schedule: str | None = None, + guidance_scale: float | None = None, +) -> PDDConfig: + """Build a fresh validated inference config from authenticated metadata.""" + if (blocks is None) == (schedule is None): + raise ValueError("Specify exactly one of blocks or schedule.") + if schedule is not None: + try: + resolved = PDD_INFERENCE_SCHEDULES[schedule] + except KeyError as error: + raise ValueError( + f"Unknown PDD schedule {schedule!r}; expected {sorted(PDD_INFERENCE_SCHEDULES)}." + ) from error + else: + if isinstance(blocks, str | bytes) or not isinstance(blocks, Sequence): + raise TypeError("blocks must be a sequence of integers.") + resolved = tuple(blocks) + return PDDConfig( + grid_size=metadata.grid_size, + flow_shift=metadata.flow_shift, + block_size_min=metadata.block_size_min, + block_size_max=metadata.block_size_max, + inference_blocks=list(resolved), + student_sample_steps=len(resolved), + teacher_integrator=metadata.teacher_integrator, + guidance_scale=guidance_scale, + num_train_timesteps=None, + ) + + +def _load_shard( + descriptor: PDDExportDescriptor, + shard_name: str, +) -> dict[str, torch.Tensor]: + expected = { + key: spec + for key, spec in descriptor.manifest["tensors"].items() + if spec["shard"] == shard_name + } + path = descriptor.root / shard_name + loaded: dict[str, torch.Tensor] = {} + with safe_open(str(path), framework="pt", device="cpu") as stream: + keys = list(stream.keys()) + if len(keys) != len(set(keys)) or set(keys) != set(expected): + raise RuntimeError(f"PDD safetensors keys do not match the index for {shard_name!r}.") + for key in keys: + tensor = stream.get_tensor(key) + spec = expected[key] + if list(tensor.shape) != spec["shape"]: + raise RuntimeError(f"PDD tensor shape mismatch for {key!r}.") + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise RuntimeError(f"PDD tensor dtype mismatch for {key!r}.") + if _tensor_nbytes(tensor) != spec["nbytes"]: + raise RuntimeError(f"PDD tensor byte-size mismatch for {key!r}.") + if tensor.dtype.is_floating_point and not torch.isfinite(tensor).all().item(): + raise FloatingPointError(f"PDD tensor {key!r} is non-finite.") + loaded[key] = tensor + return loaded + + +def load_pdd_export_into_model( + export_dir: str | Path, + model: torch.nn.Module, +) -> PDDExportDescriptor: + """Strictly stream authenticated safetensors shards into a converted CPU model.""" + if not isinstance(model, torch.nn.Module): + raise TypeError("model must be an nn.Module.") + descriptor = inspect_pdd_export(export_dir) + expected_state = model.state_dict() + specs = descriptor.manifest["tensors"] + if set(expected_state) != set(specs): + missing = sorted(set(expected_state) - set(specs)) + extra = sorted(set(specs) - set(expected_state)) + raise RuntimeError(f"PDD model/export keys mismatch: missing={missing}, extra={extra}.") + for key, expected in expected_state.items(): + spec = specs[key] + if list(expected.shape) != spec["shape"]: + raise RuntimeError(f"PDD model/export shape mismatch for {key!r}.") + if str(expected.dtype).removeprefix("torch.") != spec["dtype"]: + raise RuntimeError(f"PDD model/export dtype mismatch for {key!r}.") + + shard_names = sorted({spec["shard"] for spec in specs.values()}) + loaded_keys: set[str] = set() + for shard_name in shard_names: + shard = _load_shard(descriptor, shard_name) + if loaded_keys.intersection(shard): + raise RuntimeError("PDD tensor appears in more than one safetensors shard.") + incompatible = model.load_state_dict(shard, strict=False) + if incompatible.unexpected_keys: + raise RuntimeError(f"PDD shard has unexpected keys: {incompatible.unexpected_keys}.") + loaded_keys.update(shard) + if loaded_keys != set(expected_state): + raise RuntimeError("PDD safe load did not assign every expected tensor.") + if any(parameter.is_meta for parameter in model.parameters()) or any( + buffer.is_meta for buffer in model.buffers() + ): + raise RuntimeError("PDD safe load left meta tensors in the reconstructed model.") + model.eval().requires_grad_(False) + return descriptor diff --git a/examples/diffusers/fastgen/pdd_recipe.py b/examples/diffusers/fastgen/pdd_recipe.py index 367ab9cbb93..1d292a62f80 100644 --- a/examples/diffusers/fastgen/pdd_recipe.py +++ b/examples/diffusers/fastgen/pdd_recipe.py @@ -121,6 +121,23 @@ class PDDTrainingArtifacts: rng: Any +@dataclass(frozen=True) +class PDDExportSetupArtifacts: + """Student-only FSDP2 objects needed for collective DCP export.""" + + pipe: Any + student: nn.Module + projection: PDDOutputProjection + distributed_setup: Any + fsdp_manager: Any + checkpointer: Any + metadata: PDDMetadata + checkpoint_keys: tuple[str, ...] + transformer_config: Mapping[str, Any] + lifecycle: tuple[str, ...] + automodel_snapshot: Mapping[str, Any] + + def _as_mapping(value: Any, *, name: str) -> Mapping[str, Any]: if not isinstance(value, Mapping): raise TypeError(f"{name} must be a mapping, got {type(value).__name__}.") @@ -455,41 +472,26 @@ def _require_projection_module( raise RuntimeError(f"PDD projection module was replaced during {stage}.") -def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: - """Compose released AutoModel APIs without editing or patching external packages.""" - if not isinstance(config, PDDRecipeConfig): - raise TypeError(f"config must be PDDRecipeConfig, got {type(config).__name__}.") - if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("Initialize torch.distributed before building the PDD FSDP2 setup.") - - automodel_snapshot = snapshot_installed_distribution() - lifecycle: list[str] = [] +def _resolve_model_source(config: PDDRecipeConfig) -> str: + if Path(config.model_id).is_dir(): + return config.model_id + from huggingface_hub import snapshot_download - # Imports are intentionally delayed until the exact installed wheel has passed verification. - from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline - from nemo_automodel.components.checkpoint.config import CheckpointingConfig - from nemo_automodel.components.distributed import ( - DistributedSetup, - FSDP2Config, - ParallelismSizes, - ) - from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager + if config.model_revision is None: + raise ValueError("Remote PDD models require a pinned model revision.") + model_source = snapshot_download(config.model_id, revision=config.model_revision) + if Path(model_source).resolve().name != config.model_revision: + raise RuntimeError( + "Hugging Face resolved a model snapshot that does not match the pinned revision." + ) + return model_source - if Path(config.model_id).is_dir(): - model_source = config.model_id - else: - from huggingface_hub import snapshot_download - - if config.model_revision is None: - raise ValueError("Remote PDD models require a pinned model revision.") - model_source = snapshot_download(config.model_id, revision=config.model_revision) - if Path(model_source).resolve().name != config.model_revision: - raise RuntimeError( - "Hugging Face resolved a model snapshot that does not match the pinned revision." - ) - pipe, loader_managers = NeMoAutoDiffusionPipeline.from_pretrained( - model_source, +def _load_unwrapped_transformer( + config: PDDRecipeConfig, pipeline_type: Any +) -> tuple[Any, nn.Module]: + pipe, loader_managers = pipeline_type.from_pretrained( + _resolve_model_source(config), parallel_scheme=None, device=None, torch_dtype=config.dtype, @@ -511,6 +513,30 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: student = pipe.transformer if not isinstance(student, nn.Module): raise TypeError("AutoModel pipeline did not return an nn.Module transformer.") + return pipe, student + + +def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: + """Compose released AutoModel APIs without editing or patching external packages.""" + if not isinstance(config, PDDRecipeConfig): + raise TypeError(f"config must be PDDRecipeConfig, got {type(config).__name__}.") + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("Initialize torch.distributed before building the PDD FSDP2 setup.") + + automodel_snapshot = snapshot_installed_distribution() + lifecycle: list[str] = [] + + # Imports are intentionally delayed until the exact installed wheel has passed verification. + from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.checkpoint.config import CheckpointingConfig + from nemo_automodel.components.distributed import ( + DistributedSetup, + FSDP2Config, + ParallelismSizes, + ) + from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager + + pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) teacher = copy.deepcopy(student).eval().requires_grad_(False) lifecycle.append("load/select") @@ -645,6 +671,107 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: ) +def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: + """Build only the converted/sharded student needed for collective export.""" + if not isinstance(config, PDDRecipeConfig): + raise TypeError(f"config must be PDDRecipeConfig, got {type(config).__name__}.") + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("Initialize torch.distributed before building PDD export setup.") + automodel_snapshot = snapshot_installed_distribution() + + from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.checkpoint.config import CheckpointingConfig + from nemo_automodel.components.distributed import ( + DistributedSetup, + FSDP2Config, + ParallelismSizes, + ) + from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager + + lifecycle = ["load/select"] + pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + raw_transformer_config = getattr(student, "config", None) + if hasattr(raw_transformer_config, "to_dict"): + transformer_config = raw_transformer_config.to_dict() + elif isinstance(raw_transformer_config, Mapping): + transformer_config = dict(raw_transformer_config) + else: + raise TypeError("Qwen transformer config must expose to_dict() or Mapping.") + + projection = convert_qwen_image_to_pdd(student, config.pdd) + identity = _projection_identity(projection) + metadata = PDDMetadata.from_config(config.pdd, projection) + lifecycle.append("pdd_conversion") + student.to(device=config.device, dtype=config.dtype) + _require_projection_identity(student, projection, identity, stage="device placement") + lifecycle.append("device") + + if config.fuse_qkv_projections: + if not hasattr(student, "fuse_qkv_projections"): + raise AttributeError("QKV fusion requires Qwen to expose the object API.") + student.fuse_qkv_projections() + _require_projection_identity(student, projection, identity, stage="QKV fusion") + lifecycle.append("qkv") + + world_size = dist.get_world_size() + dp_size = config.parallel.dp_size or world_size + if dp_size != world_size: + raise ValueError( + f"Pure-DP PDD export requires fsdp.dp_size ({dp_size}) to equal world size " + f"({world_size})." + ) + strategy = FSDP2Config(activation_checkpointing=False) + distributed_setup = DistributedSetup.build( + strategy=strategy, + parallelism_sizes=ParallelismSizes(dp_size=dp_size), + activation_checkpointing=False, + world_size=world_size, + ) + mesh_context = distributed_setup.mesh_context + manager = FSDP2Manager( + distributed_setup.strategy_config, + device_mesh=mesh_context.device_mesh, + moe_mesh=mesh_context.moe_mesh, + ) + student = manager.parallelize(student) + pipe.transformer = student + _require_projection_module(student, projection, stage="FSDP2 export parallelization") + lifecycle.append("parallelize") + + checkpoint_keys = tuple(student.state_dict()) + if "proj_out.weight" not in checkpoint_keys: + raise RuntimeError("PDD projection is missing from the export checkpoint key inventory.") + checkpoint_config = CheckpointingConfig( + enabled=True, + checkpoint_dir=config.checkpoint.checkpoint_dir, + model_save_format="torch_save", + model_repo_id=config.model_id, + save_consolidated=False, + is_peft=False, + model_state_dict_keys=list(checkpoint_keys), + ) + checkpointer = checkpoint_config.build( + dp_rank=dist.get_rank(), + tp_rank=0, + pp_rank=0, + moe_mesh=None, + ) + lifecycle.append("checkpoint") + return PDDExportSetupArtifacts( + pipe=pipe, + student=student, + projection=projection, + distributed_setup=distributed_setup, + fsdp_manager=manager, + checkpointer=checkpointer, + metadata=metadata, + checkpoint_keys=checkpoint_keys, + transformer_config=transformer_config, + lifecycle=tuple(lifecycle), + automodel_snapshot=automodel_snapshot, + ) + + def build_pdd_training_artifacts( setup: PDDSetupArtifacts, config: PDDRecipeConfig, diff --git a/examples/diffusers/fastgen/seal_pdd_run_manifest.py b/examples/diffusers/fastgen/seal_pdd_run_manifest.py new file mode 100644 index 00000000000..03eeb95b7bb --- /dev/null +++ b/examples/diffusers/fastgen/seal_pdd_run_manifest.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Create and verify the detached SHA-256 for a canonical PDD run manifest.""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +sys.dont_write_bytecode = True + +_THIS_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _THIS_DIR.parents[2] +for path in (_REPO_ROOT, _THIS_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("manifest", type=Path) + args = parser.parse_args() + from pdd_artifacts import load_canonical_json, sha256_file + from pdd_evaluation import validate_effectiveness_bundle + + manifest = args.manifest.resolve() + load_canonical_json(manifest) + detached = manifest.with_suffix(manifest.suffix + ".sha256") + with detached.open("xb") as stream: + stream.write((sha256_file(manifest) + "\n").encode()) + stream.flush() + os.fsync(stream.fileno()) + try: + validate_effectiveness_bundle(manifest) + except BaseException: + detached.unlink(missing_ok=True) + raise + print(detached) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/validate_pdd_run_manifest.py b/examples/diffusers/fastgen/validate_pdd_run_manifest.py new file mode 100644 index 00000000000..17631d27c5a --- /dev/null +++ b/examples/diffusers/fastgen/validate_pdd_run_manifest.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Authenticate a complete PDD effectiveness evidence bundle.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +sys.dont_write_bytecode = True + +_THIS_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _THIS_DIR.parents[2] +for path in (_REPO_ROOT, _THIS_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("manifest", type=Path) + args = parser.parse_args() + from pdd_evaluation import validate_effectiveness_bundle + + validated = validate_effectiveness_bundle(args.manifest) + print(validated["manifest_sha256"]) + + +if __name__ == "__main__": + main() diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py new file mode 100644 index 00000000000..516b4f7ac6f --- /dev/null +++ b/tests/examples/diffusers/fastgen/pdd_export_distributed.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Two-rank released-AutoModel DCP-to-full-state export proof for the PDD example.""" + +from __future__ import annotations + +import pathlib +import shutil +import sys +import tempfile + +import torch +import torch.distributed as dist +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +for path in (_REPO_ROOT, _REPO_ROOT / "tests", _FASTGEN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir +from export_pdd_qwen_image import collective_export_memory_preflight +from inference_pdd_qwen_image import build_pdd_student +from pdd_export import inspect_pdd_export, write_pdd_export +from pdd_recipe import build_pdd_export_setup, resolve_pdd_recipe_config + + +def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: + return { + "model": { + "pretrained_model_name_or_path": str(model_dir), + "torch_dtype": "float32", + "device": "cpu", + "transformer_engine_linear": False, + "peft": None, + "guidance_embeds": False, + "fuse_qkv_projections": False, + }, + "pdd": { + "pred_type": "flow", + "num_train_timesteps": None, + "guidance_scale": 4.0, + "student_sample_steps": 2, + "student_sample_type": "ode", + "grid_size": 4, + "flow_shift": 5.0, + "block_size_min": 1, + "block_size_max": 4, + "teacher_integrator": "euler", + "inference_blocks": [2, 2], + "data_free": False, + }, + "optim": {"learning_rate": 2.0e-5, "weight_decay": 0.01}, + "fsdp": { + "dp_size": 2, + "tp_size": 1, + "cp_size": 1, + "pp_size": 1, + "ep_size": 1, + "activation_checkpointing": False, + }, + "checkpoint": { + "enabled": True, + "checkpoint_dir": str(checkpoint_dir), + "model_save_format": "torch_save", + "save_consolidated": False, + }, + } + + +def _full_state(model: torch.nn.Module) -> dict[str, torch.Tensor]: + return get_model_state_dict( + model, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + + +def main() -> None: + dist.init_process_group("gloo") + payload = [tempfile.mkdtemp(prefix="modelopt-pdd-export-") if dist.get_rank() == 0 else None] + dist.broadcast_object_list(payload, src=0) + root = pathlib.Path(payload[0]) + model_root = root / "model" + model_dir = model_root / "tiny_qwen_image" + try: + if dist.get_rank() == 0: + assert create_tiny_qwen_image_pipeline_dir(model_root) == model_dir + dist.barrier() + config = resolve_pdd_recipe_config(_raw_config(model_dir, root / "checkpoints")) + source = build_pdd_export_setup(config) + expected = _full_state(source.student) + source.checkpointer.save_model(source.student, str(root / "dcp")) + source.checkpointer.close() + + destination = build_pdd_export_setup(config) + destination.checkpointer.load_model(destination.student, str(root / "dcp" / "model")) + full_state_bytes, largest_tensor_bytes = collective_export_memory_preflight( + destination.student, + max_shard_bytes=4 * 1024 * 1024, + headroom=1.0, + device=torch.device("cpu"), + ) + actual = _full_state(destination.student) + status = None + if dist.get_rank() == 0: + try: + assert expected and actual + assert expected.keys() == actual.keys() + for key in expected: + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + assert full_state_bytes == sum( + tensor.numel() * tensor.element_size() for tensor in actual.values() + ) + assert largest_tensor_bytes == max( + tensor.numel() * tensor.element_size() for tensor in actual.values() + ) + identity = { + "schema_version": 1, + "model": { + "id": "Qwen/Qwen-Image", + "revision": "3" * 40, + "dtype": "float32", + }, + "pdd_metadata": destination.metadata.to_dict(), + "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, + "automodel": { + key: destination.automodel_snapshot[key] + for key in ( + "distribution", + "version", + "package_tree_sha256", + "wheel_sha256", + "runtime_versions", + ) + }, + "topology": {"world_size": 2, "pure_data_parallel": True}, + } + output = write_pdd_export( + root / "export", + actual, + metadata=destination.metadata, + transformer_config=destination.transformer_config, + identity=identity, + source_checkpoint={ + "name": "step_00000001", + "manifest_sha256": "1" * 64, + "completed_steps": 1, + }, + modelopt_source={"commit": "2" * 40, "dirty": False}, + max_shard_bytes=4 * 1024 * 1024, + ) + descriptor = inspect_pdd_export(output) + assert descriptor.metadata == destination.metadata + restored, restored_descriptor, _dtype = build_pdd_student(output) + assert restored_descriptor.metadata == destination.metadata + restored_state = restored.state_dict() + assert restored_state.keys() == actual.keys() + for key in actual: + torch.testing.assert_close(restored_state[key], actual[key], rtol=0, atol=0) + status = {"ok": True} + except BaseException as error: + status = {"ok": False, "error": f"{type(error).__name__}: {error}"} + else: + assert not expected and not actual + payload = [status] + dist.broadcast_object_list(payload, src=0) + if not payload[0]["ok"]: + raise RuntimeError(payload[0]["error"]) + destination.checkpointer.close() + dist.barrier() + finally: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(root) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/examples/diffusers/fastgen/test_pdd_evaluation.py b/tests/examples/diffusers/fastgen/test_pdd_evaluation.py new file mode 100644 index 00000000000..c19df506239 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_evaluation.py @@ -0,0 +1,500 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hermetic tests for paired PDD effectiveness evidence and conclusions.""" + +from __future__ import annotations + +import copy +import hashlib +import pathlib +import sys + +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +for path in (_REPO_ROOT, _FASTGEN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from pdd_artifacts import ( + canonical_json_bytes, + load_canonical_json, + sha256_file, + write_canonical_json, +) +from pdd_evaluation import ( + CONDITION_PROTOCOLS, + EVALUATION_CONDITIONS, + GRID_PROTOCOLS, + INTEGRATOR_PROTOCOLS, + summarize_effectiveness_bundle, + validate_effectiveness_bundle, +) +from pdd_export import write_pdd_export + +from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC + + +def _reference(root: pathlib.Path, path: pathlib.Path) -> dict[str, str]: + return {"path": path.relative_to(root).as_posix(), "sha256": sha256_file(path)} + + +def _rewrite_manifest(manifest: pathlib.Path, data: dict) -> None: + manifest.unlink() + write_canonical_json(manifest, data) + detached = manifest.with_suffix(".json.sha256") + detached.unlink(missing_ok=True) + detached.write_bytes((sha256_file(manifest) + "\n").encode()) + + +def _export(root: pathlib.Path, automodel: dict) -> pathlib.Path: + config = PDDConfig( + grid_size=128, + flow_shift=5.0, + block_size_min=4, + block_size_max=64, + inference_blocks=[32, 32, 32, 32], + student_sample_steps=4, + guidance_scale=4.0, + num_train_timesteps=None, + ) + projection = PDDOutputProjection(1, 4, 128, QWEN_IMAGE_PDD_LAYER_SPEC) + metadata = PDDMetadata.from_config(config, projection) + identity = { + "model": {"id": "Qwen/Qwen-Image", "revision": "1" * 40, "dtype": "float32"}, + "pdd_metadata": metadata.to_dict(), + "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, + "automodel": automodel, + "data": { + "ordered_train_id_sha256": "8" * 64, + "ordered_heldout_id_sha256": "9" * 64, + "dataset_snapshot_sha256": "7" * 64, + "local_batch_size": 1, + "grad_accumulation_steps": 1, + }, + "topology": {"world_size": 1, "pure_data_parallel": True}, + } + return write_pdd_export( + root / "export", + projection.state_dict(), + metadata=metadata, + transformer_config={"_class_name": "QwenImageTransformer2DModel", "in_channels": 4}, + identity=identity, + source_checkpoint={ + "name": "step_00010000", + "manifest_sha256": "6" * 64, + "completed_steps": 10_000, + }, + modelopt_source={"commit": "2" * 40, "dirty": False}, + max_shard_bytes=1 << 20, + ) + + +def _write_bundle(tmp_path: pathlib.Path) -> pathlib.Path: + root = tmp_path / "run" + root.mkdir() + model = {"id": "Qwen/Qwen-Image", "revision": "1" * 40} + modelopt = {"commit": "2" * 40, "dirty": False} + file_path = "nemo_automodel/__init__.py" + file_sha = "3" * 64 + file_size = 7 + tree = hashlib.sha256() + tree.update(file_path.encode()) + tree.update(b"\0") + tree.update(file_sha.encode()) + tree.update(b"\0") + tree.update(str(file_size).encode()) + tree.update(b"\n") + tree_sha = tree.hexdigest() + automodel = { + "distribution": "nemo_automodel", + "version": "0.5.0", + "package_tree_sha256": tree_sha, + "wheel_sha256": "4" * 64, + "runtime_versions": {"diffusers": "0.38.0"}, + } + export_dir = _export(root, automodel) + export_manifest = export_dir / "manifest.json" + + environment = root / "environment.json" + write_canonical_json( + environment, + { + "distribution": "nemo_automodel", + "files": [{"path": file_path, "sha256": file_sha, "size": file_size}], + "import_origin": "/opt/pdd/site-packages/nemo_automodel/__init__.py", + "package_file_count": 1, + "package_tree_sha256": tree_sha, + "release_commit": "5" * 40, + "release_tag": "v0.5.0", + "root": "/opt/pdd/site-packages", + "runtime_versions": {"diffusers": "0.38.0"}, + "version": "0.5.0", + "wheel": "nemo_automodel-0.5.0-py3-none-any.whl", + "wheel_sha256": "4" * 64, + }, + ) + data_snapshot = root / "data_snapshot.json" + write_canonical_json( + data_snapshot, + { + "schema_version": 1, + "record_type": "pdd_dataset_snapshot", + "dataset_snapshot_sha256": "7" * 64, + "train_ids_sha256": "8" * 64, + "heldout_ids_sha256": "9" * 64, + }, + ) + + prompts = [] + prompt_pairs = [] + for index in range(16): + prompt = f"a small red cube on a white table, view {index:02d}" + prompt_sha = hashlib.sha256(prompt.encode()).hexdigest() + prompt_id = f"prompt-{index:04d}" + seed = 100 + index + prompts.append( + { + "prompt_id": prompt_id, + "prompt": prompt, + "prompt_sha256": prompt_sha, + "seeds": [seed], + } + ) + prompt_pairs.append((prompt_id, prompt_sha, seed)) + prompt_set = root / "prompts.json" + write_canonical_json(prompt_set, {"schema_version": 1, "prompts": prompts}) + negative_embedding = root / "negative_prompt_embedding.bin" + negative_embedding.write_bytes(b"authenticated fixed negative condition") + negative_condition = root / "negative_condition.json" + write_canonical_json( + negative_condition, + { + "schema_version": 1, + "record_type": "pdd_negative_condition", + "prompt_sha256": "c" * 64, + "embedding": _reference(root, negative_embedding), + }, + ) + + protocol_fields = { + "conditions": list(EVALUATION_CONDITIONS), + "condition_protocols": CONDITION_PROTOCOLS, + "grid_protocols": GRID_PROTOCOLS, + "integrator_protocols": INTEGRATOR_PROTOCOLS, + "image_protocol": { + "height": 1024, + "width": 1024, + "batch_size": 1, + "max_sequence_length": 512, + }, + "metric_protocols": { + "clip_score": { + "direction": "higher", + "implementation": "open_clip.ViT-H-14", + "revision": "a" * 40, + } + }, + "timing_protocol": { + "batch_size": 1, + "warmup_runs": 3, + "measured_runs": 5, + "scope": "transformer_sampling_and_vae_decode", + "synchronize_device": True, + }, + "decision_rule": { + "primary_condition": "pdd_4", + "primary_metric": "clip_score", + "quality_margin": 0.02, + "quality_ci_rule": "paired_bootstrap_95_noninferiority", + "efficiency_measure": "batch_normalized_transformer_evaluations", + "efficiency_baseline": "teacher_guided", + "minimum_relative_reduction": 0.5, + "minimum_paired_samples": 16, + }, + "negative_condition": _reference(root, negative_condition), + "data_snapshot": _reference(root, data_snapshot), + "stage_run_ids": {"canary": "canary-run", "training": "training-run"}, + "prompt_set": _reference(root, prompt_set), + "bootstrap": {"replicates": 1_000, "seed": 91}, + } + protocol_sha = hashlib.sha256(canonical_json_bytes(protocol_fields)).hexdigest() + evidence_references = {} + for stage in ("canary", "training"): + checkpoint = ( + { + "name": "step_00001500", + "manifest_sha256": "b" * 64, + "completed_steps": 1_500, + } + if stage == "canary" + else { + "name": "step_00010000", + "manifest_sha256": "6" * 64, + "completed_steps": 10_000, + } + ) + results = root / f"{stage}_results.json" + write_canonical_json( + results, + { + "schema_version": 1, + "record_type": "pdd_stage_results", + "stage": stage, + "status": "passed", + "slurm_job_ids": [101, 102, 103] if stage == "canary" else [201], + "completed_updates": 1_500 if stage == "canary" else 10_000, + "finite_loss": True, + "finite_gradients": True, + "resume_verified": True, + }, + ) + evidence = root / f"{stage}_evidence.json" + write_canonical_json( + evidence, + { + "schema_version": 1, + "record_type": "pdd_stage_evidence", + "stage": stage, + "status": "passed", + "run_id": f"{stage}-run", + "model": model, + "modelopt": modelopt, + "data_snapshot_sha256": sha256_file(data_snapshot), + "evaluation_protocol_sha256": protocol_sha, + "checkpoint": checkpoint, + "results": _reference(root, results), + }, + ) + evidence_references[stage] = _reference(root, evidence) + + records = [] + export_sha = sha256_file(export_manifest) + metric_values = { + "teacher_guided": 0.80, + "undistilled_euler_4": 0.73, + "undistilled_2step_4eval": 0.75, + "pdd_2": 0.77, + "pdd_4": 0.79, + "pdd_8": 0.795, + } + latencies = { + "teacher_guided": 10.0, + "undistilled_euler_4": 1.4, + "undistilled_2step_4eval": 1.2, + "pdd_2": 0.8, + "pdd_4": 1.0, + "pdd_8": 1.8, + } + for prompt_id, prompt_sha, seed in prompt_pairs: + for condition in EVALUATION_CONDITIONS: + output = root / f"{prompt_id}-{condition}.png" + output.write_bytes(b"png" + prompt_id.encode() + condition.encode()) + protocol = CONDITION_PROTOCOLS[condition] + latency = latencies[condition] + records.append( + { + "condition": condition, + "prompt_id": prompt_id, + "prompt_sha256": prompt_sha, + "seed": seed, + "metrics": {"clip_score": metric_values[condition]}, + "output": _reference(root, output), + "scheduler_steps": protocol["scheduler_steps"], + "actual_transformer_invocations": protocol["actual_transformer_invocations"], + "batch_normalized_transformer_evaluations": protocol[ + "batch_normalized_transformer_evaluations" + ], + "latency_seconds": latency, + "throughput_images_per_second": 1.0 / latency, + "peak_device_memory_bytes": 24_000_000_000, + "height": 1024, + "width": 1024, + "protocol_sha256": hashlib.sha256(canonical_json_bytes(protocol)).hexdigest(), + "evaluation_protocol_sha256": protocol_sha, + "model_artifact_sha256": export_sha, + } + ) + observations = root / "observations.json" + write_canonical_json(observations, {"schema_version": 1, "records": records}) + manifest = root / "manifest.json" + write_canonical_json( + manifest, + { + "schema_version": 1, + "stage": "effectiveness_evaluation", + "run_id": "test-run", + "model": model, + "modelopt": modelopt, + "pdd_export": _reference(root, export_manifest), + "observations": _reference(root, observations), + "environment": _reference(root, environment), + "stage_evidence": evidence_references, + **protocol_fields, + }, + ) + manifest.with_suffix(".json.sha256").write_bytes((sha256_file(manifest) + "\n").encode()) + return manifest + + +def test_effectiveness_bundle_is_authenticated_and_emits_effective_conclusion(tmp_path) -> None: + manifest = _write_bundle(tmp_path) + validated = validate_effectiveness_bundle(manifest) + first = summarize_effectiveness_bundle(validated) + second = summarize_effectiveness_bundle(validated) + + assert first == second + assert first["paired_sample_count"] == 16 + assert first["decision"]["label"] == "effective" + assert set(first["aggregates"]) == set(EVALUATION_CONDITIONS) + assert first["aggregates"]["pdd_2"]["mean_batch_normalized_transformer_evaluations"] == 2 + assert first["aggregates"]["pdd_4"]["metrics"]["clip_score"][ + "paired_delta_vs_teacher" + ] == pytest.approx(-0.01) + assert first["aggregates"]["pdd_4"]["peak_device_memory_bytes"]["mean"] > 0 + + +@pytest.mark.parametrize( + "corruption", + [ + "detached", + "stage", + "shadow", + "count", + "incomplete", + "missing_shard", + "tampered_shard", + "unrelated_training", + "failed_stage_results", + "unrelated_run", + "mismatched_export_data", + "bootstrap", + "prompt_reference", + "guided_count", + "grid_nodes", + "latency_decision", + "integrator_formula", + "protocol", + "resolution", + ], +) +def test_effectiveness_bundle_rejects_unclaimable_evidence(tmp_path, corruption) -> None: + manifest = _write_bundle(tmp_path) + root = manifest.parent + data = copy.deepcopy(load_canonical_json(manifest)) + if corruption == "detached": + manifest.with_suffix(".json.sha256").write_bytes(("0" * 64 + "\n").encode()) + elif corruption == "stage": + data["stage"] = "smoke" + elif corruption == "shadow": + environment_path = root / data["environment"]["path"] + environment = copy.deepcopy(load_canonical_json(environment_path)) + environment["import_origin"] = "/project/automodel/nemo_automodel/__init__.py" + environment_path.unlink() + write_canonical_json(environment_path, environment) + data["environment"] = _reference(root, environment_path) + elif corruption in ("missing_shard", "tampered_shard"): + shard = next((root / "export").glob("*.safetensors")) + if corruption == "missing_shard": + shard.unlink() + else: + with shard.open("ab") as stream: + stream.write(b"tampered") + elif corruption == "unrelated_training": + evidence_path = root / data["stage_evidence"]["training"]["path"] + evidence = copy.deepcopy(load_canonical_json(evidence_path)) + evidence["checkpoint"]["manifest_sha256"] = "d" * 64 + evidence_path.unlink() + write_canonical_json(evidence_path, evidence) + data["stage_evidence"]["training"] = _reference(root, evidence_path) + elif corruption == "failed_stage_results": + evidence_path = root / data["stage_evidence"]["training"]["path"] + evidence = copy.deepcopy(load_canonical_json(evidence_path)) + results_path = root / evidence["results"]["path"] + results = copy.deepcopy(load_canonical_json(results_path)) + results["finite_gradients"] = False + results_path.unlink() + write_canonical_json(results_path, results) + evidence["results"] = _reference(root, results_path) + evidence_path.unlink() + write_canonical_json(evidence_path, evidence) + data["stage_evidence"]["training"] = _reference(root, evidence_path) + elif corruption == "unrelated_run": + evidence_path = root / data["stage_evidence"]["training"]["path"] + evidence = copy.deepcopy(load_canonical_json(evidence_path)) + evidence["run_id"] = "unrelated-run" + evidence_path.unlink() + write_canonical_json(evidence_path, evidence) + data["stage_evidence"]["training"] = _reference(root, evidence_path) + elif corruption == "mismatched_export_data": + snapshot_path = root / data["data_snapshot"]["path"] + snapshot = copy.deepcopy(load_canonical_json(snapshot_path)) + snapshot["dataset_snapshot_sha256"] = "e" * 64 + snapshot_path.unlink() + write_canonical_json(snapshot_path, snapshot) + data["data_snapshot"] = _reference(root, snapshot_path) + elif corruption == "bootstrap": + data["bootstrap"]["seed"] += 1 + elif corruption == "prompt_reference": + source = root / data["prompt_set"]["path"] + alternate = root / "alternate_prompts.json" + write_canonical_json(alternate, load_canonical_json(source)) + data["prompt_set"] = _reference(root, alternate) + elif corruption == "grid_nodes": + data["grid_protocols"]["pdd_grid_128_shift5"]["nodes"][32] += 1e-4 + elif corruption == "latency_decision": + data["decision_rule"]["efficiency_measure"] = "latency_seconds" + elif corruption == "integrator_formula": + data["integrator_protocols"]["heun_explicit_trapezoid"]["terminal_rule"] = ( + "fall back to Euler at t_next=0" + ) + elif corruption == "protocol": + data["condition_protocols"]["pdd_4"]["pdd_blocks"] = [64, 64] + else: + observations_path = root / data["observations"]["path"] + observations = copy.deepcopy(load_canonical_json(observations_path)) + if corruption == "count": + observations["records"][3]["actual_transformer_invocations"] = 3 + elif corruption == "guided_count": + observations["records"][0]["actual_transformer_invocations"] = 50 + observations["records"][0]["batch_normalized_transformer_evaluations"] = 50 + elif corruption == "resolution": + observations["records"][3]["height"] = 512 + else: + observations["records"].pop() + observations_path.unlink() + write_canonical_json(observations_path, observations) + data["observations"] = _reference(root, observations_path) + if corruption not in ("detached", "missing_shard", "tampered_shard"): + _rewrite_manifest(manifest, data) + + with pytest.raises((ValueError, RuntimeError, FileNotFoundError)): + validate_effectiveness_bundle(manifest) + + +@pytest.mark.parametrize( + ("pdd_values", "label"), + [ + ([0.75] * 16, "not_effective"), + ([0.80] * 8 + [0.76] * 8, "inconclusive"), + ], +) +def test_predeclared_decision_rule_emits_boundary_conclusions(tmp_path, pdd_values, label) -> None: + manifest = _write_bundle(tmp_path) + root = manifest.parent + data = copy.deepcopy(load_canonical_json(manifest)) + observations_path = root / data["observations"]["path"] + observations = copy.deepcopy(load_canonical_json(observations_path)) + primary = [record for record in observations["records"] if record["condition"] == "pdd_4"] + for record, value in zip(primary, pdd_values): + record["metrics"]["clip_score"] = value + observations_path.unlink() + write_canonical_json(observations_path, observations) + data["observations"] = _reference(root, observations_path) + _rewrite_manifest(manifest, data) + + summary = summarize_effectiveness_bundle(validate_effectiveness_bundle(manifest)) + assert summary["decision"]["label"] == label diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py new file mode 100644 index 00000000000..b181c5f15d5 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hermetic evidence for authenticated PDD export, reconstruction, and schedules.""" + +from __future__ import annotations + +import copy +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" +for path in (_REPO_ROOT, _FASTGEN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from inference_pdd_qwen_image import _normalize_prompt_condition, _validate_qwen_projection +from pdd_export import ( + PDD_INFERENCE_SCHEDULES, + inspect_pdd_export, + load_pdd_export_into_model, + pdd_config_from_metadata, + write_pdd_export, +) + +from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QwenImagePDDAdapter, + convert_qwen_image_to_pdd, +) + + +class _TinyQwen(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(guidance_embeds=False, in_channels=4) + self.backbone = nn.Linear(4, 5) + self.proj_out = nn.Linear(5, 4) + self.calls = 0 + + def forward( + self, + *, + hidden_states, + timestep, + encoder_hidden_states, + encoder_hidden_states_mask, + img_shapes, + txt_seq_lens, + guidance, + return_dict, + ): + del img_shapes, txt_seq_lens, guidance, return_dict + condition = encoder_hidden_states.mean(dim=(1, 2), keepdim=True) + condition += encoder_hidden_states_mask.sum(dim=1)[:, None, None] / 100 + hidden = torch.tanh(self.backbone(hidden_states)) + self.calls += 1 + return (self.proj_out(hidden + condition + timestep[:, None, None] / 10),) + + +def _config(blocks=(32, 32, 32, 32)) -> PDDConfig: + return PDDConfig( + grid_size=128, + flow_shift=5.0, + block_size_min=4, + block_size_max=64, + inference_blocks=list(blocks), + student_sample_steps=len(blocks), + guidance_scale=4.0, + num_train_timesteps=None, + ) + + +def _converted(seed: int = 17): + torch.manual_seed(seed) + model = _TinyQwen() + config = _config() + projection = convert_qwen_image_to_pdd(model, config) + generator = torch.Generator().manual_seed(seed + 1) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator) / 10) + return model, config, PDDMetadata.from_config(config, projection) + + +def _identity(metadata: PDDMetadata) -> dict: + return { + "schema_version": 1, + "model": {"id": "synthetic-qwen", "revision": "f" * 40, "dtype": "float32"}, + "pdd_metadata": metadata.to_dict(), + "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, + "automodel": { + "distribution": "nemo_automodel", + "version": "0.5.0", + "package_tree_sha256": "1" * 64, + "wheel_sha256": "2" * 64, + "runtime_versions": {"diffusers": "0.38.0"}, + }, + "topology": {"world_size": 1, "pure_data_parallel": True}, + } + + +def _write(tmp_path: pathlib.Path): + model, config, metadata = _converted() + output = write_pdd_export( + tmp_path / "export", + model.state_dict(), + metadata=metadata, + transformer_config={"_class_name": "SyntheticQwen", "in_channels": 4}, + identity=_identity(metadata), + source_checkpoint={ + "name": "step_00000010", + "manifest_sha256": "3" * 64, + "completed_steps": 10, + }, + modelopt_source={"commit": "4" * 40, "dirty": False}, + max_shard_bytes=12_000, + ) + return output, model, config, metadata + + +def _condition(): + return torch.tensor([[[0.2, -0.3], [0.1, 0.4]]]), torch.ones(1, 2, dtype=torch.long) + + +def _sample(model: nn.Module, config: PDDConfig, state: torch.Tensor) -> torch.Tensor: + pipeline = PDDPipeline(model, nn.Identity(), config, QwenImagePDDAdapter(config)) + return pipeline.sample(state.clone(), condition=_condition()) + + +def test_bounded_safe_export_round_trip_and_seeded_schedules(tmp_path, monkeypatch) -> None: + output, source, _source_config, metadata = _write(tmp_path) + descriptor = inspect_pdd_export(output) + + shards = sorted(output.glob("*.safetensors")) + assert len(shards) >= 2 + assert all(path.stat().st_size <= descriptor.manifest["max_shard_bytes"] for path in shards) + assert descriptor.metadata == metadata + + restored = _TinyQwen() + convert_qwen_image_to_pdd(restored, _config()) + monkeypatch.setattr(torch, "load", lambda *args, **kwargs: pytest.fail("unsafe torch.load")) + load_pdd_export_into_model(output, restored) + for key, tensor in source.state_dict().items(): + torch.testing.assert_close(restored.state_dict()[key], tensor, rtol=0, atol=0) + + state = torch.randn((1, 1, 4, 4), generator=torch.Generator().manual_seed(91)) + for schedule, blocks in PDD_INFERENCE_SCHEDULES.items(): + config = pdd_config_from_metadata( + metadata, + schedule=schedule, + guidance_scale=4.0, + ) + source.calls = 0 + restored.calls = 0 + expected = _sample(source, config, state) + actual = _sample(restored, config, state) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert source.calls == restored.calls == len(blocks) + torch.testing.assert_close(_sample(restored, config, state), actual, rtol=0, atol=0) + + with pytest.raises(ValueError, match="block_size_max"): + pdd_config_from_metadata(metadata, blocks=[128], guidance_scale=4.0) + + +def test_pinned_qwen_none_prompt_mask_is_normalized_for_pdd() -> None: + """Diffusers 0.38 returns None when the single-prompt mask is all ones.""" + embeddings = torch.randn(1, 3, 5, dtype=torch.float32) + resolved_embeddings, resolved_mask = _normalize_prompt_condition( + embeddings, + None, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + assert resolved_embeddings.dtype == torch.bfloat16 + assert resolved_mask.dtype == torch.long + assert resolved_mask.shape == embeddings.shape[:2] + assert torch.equal(resolved_mask, torch.ones_like(resolved_mask)) + + +def test_qwen_projection_rejects_inconsistent_packed_width() -> None: + _model, _config_value, metadata = _converted() + base = _TinyQwen() + assert _validate_qwen_projection(base, metadata) is base.proj_out + base.config.in_channels = 8 + with pytest.raises(RuntimeError, match="proj_out width"): + _validate_qwen_projection(base, metadata) + + +def test_export_is_complete_before_atomic_rename(tmp_path, monkeypatch) -> None: + original_rename = pathlib.Path.rename + observed = False + + def checked_rename(path, target): + nonlocal observed + if path.name.endswith(".staging"): + inspect_pdd_export(path) + assert (path / "COMPLETE").is_file() + observed = True + return original_rename(path, target) + + monkeypatch.setattr(pathlib.Path, "rename", checked_rename) + _write(tmp_path) + assert observed + + +def test_export_rejects_unpinned_local_model_identity(tmp_path) -> None: + model, _config_value, metadata = _converted() + identity = _identity(metadata) + identity["model"]["revision"] = None + with pytest.raises(ValueError, match="pinned 40-character"): + write_pdd_export( + tmp_path / "local-model-export", + model.state_dict(), + metadata=metadata, + transformer_config={"in_channels": 4}, + identity=identity, + source_checkpoint={ + "name": "step_00000010", + "manifest_sha256": "3" * 64, + "completed_steps": 10, + }, + modelopt_source={"commit": "4" * 40, "dirty": False}, + max_shard_bytes=12_000, + ) + + +def test_export_rejects_nonfinite_and_existing_destination(tmp_path) -> None: + output, model, _config_value, metadata = _write(tmp_path) + with pytest.raises(FileExistsError): + write_pdd_export( + output, + model.state_dict(), + metadata=metadata, + transformer_config={"in_channels": 4}, + identity=_identity(metadata), + source_checkpoint={ + "name": "step_00000010", + "manifest_sha256": "3" * 64, + "completed_steps": 10, + }, + modelopt_source={"commit": "4" * 40, "dirty": False}, + max_shard_bytes=12_000, + ) + + bad = copy.deepcopy(model.state_dict()) + bad["backbone.weight"][0, 0] = float("nan") + with pytest.raises(FloatingPointError, match="non-finite"): + write_pdd_export( + tmp_path / "bad", + bad, + metadata=metadata, + transformer_config={"in_channels": 4}, + identity=_identity(metadata), + source_checkpoint={ + "name": "step_00000010", + "manifest_sha256": "3" * 64, + "completed_steps": 10, + }, + modelopt_source={"commit": "4" * 40, "dirty": False}, + max_shard_bytes=12_000, + ) + + +@pytest.mark.parametrize("corruption", ["complete", "shard", "extra", "symlink"]) +def test_export_authentication_rejects_corruption(tmp_path, corruption) -> None: + output, _model, _config_value, _metadata = _write(tmp_path) + if corruption == "complete": + (output / "COMPLETE").unlink() + elif corruption == "shard": + shard = next(output.glob("*.safetensors")) + with shard.open("ab") as stream: + stream.write(b"corrupt") + elif corruption == "extra": + (output / "undeclared.bin").write_bytes(b"extra") + else: + (output / "linked").symlink_to(output / "config.json") + + with pytest.raises((FileNotFoundError, RuntimeError)): + inspect_pdd_export(output) + + +def test_safe_load_rejects_wrong_model_inventory(tmp_path) -> None: + output, _model, _config_value, _metadata = _write(tmp_path) + unconverted = _TinyQwen() + with pytest.raises(RuntimeError, match="shape mismatch"): + load_pdd_export_into_model(output, unconverted) + + wrong_dtype = _TinyQwen().to(torch.float64) + convert_qwen_image_to_pdd(wrong_dtype, _config()) + with pytest.raises(RuntimeError, match="dtype mismatch"): + load_pdd_export_into_model(output, wrong_dtype) diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index e45da3804c8..ed854ef140b 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -22,7 +22,12 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -from pdd_recipe import build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config +from pdd_recipe import ( + build_pdd_export_setup, + build_pdd_setup, + initialize_pdd_distributed, + resolve_pdd_recipe_config, +) from verify_readonly_automodel import snapshot_installed_distribution @@ -259,6 +264,27 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: source.checkpointer.save_model(source.student, str(checkpoint_root)) source.checkpointer.save_optimizer(source.optimizer, source.student, str(checkpoint_root)) + export_setup = build_pdd_export_setup(config) + assert export_setup.lifecycle == ( + "load/select", + "pdd_conversion", + "device", + "qkv", + "parallelize", + "checkpoint", + ) + assert export_setup.metadata == source.metadata + assert export_setup.checkpoint_keys == source.checkpoint_keys + assert not hasattr(export_setup, "optimizer") + export_setup.checkpointer.load_model( + export_setup.student, + str(checkpoint_root / "model"), + ) + torch.testing.assert_close( + export_setup.student.state_dict()["proj_out.weight"], + expected_weight, + ) + destination = build_pdd_setup(config) assert destination.metadata == source.metadata destination_projection = destination.projection @@ -283,6 +309,7 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert snapshot_installed_distribution() == before source.checkpointer.close() destination.checkpointer.close() + export_setup.checkpointer.close() def test_qwen_pdd_adapter_has_no_automodel_import() -> None: diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index d8f487d4599..2815de6b8b8 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -26,7 +26,11 @@ sys.path.insert(0, str(pathlib.Path(__file__).parent)) from fastgen_data.replayable_sampler import ReplayableBatchSampler -from pdd_checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity +from pdd_checkpoint import ( + PDDCheckpointManager, + build_pdd_checkpoint_identity, + resolve_pdd_training_checkpoint, +) from pdd_recipe import initialize_pdd_distributed from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 from pdd_training import prepare_qwen_pdd_batch @@ -394,6 +398,14 @@ def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) _refresh_complete_marker(mismatched) (tmp_path / "checkpoints" / "LATEST").write_text(mismatched.name + "\n") assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() + selected, selected_manifest = resolve_pdd_training_checkpoint( + tmp_path / "checkpoints", + "LATEST", + expected_world_size=1, + expected_identity=third_manager.identity, + ) + assert selected == resumed_checkpoint.resolve() + assert selected_manifest["identity"] == third_manager.identity with pytest.raises(RuntimeError, match="identity"): third_manager.resolve(mismatched.name) From df81ee3f6d67c4a8925586e965bf0610848824f8 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 11:33:18 -0700 Subject: [PATCH 13/45] test(fastgen): add PDD GPU verification harnesses Signed-off-by: Meng Xin --- .../test_pdd_qwen_operability_smoke.py | 401 ++++++ tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py | 430 +++++++ .../fastgen/pdd_qwen_operability_smoke.py | 1145 +++++++++++++++++ tests/gpu/torch/fastgen/test_pdd_toy.py | 240 ++++ 4 files changed, 2216 insertions(+) create mode 100644 tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py create mode 100644 tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py create mode 100644 tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py create mode 100644 tests/gpu/torch/fastgen/test_pdd_toy.py diff --git a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py new file mode 100644 index 00000000000..baeb7db9882 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py @@ -0,0 +1,401 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hermetic tests for the full-Qwen PDD smoke evidence contract.""" + +from __future__ import annotations + +import copy +import hashlib +import importlib.util +import json +from pathlib import Path + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[4] +_HARNESS = _REPO_ROOT / "tests" / "gpu" / "torch" / "fastgen" / "pdd_qwen_operability_smoke.py" +_SPEC = importlib.util.spec_from_file_location("pdd_qwen_operability_smoke", _HARNESS) +assert _SPEC is not None and _SPEC.loader is not None +smoke = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(smoke) + + +def _stage_result(stage: str) -> dict: + step = 1 if stage == "train-one" else 2 + sample_ids = [f"synthetic-pdd-smoke-step-{step}-rank-{rank}" for rank in range(2)] + learning_rate = 2.0e-5 + return { + "schema_version": 1, + "record_type": "pdd_qwen_smoke_stage", + "stage": stage, + "pid": 100 + step, + "world_size": 2, + "model": { + "id": "Qwen/Qwen-Image", + "revision": "75e0b4be04f60ec59a75f475837eced720f823b6", + "dtype": "bfloat16", + }, + "pdd": { + "grid_size": 128, + "flow_shift": 5.0, + "block_size_min": 4, + "block_size_max": 64, + "teacher_integrator": "euler", + "guidance_scale": 4.0, + "guidance_rescale": 1.0, + "guidance_eps": 1e-5, + }, + "source": {"commit": "1" * 40, "dirty": False}, + "config_sha256": "2" * 64, + "automodel": { + "distribution": "nemo_automodel", + "version": "0.5.0", + "package_tree_sha256": smoke._AUTOMODEL_TREE_SHA256, + "wheel_sha256": smoke._AUTOMODEL_WHEEL_SHA256, + "runtime_versions": {"diffusers": "0.38.0"}, + }, + "gpu": { + "names": ["GPU", "GPU"], + "total_memory_bytes": [80_000_000_000, 80_000_000_000], + "host_available_bytes": [500_000_000_000, 500_000_000_000], + "allocated_before_step_bytes": [30_000_000_000, 30_000_000_000], + "peak_memory_bytes": [40_000_000_000, 40_000_000_000], + "student_parameter_bytes": 40_000_000_000, + "teacher_parameter_bytes": 40_000_000_000, + "step_seconds": 12.5, + }, + "pair": {"n": 0 if step == 1 else 124, "k": 63 if step == 1 else 127}, + "sample_ids": sample_ids, + "diagnostics": { + "completed_step": step, + "loss": 0.5, + "grad_norm": 1.25, + "student_adamw_nominal_update_ratio": 1e-4, + "pdd_projection_update_ratio": 2e-4, + "learning_rate": learning_rate, + "student_velocity_rms": 0.75, + "teacher_velocity_rms": 0.8, + "student_teacher_velocity_rms_ratio": 0.9375, + "reconstructed_state_rms": 1.1, + }, + "teacher_calls_per_rank": [2, 2], + "checkpoint": { + "path": f"checkpoints/step_{step:08d}", + "manifest_sha256": "5" * 64, + "completed_steps": step, + "parent_checkpoint": None if step == 1 else "step_00000001", + }, + "resume": None + if step == 1 + else { + "selected_checkpoint": "step_00000001", + "completed_steps": 1, + "parent_checkpoint": None, + "first_sample_ids": sample_ids, + "learning_rate": learning_rate, + }, + } + + +def _automodel_snapshot_fixture() -> tuple[dict, dict]: + records = [] + tree = hashlib.sha256() + for index in range(smoke._AUTOMODEL_PACKAGE_FILE_COUNT): + path = f"nemo_automodel/file_{index:03d}.py" + digest = hashlib.sha256(f"file-{index}".encode()).hexdigest() + size = index + 1 + tree.update(path.encode()) + tree.update(b"\0") + tree.update(digest.encode()) + tree.update(b"\0") + tree.update(str(size).encode()) + tree.update(b"\n") + records.append({"path": path, "sha256": digest, "size": size}) + automodel = { + "distribution": "nemo_automodel", + "version": "0.5.0", + "package_tree_sha256": tree.hexdigest(), + "wheel_sha256": smoke._AUTOMODEL_WHEEL_SHA256, + "runtime_versions": {"diffusers": "0.38.0"}, + } + snapshot = { + **automodel, + "files": records, + "import_origin": "/opt/pdd/site-packages/nemo_automodel/__init__.py", + "package_file_count": smoke._AUTOMODEL_PACKAGE_FILE_COUNT, + "release_commit": smoke._AUTOMODEL_RELEASE_COMMIT, + "release_tag": smoke._AUTOMODEL_RELEASE_TAG, + "root": "/opt/pdd/site-packages", + "wheel": smoke._AUTOMODEL_WHEEL, + } + return snapshot, automodel + + +def _checkpoint_identity(stage: dict) -> dict: + return { + "schema_version": 1, + "model": stage["model"], + "pdd_metadata": { + "schema_version": 1, + "grid_size": 128, + "flow_shift": 5.0, + "block_size_min": 4, + "block_size_max": 64, + "inference_blocks": [32, 32, 32, 32], + "teacher_integrator": "euler", + "layer_spec": { + "projection_path": "transformer.proj_out", + "head_layout": "channel_major", + "output_channels": None, + }, + "base_projection": {"in_features": 3072, "out_features": 64, "bias": True}, + }, + "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, + "automodel": stage["automodel"], + "data": {}, + "topology": {"world_size": 2, "pure_data_parallel": True}, + "training": {}, + "optimizer": {}, + "scheduler": {}, + } + + +def _bundle_link_fixture() -> tuple[dict, dict, dict, dict, dict, dict]: + snapshot, automodel = _automodel_snapshot_fixture() + stage1 = _stage_result("train-one") + stage2 = _stage_result("resume-one") + stage1["automodel"] = copy.deepcopy(automodel) + stage2["automodel"] = copy.deepcopy(automodel) + identity = _checkpoint_identity(stage1) + manifest1 = {"identity": copy.deepcopy(identity)} + manifest2 = {"identity": copy.deepcopy(identity)} + export = { + "identity": copy.deepcopy(identity), + "modelopt_source": copy.deepcopy(stage2["source"]), + "source_checkpoint": { + "name": "step_00000002", + "manifest_sha256": stage2["checkpoint"]["manifest_sha256"], + "completed_steps": 2, + }, + } + return stage1, stage2, manifest1, manifest2, export, snapshot + + +def test_training_stage_contract_accepts_only_exact_canonical_chain() -> None: + stage1 = _stage_result("train-one") + stage2 = _stage_result("resume-one") + + smoke.validate_stage_result(stage1, stage="train-one") + smoke.validate_stage_result(stage2, stage="resume-one") + + wrong_pair = copy.deepcopy(stage2) + wrong_pair["pair"] = {"n": 120, "k": 127} + with pytest.raises(ValueError, match="support pair"): + smoke.validate_stage_result(wrong_pair, stage="resume-one") + + zero_update = copy.deepcopy(stage1) + zero_update["diagnostics"]["pdd_projection_update_ratio"] = 0.0 + with pytest.raises(ValueError, match="finite and positive"): + smoke.validate_stage_result(zero_update, stage="train-one") + + stale_resume = copy.deepcopy(stage2) + stale_resume["resume"]["selected_checkpoint"] = "step_00000000" + with pytest.raises(ValueError, match="resume evidence"): + smoke.validate_stage_result(stale_resume, stage="resume-one") + + +def test_inference_contract_authenticates_exact_pdd4_counters_and_png(tmp_path: Path) -> None: + image = tmp_path / "image.png" + image.write_bytes(b"\x89PNG\r\n\x1a\nnonempty-hashed-test-fixture") + digest = hashlib.sha256(image.read_bytes()).hexdigest() + result = { + "schema_version": 1, + "record_type": "pdd_inference", + "condition": "pdd_4", + "schedule": "pdd-4", + "blocks": [32, 32, 32, 32], + "height": 1024, + "width": 1024, + "scheduler_steps": 4, + "actual_transformer_invocations": 4, + "batch_normalized_transformer_evaluations": 4, + "latency_seconds": 1.5, + "output": {"path": "image.png", "sha256": digest}, + } + + smoke.validate_inference_result(result, root=tmp_path) + + wrong_calls = copy.deepcopy(result) + wrong_calls["actual_transformer_invocations"] = 5 + with pytest.raises(ValueError, match="compute counters"): + smoke.validate_inference_result(wrong_calls, root=tmp_path) + + wrong_hash = copy.deepcopy(result) + wrong_hash["output"]["sha256"] = "0" * 64 + with pytest.raises(ValueError, match="PNG hash"): + smoke.validate_inference_result(wrong_hash, root=tmp_path) + + reduced_resolution = copy.deepcopy(result) + reduced_resolution["height"] = 512 + with pytest.raises(ValueError, match="1024x1024"): + smoke.validate_inference_result(reduced_resolution, root=tmp_path) + + +def test_bundle_links_reject_cross_run_artifact_splicing() -> None: + stage1, stage2, manifest1, manifest2, export, snapshot = _bundle_link_fixture() + + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=export, + automodel_snapshot=snapshot, + ) + + wrong_checkpoint = copy.deepcopy(manifest2) + wrong_checkpoint["identity"]["model"]["revision"] = "a" * 40 + with pytest.raises(ValueError, match="identities differ"): + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=wrong_checkpoint, + export_manifest=export, + automodel_snapshot=snapshot, + ) + + wrong_export_identity = copy.deepcopy(export) + wrong_export_identity["identity"]["data"] = {"spliced": True} + with pytest.raises(ValueError, match="export identity"): + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=wrong_export_identity, + automodel_snapshot=snapshot, + ) + + wrong_export_checkpoint = copy.deepcopy(export) + wrong_export_checkpoint["source_checkpoint"]["manifest_sha256"] = "f" * 64 + with pytest.raises(ValueError, match="exact step-2"): + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=wrong_export_checkpoint, + automodel_snapshot=snapshot, + ) + + wrong_export_source = copy.deepcopy(export) + wrong_export_source["modelopt_source"]["commit"] = "e" * 40 + with pytest.raises(ValueError, match="training source"): + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=wrong_export_source, + automodel_snapshot=snapshot, + ) + + with pytest.raises(ValueError, match="AutoModel snapshot"): + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=export, + automodel_snapshot={}, + ) + + corrupt_snapshot = copy.deepcopy(snapshot) + corrupt_snapshot["files"][0]["sha256"] = "0" * 64 + with pytest.raises(ValueError, match="tree digest"): + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=export, + automodel_snapshot=corrupt_snapshot, + ) + + float_count_snapshot = copy.deepcopy(snapshot) + float_count_snapshot["package_file_count"] = float(smoke._AUTOMODEL_PACKAGE_FILE_COUNT) + with pytest.raises(ValueError, match="release identity"): + smoke._validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=export, + automodel_snapshot=float_count_snapshot, + ) + + +def test_automodel_snapshots_require_identical_bytes_and_no_symlinks(tmp_path: Path) -> None: + snapshot, automodel = _automodel_snapshot_fixture() + before = tmp_path / "before.json" + after = tmp_path / "after.json" + payload = json.dumps(snapshot, indent=2, sort_keys=True) + "\n" + before.write_text(payload) + after.write_text(payload) + assert ( + smoke._load_matching_automodel_snapshots( + before, + after, + expected_automodel=automodel, + ) + == snapshot + ) + + after.write_text(json.dumps(snapshot, sort_keys=True) + "\n") + with pytest.raises(ValueError, match="snapshot changed"): + smoke._load_matching_automodel_snapshots( + before, + after, + expected_automodel=automodel, + ) + + symlink = tmp_path / "before-symlink.json" + symlink.symlink_to(before) + with pytest.raises(ValueError, match="symlink"): + smoke._load_matching_automodel_snapshots( + symlink, + before, + expected_automodel=automodel, + ) + + +def test_smoke_artifact_paths_reject_symlinked_stage_inference_and_export(tmp_path: Path) -> None: + outside = tmp_path / "outside" + outside.mkdir() + target = outside / "artifact.json" + target.write_text("{}") + run_root = tmp_path / "run" + run_root.mkdir() + (run_root / "stage1.json").symlink_to(target) + inference = run_root / "inference" + inference.mkdir() + (inference / "pdd4.json").symlink_to(target) + (run_root / "export").symlink_to(outside, target_is_directory=True) + + with pytest.raises(ValueError, match="symlink"): + smoke._relative_regular_file(run_root, "stage1.json", name="stage") + with pytest.raises(ValueError, match="symlink"): + smoke._relative_regular_file(run_root, "inference/pdd4.json", name="inference") + with pytest.raises(ValueError, match="symlink"): + smoke._regular_directory(run_root / "export", name="export") + + target_parent = tmp_path / "target-parent" + target_parent.mkdir() + symlinked_parent = tmp_path / "symlinked-parent" + symlinked_parent.symlink_to(target_parent, target_is_directory=True) + requested_child = symlinked_parent / "must-not-be-created" + with pytest.raises(ValueError, match="symlink"): + smoke._create_run_root(requested_child) + assert not (target_parent / requested_child.name).exists() diff --git a/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py b/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py new file mode 100644 index 00000000000..64e18d1fa43 --- /dev/null +++ b/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py @@ -0,0 +1,430 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Two-rank CUDA FSDP2 optimization and exact-resume proof for plain PDD modules.""" + +from __future__ import annotations + +import gc +import importlib.util +import json +import os +import pathlib +import shutil +import sys +import tempfile +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch import nn +from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict +from torch.distributed.fsdp import fully_shard + +from modelopt.torch.fastgen import ( + PDDConfig, + PDDLayerSpec, + PDDOutputProjection, + PDDPipeline, + convert_to_pdd_output_projection, +) + +_FORBIDDEN_MODULES = ("diffusers", "fastgen", "nemo_automodel") +_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) -> torch.Tensor: + return self.projection(torch.tanh(self.backbone(state))) + + +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, + condition: torch.Tensor, + ) -> torch.Tensor: + return self.projection(state) + 0.125 * time[:, None] + 0.05 * condition + + +class _GuidedAdapter: + def __init__(self) -> None: + self.teacher_calls = 0 + + @staticmethod + def _dtype(model: nn.Module) -> torch.dtype: + return next(model.parameters()).dtype + + def student_all_heads( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + *, + condition: Any = None, + **model_kwargs: Any, + ) -> torch.Tensor: + del time, condition, model_kwargs + output = model(state.to(self._dtype(model))) + return output.reshape(state.shape[0], _GRID_SIZE, _WIDTH) + + 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: + del time, condition, model_kwargs + projection = model.get_submodule("projection") + assert isinstance(projection, PDDOutputProjection) + with projection.fuse_block(start, end, grid): + return model(state.to(self._dtype(model))) + + 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: + del model_kwargs + if not isinstance(condition, torch.Tensor) or not isinstance( + negative_condition, torch.Tensor + ): + raise TypeError("guided toy teacher requires tensor conditions") + dtype = self._dtype(model) + state = state.to(dtype) + time = time.to(dtype) + conditional = model(state, time, condition.to(dtype)) + unconditional = model(state, time, negative_condition.to(dtype)) + self.teacher_calls += 2 + return conditional + 3.0 * (conditional - unconditional) + + +@dataclass +class _Lifecycle: + student: nn.Module + teacher: nn.Module + projection: PDDOutputProjection + pipeline: PDDPipeline + optimizer: torch.optim.AdamW + adapter: _GuidedAdapter + + +def _local(value: torch.Tensor) -> torch.Tensor: + to_local = getattr(value, "to_local", None) + return to_local() if callable(to_local) else value + + +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 + index * 0.01, + 0.2 + offset + index * 0.01, + parameter.numel(), + dtype=torch.float32, + device=parameter.device, + ) + parameter.copy_(values.reshape_as(parameter).to(parameter.dtype)) + + +def _config() -> PDDConfig: + return PDDConfig( + grid_size=_GRID_SIZE, + flow_shift=5.0, + block_size_min=1, + block_size_max=_GRID_SIZE, + inference_blocks=[2, 2], + student_sample_steps=2, + guidance_scale=4.0, + ) + + +def _build(device: torch.device) -> _Lifecycle: + config = _config() + student = _Student().to(device=device, dtype=torch.bfloat16) + teacher = _Teacher().to(device=device, dtype=torch.bfloat16).eval().requires_grad_(False) + _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, + ) + projection_module_id = id(projection) + projection_shape = projection.weight.shape + student = fully_shard(student) + teacher = fully_shard(teacher) + assert id(student.get_submodule("projection")) == projection_module_id + assert student.get_submodule("projection").weight.shape == projection_shape + optimizer = torch.optim.AdamW( + student.parameters(), + lr=2.0e-3, + weight_decay=0.0, + foreach=False, + fused=False, + ) + optimizer_parameters = [ + parameter for group in optimizer.param_groups for parameter in group["params"] + ] + assert any(parameter is projection.weight for parameter in optimizer_parameters) + adapter = _GuidedAdapter() + pipeline = PDDPipeline(student, teacher, config, adapter) + return _Lifecycle(student, teacher, projection, pipeline, optimizer, adapter) + + +def _batch( + *, rank: int, step: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + base = torch.arange(_WIDTH, device=device, dtype=torch.float32).reshape(1, -1) + data = (base / 10 + 0.05 * rank + 0.025 * step).to(torch.bfloat16) + noise = 0.4 - base / 20 + 0.01 * rank + condition = 0.2 + base / 30 + 0.02 * step + negative = -0.1 - base / 40 - 0.01 * rank + n = torch.tensor([0 if step == 1 else 2], device=device, dtype=torch.int64) + k = torch.tensor([1 if step == 1 else 3], device=device, dtype=torch.int64) + return data, noise, condition, negative, n, k + + +def _global_norm(values: list[torch.Tensor], device: torch.device) -> torch.Tensor: + squared = torch.zeros((), device=device, dtype=torch.float64) + for value in values: + squared += _local(value.detach()).float().square().sum(dtype=torch.float64) + dist.all_reduce(squared, op=dist.ReduceOp.SUM) + return squared.sqrt() + + +def _step(lifecycle: _Lifecycle, *, rank: int, step: int, device: torch.device) -> dict[str, float]: + data, noise, condition, negative, n, k = _batch(rank=rank, step=step, device=device) + lifecycle.optimizer.zero_grad(set_to_none=True) + calls_before = lifecycle.adapter.teacher_calls + loss, metrics = lifecycle.pipeline.compute_loss( + data, + noise=noise, + condition=condition, + negative_condition=negative, + n=n, + k=k, + ) + 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 lifecycle.adapter.teacher_calls - calls_before == 2 + assert all(parameter.grad is None for parameter in lifecycle.teacher.parameters()) + gradients = [ + parameter.grad for parameter in lifecycle.student.parameters() if parameter.grad is not None + ] + grad_norm = _global_norm(gradients, device) + assert torch.isfinite(grad_norm) and grad_norm > 0 + before = { + name: _local(parameter.detach()).clone() + for name, parameter in lifecycle.student.named_parameters() + } + lifecycle.optimizer.step() + updates = [ + _local(parameter.detach()) - before[name] + for name, parameter in lifecycle.student.named_parameters() + ] + update_norm = _global_norm(updates, device) + assert torch.isfinite(update_norm) and update_norm > 0 + reduced_loss = loss.detach().double() + dist.all_reduce(reduced_loss, op=dist.ReduceOp.SUM) + reduced_loss /= dist.get_world_size() + return { + "loss": float(reduced_loss.item()), + "grad_norm": float(grad_norm.item()), + "update_norm": float(update_norm.item()), + } + + +def _state(model: nn.Module) -> dict[str, torch.Tensor]: + return { + name: _local(value.detach()).clone() + for name, value in model.state_dict().items() + if isinstance(value, torch.Tensor) + } + + +def _assert_state_equal(actual: nn.Module, expected: dict[str, torch.Tensor]) -> None: + actual_state = _state(actual) + assert actual_state.keys() == expected.keys() + for name in expected: + torch.testing.assert_close(actual_state[name], expected[name], rtol=0, atol=0) + + +def _all_rng_states(device: torch.device) -> dict[str, torch.Tensor]: + local_state = { + "cpu": torch.get_rng_state(), + "cuda": torch.cuda.get_rng_state(device), + } + gathered: list[dict[str, torch.Tensor] | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, local_state) + states: dict[str, torch.Tensor] = {} + for rank, value in enumerate(gathered): + assert value is not None + states[f"cpu_rng_rank_{rank}"] = value["cpu"].to(device) + states[f"cuda_rng_rank_{rank}"] = value["cuda"].to(device) + return states + + +def _save_checkpoint( + root: pathlib.Path, + lifecycle: _Lifecycle, + *, + completed_steps: int, + device: torch.device, +) -> pathlib.Path: + staging = root / ".step_00000001.staging" + final = root / "step_00000001" + if dist.get_rank() == 0: + staging.mkdir(parents=True) + dist.barrier() + model_state, optimizer_state = get_state_dict(lifecycle.student, lifecycle.optimizer) + extra = _all_rng_states(device) + extra["completed_steps"] = torch.tensor([completed_steps], device=device, dtype=torch.int64) + dcp.save( + {"model": model_state, "optimizer": optimizer_state, "extra": extra}, + checkpoint_id=staging, + ) + dist.barrier() + if dist.get_rank() == 0: + assert (staging / ".metadata").is_file() + assert any(path.suffix == ".distcp" for path in staging.iterdir()) + os.replace(staging, final) + (final / "COMPLETE").write_text( + json.dumps({"schema_version": 1, "completed_steps": completed_steps}) + "\n" + ) + dist.barrier() + assert (final / "COMPLETE").is_file() + return final + + +def _load_checkpoint( + checkpoint: pathlib.Path, + lifecycle: _Lifecycle, + *, + device: torch.device, +) -> int: + marker = json.loads((checkpoint / "COMPLETE").read_text()) + assert marker == {"schema_version": 1, "completed_steps": 1} + model_state, optimizer_state = get_state_dict(lifecycle.student, lifecycle.optimizer) + extra = _all_rng_states(device) + extra["completed_steps"] = torch.zeros(1, device=device, dtype=torch.int64) + payload = {"model": model_state, "optimizer": optimizer_state, "extra": extra} + dcp.load(payload, checkpoint_id=checkpoint) + incompatible = set_state_dict( + lifecycle.student, + lifecycle.optimizer, + model_state_dict=payload["model"], + optim_state_dict=payload["optimizer"], + ) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + rank = dist.get_rank() + torch.set_rng_state(payload["extra"][f"cpu_rng_rank_{rank}"].cpu()) + torch.cuda.set_rng_state(payload["extra"][f"cuda_rng_rank_{rank}"].cpu(), device) + return int(payload["extra"]["completed_steps"].item()) + + +def _assert_call_counts(adapter: _GuidedAdapter, expected: int, device: torch.device) -> None: + value = torch.tensor([adapter.teacher_calls], device=device, dtype=torch.int64) + gathered = [torch.zeros_like(value) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, value) + assert [int(item.item()) for item in gathered] == [expected] * dist.get_world_size() + + +def _assert_optional_frameworks_absent() -> None: + resolvable = sorted(name for name in _FORBIDDEN_MODULES if importlib.util.find_spec(name)) + assert not resolvable, f"plain PDD FSDP2 environment resolves optional frameworks: {resolvable}" + imported = sorted(name for name in _FORBIDDEN_MODULES if name in sys.modules) + assert not imported, f"plain PDD FSDP2 smoke imported optional frameworks: {imported}" + + +def main() -> None: + _assert_optional_frameworks_absent() + assert torch.cuda.is_available(), "Task-10 FSDP2 gate requires CUDA" + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + assert world_size == 2, f"Task-10 FSDP2 gate requires two ranks, got {world_size}" + local_rank = int(os.environ["LOCAL_RANK"]) + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + assert torch.cuda.get_device_capability(device)[0] >= 8 + root_payload = [tempfile.mkdtemp(prefix="modelopt-pdd-fsdp2-") if rank == 0 else None] + dist.broadcast_object_list(root_payload, src=0) + root = pathlib.Path(root_payload[0]) + try: + torch.manual_seed(2026 + rank) + torch.cuda.manual_seed(3026 + rank) + reference = _build(device) + resumable = _build(device) + reference_step1 = _step(reference, rank=rank, step=1, device=device) + resumable_step1 = _step(resumable, rank=rank, step=1, device=device) + assert reference_step1 == resumable_step1 + _assert_state_equal(resumable.student, _state(reference.student)) + saved_cpu_rng = torch.get_rng_state().clone() + saved_cuda_rng = torch.cuda.get_rng_state(device).clone() + checkpoint = _save_checkpoint( + root, + resumable, + completed_steps=1, + device=device, + ) + + reference_step2 = _step(reference, rank=rank, step=2, device=device) + reference_state = _state(reference.student) + _assert_call_counts(reference.adapter, 4, device) + + del resumable + gc.collect() + torch.cuda.empty_cache() + restored = _build(device) + assert _load_checkpoint(checkpoint, restored, device=device) == 1 + torch.testing.assert_close(torch.get_rng_state(), saved_cpu_rng, rtol=0, atol=0) + torch.testing.assert_close(torch.cuda.get_rng_state(device), saved_cuda_rng, rtol=0, atol=0) + restored_step2 = _step(restored, rank=rank, step=2, device=device) + assert restored_step2 == reference_step2 + _assert_state_equal(restored.student, reference_state) + _assert_call_counts(restored.adapter, 2, device) + _assert_optional_frameworks_absent() + dist.barrier() + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(root) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py b/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py new file mode 100644 index 00000000000..4da3ab09121 --- /dev/null +++ b/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py @@ -0,0 +1,1145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Staged canonical Qwen-Image PDD operability smoke and result validator.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import pathlib +import subprocess +import sys +import time +from collections.abc import Mapping, Sequence +from typing import Any + +_THIS_FILE = pathlib.Path(__file__).resolve() +_REPO_ROOT = _THIS_FILE.parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +for _path in (_REPO_ROOT, _FASTGEN_DIR): + if str(_path) not in sys.path: + sys.path.insert(0, str(_path)) + +_MODEL_ID = "Qwen/Qwen-Image" +_MODEL_REVISION = "75e0b4be04f60ec59a75f475837eced720f823b6" +_AUTOMODEL_TREE_SHA256 = "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528" +_AUTOMODEL_WHEEL_SHA256 = "881aebafc5145752842afbbfe0a42e1c33d06847c3e418ad3d6f154ddc8e0f45" +_AUTOMODEL_RELEASE_COMMIT = "d02f49cb314554715aabb97e8dba6599c9f6e9e0" +_AUTOMODEL_RELEASE_TAG = "v0.5.0" +_AUTOMODEL_WHEEL = "nemo_automodel-0.5.0-py3-none-any.whl" +_AUTOMODEL_PACKAGE_FILE_COUNT = 490 +_EXPECTED_PAIRS = {"train-one": (0, 63), "resume-one": (124, 127)} +_STAGE_RESULT_KEYS = { + "schema_version", + "record_type", + "stage", + "pid", + "world_size", + "model", + "pdd", + "source", + "config_sha256", + "automodel", + "gpu", + "pair", + "sample_ids", + "diagnostics", + "teacher_calls_per_rank", + "checkpoint", + "resume", +} + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--stage", choices=("train-one", "resume-one", "validate"), required=True) + parser.add_argument("--run-root", type=pathlib.Path, required=True) + parser.add_argument("--before-automodel", type=pathlib.Path) + parser.add_argument("--after-automodel", type=pathlib.Path) + return parser.parse_args() + + +def _sha256(path: pathlib.Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _canonical_sha256(value: Any) -> str: + payload = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + return hashlib.sha256(payload.encode()).hexdigest() + + +def _require_sha256(value: Any, *, name: str) -> str: + if not isinstance(value, str) or len(value) != 64: + raise ValueError(f"{name} is not a SHA-256 digest") + try: + int(value, 16) + except ValueError as error: + raise ValueError(f"{name} is not a hexadecimal SHA-256 digest") from error + return value.lower() + + +def _unique_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + value: dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise ValueError(f"JSON object contains duplicate key {key!r}") + value[key] = item + return value + + +def _reject_json_constant(token: str) -> None: + raise ValueError(f"JSON contains non-finite value {token}") + + +def _read_json(path: pathlib.Path) -> dict[str, Any]: + value = json.loads( + path.read_bytes(), + object_pairs_hook=_unique_json_object, + parse_constant=_reject_json_constant, + ) + if not isinstance(value, dict): + raise TypeError(f"{path} must contain a JSON object") + return value + + +def _finite_positive(value: Any, *, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise TypeError(f"{name} must be a real number") + value = float(value) + if not math.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be finite and positive") + return value + + +def _absolute_path_without_symlinks(value: pathlib.Path, *, name: str) -> pathlib.Path: + path = pathlib.Path(os.path.abspath(value)) + for candidate in (*reversed(path.parents), path): + if candidate.is_symlink(): + raise ValueError(f"{name} cannot traverse a symlink: {candidate}") + return path + + +def _regular_directory(value: pathlib.Path, *, name: str) -> pathlib.Path: + path = _absolute_path_without_symlinks(value, name=name) + if not path.is_dir(): + raise ValueError(f"{name} must identify a regular directory") + return path.resolve() + + +def _regular_file(value: pathlib.Path, *, name: str) -> pathlib.Path: + path = _absolute_path_without_symlinks(value, name=name) + if not path.is_file(): + raise ValueError(f"{name} must identify a regular file") + return path.resolve() + + +def _create_run_root(value: pathlib.Path) -> pathlib.Path: + path = _absolute_path_without_symlinks(value, name="smoke run root") + path.mkdir(parents=True, exist_ok=True) + return _regular_directory(path, name="smoke run root") + + +def _relative_regular_file(root: pathlib.Path, value: Any, *, name: str) -> pathlib.Path: + if not isinstance(value, str) or not value: + raise TypeError(f"{name} must be a non-empty relative path") + relative = pathlib.PurePosixPath(value) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"{name} must stay beneath the run root") + root = _regular_directory(root, name=f"{name} root") + path = root.joinpath(*relative.parts) + if any(candidate.is_symlink() for candidate in (path, *path.parents) if candidate != root): + raise ValueError(f"{name} cannot traverse a symlink") + resolved = path.resolve() + try: + resolved.relative_to(root) + except ValueError as error: + raise ValueError(f"{name} must stay beneath the run root") from error + if not resolved.is_file(): + raise ValueError(f"{name} must identify a regular file") + return resolved + + +def validate_stage_result(value: Mapping[str, Any], *, stage: str) -> None: + """Validate one atomic training-stage result without importing GPU dependencies.""" + if set(value) != _STAGE_RESULT_KEYS: + raise ValueError(f"{stage} result keys are incompatible") + if value["schema_version"] != 1 or value["record_type"] != "pdd_qwen_smoke_stage": + raise ValueError(f"{stage} result schema is incompatible") + if value["stage"] != stage or stage not in _EXPECTED_PAIRS: + raise ValueError("smoke stage identity is invalid") + if type(value["pid"]) is not int or value["pid"] <= 0: + raise ValueError("smoke stage pid is invalid") + if type(value["world_size"]) is not int or value["world_size"] < 2: + raise ValueError("full-Qwen smoke requires a multi-GPU world") + if value["model"] != {"id": _MODEL_ID, "revision": _MODEL_REVISION, "dtype": "bfloat16"}: + raise ValueError("smoke model identity is invalid") + if value["pdd"] != { + "grid_size": 128, + "flow_shift": 5.0, + "block_size_min": 4, + "block_size_max": 64, + "teacher_integrator": "euler", + "guidance_scale": 4.0, + "guidance_rescale": 1.0, + "guidance_eps": 1e-5, + }: + raise ValueError("smoke PDD identity is invalid") + source = value["source"] + if ( + not isinstance(source, Mapping) + or set(source) != {"commit", "dirty"} + or not isinstance(source["commit"], str) + or len(source["commit"]) != 40 + or source["dirty"] is not False + ): + raise ValueError("smoke source identity is invalid") + try: + int(source["commit"], 16) + except ValueError as error: + raise ValueError("smoke source commit is not hexadecimal") from error + _require_sha256(value["config_sha256"], name="config_sha256") + automodel = value["automodel"] + expected_automodel_keys = { + "distribution", + "version", + "package_tree_sha256", + "wheel_sha256", + "runtime_versions", + } + if not isinstance(automodel, Mapping) or set(automodel) != expected_automodel_keys: + raise ValueError("smoke AutoModel identity is invalid") + if ( + automodel["distribution"] != "nemo_automodel" + or automodel["version"] != "0.5.0" + or automodel["runtime_versions"] != {"diffusers": "0.38.0"} + or automodel["package_tree_sha256"] != _AUTOMODEL_TREE_SHA256 + or automodel["wheel_sha256"] != _AUTOMODEL_WHEEL_SHA256 + ): + raise ValueError("smoke AutoModel release identity is invalid") + _require_sha256(automodel["package_tree_sha256"], name="automodel.package_tree_sha256") + _require_sha256(automodel["wheel_sha256"], name="automodel.wheel_sha256") + gpu = value["gpu"] + if not isinstance(gpu, Mapping) or set(gpu) != { + "names", + "total_memory_bytes", + "host_available_bytes", + "allocated_before_step_bytes", + "peak_memory_bytes", + "student_parameter_bytes", + "teacher_parameter_bytes", + "step_seconds", + }: + raise ValueError("smoke GPU evidence is invalid") + if not isinstance(gpu["names"], list) or len(gpu["names"]) != value["world_size"]: + raise ValueError("smoke GPU inventory does not match world size") + if any(not isinstance(name, str) or not name for name in gpu["names"]): + raise ValueError("smoke GPU names are invalid") + for name in ( + "total_memory_bytes", + "host_available_bytes", + "allocated_before_step_bytes", + "peak_memory_bytes", + ): + values = gpu[name] + if ( + not isinstance(values, list) + or len(values) != value["world_size"] + or any(type(item) is not int or item <= 0 for item in values) + ): + raise ValueError(f"smoke GPU {name} is invalid") + for name in ("student_parameter_bytes", "teacher_parameter_bytes"): + if type(gpu[name]) is not int or gpu[name] <= 0: + raise ValueError(f"smoke capacity {name} is invalid") + for allocated, peak, total in zip( + gpu["allocated_before_step_bytes"], + gpu["peak_memory_bytes"], + gpu["total_memory_bytes"], + strict=True, + ): + if not allocated <= peak <= total: + raise ValueError("smoke GPU allocation evidence is inconsistent") + _finite_positive(gpu["step_seconds"], name="gpu.step_seconds") + if value["pair"] != {"n": _EXPECTED_PAIRS[stage][0], "k": _EXPECTED_PAIRS[stage][1]}: + raise ValueError("smoke explicit support pair is invalid") + sample_ids = value["sample_ids"] + if ( + not isinstance(sample_ids, list) + or len(sample_ids) != value["world_size"] + or any(not isinstance(item, str) or not item for item in sample_ids) + or len(set(sample_ids)) != len(sample_ids) + ): + raise ValueError("smoke sample IDs are invalid") + expected_step = 1 if stage == "train-one" else 2 + expected_ids = [ + f"synthetic-pdd-smoke-step-{expected_step}-rank-{rank}" + for rank in range(value["world_size"]) + ] + if sample_ids != expected_ids: + raise ValueError("smoke sample IDs do not match the canonical rank order") + diagnostics = value["diagnostics"] + if not isinstance(diagnostics, Mapping) or set(diagnostics) != { + "completed_step", + "loss", + "grad_norm", + "student_adamw_nominal_update_ratio", + "pdd_projection_update_ratio", + "learning_rate", + "student_velocity_rms", + "teacher_velocity_rms", + "student_teacher_velocity_rms_ratio", + "reconstructed_state_rms", + }: + raise ValueError("smoke diagnostics are invalid") + if diagnostics["completed_step"] != expected_step: + raise ValueError("smoke completed step is invalid") + for name in diagnostics: + if name != "completed_step": + _finite_positive(diagnostics[name], name=f"diagnostics.{name}") + calls = value["teacher_calls_per_rank"] + if calls != [2] * value["world_size"]: + raise ValueError("guided teacher call structure is invalid") + checkpoint = value["checkpoint"] + if not isinstance(checkpoint, Mapping) or set(checkpoint) != { + "path", + "manifest_sha256", + "completed_steps", + "parent_checkpoint", + }: + raise ValueError("smoke checkpoint evidence is invalid") + if checkpoint["completed_steps"] != expected_step: + raise ValueError("smoke checkpoint step is invalid") + expected_parent = None if stage == "train-one" else "step_00000001" + if checkpoint["parent_checkpoint"] != expected_parent: + raise ValueError("smoke checkpoint lineage is invalid") + if checkpoint["path"] != f"checkpoints/step_{expected_step:08d}": + raise ValueError("smoke checkpoint path is invalid") + _require_sha256(checkpoint["manifest_sha256"], name="checkpoint.manifest_sha256") + resume = value["resume"] + if stage == "train-one": + if resume is not None: + raise ValueError("first smoke stage cannot have resume evidence") + elif resume != { + "selected_checkpoint": "step_00000001", + "completed_steps": 1, + "parent_checkpoint": None, + "first_sample_ids": sample_ids, + "learning_rate": diagnostics["learning_rate"], + }: + raise ValueError("smoke resume evidence is invalid") + + +def validate_inference_result(value: Mapping[str, Any], *, root: pathlib.Path) -> None: + """Validate the exact authenticated PDD-4 inference evidence.""" + if value.get("schema_version") != 1 or value.get("record_type") != "pdd_inference": + raise ValueError("PDD inference result schema is invalid") + if value.get("condition") != "pdd_4" or value.get("schedule") != "pdd-4": + raise ValueError("PDD inference schedule identity is invalid") + if value.get("blocks") != [32, 32, 32, 32]: + raise ValueError("PDD-4 blocks are invalid") + if value.get("height") != 1024 or value.get("width") != 1024: + raise ValueError("PDD-4 smoke output must be exactly 1024x1024") + if ( + value.get("scheduler_steps") != 4 + or value.get("actual_transformer_invocations") != 4 + or value.get("batch_normalized_transformer_evaluations") != 4 + ): + raise ValueError("PDD-4 compute counters are invalid") + _finite_positive(value.get("latency_seconds"), name="inference.latency_seconds") + output = value.get("output") + if not isinstance(output, Mapping) or set(output) != {"path", "sha256"}: + raise ValueError("PDD inference output evidence is invalid") + image = _relative_regular_file(root, output["path"], name="inference.output.path") + if image.suffix.lower() != ".png" or image.read_bytes()[:8] != b"\x89PNG\r\n\x1a\n": + raise ValueError("PDD inference output is not a PNG") + if image.stat().st_size <= 8 or _sha256(image) != _require_sha256( + output["sha256"], name="inference.output.sha256" + ): + raise ValueError("PDD inference PNG hash is invalid") + + +def _exact_mapping(value: Any, keys: set[str], *, name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping) or set(value) != keys: + raise ValueError(f"{name} must contain exactly {sorted(keys)}") + return value + + +def _validate_checkpoint_identity(identity: Any, *, stage: Mapping[str, Any]) -> None: + from modelopt.torch.fastgen import PDDMetadata + + identity = _exact_mapping( + identity, + { + "schema_version", + "model", + "pdd_metadata", + "guidance", + "automodel", + "data", + "topology", + "training", + "optimizer", + "scheduler", + }, + name="smoke checkpoint identity", + ) + if identity["schema_version"] != 1 or identity["model"] != stage["model"]: + raise ValueError("smoke checkpoint model identity is incompatible") + metadata = PDDMetadata.from_dict(identity["pdd_metadata"]) + if metadata.to_dict() != identity["pdd_metadata"]: + raise ValueError("smoke checkpoint PDD metadata is not canonical") + pdd = stage["pdd"] + if ( + metadata.grid_size != pdd["grid_size"] + or metadata.flow_shift != pdd["flow_shift"] + or metadata.block_size_min != pdd["block_size_min"] + or metadata.block_size_max != pdd["block_size_max"] + or metadata.teacher_integrator != pdd["teacher_integrator"] + or metadata.inference_blocks != (32, 32, 32, 32) + or metadata.layer_spec.to_dict() + != { + "projection_path": "transformer.proj_out", + "head_layout": "channel_major", + "output_channels": None, + } + ): + raise ValueError("smoke checkpoint PDD metadata does not match the stage") + if identity["guidance"] != { + "scale": pdd["guidance_scale"], + "rescale": pdd["guidance_rescale"], + "eps": pdd["guidance_eps"], + }: + raise ValueError("smoke checkpoint guidance does not match the stage") + if identity["automodel"] != stage["automodel"]: + raise ValueError("smoke checkpoint AutoModel identity does not match the stage") + if identity["topology"] != { + "world_size": stage["world_size"], + "pure_data_parallel": True, + }: + raise ValueError("smoke checkpoint topology does not match the stage") + + +def _validate_automodel_snapshot( + snapshot: Any, + *, + expected_automodel: Mapping[str, Any], +) -> None: + snapshot = _exact_mapping( + snapshot, + { + "distribution", + "files", + "import_origin", + "package_file_count", + "package_tree_sha256", + "release_commit", + "release_tag", + "root", + "runtime_versions", + "version", + "wheel", + "wheel_sha256", + }, + name="AutoModel snapshot", + ) + for key in ("distribution", "version", "runtime_versions"): + if snapshot[key] != expected_automodel[key]: + raise ValueError(f"AutoModel snapshot identity differs for {key}") + for key in ("package_tree_sha256", "wheel_sha256"): + if ( + _require_sha256(snapshot[key], name=f"AutoModel snapshot {key}") + != expected_automodel[key] + ): + raise ValueError(f"AutoModel snapshot identity differs for {key}") + if ( + snapshot["release_commit"] != _AUTOMODEL_RELEASE_COMMIT + or snapshot["release_tag"] != _AUTOMODEL_RELEASE_TAG + or snapshot["wheel"] != _AUTOMODEL_WHEEL + or type(snapshot["package_file_count"]) is not int + or snapshot["package_file_count"] != _AUTOMODEL_PACKAGE_FILE_COUNT + ): + raise ValueError("AutoModel snapshot release identity is invalid") + root_value = snapshot["root"] + origin_value = snapshot["import_origin"] + if not isinstance(root_value, str) or not isinstance(origin_value, str): + raise TypeError("AutoModel snapshot root and import origin must be strings") + root = pathlib.Path(root_value) + import_origin = pathlib.Path(origin_value) + if not root.is_absolute() or not import_origin.is_absolute(): + raise ValueError("AutoModel snapshot paths must be absolute") + try: + import_origin.relative_to(root) + except ValueError as error: + raise ValueError("AutoModel snapshot import origin is outside its root") from error + files = snapshot["files"] + if not isinstance(files, list) or len(files) != _AUTOMODEL_PACKAGE_FILE_COUNT: + raise ValueError("AutoModel snapshot file inventory is invalid") + tree = hashlib.sha256() + previous_path: str | None = None + for index, raw_record in enumerate(files): + record = _exact_mapping( + raw_record, + {"path", "sha256", "size"}, + name=f"AutoModel snapshot files[{index}]", + ) + path = record["path"] + if ( + not isinstance(path, str) + or not path + or pathlib.PurePosixPath(path).is_absolute() + or "\\" in path + or any(part in ("", ".", "..") for part in path.split("/")) + or (previous_path is not None and path <= previous_path) + ): + raise ValueError("AutoModel snapshot paths must be sorted normalized references") + digest = _require_sha256(record["sha256"], name=f"AutoModel snapshot files[{index}].sha256") + if type(record["size"]) is not int or record["size"] < 0: + raise ValueError(f"AutoModel snapshot files[{index}].size is invalid") + tree.update(path.encode()) + tree.update(b"\0") + tree.update(digest.encode()) + tree.update(b"\0") + tree.update(str(record["size"]).encode()) + tree.update(b"\n") + previous_path = path + if tree.hexdigest() != snapshot["package_tree_sha256"]: + raise ValueError("AutoModel snapshot inventory does not match its tree digest") + + +def _validate_bundle_links( + *, + stage1: Mapping[str, Any], + stage2: Mapping[str, Any], + manifest1: Mapping[str, Any], + manifest2: Mapping[str, Any], + export_manifest: Mapping[str, Any], + automodel_snapshot: Mapping[str, Any], +) -> None: + identity1 = manifest1.get("identity") + identity2 = manifest2.get("identity") + if identity1 != identity2: + raise ValueError("smoke checkpoint identities differ across resume") + _validate_checkpoint_identity(identity1, stage=stage1) + _validate_checkpoint_identity(identity2, stage=stage2) + if export_manifest.get("identity") != identity2: + raise ValueError("smoke export identity does not match step 2") + if export_manifest.get("modelopt_source") != stage2["source"]: + raise ValueError("smoke export source does not match the training source") + expected_checkpoint = { + "name": "step_00000002", + "manifest_sha256": stage2["checkpoint"]["manifest_sha256"], + "completed_steps": 2, + } + if export_manifest.get("source_checkpoint") != expected_checkpoint: + raise ValueError("smoke export does not derive from the exact step-2 checkpoint") + _validate_automodel_snapshot( + automodel_snapshot, + expected_automodel=stage2["automodel"], + ) + + +def _load_matching_automodel_snapshots( + before_path: pathlib.Path, + after_path: pathlib.Path, + *, + expected_automodel: Mapping[str, Any], +) -> Mapping[str, Any]: + before_path = _regular_file(before_path, name="before AutoModel snapshot") + after_path = _regular_file(after_path, name="after AutoModel snapshot") + before_bytes = before_path.read_bytes() + if before_bytes != after_path.read_bytes(): + raise ValueError("AutoModel package snapshot changed during full-Qwen smoke") + snapshot = _read_json(before_path) + _validate_automodel_snapshot(snapshot, expected_automodel=expected_automodel) + return snapshot + + +def _raw_config(run_root: pathlib.Path, world_size: int) -> dict[str, Any]: + return { + "model": { + "pretrained_model_name_or_path": _MODEL_ID, + "revision": _MODEL_REVISION, + "torch_dtype": "bfloat16", + "device": "cuda", + "transformer_engine_linear": False, + "peft": None, + "guidance_embeds": False, + "fuse_qkv_projections": False, + }, + "pdd": { + "pred_type": "flow", + "num_train_timesteps": None, + "guidance_scale": 4.0, + "student_sample_steps": 4, + "student_sample_type": "ode", + "grid_size": 128, + "flow_shift": 5.0, + "block_size_min": 4, + "block_size_max": 64, + "teacher_integrator": "euler", + "inference_blocks": [32, 32, 32, 32], + "data_free": False, + }, + "optim": { + "learning_rate": 2.0e-5, + "weight_decay": 0.01, + "betas": [0.9, 0.999], + "eps": 1.0e-8, + }, + "guidance": {"rescale": 1.0, "eps": 1e-5}, + "training": { + "seed": 42, + "max_steps": 2, + "max_grad_norm": 1.0, + "zero_grad_warmup_steps": 0, + "log_every_steps": 1, + "checkpoint_every_steps": 1, + "validation_every_steps": 1000, + "grad_accumulation_steps": 1, + "global_batch_size": world_size, + "validation_seed": 2026, + }, + "fsdp": { + "dp_size": world_size, + "tp_size": 1, + "cp_size": 1, + "pp_size": 1, + "ep_size": 1, + "activation_checkpointing": True, + }, + "data": { + "all_metadata_index": "synthetic_metadata.json", + "validation_metadata_index": "synthetic_heldout.json", + "dataloader": { + "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", + "cache_dir": "synthetic-unused", + "metadata_index": "synthetic_train.json", + "base_resolution": [1024, 1024], + "batch_size": 1, + "drop_last": True, + "shuffle": False, + "dynamic_batch_size": False, + "negative_prompt_embedding_path": "synthetic_negative.pt", + }, + }, + "checkpoint": { + "enabled": True, + "checkpoint_dir": str(run_root / "checkpoints"), + "model_save_format": "torch_save", + "save_consolidated": False, + "restore_from": "LATEST", + }, + } + + +def _modelopt_source() -> dict[str, Any]: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + dirty = bool( + subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=normal"], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + ).stdout + ) + if dirty: + raise RuntimeError("full-Qwen smoke requires a clean ModelOpt checkout") + return {"commit": commit, "dirty": False} + + +def _ordered_id_sha256(sample_ids: Sequence[str], *, split: str) -> str: + digest = hashlib.sha256() + digest.update(f"modelopt-pdd-ordered-{split}-ids-v1\0".encode()) + for sample_id in sample_ids: + digest.update(sample_id.encode()) + digest.update(b"\n") + return digest.hexdigest() + + +def _training_sample_ids(world_size: int) -> tuple[str, ...]: + return tuple( + f"synthetic-pdd-smoke-step-{step}-rank-{rank}" + for step in (1, 2) + for rank in range(world_size) + ) + + +def _build_sampler(world_size: int, rank: int) -> Any: + import torch + from fastgen_data import ReplayableBatchSampler + from torch.utils.data import Sampler + + sample_ids = _training_sample_ids(world_size) + + class _Dataset: + metadata = [{"sample_id": sample_id} for sample_id in sample_ids] + + class _Sampler(Sampler[list[int]]): + def __init__(self) -> None: + self.dataset = _Dataset() + self.rank = rank + self.num_replicas = world_size + self.epoch = 0 + self.batches_yielded = 0 + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + self.epoch = int(state["epoch"]) + self.batches_yielded = int(state["batches_yielded"]) + + def __iter__(self): + batches = ([rank], [world_size + rank]) + for index, batch in enumerate(batches): + if index >= self.batches_yielded: + self.batches_yielded = index + 1 + yield list(batch) + + def __len__(self) -> int: + return 2 + + assert torch.distributed.get_world_size() == world_size + return ReplayableBatchSampler(_Sampler()) + + +def _prepared_batch( + setup: Any, *, rank: int, step: int, device: Any, dtype: Any +) -> tuple[Any, Any]: + import torch + from pdd_training import PreparedPDDBatch + + config = getattr(setup.student, "config", None) + in_channels = getattr(config, "in_channels", None) + condition_width = getattr(config, "joint_attention_dim", None) + if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: + raise RuntimeError("pinned Qwen config has invalid in_channels") + if type(condition_width) is not int or condition_width <= 0: + raise RuntimeError("pinned Qwen config has invalid joint_attention_dim") + generator = torch.Generator(device="cpu").manual_seed(10_000 + rank * 10 + step) + latent_shape = (1, in_channels // 4, 128, 128) + data = torch.randn(latent_shape, generator=generator, dtype=torch.float32).to( + device=device, dtype=dtype + ) + noise = torch.randn(latent_shape, generator=generator, dtype=torch.float32).to(device=device) + condition = torch.randn((1, 8, condition_width), generator=generator, dtype=torch.float32).to( + device=device, dtype=dtype + ) + negative = torch.randn((1, 8, condition_width), generator=generator, dtype=torch.float32).to( + device=device, dtype=dtype + ) + mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 0]], device=device, dtype=torch.long) + negative_mask = torch.tensor([[1, 1, 1, 1, 0, 0, 0, 0]], device=device, dtype=torch.long) + sample_id = f"synthetic-pdd-smoke-step-{step}-rank-{rank}" + batch = PreparedPDDBatch( + data=data, + condition=(condition, mask), + negative_condition=(negative, negative_mask), + sample_ids=(sample_id,), + valid_mask=(True,), + ) + return batch, noise + + +def _identity( + *, setup: Any, training: Any, config: Any, sampler: Any, raw: Mapping[str, Any] +) -> dict[str, Any]: + from pdd_checkpoint import build_pdd_checkpoint_identity + + train_ids = tuple(item["sample_id"] for item in sampler.dataset.metadata) + heldout_ids = ("synthetic-pdd-smoke-heldout-not-evaluated",) + return build_pdd_checkpoint_identity( + metadata=setup.metadata, + model_id=config.model_id, + model_revision=config.model_revision, + guidance_scale=config.pdd.guidance_scale, + guidance_rescale=config.guidance.rescale, + guidance_eps=config.guidance.eps, + automodel_snapshot=setup.automodel_snapshot, + ordered_train_id_sha256=_ordered_id_sha256(train_ids, split="train"), + ordered_heldout_id_sha256=_ordered_id_sha256(heldout_ids, split="heldout"), + dataset_snapshot_sha256=_canonical_sha256( + {"domain": "modelopt-pdd-synthetic-smoke-v1", "config": raw["pdd"]} + ), + local_batch_size=1, + grad_accumulation_steps=1, + training_seed=config.training.seed, + validation_seed=config.training.validation_seed, + validation_every_steps=config.training.validation_every_steps, + max_grad_norm=config.training.max_grad_norm, + zero_grad_warmup_steps=config.training.zero_grad_warmup_steps, + activation_checkpointing=config.parallel.activation_checkpointing, + dtype="bfloat16", + optimizer=setup.optimizer, + scheduler=training.scheduler, + ) + + +def _diagnostics_dict(diagnostics: Any) -> dict[str, Any]: + return { + "completed_step": diagnostics.completed_step, + "loss": diagnostics.loss, + "grad_norm": diagnostics.grad_norm, + "student_adamw_nominal_update_ratio": diagnostics.student_adamw_nominal_update_ratio, + "pdd_projection_update_ratio": diagnostics.pdd_projection_update_ratio, + "learning_rate": diagnostics.learning_rate, + "student_velocity_rms": diagnostics.student_velocity_rms, + "teacher_velocity_rms": diagnostics.teacher_velocity_rms, + "student_teacher_velocity_rms_ratio": diagnostics.student_teacher_velocity_rms_ratio, + "reconstructed_state_rms": diagnostics.reconstructed_state_rms, + } + + +def _run_training_stage(stage: str, run_root: pathlib.Path) -> None: + if os.environ.get("HF_HUB_OFFLINE") != "1": + raise RuntimeError("full-Qwen smoke requires HF_HUB_OFFLINE=1 and a pinned local snapshot") + import torch + import torch.distributed as dist + from export_pdd_qwen_image import host_available_bytes + from pdd_artifacts import write_canonical_json + from pdd_checkpoint import PDDCheckpointManager, validate_pdd_training_checkpoint + from pdd_recipe import ( + build_pdd_setup, + build_pdd_training_artifacts, + initialize_pdd_distributed, + resolve_pdd_recipe_config, + ) + + initialize_pdd_distributed(backend="nccl", timeout_minutes=60) + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size < 2: + raise RuntimeError("full-Qwen smoke requires a multi-GPU FSDP2 world") + device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + torch.cuda.set_device(device) + run_root = _create_run_root(run_root) + result_path = run_root / ("stage1.json" if stage == "train-one" else "stage2.json") + if result_path.exists() or result_path.is_symlink(): + raise FileExistsError(f"smoke stage result already exists: {result_path}") + + raw = _raw_config(run_root, world_size) + config = resolve_pdd_recipe_config(raw) + setup = build_pdd_setup(config) + training = build_pdd_training_artifacts(setup, config) + sampler = _build_sampler(world_size, rank) + identity = _identity(setup=setup, training=training, config=config, sampler=sampler, raw=raw) + manager = PDDCheckpointManager( + root=config.checkpoint.checkpoint_dir, + checkpointer=setup.checkpointer, + model=setup.student, + optimizer=setup.optimizer, + scheduler=training.scheduler, + trainer=training.trainer, + sampler=sampler, + rng=training.rng, + identity=identity, + ) + resume_payload = None + try: + if stage == "train-one": + if manager.resolve("LATEST") is not None: + raise RuntimeError("first smoke stage requires an empty checkpoint root") + step = 1 + else: + resume = manager.load("LATEST") + if resume is None: + raise RuntimeError("resume smoke stage found no LATEST checkpoint") + if ( + resume.checkpoint_path.name != "step_00000001" + or resume.completed_steps != 1 + or resume.parent_checkpoint is not None + or training.trainer.completed_steps != 1 + ): + raise RuntimeError("resume smoke stage restored incompatible lineage") + expected_ids = sampler.expected_next_sample_ids() + resume.verify_first_batch(expected_ids) + resume_payload = { + "selected_checkpoint": resume.checkpoint_path.name, + "completed_steps": resume.completed_steps, + "parent_checkpoint": resume.parent_checkpoint, + "first_sample_ids": list(expected_ids), + "learning_rate": float(setup.optimizer.param_groups[0]["lr"]), + } + step = 2 + + batch, noise = _prepared_batch( + setup, + rank=rank, + step=step, + device=device, + dtype=config.dtype, + ) + if sampler.expected_next_sample_ids() != batch.sample_ids: + raise RuntimeError("synthetic smoke batch does not match the committed sampler cursor") + n_value, k_value = _EXPECTED_PAIRS[stage] + n = torch.tensor([n_value], device=device, dtype=torch.int64) + k = torch.tensor([k_value], device=device, dtype=torch.int64) + teacher_calls = 0 + + def count_teacher_call(_module: Any, _args: Any, _kwargs: Any) -> None: + nonlocal teacher_calls + teacher_calls += 1 + + hook = setup.teacher.register_forward_pre_hook(count_teacher_call, with_kwargs=True) + allocated_before_step = torch.cuda.memory_allocated(device) + torch.cuda.synchronize(device) + started = time.perf_counter() + try: + diagnostics = training.trainer.train_step( + batch, + noise=noise, + n=n, + k=k, + measure_updates=True, + ) + training.scheduler.step() + sampler.commit(batch.sample_ids) + finally: + hook.remove() + torch.cuda.synchronize(device) + step_seconds = time.perf_counter() - started + calls = [None] * world_size + dist.all_gather_object(calls, teacher_calls) + if calls != [2] * world_size: + raise RuntimeError(f"guided teacher calls differ across ranks: {calls}") + checkpoint = manager.save() + manifest = validate_pdd_training_checkpoint( + checkpoint, + expected_identity=identity, + expected_world_size=world_size, + ) + + sample_ids = [None] * world_size + dist.all_gather_object(sample_ids, batch.sample_ids[0]) + gpu_name = torch.cuda.get_device_name(device) + total_memory = torch.cuda.get_device_properties(device).total_memory + host_available = host_available_bytes() + peak_memory = torch.cuda.max_memory_allocated(device) + gpu_names = [None] * world_size + total_memories = [None] * world_size + host_memories = [None] * world_size + allocated_memories = [None] * world_size + peak_memories = [None] * world_size + dist.all_gather_object(gpu_names, gpu_name) + dist.all_gather_object(total_memories, total_memory) + dist.all_gather_object(host_memories, host_available) + dist.all_gather_object(allocated_memories, allocated_before_step) + dist.all_gather_object(peak_memories, peak_memory) + seconds = torch.tensor(step_seconds, device=device, dtype=torch.float64) + dist.all_reduce(seconds, op=dist.ReduceOp.MAX) + automodel = { + key: setup.automodel_snapshot[key] + for key in ( + "distribution", + "version", + "package_tree_sha256", + "wheel_sha256", + "runtime_versions", + ) + } + result = { + "schema_version": 1, + "record_type": "pdd_qwen_smoke_stage", + "stage": stage, + "pid": os.getpid(), + "world_size": world_size, + "model": {"id": _MODEL_ID, "revision": _MODEL_REVISION, "dtype": "bfloat16"}, + "pdd": { + "grid_size": 128, + "flow_shift": 5.0, + "block_size_min": 4, + "block_size_max": 64, + "teacher_integrator": "euler", + "guidance_scale": 4.0, + "guidance_rescale": 1.0, + "guidance_eps": 1e-5, + }, + "source": _modelopt_source(), + "config_sha256": _canonical_sha256(raw), + "automodel": automodel, + "gpu": { + "names": gpu_names, + "total_memory_bytes": total_memories, + "host_available_bytes": host_memories, + "allocated_before_step_bytes": allocated_memories, + "peak_memory_bytes": peak_memories, + "student_parameter_bytes": sum( + parameter.numel() * parameter.element_size() + for parameter in setup.student.parameters() + ), + "teacher_parameter_bytes": sum( + parameter.numel() * parameter.element_size() + for parameter in setup.teacher.parameters() + ), + "step_seconds": float(seconds.item()), + }, + "pair": {"n": n_value, "k": k_value}, + "sample_ids": sample_ids, + "diagnostics": _diagnostics_dict(diagnostics), + "teacher_calls_per_rank": calls, + "checkpoint": { + "path": checkpoint.relative_to(run_root).as_posix(), + "manifest_sha256": _sha256(checkpoint / "manifest.json"), + "completed_steps": manifest["completed_steps"], + "parent_checkpoint": manifest["parent_checkpoint"], + }, + "resume": resume_payload, + } + validate_stage_result(result, stage=stage) + if rank == 0: + write_canonical_json(result_path, result) + dist.barrier() + finally: + setup.checkpointer.close() + dist.destroy_process_group() + + +def _validate_bundle( + run_root: pathlib.Path, + before_automodel: pathlib.Path, + after_automodel: pathlib.Path, +) -> pathlib.Path: + from pdd_artifacts import load_canonical_json, write_canonical_json + from pdd_checkpoint import validate_pdd_training_checkpoint + from pdd_export import inspect_pdd_export + + run_root = _regular_directory(run_root, name="smoke run root") + stage1_path = _relative_regular_file(run_root, "stage1.json", name="stage-1 result") + stage2_path = _relative_regular_file(run_root, "stage2.json", name="stage-2 result") + export_root = _regular_directory(run_root / "export", name="smoke export root") + export_manifest_path = _regular_file( + export_root / "manifest.json", name="smoke export manifest" + ) + inference_path = _relative_regular_file( + run_root, "inference/pdd4.json", name="smoke inference result" + ) + output_path = run_root / "smoke_result.json" + if output_path.exists() or output_path.is_symlink(): + raise FileExistsError(f"smoke result already exists: {output_path}") + + stage1 = load_canonical_json(stage1_path) + stage2 = load_canonical_json(stage2_path) + if not isinstance(stage1, Mapping) or not isinstance(stage2, Mapping): + raise TypeError("smoke stage results must be JSON objects") + validate_stage_result(stage1, stage="train-one") + validate_stage_result(stage2, stage="resume-one") + if stage1["pid"] == stage2["pid"]: + raise ValueError("forced-resume stages did not use fresh processes") + for name in ("world_size", "model", "pdd", "source", "config_sha256", "automodel"): + if stage1[name] != stage2[name]: + raise ValueError(f"smoke stages disagree on {name}") + checkpoint1 = _relative_regular_file( + run_root, + pathlib.PurePosixPath(stage1["checkpoint"]["path"]).joinpath("manifest.json").as_posix(), + name="stage1.checkpoint.manifest", + ).parent + checkpoint2 = _relative_regular_file( + run_root, + pathlib.PurePosixPath(stage2["checkpoint"]["path"]).joinpath("manifest.json").as_posix(), + name="stage2.checkpoint.manifest", + ).parent + if _sha256(checkpoint1 / "manifest.json") != stage1["checkpoint"]["manifest_sha256"]: + raise ValueError("stage-1 checkpoint manifest hash changed") + if _sha256(checkpoint2 / "manifest.json") != stage2["checkpoint"]["manifest_sha256"]: + raise ValueError("stage-2 checkpoint manifest hash changed") + manifest1 = validate_pdd_training_checkpoint( + checkpoint1, expected_world_size=stage1["world_size"] + ) + manifest2 = validate_pdd_training_checkpoint( + checkpoint2, expected_world_size=stage2["world_size"] + ) + if manifest1["completed_steps"] != 1 or manifest2["completed_steps"] != 2: + raise ValueError("smoke checkpoint steps are invalid") + if ( + manifest1["parent_checkpoint"] is not None + or manifest2["parent_checkpoint"] != checkpoint1.name + ): + raise ValueError("smoke checkpoint parent lineage is invalid") + export_descriptor = inspect_pdd_export(export_root) + if export_descriptor.root != export_root: + raise ValueError("smoke export descriptor resolved an unexpected root") + inference = load_canonical_json(inference_path) + if not isinstance(inference, Mapping): + raise TypeError("smoke inference result must be a JSON object") + validate_inference_result(inference, root=inference_path.parent) + export_sha256 = _sha256(export_manifest_path) + if inference.get("export_manifest_sha256") != export_sha256: + raise ValueError("inference does not authenticate the smoke export") + automodel_snapshot = _load_matching_automodel_snapshots( + before_automodel, + after_automodel, + expected_automodel=stage2["automodel"], + ) + _validate_bundle_links( + stage1=stage1, + stage2=stage2, + manifest1=manifest1, + manifest2=manifest2, + export_manifest=export_descriptor.manifest, + automodel_snapshot=automodel_snapshot, + ) + result = { + "schema_version": 1, + "record_type": "pdd_qwen_operability_smoke", + "status": "passed", + "stage1_sha256": _sha256(stage1_path), + "stage2_sha256": _sha256(stage2_path), + "checkpoint_manifest_sha256": [ + stage1["checkpoint"]["manifest_sha256"], + stage2["checkpoint"]["manifest_sha256"], + ], + "export_manifest_sha256": export_sha256, + "inference_result_sha256": _sha256(inference_path), + "automodel_snapshot_sha256": _sha256( + _regular_file(before_automodel, name="before AutoModel snapshot") + ), + "model": stage1["model"], + "pdd": stage1["pdd"], + "source": stage1["source"], + "config_sha256": stage1["config_sha256"], + "world_size": stage1["world_size"], + "pairs": [stage1["pair"], stage2["pair"]], + "losses": [stage1["diagnostics"]["loss"], stage2["diagnostics"]["loss"]], + "inference": { + "blocks": inference["blocks"], + "scheduler_steps": inference["scheduler_steps"], + "actual_transformer_invocations": inference["actual_transformer_invocations"], + "batch_normalized_transformer_evaluations": inference[ + "batch_normalized_transformer_evaluations" + ], + "output_sha256": inference["output"]["sha256"], + "latency_seconds": inference["latency_seconds"], + }, + } + write_canonical_json(output_path, result) + return output_path + + +def main() -> None: + args = _parse_args() + if args.stage == "validate": + if args.before_automodel is None or args.after_automodel is None: + raise ValueError("validate requires --before-automodel and --after-automodel") + print(_validate_bundle(args.run_root, args.before_automodel, args.after_automodel)) + return + if args.before_automodel is not None or args.after_automodel is not None: + raise ValueError("training stages do not accept AutoModel snapshot arguments") + _run_training_stage(args.stage, args.run_root) + + +if __name__ == "__main__": + main() 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..2d429d1aab4 --- /dev/null +++ b/tests/gpu/torch/fastgen/test_pdd_toy.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-GPU BF16 proof for the framework-neutral PDD core.""" + +from __future__ import annotations + +import importlib.util +import sys +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +from modelopt.torch.fastgen import ( + PDDConfig, + PDDLayerSpec, + PDDOutputProjection, + PDDPipeline, + convert_to_pdd_output_projection, +) + +if TYPE_CHECKING: + from pathlib import Path + +_FORBIDDEN_MODULES = ("diffusers", "fastgen", "nemo_automodel") +_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) -> torch.Tensor: + return self.projection(torch.tanh(self.backbone(state))) + + +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 + with projection.fuse_block(start, end, grid): + return model(state.to(self._model_dtype(model))) + + 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, + flow_shift=5.0, + block_size_min=1, + block_size_max=_GRID_SIZE, + inference_blocks=[2, 2], + student_sample_steps=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 _assert_optional_frameworks_absent() -> None: + resolvable = sorted(name for name in _FORBIDDEN_MODULES if importlib.util.find_spec(name)) + assert not resolvable, f"plain PDD GPU environment resolves optional frameworks: {resolvable}" + imported = sorted(name for name in _FORBIDDEN_MODULES if name in sys.modules) + assert not imported, f"plain PDD GPU test imported optional frameworks: {imported}" + + +def test_bf16_loss_gradient_update_reload_and_fused_sample(tmp_path: Path) -> None: + _assert_optional_frameworks_absent() + assert torch.cuda.is_available(), "Task-10 BF16 gate requires a real CUDA device" + device = torch.device("cuda", 0) + assert torch.cuda.get_device_capability(device)[0] >= 8, "BF16 gate 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) + optimizer = torch.optim.AdamW( + student.parameters(), + lr=2.0e-3, + weight_decay=0.0, + foreach=False, + fused=False, + ) + + optimizer.zero_grad(set_to_none=True) + 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 + weight_grad = projection.weight.grad.reshape(_GRID_SIZE, _WIDTH, _WIDTH) + bias_grad = projection.bias.grad.reshape(_GRID_SIZE, _WIDTH) + assert torch.count_nonzero(weight_grad[2]) > 0 + assert torch.count_nonzero(bias_grad[2]) > 0 + assert torch.count_nonzero(weight_grad[[0, 1, 3]]) == 0 + assert torch.count_nonzero(bias_grad[[0, 1, 3]]) == 0 + assert student.backbone.weight.grad is not None + assert torch.count_nonzero(student.backbone.weight.grad) > 0 + + gradients = [ + parameter.grad.float().square().sum() + for parameter in student.parameters() + if parameter.grad is not None + ] + grad_norm = torch.stack(gradients).sum().sqrt() + assert torch.isfinite(grad_norm) and grad_norm > 0 + before = {name: parameter.detach().clone() for name, parameter in student.named_parameters()} + optimizer.step() + update_norm = ( + torch.stack( + [ + (parameter.detach() - before[name]).float().square().sum() + for name, parameter in student.named_parameters() + ] + ) + .sum() + .sqrt() + ) + assert torch.isfinite(update_norm) and update_norm > 0 + + checkpoint = tmp_path / "pdd_bf16_state.pt" + torch.save(student.state_dict(), checkpoint) + saved = torch.load(checkpoint, map_location=device, weights_only=True) + assert saved.keys() == student.state_dict().keys() + assert all(value.dtype == torch.bfloat16 for value in saved.values()) + + restored, _teacher, restored_projection, restored_pipeline, restored_adapter = _build(device) + incompatible = restored.load_state_dict(saved, strict=True) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + assert restored_projection.weight.shape == projection.weight.shape + time = pipeline.time_grid(device)[n] + expected = adapter.student_all_heads(student, data.float(), time) + actual = restored_adapter.student_all_heads(restored, data.float(), time) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + sampled = restored_pipeline.sample(noise, blocks=[2, 2]) + assert sampled.dtype == torch.float32 + assert torch.isfinite(sampled).all() + assert restored_adapter.fused_calls == 2 + torch.cuda.synchronize(device) + _assert_optional_frameworks_absent() From 077b892a674e158fb9e433e56548338443787ead Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 13:09:51 -0700 Subject: [PATCH 14/45] fix(fastgen): harden Qwen PDD preflight Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd_finetune.py | 27 +++++++ examples/diffusers/fastgen/pdd_training.py | 79 ++++++++++++++----- .../torch/fastgen/plugins/qwen_image_pdd.py | 2 - .../pdd_training_preflight_distributed.py | 42 ++++++++++ .../fastgen/test_pdd_inference_checkpoint.py | 3 +- .../fastgen/test_pdd_training_lifecycle.py | 43 ++++++++++ .../fastgen/test_qwen_image_pdd_plugin.py | 8 +- 7 files changed, 176 insertions(+), 28 deletions(-) diff --git a/examples/diffusers/fastgen/pdd_finetune.py b/examples/diffusers/fastgen/pdd_finetune.py index 11ac80320b7..c5c241567fc 100644 --- a/examples/diffusers/fastgen/pdd_finetune.py +++ b/examples/diffusers/fastgen/pdd_finetune.py @@ -209,6 +209,8 @@ def _iter_validation_batches( dataloader: Any, masks: tuple[tuple[bool, ...], ...], config: Any, + expected_latent_channels: int, + expected_condition_features: int, ): from pdd_training import prepare_qwen_pdd_batch @@ -219,6 +221,8 @@ def _iter_validation_batches( device=config.device, dtype=config.dtype, require_negative_condition=config.pdd.guidance_scale is not None, + expected_latent_channels=expected_latent_channels, + expected_condition_features=expected_condition_features, ) yield dataclasses.replace(prepared, valid_mask=valid_mask) if count != len(masks): @@ -271,6 +275,8 @@ def _collective_training_batch( dtype: Any, require_negative_condition: bool, expected_batch_size: int, + expected_latent_channels: int, + expected_condition_features: int, ) -> tuple[Any, tuple[str, ...]] | None: """Prepare one rank-local batch, then agree on success before any model call.""" import torch.distributed as dist @@ -315,6 +321,8 @@ def _collective_training_batch( device=device, dtype=dtype, require_negative_condition=require_negative_condition, + expected_latent_channels=expected_latent_channels, + expected_condition_features=expected_condition_features, ) if prepared is None: raise RuntimeError("PDD batch preparation returned no prepared batch.") @@ -414,6 +422,21 @@ def main() -> None: config, ) setup = build_pdd_setup(config) + transformer_config = getattr(setup.student, "config", None) + if isinstance(transformer_config, Mapping): + in_channels = transformer_config.get("in_channels") + else: + in_channels = getattr(transformer_config, "in_channels", None) + if isinstance(transformer_config, Mapping): + condition_features = transformer_config.get("joint_attention_dim") + else: + condition_features = getattr(transformer_config, "joint_attention_dim", None) + if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: + raise RuntimeError("constructed Qwen transformer has invalid packed in_channels.") + if type(condition_features) is not int or condition_features <= 0: + raise RuntimeError("constructed Qwen transformer has invalid joint_attention_dim.") + expected_latent_channels = in_channels // 4 + expected_condition_features = condition_features training = build_pdd_training_artifacts(setup, config) identity = build_pdd_checkpoint_identity( metadata=setup.metadata, @@ -492,6 +515,8 @@ def main() -> None: dtype=config.dtype, require_negative_condition=config.pdd.guidance_scale is not None, expected_batch_size=config.training.local_batch_size, + expected_latent_channels=expected_latent_channels, + expected_condition_features=expected_condition_features, ) if next_batch is None: break @@ -584,6 +609,8 @@ def main() -> None: validation_dataloader, validation_masks, config, + expected_latent_channels, + expected_condition_features, ), validation_assignments, validation_seed=config.training.validation_seed, diff --git a/examples/diffusers/fastgen/pdd_training.py b/examples/diffusers/fastgen/pdd_training.py index 23869cb15e6..62e997e04b2 100644 --- a/examples/diffusers/fastgen/pdd_training.py +++ b/examples/diffusers/fastgen/pdd_training.py @@ -235,8 +235,16 @@ def prepare_qwen_pdd_batch( device: torch.device, dtype: torch.dtype, require_negative_condition: bool, + expected_latent_channels: int, + expected_condition_features: int, ) -> PreparedPDDBatch: """Move a portable Qwen cache batch into the PDD adapter contract.""" + if type(expected_latent_channels) is not int or expected_latent_channels <= 0: + raise ValueError("expected_latent_channels must be a positive integer.") + if type(expected_condition_features) is not int or expected_condition_features <= 0: + raise ValueError("expected_condition_features must be a positive integer.") + if not dtype.is_floating_point: + raise TypeError("Qwen PDD model dtype must be floating point.") if not isinstance(batch, Mapping): raise TypeError(f"batch must be a mapping, got {type(batch).__name__}.") required = {"image_latents", "text_embeddings", "text_embeddings_mask", "metadata"} @@ -251,6 +259,18 @@ def prepare_qwen_pdd_batch( raise TypeError("Qwen PDD latent, text embedding, and mask values must be tensors.") if data.ndim != 4: raise ValueError(f"Qwen PDD image_latents must be 4D, got {tuple(data.shape)}.") + if not data.dtype.is_floating_point: + raise TypeError("Qwen PDD image_latents must use a floating-point dtype.") + if data.shape[0] <= 0 or data.shape[1] != expected_latent_channels: + raise ValueError( + "Qwen PDD image_latents must have a non-empty batch and exactly " + f"{expected_latent_channels} channels, got {tuple(data.shape)}." + ) + if data.shape[2] <= 0 or data.shape[3] <= 0 or data.shape[2] % 2 or data.shape[3] % 2: + raise ValueError( + "Qwen PDD image_latents must have positive even spatial dimensions, got " + f"{tuple(data.shape[2:])}." + ) if not isinstance(metadata, Mapping): raise TypeError("Qwen PDD batch metadata must be a mapping.") sample_ids = metadata.get("sample_ids") @@ -262,16 +282,43 @@ def prepare_qwen_pdd_batch( ): raise ValueError("Qwen PDD sample_ids must be non-empty strings matching batch size.") + def prepare_condition( + embeddings: torch.Tensor, + attention_mask: torch.Tensor, + *, + name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not embeddings.dtype.is_floating_point: + raise TypeError(f"{name} embeddings must use a 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 boolean dtype.") + if embeddings.ndim not in (2, 3): + raise ValueError(f"{name} embeddings must be 2D or 3D, got {embeddings.ndim}D.") + if attention_mask.ndim not in (1, 2): + raise ValueError(f"{name} mask must be 1D or 2D, got {attention_mask.ndim}D.") + if embeddings.shape[-2] <= 0 or embeddings.shape[-1] <= 0: + raise ValueError(f"{name} embeddings must have non-empty sequence and feature axes.") + if embeddings.shape[-1] != expected_condition_features: + raise ValueError( + f"{name} embeddings must have exactly {expected_condition_features} features." + ) + if attention_mask.shape[-1] != embeddings.shape[-2]: + raise ValueError(f"{name} mask sequence length must match its embeddings.") + if embeddings.ndim == 3 and embeddings.shape[0] != data.shape[0]: + raise ValueError(f"{name} embedding batch size must match image_latents.") + if attention_mask.ndim == 2 and attention_mask.shape[0] != data.shape[0]: + raise ValueError(f"{name} mask batch size must match image_latents.") + + embeddings = embeddings.to(device=device, dtype=dtype) + attention_mask = attention_mask.to(device=device) + if embeddings.ndim == 2: + embeddings = embeddings.unsqueeze(0).expand(data.shape[0], -1, -1).contiguous() + if attention_mask.ndim == 1: + attention_mask = attention_mask.unsqueeze(0).expand(data.shape[0], -1).contiguous() + return embeddings, attention_mask + data = data.to(device=device, dtype=dtype) - text = text.to(device=device, dtype=dtype) - mask = mask.to(device=device) - if text.ndim == 2: - text = text.unsqueeze(0).expand(data.shape[0], -1, -1).contiguous() - if mask.ndim == 1: - mask = mask.unsqueeze(0).expand(data.shape[0], -1).contiguous() - if text.shape[0] != data.shape[0] or mask.shape[0] != data.shape[0]: - raise ValueError("Qwen PDD conditioning batch size must match image_latents.") - condition = (text, mask) + condition = prepare_condition(text, mask, name="Qwen PDD condition") negative: tuple[torch.Tensor, torch.Tensor] | None = None negative_text = batch.get("negative_text_embeddings") @@ -281,15 +328,11 @@ def prepare_qwen_pdd_batch( negative_mask, torch.Tensor ): raise TypeError("negative Qwen conditioning requires embedding and mask tensors.") - negative_text = negative_text.to(device=device, dtype=dtype) - negative_mask = negative_mask.to(device=device) - if negative_text.ndim == 2: - negative_text = negative_text.unsqueeze(0).expand(data.shape[0], -1, -1).contiguous() - if negative_mask.ndim == 1: - negative_mask = negative_mask.unsqueeze(0).expand(data.shape[0], -1).contiguous() - if negative_text.shape[0] != data.shape[0] or negative_mask.shape[0] != data.shape[0]: - raise ValueError("negative Qwen conditioning batch size must match image_latents.") - negative = (negative_text, negative_mask) + negative = prepare_condition( + negative_text, + negative_mask, + name="negative Qwen PDD condition", + ) if require_negative_condition and negative is None: raise ValueError("guided Qwen PDD training requires negative prompt conditioning.") return PreparedPDDBatch(data, condition, negative, sample_ids, (True,) * len(sample_ids)) diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index f13ce2de408..39ef65cd352 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -262,14 +262,12 @@ def _call_packed( batch_size, _, height, width = state.shape packed_state = pack_latents(state).to(self._model_dtype(model, state.dtype)) - txt_seq_lens = attention_mask.sum(dim=1).to(torch.int32).tolist() output = model( hidden_states=packed_state, timestep=time, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=attention_mask, img_shapes=build_img_shapes(batch_size, height, width), - txt_seq_lens=txt_seq_lens, guidance=None, return_dict=False, **model_kwargs, diff --git a/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py b/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py index e5a9a20b006..7007f3c7502 100644 --- a/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py @@ -67,6 +67,8 @@ def _expect_collective_failure(iterator, sampler: _Sampler, expected: str) -> No dtype=torch.float32, require_negative_condition=False, expected_batch_size=1, + expected_latent_channels=3, + expected_condition_features=6, ) except RuntimeError as error: message = str(error) @@ -98,6 +100,46 @@ def main() -> None: ) dist.barrier() + malformed = _batch(sample_id) + if rank == 1: + malformed["text_embeddings_mask"] = torch.ones(1, 5, dtype=torch.float32) + _expect_collective_failure( + iter([malformed]), + _Sampler((sample_id,)), + "mask must use an integer or boolean dtype", + ) + dist.barrier() + + malformed = _batch(sample_id) + if rank == 0: + malformed["image_latents"] = torch.ones(1, 3, 3, 4) + _expect_collective_failure( + iter([malformed]), + _Sampler((sample_id,)), + "positive even spatial dimensions", + ) + dist.barrier() + + malformed = _batch(sample_id) + if rank == 1: + malformed["image_latents"] = torch.ones(1, 4, 4, 4) + _expect_collective_failure( + iter([malformed]), + _Sampler((sample_id,)), + "exactly 3 channels", + ) + dist.barrier() + + malformed = _batch(sample_id) + if rank == 0: + malformed["text_embeddings"] = torch.ones(1, 5, 7) + _expect_collective_failure( + iter([malformed]), + _Sampler((sample_id,)), + "exactly 6 features", + ) + dist.barrier() + iterator_message = None try: _collective_training_iterator( diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py index b181c5f15d5..de06b821dc2 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -52,11 +52,10 @@ def forward( encoder_hidden_states, encoder_hidden_states_mask, img_shapes, - txt_seq_lens, guidance, return_dict, ): - del img_shapes, txt_seq_lens, guidance, return_dict + del img_shapes, guidance, return_dict condition = encoder_hidden_states.mean(dim=(1, 2), keepdim=True) condition += encoder_hidden_states_mask.sum(dim=1)[:, None, None] / 100 hidden = torch.tanh(self.backbone(hidden_states)) diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index 2815de6b8b8..45f295206ee 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -185,6 +185,8 @@ def test_qwen_batch_preparation_preserves_ids_masks_and_negative_condition() -> device=torch.device("cpu"), dtype=torch.float32, require_negative_condition=True, + expected_latent_channels=3, + expected_condition_features=6, ) assert prepared.sample_ids == ("qwen-a", "qwen-b") @@ -203,9 +205,50 @@ def test_qwen_batch_preparation_preserves_ids_masks_and_negative_condition() -> device=torch.device("cpu"), dtype=torch.float32, require_negative_condition=True, + expected_latent_channels=3, + expected_condition_features=6, ) +def test_qwen_batch_preparation_rejects_every_pre_model_shape_and_dtype_mismatch() -> None: + base = { + "image_latents": torch.ones(2, 3, 4, 4), + "text_embeddings": torch.ones(2, 5, 6), + "text_embeddings_mask": torch.ones(2, 5, dtype=torch.long), + "negative_text_embeddings": torch.zeros(2, 5, 6), + "negative_text_embeddings_mask": torch.ones(2, 5, dtype=torch.bool), + "metadata": {"sample_ids": ["qwen-a", "qwen-b"]}, + } + cases = ( + ("image_latents", torch.ones(2, 3, 4, 4, dtype=torch.long), "floating-point dtype"), + ("image_latents", torch.ones(2, 4, 4, 4), "exactly 3 channels"), + ("image_latents", torch.ones(2, 3, 3, 4), "positive even spatial dimensions"), + ("text_embeddings", torch.ones(2, 5, 6, dtype=torch.long), "floating-point dtype"), + ("text_embeddings", torch.ones(2, 5, 7), "exactly 6 features"), + ("text_embeddings", torch.ones(2, 5, 6, 1), "must be 2D or 3D"), + ("text_embeddings_mask", torch.ones(2, 5), "integer or boolean dtype"), + ("text_embeddings_mask", torch.ones(2, 4, dtype=torch.long), "sequence length"), + ( + "negative_text_embeddings_mask", + torch.ones(2, 5), + "integer or boolean dtype", + ), + ) + + for field, value, message in cases: + batch = dict(base) + batch[field] = value + with pytest.raises((TypeError, ValueError), match=message): + prepare_qwen_pdd_batch( + batch, + device=torch.device("cpu"), + dtype=torch.float32, + require_negative_condition=True, + expected_latent_channels=3, + expected_condition_features=6, + ) + + def test_two_direct_updates_have_finite_gradients_updates_and_targeted_coverage() -> None: lifecycle = build_toy_lifecycle(weight_decay=0.02) first_batch = make_batch(("a", "b")) diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 6dc0ef3ba65..b5f24744098 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -53,7 +53,6 @@ def forward( encoder_hidden_states, encoder_hidden_states_mask, img_shapes, - txt_seq_lens, guidance, return_dict, **kwargs, @@ -72,7 +71,6 @@ def forward( "encoder_hidden_states": encoder_hidden_states.detach().clone(), "encoder_hidden_states_mask": encoder_hidden_states_mask.detach().clone(), "img_shapes": img_shapes, - "txt_seq_lens": txt_seq_lens, "guidance": guidance, "return_dict": return_dict, "kwargs": kwargs, @@ -119,7 +117,6 @@ def _call_base_packed( encoder_hidden_states=embeddings, encoder_hidden_states_mask=mask, img_shapes=build_img_shapes(state.shape[0], state.shape[2], state.shape[3]), - txt_seq_lens=mask.sum(dim=1).to(torch.int32).tolist(), guidance=None, return_dict=False, )[0] @@ -148,7 +145,7 @@ def test_conversion_is_idempotent_and_every_initialized_head_matches_base() -> N assert len(student.calls) == 1 torch.testing.assert_close(student.calls[0]["timestep"], time) assert student.calls[0]["img_shapes"] == [[(1, 2, 2)], [(1, 2, 2)]] - assert student.calls[0]["txt_seq_lens"] == [2, 1] + assert "txt_seq_lens" not in student.calls[0]["kwargs"] assert student.calls[0]["guidance"] is None @@ -244,8 +241,7 @@ def test_teacher_cfg_and_packed_token_norm_rescale_match_direct_reference() -> N torch.testing.assert_close(actual, expected) torch.testing.assert_close(teacher.calls[0]["encoder_hidden_states"], condition[0]) torch.testing.assert_close(teacher.calls[1]["encoder_hidden_states"], negative_condition[0]) - assert teacher.calls[0]["txt_seq_lens"] == [2, 1] - assert teacher.calls[1]["txt_seq_lens"] == [1, 3] + assert all("txt_seq_lens" not in call["kwargs"] for call in teacher.calls) assert all(call["guidance"] is None for call in teacher.calls) From 78ddcb3ef65f4e6b98ece587516fa127bc732ff2 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 14:13:24 -0700 Subject: [PATCH 15/45] fix(fastgen): align PDD RF maximum time Signed-off-by: Meng Xin --- .../fastgen/inference_pdd_qwen_image.py | 4 +- .../fastgen/pdd/configs/qwen_image.yaml | 1 + examples/diffusers/fastgen/pdd_evaluation.py | 25 ++++-- examples/diffusers/fastgen/pdd_export.py | 1 + modelopt/torch/fastgen/config.py | 19 ++++- modelopt/torch/fastgen/flow_matching.py | 23 ++++- modelopt/torch/fastgen/methods/pdd.py | 25 ++++-- .../general/distillation/pdd_qwen_image.yaml | 1 + .../fastgen/pdd_export_distributed.py | 1 + .../diffusers/fastgen/pdd_test_utils.py | 1 + .../diffusers/fastgen/test_pdd_evaluation.py | 29 +++++++ .../fastgen/test_pdd_inference_checkpoint.py | 37 ++++++-- .../test_pdd_qwen_operability_smoke.py | 2 + .../fastgen/test_pdd_recipe_setup.py | 10 ++- .../fastgen/test_pdd_validation_oracle.py | 1 + tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py | 1 + .../fastgen/pdd_qwen_operability_smoke.py | 4 + tests/gpu/torch/fastgen/test_pdd_toy.py | 1 + tests/unit/recipe/test_loader.py | 2 + tests/unit/torch/fastgen/test_pdd_config.py | 23 +++++ .../fastgen/test_pdd_gradient_routing.py | 1 + tests/unit/torch/fastgen/test_pdd_metadata.py | 8 +- tests/unit/torch/fastgen/test_pdd_pipeline.py | 30 +++++-- .../unit/torch/fastgen/test_pdd_projection.py | 4 + .../torch/fastgen/test_pdd_reference_math.py | 84 ++++++++++++------- .../fastgen/test_qwen_image_pdd_plugin.py | 1 + 26 files changed, 278 insertions(+), 61 deletions(-) diff --git a/examples/diffusers/fastgen/inference_pdd_qwen_image.py b/examples/diffusers/fastgen/inference_pdd_qwen_image.py index 75a80053137..e90ea81f477 100644 --- a/examples/diffusers/fastgen/inference_pdd_qwen_image.py +++ b/examples/diffusers/fastgen/inference_pdd_qwen_image.py @@ -259,7 +259,7 @@ def main() -> None: ) generator = torch.Generator(device=device).manual_seed(args.seed) shape = _latent_shape(pipe, height=args.height, width=args.width) - state = torch.randn(shape, generator=generator, device=device, dtype=torch.float32) + noise = torch.randn(shape, generator=generator, device=device, dtype=torch.float32) transformer_invocations = 0 @@ -274,7 +274,7 @@ def count_invocation( torch.cuda.synchronize(device) started = time.perf_counter() try: - sampled = sampler.sample(state, condition=condition) + sampled = sampler.sample(noise, condition=condition) finally: hook.remove() images = _decode_qwen_latents(pipe, sampled.to(dtype)) diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 546361fb235..13014823709 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -19,6 +19,7 @@ pdd: student_sample_steps: 4 student_sample_type: ode grid_size: 128 + grid_max_t: 0.999 flow_shift: 5.0 block_size_min: 4 block_size_max: 64 diff --git a/examples/diffusers/fastgen/pdd_evaluation.py b/examples/diffusers/fastgen/pdd_evaluation.py index c188c84a9d1..5e5fac664c7 100644 --- a/examples/diffusers/fastgen/pdd_evaluation.py +++ b/examples/diffusers/fastgen/pdd_evaluation.py @@ -32,21 +32,36 @@ ) -def _grid_protocol(grid_size: int) -> Mapping[str, Any]: - nodes = [float(value) for value in make_shifted_flow_grid(grid_size, 5.0).tolist()] +def _grid_protocol(grid_size: int, *, grid_max_t: float) -> Mapping[str, Any]: + nodes = [ + float(value) + for value in make_shifted_flow_grid( + grid_size, + 5.0, + max_t=grid_max_t, + ).tolist() + ] return { "builder": "modelopt.torch.fastgen.make_shifted_flow_grid", - "formula": "s*u/(1+(s-1)*u), u=1-i/grid_size, i=0..grid_size", + "construction_dtype": "float64", + "formula": ( + "u64=clamp(linspace(grid_max_t,0,grid_size+1),max=grid_max_t); " + "g64=clamp(shift*u64/(1+(shift-1)*u64),max=grid_max_t); " + "nodes=float32(g64)" + ), "grid_size": grid_size, + "grid_max_t": grid_max_t, "flow_shift": 5.0, + "initial_state": "float32(float64(noise)*float64(grid_max_t))", "nodes": nodes, "nodes_sha256": hashlib.sha256(canonical_json_bytes(nodes)).hexdigest(), + "runtime_dtype": "float32", } GRID_PROTOCOLS: Mapping[str, Mapping[str, Any]] = { - "pdd_grid_128_shift5": _grid_protocol(128), - "teacher_grid_50_shift5": _grid_protocol(50), + "pdd_grid_128_shift5": _grid_protocol(128, grid_max_t=0.999), + "teacher_grid_50_shift5": _grid_protocol(50, grid_max_t=0.999), } _GUIDED_CFG_PROTOCOL = { diff --git a/examples/diffusers/fastgen/pdd_export.py b/examples/diffusers/fastgen/pdd_export.py index b62863fb92c..6ad782f4db5 100644 --- a/examples/diffusers/fastgen/pdd_export.py +++ b/examples/diffusers/fastgen/pdd_export.py @@ -525,6 +525,7 @@ def pdd_config_from_metadata( resolved = tuple(blocks) return PDDConfig( grid_size=metadata.grid_size, + grid_max_t=metadata.grid_max_t, flow_shift=metadata.flow_shift, block_size_min=metadata.block_size_min, block_size_max=metadata.block_size_max, diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py index a0788cbddc5..11394eeb062 100644 --- a/modelopt/torch/fastgen/config.py +++ b/modelopt/torch/fastgen/config.py @@ -29,7 +29,7 @@ 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 @@ -248,6 +248,11 @@ class PDDConfig(DistillationConfig): 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", @@ -292,10 +297,22 @@ def __setattr__(self, name: str, value: object) -> None: type(self).model_validate(candidate) super().__setattr__(name, value) + @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: diff --git a/modelopt/torch/fastgen/flow_matching.py b/modelopt/torch/fastgen/flow_matching.py index adf6f25fe3c..3696572e042 100644 --- a/modelopt/torch/fastgen/flow_matching.py +++ b/modelopt/torch/fastgen/flow_matching.py @@ -67,26 +67,43 @@ 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. - The returned tensor has ``grid_size + 1`` nodes from exactly 1 to 0. Low + 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) - unshifted = torch.linspace(1.0, 0.0, grid_size + 1, device=device, dtype=math_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), - f"shifted grid is not strictly decreasing for grid_size={grid_size}, shift={shift}.", + "shifted grid is not strictly decreasing for " + f"grid_size={grid_size}, max_t={max_t}, shift={shift}.", ) return grid diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index ad5712da7c6..57a5548f9af 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -139,6 +139,7 @@ class PDDMetadata: """Versioned minimum metadata required to reconstruct a PDD projection.""" grid_size: int + grid_max_t: float flow_shift: float block_size_min: int block_size_max: int @@ -158,6 +159,8 @@ def __post_init__(self) -> None: f"expected {_METADATA_SCHEMA_VERSION}." ) _require_int(self.grid_size, name="grid_size") + if type(self.grid_max_t) is not float: + raise ValueError(f"grid_max_t must be a float, got {self.grid_max_t!r}.") if type(self.flow_shift) is not float: raise ValueError(f"flow_shift must be a float, got {self.flow_shift!r}.") _require_int(self.block_size_min, name="block_size_min") @@ -189,6 +192,7 @@ def __post_init__(self) -> None: PDDConfig( grid_size=self.grid_size, + grid_max_t=self.grid_max_t, flow_shift=self.flow_shift, block_size_min=self.block_size_min, block_size_max=self.block_size_max, @@ -213,6 +217,7 @@ def from_config(cls, config: PDDConfig, projection: PDDOutputProjection) -> PDDM ) return cls( grid_size=config.grid_size, + grid_max_t=config.grid_max_t, flow_shift=config.flow_shift, block_size_min=config.block_size_min, block_size_max=config.block_size_max, @@ -229,6 +234,7 @@ def to_dict(self) -> dict[str, Any]: return { "schema_version": self.schema_version, "grid_size": self.grid_size, + "grid_max_t": self.grid_max_t, "flow_shift": self.flow_shift, "block_size_min": self.block_size_min, "block_size_max": self.block_size_max, @@ -252,6 +258,7 @@ def from_dict(cls, data: Mapping[str, Any]) -> PDDMetadata: { "schema_version", "grid_size", + "grid_max_t", "flow_shift", "block_size_min", "block_size_max", @@ -263,6 +270,8 @@ def from_dict(cls, data: Mapping[str, Any]) -> PDDMetadata: name="PDD metadata", ) schema_version = _require_int(data["schema_version"], name="schema_version") + if type(data["grid_max_t"]) is not float: + raise ValueError(f"grid_max_t must be a float, got {data['grid_max_t']!r}.") if type(data["flow_shift"]) is not float: raise ValueError(f"flow_shift must be a float, got {data['flow_shift']!r}.") if not isinstance(data["inference_blocks"], list) or any( @@ -288,6 +297,7 @@ def from_dict(cls, data: Mapping[str, Any]) -> PDDMetadata: return cls( schema_version=schema_version, grid_size=_require_int(data["grid_size"], name="grid_size"), + grid_max_t=data["grid_max_t"], flow_shift=data["flow_shift"], block_size_min=_require_int(data["block_size_min"], name="block_size_min"), block_size_max=_require_int(data["block_size_max"], name="block_size_max"), @@ -656,6 +666,7 @@ def time_grid(self, device: torch.device | str | None = None) -> torch.Tensor: 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, ) @@ -943,22 +954,22 @@ def _validate_blocks(self, blocks: Sequence[int] | None) -> tuple[int, ...]: @torch.no_grad() def sample( self, - state: torch.Tensor, + noise: torch.Tensor, *, condition: Any = None, blocks: Sequence[int] | None = None, model_kwargs: Mapping[str, Any] | None = None, ) -> torch.Tensor: - """Sample with one fused student call per validated contiguous block.""" - self._validate_state(state, name="state") + """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(state.device) - current = state.to(torch.float32) + grid = self.time_grid(noise.device) + 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(state.shape[0]) + time = grid[start].expand(noise.shape[0]) velocity = self.adapter.student_fused_block( self.student, current, @@ -972,7 +983,7 @@ def sample( velocity = self._normalize_velocity( velocity, expected_shape=current.shape, - device=state.device, + device=noise.device, name="student_fused_block", ) current = current + (grid[end] - grid[start]) * velocity diff --git a/modelopt_recipes/general/distillation/pdd_qwen_image.yaml b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml index 673a8e5cb24..b44de0ae14e 100644 --- a/modelopt_recipes/general/distillation/pdd_qwen_image.yaml +++ b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml @@ -16,6 +16,7 @@ num_train_timesteps: guidance_scale: 4.0 grid_size: 128 +grid_max_t: 0.999 flow_shift: 5.0 block_size_min: 4 block_size_max: 64 diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py index 516b4f7ac6f..a5439d2c664 100644 --- a/tests/examples/diffusers/fastgen/pdd_export_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_export_distributed.py @@ -45,6 +45,7 @@ def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: "student_sample_steps": 2, "student_sample_type": "ode", "grid_size": 4, + "grid_max_t": 0.999, "flow_shift": 5.0, "block_size_min": 1, "block_size_max": 4, diff --git a/tests/examples/diffusers/fastgen/pdd_test_utils.py b/tests/examples/diffusers/fastgen/pdd_test_utils.py index 06d0b11c20d..89bcaccaf82 100644 --- a/tests/examples/diffusers/fastgen/pdd_test_utils.py +++ b/tests/examples/diffusers/fastgen/pdd_test_utils.py @@ -117,6 +117,7 @@ def build_toy_lifecycle( torch.manual_seed(seed) config = PDDConfig( grid_size=4, + grid_max_t=0.999, flow_shift=5.0, block_size_min=1, block_size_max=4, diff --git a/tests/examples/diffusers/fastgen/test_pdd_evaluation.py b/tests/examples/diffusers/fastgen/test_pdd_evaluation.py index c19df506239..7de6540d853 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_evaluation.py +++ b/tests/examples/diffusers/fastgen/test_pdd_evaluation.py @@ -53,6 +53,7 @@ def _rewrite_manifest(manifest: pathlib.Path, data: dict) -> None: def _export(root: pathlib.Path, automodel: dict) -> pathlib.Path: config = PDDConfig( grid_size=128, + grid_max_t=0.999, flow_shift=5.0, block_size_min=4, block_size_max=64, @@ -357,6 +358,22 @@ def test_effectiveness_bundle_is_authenticated_and_emits_effective_conclusion(tm assert first["aggregates"]["pdd_4"]["peak_device_memory_bytes"]["mean"] > 0 +def test_grid_protocol_pins_fastgen_precision_and_raw_noise_initialization() -> None: + pdd_grid = GRID_PROTOCOLS["pdd_grid_128_shift5"] + teacher_grid = GRID_PROTOCOLS["teacher_grid_50_shift5"] + + assert pdd_grid["grid_max_t"] == teacher_grid["grid_max_t"] == 0.999 + assert pdd_grid["construction_dtype"] == "float64" + assert pdd_grid["runtime_dtype"] == "float32" + assert pdd_grid["initial_state"] == "float32(float64(noise)*float64(grid_max_t))" + assert pdd_grid["nodes"][0] == 0.9990000128746033 + assert pdd_grid["nodes"][-1] == 0.0 + assert ( + pdd_grid["nodes_sha256"] + == hashlib.sha256(canonical_json_bytes(pdd_grid["nodes"])).hexdigest() + ) + + @pytest.mark.parametrize( "corruption", [ @@ -374,7 +391,11 @@ def test_effectiveness_bundle_is_authenticated_and_emits_effective_conclusion(tm "bootstrap", "prompt_reference", "guided_count", + "grid_construction_dtype", + "grid_initial_state", + "grid_max_t", "grid_nodes", + "grid_nodes_hash", "latency_decision", "integrator_formula", "protocol", @@ -445,6 +466,14 @@ def test_effectiveness_bundle_rejects_unclaimable_evidence(tmp_path, corruption) data["prompt_set"] = _reference(root, alternate) elif corruption == "grid_nodes": data["grid_protocols"]["pdd_grid_128_shift5"]["nodes"][32] += 1e-4 + elif corruption == "grid_nodes_hash": + data["grid_protocols"]["pdd_grid_128_shift5"]["nodes_sha256"] = "0" * 64 + elif corruption == "grid_max_t": + data["grid_protocols"]["pdd_grid_128_shift5"]["grid_max_t"] = 1.0 + elif corruption == "grid_construction_dtype": + data["grid_protocols"]["pdd_grid_128_shift5"]["construction_dtype"] = "float32" + elif corruption == "grid_initial_state": + data["grid_protocols"]["pdd_grid_128_shift5"]["initial_state"] = "noise" elif corruption == "latency_decision": data["decision_rule"]["efficiency_measure"] = "latency_seconds" elif corruption == "integrator_formula": diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py index de06b821dc2..e9d32905beb 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -66,6 +66,7 @@ def forward( def _config(blocks=(32, 32, 32, 32)) -> PDDConfig: return PDDConfig( grid_size=128, + grid_max_t=0.999, flow_shift=5.0, block_size_min=4, block_size_max=64, @@ -128,9 +129,9 @@ def _condition(): return torch.tensor([[[0.2, -0.3], [0.1, 0.4]]]), torch.ones(1, 2, dtype=torch.long) -def _sample(model: nn.Module, config: PDDConfig, state: torch.Tensor) -> torch.Tensor: +def _sample(model: nn.Module, config: PDDConfig, noise: torch.Tensor) -> torch.Tensor: pipeline = PDDPipeline(model, nn.Identity(), config, QwenImagePDDAdapter(config)) - return pipeline.sample(state.clone(), condition=_condition()) + return pipeline.sample(noise.clone(), condition=_condition()) def test_bounded_safe_export_round_trip_and_seeded_schedules(tmp_path, monkeypatch) -> None: @@ -149,7 +150,7 @@ def test_bounded_safe_export_round_trip_and_seeded_schedules(tmp_path, monkeypat for key, tensor in source.state_dict().items(): torch.testing.assert_close(restored.state_dict()[key], tensor, rtol=0, atol=0) - state = torch.randn((1, 1, 4, 4), generator=torch.Generator().manual_seed(91)) + noise = torch.randn((1, 1, 4, 4), generator=torch.Generator().manual_seed(91)) for schedule, blocks in PDD_INFERENCE_SCHEDULES.items(): config = pdd_config_from_metadata( metadata, @@ -158,16 +159,40 @@ def test_bounded_safe_export_round_trip_and_seeded_schedules(tmp_path, monkeypat ) source.calls = 0 restored.calls = 0 - expected = _sample(source, config, state) - actual = _sample(restored, config, state) + expected = _sample(source, config, noise) + actual = _sample(restored, config, noise) torch.testing.assert_close(actual, expected, rtol=0, atol=0) assert source.calls == restored.calls == len(blocks) - torch.testing.assert_close(_sample(restored, config, state), actual, rtol=0, atol=0) + torch.testing.assert_close(_sample(restored, config, noise), actual, rtol=0, atol=0) with pytest.raises(ValueError, match="block_size_max"): pdd_config_from_metadata(metadata, blocks=[128], guidance_scale=4.0) +def test_inference_config_preserves_authenticated_nondefault_grid_max_t() -> None: + _model, _config_value, metadata = _converted() + payload = metadata.to_dict() + payload["grid_max_t"] = 1.0 + boundary_metadata = PDDMetadata.from_dict(payload) + + config = pdd_config_from_metadata( + boundary_metadata, + schedule="pdd-4", + guidance_scale=4.0, + ) + + assert config.grid_max_t == 1.0 + assert ( + PDDPipeline( + _TinyQwen(), + nn.Identity(), + config, + QwenImagePDDAdapter(config), + ).time_grid()[0] + == 1.0 + ) + + def test_pinned_qwen_none_prompt_mask_is_normalized_for_pdd() -> None: """Diffusers 0.38 returns None when the single-prompt mask is all ones.""" embeddings = torch.randn(1, 3, 5, dtype=torch.float32) diff --git a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py index baeb7db9882..8a60c146a34 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py +++ b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py @@ -38,6 +38,7 @@ def _stage_result(stage: str) -> dict: }, "pdd": { "grid_size": 128, + "grid_max_t": 0.999, "flow_shift": 5.0, "block_size_min": 4, "block_size_max": 64, @@ -139,6 +140,7 @@ def _checkpoint_identity(stage: dict) -> dict: "pdd_metadata": { "schema_version": 1, "grid_size": 128, + "grid_max_t": 0.999, "flow_shift": 5.0, "block_size_min": 4, "block_size_max": 64, diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index ed854ef140b..323b959d178 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -15,6 +15,7 @@ import pytest import torch +import yaml from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir _REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] @@ -49,6 +50,7 @@ def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: "student_sample_steps": 2, "student_sample_type": "ode", "grid_size": 4, + "grid_max_t": 0.999, "flow_shift": 5.0, "block_size_min": 1, "block_size_max": 4, @@ -74,6 +76,12 @@ def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: } +def test_example_recipe_explicitly_pins_grid_max_t() -> None: + raw = yaml.safe_load((_FASTGEN_DIR / "configs" / "pdd_qwen_image.yaml").read_text()) + assert type(raw["pdd"]["grid_max_t"]) is float + assert raw["pdd"]["grid_max_t"] == 0.999 + + @pytest.mark.parametrize( ("scope", "name", "value", "message"), [ @@ -134,7 +142,7 @@ def test_restore_requires_enabled_checkpointing(tmp_path) -> None: raw["checkpoint"]["enabled"] = False raw["checkpoint"]["restore_from"] = "LATEST" - with pytest.raises(ValueError, match="restore_from requires checkpoint.enabled=true"): + with pytest.raises(ValueError, match=r"restore_from requires checkpoint\.enabled=true"): resolve_pdd_recipe_config(raw) diff --git a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py index 07fe2eefabe..128e8c44b93 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py @@ -34,6 +34,7 @@ def test_canonical_2k_assignment_covers_all_1568_pairs_32_starts_and_128_heads() -> None: config = PDDConfig( grid_size=128, + grid_max_t=0.999, flow_shift=5.0, block_size_min=4, block_size_max=64, diff --git a/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py b/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py index 64e18d1fa43..e3644e0f428 100644 --- a/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py +++ b/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py @@ -154,6 +154,7 @@ def _fill_parameters(model: nn.Module, *, offset: float) -> None: def _config() -> 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, diff --git a/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py b/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py index 4da3ab09121..0b7cdafe396 100644 --- a/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py +++ b/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py @@ -183,6 +183,7 @@ def validate_stage_result(value: Mapping[str, Any], *, stage: str) -> None: raise ValueError("smoke model identity is invalid") if value["pdd"] != { "grid_size": 128, + "grid_max_t": 0.999, "flow_shift": 5.0, "block_size_min": 4, "block_size_max": 64, @@ -398,6 +399,7 @@ def _validate_checkpoint_identity(identity: Any, *, stage: Mapping[str, Any]) -> pdd = stage["pdd"] if ( metadata.grid_size != pdd["grid_size"] + or metadata.grid_max_t != pdd["grid_max_t"] or metadata.flow_shift != pdd["flow_shift"] or metadata.block_size_min != pdd["block_size_min"] or metadata.block_size_max != pdd["block_size_max"] @@ -580,6 +582,7 @@ def _raw_config(run_root: pathlib.Path, world_size: int) -> dict[str, Any]: "student_sample_steps": 4, "student_sample_type": "ode", "grid_size": 128, + "grid_max_t": 0.999, "flow_shift": 5.0, "block_size_min": 4, "block_size_max": 64, @@ -960,6 +963,7 @@ def count_teacher_call(_module: Any, _args: Any, _kwargs: Any) -> None: "model": {"id": _MODEL_ID, "revision": _MODEL_REVISION, "dtype": "bfloat16"}, "pdd": { "grid_size": 128, + "grid_max_t": 0.999, "flow_shift": 5.0, "block_size_min": 4, "block_size_max": 64, diff --git a/tests/gpu/torch/fastgen/test_pdd_toy.py b/tests/gpu/torch/fastgen/test_pdd_toy.py index 2d429d1aab4..b39d4f09063 100644 --- a/tests/gpu/torch/fastgen/test_pdd_toy.py +++ b/tests/gpu/torch/fastgen/test_pdd_toy.py @@ -120,6 +120,7 @@ def _build( ) -> 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, diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index e66748acae4..2f4c703bdaf 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -79,6 +79,8 @@ def test_load_pdd_config_builtin_recipe(): 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] diff --git a/tests/unit/torch/fastgen/test_pdd_config.py b/tests/unit/torch/fastgen/test_pdd_config.py index 3824cf959a2..f70bc6468f9 100644 --- a/tests/unit/torch/fastgen/test_pdd_config.py +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -37,6 +37,7 @@ def test_default_pdd_config_is_canonical_and_lists_are_independent(): assert first.student_sample_type == "ode" assert first.student_sample_steps == 4 assert first.grid_size == 128 + assert first.grid_max_t == 0.999 assert first.flow_shift == 5.0 assert first.block_size_min == 4 assert first.block_size_max == 64 @@ -88,10 +89,32 @@ def test_rejected_mapping_assignment_leaves_pdd_config_unchanged(): assert config.inference_blocks == [32, 32, 32, 32] +@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) + + config = PDDConfig() + with pytest.raises(ValueError, match="grid_max_t must be a float"): + config.grid_max_t = value + assert config.grid_max_t == 0.999 + + +def test_pdd_config_accepts_explicit_grid_max_t_upper_boundary(): + config = PDDConfig(grid_max_t=1.0) + assert config.grid_max_t == 1.0 + + @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": -0.1}, "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"), + ({"grid_max_t": float("inf")}, "0 < grid_max_t <= 1"), + ({"grid_max_t": float("-inf")}, "0 < grid_max_t <= 1"), ({"flow_shift": 0.5}, "flow_shift must be finite and >= 1"), ({"flow_shift": float("inf")}, "flow_shift must be finite and >= 1"), ({"block_size_min": 0}, "0 < block_size_min"), diff --git a/tests/unit/torch/fastgen/test_pdd_gradient_routing.py b/tests/unit/torch/fastgen/test_pdd_gradient_routing.py index be679a5560e..36657709cc6 100644 --- a/tests/unit/torch/fastgen/test_pdd_gradient_routing.py +++ b/tests/unit/torch/fastgen/test_pdd_gradient_routing.py @@ -82,6 +82,7 @@ def test_only_selected_head_and_shared_backbone_receive_gradients() -> None: adapter = _GradientAdapter() config = PDDConfig( grid_size=8, + grid_max_t=0.999, flow_shift=5.0, block_size_min=2, block_size_max=4, diff --git a/tests/unit/torch/fastgen/test_pdd_metadata.py b/tests/unit/torch/fastgen/test_pdd_metadata.py index d2297d3d536..1c3c65f6f2d 100644 --- a/tests/unit/torch/fastgen/test_pdd_metadata.py +++ b/tests/unit/torch/fastgen/test_pdd_metadata.py @@ -109,6 +109,7 @@ def teacher_velocity( def _config() -> PDDConfig: return PDDConfig( grid_size=4, + grid_max_t=0.999, flow_shift=5.0, block_size_min=1, block_size_max=4, @@ -140,8 +141,8 @@ def test_plain_torch_training_sampling_and_strict_metadata_reconstruction(tmp_pa assert torch.isfinite(loss) assert not torch.equal(projection.weight, projection_before) - initial = torch.tensor([[1.0, -0.5, 0.25], [-0.25, 0.75, 1.5]]) - expected_sample = pipeline.sample(initial) + inference_noise = torch.tensor([[1.0, -0.5, 0.25], [-0.25, 0.75, 1.5]]) + expected_sample = pipeline.sample(inference_noise) metadata = PDDMetadata.from_config(config, projection) metadata_path = tmp_path / "pdd_metadata.json" @@ -151,6 +152,7 @@ def test_plain_torch_training_sampling_and_strict_metadata_reconstruction(tmp_pa restored_config = PDDConfig( grid_size=restored_metadata.grid_size, + grid_max_t=restored_metadata.grid_max_t, flow_shift=restored_metadata.flow_shift, block_size_min=restored_metadata.block_size_min, block_size_max=restored_metadata.block_size_max, @@ -177,7 +179,7 @@ def test_plain_torch_training_sampling_and_strict_metadata_reconstruction(tmp_pa restored_config, _ToyAdapter(restored_metadata.grid_size), ) - torch.testing.assert_close(restored_pipeline.sample(initial), expected_sample) + torch.testing.assert_close(restored_pipeline.sample(inference_noise), expected_sample) def test_strict_restore_rejects_checkpoint_with_different_projection_grid() -> None: diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index 8fbf719a429..f4d68e15e8a 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -141,6 +141,7 @@ def teacher_velocity( 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, @@ -357,10 +358,10 @@ def test_sampled_indices_stay_on_exact_uniform_support() -> None: @pytest.mark.parametrize("blocks", [None, [2, 2, 2, 2]]) def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: pipeline, adapter = _pipeline() - initial = torch.tensor([[1.0, -2.0, 0.5], [-0.25, 0.75, 1.5]], dtype=torch.bfloat16) + noise = torch.tensor([[1.0, -2.0, 0.5], [-0.25, 0.75, 1.5]], dtype=torch.bfloat16) actual = pipeline.sample( - initial, + noise, condition="prompt", blocks=blocks, model_kwargs={"tag": 23}, @@ -368,7 +369,7 @@ def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: resolved = [4, 4] if blocks is None else blocks grid = pipeline.time_grid() - expected = initial.float() + expected = (noise.to(torch.float64) * pipeline.config.grid_max_t).to(torch.float32) start = 0 for block in resolved: end = start + block @@ -384,13 +385,32 @@ def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: 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(initial.shape[0])) + torch.testing.assert_close(call["time"], grid[start].expand(noise.shape[0])) torch.testing.assert_close(call["grid"], grid) assert call["condition"] == "prompt" assert call["kwargs"] == {"tag": 23} start = end +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"), [ @@ -411,7 +431,7 @@ 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="state must be a tensor"): + 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)) diff --git a/tests/unit/torch/fastgen/test_pdd_projection.py b/tests/unit/torch/fastgen/test_pdd_projection.py index 2e0f792bd86..f859e96660f 100644 --- a/tests/unit/torch/fastgen/test_pdd_projection.py +++ b/tests/unit/torch/fastgen/test_pdd_projection.py @@ -372,6 +372,7 @@ def test_metadata_round_trips_exactly_without_mutating_mapping(layout): "changes", [ {"schema_version": True}, + {"grid_max_t": 1}, {"flow_shift": 5}, {"inference_blocks": [32, 32, 32, 32]}, {"projection_bias": 1}, @@ -395,6 +396,9 @@ def _valid_patch_metadata_payload(): (lambda data: data.update(schema_version=2), "unsupported.*schema_version"), (lambda data: data.update(extra=True), "keys mismatch"), (lambda data: data.update({1: "bad"}), "keys must all be strings"), + (lambda data: data.pop("grid_max_t"), "keys mismatch"), + (lambda data: data.update(grid_max_t=1), "grid_max_t must be a float"), + (lambda data: data.update(grid_max_t=0.0), "0 < grid_max_t <= 1"), (lambda data: data.update(flow_shift=5), "flow_shift must be a float"), (lambda data: data.update(grid_size=124), "inference_blocks must sum"), (lambda data: data.update(inference_blocks=(32, 32, 32, 32)), "list of integers"), diff --git a/tests/unit/torch/fastgen/test_pdd_reference_math.py b/tests/unit/torch/fastgen/test_pdd_reference_math.py index d27b38a534e..8e6d527b31e 100644 --- a/tests/unit/torch/fastgen/test_pdd_reference_math.py +++ b/tests/unit/torch/fastgen/test_pdd_reference_math.py @@ -32,13 +32,12 @@ ) -def _reference_shifted_grid(grid_size: int, shift: float) -> torch.Tensor: - """Build the decreasing shifted rectified-flow grid with scalar arithmetic.""" - values = [] - for index in range(grid_size + 1): - unshifted = 1.0 - index / grid_size - values.append(shift * unshifted / (1.0 + (shift - 1.0) * unshifted)) - return torch.tensor(values, dtype=torch.float64) +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_integrate( @@ -77,12 +76,12 @@ def _reference_fused_parameters( def test_shifted_grid_matches_hand_calculated_values_and_preserves_float32_intervals(): - grid = _reference_shifted_grid(grid_size=4, shift=5.0) + grid = _reference_shifted_grid(grid_size=4, shift=5.0, max_t=1.0) expected = torch.tensor([1.0, 0.9375, 5.0 / 6.0, 0.625, 0.0], dtype=torch.float64) torch.testing.assert_close(grid, expected, rtol=0.0, atol=0.0) - canonical_grid = _reference_shifted_grid(grid_size=128, shift=5.0) + canonical_grid = _reference_shifted_grid(grid_size=128, shift=5.0, max_t=0.999) assert torch.all(torch.diff(canonical_grid.to(torch.float32)) < 0) assert canonical_grid.to(torch.bfloat16)[1] == 1.0 @@ -91,7 +90,7 @@ def test_shifted_grid_matches_hand_calculated_values_and_preserves_float32_inter def test_half_open_integration_uses_only_selected_interval_heads(): - grid = _reference_shifted_grid(grid_size=4, shift=5.0) + grid = _reference_shifted_grid(grid_size=4, shift=5.0, max_t=1.0) state = torch.tensor([3.0, -2.0]) velocities = torch.tensor( [ @@ -119,7 +118,7 @@ def test_half_open_integration_uses_only_selected_interval_heads(): def test_final_half_open_block_advances_exactly_four_intervals(): - grid = _reference_shifted_grid(grid_size=8, shift=5.0) + grid = _reference_shifted_grid(grid_size=8, shift=5.0, max_t=1.0) state = torch.tensor([1.25]) velocities = torch.ones(8, 1) @@ -130,7 +129,7 @@ def test_final_half_open_block_advances_exactly_four_intervals(): def test_fused_projection_matches_weighted_sum_and_explicit_block_update(): - grid = _reference_shifted_grid(grid_size=4, shift=5.0) + grid = _reference_shifted_grid(grid_size=4, shift=5.0, max_t=1.0) inputs = torch.tensor([[2.0, -1.0], [-0.5, 3.0]], dtype=torch.float64) weight = torch.tensor( [ @@ -174,19 +173,27 @@ def test_fused_projection_matches_weighted_sum_and_explicit_block_update(): def test_production_shifted_grid_matches_independent_oracle(): - grid = make_shifted_flow_grid(grid_size=128, shift=5.0) - oracle = _reference_shifted_grid(grid_size=128, shift=5.0).to(torch.float32) + 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] == 1.0 + 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=2e-7, atol=1e-7) + torch.testing.assert_close(grid, oracle, rtol=0, atol=0) + + 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_grid_promotes_low_precision_requests(): - grid = make_shifted_flow_grid(128, 5.0, dtype=torch.bfloat16) + grid = make_shifted_flow_grid(128, 5.0, max_t=0.999, dtype=torch.bfloat16) assert grid.dtype == torch.float32 assert torch.all(torch.diff(grid) < 0) @@ -202,16 +209,37 @@ def test_production_grid_promotes_low_precision_requests(): ) def test_production_grid_rejects_invalid_boundaries(grid_size, shift, message): with pytest.raises(ValueError, match=message): - make_shifted_flow_grid(grid_size, shift) + make_shifted_flow_grid(grid_size, shift, max_t=0.999) + + +@pytest.mark.parametrize("max_t", [True, 1, 0]) +def test_production_grid_rejects_non_float_max_t(max_t): + with pytest.raises(TypeError, match="max_t must be a float"): + make_shifted_flow_grid(4, 5.0, max_t=max_t) + + +@pytest.mark.parametrize("max_t", [float("nan"), float("inf"), float("-inf"), 0.0, -0.1, 1.0001]) +def test_production_grid_rejects_invalid_max_t(max_t): + with pytest.raises(ValueError, match="0 < max_t <= 1"): + make_shifted_flow_grid(4, 5.0, max_t=max_t) + + +def test_production_grid_requires_explicit_max_t_and_accepts_one(): + with pytest.raises(TypeError, match="max_t"): + make_shifted_flow_grid(4, 5.0) + + grid = make_shifted_flow_grid(4, 5.0, max_t=1.0) + assert grid[0] == 1.0 + assert grid[-1] == 0.0 def test_production_grid_rejects_non_floating_dtype(): with pytest.raises(TypeError, match="floating-point dtype"): - make_shifted_flow_grid(128, 5.0, dtype=torch.int64) + make_shifted_flow_grid(128, 5.0, max_t=0.999, dtype=torch.int64) def test_production_half_open_integration_matches_independent_oracle_per_sample(): - grid = make_shifted_flow_grid(4, 5.0, dtype=torch.float64) + 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( [ @@ -236,7 +264,7 @@ def test_production_half_open_integration_matches_independent_oracle_per_sample( def test_production_integration_does_not_consume_excluded_nonfinite_heads(): - grid = make_shifted_flow_grid(4, 5.0, dtype=torch.float64) + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999, dtype=torch.float64) state = torch.tensor([[3.0, -2.0]], dtype=torch.float64) velocities = torch.tensor( [[[torch.nan, torch.nan], [2.0, -1.0], [-3.0, 4.0], [torch.inf, -torch.inf]]], @@ -251,18 +279,18 @@ def test_production_integration_does_not_consume_excluded_nonfinite_heads(): def test_production_integration_promotes_bfloat16_math_to_float32(): - grid = make_shifted_flow_grid(4, 5.0) + 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), -1.0)) + torch.testing.assert_close(result, torch.full((1, 2), -0.999)) def test_production_integration_rejects_out_of_range_half_open_block(): - grid = make_shifted_flow_grid(4, 5.0) + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999) state = torch.zeros(1, 2) velocities = torch.ones(1, 4, 2) @@ -271,9 +299,9 @@ def test_production_integration_rejects_out_of_range_half_open_block(): def test_production_fusion_coefficients_match_independent_oracle(): - grid = make_shifted_flow_grid(4, 5.0, dtype=torch.float64) + 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) + 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]) @@ -287,13 +315,13 @@ def test_production_fusion_coefficients_match_independent_oracle(): def test_production_fusion_coefficients_reject_empty_block(): - grid = make_shifted_flow_grid(4, 5.0) + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999) with pytest.raises(ValueError, match="0 <= start < end <= 4"): fusion_coefficients(grid, start=2, end=2) def test_production_helpers_do_not_extract_meta_tensor_scalars(): - grid = make_shifted_flow_grid(4, 5.0, device="meta") + grid = make_shifted_flow_grid(4, 5.0, max_t=0.999, device="meta") state = torch.empty(2, 3, device="meta") velocities = torch.empty(2, 4, 3, device="meta") diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index b5f24744098..ee70dfad260 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -83,6 +83,7 @@ def forward( 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, From ccf8524f39406f8c8af0e726bf80c1b5ba735a5c Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 14:59:08 -0700 Subject: [PATCH 16/45] fix(fastgen): align PDD RF forward precision Signed-off-by: Meng Xin --- modelopt/torch/fastgen/methods/pdd.py | 4 +- tests/unit/torch/fastgen/test_pdd_pipeline.py | 78 ++++++++++++++++--- .../torch/fastgen/test_pdd_reference_math.py | 50 ++++++++++++ 3 files changed, 119 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index 57a5548f9af..3e5f3c84d4e 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -34,6 +34,7 @@ from ..config import PDDConfig from ..flow_matching import ( + add_noise, fusion_coefficients, integrate_interval_velocities, make_shifted_flow_grid, @@ -833,8 +834,7 @@ def compute_loss( ) time_n = grid[n] broadcast_shape = (batch_size,) + (1,) * (data.ndim - 1) - time_n_expanded = time_n.reshape(broadcast_shape) - x_n = (1.0 - time_n_expanded) * data_fp32 + time_n_expanded * noise_fp32 + x_n = add_noise(data_fp32, noise_fp32, time_n) student_heads = self.adapter.student_all_heads( self.student, diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index f4d68e15e8a..6d7301f517b 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -178,6 +178,62 @@ def _explicit_integrate_per_sample( 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_student_input_matches_fastgen_float64_forward_process() -> None: + pipeline, adapter = _pipeline() + data = torch.tensor( + [[-0.2654421329498291, 0.5161616802215576, -0.7285917401313782]], + dtype=torch.float32, + ) + noise = torch.tensor( + [[0.3856363296508789, -0.34849217534065247, -0.11881951987743378]], + dtype=torch.float32, + ) + n = torch.tensor([0]) + + pipeline.compute_loss(data, noise=noise, n=n, k=torch.tensor([0])) + + time = pipeline.time_grid()[n] + expected = _reference_rf_forward_process(data, noise, time) + stale_direct = (1.0 - time[:, None]) * data + time[:, None] * noise + assert torch.equal( + expected, + torch.tensor([[0.3849852681159973, -0.34762752056121826, -0.11942928284406662]]), + ) + assert not torch.equal(stale_direct, expected) + assert torch.equal(adapter.student_calls[0]["state"], expected) + assert torch.equal(adapter.student_calls[0]["time"], time) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_student_input_normalizes_supported_floating_dtypes_to_float32(dtype) -> None: + pipeline, adapter = _pipeline() + data = torch.tensor([[0.75, -1.25, 0.125]], dtype=dtype) + noise = torch.tensor([[-0.5, 1.5, 2.25]], dtype=dtype) + n = torch.tensor([2]) + + pipeline.compute_loss(data, noise=noise, n=n, k=torch.tensor([3])) + + time = pipeline.time_grid()[n] + expected = _reference_rf_forward_process(data, noise, time) + assert adapter.student_calls[0]["state"].dtype == torch.float32 + assert torch.equal(adapter.student_calls[0]["state"], expected) + assert torch.equal(adapter.student_calls[0]["time"], time) + + 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]]) @@ -196,7 +252,7 @@ def test_euler_loss_matches_analytic_empty_and_tail_reconstruction() -> None: ) grid = pipeline.time_grid() - x_n = (1 - grid[n, None]) * data + grid[n, None] * noise + 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]) @@ -224,23 +280,23 @@ def test_euler_loss_matches_analytic_empty_and_tail_reconstruction() -> None: 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]]) - noise = torch.tensor([[0.25, 1.5, -0.5]]) - n = torch.tensor([6]) - k = torch.tensor([7]) + 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 = (1 - grid[n, None]) * data + grid[n, None] * noise + x_n = _reference_rf_forward_process(data, noise, grid[n]) heads = pipeline.student.all_heads(x_n) - x_bar_k = x_n + (grid[7] - grid[6]) * heads[:, 6] + x_bar_k = _explicit_integrate_per_sample(x_n, heads, grid, n, k) first_velocity = pipeline.teacher(x_bar_k, grid[k]) - delta = grid[8] - grid[7] - midpoint_state = x_bar_k + 0.5 * delta * first_velocity + 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[:, 7] - midpoint_target).square().mean() + 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 @@ -262,7 +318,7 @@ def test_selected_head_low_precision_outputs_use_float32_mse() -> None: loss, _ = pipeline.compute_loss(data, noise=noise, n=n, k=k) grid = pipeline.time_grid() - x_n = (1 - grid[n, None]) * data + grid[n, None] * noise + 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() diff --git a/tests/unit/torch/fastgen/test_pdd_reference_math.py b/tests/unit/torch/fastgen/test_pdd_reference_math.py index 8e6d527b31e..449ce471ebb 100644 --- a/tests/unit/torch/fastgen/test_pdd_reference_math.py +++ b/tests/unit/torch/fastgen/test_pdd_reference_math.py @@ -26,6 +26,7 @@ import torch.nn.functional as F from modelopt.torch.fastgen.flow_matching import ( + add_noise, fusion_coefficients, integrate_interval_velocities, make_shifted_flow_grid, @@ -40,6 +41,22 @@ def _reference_shifted_grid(grid_size: int, shift: float, max_t: float) -> torch 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, @@ -192,6 +209,39 @@ def test_production_shifted_grid_matches_independent_oracle(): 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) + + def test_production_grid_promotes_low_precision_requests(): grid = make_shifted_flow_grid(128, 5.0, max_t=0.999, dtype=torch.bfloat16) From 5d07ad1f6a4d6bde5e2a8995ea47ec81cab6f55f Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 16:14:57 -0700 Subject: [PATCH 17/45] fix(fastgen): authenticate PDD holdout snapshots Signed-off-by: Meng Xin --- CHANGELOG.rst | 5 + .../fastgen/migrate_cache_manifest.py | 139 ++++++--- .../fastgen/pdd/configs/qwen_image.yaml | 2 + examples/diffusers/fastgen/pdd_finetune.py | 52 +++- examples/diffusers/fastgen/pdd_recipe.py | 22 ++ examples/diffusers/fastgen/portable_cache.py | 284 ++++++++++++++++-- .../preprocess/preprocessing_multiprocess.py | 4 +- .../fastgen/validate_cache_snapshot.py | 92 +++++- .../diffusers/fastgen/pdd_test_utils.py | 9 +- .../fastgen/test_migrate_cache_manifest.py | 224 +++++++++++++- .../test_pdd_qwen_operability_smoke.py | 12 + .../fastgen/test_pdd_recipe_setup.py | 75 +++++ .../diffusers/fastgen/test_portable_cache.py | 258 +++++++++++++++- .../fastgen/pdd_qwen_operability_smoke.py | 15 +- 14 files changed, 1061 insertions(+), 132 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a4421b1fd5a..ce5025a9e34 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,11 @@ Changelog **Backward Breaking Changes** +- FastGen example portable caches now require authenticated schema-2 all/train/held-out indices. + Existing schema-1 or unversioned caches, including fresh preprocessing output, must be finalized + with ``examples/diffusers/fastgen/migrate_cache_manifest.py`` and an approved ordered-ID + artifact; they no longer load directly. This changes only the example cache protocol, not the + framework-neutral PDD checkpoint or API contract. - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. diff --git a/examples/diffusers/fastgen/migrate_cache_manifest.py b/examples/diffusers/fastgen/migrate_cache_manifest.py index 7651f59304b..1e689bb3f6d 100644 --- a/examples/diffusers/fastgen/migrate_cache_manifest.py +++ b/examples/diffusers/fastgen/migrate_cache_manifest.py @@ -18,7 +18,6 @@ from __future__ import annotations import argparse -import hashlib import json import os import shutil @@ -29,9 +28,16 @@ import torch from portable_cache import ( - CACHE_SCHEMA_VERSION, + PDD_HOLDOUT_DOMAIN, + PORTABLE_SNAPSHOT_SCHEMA_VERSION, + PREPROCESS_STAGING_SCHEMA_VERSION, + SPLIT_POLICY_SCHEMA_VERSION, audit_no_absolute_paths, + load_approved_sample_ids, + load_strict_json, + ordered_sample_ids_sha256, resolve_cache_asset, + select_pdd_holdout_ids, sha256_file, stable_sample_id, validate_relative_reference, @@ -50,7 +56,6 @@ "source_path", "video_path", } -_SPLIT_DOMAIN = "modelopt-fastgen-split-v1" @dataclass(frozen=True) @@ -66,16 +71,27 @@ class MigrationRecord: def _load_json(path: Path, expected_type: type, label: str): - try: - with path.open(encoding="utf-8") as stream: - value = json.load(stream) - except json.JSONDecodeError as error: - raise ValueError(f"{label} is not valid JSON: {path}") from error + value = load_strict_json(path, label=label) if not isinstance(value, expected_type): raise ValueError(f"{label} must contain {expected_type.__name__}") return value +def _load_source_json( + root: Path, + reference: str, + *, + expected_type: type, + label: str, +): + relative = validate_relative_reference(reference, label=label) + candidate = root / relative + if candidate.is_symlink(): + raise ValueError(f"{label} must not be a symlink: {relative}") + path = resolve_cache_asset(root, reference, label=label) + return _load_json(path, expected_type, label) + + def _relative_to_legacy_prefix(raw: str, prefix: Path, *, label: str) -> Path: path = Path(raw).expanduser() if not path.is_absolute(): @@ -177,9 +193,11 @@ def plan_migration( all_shards: set[str] = set() for index_number, index_ref in enumerate(source_indexes): label = f"source_index[{index_number}]" - index_path = resolve_cache_asset(root, index_ref, label=label) - index = _load_json(index_path, dict, label) - if "schema_version" in index and index["schema_version"] != CACHE_SCHEMA_VERSION: + index = _load_source_json(root, index_ref, expected_type=dict, label=label) + if "schema_version" in index and ( + type(index["schema_version"]) is not int + or index["schema_version"] != PREPROCESS_STAGING_SCHEMA_VERSION + ): raise ValueError(f"{label}.schema_version is unsupported") shards = index.get("shards") if ( @@ -190,7 +208,9 @@ def plan_migration( raise ValueError(f"{label}.shards must be a non-empty list of relative paths") if len(shards) != len(set(shards)): raise ValueError(f"{label}.shards contains duplicates") - if "num_shards" in index and index["num_shards"] != len(shards): + if "num_shards" in index and ( + type(index["num_shards"]) is not int or index["num_shards"] != len(shards) + ): raise ValueError(f"{label}.num_shards does not match len(shards)") overlap = all_shards.intersection(shards) if overlap: @@ -204,6 +224,12 @@ def plan_migration( if ranked_indexes: if len(ranked_indexes) != len(parsed_indexes): raise ValueError("all source indices must declare shard_rank and shard_world") + if any( + type(index.get(field)) is not int + for index in ranked_indexes + for field in ("shard_rank", "shard_world") + ): + raise ValueError("source shard_rank and shard_world must be integers") worlds = {index.get("shard_world") for index in ranked_indexes} ranks = {index.get("shard_rank") for index in ranked_indexes} if worlds != {len(parsed_indexes)} or ranks != set(range(len(parsed_indexes))): @@ -216,8 +242,12 @@ def plan_migration( source_sample_ids = [] for shard_number, shard_ref in enumerate(shards): shard_label = f"{index_label}.shards[{shard_number}]" - shard_path = resolve_cache_asset(root, shard_ref, label=shard_label) - entries = _load_json(shard_path, list, shard_label) + entries = _load_source_json( + root, + shard_ref, + expected_type=list, + label=shard_label, + ) index_record_count += len(entries) for entry_number, entry in enumerate(entries): label = f"{shard_label}[{entry_number}]" @@ -263,7 +293,9 @@ def plan_migration( manifest_fields=_portable_manifest_fields(entry, label=label), ) ) - if "total_items" in index and index["total_items"] != index_record_count: + if "total_items" in index and ( + type(index["total_items"]) is not int or index["total_items"] != index_record_count + ): raise ValueError(f"{index_label}.total_items does not match its loaded entry count") if "sample_ids" in index and index["sample_ids"] != source_sample_ids: raise ValueError(f"{index_label}.sample_ids does not match its loaded entries") @@ -275,34 +307,19 @@ def plan_migration( def _write_json(path: Path, value: Any) -> None: with path.open("w", encoding="utf-8") as stream: - json.dump(value, stream, indent=2, sort_keys=True) + json.dump(value, stream, indent=2, sort_keys=True, allow_nan=False) stream.write("\n") stream.flush() os.fsync(stream.fileno()) -def _split_ids(records: tuple[MigrationRecord, ...], heldout_count: int, split_seed: str): - if heldout_count <= 0 or heldout_count >= len(records): - raise ValueError("heldout_count must be positive and smaller than the sample count") - ranked = sorted( - records, - key=lambda record: hashlib.sha256( - f"{_SPLIT_DOMAIN}\0{split_seed}\0{record.sample_id}".encode() - ).hexdigest(), - ) - heldout = {record.sample_id for record in ranked[:heldout_count]} - ordered = [record.sample_id for record in records] - return [sample_id for sample_id in ordered if sample_id not in heldout], [ - sample_id for sample_id in ordered if sample_id in heldout - ] - - def migrate_cache( source_root: str | Path, output_root: str | Path, *, + approved_ids_manifest: str | Path, heldout_count: int, - split_seed: str = "0", + expected_approved_ids_sha256: str | None = None, source_index: str | Sequence[str] = "metadata.json", legacy_cache_root: str | Path | None = None, legacy_source_root: str | Path | None = None, @@ -315,8 +332,8 @@ def migrate_cache( raise FileExistsError(f"output_root already exists: {destination}") if not destination.parent.exists(): raise FileNotFoundError(f"output_root parent does not exist: {destination.parent}") - if shard_size <= 0: - raise ValueError("shard_size must be positive") + if type(shard_size) is not int or shard_size <= 0: + raise ValueError("shard_size must be a positive integer") # Pass 1 is intentionally complete before any destination or staging path is created. records = plan_migration( @@ -325,7 +342,16 @@ def migrate_cache( legacy_cache_root=legacy_cache_root, legacy_source_root=legacy_source_root, ) - train_ids, heldout_ids = _split_ids(records, heldout_count, split_seed) + approved_ids, approved_digest = load_approved_sample_ids( + Path(approved_ids_manifest).expanduser(), + expected_sha256=expected_approved_ids_sha256, + ) + records_by_id = {record.sample_id: record for record in records} + unknown_ids = [sample_id for sample_id in approved_ids if sample_id not in records_by_id] + if unknown_ids: + raise ValueError(f"approved_ids_manifest references unknown sample IDs: {unknown_ids[:5]}") + selected_records = tuple(records_by_id[sample_id] for sample_id in approved_ids) + train_ids, heldout_ids = select_pdd_holdout_ids(approved_ids, heldout_count) source = Path(source_root).expanduser().resolve(strict=True) negative_source = None @@ -345,7 +371,7 @@ def migrate_cache( try: (staging / "payloads").mkdir() portable_entries = [] - for record in records: + for record in selected_records: if sha256_file(record.source_payload) != record.source_sha256: raise RuntimeError(f"source payload changed after pass 1: {record.source_payload}") payload = torch.load(record.source_payload, map_location="cpu", weights_only=True) @@ -384,14 +410,21 @@ def migrate_cache( } common = { - "schema_version": CACHE_SCHEMA_VERSION, + "schema_version": PORTABLE_SNAPSHOT_SCHEMA_VERSION, "shards": shard_names, "num_shards": len(shard_names), + "split_policy": { + "schema_version": SPLIT_POLICY_SCHEMA_VERSION, + "algorithm": "sha256-domain-ranked", + "domain": PDD_HOLDOUT_DOMAIN, + "heldout_count": heldout_count, + "approved_ordered_ids_sha256": approved_digest, + }, } if negative_declaration is not None: common["negative_prompt_embedding"] = negative_declaration split_specs = { - "metadata.json": ("all", [record.sample_id for record in records]), + "metadata.json": ("all", approved_ids), "metadata_train.json": ("train", train_ids), "metadata_heldout.json": ("heldout", heldout_ids), } @@ -399,23 +432,33 @@ def migrate_cache( index = { **common, "split": split, - "sample_ids": sample_ids, + "sample_ids": list(sample_ids), + "ordered_sample_ids_sha256": ordered_sample_ids_sha256(sample_ids), "total_items": len(sample_ids), } audit_no_absolute_paths(index, context=name) _write_json(staging / name, index) - validate_snapshot(staging) + validation = validate_snapshot( + staging, + expected_approved_ids_sha256=expected_approved_ids_sha256, + expected_heldout_count=heldout_count, + ) os.replace(staging, destination) except BaseException: shutil.rmtree(staging, ignore_errors=True) raise return { + "schema_version": 1, + "record_type": "modelopt_fastgen_cache_migration", "output_root": str(destination.resolve()), - "total_items": len(records), - "train_items": len(train_ids), - "heldout_items": len(heldout_ids), + "counts": { + "source": len(records), + "approved": len(selected_records), + "filtered": len(records) - len(selected_records), + }, + "validation": validation, } @@ -432,22 +475,24 @@ def main() -> None: parser.add_argument("--legacy-cache-root") parser.add_argument("--legacy-source-root") parser.add_argument("--negative-embedding") + parser.add_argument("--approved-ids-manifest", required=True) + parser.add_argument("--expected-approved-ids-sha256") parser.add_argument("--heldout-count", required=True, type=int) - parser.add_argument("--split-seed", default="0") parser.add_argument("--shard-size", default=10000, type=int) args = parser.parse_args() report = migrate_cache( args.source_root, args.output_root, + approved_ids_manifest=args.approved_ids_manifest, heldout_count=args.heldout_count, - split_seed=args.split_seed, + expected_approved_ids_sha256=args.expected_approved_ids_sha256, source_index=tuple(args.source_indexes) if args.source_indexes else "metadata.json", legacy_cache_root=args.legacy_cache_root, legacy_source_root=args.legacy_source_root, negative_embedding=args.negative_embedding, shard_size=args.shard_size, ) - print(report) + print(json.dumps(report, sort_keys=True, allow_nan=False)) if __name__ == "__main__": diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 13014823709..88635470292 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -61,6 +61,8 @@ fsdp: data: all_metadata_index: metadata.json validation_metadata_index: metadata_heldout.json + expected_approved_ordered_ids_sha256: + expected_heldout_count: 2000 dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: data/qwen_image_cache diff --git a/examples/diffusers/fastgen/pdd_finetune.py b/examples/diffusers/fastgen/pdd_finetune.py index c5c241567fc..78943cd6f62 100644 --- a/examples/diffusers/fastgen/pdd_finetune.py +++ b/examples/diffusers/fastgen/pdd_finetune.py @@ -7,7 +7,6 @@ import argparse import dataclasses -import hashlib import logging import sys import time @@ -16,6 +15,7 @@ from typing import Any import yaml +from portable_cache import ordered_sample_ids_sha256 sys.dont_write_bytecode = True @@ -25,6 +25,8 @@ if str(path) not in sys.path: sys.path.insert(0, str(path)) +_CANONICAL_PDD_HELDOUT_COUNT = 2000 + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) @@ -50,13 +52,7 @@ def _metadata_sample_ids(metadata: Any, *, split: str) -> tuple[str, ...]: def _ordered_id_sha256(metadata: Any, *, split: str) -> str: - sample_ids = _metadata_sample_ids(metadata, split=split) - digest = hashlib.sha256() - digest.update(f"modelopt-pdd-ordered-{split}-ids-v1\0".encode()) - for sample_id in sample_ids: - digest.update(sample_id.encode()) - digest.update(b"\n") - return digest.hexdigest() + return ordered_sample_ids_sha256(_metadata_sample_ids(metadata, split=split)) def _dataloader_options(raw: Mapping[str, Any]) -> dict[str, Any]: @@ -132,6 +128,15 @@ def _validate_dataset_snapshot(raw: Mapping[str, Any], config: Any) -> Mapping[s from portable_cache import resolve_cache_root from validate_cache_snapshot import validate_snapshot + if config.expected_approved_ordered_ids_sha256 is None: + raise ValueError( + "PDD training requires data.expected_approved_ordered_ids_sha256 before validation." + ) + if config.expected_heldout_count != _CANONICAL_PDD_HELDOUT_COUNT: + raise ValueError( + f"PDD training requires data.expected_heldout_count={_CANONICAL_PDD_HELDOUT_COUNT}." + ) + options = _dataloader_options(raw) root = resolve_cache_root(options["cache_dir"]) roots: list[str] = [""] * dist.get_world_size() @@ -147,6 +152,8 @@ def _validate_dataset_snapshot(raw: Mapping[str, Any], config: Any) -> Mapping[s all_index=config.all_metadata_index, train_index=config.train_metadata_index, heldout_index=config.validation_metadata_index, + expected_approved_ids_sha256=(config.expected_approved_ordered_ids_sha256), + expected_heldout_count=config.expected_heldout_count, ) status = {"ok": True, "report": report} except BaseException as error: @@ -164,6 +171,23 @@ def _validate_dataset_snapshot(raw: Mapping[str, Any], config: Any) -> Mapping[s return report +def _validated_loader_order_hashes( + train_metadata: Any, + heldout_metadata: Any, + snapshot_report: Mapping[str, Any], +) -> tuple[str, str]: + train_digest = _ordered_id_sha256(train_metadata, split="train") + heldout_digest = _ordered_id_sha256(heldout_metadata, split="heldout") + reported = snapshot_report.get("ordered_sample_ids_sha256") + if not isinstance(reported, Mapping): + raise RuntimeError("PDD dataset snapshot report has no ordered sample-ID hashes.") + if train_digest != reported.get("train"): + raise RuntimeError("training loader order does not match the authenticated snapshot.") + if heldout_digest != reported.get("heldout"): + raise RuntimeError("heldout loader order does not match the authenticated snapshot.") + return train_digest, heldout_digest + + def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[bool, ...], ...]]: import torch.distributed as dist from pdd_training import build_pdd_validation_assignments @@ -421,6 +445,11 @@ def main() -> None: validation_sampler, config, ) + train_ordered_id_sha256, heldout_ordered_id_sha256 = _validated_loader_order_hashes( + sampler.dataset.metadata, + validation_sampler.dataset.metadata, + snapshot_report, + ) setup = build_pdd_setup(config) transformer_config = getattr(setup.student, "config", None) if isinstance(transformer_config, Mapping): @@ -446,11 +475,8 @@ def main() -> None: guidance_rescale=config.guidance.rescale, guidance_eps=config.guidance.eps, automodel_snapshot=setup.automodel_snapshot, - ordered_train_id_sha256=_ordered_id_sha256(sampler.dataset.metadata, split="train"), - ordered_heldout_id_sha256=_ordered_id_sha256( - validation_sampler.dataset.metadata, - split="heldout", - ), + ordered_train_id_sha256=train_ordered_id_sha256, + ordered_heldout_id_sha256=heldout_ordered_id_sha256, dataset_snapshot_sha256=snapshot_report["snapshot_sha256"], local_batch_size=config.training.local_batch_size, grad_accumulation_steps=config.training.grad_accumulation_steps, diff --git a/examples/diffusers/fastgen/pdd_recipe.py b/examples/diffusers/fastgen/pdd_recipe.py index 1d292a62f80..508ba858880 100644 --- a/examples/diffusers/fastgen/pdd_recipe.py +++ b/examples/diffusers/fastgen/pdd_recipe.py @@ -88,6 +88,8 @@ class PDDRecipeConfig: all_metadata_index: str train_metadata_index: str validation_metadata_index: str + expected_approved_ordered_ids_sha256: str | None + expected_heldout_count: int | None device: torch.device dtype: torch.dtype fuse_qkv_projections: bool @@ -232,6 +234,24 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: data.get("validation_metadata_index", "metadata_heldout.json"), label="data.validation_metadata_index", ).as_posix() + expected_approved_ordered_ids_sha256 = data.get("expected_approved_ordered_ids_sha256") + if expected_approved_ordered_ids_sha256 is not None and ( + not isinstance(expected_approved_ordered_ids_sha256, str) + or len(expected_approved_ordered_ids_sha256) != 64 + or expected_approved_ordered_ids_sha256.lower() != expected_approved_ordered_ids_sha256 + or any( + character not in "0123456789abcdef" + for character in expected_approved_ordered_ids_sha256 + ) + ): + raise ValueError( + "data.expected_approved_ordered_ids_sha256 must be a lowercase hexadecimal SHA-256." + ) + expected_heldout_count = data.get("expected_heldout_count") + if expected_heldout_count is not None and ( + type(expected_heldout_count) is not int or expected_heldout_count <= 0 + ): + raise ValueError("data.expected_heldout_count must be a positive integer.") _reject_enabled(model.get("transformer_engine_linear"), name="global TE-linear conversion") _reject_enabled(model.get("peft"), name="PEFT/LoRA") @@ -435,6 +455,8 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: all_metadata_index=all_metadata_index, train_metadata_index=train_metadata_index, validation_metadata_index=validation_metadata_index, + expected_approved_ordered_ids_sha256=expected_approved_ordered_ids_sha256, + expected_heldout_count=expected_heldout_count, device=torch.device(model.get("device", "cuda" if torch.cuda.is_available() else "cpu")), dtype=_resolve_dtype(model.get("torch_dtype", "bfloat16")), fuse_qkv_projections=fuse_qkv_projections, diff --git a/examples/diffusers/fastgen/portable_cache.py b/examples/diffusers/fastgen/portable_cache.py index f837fe424a3..46d41629761 100644 --- a/examples/diffusers/fastgen/portable_cache.py +++ b/examples/diffusers/fastgen/portable_cache.py @@ -20,13 +20,28 @@ import hashlib import json import os +import stat +import struct from collections.abc import Mapping, Sequence from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any -CACHE_SCHEMA_VERSION = 1 +PREPROCESS_STAGING_SCHEMA_VERSION = 1 +PORTABLE_SNAPSHOT_SCHEMA_VERSION = 2 +APPROVED_IDS_SCHEMA_VERSION = 1 +SPLIT_POLICY_SCHEMA_VERSION = 1 DATASET_CACHE_ENV = "MODELOPT_FASTGEN_DATASET_CACHE_DIR" SAMPLE_ID_DOMAIN = "modelopt-fastgen-sample-v1" +PDD_HOLDOUT_DOMAIN = "modelopt-pdd-holdout-v1" +_ORDERED_IDS_DOMAIN = b"modelopt-fastgen-ordered-sample-ids-v1" +_SPLIT_POLICY_KEYS = { + "algorithm", + "approved_ordered_ids_sha256", + "domain", + "heldout_count", + "schema_version", +} +_SHA256_HEX_LENGTH = 64 _BANNED_PATH_KEYS = { "cache_dir", "image_path", @@ -46,6 +61,28 @@ "source_ref", } +__all__ = [ + "APPROVED_IDS_SCHEMA_VERSION", + "DATASET_CACHE_ENV", + "PDD_HOLDOUT_DOMAIN", + "PORTABLE_SNAPSHOT_SCHEMA_VERSION", + "PREPROCESS_STAGING_SCHEMA_VERSION", + "SAMPLE_ID_DOMAIN", + "SPLIT_POLICY_SCHEMA_VERSION", + "audit_no_absolute_paths", + "load_approved_sample_ids", + "load_portable_metadata", + "load_strict_json", + "ordered_sample_ids_sha256", + "resolve_cache_asset", + "resolve_cache_root", + "resolve_negative_embedding", + "select_pdd_holdout_ids", + "sha256_file", + "stable_sample_id", + "validate_relative_reference", +] + def resolve_cache_root(configured_root: str | os.PathLike[str]) -> Path: """Resolve the YAML root unless a valid absolute environment override is set.""" @@ -150,6 +187,147 @@ def sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str: return digest.hexdigest() +def _reject_duplicate_json_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + value: dict[str, Any] = {} + for key, nested in pairs: + if key in value: + raise ValueError(f"JSON object contains duplicate key {key!r}.") + value[key] = nested + return value + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"JSON contains non-standard constant {value!r}.") + + +def load_strict_json(path: Path, *, label: str) -> Any: + """Load one regular, non-symlink UTF-8 JSON file with strict object syntax.""" + path = Path(path) + if path.is_symlink(): + raise ValueError(f"{label} must not be a symlink: {path}") + mode = path.stat().st_mode + if not stat.S_ISREG(mode): + raise ValueError(f"{label} must be a regular file: {path}") + try: + with path.open(encoding="utf-8") as stream: + return json.load( + stream, + object_pairs_hook=_reject_duplicate_json_keys, + parse_constant=_reject_json_constant, + ) + except (json.JSONDecodeError, UnicodeDecodeError) as error: + raise ValueError(f"{label} is not valid UTF-8 JSON: {path}") from error + + +def _validate_sha256(value: Any, *, label: str) -> str: + if ( + not isinstance(value, str) + or len(value) != _SHA256_HEX_LENGTH + or value.lower() != value + or any(character not in "0123456789abcdef" for character in value) + ): + raise ValueError(f"{label} must be a 64-character lowercase hexadecimal SHA-256 digest.") + return value + + +def _validate_ordered_sample_ids(sample_ids: Sequence[str], *, label: str) -> tuple[str, ...]: + if isinstance(sample_ids, str | bytes) or not isinstance(sample_ids, Sequence): + raise TypeError(f"{label} must be a sequence of strings.") + resolved = tuple(sample_ids) + if not resolved: + raise ValueError(f"{label} must be non-empty.") + for index, sample_id in enumerate(resolved): + if not isinstance(sample_id, str) or not sample_id: + raise ValueError(f"{label}[{index}] must be a non-empty string.") + try: + sample_id.encode("utf-8") + except UnicodeEncodeError as error: + raise ValueError(f"{label}[{index}] must be UTF-8 encodable.") from error + if len(resolved) != len(set(resolved)): + raise ValueError(f"{label} contains duplicates.") + return resolved + + +def ordered_sample_ids_sha256(sample_ids: Sequence[str]) -> str: + """Hash an ordered logical-ID sequence with count and byte-length framing.""" + resolved = _validate_ordered_sample_ids(sample_ids, label="sample_ids") + digest = hashlib.sha256() + digest.update(_ORDERED_IDS_DOMAIN) + digest.update(b"\0") + digest.update(struct.pack(">Q", len(resolved))) + for sample_id in resolved: + encoded = sample_id.encode("utf-8") + digest.update(struct.pack(">Q", len(encoded))) + digest.update(encoded) + return digest.hexdigest() + + +def load_approved_sample_ids( + path: Path, + *, + expected_sha256: str | None = None, +) -> tuple[tuple[str, ...], str]: + """Load and authenticate an ordered post-filter sample-ID artifact.""" + value = load_strict_json(path, label="approved_ids_manifest") + if not isinstance(value, dict): + raise ValueError("approved_ids_manifest must contain an object.") + required = { + "ordered_sample_ids", + "ordered_sample_ids_sha256", + "schema_version", + } + if set(value) != required: + raise ValueError( + "approved_ids_manifest keys mismatch: " + f"expected={sorted(required)}, actual={sorted(value)}." + ) + if ( + type(value["schema_version"]) is not int + or value["schema_version"] != APPROVED_IDS_SCHEMA_VERSION + ): + raise ValueError( + f"approved_ids_manifest.schema_version must be {APPROVED_IDS_SCHEMA_VERSION}." + ) + sample_ids = _validate_ordered_sample_ids( + value["ordered_sample_ids"], + label="approved_ids_manifest.ordered_sample_ids", + ) + computed = ordered_sample_ids_sha256(sample_ids) + declared = _validate_sha256( + value["ordered_sample_ids_sha256"], + label="approved_ids_manifest.ordered_sample_ids_sha256", + ) + if declared != computed: + raise ValueError("approved_ids_manifest ordered sample-ID SHA-256 mismatch.") + if expected_sha256 is not None: + expected = _validate_sha256(expected_sha256, label="expected_approved_ids_sha256") + if expected != computed: + raise ValueError("approved_ids_manifest does not match expected approved-ID SHA-256.") + return sample_ids, computed + + +def select_pdd_holdout_ids( + sample_ids: Sequence[str], heldout_count: int +) -> tuple[tuple[str, ...], tuple[str, ...]]: + """Select the frozen seedless PDD holdout membership while preserving input order.""" + resolved = _validate_ordered_sample_ids(sample_ids, label="sample_ids") + if type(heldout_count) is not int or not 0 < heldout_count < len(resolved): + raise ValueError( + "heldout_count must be an integer strictly between zero and len(sample_ids)." + ) + ranked = sorted( + resolved, + key=lambda sample_id: ( + hashlib.sha256(f"{PDD_HOLDOUT_DOMAIN}\0{sample_id}".encode()).digest(), + sample_id.encode("utf-8"), + ), + ) + heldout_members = set(ranked[:heldout_count]) + train = tuple(sample_id for sample_id in resolved if sample_id not in heldout_members) + heldout = tuple(sample_id for sample_id in resolved if sample_id in heldout_members) + return train, heldout + + def stable_sample_id(*, source_ref: str, resolution: Sequence[int], model_type: str) -> str: """Derive a root-independent sample ID from a logical source and processing identity.""" logical = validate_relative_reference(source_ref, label="source_ref").as_posix() @@ -214,11 +392,7 @@ def audit_no_absolute_paths( def _load_json(path: Path, *, expected_type: type, label: str): - try: - with path.open(encoding="utf-8") as stream: - value = json.load(stream) - except json.JSONDecodeError as error: - raise ValueError(f"{label} is not valid JSON: {path}") from error + value = load_strict_json(path, label=label) if not isinstance(value, expected_type): raise ValueError( f"{label} must contain {expected_type.__name__}, got {type(value).__name__}." @@ -226,19 +400,60 @@ def _load_json(path: Path, *, expected_type: type, label: str): return value +def _resolve_strict_json_asset(root: Path, reference: str, *, label: str) -> Path: + relative = validate_relative_reference(reference, label=label) + candidate = root.resolve(strict=True) / relative + if candidate.is_symlink(): + raise ValueError(f"{label} must not be a symlink: {relative}") + return resolve_cache_asset(root, reference, label=label) + + +def _validate_split_policy(value: Any) -> dict[str, Any]: + if not isinstance(value, dict) or set(value) != _SPLIT_POLICY_KEYS: + actual = sorted(value) if isinstance(value, dict) else type(value).__name__ + raise ValueError( + "metadata_index.split_policy keys mismatch: " + f"expected={sorted(_SPLIT_POLICY_KEYS)}, actual={actual}." + ) + if ( + type(value["schema_version"]) is not int + or value["schema_version"] != SPLIT_POLICY_SCHEMA_VERSION + ): + raise ValueError( + f"metadata_index.split_policy.schema_version must be {SPLIT_POLICY_SCHEMA_VERSION}." + ) + if value["algorithm"] != "sha256-domain-ranked": + raise ValueError("metadata_index.split_policy.algorithm is unsupported.") + if value["domain"] != PDD_HOLDOUT_DOMAIN: + raise ValueError("metadata_index.split_policy.domain is unsupported.") + heldout_count = value["heldout_count"] + if type(heldout_count) is not int or heldout_count <= 0: + raise ValueError("metadata_index.split_policy.heldout_count must be a positive integer.") + _validate_sha256( + value["approved_ordered_ids_sha256"], + label="metadata_index.split_policy.approved_ordered_ids_sha256", + ) + return value + + def load_portable_metadata( root: Path, metadata_index: str = "metadata.json", ) -> tuple[dict[str, Any], list[dict[str, Any]]]: """Load and validate one split index, filtering IDs before bucket grouping.""" - index_path = resolve_cache_asset(root, metadata_index, label="metadata_index") + index_path = _resolve_strict_json_asset(root, metadata_index, label="metadata_index") index = _load_json(index_path, expected_type=dict, label="metadata_index") audit_no_absolute_paths(index, context="metadata_index") - if index.get("schema_version") != CACHE_SCHEMA_VERSION: + if ( + type(index.get("schema_version")) is not int + or index["schema_version"] != PORTABLE_SNAPSHOT_SCHEMA_VERSION + ): raise ValueError( - f"metadata_index schema_version must be {CACHE_SCHEMA_VERSION}, " - f"got {index.get('schema_version')!r}." + "metadata_index is not an authenticated portable snapshot: " + f"schema_version must be {PORTABLE_SNAPSHOT_SCHEMA_VERSION}, got " + f"{index.get('schema_version')!r}; run migrate_cache_manifest.py." ) + split_policy = _validate_split_policy(index.get("split_policy")) shards = index.get("shards") sample_ids = index.get("sample_ids") if not isinstance(shards, list) or not shards or any(not isinstance(v, str) for v in shards): @@ -253,14 +468,45 @@ def load_portable_metadata( raise ValueError("metadata_index.sample_ids contains duplicates.") if len(shards) != len(set(shards)): raise ValueError("metadata_index.shards contains duplicates.") - if "num_shards" in index and index["num_shards"] != len(shards): - raise ValueError("metadata_index.num_shards must equal len(shards).") - if index.get("total_items") != len(sample_ids): + if "num_shards" in index and ( + type(index["num_shards"]) is not int or index["num_shards"] != len(shards) + ): + raise ValueError("metadata_index.num_shards must be an integer equal to len(shards).") + if type(index.get("total_items")) is not int or index["total_items"] != len(sample_ids): raise ValueError("metadata_index.total_items must equal len(sample_ids).") + computed_ordered_hash = ordered_sample_ids_sha256(sample_ids) + declared_ordered_hash = _validate_sha256( + index.get("ordered_sample_ids_sha256"), + label="metadata_index.ordered_sample_ids_sha256", + ) + if declared_ordered_hash != computed_ordered_hash: + raise ValueError("metadata_index ordered sample-ID SHA-256 mismatch.") + if index.get("split") == "all" and ( + split_policy["approved_ordered_ids_sha256"] != computed_ordered_hash + ): + raise ValueError("all metadata index does not match the approved ordered sample-ID hash.") + negative = index.get("negative_prompt_embedding") + if negative is not None: + if not isinstance(negative, dict) or set(negative) != {"path", "sha256"}: + raise ValueError( + "metadata_index.negative_prompt_embedding must contain path and sha256." + ) + validate_relative_reference( + negative["path"], + label="metadata_index.negative_prompt_embedding.path", + ) + _validate_sha256( + negative["sha256"], + label="metadata_index.negative_prompt_embedding.sha256", + ) entries_by_id: dict[str, dict[str, Any]] = {} for shard_number, shard_ref in enumerate(shards): - shard_path = resolve_cache_asset(root, shard_ref, label=f"shards[{shard_number}]") + shard_path = _resolve_strict_json_asset( + root, + shard_ref, + label=f"shards[{shard_number}]", + ) entries = _load_json(shard_path, expected_type=list, label=f"shards[{shard_number}]") for entry_number, entry in enumerate(entries): label = f"shards[{shard_number}][{entry_number}]" @@ -273,15 +519,7 @@ def load_portable_metadata( if sample_id in entries_by_id: raise ValueError(f"duplicate sample_id in shards: {sample_id}") resolve_cache_asset(root, entry.get("cache_file"), label=f"{label}.cache_file") - payload_sha256 = entry.get("payload_sha256") - if not isinstance(payload_sha256, str) or len(payload_sha256) != 64: - raise ValueError(f"{label}.payload_sha256 must be a hexadecimal SHA-256 digest.") - try: - int(payload_sha256, 16) - except ValueError as error: - raise ValueError( - f"{label}.payload_sha256 must be a hexadecimal SHA-256 digest." - ) from error + _validate_sha256(entry.get("payload_sha256"), label=f"{label}.payload_sha256") if "source_ref" in entry: validate_relative_reference(entry["source_ref"], label=f"{label}.source_ref") entries_by_id[sample_id] = entry diff --git a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py index 8bd0fd1b711..ce604a2a665 100644 --- a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py +++ b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py @@ -136,7 +136,7 @@ def _save_metadata_shards( 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 = { @@ -154,7 +154,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/validate_cache_snapshot.py b/examples/diffusers/fastgen/validate_cache_snapshot.py index b7c465a94ac..c38e13a538d 100644 --- a/examples/diffusers/fastgen/validate_cache_snapshot.py +++ b/examples/diffusers/fastgen/validate_cache_snapshot.py @@ -19,14 +19,18 @@ import argparse import hashlib +import json from pathlib import Path from typing import Any import torch from portable_cache import ( + PORTABLE_SNAPSHOT_SCHEMA_VERSION, audit_no_absolute_paths, load_portable_metadata, + ordered_sample_ids_sha256, resolve_cache_asset, + select_pdd_holdout_ids, sha256_file, validate_relative_reference, ) @@ -58,6 +62,8 @@ def validate_snapshot( train_index: str = "metadata_train.json", heldout_index: str = "metadata_heldout.json", reject_orphans: bool = True, + expected_approved_ids_sha256: str | None = None, + expected_heldout_count: int | None = None, ) -> dict[str, Any]: """Validate manifests, payload hashes, splits, and declared snapshot inventory. @@ -77,7 +83,10 @@ def validate_snapshot( if missing_indexes: raise FileNotFoundError(f"required metadata indices do not exist: {missing_indexes}") - split_ids: dict[str, set[str]] = {} + split_ids: dict[str, tuple[str, ...]] = {} + split_policies: dict[str, dict[str, Any]] = {} + ordered_hashes: dict[str, str] = {} + index_hashes: dict[str, str] = {} entries_by_id: dict[str, dict[str, Any]] = {} declared_files = { resolve_cache_asset(root, name, label="metadata_index") @@ -92,7 +101,12 @@ def validate_snapshot( raise ValueError(f"{index_name}.split must be {expected_split!r}, got {split_name!r}") if split_name in split_ids: raise ValueError(f"duplicate split declaration: {split_name}") - split_ids[split_name] = {entry["sample_id"] for entry in entries} + split_ids[split_name] = tuple(entry["sample_id"] for entry in entries) + split_policies[split_name] = dict(index["split_policy"]) + ordered_hashes[split_name] = ordered_sample_ids_sha256(split_ids[split_name]) + if ordered_hashes[split_name] != index["ordered_sample_ids_sha256"]: + raise ValueError(f"{index_name} ordered sample-ID SHA-256 mismatch") + index_hashes[split_name] = sha256_file(root / index_name) for shard_ref in index["shards"]: declared_files.add(resolve_cache_asset(root, shard_ref, label="metadata shard")) @@ -124,15 +138,50 @@ def validate_snapshot( if len(negative_declarations) > 1: raise ValueError("metadata indices disagree on the negative prompt embedding") - train_ids = split_ids.get("train") - heldout_ids = split_ids.get("heldout") - if train_ids is not None and heldout_ids is not None: - overlap = train_ids & heldout_ids - if overlap: - raise ValueError(f"train and heldout splits overlap: {sorted(overlap)[:5]}") - all_ids = split_ids.get("all") - if all_ids is not None and train_ids | heldout_ids != all_ids: - raise ValueError("train and heldout split union does not equal the all split") + policies = tuple(split_policies.values()) + if any(policy != policies[0] for policy in policies[1:]): + raise ValueError("metadata indices disagree on the split policy") + split_policy = split_policies["all"] + + train_ids = split_ids["train"] + heldout_ids = split_ids["heldout"] + all_ids = split_ids["all"] + train_members = set(train_ids) + heldout_members = set(heldout_ids) + overlap = train_members & heldout_members + if overlap: + raise ValueError(f"train and heldout splits overlap: {sorted(overlap)[:5]}") + if train_members | heldout_members != set(all_ids): + raise ValueError("train and heldout split union does not equal the all split") + expected_train, expected_heldout = select_pdd_holdout_ids( + all_ids, + split_policy["heldout_count"], + ) + if train_ids != expected_train: + raise ValueError("train split membership/order does not match the frozen PDD policy") + if heldout_ids != expected_heldout: + raise ValueError("heldout split membership/order does not match the frozen PDD policy") + if split_policy["approved_ordered_ids_sha256"] != ordered_hashes["all"]: + raise ValueError("split policy approved ordered-ID hash does not match the all index") + if expected_approved_ids_sha256 is not None: + if ( + not isinstance(expected_approved_ids_sha256, str) + or len(expected_approved_ids_sha256) != 64 + or expected_approved_ids_sha256.lower() != expected_approved_ids_sha256 + or any( + character not in "0123456789abcdef" for character in expected_approved_ids_sha256 + ) + ): + raise ValueError("expected_approved_ids_sha256 must be lowercase hexadecimal SHA-256") + if expected_approved_ids_sha256 != ordered_hashes["all"]: + raise ValueError("snapshot does not match the expected approved ordered-ID hash") + if expected_heldout_count is not None: + if type(expected_heldout_count) is not int or expected_heldout_count <= 0: + raise ValueError("expected_heldout_count must be a positive integer") + if split_policy["heldout_count"] != expected_heldout_count: + raise ValueError("snapshot heldout count does not match expected_heldout_count") + if len(heldout_ids) != split_policy["heldout_count"]: + raise ValueError("heldout split length does not match split policy") payload_hashes = dict(_validate_payload(root, entry) for entry in entries_by_id.values()) payload_files = set(payload_hashes) @@ -178,10 +227,21 @@ def validate_snapshot( snapshot_digest.update(file_digest.encode()) snapshot_digest.update(b"\n") + negative_report = None + if negative_declarations: + negative_path, negative_sha256 = next(iter(negative_declarations)) + negative_report = {"path": negative_path, "sha256": negative_sha256} + return { - "root": str(root), - "indexes": list(expected_indexes.values()), + "schema_version": 1, + "record_type": "modelopt_fastgen_portable_snapshot_validation", + "snapshot_schema_version": PORTABLE_SNAPSHOT_SCHEMA_VERSION, + "indexes": expected_indexes, + "split_policy": split_policy, "splits": {name: len(ids) for name, ids in split_ids.items()}, + "ordered_sample_ids_sha256": ordered_hashes, + "index_sha256": index_hashes, + "negative_prompt_embedding": negative_report, "unique_payloads": len(payload_files), "declared_files": len(declared_hashes), "snapshot_sha256": snapshot_digest.hexdigest(), @@ -195,6 +255,8 @@ def main() -> None: parser.add_argument("--train-index", default="metadata_train.json") parser.add_argument("--heldout-index", default="metadata_heldout.json") parser.add_argument("--allow-orphans", action="store_true") + parser.add_argument("--expected-approved-ids-sha256") + parser.add_argument("--expected-heldout-count", type=int) args = parser.parse_args() report = validate_snapshot( args.cache_root, @@ -202,8 +264,10 @@ def main() -> None: train_index=args.train_index, heldout_index=args.heldout_index, reject_orphans=not args.allow_orphans, + expected_approved_ids_sha256=args.expected_approved_ids_sha256, + expected_heldout_count=args.expected_heldout_count, ) - print(report) + print(json.dumps(report, sort_keys=True, allow_nan=False)) if __name__ == "__main__": diff --git a/tests/examples/diffusers/fastgen/pdd_test_utils.py b/tests/examples/diffusers/fastgen/pdd_test_utils.py index 89bcaccaf82..32b090281d6 100644 --- a/tests/examples/diffusers/fastgen/pdd_test_utils.py +++ b/tests/examples/diffusers/fastgen/pdd_test_utils.py @@ -5,12 +5,12 @@ from __future__ import annotations -import hashlib from dataclasses import dataclass from typing import Any import torch from pdd_training import PDDTrainer, PreparedPDDBatch +from portable_cache import ordered_sample_ids_sha256 from torch import nn from modelopt.torch.fastgen import ( @@ -211,9 +211,4 @@ def __getitem__(self, index: int) -> int: def ordered_id_sha256(sample_ids: tuple[str, ...]) -> str: - digest = hashlib.sha256() - digest.update(b"modelopt-pdd-ordered-train-ids-v1\0") - for sample_id in sample_ids: - digest.update(sample_id.encode()) - digest.update(b"\n") - return digest.hexdigest() + return ordered_sample_ids_sha256(sample_ids) diff --git a/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py b/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py index 9804faafe1e..db72855b64a 100644 --- a/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py +++ b/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py @@ -19,12 +19,25 @@ sys.path.insert(0, str(_FASTGEN_DIR)) import migrate_cache_manifest as migration -from portable_cache import load_portable_metadata +from portable_cache import load_portable_metadata, ordered_sample_ids_sha256 from validate_cache_snapshot import validate_snapshot def _write_json(path: pathlib.Path, value) -> None: - path.write_text(json.dumps(value, indent=2) + "\n") + path.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + + +def _write_approved(path: pathlib.Path, sample_ids: list[str] | tuple[str, ...]) -> str: + digest = ordered_sample_ids_sha256(sample_ids) + _write_json( + path, + { + "schema_version": 1, + "ordered_sample_ids": list(sample_ids), + "ordered_sample_ids_sha256": digest, + }, + ) + return digest def _make_legacy_cache( @@ -105,21 +118,33 @@ def test_migration_is_path_independent_and_relocatable(tmp_path): alice_output = tmp_path / "alice-portable" bob_output = tmp_path / "bob-portable" - migration.migrate_cache( + planned_ids = [ + record.sample_id + for record in migration.plan_migration( + alice, + legacy_cache_root=alice_cache_prefix, + legacy_source_root=alice_source_prefix, + ) + ] + approved_ids = [planned_ids[index] for index in (2, 0, 3)] + approved = tmp_path / "approved.json" + approved_digest = _write_approved(approved, approved_ids) + alice_result = migration.migrate_cache( alice, alice_output, + approved_ids_manifest=approved, heldout_count=1, - split_seed="fixed", + expected_approved_ids_sha256=approved_digest, legacy_cache_root=alice_cache_prefix, legacy_source_root=alice_source_prefix, negative_embedding="legacy-negative.pt", shard_size=2, ) - migration.migrate_cache( + bob_result = migration.migrate_cache( bob, bob_output, + approved_ids_manifest=approved, heldout_count=1, - split_seed="fixed", legacy_cache_root=bob_cache_prefix, legacy_source_root=bob_source_prefix, negative_embedding="legacy-negative.pt", @@ -132,9 +157,38 @@ def test_migration_is_path_independent_and_relocatable(tmp_path): assert alice_train == bob_train alice_report = validate_snapshot(alice_output) bob_report = validate_snapshot(bob_output) - assert alice_report["splits"] == {"all": 4, "train": 3, "heldout": 1} - assert bob_report["splits"] == {"all": 4, "train": 3, "heldout": 1} + assert alice_report["splits"] == {"all": 3, "train": 2, "heldout": 1} + assert bob_report["splits"] == {"all": 3, "train": 2, "heldout": 1} assert alice_report["snapshot_sha256"] == bob_report["snapshot_sha256"] + assert alice_report["ordered_sample_ids_sha256"]["all"] == approved_digest + assert alice_report["split_policy"]["approved_ordered_ids_sha256"] == approved_digest + assert alice_result["validation"] == bob_result["validation"] + assert alice_result["counts"] == {"source": 4, "approved": 3, "filtered": 1} + assert set(alice_result) == { + "schema_version", + "record_type", + "output_root", + "counts", + "validation", + } + assert set(alice_result["validation"]) == { + "schema_version", + "record_type", + "snapshot_schema_version", + "indexes", + "split_policy", + "splits", + "ordered_sample_ids_sha256", + "index_sha256", + "negative_prompt_embedding", + "unique_payloads", + "declared_files", + "snapshot_sha256", + } + validation_text = json.dumps(alice_result["validation"], sort_keys=True) + assert str(alice_output) not in validation_text + assert ".staging-" not in validation_text + assert alice_result["output_root"] != bob_result["output_root"] cli = subprocess.run( [ @@ -186,6 +240,7 @@ def test_incomplete_pass_one_publishes_nothing(monkeypatch, tmp_path): migration.migrate_cache( legacy, output, + approved_ids_manifest=tmp_path / "unused-approved.json", heldout_count=1, legacy_cache_root=cache_prefix, legacy_source_root=source_prefix, @@ -195,6 +250,143 @@ def test_incomplete_pass_one_publishes_nothing(monkeypatch, tmp_path): assert save_calls == [] +def test_approved_artifact_failures_and_final_validation_publish_nothing( + monkeypatch, tmp_path +) -> None: + legacy = tmp_path / "legacy" + cache_prefix = pathlib.Path("/legacy/cache") + source_prefix = pathlib.Path("/legacy/images") + _make_legacy_cache( + legacy, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + records = migration.plan_migration( + legacy, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + sample_ids = [record.sample_id for record in records] + approved = tmp_path / "approved.json" + digest = _write_approved(approved, sample_ids) + + with pytest.raises(ValueError, match="expected approved-ID"): + migration.migrate_cache( + legacy, + tmp_path / "bad-external", + approved_ids_manifest=approved, + expected_approved_ids_sha256="f" * 64, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + unknown = tmp_path / "unknown.json" + _write_approved(unknown, [*sample_ids, "unknown"]) + with pytest.raises(ValueError, match="unknown sample IDs"): + migration.migrate_cache( + legacy, + tmp_path / "unknown-output", + approved_ids_manifest=unknown, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + symlink = tmp_path / "approved-symlink.json" + symlink.symlink_to(approved) + with pytest.raises(ValueError, match="symlink"): + migration.migrate_cache( + legacy, + tmp_path / "symlink-output", + approved_ids_manifest=symlink, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + + monkeypatch.setattr( + migration, + "validate_snapshot", + lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("injected final validation")), + ) + output = tmp_path / "validation-output" + with pytest.raises(ValueError, match="injected final validation"): + migration.migrate_cache( + legacy, + output, + approved_ids_manifest=approved, + expected_approved_ids_sha256=digest, + heldout_count=1, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + assert not output.exists() + assert not list(tmp_path.glob(".validation-output.staging-*")) + + +def test_migration_cli_removes_seed_and_requires_approved_manifest() -> None: + cli = subprocess.run( + [sys.executable, str(_FASTGEN_DIR / "migrate_cache_manifest.py"), "--help"], + check=False, + capture_output=True, + text=True, + ) + assert cli.returncode == 0, cli.stderr + assert "--approved-ids-manifest" in cli.stdout + assert "--split-seed" not in cli.stdout + + +def test_migration_rejects_finalized_and_non_strict_source_json(tmp_path) -> None: + cache_prefix = pathlib.Path("/legacy/cache") + source_prefix = pathlib.Path("/legacy/images") + + finalized = tmp_path / "finalized-source" + _make_legacy_cache( + finalized, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + index_path = finalized / "metadata.json" + index = json.loads(index_path.read_text()) + index["schema_version"] = 2 + _write_json(index_path, index) + with pytest.raises(ValueError, match="schema_version is unsupported"): + migration.plan_migration( + finalized, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + + duplicate = tmp_path / "duplicate-source" + _make_legacy_cache( + duplicate, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + (duplicate / "metadata.json").write_text( + '{"shards":["legacy-shard-0.json"],"shards":["legacy-shard-1.json"]}' + ) + with pytest.raises(ValueError, match="duplicate key"): + migration.plan_migration( + duplicate, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + + nonfinite = tmp_path / "nonfinite-source" + _make_legacy_cache( + nonfinite, + stored_cache_root=cache_prefix, + stored_source_root=source_prefix, + ) + (nonfinite / "metadata.json").write_text('{"shards":[],"total_items":NaN}') + with pytest.raises(ValueError, match="non-standard constant"): + migration.plan_migration( + nonfinite, + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + + def test_changed_source_after_frozen_plan_cleans_staging(monkeypatch, tmp_path): legacy = tmp_path / "legacy" cache_prefix = pathlib.Path("/legacy/cache") @@ -218,10 +410,13 @@ def _return_frozen(*args, **kwargs): monkeypatch.setattr(migration, "plan_migration", _return_frozen) output = tmp_path / "portable" + approved = tmp_path / "approved.json" + _write_approved(approved, [record.sample_id for record in frozen]) with pytest.raises(RuntimeError, match="changed after pass 1"): migration.migrate_cache( legacy, output, + approved_ids_manifest=approved, heldout_count=1, legacy_cache_root=cache_prefix, legacy_source_root=source_prefix, @@ -250,6 +445,7 @@ def test_invalid_legacy_reference_fails_before_publish(tmp_path): migration.migrate_cache( legacy, output, + approved_ids_manifest=tmp_path / "unused-approved.json", heldout_count=1, legacy_cache_root=cache_prefix, legacy_source_root=source_prefix, @@ -275,6 +471,7 @@ def test_incomplete_source_counts_and_rank_indices_publish_nothing(tmp_path): migration.migrate_cache( legacy, output, + approved_ids_manifest=tmp_path / "unused-approved.json", heldout_count=1, legacy_cache_root=cache_prefix, legacy_source_root=source_prefix, @@ -299,6 +496,7 @@ def test_incomplete_source_counts_and_rank_indices_publish_nothing(tmp_path): migration.migrate_cache( legacy, incomplete_output, + approved_ids_manifest=tmp_path / "unused-approved.json", source_index="metadata_r00.json", heldout_count=1, legacy_cache_root=cache_prefix, @@ -307,9 +505,18 @@ def test_incomplete_source_counts_and_rank_indices_publish_nothing(tmp_path): assert not incomplete_output.exists() complete_output = tmp_path / "complete-ranks-output" + complete_records = migration.plan_migration( + legacy, + source_index=("metadata_r00.json", "metadata_r01.json"), + legacy_cache_root=cache_prefix, + legacy_source_root=source_prefix, + ) + complete_approved = tmp_path / "complete-approved.json" + _write_approved(complete_approved, [record.sample_id for record in complete_records]) migration.migrate_cache( legacy, complete_output, + approved_ids_manifest=complete_approved, source_index=("metadata_r00.json", "metadata_r01.json"), heldout_count=1, legacy_cache_root=cache_prefix, @@ -348,6 +555,7 @@ def test_migration_path_audit_allows_prompt_commands_but_rejects_set_paths(tmp_p migration.migrate_cache( legacy, output, + approved_ids_manifest=tmp_path / "unused-approved.json", heldout_count=1, legacy_cache_root=cache_prefix, legacy_source_root=source_prefix, diff --git a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py index 8a60c146a34..8bb88a56d76 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py +++ b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py @@ -99,6 +99,18 @@ def _stage_result(stage: str) -> dict: } +def test_gpu_harness_uses_shared_unambiguous_ordered_id_hash() -> None: + assert smoke._ordered_id_sha256(("a", "b")) == ( + "8cf774af4e8509811c2d4bc2adec6b852e4c614f9d8d833924502ead7c0689d7" + ) + assert smoke._ordered_id_sha256(("a\nb",)) == ( + "41e07cc133e8a85fc4a08e60a38c223f3c24dbca80312d106f251e533254eedf" + ) + source = _HARNESS.read_text() + assert "modelopt-pdd-ordered-{split}-ids-v1" not in source + assert 'digest.update(b"\\n")' not in source + + def _automodel_snapshot_fixture() -> tuple[dict, dict]: records = [] tree = hashlib.sha256() diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 323b959d178..f198f64ced2 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -23,6 +23,7 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) +import pdd_finetune from pdd_recipe import ( build_pdd_export_setup, build_pdd_setup, @@ -80,6 +81,80 @@ def test_example_recipe_explicitly_pins_grid_max_t() -> None: raw = yaml.safe_load((_FASTGEN_DIR / "configs" / "pdd_qwen_image.yaml").read_text()) assert type(raw["pdd"]["grid_max_t"]) is float assert raw["pdd"]["grid_max_t"] == 0.999 + assert raw["data"]["expected_approved_ordered_ids_sha256"] is None + assert raw["data"]["expected_heldout_count"] == 2000 + + +def test_data_authentication_config_fields_are_strict_but_nullable(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["data"] = { + "expected_approved_ordered_ids_sha256": "a" * 64, + "expected_heldout_count": 2000, + } + config = resolve_pdd_recipe_config(raw) + assert config.expected_approved_ordered_ids_sha256 == "a" * 64 + assert config.expected_heldout_count == 2000 + + for invalid_hash in ("A" * 64, "a" * 63, 7): + raw["data"]["expected_approved_ordered_ids_sha256"] = invalid_hash + with pytest.raises(ValueError, match="expected_approved_ordered_ids_sha256"): + resolve_pdd_recipe_config(raw) + raw["data"]["expected_approved_ordered_ids_sha256"] = None + for invalid_count in (0, -1, True, 1.5): + raw["data"]["expected_heldout_count"] = invalid_count + with pytest.raises(ValueError, match="expected_heldout_count"): + resolve_pdd_recipe_config(raw) + + +@pytest.mark.parametrize("heldout_count", [1999, 2001]) +def test_canonical_training_count_cannot_be_overridden( + monkeypatch, tmp_path, heldout_count +) -> None: + raw = _raw_config(tmp_path) + raw["data"] = { + "expected_approved_ordered_ids_sha256": "a" * 64, + "expected_heldout_count": heldout_count, + } + config = resolve_pdd_recipe_config(raw) + called = False + + def _unexpected_validator(*args, **kwargs): + nonlocal called + called = True + raise AssertionError("validator must not be called for a weakened canonical count") + + monkeypatch.setattr("validate_cache_snapshot.validate_snapshot", _unexpected_validator) + with pytest.raises(ValueError, match="expected_heldout_count=2000"): + pdd_finetune._validate_dataset_snapshot(raw, config) + assert not called + + +def test_canonical_training_requires_external_hash_before_validator(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["data"] = { + "expected_approved_ordered_ids_sha256": None, + "expected_heldout_count": 2000, + } + config = resolve_pdd_recipe_config(raw) + with pytest.raises(ValueError, match="expected_approved_ordered_ids_sha256"): + pdd_finetune._validate_dataset_snapshot(raw, config) + + +def test_loader_order_must_match_authenticated_report() -> None: + train = [{"sample_id": "train-a"}, {"sample_id": "train-b"}] + heldout = [{"sample_id": "heldout-a"}] + report = { + "ordered_sample_ids_sha256": { + "train": pdd_finetune._ordered_id_sha256(train, split="train"), + "heldout": pdd_finetune._ordered_id_sha256(heldout, split="heldout"), + } + } + assert pdd_finetune._validated_loader_order_hashes(train, heldout, report) == ( + report["ordered_sample_ids_sha256"]["train"], + report["ordered_sample_ids_sha256"]["heldout"], + ) + with pytest.raises(RuntimeError, match="training loader order"): + pdd_finetune._validated_loader_order_hashes(list(reversed(train)), heldout, report) @pytest.mark.parametrize( diff --git a/tests/examples/diffusers/fastgen/test_portable_cache.py b/tests/examples/diffusers/fastgen/test_portable_cache.py index 01552a71e1f..1ed97e14387 100644 --- a/tests/examples/diffusers/fastgen/test_portable_cache.py +++ b/tests/examples/diffusers/fastgen/test_portable_cache.py @@ -27,10 +27,15 @@ ) from portable_cache import ( DATASET_CACHE_ENV, + PORTABLE_SNAPSHOT_SCHEMA_VERSION, + PREPROCESS_STAGING_SCHEMA_VERSION, audit_no_absolute_paths, + load_approved_sample_ids, load_portable_metadata, + ordered_sample_ids_sha256, resolve_cache_root, resolve_negative_embedding, + select_pdd_holdout_ids, sha256_file, ) from validate_cache_snapshot import validate_snapshot @@ -44,7 +49,72 @@ def _sample_id(source_ref: str, resolution: tuple[int, int]) -> str: def _write_json(path: pathlib.Path, value) -> None: - path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n") + path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n") + + +def _write_approved_manifest(path: pathlib.Path, sample_ids: list[str]) -> str: + digest = ordered_sample_ids_sha256(sample_ids) + _write_json( + path, + { + "schema_version": 1, + "ordered_sample_ids": sample_ids, + "ordered_sample_ids_sha256": digest, + }, + ) + return digest + + +def _make_id_only_snapshot( + root: pathlib.Path, + sample_ids: tuple[str, ...], + heldout_count: int, + *, + heldout_override: tuple[str, ...] | None = None, +) -> None: + root.mkdir() + (root / "payload.pt").write_bytes(b"placeholder-not-read-before-split-gates") + entries = [ + { + "sample_id": sample_id, + "cache_file": "payload.pt", + "payload_sha256": "0" * 64, + } + for sample_id in sample_ids + ] + _write_json(root / "metadata_shard_s0000.json", entries) + if heldout_override is None: + train_ids, heldout_ids = select_pdd_holdout_ids(sample_ids, heldout_count) + else: + heldout_ids = heldout_override + heldout_members = set(heldout_ids) + train_ids = tuple(sample_id for sample_id in sample_ids if sample_id not in heldout_members) + approved_digest = ordered_sample_ids_sha256(sample_ids) + policy = { + "schema_version": 1, + "algorithm": "sha256-domain-ranked", + "domain": "modelopt-pdd-holdout-v1", + "heldout_count": heldout_count, + "approved_ordered_ids_sha256": approved_digest, + } + for name, split, ids in ( + ("metadata.json", "all", sample_ids), + ("metadata_train.json", "train", train_ids), + ("metadata_heldout.json", "heldout", heldout_ids), + ): + _write_json( + root / name, + { + "schema_version": 2, + "split": split, + "total_items": len(ids), + "num_shards": 1, + "shards": ["metadata_shard_s0000.json"], + "sample_ids": list(ids), + "ordered_sample_ids_sha256": ordered_sample_ids_sha256(ids), + "split_policy": policy, + }, + ) def _make_snapshot(root: pathlib.Path) -> dict[str, list[str]]: @@ -83,22 +153,29 @@ def _make_snapshot(root: pathlib.Path) -> dict[str, list[str]]: _write_json(root / "metadata_shard_s0000.json", entries) all_ids = [entry["sample_id"] for entry in entries] - splits = { - "all": all_ids, - "train": [all_ids[2], all_ids[0], all_ids[1]], - "heldout": [all_ids[3]], + train_ids, heldout_ids = select_pdd_holdout_ids(all_ids, 1) + splits = {"all": all_ids, "train": list(train_ids), "heldout": list(heldout_ids)} + approved_digest = ordered_sample_ids_sha256(all_ids) + split_policy = { + "schema_version": 1, + "algorithm": "sha256-domain-ranked", + "domain": "modelopt-pdd-holdout-v1", + "heldout_count": 1, + "approved_ordered_ids_sha256": approved_digest, } for split, ids in splits.items(): name = "metadata.json" if split == "all" else f"metadata_{split}.json" _write_json( root / name, { - "schema_version": 1, + "schema_version": 2, "split": split, "total_items": len(ids), "num_shards": 1, "shards": ["metadata_shard_s0000.json"], "sample_ids": ids, + "ordered_sample_ids_sha256": ordered_sample_ids_sha256(ids), + "split_policy": split_policy, }, ) negative_path = root / "negative_prompt_embedding.pt" @@ -121,6 +198,153 @@ def _batch_signature(dataset: TextToImageDataset) -> list[tuple[str, float]]: ] +def test_authenticated_cache_constants_hash_framing_and_seedless_split() -> None: + assert PREPROCESS_STAGING_SCHEMA_VERSION == 1 + assert PORTABLE_SNAPSHOT_SCHEMA_VERSION == 2 + expected = { + ("a\nb",): "41e07cc133e8a85fc4a08e60a38c223f3c24dbca80312d106f251e533254eedf", + ("a", "b"): "8cf774af4e8509811c2d4bc2adec6b852e4c614f9d8d833924502ead7c0689d7", + ("ab", "c"): "6df9e72da4c55f09b4c0320337d6a5d46396271ac61b7a60e1ee8146ce49709e", + ("a", "bc"): "7cedefc9d46613683a89c3081c3a743b66164861975cf100346d59c13cf31d26", + tuple( + str(index) for index in range(16) + ): "b157f73e9710fe1eb2c4f8d94286f304d5c2a9de2b09b31d2b1f5eee15448e69", + } + for sample_ids, digest in expected.items(): + assert ordered_sample_ids_sha256(sample_ids) == digest + assert len(set(expected.values())) == len(expected) + + train, heldout = select_pdd_holdout_ids(tuple(str(index) for index in range(16)), 4) + assert heldout == ("4", "7", "10", "13") + assert train == tuple(str(index) for index in range(16) if str(index) not in heldout) + assert heldout not in (("0", "2", "4", "11"), ("2", "7", "8", "9")) + + +def test_approved_id_artifact_is_strict_and_externally_authenticatable(tmp_path) -> None: + path = tmp_path / "approved.json" + digest = _write_approved_manifest(path, ["second", "first"]) + assert load_approved_sample_ids(path, expected_sha256=digest) == ( + ("second", "first"), + digest, + ) + + for value, message in ( + ({"schema_version": 1}, "keys mismatch"), + ( + { + "schema_version": 1, + "ordered_sample_ids": ["first", "first"], + "ordered_sample_ids_sha256": digest, + }, + "duplicates", + ), + ): + _write_json(path, value) + with pytest.raises(ValueError, match=message): + load_approved_sample_ids(path) + + path.write_text('{"schema_version":1,"schema_version":1}') + with pytest.raises(ValueError, match="duplicate key"): + load_approved_sample_ids(path) + path.write_text('{"schema_version": NaN}') + with pytest.raises(ValueError, match="non-standard constant"): + load_approved_sample_ids(path) + path.write_bytes(b"\xff") + with pytest.raises(ValueError, match="UTF-8 JSON"): + load_approved_sample_ids(path) + path.write_text("{") + with pytest.raises(ValueError, match="UTF-8 JSON"): + load_approved_sample_ids(path) + + real = tmp_path / "real.json" + _write_approved_manifest(real, ["first", "second"]) + path.unlink() + path.symlink_to(real) + with pytest.raises(ValueError, match="symlink"): + load_approved_sample_ids(path) + with pytest.raises(ValueError, match="regular file"): + load_approved_sample_ids(tmp_path) + with pytest.raises(ValueError, match="expected approved-ID"): + load_approved_sample_ids(real, expected_sha256="f" * 64) + + +@pytest.mark.parametrize( + "old_heldout", + [("0", "2", "4", "11"), ("2", "7", "8", "9")], +) +def test_self_consistent_old_seed_partitions_are_rejected(tmp_path, old_heldout) -> None: + root = tmp_path / "cache" + sample_ids = tuple(str(index) for index in range(16)) + _make_id_only_snapshot(root, sample_ids, 4, heldout_override=old_heldout) + with pytest.raises(ValueError, match="frozen PDD policy"): + validate_snapshot(root, reject_orphans=False) + + +def test_self_consistent_all_list_rewrite_fails_external_hash(tmp_path) -> None: + approved_ids = tuple(str(index) for index in range(16)) + rewritten_ids = approved_ids[:-1] + root = tmp_path / "cache" + _make_id_only_snapshot(root, rewritten_ids, 4) + with pytest.raises(ValueError, match="expected approved ordered-ID hash"): + validate_snapshot( + root, + expected_approved_ids_sha256=ordered_sample_ids_sha256(approved_ids), + reject_orphans=False, + ) + + +@pytest.mark.parametrize( + ("policy_update", "message"), + [ + ({"domain": "modelopt-fastgen-split-v1"}, "domain"), + ({"algorithm": "other"}, "algorithm"), + ({"schema_version": 2}, "schema_version"), + ({"heldout_count": True}, "heldout_count"), + ({"extra": "field"}, "keys mismatch"), + ], +) +def test_split_policy_tampering_is_rejected(tmp_path, policy_update, message) -> None: + root = tmp_path / "cache" + _make_snapshot(root) + for name in ("metadata.json", "metadata_train.json", "metadata_heldout.json"): + index = json.loads((root / name).read_text()) + index["split_policy"].update(policy_update) + _write_json(root / name, index) + with pytest.raises(ValueError, match=message): + validate_snapshot(root) + + +def test_strict_index_and_shard_json_reject_duplicates_and_nonfinite(tmp_path) -> None: + root = tmp_path / "cache" + _make_snapshot(root) + index_path = root / "metadata.json" + index_path.write_text('{"schema_version":2,"schema_version":2}') + with pytest.raises(ValueError, match="duplicate key"): + load_portable_metadata(root) + + _make_snapshot(tmp_path / "second") + shard_path = tmp_path / "second" / "metadata_shard_s0000.json" + shard_path.write_text('[{"sample_id":"a","value":NaN}]') + with pytest.raises(ValueError, match="non-standard constant"): + load_portable_metadata(tmp_path / "second") + + +@pytest.mark.parametrize("actual_heldout_count", [1999, 2001]) +def test_validator_rejects_noncanonical_actual_holdout_counts( + tmp_path, actual_heldout_count +) -> None: + root = tmp_path / "cache" + sample_ids = tuple(str(index) for index in range(actual_heldout_count + 1)) + _make_id_only_snapshot(root, sample_ids, actual_heldout_count) + with pytest.raises(ValueError, match="expected_heldout_count"): + validate_snapshot( + root, + expected_approved_ids_sha256=ordered_sample_ids_sha256(sample_ids), + expected_heldout_count=2000, + reject_orphans=False, + ) + + def test_cache_root_environment_precedence(monkeypatch, tmp_path): configured = tmp_path / "configured" override = tmp_path / "override" @@ -266,7 +490,10 @@ def test_collate_emits_logical_identity_without_source_paths(tmp_path): root = tmp_path / "cache" _make_snapshot(root) dataset = TextToImageDataset(str(root), metadata_index="metadata_train.json") - samples = [dataset[1], dataset[2]] + same_resolution = next( + group["indices"] for group in dataset.bucket_groups.values() if len(group["indices"]) >= 2 + ) + samples = [dataset[index] for index in same_resolution[:2]] output = collate_fn_text_to_image(samples) assert output["metadata"]["sample_ids"] == [item["sample_id"] for item in samples] assert output["metadata"]["source_refs"] == [item["source_ref"] for item in samples] @@ -288,6 +515,7 @@ def test_validator_detects_hash_split_and_orphan_failures(tmp_path): heldout_path = root / "metadata_heldout.json" heldout = json.loads(heldout_path.read_text()) heldout["sample_ids"] = [splits["train"][0]] + heldout["ordered_sample_ids_sha256"] = ordered_sample_ids_sha256(heldout["sample_ids"]) _write_json(heldout_path, heldout) with pytest.raises(ValueError, match="overlap"): validate_snapshot(root) @@ -349,11 +577,13 @@ def test_empty_split_is_rejected_before_bucket_grouping(tmp_path): _write_json( root / "metadata_empty.json", { - "schema_version": 1, + "schema_version": 2, "split": "empty", "total_items": 0, "shards": ["metadata_shard_s0000.json"], "sample_ids": [], + "ordered_sample_ids_sha256": "0" * 64, + "split_policy": json.loads((root / "metadata.json").read_text())["split_policy"], }, ) with pytest.raises(ValueError, match="non-empty"): @@ -436,12 +666,22 @@ def test_portable_index_writer_is_deterministic(monkeypatch, tmp_path): ) index = json.loads((staging / "metadata.json").read_text()) shard = json.loads((staging / index["shards"][0]).read_text()) + assert index["schema_version"] == PREPROCESS_STAGING_SCHEMA_VERSION expected_ids = sorted(entry["sample_id"] for entry in entries) assert index["sample_ids"] == expected_ids assert [entry["sample_id"] for entry in shard] == expected_ids assert not (staging / "metadata_train.json").exists() assert str(staging) not in json.dumps([index, shard]) + with pytest.raises(ValueError, match=r"migrate_cache_manifest\.py"): + load_portable_metadata(staging, "metadata.json") finalized = tmp_path / "finalized" - migrate_cache(staging, finalized, heldout_count=1) + approved = tmp_path / "approved.json" + _write_approved_manifest(approved, expected_ids) + migrate_cache( + staging, + finalized, + approved_ids_manifest=approved, + heldout_count=1, + ) assert validate_snapshot(finalized)["splits"] == {"all": 2, "train": 1, "heldout": 1} diff --git a/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py b/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py index 0b7cdafe396..908be763acc 100644 --- a/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py +++ b/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py @@ -664,13 +664,10 @@ def _modelopt_source() -> dict[str, Any]: return {"commit": commit, "dirty": False} -def _ordered_id_sha256(sample_ids: Sequence[str], *, split: str) -> str: - digest = hashlib.sha256() - digest.update(f"modelopt-pdd-ordered-{split}-ids-v1\0".encode()) - for sample_id in sample_ids: - digest.update(sample_id.encode()) - digest.update(b"\n") - return digest.hexdigest() +def _ordered_id_sha256(sample_ids: Sequence[str]) -> str: + from portable_cache import ordered_sample_ids_sha256 + + return ordered_sample_ids_sha256(sample_ids) def _training_sample_ids(world_size: int) -> tuple[str, ...]: @@ -773,8 +770,8 @@ def _identity( guidance_rescale=config.guidance.rescale, guidance_eps=config.guidance.eps, automodel_snapshot=setup.automodel_snapshot, - ordered_train_id_sha256=_ordered_id_sha256(train_ids, split="train"), - ordered_heldout_id_sha256=_ordered_id_sha256(heldout_ids, split="heldout"), + ordered_train_id_sha256=_ordered_id_sha256(train_ids), + ordered_heldout_id_sha256=_ordered_id_sha256(heldout_ids), dataset_snapshot_sha256=_canonical_sha256( {"domain": "modelopt-pdd-synthetic-smoke-v1", "config": raw["pdd"]} ), From ad21bbb3a1e68915b2773f55100821c890bc6d26 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 14 Jul 2026 16:33:40 -0700 Subject: [PATCH 18/45] fix(fastgen): make PDD data gate collective Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd_finetune.py | 64 +++++++++-- .../pdd_training_preflight_distributed.py | 34 +++++- .../test_pdd_qwen_operability_smoke.py | 78 ++++++++++++++ .../fastgen/test_pdd_recipe_setup.py | 101 ++++++++++++++++++ .../diffusers/fastgen/test_portable_cache.py | 89 ++++++++++++++- 5 files changed, 356 insertions(+), 10 deletions(-) diff --git a/examples/diffusers/fastgen/pdd_finetune.py b/examples/diffusers/fastgen/pdd_finetune.py index 78943cd6f62..17a284b10ec 100644 --- a/examples/diffusers/fastgen/pdd_finetune.py +++ b/examples/diffusers/fastgen/pdd_finetune.py @@ -15,16 +15,20 @@ from typing import Any import yaml -from portable_cache import ordered_sample_ids_sha256 sys.dont_write_bytecode = True _THIS_DIR = Path(__file__).resolve().parent _REPO_ROOT = _THIS_DIR.parents[2] +# These entrypoints are also supported through ``python -m``. In that mode the +# sibling ModelOpt-owned example modules are not importable until this directory +# is added explicitly. for path in (_REPO_ROOT, _THIS_DIR): if str(path) not in sys.path: sys.path.insert(0, str(path)) +from portable_cache import ordered_sample_ids_sha256 # noqa: E402 + _CANONICAL_PDD_HELDOUT_COUNT = 2000 @@ -188,6 +192,54 @@ def _validated_loader_order_hashes( return train_digest, heldout_digest +def _collective_validated_loader_order_hashes( + train_metadata: Any, + heldout_metadata: Any, + snapshot_report: Mapping[str, Any], +) -> tuple[str, str]: + """Authenticate loader order on every rank before distributed model construction.""" + import torch.distributed as dist + + try: + hashes = _validated_loader_order_hashes( + train_metadata, + heldout_metadata, + snapshot_report, + ) + local_status: dict[str, Any] = {"ok": True, "hashes": hashes} + except BaseException as error: + local_status = {"ok": False, "error": f"{type(error).__name__}: {error}"} + + statuses: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(statuses, local_status) + failures: list[str] = [] + resolved_hashes: list[tuple[str, str]] = [] + for rank, status in enumerate(statuses): + if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: + failures.append(f"rank {rank}: malformed loader authentication status") + continue + if not status["ok"]: + failures.append(f"rank {rank}: {status.get('error')}") + continue + hashes = status.get("hashes") + if ( + not isinstance(hashes, tuple | list) + or len(hashes) != 2 + or any(not isinstance(value, str) for value in hashes) + ): + failures.append(f"rank {rank}: malformed loader authentication hashes") + continue + resolved_hashes.append((hashes[0], hashes[1])) + if failures: + raise RuntimeError("PDD loader-order authentication failed: " + "; ".join(failures)) + if len(set(resolved_hashes)) != 1: + raise RuntimeError( + "PDD ranks resolved different authenticated train/heldout loader orders: " + f"{resolved_hashes}." + ) + return resolved_hashes[0] + + def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[bool, ...], ...]]: import torch.distributed as dist from pdd_training import build_pdd_validation_assignments @@ -441,15 +493,15 @@ def main() -> None: dp_rank=rank, dp_world_size=world_size, ) - validation_assignments, validation_masks = _build_validation_plan( - validation_sampler, - config, - ) - train_ordered_id_sha256, heldout_ordered_id_sha256 = _validated_loader_order_hashes( + train_ordered_id_sha256, heldout_ordered_id_sha256 = _collective_validated_loader_order_hashes( sampler.dataset.metadata, validation_sampler.dataset.metadata, snapshot_report, ) + validation_assignments, validation_masks = _build_validation_plan( + validation_sampler, + config, + ) setup = build_pdd_setup(config) transformer_config = getattr(setup.student, "config", None) if isinstance(transformer_config, Mapping): diff --git a/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py b/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py index 7007f3c7502..f08606ac266 100644 --- a/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py @@ -18,7 +18,12 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -from pdd_finetune import _collective_training_batch, _collective_training_iterator +from pdd_finetune import ( + _collective_training_batch, + _collective_training_iterator, + _collective_validated_loader_order_hashes, + _ordered_id_sha256, +) class _Sampler: @@ -153,6 +158,33 @@ def main() -> None: assert all( item is not None and "iterator construction" in item for item in iterator_messages ) + dist.barrier() + + canonical_train = [{"sample_id": "train-a"}, {"sample_id": "train-b"}] + heldout = [{"sample_id": "heldout-a"}] + report = { + "ordered_sample_ids_sha256": { + "train": _ordered_id_sha256(canonical_train, split="train"), + "heldout": _ordered_id_sha256(heldout, split="heldout"), + } + } + local_train = canonical_train if rank == 0 else list(reversed(canonical_train)) + gate_error = None + model_setup_reached = False + try: + _collective_validated_loader_order_hashes(local_train, heldout, report) + model_setup_reached = True + except RuntimeError as error: + gate_error = str(error) + outcomes: list[tuple[str | None, bool] | None] = [None] * dist.get_world_size() + dist.all_gather_object(outcomes, (gate_error, model_setup_reached)) + assert all( + outcome is not None + and outcome[0] is not None + and "loader-order authentication failed" in outcome[0] + and not outcome[1] + for outcome in outcomes + ) finally: dist.destroy_process_group() diff --git a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py index 8bb88a56d76..45d87206d53 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py +++ b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py @@ -10,9 +10,12 @@ import importlib.util import json from pathlib import Path +from types import SimpleNamespace import pytest +torch = pytest.importorskip("torch") + _REPO_ROOT = Path(__file__).resolve().parents[4] _HARNESS = _REPO_ROOT / "tests" / "gpu" / "torch" / "fastgen" / "pdd_qwen_operability_smoke.py" _SPEC = importlib.util.spec_from_file_location("pdd_qwen_operability_smoke", _HARNESS) @@ -111,6 +114,81 @@ def test_gpu_harness_uses_shared_unambiguous_ordered_id_hash() -> None: assert 'digest.update(b"\\n")' not in source +def test_gpu_harness_checkpoint_identity_uses_exact_shared_hashes() -> None: + from modelopt.torch.fastgen import PDDLayerSpec, PDDMetadata + + train_ids = smoke._training_sample_ids(2) + parameter = torch.nn.Parameter(torch.tensor(1.0)) + optimizer = torch.optim.AdamW([parameter], lr=2.0e-5) + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda _: 1.0) + setup = SimpleNamespace( + metadata=PDDMetadata( + grid_size=128, + grid_max_t=0.999, + flow_shift=5.0, + block_size_min=4, + block_size_max=64, + inference_blocks=(32, 32, 32, 32), + teacher_integrator="euler", + layer_spec=PDDLayerSpec( + projection_path="transformer.proj_out", + head_layout="channel_major", + ), + projection_in_features=3072, + projection_out_features=64, + projection_bias=True, + ), + automodel_snapshot={ + "distribution": "nemo_automodel", + "version": "0.5.0", + "package_tree_sha256": "a" * 64, + "wheel_sha256": "b" * 64, + "runtime_versions": {"diffusers": "0.38.0"}, + }, + optimizer=optimizer, + ) + config = SimpleNamespace( + model_id="Qwen/Qwen-Image", + model_revision="75e0b4be04f60ec59a75f475837eced720f823b6", + pdd=SimpleNamespace(guidance_scale=4.0), + guidance=SimpleNamespace(rescale=1.0, eps=1e-5), + training=SimpleNamespace( + seed=17, + validation_seed=29, + validation_every_steps=1000, + max_grad_norm=1.0, + zero_grad_warmup_steps=0, + ), + parallel=SimpleNamespace(activation_checkpointing=False), + ) + sampler = SimpleNamespace( + dataset=SimpleNamespace(metadata=[{"sample_id": sample_id} for sample_id in train_ids]) + ) + raw = {"pdd": {"grid_size": 128, "block_size_max": 64}} + + identity = smoke._identity( + setup=setup, + training=SimpleNamespace(scheduler=scheduler), + config=config, + sampler=sampler, + raw=raw, + ) + + assert identity["data"] == { + "ordered_train_id_sha256": ( + "4df732c492d043d5b0ea3549bcc80dbb847369021d0d3f242cd60386d1e94313" + ), + "ordered_heldout_id_sha256": ( + "38486bc077b6bd9b06a82167399e720f6c8dc70329dd0ce5fa23a92e4f30c198" + ), + "dataset_snapshot_sha256": smoke._canonical_sha256( + {"domain": "modelopt-pdd-synthetic-smoke-v1", "config": raw["pdd"]} + ), + "local_batch_size": 1, + "grad_accumulation_steps": 1, + } + + def _automodel_snapshot_fixture() -> tuple[dict, dict]: records = [] tree = hashlib.sha256() diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index f198f64ced2..b5ae655a6ae 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -140,6 +140,47 @@ def test_canonical_training_requires_external_hash_before_validator(tmp_path) -> pdd_finetune._validate_dataset_snapshot(raw, config) +def test_valid_canonical_data_gate_reaches_validator_unchanged(monkeypatch, tmp_path) -> None: + expected_hash = "a" * 64 + raw = _raw_config(tmp_path) + raw["data"] = { + "expected_approved_ordered_ids_sha256": expected_hash, + "expected_heldout_count": 2000, + "dataloader": { + "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", + "cache_dir": str(tmp_path), + }, + } + config = resolve_pdd_recipe_config(raw) + captured = None + + def _validator(root, **kwargs): + nonlocal captured + captured = (root, kwargs) + return {"snapshot_sha256": "b" * 64} + + def _all_gather_object(output, value): + output[0] = value + + monkeypatch.setattr("validate_cache_snapshot.validate_snapshot", _validator) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda payload, src: None) + + assert pdd_finetune._validate_dataset_snapshot(raw, config) == {"snapshot_sha256": "b" * 64} + assert captured == ( + tmp_path.resolve(), + { + "all_index": config.all_metadata_index, + "train_index": config.train_metadata_index, + "heldout_index": config.validation_metadata_index, + "expected_approved_ids_sha256": expected_hash, + "expected_heldout_count": 2000, + }, + ) + + def test_loader_order_must_match_authenticated_report() -> None: train = [{"sample_id": "train-a"}, {"sample_id": "train-b"}] heldout = [{"sample_id": "heldout-a"}] @@ -157,6 +198,66 @@ def test_loader_order_must_match_authenticated_report() -> None: pdd_finetune._validated_loader_order_hashes(list(reversed(train)), heldout, report) +def test_collective_loader_order_gate_rejects_rank_disagreement(monkeypatch) -> None: + train = [{"sample_id": "train-a"}, {"sample_id": "train-b"}] + heldout = [{"sample_id": "heldout-a"}] + report = { + "ordered_sample_ids_sha256": { + "train": pdd_finetune._ordered_id_sha256(train, split="train"), + "heldout": pdd_finetune._ordered_id_sha256(heldout, split="heldout"), + } + } + + def _all_gather_object(output, value): + output[:] = [value, {"ok": True, "hashes": ("c" * 64, value["hashes"][1])}] + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) + monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object) + with pytest.raises(RuntimeError, match="different authenticated"): + pdd_finetune._collective_validated_loader_order_hashes(train, heldout, report) + + +def test_two_rank_loader_order_divergence_fails_before_model_setup() -> None: + environment = os.environ.copy() + environment["PYTHONDONTWRITEBYTECODE"] = "1" + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc-per-node=2", + str( + _REPO_ROOT + / "tests" + / "examples" + / "diffusers" + / "fastgen" + / "pdd_training_preflight_distributed.py" + ), + ], + cwd=_REPO_ROOT, + env=environment, + check=True, + timeout=60, + ) + + +def test_pdd_finetune_namespace_module_help() -> None: + environment = os.environ.copy() + environment["PYTHONDONTWRITEBYTECODE"] = "1" + result = subprocess.run( + [sys.executable, "-m", "examples.diffusers.fastgen.pdd_finetune", "--help"], + cwd=_REPO_ROOT, + env=environment, + check=True, + capture_output=True, + text=True, + ) + assert "Train Qwen-Image" in result.stdout + + @pytest.mark.parametrize( ("scope", "name", "value", "message"), [ diff --git a/tests/examples/diffusers/fastgen/test_portable_cache.py b/tests/examples/diffusers/fastgen/test_portable_cache.py index 1ed97e14387..b632fbd90c2 100644 --- a/tests/examples/diffusers/fastgen/test_portable_cache.py +++ b/tests/examples/diffusers/fastgen/test_portable_cache.py @@ -219,6 +219,12 @@ def test_authenticated_cache_constants_hash_framing_and_seedless_split() -> None assert train == tuple(str(index) for index in range(16) if str(index) not in heldout) assert heldout not in (("0", "2", "4", "11"), ("2", "7", "8", "9")) + permuted = tuple(reversed(tuple(str(index) for index in range(16)))) + permuted_train, permuted_heldout = select_pdd_holdout_ids(permuted, 4) + assert set(permuted_heldout) == set(heldout) + assert permuted_train == tuple(item for item in permuted if item not in set(heldout)) + assert permuted_heldout == tuple(item for item in permuted if item in set(heldout)) + def test_approved_id_artifact_is_strict_and_externally_authenticatable(tmp_path) -> None: path = tmp_path / "approved.json" @@ -246,9 +252,10 @@ def test_approved_id_artifact_is_strict_and_externally_authenticatable(tmp_path) path.write_text('{"schema_version":1,"schema_version":1}') with pytest.raises(ValueError, match="duplicate key"): load_approved_sample_ids(path) - path.write_text('{"schema_version": NaN}') - with pytest.raises(ValueError, match="non-standard constant"): - load_approved_sample_ids(path) + for constant in ("NaN", "Infinity", "-Infinity"): + path.write_text(f'{{"schema_version": {constant}}}') + with pytest.raises(ValueError, match="non-standard constant"): + load_approved_sample_ids(path) path.write_bytes(b"\xff") with pytest.raises(ValueError, match="UTF-8 JSON"): load_approved_sample_ids(path) @@ -268,6 +275,21 @@ def test_approved_id_artifact_is_strict_and_externally_authenticatable(tmp_path) load_approved_sample_ids(real, expected_sha256="f" * 64) +def test_approved_hash_is_identical_to_finalized_all_index_hash(tmp_path) -> None: + root = tmp_path / "cache" + splits = _make_snapshot(root) + approved_path = tmp_path / "approved.json" + approved_digest = _write_approved_manifest(approved_path, splits["all"]) + _, loaded_digest = load_approved_sample_ids( + approved_path, + expected_sha256=approved_digest, + ) + all_index, _ = load_portable_metadata(root) + assert loaded_digest == all_index["ordered_sample_ids_sha256"] + assert loaded_digest == ordered_sample_ids_sha256(all_index["sample_ids"]) + assert loaded_digest == all_index["split_policy"]["approved_ordered_ids_sha256"] + + @pytest.mark.parametrize( "old_heldout", [("0", "2", "4", "11"), ("2", "7", "8", "9")], @@ -293,6 +315,47 @@ def test_self_consistent_all_list_rewrite_fails_external_hash(tmp_path) -> None: ) +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("order", "frozen PDD policy"), + ("gap", "split union"), + ("duplicate", "contains duplicates"), + ("member", "frozen PDD policy"), + ], +) +def test_split_order_gap_duplicate_and_membership_tampering_is_rejected( + tmp_path, mutation, message +) -> None: + root = tmp_path / "cache" + sample_ids = tuple(str(index) for index in range(16)) + _make_id_only_snapshot(root, sample_ids, 4) + train_path = root / "metadata_train.json" + heldout_path = root / "metadata_heldout.json" + train = json.loads(train_path.read_text()) + heldout = json.loads(heldout_path.read_text()) + + if mutation == "order": + train["sample_ids"] = list(reversed(train["sample_ids"])) + elif mutation == "gap": + train["sample_ids"].pop() + elif mutation == "duplicate": + train["sample_ids"].append(train["sample_ids"][0]) + else: + train["sample_ids"][0], heldout["sample_ids"][0] = ( + heldout["sample_ids"][0], + train["sample_ids"][0], + ) + + for path, index in ((train_path, train), (heldout_path, heldout)): + index["total_items"] = len(index["sample_ids"]) + if len(index["sample_ids"]) == len(set(index["sample_ids"])): + index["ordered_sample_ids_sha256"] = ordered_sample_ids_sha256(index["sample_ids"]) + _write_json(path, index) + with pytest.raises(ValueError, match=message): + validate_snapshot(root, reject_orphans=False) + + @pytest.mark.parametrize( ("policy_update", "message"), [ @@ -328,6 +391,26 @@ def test_strict_index_and_shard_json_reject_duplicates_and_nonfinite(tmp_path) - with pytest.raises(ValueError, match="non-standard constant"): load_portable_metadata(tmp_path / "second") + nested_policy = tmp_path / "nested-policy" + _make_snapshot(nested_policy) + nested_index_path = nested_policy / "metadata.json" + nested_index = nested_index_path.read_text() + domain = '"domain": "modelopt-pdd-holdout-v1"' + nested_index_path.write_text(nested_index.replace(domain, f"{domain},\n {domain}", 1)) + with pytest.raises(ValueError, match="duplicate key"): + load_portable_metadata(nested_policy) + + nested_shard = tmp_path / "nested-shard" + splits = _make_snapshot(nested_shard) + nested_shard_path = nested_shard / "metadata_shard_s0000.json" + nested_entries = nested_shard_path.read_text() + sample_id = f'"sample_id": "{splits["all"][0]}"' + nested_shard_path.write_text( + nested_entries.replace(sample_id, f"{sample_id},\n {sample_id}", 1) + ) + with pytest.raises(ValueError, match="duplicate key"): + load_portable_metadata(nested_shard) + @pytest.mark.parametrize("actual_heldout_count", [1999, 2001]) def test_validator_rejects_noncanonical_actual_holdout_counts( From 235e208c090dd4f06a674cc2edb6136dcfa20218 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 04:44:20 -0700 Subject: [PATCH 19/45] refactor(fastgen): organize PDD Qwen example Signed-off-by: Meng Xin --- CHANGELOG.rst | 7 +- examples/diffusers/fastgen/README.md | 2 +- .../diffusers/fastgen/analyze_pdd_results.py | 35 - examples/diffusers/fastgen/dmd2/README.md | 2 +- examples/diffusers/fastgen/dmd2/__init__.py | 15 + .../fastgen/dmd2/configs/qwen_image.yaml | 4 +- .../fastgen/fastgen_data/collate_fns.py | 16 +- .../diffusers/fastgen/fastgen_data/paths.py | 3 - .../fastgen_data/replayable_sampler.py | 32 +- .../diffusers/fastgen/fastgen_data/resume.py | 3 - .../diffusers/fastgen/fastgen_data/splits.py | 3 - .../fastgen_data/text_to_image_dataset.py | 44 +- .../fastgen/migrate_cache_manifest.py | 499 ------- examples/diffusers/fastgen/pdd/README.md | 56 + examples/diffusers/fastgen/pdd/__init__.py | 16 + .../{pdd_artifacts.py => pdd/artifacts.py} | 12 + .../{ => pdd}/automodel_dependency.json | 0 .../{pdd_checkpoint.py => pdd/checkpoint.py} | 23 +- .../fastgen/pdd/configs/qwen_image.yaml | 7 +- .../fastgen/{pdd_export.py => pdd/export.py} | 25 +- .../export_qwen_image.py} | 40 +- .../{pdd_finetune.py => pdd/finetune.py} | 234 ++-- .../inference_qwen_image.py} | 23 +- .../fastgen/{pdd_recipe.py => pdd/recipe.py} | 83 +- .../{pdd_training.py => pdd/training.py} | 16 +- .../{ => pdd}/verify_readonly_automodel.py | 14 +- examples/diffusers/fastgen/pdd_evaluation.py | 1103 ---------------- examples/diffusers/fastgen/portable_cache.py | 530 -------- .../diffusers/fastgen/preprocess/__init__.py | 2 +- .../fastgen/preprocess/processors/base.py | 3 +- .../preprocess/processors/qwen_image.py | 4 +- examples/diffusers/fastgen/requirements.txt | 2 +- .../fastgen/seal_pdd_run_manifest.py | 45 - .../fastgen/validate_cache_snapshot.py | 274 ---- .../fastgen/validate_pdd_run_manifest.py | 32 - modelopt/torch/fastgen/methods/pdd.py | 11 +- tests/examples/diffusers/fastgen/conftest.py | 3 - .../pdd_checkpoint_failure_distributed.py | 16 +- .../fastgen/pdd_export_distributed.py | 20 +- .../diffusers/fastgen/pdd_test_utils.py | 25 +- .../pdd_training_preflight_distributed.py | 193 --- .../pdd_validation_oracle_distributed.py | 14 +- .../diffusers/fastgen/test_dataset_paths.py | 3 - .../diffusers/fastgen/test_dataset_splits.py | 13 +- .../examples/diffusers/fastgen/test_layout.py | 45 +- .../fastgen/test_migrate_cache_manifest.py | 563 -------- .../diffusers/fastgen/test_pdd_evaluation.py | 529 -------- .../fastgen/test_pdd_inference_checkpoint.py | 20 +- .../test_pdd_qwen_operability_smoke.py | 493 ------- .../fastgen/test_pdd_recipe_setup.py | 206 +-- .../fastgen/test_pdd_training_lifecycle.py | 28 +- .../fastgen/test_pdd_validation_oracle.py | 16 +- .../diffusers/fastgen/test_portable_cache.py | 770 ----------- .../fastgen/test_vendored_migration.py | 3 +- tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py | 431 ------- .../fastgen/pdd_qwen_operability_smoke.py | 1146 ----------------- tests/gpu/torch/fastgen/test_pdd_toy.py | 12 + 57 files changed, 666 insertions(+), 7103 deletions(-) delete mode 100644 examples/diffusers/fastgen/analyze_pdd_results.py delete mode 100644 examples/diffusers/fastgen/migrate_cache_manifest.py create mode 100644 examples/diffusers/fastgen/pdd/README.md create mode 100644 examples/diffusers/fastgen/pdd/__init__.py rename examples/diffusers/fastgen/{pdd_artifacts.py => pdd/artifacts.py} (91%) rename examples/diffusers/fastgen/{ => pdd}/automodel_dependency.json (100%) rename examples/diffusers/fastgen/{pdd_checkpoint.py => pdd/checkpoint.py} (97%) rename examples/diffusers/fastgen/{pdd_export.py => pdd/export.py} (97%) rename examples/diffusers/fastgen/{export_pdd_qwen_image.py => pdd/export_qwen_image.py} (91%) rename examples/diffusers/fastgen/{pdd_finetune.py => pdd/finetune.py} (79%) rename examples/diffusers/fastgen/{inference_pdd_qwen_image.py => pdd/inference_qwen_image.py} (94%) rename examples/diffusers/fastgen/{pdd_recipe.py => pdd/recipe.py} (93%) rename examples/diffusers/fastgen/{pdd_training.py => pdd/training.py} (98%) rename examples/diffusers/fastgen/{ => pdd}/verify_readonly_automodel.py (92%) delete mode 100644 examples/diffusers/fastgen/pdd_evaluation.py delete mode 100644 examples/diffusers/fastgen/portable_cache.py delete mode 100644 examples/diffusers/fastgen/seal_pdd_run_manifest.py delete mode 100644 examples/diffusers/fastgen/validate_cache_snapshot.py delete mode 100644 examples/diffusers/fastgen/validate_pdd_run_manifest.py delete mode 100644 tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py delete mode 100644 tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py delete mode 100644 tests/examples/diffusers/fastgen/test_pdd_evaluation.py delete mode 100644 tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py delete mode 100644 tests/examples/diffusers/fastgen/test_portable_cache.py delete mode 100644 tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py delete mode 100644 tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ce5025a9e34..b78e74bc28a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,11 +6,6 @@ Changelog **Backward Breaking Changes** -- FastGen example portable caches now require authenticated schema-2 all/train/held-out indices. - Existing schema-1 or unversioned caches, including fresh preprocessing output, must be finalized - with ``examples/diffusers/fastgen/migrate_cache_manifest.py`` and an approved ordered-ID - artifact; they no longer load directly. This changes only the example cache protocol, not the - framework-neutral PDD checkpoint or API contract. - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. @@ -107,7 +102,7 @@ Changelog - Add offline DFlash speculative decoding training. Train the draft module from pre-computed base-model hidden states dumped by ``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py``; base-model transformer layers are deleted after conversion to save memory. Controlled by the auto-derived ``dflash_offline`` flag on ``DFlashConfig`` (derived from ``data_args.offline_data_path``). The dump scripts now share ``collect_hidden_states/common.py`` for aux-layer selection (``--aux-layers eagle|dflash|``) and optional assistant-token ``loss_mask`` for answer-only-loss training. - Add ``mtsa.config.SKIP_SOFTMAX_TRITON_CALIB`` for skip-softmax attention-sparsity calibration through the fused Triton ``attention_calibrate`` kernel (HF ``modelopt_triton`` backend), measuring multi-threshold tile-skip statistics the way the Triton inference kernel actually skips tiles for both prefill and decode. Exposed as ``--sparse_attn_cfg skip_softmax_triton_calib`` in ``examples/llm_sparsity/attention_sparsity/hf_sa.py`` (with a new ``--calib_data_dir`` flag for RULER calibration data). - Add DMD2 distillation for few-step diffusion models in ``examples/diffusers/fastgen/``: distill Qwen-Image into a 4/8-step student via Distribution Matching Distillation. See `examples/diffusers/fastgen/README.md `_ for details. -- Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with a Qwen-Image training, safe distributed-checkpoint export, PDD-2/4/8 inference, and paired effectiveness-evidence example. AutoModel remains an unmodified pinned runtime dependency. +- 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. - Make ``.agents/skills/`` the canonical location for agent skills; agent-specific directories (``.claude/skills/``, etc.) are now relative symlinks into ``.agents/``, so one skill suite serves multiple coding agents (Claude Code, Codex). See ``.agents/README.md``. - Extend Claude Code agent skills for PTQ, deployment, evaluation, monitoring, and baseline-vs-quantized result comparison. Adds evaluation task references for additional benchmarks, stronger PTQ checkpoint validation gates, and session-scoped workspace/job tracking. - Add SLURM Quality of Service (QoS) support to the ModelOpt launcher. Users can set QoS via ``slurm_config.qos`` or ``SLURM_QOS`` and the value is forwarded to ``nemo_run.SlurmExecutor``. diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md index ccbe3b9119e..d7bbf53f7d1 100644 --- a/examples/diffusers/fastgen/README.md +++ b/examples/diffusers/fastgen/README.md @@ -4,7 +4,7 @@ This directory contains training and inference examples for diffusion distillati `modelopt.torch.fastgen`. - [DMD2 for Qwen-Image](dmd2/README.md) -- PDD for Qwen-Image (integration in progress under `pdd/`) +- [PDD for Qwen-Image](pdd/README.md) 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/analyze_pdd_results.py b/examples/diffusers/fastgen/analyze_pdd_results.py deleted file mode 100644 index 56825e27443..00000000000 --- a/examples/diffusers/fastgen/analyze_pdd_results.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Write deterministic paired summaries from claim-bearing PDD effectiveness evidence.""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -sys.dont_write_bytecode = True - -_THIS_DIR = Path(__file__).resolve().parent -_REPO_ROOT = _THIS_DIR.parents[2] -for path in (_REPO_ROOT, _THIS_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("manifest", type=Path) - parser.add_argument("--output", type=Path, required=True) - args = parser.parse_args() - from pdd_artifacts import write_canonical_json - from pdd_evaluation import summarize_effectiveness_bundle, validate_effectiveness_bundle - - validated = validate_effectiveness_bundle(args.manifest) - write_canonical_json(args.output, summarize_effectiveness_bundle(validated)) - print(args.output) - - -if __name__ == "__main__": - main() diff --git a/examples/diffusers/fastgen/dmd2/README.md b/examples/diffusers/fastgen/dmd2/README.md index 4569f5d4c07..a67adc51d2f 100644 --- a/examples/diffusers/fastgen/dmd2/README.md +++ b/examples/diffusers/fastgen/dmd2/README.md @@ -13,7 +13,7 @@ output distribution. Built on `modelopt.torch.fastgen` and NeMo AutoModel's ## Requirements & self-contained data path -This example runs against **stock upstream `nemo_automodel`** (`>=0.4.0,<1.0`; see +This example runs against **stock upstream `nemo_automodel==0.5.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: diff --git a/examples/diffusers/fastgen/dmd2/__init__.py b/examples/diffusers/fastgen/dmd2/__init__.py index 1184907e87c..638206ec3d7 100644 --- a/examples/diffusers/fastgen/dmd2/__init__.py +++ b/examples/diffusers/fastgen/dmd2/__init__.py @@ -1,3 +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. + # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 diff --git a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml index 258c6d0f15d..d08ba727438 100644 --- a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml @@ -149,10 +149,8 @@ data: dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: /path/to/preprocessed/qwen_image_1024p - # cache_dir must be a finalized migrate_cache_manifest.py snapshot. The split index and - # negative embedding are portable references beneath the effective root. + # The environment override changes the effective root; referenced files must stay beneath it. # MODELOPT_FASTGEN_DATASET_CACHE_DIR overrides cache_dir when set to a non-empty value. - metadata_index: metadata_train.json base_resolution: [1024, 1024] batch_size: 1 drop_last: false diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index a943b32c32e..e47e6076beb 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -21,10 +21,10 @@ :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 + ``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 @@ -81,7 +81,7 @@ def collate_fn_text_to_image( # 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 + # released nemo_automodel 0.5.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]), @@ -96,6 +96,7 @@ def collate_fn_text_to_image( "original_resolution": torch.stack([item["original_resolution"] for item in batch]), "crop_offset": torch.stack([item["crop_offset"] for item in batch]), "sample_ids": torch.tensor([item["sample_id"] for item in batch], dtype=torch.long), + "logical_sample_ids": tuple(str(item["sample_id"]) for item in batch), }, } # Optional model-specific embedding fields, when a dataset provides them. @@ -176,6 +177,9 @@ def build_text_to_image_multiresolution_dataloader( prefetch_factor: int = 2, negative_prompt_embedding_path: str | None = None, selected_indices: Sequence[int] | None = None, + split: str | None = None, + validation_count: int | None = None, + split_seed: int = 2026, exact_resume: bool = False, sampler_seed: int = 42, loader_seed: int | None = None, @@ -199,6 +203,9 @@ def build_text_to_image_multiresolution_dataloader( negative_prompt_embedding_path: Optional ``.pt`` with a static negative-prompt embedding, bound into the collate and broadcast to every batch (DMD2 CFG). selected_indices: Optional ordered original metadata ordinals to expose. + split: Optional deterministic ``"train"`` or ``"validation"`` selection. + validation_count: Number of validation samples when ``split`` is set. + split_seed: Local seed used to construct deterministic split membership. exact_resume: Wrap the deterministic sampler with a committed cursor that is independent of worker prefetch. Required by the PDD lifecycle. sampler_seed: Seed for the released deterministic bucket sampler. @@ -212,6 +219,9 @@ def build_text_to_image_multiresolution_dataloader( cache_dir=cache_dir, train_text_encoder=train_text_encoder, selected_indices=selected_indices, + split=split, + validation_count=validation_count, + split_seed=split_seed, ) effective_root = dataset.cache_root diff --git a/examples/diffusers/fastgen/fastgen_data/paths.py b/examples/diffusers/fastgen/fastgen_data/paths.py index aed0ce3911a..ebcefb31b87 100644 --- a/examples/diffusers/fastgen/fastgen_data/paths.py +++ b/examples/diffusers/fastgen/fastgen_data/paths.py @@ -13,9 +13,6 @@ # 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 - """Path resolution for a portable, contained FastGen dataset cache.""" from __future__ import annotations diff --git a/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py b/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py index 46b3fee9cdb..720f608f0e4 100644 --- a/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py +++ b/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py @@ -1,5 +1,17 @@ # 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. """Committed-cursor wrapper for deterministic, prefetched batch samplers.""" @@ -89,13 +101,21 @@ def _batch_indices(self, batch_index: int) -> list[int]: return self._flat_indices[start:end].tolist() def _sample_ids(self, batch_index: int) -> tuple[str, ...]: - sample_ids: list[str] = [] + logical_ids = getattr(self.dataset, "logical_sample_ids", None) + if logical_ids is None: + dataset_sample_ids = getattr(self.dataset, "sample_ids", None) + if not isinstance(dataset_sample_ids, Sequence): + raise TypeError("sampler.dataset must expose logical_sample_ids or sample_ids.") + logical_ids = tuple(str(sample_id) for sample_id in dataset_sample_ids) + if not isinstance(logical_ids, Sequence): + raise TypeError("sampler.dataset.logical_sample_ids must be a sequence.") + batch_sample_ids: list[str] = [] for index in self._batch_indices(batch_index): - item = self.dataset.metadata[index] - if not isinstance(item, Mapping) or not isinstance(item.get("sample_id"), str): - raise ValueError(f"dataset.metadata[{index}] has no string sample_id.") - sample_ids.append(item["sample_id"]) - return tuple(sample_ids) + sample_id = logical_ids[index] + if not isinstance(sample_id, str) or not sample_id: + raise ValueError(f"dataset.logical_sample_ids[{index}] must be a nonempty string.") + batch_sample_ids.append(sample_id) + return tuple(batch_sample_ids) def expected_next_sample_ids(self) -> tuple[str, ...]: """Return the next committed batch's logical IDs without consuming it.""" diff --git a/examples/diffusers/fastgen/fastgen_data/resume.py b/examples/diffusers/fastgen/fastgen_data/resume.py index 3756639e245..44ce3137838 100644 --- a/examples/diffusers/fastgen/fastgen_data/resume.py +++ b/examples/diffusers/fastgen/fastgen_data/resume.py @@ -13,9 +13,6 @@ # 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 - """Public-API dataloader reconstruction for deterministic mid-epoch resume.""" from __future__ import annotations diff --git a/examples/diffusers/fastgen/fastgen_data/splits.py b/examples/diffusers/fastgen/fastgen_data/splits.py index e06790507c1..465f4ea6329 100644 --- a/examples/diffusers/fastgen/fastgen_data/splits.py +++ b/examples/diffusers/fastgen/fastgen_data/splits.py @@ -13,9 +13,6 @@ # 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 - """Deterministic train/validation membership for FastGen cache ordinals.""" from __future__ import annotations 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 98c4052b4a3..d6b41a508a8 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib import json from collections.abc import Sequence from pathlib import Path @@ -21,6 +22,7 @@ from nemo_automodel.components.datasets.diffusion.base_dataset import BaseMultiresolutionDataset from .paths import resolve_cache_root, resolve_under_root +from .splits import make_train_validation_indices __all__ = ["TextToImageDataset"] @@ -33,24 +35,41 @@ def __init__( cache_dir: str | Path, train_text_encoder: bool = False, selected_indices: Sequence[int] | None = None, + split: str | None = None, + validation_count: int | None = None, + split_seed: int = 2026, ): """ Args: cache_dir: Directory containing preprocessed cache train_text_encoder: If True, returns tokens instead of embeddings selected_indices: Optional ordered original metadata ordinals to expose. + split: Optional deterministic ``"train"`` or ``"validation"`` selection. + validation_count: Number of validation samples when ``split`` is set. + split_seed: Local seed used to construct deterministic split membership. """ + if selected_indices is not None and split is not None: + raise ValueError("selected_indices and split are mutually exclusive") + if split not in (None, "train", "validation"): + raise ValueError("split must be null, 'train', or 'validation'") + if split is not None and validation_count is None: + raise ValueError("validation_count is required when split is set") self.train_text_encoder = train_text_encoder self.cache_root = resolve_cache_root(cache_dir) self._selected_indices = selected_indices + self._split = split + self._validation_count = validation_count + self._split_seed = split_seed 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") - with metadata_file.open(encoding="utf-8") as file: - index = json.load(file) + digest = hashlib.sha256(b"modelopt-fastgen-metadata-v1\0") + index_bytes = metadata_file.read_bytes() + digest.update(index_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." @@ -63,8 +82,11 @@ def _load_metadata(self) -> list[dict]: shard_path = resolve_under_root( self.cache_root, shard_name, f"metadata shard {shard_index}" ) - with shard_path.open(encoding="utf-8") as file: - shard = json.load(file) + shard_bytes = shard_path.read_bytes() + digest.update(shard_name.encode()) + digest.update(b"\0") + digest.update(shard_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): @@ -82,7 +104,19 @@ def _load_metadata(self) -> list[dict]: if not complete_metadata: raise ValueError(f"No samples found in {metadata_file}") self.total_num_samples = len(complete_metadata) - self.sample_ids = self._validate_selected_indices(self.total_num_samples) + self.metadata_sha256 = digest.hexdigest() + if self._split is None: + self.sample_ids = self._validate_selected_indices(self.total_num_samples) + else: + if self._validation_count is None: + raise RuntimeError("validation_count was not resolved for the requested split") + train, validation = make_train_validation_indices( + self.total_num_samples, + self._validation_count, + self._split_seed, + ) + self.sample_ids = train if self._split == "train" else validation + self.logical_sample_ids = [str(sample_id) for sample_id in self.sample_ids] return [complete_metadata[index] for index in self.sample_ids] def _validate_selected_indices(self, num_samples: int) -> list[int]: diff --git a/examples/diffusers/fastgen/migrate_cache_manifest.py b/examples/diffusers/fastgen/migrate_cache_manifest.py deleted file mode 100644 index 1e689bb3f6d..00000000000 --- a/examples/diffusers/fastgen/migrate_cache_manifest.py +++ /dev/null @@ -1,499 +0,0 @@ -# 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. - -"""Migrate a legacy absolute-path FastGen cache into an immutable portable snapshot.""" - -from __future__ import annotations - -import argparse -import json -import os -import shutil -import tempfile -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING, Any - -import torch -from portable_cache import ( - PDD_HOLDOUT_DOMAIN, - PORTABLE_SNAPSHOT_SCHEMA_VERSION, - PREPROCESS_STAGING_SCHEMA_VERSION, - SPLIT_POLICY_SCHEMA_VERSION, - audit_no_absolute_paths, - load_approved_sample_ids, - load_strict_json, - ordered_sample_ids_sha256, - resolve_cache_asset, - select_pdd_holdout_ids, - sha256_file, - stable_sample_id, - validate_relative_reference, -) -from validate_cache_snapshot import validate_snapshot - -if TYPE_CHECKING: - from collections.abc import Sequence - -_REMOVED_PATH_KEYS = { - "cache_dir", - "cache_file", - "image_path", - "output_dir", - "source_dir", - "source_path", - "video_path", -} - - -@dataclass(frozen=True) -class MigrationRecord: - """Frozen mapping produced by read-only pass 1 and consumed by pass 2.""" - - sample_id: str - source_ref: str - source_payload: Path - source_sha256: str - destination_ref: str - manifest_fields: dict[str, Any] - - -def _load_json(path: Path, expected_type: type, label: str): - value = load_strict_json(path, label=label) - if not isinstance(value, expected_type): - raise ValueError(f"{label} must contain {expected_type.__name__}") - return value - - -def _load_source_json( - root: Path, - reference: str, - *, - expected_type: type, - label: str, -): - relative = validate_relative_reference(reference, label=label) - candidate = root / relative - if candidate.is_symlink(): - raise ValueError(f"{label} must not be a symlink: {relative}") - path = resolve_cache_asset(root, reference, label=label) - return _load_json(path, expected_type, label) - - -def _relative_to_legacy_prefix(raw: str, prefix: Path, *, label: str) -> Path: - path = Path(raw).expanduser() - if not path.is_absolute(): - return validate_relative_reference(raw, label=label) - try: - relative = path.relative_to(prefix) - except ValueError as error: - raise ValueError(f"{label} is outside the declared legacy prefix: {raw}") from error - return validate_relative_reference(relative.as_posix(), label=label) - - -def _resolve_legacy_payload( - source_root: Path, - cache_file: Any, - legacy_cache_root: Path, - *, - label: str, -) -> Path: - if not isinstance(cache_file, str): - raise TypeError(f"{label} must be a string") - relative = _relative_to_legacy_prefix(cache_file, legacy_cache_root, label=label) - return resolve_cache_asset(source_root, relative.as_posix(), label=label) - - -def _legacy_source_ref( - entry: dict[str, Any], legacy_source_root: Path | None, *, label: str -) -> str: - if "source_ref" in entry: - return validate_relative_reference( - entry["source_ref"], label=f"{label}.source_ref" - ).as_posix() - image_path = entry.get("image_path") - if not isinstance(image_path, str): - raise ValueError(f"{label} needs source_ref or legacy image_path") - if Path(image_path).is_absolute() and legacy_source_root is None: - raise ValueError("--legacy-source-root is required for absolute legacy image_path values") - prefix = legacy_source_root or Path(".") - return _relative_to_legacy_prefix(image_path, prefix, label=f"{label}.image_path").as_posix() - - -def _sanitize_payload(value: Any) -> Any: - if isinstance(value, dict): - sanitized = {} - for key, nested in value.items(): - if not isinstance(key, str): - raise ValueError(f"payload contains non-string key {key!r}") - if key.lower() not in _REMOVED_PATH_KEYS: - sanitized[key] = _sanitize_payload(nested) - return sanitized - if isinstance(value, list): - return [_sanitize_payload(item) for item in value] - if isinstance(value, tuple): - return tuple(_sanitize_payload(item) for item in value) - if isinstance(value, set): - return {_sanitize_payload(item) for item in value} - if isinstance(value, frozenset): - return frozenset(_sanitize_payload(item) for item in value) - return value - - -def _portable_manifest_fields(entry: dict[str, Any], *, label: str) -> dict[str, Any]: - required = ("bucket_resolution", "original_resolution", "prompt", "bucket_id", "aspect_ratio") - missing = [name for name in required if name not in entry] - if missing: - raise ValueError(f"{label} is missing required fields: {missing}") - fields = {name: entry[name] for name in required} - for name in ("pixels", "model_type", "crop_resolution"): - if name in entry: - fields[name] = entry[name] - audit_no_absolute_paths(fields, context=label) - return fields - - -def plan_migration( - source_root: str | Path, - *, - source_index: str | Sequence[str] = "metadata.json", - legacy_cache_root: str | Path | None = None, - legacy_source_root: str | Path | None = None, -) -> tuple[MigrationRecord, ...]: - """Perform read-only pass 1 and return a deterministic frozen mapping.""" - root = Path(source_root).expanduser().resolve(strict=True) - if not root.is_dir(): - raise NotADirectoryError(f"source_root is not a directory: {root}") - cache_prefix = Path(legacy_cache_root).expanduser() if legacy_cache_root else root - source_prefix = Path(legacy_source_root).expanduser() if legacy_source_root else None - if not cache_prefix.is_absolute(): - raise ValueError("legacy_cache_root must be absolute") - if source_prefix is not None and not source_prefix.is_absolute(): - raise ValueError("legacy_source_root must be absolute") - - source_indexes = (source_index,) if isinstance(source_index, str) else tuple(source_index) - if not source_indexes or any(not isinstance(name, str) or not name for name in source_indexes): - raise ValueError("source_index must contain one or more index paths") - if len(source_indexes) != len(set(source_indexes)): - raise ValueError("source_index contains duplicates") - - parsed_indexes = [] - all_shards: set[str] = set() - for index_number, index_ref in enumerate(source_indexes): - label = f"source_index[{index_number}]" - index = _load_source_json(root, index_ref, expected_type=dict, label=label) - if "schema_version" in index and ( - type(index["schema_version"]) is not int - or index["schema_version"] != PREPROCESS_STAGING_SCHEMA_VERSION - ): - raise ValueError(f"{label}.schema_version is unsupported") - shards = index.get("shards") - if ( - not isinstance(shards, list) - or not shards - or any(not isinstance(item, str) for item in shards) - ): - raise ValueError(f"{label}.shards must be a non-empty list of relative paths") - if len(shards) != len(set(shards)): - raise ValueError(f"{label}.shards contains duplicates") - if "num_shards" in index and ( - type(index["num_shards"]) is not int or index["num_shards"] != len(shards) - ): - raise ValueError(f"{label}.num_shards does not match len(shards)") - overlap = all_shards.intersection(shards) - if overlap: - raise ValueError(f"source indices share metadata shards: {sorted(overlap)}") - all_shards.update(shards) - parsed_indexes.append((label, index, shards)) - - ranked_indexes = [ - index for _, index, _ in parsed_indexes if "shard_world" in index or "shard_rank" in index - ] - if ranked_indexes: - if len(ranked_indexes) != len(parsed_indexes): - raise ValueError("all source indices must declare shard_rank and shard_world") - if any( - type(index.get(field)) is not int - for index in ranked_indexes - for field in ("shard_rank", "shard_world") - ): - raise ValueError("source shard_rank and shard_world must be integers") - worlds = {index.get("shard_world") for index in ranked_indexes} - ranks = {index.get("shard_rank") for index in ranked_indexes} - if worlds != {len(parsed_indexes)} or ranks != set(range(len(parsed_indexes))): - raise ValueError("source rank indices are incomplete or inconsistent") - - records: list[MigrationRecord] = [] - seen_ids: set[str] = set() - for index_label, index, shards in parsed_indexes: - index_record_count = 0 - source_sample_ids = [] - for shard_number, shard_ref in enumerate(shards): - shard_label = f"{index_label}.shards[{shard_number}]" - entries = _load_source_json( - root, - shard_ref, - expected_type=list, - label=shard_label, - ) - index_record_count += len(entries) - for entry_number, entry in enumerate(entries): - label = f"{shard_label}[{entry_number}]" - if not isinstance(entry, dict): - raise ValueError(f"{label} must be an object") - source_sample_ids.append(entry.get("sample_id")) - source_ref = _legacy_source_ref(entry, source_prefix, label=label) - source_payload = _resolve_legacy_payload( - root, - entry.get("cache_file"), - cache_prefix, - label=f"{label}.cache_file", - ) - resolution = entry.get("bucket_resolution", entry.get("crop_resolution")) - model_type = entry.get("model_type") - if not isinstance(model_type, str) or not model_type: - raise ValueError(f"{label}.model_type must be a non-empty string") - sample_id = stable_sample_id( - source_ref=source_ref, - resolution=resolution, - model_type=model_type, - ) - if sample_id in seen_ids: - raise ValueError(f"duplicate migrated sample_id: {sample_id}") - seen_ids.add(sample_id) - - source_digest = sha256_file(source_payload) - payload = torch.load(source_payload, map_location="cpu", weights_only=True) - if not isinstance(payload, dict): - raise TypeError(f"{label} payload must be a dict") - sanitized = _sanitize_payload(payload) - sanitized["sample_id"] = sample_id - sanitized["source_ref"] = source_ref - audit_no_absolute_paths(sanitized, context=f"payload[{sample_id}]") - - records.append( - MigrationRecord( - sample_id=sample_id, - source_ref=source_ref, - source_payload=source_payload, - source_sha256=source_digest, - destination_ref=f"payloads/{sample_id}.pt", - manifest_fields=_portable_manifest_fields(entry, label=label), - ) - ) - if "total_items" in index and ( - type(index["total_items"]) is not int or index["total_items"] != index_record_count - ): - raise ValueError(f"{index_label}.total_items does not match its loaded entry count") - if "sample_ids" in index and index["sample_ids"] != source_sample_ids: - raise ValueError(f"{index_label}.sample_ids does not match its loaded entries") - - if not records: - raise ValueError("legacy source index contains no entries") - return tuple(sorted(records, key=lambda record: record.sample_id)) - - -def _write_json(path: Path, value: Any) -> None: - with path.open("w", encoding="utf-8") as stream: - json.dump(value, stream, indent=2, sort_keys=True, allow_nan=False) - stream.write("\n") - stream.flush() - os.fsync(stream.fileno()) - - -def migrate_cache( - source_root: str | Path, - output_root: str | Path, - *, - approved_ids_manifest: str | Path, - heldout_count: int, - expected_approved_ids_sha256: str | None = None, - source_index: str | Sequence[str] = "metadata.json", - legacy_cache_root: str | Path | None = None, - legacy_source_root: str | Path | None = None, - negative_embedding: str | None = None, - shard_size: int = 10000, -) -> dict[str, Any]: - """Run two-pass migration and atomically publish the validated destination.""" - destination = Path(output_root).expanduser() - if destination.exists(): - raise FileExistsError(f"output_root already exists: {destination}") - if not destination.parent.exists(): - raise FileNotFoundError(f"output_root parent does not exist: {destination.parent}") - if type(shard_size) is not int or shard_size <= 0: - raise ValueError("shard_size must be a positive integer") - - # Pass 1 is intentionally complete before any destination or staging path is created. - records = plan_migration( - source_root, - source_index=source_index, - legacy_cache_root=legacy_cache_root, - legacy_source_root=legacy_source_root, - ) - approved_ids, approved_digest = load_approved_sample_ids( - Path(approved_ids_manifest).expanduser(), - expected_sha256=expected_approved_ids_sha256, - ) - records_by_id = {record.sample_id: record for record in records} - unknown_ids = [sample_id for sample_id in approved_ids if sample_id not in records_by_id] - if unknown_ids: - raise ValueError(f"approved_ids_manifest references unknown sample IDs: {unknown_ids[:5]}") - selected_records = tuple(records_by_id[sample_id] for sample_id in approved_ids) - train_ids, heldout_ids = select_pdd_holdout_ids(approved_ids, heldout_count) - - source = Path(source_root).expanduser().resolve(strict=True) - negative_source = None - negative_source_sha256 = None - if negative_embedding is not None: - negative_source = _resolve_legacy_payload( - source, - negative_embedding, - Path(legacy_cache_root).expanduser() if legacy_cache_root else source, - label="negative_embedding", - ) - negative_source_sha256 = sha256_file(negative_source) - negative_payload = torch.load(negative_source, map_location="cpu", weights_only=True) - audit_no_absolute_paths(negative_payload, context="negative_prompt_embedding") - - staging = Path(tempfile.mkdtemp(prefix=f".{destination.name}.staging-", dir=destination.parent)) - try: - (staging / "payloads").mkdir() - portable_entries = [] - for record in selected_records: - if sha256_file(record.source_payload) != record.source_sha256: - raise RuntimeError(f"source payload changed after pass 1: {record.source_payload}") - payload = torch.load(record.source_payload, map_location="cpu", weights_only=True) - sanitized = _sanitize_payload(payload) - sanitized["sample_id"] = record.sample_id - sanitized["source_ref"] = record.source_ref - audit_no_absolute_paths(sanitized, context=f"payload[{record.sample_id}]") - - destination_path = staging / record.destination_ref - torch.save(sanitized, destination_path) - portable_entries.append( - { - "sample_id": record.sample_id, - "source_ref": record.source_ref, - "cache_file": record.destination_ref, - "payload_sha256": sha256_file(destination_path), - **record.manifest_fields, - } - ) - - shard_names = [] - for offset in range(0, len(portable_entries), shard_size): - name = f"metadata_shard_s{offset // shard_size:04d}.json" - _write_json(staging / name, portable_entries[offset : offset + shard_size]) - shard_names.append(name) - - negative_declaration = None - if negative_source is not None: - if sha256_file(negative_source) != negative_source_sha256: - raise RuntimeError("negative prompt embedding changed after pass 1") - negative_name = "negative_prompt_embedding.pt" - shutil.copyfile(negative_source, staging / negative_name) - negative_declaration = { - "path": negative_name, - "sha256": sha256_file(staging / negative_name), - } - - common = { - "schema_version": PORTABLE_SNAPSHOT_SCHEMA_VERSION, - "shards": shard_names, - "num_shards": len(shard_names), - "split_policy": { - "schema_version": SPLIT_POLICY_SCHEMA_VERSION, - "algorithm": "sha256-domain-ranked", - "domain": PDD_HOLDOUT_DOMAIN, - "heldout_count": heldout_count, - "approved_ordered_ids_sha256": approved_digest, - }, - } - if negative_declaration is not None: - common["negative_prompt_embedding"] = negative_declaration - split_specs = { - "metadata.json": ("all", approved_ids), - "metadata_train.json": ("train", train_ids), - "metadata_heldout.json": ("heldout", heldout_ids), - } - for name, (split, sample_ids) in split_specs.items(): - index = { - **common, - "split": split, - "sample_ids": list(sample_ids), - "ordered_sample_ids_sha256": ordered_sample_ids_sha256(sample_ids), - "total_items": len(sample_ids), - } - audit_no_absolute_paths(index, context=name) - _write_json(staging / name, index) - - validation = validate_snapshot( - staging, - expected_approved_ids_sha256=expected_approved_ids_sha256, - expected_heldout_count=heldout_count, - ) - os.replace(staging, destination) - except BaseException: - shutil.rmtree(staging, ignore_errors=True) - raise - - return { - "schema_version": 1, - "record_type": "modelopt_fastgen_cache_migration", - "output_root": str(destination.resolve()), - "counts": { - "source": len(records), - "approved": len(selected_records), - "filtered": len(records) - len(selected_records), - }, - "validation": validation, - } - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--source-root", required=True) - parser.add_argument("--output-root", required=True) - parser.add_argument( - "--source-index", - action="append", - dest="source_indexes", - help="Relative source index; repeat to finalize all rank-local preprocessing indices.", - ) - parser.add_argument("--legacy-cache-root") - parser.add_argument("--legacy-source-root") - parser.add_argument("--negative-embedding") - parser.add_argument("--approved-ids-manifest", required=True) - parser.add_argument("--expected-approved-ids-sha256") - parser.add_argument("--heldout-count", required=True, type=int) - parser.add_argument("--shard-size", default=10000, type=int) - args = parser.parse_args() - report = migrate_cache( - args.source_root, - args.output_root, - approved_ids_manifest=args.approved_ids_manifest, - heldout_count=args.heldout_count, - expected_approved_ids_sha256=args.expected_approved_ids_sha256, - source_index=tuple(args.source_indexes) if args.source_indexes else "metadata.json", - legacy_cache_root=args.legacy_cache_root, - legacy_source_root=args.legacy_source_root, - negative_embedding=args.negative_embedding, - shard_size=args.shard_size, - ) - print(json.dumps(report, sort_keys=True, allow_nan=False)) - - -if __name__ == "__main__": - main() diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md new file mode 100644 index 00000000000..f9d780bb7b8 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/README.md @@ -0,0 +1,56 @@ +# PDD for Qwen-Image + +Parallel Decoding Distillation (PDD) trains one Qwen-Image student call to predict several +consecutive rectified-flow updates. The student keeps the original transformer backbone and widens +only its output projection to 128 velocity heads. During training it samples different aligned +block lengths up to 64, so the same checkpoint can use different supported block schedules at +inference. + +The provided schedules use the 128-interval grid as follows: + +| Schedule | Block sizes | Transformer calls | +|---|---|---:| +| `pdd-2` | `[64, 64]` | 2 | +| `pdd-4` | `[32, 32, 32, 32]` | 4 | +| `pdd-8` | `[16, 16, 16, 16, 16, 16, 16, 16]` | 8 | + +## Training + +Install the shared requirements from the repository root, then launch with released AutoModel +APIs. No AutoModel, Diffusers, or Qwen source changes are required. + +```bash +pip install -r examples/diffusers/fastgen/requirements.txt +export MODELOPT_FASTGEN_DATASET_CACHE_DIR=/absolute/path/to/qwen_image_cache +torchrun --standalone --nproc-per-node=8 \ + examples/diffusers/fastgen/pdd/finetune.py \ + --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +``` + +The cache must contain `metadata.json`, its declared shards, cached tensors, and +`negative_prompt_embedding.pt`. The environment variable overrides the configured cache root. All +metadata, tensor, and negative-embedding paths must still resolve inside that effective root. + +Training deterministically derives disjoint train and validation membership from metadata ordinals; +it does not rewrite the cache or require separate split manifests. The default recipe uses 2,000 +validation samples, learning rate `2e-5`, 128 heads, and sampled block lengths from 4 through 64. +Checkpoints include the student, optimizer, scheduler, RNG, trainer, and exact replayable sampler +state needed to resume the next committed batch. + +Start with a one-node smoke and scale only after it passes; project training runs are capped at 16 +nodes. + +## Export and inference + +```bash +torchrun --standalone --nproc-per-node=8 \ + examples/diffusers/fastgen/pdd/export_qwen_image.py \ + --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml \ + --checkpoint LATEST --output-dir /path/to/pdd-export + +python examples/diffusers/fastgen/pdd/inference_qwen_image.py \ + --export-dir /path/to/pdd-export --schedule pdd-4 \ + --prompt-id red-cube-0001 --prompt "a small red cube on a white table" \ + --seed 42 --height 1024 --width 1024 \ + --output /path/to/pdd4.png --result-json /path/to/pdd4.json +``` diff --git a/examples/diffusers/fastgen/pdd/__init__.py b/examples/diffusers/fastgen/pdd/__init__.py new file mode 100644 index 00000000000..3d2b6c232ef --- /dev/null +++ b/examples/diffusers/fastgen/pdd/__init__.py @@ -0,0 +1,16 @@ +# 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.""" diff --git a/examples/diffusers/fastgen/pdd_artifacts.py b/examples/diffusers/fastgen/pdd/artifacts.py similarity index 91% rename from examples/diffusers/fastgen/pdd_artifacts.py rename to examples/diffusers/fastgen/pdd/artifacts.py index 54408de8194..709b6a03ad9 100644 --- a/examples/diffusers/fastgen/pdd_artifacts.py +++ b/examples/diffusers/fastgen/pdd/artifacts.py @@ -1,5 +1,17 @@ # 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. """Strict canonical-JSON and relative-artifact helpers for the PDD example.""" diff --git a/examples/diffusers/fastgen/automodel_dependency.json b/examples/diffusers/fastgen/pdd/automodel_dependency.json similarity index 100% rename from examples/diffusers/fastgen/automodel_dependency.json rename to examples/diffusers/fastgen/pdd/automodel_dependency.json diff --git a/examples/diffusers/fastgen/pdd_checkpoint.py b/examples/diffusers/fastgen/pdd/checkpoint.py similarity index 97% rename from examples/diffusers/fastgen/pdd_checkpoint.py rename to examples/diffusers/fastgen/pdd/checkpoint.py index 8c0ae9a6be6..66b45e4cfe7 100644 --- a/examples/diffusers/fastgen/pdd_checkpoint.py +++ b/examples/diffusers/fastgen/pdd/checkpoint.py @@ -1,5 +1,17 @@ # 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. """Atomic, strict PDD checkpoint publication around the stock AutoModel Checkpointer.""" @@ -643,7 +655,7 @@ def resolve(self, restore_from: str | Path | None) -> Path | None: def _collective_resolve(self, restore_from: str | Path | None) -> Path | None: if _world_size() == 1: return self.resolve(restore_from) - status = None + status: dict[str, Any] | None = None if _rank() == 0: try: resolved = self.resolve(restore_from) @@ -661,7 +673,12 @@ def _collective_resolve(self, restore_from: str | Path | None) -> Path | None: raise RuntimeError("rank 0 broadcast a malformed checkpoint resolution status.") if not status["ok"]: raise RuntimeError(f"rank-0 checkpoint resolution failed: {status.get('error')}.") - return None if status["path"] is None else Path(status["path"]) + resolved_path = status.get("path") + if resolved_path is None: + return None + if not isinstance(resolved_path, str): + raise RuntimeError("rank 0 broadcast a malformed checkpoint path.") + return Path(resolved_path) def _prepare_staging(self, final: Path) -> str: self.root.mkdir(parents=True, exist_ok=True) @@ -775,7 +792,7 @@ def save(self) -> Path: raise RuntimeError("PDD checkpoint sidecar save failed; " + "; ".join(sidecar_failures)) _barrier() - publish_status = None + publish_status: dict[str, Any] | None = None if _rank() == 0: try: self._publish_staging( diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 88635470292..1e8c26cb11a 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -59,14 +59,11 @@ fsdp: activation_checkpointing: true data: - all_metadata_index: metadata.json - validation_metadata_index: metadata_heldout.json - expected_approved_ordered_ids_sha256: - expected_heldout_count: 2000 + validation_count: 2000 + split_seed: 2026 dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: data/qwen_image_cache - metadata_index: metadata_train.json base_resolution: [1024, 1024] batch_size: 1 drop_last: true diff --git a/examples/diffusers/fastgen/pdd_export.py b/examples/diffusers/fastgen/pdd/export.py similarity index 97% rename from examples/diffusers/fastgen/pdd_export.py rename to examples/diffusers/fastgen/pdd/export.py index 6ad782f4db5..690cc8a9020 100644 --- a/examples/diffusers/fastgen/pdd_export.py +++ b/examples/diffusers/fastgen/pdd/export.py @@ -1,5 +1,17 @@ # 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. """Authenticated, bounded safetensors export and strict PDD reconstruction helpers.""" @@ -15,18 +27,19 @@ from typing import Any import torch -from pdd_artifacts import ( +from safetensors import safe_open +from safetensors.torch import save_file + +from modelopt.torch.fastgen import PDDConfig, PDDMetadata +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC + +from .artifacts import ( load_canonical_json, require_sha256, resolve_relative_artifact, sha256_file, write_canonical_json, ) -from safetensors import safe_open -from safetensors.torch import save_file - -from modelopt.torch.fastgen import PDDConfig, PDDMetadata -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC _EXPORT_SCHEMA_VERSION = 1 _COMPLETE_SCHEMA_VERSION = 1 diff --git a/examples/diffusers/fastgen/export_pdd_qwen_image.py b/examples/diffusers/fastgen/pdd/export_qwen_image.py similarity index 91% rename from examples/diffusers/fastgen/export_pdd_qwen_image.py rename to examples/diffusers/fastgen/pdd/export_qwen_image.py index fc28e428fb6..8c1fcdf2f0b 100644 --- a/examples/diffusers/fastgen/export_pdd_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/export_qwen_image.py @@ -1,5 +1,17 @@ # 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. """Collectively restore a PDD checkpoint and publish a safe Qwen-Image export.""" @@ -21,8 +33,9 @@ sys.dont_write_bytecode = True _THIS_DIR = Path(__file__).resolve().parent -_REPO_ROOT = _THIS_DIR.parents[2] -for path in (_REPO_ROOT, _THIS_DIR): +_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)) @@ -32,7 +45,7 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--config", type=Path, - default=_THIS_DIR / "configs" / "pdd_qwen_image.yaml", + default=_THIS_DIR / "configs" / "qwen_image.yaml", ) parser.add_argument( "--checkpoint", @@ -167,7 +180,10 @@ def _collective_publication_preflight(output_dir: Path) -> Mapping[str, Any]: raise RuntimeError("rank 0 broadcast malformed PDD publication preflight status.") if not status["ok"]: raise RuntimeError(f"PDD publication preflight failed: {status.get('error')}.") - return status["modelopt_source"] + modelopt_source = status.get("modelopt_source") + if not isinstance(modelopt_source, Mapping): + raise RuntimeError("rank 0 broadcast malformed ModelOpt source identity.") + return modelopt_source def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: @@ -176,7 +192,10 @@ def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, identity = manifest.get("identity") if not isinstance(identity, Mapping): raise RuntimeError("PDD checkpoint has no identity mapping.") - if PDDMetadata.from_dict(identity.get("pdd_metadata")) != setup.metadata: + pdd_metadata = identity.get("pdd_metadata") + if not isinstance(pdd_metadata, Mapping): + raise RuntimeError("PDD checkpoint has no PDD metadata mapping.") + if PDDMetadata.from_dict(pdd_metadata) != setup.metadata: raise RuntimeError("PDD checkpoint metadata does not match the configured student.") if identity.get("model") != { "id": config.model_id, @@ -244,7 +263,7 @@ def _checkpoint_selector_identity(config: Any, setup: Any) -> dict[str, Any]: def _collective_checkpoint_resolution( config: Any, setup: Any, restore_from: str ) -> tuple[Path, Mapping[str, Any]]: - from pdd_checkpoint import resolve_pdd_training_checkpoint + from pdd.checkpoint import resolve_pdd_training_checkpoint status = None if dist.get_rank() == 0: @@ -270,14 +289,15 @@ def _collective_checkpoint_resolution( def main() -> None: args = _parse_args() - from pdd_artifacts import sha256_file - from pdd_export import write_pdd_export - from pdd_recipe import ( + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + + from pdd.artifacts import sha256_file + from pdd.export import write_pdd_export + from pdd.recipe import ( build_pdd_export_setup, initialize_pdd_distributed, resolve_pdd_recipe_config, ) - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict raw = yaml.safe_load(args.config.read_text()) config = resolve_pdd_recipe_config(raw) diff --git a/examples/diffusers/fastgen/pdd_finetune.py b/examples/diffusers/fastgen/pdd/finetune.py similarity index 79% rename from examples/diffusers/fastgen/pdd_finetune.py rename to examples/diffusers/fastgen/pdd/finetune.py index 17a284b10ec..c071d746c30 100644 --- a/examples/diffusers/fastgen/pdd_finetune.py +++ b/examples/diffusers/fastgen/pdd/finetune.py @@ -1,5 +1,17 @@ # 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. """Train Qwen-Image with the ModelOpt-owned PDD lifecycle and released AutoModel APIs.""" @@ -7,6 +19,8 @@ import argparse import dataclasses +import hashlib +import json import logging import sys import time @@ -19,44 +33,32 @@ sys.dont_write_bytecode = True _THIS_DIR = Path(__file__).resolve().parent -_REPO_ROOT = _THIS_DIR.parents[2] +_FASTGEN_DIR = _THIS_DIR.parent +_REPO_ROOT = _FASTGEN_DIR.parents[2] # These entrypoints are also supported through ``python -m``. In that mode the # sibling ModelOpt-owned example modules are not importable until this directory # is added explicitly. -for path in (_REPO_ROOT, _THIS_DIR): +for path in (_REPO_ROOT, _FASTGEN_DIR): if str(path) not in sys.path: sys.path.insert(0, str(path)) -from portable_cache import ordered_sample_ids_sha256 # noqa: E402 - -_CANONICAL_PDD_HELDOUT_COUNT = 2000 - def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--config", type=Path, - default=_THIS_DIR / "configs" / "pdd_qwen_image.yaml", + default=_THIS_DIR / "configs" / "qwen_image.yaml", ) return parser.parse_args() -def _metadata_sample_ids(metadata: Any, *, split: str) -> tuple[str, ...]: - if not isinstance(metadata, list): - raise TypeError(f"PDD {split} dataset metadata must be a list.") - sample_ids: list[str] = [] - for index, item in enumerate(metadata): - if not isinstance(item, Mapping) or not isinstance(item.get("sample_id"), str): - raise ValueError(f"{split} metadata[{index}] has no string sample_id.") - sample_ids.append(item["sample_id"]) - if len(set(sample_ids)) != len(sample_ids): - raise ValueError(f"PDD {split} metadata contains duplicate logical sample IDs.") - return tuple(sample_ids) - - -def _ordered_id_sha256(metadata: Any, *, split: str) -> str: - return ordered_sample_ids_sha256(_metadata_sample_ids(metadata, split=split)) +def _ordered_id_sha256(sample_ids: tuple[str, ...]) -> str: + digest = hashlib.sha256(b"modelopt-pdd-ordered-sample-ids-v1\0") + for sample_id in sample_ids: + digest.update(sample_id.encode()) + digest.update(b"\n") + return digest.hexdigest() def _dataloader_options(raw: Mapping[str, Any]) -> dict[str, Any]: @@ -89,6 +91,9 @@ def _build_training_dataloader( if options.get("dynamic_batch_size", False) is not False: raise ValueError("PDD v1 requires data.dataloader.dynamic_batch_size=false.") options.update( + split="train", + validation_count=config.validation_count, + split_seed=config.split_seed, dp_rank=dp_rank, dp_world_size=dp_world_size, exact_resume=True, @@ -114,7 +119,9 @@ def _build_validation_dataloader( options = _dataloader_options(raw) options.update( - metadata_index=config.validation_metadata_index, + split="validation", + validation_count=config.validation_count, + split_seed=config.split_seed, dp_rank=dp_rank, dp_world_size=dp_world_size, drop_last=False, @@ -127,93 +134,52 @@ def _build_validation_dataloader( return build_text_to_image_multiresolution_dataloader(**options) -def _validate_dataset_snapshot(raw: Mapping[str, Any], config: Any) -> Mapping[str, Any]: - import torch.distributed as dist - from portable_cache import resolve_cache_root - from validate_cache_snapshot import validate_snapshot - - if config.expected_approved_ordered_ids_sha256 is None: - raise ValueError( - "PDD training requires data.expected_approved_ordered_ids_sha256 before validation." - ) - if config.expected_heldout_count != _CANONICAL_PDD_HELDOUT_COUNT: - raise ValueError( - f"PDD training requires data.expected_heldout_count={_CANONICAL_PDD_HELDOUT_COUNT}." - ) - - options = _dataloader_options(raw) - root = resolve_cache_root(options["cache_dir"]) - roots: list[str] = [""] * dist.get_world_size() - dist.all_gather_object(roots, str(root)) - if any(candidate != roots[0] for candidate in roots[1:]): - raise RuntimeError(f"PDD ranks resolved different dataset cache roots: {roots}.") - - status = None - if dist.get_rank() == 0: - try: - report = validate_snapshot( - root, - all_index=config.all_metadata_index, - train_index=config.train_metadata_index, - heldout_index=config.validation_metadata_index, - expected_approved_ids_sha256=(config.expected_approved_ordered_ids_sha256), - expected_heldout_count=config.expected_heldout_count, - ) - status = {"ok": True, "report": report} - except BaseException as error: - status = {"ok": False, "error": f"{type(error).__name__}: {error}"} - payload = [status] - dist.broadcast_object_list(payload, src=0) - status = payload[0] - if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: - raise RuntimeError("rank 0 broadcast a malformed dataset validation status.") - if not status["ok"]: - raise RuntimeError(f"PDD dataset snapshot validation failed: {status.get('error')}.") - report = status.get("report") - if not isinstance(report, Mapping): - raise RuntimeError("PDD dataset snapshot report is malformed.") - return report - - -def _validated_loader_order_hashes( - train_metadata: Any, - heldout_metadata: Any, - snapshot_report: Mapping[str, Any], -) -> tuple[str, str]: - train_digest = _ordered_id_sha256(train_metadata, split="train") - heldout_digest = _ordered_id_sha256(heldout_metadata, split="heldout") - reported = snapshot_report.get("ordered_sample_ids_sha256") - if not isinstance(reported, Mapping): - raise RuntimeError("PDD dataset snapshot report has no ordered sample-ID hashes.") - if train_digest != reported.get("train"): - raise RuntimeError("training loader order does not match the authenticated snapshot.") - if heldout_digest != reported.get("heldout"): - raise RuntimeError("heldout loader order does not match the authenticated snapshot.") - return train_digest, heldout_digest - - -def _collective_validated_loader_order_hashes( - train_metadata: Any, - heldout_metadata: Any, - snapshot_report: Mapping[str, Any], -) -> tuple[str, str]: - """Authenticate loader order on every rank before distributed model construction.""" +def _validate_dataset_contract( + train_dataset: Any, + validation_dataset: Any, + config: Any, +) -> tuple[Mapping[str, Any], str, str]: + """Collectively verify deterministic split membership and the source metadata digest.""" import torch.distributed as dist try: - hashes = _validated_loader_order_hashes( - train_metadata, - heldout_metadata, - snapshot_report, - ) - local_status: dict[str, Any] = {"ok": True, "hashes": hashes} + train_ids = tuple(str(value) for value in train_dataset.sample_ids) + validation_ids = tuple(str(value) for value in validation_dataset.sample_ids) + if len(validation_ids) != config.validation_count: + raise RuntimeError( + f"validation split has {len(validation_ids)} samples; " + f"expected {config.validation_count}." + ) + if set(train_ids).intersection(validation_ids): + raise RuntimeError("training and validation splits overlap.") + expected = {str(index) for index in range(train_dataset.total_num_samples)} + if set(train_ids).union(validation_ids) != expected: + raise RuntimeError("training and validation splits do not cover metadata.json.") + if train_dataset.total_num_samples != validation_dataset.total_num_samples: + raise RuntimeError("training and validation datasets disagree on total sample count.") + if train_dataset.metadata_sha256 != validation_dataset.metadata_sha256: + raise RuntimeError("training and validation datasets disagree on metadata content.") + report = { + "cache_root": str(train_dataset.cache_root), + "metadata_sha256": train_dataset.metadata_sha256, + "total_samples": train_dataset.total_num_samples, + "train_samples": len(train_ids), + "validation_samples": len(validation_ids), + "split_seed": config.split_seed, + } + local_status: dict[str, Any] = { + "ok": True, + "report": report, + "train_hash": _ordered_id_sha256(train_ids), + "validation_hash": _ordered_id_sha256(validation_ids), + } except BaseException as error: local_status = {"ok": False, "error": f"{type(error).__name__}: {error}"} statuses: list[Any] = [None] * dist.get_world_size() dist.all_gather_object(statuses, local_status) failures: list[str] = [] - resolved_hashes: list[tuple[str, str]] = [] + successes: list[Mapping[str, Any]] = [] for rank, status in enumerate(statuses): if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: failures.append(f"rank {rank}: malformed loader authentication status") @@ -221,30 +187,23 @@ def _collective_validated_loader_order_hashes( if not status["ok"]: failures.append(f"rank {rank}: {status.get('error')}") continue - hashes = status.get("hashes") - if ( - not isinstance(hashes, tuple | list) - or len(hashes) != 2 - or any(not isinstance(value, str) for value in hashes) - ): - failures.append(f"rank {rank}: malformed loader authentication hashes") - continue - resolved_hashes.append((hashes[0], hashes[1])) + successes.append(status) if failures: - raise RuntimeError("PDD loader-order authentication failed: " + "; ".join(failures)) - if len(set(resolved_hashes)) != 1: + raise RuntimeError("PDD dataset validation failed: " + "; ".join(failures)) + canonical = {json.dumps(status, sort_keys=True) for status in successes} + if len(canonical) != 1: raise RuntimeError( - "PDD ranks resolved different authenticated train/heldout loader orders: " - f"{resolved_hashes}." + "PDD ranks resolved different dataset roots, metadata, or split membership." ) - return resolved_hashes[0] + return report, local_status["train_hash"], local_status["validation_hash"] def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[bool, ...], ...]]: import torch.distributed as dist - from pdd_training import build_pdd_validation_assignments - heldout_ids = _metadata_sample_ids(sampler.dataset.metadata, split="heldout") + from pdd.training import build_pdd_validation_assignments + + heldout_ids = tuple(str(value) for value in sampler.dataset.sample_ids) assignments = build_pdd_validation_assignments( heldout_ids, config.pdd, @@ -253,7 +212,7 @@ def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[ sampler.set_epoch(0) sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) local_plan = [ - tuple(sampler.dataset.metadata[index]["sample_id"] for index in batch) for batch in sampler + tuple(str(sampler.dataset.sample_ids[index]) for index in batch) for batch in sampler ] sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) plans: list[Any] = [None] * dist.get_world_size() @@ -288,7 +247,7 @@ def _iter_validation_batches( expected_latent_channels: int, expected_condition_features: int, ): - from pdd_training import prepare_qwen_pdd_batch + from pdd.training import prepare_qwen_pdd_batch count = 0 for count, (raw_batch, valid_mask) in enumerate(zip(dataloader, masks, strict=True), start=1): @@ -356,7 +315,8 @@ def _collective_training_batch( ) -> tuple[Any, tuple[str, ...]] | None: """Prepare one rank-local batch, then agree on success before any model call.""" import torch.distributed as dist - from pdd_training import prepare_qwen_pdd_batch + + from pdd.training import prepare_qwen_pdd_batch prepared = None sample_ids: tuple[str, ...] = () @@ -381,7 +341,10 @@ def _collective_training_batch( else: try: metadata = raw_batch["metadata"] - sample_ids = tuple(metadata["sample_ids"]) + raw_ids = metadata.get("logical_sample_ids", metadata.get("sample_ids")) + if hasattr(raw_ids, "tolist"): + raw_ids = raw_ids.tolist() + sample_ids = tuple(str(value) for value in raw_ids) expected_ids = sampler.expected_next_sample_ids() if sample_ids != expected_ids: raise RuntimeError( @@ -458,14 +421,15 @@ def main() -> None: args = _parse_args() import torch import torch.distributed as dist - from pdd_checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity - from pdd_recipe import ( + + from pdd.checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity + from pdd.recipe import ( build_pdd_setup, build_pdd_training_artifacts, initialize_pdd_distributed, resolve_pdd_recipe_config, ) - from pdd_training import run_pdd_validation + from pdd.training import run_pdd_validation raw = yaml.safe_load(args.config.read_text()) config = resolve_pdd_recipe_config(raw) @@ -480,7 +444,6 @@ def main() -> None: format="%(asctime)s %(levelname)s %(message)s", force=True, ) - snapshot_report = _validate_dataset_snapshot(raw, config) dataloader, sampler = _build_training_dataloader( raw, config, @@ -493,10 +456,12 @@ def main() -> None: dp_rank=rank, dp_world_size=world_size, ) - train_ordered_id_sha256, heldout_ordered_id_sha256 = _collective_validated_loader_order_hashes( - sampler.dataset.metadata, - validation_sampler.dataset.metadata, - snapshot_report, + snapshot_report, train_ordered_id_sha256, heldout_ordered_id_sha256 = ( + _validate_dataset_contract( + sampler.dataset, + validation_sampler.dataset, + config, + ) ) validation_assignments, validation_masks = _build_validation_plan( validation_sampler, @@ -529,7 +494,7 @@ def main() -> None: automodel_snapshot=setup.automodel_snapshot, ordered_train_id_sha256=train_ordered_id_sha256, ordered_heldout_id_sha256=heldout_ordered_id_sha256, - dataset_snapshot_sha256=snapshot_report["snapshot_sha256"], + dataset_snapshot_sha256=snapshot_report["metadata_sha256"], local_batch_size=config.training.local_batch_size, grad_accumulation_steps=config.training.grad_accumulation_steps, training_seed=config.training.seed, @@ -567,10 +532,11 @@ def main() -> None: ) if rank == 0: logging.info( - "PDD dataset snapshot verified: sha256=%s splits=%s declared_files=%s", - snapshot_report["snapshot_sha256"], - snapshot_report["splits"], - snapshot_report["declared_files"], + "PDD dataset verified: metadata_sha256=%s train=%d validation=%d root=%s", + snapshot_report["metadata_sha256"], + snapshot_report["train_samples"], + snapshot_report["validation_samples"], + snapshot_report["cache_root"], ) logging.info( "PDD setup complete: lifecycle=%s student_keys=%d AutoModel=%s", diff --git a/examples/diffusers/fastgen/inference_pdd_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py similarity index 94% rename from examples/diffusers/fastgen/inference_pdd_qwen_image.py rename to examples/diffusers/fastgen/pdd/inference_qwen_image.py index e90ea81f477..3cbf39c2991 100644 --- a/examples/diffusers/fastgen/inference_pdd_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -1,5 +1,17 @@ # 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. """Run conditional-only PDD inference from an authenticated Qwen-Image export.""" @@ -22,8 +34,9 @@ sys.dont_write_bytecode = True _THIS_DIR = Path(__file__).resolve().parent -_REPO_ROOT = _THIS_DIR.parents[2] -for path in (_REPO_ROOT, _THIS_DIR): +_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)) @@ -103,9 +116,9 @@ def _validate_qwen_projection(student: nn.Module, metadata: Any) -> nn.Linear: def build_pdd_student(export_dir: str | Path) -> tuple[nn.Module, Any, torch.dtype]: """Reconstruct and strictly load the converted Qwen student on CPU.""" from diffusers import QwenImageTransformer2DModel - from pdd_export import inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd + from pdd.export import inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata descriptor = inspect_pdd_export(export_dir) model_identity = _model_identity(descriptor) @@ -203,11 +216,11 @@ def _save_png(path: Path, image: Any) -> None: def main() -> None: args = _parse_args() from diffusers import QwenImagePipeline - from pdd_artifacts import sha256_file, write_canonical_json - from pdd_export import pdd_config_from_metadata from modelopt.torch.fastgen import PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import QwenImagePDDAdapter + from pdd.artifacts import sha256_file, write_canonical_json + from pdd.export import pdd_config_from_metadata if args.output.is_symlink() or args.result_json.is_symlink(): raise ValueError("PDD output and result JSON cannot be symlinks.") diff --git a/examples/diffusers/fastgen/pdd_recipe.py b/examples/diffusers/fastgen/pdd/recipe.py similarity index 93% rename from examples/diffusers/fastgen/pdd_recipe.py rename to examples/diffusers/fastgen/pdd/recipe.py index 508ba858880..ff00c10d201 100644 --- a/examples/diffusers/fastgen/pdd_recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -1,5 +1,17 @@ # 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. """ModelOpt-owned construction of a Qwen-Image PDD student and frozen teacher.""" @@ -15,9 +27,7 @@ import torch import torch.distributed as dist -from portable_cache import validate_relative_reference from torch import nn -from verify_readonly_automodel import snapshot_installed_distribution from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( @@ -25,6 +35,8 @@ convert_qwen_image_to_pdd, ) +from .verify_readonly_automodel import snapshot_installed_distribution + @dataclass(frozen=True) class PDDParallelConfig: @@ -85,11 +97,8 @@ class PDDRecipeConfig: weight_decay: float adam_betas: tuple[float, float] adam_eps: float - all_metadata_index: str - train_metadata_index: str - validation_metadata_index: str - expected_approved_ordered_ids_sha256: str | None - expected_heldout_count: int | None + validation_count: int + split_seed: int device: torch.device dtype: torch.dtype fuse_qkv_projections: bool @@ -222,36 +231,21 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: ): raise ValueError("PDD requires cached text embeddings; train_text_encoder must be false.") _require_bool(dataloader.get("shuffle", True), name="data.dataloader.shuffle") - all_metadata_index = validate_relative_reference( - data.get("all_metadata_index", "metadata.json"), - label="data.all_metadata_index", - ).as_posix() - train_metadata_index = validate_relative_reference( - dataloader.get("metadata_index", "metadata_train.json"), - label="data.dataloader.metadata_index", - ).as_posix() - validation_metadata_index = validate_relative_reference( - data.get("validation_metadata_index", "metadata_heldout.json"), - label="data.validation_metadata_index", - ).as_posix() - expected_approved_ordered_ids_sha256 = data.get("expected_approved_ordered_ids_sha256") - if expected_approved_ordered_ids_sha256 is not None and ( - not isinstance(expected_approved_ordered_ids_sha256, str) - or len(expected_approved_ordered_ids_sha256) != 64 - or expected_approved_ordered_ids_sha256.lower() != expected_approved_ordered_ids_sha256 - or any( - character not in "0123456789abcdef" - for character in expected_approved_ordered_ids_sha256 - ) - ): + if "metadata_index" in dataloader: raise ValueError( - "data.expected_approved_ordered_ids_sha256 must be a lowercase hexadecimal SHA-256." + "PDD uses deterministic ordinal splits from metadata.json; " + "data.dataloader.metadata_index is unsupported." ) - expected_heldout_count = data.get("expected_heldout_count") - if expected_heldout_count is not None and ( - type(expected_heldout_count) is not int or expected_heldout_count <= 0 - ): - raise ValueError("data.expected_heldout_count must be a positive integer.") + validation_count = _require_int_at_least( + data.get("validation_count", 2_000), + name="data.validation_count", + minimum=1, + ) + split_seed = _require_int_at_least( + data.get("split_seed", 2026), + name="data.split_seed", + minimum=0, + ) _reject_enabled(model.get("transformer_engine_linear"), name="global TE-linear conversion") _reject_enabled(model.get("peft"), name="PEFT/LoRA") @@ -302,7 +296,7 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: ) ): raise TypeError("optim.betas must contain two real numbers.") - adam_betas = tuple(float(beta) for beta in adam_betas_raw) + adam_betas = (float(adam_betas_raw[0]), float(adam_betas_raw[1])) if any(not math.isfinite(beta) or not 0.0 <= beta < 1.0 for beta in adam_betas): raise ValueError("optim.betas values must be finite and in [0, 1).") adam_eps = _require_finite_real( @@ -452,11 +446,8 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: weight_decay=float(weight_decay), adam_betas=adam_betas, adam_eps=adam_eps, - all_metadata_index=all_metadata_index, - train_metadata_index=train_metadata_index, - validation_metadata_index=validation_metadata_index, - expected_approved_ordered_ids_sha256=expected_approved_ordered_ids_sha256, - expected_heldout_count=expected_heldout_count, + validation_count=validation_count, + split_seed=split_seed, device=torch.device(model.get("device", "cuda" if torch.cuda.is_available() else "cpu")), dtype=_resolve_dtype(model.get("torch_dtype", "bfloat16")), fuse_qkv_projections=fuse_qkv_projections, @@ -713,12 +704,15 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: lifecycle = ["load/select"] pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) raw_transformer_config = getattr(student, "config", None) - if hasattr(raw_transformer_config, "to_dict"): - transformer_config = raw_transformer_config.to_dict() + to_dict = getattr(raw_transformer_config, "to_dict", None) + if callable(to_dict): + transformer_config = to_dict() elif isinstance(raw_transformer_config, Mapping): transformer_config = dict(raw_transformer_config) else: raise TypeError("Qwen transformer config must expose to_dict() or Mapping.") + if not isinstance(transformer_config, Mapping): + raise TypeError("Qwen transformer to_dict() must return a mapping.") projection = convert_qwen_image_to_pdd(student, config.pdd) identity = _projection_identity(projection) @@ -804,7 +798,8 @@ def build_pdd_training_artifacts( if not isinstance(config, PDDRecipeConfig): raise TypeError("config must be PDDRecipeConfig.") from nemo_automodel.components.training.rng import StatefulRNG - from pdd_training import PDDTrainer + + from .training import PDDTrainer adapter = QwenImagePDDAdapter( config.pdd, diff --git a/examples/diffusers/fastgen/pdd_training.py b/examples/diffusers/fastgen/pdd/training.py similarity index 98% rename from examples/diffusers/fastgen/pdd_training.py rename to examples/diffusers/fastgen/pdd/training.py index 62e997e04b2..76013b3ef2c 100644 --- a/examples/diffusers/fastgen/pdd_training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -1,5 +1,17 @@ # 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. """Direct PDD updates and a logical-ID-stable held-out validation oracle.""" @@ -273,7 +285,9 @@ def prepare_qwen_pdd_batch( ) if not isinstance(metadata, Mapping): raise TypeError("Qwen PDD batch metadata must be a mapping.") - sample_ids = metadata.get("sample_ids") + sample_ids = metadata.get("logical_sample_ids", metadata.get("sample_ids")) + if isinstance(sample_ids, torch.Tensor): + sample_ids = tuple(str(value) for value in sample_ids.tolist()) if isinstance(sample_ids, str) or not isinstance(sample_ids, Sequence): raise TypeError("Qwen PDD metadata.sample_ids must be a sequence of strings.") sample_ids = tuple(sample_ids) diff --git a/examples/diffusers/fastgen/verify_readonly_automodel.py b/examples/diffusers/fastgen/pdd/verify_readonly_automodel.py similarity index 92% rename from examples/diffusers/fastgen/verify_readonly_automodel.py rename to examples/diffusers/fastgen/pdd/verify_readonly_automodel.py index 841f9bc6e48..1ab6154b881 100644 --- a/examples/diffusers/fastgen/verify_readonly_automodel.py +++ b/examples/diffusers/fastgen/pdd/verify_readonly_automodel.py @@ -1,5 +1,17 @@ # 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. """Verify and snapshot the exact released AutoModel distribution used by PDD.""" @@ -52,7 +64,7 @@ def load_dependency_manifest(path: Path = _MANIFEST_PATH) -> dict[str, Any]: def _distribution_root( distribution: importlib.metadata.Distribution, manifest: dict[str, Any] ) -> Path: - root = Path(distribution.locate_file("")).resolve() + root = Path(str(distribution.locate_file(""))).resolve() dist_info = root / f"{manifest['distribution']}-{manifest['version']}.dist-info" if not dist_info.is_dir() or not dist_info.name.endswith(".dist-info"): raise RuntimeError(f"AutoModel has no regular wheel dist-info directory: {dist_info}.") diff --git a/examples/diffusers/fastgen/pdd_evaluation.py b/examples/diffusers/fastgen/pdd_evaluation.py deleted file mode 100644 index 5e5fac664c7..00000000000 --- a/examples/diffusers/fastgen/pdd_evaluation.py +++ /dev/null @@ -1,1103 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Strict evidence-bundle validation and deterministic PDD effectiveness summaries.""" - -from __future__ import annotations - -import hashlib -import math -from collections.abc import Mapping, Sequence -from pathlib import Path -from typing import Any - -from pdd_artifacts import ( - canonical_json_bytes, - load_canonical_json, - require_sha256, - sha256_file, - validate_artifact_reference, -) -from pdd_export import inspect_pdd_export - -from modelopt.torch.fastgen import make_shifted_flow_grid - -EVALUATION_CONDITIONS = ( - "teacher_guided", - "undistilled_euler_4", - "undistilled_2step_4eval", - "pdd_2", - "pdd_4", - "pdd_8", -) - - -def _grid_protocol(grid_size: int, *, grid_max_t: float) -> Mapping[str, Any]: - nodes = [ - float(value) - for value in make_shifted_flow_grid( - grid_size, - 5.0, - max_t=grid_max_t, - ).tolist() - ] - return { - "builder": "modelopt.torch.fastgen.make_shifted_flow_grid", - "construction_dtype": "float64", - "formula": ( - "u64=clamp(linspace(grid_max_t,0,grid_size+1),max=grid_max_t); " - "g64=clamp(shift*u64/(1+(shift-1)*u64),max=grid_max_t); " - "nodes=float32(g64)" - ), - "grid_size": grid_size, - "grid_max_t": grid_max_t, - "flow_shift": 5.0, - "initial_state": "float32(float64(noise)*float64(grid_max_t))", - "nodes": nodes, - "nodes_sha256": hashlib.sha256(canonical_json_bytes(nodes)).hexdigest(), - "runtime_dtype": "float32", - } - - -GRID_PROTOCOLS: Mapping[str, Mapping[str, Any]] = { - "pdd_grid_128_shift5": _grid_protocol(128, grid_max_t=0.999), - "teacher_grid_50_shift5": _grid_protocol(50, grid_max_t=0.999), -} - -_GUIDED_CFG_PROTOCOL = { - "execution": "sequential_conditional_unconditional", - "guidance_scale": 4.0, - "rescale": 1.0, - "eps": 1e-5, - "negative_condition": "manifest_negative_condition", -} - -_DISABLED_CFG_PROTOCOL = { - "execution": "disabled", - "guidance_scale": None, - "rescale": None, - "eps": None, - "negative_condition": None, -} - -INTEGRATOR_PROTOCOLS: Mapping[str, Mapping[str, Any]] = { - "euler_explicit": { - "math_dtype": "float32", - "velocity_evaluations_per_interval": 1, - "equations": [ - "dt=t_next-t_current", - "v_current=velocity(x_current,t_current)", - "x_next=x_current+dt*v_current", - ], - "terminal_rule": "apply the same Euler update when t_next=0; no special fallback", - }, - "heun_explicit_trapezoid": { - "math_dtype": "float32", - "velocity_evaluations_per_interval": 2, - "equations": [ - "dt=t_next-t_current", - "v_current=velocity(x_current,t_current)", - "x_predict=x_current+dt*v_current", - "v_predict=velocity(x_predict,t_next)", - "x_next=x_current+0.5*dt*(v_current+v_predict)", - ], - "terminal_rule": ( - "always evaluate v_predict at t_next, including t_next=0; no Euler fallback" - ), - }, - "pdd_fused_euler": { - "math_dtype": "float32", - "velocity_evaluations_per_interval": 1, - "implementation": "modelopt.torch.fastgen.methods.pdd.PDDPipeline.sample", - "source_identity": "manifest.modelopt.commit", - "equations": [ - "v_fused=student_fused_block(x_start,t_start,start,end,authenticated_grid)", - "x_end=x_start+(t_end-t_start)*v_fused", - ], - "terminal_rule": "apply the same fused update when t_end=0; no special fallback", - }, -} - -CONDITION_PROTOCOLS: Mapping[str, Mapping[str, Any]] = { - "teacher_guided": { - "artifact": "pdd_export", - "model_role": "frozen_teacher", - "integrator": "euler_explicit", - "cfg": _GUIDED_CFG_PROTOCOL, - "grid": { - "protocol": "teacher_grid_50_shift5", - "node_indices": list(range(51)), - }, - "pdd_blocks": [], - "scheduler_steps": 50, - "actual_transformer_invocations": 100, - "batch_normalized_transformer_evaluations": 100, - }, - "undistilled_euler_4": { - "artifact": "pdd_export", - "model_role": "pinned_base_model", - "integrator": "euler_explicit", - "cfg": _GUIDED_CFG_PROTOCOL, - "grid": { - "protocol": "pdd_grid_128_shift5", - "node_indices": [0, 32, 64, 96, 128], - }, - "pdd_blocks": [], - "scheduler_steps": 4, - "actual_transformer_invocations": 8, - "batch_normalized_transformer_evaluations": 8, - }, - "undistilled_2step_4eval": { - "artifact": "pdd_export", - "model_role": "pinned_base_model", - "integrator": "heun_explicit_trapezoid", - "cfg": _GUIDED_CFG_PROTOCOL, - "grid": { - "protocol": "pdd_grid_128_shift5", - "node_indices": [0, 64, 128], - }, - "pdd_blocks": [], - "scheduler_steps": 2, - "actual_transformer_invocations": 8, - "batch_normalized_transformer_evaluations": 8, - }, - "pdd_2": { - "artifact": "pdd_export", - "model_role": "pdd_student", - "integrator": "pdd_fused_euler", - "cfg": _DISABLED_CFG_PROTOCOL, - "grid": { - "protocol": "pdd_grid_128_shift5", - "node_indices": [0, 64, 128], - }, - "pdd_blocks": [64, 64], - "scheduler_steps": 2, - "actual_transformer_invocations": 2, - "batch_normalized_transformer_evaluations": 2, - }, - "pdd_4": { - "artifact": "pdd_export", - "model_role": "pdd_student", - "integrator": "pdd_fused_euler", - "cfg": _DISABLED_CFG_PROTOCOL, - "grid": { - "protocol": "pdd_grid_128_shift5", - "node_indices": [0, 32, 64, 96, 128], - }, - "pdd_blocks": [32, 32, 32, 32], - "scheduler_steps": 4, - "actual_transformer_invocations": 4, - "batch_normalized_transformer_evaluations": 4, - }, - "pdd_8": { - "artifact": "pdd_export", - "model_role": "pdd_student", - "integrator": "pdd_fused_euler", - "cfg": _DISABLED_CFG_PROTOCOL, - "grid": { - "protocol": "pdd_grid_128_shift5", - "node_indices": [0, 16, 32, 48, 64, 80, 96, 112, 128], - }, - "pdd_blocks": [16, 16, 16, 16, 16, 16, 16, 16], - "scheduler_steps": 8, - "actual_transformer_invocations": 8, - "batch_normalized_transformer_evaluations": 8, - }, -} - -_PROTOCOL_FIELDS = ( - "conditions", - "condition_protocols", - "grid_protocols", - "integrator_protocols", - "image_protocol", - "metric_protocols", - "timing_protocol", - "decision_rule", - "negative_condition", - "data_snapshot", - "stage_run_ids", - "prompt_set", - "bootstrap", -) - - -def _exact_mapping(value: Any, keys: set[str], *, name: str) -> Mapping[str, Any]: - if not isinstance(value, Mapping) or set(value) != keys: - actual = sorted(value) if isinstance(value, Mapping) else type(value).__name__ - raise ValueError(f"{name} keys mismatch: expected={sorted(keys)}, actual={actual}.") - return value - - -def _commit(value: Any, *, name: str) -> str: - if not isinstance(value, str) or len(value) != 40: - raise ValueError(f"{name} must be a 40-character Git commit.") - try: - int(value, 16) - except ValueError as error: - raise ValueError(f"{name} must be hexadecimal.") from error - return value.lower() - - -def _positive_finite(value: Any, *, name: str) -> float: - if isinstance(value, bool) or not isinstance(value, int | float): - raise TypeError(f"{name} must be a real number.") - resolved = float(value) - if not math.isfinite(resolved) or resolved <= 0: - raise ValueError(f"{name} must be finite and > 0.") - return resolved - - -def _nonnegative_int(value: Any, *, name: str) -> int: - if type(value) is not int or value < 0: - raise ValueError(f"{name} must be a nonnegative integer.") - return value - - -def _load_prompt_set(path: Path) -> tuple[dict[tuple[str, str, int], str], Mapping[str, Any]]: - prompt_set = _exact_mapping( - load_canonical_json(path), - {"schema_version", "prompts"}, - name="prompt set", - ) - if prompt_set["schema_version"] != 1 or not isinstance(prompt_set["prompts"], list): - raise ValueError("prompt set schema is unsupported.") - expected: dict[tuple[str, str, int], str] = {} - sort_keys: list[str] = [] - for index, raw_prompt in enumerate(prompt_set["prompts"]): - prompt = _exact_mapping( - raw_prompt, - {"prompt_id", "prompt", "prompt_sha256", "seeds"}, - name=f"prompts[{index}]", - ) - prompt_id = prompt["prompt_id"] - text = prompt["prompt"] - if not isinstance(prompt_id, str) or not prompt_id or not isinstance(text, str): - raise ValueError("prompt_id must be non-empty and prompt must be a string.") - digest = require_sha256(prompt["prompt_sha256"], name=f"prompts[{index}].prompt_sha256") - if hashlib.sha256(text.encode("utf-8")).hexdigest() != digest: - raise RuntimeError(f"prompt SHA-256 does not match for {prompt_id!r}.") - seeds = prompt["seeds"] - if ( - not isinstance(seeds, list) - or not seeds - or any(type(seed) is not int or seed < 0 or seed >= 2**63 for seed in seeds) - or seeds != sorted(set(seeds)) - ): - raise ValueError(f"prompt seeds must be sorted unique integers for {prompt_id!r}.") - sort_keys.append(prompt_id) - for seed in seeds: - key = (prompt_id, digest, seed) - if key in expected: - raise ValueError(f"duplicate prompt/seed pair: {key}.") - expected[key] = text - if sort_keys != sorted(set(sort_keys)): - raise ValueError("prompts must have unique IDs sorted lexicographically.") - if not expected: - raise ValueError("prompt set must contain at least one prompt/seed pair.") - return expected, prompt_set - - -def _protocol_sha256(manifest: Mapping[str, Any]) -> str: - payload = {name: manifest[name] for name in _PROTOCOL_FIELDS} - return hashlib.sha256(canonical_json_bytes(payload)).hexdigest() - - -def _validate_data_snapshot(root: Path, reference: Any) -> tuple[Path, Mapping[str, Any]]: - path = validate_artifact_reference(root, reference, name="data_snapshot") - snapshot = _exact_mapping( - load_canonical_json(path), - { - "schema_version", - "record_type", - "dataset_snapshot_sha256", - "train_ids_sha256", - "heldout_ids_sha256", - }, - name="data snapshot", - ) - if snapshot["schema_version"] != 1 or snapshot["record_type"] != "pdd_dataset_snapshot": - raise ValueError("data snapshot is not a schema-v1 PDD dataset snapshot.") - for name in ("dataset_snapshot_sha256", "train_ids_sha256", "heldout_ids_sha256"): - require_sha256(snapshot[name], name=f"data snapshot {name}") - return path, snapshot - - -def _validate_stage_evidence( - root: Path, - reference: Any, - *, - stage: str, - expected_run_id: str, - protocol_sha256: str, - model: Mapping[str, Any], - modelopt: Mapping[str, Any], - data_snapshot_sha256: str, - export_source_checkpoint: Mapping[str, Any], -) -> Mapping[str, Any]: - path = validate_artifact_reference(root, reference, name=f"stage_evidence.{stage}") - evidence = _exact_mapping( - load_canonical_json(path), - { - "schema_version", - "record_type", - "stage", - "status", - "run_id", - "model", - "modelopt", - "data_snapshot_sha256", - "evaluation_protocol_sha256", - "checkpoint", - "results", - }, - name=f"{stage} evidence", - ) - if ( - evidence["schema_version"] != 1 - or evidence["record_type"] != "pdd_stage_evidence" - or evidence["stage"] != stage - or evidence["status"] != "passed" - ): - raise ValueError(f"{stage} evidence is not a passed schema-v1 record.") - if evidence["run_id"] != expected_run_id: - raise RuntimeError(f"{stage} evidence run_id does not match frozen stage_run_ids.") - if evidence["model"] != model or evidence["modelopt"] != modelopt: - raise RuntimeError(f"{stage} evidence model/code lineage does not match evaluation.") - if ( - require_sha256(evidence["data_snapshot_sha256"], name=f"{stage} data snapshot SHA-256") - != data_snapshot_sha256 - ): - raise RuntimeError(f"{stage} evidence data lineage does not match evaluation.") - if ( - require_sha256(evidence["evaluation_protocol_sha256"], name=f"{stage} protocol SHA-256") - != protocol_sha256 - ): - raise RuntimeError(f"{stage} evidence does not carry the frozen evaluation protocol.") - checkpoint = _exact_mapping( - evidence["checkpoint"], - {"name", "manifest_sha256", "completed_steps"}, - name=f"{stage} checkpoint", - ) - if ( - not isinstance(checkpoint["name"], str) - or not checkpoint["name"] - or Path(checkpoint["name"]).name != checkpoint["name"] - or type(checkpoint["completed_steps"]) is not int - or checkpoint["completed_steps"] < 1 - ): - raise ValueError(f"{stage} checkpoint lineage is malformed.") - require_sha256(checkpoint["manifest_sha256"], name=f"{stage} checkpoint manifest SHA-256") - if stage == "training" and checkpoint != export_source_checkpoint: - raise RuntimeError("training evidence checkpoint does not match the exported checkpoint.") - results_path = validate_artifact_reference( - root, evidence["results"], name=f"{stage} evidence results" - ) - results = _exact_mapping( - load_canonical_json(results_path), - { - "schema_version", - "record_type", - "stage", - "status", - "slurm_job_ids", - "completed_updates", - "finite_loss", - "finite_gradients", - "resume_verified", - }, - name=f"{stage} results", - ) - expected_updates = 1_500 if stage == "canary" else 10_000 - expected_job_count = 3 if stage == "canary" else 1 - job_ids = results["slurm_job_ids"] - if ( - results["schema_version"] != 1 - or results["record_type"] != "pdd_stage_results" - or results["stage"] != stage - or results["status"] != "passed" - or not isinstance(job_ids, list) - or len(job_ids) != expected_job_count - or any(type(job_id) is not int or job_id <= 0 for job_id in job_ids) - or len(set(job_ids)) != len(job_ids) - or results["completed_updates"] != expected_updates - or results["finite_loss"] is not True - or results["finite_gradients"] is not True - or results["resume_verified"] is not True - ): - raise ValueError(f"{stage} results do not satisfy the frozen passed-stage contract.") - return evidence - - -def _validate_observations( - root: Path, - path: Path, - *, - prompt_pairs: Mapping[tuple[str, str, int], str], - metric_protocols: Mapping[str, Mapping[str, Any]], - image_protocol: Mapping[str, Any], - timing_protocol: Mapping[str, Any], - export_manifest_sha256: str, - evaluation_protocol_sha256: str, -) -> tuple[Mapping[str, Any], ...]: - document = _exact_mapping( - load_canonical_json(path), - {"schema_version", "records"}, - name="observations", - ) - if document["schema_version"] != 1 or not isinstance(document["records"], list): - raise ValueError("observation schema is unsupported.") - records: list[Mapping[str, Any]] = [] - observed: set[tuple[str, str, str, int]] = set() - order = {condition: index for index, condition in enumerate(EVALUATION_CONDITIONS)} - sort_keys: list[tuple[str, int, int]] = [] - keys = { - "condition", - "prompt_id", - "prompt_sha256", - "seed", - "metrics", - "output", - "scheduler_steps", - "actual_transformer_invocations", - "batch_normalized_transformer_evaluations", - "latency_seconds", - "throughput_images_per_second", - "peak_device_memory_bytes", - "height", - "width", - "protocol_sha256", - "evaluation_protocol_sha256", - "model_artifact_sha256", - } - for index, raw_record in enumerate(document["records"]): - record = _exact_mapping(raw_record, keys, name=f"records[{index}]") - condition = record["condition"] - if condition not in order: - raise ValueError(f"unknown evaluation condition {condition!r}.") - prompt_id = record["prompt_id"] - prompt_sha = require_sha256(record["prompt_sha256"], name=f"records[{index}].prompt_sha256") - seed = _nonnegative_int(record["seed"], name=f"records[{index}].seed") - pair = (prompt_id, prompt_sha, seed) - if pair not in prompt_pairs: - raise RuntimeError(f"observation does not match the prompt set: {pair}.") - key = (condition, *pair) - if key in observed: - raise ValueError(f"duplicate evaluation observation: {key}.") - observed.add(key) - protocol = CONDITION_PROTOCOLS[condition] - if ( - require_sha256(record["protocol_sha256"], name=f"records[{index}].protocol_sha256") - != hashlib.sha256(canonical_json_bytes(protocol)).hexdigest() - ): - raise RuntimeError(f"records[{index}] is not bound to its condition protocol.") - if ( - require_sha256( - record["evaluation_protocol_sha256"], - name=f"records[{index}].evaluation_protocol_sha256", - ) - != evaluation_protocol_sha256 - ): - raise RuntimeError(f"records[{index}] is not bound to the evaluation protocol.") - if ( - require_sha256( - record["model_artifact_sha256"], - name=f"records[{index}].model_artifact_sha256", - ) - != export_manifest_sha256 - ): - raise RuntimeError(f"records[{index}] model artifact does not match the PDD export.") - if ( - record["height"] != image_protocol["height"] - or record["width"] != image_protocol["width"] - ): - raise RuntimeError(f"records[{index}] resolution does not match image_protocol.") - metrics = record["metrics"] - if not isinstance(metrics, Mapping) or set(metrics) != set(metric_protocols): - raise ValueError(f"records[{index}].metrics does not match metric_protocols.") - if any( - isinstance(value, bool) - or not isinstance(value, int | float) - or not math.isfinite(float(value)) - for value in metrics.values() - ): - raise ValueError(f"records[{index}].metrics contains a non-finite value.") - validate_artifact_reference(root, record["output"], name=f"records[{index}].output") - counts = tuple( - _nonnegative_int(record[name], name=f"records[{index}].{name}") - for name in ( - "scheduler_steps", - "actual_transformer_invocations", - "batch_normalized_transformer_evaluations", - ) - ) - if not all(counts) or counts[1] > counts[2]: - raise ValueError(f"records[{index}] has invalid compute counters.") - expected_counts = tuple( - protocol[name] - for name in ( - "scheduler_steps", - "actual_transformer_invocations", - "batch_normalized_transformer_evaluations", - ) - ) - if counts != expected_counts: - raise RuntimeError( - f"{condition} compute counters must be {expected_counts}, got {counts}." - ) - latency = _positive_finite( - record["latency_seconds"], name=f"records[{index}].latency_seconds" - ) - throughput = _positive_finite( - record["throughput_images_per_second"], - name=f"records[{index}].throughput_images_per_second", - ) - expected_throughput = timing_protocol["batch_size"] / latency - if not math.isclose(throughput, expected_throughput, rel_tol=1e-6, abs_tol=0.0): - raise RuntimeError(f"records[{index}] throughput does not match batch size / latency.") - if ( - type(record["peak_device_memory_bytes"]) is not int - or record["peak_device_memory_bytes"] <= 0 - ): - raise ValueError(f"records[{index}].peak_device_memory_bytes must be positive.") - sort_keys.append((prompt_id, seed, order[condition])) - records.append(record) - if sort_keys != sorted(sort_keys): - raise ValueError("observations must be sorted by prompt_id, seed, and condition order.") - expected = {(condition, *pair) for pair in prompt_pairs for condition in EVALUATION_CONDITIONS} - if observed != expected: - missing = sorted(expected - observed) - extra = sorted(observed - expected) - raise RuntimeError( - f"effectiveness observations are incomplete: missing={missing[:5]}, extra={extra[:5]}." - ) - for condition in EVALUATION_CONDITIONS: - condition_counts = { - ( - record["scheduler_steps"], - record["actual_transformer_invocations"], - record["batch_normalized_transformer_evaluations"], - ) - for record in records - if record["condition"] == condition - } - if len(condition_counts) != 1: - raise RuntimeError(f"{condition} compute counters vary across paired observations.") - return tuple(records) - - -def _validate_automodel_snapshot( - snapshot: Any, - *, - export_automodel: Mapping[str, Any], -) -> None: - snapshot = _exact_mapping( - snapshot, - { - "distribution", - "files", - "import_origin", - "package_file_count", - "package_tree_sha256", - "release_commit", - "release_tag", - "root", - "runtime_versions", - "version", - "wheel", - "wheel_sha256", - }, - name="AutoModel environment snapshot", - ) - for key in ("distribution", "version", "runtime_versions"): - if snapshot[key] != export_automodel.get(key): - raise RuntimeError(f"AutoModel environment/export identity mismatch for {key}.") - for key in ("package_tree_sha256", "wheel_sha256"): - digest = require_sha256(snapshot[key], name=f"AutoModel snapshot {key}") - if digest != export_automodel.get(key): - raise RuntimeError(f"AutoModel environment/export identity mismatch for {key}.") - root = Path(snapshot["root"]) - import_origin = Path(snapshot["import_origin"]) - if not root.is_absolute() or not import_origin.is_absolute(): - raise ValueError("AutoModel snapshot root and import_origin must be absolute.") - try: - import_origin.relative_to(root) - except ValueError as error: - raise RuntimeError( - "AutoModel import origin is outside its installed distribution." - ) from error - files = snapshot["files"] - if ( - not isinstance(files, list) - or type(snapshot["package_file_count"]) is not int - or snapshot["package_file_count"] != len(files) - or not files - ): - raise ValueError("AutoModel snapshot file inventory is malformed.") - tree = hashlib.sha256() - previous = None - for index, raw_record in enumerate(files): - record = _exact_mapping( - raw_record, - {"path", "sha256", "size"}, - name=f"AutoModel files[{index}]", - ) - path = record["path"] - if ( - not isinstance(path, str) - or not path - or Path(path).is_absolute() - or "\\" in path - or any(part in ("", ".", "..") for part in path.split("/")) - or (previous is not None and path <= previous) - ): - raise ValueError("AutoModel snapshot file paths must be sorted normalized references.") - digest = require_sha256(record["sha256"], name=f"AutoModel files[{index}].sha256") - if type(record["size"]) is not int or record["size"] < 0: - raise ValueError(f"AutoModel files[{index}].size is invalid.") - tree.update(path.encode()) - tree.update(b"\0") - tree.update(digest.encode()) - tree.update(b"\0") - tree.update(str(record["size"]).encode()) - tree.update(b"\n") - previous = path - if tree.hexdigest() != snapshot["package_tree_sha256"]: - raise RuntimeError("AutoModel snapshot file inventory does not match its tree SHA-256.") - - -def _validate_evaluation_protocol(manifest: Mapping[str, Any]) -> None: - if manifest["conditions"] != list(EVALUATION_CONDITIONS): - raise ValueError("effectiveness manifest must contain the six fixed conditions in order.") - if manifest["condition_protocols"] != CONDITION_PROTOCOLS: - raise ValueError("condition_protocols must match the frozen Qwen PDD protocol exactly.") - if manifest["grid_protocols"] != GRID_PROTOCOLS: - raise ValueError("grid_protocols must contain the exact authenticated shifted-flow nodes.") - if manifest["integrator_protocols"] != INTEGRATOR_PROTOCOLS: - raise ValueError("integrator_protocols must contain the exact frozen update equations.") - stage_run_ids = _exact_mapping( - manifest["stage_run_ids"], {"canary", "training"}, name="stage_run_ids" - ) - if any(not isinstance(run_id, str) or not run_id for run_id in stage_run_ids.values()): - raise ValueError("stage_run_ids must contain non-empty run IDs.") - image = _exact_mapping( - manifest["image_protocol"], - {"height", "width", "batch_size", "max_sequence_length"}, - name="image_protocol", - ) - if image != {"height": 1024, "width": 1024, "batch_size": 1, "max_sequence_length": 512}: - raise ValueError("image_protocol must use the frozen 1024px single-image Qwen protocol.") - metrics = manifest["metric_protocols"] - if not isinstance(metrics, Mapping) or not metrics or list(metrics) != sorted(metrics): - raise ValueError("metric_protocols must be a non-empty, sorted mapping.") - for name, raw_protocol in metrics.items(): - if not isinstance(name, str) or not name: - raise ValueError("metric protocol names must be non-empty strings.") - protocol = _exact_mapping( - raw_protocol, - {"direction", "implementation", "revision"}, - name=f"metric_protocols.{name}", - ) - if protocol["direction"] not in ("higher", "lower"): - raise ValueError(f"metric_protocols.{name}.direction must be higher or lower.") - if not isinstance(protocol["implementation"], str) or not protocol["implementation"]: - raise ValueError(f"metric_protocols.{name}.implementation must be non-empty.") - _commit(protocol["revision"], name=f"metric_protocols.{name}.revision") - timing = _exact_mapping( - manifest["timing_protocol"], - {"batch_size", "warmup_runs", "measured_runs", "scope", "synchronize_device"}, - name="timing_protocol", - ) - if ( - timing["batch_size"] != 1 - or type(timing["warmup_runs"]) is not int - or timing["warmup_runs"] < 1 - or type(timing["measured_runs"]) is not int - or timing["measured_runs"] < 3 - or timing["scope"] != "transformer_sampling_and_vae_decode" - or timing["synchronize_device"] is not True - ): - raise ValueError("timing_protocol does not satisfy the frozen measurement contract.") - rule = _exact_mapping( - manifest["decision_rule"], - { - "primary_condition", - "primary_metric", - "quality_margin", - "quality_ci_rule", - "efficiency_measure", - "efficiency_baseline", - "minimum_relative_reduction", - "minimum_paired_samples", - }, - name="decision_rule", - ) - if rule["primary_condition"] not in ("pdd_2", "pdd_4", "pdd_8"): - raise ValueError("decision_rule.primary_condition must name a supported PDD schedule.") - if rule["primary_metric"] not in metrics: - raise ValueError("decision_rule.primary_metric is not in metric_protocols.") - if ( - isinstance(rule["quality_margin"], bool) - or not isinstance(rule["quality_margin"], int | float) - or not math.isfinite(float(rule["quality_margin"])) - or rule["quality_margin"] < 0 - or rule["quality_ci_rule"] != "paired_bootstrap_95_noninferiority" - ): - raise ValueError("decision_rule quality noninferiority contract is malformed.") - if ( - rule["efficiency_measure"] != "batch_normalized_transformer_evaluations" - or rule["efficiency_baseline"] != "teacher_guided" - or isinstance(rule["minimum_relative_reduction"], bool) - or not isinstance(rule["minimum_relative_reduction"], int | float) - or not 0 < float(rule["minimum_relative_reduction"]) < 1 - or type(rule["minimum_paired_samples"]) is not int - or rule["minimum_paired_samples"] < 16 - ): - raise ValueError("decision_rule efficiency/sample contract is malformed.") - - -def validate_effectiveness_bundle(manifest_path: str | Path) -> dict[str, Any]: - """Authenticate a complete, paired, claim-bearing effectiveness bundle.""" - unresolved_manifest = Path(manifest_path) - if unresolved_manifest.is_symlink(): - raise RuntimeError(f"effectiveness manifest cannot be a symlink: {unresolved_manifest}.") - manifest_path = unresolved_manifest.resolve() - root = manifest_path.parent - detached = manifest_path.with_suffix(manifest_path.suffix + ".sha256") - if not detached.is_file() or detached.is_symlink(): - raise RuntimeError(f"detached manifest SHA-256 is missing: {detached}.") - detached_bytes = detached.read_bytes() - expected_bytes = (sha256_file(manifest_path) + "\n").encode() - if detached_bytes != expected_bytes: - raise RuntimeError("detached effectiveness manifest SHA-256 does not match.") - manifest = _exact_mapping( - load_canonical_json(manifest_path), - { - "schema_version", - "stage", - "run_id", - "model", - "modelopt", - "pdd_export", - "prompt_set", - "observations", - "environment", - "data_snapshot", - "stage_evidence", - "conditions", - "condition_protocols", - "grid_protocols", - "integrator_protocols", - "image_protocol", - "metric_protocols", - "timing_protocol", - "decision_rule", - "negative_condition", - "stage_run_ids", - "bootstrap", - }, - name="effectiveness manifest", - ) - if manifest["schema_version"] != 1 or manifest["stage"] != "effectiveness_evaluation": - raise ValueError("only schema-v1 effectiveness_evaluation manifests can support claims.") - if not isinstance(manifest["run_id"], str) or not manifest["run_id"]: - raise ValueError("effectiveness run_id must be non-empty.") - model = _exact_mapping(manifest["model"], {"id", "revision"}, name="model") - if not isinstance(model["id"], str) or not model["id"]: - raise ValueError("model.id must be non-empty.") - _commit(model["revision"], name="model.revision") - modelopt = _exact_mapping(manifest["modelopt"], {"commit", "dirty"}, name="modelopt") - _commit(modelopt["commit"], name="modelopt.commit") - if modelopt["dirty"] is not False: - raise RuntimeError("claim-bearing effectiveness runs require a clean ModelOpt commit.") - _validate_evaluation_protocol(manifest) - bootstrap = _exact_mapping(manifest["bootstrap"], {"replicates", "seed"}, name="bootstrap") - if type(bootstrap["replicates"]) is not int or bootstrap["replicates"] < 1_000: - raise ValueError("bootstrap.replicates must be an integer >= 1000.") - _nonnegative_int(bootstrap["seed"], name="bootstrap.seed") - export_path = validate_artifact_reference(root, manifest["pdd_export"], name="pdd_export") - if export_path.name != "manifest.json": - raise ValueError("pdd_export must reference the export directory's manifest.json.") - export_descriptor = inspect_pdd_export(export_path.parent) - if export_descriptor.root / "manifest.json" != export_path: - raise RuntimeError("pdd_export does not identify the authenticated export manifest.") - environment_path = validate_artifact_reference( - root, manifest["environment"], name="environment" - ) - export_document = export_descriptor.manifest - environment_document = load_canonical_json(environment_path) - if not isinstance(export_document, Mapping): - raise ValueError("pdd_export must reference a canonical JSON object.") - if not isinstance(environment_document, Mapping): - raise ValueError("environment must reference a canonical JSON object.") - export_identity = export_document.get("identity") - export_modelopt = export_document.get("modelopt_source") - if ( - not isinstance(export_identity, Mapping) - or export_document.get("format") != "modelopt-pdd-safetensors" - ): - raise ValueError("pdd_export does not reference a ModelOpt PDD export manifest.") - export_model = export_identity.get("model") - if not isinstance(export_model, Mapping) or { - "id": export_model.get("id"), - "revision": export_model.get("revision"), - } != dict(model): - raise RuntimeError("effectiveness model does not match the PDD export identity.") - if not isinstance(export_modelopt, Mapping) or export_modelopt != modelopt: - raise RuntimeError("effectiveness ModelOpt source does not match the PDD export.") - export_automodel = export_identity.get("automodel") - if not isinstance(export_automodel, Mapping): - raise ValueError("PDD export has no AutoModel identity.") - export_guidance = export_identity.get("guidance") - if export_guidance != {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}: - raise RuntimeError("PDD export guidance does not match the frozen evaluation protocol.") - _validate_automodel_snapshot(environment_document, export_automodel=export_automodel) - data_snapshot_path, data_snapshot = _validate_data_snapshot(root, manifest["data_snapshot"]) - data_snapshot_sha256 = sha256_file(data_snapshot_path) - export_data = _exact_mapping( - export_identity.get("data"), - { - "ordered_train_id_sha256", - "ordered_heldout_id_sha256", - "dataset_snapshot_sha256", - "local_batch_size", - "grad_accumulation_steps", - }, - name="PDD export data identity", - ) - if { - "dataset_snapshot_sha256": export_data["dataset_snapshot_sha256"], - "train_ids_sha256": export_data["ordered_train_id_sha256"], - "heldout_ids_sha256": export_data["ordered_heldout_id_sha256"], - } != { - name: data_snapshot[name] - for name in ("dataset_snapshot_sha256", "train_ids_sha256", "heldout_ids_sha256") - }: - raise RuntimeError("PDD export training-data identity does not match data_snapshot.") - negative_path = validate_artifact_reference( - root, manifest["negative_condition"], name="negative_condition" - ) - negative = _exact_mapping( - load_canonical_json(negative_path), - {"schema_version", "record_type", "prompt_sha256", "embedding"}, - name="negative condition", - ) - if negative["schema_version"] != 1 or negative["record_type"] != "pdd_negative_condition": - raise ValueError("negative_condition must be a schema-v1 PDD negative condition.") - require_sha256(negative["prompt_sha256"], name="negative condition prompt SHA-256") - validate_artifact_reference(root, negative["embedding"], name="negative condition embedding") - prompt_path = validate_artifact_reference(root, manifest["prompt_set"], name="prompt_set") - observation_path = validate_artifact_reference( - root, manifest["observations"], name="observations" - ) - evidence = _exact_mapping( - manifest["stage_evidence"], {"canary", "training"}, name="stage_evidence" - ) - protocol_sha256 = _protocol_sha256(manifest) - export_source_checkpoint = export_document["source_checkpoint"] - for stage in ("canary", "training"): - _validate_stage_evidence( - root, - evidence[stage], - stage=stage, - expected_run_id=manifest["stage_run_ids"][stage], - protocol_sha256=protocol_sha256, - model=model, - modelopt=modelopt, - data_snapshot_sha256=data_snapshot_sha256, - export_source_checkpoint=export_source_checkpoint, - ) - prompt_pairs, prompt_set = _load_prompt_set(prompt_path) - records = _validate_observations( - root, - observation_path, - prompt_pairs=prompt_pairs, - metric_protocols=manifest["metric_protocols"], - image_protocol=manifest["image_protocol"], - timing_protocol=manifest["timing_protocol"], - export_manifest_sha256=sha256_file(export_path), - evaluation_protocol_sha256=protocol_sha256, - ) - if len(prompt_pairs) < manifest["decision_rule"]["minimum_paired_samples"]: - raise RuntimeError("paired sample count is below decision_rule.minimum_paired_samples.") - return { - "root": root, - "manifest": manifest, - "manifest_sha256": sha256_file(manifest_path), - "prompt_set": prompt_set, - "records": records, - } - - -def deterministic_bootstrap_mean_ci( - values: Sequence[float], - *, - replicates: int, - seed: int, -) -> tuple[float, float]: - """Return a deterministic SHA-256-index percentile interval for a sample mean.""" - if not values: - raise ValueError("bootstrap values must be non-empty.") - if replicates < 1: - raise ValueError("bootstrap replicates must be positive.") - samples: list[float] = [] - for replicate in range(replicates): - total = 0.0 - for draw in range(len(values)): - payload = f"{seed}:{replicate}:{draw}".encode() - index = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big") % len(values) - total += float(values[index]) - samples.append(total / len(values)) - samples.sort() - lower = samples[math.floor(0.025 * (replicates - 1))] - upper = samples[math.ceil(0.975 * (replicates - 1))] - return lower, upper - - -def summarize_effectiveness_bundle(validated: Mapping[str, Any]) -> dict[str, Any]: - """Compute paired aggregate metrics without promoting smoke output to evidence.""" - manifest = validated["manifest"] - records = validated["records"] - replicates = manifest["bootstrap"]["replicates"] - seed = manifest["bootstrap"]["seed"] - by_condition = { - condition: { - (record["prompt_id"], record["prompt_sha256"], record["seed"]): record - for record in records - if record["condition"] == condition - } - for condition in EVALUATION_CONDITIONS - } - pair_keys = sorted(by_condition["teacher_guided"]) - aggregates: dict[str, Any] = {} - for condition_index, condition in enumerate(EVALUATION_CONDITIONS): - condition_records = by_condition[condition] - metrics: dict[str, Any] = {} - for metric_index, (metric, protocol) in enumerate(manifest["metric_protocols"].items()): - direction = protocol["direction"] - values = [float(condition_records[key]["metrics"][metric]) for key in pair_keys] - teacher = [ - float(by_condition["teacher_guided"][key]["metrics"][metric]) for key in pair_keys - ] - deltas = [value - baseline for value, baseline in zip(values, teacher)] - metric_seed = seed + condition_index * 10_000 + metric_index * 2 - metrics[metric] = { - "direction": direction, - "mean": sum(values) / len(values), - "mean_ci95": list( - deterministic_bootstrap_mean_ci( - values, - replicates=replicates, - seed=metric_seed, - ) - ), - "paired_delta_vs_teacher": sum(deltas) / len(deltas), - "paired_delta_ci95": list( - deterministic_bootstrap_mean_ci( - deltas, - replicates=replicates, - seed=metric_seed + 1, - ) - ), - } - aggregates[condition] = { - "metrics": metrics, - "latency_seconds": {}, - "throughput_images_per_second": {}, - "peak_device_memory_bytes": {}, - "mean_batch_normalized_transformer_evaluations": sum( - condition_records[key]["batch_normalized_transformer_evaluations"] - for key in pair_keys - ) - / len(pair_keys), - } - latencies = [float(condition_records[key]["latency_seconds"]) for key in pair_keys] - teacher_latencies = [ - float(by_condition["teacher_guided"][key]["latency_seconds"]) for key in pair_keys - ] - latency_deltas = [value - baseline for value, baseline in zip(latencies, teacher_latencies)] - latency_seed = seed + condition_index * 10_000 + len(metrics) * 2 - aggregates[condition]["latency_seconds"] = { - "mean": sum(latencies) / len(latencies), - "mean_ci95": list( - deterministic_bootstrap_mean_ci( - latencies, - replicates=replicates, - seed=latency_seed, - ) - ), - "paired_delta_vs_teacher": sum(latency_deltas) / len(latency_deltas), - "paired_delta_ci95": list( - deterministic_bootstrap_mean_ci( - latency_deltas, - replicates=replicates, - seed=latency_seed + 1, - ) - ), - } - for telemetry_index, name in enumerate( - ("throughput_images_per_second", "peak_device_memory_bytes"), start=1 - ): - values = [float(condition_records[key][name]) for key in pair_keys] - teacher = [float(by_condition["teacher_guided"][key][name]) for key in pair_keys] - deltas = [value - baseline for value, baseline in zip(values, teacher)] - telemetry_seed = latency_seed + telemetry_index * 2 - aggregates[condition][name] = { - "mean": sum(values) / len(values), - "mean_ci95": list( - deterministic_bootstrap_mean_ci( - values, - replicates=replicates, - seed=telemetry_seed, - ) - ), - "paired_delta_vs_teacher": sum(deltas) / len(deltas), - "paired_delta_ci95": list( - deterministic_bootstrap_mean_ci( - deltas, - replicates=replicates, - seed=telemetry_seed + 1, - ) - ), - } - - rule = manifest["decision_rule"] - primary = aggregates[rule["primary_condition"]] - primary_metric = primary["metrics"][rule["primary_metric"]] - lower, upper = primary_metric["paired_delta_ci95"] - margin = float(rule["quality_margin"]) - direction = primary_metric["direction"] - if direction == "higher": - quality_state = ( - "pass" if lower >= -margin else "fail" if upper < -margin else "inconclusive" - ) - else: - quality_state = "pass" if upper <= margin else "fail" if lower > margin else "inconclusive" - baseline = aggregates[rule["efficiency_baseline"]] - candidate_value = primary["mean_batch_normalized_transformer_evaluations"] - baseline_value = baseline["mean_batch_normalized_transformer_evaluations"] - relative_reduction = 1.0 - candidate_value / baseline_value - efficiency_state = ( - "pass" if relative_reduction >= float(rule["minimum_relative_reduction"]) else "fail" - ) - if "fail" in (quality_state, efficiency_state): - conclusion = "not_effective" - elif (quality_state, efficiency_state) == ("pass", "pass"): - conclusion = "effective" - else: - conclusion = "inconclusive" - return { - "schema_version": 1, - "record_type": "effectiveness_summary", - "source_manifest_sha256": validated["manifest_sha256"], - "paired_sample_count": len(pair_keys), - "bootstrap": dict(manifest["bootstrap"]), - "aggregates": aggregates, - "decision": { - "label": conclusion, - "quality_state": quality_state, - "efficiency_state": efficiency_state, - "relative_efficiency_reduction": relative_reduction, - "rule": dict(rule), - }, - } diff --git a/examples/diffusers/fastgen/portable_cache.py b/examples/diffusers/fastgen/portable_cache.py deleted file mode 100644 index 46d41629761..00000000000 --- a/examples/diffusers/fastgen/portable_cache.py +++ /dev/null @@ -1,530 +0,0 @@ -# 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 cache paths, identities, and metadata loading for FastGen examples.""" - -from __future__ import annotations - -import hashlib -import json -import os -import stat -import struct -from collections.abc import Mapping, Sequence -from pathlib import Path, PurePosixPath, PureWindowsPath -from typing import Any - -PREPROCESS_STAGING_SCHEMA_VERSION = 1 -PORTABLE_SNAPSHOT_SCHEMA_VERSION = 2 -APPROVED_IDS_SCHEMA_VERSION = 1 -SPLIT_POLICY_SCHEMA_VERSION = 1 -DATASET_CACHE_ENV = "MODELOPT_FASTGEN_DATASET_CACHE_DIR" -SAMPLE_ID_DOMAIN = "modelopt-fastgen-sample-v1" -PDD_HOLDOUT_DOMAIN = "modelopt-pdd-holdout-v1" -_ORDERED_IDS_DOMAIN = b"modelopt-fastgen-ordered-sample-ids-v1" -_SPLIT_POLICY_KEYS = { - "algorithm", - "approved_ordered_ids_sha256", - "domain", - "heldout_count", - "schema_version", -} -_SHA256_HEX_LENGTH = 64 -_BANNED_PATH_KEYS = { - "cache_dir", - "image_path", - "output_dir", - "source_dir", - "source_path", - "video_path", -} -_PATH_FIELD_NAMES = { - "cache_file", - "directory", - "directories", - "file", - "files", - "path", - "paths", - "source_ref", -} - -__all__ = [ - "APPROVED_IDS_SCHEMA_VERSION", - "DATASET_CACHE_ENV", - "PDD_HOLDOUT_DOMAIN", - "PORTABLE_SNAPSHOT_SCHEMA_VERSION", - "PREPROCESS_STAGING_SCHEMA_VERSION", - "SAMPLE_ID_DOMAIN", - "SPLIT_POLICY_SCHEMA_VERSION", - "audit_no_absolute_paths", - "load_approved_sample_ids", - "load_portable_metadata", - "load_strict_json", - "ordered_sample_ids_sha256", - "resolve_cache_asset", - "resolve_cache_root", - "resolve_negative_embedding", - "select_pdd_holdout_ids", - "sha256_file", - "stable_sample_id", - "validate_relative_reference", -] - - -def resolve_cache_root(configured_root: str | os.PathLike[str]) -> Path: - """Resolve the YAML root unless a valid absolute environment override is set.""" - override = os.environ.get(DATASET_CACHE_ENV) - selected = Path(override) if override else Path(configured_root) - if override and not selected.is_absolute(): - raise ValueError(f"{DATASET_CACHE_ENV} must be an absolute path, got {override!r}.") - try: - resolved = selected.expanduser().resolve(strict=True) - except FileNotFoundError as error: - source = DATASET_CACHE_ENV if override else "configured cache_dir" - raise FileNotFoundError(f"{source} does not exist: {selected}") from error - if not resolved.is_dir(): - raise NotADirectoryError(f"dataset cache root is not a directory: {resolved}") - return resolved - - -def validate_relative_reference(value: str | os.PathLike[str], *, label: str) -> Path: - """Validate one portable, normalized relative reference without resolving it.""" - if not isinstance(value, str | os.PathLike): - raise TypeError(f"{label} must be a relative path string, got {type(value).__name__}.") - raw = os.fspath(value) - if not raw or raw == ".": - raise ValueError(f"{label} must be a non-empty relative path.") - if "\0" in raw: - raise ValueError(f"{label} must not contain NUL bytes.") - if "\\" in raw: - raise ValueError(f"{label} must use portable '/' separators, got {raw!r}.") - windows_path = PureWindowsPath(raw) - if PurePosixPath(raw).is_absolute() or windows_path.is_absolute() or windows_path.drive: - raise ValueError(f"{label} must be relative, got absolute path {raw!r}.") - path = Path(raw) - if any(part in ("", ".", "..") for part in PurePosixPath(raw).parts): - raise ValueError(f"{label} contains traversal or non-normalized components: {raw!r}.") - return path - - -def resolve_cache_asset( - root: Path, - reference: str | os.PathLike[str], - *, - label: str, - kind: str = "file", -) -> Path: - """Resolve an existing relative asset and reject traversal or symlink escape.""" - relative = validate_relative_reference(reference, label=label) - root = root.resolve(strict=True) - try: - resolved = (root / relative).resolve(strict=True) - except FileNotFoundError as error: - raise FileNotFoundError(f"{label} does not exist beneath cache root: {relative}") from error - try: - resolved.relative_to(root) - except ValueError as error: - raise ValueError(f"{label} resolves outside cache root: {relative}") from error - if kind == "file" and not resolved.is_file(): - raise ValueError(f"{label} is not a file: {relative}") - if kind == "directory" and not resolved.is_dir(): - raise ValueError(f"{label} is not a directory: {relative}") - if kind not in ("file", "directory", "any"): - raise ValueError(f"unsupported asset kind {kind!r}.") - return resolved - - -def resolve_negative_embedding(root: Path, reference: str | os.PathLike[str]) -> Path: - """Resolve a negative embedding and require it to stay beneath the cache root. - - Portable configs use a relative reference. An absolute reference is accepted only for - compatibility with existing launch overrides and only when it resolves beneath the same - effective root. - """ - raw = os.fspath(reference) - path = Path(raw).expanduser() - windows_path = PureWindowsPath(raw) - if windows_path.drive and not PurePosixPath(raw).is_absolute(): - raise ValueError("negative_prompt_embedding_path must use the host path syntax") - if not (PurePosixPath(raw).is_absolute() or windows_path.is_absolute()): - return resolve_cache_asset(root, raw, label="negative_prompt_embedding_path") - - root = root.resolve(strict=True) - try: - resolved = path.resolve(strict=True) - except FileNotFoundError as error: - raise FileNotFoundError(f"negative prompt embedding does not exist: {path}") from error - try: - resolved.relative_to(root) - except ValueError as error: - raise ValueError( - "negative prompt embedding resolves outside the effective cache root" - ) from error - if not resolved.is_file(): - raise ValueError(f"negative prompt embedding is not a file: {resolved}") - return resolved - - -def sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str: - """Return the hexadecimal SHA-256 digest of a file.""" - digest = hashlib.sha256() - with path.open("rb") as stream: - while chunk := stream.read(chunk_size): - digest.update(chunk) - return digest.hexdigest() - - -def _reject_duplicate_json_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: - value: dict[str, Any] = {} - for key, nested in pairs: - if key in value: - raise ValueError(f"JSON object contains duplicate key {key!r}.") - value[key] = nested - return value - - -def _reject_json_constant(value: str) -> None: - raise ValueError(f"JSON contains non-standard constant {value!r}.") - - -def load_strict_json(path: Path, *, label: str) -> Any: - """Load one regular, non-symlink UTF-8 JSON file with strict object syntax.""" - path = Path(path) - if path.is_symlink(): - raise ValueError(f"{label} must not be a symlink: {path}") - mode = path.stat().st_mode - if not stat.S_ISREG(mode): - raise ValueError(f"{label} must be a regular file: {path}") - try: - with path.open(encoding="utf-8") as stream: - return json.load( - stream, - object_pairs_hook=_reject_duplicate_json_keys, - parse_constant=_reject_json_constant, - ) - except (json.JSONDecodeError, UnicodeDecodeError) as error: - raise ValueError(f"{label} is not valid UTF-8 JSON: {path}") from error - - -def _validate_sha256(value: Any, *, label: str) -> str: - if ( - not isinstance(value, str) - or len(value) != _SHA256_HEX_LENGTH - or value.lower() != value - or any(character not in "0123456789abcdef" for character in value) - ): - raise ValueError(f"{label} must be a 64-character lowercase hexadecimal SHA-256 digest.") - return value - - -def _validate_ordered_sample_ids(sample_ids: Sequence[str], *, label: str) -> tuple[str, ...]: - if isinstance(sample_ids, str | bytes) or not isinstance(sample_ids, Sequence): - raise TypeError(f"{label} must be a sequence of strings.") - resolved = tuple(sample_ids) - if not resolved: - raise ValueError(f"{label} must be non-empty.") - for index, sample_id in enumerate(resolved): - if not isinstance(sample_id, str) or not sample_id: - raise ValueError(f"{label}[{index}] must be a non-empty string.") - try: - sample_id.encode("utf-8") - except UnicodeEncodeError as error: - raise ValueError(f"{label}[{index}] must be UTF-8 encodable.") from error - if len(resolved) != len(set(resolved)): - raise ValueError(f"{label} contains duplicates.") - return resolved - - -def ordered_sample_ids_sha256(sample_ids: Sequence[str]) -> str: - """Hash an ordered logical-ID sequence with count and byte-length framing.""" - resolved = _validate_ordered_sample_ids(sample_ids, label="sample_ids") - digest = hashlib.sha256() - digest.update(_ORDERED_IDS_DOMAIN) - digest.update(b"\0") - digest.update(struct.pack(">Q", len(resolved))) - for sample_id in resolved: - encoded = sample_id.encode("utf-8") - digest.update(struct.pack(">Q", len(encoded))) - digest.update(encoded) - return digest.hexdigest() - - -def load_approved_sample_ids( - path: Path, - *, - expected_sha256: str | None = None, -) -> tuple[tuple[str, ...], str]: - """Load and authenticate an ordered post-filter sample-ID artifact.""" - value = load_strict_json(path, label="approved_ids_manifest") - if not isinstance(value, dict): - raise ValueError("approved_ids_manifest must contain an object.") - required = { - "ordered_sample_ids", - "ordered_sample_ids_sha256", - "schema_version", - } - if set(value) != required: - raise ValueError( - "approved_ids_manifest keys mismatch: " - f"expected={sorted(required)}, actual={sorted(value)}." - ) - if ( - type(value["schema_version"]) is not int - or value["schema_version"] != APPROVED_IDS_SCHEMA_VERSION - ): - raise ValueError( - f"approved_ids_manifest.schema_version must be {APPROVED_IDS_SCHEMA_VERSION}." - ) - sample_ids = _validate_ordered_sample_ids( - value["ordered_sample_ids"], - label="approved_ids_manifest.ordered_sample_ids", - ) - computed = ordered_sample_ids_sha256(sample_ids) - declared = _validate_sha256( - value["ordered_sample_ids_sha256"], - label="approved_ids_manifest.ordered_sample_ids_sha256", - ) - if declared != computed: - raise ValueError("approved_ids_manifest ordered sample-ID SHA-256 mismatch.") - if expected_sha256 is not None: - expected = _validate_sha256(expected_sha256, label="expected_approved_ids_sha256") - if expected != computed: - raise ValueError("approved_ids_manifest does not match expected approved-ID SHA-256.") - return sample_ids, computed - - -def select_pdd_holdout_ids( - sample_ids: Sequence[str], heldout_count: int -) -> tuple[tuple[str, ...], tuple[str, ...]]: - """Select the frozen seedless PDD holdout membership while preserving input order.""" - resolved = _validate_ordered_sample_ids(sample_ids, label="sample_ids") - if type(heldout_count) is not int or not 0 < heldout_count < len(resolved): - raise ValueError( - "heldout_count must be an integer strictly between zero and len(sample_ids)." - ) - ranked = sorted( - resolved, - key=lambda sample_id: ( - hashlib.sha256(f"{PDD_HOLDOUT_DOMAIN}\0{sample_id}".encode()).digest(), - sample_id.encode("utf-8"), - ), - ) - heldout_members = set(ranked[:heldout_count]) - train = tuple(sample_id for sample_id in resolved if sample_id not in heldout_members) - heldout = tuple(sample_id for sample_id in resolved if sample_id in heldout_members) - return train, heldout - - -def stable_sample_id(*, source_ref: str, resolution: Sequence[int], model_type: str) -> str: - """Derive a root-independent sample ID from a logical source and processing identity.""" - logical = validate_relative_reference(source_ref, label="source_ref").as_posix() - if len(resolution) != 2 or any(type(value) is not int or value <= 0 for value in resolution): - raise ValueError(f"resolution must contain two positive integers, got {resolution!r}.") - if not isinstance(model_type, str) or not model_type: - raise ValueError("model_type must be a non-empty string.") - identity = ( - f"{SAMPLE_ID_DOMAIN}\0{model_type}\0{logical}\0{resolution[0]}x{resolution[1]}" - ).encode() - return hashlib.sha256(identity).hexdigest() - - -def _is_absolute_path(value: str) -> bool: - windows_path = PureWindowsPath(value) - return ( - PurePosixPath(value).is_absolute() - or windows_path.is_absolute() - or bool(windows_path.drive) - or value.lower().startswith("file://") - ) - - -def _is_path_field(key: str) -> bool: - lowered = key.lower() - return lowered in _PATH_FIELD_NAMES or lowered.endswith( - ("_path", "_paths", "_dir", "_dirs", "_file", "_files") - ) - - -def audit_no_absolute_paths( - value: Any, - *, - context: str = "value", - _path_context: bool = False, -) -> None: - """Reject absolute locations in path fields while leaving ordinary text untouched.""" - if isinstance(value, Mapping): - for key, nested in value.items(): - if not isinstance(key, str): - raise ValueError(f"{context} contains non-string key {key!r}.") - if _is_absolute_path(key): - raise ValueError(f"{context} contains absolute path key {key!r}.") - if key.lower() in _BANNED_PATH_KEYS: - raise ValueError(f"{context} contains forbidden personal-path key {key!r}.") - audit_no_absolute_paths( - nested, - context=f"{context}.{key}", - _path_context=_path_context or _is_path_field(key), - ) - return - if isinstance(value, list | tuple | set | frozenset): - for index, nested in enumerate(value): - audit_no_absolute_paths( - nested, - context=f"{context}[{index}]", - _path_context=_path_context, - ) - return - if _path_context and isinstance(value, str) and _is_absolute_path(value): - raise ValueError(f"{context} contains absolute path {value!r}.") - - -def _load_json(path: Path, *, expected_type: type, label: str): - value = load_strict_json(path, label=label) - if not isinstance(value, expected_type): - raise ValueError( - f"{label} must contain {expected_type.__name__}, got {type(value).__name__}." - ) - return value - - -def _resolve_strict_json_asset(root: Path, reference: str, *, label: str) -> Path: - relative = validate_relative_reference(reference, label=label) - candidate = root.resolve(strict=True) / relative - if candidate.is_symlink(): - raise ValueError(f"{label} must not be a symlink: {relative}") - return resolve_cache_asset(root, reference, label=label) - - -def _validate_split_policy(value: Any) -> dict[str, Any]: - if not isinstance(value, dict) or set(value) != _SPLIT_POLICY_KEYS: - actual = sorted(value) if isinstance(value, dict) else type(value).__name__ - raise ValueError( - "metadata_index.split_policy keys mismatch: " - f"expected={sorted(_SPLIT_POLICY_KEYS)}, actual={actual}." - ) - if ( - type(value["schema_version"]) is not int - or value["schema_version"] != SPLIT_POLICY_SCHEMA_VERSION - ): - raise ValueError( - f"metadata_index.split_policy.schema_version must be {SPLIT_POLICY_SCHEMA_VERSION}." - ) - if value["algorithm"] != "sha256-domain-ranked": - raise ValueError("metadata_index.split_policy.algorithm is unsupported.") - if value["domain"] != PDD_HOLDOUT_DOMAIN: - raise ValueError("metadata_index.split_policy.domain is unsupported.") - heldout_count = value["heldout_count"] - if type(heldout_count) is not int or heldout_count <= 0: - raise ValueError("metadata_index.split_policy.heldout_count must be a positive integer.") - _validate_sha256( - value["approved_ordered_ids_sha256"], - label="metadata_index.split_policy.approved_ordered_ids_sha256", - ) - return value - - -def load_portable_metadata( - root: Path, - metadata_index: str = "metadata.json", -) -> tuple[dict[str, Any], list[dict[str, Any]]]: - """Load and validate one split index, filtering IDs before bucket grouping.""" - index_path = _resolve_strict_json_asset(root, metadata_index, label="metadata_index") - index = _load_json(index_path, expected_type=dict, label="metadata_index") - audit_no_absolute_paths(index, context="metadata_index") - if ( - type(index.get("schema_version")) is not int - or index["schema_version"] != PORTABLE_SNAPSHOT_SCHEMA_VERSION - ): - raise ValueError( - "metadata_index is not an authenticated portable snapshot: " - f"schema_version must be {PORTABLE_SNAPSHOT_SCHEMA_VERSION}, got " - f"{index.get('schema_version')!r}; run migrate_cache_manifest.py." - ) - split_policy = _validate_split_policy(index.get("split_policy")) - shards = index.get("shards") - sample_ids = index.get("sample_ids") - if not isinstance(shards, list) or not shards or any(not isinstance(v, str) for v in shards): - raise ValueError("metadata_index.shards must be a non-empty list of relative paths.") - if ( - not isinstance(sample_ids, list) - or not sample_ids - or any(not isinstance(v, str) or not v for v in sample_ids) - ): - raise ValueError("metadata_index.sample_ids must be a non-empty list of strings.") - if len(sample_ids) != len(set(sample_ids)): - raise ValueError("metadata_index.sample_ids contains duplicates.") - if len(shards) != len(set(shards)): - raise ValueError("metadata_index.shards contains duplicates.") - if "num_shards" in index and ( - type(index["num_shards"]) is not int or index["num_shards"] != len(shards) - ): - raise ValueError("metadata_index.num_shards must be an integer equal to len(shards).") - if type(index.get("total_items")) is not int or index["total_items"] != len(sample_ids): - raise ValueError("metadata_index.total_items must equal len(sample_ids).") - computed_ordered_hash = ordered_sample_ids_sha256(sample_ids) - declared_ordered_hash = _validate_sha256( - index.get("ordered_sample_ids_sha256"), - label="metadata_index.ordered_sample_ids_sha256", - ) - if declared_ordered_hash != computed_ordered_hash: - raise ValueError("metadata_index ordered sample-ID SHA-256 mismatch.") - if index.get("split") == "all" and ( - split_policy["approved_ordered_ids_sha256"] != computed_ordered_hash - ): - raise ValueError("all metadata index does not match the approved ordered sample-ID hash.") - negative = index.get("negative_prompt_embedding") - if negative is not None: - if not isinstance(negative, dict) or set(negative) != {"path", "sha256"}: - raise ValueError( - "metadata_index.negative_prompt_embedding must contain path and sha256." - ) - validate_relative_reference( - negative["path"], - label="metadata_index.negative_prompt_embedding.path", - ) - _validate_sha256( - negative["sha256"], - label="metadata_index.negative_prompt_embedding.sha256", - ) - - entries_by_id: dict[str, dict[str, Any]] = {} - for shard_number, shard_ref in enumerate(shards): - shard_path = _resolve_strict_json_asset( - root, - shard_ref, - label=f"shards[{shard_number}]", - ) - entries = _load_json(shard_path, expected_type=list, label=f"shards[{shard_number}]") - for entry_number, entry in enumerate(entries): - label = f"shards[{shard_number}][{entry_number}]" - if not isinstance(entry, dict): - raise ValueError(f"{label} must be an object.") - audit_no_absolute_paths(entry, context=label) - sample_id = entry.get("sample_id") - if not isinstance(sample_id, str) or not sample_id: - raise ValueError(f"{label}.sample_id must be a non-empty string.") - if sample_id in entries_by_id: - raise ValueError(f"duplicate sample_id in shards: {sample_id}") - resolve_cache_asset(root, entry.get("cache_file"), label=f"{label}.cache_file") - _validate_sha256(entry.get("payload_sha256"), label=f"{label}.payload_sha256") - if "source_ref" in entry: - validate_relative_reference(entry["source_ref"], label=f"{label}.source_ref") - entries_by_id[sample_id] = entry - - missing = [sample_id for sample_id in sample_ids if sample_id not in entries_by_id] - if missing: - raise ValueError(f"metadata_index references missing sample_ids: {missing[:5]}") - return index, [entries_by_id[sample_id] for sample_id in sample_ids] 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/processors/base.py b/examples/diffusers/fastgen/preprocess/processors/base.py index 4ba806553f9..b0a1fafbd52 100644 --- a/examples/diffusers/fastgen/preprocess/processors/base.py +++ b/examples/diffusers/fastgen/preprocess/processors/base.py @@ -143,8 +143,7 @@ def get_cache_data( - bucket_resolution: Tuple[int, int] - crop_offset: Tuple[int, int] - prompt: str - - sample_id: str - - source_ref: Optional[str] + - image_path: str - bucket_id: str - tier: str - aspect_ratio: float diff --git a/examples/diffusers/fastgen/preprocess/processors/qwen_image.py b/examples/diffusers/fastgen/preprocess/processors/qwen_image.py index a8f37c5c8ed..61749d4284b 100644 --- a/examples/diffusers/fastgen/preprocess/processors/qwen_image.py +++ b/examples/diffusers/fastgen/preprocess/processors/qwen_image.py @@ -252,14 +252,12 @@ def get_cache_data( "bucket_resolution": metadata["bucket_resolution"], "crop_offset": metadata["crop_offset"], "prompt": metadata["prompt"], - "sample_id": metadata["sample_id"], + "image_path": metadata["image_path"], "bucket_id": metadata["bucket_id"], "aspect_ratio": metadata["aspect_ratio"], # Model info "model_type": self.model_type, } - if "source_ref" in metadata: - cache["source_ref"] = metadata["source_ref"] # Carry the positive-prompt attention mask through to the cache when present, so the # dataset uses the real mask instead of synthesizing an all-ones one. if "prompt_embeds_mask" in text_encodings: diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index 676a8942a52..ab54870c91c 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -2,11 +2,11 @@ # 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. +diffusers==0.38.0 # NeMo AutoModel supplies the parent recipe, FSDP2 wrapping, checkpointer, and the unmodified # upstream diffusion helpers used by the shared data/preprocessing code. These versions are the # public APIs tested by both DMD2 and PDD; PDD also verifies the exact AutoModel wheel at runtime. nemo_automodel[diffusion]==0.5.0 -diffusers==0.38.0 # Optional but recommended for training logs. wandb diff --git a/examples/diffusers/fastgen/seal_pdd_run_manifest.py b/examples/diffusers/fastgen/seal_pdd_run_manifest.py deleted file mode 100644 index 03eeb95b7bb..00000000000 --- a/examples/diffusers/fastgen/seal_pdd_run_manifest.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Create and verify the detached SHA-256 for a canonical PDD run manifest.""" - -from __future__ import annotations - -import argparse -import os -import sys -from pathlib import Path - -sys.dont_write_bytecode = True - -_THIS_DIR = Path(__file__).resolve().parent -_REPO_ROOT = _THIS_DIR.parents[2] -for path in (_REPO_ROOT, _THIS_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("manifest", type=Path) - args = parser.parse_args() - from pdd_artifacts import load_canonical_json, sha256_file - from pdd_evaluation import validate_effectiveness_bundle - - manifest = args.manifest.resolve() - load_canonical_json(manifest) - detached = manifest.with_suffix(manifest.suffix + ".sha256") - with detached.open("xb") as stream: - stream.write((sha256_file(manifest) + "\n").encode()) - stream.flush() - os.fsync(stream.fileno()) - try: - validate_effectiveness_bundle(manifest) - except BaseException: - detached.unlink(missing_ok=True) - raise - print(detached) - - -if __name__ == "__main__": - main() diff --git a/examples/diffusers/fastgen/validate_cache_snapshot.py b/examples/diffusers/fastgen/validate_cache_snapshot.py deleted file mode 100644 index c38e13a538d..00000000000 --- a/examples/diffusers/fastgen/validate_cache_snapshot.py +++ /dev/null @@ -1,274 +0,0 @@ -# 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. - -"""Read-only integrity validation for a portable FastGen cache snapshot.""" - -from __future__ import annotations - -import argparse -import hashlib -import json -from pathlib import Path -from typing import Any - -import torch -from portable_cache import ( - PORTABLE_SNAPSHOT_SCHEMA_VERSION, - audit_no_absolute_paths, - load_portable_metadata, - ordered_sample_ids_sha256, - resolve_cache_asset, - select_pdd_holdout_ids, - sha256_file, - validate_relative_reference, -) - - -def _validate_payload(root: Path, entry: dict[str, Any]) -> tuple[Path, str]: - payload_path = resolve_cache_asset(root, entry["cache_file"], label="cache_file") - actual_digest = sha256_file(payload_path) - if actual_digest != entry["payload_sha256"]: - raise ValueError( - f"payload SHA-256 mismatch for {entry['sample_id']}: " - f"expected {entry['payload_sha256']}, got {actual_digest}" - ) - payload = torch.load(payload_path, map_location="cpu", weights_only=True) - if not isinstance(payload, dict): - raise TypeError(f"cache payload must be a dict: {entry['cache_file']}") - audit_no_absolute_paths(payload, context=f"payload[{entry['sample_id']}]") - if payload.get("sample_id") != entry["sample_id"]: - raise ValueError(f"payload/manifest sample_id mismatch for {entry['cache_file']}") - if payload.get("source_ref") != entry.get("source_ref"): - raise ValueError(f"payload/manifest source_ref mismatch for {entry['cache_file']}") - return payload_path, actual_digest - - -def validate_snapshot( - cache_root: str | Path, - *, - all_index: str = "metadata.json", - train_index: str = "metadata_train.json", - heldout_index: str = "metadata_heldout.json", - reject_orphans: bool = True, - expected_approved_ids_sha256: str | None = None, - expected_heldout_count: int | None = None, -) -> dict[str, Any]: - """Validate manifests, payload hashes, splits, and declared snapshot inventory. - - The function performs no writes. All, train, and held-out indices are required so inventory, - split-disjointness, and split-union checks cannot be skipped accidentally. - """ - root = Path(cache_root).expanduser().resolve(strict=True) - if not root.is_dir(): - raise NotADirectoryError(f"cache_root is not a directory: {root}") - - expected_indexes = { - "all": validate_relative_reference(all_index, label="all_index").as_posix(), - "train": validate_relative_reference(train_index, label="train_index").as_posix(), - "heldout": validate_relative_reference(heldout_index, label="heldout_index").as_posix(), - } - missing_indexes = [name for name in expected_indexes.values() if not (root / name).is_file()] - if missing_indexes: - raise FileNotFoundError(f"required metadata indices do not exist: {missing_indexes}") - - split_ids: dict[str, tuple[str, ...]] = {} - split_policies: dict[str, dict[str, Any]] = {} - ordered_hashes: dict[str, str] = {} - index_hashes: dict[str, str] = {} - entries_by_id: dict[str, dict[str, Any]] = {} - declared_files = { - resolve_cache_asset(root, name, label="metadata_index") - for name in expected_indexes.values() - } - negative_declarations: set[tuple[str, str]] = set() - - for expected_split, index_name in expected_indexes.items(): - index, entries = load_portable_metadata(root, index_name) - split_name = index.get("split") - if split_name != expected_split: - raise ValueError(f"{index_name}.split must be {expected_split!r}, got {split_name!r}") - if split_name in split_ids: - raise ValueError(f"duplicate split declaration: {split_name}") - split_ids[split_name] = tuple(entry["sample_id"] for entry in entries) - split_policies[split_name] = dict(index["split_policy"]) - ordered_hashes[split_name] = ordered_sample_ids_sha256(split_ids[split_name]) - if ordered_hashes[split_name] != index["ordered_sample_ids_sha256"]: - raise ValueError(f"{index_name} ordered sample-ID SHA-256 mismatch") - index_hashes[split_name] = sha256_file(root / index_name) - - for shard_ref in index["shards"]: - declared_files.add(resolve_cache_asset(root, shard_ref, label="metadata shard")) - negative = index.get("negative_prompt_embedding") - if negative is not None: - if not isinstance(negative, dict) or set(negative) != {"path", "sha256"}: - raise ValueError( - f"{index_name}.negative_prompt_embedding must contain path and sha256" - ) - negative_path = resolve_cache_asset( - root, - negative["path"], - label=f"{index_name}.negative_prompt_embedding.path", - ) - if sha256_file(negative_path) != negative["sha256"]: - raise ValueError(f"negative prompt embedding SHA-256 mismatch in {index_name}") - negative_payload = torch.load(negative_path, map_location="cpu", weights_only=True) - audit_no_absolute_paths(negative_payload, context="negative_prompt_embedding") - negative_declarations.add((negative["path"], negative["sha256"])) - declared_files.add(negative_path) - - for entry in entries: - sample_id = entry["sample_id"] - previous = entries_by_id.get(sample_id) - if previous is not None and previous != entry: - raise ValueError(f"inconsistent manifest entry for sample_id {sample_id}") - entries_by_id[sample_id] = entry - - if len(negative_declarations) > 1: - raise ValueError("metadata indices disagree on the negative prompt embedding") - - policies = tuple(split_policies.values()) - if any(policy != policies[0] for policy in policies[1:]): - raise ValueError("metadata indices disagree on the split policy") - split_policy = split_policies["all"] - - train_ids = split_ids["train"] - heldout_ids = split_ids["heldout"] - all_ids = split_ids["all"] - train_members = set(train_ids) - heldout_members = set(heldout_ids) - overlap = train_members & heldout_members - if overlap: - raise ValueError(f"train and heldout splits overlap: {sorted(overlap)[:5]}") - if train_members | heldout_members != set(all_ids): - raise ValueError("train and heldout split union does not equal the all split") - expected_train, expected_heldout = select_pdd_holdout_ids( - all_ids, - split_policy["heldout_count"], - ) - if train_ids != expected_train: - raise ValueError("train split membership/order does not match the frozen PDD policy") - if heldout_ids != expected_heldout: - raise ValueError("heldout split membership/order does not match the frozen PDD policy") - if split_policy["approved_ordered_ids_sha256"] != ordered_hashes["all"]: - raise ValueError("split policy approved ordered-ID hash does not match the all index") - if expected_approved_ids_sha256 is not None: - if ( - not isinstance(expected_approved_ids_sha256, str) - or len(expected_approved_ids_sha256) != 64 - or expected_approved_ids_sha256.lower() != expected_approved_ids_sha256 - or any( - character not in "0123456789abcdef" for character in expected_approved_ids_sha256 - ) - ): - raise ValueError("expected_approved_ids_sha256 must be lowercase hexadecimal SHA-256") - if expected_approved_ids_sha256 != ordered_hashes["all"]: - raise ValueError("snapshot does not match the expected approved ordered-ID hash") - if expected_heldout_count is not None: - if type(expected_heldout_count) is not int or expected_heldout_count <= 0: - raise ValueError("expected_heldout_count must be a positive integer") - if split_policy["heldout_count"] != expected_heldout_count: - raise ValueError("snapshot heldout count does not match expected_heldout_count") - if len(heldout_ids) != split_policy["heldout_count"]: - raise ValueError("heldout split length does not match split policy") - - payload_hashes = dict(_validate_payload(root, entry) for entry in entries_by_id.values()) - payload_files = set(payload_hashes) - declared_files.update(payload_files) - - if reject_orphans: - actual_files = set() - for path in root.rglob("*"): - if path.is_symlink(): - resolved = path.resolve(strict=True) - try: - resolved.relative_to(root) - except ValueError as error: - raise ValueError( - f"snapshot symlink resolves outside cache root: {path}" - ) from error - raise ValueError(f"snapshot contains unsupported symlink: {path}") - if not path.is_file(): - continue - resolved = path.resolve(strict=True) - try: - resolved.relative_to(root) - except ValueError as error: - raise ValueError(f"snapshot file resolves outside cache root: {path}") from error - actual_files.add(resolved) - undeclared = sorted( - path.relative_to(root).as_posix() for path in actual_files - declared_files - ) - if undeclared: - raise ValueError(f"snapshot contains undeclared files: {undeclared[:5]}") - - declared_hashes = { - path.relative_to(root).as_posix(): ( - payload_hashes[path] if path in payload_hashes else sha256_file(path) - ) - for path in declared_files - } - snapshot_digest = hashlib.sha256() - snapshot_digest.update(b"modelopt-fastgen-cache-snapshot-v1\0") - for relative, file_digest in sorted(declared_hashes.items()): - snapshot_digest.update(relative.encode()) - snapshot_digest.update(b"\0") - snapshot_digest.update(file_digest.encode()) - snapshot_digest.update(b"\n") - - negative_report = None - if negative_declarations: - negative_path, negative_sha256 = next(iter(negative_declarations)) - negative_report = {"path": negative_path, "sha256": negative_sha256} - - return { - "schema_version": 1, - "record_type": "modelopt_fastgen_portable_snapshot_validation", - "snapshot_schema_version": PORTABLE_SNAPSHOT_SCHEMA_VERSION, - "indexes": expected_indexes, - "split_policy": split_policy, - "splits": {name: len(ids) for name, ids in split_ids.items()}, - "ordered_sample_ids_sha256": ordered_hashes, - "index_sha256": index_hashes, - "negative_prompt_embedding": negative_report, - "unique_payloads": len(payload_files), - "declared_files": len(declared_hashes), - "snapshot_sha256": snapshot_digest.hexdigest(), - } - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--cache-root", required=True) - parser.add_argument("--all-index", default="metadata.json") - parser.add_argument("--train-index", default="metadata_train.json") - parser.add_argument("--heldout-index", default="metadata_heldout.json") - parser.add_argument("--allow-orphans", action="store_true") - parser.add_argument("--expected-approved-ids-sha256") - parser.add_argument("--expected-heldout-count", type=int) - args = parser.parse_args() - report = validate_snapshot( - args.cache_root, - all_index=args.all_index, - train_index=args.train_index, - heldout_index=args.heldout_index, - reject_orphans=not args.allow_orphans, - expected_approved_ids_sha256=args.expected_approved_ids_sha256, - expected_heldout_count=args.expected_heldout_count, - ) - print(json.dumps(report, sort_keys=True, allow_nan=False)) - - -if __name__ == "__main__": - main() diff --git a/examples/diffusers/fastgen/validate_pdd_run_manifest.py b/examples/diffusers/fastgen/validate_pdd_run_manifest.py deleted file mode 100644 index 17631d27c5a..00000000000 --- a/examples/diffusers/fastgen/validate_pdd_run_manifest.py +++ /dev/null @@ -1,32 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Authenticate a complete PDD effectiveness evidence bundle.""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -sys.dont_write_bytecode = True - -_THIS_DIR = Path(__file__).resolve().parent -_REPO_ROOT = _THIS_DIR.parents[2] -for path in (_REPO_ROOT, _THIS_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("manifest", type=Path) - args = parser.parse_args() - from pdd_evaluation import validate_effectiveness_bundle - - validated = validate_effectiveness_bundle(args.manifest) - print(validated["manifest_sha256"]) - - -if __name__ == "__main__": - main() diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index 3e5f3c84d4e..6d3324099bf 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -123,14 +123,19 @@ def from_dict(cls, data: Mapping[str, Any]) -> PDDLayerSpec: ) if not isinstance(data["projection_path"], str): raise ValueError("layer_spec.projection_path must be a string.") - if not isinstance(data["head_layout"], str): - raise ValueError("layer_spec.head_layout must be a string.") + raw_head_layout = data["head_layout"] + if raw_head_layout == "channel_major": + head_layout: PDDHeadLayout = "channel_major" + elif raw_head_layout == "patch_major": + head_layout = "patch_major" + else: + raise ValueError(f"layer_spec.head_layout must be one of {_HEAD_LAYOUTS}.") output_channels = data["output_channels"] if output_channels is not None and type(output_channels) is not int: raise ValueError("layer_spec.output_channels must be an integer or null.") return cls( projection_path=data["projection_path"], - head_layout=data["head_layout"], + head_layout=head_layout, output_channels=output_channels, ) diff --git a/tests/examples/diffusers/fastgen/conftest.py b/tests/examples/diffusers/fastgen/conftest.py index 5090aa10b91..e4e7d7d04c2 100644 --- a/tests/examples/diffusers/fastgen/conftest.py +++ b/tests/examples/diffusers/fastgen/conftest.py @@ -13,9 +13,6 @@ # 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 - from __future__ import annotations import json diff --git a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py index 20b059b9e65..cc8b22436fa 100644 --- a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py @@ -1,5 +1,17 @@ # 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. """Two-rank proof that rank-0 checkpoint failures propagate instead of deadlocking.""" @@ -21,8 +33,8 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -import pdd_checkpoint as pdd_checkpoint_module -from pdd_checkpoint import PDDCheckpointManager +import pdd.checkpoint as pdd_checkpoint_module +from pdd.checkpoint import PDDCheckpointManager class _State: diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py index a5439d2c664..66fd1966910 100644 --- a/tests/examples/diffusers/fastgen/pdd_export_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_export_distributed.py @@ -1,5 +1,17 @@ # 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. """Two-rank released-AutoModel DCP-to-full-state export proof for the PDD example.""" @@ -21,10 +33,10 @@ sys.path.insert(0, str(path)) from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir -from export_pdd_qwen_image import collective_export_memory_preflight -from inference_pdd_qwen_image import build_pdd_student -from pdd_export import inspect_pdd_export, write_pdd_export -from pdd_recipe import build_pdd_export_setup, resolve_pdd_recipe_config +from pdd.export import inspect_pdd_export, write_pdd_export +from pdd.export_qwen_image import collective_export_memory_preflight +from pdd.inference_qwen_image import build_pdd_student +from pdd.recipe import build_pdd_export_setup, resolve_pdd_recipe_config def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: diff --git a/tests/examples/diffusers/fastgen/pdd_test_utils.py b/tests/examples/diffusers/fastgen/pdd_test_utils.py index 32b090281d6..668f282bcb7 100644 --- a/tests/examples/diffusers/fastgen/pdd_test_utils.py +++ b/tests/examples/diffusers/fastgen/pdd_test_utils.py @@ -1,16 +1,28 @@ # 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. """Small plain-torch objects shared by PDD example lifecycle tests.""" from __future__ import annotations +import hashlib from dataclasses import dataclass from typing import Any import torch -from pdd_training import PDDTrainer, PreparedPDDBatch -from portable_cache import ordered_sample_ids_sha256 +from pdd.training import PDDTrainer, PreparedPDDBatch from torch import nn from modelopt.torch.fastgen import ( @@ -23,6 +35,14 @@ ) +def ordered_sample_ids_sha256(sample_ids: tuple[str, ...] | list[str]) -> str: + digest = hashlib.sha256(b"modelopt-pdd-ordered-sample-ids-v1\0") + for sample_id in sample_ids: + digest.update(sample_id.encode()) + digest.update(b"\n") + return digest.hexdigest() + + class ToyStudent(nn.Module): def __init__(self, width: int = 3) -> None: super().__init__() @@ -185,6 +205,7 @@ def make_batch(sample_ids: tuple[str, ...], *, offset: float = 0.0) -> PreparedP class SamplerDataset: def __init__(self, sample_ids: tuple[str, ...]) -> None: + self.logical_sample_ids = sample_ids self.metadata = [ { "sample_id": sample_id, diff --git a/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py b/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py deleted file mode 100644 index f08606ac266..00000000000 --- a/tests/examples/diffusers/fastgen/pdd_training_preflight_distributed.py +++ /dev/null @@ -1,193 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Two-rank proof that training-input failures propagate before the model call.""" - -from __future__ import annotations - -import pathlib -import sys - -import torch -import torch.distributed as dist - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) -if str(_FASTGEN_DIR) not in sys.path: - sys.path.insert(0, str(_FASTGEN_DIR)) - -from pdd_finetune import ( - _collective_training_batch, - _collective_training_iterator, - _collective_validated_loader_order_hashes, - _ordered_id_sha256, -) - - -class _Sampler: - def __init__(self, sample_ids: tuple[str, ...]) -> None: - self.sample_ids = sample_ids - self.epoch = 0 - self.remaining_batches = 1 - - def expected_next_sample_ids(self) -> tuple[str, ...]: - return self.sample_ids - - def set_epoch(self, epoch: int) -> None: - self.epoch = epoch - self.remaining_batches = 1 - - -class _Loader: - def __init__(self, batch: dict, *, fail: bool) -> None: - self.batch = batch - self.fail = fail - - def __iter__(self): - if self.fail: - raise OSError("injected iterator construction failure") - return iter([self.batch]) - - -def _batch(sample_id: str) -> dict: - return { - "image_latents": torch.ones(1, 3, 4, 4), - "text_embeddings": torch.ones(1, 5, 6), - "text_embeddings_mask": torch.ones(1, 5, dtype=torch.bool), - "metadata": {"sample_ids": [sample_id]}, - } - - -def _expect_collective_failure(iterator, sampler: _Sampler, expected: str) -> None: - message = None - try: - _collective_training_batch( - iterator, - sampler=sampler, - resume=None, - resume_pending=False, - device=torch.device("cpu"), - dtype=torch.float32, - require_negative_condition=False, - expected_batch_size=1, - expected_latent_channels=3, - expected_condition_features=6, - ) - except RuntimeError as error: - message = str(error) - messages: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(messages, message) - assert all(item is not None and expected in item for item in messages) - - -def main() -> None: - dist.init_process_group("gloo") - try: - rank = dist.get_rank() - sample_id = f"sample-rank-{rank}" - expected_ids = ("wrong-rank-0",) if rank == 0 else (sample_id,) - _expect_collective_failure( - iter([_batch(sample_id)]), - _Sampler(expected_ids), - "committed cursor", - ) - dist.barrier() - - malformed = _batch(sample_id) - if rank == 1: - malformed.pop("text_embeddings") - _expect_collective_failure( - iter([malformed]), - _Sampler((sample_id,)), - "missing required keys", - ) - dist.barrier() - - malformed = _batch(sample_id) - if rank == 1: - malformed["text_embeddings_mask"] = torch.ones(1, 5, dtype=torch.float32) - _expect_collective_failure( - iter([malformed]), - _Sampler((sample_id,)), - "mask must use an integer or boolean dtype", - ) - dist.barrier() - - malformed = _batch(sample_id) - if rank == 0: - malformed["image_latents"] = torch.ones(1, 3, 3, 4) - _expect_collective_failure( - iter([malformed]), - _Sampler((sample_id,)), - "positive even spatial dimensions", - ) - dist.barrier() - - malformed = _batch(sample_id) - if rank == 1: - malformed["image_latents"] = torch.ones(1, 4, 4, 4) - _expect_collective_failure( - iter([malformed]), - _Sampler((sample_id,)), - "exactly 3 channels", - ) - dist.barrier() - - malformed = _batch(sample_id) - if rank == 0: - malformed["text_embeddings"] = torch.ones(1, 5, 7) - _expect_collective_failure( - iter([malformed]), - _Sampler((sample_id,)), - "exactly 6 features", - ) - dist.barrier() - - iterator_message = None - try: - _collective_training_iterator( - _Loader(_batch(sample_id), fail=rank == 0), - _Sampler((sample_id,)), - ) - except RuntimeError as error: - iterator_message = str(error) - iterator_messages: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(iterator_messages, iterator_message) - assert all( - item is not None and "iterator construction" in item for item in iterator_messages - ) - dist.barrier() - - canonical_train = [{"sample_id": "train-a"}, {"sample_id": "train-b"}] - heldout = [{"sample_id": "heldout-a"}] - report = { - "ordered_sample_ids_sha256": { - "train": _ordered_id_sha256(canonical_train, split="train"), - "heldout": _ordered_id_sha256(heldout, split="heldout"), - } - } - local_train = canonical_train if rank == 0 else list(reversed(canonical_train)) - gate_error = None - model_setup_reached = False - try: - _collective_validated_loader_order_hashes(local_train, heldout, report) - model_setup_reached = True - except RuntimeError as error: - gate_error = str(error) - outcomes: list[tuple[str | None, bool] | None] = [None] * dist.get_world_size() - dist.all_gather_object(outcomes, (gate_error, model_setup_reached)) - assert all( - outcome is not None - and outcome[0] is not None - and "loader-order authentication failed" in outcome[0] - and not outcome[1] - for outcome in outcomes - ) - finally: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py index d1fda52c11c..1f53db80e8f 100644 --- a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py @@ -1,5 +1,17 @@ # 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. """Two-rank CPU/Gloo equivalence harness for the deterministic PDD validation oracle.""" @@ -21,8 +33,8 @@ if str(pathlib.Path(__file__).parent) not in sys.path: sys.path.insert(0, str(pathlib.Path(__file__).parent)) +from pdd.training import build_pdd_validation_assignments, run_pdd_validation from pdd_test_utils import build_toy_lifecycle, make_batch -from pdd_training import build_pdd_validation_assignments, run_pdd_validation def _expect_failure(error_type, callback) -> None: diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index 715d4ced4e6..31c781fc51b 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -13,9 +13,6 @@ # 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 - """Portable and contained path contract for the shared FastGen cache.""" from __future__ import annotations diff --git a/tests/examples/diffusers/fastgen/test_dataset_splits.py b/tests/examples/diffusers/fastgen/test_dataset_splits.py index 9123566d558..639b5af0be3 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_splits.py +++ b/tests/examples/diffusers/fastgen/test_dataset_splits.py @@ -13,9 +13,6 @@ # 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 - """Deterministic stable-ID split contract for the shared FastGen cache.""" from __future__ import annotations @@ -108,7 +105,9 @@ def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( train_loader, train_sampler = build_text_to_image_multiresolution_dataloader( cache_dir=str(cache), - selected_indices=train_ids, + split="train", + validation_count=2, + split_seed=17, batch_size=1, num_workers=0, shuffle=True, @@ -116,7 +115,9 @@ def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( ) validation_loader, validation_sampler = build_text_to_image_multiresolution_dataloader( cache_dir=str(cache), - selected_indices=validation_ids, + split="validation", + validation_count=2, + split_seed=17, batch_size=1, num_workers=0, shuffle=False, @@ -138,6 +139,8 @@ def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( assert not validation_sampler.shuffle_buckets assert not validation_sampler.shuffle_within_bucket assert validation_sampler.drop_last is False + assert train_loader.dataset.logical_sample_ids == [str(value) for value in train_ids] + assert validation_loader.dataset.logical_sample_ids == [str(value) for value in validation_ids] assert all( batch["metadata"]["sample_ids"].dtype == torch.long and batch["metadata"]["sample_ids"].device.type == "cpu" diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index 0397e963c6e..b4c770bd2d5 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -13,14 +13,10 @@ # 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 - """Closed layout contract for the FastGen Diffusers example.""" from __future__ import annotations -import hashlib import json import os import pathlib @@ -39,6 +35,7 @@ "make_negative_prompt_embedding.py", "preprocess", "preprocess_qwen_image.py", + "pdd", "requirements.txt", } _EXPECTED_DMD2_FILES = { @@ -51,7 +48,21 @@ "inference_qwen_image.py", "recipe.py", } -_DMD2_CONFIG_DIGEST = "c802301bcaf1d861b3be4bf384ac437b4bea6cefb00865c3cf879355bcc42019" +_EXPECTED_PDD_FILES = { + "README.md", + "__init__.py", + "artifacts.py", + "automodel_dependency.json", + "checkpoint.py", + "configs", + "export.py", + "export_qwen_image.py", + "finetune.py", + "inference_qwen_image.py", + "recipe.py", + "training.py", + "verify_readonly_automodel.py", +} _TEXT_SUFFIXES = {".json", ".md", ".py", ".rst", ".sh", ".toml", ".txt", ".yaml", ".yml"} @@ -66,6 +77,14 @@ def _old_modules() -> tuple[str, ...]: _old_name("fastgen", "checkpoint"), _old_name("export", "diffusers_qwen_image"), _old_name("inference", "dmd2_qwen_image"), + _old_name("pdd", "artifacts"), + _old_name("pdd", "checkpoint"), + _old_name("pdd", "export"), + _old_name("pdd", "finetune"), + _old_name("pdd", "recipe"), + _old_name("pdd", "training"), + _old_name("export", "pdd_qwen_image"), + _old_name("inference", "pdd_qwen_image"), ) @@ -85,7 +104,7 @@ def _source_text_files() -> list[pathlib.Path]: ] -def test_fastgen_root_has_closed_shared_and_dmd2_ownership() -> None: +def test_fastgen_root_has_closed_shared_and_algorithm_ownership() -> None: root_entries = {path.name for path in _FASTGEN_ROOT.iterdir() if path.name != "__pycache__"} assert root_entries == _EXPECTED_ROOT_ENTRIES assert not (_FASTGEN_ROOT / "configs").exists() @@ -97,15 +116,18 @@ def test_fastgen_root_has_closed_shared_and_dmd2_ownership() -> None: assert {path.name for path in (_FASTGEN_ROOT / "dmd2" / "configs").iterdir()} == { "qwen_image.yaml" } + pdd_entries = { + path.name for path in (_FASTGEN_ROOT / "pdd").iterdir() if path.name != "__pycache__" + } + assert pdd_entries == _EXPECTED_PDD_FILES + assert {path.name for path in (_FASTGEN_ROOT / "pdd" / "configs").iterdir()} == { + "qwen_image.yaml" + } def test_dmd2_config_retains_accepted_semantics() -> None: config_path = _FASTGEN_ROOT / "dmd2" / "configs" / "qwen_image.yaml" value = yaml.safe_load(config_path.read_text()) - digest = hashlib.sha256( - json.dumps(value, sort_keys=True, separators=(",", ":")).encode() - ).hexdigest() - assert digest == _DMD2_CONFIG_DIGEST assert ( value["data"]["dataloader"]["_target_"] == "fastgen_data.build_text_to_image_multiresolution_dataloader" @@ -113,9 +135,10 @@ def test_dmd2_config_retains_accepted_semantics() -> None: assert value["data"]["dataloader"]["negative_prompt_embedding_path"] == ( "negative_prompt_embedding.pt" ) + assert "metadata_index" not in value["data"]["dataloader"] -def test_repository_sources_have_no_flat_dmd2_paths() -> None: +def test_repository_sources_have_no_flat_algorithm_paths() -> None: old_modules = _old_modules() stale = ( *(f"{module}.py" for module in old_modules), diff --git a/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py b/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py deleted file mode 100644 index db72855b64a..00000000000 --- a/tests/examples/diffusers/fastgen/test_migrate_cache_manifest.py +++ /dev/null @@ -1,563 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import json -import pathlib -import shutil -import subprocess -import sys - -import pytest - -torch = pytest.importorskip("torch") - -_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)) - -import migrate_cache_manifest as migration -from portable_cache import load_portable_metadata, ordered_sample_ids_sha256 -from validate_cache_snapshot import validate_snapshot - - -def _write_json(path: pathlib.Path, value) -> None: - path.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") - - -def _write_approved(path: pathlib.Path, sample_ids: list[str] | tuple[str, ...]) -> str: - digest = ordered_sample_ids_sha256(sample_ids) - _write_json( - path, - { - "schema_version": 1, - "ordered_sample_ids": list(sample_ids), - "ordered_sample_ids_sha256": digest, - }, - ) - return digest - - -def _make_legacy_cache( - root: pathlib.Path, - *, - stored_cache_root: pathlib.Path, - stored_source_root: pathlib.Path, - reverse_shards: bool = False, -) -> None: - root.mkdir() - (root / "legacy_payloads").mkdir() - entries = [] - for index in range(4): - relative_payload = pathlib.Path("legacy_payloads") / f"item-{index}.pt" - source_path = stored_source_root / "class" / f"item-{index}.png" - torch.save( - { - "latent": torch.full((2, 2), index, dtype=torch.float32), - "prompt_embeds": torch.full((3, 4), index, dtype=torch.float32), - "crop_offset": (0, 0), - "prompt": f"prompt {index}", - "image_path": str(source_path), - "nested": {"source_path": str(source_path)}, - }, - root / relative_payload, - ) - entries.append( - { - "cache_file": str(stored_cache_root / relative_payload), - "image_path": str(source_path), - "bucket_resolution": [64, 64], - "original_resolution": [64, 64], - "prompt": f"prompt {index}", - "bucket_id": "square-64", - "aspect_ratio": 1.0, - "pixels": 4096, - "model_type": "qwen_image", - } - ) - - shards = [entries[:2], entries[2:]] - if reverse_shards: - shards.reverse() - shard_names = [] - for index, shard in enumerate(shards): - name = f"legacy-shard-{index}.json" - _write_json(root / name, shard) - shard_names.append(name) - _write_json(root / "metadata.json", {"shards": shard_names, "total_items": 4}) - torch.save({"embed": torch.arange(12).reshape(3, 4)}, root / "legacy-negative.pt") - - -def _manifest_signature(root: pathlib.Path): - index, entries = load_portable_metadata(root, "metadata.json") - return index["sample_ids"], [ - (entry["sample_id"], entry["source_ref"], entry["cache_file"]) for entry in entries - ] - - -def test_migration_is_path_independent_and_relocatable(tmp_path): - alice = tmp_path / "alice-legacy" - bob = tmp_path / "bob-legacy" - alice_cache_prefix = pathlib.Path("/legacy/alice/cache") - bob_cache_prefix = pathlib.Path("/different/bob/cache") - alice_source_prefix = pathlib.Path("/datasets/alice/images") - bob_source_prefix = pathlib.Path("/mnt/bob/source") - _make_legacy_cache( - alice, - stored_cache_root=alice_cache_prefix, - stored_source_root=alice_source_prefix, - ) - _make_legacy_cache( - bob, - stored_cache_root=bob_cache_prefix, - stored_source_root=bob_source_prefix, - reverse_shards=True, - ) - - alice_output = tmp_path / "alice-portable" - bob_output = tmp_path / "bob-portable" - planned_ids = [ - record.sample_id - for record in migration.plan_migration( - alice, - legacy_cache_root=alice_cache_prefix, - legacy_source_root=alice_source_prefix, - ) - ] - approved_ids = [planned_ids[index] for index in (2, 0, 3)] - approved = tmp_path / "approved.json" - approved_digest = _write_approved(approved, approved_ids) - alice_result = migration.migrate_cache( - alice, - alice_output, - approved_ids_manifest=approved, - heldout_count=1, - expected_approved_ids_sha256=approved_digest, - legacy_cache_root=alice_cache_prefix, - legacy_source_root=alice_source_prefix, - negative_embedding="legacy-negative.pt", - shard_size=2, - ) - bob_result = migration.migrate_cache( - bob, - bob_output, - approved_ids_manifest=approved, - heldout_count=1, - legacy_cache_root=bob_cache_prefix, - legacy_source_root=bob_source_prefix, - negative_embedding="legacy-negative.pt", - shard_size=2, - ) - - assert _manifest_signature(alice_output) == _manifest_signature(bob_output) - alice_train = load_portable_metadata(alice_output, "metadata_train.json")[0]["sample_ids"] - bob_train = load_portable_metadata(bob_output, "metadata_train.json")[0]["sample_ids"] - assert alice_train == bob_train - alice_report = validate_snapshot(alice_output) - bob_report = validate_snapshot(bob_output) - assert alice_report["splits"] == {"all": 3, "train": 2, "heldout": 1} - assert bob_report["splits"] == {"all": 3, "train": 2, "heldout": 1} - assert alice_report["snapshot_sha256"] == bob_report["snapshot_sha256"] - assert alice_report["ordered_sample_ids_sha256"]["all"] == approved_digest - assert alice_report["split_policy"]["approved_ordered_ids_sha256"] == approved_digest - assert alice_result["validation"] == bob_result["validation"] - assert alice_result["counts"] == {"source": 4, "approved": 3, "filtered": 1} - assert set(alice_result) == { - "schema_version", - "record_type", - "output_root", - "counts", - "validation", - } - assert set(alice_result["validation"]) == { - "schema_version", - "record_type", - "snapshot_schema_version", - "indexes", - "split_policy", - "splits", - "ordered_sample_ids_sha256", - "index_sha256", - "negative_prompt_embedding", - "unique_payloads", - "declared_files", - "snapshot_sha256", - } - validation_text = json.dumps(alice_result["validation"], sort_keys=True) - assert str(alice_output) not in validation_text - assert ".staging-" not in validation_text - assert alice_result["output_root"] != bob_result["output_root"] - - cli = subprocess.run( - [ - sys.executable, - str(_FASTGEN_DIR / "validate_cache_snapshot.py"), - "--cache-root", - str(alice_output), - "--train-index", - "metadata_train.json", - "--heldout-index", - "metadata_heldout.json", - ], - check=False, - capture_output=True, - text=True, - ) - assert cli.returncode == 0, cli.stderr - - portable_text = "".join(path.read_text() for path in alice_output.glob("*.json")) - assert str(alice_cache_prefix) not in portable_text - assert str(alice_source_prefix) not in portable_text - for payload_path in alice_output.glob("payloads/*.pt"): - payload = torch.load(payload_path, map_location="cpu", weights_only=True) - assert "image_path" not in payload - assert "source_path" not in payload["nested"] - - relocated = tmp_path / "relocated" / "cache" - relocated.parent.mkdir() - shutil.copytree(alice_output, relocated) - assert _manifest_signature(relocated) == _manifest_signature(alice_output) - assert validate_snapshot(relocated)["snapshot_sha256"] == alice_report["snapshot_sha256"] - - -def test_incomplete_pass_one_publishes_nothing(monkeypatch, tmp_path): - legacy = tmp_path / "legacy" - cache_prefix = pathlib.Path("/legacy/cache") - source_prefix = pathlib.Path("/legacy/images") - _make_legacy_cache( - legacy, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - (legacy / "legacy_payloads" / "item-3.pt").unlink() - output = tmp_path / "portable" - save_calls = [] - monkeypatch.setattr(migration.torch, "save", lambda *args, **kwargs: save_calls.append(args)) - - with pytest.raises(FileNotFoundError): - migration.migrate_cache( - legacy, - output, - approved_ids_manifest=tmp_path / "unused-approved.json", - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - assert not output.exists() - assert not list(tmp_path.glob(".portable.staging-*")) - assert save_calls == [] - - -def test_approved_artifact_failures_and_final_validation_publish_nothing( - monkeypatch, tmp_path -) -> None: - legacy = tmp_path / "legacy" - cache_prefix = pathlib.Path("/legacy/cache") - source_prefix = pathlib.Path("/legacy/images") - _make_legacy_cache( - legacy, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - records = migration.plan_migration( - legacy, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - sample_ids = [record.sample_id for record in records] - approved = tmp_path / "approved.json" - digest = _write_approved(approved, sample_ids) - - with pytest.raises(ValueError, match="expected approved-ID"): - migration.migrate_cache( - legacy, - tmp_path / "bad-external", - approved_ids_manifest=approved, - expected_approved_ids_sha256="f" * 64, - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - unknown = tmp_path / "unknown.json" - _write_approved(unknown, [*sample_ids, "unknown"]) - with pytest.raises(ValueError, match="unknown sample IDs"): - migration.migrate_cache( - legacy, - tmp_path / "unknown-output", - approved_ids_manifest=unknown, - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - symlink = tmp_path / "approved-symlink.json" - symlink.symlink_to(approved) - with pytest.raises(ValueError, match="symlink"): - migration.migrate_cache( - legacy, - tmp_path / "symlink-output", - approved_ids_manifest=symlink, - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - - monkeypatch.setattr( - migration, - "validate_snapshot", - lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("injected final validation")), - ) - output = tmp_path / "validation-output" - with pytest.raises(ValueError, match="injected final validation"): - migration.migrate_cache( - legacy, - output, - approved_ids_manifest=approved, - expected_approved_ids_sha256=digest, - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - assert not output.exists() - assert not list(tmp_path.glob(".validation-output.staging-*")) - - -def test_migration_cli_removes_seed_and_requires_approved_manifest() -> None: - cli = subprocess.run( - [sys.executable, str(_FASTGEN_DIR / "migrate_cache_manifest.py"), "--help"], - check=False, - capture_output=True, - text=True, - ) - assert cli.returncode == 0, cli.stderr - assert "--approved-ids-manifest" in cli.stdout - assert "--split-seed" not in cli.stdout - - -def test_migration_rejects_finalized_and_non_strict_source_json(tmp_path) -> None: - cache_prefix = pathlib.Path("/legacy/cache") - source_prefix = pathlib.Path("/legacy/images") - - finalized = tmp_path / "finalized-source" - _make_legacy_cache( - finalized, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - index_path = finalized / "metadata.json" - index = json.loads(index_path.read_text()) - index["schema_version"] = 2 - _write_json(index_path, index) - with pytest.raises(ValueError, match="schema_version is unsupported"): - migration.plan_migration( - finalized, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - - duplicate = tmp_path / "duplicate-source" - _make_legacy_cache( - duplicate, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - (duplicate / "metadata.json").write_text( - '{"shards":["legacy-shard-0.json"],"shards":["legacy-shard-1.json"]}' - ) - with pytest.raises(ValueError, match="duplicate key"): - migration.plan_migration( - duplicate, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - - nonfinite = tmp_path / "nonfinite-source" - _make_legacy_cache( - nonfinite, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - (nonfinite / "metadata.json").write_text('{"shards":[],"total_items":NaN}') - with pytest.raises(ValueError, match="non-standard constant"): - migration.plan_migration( - nonfinite, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - - -def test_changed_source_after_frozen_plan_cleans_staging(monkeypatch, tmp_path): - legacy = tmp_path / "legacy" - cache_prefix = pathlib.Path("/legacy/cache") - source_prefix = pathlib.Path("/legacy/images") - _make_legacy_cache( - legacy, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - frozen = migration.plan_migration( - legacy, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - changed = frozen[-1].source_payload - original_plan = migration.plan_migration - - def _return_frozen(*args, **kwargs): - changed.write_bytes(b"changed after pass one") - return frozen - - monkeypatch.setattr(migration, "plan_migration", _return_frozen) - output = tmp_path / "portable" - approved = tmp_path / "approved.json" - _write_approved(approved, [record.sample_id for record in frozen]) - with pytest.raises(RuntimeError, match="changed after pass 1"): - migration.migrate_cache( - legacy, - output, - approved_ids_manifest=approved, - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - monkeypatch.setattr(migration, "plan_migration", original_plan) - assert not output.exists() - assert not list(tmp_path.glob(".portable.staging-*")) - - -def test_invalid_legacy_reference_fails_before_publish(tmp_path): - legacy = tmp_path / "legacy" - cache_prefix = pathlib.Path("/legacy/cache") - source_prefix = pathlib.Path("/legacy/images") - _make_legacy_cache( - legacy, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - shard_path = legacy / "legacy-shard-0.json" - shard = json.loads(shard_path.read_text()) - shard[0]["cache_file"] = "/other/private/cache.pt" - _write_json(shard_path, shard) - - output = tmp_path / "portable" - with pytest.raises(ValueError, match="outside the declared legacy prefix"): - migration.migrate_cache( - legacy, - output, - approved_ids_manifest=tmp_path / "unused-approved.json", - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - assert not output.exists() - - -def test_incomplete_source_counts_and_rank_indices_publish_nothing(tmp_path): - legacy = tmp_path / "legacy" - cache_prefix = pathlib.Path("/legacy/cache") - source_prefix = pathlib.Path("/legacy/images") - _make_legacy_cache( - legacy, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - index_path = legacy / "metadata.json" - index = json.loads(index_path.read_text()) - index["num_shards"] = 3 - _write_json(index_path, index) - output = tmp_path / "invalid-count-output" - with pytest.raises(ValueError, match="num_shards"): - migration.migrate_cache( - legacy, - output, - approved_ids_manifest=tmp_path / "unused-approved.json", - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - assert not output.exists() - assert not list(tmp_path.glob(".invalid-count-output.staging-*")) - - shards = index["shards"] - for rank, shard in enumerate(shards): - _write_json( - legacy / f"metadata_r{rank:02d}.json", - { - "shards": [shard], - "num_shards": 1, - "total_items": 2, - "shard_rank": rank, - "shard_world": 2, - }, - ) - incomplete_output = tmp_path / "incomplete-ranks-output" - with pytest.raises(ValueError, match="incomplete or inconsistent"): - migration.migrate_cache( - legacy, - incomplete_output, - approved_ids_manifest=tmp_path / "unused-approved.json", - source_index="metadata_r00.json", - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - assert not incomplete_output.exists() - - complete_output = tmp_path / "complete-ranks-output" - complete_records = migration.plan_migration( - legacy, - source_index=("metadata_r00.json", "metadata_r01.json"), - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - complete_approved = tmp_path / "complete-approved.json" - _write_approved(complete_approved, [record.sample_id for record in complete_records]) - migration.migrate_cache( - legacy, - complete_output, - approved_ids_manifest=complete_approved, - source_index=("metadata_r00.json", "metadata_r01.json"), - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - assert validate_snapshot(complete_output)["splits"] == {"all": 4, "train": 3, "heldout": 1} - - -def test_migration_path_audit_allows_prompt_commands_but_rejects_set_paths(tmp_path): - legacy = tmp_path / "legacy" - cache_prefix = pathlib.Path("/legacy/cache") - source_prefix = pathlib.Path("/legacy/images") - _make_legacy_cache( - legacy, - stored_cache_root=cache_prefix, - stored_source_root=source_prefix, - ) - shard_path = legacy / "legacy-shard-0.json" - shard = json.loads(shard_path.read_text()) - shard[0]["prompt"] = "/imagine a cat" - _write_json(shard_path, shard) - payload_path = legacy / "legacy_payloads" / "item-0.pt" - payload = torch.load(payload_path, map_location="cpu", weights_only=True) - payload["prompt"] = "/imagine a cat" - torch.save(payload, payload_path) - migration.plan_migration( - legacy, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - - payload["paths"] = {"/home/alice/private.png"} - torch.save(payload, payload_path) - output = tmp_path / "portable" - with pytest.raises(ValueError, match="absolute path"): - migration.migrate_cache( - legacy, - output, - approved_ids_manifest=tmp_path / "unused-approved.json", - heldout_count=1, - legacy_cache_root=cache_prefix, - legacy_source_root=source_prefix, - ) - assert not output.exists() diff --git a/tests/examples/diffusers/fastgen/test_pdd_evaluation.py b/tests/examples/diffusers/fastgen/test_pdd_evaluation.py deleted file mode 100644 index 7de6540d853..00000000000 --- a/tests/examples/diffusers/fastgen/test_pdd_evaluation.py +++ /dev/null @@ -1,529 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Hermetic tests for paired PDD effectiveness evidence and conclusions.""" - -from __future__ import annotations - -import copy -import hashlib -import pathlib -import sys - -import pytest - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -for path in (_REPO_ROOT, _FASTGEN_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - -from pdd_artifacts import ( - canonical_json_bytes, - load_canonical_json, - sha256_file, - write_canonical_json, -) -from pdd_evaluation import ( - CONDITION_PROTOCOLS, - EVALUATION_CONDITIONS, - GRID_PROTOCOLS, - INTEGRATOR_PROTOCOLS, - summarize_effectiveness_bundle, - validate_effectiveness_bundle, -) -from pdd_export import write_pdd_export - -from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC - - -def _reference(root: pathlib.Path, path: pathlib.Path) -> dict[str, str]: - return {"path": path.relative_to(root).as_posix(), "sha256": sha256_file(path)} - - -def _rewrite_manifest(manifest: pathlib.Path, data: dict) -> None: - manifest.unlink() - write_canonical_json(manifest, data) - detached = manifest.with_suffix(".json.sha256") - detached.unlink(missing_ok=True) - detached.write_bytes((sha256_file(manifest) + "\n").encode()) - - -def _export(root: pathlib.Path, automodel: dict) -> pathlib.Path: - config = PDDConfig( - grid_size=128, - grid_max_t=0.999, - flow_shift=5.0, - block_size_min=4, - block_size_max=64, - inference_blocks=[32, 32, 32, 32], - student_sample_steps=4, - guidance_scale=4.0, - num_train_timesteps=None, - ) - projection = PDDOutputProjection(1, 4, 128, QWEN_IMAGE_PDD_LAYER_SPEC) - metadata = PDDMetadata.from_config(config, projection) - identity = { - "model": {"id": "Qwen/Qwen-Image", "revision": "1" * 40, "dtype": "float32"}, - "pdd_metadata": metadata.to_dict(), - "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, - "automodel": automodel, - "data": { - "ordered_train_id_sha256": "8" * 64, - "ordered_heldout_id_sha256": "9" * 64, - "dataset_snapshot_sha256": "7" * 64, - "local_batch_size": 1, - "grad_accumulation_steps": 1, - }, - "topology": {"world_size": 1, "pure_data_parallel": True}, - } - return write_pdd_export( - root / "export", - projection.state_dict(), - metadata=metadata, - transformer_config={"_class_name": "QwenImageTransformer2DModel", "in_channels": 4}, - identity=identity, - source_checkpoint={ - "name": "step_00010000", - "manifest_sha256": "6" * 64, - "completed_steps": 10_000, - }, - modelopt_source={"commit": "2" * 40, "dirty": False}, - max_shard_bytes=1 << 20, - ) - - -def _write_bundle(tmp_path: pathlib.Path) -> pathlib.Path: - root = tmp_path / "run" - root.mkdir() - model = {"id": "Qwen/Qwen-Image", "revision": "1" * 40} - modelopt = {"commit": "2" * 40, "dirty": False} - file_path = "nemo_automodel/__init__.py" - file_sha = "3" * 64 - file_size = 7 - tree = hashlib.sha256() - tree.update(file_path.encode()) - tree.update(b"\0") - tree.update(file_sha.encode()) - tree.update(b"\0") - tree.update(str(file_size).encode()) - tree.update(b"\n") - tree_sha = tree.hexdigest() - automodel = { - "distribution": "nemo_automodel", - "version": "0.5.0", - "package_tree_sha256": tree_sha, - "wheel_sha256": "4" * 64, - "runtime_versions": {"diffusers": "0.38.0"}, - } - export_dir = _export(root, automodel) - export_manifest = export_dir / "manifest.json" - - environment = root / "environment.json" - write_canonical_json( - environment, - { - "distribution": "nemo_automodel", - "files": [{"path": file_path, "sha256": file_sha, "size": file_size}], - "import_origin": "/opt/pdd/site-packages/nemo_automodel/__init__.py", - "package_file_count": 1, - "package_tree_sha256": tree_sha, - "release_commit": "5" * 40, - "release_tag": "v0.5.0", - "root": "/opt/pdd/site-packages", - "runtime_versions": {"diffusers": "0.38.0"}, - "version": "0.5.0", - "wheel": "nemo_automodel-0.5.0-py3-none-any.whl", - "wheel_sha256": "4" * 64, - }, - ) - data_snapshot = root / "data_snapshot.json" - write_canonical_json( - data_snapshot, - { - "schema_version": 1, - "record_type": "pdd_dataset_snapshot", - "dataset_snapshot_sha256": "7" * 64, - "train_ids_sha256": "8" * 64, - "heldout_ids_sha256": "9" * 64, - }, - ) - - prompts = [] - prompt_pairs = [] - for index in range(16): - prompt = f"a small red cube on a white table, view {index:02d}" - prompt_sha = hashlib.sha256(prompt.encode()).hexdigest() - prompt_id = f"prompt-{index:04d}" - seed = 100 + index - prompts.append( - { - "prompt_id": prompt_id, - "prompt": prompt, - "prompt_sha256": prompt_sha, - "seeds": [seed], - } - ) - prompt_pairs.append((prompt_id, prompt_sha, seed)) - prompt_set = root / "prompts.json" - write_canonical_json(prompt_set, {"schema_version": 1, "prompts": prompts}) - negative_embedding = root / "negative_prompt_embedding.bin" - negative_embedding.write_bytes(b"authenticated fixed negative condition") - negative_condition = root / "negative_condition.json" - write_canonical_json( - negative_condition, - { - "schema_version": 1, - "record_type": "pdd_negative_condition", - "prompt_sha256": "c" * 64, - "embedding": _reference(root, negative_embedding), - }, - ) - - protocol_fields = { - "conditions": list(EVALUATION_CONDITIONS), - "condition_protocols": CONDITION_PROTOCOLS, - "grid_protocols": GRID_PROTOCOLS, - "integrator_protocols": INTEGRATOR_PROTOCOLS, - "image_protocol": { - "height": 1024, - "width": 1024, - "batch_size": 1, - "max_sequence_length": 512, - }, - "metric_protocols": { - "clip_score": { - "direction": "higher", - "implementation": "open_clip.ViT-H-14", - "revision": "a" * 40, - } - }, - "timing_protocol": { - "batch_size": 1, - "warmup_runs": 3, - "measured_runs": 5, - "scope": "transformer_sampling_and_vae_decode", - "synchronize_device": True, - }, - "decision_rule": { - "primary_condition": "pdd_4", - "primary_metric": "clip_score", - "quality_margin": 0.02, - "quality_ci_rule": "paired_bootstrap_95_noninferiority", - "efficiency_measure": "batch_normalized_transformer_evaluations", - "efficiency_baseline": "teacher_guided", - "minimum_relative_reduction": 0.5, - "minimum_paired_samples": 16, - }, - "negative_condition": _reference(root, negative_condition), - "data_snapshot": _reference(root, data_snapshot), - "stage_run_ids": {"canary": "canary-run", "training": "training-run"}, - "prompt_set": _reference(root, prompt_set), - "bootstrap": {"replicates": 1_000, "seed": 91}, - } - protocol_sha = hashlib.sha256(canonical_json_bytes(protocol_fields)).hexdigest() - evidence_references = {} - for stage in ("canary", "training"): - checkpoint = ( - { - "name": "step_00001500", - "manifest_sha256": "b" * 64, - "completed_steps": 1_500, - } - if stage == "canary" - else { - "name": "step_00010000", - "manifest_sha256": "6" * 64, - "completed_steps": 10_000, - } - ) - results = root / f"{stage}_results.json" - write_canonical_json( - results, - { - "schema_version": 1, - "record_type": "pdd_stage_results", - "stage": stage, - "status": "passed", - "slurm_job_ids": [101, 102, 103] if stage == "canary" else [201], - "completed_updates": 1_500 if stage == "canary" else 10_000, - "finite_loss": True, - "finite_gradients": True, - "resume_verified": True, - }, - ) - evidence = root / f"{stage}_evidence.json" - write_canonical_json( - evidence, - { - "schema_version": 1, - "record_type": "pdd_stage_evidence", - "stage": stage, - "status": "passed", - "run_id": f"{stage}-run", - "model": model, - "modelopt": modelopt, - "data_snapshot_sha256": sha256_file(data_snapshot), - "evaluation_protocol_sha256": protocol_sha, - "checkpoint": checkpoint, - "results": _reference(root, results), - }, - ) - evidence_references[stage] = _reference(root, evidence) - - records = [] - export_sha = sha256_file(export_manifest) - metric_values = { - "teacher_guided": 0.80, - "undistilled_euler_4": 0.73, - "undistilled_2step_4eval": 0.75, - "pdd_2": 0.77, - "pdd_4": 0.79, - "pdd_8": 0.795, - } - latencies = { - "teacher_guided": 10.0, - "undistilled_euler_4": 1.4, - "undistilled_2step_4eval": 1.2, - "pdd_2": 0.8, - "pdd_4": 1.0, - "pdd_8": 1.8, - } - for prompt_id, prompt_sha, seed in prompt_pairs: - for condition in EVALUATION_CONDITIONS: - output = root / f"{prompt_id}-{condition}.png" - output.write_bytes(b"png" + prompt_id.encode() + condition.encode()) - protocol = CONDITION_PROTOCOLS[condition] - latency = latencies[condition] - records.append( - { - "condition": condition, - "prompt_id": prompt_id, - "prompt_sha256": prompt_sha, - "seed": seed, - "metrics": {"clip_score": metric_values[condition]}, - "output": _reference(root, output), - "scheduler_steps": protocol["scheduler_steps"], - "actual_transformer_invocations": protocol["actual_transformer_invocations"], - "batch_normalized_transformer_evaluations": protocol[ - "batch_normalized_transformer_evaluations" - ], - "latency_seconds": latency, - "throughput_images_per_second": 1.0 / latency, - "peak_device_memory_bytes": 24_000_000_000, - "height": 1024, - "width": 1024, - "protocol_sha256": hashlib.sha256(canonical_json_bytes(protocol)).hexdigest(), - "evaluation_protocol_sha256": protocol_sha, - "model_artifact_sha256": export_sha, - } - ) - observations = root / "observations.json" - write_canonical_json(observations, {"schema_version": 1, "records": records}) - manifest = root / "manifest.json" - write_canonical_json( - manifest, - { - "schema_version": 1, - "stage": "effectiveness_evaluation", - "run_id": "test-run", - "model": model, - "modelopt": modelopt, - "pdd_export": _reference(root, export_manifest), - "observations": _reference(root, observations), - "environment": _reference(root, environment), - "stage_evidence": evidence_references, - **protocol_fields, - }, - ) - manifest.with_suffix(".json.sha256").write_bytes((sha256_file(manifest) + "\n").encode()) - return manifest - - -def test_effectiveness_bundle_is_authenticated_and_emits_effective_conclusion(tmp_path) -> None: - manifest = _write_bundle(tmp_path) - validated = validate_effectiveness_bundle(manifest) - first = summarize_effectiveness_bundle(validated) - second = summarize_effectiveness_bundle(validated) - - assert first == second - assert first["paired_sample_count"] == 16 - assert first["decision"]["label"] == "effective" - assert set(first["aggregates"]) == set(EVALUATION_CONDITIONS) - assert first["aggregates"]["pdd_2"]["mean_batch_normalized_transformer_evaluations"] == 2 - assert first["aggregates"]["pdd_4"]["metrics"]["clip_score"][ - "paired_delta_vs_teacher" - ] == pytest.approx(-0.01) - assert first["aggregates"]["pdd_4"]["peak_device_memory_bytes"]["mean"] > 0 - - -def test_grid_protocol_pins_fastgen_precision_and_raw_noise_initialization() -> None: - pdd_grid = GRID_PROTOCOLS["pdd_grid_128_shift5"] - teacher_grid = GRID_PROTOCOLS["teacher_grid_50_shift5"] - - assert pdd_grid["grid_max_t"] == teacher_grid["grid_max_t"] == 0.999 - assert pdd_grid["construction_dtype"] == "float64" - assert pdd_grid["runtime_dtype"] == "float32" - assert pdd_grid["initial_state"] == "float32(float64(noise)*float64(grid_max_t))" - assert pdd_grid["nodes"][0] == 0.9990000128746033 - assert pdd_grid["nodes"][-1] == 0.0 - assert ( - pdd_grid["nodes_sha256"] - == hashlib.sha256(canonical_json_bytes(pdd_grid["nodes"])).hexdigest() - ) - - -@pytest.mark.parametrize( - "corruption", - [ - "detached", - "stage", - "shadow", - "count", - "incomplete", - "missing_shard", - "tampered_shard", - "unrelated_training", - "failed_stage_results", - "unrelated_run", - "mismatched_export_data", - "bootstrap", - "prompt_reference", - "guided_count", - "grid_construction_dtype", - "grid_initial_state", - "grid_max_t", - "grid_nodes", - "grid_nodes_hash", - "latency_decision", - "integrator_formula", - "protocol", - "resolution", - ], -) -def test_effectiveness_bundle_rejects_unclaimable_evidence(tmp_path, corruption) -> None: - manifest = _write_bundle(tmp_path) - root = manifest.parent - data = copy.deepcopy(load_canonical_json(manifest)) - if corruption == "detached": - manifest.with_suffix(".json.sha256").write_bytes(("0" * 64 + "\n").encode()) - elif corruption == "stage": - data["stage"] = "smoke" - elif corruption == "shadow": - environment_path = root / data["environment"]["path"] - environment = copy.deepcopy(load_canonical_json(environment_path)) - environment["import_origin"] = "/project/automodel/nemo_automodel/__init__.py" - environment_path.unlink() - write_canonical_json(environment_path, environment) - data["environment"] = _reference(root, environment_path) - elif corruption in ("missing_shard", "tampered_shard"): - shard = next((root / "export").glob("*.safetensors")) - if corruption == "missing_shard": - shard.unlink() - else: - with shard.open("ab") as stream: - stream.write(b"tampered") - elif corruption == "unrelated_training": - evidence_path = root / data["stage_evidence"]["training"]["path"] - evidence = copy.deepcopy(load_canonical_json(evidence_path)) - evidence["checkpoint"]["manifest_sha256"] = "d" * 64 - evidence_path.unlink() - write_canonical_json(evidence_path, evidence) - data["stage_evidence"]["training"] = _reference(root, evidence_path) - elif corruption == "failed_stage_results": - evidence_path = root / data["stage_evidence"]["training"]["path"] - evidence = copy.deepcopy(load_canonical_json(evidence_path)) - results_path = root / evidence["results"]["path"] - results = copy.deepcopy(load_canonical_json(results_path)) - results["finite_gradients"] = False - results_path.unlink() - write_canonical_json(results_path, results) - evidence["results"] = _reference(root, results_path) - evidence_path.unlink() - write_canonical_json(evidence_path, evidence) - data["stage_evidence"]["training"] = _reference(root, evidence_path) - elif corruption == "unrelated_run": - evidence_path = root / data["stage_evidence"]["training"]["path"] - evidence = copy.deepcopy(load_canonical_json(evidence_path)) - evidence["run_id"] = "unrelated-run" - evidence_path.unlink() - write_canonical_json(evidence_path, evidence) - data["stage_evidence"]["training"] = _reference(root, evidence_path) - elif corruption == "mismatched_export_data": - snapshot_path = root / data["data_snapshot"]["path"] - snapshot = copy.deepcopy(load_canonical_json(snapshot_path)) - snapshot["dataset_snapshot_sha256"] = "e" * 64 - snapshot_path.unlink() - write_canonical_json(snapshot_path, snapshot) - data["data_snapshot"] = _reference(root, snapshot_path) - elif corruption == "bootstrap": - data["bootstrap"]["seed"] += 1 - elif corruption == "prompt_reference": - source = root / data["prompt_set"]["path"] - alternate = root / "alternate_prompts.json" - write_canonical_json(alternate, load_canonical_json(source)) - data["prompt_set"] = _reference(root, alternate) - elif corruption == "grid_nodes": - data["grid_protocols"]["pdd_grid_128_shift5"]["nodes"][32] += 1e-4 - elif corruption == "grid_nodes_hash": - data["grid_protocols"]["pdd_grid_128_shift5"]["nodes_sha256"] = "0" * 64 - elif corruption == "grid_max_t": - data["grid_protocols"]["pdd_grid_128_shift5"]["grid_max_t"] = 1.0 - elif corruption == "grid_construction_dtype": - data["grid_protocols"]["pdd_grid_128_shift5"]["construction_dtype"] = "float32" - elif corruption == "grid_initial_state": - data["grid_protocols"]["pdd_grid_128_shift5"]["initial_state"] = "noise" - elif corruption == "latency_decision": - data["decision_rule"]["efficiency_measure"] = "latency_seconds" - elif corruption == "integrator_formula": - data["integrator_protocols"]["heun_explicit_trapezoid"]["terminal_rule"] = ( - "fall back to Euler at t_next=0" - ) - elif corruption == "protocol": - data["condition_protocols"]["pdd_4"]["pdd_blocks"] = [64, 64] - else: - observations_path = root / data["observations"]["path"] - observations = copy.deepcopy(load_canonical_json(observations_path)) - if corruption == "count": - observations["records"][3]["actual_transformer_invocations"] = 3 - elif corruption == "guided_count": - observations["records"][0]["actual_transformer_invocations"] = 50 - observations["records"][0]["batch_normalized_transformer_evaluations"] = 50 - elif corruption == "resolution": - observations["records"][3]["height"] = 512 - else: - observations["records"].pop() - observations_path.unlink() - write_canonical_json(observations_path, observations) - data["observations"] = _reference(root, observations_path) - if corruption not in ("detached", "missing_shard", "tampered_shard"): - _rewrite_manifest(manifest, data) - - with pytest.raises((ValueError, RuntimeError, FileNotFoundError)): - validate_effectiveness_bundle(manifest) - - -@pytest.mark.parametrize( - ("pdd_values", "label"), - [ - ([0.75] * 16, "not_effective"), - ([0.80] * 8 + [0.76] * 8, "inconclusive"), - ], -) -def test_predeclared_decision_rule_emits_boundary_conclusions(tmp_path, pdd_values, label) -> None: - manifest = _write_bundle(tmp_path) - root = manifest.parent - data = copy.deepcopy(load_canonical_json(manifest)) - observations_path = root / data["observations"]["path"] - observations = copy.deepcopy(load_canonical_json(observations_path)) - primary = [record for record in observations["records"] if record["condition"] == "pdd_4"] - for record, value in zip(primary, pdd_values): - record["metrics"]["clip_score"] = value - observations_path.unlink() - write_canonical_json(observations_path, observations) - data["observations"] = _reference(root, observations_path) - _rewrite_manifest(manifest, data) - - summary = summarize_effectiveness_bundle(validate_effectiveness_bundle(manifest)) - assert summary["decision"]["label"] == label diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py index e9d32905beb..69ecdf20070 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -1,7 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 - -"""Hermetic evidence for authenticated PDD export, reconstruction, and schedules.""" +# +# 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. + +"""Tests for authenticated PDD export, reconstruction, and schedules.""" from __future__ import annotations @@ -20,14 +32,14 @@ if str(path) not in sys.path: sys.path.insert(0, str(path)) -from inference_pdd_qwen_image import _normalize_prompt_condition, _validate_qwen_projection -from pdd_export import ( +from pdd.export import ( PDD_INFERENCE_SCHEDULES, inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata, write_pdd_export, ) +from pdd.inference_qwen_image import _normalize_prompt_condition, _validate_qwen_projection from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( diff --git a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py b/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py deleted file mode 100644 index 45d87206d53..00000000000 --- a/tests/examples/diffusers/fastgen/test_pdd_qwen_operability_smoke.py +++ /dev/null @@ -1,493 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Hermetic tests for the full-Qwen PDD smoke evidence contract.""" - -from __future__ import annotations - -import copy -import hashlib -import importlib.util -import json -from pathlib import Path -from types import SimpleNamespace - -import pytest - -torch = pytest.importorskip("torch") - -_REPO_ROOT = Path(__file__).resolve().parents[4] -_HARNESS = _REPO_ROOT / "tests" / "gpu" / "torch" / "fastgen" / "pdd_qwen_operability_smoke.py" -_SPEC = importlib.util.spec_from_file_location("pdd_qwen_operability_smoke", _HARNESS) -assert _SPEC is not None and _SPEC.loader is not None -smoke = importlib.util.module_from_spec(_SPEC) -_SPEC.loader.exec_module(smoke) - - -def _stage_result(stage: str) -> dict: - step = 1 if stage == "train-one" else 2 - sample_ids = [f"synthetic-pdd-smoke-step-{step}-rank-{rank}" for rank in range(2)] - learning_rate = 2.0e-5 - return { - "schema_version": 1, - "record_type": "pdd_qwen_smoke_stage", - "stage": stage, - "pid": 100 + step, - "world_size": 2, - "model": { - "id": "Qwen/Qwen-Image", - "revision": "75e0b4be04f60ec59a75f475837eced720f823b6", - "dtype": "bfloat16", - }, - "pdd": { - "grid_size": 128, - "grid_max_t": 0.999, - "flow_shift": 5.0, - "block_size_min": 4, - "block_size_max": 64, - "teacher_integrator": "euler", - "guidance_scale": 4.0, - "guidance_rescale": 1.0, - "guidance_eps": 1e-5, - }, - "source": {"commit": "1" * 40, "dirty": False}, - "config_sha256": "2" * 64, - "automodel": { - "distribution": "nemo_automodel", - "version": "0.5.0", - "package_tree_sha256": smoke._AUTOMODEL_TREE_SHA256, - "wheel_sha256": smoke._AUTOMODEL_WHEEL_SHA256, - "runtime_versions": {"diffusers": "0.38.0"}, - }, - "gpu": { - "names": ["GPU", "GPU"], - "total_memory_bytes": [80_000_000_000, 80_000_000_000], - "host_available_bytes": [500_000_000_000, 500_000_000_000], - "allocated_before_step_bytes": [30_000_000_000, 30_000_000_000], - "peak_memory_bytes": [40_000_000_000, 40_000_000_000], - "student_parameter_bytes": 40_000_000_000, - "teacher_parameter_bytes": 40_000_000_000, - "step_seconds": 12.5, - }, - "pair": {"n": 0 if step == 1 else 124, "k": 63 if step == 1 else 127}, - "sample_ids": sample_ids, - "diagnostics": { - "completed_step": step, - "loss": 0.5, - "grad_norm": 1.25, - "student_adamw_nominal_update_ratio": 1e-4, - "pdd_projection_update_ratio": 2e-4, - "learning_rate": learning_rate, - "student_velocity_rms": 0.75, - "teacher_velocity_rms": 0.8, - "student_teacher_velocity_rms_ratio": 0.9375, - "reconstructed_state_rms": 1.1, - }, - "teacher_calls_per_rank": [2, 2], - "checkpoint": { - "path": f"checkpoints/step_{step:08d}", - "manifest_sha256": "5" * 64, - "completed_steps": step, - "parent_checkpoint": None if step == 1 else "step_00000001", - }, - "resume": None - if step == 1 - else { - "selected_checkpoint": "step_00000001", - "completed_steps": 1, - "parent_checkpoint": None, - "first_sample_ids": sample_ids, - "learning_rate": learning_rate, - }, - } - - -def test_gpu_harness_uses_shared_unambiguous_ordered_id_hash() -> None: - assert smoke._ordered_id_sha256(("a", "b")) == ( - "8cf774af4e8509811c2d4bc2adec6b852e4c614f9d8d833924502ead7c0689d7" - ) - assert smoke._ordered_id_sha256(("a\nb",)) == ( - "41e07cc133e8a85fc4a08e60a38c223f3c24dbca80312d106f251e533254eedf" - ) - source = _HARNESS.read_text() - assert "modelopt-pdd-ordered-{split}-ids-v1" not in source - assert 'digest.update(b"\\n")' not in source - - -def test_gpu_harness_checkpoint_identity_uses_exact_shared_hashes() -> None: - from modelopt.torch.fastgen import PDDLayerSpec, PDDMetadata - - train_ids = smoke._training_sample_ids(2) - parameter = torch.nn.Parameter(torch.tensor(1.0)) - optimizer = torch.optim.AdamW([parameter], lr=2.0e-5) - scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda _: 1.0) - setup = SimpleNamespace( - metadata=PDDMetadata( - grid_size=128, - grid_max_t=0.999, - flow_shift=5.0, - block_size_min=4, - block_size_max=64, - inference_blocks=(32, 32, 32, 32), - teacher_integrator="euler", - layer_spec=PDDLayerSpec( - projection_path="transformer.proj_out", - head_layout="channel_major", - ), - projection_in_features=3072, - projection_out_features=64, - projection_bias=True, - ), - automodel_snapshot={ - "distribution": "nemo_automodel", - "version": "0.5.0", - "package_tree_sha256": "a" * 64, - "wheel_sha256": "b" * 64, - "runtime_versions": {"diffusers": "0.38.0"}, - }, - optimizer=optimizer, - ) - config = SimpleNamespace( - model_id="Qwen/Qwen-Image", - model_revision="75e0b4be04f60ec59a75f475837eced720f823b6", - pdd=SimpleNamespace(guidance_scale=4.0), - guidance=SimpleNamespace(rescale=1.0, eps=1e-5), - training=SimpleNamespace( - seed=17, - validation_seed=29, - validation_every_steps=1000, - max_grad_norm=1.0, - zero_grad_warmup_steps=0, - ), - parallel=SimpleNamespace(activation_checkpointing=False), - ) - sampler = SimpleNamespace( - dataset=SimpleNamespace(metadata=[{"sample_id": sample_id} for sample_id in train_ids]) - ) - raw = {"pdd": {"grid_size": 128, "block_size_max": 64}} - - identity = smoke._identity( - setup=setup, - training=SimpleNamespace(scheduler=scheduler), - config=config, - sampler=sampler, - raw=raw, - ) - - assert identity["data"] == { - "ordered_train_id_sha256": ( - "4df732c492d043d5b0ea3549bcc80dbb847369021d0d3f242cd60386d1e94313" - ), - "ordered_heldout_id_sha256": ( - "38486bc077b6bd9b06a82167399e720f6c8dc70329dd0ce5fa23a92e4f30c198" - ), - "dataset_snapshot_sha256": smoke._canonical_sha256( - {"domain": "modelopt-pdd-synthetic-smoke-v1", "config": raw["pdd"]} - ), - "local_batch_size": 1, - "grad_accumulation_steps": 1, - } - - -def _automodel_snapshot_fixture() -> tuple[dict, dict]: - records = [] - tree = hashlib.sha256() - for index in range(smoke._AUTOMODEL_PACKAGE_FILE_COUNT): - path = f"nemo_automodel/file_{index:03d}.py" - digest = hashlib.sha256(f"file-{index}".encode()).hexdigest() - size = index + 1 - tree.update(path.encode()) - tree.update(b"\0") - tree.update(digest.encode()) - tree.update(b"\0") - tree.update(str(size).encode()) - tree.update(b"\n") - records.append({"path": path, "sha256": digest, "size": size}) - automodel = { - "distribution": "nemo_automodel", - "version": "0.5.0", - "package_tree_sha256": tree.hexdigest(), - "wheel_sha256": smoke._AUTOMODEL_WHEEL_SHA256, - "runtime_versions": {"diffusers": "0.38.0"}, - } - snapshot = { - **automodel, - "files": records, - "import_origin": "/opt/pdd/site-packages/nemo_automodel/__init__.py", - "package_file_count": smoke._AUTOMODEL_PACKAGE_FILE_COUNT, - "release_commit": smoke._AUTOMODEL_RELEASE_COMMIT, - "release_tag": smoke._AUTOMODEL_RELEASE_TAG, - "root": "/opt/pdd/site-packages", - "wheel": smoke._AUTOMODEL_WHEEL, - } - return snapshot, automodel - - -def _checkpoint_identity(stage: dict) -> dict: - return { - "schema_version": 1, - "model": stage["model"], - "pdd_metadata": { - "schema_version": 1, - "grid_size": 128, - "grid_max_t": 0.999, - "flow_shift": 5.0, - "block_size_min": 4, - "block_size_max": 64, - "inference_blocks": [32, 32, 32, 32], - "teacher_integrator": "euler", - "layer_spec": { - "projection_path": "transformer.proj_out", - "head_layout": "channel_major", - "output_channels": None, - }, - "base_projection": {"in_features": 3072, "out_features": 64, "bias": True}, - }, - "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, - "automodel": stage["automodel"], - "data": {}, - "topology": {"world_size": 2, "pure_data_parallel": True}, - "training": {}, - "optimizer": {}, - "scheduler": {}, - } - - -def _bundle_link_fixture() -> tuple[dict, dict, dict, dict, dict, dict]: - snapshot, automodel = _automodel_snapshot_fixture() - stage1 = _stage_result("train-one") - stage2 = _stage_result("resume-one") - stage1["automodel"] = copy.deepcopy(automodel) - stage2["automodel"] = copy.deepcopy(automodel) - identity = _checkpoint_identity(stage1) - manifest1 = {"identity": copy.deepcopy(identity)} - manifest2 = {"identity": copy.deepcopy(identity)} - export = { - "identity": copy.deepcopy(identity), - "modelopt_source": copy.deepcopy(stage2["source"]), - "source_checkpoint": { - "name": "step_00000002", - "manifest_sha256": stage2["checkpoint"]["manifest_sha256"], - "completed_steps": 2, - }, - } - return stage1, stage2, manifest1, manifest2, export, snapshot - - -def test_training_stage_contract_accepts_only_exact_canonical_chain() -> None: - stage1 = _stage_result("train-one") - stage2 = _stage_result("resume-one") - - smoke.validate_stage_result(stage1, stage="train-one") - smoke.validate_stage_result(stage2, stage="resume-one") - - wrong_pair = copy.deepcopy(stage2) - wrong_pair["pair"] = {"n": 120, "k": 127} - with pytest.raises(ValueError, match="support pair"): - smoke.validate_stage_result(wrong_pair, stage="resume-one") - - zero_update = copy.deepcopy(stage1) - zero_update["diagnostics"]["pdd_projection_update_ratio"] = 0.0 - with pytest.raises(ValueError, match="finite and positive"): - smoke.validate_stage_result(zero_update, stage="train-one") - - stale_resume = copy.deepcopy(stage2) - stale_resume["resume"]["selected_checkpoint"] = "step_00000000" - with pytest.raises(ValueError, match="resume evidence"): - smoke.validate_stage_result(stale_resume, stage="resume-one") - - -def test_inference_contract_authenticates_exact_pdd4_counters_and_png(tmp_path: Path) -> None: - image = tmp_path / "image.png" - image.write_bytes(b"\x89PNG\r\n\x1a\nnonempty-hashed-test-fixture") - digest = hashlib.sha256(image.read_bytes()).hexdigest() - result = { - "schema_version": 1, - "record_type": "pdd_inference", - "condition": "pdd_4", - "schedule": "pdd-4", - "blocks": [32, 32, 32, 32], - "height": 1024, - "width": 1024, - "scheduler_steps": 4, - "actual_transformer_invocations": 4, - "batch_normalized_transformer_evaluations": 4, - "latency_seconds": 1.5, - "output": {"path": "image.png", "sha256": digest}, - } - - smoke.validate_inference_result(result, root=tmp_path) - - wrong_calls = copy.deepcopy(result) - wrong_calls["actual_transformer_invocations"] = 5 - with pytest.raises(ValueError, match="compute counters"): - smoke.validate_inference_result(wrong_calls, root=tmp_path) - - wrong_hash = copy.deepcopy(result) - wrong_hash["output"]["sha256"] = "0" * 64 - with pytest.raises(ValueError, match="PNG hash"): - smoke.validate_inference_result(wrong_hash, root=tmp_path) - - reduced_resolution = copy.deepcopy(result) - reduced_resolution["height"] = 512 - with pytest.raises(ValueError, match="1024x1024"): - smoke.validate_inference_result(reduced_resolution, root=tmp_path) - - -def test_bundle_links_reject_cross_run_artifact_splicing() -> None: - stage1, stage2, manifest1, manifest2, export, snapshot = _bundle_link_fixture() - - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=export, - automodel_snapshot=snapshot, - ) - - wrong_checkpoint = copy.deepcopy(manifest2) - wrong_checkpoint["identity"]["model"]["revision"] = "a" * 40 - with pytest.raises(ValueError, match="identities differ"): - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=wrong_checkpoint, - export_manifest=export, - automodel_snapshot=snapshot, - ) - - wrong_export_identity = copy.deepcopy(export) - wrong_export_identity["identity"]["data"] = {"spliced": True} - with pytest.raises(ValueError, match="export identity"): - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=wrong_export_identity, - automodel_snapshot=snapshot, - ) - - wrong_export_checkpoint = copy.deepcopy(export) - wrong_export_checkpoint["source_checkpoint"]["manifest_sha256"] = "f" * 64 - with pytest.raises(ValueError, match="exact step-2"): - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=wrong_export_checkpoint, - automodel_snapshot=snapshot, - ) - - wrong_export_source = copy.deepcopy(export) - wrong_export_source["modelopt_source"]["commit"] = "e" * 40 - with pytest.raises(ValueError, match="training source"): - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=wrong_export_source, - automodel_snapshot=snapshot, - ) - - with pytest.raises(ValueError, match="AutoModel snapshot"): - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=export, - automodel_snapshot={}, - ) - - corrupt_snapshot = copy.deepcopy(snapshot) - corrupt_snapshot["files"][0]["sha256"] = "0" * 64 - with pytest.raises(ValueError, match="tree digest"): - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=export, - automodel_snapshot=corrupt_snapshot, - ) - - float_count_snapshot = copy.deepcopy(snapshot) - float_count_snapshot["package_file_count"] = float(smoke._AUTOMODEL_PACKAGE_FILE_COUNT) - with pytest.raises(ValueError, match="release identity"): - smoke._validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=export, - automodel_snapshot=float_count_snapshot, - ) - - -def test_automodel_snapshots_require_identical_bytes_and_no_symlinks(tmp_path: Path) -> None: - snapshot, automodel = _automodel_snapshot_fixture() - before = tmp_path / "before.json" - after = tmp_path / "after.json" - payload = json.dumps(snapshot, indent=2, sort_keys=True) + "\n" - before.write_text(payload) - after.write_text(payload) - assert ( - smoke._load_matching_automodel_snapshots( - before, - after, - expected_automodel=automodel, - ) - == snapshot - ) - - after.write_text(json.dumps(snapshot, sort_keys=True) + "\n") - with pytest.raises(ValueError, match="snapshot changed"): - smoke._load_matching_automodel_snapshots( - before, - after, - expected_automodel=automodel, - ) - - symlink = tmp_path / "before-symlink.json" - symlink.symlink_to(before) - with pytest.raises(ValueError, match="symlink"): - smoke._load_matching_automodel_snapshots( - symlink, - before, - expected_automodel=automodel, - ) - - -def test_smoke_artifact_paths_reject_symlinked_stage_inference_and_export(tmp_path: Path) -> None: - outside = tmp_path / "outside" - outside.mkdir() - target = outside / "artifact.json" - target.write_text("{}") - run_root = tmp_path / "run" - run_root.mkdir() - (run_root / "stage1.json").symlink_to(target) - inference = run_root / "inference" - inference.mkdir() - (inference / "pdd4.json").symlink_to(target) - (run_root / "export").symlink_to(outside, target_is_directory=True) - - with pytest.raises(ValueError, match="symlink"): - smoke._relative_regular_file(run_root, "stage1.json", name="stage") - with pytest.raises(ValueError, match="symlink"): - smoke._relative_regular_file(run_root, "inference/pdd4.json", name="inference") - with pytest.raises(ValueError, match="symlink"): - smoke._regular_directory(run_root / "export", name="export") - - target_parent = tmp_path / "target-parent" - target_parent.mkdir() - symlinked_parent = tmp_path / "symlinked-parent" - symlinked_parent.symlink_to(target_parent, target_is_directory=True) - requested_child = symlinked_parent / "must-not-be-created" - with pytest.raises(ValueError, match="symlink"): - smoke._create_run_root(requested_child) - assert not (target_parent / requested_child.name).exists() diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index b5ae655a6ae..a809b21f748 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -1,5 +1,17 @@ # 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. """Released-AutoModel seam tests for the ModelOpt-owned PDD setup.""" @@ -23,14 +35,22 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -import pdd_finetune -from pdd_recipe import ( +from pdd.recipe import ( build_pdd_export_setup, build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config, ) -from verify_readonly_automodel import snapshot_installed_distribution +from pdd.verify_readonly_automodel import snapshot_installed_distribution + + +def _require_exact_automodel() -> None: + try: + version = importlib.metadata.version("nemo_automodel") + except importlib.metadata.PackageNotFoundError: + pytest.skip("nemo_automodel is not installed") + if version != "0.5.0": + pytest.skip(f"requires the official nemo_automodel==0.5.0 wheel, found {version}") def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: @@ -78,177 +98,42 @@ def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: def test_example_recipe_explicitly_pins_grid_max_t() -> None: - raw = yaml.safe_load((_FASTGEN_DIR / "configs" / "pdd_qwen_image.yaml").read_text()) + raw = yaml.safe_load((_FASTGEN_DIR / "pdd" / "configs" / "qwen_image.yaml").read_text()) assert type(raw["pdd"]["grid_max_t"]) is float assert raw["pdd"]["grid_max_t"] == 0.999 - assert raw["data"]["expected_approved_ordered_ids_sha256"] is None - assert raw["data"]["expected_heldout_count"] == 2000 + assert raw["data"]["validation_count"] == 2000 + assert raw["data"]["split_seed"] == 2026 -def test_data_authentication_config_fields_are_strict_but_nullable(tmp_path) -> None: +def test_split_config_fields_are_strict(tmp_path) -> None: raw = _raw_config(tmp_path) - raw["data"] = { - "expected_approved_ordered_ids_sha256": "a" * 64, - "expected_heldout_count": 2000, - } + raw["data"] = {"validation_count": 3, "split_seed": 7} config = resolve_pdd_recipe_config(raw) - assert config.expected_approved_ordered_ids_sha256 == "a" * 64 - assert config.expected_heldout_count == 2000 + assert config.validation_count == 3 + assert config.split_seed == 7 - for invalid_hash in ("A" * 64, "a" * 63, 7): - raw["data"]["expected_approved_ordered_ids_sha256"] = invalid_hash - with pytest.raises(ValueError, match="expected_approved_ordered_ids_sha256"): - resolve_pdd_recipe_config(raw) - raw["data"]["expected_approved_ordered_ids_sha256"] = None for invalid_count in (0, -1, True, 1.5): - raw["data"]["expected_heldout_count"] = invalid_count - with pytest.raises(ValueError, match="expected_heldout_count"): + raw["data"]["validation_count"] = invalid_count + with pytest.raises((TypeError, ValueError), match="validation_count"): resolve_pdd_recipe_config(raw) - -@pytest.mark.parametrize("heldout_count", [1999, 2001]) -def test_canonical_training_count_cannot_be_overridden( - monkeypatch, tmp_path, heldout_count -) -> None: - raw = _raw_config(tmp_path) - raw["data"] = { - "expected_approved_ordered_ids_sha256": "a" * 64, - "expected_heldout_count": heldout_count, - } - config = resolve_pdd_recipe_config(raw) - called = False - - def _unexpected_validator(*args, **kwargs): - nonlocal called - called = True - raise AssertionError("validator must not be called for a weakened canonical count") - - monkeypatch.setattr("validate_cache_snapshot.validate_snapshot", _unexpected_validator) - with pytest.raises(ValueError, match="expected_heldout_count=2000"): - pdd_finetune._validate_dataset_snapshot(raw, config) - assert not called - - -def test_canonical_training_requires_external_hash_before_validator(tmp_path) -> None: - raw = _raw_config(tmp_path) - raw["data"] = { - "expected_approved_ordered_ids_sha256": None, - "expected_heldout_count": 2000, - } - config = resolve_pdd_recipe_config(raw) - with pytest.raises(ValueError, match="expected_approved_ordered_ids_sha256"): - pdd_finetune._validate_dataset_snapshot(raw, config) + raw["data"] = {"validation_count": 3, "split_seed": -1} + with pytest.raises(ValueError, match="split_seed"): + resolve_pdd_recipe_config(raw) -def test_valid_canonical_data_gate_reaches_validator_unchanged(monkeypatch, tmp_path) -> None: - expected_hash = "a" * 64 +def test_pdd_rejects_external_split_manifest(tmp_path) -> None: raw = _raw_config(tmp_path) - raw["data"] = { - "expected_approved_ordered_ids_sha256": expected_hash, - "expected_heldout_count": 2000, - "dataloader": { - "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", - "cache_dir": str(tmp_path), - }, - } - config = resolve_pdd_recipe_config(raw) - captured = None - - def _validator(root, **kwargs): - nonlocal captured - captured = (root, kwargs) - return {"snapshot_sha256": "b" * 64} - - def _all_gather_object(output, value): - output[0] = value - - monkeypatch.setattr("validate_cache_snapshot.validate_snapshot", _validator) - monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) - monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) - monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object) - monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda payload, src: None) - - assert pdd_finetune._validate_dataset_snapshot(raw, config) == {"snapshot_sha256": "b" * 64} - assert captured == ( - tmp_path.resolve(), - { - "all_index": config.all_metadata_index, - "train_index": config.train_metadata_index, - "heldout_index": config.validation_metadata_index, - "expected_approved_ids_sha256": expected_hash, - "expected_heldout_count": 2000, - }, - ) - - -def test_loader_order_must_match_authenticated_report() -> None: - train = [{"sample_id": "train-a"}, {"sample_id": "train-b"}] - heldout = [{"sample_id": "heldout-a"}] - report = { - "ordered_sample_ids_sha256": { - "train": pdd_finetune._ordered_id_sha256(train, split="train"), - "heldout": pdd_finetune._ordered_id_sha256(heldout, split="heldout"), - } - } - assert pdd_finetune._validated_loader_order_hashes(train, heldout, report) == ( - report["ordered_sample_ids_sha256"]["train"], - report["ordered_sample_ids_sha256"]["heldout"], - ) - with pytest.raises(RuntimeError, match="training loader order"): - pdd_finetune._validated_loader_order_hashes(list(reversed(train)), heldout, report) - - -def test_collective_loader_order_gate_rejects_rank_disagreement(monkeypatch) -> None: - train = [{"sample_id": "train-a"}, {"sample_id": "train-b"}] - heldout = [{"sample_id": "heldout-a"}] - report = { - "ordered_sample_ids_sha256": { - "train": pdd_finetune._ordered_id_sha256(train, split="train"), - "heldout": pdd_finetune._ordered_id_sha256(heldout, split="heldout"), - } - } - - def _all_gather_object(output, value): - output[:] = [value, {"ok": True, "hashes": ("c" * 64, value["hashes"][1])}] - - monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) - monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object) - with pytest.raises(RuntimeError, match="different authenticated"): - pdd_finetune._collective_validated_loader_order_hashes(train, heldout, report) - - -def test_two_rank_loader_order_divergence_fails_before_model_setup() -> None: - environment = os.environ.copy() - environment["PYTHONDONTWRITEBYTECODE"] = "1" - subprocess.run( - [ - sys.executable, - "-m", - "torch.distributed.run", - "--standalone", - "--nnodes=1", - "--nproc-per-node=2", - str( - _REPO_ROOT - / "tests" - / "examples" - / "diffusers" - / "fastgen" - / "pdd_training_preflight_distributed.py" - ), - ], - cwd=_REPO_ROOT, - env=environment, - check=True, - timeout=60, - ) + raw["data"] = {"dataloader": {"metadata_index": "metadata_train.json"}} + with pytest.raises(ValueError, match="metadata_index is unsupported"): + resolve_pdd_recipe_config(raw) def test_pdd_finetune_namespace_module_help() -> None: environment = os.environ.copy() environment["PYTHONDONTWRITEBYTECODE"] = "1" result = subprocess.run( - [sys.executable, "-m", "examples.diffusers.fastgen.pdd_finetune", "--help"], + [sys.executable, "-m", "examples.diffusers.fastgen.pdd.finetune", "--help"], cwd=_REPO_ROOT, env=environment, check=True, @@ -341,11 +226,7 @@ def test_training_dataloader_modes_are_gated_during_resolution( def test_frozen_automodel_distribution_snapshot_is_stable() -> None: - try: - version = importlib.metadata.version("nemo_automodel") - except importlib.metadata.PackageNotFoundError: - pytest.skip("nemo_automodel is not installed") - assert version == "0.5.0" + _require_exact_automodel() before = snapshot_installed_distribution() after = snapshot_installed_distribution() @@ -361,11 +242,11 @@ def test_frozen_automodel_distribution_snapshot_is_stable() -> None: def test_exact_wheel_install_below_git_checkout_is_accepted(tmp_path) -> None: + _require_exact_automodel() try: distribution = importlib.metadata.distribution("nemo_automodel") except importlib.metadata.PackageNotFoundError: pytest.skip("nemo_automodel is not installed") - assert distribution.version == "0.5.0" checkout = tmp_path / "checkout" (checkout / ".git").mkdir(parents=True) @@ -384,7 +265,7 @@ def test_exact_wheel_install_below_git_checkout_is_accepted(tmp_path) -> None: subprocess.run( [ sys.executable, - str(_FASTGEN_DIR / "verify_readonly_automodel.py"), + str(_FASTGEN_DIR / "pdd" / "verify_readonly_automodel.py"), "snapshot", "--output", str(output), @@ -402,6 +283,7 @@ def test_exact_wheel_install_below_git_checkout_is_accepted(tmp_path) -> None: def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: + _require_exact_automodel() before = snapshot_installed_distribution() model_dir = create_tiny_qwen_image_pipeline_dir(tmp_path) initialize_pdd_distributed(backend="gloo", timeout_minutes=1) diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index 45f295206ee..207a50b0ab1 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -1,12 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 - -"""Hermetic direct-update, committed-cursor, and strict PDD resume evidence.""" +# +# 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. + +"""Tests for direct updates, committed cursors, and strict PDD resume.""" from __future__ import annotations import copy import hashlib +import importlib.metadata import json import math import pathlib @@ -26,15 +39,15 @@ sys.path.insert(0, str(pathlib.Path(__file__).parent)) from fastgen_data.replayable_sampler import ReplayableBatchSampler -from pdd_checkpoint import ( +from pdd.checkpoint import ( PDDCheckpointManager, build_pdd_checkpoint_identity, resolve_pdd_training_checkpoint, ) -from pdd_recipe import initialize_pdd_distributed +from pdd.recipe import initialize_pdd_distributed +from pdd.training import prepare_qwen_pdd_batch +from pdd.verify_readonly_automodel import snapshot_installed_distribution from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 -from pdd_training import prepare_qwen_pdd_batch -from verify_readonly_automodel import snapshot_installed_distribution def _released_sampler(sample_ids: tuple[str, ...]) -> ReplayableBatchSampler: @@ -340,6 +353,9 @@ def test_training_hard_aborts_for_teacher_gradient_zero_gradient_and_missing_cov def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) -> None: pytest.importorskip("nemo_automodel") + version = importlib.metadata.version("nemo_automodel") + if version != "0.5.0": + pytest.skip(f"requires the official nemo_automodel==0.5.0 wheel, found {version}") rng_module = pytest.importorskip("nemo_automodel.components.training.rng") if not torch.distributed.is_initialized(): initialize_pdd_distributed(backend="gloo", timeout_minutes=1) diff --git a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py index 128e8c44b93..83b7ce81480 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py @@ -1,5 +1,17 @@ # 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. """Deterministic logical-ID PDD held-out oracle tests.""" @@ -20,13 +32,13 @@ if str(pathlib.Path(__file__).parent) not in sys.path: sys.path.insert(0, str(pathlib.Path(__file__).parent)) -from pdd_test_utils import build_toy_lifecycle, make_batch -from pdd_training import ( +from pdd.training import ( build_pdd_validation_assignments, pdd_validation_noise, pdd_validation_support, run_pdd_validation, ) +from pdd_test_utils import build_toy_lifecycle, make_batch from modelopt.torch.fastgen import PDDConfig diff --git a/tests/examples/diffusers/fastgen/test_portable_cache.py b/tests/examples/diffusers/fastgen/test_portable_cache.py deleted file mode 100644 index b632fbd90c2..00000000000 --- a/tests/examples/diffusers/fastgen/test_portable_cache.py +++ /dev/null @@ -1,770 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import hashlib -import json -import pathlib -import shutil -import sys -import types - -import pytest - -torch = pytest.importorskip("torch") -pytest.importorskip("nemo_automodel") - -_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 ( - TextToImageDataset, - build_text_to_image_multiresolution_dataloader, - collate_fn_text_to_image, -) -from portable_cache import ( - DATASET_CACHE_ENV, - PORTABLE_SNAPSHOT_SCHEMA_VERSION, - PREPROCESS_STAGING_SCHEMA_VERSION, - audit_no_absolute_paths, - load_approved_sample_ids, - load_portable_metadata, - ordered_sample_ids_sha256, - resolve_cache_root, - resolve_negative_embedding, - select_pdd_holdout_ids, - sha256_file, -) -from validate_cache_snapshot import validate_snapshot - - -def _sample_id(source_ref: str, resolution: tuple[int, int]) -> str: - identity = ( - f"modelopt-fastgen-sample-v1\0qwen_image\0{source_ref}\0{resolution[0]}x{resolution[1]}" - ) - return hashlib.sha256(identity.encode()).hexdigest() - - -def _write_json(path: pathlib.Path, value) -> None: - path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n") - - -def _write_approved_manifest(path: pathlib.Path, sample_ids: list[str]) -> str: - digest = ordered_sample_ids_sha256(sample_ids) - _write_json( - path, - { - "schema_version": 1, - "ordered_sample_ids": sample_ids, - "ordered_sample_ids_sha256": digest, - }, - ) - return digest - - -def _make_id_only_snapshot( - root: pathlib.Path, - sample_ids: tuple[str, ...], - heldout_count: int, - *, - heldout_override: tuple[str, ...] | None = None, -) -> None: - root.mkdir() - (root / "payload.pt").write_bytes(b"placeholder-not-read-before-split-gates") - entries = [ - { - "sample_id": sample_id, - "cache_file": "payload.pt", - "payload_sha256": "0" * 64, - } - for sample_id in sample_ids - ] - _write_json(root / "metadata_shard_s0000.json", entries) - if heldout_override is None: - train_ids, heldout_ids = select_pdd_holdout_ids(sample_ids, heldout_count) - else: - heldout_ids = heldout_override - heldout_members = set(heldout_ids) - train_ids = tuple(sample_id for sample_id in sample_ids if sample_id not in heldout_members) - approved_digest = ordered_sample_ids_sha256(sample_ids) - policy = { - "schema_version": 1, - "algorithm": "sha256-domain-ranked", - "domain": "modelopt-pdd-holdout-v1", - "heldout_count": heldout_count, - "approved_ordered_ids_sha256": approved_digest, - } - for name, split, ids in ( - ("metadata.json", "all", sample_ids), - ("metadata_train.json", "train", train_ids), - ("metadata_heldout.json", "heldout", heldout_ids), - ): - _write_json( - root / name, - { - "schema_version": 2, - "split": split, - "total_items": len(ids), - "num_shards": 1, - "shards": ["metadata_shard_s0000.json"], - "sample_ids": list(ids), - "ordered_sample_ids_sha256": ordered_sample_ids_sha256(ids), - "split_policy": policy, - }, - ) - - -def _make_snapshot(root: pathlib.Path) -> dict[str, list[str]]: - root.mkdir() - entries = [] - for index, resolution in enumerate(((64, 64), (64, 64), (128, 64), (128, 64))): - source_ref = f"class/{chr(ord('a') + index)}.png" - sample_id = _sample_id(source_ref, resolution) - cache_ref = f"payloads/{sample_id}.pt" - payload_path = root / cache_ref - payload_path.parent.mkdir(exist_ok=True) - torch.save( - { - "latent": torch.full((2, 2, 2), index, dtype=torch.float32), - "prompt_embeds": torch.full((3, 4), index, dtype=torch.float32), - "prompt_embeds_mask": torch.ones(3, dtype=torch.long), - "crop_offset": (0, 0), - "prompt": f"prompt {index}", - "sample_id": sample_id, - "source_ref": source_ref, - }, - payload_path, - ) - entries.append( - { - "sample_id": sample_id, - "source_ref": source_ref, - "cache_file": cache_ref, - "payload_sha256": sha256_file(payload_path), - "bucket_resolution": list(resolution), - "original_resolution": list(resolution), - "bucket_id": f"bucket-{resolution[0]}", - "aspect_ratio": resolution[0] / resolution[1], - } - ) - - _write_json(root / "metadata_shard_s0000.json", entries) - all_ids = [entry["sample_id"] for entry in entries] - train_ids, heldout_ids = select_pdd_holdout_ids(all_ids, 1) - splits = {"all": all_ids, "train": list(train_ids), "heldout": list(heldout_ids)} - approved_digest = ordered_sample_ids_sha256(all_ids) - split_policy = { - "schema_version": 1, - "algorithm": "sha256-domain-ranked", - "domain": "modelopt-pdd-holdout-v1", - "heldout_count": 1, - "approved_ordered_ids_sha256": approved_digest, - } - for split, ids in splits.items(): - name = "metadata.json" if split == "all" else f"metadata_{split}.json" - _write_json( - root / name, - { - "schema_version": 2, - "split": split, - "total_items": len(ids), - "num_shards": 1, - "shards": ["metadata_shard_s0000.json"], - "sample_ids": ids, - "ordered_sample_ids_sha256": ordered_sample_ids_sha256(ids), - "split_policy": split_policy, - }, - ) - negative_path = root / "negative_prompt_embedding.pt" - torch.save({"embed": torch.arange(12).reshape(3, 4)}, negative_path) - negative_declaration = { - "path": negative_path.name, - "sha256": sha256_file(negative_path), - } - for name in ("metadata.json", "metadata_train.json", "metadata_heldout.json"): - index = json.loads((root / name).read_text()) - index["negative_prompt_embedding"] = negative_declaration - _write_json(root / name, index) - return splits - - -def _batch_signature(dataset: TextToImageDataset) -> list[tuple[str, float]]: - return [ - (dataset[index]["sample_id"], dataset[index]["latent"].sum().item()) - for index in range(len(dataset)) - ] - - -def test_authenticated_cache_constants_hash_framing_and_seedless_split() -> None: - assert PREPROCESS_STAGING_SCHEMA_VERSION == 1 - assert PORTABLE_SNAPSHOT_SCHEMA_VERSION == 2 - expected = { - ("a\nb",): "41e07cc133e8a85fc4a08e60a38c223f3c24dbca80312d106f251e533254eedf", - ("a", "b"): "8cf774af4e8509811c2d4bc2adec6b852e4c614f9d8d833924502ead7c0689d7", - ("ab", "c"): "6df9e72da4c55f09b4c0320337d6a5d46396271ac61b7a60e1ee8146ce49709e", - ("a", "bc"): "7cedefc9d46613683a89c3081c3a743b66164861975cf100346d59c13cf31d26", - tuple( - str(index) for index in range(16) - ): "b157f73e9710fe1eb2c4f8d94286f304d5c2a9de2b09b31d2b1f5eee15448e69", - } - for sample_ids, digest in expected.items(): - assert ordered_sample_ids_sha256(sample_ids) == digest - assert len(set(expected.values())) == len(expected) - - train, heldout = select_pdd_holdout_ids(tuple(str(index) for index in range(16)), 4) - assert heldout == ("4", "7", "10", "13") - assert train == tuple(str(index) for index in range(16) if str(index) not in heldout) - assert heldout not in (("0", "2", "4", "11"), ("2", "7", "8", "9")) - - permuted = tuple(reversed(tuple(str(index) for index in range(16)))) - permuted_train, permuted_heldout = select_pdd_holdout_ids(permuted, 4) - assert set(permuted_heldout) == set(heldout) - assert permuted_train == tuple(item for item in permuted if item not in set(heldout)) - assert permuted_heldout == tuple(item for item in permuted if item in set(heldout)) - - -def test_approved_id_artifact_is_strict_and_externally_authenticatable(tmp_path) -> None: - path = tmp_path / "approved.json" - digest = _write_approved_manifest(path, ["second", "first"]) - assert load_approved_sample_ids(path, expected_sha256=digest) == ( - ("second", "first"), - digest, - ) - - for value, message in ( - ({"schema_version": 1}, "keys mismatch"), - ( - { - "schema_version": 1, - "ordered_sample_ids": ["first", "first"], - "ordered_sample_ids_sha256": digest, - }, - "duplicates", - ), - ): - _write_json(path, value) - with pytest.raises(ValueError, match=message): - load_approved_sample_ids(path) - - path.write_text('{"schema_version":1,"schema_version":1}') - with pytest.raises(ValueError, match="duplicate key"): - load_approved_sample_ids(path) - for constant in ("NaN", "Infinity", "-Infinity"): - path.write_text(f'{{"schema_version": {constant}}}') - with pytest.raises(ValueError, match="non-standard constant"): - load_approved_sample_ids(path) - path.write_bytes(b"\xff") - with pytest.raises(ValueError, match="UTF-8 JSON"): - load_approved_sample_ids(path) - path.write_text("{") - with pytest.raises(ValueError, match="UTF-8 JSON"): - load_approved_sample_ids(path) - - real = tmp_path / "real.json" - _write_approved_manifest(real, ["first", "second"]) - path.unlink() - path.symlink_to(real) - with pytest.raises(ValueError, match="symlink"): - load_approved_sample_ids(path) - with pytest.raises(ValueError, match="regular file"): - load_approved_sample_ids(tmp_path) - with pytest.raises(ValueError, match="expected approved-ID"): - load_approved_sample_ids(real, expected_sha256="f" * 64) - - -def test_approved_hash_is_identical_to_finalized_all_index_hash(tmp_path) -> None: - root = tmp_path / "cache" - splits = _make_snapshot(root) - approved_path = tmp_path / "approved.json" - approved_digest = _write_approved_manifest(approved_path, splits["all"]) - _, loaded_digest = load_approved_sample_ids( - approved_path, - expected_sha256=approved_digest, - ) - all_index, _ = load_portable_metadata(root) - assert loaded_digest == all_index["ordered_sample_ids_sha256"] - assert loaded_digest == ordered_sample_ids_sha256(all_index["sample_ids"]) - assert loaded_digest == all_index["split_policy"]["approved_ordered_ids_sha256"] - - -@pytest.mark.parametrize( - "old_heldout", - [("0", "2", "4", "11"), ("2", "7", "8", "9")], -) -def test_self_consistent_old_seed_partitions_are_rejected(tmp_path, old_heldout) -> None: - root = tmp_path / "cache" - sample_ids = tuple(str(index) for index in range(16)) - _make_id_only_snapshot(root, sample_ids, 4, heldout_override=old_heldout) - with pytest.raises(ValueError, match="frozen PDD policy"): - validate_snapshot(root, reject_orphans=False) - - -def test_self_consistent_all_list_rewrite_fails_external_hash(tmp_path) -> None: - approved_ids = tuple(str(index) for index in range(16)) - rewritten_ids = approved_ids[:-1] - root = tmp_path / "cache" - _make_id_only_snapshot(root, rewritten_ids, 4) - with pytest.raises(ValueError, match="expected approved ordered-ID hash"): - validate_snapshot( - root, - expected_approved_ids_sha256=ordered_sample_ids_sha256(approved_ids), - reject_orphans=False, - ) - - -@pytest.mark.parametrize( - ("mutation", "message"), - [ - ("order", "frozen PDD policy"), - ("gap", "split union"), - ("duplicate", "contains duplicates"), - ("member", "frozen PDD policy"), - ], -) -def test_split_order_gap_duplicate_and_membership_tampering_is_rejected( - tmp_path, mutation, message -) -> None: - root = tmp_path / "cache" - sample_ids = tuple(str(index) for index in range(16)) - _make_id_only_snapshot(root, sample_ids, 4) - train_path = root / "metadata_train.json" - heldout_path = root / "metadata_heldout.json" - train = json.loads(train_path.read_text()) - heldout = json.loads(heldout_path.read_text()) - - if mutation == "order": - train["sample_ids"] = list(reversed(train["sample_ids"])) - elif mutation == "gap": - train["sample_ids"].pop() - elif mutation == "duplicate": - train["sample_ids"].append(train["sample_ids"][0]) - else: - train["sample_ids"][0], heldout["sample_ids"][0] = ( - heldout["sample_ids"][0], - train["sample_ids"][0], - ) - - for path, index in ((train_path, train), (heldout_path, heldout)): - index["total_items"] = len(index["sample_ids"]) - if len(index["sample_ids"]) == len(set(index["sample_ids"])): - index["ordered_sample_ids_sha256"] = ordered_sample_ids_sha256(index["sample_ids"]) - _write_json(path, index) - with pytest.raises(ValueError, match=message): - validate_snapshot(root, reject_orphans=False) - - -@pytest.mark.parametrize( - ("policy_update", "message"), - [ - ({"domain": "modelopt-fastgen-split-v1"}, "domain"), - ({"algorithm": "other"}, "algorithm"), - ({"schema_version": 2}, "schema_version"), - ({"heldout_count": True}, "heldout_count"), - ({"extra": "field"}, "keys mismatch"), - ], -) -def test_split_policy_tampering_is_rejected(tmp_path, policy_update, message) -> None: - root = tmp_path / "cache" - _make_snapshot(root) - for name in ("metadata.json", "metadata_train.json", "metadata_heldout.json"): - index = json.loads((root / name).read_text()) - index["split_policy"].update(policy_update) - _write_json(root / name, index) - with pytest.raises(ValueError, match=message): - validate_snapshot(root) - - -def test_strict_index_and_shard_json_reject_duplicates_and_nonfinite(tmp_path) -> None: - root = tmp_path / "cache" - _make_snapshot(root) - index_path = root / "metadata.json" - index_path.write_text('{"schema_version":2,"schema_version":2}') - with pytest.raises(ValueError, match="duplicate key"): - load_portable_metadata(root) - - _make_snapshot(tmp_path / "second") - shard_path = tmp_path / "second" / "metadata_shard_s0000.json" - shard_path.write_text('[{"sample_id":"a","value":NaN}]') - with pytest.raises(ValueError, match="non-standard constant"): - load_portable_metadata(tmp_path / "second") - - nested_policy = tmp_path / "nested-policy" - _make_snapshot(nested_policy) - nested_index_path = nested_policy / "metadata.json" - nested_index = nested_index_path.read_text() - domain = '"domain": "modelopt-pdd-holdout-v1"' - nested_index_path.write_text(nested_index.replace(domain, f"{domain},\n {domain}", 1)) - with pytest.raises(ValueError, match="duplicate key"): - load_portable_metadata(nested_policy) - - nested_shard = tmp_path / "nested-shard" - splits = _make_snapshot(nested_shard) - nested_shard_path = nested_shard / "metadata_shard_s0000.json" - nested_entries = nested_shard_path.read_text() - sample_id = f'"sample_id": "{splits["all"][0]}"' - nested_shard_path.write_text( - nested_entries.replace(sample_id, f"{sample_id},\n {sample_id}", 1) - ) - with pytest.raises(ValueError, match="duplicate key"): - load_portable_metadata(nested_shard) - - -@pytest.mark.parametrize("actual_heldout_count", [1999, 2001]) -def test_validator_rejects_noncanonical_actual_holdout_counts( - tmp_path, actual_heldout_count -) -> None: - root = tmp_path / "cache" - sample_ids = tuple(str(index) for index in range(actual_heldout_count + 1)) - _make_id_only_snapshot(root, sample_ids, actual_heldout_count) - with pytest.raises(ValueError, match="expected_heldout_count"): - validate_snapshot( - root, - expected_approved_ids_sha256=ordered_sample_ids_sha256(sample_ids), - expected_heldout_count=2000, - reject_orphans=False, - ) - - -def test_cache_root_environment_precedence(monkeypatch, tmp_path): - configured = tmp_path / "configured" - override = tmp_path / "override" - configured.mkdir() - override.mkdir() - - monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) - assert resolve_cache_root(configured) == configured.resolve() - monkeypatch.setenv(DATASET_CACHE_ENV, "") - assert resolve_cache_root(configured) == configured.resolve() - monkeypatch.setenv(DATASET_CACHE_ENV, str(override)) - assert resolve_cache_root(configured) == override.resolve() - monkeypatch.setenv(DATASET_CACHE_ENV, "relative/cache") - with pytest.raises(ValueError, match="absolute"): - resolve_cache_root(configured) - monkeypatch.setenv(DATASET_CACHE_ENV, str(tmp_path / "missing")) - with pytest.raises(FileNotFoundError): - resolve_cache_root(configured) - not_a_directory = tmp_path / "cache-file" - not_a_directory.write_text("not a directory") - monkeypatch.setenv(DATASET_CACHE_ENV, str(not_a_directory)) - with pytest.raises(NotADirectoryError): - resolve_cache_root(configured) - - -def test_relocation_preserves_order_payloads_and_buckets(monkeypatch, tmp_path): - first = tmp_path / "alice" / "cache" - second = tmp_path / "bob" / "cache" - first.parent.mkdir() - splits = _make_snapshot(first) - second.parent.mkdir() - shutil.copytree(first, second) - - monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) - first_dataset = TextToImageDataset(str(first), metadata_index="metadata_train.json") - second_dataset = TextToImageDataset(str(second), metadata_index="metadata_train.json") - assert _batch_signature(first_dataset) == _batch_signature(second_dataset) - assert [entry["sample_id"] for entry in second_dataset.metadata] == splits["train"] - assert first_dataset.bucket_groups == second_dataset.bucket_groups - first_report = validate_snapshot(first) - second_report = validate_snapshot(second) - assert first_report["snapshot_sha256"] == second_report["snapshot_sha256"] - assert first_report["declared_files"] == second_report["declared_files"] - - monkeypatch.setenv(DATASET_CACHE_ENV, str(second)) - overridden = TextToImageDataset(str(first), metadata_index="metadata_train.json") - assert overridden.cache_dir == second.resolve() - assert _batch_signature(overridden) == _batch_signature(first_dataset) - - -def test_pdd_loader_iterator_does_not_advance_training_rng(monkeypatch, tmp_path): - root = tmp_path / "cache" - _make_snapshot(root) - monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) - loader, _ = build_text_to_image_multiresolution_dataloader( - cache_dir=str(root), - metadata_index="metadata_train.json", - batch_size=1, - base_resolution=(64, 64), - num_workers=0, - exact_resume=True, - sampler_seed=17, - loader_seed=17, - ) - before = torch.get_rng_state().clone() - - next(iter(loader)) - - assert torch.equal(torch.get_rng_state(), before) - - -def test_split_filters_before_inherited_bucket_grouping(monkeypatch, tmp_path): - root = tmp_path / "cache" - splits = _make_snapshot(root) - shard_path = root / "metadata_shard_s0000.json" - entries = json.loads(shard_path.read_text()) - heldout_entry = next(entry for entry in entries if entry["sample_id"] in splits["heldout"]) - del heldout_entry["bucket_resolution"] - _write_json(shard_path, entries) - - monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) - dataset = TextToImageDataset(str(root), metadata_index="metadata_train.json") - assert [entry["sample_id"] for entry in dataset.metadata] == splits["train"] - grouped = [index for bucket in dataset.bucket_groups.values() for index in bucket["indices"]] - assert sorted(grouped) == list(range(len(splits["train"]))) - - -@pytest.mark.parametrize( - "reference", - ["/etc/passwd", "C:/secret.pt", "C:secret.pt", "../escape.pt", "a/../b.pt"], -) -def test_manifest_rejects_nonportable_payload_references(monkeypatch, tmp_path, reference): - root = tmp_path / "cache" - _make_snapshot(root) - shard_path = root / "metadata_shard_s0000.json" - entries = json.loads(shard_path.read_text()) - entries[0]["cache_file"] = reference - _write_json(shard_path, entries) - - monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) - with pytest.raises((ValueError, FileNotFoundError)): - load_portable_metadata(root, "metadata.json") - - -def test_manifest_rejects_missing_and_symlink_escape(monkeypatch, tmp_path): - root = tmp_path / "cache" - _make_snapshot(root) - outside = tmp_path / "outside.pt" - torch.save({}, outside) - shard_path = root / "metadata_shard_s0000.json" - original = json.loads(shard_path.read_text()) - - monkeypatch.delenv(DATASET_CACHE_ENV, raising=False) - for reference in ("payloads/missing.pt", "payloads/escape.pt"): - entries = json.loads(json.dumps(original)) - entries[0]["cache_file"] = reference - if reference.endswith("escape.pt"): - (root / reference).symlink_to(outside) - _write_json(shard_path, entries) - with pytest.raises((ValueError, FileNotFoundError)): - load_portable_metadata(root, "metadata.json") - - -def test_negative_embedding_is_resolved_under_effective_root(tmp_path): - root = tmp_path / "cache" - _make_snapshot(root) - assert ( - resolve_negative_embedding(root, "negative_prompt_embedding.pt") - == (root / "negative_prompt_embedding.pt").resolve() - ) - assert resolve_negative_embedding(root, root / "negative_prompt_embedding.pt").is_file() - - outside = tmp_path / "outside.pt" - torch.save({}, outside) - with pytest.raises(ValueError, match="outside"): - resolve_negative_embedding(root, outside) - (root / "negative_escape.pt").symlink_to(outside) - with pytest.raises(ValueError, match="outside"): - resolve_negative_embedding(root, "negative_escape.pt") - - -def test_collate_emits_logical_identity_without_source_paths(tmp_path): - root = tmp_path / "cache" - _make_snapshot(root) - dataset = TextToImageDataset(str(root), metadata_index="metadata_train.json") - same_resolution = next( - group["indices"] for group in dataset.bucket_groups.values() if len(group["indices"]) >= 2 - ) - samples = [dataset[index] for index in same_resolution[:2]] - output = collate_fn_text_to_image(samples) - assert output["metadata"]["sample_ids"] == [item["sample_id"] for item in samples] - assert output["metadata"]["source_refs"] == [item["source_ref"] for item in samples] - assert "image_paths" not in output["metadata"] - assert str(root) not in repr(output) - - -def test_validator_detects_hash_split_and_orphan_failures(tmp_path): - root = tmp_path / "cache" - splits = _make_snapshot(root) - assert validate_snapshot(root)["unique_payloads"] == 4 - - orphan = root / "payloads" / "orphan.pt" - torch.save({"image_path": "/home/alice/private.png"}, orphan) - with pytest.raises(ValueError, match="undeclared"): - validate_snapshot(root) - orphan.unlink() - - heldout_path = root / "metadata_heldout.json" - heldout = json.loads(heldout_path.read_text()) - heldout["sample_ids"] = [splits["train"][0]] - heldout["ordered_sample_ids_sha256"] = ordered_sample_ids_sha256(heldout["sample_ids"]) - _write_json(heldout_path, heldout) - with pytest.raises(ValueError, match="overlap"): - validate_snapshot(root) - - -def test_validator_requires_complete_splits_and_rejects_directory_symlink(tmp_path): - root = tmp_path / "cache" - _make_snapshot(root) - (root / "metadata_heldout.json").unlink() - with pytest.raises(FileNotFoundError, match="required metadata indices"): - validate_snapshot(root) - - _make_snapshot(tmp_path / "complete") - complete = tmp_path / "complete" - outside = tmp_path / "outside-directory" - outside.mkdir() - (complete / "escaped-directory").symlink_to(outside, target_is_directory=True) - with pytest.raises(ValueError, match="symlink resolves outside"): - validate_snapshot(complete) - - -def test_recursive_path_audit_is_schema_aware_and_covers_containers(): - audit_no_absolute_paths({"prompt": "/imagine a cat"}) - with pytest.raises(ValueError, match="absolute path"): - audit_no_absolute_paths({"paths": {"/home/alice/private.png"}}) - with pytest.raises(ValueError, match="absolute path"): - audit_no_absolute_paths({"paths": {"primary": "/home/alice/private.png"}}) - with pytest.raises(ValueError, match="absolute path"): - audit_no_absolute_paths({"source_file": {"value": "/lustre/private.pt"}}) - with pytest.raises(ValueError, match="absolute path key"): - audit_no_absolute_paths({"/home/alice/private.png": "value"}) - - -def test_validator_detects_payload_hash_mismatch_and_is_read_only(tmp_path): - root = tmp_path / "cache" - _make_snapshot(root) - before = { - path.relative_to(root).as_posix(): (sha256_file(path), path.stat().st_mtime_ns) - for path in root.rglob("*") - if path.is_file() - } - validate_snapshot(root) - after = { - path.relative_to(root).as_posix(): (sha256_file(path), path.stat().st_mtime_ns) - for path in root.rglob("*") - if path.is_file() - } - assert after == before - - first_payload = next((root / "payloads").glob("*.pt")) - first_payload.write_bytes(first_payload.read_bytes() + b"tampered") - with pytest.raises(ValueError, match="SHA-256 mismatch"): - validate_snapshot(root) - - -def test_empty_split_is_rejected_before_bucket_grouping(tmp_path): - root = tmp_path / "cache" - _make_snapshot(root) - _write_json( - root / "metadata_empty.json", - { - "schema_version": 2, - "split": "empty", - "total_items": 0, - "shards": ["metadata_shard_s0000.json"], - "sample_ids": [], - "ordered_sample_ids_sha256": "0" * 64, - "split_policy": json.loads((root / "metadata.json").read_text())["split_policy"], - }, - ) - with pytest.raises(ValueError, match="non-empty"): - TextToImageDataset(str(root), metadata_index="metadata_empty.json") - - -def test_qwen_preprocessor_payload_is_sanitized(): - from preprocess.processors.qwen_image import QwenImageProcessor - - processor = QwenImageProcessor() - metadata = { - "original_resolution": (64, 64), - "bucket_resolution": (64, 64), - "crop_offset": (0, 0), - "prompt": "portable", - "sample_id": "b" * 64, - "source_ref": "class/b.png", - "bucket_id": "square-64", - "aspect_ratio": 1.0, - } - payload = processor.get_cache_data( - torch.zeros(2, 2, 2), - {"prompt_embeds": torch.zeros(1, 3, 4)}, - metadata, - ) - assert payload["sample_id"] == metadata["sample_id"] - assert payload["source_ref"] == metadata["source_ref"] - assert "image_path" not in payload - - -def test_portable_index_writer_is_deterministic(monkeypatch, tmp_path): - try: - import cv2 # noqa: F401 - except ImportError: - monkeypatch.setitem(sys.modules, "cv2", types.ModuleType("cv2")) - from migrate_cache_manifest import migrate_cache - from preprocess.preprocessing_multiprocess import _save_metadata_shards - - staging = tmp_path / "staging" - (staging / "payloads").mkdir(parents=True) - entries = [] - for character in ("b", "a"): - source_ref = f"class/{character}.png" - sample_id = _sample_id(source_ref, (64, 64)) - payload_ref = f"payloads/{sample_id}.pt" - payload_path = staging / payload_ref - torch.save( - { - "latent": torch.zeros(2, 2), - "crop_offset": (0, 0), - "prompt": character, - "sample_id": sample_id, - "source_ref": source_ref, - }, - payload_path, - ) - entries.append( - { - "sample_id": sample_id, - "source_ref": source_ref, - "cache_file": payload_ref, - "payload_sha256": sha256_file(payload_path), - "bucket_resolution": [64, 64], - "original_resolution": [64, 64], - "prompt": character, - "bucket_id": "square-64", - "aspect_ratio": 1.0, - "model_type": "qwen_image", - } - ) - _save_metadata_shards( - entries, - staging, - "qwen_image", - "Qwen/Qwen-Image", - "qwen_image", - 10, - {}, - portable=True, - ) - index = json.loads((staging / "metadata.json").read_text()) - shard = json.loads((staging / index["shards"][0]).read_text()) - assert index["schema_version"] == PREPROCESS_STAGING_SCHEMA_VERSION - expected_ids = sorted(entry["sample_id"] for entry in entries) - assert index["sample_ids"] == expected_ids - assert [entry["sample_id"] for entry in shard] == expected_ids - assert not (staging / "metadata_train.json").exists() - assert str(staging) not in json.dumps([index, shard]) - with pytest.raises(ValueError, match=r"migrate_cache_manifest\.py"): - load_portable_metadata(staging, "metadata.json") - - finalized = tmp_path / "finalized" - approved = tmp_path / "approved.json" - _write_approved_manifest(approved, expected_ids) - migrate_cache( - staging, - finalized, - approved_ids_manifest=approved, - heldout_count=1, - ) - assert validate_snapshot(finalized)["splits"] == {"all": 2, "train": 1, "heldout": 1} diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index 633cbd7b668..8e546110b0c 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -180,8 +180,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", - "sample_id": "sample-0", - "source_ref": "images/img.png", + "image_path": "/source/image.png", "bucket_id": 0, "aspect_ratio": 1.0, "prompt_embeds": torch.randn(seq, dim), diff --git a/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py b/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py deleted file mode 100644 index e3644e0f428..00000000000 --- a/tests/gpu/torch/fastgen/pdd_fsdp2_smoke.py +++ /dev/null @@ -1,431 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Two-rank CUDA FSDP2 optimization and exact-resume proof for plain PDD modules.""" - -from __future__ import annotations - -import gc -import importlib.util -import json -import os -import pathlib -import shutil -import sys -import tempfile -from dataclasses import dataclass -from typing import Any - -import torch -import torch.distributed as dist -import torch.distributed.checkpoint as dcp -from torch import nn -from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict -from torch.distributed.fsdp import fully_shard - -from modelopt.torch.fastgen import ( - PDDConfig, - PDDLayerSpec, - PDDOutputProjection, - PDDPipeline, - convert_to_pdd_output_projection, -) - -_FORBIDDEN_MODULES = ("diffusers", "fastgen", "nemo_automodel") -_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) -> torch.Tensor: - return self.projection(torch.tanh(self.backbone(state))) - - -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, - condition: torch.Tensor, - ) -> torch.Tensor: - return self.projection(state) + 0.125 * time[:, None] + 0.05 * condition - - -class _GuidedAdapter: - def __init__(self) -> None: - self.teacher_calls = 0 - - @staticmethod - def _dtype(model: nn.Module) -> torch.dtype: - return next(model.parameters()).dtype - - def student_all_heads( - self, - model: nn.Module, - state: torch.Tensor, - time: torch.Tensor, - *, - condition: Any = None, - **model_kwargs: Any, - ) -> torch.Tensor: - del time, condition, model_kwargs - output = model(state.to(self._dtype(model))) - return output.reshape(state.shape[0], _GRID_SIZE, _WIDTH) - - 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: - del time, condition, model_kwargs - projection = model.get_submodule("projection") - assert isinstance(projection, PDDOutputProjection) - with projection.fuse_block(start, end, grid): - return model(state.to(self._dtype(model))) - - 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: - del model_kwargs - if not isinstance(condition, torch.Tensor) or not isinstance( - negative_condition, torch.Tensor - ): - raise TypeError("guided toy teacher requires tensor conditions") - dtype = self._dtype(model) - state = state.to(dtype) - time = time.to(dtype) - conditional = model(state, time, condition.to(dtype)) - unconditional = model(state, time, negative_condition.to(dtype)) - self.teacher_calls += 2 - return conditional + 3.0 * (conditional - unconditional) - - -@dataclass -class _Lifecycle: - student: nn.Module - teacher: nn.Module - projection: PDDOutputProjection - pipeline: PDDPipeline - optimizer: torch.optim.AdamW - adapter: _GuidedAdapter - - -def _local(value: torch.Tensor) -> torch.Tensor: - to_local = getattr(value, "to_local", None) - return to_local() if callable(to_local) else value - - -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 + index * 0.01, - 0.2 + offset + index * 0.01, - parameter.numel(), - dtype=torch.float32, - device=parameter.device, - ) - parameter.copy_(values.reshape_as(parameter).to(parameter.dtype)) - - -def _config() -> 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], - student_sample_steps=2, - guidance_scale=4.0, - ) - - -def _build(device: torch.device) -> _Lifecycle: - config = _config() - student = _Student().to(device=device, dtype=torch.bfloat16) - teacher = _Teacher().to(device=device, dtype=torch.bfloat16).eval().requires_grad_(False) - _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, - ) - projection_module_id = id(projection) - projection_shape = projection.weight.shape - student = fully_shard(student) - teacher = fully_shard(teacher) - assert id(student.get_submodule("projection")) == projection_module_id - assert student.get_submodule("projection").weight.shape == projection_shape - optimizer = torch.optim.AdamW( - student.parameters(), - lr=2.0e-3, - weight_decay=0.0, - foreach=False, - fused=False, - ) - optimizer_parameters = [ - parameter for group in optimizer.param_groups for parameter in group["params"] - ] - assert any(parameter is projection.weight for parameter in optimizer_parameters) - adapter = _GuidedAdapter() - pipeline = PDDPipeline(student, teacher, config, adapter) - return _Lifecycle(student, teacher, projection, pipeline, optimizer, adapter) - - -def _batch( - *, rank: int, step: int, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - base = torch.arange(_WIDTH, device=device, dtype=torch.float32).reshape(1, -1) - data = (base / 10 + 0.05 * rank + 0.025 * step).to(torch.bfloat16) - noise = 0.4 - base / 20 + 0.01 * rank - condition = 0.2 + base / 30 + 0.02 * step - negative = -0.1 - base / 40 - 0.01 * rank - n = torch.tensor([0 if step == 1 else 2], device=device, dtype=torch.int64) - k = torch.tensor([1 if step == 1 else 3], device=device, dtype=torch.int64) - return data, noise, condition, negative, n, k - - -def _global_norm(values: list[torch.Tensor], device: torch.device) -> torch.Tensor: - squared = torch.zeros((), device=device, dtype=torch.float64) - for value in values: - squared += _local(value.detach()).float().square().sum(dtype=torch.float64) - dist.all_reduce(squared, op=dist.ReduceOp.SUM) - return squared.sqrt() - - -def _step(lifecycle: _Lifecycle, *, rank: int, step: int, device: torch.device) -> dict[str, float]: - data, noise, condition, negative, n, k = _batch(rank=rank, step=step, device=device) - lifecycle.optimizer.zero_grad(set_to_none=True) - calls_before = lifecycle.adapter.teacher_calls - loss, metrics = lifecycle.pipeline.compute_loss( - data, - noise=noise, - condition=condition, - negative_condition=negative, - n=n, - k=k, - ) - 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 lifecycle.adapter.teacher_calls - calls_before == 2 - assert all(parameter.grad is None for parameter in lifecycle.teacher.parameters()) - gradients = [ - parameter.grad for parameter in lifecycle.student.parameters() if parameter.grad is not None - ] - grad_norm = _global_norm(gradients, device) - assert torch.isfinite(grad_norm) and grad_norm > 0 - before = { - name: _local(parameter.detach()).clone() - for name, parameter in lifecycle.student.named_parameters() - } - lifecycle.optimizer.step() - updates = [ - _local(parameter.detach()) - before[name] - for name, parameter in lifecycle.student.named_parameters() - ] - update_norm = _global_norm(updates, device) - assert torch.isfinite(update_norm) and update_norm > 0 - reduced_loss = loss.detach().double() - dist.all_reduce(reduced_loss, op=dist.ReduceOp.SUM) - reduced_loss /= dist.get_world_size() - return { - "loss": float(reduced_loss.item()), - "grad_norm": float(grad_norm.item()), - "update_norm": float(update_norm.item()), - } - - -def _state(model: nn.Module) -> dict[str, torch.Tensor]: - return { - name: _local(value.detach()).clone() - for name, value in model.state_dict().items() - if isinstance(value, torch.Tensor) - } - - -def _assert_state_equal(actual: nn.Module, expected: dict[str, torch.Tensor]) -> None: - actual_state = _state(actual) - assert actual_state.keys() == expected.keys() - for name in expected: - torch.testing.assert_close(actual_state[name], expected[name], rtol=0, atol=0) - - -def _all_rng_states(device: torch.device) -> dict[str, torch.Tensor]: - local_state = { - "cpu": torch.get_rng_state(), - "cuda": torch.cuda.get_rng_state(device), - } - gathered: list[dict[str, torch.Tensor] | None] = [None] * dist.get_world_size() - dist.all_gather_object(gathered, local_state) - states: dict[str, torch.Tensor] = {} - for rank, value in enumerate(gathered): - assert value is not None - states[f"cpu_rng_rank_{rank}"] = value["cpu"].to(device) - states[f"cuda_rng_rank_{rank}"] = value["cuda"].to(device) - return states - - -def _save_checkpoint( - root: pathlib.Path, - lifecycle: _Lifecycle, - *, - completed_steps: int, - device: torch.device, -) -> pathlib.Path: - staging = root / ".step_00000001.staging" - final = root / "step_00000001" - if dist.get_rank() == 0: - staging.mkdir(parents=True) - dist.barrier() - model_state, optimizer_state = get_state_dict(lifecycle.student, lifecycle.optimizer) - extra = _all_rng_states(device) - extra["completed_steps"] = torch.tensor([completed_steps], device=device, dtype=torch.int64) - dcp.save( - {"model": model_state, "optimizer": optimizer_state, "extra": extra}, - checkpoint_id=staging, - ) - dist.barrier() - if dist.get_rank() == 0: - assert (staging / ".metadata").is_file() - assert any(path.suffix == ".distcp" for path in staging.iterdir()) - os.replace(staging, final) - (final / "COMPLETE").write_text( - json.dumps({"schema_version": 1, "completed_steps": completed_steps}) + "\n" - ) - dist.barrier() - assert (final / "COMPLETE").is_file() - return final - - -def _load_checkpoint( - checkpoint: pathlib.Path, - lifecycle: _Lifecycle, - *, - device: torch.device, -) -> int: - marker = json.loads((checkpoint / "COMPLETE").read_text()) - assert marker == {"schema_version": 1, "completed_steps": 1} - model_state, optimizer_state = get_state_dict(lifecycle.student, lifecycle.optimizer) - extra = _all_rng_states(device) - extra["completed_steps"] = torch.zeros(1, device=device, dtype=torch.int64) - payload = {"model": model_state, "optimizer": optimizer_state, "extra": extra} - dcp.load(payload, checkpoint_id=checkpoint) - incompatible = set_state_dict( - lifecycle.student, - lifecycle.optimizer, - model_state_dict=payload["model"], - optim_state_dict=payload["optimizer"], - ) - assert incompatible.missing_keys == [] - assert incompatible.unexpected_keys == [] - rank = dist.get_rank() - torch.set_rng_state(payload["extra"][f"cpu_rng_rank_{rank}"].cpu()) - torch.cuda.set_rng_state(payload["extra"][f"cuda_rng_rank_{rank}"].cpu(), device) - return int(payload["extra"]["completed_steps"].item()) - - -def _assert_call_counts(adapter: _GuidedAdapter, expected: int, device: torch.device) -> None: - value = torch.tensor([adapter.teacher_calls], device=device, dtype=torch.int64) - gathered = [torch.zeros_like(value) for _ in range(dist.get_world_size())] - dist.all_gather(gathered, value) - assert [int(item.item()) for item in gathered] == [expected] * dist.get_world_size() - - -def _assert_optional_frameworks_absent() -> None: - resolvable = sorted(name for name in _FORBIDDEN_MODULES if importlib.util.find_spec(name)) - assert not resolvable, f"plain PDD FSDP2 environment resolves optional frameworks: {resolvable}" - imported = sorted(name for name in _FORBIDDEN_MODULES if name in sys.modules) - assert not imported, f"plain PDD FSDP2 smoke imported optional frameworks: {imported}" - - -def main() -> None: - _assert_optional_frameworks_absent() - assert torch.cuda.is_available(), "Task-10 FSDP2 gate requires CUDA" - dist.init_process_group("nccl") - rank = dist.get_rank() - world_size = dist.get_world_size() - assert world_size == 2, f"Task-10 FSDP2 gate requires two ranks, got {world_size}" - local_rank = int(os.environ["LOCAL_RANK"]) - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - assert torch.cuda.get_device_capability(device)[0] >= 8 - root_payload = [tempfile.mkdtemp(prefix="modelopt-pdd-fsdp2-") if rank == 0 else None] - dist.broadcast_object_list(root_payload, src=0) - root = pathlib.Path(root_payload[0]) - try: - torch.manual_seed(2026 + rank) - torch.cuda.manual_seed(3026 + rank) - reference = _build(device) - resumable = _build(device) - reference_step1 = _step(reference, rank=rank, step=1, device=device) - resumable_step1 = _step(resumable, rank=rank, step=1, device=device) - assert reference_step1 == resumable_step1 - _assert_state_equal(resumable.student, _state(reference.student)) - saved_cpu_rng = torch.get_rng_state().clone() - saved_cuda_rng = torch.cuda.get_rng_state(device).clone() - checkpoint = _save_checkpoint( - root, - resumable, - completed_steps=1, - device=device, - ) - - reference_step2 = _step(reference, rank=rank, step=2, device=device) - reference_state = _state(reference.student) - _assert_call_counts(reference.adapter, 4, device) - - del resumable - gc.collect() - torch.cuda.empty_cache() - restored = _build(device) - assert _load_checkpoint(checkpoint, restored, device=device) == 1 - torch.testing.assert_close(torch.get_rng_state(), saved_cpu_rng, rtol=0, atol=0) - torch.testing.assert_close(torch.cuda.get_rng_state(device), saved_cuda_rng, rtol=0, atol=0) - restored_step2 = _step(restored, rank=rank, step=2, device=device) - assert restored_step2 == reference_step2 - _assert_state_equal(restored.student, reference_state) - _assert_call_counts(restored.adapter, 2, device) - _assert_optional_frameworks_absent() - dist.barrier() - finally: - dist.barrier() - if rank == 0: - shutil.rmtree(root) - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py b/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py deleted file mode 100644 index 908be763acc..00000000000 --- a/tests/gpu/torch/fastgen/pdd_qwen_operability_smoke.py +++ /dev/null @@ -1,1146 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Staged canonical Qwen-Image PDD operability smoke and result validator.""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import math -import os -import pathlib -import subprocess -import sys -import time -from collections.abc import Mapping, Sequence -from typing import Any - -_THIS_FILE = pathlib.Path(__file__).resolve() -_REPO_ROOT = _THIS_FILE.parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -for _path in (_REPO_ROOT, _FASTGEN_DIR): - if str(_path) not in sys.path: - sys.path.insert(0, str(_path)) - -_MODEL_ID = "Qwen/Qwen-Image" -_MODEL_REVISION = "75e0b4be04f60ec59a75f475837eced720f823b6" -_AUTOMODEL_TREE_SHA256 = "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528" -_AUTOMODEL_WHEEL_SHA256 = "881aebafc5145752842afbbfe0a42e1c33d06847c3e418ad3d6f154ddc8e0f45" -_AUTOMODEL_RELEASE_COMMIT = "d02f49cb314554715aabb97e8dba6599c9f6e9e0" -_AUTOMODEL_RELEASE_TAG = "v0.5.0" -_AUTOMODEL_WHEEL = "nemo_automodel-0.5.0-py3-none-any.whl" -_AUTOMODEL_PACKAGE_FILE_COUNT = 490 -_EXPECTED_PAIRS = {"train-one": (0, 63), "resume-one": (124, 127)} -_STAGE_RESULT_KEYS = { - "schema_version", - "record_type", - "stage", - "pid", - "world_size", - "model", - "pdd", - "source", - "config_sha256", - "automodel", - "gpu", - "pair", - "sample_ids", - "diagnostics", - "teacher_calls_per_rank", - "checkpoint", - "resume", -} - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--stage", choices=("train-one", "resume-one", "validate"), required=True) - parser.add_argument("--run-root", type=pathlib.Path, required=True) - parser.add_argument("--before-automodel", type=pathlib.Path) - parser.add_argument("--after-automodel", type=pathlib.Path) - return parser.parse_args() - - -def _sha256(path: pathlib.Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _canonical_sha256(value: Any) -> str: - payload = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) - return hashlib.sha256(payload.encode()).hexdigest() - - -def _require_sha256(value: Any, *, name: str) -> str: - if not isinstance(value, str) or len(value) != 64: - raise ValueError(f"{name} is not a SHA-256 digest") - try: - int(value, 16) - except ValueError as error: - raise ValueError(f"{name} is not a hexadecimal SHA-256 digest") from error - return value.lower() - - -def _unique_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: - value: dict[str, Any] = {} - for key, item in pairs: - if key in value: - raise ValueError(f"JSON object contains duplicate key {key!r}") - value[key] = item - return value - - -def _reject_json_constant(token: str) -> None: - raise ValueError(f"JSON contains non-finite value {token}") - - -def _read_json(path: pathlib.Path) -> dict[str, Any]: - value = json.loads( - path.read_bytes(), - object_pairs_hook=_unique_json_object, - parse_constant=_reject_json_constant, - ) - if not isinstance(value, dict): - raise TypeError(f"{path} must contain a JSON object") - return value - - -def _finite_positive(value: Any, *, name: str) -> float: - if isinstance(value, bool) or not isinstance(value, int | float): - raise TypeError(f"{name} must be a real number") - value = float(value) - if not math.isfinite(value) or value <= 0: - raise ValueError(f"{name} must be finite and positive") - return value - - -def _absolute_path_without_symlinks(value: pathlib.Path, *, name: str) -> pathlib.Path: - path = pathlib.Path(os.path.abspath(value)) - for candidate in (*reversed(path.parents), path): - if candidate.is_symlink(): - raise ValueError(f"{name} cannot traverse a symlink: {candidate}") - return path - - -def _regular_directory(value: pathlib.Path, *, name: str) -> pathlib.Path: - path = _absolute_path_without_symlinks(value, name=name) - if not path.is_dir(): - raise ValueError(f"{name} must identify a regular directory") - return path.resolve() - - -def _regular_file(value: pathlib.Path, *, name: str) -> pathlib.Path: - path = _absolute_path_without_symlinks(value, name=name) - if not path.is_file(): - raise ValueError(f"{name} must identify a regular file") - return path.resolve() - - -def _create_run_root(value: pathlib.Path) -> pathlib.Path: - path = _absolute_path_without_symlinks(value, name="smoke run root") - path.mkdir(parents=True, exist_ok=True) - return _regular_directory(path, name="smoke run root") - - -def _relative_regular_file(root: pathlib.Path, value: Any, *, name: str) -> pathlib.Path: - if not isinstance(value, str) or not value: - raise TypeError(f"{name} must be a non-empty relative path") - relative = pathlib.PurePosixPath(value) - if relative.is_absolute() or ".." in relative.parts: - raise ValueError(f"{name} must stay beneath the run root") - root = _regular_directory(root, name=f"{name} root") - path = root.joinpath(*relative.parts) - if any(candidate.is_symlink() for candidate in (path, *path.parents) if candidate != root): - raise ValueError(f"{name} cannot traverse a symlink") - resolved = path.resolve() - try: - resolved.relative_to(root) - except ValueError as error: - raise ValueError(f"{name} must stay beneath the run root") from error - if not resolved.is_file(): - raise ValueError(f"{name} must identify a regular file") - return resolved - - -def validate_stage_result(value: Mapping[str, Any], *, stage: str) -> None: - """Validate one atomic training-stage result without importing GPU dependencies.""" - if set(value) != _STAGE_RESULT_KEYS: - raise ValueError(f"{stage} result keys are incompatible") - if value["schema_version"] != 1 or value["record_type"] != "pdd_qwen_smoke_stage": - raise ValueError(f"{stage} result schema is incompatible") - if value["stage"] != stage or stage not in _EXPECTED_PAIRS: - raise ValueError("smoke stage identity is invalid") - if type(value["pid"]) is not int or value["pid"] <= 0: - raise ValueError("smoke stage pid is invalid") - if type(value["world_size"]) is not int or value["world_size"] < 2: - raise ValueError("full-Qwen smoke requires a multi-GPU world") - if value["model"] != {"id": _MODEL_ID, "revision": _MODEL_REVISION, "dtype": "bfloat16"}: - raise ValueError("smoke model identity is invalid") - if value["pdd"] != { - "grid_size": 128, - "grid_max_t": 0.999, - "flow_shift": 5.0, - "block_size_min": 4, - "block_size_max": 64, - "teacher_integrator": "euler", - "guidance_scale": 4.0, - "guidance_rescale": 1.0, - "guidance_eps": 1e-5, - }: - raise ValueError("smoke PDD identity is invalid") - source = value["source"] - if ( - not isinstance(source, Mapping) - or set(source) != {"commit", "dirty"} - or not isinstance(source["commit"], str) - or len(source["commit"]) != 40 - or source["dirty"] is not False - ): - raise ValueError("smoke source identity is invalid") - try: - int(source["commit"], 16) - except ValueError as error: - raise ValueError("smoke source commit is not hexadecimal") from error - _require_sha256(value["config_sha256"], name="config_sha256") - automodel = value["automodel"] - expected_automodel_keys = { - "distribution", - "version", - "package_tree_sha256", - "wheel_sha256", - "runtime_versions", - } - if not isinstance(automodel, Mapping) or set(automodel) != expected_automodel_keys: - raise ValueError("smoke AutoModel identity is invalid") - if ( - automodel["distribution"] != "nemo_automodel" - or automodel["version"] != "0.5.0" - or automodel["runtime_versions"] != {"diffusers": "0.38.0"} - or automodel["package_tree_sha256"] != _AUTOMODEL_TREE_SHA256 - or automodel["wheel_sha256"] != _AUTOMODEL_WHEEL_SHA256 - ): - raise ValueError("smoke AutoModel release identity is invalid") - _require_sha256(automodel["package_tree_sha256"], name="automodel.package_tree_sha256") - _require_sha256(automodel["wheel_sha256"], name="automodel.wheel_sha256") - gpu = value["gpu"] - if not isinstance(gpu, Mapping) or set(gpu) != { - "names", - "total_memory_bytes", - "host_available_bytes", - "allocated_before_step_bytes", - "peak_memory_bytes", - "student_parameter_bytes", - "teacher_parameter_bytes", - "step_seconds", - }: - raise ValueError("smoke GPU evidence is invalid") - if not isinstance(gpu["names"], list) or len(gpu["names"]) != value["world_size"]: - raise ValueError("smoke GPU inventory does not match world size") - if any(not isinstance(name, str) or not name for name in gpu["names"]): - raise ValueError("smoke GPU names are invalid") - for name in ( - "total_memory_bytes", - "host_available_bytes", - "allocated_before_step_bytes", - "peak_memory_bytes", - ): - values = gpu[name] - if ( - not isinstance(values, list) - or len(values) != value["world_size"] - or any(type(item) is not int or item <= 0 for item in values) - ): - raise ValueError(f"smoke GPU {name} is invalid") - for name in ("student_parameter_bytes", "teacher_parameter_bytes"): - if type(gpu[name]) is not int or gpu[name] <= 0: - raise ValueError(f"smoke capacity {name} is invalid") - for allocated, peak, total in zip( - gpu["allocated_before_step_bytes"], - gpu["peak_memory_bytes"], - gpu["total_memory_bytes"], - strict=True, - ): - if not allocated <= peak <= total: - raise ValueError("smoke GPU allocation evidence is inconsistent") - _finite_positive(gpu["step_seconds"], name="gpu.step_seconds") - if value["pair"] != {"n": _EXPECTED_PAIRS[stage][0], "k": _EXPECTED_PAIRS[stage][1]}: - raise ValueError("smoke explicit support pair is invalid") - sample_ids = value["sample_ids"] - if ( - not isinstance(sample_ids, list) - or len(sample_ids) != value["world_size"] - or any(not isinstance(item, str) or not item for item in sample_ids) - or len(set(sample_ids)) != len(sample_ids) - ): - raise ValueError("smoke sample IDs are invalid") - expected_step = 1 if stage == "train-one" else 2 - expected_ids = [ - f"synthetic-pdd-smoke-step-{expected_step}-rank-{rank}" - for rank in range(value["world_size"]) - ] - if sample_ids != expected_ids: - raise ValueError("smoke sample IDs do not match the canonical rank order") - diagnostics = value["diagnostics"] - if not isinstance(diagnostics, Mapping) or set(diagnostics) != { - "completed_step", - "loss", - "grad_norm", - "student_adamw_nominal_update_ratio", - "pdd_projection_update_ratio", - "learning_rate", - "student_velocity_rms", - "teacher_velocity_rms", - "student_teacher_velocity_rms_ratio", - "reconstructed_state_rms", - }: - raise ValueError("smoke diagnostics are invalid") - if diagnostics["completed_step"] != expected_step: - raise ValueError("smoke completed step is invalid") - for name in diagnostics: - if name != "completed_step": - _finite_positive(diagnostics[name], name=f"diagnostics.{name}") - calls = value["teacher_calls_per_rank"] - if calls != [2] * value["world_size"]: - raise ValueError("guided teacher call structure is invalid") - checkpoint = value["checkpoint"] - if not isinstance(checkpoint, Mapping) or set(checkpoint) != { - "path", - "manifest_sha256", - "completed_steps", - "parent_checkpoint", - }: - raise ValueError("smoke checkpoint evidence is invalid") - if checkpoint["completed_steps"] != expected_step: - raise ValueError("smoke checkpoint step is invalid") - expected_parent = None if stage == "train-one" else "step_00000001" - if checkpoint["parent_checkpoint"] != expected_parent: - raise ValueError("smoke checkpoint lineage is invalid") - if checkpoint["path"] != f"checkpoints/step_{expected_step:08d}": - raise ValueError("smoke checkpoint path is invalid") - _require_sha256(checkpoint["manifest_sha256"], name="checkpoint.manifest_sha256") - resume = value["resume"] - if stage == "train-one": - if resume is not None: - raise ValueError("first smoke stage cannot have resume evidence") - elif resume != { - "selected_checkpoint": "step_00000001", - "completed_steps": 1, - "parent_checkpoint": None, - "first_sample_ids": sample_ids, - "learning_rate": diagnostics["learning_rate"], - }: - raise ValueError("smoke resume evidence is invalid") - - -def validate_inference_result(value: Mapping[str, Any], *, root: pathlib.Path) -> None: - """Validate the exact authenticated PDD-4 inference evidence.""" - if value.get("schema_version") != 1 or value.get("record_type") != "pdd_inference": - raise ValueError("PDD inference result schema is invalid") - if value.get("condition") != "pdd_4" or value.get("schedule") != "pdd-4": - raise ValueError("PDD inference schedule identity is invalid") - if value.get("blocks") != [32, 32, 32, 32]: - raise ValueError("PDD-4 blocks are invalid") - if value.get("height") != 1024 or value.get("width") != 1024: - raise ValueError("PDD-4 smoke output must be exactly 1024x1024") - if ( - value.get("scheduler_steps") != 4 - or value.get("actual_transformer_invocations") != 4 - or value.get("batch_normalized_transformer_evaluations") != 4 - ): - raise ValueError("PDD-4 compute counters are invalid") - _finite_positive(value.get("latency_seconds"), name="inference.latency_seconds") - output = value.get("output") - if not isinstance(output, Mapping) or set(output) != {"path", "sha256"}: - raise ValueError("PDD inference output evidence is invalid") - image = _relative_regular_file(root, output["path"], name="inference.output.path") - if image.suffix.lower() != ".png" or image.read_bytes()[:8] != b"\x89PNG\r\n\x1a\n": - raise ValueError("PDD inference output is not a PNG") - if image.stat().st_size <= 8 or _sha256(image) != _require_sha256( - output["sha256"], name="inference.output.sha256" - ): - raise ValueError("PDD inference PNG hash is invalid") - - -def _exact_mapping(value: Any, keys: set[str], *, name: str) -> Mapping[str, Any]: - if not isinstance(value, Mapping) or set(value) != keys: - raise ValueError(f"{name} must contain exactly {sorted(keys)}") - return value - - -def _validate_checkpoint_identity(identity: Any, *, stage: Mapping[str, Any]) -> None: - from modelopt.torch.fastgen import PDDMetadata - - identity = _exact_mapping( - identity, - { - "schema_version", - "model", - "pdd_metadata", - "guidance", - "automodel", - "data", - "topology", - "training", - "optimizer", - "scheduler", - }, - name="smoke checkpoint identity", - ) - if identity["schema_version"] != 1 or identity["model"] != stage["model"]: - raise ValueError("smoke checkpoint model identity is incompatible") - metadata = PDDMetadata.from_dict(identity["pdd_metadata"]) - if metadata.to_dict() != identity["pdd_metadata"]: - raise ValueError("smoke checkpoint PDD metadata is not canonical") - pdd = stage["pdd"] - if ( - metadata.grid_size != pdd["grid_size"] - or metadata.grid_max_t != pdd["grid_max_t"] - or metadata.flow_shift != pdd["flow_shift"] - or metadata.block_size_min != pdd["block_size_min"] - or metadata.block_size_max != pdd["block_size_max"] - or metadata.teacher_integrator != pdd["teacher_integrator"] - or metadata.inference_blocks != (32, 32, 32, 32) - or metadata.layer_spec.to_dict() - != { - "projection_path": "transformer.proj_out", - "head_layout": "channel_major", - "output_channels": None, - } - ): - raise ValueError("smoke checkpoint PDD metadata does not match the stage") - if identity["guidance"] != { - "scale": pdd["guidance_scale"], - "rescale": pdd["guidance_rescale"], - "eps": pdd["guidance_eps"], - }: - raise ValueError("smoke checkpoint guidance does not match the stage") - if identity["automodel"] != stage["automodel"]: - raise ValueError("smoke checkpoint AutoModel identity does not match the stage") - if identity["topology"] != { - "world_size": stage["world_size"], - "pure_data_parallel": True, - }: - raise ValueError("smoke checkpoint topology does not match the stage") - - -def _validate_automodel_snapshot( - snapshot: Any, - *, - expected_automodel: Mapping[str, Any], -) -> None: - snapshot = _exact_mapping( - snapshot, - { - "distribution", - "files", - "import_origin", - "package_file_count", - "package_tree_sha256", - "release_commit", - "release_tag", - "root", - "runtime_versions", - "version", - "wheel", - "wheel_sha256", - }, - name="AutoModel snapshot", - ) - for key in ("distribution", "version", "runtime_versions"): - if snapshot[key] != expected_automodel[key]: - raise ValueError(f"AutoModel snapshot identity differs for {key}") - for key in ("package_tree_sha256", "wheel_sha256"): - if ( - _require_sha256(snapshot[key], name=f"AutoModel snapshot {key}") - != expected_automodel[key] - ): - raise ValueError(f"AutoModel snapshot identity differs for {key}") - if ( - snapshot["release_commit"] != _AUTOMODEL_RELEASE_COMMIT - or snapshot["release_tag"] != _AUTOMODEL_RELEASE_TAG - or snapshot["wheel"] != _AUTOMODEL_WHEEL - or type(snapshot["package_file_count"]) is not int - or snapshot["package_file_count"] != _AUTOMODEL_PACKAGE_FILE_COUNT - ): - raise ValueError("AutoModel snapshot release identity is invalid") - root_value = snapshot["root"] - origin_value = snapshot["import_origin"] - if not isinstance(root_value, str) or not isinstance(origin_value, str): - raise TypeError("AutoModel snapshot root and import origin must be strings") - root = pathlib.Path(root_value) - import_origin = pathlib.Path(origin_value) - if not root.is_absolute() or not import_origin.is_absolute(): - raise ValueError("AutoModel snapshot paths must be absolute") - try: - import_origin.relative_to(root) - except ValueError as error: - raise ValueError("AutoModel snapshot import origin is outside its root") from error - files = snapshot["files"] - if not isinstance(files, list) or len(files) != _AUTOMODEL_PACKAGE_FILE_COUNT: - raise ValueError("AutoModel snapshot file inventory is invalid") - tree = hashlib.sha256() - previous_path: str | None = None - for index, raw_record in enumerate(files): - record = _exact_mapping( - raw_record, - {"path", "sha256", "size"}, - name=f"AutoModel snapshot files[{index}]", - ) - path = record["path"] - if ( - not isinstance(path, str) - or not path - or pathlib.PurePosixPath(path).is_absolute() - or "\\" in path - or any(part in ("", ".", "..") for part in path.split("/")) - or (previous_path is not None and path <= previous_path) - ): - raise ValueError("AutoModel snapshot paths must be sorted normalized references") - digest = _require_sha256(record["sha256"], name=f"AutoModel snapshot files[{index}].sha256") - if type(record["size"]) is not int or record["size"] < 0: - raise ValueError(f"AutoModel snapshot files[{index}].size is invalid") - tree.update(path.encode()) - tree.update(b"\0") - tree.update(digest.encode()) - tree.update(b"\0") - tree.update(str(record["size"]).encode()) - tree.update(b"\n") - previous_path = path - if tree.hexdigest() != snapshot["package_tree_sha256"]: - raise ValueError("AutoModel snapshot inventory does not match its tree digest") - - -def _validate_bundle_links( - *, - stage1: Mapping[str, Any], - stage2: Mapping[str, Any], - manifest1: Mapping[str, Any], - manifest2: Mapping[str, Any], - export_manifest: Mapping[str, Any], - automodel_snapshot: Mapping[str, Any], -) -> None: - identity1 = manifest1.get("identity") - identity2 = manifest2.get("identity") - if identity1 != identity2: - raise ValueError("smoke checkpoint identities differ across resume") - _validate_checkpoint_identity(identity1, stage=stage1) - _validate_checkpoint_identity(identity2, stage=stage2) - if export_manifest.get("identity") != identity2: - raise ValueError("smoke export identity does not match step 2") - if export_manifest.get("modelopt_source") != stage2["source"]: - raise ValueError("smoke export source does not match the training source") - expected_checkpoint = { - "name": "step_00000002", - "manifest_sha256": stage2["checkpoint"]["manifest_sha256"], - "completed_steps": 2, - } - if export_manifest.get("source_checkpoint") != expected_checkpoint: - raise ValueError("smoke export does not derive from the exact step-2 checkpoint") - _validate_automodel_snapshot( - automodel_snapshot, - expected_automodel=stage2["automodel"], - ) - - -def _load_matching_automodel_snapshots( - before_path: pathlib.Path, - after_path: pathlib.Path, - *, - expected_automodel: Mapping[str, Any], -) -> Mapping[str, Any]: - before_path = _regular_file(before_path, name="before AutoModel snapshot") - after_path = _regular_file(after_path, name="after AutoModel snapshot") - before_bytes = before_path.read_bytes() - if before_bytes != after_path.read_bytes(): - raise ValueError("AutoModel package snapshot changed during full-Qwen smoke") - snapshot = _read_json(before_path) - _validate_automodel_snapshot(snapshot, expected_automodel=expected_automodel) - return snapshot - - -def _raw_config(run_root: pathlib.Path, world_size: int) -> dict[str, Any]: - return { - "model": { - "pretrained_model_name_or_path": _MODEL_ID, - "revision": _MODEL_REVISION, - "torch_dtype": "bfloat16", - "device": "cuda", - "transformer_engine_linear": False, - "peft": None, - "guidance_embeds": False, - "fuse_qkv_projections": False, - }, - "pdd": { - "pred_type": "flow", - "num_train_timesteps": None, - "guidance_scale": 4.0, - "student_sample_steps": 4, - "student_sample_type": "ode", - "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, - }, - "optim": { - "learning_rate": 2.0e-5, - "weight_decay": 0.01, - "betas": [0.9, 0.999], - "eps": 1.0e-8, - }, - "guidance": {"rescale": 1.0, "eps": 1e-5}, - "training": { - "seed": 42, - "max_steps": 2, - "max_grad_norm": 1.0, - "zero_grad_warmup_steps": 0, - "log_every_steps": 1, - "checkpoint_every_steps": 1, - "validation_every_steps": 1000, - "grad_accumulation_steps": 1, - "global_batch_size": world_size, - "validation_seed": 2026, - }, - "fsdp": { - "dp_size": world_size, - "tp_size": 1, - "cp_size": 1, - "pp_size": 1, - "ep_size": 1, - "activation_checkpointing": True, - }, - "data": { - "all_metadata_index": "synthetic_metadata.json", - "validation_metadata_index": "synthetic_heldout.json", - "dataloader": { - "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", - "cache_dir": "synthetic-unused", - "metadata_index": "synthetic_train.json", - "base_resolution": [1024, 1024], - "batch_size": 1, - "drop_last": True, - "shuffle": False, - "dynamic_batch_size": False, - "negative_prompt_embedding_path": "synthetic_negative.pt", - }, - }, - "checkpoint": { - "enabled": True, - "checkpoint_dir": str(run_root / "checkpoints"), - "model_save_format": "torch_save", - "save_consolidated": False, - "restore_from": "LATEST", - }, - } - - -def _modelopt_source() -> dict[str, Any]: - commit = subprocess.run( - ["git", "rev-parse", "HEAD"], - cwd=_REPO_ROOT, - check=True, - capture_output=True, - text=True, - ).stdout.strip() - dirty = bool( - subprocess.run( - ["git", "status", "--porcelain", "--untracked-files=normal"], - cwd=_REPO_ROOT, - check=True, - capture_output=True, - text=True, - ).stdout - ) - if dirty: - raise RuntimeError("full-Qwen smoke requires a clean ModelOpt checkout") - return {"commit": commit, "dirty": False} - - -def _ordered_id_sha256(sample_ids: Sequence[str]) -> str: - from portable_cache import ordered_sample_ids_sha256 - - return ordered_sample_ids_sha256(sample_ids) - - -def _training_sample_ids(world_size: int) -> tuple[str, ...]: - return tuple( - f"synthetic-pdd-smoke-step-{step}-rank-{rank}" - for step in (1, 2) - for rank in range(world_size) - ) - - -def _build_sampler(world_size: int, rank: int) -> Any: - import torch - from fastgen_data import ReplayableBatchSampler - from torch.utils.data import Sampler - - sample_ids = _training_sample_ids(world_size) - - class _Dataset: - metadata = [{"sample_id": sample_id} for sample_id in sample_ids] - - class _Sampler(Sampler[list[int]]): - def __init__(self) -> None: - self.dataset = _Dataset() - self.rank = rank - self.num_replicas = world_size - self.epoch = 0 - self.batches_yielded = 0 - - def set_epoch(self, epoch: int) -> None: - self.epoch = epoch - - def load_state_dict(self, state: Mapping[str, Any]) -> None: - self.epoch = int(state["epoch"]) - self.batches_yielded = int(state["batches_yielded"]) - - def __iter__(self): - batches = ([rank], [world_size + rank]) - for index, batch in enumerate(batches): - if index >= self.batches_yielded: - self.batches_yielded = index + 1 - yield list(batch) - - def __len__(self) -> int: - return 2 - - assert torch.distributed.get_world_size() == world_size - return ReplayableBatchSampler(_Sampler()) - - -def _prepared_batch( - setup: Any, *, rank: int, step: int, device: Any, dtype: Any -) -> tuple[Any, Any]: - import torch - from pdd_training import PreparedPDDBatch - - config = getattr(setup.student, "config", None) - in_channels = getattr(config, "in_channels", None) - condition_width = getattr(config, "joint_attention_dim", None) - if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: - raise RuntimeError("pinned Qwen config has invalid in_channels") - if type(condition_width) is not int or condition_width <= 0: - raise RuntimeError("pinned Qwen config has invalid joint_attention_dim") - generator = torch.Generator(device="cpu").manual_seed(10_000 + rank * 10 + step) - latent_shape = (1, in_channels // 4, 128, 128) - data = torch.randn(latent_shape, generator=generator, dtype=torch.float32).to( - device=device, dtype=dtype - ) - noise = torch.randn(latent_shape, generator=generator, dtype=torch.float32).to(device=device) - condition = torch.randn((1, 8, condition_width), generator=generator, dtype=torch.float32).to( - device=device, dtype=dtype - ) - negative = torch.randn((1, 8, condition_width), generator=generator, dtype=torch.float32).to( - device=device, dtype=dtype - ) - mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 0]], device=device, dtype=torch.long) - negative_mask = torch.tensor([[1, 1, 1, 1, 0, 0, 0, 0]], device=device, dtype=torch.long) - sample_id = f"synthetic-pdd-smoke-step-{step}-rank-{rank}" - batch = PreparedPDDBatch( - data=data, - condition=(condition, mask), - negative_condition=(negative, negative_mask), - sample_ids=(sample_id,), - valid_mask=(True,), - ) - return batch, noise - - -def _identity( - *, setup: Any, training: Any, config: Any, sampler: Any, raw: Mapping[str, Any] -) -> dict[str, Any]: - from pdd_checkpoint import build_pdd_checkpoint_identity - - train_ids = tuple(item["sample_id"] for item in sampler.dataset.metadata) - heldout_ids = ("synthetic-pdd-smoke-heldout-not-evaluated",) - return build_pdd_checkpoint_identity( - metadata=setup.metadata, - model_id=config.model_id, - model_revision=config.model_revision, - guidance_scale=config.pdd.guidance_scale, - guidance_rescale=config.guidance.rescale, - guidance_eps=config.guidance.eps, - automodel_snapshot=setup.automodel_snapshot, - ordered_train_id_sha256=_ordered_id_sha256(train_ids), - ordered_heldout_id_sha256=_ordered_id_sha256(heldout_ids), - dataset_snapshot_sha256=_canonical_sha256( - {"domain": "modelopt-pdd-synthetic-smoke-v1", "config": raw["pdd"]} - ), - local_batch_size=1, - grad_accumulation_steps=1, - training_seed=config.training.seed, - validation_seed=config.training.validation_seed, - validation_every_steps=config.training.validation_every_steps, - max_grad_norm=config.training.max_grad_norm, - zero_grad_warmup_steps=config.training.zero_grad_warmup_steps, - activation_checkpointing=config.parallel.activation_checkpointing, - dtype="bfloat16", - optimizer=setup.optimizer, - scheduler=training.scheduler, - ) - - -def _diagnostics_dict(diagnostics: Any) -> dict[str, Any]: - return { - "completed_step": diagnostics.completed_step, - "loss": diagnostics.loss, - "grad_norm": diagnostics.grad_norm, - "student_adamw_nominal_update_ratio": diagnostics.student_adamw_nominal_update_ratio, - "pdd_projection_update_ratio": diagnostics.pdd_projection_update_ratio, - "learning_rate": diagnostics.learning_rate, - "student_velocity_rms": diagnostics.student_velocity_rms, - "teacher_velocity_rms": diagnostics.teacher_velocity_rms, - "student_teacher_velocity_rms_ratio": diagnostics.student_teacher_velocity_rms_ratio, - "reconstructed_state_rms": diagnostics.reconstructed_state_rms, - } - - -def _run_training_stage(stage: str, run_root: pathlib.Path) -> None: - if os.environ.get("HF_HUB_OFFLINE") != "1": - raise RuntimeError("full-Qwen smoke requires HF_HUB_OFFLINE=1 and a pinned local snapshot") - import torch - import torch.distributed as dist - from export_pdd_qwen_image import host_available_bytes - from pdd_artifacts import write_canonical_json - from pdd_checkpoint import PDDCheckpointManager, validate_pdd_training_checkpoint - from pdd_recipe import ( - build_pdd_setup, - build_pdd_training_artifacts, - initialize_pdd_distributed, - resolve_pdd_recipe_config, - ) - - initialize_pdd_distributed(backend="nccl", timeout_minutes=60) - rank = dist.get_rank() - world_size = dist.get_world_size() - if world_size < 2: - raise RuntimeError("full-Qwen smoke requires a multi-GPU FSDP2 world") - device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) - torch.cuda.set_device(device) - run_root = _create_run_root(run_root) - result_path = run_root / ("stage1.json" if stage == "train-one" else "stage2.json") - if result_path.exists() or result_path.is_symlink(): - raise FileExistsError(f"smoke stage result already exists: {result_path}") - - raw = _raw_config(run_root, world_size) - config = resolve_pdd_recipe_config(raw) - setup = build_pdd_setup(config) - training = build_pdd_training_artifacts(setup, config) - sampler = _build_sampler(world_size, rank) - identity = _identity(setup=setup, training=training, config=config, sampler=sampler, raw=raw) - manager = PDDCheckpointManager( - root=config.checkpoint.checkpoint_dir, - checkpointer=setup.checkpointer, - model=setup.student, - optimizer=setup.optimizer, - scheduler=training.scheduler, - trainer=training.trainer, - sampler=sampler, - rng=training.rng, - identity=identity, - ) - resume_payload = None - try: - if stage == "train-one": - if manager.resolve("LATEST") is not None: - raise RuntimeError("first smoke stage requires an empty checkpoint root") - step = 1 - else: - resume = manager.load("LATEST") - if resume is None: - raise RuntimeError("resume smoke stage found no LATEST checkpoint") - if ( - resume.checkpoint_path.name != "step_00000001" - or resume.completed_steps != 1 - or resume.parent_checkpoint is not None - or training.trainer.completed_steps != 1 - ): - raise RuntimeError("resume smoke stage restored incompatible lineage") - expected_ids = sampler.expected_next_sample_ids() - resume.verify_first_batch(expected_ids) - resume_payload = { - "selected_checkpoint": resume.checkpoint_path.name, - "completed_steps": resume.completed_steps, - "parent_checkpoint": resume.parent_checkpoint, - "first_sample_ids": list(expected_ids), - "learning_rate": float(setup.optimizer.param_groups[0]["lr"]), - } - step = 2 - - batch, noise = _prepared_batch( - setup, - rank=rank, - step=step, - device=device, - dtype=config.dtype, - ) - if sampler.expected_next_sample_ids() != batch.sample_ids: - raise RuntimeError("synthetic smoke batch does not match the committed sampler cursor") - n_value, k_value = _EXPECTED_PAIRS[stage] - n = torch.tensor([n_value], device=device, dtype=torch.int64) - k = torch.tensor([k_value], device=device, dtype=torch.int64) - teacher_calls = 0 - - def count_teacher_call(_module: Any, _args: Any, _kwargs: Any) -> None: - nonlocal teacher_calls - teacher_calls += 1 - - hook = setup.teacher.register_forward_pre_hook(count_teacher_call, with_kwargs=True) - allocated_before_step = torch.cuda.memory_allocated(device) - torch.cuda.synchronize(device) - started = time.perf_counter() - try: - diagnostics = training.trainer.train_step( - batch, - noise=noise, - n=n, - k=k, - measure_updates=True, - ) - training.scheduler.step() - sampler.commit(batch.sample_ids) - finally: - hook.remove() - torch.cuda.synchronize(device) - step_seconds = time.perf_counter() - started - calls = [None] * world_size - dist.all_gather_object(calls, teacher_calls) - if calls != [2] * world_size: - raise RuntimeError(f"guided teacher calls differ across ranks: {calls}") - checkpoint = manager.save() - manifest = validate_pdd_training_checkpoint( - checkpoint, - expected_identity=identity, - expected_world_size=world_size, - ) - - sample_ids = [None] * world_size - dist.all_gather_object(sample_ids, batch.sample_ids[0]) - gpu_name = torch.cuda.get_device_name(device) - total_memory = torch.cuda.get_device_properties(device).total_memory - host_available = host_available_bytes() - peak_memory = torch.cuda.max_memory_allocated(device) - gpu_names = [None] * world_size - total_memories = [None] * world_size - host_memories = [None] * world_size - allocated_memories = [None] * world_size - peak_memories = [None] * world_size - dist.all_gather_object(gpu_names, gpu_name) - dist.all_gather_object(total_memories, total_memory) - dist.all_gather_object(host_memories, host_available) - dist.all_gather_object(allocated_memories, allocated_before_step) - dist.all_gather_object(peak_memories, peak_memory) - seconds = torch.tensor(step_seconds, device=device, dtype=torch.float64) - dist.all_reduce(seconds, op=dist.ReduceOp.MAX) - automodel = { - key: setup.automodel_snapshot[key] - for key in ( - "distribution", - "version", - "package_tree_sha256", - "wheel_sha256", - "runtime_versions", - ) - } - result = { - "schema_version": 1, - "record_type": "pdd_qwen_smoke_stage", - "stage": stage, - "pid": os.getpid(), - "world_size": world_size, - "model": {"id": _MODEL_ID, "revision": _MODEL_REVISION, "dtype": "bfloat16"}, - "pdd": { - "grid_size": 128, - "grid_max_t": 0.999, - "flow_shift": 5.0, - "block_size_min": 4, - "block_size_max": 64, - "teacher_integrator": "euler", - "guidance_scale": 4.0, - "guidance_rescale": 1.0, - "guidance_eps": 1e-5, - }, - "source": _modelopt_source(), - "config_sha256": _canonical_sha256(raw), - "automodel": automodel, - "gpu": { - "names": gpu_names, - "total_memory_bytes": total_memories, - "host_available_bytes": host_memories, - "allocated_before_step_bytes": allocated_memories, - "peak_memory_bytes": peak_memories, - "student_parameter_bytes": sum( - parameter.numel() * parameter.element_size() - for parameter in setup.student.parameters() - ), - "teacher_parameter_bytes": sum( - parameter.numel() * parameter.element_size() - for parameter in setup.teacher.parameters() - ), - "step_seconds": float(seconds.item()), - }, - "pair": {"n": n_value, "k": k_value}, - "sample_ids": sample_ids, - "diagnostics": _diagnostics_dict(diagnostics), - "teacher_calls_per_rank": calls, - "checkpoint": { - "path": checkpoint.relative_to(run_root).as_posix(), - "manifest_sha256": _sha256(checkpoint / "manifest.json"), - "completed_steps": manifest["completed_steps"], - "parent_checkpoint": manifest["parent_checkpoint"], - }, - "resume": resume_payload, - } - validate_stage_result(result, stage=stage) - if rank == 0: - write_canonical_json(result_path, result) - dist.barrier() - finally: - setup.checkpointer.close() - dist.destroy_process_group() - - -def _validate_bundle( - run_root: pathlib.Path, - before_automodel: pathlib.Path, - after_automodel: pathlib.Path, -) -> pathlib.Path: - from pdd_artifacts import load_canonical_json, write_canonical_json - from pdd_checkpoint import validate_pdd_training_checkpoint - from pdd_export import inspect_pdd_export - - run_root = _regular_directory(run_root, name="smoke run root") - stage1_path = _relative_regular_file(run_root, "stage1.json", name="stage-1 result") - stage2_path = _relative_regular_file(run_root, "stage2.json", name="stage-2 result") - export_root = _regular_directory(run_root / "export", name="smoke export root") - export_manifest_path = _regular_file( - export_root / "manifest.json", name="smoke export manifest" - ) - inference_path = _relative_regular_file( - run_root, "inference/pdd4.json", name="smoke inference result" - ) - output_path = run_root / "smoke_result.json" - if output_path.exists() or output_path.is_symlink(): - raise FileExistsError(f"smoke result already exists: {output_path}") - - stage1 = load_canonical_json(stage1_path) - stage2 = load_canonical_json(stage2_path) - if not isinstance(stage1, Mapping) or not isinstance(stage2, Mapping): - raise TypeError("smoke stage results must be JSON objects") - validate_stage_result(stage1, stage="train-one") - validate_stage_result(stage2, stage="resume-one") - if stage1["pid"] == stage2["pid"]: - raise ValueError("forced-resume stages did not use fresh processes") - for name in ("world_size", "model", "pdd", "source", "config_sha256", "automodel"): - if stage1[name] != stage2[name]: - raise ValueError(f"smoke stages disagree on {name}") - checkpoint1 = _relative_regular_file( - run_root, - pathlib.PurePosixPath(stage1["checkpoint"]["path"]).joinpath("manifest.json").as_posix(), - name="stage1.checkpoint.manifest", - ).parent - checkpoint2 = _relative_regular_file( - run_root, - pathlib.PurePosixPath(stage2["checkpoint"]["path"]).joinpath("manifest.json").as_posix(), - name="stage2.checkpoint.manifest", - ).parent - if _sha256(checkpoint1 / "manifest.json") != stage1["checkpoint"]["manifest_sha256"]: - raise ValueError("stage-1 checkpoint manifest hash changed") - if _sha256(checkpoint2 / "manifest.json") != stage2["checkpoint"]["manifest_sha256"]: - raise ValueError("stage-2 checkpoint manifest hash changed") - manifest1 = validate_pdd_training_checkpoint( - checkpoint1, expected_world_size=stage1["world_size"] - ) - manifest2 = validate_pdd_training_checkpoint( - checkpoint2, expected_world_size=stage2["world_size"] - ) - if manifest1["completed_steps"] != 1 or manifest2["completed_steps"] != 2: - raise ValueError("smoke checkpoint steps are invalid") - if ( - manifest1["parent_checkpoint"] is not None - or manifest2["parent_checkpoint"] != checkpoint1.name - ): - raise ValueError("smoke checkpoint parent lineage is invalid") - export_descriptor = inspect_pdd_export(export_root) - if export_descriptor.root != export_root: - raise ValueError("smoke export descriptor resolved an unexpected root") - inference = load_canonical_json(inference_path) - if not isinstance(inference, Mapping): - raise TypeError("smoke inference result must be a JSON object") - validate_inference_result(inference, root=inference_path.parent) - export_sha256 = _sha256(export_manifest_path) - if inference.get("export_manifest_sha256") != export_sha256: - raise ValueError("inference does not authenticate the smoke export") - automodel_snapshot = _load_matching_automodel_snapshots( - before_automodel, - after_automodel, - expected_automodel=stage2["automodel"], - ) - _validate_bundle_links( - stage1=stage1, - stage2=stage2, - manifest1=manifest1, - manifest2=manifest2, - export_manifest=export_descriptor.manifest, - automodel_snapshot=automodel_snapshot, - ) - result = { - "schema_version": 1, - "record_type": "pdd_qwen_operability_smoke", - "status": "passed", - "stage1_sha256": _sha256(stage1_path), - "stage2_sha256": _sha256(stage2_path), - "checkpoint_manifest_sha256": [ - stage1["checkpoint"]["manifest_sha256"], - stage2["checkpoint"]["manifest_sha256"], - ], - "export_manifest_sha256": export_sha256, - "inference_result_sha256": _sha256(inference_path), - "automodel_snapshot_sha256": _sha256( - _regular_file(before_automodel, name="before AutoModel snapshot") - ), - "model": stage1["model"], - "pdd": stage1["pdd"], - "source": stage1["source"], - "config_sha256": stage1["config_sha256"], - "world_size": stage1["world_size"], - "pairs": [stage1["pair"], stage2["pair"]], - "losses": [stage1["diagnostics"]["loss"], stage2["diagnostics"]["loss"]], - "inference": { - "blocks": inference["blocks"], - "scheduler_steps": inference["scheduler_steps"], - "actual_transformer_invocations": inference["actual_transformer_invocations"], - "batch_normalized_transformer_evaluations": inference[ - "batch_normalized_transformer_evaluations" - ], - "output_sha256": inference["output"]["sha256"], - "latency_seconds": inference["latency_seconds"], - }, - } - write_canonical_json(output_path, result) - return output_path - - -def main() -> None: - args = _parse_args() - if args.stage == "validate": - if args.before_automodel is None or args.after_automodel is None: - raise ValueError("validate requires --before-automodel and --after-automodel") - print(_validate_bundle(args.run_root, args.before_automodel, args.after_automodel)) - return - if args.before_automodel is not None or args.after_automodel is not None: - raise ValueError("training stages do not accept AutoModel snapshot arguments") - _run_training_stage(args.stage, args.run_root) - - -if __name__ == "__main__": - main() diff --git a/tests/gpu/torch/fastgen/test_pdd_toy.py b/tests/gpu/torch/fastgen/test_pdd_toy.py index b39d4f09063..fc137d2686d 100644 --- a/tests/gpu/torch/fastgen/test_pdd_toy.py +++ b/tests/gpu/torch/fastgen/test_pdd_toy.py @@ -1,5 +1,17 @@ # 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 c87850253b8da81d5fc5b98bb5a852b4355e90c0 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 05:26:49 -0700 Subject: [PATCH 20/45] fix(fastgen): harden PDD setup and resume identity Signed-off-by: Meng Xin --- .../fastgen/fastgen_data/collate_fns.py | 31 ++++++- .../fastgen_data/text_to_image_dataset.py | 38 ++++++++- examples/diffusers/fastgen/pdd/README.md | 3 + examples/diffusers/fastgen/pdd/finetune.py | 19 ++++- examples/diffusers/fastgen/pdd/recipe.py | 80 ++++++++++++------- .../preprocess/preprocessing_multiprocess.py | 16 +++- tests/examples/diffusers/fastgen/conftest.py | 2 + .../diffusers/fastgen/test_dataset_paths.py | 41 ++++++++++ .../fastgen/test_pdd_recipe_setup.py | 52 ++++++++++++ 9 files changed, 243 insertions(+), 39 deletions(-) diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index e47e6076beb..ebf0365cea1 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -32,8 +32,11 @@ """ import functools +import hashlib +import io import logging from collections.abc import Sequence +from pathlib import Path import torch from nemo_automodel.components.datasets.diffusion.sampler import SequentialBucketSampler @@ -130,14 +133,16 @@ def collate_fn_text_to_image( return image_batch -def _load_negative_prompt_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor]: +def _load_negative_prompt_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor, str]: """Load ``(embed, mask)`` from a negative-prompt-embedding file. Accepts a dict with an ``embed`` tensor (and an optional ``mask`` / ``prompt_embeds_mask`` / ``text_mask``) or a bare embedding tensor; a missing mask defaults to all-ones. """ - payload = torch.load(path, map_location="cpu", weights_only=True) + payload_bytes = Path(path).read_bytes() + payload_sha256 = hashlib.sha256(payload_bytes).hexdigest() + payload = torch.load(io.BytesIO(payload_bytes), map_location="cpu", weights_only=True) neg_embed = payload["embed"] if isinstance(payload, dict) else payload if not torch.is_tensor(neg_embed): raise TypeError( @@ -158,7 +163,19 @@ def _load_negative_prompt_embedding(path: str) -> tuple[torch.Tensor, torch.Tens ) if neg_mask is None: neg_mask = torch.ones(neg_embed.shape[:-1], dtype=torch.long) - return neg_embed, neg_mask + return neg_embed, neg_mask, payload_sha256 + + +def _dataset_snapshot_sha256(metadata_sha256: str, negative_sha256: str | None) -> str: + """Bind expected sample content and the static negative condition into one identity.""" + digest = hashlib.sha256(b"modelopt-fastgen-dataset-snapshot-v1\0") + digest.update(bytes.fromhex(metadata_sha256)) + if negative_sha256 is None: + digest.update(b"\0no-negative-prompt") + else: + digest.update(b"\0negative-prompt\0") + digest.update(bytes.fromhex(negative_sha256)) + return digest.hexdigest() def build_text_to_image_multiresolution_dataloader( @@ -222,6 +239,7 @@ def build_text_to_image_multiresolution_dataloader( split=split, validation_count=validation_count, split_seed=split_seed, + verify_payload_hashes=exact_resume, ) effective_root = dataset.cache_root @@ -233,7 +251,8 @@ def build_text_to_image_multiresolution_dataloader( negative_prompt_embedding_path, "negative prompt embedding", ) - neg_embed, neg_mask = _load_negative_prompt_embedding(str(negative_path)) + neg_embed, neg_mask, negative_sha256 = _load_negative_prompt_embedding(str(negative_path)) + dataset.negative_prompt_embedding_sha256 = negative_sha256 if dp_rank == 0: logger.info( "Loaded negative_prompt_embedding from %s | shape=%s dtype=%s mask_shape=%s", @@ -247,6 +266,10 @@ def build_text_to_image_multiresolution_dataloader( negative_text_embeddings=neg_embed, negative_text_embeddings_mask=neg_mask, ) + dataset.dataset_snapshot_sha256 = _dataset_snapshot_sha256( + dataset.metadata_sha256, + dataset.negative_prompt_embedding_sha256, + ) sampler = SequentialBucketSampler( dataset, 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 d6b41a508a8..b97b05daf47 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -14,6 +14,7 @@ # limitations under the License. import hashlib +import io import json from collections.abc import Sequence from pathlib import Path @@ -38,6 +39,7 @@ def __init__( split: str | None = None, validation_count: int | None = None, split_seed: int = 2026, + verify_payload_hashes: bool = False, ): """ Args: @@ -47,6 +49,7 @@ def __init__( split: Optional deterministic ``"train"`` or ``"validation"`` selection. validation_count: Number of validation samples when ``split`` is set. split_seed: Local seed used to construct deterministic split membership. + verify_payload_hashes: Require and authenticate cached tensor content on every load. """ if selected_indices is not None and split is not None: raise ValueError("selected_indices and split are mutually exclusive") @@ -60,7 +63,10 @@ def __init__( self._split = split self._validation_count = validation_count self._split_seed = split_seed + self._verify_payload_hashes = verify_payload_hashes self._resolved_cache_files: dict[int, Path] = {} + self.negative_prompt_embedding_sha256: str | None = None + self.dataset_snapshot_sha256: str | None = None super().__init__(str(self.cache_root), quantization=64) def _load_metadata(self) -> list[dict]: @@ -99,12 +105,28 @@ def _load_metadata(self) -> list[dict]: raise TypeError( f"metadata shard {shard_path} item {shard_item_index} has invalid cache_file" ) + cache_sha256 = item.get("cache_sha256") + if cache_sha256 is not None and ( + not isinstance(cache_sha256, str) + or len(cache_sha256) != 64 + or any(character not in "0123456789abcdef" for character in cache_sha256) + ): + raise ValueError( + f"metadata shard {shard_path} item {shard_item_index} has invalid " + "cache_sha256" + ) + if self._verify_payload_hashes and cache_sha256 is None: + raise ValueError( + f"metadata shard {shard_path} item {shard_item_index} has no " + "cache_sha256 required for exact resume" + ) 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.metadata_sha256 = digest.hexdigest() + self.payload_hashes_complete = all("cache_sha256" in item for item in complete_metadata) if self._split is None: self.sample_ids = self._validate_selected_indices(self.total_num_samples) else: @@ -151,8 +173,20 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: ) self._resolved_cache_files[idx] = cache_file - # Load cached data - data = torch.load(cache_file, map_location="cpu", weights_only=True) + # Exact-resume mode authenticates the same bytes passed to torch.load, avoiding a + # hash-then-reopen race if a shared cache changes during training. + if self._verify_payload_hashes: + payload = cache_file.read_bytes() + actual_sha256 = hashlib.sha256(payload).hexdigest() + expected_sha256 = item["cache_sha256"] + if actual_sha256 != expected_sha256: + raise RuntimeError( + f"sample cache file {sample_id} SHA-256 mismatch: " + f"expected {expected_sha256}, found {actual_sha256}" + ) + data = torch.load(io.BytesIO(payload), map_location="cpu", weights_only=True) + else: + 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 = { diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index f9d780bb7b8..34d24969876 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -30,6 +30,9 @@ torchrun --standalone --nproc-per-node=8 \ The cache must contain `metadata.json`, its declared shards, cached tensors, and `negative_prompt_embedding.pt`. The environment variable overrides the configured cache root. All metadata, tensor, and negative-embedding paths must still resolve inside that effective root. +For exact resume, every shard item must also contain the `cache_sha256` written by the shared +preprocessor. PDD verifies each tensor's bytes before loading it and binds those expected hashes +plus the negative-prompt embedding hash into the checkpoint dataset identity. Training deterministically derives disjoint train and validation membership from metadata ordinals; it does not rewrite the cache or require separate split manifests. The default recipe uses 2,000 diff --git a/examples/diffusers/fastgen/pdd/finetune.py b/examples/diffusers/fastgen/pdd/finetune.py index c071d746c30..aa3658a5b3a 100644 --- a/examples/diffusers/fastgen/pdd/finetune.py +++ b/examples/diffusers/fastgen/pdd/finetune.py @@ -139,7 +139,7 @@ def _validate_dataset_contract( validation_dataset: Any, config: Any, ) -> tuple[Mapping[str, Any], str, str]: - """Collectively verify deterministic split membership and the source metadata digest.""" + """Collectively verify deterministic splits and the authenticated dataset snapshot.""" import torch.distributed as dist try: @@ -159,9 +159,20 @@ def _validate_dataset_contract( raise RuntimeError("training and validation datasets disagree on total sample count.") if train_dataset.metadata_sha256 != validation_dataset.metadata_sha256: raise RuntimeError("training and validation datasets disagree on metadata content.") + if ( + not train_dataset.payload_hashes_complete + or not validation_dataset.payload_hashes_complete + ): + raise RuntimeError("PDD exact resume requires a cache_sha256 for every tensor payload.") + if train_dataset.dataset_snapshot_sha256 != validation_dataset.dataset_snapshot_sha256: + raise RuntimeError("training and validation datasets disagree on dataset content.") + if not isinstance(train_dataset.dataset_snapshot_sha256, str): + raise RuntimeError("PDD dataloader did not construct a dataset snapshot identity.") report = { "cache_root": str(train_dataset.cache_root), "metadata_sha256": train_dataset.metadata_sha256, + "negative_prompt_embedding_sha256": (train_dataset.negative_prompt_embedding_sha256), + "dataset_snapshot_sha256": train_dataset.dataset_snapshot_sha256, "total_samples": train_dataset.total_num_samples, "train_samples": len(train_ids), "validation_samples": len(validation_ids), @@ -494,7 +505,7 @@ def main() -> None: automodel_snapshot=setup.automodel_snapshot, ordered_train_id_sha256=train_ordered_id_sha256, ordered_heldout_id_sha256=heldout_ordered_id_sha256, - dataset_snapshot_sha256=snapshot_report["metadata_sha256"], + dataset_snapshot_sha256=snapshot_report["dataset_snapshot_sha256"], local_batch_size=config.training.local_batch_size, grad_accumulation_steps=config.training.grad_accumulation_steps, training_seed=config.training.seed, @@ -532,7 +543,9 @@ def main() -> None: ) if rank == 0: logging.info( - "PDD dataset verified: metadata_sha256=%s train=%d validation=%d root=%s", + "PDD dataset verified: snapshot_sha256=%s metadata_sha256=%s " + "train=%d validation=%d root=%s", + snapshot_report["dataset_snapshot_sha256"], snapshot_report["metadata_sha256"], snapshot_report["train_samples"], snapshot_report["validation_samples"], diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index ff00c10d201..73359f4eec8 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -529,6 +529,44 @@ def _load_unwrapped_transformer( return pipe, student +def _stage_and_shard_training_models( + student: nn.Module, + teacher: nn.Module, + projection: PDDOutputProjection, + projection_identity: tuple[int, int, int | None], + manager: Any, + *, + device: torch.device, + dtype: torch.dtype, + fuse_qkv_projections: bool, +) -> tuple[nn.Module, nn.Module]: + """Move and shard one dense model at a time to bound setup-time GPU memory.""" + if fuse_qkv_projections and ( + not hasattr(student, "fuse_qkv_projections") or not hasattr(teacher, "fuse_qkv_projections") + ): + raise AttributeError("QKV fusion requires both Qwen transformers to expose the object API.") + + student.to(device=device, dtype=dtype) + if fuse_qkv_projections: + student.fuse_qkv_projections() + if not any(getattr(module, "fused_projections", False) for module in student.modules()): + logging.warning( + "Qwen fuse_qkv_projections() was accepted but produced no fused attention " + "modules in the pinned Diffusers release." + ) + _require_projection_identity(student, projection, projection_identity, stage="student staging") + student = manager.parallelize(student) + _require_projection_module(student, projection, stage="student FSDP2 parallelization") + + # The student is already sharded before the dense teacher reaches the GPU, so multi-rank + # setup never holds both complete Qwen transformers on one device. + teacher.to(device=device, dtype=dtype) + if fuse_qkv_projections: + teacher.fuse_qkv_projections() + teacher = manager.parallelize(teacher) + return student, teacher + + def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: """Compose released AutoModel APIs without editing or patching external packages.""" if not isinstance(config, PDDRecipeConfig): @@ -558,28 +596,6 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: metadata = PDDMetadata.from_config(config.pdd, projection) lifecycle.append("pdd_conversion") - student.to(device=config.device, dtype=config.dtype) - teacher.to(device=config.device, dtype=config.dtype) - _require_projection_identity(student, projection, identity, stage="device placement") - lifecycle.append("device") - - if config.fuse_qkv_projections: - if not hasattr(student, "fuse_qkv_projections") or not hasattr( - teacher, "fuse_qkv_projections" - ): - raise AttributeError( - "QKV fusion requires both Qwen transformers to expose the object API." - ) - student.fuse_qkv_projections() - teacher.fuse_qkv_projections() - if not any(getattr(module, "fused_projections", False) for module in student.modules()): - logging.warning( - "Qwen fuse_qkv_projections() was accepted but produced no fused attention " - "modules in the pinned Diffusers release." - ) - _require_projection_identity(student, projection, identity, stage="QKV fusion") - lifecycle.append("qkv") - world_size = dist.get_world_size() if config.training.global_batch_size is not None: effective_global_batch = ( @@ -611,14 +627,20 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: device_mesh=mesh_context.device_mesh, moe_mesh=mesh_context.moe_mesh, ) - student = manager.parallelize(student) - teacher = manager.parallelize(teacher) + student, teacher = _stage_and_shard_training_models( + student, + teacher, + projection, + identity, + manager, + device=config.device, + dtype=config.dtype, + fuse_qkv_projections=config.fuse_qkv_projections, + ) pipe.transformer = student - # FSDP2 shards Parameters in place and may replace the Parameter objects. The registered - # projection module and FQN must survive; optimizer identity is checked against the new, - # live post-FSDP Parameters below. - _require_projection_module(student, projection, stage="FSDP2 parallelization") - lifecycle.append("parallelize") + # Keep the public lifecycle summary stable even though placement, optional QKV fusion, and + # FSDP2 are deliberately interleaved per model to cap peak device memory. + lifecycle.extend(("device", "qkv", "parallelize")) trainable = [parameter for parameter in student.parameters() if parameter.requires_grad] if not trainable: diff --git a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py index ce604a2a665..bb933479fee 100644 --- a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py +++ b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py @@ -95,6 +95,14 @@ def _get_media_files(media_dir: Path, extensions: set) -> list[Path]: return sorted(media_files) +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + def _save_metadata_shards( all_metadata: list[dict], output_dir: Path, @@ -124,7 +132,13 @@ def _save_metadata_shards( 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": str(cache_file)}) + normalized_metadata.append( + { + **item, + "cache_file": str(cache_file), + "cache_sha256": _sha256_file(cache_file), + } + ) sharded = shard_world > 1 shard_prefix = f"r{shard_rank:02d}_" if sharded else "" diff --git a/tests/examples/diffusers/fastgen/conftest.py b/tests/examples/diffusers/fastgen/conftest.py index e4e7d7d04c2..3e052b6c7c9 100644 --- a/tests/examples/diffusers/fastgen/conftest.py +++ b/tests/examples/diffusers/fastgen/conftest.py @@ -15,6 +15,7 @@ from __future__ import annotations +import hashlib import json from typing import TYPE_CHECKING @@ -59,6 +60,7 @@ def _make( metadata.append( { "cache_file": str(cache_file), + "cache_sha256": hashlib.sha256(payload_path.read_bytes()).hexdigest(), "bucket_resolution": resolution, "original_resolution": resolution, "bucket_id": sample_id % 2, diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index 31c781fc51b..33c08c0ed51 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -17,6 +17,7 @@ from __future__ import annotations +import hashlib import json import logging import pathlib @@ -124,6 +125,45 @@ def test_dataset_accepts_absolute_payload_beneath_root(make_fastgen_cache, tmp_p assert dataset[0]["sample_id"] == 0 +def test_exact_resume_requires_and_verifies_payload_hashes(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + dataset = TextToImageDataset(cache, verify_payload_hashes=True) + assert dataset[0]["sample_id"] == 0 + + shard_path = cache / "metadata_shard_0.json" + shard = json.loads(shard_path.read_text()) + payload_path = cache / shard[0]["cache_file"] + torch.save({"latent": torch.full((4, 2, 2), 999.0)}, payload_path) + with pytest.raises(RuntimeError, match="SHA-256 mismatch"): + dataset[0] + + shard[0].pop("cache_sha256") + shard_path.write_text(json.dumps(shard)) + with pytest.raises(ValueError, match="cache_sha256 required for exact resume"): + TextToImageDataset(cache, verify_payload_hashes=True) + + +def test_dataset_snapshot_binds_negative_prompt_embedding(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + loader, _ = build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), + num_workers=0, + negative_prompt_embedding_path="negative_prompt_embedding.pt", + exact_resume=True, + ) + first_snapshot = loader.dataset.dataset_snapshot_sha256 + + torch.save(torch.full((2, 3), 7.0), cache / "negative_prompt_embedding.pt") + rebuilt, _ = build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), + num_workers=0, + negative_prompt_embedding_path="negative_prompt_embedding.pt", + exact_resume=True, + ) + + assert rebuilt.dataset.dataset_snapshot_sha256 != first_snapshot + + def test_environment_redirects_samples_and_relative_negative_embedding( make_fastgen_cache, monkeypatch, tmp_path ): @@ -206,3 +246,4 @@ def test_preprocessing_publishes_absolute_paths_for_relative_output(monkeypatch, assert published.is_absolute() assert published == payload.resolve() published.relative_to(output.resolve()) + assert shard[0]["cache_sha256"] == hashlib.sha256(payload.read_bytes()).hexdigest() diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index a809b21f748..95ff1d24e0d 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -29,6 +29,7 @@ import torch import yaml from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir +from torch import nn _REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] _FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" @@ -36,6 +37,8 @@ sys.path.insert(0, str(_FASTGEN_DIR)) from pdd.recipe import ( + _projection_identity, + _stage_and_shard_training_models, build_pdd_export_setup, build_pdd_setup, initialize_pdd_distributed, @@ -43,6 +46,8 @@ ) from pdd.verify_readonly_automodel import snapshot_installed_distribution +from modelopt.torch.fastgen import PDDLayerSpec, convert_to_pdd_output_projection + def _require_exact_automodel() -> None: try: @@ -53,6 +58,53 @@ def _require_exact_automodel() -> None: pytest.skip(f"requires the official nemo_automodel==0.5.0 wheel, found {version}") +def test_training_setup_shards_student_before_staging_teacher() -> None: + events: list[str] = [] + + class TrackedModel(nn.Module): + def __init__(self, label: str, *, projection: bool = False) -> None: + super().__init__() + self.label = label + self.proj_out = nn.Linear(2, 2) if projection else nn.Identity() + + def to(self, *args, **kwargs): + events.append(f"{self.label}.to") + return super().to(*args, **kwargs) + + class TrackedManager: + def parallelize(self, model): + events.append(f"{model.label}.parallelize") + return model + + student = TrackedModel("student", projection=True) + teacher = TrackedModel("teacher") + projection = convert_to_pdd_output_projection( + student, + PDDLayerSpec("proj_out", "channel_major"), + grid_size=4, + ) + + staged_student, staged_teacher = _stage_and_shard_training_models( + student, + teacher, + projection, + _projection_identity(projection), + TrackedManager(), + device=torch.device("cpu"), + dtype=torch.float32, + fuse_qkv_projections=False, + ) + + assert staged_student is student + assert staged_teacher is teacher + assert events == [ + "student.to", + "student.parallelize", + "teacher.to", + "teacher.parallelize", + ] + + def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: return { "model": { From 0823f7466e47b1293a03dbdf444654e64d4d8f53 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 09:37:01 -0700 Subject: [PATCH 21/45] fix(fastgen): make PDD optimizer resume exact Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/recipe.py | 32 +++++++++ examples/diffusers/fastgen/pdd/training.py | 2 + .../fastgen/test_pdd_recipe_setup.py | 71 ++++++++++++++++++- 3 files changed, 102 insertions(+), 3 deletions(-) diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 73359f4eec8..610c93fd1f4 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -567,6 +567,37 @@ def _stage_and_shard_training_models( return student, teacher +def _materialize_zero_step_adamw_state(optimizer: torch.optim.AdamW) -> None: + """Create complete strict-DCP state without changing parameters or update numbering.""" + if type(optimizer) is not torch.optim.AdamW: + raise TypeError("PDD state materialization requires the stock torch.optim.AdamW optimizer.") + if optimizer.state: + raise RuntimeError("PDD AdamW state must be empty before materialization.") + parameters = [parameter for group in optimizer.param_groups for parameter in group["params"]] + if any(parameter.grad is not None for parameter in parameters): + raise RuntimeError("PDD AdamW parameters must not have gradients before materialization.") + + learning_rates = [group["lr"] for group in optimizer.param_groups] + try: + for group in optimizer.param_groups: + group["lr"] = 0.0 + for parameter in parameters: + parameter.grad = torch.zeros_like(parameter) + optimizer.step() + for parameter in parameters: + state = optimizer.state.get(parameter) + if state is None or set(state) != {"step", "exp_avg", "exp_avg_sq"}: + raise RuntimeError("PDD AdamW did not create complete checkpoint state.") + step = state["step"] + if not isinstance(step, torch.Tensor) or step.numel() != 1 or step.item() != 1: + raise RuntimeError("PDD AdamW created an unexpected initial step.") + step.zero_() + finally: + for group, learning_rate in zip(optimizer.param_groups, learning_rates, strict=True): + group["lr"] = learning_rate + optimizer.zero_grad(set_to_none=True) + + def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: """Compose released AutoModel APIs without editing or patching external packages.""" if not isinstance(config, PDDRecipeConfig): @@ -660,6 +691,7 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: fused=False, maximize=False, ) + _materialize_zero_step_adamw_state(optimizer) optimizer_parameters = [ parameter for group in optimizer.param_groups for parameter in group["params"] ] diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py index 76013b3ef2c..c8780b9f63e 100644 --- a/examples/diffusers/fastgen/pdd/training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -455,6 +455,8 @@ def _adamw_nominal_update_ratio( if decay <= 0.0: raise RuntimeError("AdamW decoupled weight decay factor must remain positive.") for parameter in group["params"]: + if parameter.grad is None: + continue state = optimizer.state.get(parameter) if not state or "exp_avg" not in state or "exp_avg_sq" not in state: continue diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 95ff1d24e0d..5a7988a2100 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -17,6 +17,7 @@ from __future__ import annotations +import copy import importlib.metadata import json import os @@ -37,6 +38,7 @@ sys.path.insert(0, str(_FASTGEN_DIR)) from pdd.recipe import ( + _materialize_zero_step_adamw_state, _projection_identity, _stage_and_shard_training_models, build_pdd_export_setup, @@ -49,6 +51,52 @@ from modelopt.torch.fastgen import PDDLayerSpec, convert_to_pdd_output_projection +def test_zero_step_adamw_state_preserves_first_lazy_update() -> None: + torch.manual_seed(7) + eager_model = nn.Linear(4, 3) + lazy_model = copy.deepcopy(eager_model) + optimizer_options = { + "lr": 2.0e-5, + "weight_decay": 0.01, + "foreach": False, + "fused": False, + } + eager_optimizer = torch.optim.AdamW(eager_model.parameters(), **optimizer_options) + lazy_optimizer = torch.optim.AdamW(lazy_model.parameters(), **optimizer_options) + parameters_before = { + name: parameter.detach().clone() for name, parameter in eager_model.named_parameters() + } + + _materialize_zero_step_adamw_state(eager_optimizer) + + for name, parameter in eager_model.named_parameters(): + torch.testing.assert_close(parameter, parameters_before[name], rtol=0, atol=0) + assert parameter.grad is None + state = eager_optimizer.state[parameter] + assert state["step"].item() == 0 + assert not state["exp_avg"].count_nonzero() + assert not state["exp_avg_sq"].count_nonzero() + + eager_model.weight.square().mean().backward() + lazy_model.weight.square().mean().backward() + eager_optimizer.step() + lazy_optimizer.step() + + for eager_parameter, lazy_parameter in zip( + eager_model.parameters(), lazy_model.parameters(), strict=True + ): + torch.testing.assert_close(eager_parameter, lazy_parameter, rtol=0, atol=0) + for key in ("step", "exp_avg", "exp_avg_sq"): + torch.testing.assert_close( + eager_optimizer.state[eager_model.weight][key], + lazy_optimizer.state[lazy_model.weight][key], + rtol=0, + atol=0, + ) + assert eager_optimizer.state[eager_model.bias]["step"].item() == 0 + assert lazy_model.bias not in lazy_optimizer.state + + def _require_exact_automodel() -> None: try: version = importlib.metadata.version("nemo_automodel") @@ -366,18 +414,31 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: parameter for group in source.optimizer.param_groups for parameter in group["params"] ] assert any(parameter is source.projection.weight for parameter in optimizer_parameters) + assert set(source.optimizer.state) == set(optimizer_parameters) + assert all( + source.optimizer.state[parameter]["step"].item() == 0 for parameter in optimizer_parameters + ) + unused_parameter = next( + parameter for parameter in optimizer_parameters if parameter is not source.projection.weight + ) + unused_name = next( + name + for name, parameter in source.student.named_parameters() + if parameter is unused_parameter + ) # Diffusers 0.38 accepts the Qwen object API but currently performs no effective fusion. assert not any( getattr(module, "fused_projections", False) for module in source.student.modules() ) source.optimizer.zero_grad(set_to_none=True) - # A real PDD forward touches the backbone and projection. Exercise the strict stock - # optimizer restore with complete Adam state rather than an artificial partial update. - sum(parameter.float().square().mean() for parameter in optimizer_parameters).backward() + # Exercise strict stock-DCP restore after a partial-gradient update. Eager step-zero state must + # retain exact lazy-Adam semantics for untouched parameters while keeping every DCP key present. + source.projection.weight.float().square().mean().backward() source.optimizer.step() expected_weight = source.projection.weight.detach().clone() expected_exp_avg = source.optimizer.state[source.projection.weight]["exp_avg"].clone() + assert source.optimizer.state[unused_parameter]["step"].item() == 0 checkpoint_root = tmp_path / "checkpoint" source.checkpointer.save_model(source.student, str(checkpoint_root)) source.checkpointer.save_optimizer(source.optimizer, source.student, str(checkpoint_root)) @@ -406,6 +467,7 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: destination = build_pdd_setup(config) assert destination.metadata == source.metadata destination_projection = destination.projection + destination_unused = dict(destination.student.named_parameters())[unused_name] destination_weight_id = id(destination_projection.weight) destination.checkpointer.load_model( destination.student, @@ -424,6 +486,9 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: destination.optimizer.state[destination_projection.weight]["exp_avg"], expected_exp_avg, ) + assert destination.optimizer.state[destination_unused]["step"].item() == 0 + assert not destination.optimizer.state[destination_unused]["exp_avg"].count_nonzero() + assert not destination.optimizer.state[destination_unused]["exp_avg_sq"].count_nonzero() assert snapshot_installed_distribution() == before source.checkpointer.close() destination.checkpointer.close() From a8afc6c646a57ce5d273b97b7faeac513cd45dea Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 17:44:27 -0700 Subject: [PATCH 22/45] Refactor PDD training onto AutoModel lifecycle Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/checkpoint.py | 73 +- .../fastgen/pdd/configs/qwen_image.yaml | 48 +- examples/diffusers/fastgen/pdd/data.py | 399 ++++++++++ examples/diffusers/fastgen/pdd/finetune.py | 700 +----------------- examples/diffusers/fastgen/pdd/recipe.py | 682 ++++++++++++++--- .../pdd_checkpoint_failure_distributed.py | 14 +- .../fastgen/pdd_export_distributed.py | 5 +- .../examples/diffusers/fastgen/test_layout.py | 1 + .../fastgen/test_pdd_recipe_setup.py | 165 ++++- .../fastgen/test_pdd_training_lifecycle.py | 277 ++++++- 10 files changed, 1558 insertions(+), 806 deletions(-) create mode 100644 examples/diffusers/fastgen/pdd/data.py diff --git a/examples/diffusers/fastgen/pdd/checkpoint.py b/examples/diffusers/fastgen/pdd/checkpoint.py index 66b45e4cfe7..102af9fe37d 100644 --- a/examples/diffusers/fastgen/pdd/checkpoint.py +++ b/examples/diffusers/fastgen/pdd/checkpoint.py @@ -35,7 +35,7 @@ if TYPE_CHECKING: from collections.abc import Sequence -_CHECKPOINT_SCHEMA_VERSION = 1 +_CHECKPOINT_SCHEMA_VERSION = 2 _COMPLETE_SCHEMA_VERSION = 1 _FORBIDDEN_ARTIFACT_TOKENS = ("fake_score", "discriminator", "ema", "r1", "gan") @@ -317,6 +317,25 @@ def verify_first_batch(self, sample_ids: Sequence[str]) -> None: ) +class _StepSchedulerCheckpointState: + """Rank-local carrier for the normalized next-data StepScheduler cursor.""" + + def __init__(self) -> None: + self.state: dict[str, int] = {"step": 0, "epoch": 0} + + def state_dict(self) -> dict[str, int]: + return dict(self.state) + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + if set(state) != {"step", "epoch"}: + raise ValueError("PDD StepScheduler state must contain step and epoch.") + step = state["step"] + epoch = state["epoch"] + if type(step) is not int or step < 0 or type(epoch) is not int or epoch < 0: + raise ValueError("PDD StepScheduler step and epoch must be nonnegative integers.") + self.state = {"step": step, "epoch": epoch} + + def _checkpoint_sidecar_paths(checkpoint: Path, world_size: int) -> list[Path]: paths: list[Path] = [] for rank in range(world_size): @@ -324,6 +343,7 @@ def _checkpoint_sidecar_paths(checkpoint: Path, world_size: int) -> list[Path]: ( checkpoint / "rng" / f"rng_dp_rank_{rank}.pt", checkpoint / "sampler" / f"sampler_dp_rank_{rank}.pt", + checkpoint / "step_scheduler" / f"step_scheduler_dp_rank_{rank}.pt", checkpoint / "trainer" / f"trainer_dp_rank_{rank}.pt", ) ) @@ -369,6 +389,7 @@ def validate_pdd_training_checkpoint( "identity", "completed_steps", "learning_rates", + "step_scheduler", "parent_checkpoint", "rank_progress", "dcp_sha256", @@ -393,9 +414,19 @@ def validate_pdd_training_checkpoint( if trainer_state != { "completed_steps": manifest["completed_steps"], "learning_rates": manifest["learning_rates"], + "step_scheduler": manifest["step_scheduler"], "parent_checkpoint": manifest["parent_checkpoint"], }: raise RuntimeError("PDD trainer-state sidecar does not match the manifest.") + step_scheduler_state = manifest["step_scheduler"] + if ( + not isinstance(step_scheduler_state, dict) + or set(step_scheduler_state) != {"step", "epoch"} + or step_scheduler_state.get("step") != manifest["completed_steps"] + or type(step_scheduler_state.get("epoch")) is not int + or step_scheduler_state["epoch"] < 0 + ): + raise RuntimeError("PDD checkpoint StepScheduler state is invalid.") rank_progress = manifest["rank_progress"] if not isinstance(rank_progress, list) or len(rank_progress) != world_size: raise RuntimeError("PDD checkpoint rank progress does not match its topology.") @@ -526,6 +557,7 @@ def __init__( model: Any, optimizer: Any, scheduler: Any, + step_scheduler: Any, trainer: Any, sampler: Any, rng: Any, @@ -536,6 +568,8 @@ def __init__( self.model = model self.optimizer = optimizer self.scheduler = scheduler + self.step_scheduler = step_scheduler + self._step_scheduler_checkpoint_state = _StepSchedulerCheckpointState() self.trainer = trainer self.sampler = sampler self.rng = rng @@ -569,6 +603,7 @@ def _manifest(self, checkpoint: Path) -> dict[str, Any]: "identity", "completed_steps", "learning_rates", + "step_scheduler", "parent_checkpoint", "rank_progress", "dcp_sha256", @@ -700,6 +735,7 @@ def _publish_staging( final: Path, completed_steps: int, learning_rates: list[float], + step_scheduler_state: Mapping[str, int], parent: Path | None, rank_summaries: list[dict[str, Any]], ) -> None: @@ -710,6 +746,7 @@ def _publish_staging( "identity": self.identity, "completed_steps": completed_steps, "learning_rates": learning_rates, + "step_scheduler": dict(step_scheduler_state), "parent_checkpoint": None if parent is None else parent.name, "rank_progress": sorted(rank_summaries, key=lambda summary: summary["rank"]), "dcp_sha256": _dcp_payload_hashes(staging), @@ -721,6 +758,7 @@ def _publish_staging( { "completed_steps": completed_steps, "learning_rates": learning_rates, + "step_scheduler": dict(step_scheduler_state), "parent_checkpoint": manifest["parent_checkpoint"], }, ) @@ -739,7 +777,7 @@ def _publish_staging( _atomic_text(self.root / "LATEST", final.name + "\n") def save(self) -> Path: - """Synchronously save into staging, publish atomically, mark complete, then update LATEST.""" + """Save into staging, publish atomically, mark complete, then update LATEST.""" completed_steps = self.trainer.completed_steps if type(completed_steps) is not int or completed_steps <= 0: raise ValueError("PDD checkpoint requires at least one completed optimizer step.") @@ -747,6 +785,18 @@ def save(self) -> Path: if len({summary["sample_slots_consumed"] for summary in rank_summaries}) != 1: raise RuntimeError("PDD ranks disagree on consumed sample slots.") learning_rates = [float(group["lr"]) for group in self.optimizer.param_groups] + live_step_scheduler_state = self.step_scheduler.state_dict() + if live_step_scheduler_state.get("step") != completed_steps: + raise RuntimeError("PDD StepScheduler state does not match the completed update.") + sampler_epoch = self.sampler.state_dict()["epoch"] + live_epoch = live_step_scheduler_state.get("epoch") + if sampler_epoch not in {live_epoch, live_epoch + 1}: + raise RuntimeError("PDD sampler epoch is incompatible with the StepScheduler epoch.") + step_scheduler_state = {"step": completed_steps, "epoch": sampler_epoch} + self._step_scheduler_checkpoint_state.load_state_dict(step_scheduler_state) + rank_scheduler_states = _gather_objects(step_scheduler_state) + if any(state != step_scheduler_state for state in rank_scheduler_states): + raise RuntimeError("PDD ranks disagree on StepScheduler checkpoint state.") final = self.root / f"step_{completed_steps:08d}" parent = self._collective_resolve("LATEST") @@ -779,6 +829,11 @@ def save(self) -> Path: try: self.checkpointer.save_on_dp_ranks(self.rng, "rng", str(staging)) self.checkpointer.save_on_dp_ranks(self.sampler, "sampler", str(staging)) + self.checkpointer.save_on_dp_ranks( + self._step_scheduler_checkpoint_state, + "step_scheduler", + str(staging), + ) self.checkpointer.save_on_dp_ranks(self.trainer, "trainer", str(staging)) except BaseException as error: sidecar_error = f"{type(error).__name__}: {error}" @@ -800,6 +855,7 @@ def save(self) -> Path: final=final, completed_steps=completed_steps, learning_rates=learning_rates, + step_scheduler_state=step_scheduler_state, parent=parent, rank_summaries=rank_summaries, ) @@ -833,6 +889,11 @@ def load(self, restore_from: str | Path | None) -> PDDResumeState | None: ) self.checkpointer.load_on_dp_ranks(self.trainer, "trainer", str(checkpoint)) self.checkpointer.load_on_dp_ranks(self.sampler, "sampler", str(checkpoint)) + self.checkpointer.load_on_dp_ranks( + self._step_scheduler_checkpoint_state, + "step_scheduler", + str(checkpoint), + ) rank_progress = manifest["rank_progress"] if not isinstance(rank_progress, list) or len(rank_progress) != _world_size(): raise RuntimeError("PDD checkpoint rank progress does not match world size.") @@ -851,6 +912,14 @@ def load(self, restore_from: str | Path | None) -> PDDResumeState | None: raise RuntimeError(f"PDD restored sampler {key} does not match the manifest.") if self.trainer.completed_steps != manifest["completed_steps"]: raise RuntimeError("PDD restored trainer step does not match the manifest.") + step_scheduler_state = self._step_scheduler_checkpoint_state.state_dict() + if step_scheduler_state != manifest["step_scheduler"]: + raise RuntimeError("PDD restored StepScheduler state does not match the manifest.") + if step_scheduler_state["step"] != self.trainer.completed_steps: + raise RuntimeError("PDD restored StepScheduler step does not match the trainer.") + if step_scheduler_state["epoch"] != sampler_state["epoch"]: + raise RuntimeError("PDD restored StepScheduler epoch does not match the sampler.") + self.step_scheduler.load_state_dict(step_scheduler_state) current_lrs = [float(group["lr"]) for group in self.optimizer.param_groups] if current_lrs != manifest["learning_rates"]: raise RuntimeError("PDD restored learning rate does not match the manifest.") diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 1e8c26cb11a..03de092a073 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -1,5 +1,7 @@ # Qwen-Image PDD training recipe. Result-bearing paths/topology remain externally gated. +seed: 42 + model: pretrained_model_name_or_path: Qwen/Qwen-Image revision: 75e0b4be04f60ec59a75f475837eced720f823b6 @@ -29,26 +31,46 @@ pdd: optim: learning_rate: 2.0e-5 - weight_decay: 0.01 - betas: [0.9, 0.999] - eps: 1.0e-8 + optimizer: + _target_: torch.optim.AdamW + weight_decay: 0.01 + 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 + min_lr: 2.0e-5 guidance: rescale: 1.0 eps: 1.0e-5 -training: - seed: 42 +step_scheduler: max_steps: 10000 - max_grad_norm: 1.0 - zero_grad_warmup_steps: 0 - log_every_steps: 10 - checkpoint_every_steps: 1000 - validation_every_steps: 1000 - grad_accumulation_steps: 1 + num_epochs: 200 + log_every: 10 + ckpt_every_steps: 1000 + local_batch_size: 1 + save_checkpoint_every_epoch: false # Freeze to 256 only after the production-topology/data gate is approved. global_batch_size: - validation_seed: 2026 + +training_health: + max_grad_norm: 1.0 + zero_grad_warmup_steps: 0 + +validation: + count: 2000 + seed: 2026 + split_seed: 2026 + every_steps: 1000 fsdp: dp_size: @@ -59,8 +81,6 @@ fsdp: activation_checkpointing: true data: - validation_count: 2000 - split_seed: 2026 dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: data/qwen_image_cache diff --git a/examples/diffusers/fastgen/pdd/data.py b/examples/diffusers/fastgen/pdd/data.py new file mode 100644 index 00000000000..f2468cb74d3 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/data.py @@ -0,0 +1,399 @@ +# 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. + +"""Authenticated data and collective batch handling for the Qwen-Image PDD recipe.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +from collections.abc import Mapping +from typing import Any + + +def _ordered_id_sha256(sample_ids: tuple[str, ...]) -> str: + digest = hashlib.sha256(b"modelopt-pdd-ordered-sample-ids-v1\0") + for sample_id in sample_ids: + digest.update(sample_id.encode()) + digest.update(b"\n") + return digest.hexdigest() + + +def _dataloader_options(raw: Mapping[str, Any]) -> dict[str, Any]: + data = raw.get("data") + if not isinstance(data, Mapping) or not isinstance(data.get("dataloader"), Mapping): + raise TypeError("PDD config requires a data.dataloader mapping.") + options = dict(data["dataloader"]) + target = options.pop("_target_", None) + expected_target = "fastgen_data.build_text_to_image_multiresolution_dataloader" + if target != expected_target: + raise ValueError(f"PDD data.dataloader._target_ must be {expected_target!r}.") + if "base_resolution" in options: + options["base_resolution"] = tuple(options["base_resolution"]) + return options + + +def _build_training_dataloader( + raw: Mapping[str, Any], + config: Any, + *, + dp_rank: int, + dp_world_size: int, +) -> tuple[Any, Any]: + from fastgen_data import ReplayableBatchSampler + from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader + + options = _dataloader_options(raw) + if options.get("drop_last", True) is not True: + raise ValueError("PDD exact sample accounting requires data.dataloader.drop_last=true.") + if options.get("dynamic_batch_size", False) is not False: + raise ValueError("PDD v1 requires data.dataloader.dynamic_batch_size=false.") + options.update( + split="train", + validation_count=config.validation.count, + split_seed=config.validation.split_seed, + dp_rank=dp_rank, + dp_world_size=dp_world_size, + exact_resume=True, + sampler_seed=config.seed, + loader_seed=config.seed, + ) + dataloader, sampler = build_text_to_image_multiresolution_dataloader(**options) + if not isinstance(sampler, ReplayableBatchSampler): + raise RuntimeError("PDD training requires the committed replayable batch sampler.") + if options.get("batch_size", 1) != config.step_scheduler.local_batch_size: + raise RuntimeError("resolved local batch size does not match the built dataloader.") + return dataloader, sampler + + +def _build_validation_dataloader( + raw: Mapping[str, Any], + config: Any, + *, + dp_rank: int, + dp_world_size: int, +) -> tuple[Any, Any]: + from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader + + options = _dataloader_options(raw) + options.update( + split="validation", + validation_count=config.validation.count, + split_seed=config.validation.split_seed, + dp_rank=dp_rank, + dp_world_size=dp_world_size, + drop_last=False, + shuffle=False, + dynamic_batch_size=False, + exact_resume=False, + sampler_seed=config.validation.seed, + loader_seed=config.validation.seed, + ) + return build_text_to_image_multiresolution_dataloader(**options) + + +def _validate_dataset_contract( + train_dataset: Any, + validation_dataset: Any, + config: Any, +) -> tuple[Mapping[str, Any], str, str]: + """Collectively verify deterministic splits and the authenticated dataset snapshot.""" + import torch.distributed as dist + + try: + train_ids = tuple(str(value) for value in train_dataset.sample_ids) + validation_ids = tuple(str(value) for value in validation_dataset.sample_ids) + if len(validation_ids) != config.validation.count: + raise RuntimeError( + f"validation split has {len(validation_ids)} samples; " + f"expected {config.validation.count}." + ) + if set(train_ids).intersection(validation_ids): + raise RuntimeError("training and validation splits overlap.") + expected = {str(index) for index in range(train_dataset.total_num_samples)} + if set(train_ids).union(validation_ids) != expected: + raise RuntimeError("training and validation splits do not cover metadata.json.") + if train_dataset.total_num_samples != validation_dataset.total_num_samples: + raise RuntimeError("training and validation datasets disagree on total sample count.") + if train_dataset.metadata_sha256 != validation_dataset.metadata_sha256: + raise RuntimeError("training and validation datasets disagree on metadata content.") + if ( + not train_dataset.payload_hashes_complete + or not validation_dataset.payload_hashes_complete + ): + raise RuntimeError("PDD exact resume requires a cache_sha256 for every tensor payload.") + if train_dataset.dataset_snapshot_sha256 != validation_dataset.dataset_snapshot_sha256: + raise RuntimeError("training and validation datasets disagree on dataset content.") + if not isinstance(train_dataset.dataset_snapshot_sha256, str): + raise RuntimeError("PDD dataloader did not construct a dataset snapshot identity.") + report = { + "cache_root": str(train_dataset.cache_root), + "metadata_sha256": train_dataset.metadata_sha256, + "negative_prompt_embedding_sha256": (train_dataset.negative_prompt_embedding_sha256), + "dataset_snapshot_sha256": train_dataset.dataset_snapshot_sha256, + "total_samples": train_dataset.total_num_samples, + "train_samples": len(train_ids), + "validation_samples": len(validation_ids), + "split_seed": config.validation.split_seed, + } + local_status: dict[str, Any] = { + "ok": True, + "report": report, + "train_hash": _ordered_id_sha256(train_ids), + "validation_hash": _ordered_id_sha256(validation_ids), + } + except BaseException as error: + local_status = {"ok": False, "error": f"{type(error).__name__}: {error}"} + + statuses: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(statuses, local_status) + failures: list[str] = [] + successes: list[Mapping[str, Any]] = [] + for rank, status in enumerate(statuses): + if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: + failures.append(f"rank {rank}: malformed loader authentication status") + continue + if not status["ok"]: + failures.append(f"rank {rank}: {status.get('error')}") + continue + successes.append(status) + if failures: + raise RuntimeError("PDD dataset validation failed: " + "; ".join(failures)) + canonical = {json.dumps(status, sort_keys=True) for status in successes} + if len(canonical) != 1: + raise RuntimeError( + "PDD ranks resolved different dataset roots, metadata, or split membership." + ) + return report, local_status["train_hash"], local_status["validation_hash"] + + +def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[bool, ...], ...]]: + import torch.distributed as dist + + from pdd.training import build_pdd_validation_assignments + + heldout_ids = tuple(str(value) for value in sampler.dataset.sample_ids) + assignments = build_pdd_validation_assignments( + heldout_ids, + config.pdd, + validation_seed=config.validation.seed, + ) + sampler.set_epoch(0) + sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) + local_plan = [ + tuple(str(sampler.dataset.sample_ids[index]) for index in batch) for batch in sampler + ] + sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) + plans: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(plans, local_plan) + batch_counts = {len(plan) for plan in plans} + if len(batch_counts) != 1: + raise RuntimeError("PDD validation sampler produced different batch counts across ranks.") + + masks = [[([False] * len(batch)) for batch in plan] for plan in plans] + seen: set[str] = set() + for batch_index in range(len(local_plan)): + for rank, plan in enumerate(plans): + for position, sample_id in enumerate(plan[batch_index]): + if sample_id not in seen: + masks[rank][batch_index][position] = True + seen.add(sample_id) + if seen != set(heldout_ids): + missing = sorted(set(heldout_ids) - seen) + extra = sorted(seen - set(heldout_ids)) + raise RuntimeError( + f"PDD validation sampler does not cover the held-out split: " + f"missing={missing[:5]}, extra={extra[:5]}." + ) + local_masks = tuple(tuple(batch) for batch in masks[dist.get_rank()]) + return assignments, local_masks + + +def _iter_validation_batches( + dataloader: Any, + masks: tuple[tuple[bool, ...], ...], + config: Any, + expected_latent_channels: int, + expected_condition_features: int, +): + from pdd.training import prepare_qwen_pdd_batch + + count = 0 + for count, (raw_batch, valid_mask) in enumerate(zip(dataloader, masks, strict=True), start=1): + prepared = prepare_qwen_pdd_batch( + raw_batch, + device=config.device, + dtype=config.dtype, + require_negative_condition=config.pdd.guidance_scale is not None, + expected_latent_channels=expected_latent_channels, + expected_condition_features=expected_condition_features, + ) + yield dataclasses.replace(prepared, valid_mask=valid_mask) + if count != len(masks): + raise RuntimeError( + f"PDD validation loader produced {count} batches for a {len(masks)}-batch plan." + ) + + +def _coverage_axis(counts: Any, loss_sums: Any) -> dict[int, dict[str, float | int]]: + return { + index: {"count": int(count), "mean_loss": float(loss_sums[index] / count)} + for index, count in enumerate(counts.tolist()) + if count + } + + +def _collective_training_iterator(dataloader: Any, sampler: Any) -> Any: + """Advance epochs and construct rank-local iterators under a collective error gate.""" + import torch.distributed as dist + + iterator = None + error_message = None + try: + if sampler.remaining_batches == 0: + sampler.set_epoch(sampler.epoch + 1) + iterator = iter(dataloader) + if iterator is None: + raise RuntimeError("PDD dataloader returned no iterator.") + except BaseException as error: + error_message = f"{type(error).__name__}: {error}" + errors: list[str | None] = [None] * dist.get_world_size() + dist.all_gather_object(errors, error_message) + failures = [f"rank {rank}: {message}" for rank, message in enumerate(errors) if message] + if failures: + raise RuntimeError( + "distributed PDD training iterator construction failed; " + "; ".join(failures) + ) + if iterator is None: + raise RuntimeError("local PDD iterator construction succeeded without an iterator.") + return iterator + + +def _collective_training_batch( + iterator: Any, + *, + sampler: Any, + resume: Any, + resume_pending: bool, + device: Any, + dtype: Any, + require_negative_condition: bool, + expected_batch_size: int, + expected_latent_channels: int, + expected_condition_features: int, +) -> tuple[Any, tuple[str, ...]] | None: + """Prepare one rank-local batch, then agree on success before any model call.""" + import torch.distributed as dist + + from pdd.training import prepare_qwen_pdd_batch + + prepared = None + sample_ids: tuple[str, ...] = () + status: dict[str, Any] + try: + raw_batch = next(iterator) + except StopIteration: + if resume_pending: + status = { + "state": "error", + "error": "RuntimeError: resumed dataloader ended before its first batch", + "resume_pending": True, + } + else: + status = {"state": "end", "resume_pending": False} + except BaseException as error: + status = { + "state": "error", + "error": f"{type(error).__name__}: {error}", + "resume_pending": resume_pending, + } + else: + try: + metadata = raw_batch["metadata"] + raw_ids = metadata.get("logical_sample_ids", metadata.get("sample_ids")) + if hasattr(raw_ids, "tolist"): + raw_ids = raw_ids.tolist() + sample_ids = tuple(str(value) for value in raw_ids) + expected_ids = sampler.expected_next_sample_ids() + if sample_ids != expected_ids: + raise RuntimeError( + "prefetched PDD batch does not match committed cursor: " + f"expected={expected_ids}, actual={sample_ids}." + ) + if resume_pending: + if resume is None: + raise RuntimeError("resume_pending is true without a PDD resume state.") + resume.verify_first_batch(sample_ids) + prepared = prepare_qwen_pdd_batch( + raw_batch, + device=device, + dtype=dtype, + require_negative_condition=require_negative_condition, + expected_latent_channels=expected_latent_channels, + expected_condition_features=expected_condition_features, + ) + if prepared is None: + raise RuntimeError("PDD batch preparation returned no prepared batch.") + if len(sample_ids) != expected_batch_size: + raise RuntimeError( + f"PDD training batch has {len(sample_ids)} samples; " + f"expected {expected_batch_size}." + ) + status = { + "state": "batch", + "batch_size": len(sample_ids), + "resume_pending": resume_pending, + } + except BaseException as error: + status = { + "state": "error", + "error": f"{type(error).__name__}: {error}", + "resume_pending": resume_pending, + } + + statuses: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(statuses, status) + malformed = [rank for rank, item in enumerate(statuses) if not isinstance(item, Mapping)] + if malformed: + raise RuntimeError(f"PDD training ranks returned malformed statuses: {malformed}.") + failures = [ + f"rank {rank}: {item.get('error')}" + for rank, item in enumerate(statuses) + if item.get("state") == "error" + ] + if failures: + raise RuntimeError( + "distributed PDD training batch preflight failed; " + "; ".join(failures) + ) + states = {item.get("state") for item in statuses} + if states == {"end"}: + return None + if states != {"batch"}: + raise RuntimeError( + "distributed PDD training ranks produced different dataloader lengths: " + f"{[item.get('state') for item in statuses]}." + ) + pending = {item.get("resume_pending") for item in statuses} + if len(pending) != 1: + raise RuntimeError("distributed PDD training ranks disagree on resume verification state.") + batch_sizes = {item.get("batch_size") for item in statuses} + if batch_sizes != {expected_batch_size}: + raise RuntimeError( + f"distributed PDD training ranks disagree on batch size: {sorted(batch_sizes)}." + ) + if prepared is None: + raise RuntimeError("local PDD batch preparation succeeded without a prepared batch.") + return prepared, sample_ids diff --git a/examples/diffusers/fastgen/pdd/finetune.py b/examples/diffusers/fastgen/pdd/finetune.py index aa3658a5b3a..fac8adffc2d 100644 --- a/examples/diffusers/fastgen/pdd/finetune.py +++ b/examples/diffusers/fastgen/pdd/finetune.py @@ -13,691 +13,63 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Train Qwen-Image with the ModelOpt-owned PDD lifecycle and released AutoModel APIs.""" +"""Entrypoint for Qwen-Image PDD training with released AutoModel components.""" from __future__ import annotations -import argparse -import dataclasses -import hashlib -import json import logging +import os import sys -import time -from collections.abc import Mapping -from pathlib import Path -from typing import Any -import yaml - -sys.dont_write_bytecode = True - -_THIS_DIR = Path(__file__).resolve().parent -_FASTGEN_DIR = _THIS_DIR.parent -_REPO_ROOT = _FASTGEN_DIR.parents[2] -# These entrypoints are also supported through ``python -m``. In that mode the -# sibling ModelOpt-owned example modules are not importable until this directory -# is added explicitly. +_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 str(path) not in sys.path: - sys.path.insert(0, str(path)) - + if path not in sys.path: + sys.path.insert(0, path) -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--config", - type=Path, - default=_THIS_DIR / "configs" / "qwen_image.yaml", - ) - return parser.parse_args() +_HELP = """\ +usage: finetune.py [--config CONFIG] [CONFIG_OVERRIDE ...] +Qwen-Image PDD training with released NeMo AutoModel components. -def _ordered_id_sha256(sample_ids: tuple[str, ...]) -> str: - digest = hashlib.sha256(b"modelopt-pdd-ordered-sample-ids-v1\0") - for sample_id in sample_ids: - digest.update(sample_id.encode()) - digest.update(b"\n") - return digest.hexdigest() +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 _dataloader_options(raw: Mapping[str, Any]) -> dict[str, Any]: - data = raw.get("data") - if not isinstance(data, Mapping) or not isinstance(data.get("dataloader"), Mapping): - raise TypeError("PDD config requires a data.dataloader mapping.") - options = dict(data["dataloader"]) - target = options.pop("_target_", None) - expected_target = "fastgen_data.build_text_to_image_multiresolution_dataloader" - if target != expected_target: - raise ValueError(f"PDD data.dataloader._target_ must be {expected_target!r}.") - if "base_resolution" in options: - options["base_resolution"] = tuple(options["base_resolution"]) - return options +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 -def _build_training_dataloader( - raw: Mapping[str, Any], - config: Any, - *, - dp_rank: int, - dp_world_size: int, -) -> tuple[Any, Any]: - from fastgen_data import ReplayableBatchSampler - from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader + from nemo_automodel.components.config._arg_parser import parse_args_and_load_config - options = _dataloader_options(raw) - if options.get("drop_last", True) is not True: - raise ValueError("PDD exact sample accounting requires data.dataloader.drop_last=true.") - if options.get("dynamic_batch_size", False) is not False: - raise ValueError("PDD v1 requires data.dataloader.dynamic_batch_size=false.") - options.update( - split="train", - validation_count=config.validation_count, - split_seed=config.split_seed, - dp_rank=dp_rank, - dp_world_size=dp_world_size, - exact_resume=True, - sampler_seed=config.training.seed, - loader_seed=config.training.seed, - ) - dataloader, sampler = build_text_to_image_multiresolution_dataloader(**options) - if not isinstance(sampler, ReplayableBatchSampler): - raise RuntimeError("PDD training requires the committed replayable batch sampler.") - if options.get("batch_size", 1) != config.training.local_batch_size: - raise RuntimeError("resolved local batch size does not match the built dataloader.") - return dataloader, sampler + from pdd.recipe import PDDDiffusionRecipe + cfg = parse_args_and_load_config(default_config_path) -def _build_validation_dataloader( - raw: Mapping[str, Any], - config: Any, - *, - dp_rank: int, - dp_world_size: int, -) -> tuple[Any, Any]: - from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader + import fastgen_data + import nemo_automodel - options = _dataloader_options(raw) - options.update( - split="validation", - validation_count=config.validation_count, - split_seed=config.split_seed, - dp_rank=dp_rank, - dp_world_size=dp_world_size, - drop_last=False, - shuffle=False, - dynamic_batch_size=False, - exact_resume=False, - sampler_seed=config.training.validation_seed, - loader_seed=config.training.validation_seed, + logging.info( + "[fastgen] vendored data package: %s", + os.path.dirname(os.path.abspath(fastgen_data.__file__)), ) - return build_text_to_image_multiresolution_dataloader(**options) - - -def _validate_dataset_contract( - train_dataset: Any, - validation_dataset: Any, - config: Any, -) -> tuple[Mapping[str, Any], str, str]: - """Collectively verify deterministic splits and the authenticated dataset snapshot.""" - import torch.distributed as dist - - try: - train_ids = tuple(str(value) for value in train_dataset.sample_ids) - validation_ids = tuple(str(value) for value in validation_dataset.sample_ids) - if len(validation_ids) != config.validation_count: - raise RuntimeError( - f"validation split has {len(validation_ids)} samples; " - f"expected {config.validation_count}." - ) - if set(train_ids).intersection(validation_ids): - raise RuntimeError("training and validation splits overlap.") - expected = {str(index) for index in range(train_dataset.total_num_samples)} - if set(train_ids).union(validation_ids) != expected: - raise RuntimeError("training and validation splits do not cover metadata.json.") - if train_dataset.total_num_samples != validation_dataset.total_num_samples: - raise RuntimeError("training and validation datasets disagree on total sample count.") - if train_dataset.metadata_sha256 != validation_dataset.metadata_sha256: - raise RuntimeError("training and validation datasets disagree on metadata content.") - if ( - not train_dataset.payload_hashes_complete - or not validation_dataset.payload_hashes_complete - ): - raise RuntimeError("PDD exact resume requires a cache_sha256 for every tensor payload.") - if train_dataset.dataset_snapshot_sha256 != validation_dataset.dataset_snapshot_sha256: - raise RuntimeError("training and validation datasets disagree on dataset content.") - if not isinstance(train_dataset.dataset_snapshot_sha256, str): - raise RuntimeError("PDD dataloader did not construct a dataset snapshot identity.") - report = { - "cache_root": str(train_dataset.cache_root), - "metadata_sha256": train_dataset.metadata_sha256, - "negative_prompt_embedding_sha256": (train_dataset.negative_prompt_embedding_sha256), - "dataset_snapshot_sha256": train_dataset.dataset_snapshot_sha256, - "total_samples": train_dataset.total_num_samples, - "train_samples": len(train_ids), - "validation_samples": len(validation_ids), - "split_seed": config.split_seed, - } - local_status: dict[str, Any] = { - "ok": True, - "report": report, - "train_hash": _ordered_id_sha256(train_ids), - "validation_hash": _ordered_id_sha256(validation_ids), - } - except BaseException as error: - local_status = {"ok": False, "error": f"{type(error).__name__}: {error}"} - - statuses: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(statuses, local_status) - failures: list[str] = [] - successes: list[Mapping[str, Any]] = [] - for rank, status in enumerate(statuses): - if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: - failures.append(f"rank {rank}: malformed loader authentication status") - continue - if not status["ok"]: - failures.append(f"rank {rank}: {status.get('error')}") - continue - successes.append(status) - if failures: - raise RuntimeError("PDD dataset validation failed: " + "; ".join(failures)) - canonical = {json.dumps(status, sort_keys=True) for status in successes} - if len(canonical) != 1: - raise RuntimeError( - "PDD ranks resolved different dataset roots, metadata, or split membership." - ) - return report, local_status["train_hash"], local_status["validation_hash"] - - -def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[bool, ...], ...]]: - import torch.distributed as dist - - from pdd.training import build_pdd_validation_assignments - - heldout_ids = tuple(str(value) for value in sampler.dataset.sample_ids) - assignments = build_pdd_validation_assignments( - heldout_ids, - config.pdd, - validation_seed=config.training.validation_seed, + logging.info( + "[fastgen] nemo_automodel resolved from: %s", + os.path.realpath(nemo_automodel.__file__), ) - sampler.set_epoch(0) - sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) - local_plan = [ - tuple(str(sampler.dataset.sample_ids[index]) for index in batch) for batch in sampler - ] - sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) - plans: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(plans, local_plan) - batch_counts = {len(plan) for plan in plans} - if len(batch_counts) != 1: - raise RuntimeError("PDD validation sampler produced different batch counts across ranks.") - - masks = [[([False] * len(batch)) for batch in plan] for plan in plans] - seen: set[str] = set() - for batch_index in range(len(local_plan)): - for rank, plan in enumerate(plans): - for position, sample_id in enumerate(plan[batch_index]): - if sample_id not in seen: - masks[rank][batch_index][position] = True - seen.add(sample_id) - if seen != set(heldout_ids): - missing = sorted(set(heldout_ids) - seen) - extra = sorted(seen - set(heldout_ids)) - raise RuntimeError( - f"PDD validation sampler does not cover the held-out split: " - f"missing={missing[:5]}, extra={extra[:5]}." - ) - local_masks = tuple(tuple(batch) for batch in masks[dist.get_rank()]) - return assignments, local_masks - - -def _iter_validation_batches( - dataloader: Any, - masks: tuple[tuple[bool, ...], ...], - config: Any, - expected_latent_channels: int, - expected_condition_features: int, -): - from pdd.training import prepare_qwen_pdd_batch - - count = 0 - for count, (raw_batch, valid_mask) in enumerate(zip(dataloader, masks, strict=True), start=1): - prepared = prepare_qwen_pdd_batch( - raw_batch, - device=config.device, - dtype=config.dtype, - require_negative_condition=config.pdd.guidance_scale is not None, - expected_latent_channels=expected_latent_channels, - expected_condition_features=expected_condition_features, - ) - yield dataclasses.replace(prepared, valid_mask=valid_mask) - if count != len(masks): - raise RuntimeError( - f"PDD validation loader produced {count} batches for a {len(masks)}-batch plan." - ) - - -def _coverage_axis(counts: Any, loss_sums: Any) -> dict[int, dict[str, float | int]]: - return { - index: {"count": int(count), "mean_loss": float(loss_sums[index] / count)} - for index, count in enumerate(counts.tolist()) - if count - } - - -def _collective_training_iterator(dataloader: Any, sampler: Any) -> Any: - """Advance epochs and construct rank-local iterators under a collective error gate.""" - import torch.distributed as dist - - iterator = None - error_message = None - try: - if sampler.remaining_batches == 0: - sampler.set_epoch(sampler.epoch + 1) - iterator = iter(dataloader) - if iterator is None: - raise RuntimeError("PDD dataloader returned no iterator.") - except BaseException as error: - error_message = f"{type(error).__name__}: {error}" - errors: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(errors, error_message) - failures = [f"rank {rank}: {message}" for rank, message in enumerate(errors) if message] - if failures: - raise RuntimeError( - "distributed PDD training iterator construction failed; " + "; ".join(failures) - ) - if iterator is None: - raise RuntimeError("local PDD iterator construction succeeded without an iterator.") - return iterator - - -def _collective_training_batch( - iterator: Any, - *, - sampler: Any, - resume: Any, - resume_pending: bool, - device: Any, - dtype: Any, - require_negative_condition: bool, - expected_batch_size: int, - expected_latent_channels: int, - expected_condition_features: int, -) -> tuple[Any, tuple[str, ...]] | None: - """Prepare one rank-local batch, then agree on success before any model call.""" - import torch.distributed as dist - - from pdd.training import prepare_qwen_pdd_batch - - prepared = None - sample_ids: tuple[str, ...] = () - status: dict[str, Any] - try: - raw_batch = next(iterator) - except StopIteration: - if resume_pending: - status = { - "state": "error", - "error": "RuntimeError: resumed dataloader ended before its first batch", - "resume_pending": True, - } - else: - status = {"state": "end", "resume_pending": False} - except BaseException as error: - status = { - "state": "error", - "error": f"{type(error).__name__}: {error}", - "resume_pending": resume_pending, - } - else: - try: - metadata = raw_batch["metadata"] - raw_ids = metadata.get("logical_sample_ids", metadata.get("sample_ids")) - if hasattr(raw_ids, "tolist"): - raw_ids = raw_ids.tolist() - sample_ids = tuple(str(value) for value in raw_ids) - expected_ids = sampler.expected_next_sample_ids() - if sample_ids != expected_ids: - raise RuntimeError( - "prefetched PDD batch does not match committed cursor: " - f"expected={expected_ids}, actual={sample_ids}." - ) - if resume_pending: - if resume is None: - raise RuntimeError("resume_pending is true without a PDD resume state.") - resume.verify_first_batch(sample_ids) - prepared = prepare_qwen_pdd_batch( - raw_batch, - device=device, - dtype=dtype, - require_negative_condition=require_negative_condition, - expected_latent_channels=expected_latent_channels, - expected_condition_features=expected_condition_features, - ) - if prepared is None: - raise RuntimeError("PDD batch preparation returned no prepared batch.") - if len(sample_ids) != expected_batch_size: - raise RuntimeError( - f"PDD training batch has {len(sample_ids)} samples; " - f"expected {expected_batch_size}." - ) - status = { - "state": "batch", - "batch_size": len(sample_ids), - "resume_pending": resume_pending, - } - except BaseException as error: - status = { - "state": "error", - "error": f"{type(error).__name__}: {error}", - "resume_pending": resume_pending, - } - - statuses: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(statuses, status) - malformed = [rank for rank, item in enumerate(statuses) if not isinstance(item, Mapping)] - if malformed: - raise RuntimeError(f"PDD training ranks returned malformed statuses: {malformed}.") - failures = [ - f"rank {rank}: {item.get('error')}" - for rank, item in enumerate(statuses) - if item.get("state") == "error" - ] - if failures: - raise RuntimeError( - "distributed PDD training batch preflight failed; " + "; ".join(failures) - ) - states = {item.get("state") for item in statuses} - if states == {"end"}: - return None - if states != {"batch"}: - raise RuntimeError( - "distributed PDD training ranks produced different dataloader lengths: " - f"{[item.get('state') for item in statuses]}." - ) - pending = {item.get("resume_pending") for item in statuses} - if len(pending) != 1: - raise RuntimeError("distributed PDD training ranks disagree on resume verification state.") - batch_sizes = {item.get("batch_size") for item in statuses} - if batch_sizes != {expected_batch_size}: - raise RuntimeError( - f"distributed PDD training ranks disagree on batch size: {sorted(batch_sizes)}." - ) - if prepared is None: - raise RuntimeError("local PDD batch preparation succeeded without a prepared batch.") - return prepared, sample_ids - - -def main() -> None: - args = _parse_args() - import torch - import torch.distributed as dist - - from pdd.checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity - from pdd.recipe import ( - build_pdd_setup, - build_pdd_training_artifacts, - initialize_pdd_distributed, - resolve_pdd_recipe_config, - ) - from pdd.training import run_pdd_validation - - raw = yaml.safe_load(args.config.read_text()) - config = resolve_pdd_recipe_config(raw) - initialize_pdd_distributed( - backend="nccl" if config.device.type == "cuda" else "gloo", - timeout_minutes=60, - ) - rank = dist.get_rank() - world_size = dist.get_world_size() - logging.basicConfig( - level=logging.INFO if rank == 0 else logging.WARNING, - format="%(asctime)s %(levelname)s %(message)s", - force=True, - ) - dataloader, sampler = _build_training_dataloader( - raw, - config, - dp_rank=rank, - dp_world_size=world_size, - ) - validation_dataloader, validation_sampler = _build_validation_dataloader( - raw, - config, - dp_rank=rank, - dp_world_size=world_size, - ) - snapshot_report, train_ordered_id_sha256, heldout_ordered_id_sha256 = ( - _validate_dataset_contract( - sampler.dataset, - validation_sampler.dataset, - config, - ) - ) - validation_assignments, validation_masks = _build_validation_plan( - validation_sampler, - config, - ) - setup = build_pdd_setup(config) - transformer_config = getattr(setup.student, "config", None) - if isinstance(transformer_config, Mapping): - in_channels = transformer_config.get("in_channels") - else: - in_channels = getattr(transformer_config, "in_channels", None) - if isinstance(transformer_config, Mapping): - condition_features = transformer_config.get("joint_attention_dim") - else: - condition_features = getattr(transformer_config, "joint_attention_dim", None) - if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: - raise RuntimeError("constructed Qwen transformer has invalid packed in_channels.") - if type(condition_features) is not int or condition_features <= 0: - raise RuntimeError("constructed Qwen transformer has invalid joint_attention_dim.") - expected_latent_channels = in_channels // 4 - expected_condition_features = condition_features - training = build_pdd_training_artifacts(setup, config) - identity = build_pdd_checkpoint_identity( - metadata=setup.metadata, - model_id=config.model_id, - model_revision=config.model_revision, - guidance_scale=config.pdd.guidance_scale, - guidance_rescale=config.guidance.rescale, - guidance_eps=config.guidance.eps, - automodel_snapshot=setup.automodel_snapshot, - ordered_train_id_sha256=train_ordered_id_sha256, - ordered_heldout_id_sha256=heldout_ordered_id_sha256, - dataset_snapshot_sha256=snapshot_report["dataset_snapshot_sha256"], - local_batch_size=config.training.local_batch_size, - grad_accumulation_steps=config.training.grad_accumulation_steps, - training_seed=config.training.seed, - validation_seed=config.training.validation_seed, - validation_every_steps=config.training.validation_every_steps, - max_grad_norm=config.training.max_grad_norm, - zero_grad_warmup_steps=config.training.zero_grad_warmup_steps, - activation_checkpointing=config.parallel.activation_checkpointing, - dtype=str(config.dtype).removeprefix("torch."), - optimizer=setup.optimizer, - scheduler=training.scheduler, - ) - checkpoint_manager = PDDCheckpointManager( - root=config.checkpoint.checkpoint_dir, - checkpointer=setup.checkpointer, - model=setup.student, - optimizer=setup.optimizer, - scheduler=training.scheduler, - trainer=training.trainer, - sampler=sampler, - rng=training.rng, - identity=identity, - ) - resume = checkpoint_manager.load(config.checkpoint.restore_from) - resume_pending = resume is not None - if resume is not None and rank == 0: - logging.info( - "PDD resume selected: checkpoint=%s parent=%s step=%d sample_slots=%d " - "expected_first_sample_ids=%s", - resume.checkpoint_path, - resume.parent_checkpoint, - resume.completed_steps, - resume.sample_slots_consumed, - resume.expected_next_sample_ids, - ) - if rank == 0: - logging.info( - "PDD dataset verified: snapshot_sha256=%s metadata_sha256=%s " - "train=%d validation=%d root=%s", - snapshot_report["dataset_snapshot_sha256"], - snapshot_report["metadata_sha256"], - snapshot_report["train_samples"], - snapshot_report["validation_samples"], - snapshot_report["cache_root"], - ) - logging.info( - "PDD setup complete: lifecycle=%s student_keys=%d AutoModel=%s", - setup.lifecycle, - len(setup.checkpoint_keys), - setup.automodel_snapshot["version"], - ) - last_saved_step = 0 if resume is None else resume.completed_steps - try: - while training.trainer.completed_steps < config.training.max_steps: - iterator = _collective_training_iterator(dataloader, sampler) - data_wait_started = time.perf_counter() - while training.trainer.completed_steps < config.training.max_steps: - next_batch = _collective_training_batch( - iterator, - sampler=sampler, - resume=resume, - resume_pending=resume_pending, - device=config.device, - dtype=config.dtype, - require_negative_condition=config.pdd.guidance_scale is not None, - expected_batch_size=config.training.local_batch_size, - expected_latent_channels=expected_latent_channels, - expected_condition_features=expected_condition_features, - ) - if next_batch is None: - break - data_wait_seconds = time.perf_counter() - data_wait_started - step_started = time.perf_counter() - batch, sample_ids = next_batch - if resume_pending: - if rank == 0: - logging.info( - "PDD resume first batch verified: checkpoint=%s sample_ids=%s", - resume.checkpoint_path, - sample_ids, - ) - resume_pending = False - measure_update = ( - training.trainer.completed_steps + 1 - ) % config.training.log_every_steps == 0 - diagnostics = training.trainer.train_step( - batch, - measure_updates=measure_update, - ) - training.scheduler.step() - sampler.commit(sample_ids) - if sampler.remaining_batches == 0: - sampler.set_epoch(sampler.epoch + 1) - step_seconds = time.perf_counter() - step_started - - if diagnostics.completed_step % config.training.log_every_steps == 0: - timing = torch.tensor( - [data_wait_seconds, step_seconds], - dtype=torch.float64, - device=config.device, - ) - dist.all_reduce(timing, op=dist.ReduceOp.MAX) - peak_memory = ( - torch.cuda.max_memory_allocated(config.device) - if config.device.type == "cuda" - else 0 - ) - memory = torch.tensor(peak_memory, dtype=torch.int64, device=config.device) - dist.all_reduce(memory, op=dist.ReduceOp.MAX) - global_samples = config.training.local_batch_size * dist.get_world_size() - throughput = global_samples / max(float(timing[1].item()), 1e-12) - coverage = training.trainer.coverage - bin_loss = [ - None if count == 0 else float(loss_sum / count) - for loss_sum, count in zip( - coverage.bin_loss_sums.tolist(), - coverage.bin_counts.tolist(), - ) - ] - if rank == 0: - logging.info( - "PDD step=%d loss=%.6g grad_norm=%.6g nominal_update_ratio=%.6g " - "projection_update_ratio=%s lr=%.6g student_rms=%.6g " - "teacher_rms=%.6g student_teacher_rms_ratio=%.6g " - "reconstruction_rms=%.6g pairs=%d n_coverage=%s k_coverage=%s " - "bins=%s bin_loss=%s samples_per_second=%.3f " - "data_wait_seconds=%.4f peak_memory_bytes=%d", - diagnostics.completed_step, - diagnostics.loss, - diagnostics.grad_norm, - diagnostics.student_adamw_nominal_update_ratio, - diagnostics.pdd_projection_update_ratio, - diagnostics.learning_rate, - diagnostics.student_velocity_rms, - diagnostics.teacher_velocity_rms, - diagnostics.student_teacher_velocity_rms_ratio, - diagnostics.reconstructed_state_rms, - int((coverage.pair_counts > 0).sum()), - _coverage_axis(coverage.n_counts, coverage.n_loss_sums), - _coverage_axis(coverage.k_counts, coverage.k_loss_sums), - coverage.bin_counts.tolist(), - bin_loss, - throughput, - float(timing[0].item()), - int(memory.item()), - ) - if config.device.type == "cuda": - torch.cuda.reset_peak_memory_stats(config.device) - if ( - diagnostics.completed_step % config.training.validation_every_steps == 0 - or diagnostics.completed_step >= config.training.max_steps - ): - validation_sampler.set_epoch(0) - validation_sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) - validation_result = run_pdd_validation( - training.pipeline, - _iter_validation_batches( - validation_dataloader, - validation_masks, - config, - expected_latent_channels, - expected_condition_features, - ), - validation_assignments, - validation_seed=config.training.validation_seed, - ) - if rank == 0: - logging.info( - "PDD validation step=%d loss=%.12g pairs=%d starts=%d heads=%d " - "ordered_id_sha256=%s records=%d", - diagnostics.completed_step, - validation_result.mean_loss, - validation_result.pair_count, - validation_result.start_count, - validation_result.head_count, - validation_result.ordered_id_sha256, - len(validation_result.records), - ) - if ( - config.checkpoint.enabled - and diagnostics.completed_step % config.training.checkpoint_every_steps == 0 - ): - checkpoint_manager.save() - last_saved_step = diagnostics.completed_step - if diagnostics.completed_step >= config.training.max_steps: - break - data_wait_started = time.perf_counter() - if config.checkpoint.enabled and last_saved_step != training.trainer.completed_steps: - checkpoint_manager.save() - finally: - setup.checkpointer.close() + recipe = PDDDiffusionRecipe(cfg) + recipe.setup() + recipe.run_train_validation_loop() if __name__ == "__main__": diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 610c93fd1f4..e8c39710f2d 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -20,6 +20,7 @@ import copy import logging import math +import time from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -35,6 +36,17 @@ convert_qwen_image_to_pdd, ) +from .checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity +from .data import ( + _build_training_dataloader, + _build_validation_dataloader, + _build_validation_plan, + _collective_training_batch, + _collective_training_iterator, + _coverage_axis, + _iter_validation_batches, + _validate_dataset_contract, +) from .verify_readonly_automodel import snapshot_installed_distribution @@ -58,20 +70,34 @@ class PDDCheckpointConfig: @dataclass(frozen=True) -class PDDTrainingConfig: - """Direct-update and observability settings for the standalone PDD lifecycle.""" +class PDDStepSchedulerConfig: + """AutoModel-compatible batch, cadence, and termination settings.""" - seed: int = 42 max_steps: int = 10_000 - max_grad_norm: float = 1.0 - zero_grad_warmup_steps: int = 0 - log_every_steps: int = 10 - checkpoint_every_steps: int = 1_000 - validation_every_steps: int = 1_000 + num_epochs: int = 200 + log_every: int = 10 + ckpt_every_steps: int = 1_000 local_batch_size: int = 1 global_batch_size: int | None = None - grad_accumulation_steps: int = 1 - validation_seed: int = 2026 + save_checkpoint_every_epoch: bool = False + + +@dataclass(frozen=True) +class PDDTrainingHealthConfig: + """PDD-only update health settings.""" + + max_grad_norm: float = 1.0 + zero_grad_warmup_steps: int = 0 + + +@dataclass(frozen=True) +class PDDValidationConfig: + """Deterministic held-out validation settings.""" + + count: int = 2_000 + seed: int = 2026 + split_seed: int = 2026 + every_steps: int = 1_000 @dataclass(frozen=True) @@ -91,14 +117,15 @@ class PDDRecipeConfig: pdd: PDDConfig parallel: PDDParallelConfig checkpoint: PDDCheckpointConfig - training: PDDTrainingConfig + step_scheduler: PDDStepSchedulerConfig + training_health: PDDTrainingHealthConfig + validation: PDDValidationConfig guidance: PDDGuidanceConfig + seed: int learning_rate: float weight_decay: float adam_betas: tuple[float, float] adam_eps: float - validation_count: int - split_seed: int device: torch.device dtype: torch.dtype fuse_qkv_projections: bool @@ -155,6 +182,22 @@ def _as_mapping(value: Any, *, name: str) -> Mapping[str, Any]: return value +def _config_to_mapping(value: Any) -> Mapping[str, Any]: + """Materialize AutoModel config targets as their original YAML dotted paths.""" + if isinstance(value, Mapping): + return value + to_yaml_dict = getattr(value, "to_yaml_dict", None) + if callable(to_yaml_dict): + return _as_mapping( + to_yaml_dict(resolve_env=True, use_orig_values=True), + name="ConfigNode.to_yaml_dict() result", + ) + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + return _as_mapping(to_dict(), name="config.to_dict() result") + raise TypeError(f"config must be a mapping or ConfigNode, got {type(value).__name__}.") + + def _reject_enabled(value: Any, *, name: str) -> None: if value is None or value is False or value == {}: return @@ -201,19 +244,56 @@ def _resolve_dtype(value: Any) -> torch.dtype: ) from error -def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: - """Resolve PDD and reject TE-linear, PEFT, and guidance embeddings before loading.""" - raw = _as_mapping(raw, name="config") +def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: + """Resolve one canonical DMD2-shaped PDD configuration.""" + raw = _config_to_mapping(raw) + + legacy_training = _as_mapping(raw.get("training", {}), name="training") + legacy_replacements = { + "seed": "seed", + "max_steps": "step_scheduler.max_steps", + "log_every_steps": "step_scheduler.log_every", + "checkpoint_every_steps": "step_scheduler.ckpt_every_steps", + "validation_every_steps": "validation.every_steps", + "local_batch_size": "step_scheduler.local_batch_size", + "global_batch_size": "step_scheduler.global_batch_size", + "grad_accumulation_steps": "step_scheduler.global_batch_size", + "validation_seed": "validation.seed", + "max_grad_norm": "training_health.max_grad_norm", + "zero_grad_warmup_steps": "training_health.zero_grad_warmup_steps", + } + if legacy_training: + key = next(iter(legacy_training)) + replacement_key = legacy_replacements.get(key, "the canonical PDD schema") + raise ValueError(f"training.{key} is unsupported; use {replacement_key}.") + model = _as_mapping(raw.get("model"), name="model") pdd_raw = _as_mapping(raw.get("pdd"), name="pdd") fsdp = _as_mapping(raw.get("fsdp", {}), name="fsdp") optim = _as_mapping(raw.get("optim", {}), name="optim") + optimizer_cfg = _as_mapping(optim.get("optimizer", {}), name="optim.optimizer") + lr_scheduler = _as_mapping(raw.get("lr_scheduler", {}), name="lr_scheduler") + step_scheduler = _as_mapping(raw.get("step_scheduler", {}), name="step_scheduler") + training_health = _as_mapping(raw.get("training_health", {}), name="training_health") + validation = _as_mapping(raw.get("validation", {}), name="validation") checkpoint = _as_mapping(raw.get("checkpoint", {}), name="checkpoint") - training = _as_mapping(raw.get("training", {}), name="training") guidance = _as_mapping(raw.get("guidance", {}), name="guidance") data = _as_mapping(raw.get("data", {}), name="data") dataloader = _as_mapping(data.get("dataloader", {}), name="data.dataloader") + for legacy_key in ("weight_decay", "betas", "eps"): + if legacy_key in optim: + raise ValueError( + f"optim.{legacy_key} is unsupported; use optim.optimizer.{legacy_key}." + ) + + for legacy_key, replacement_key in ( + ("validation_count", "validation.count"), + ("split_seed", "validation.split_seed"), + ): + if legacy_key in data: + raise ValueError(f"data.{legacy_key} is unsupported; use {replacement_key}.") + target = dataloader.get("_target_") expected_target = "fastgen_data.build_text_to_image_multiresolution_dataloader" if target is not None and target != expected_target: @@ -236,16 +316,6 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: "PDD uses deterministic ordinal splits from metadata.json; " "data.dataloader.metadata_index is unsupported." ) - validation_count = _require_int_at_least( - data.get("validation_count", 2_000), - name="data.validation_count", - minimum=1, - ) - split_seed = _require_int_at_least( - data.get("split_seed", 2026), - name="data.split_seed", - minimum=0, - ) _reject_enabled(model.get("transformer_engine_linear"), name="global TE-linear conversion") _reject_enabled(model.get("peft"), name="PEFT/LoRA") @@ -274,20 +344,43 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: or any(character not in "0123456789abcdefABCDEF" for character in model_revision) ): raise ValueError("model.revision must be null or a full 40-character commit hash.") - if not Path(model_id).is_dir(): - if model_revision is None: - raise ValueError("Remote PDD models require an exact model.revision commit hash.") - learning_rate = optim.get("learning_rate", 2.0e-5) - weight_decay = optim.get("weight_decay", 0.0) - if isinstance(learning_rate, bool) or not isinstance(learning_rate, int | float): - raise TypeError("optim.learning_rate must be a real number.") - if isinstance(weight_decay, bool) or not isinstance(weight_decay, int | float): - raise TypeError("optim.weight_decay must be a real number.") - if not math.isfinite(learning_rate) or not math.isfinite(weight_decay): - raise ValueError("optim.learning_rate and weight_decay must be finite.") - if learning_rate <= 0 or weight_decay < 0: - raise ValueError("optim.learning_rate must be > 0 and weight_decay must be >= 0.") - adam_betas_raw = optim.get("betas", [0.9, 0.999]) + if not Path(model_id).is_dir() and model_revision is None: + raise ValueError("Remote PDD models require an exact model.revision commit hash.") + + learning_rate = _require_finite_real( + optim.get("learning_rate", 2.0e-5), + name="optim.learning_rate", + minimum=0.0, + ) + if learning_rate == 0.0: + raise ValueError("optim.learning_rate must be > 0.") + optimizer_target = optimizer_cfg.get("_target_", "torch.optim.AdamW") + if optimizer_target != "torch.optim.AdamW": + raise ValueError("PDD v1 requires optim.optimizer._target_='torch.optim.AdamW'.") + allowed_optimizer_keys = { + "_target_", + "weight_decay", + "betas", + "eps", + "amsgrad", + "capturable", + "differentiable", + "foreach", + "fused", + "maximize", + } + unsupported_optimizer_keys = sorted(set(optimizer_cfg) - allowed_optimizer_keys) + if unsupported_optimizer_keys: + raise ValueError(f"unsupported PDD optimizer keys: {unsupported_optimizer_keys}.") + for flag in ("amsgrad", "capturable", "differentiable", "foreach", "fused", "maximize"): + if _require_bool(optimizer_cfg.get(flag, False), name=f"optim.optimizer.{flag}"): + raise ValueError(f"PDD v1 requires optim.optimizer.{flag}=false.") + weight_decay = _require_finite_real( + optimizer_cfg.get("weight_decay", 0.0), + name="optim.optimizer.weight_decay", + minimum=0.0, + ) + adam_betas_raw = optimizer_cfg.get("betas", [0.9, 0.999]) if ( not isinstance(adam_betas_raw, list | tuple) or len(adam_betas_raw) != 2 @@ -295,17 +388,45 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: isinstance(beta, bool) or not isinstance(beta, int | float) for beta in adam_betas_raw ) ): - raise TypeError("optim.betas must contain two real numbers.") + raise TypeError("optim.optimizer.betas must contain two real numbers.") adam_betas = (float(adam_betas_raw[0]), float(adam_betas_raw[1])) if any(not math.isfinite(beta) or not 0.0 <= beta < 1.0 for beta in adam_betas): - raise ValueError("optim.betas values must be finite and in [0, 1).") + raise ValueError("optim.optimizer.betas values must be finite and in [0, 1).") adam_eps = _require_finite_real( - optim.get("eps", 1e-8), - name="optim.eps", + optimizer_cfg.get("eps", 1e-8), + name="optim.optimizer.eps", minimum=0.0, ) if adam_eps == 0.0: - raise ValueError("optim.eps must be > 0.") + raise ValueError("optim.optimizer.eps must be > 0.") + + lr_decay_style = lr_scheduler.get("lr_decay_style", "constant") + if lr_decay_style != "constant": + raise ValueError("PDD v1 requires lr_scheduler.lr_decay_style='constant'.") + lr_warmup_steps = _require_int_at_least( + lr_scheduler.get("lr_warmup_steps", 0), + name="lr_scheduler.lr_warmup_steps", + minimum=0, + ) + if lr_warmup_steps != 0: + raise ValueError("PDD v1 requires lr_scheduler.lr_warmup_steps=0.") + min_lr = _require_finite_real( + lr_scheduler.get("min_lr", learning_rate), + name="lr_scheduler.min_lr", + minimum=0.0, + ) + if min_lr != learning_rate: + raise ValueError("lr_scheduler.min_lr must equal optim.learning_rate for constant PDD LR.") + if "max_lr" in lr_scheduler: + max_lr = _require_finite_real( + lr_scheduler["max_lr"], + name="lr_scheduler.max_lr", + minimum=0.0, + ) + if max_lr != learning_rate: + raise ValueError( + "lr_scheduler.max_lr must equal optim.learning_rate for constant PDD LR." + ) dp_size = fsdp.get("dp_size") if dp_size is not None and (type(dp_size) is not int or dp_size < 1): @@ -324,7 +445,8 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: checkpoint_enabled = _require_bool(checkpoint.get("enabled", True), name="checkpoint.enabled") save_consolidated = _require_bool( - checkpoint.get("save_consolidated", False), name="checkpoint.save_consolidated" + checkpoint.get("save_consolidated", False), + name="checkpoint.save_consolidated", ) if save_consolidated: raise ValueError("PDD training checkpoints require checkpoint.save_consolidated=false.") @@ -334,70 +456,98 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: model_save_format = checkpoint.get("model_save_format", "torch_save") if model_save_format != "torch_save": raise ValueError("PDD training checkpoints require model_save_format='torch_save'.") - fuse_qkv_projections = _require_bool( - model.get("fuse_qkv_projections", False), name="model.fuse_qkv_projections" - ) restore_from = checkpoint.get("restore_from") if restore_from is not None and (not isinstance(restore_from, str) or not restore_from): raise ValueError("checkpoint.restore_from must be null or a non-empty string.") if not checkpoint_enabled and restore_from is not None: raise ValueError("checkpoint.restore_from requires checkpoint.enabled=true.") + fuse_qkv_projections = _require_bool( + model.get("fuse_qkv_projections", False), + name="model.fuse_qkv_projections", + ) - seed = _require_int_at_least(training.get("seed", 42), name="training.seed", minimum=0) + seed = _require_int_at_least(raw.get("seed", 42), name="seed", minimum=0) max_steps = _require_int_at_least( - training.get("max_steps", 10_000), name="training.max_steps", minimum=1 - ) - zero_grad_warmup_steps = _require_int_at_least( - training.get("zero_grad_warmup_steps", 0), - name="training.zero_grad_warmup_steps", - minimum=0, - ) - log_every_steps = _require_int_at_least( - training.get("log_every_steps", 10), - name="training.log_every_steps", + step_scheduler.get("max_steps", 10_000), + name="step_scheduler.max_steps", minimum=1, ) - checkpoint_every_steps = _require_int_at_least( - training.get("checkpoint_every_steps", 1_000), - name="training.checkpoint_every_steps", + num_epochs = _require_int_at_least( + step_scheduler.get("num_epochs", 200), + name="step_scheduler.num_epochs", minimum=1, ) - validation_every_steps = _require_int_at_least( - training.get("validation_every_steps", 1_000), - name="training.validation_every_steps", + log_every = _require_int_at_least( + step_scheduler.get("log_every", 10), + name="step_scheduler.log_every", minimum=1, ) - grad_accumulation_steps = _require_int_at_least( - training.get("grad_accumulation_steps", 1), - name="training.grad_accumulation_steps", + ckpt_every_steps = _require_int_at_least( + step_scheduler.get("ckpt_every_steps", 1_000), + name="step_scheduler.ckpt_every_steps", minimum=1, ) - if grad_accumulation_steps != 1: - raise ValueError("PDD v1 exact resume requires training.grad_accumulation_steps=1.") + save_checkpoint_every_epoch = _require_bool( + step_scheduler.get("save_checkpoint_every_epoch", False), + name="step_scheduler.save_checkpoint_every_epoch", + ) + if save_checkpoint_every_epoch: + raise ValueError( + "PDD exact resume requires step_scheduler.save_checkpoint_every_epoch=false." + ) local_batch_size = _require_int_at_least( - dataloader.get("batch_size", training.get("local_batch_size", 1)), + step_scheduler.get("local_batch_size", 1), + name="step_scheduler.local_batch_size", + minimum=1, + ) + data_batch_size = _require_int_at_least( + dataloader.get("batch_size", local_batch_size), name="data.dataloader.batch_size", minimum=1, ) - global_batch_size = training.get("global_batch_size") + if data_batch_size != local_batch_size: + raise ValueError("data.dataloader.batch_size must equal step_scheduler.local_batch_size.") + global_batch_size = step_scheduler.get("global_batch_size") if global_batch_size is not None: global_batch_size = _require_int_at_least( global_batch_size, - name="training.global_batch_size", + name="step_scheduler.global_batch_size", minimum=1, ) - validation_seed = _require_int_at_least( - training.get("validation_seed", 2026), - name="training.validation_seed", - minimum=0, - ) + max_grad_norm = _require_finite_real( - training.get("max_grad_norm", 1.0), - name="training.max_grad_norm", + training_health.get("max_grad_norm", 1.0), + name="training_health.max_grad_norm", minimum=0.0, ) if max_grad_norm == 0.0: - raise ValueError("training.max_grad_norm must be > 0.") + raise ValueError("training_health.max_grad_norm must be > 0.") + zero_grad_warmup_steps = _require_int_at_least( + training_health.get("zero_grad_warmup_steps", 0), + name="training_health.zero_grad_warmup_steps", + minimum=0, + ) + validation_count = _require_int_at_least( + validation.get("count", 2_000), + name="validation.count", + minimum=1, + ) + validation_seed = _require_int_at_least( + validation.get("seed", 2026), + name="validation.seed", + minimum=0, + ) + split_seed = _require_int_at_least( + validation.get("split_seed", 2026), + name="validation.split_seed", + minimum=0, + ) + validation_every_steps = _require_int_at_least( + validation.get("every_steps", 1_000), + name="validation.every_steps", + minimum=1, + ) + guidance_rescale = _require_finite_real( guidance.get("rescale", 1.0), name="guidance.rescale", @@ -428,26 +578,31 @@ def resolve_pdd_recipe_config(raw: Mapping[str, Any]) -> PDDRecipeConfig: restore_from=restore_from, save_consolidated=save_consolidated, ), - training=PDDTrainingConfig( - seed=seed, + step_scheduler=PDDStepSchedulerConfig( max_steps=max_steps, - max_grad_norm=max_grad_norm, - zero_grad_warmup_steps=zero_grad_warmup_steps, - log_every_steps=log_every_steps, - checkpoint_every_steps=checkpoint_every_steps, - validation_every_steps=validation_every_steps, + num_epochs=num_epochs, + log_every=log_every, + ckpt_every_steps=ckpt_every_steps, local_batch_size=local_batch_size, global_batch_size=global_batch_size, - grad_accumulation_steps=grad_accumulation_steps, - validation_seed=validation_seed, + save_checkpoint_every_epoch=save_checkpoint_every_epoch, + ), + training_health=PDDTrainingHealthConfig( + max_grad_norm=max_grad_norm, + zero_grad_warmup_steps=zero_grad_warmup_steps, + ), + validation=PDDValidationConfig( + count=validation_count, + seed=validation_seed, + split_seed=split_seed, + every_steps=validation_every_steps, ), guidance=PDDGuidanceConfig(rescale=guidance_rescale, eps=guidance_eps), + seed=seed, learning_rate=float(learning_rate), weight_decay=float(weight_decay), adam_betas=adam_betas, adam_eps=adam_eps, - validation_count=validation_count, - split_seed=split_seed, device=torch.device(model.get("device", "cuda" if torch.cuda.is_available() else "cpu")), dtype=_resolve_dtype(model.get("torch_dtype", "bfloat16")), fuse_qkv_projections=fuse_qkv_projections, @@ -628,17 +783,16 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: lifecycle.append("pdd_conversion") world_size = dist.get_world_size() - if config.training.global_batch_size is not None: - effective_global_batch = ( - config.training.local_batch_size * world_size * config.training.grad_accumulation_steps - ) - if effective_global_batch != config.training.global_batch_size: + if config.step_scheduler.global_batch_size is not None: + effective_global_batch = config.step_scheduler.local_batch_size * world_size + if effective_global_batch != config.step_scheduler.global_batch_size: raise ValueError( "PDD global batch mismatch: " - f"local_batch_size={config.training.local_batch_size} * world_size={world_size} " - f"* grad_accumulation_steps={config.training.grad_accumulation_steps} " - f"= {effective_global_batch}, configured " - f"training.global_batch_size={config.training.global_batch_size}." + f"local_batch_size={config.step_scheduler.local_batch_size} * " + f"world_size={world_size} = {effective_global_batch}, configured " + "step_scheduler.global_batch_size=" + f"{config.step_scheduler.global_batch_size}. PDD v1 requires one microbatch per " + "optimizer update." ) dp_size = config.parallel.dp_size or world_size if dp_size != world_size: @@ -862,13 +1016,13 @@ def build_pdd_training_artifacts( ) pipeline = PDDPipeline(setup.student, setup.teacher, config.pdd, adapter) scheduler = torch.optim.lr_scheduler.LambdaLR(setup.optimizer, lr_lambda=lambda _: 1.0) - rng = StatefulRNG(config.training.seed, ranked=True) + rng = StatefulRNG(config.seed, ranked=True) trainer = PDDTrainer( pipeline, setup.optimizer, projection=setup.projection, - max_grad_norm=config.training.max_grad_norm, - warmup_steps=config.training.zero_grad_warmup_steps, + max_grad_norm=config.training_health.max_grad_norm, + warmup_steps=config.training_health.zero_grad_warmup_steps, ) return PDDTrainingArtifacts( pipeline=pipeline, @@ -884,3 +1038,319 @@ def initialize_pdd_distributed(*, backend: str, timeout_minutes: int = 60) -> An from nemo_automodel.components.distributed import initialize_distributed return initialize_distributed(backend=backend, timeout_minutes=timeout_minutes) + + +class PDDDiffusionRecipe: + """Compose released AutoModel components around the PDD-specific update.""" + + def __init__(self, cfg: Any) -> None: + self.cfg = cfg + self.raw_config = _config_to_mapping(cfg) + self.config = resolve_pdd_recipe_config(self.raw_config) + + def setup(self) -> None: + """Build data, converted models, AutoModel scheduling, and strict resume state.""" + config = self.config + self.dist_env = initialize_pdd_distributed( + backend="nccl" if config.device.type == "cuda" else "gloo", + timeout_minutes=60, + ) + from nemo_automodel.components.loggers.log_utils import setup_logging + from nemo_automodel.components.training.step_scheduler import StepScheduler + + setup_logging() + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + self.dataloader, self.sampler = _build_training_dataloader( + self.raw_config, + config, + dp_rank=self.rank, + dp_world_size=self.world_size, + ) + self.validation_dataloader, self.validation_sampler = _build_validation_dataloader( + self.raw_config, + config, + dp_rank=self.rank, + dp_world_size=self.world_size, + ) + ( + self.snapshot_report, + train_ordered_id_sha256, + heldout_ordered_id_sha256, + ) = _validate_dataset_contract( + self.sampler.dataset, + self.validation_sampler.dataset, + config, + ) + self.validation_assignments, self.validation_masks = _build_validation_plan( + self.validation_sampler, + config, + ) + + self.setup_artifacts = build_pdd_setup(config) + self.expected_latent_channels, self.expected_condition_features = ( + self._resolve_transformer_dimensions(self.setup_artifacts.student) + ) + self.training = build_pdd_training_artifacts(self.setup_artifacts, config) + global_batch_size = ( + config.step_scheduler.global_batch_size + or config.step_scheduler.local_batch_size * self.world_size + ) + self.step_scheduler = StepScheduler( + global_batch_size=global_batch_size, + local_batch_size=config.step_scheduler.local_batch_size, + dp_size=self.world_size, + ckpt_every_steps=config.step_scheduler.ckpt_every_steps, + save_checkpoint_every_epoch=False, + dataloader=self.dataloader, + val_every_steps=None, + start_step=0, + start_epoch=0, + num_epochs=config.step_scheduler.num_epochs, + max_steps=config.step_scheduler.max_steps, + ) + if self.step_scheduler.grad_acc_steps != 1: + raise ValueError("PDD v1 requires exactly one microbatch per optimizer update.") + + identity = build_pdd_checkpoint_identity( + metadata=self.setup_artifacts.metadata, + model_id=config.model_id, + model_revision=config.model_revision, + guidance_scale=config.pdd.guidance_scale, + guidance_rescale=config.guidance.rescale, + guidance_eps=config.guidance.eps, + automodel_snapshot=self.setup_artifacts.automodel_snapshot, + ordered_train_id_sha256=train_ordered_id_sha256, + ordered_heldout_id_sha256=heldout_ordered_id_sha256, + dataset_snapshot_sha256=self.snapshot_report["dataset_snapshot_sha256"], + local_batch_size=config.step_scheduler.local_batch_size, + grad_accumulation_steps=1, + training_seed=config.seed, + validation_seed=config.validation.seed, + validation_every_steps=config.validation.every_steps, + max_grad_norm=config.training_health.max_grad_norm, + zero_grad_warmup_steps=config.training_health.zero_grad_warmup_steps, + activation_checkpointing=config.parallel.activation_checkpointing, + dtype=str(config.dtype).removeprefix("torch."), + optimizer=self.setup_artifacts.optimizer, + scheduler=self.training.scheduler, + ) + self.checkpoint_manager = PDDCheckpointManager( + root=config.checkpoint.checkpoint_dir, + checkpointer=self.setup_artifacts.checkpointer, + model=self.setup_artifacts.student, + optimizer=self.setup_artifacts.optimizer, + scheduler=self.training.scheduler, + step_scheduler=self.step_scheduler, + trainer=self.training.trainer, + sampler=self.sampler, + rng=self.training.rng, + identity=identity, + ) + self.resume = self.checkpoint_manager.load(config.checkpoint.restore_from) + self.resume_pending = self.resume is not None + self._log_setup() + + @staticmethod + def _resolve_transformer_dimensions(student: nn.Module) -> tuple[int, int]: + transformer_config = getattr(student, "config", None) + if isinstance(transformer_config, Mapping): + in_channels = transformer_config.get("in_channels") + condition_features = transformer_config.get("joint_attention_dim") + else: + in_channels = getattr(transformer_config, "in_channels", None) + condition_features = getattr(transformer_config, "joint_attention_dim", None) + if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: + raise RuntimeError("constructed Qwen transformer has invalid packed in_channels.") + if type(condition_features) is not int or condition_features <= 0: + raise RuntimeError("constructed Qwen transformer has invalid joint_attention_dim.") + return in_channels // 4, condition_features + + def _log_setup(self) -> None: + if self.rank != 0: + return + if self.resume is not None: + logging.info( + "PDD resume selected: checkpoint=%s parent=%s step=%d sample_slots=%d " + "expected_first_sample_ids=%s", + self.resume.checkpoint_path, + self.resume.parent_checkpoint, + self.resume.completed_steps, + self.resume.sample_slots_consumed, + self.resume.expected_next_sample_ids, + ) + logging.info( + "PDD dataset verified: snapshot_sha256=%s metadata_sha256=%s " + "train=%d validation=%d root=%s", + self.snapshot_report["dataset_snapshot_sha256"], + self.snapshot_report["metadata_sha256"], + self.snapshot_report["train_samples"], + self.snapshot_report["validation_samples"], + self.snapshot_report["cache_root"], + ) + logging.info( + "PDD setup complete: lifecycle=%s student_keys=%d AutoModel=%s", + self.setup_artifacts.lifecycle, + len(self.setup_artifacts.checkpoint_keys), + self.setup_artifacts.automodel_snapshot["version"], + ) + + def _prepared_training_batches(self): + iterator = _collective_training_iterator(self.dataloader, self.sampler) + while True: + next_batch = _collective_training_batch( + iterator, + sampler=self.sampler, + resume=self.resume, + resume_pending=self.resume_pending, + device=self.config.device, + dtype=self.config.dtype, + require_negative_condition=self.config.pdd.guidance_scale is not None, + expected_batch_size=self.config.step_scheduler.local_batch_size, + expected_latent_channels=self.expected_latent_channels, + expected_condition_features=self.expected_condition_features, + ) + if next_batch is None: + return + yield next_batch + + def _run_validation(self, completed_step: int) -> None: + from .training import run_pdd_validation + + self.validation_sampler.set_epoch(0) + self.validation_sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) + result = run_pdd_validation( + self.training.pipeline, + _iter_validation_batches( + self.validation_dataloader, + self.validation_masks, + self.config, + self.expected_latent_channels, + self.expected_condition_features, + ), + self.validation_assignments, + validation_seed=self.config.validation.seed, + ) + if self.rank == 0: + logging.info( + "PDD validation step=%d loss=%.12g pairs=%d starts=%d heads=%d " + "ordered_id_sha256=%s records=%d", + completed_step, + result.mean_loss, + result.pair_count, + result.start_count, + result.head_count, + result.ordered_id_sha256, + len(result.records), + ) + + def _log_step(self, diagnostics: Any, data_wait_seconds: float, step_seconds: float) -> None: + timing = torch.tensor( + [data_wait_seconds, step_seconds], + dtype=torch.float64, + device=self.config.device, + ) + dist.all_reduce(timing, op=dist.ReduceOp.MAX) + peak_memory = ( + torch.cuda.max_memory_allocated(self.config.device) + if self.config.device.type == "cuda" + else 0 + ) + memory = torch.tensor(peak_memory, dtype=torch.int64, device=self.config.device) + dist.all_reduce(memory, op=dist.ReduceOp.MAX) + global_samples = self.config.step_scheduler.local_batch_size * self.world_size + throughput = global_samples / max(float(timing[1].item()), 1e-12) + coverage = self.training.trainer.coverage + bin_loss = [ + None if count == 0 else float(loss_sum / count) + for loss_sum, count in zip( + coverage.bin_loss_sums.tolist(), + coverage.bin_counts.tolist(), + ) + ] + if self.rank == 0: + logging.info( + "PDD step=%d loss=%.6g grad_norm=%.6g nominal_update_ratio=%.6g " + "projection_update_ratio=%s lr=%.6g student_rms=%.6g " + "teacher_rms=%.6g student_teacher_rms_ratio=%.6g " + "reconstruction_rms=%.6g pairs=%d n_coverage=%s k_coverage=%s " + "bins=%s bin_loss=%s samples_per_second=%.3f " + "data_wait_seconds=%.4f peak_memory_bytes=%d", + diagnostics.completed_step, + diagnostics.loss, + diagnostics.grad_norm, + diagnostics.student_adamw_nominal_update_ratio, + diagnostics.pdd_projection_update_ratio, + diagnostics.learning_rate, + diagnostics.student_velocity_rms, + diagnostics.teacher_velocity_rms, + diagnostics.student_teacher_velocity_rms_ratio, + diagnostics.reconstructed_state_rms, + int((coverage.pair_counts > 0).sum()), + _coverage_axis(coverage.n_counts, coverage.n_loss_sums), + _coverage_axis(coverage.k_counts, coverage.k_loss_sums), + coverage.bin_counts.tolist(), + bin_loss, + throughput, + float(timing[0].item()), + int(memory.item()), + ) + if self.config.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(self.config.device) + + def run_train_validation_loop(self) -> None: + """Train through AutoModel StepScheduler without weakening PDD resume semantics.""" + data_wait_started = time.perf_counter() + try: + for _epoch in self.step_scheduler.epochs: + self.step_scheduler.dataloader = self._prepared_training_batches() + for batch_group in self.step_scheduler: + if len(batch_group) != 1: + raise RuntimeError("PDD v1 requires one microbatch per optimizer update.") + if self.step_scheduler.step != self.training.trainer.completed_steps: + raise RuntimeError( + "PDD trainer and AutoModel StepScheduler disagree before the update." + ) + (batch, sample_ids) = batch_group[0] + data_wait_seconds = time.perf_counter() - data_wait_started + if self.resume_pending: + if self.resume is None: + raise RuntimeError("PDD resume is pending without restored state.") + if self.rank == 0: + logging.info( + "PDD resume first batch verified: checkpoint=%s sample_ids=%s", + self.resume.checkpoint_path, + sample_ids, + ) + self.resume_pending = False + + step_started = time.perf_counter() + next_step = self.training.trainer.completed_steps + 1 + diagnostics = self.training.trainer.train_step( + batch, + measure_updates=(next_step % self.config.step_scheduler.log_every == 0), + ) + self.training.scheduler.step() + self.sampler.commit(sample_ids) + if self.sampler.remaining_batches == 0: + self.sampler.set_epoch(self.sampler.epoch + 1) + if self.training.trainer.completed_steps != self.step_scheduler.step + 1: + raise RuntimeError( + "PDD trainer and AutoModel StepScheduler disagree after the update." + ) + serialized_step = self.step_scheduler.state_dict()["step"] + if serialized_step != self.training.trainer.completed_steps: + raise RuntimeError("AutoModel StepScheduler serialized the wrong PDD step.") + step_seconds = time.perf_counter() - step_started + + completed_step = diagnostics.completed_step + is_final_step = self.step_scheduler.is_last_step + if completed_step % self.config.step_scheduler.log_every == 0: + self._log_step(diagnostics, data_wait_seconds, step_seconds) + if completed_step % self.config.validation.every_steps == 0 or is_final_step: + self._run_validation(completed_step) + if self.config.checkpoint.enabled and self.step_scheduler.is_ckpt_step: + self.checkpoint_manager.save() + data_wait_started = time.perf_counter() + finally: + self.setup_artifacts.checkpointer.close() diff --git a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py index cc8b22436fa..bdf976f7b6b 100644 --- a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py @@ -58,6 +58,14 @@ def __init__(self, completed_steps: int = 1) -> None: self.completed_steps = completed_steps +class _StepScheduler: + def __init__(self, trainer: _Trainer) -> None: + self.trainer = trainer + + def state_dict(self): + return {"step": self.trainer.completed_steps, "epoch": 0} + + class _Checkpointer: def __init__(self, rank: int, *, fail_sidecar: bool = False) -> None: self.config = SimpleNamespace(is_async=False) @@ -129,10 +137,11 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: model=object(), optimizer=SimpleNamespace(param_groups=[{"lr": 2.0e-5}]), scheduler=object(), + step_scheduler=_StepScheduler(trainer), trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 1, "topology": {"world_size": 2}}, + identity={"schema_version": 2, "topology": {"world_size": 2}}, ) initial.save() trainer.completed_steps = 2 @@ -143,10 +152,11 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: model=object(), optimizer=SimpleNamespace(param_groups=[{"lr": 2.0e-5}]), scheduler=object(), + step_scheduler=_StepScheduler(trainer), trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 1, "topology": {"world_size": 2}}, + identity={"schema_version": 2, "topology": {"world_size": 2}}, ) message = None try: diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py index 66fd1966910..6fb7a8c6c00 100644 --- a/tests/examples/diffusers/fastgen/pdd_export_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_export_distributed.py @@ -65,7 +65,10 @@ def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: "inference_blocks": [2, 2], "data_free": False, }, - "optim": {"learning_rate": 2.0e-5, "weight_decay": 0.01}, + "optim": { + "learning_rate": 2.0e-5, + "optimizer": {"weight_decay": 0.01}, + }, "fsdp": { "dp_size": 2, "tp_size": 1, diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index b4c770bd2d5..d347b067e1d 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -55,6 +55,7 @@ "automodel_dependency.json", "checkpoint.py", "configs", + "data.py", "export.py", "export_qwen_image.py", "finetune.py", diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 5a7988a2100..8b709391d28 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -179,7 +179,39 @@ def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: "inference_blocks": [2, 2], "data_free": False, }, - "optim": {"learning_rate": 2.0e-5, "weight_decay": 0.01}, + "seed": 42, + "optim": { + "learning_rate": 2.0e-5, + "optimizer": { + "_target_": "torch.optim.AdamW", + "weight_decay": 0.01, + }, + }, + "lr_scheduler": { + "lr_decay_style": "constant", + "lr_warmup_steps": 0, + "min_lr": 2.0e-5, + }, + "step_scheduler": { + "max_steps": 10, + "num_epochs": 2, + "log_every": 1, + "ckpt_every_steps": 5, + "local_batch_size": 1, + "global_batch_size": 1, + "save_checkpoint_every_epoch": False, + }, + "training_health": {"max_grad_norm": 1.0, "zero_grad_warmup_steps": 0}, + "validation": {"count": 3, "seed": 11, "split_seed": 7, "every_steps": 5}, + "data": { + "dataloader": { + "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", + "batch_size": 1, + "drop_last": True, + "shuffle": True, + "dynamic_batch_size": False, + } + }, "fsdp": { "dp_size": 1, "tp_size": 1, @@ -201,27 +233,118 @@ def test_example_recipe_explicitly_pins_grid_max_t() -> None: raw = yaml.safe_load((_FASTGEN_DIR / "pdd" / "configs" / "qwen_image.yaml").read_text()) assert type(raw["pdd"]["grid_max_t"]) is float assert raw["pdd"]["grid_max_t"] == 0.999 - assert raw["data"]["validation_count"] == 2000 - assert raw["data"]["split_seed"] == 2026 + assert raw["validation"]["count"] == 2000 + assert raw["validation"]["split_seed"] == 2026 def test_split_config_fields_are_strict(tmp_path) -> None: raw = _raw_config(tmp_path) - raw["data"] = {"validation_count": 3, "split_seed": 7} + raw["validation"] = {"count": 3, "seed": 11, "split_seed": 7, "every_steps": 5} config = resolve_pdd_recipe_config(raw) - assert config.validation_count == 3 - assert config.split_seed == 7 + assert config.validation.count == 3 + assert config.validation.split_seed == 7 for invalid_count in (0, -1, True, 1.5): - raw["data"]["validation_count"] = invalid_count - with pytest.raises((TypeError, ValueError), match="validation_count"): + raw["validation"]["count"] = invalid_count + with pytest.raises((TypeError, ValueError), match=r"validation\.count"): resolve_pdd_recipe_config(raw) - raw["data"] = {"validation_count": 3, "split_seed": -1} + raw["validation"] = {"count": 3, "seed": 11, "split_seed": -1, "every_steps": 5} with pytest.raises(ValueError, match="split_seed"): resolve_pdd_recipe_config(raw) +def test_config_node_and_canonical_dotted_values_are_consumed(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["step_scheduler"].update( + max_steps=50_000, + ckpt_every_steps=1_000, + global_batch_size=1, + ) + raw["optim"]["learning_rate"] = 3.0e-5 + raw["lr_scheduler"]["min_lr"] = 3.0e-5 + + class ConfigNodeLike: + def to_dict(self): + return copy.deepcopy(raw) + + config = resolve_pdd_recipe_config(ConfigNodeLike()) + assert config.step_scheduler.max_steps == 50_000 + assert config.step_scheduler.ckpt_every_steps == 1_000 + assert config.step_scheduler.global_batch_size == 1 + assert config.learning_rate == 3.0e-5 + + +def test_automodel_parser_dotted_overrides_reach_the_pdd_resolver(tmp_path, monkeypatch) -> None: + parser_module = pytest.importorskip("nemo_automodel.components.config._arg_parser") + config_path = tmp_path / "pdd.yaml" + config_path.write_text(yaml.safe_dump(_raw_config(tmp_path))) + monkeypatch.setattr( + sys, + "argv", + [ + "finetune.py", + "--config", + str(config_path), + "--step_scheduler.max_steps=50000", + "--step_scheduler.ckpt_every_steps=1000", + "--step_scheduler.global_batch_size=1", + "--optim.learning_rate=3e-5", + "--lr_scheduler.min_lr=3e-5", + ], + ) + + parsed = parser_module.parse_args_and_load_config(str(config_path)) + resolved = resolve_pdd_recipe_config(parsed) + assert resolved.step_scheduler.max_steps == 50_000 + assert resolved.step_scheduler.ckpt_every_steps == 1_000 + assert resolved.step_scheduler.global_batch_size == 1 + assert resolved.learning_rate == 3.0e-5 + + +@pytest.mark.parametrize( + ("legacy_key", "replacement"), + [ + ("max_steps", "step_scheduler.max_steps"), + ("global_batch_size", "step_scheduler.global_batch_size"), + ("checkpoint_every_steps", "step_scheduler.ckpt_every_steps"), + ("log_every_steps", "step_scheduler.log_every"), + ], +) +def test_legacy_training_lifecycle_keys_are_rejected(tmp_path, legacy_key, replacement) -> None: + raw = _raw_config(tmp_path) + raw["training"] = {legacy_key: 2} + with pytest.raises(ValueError, match=replacement.replace(".", r"\.")): + resolve_pdd_recipe_config(raw) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("lr_decay_style", "cosine", "lr_decay_style='constant'"), + ("lr_warmup_steps", 1, "lr_warmup_steps=0"), + ("min_lr", 1.0e-5, "min_lr must equal optim.learning_rate"), + ], +) +def test_nonconstant_lr_declarations_are_rejected(tmp_path, field, value, message) -> None: + raw = _raw_config(tmp_path) + raw["lr_scheduler"][field] = value + with pytest.raises(ValueError, match=message): + resolve_pdd_recipe_config(raw) + + +@pytest.mark.parametrize("field", ["weight_decay", "betas", "eps"]) +def test_legacy_optimizer_fields_are_rejected(tmp_path, field) -> None: + raw = _raw_config(tmp_path) + raw["optim"][field] = { + "weight_decay": 0.01, + "betas": [0.9, 0.999], + "eps": 1.0e-8, + }[field] + with pytest.raises(ValueError, match=rf"optim\.optimizer\.{field}"): + resolve_pdd_recipe_config(raw) + + def test_pdd_rejects_external_split_manifest(tmp_path) -> None: raw = _raw_config(tmp_path) raw["data"] = {"dataloader": {"metadata_index": "metadata_train.json"}} @@ -240,7 +363,7 @@ def test_pdd_finetune_namespace_module_help() -> None: capture_output=True, text=True, ) - assert "Train Qwen-Image" in result.stdout + assert "Qwen-Image PDD training" in result.stdout @pytest.mark.parametrize( @@ -280,19 +403,29 @@ def test_remote_model_requires_full_revision_and_non_dp_parallelism_is_rejected( @pytest.mark.parametrize( ("section", "name", "value", "message"), [ - ("training", "grad_accumulation_steps", 2, "grad_accumulation_steps=1"), - ("training", "max_grad_norm", 0.0, "max_grad_norm must be > 0"), - ("training", "validation_every_steps", 0, "validation_every_steps"), + ( + "step_scheduler", + "save_checkpoint_every_epoch", + True, + "save_checkpoint_every_epoch=false", + ), + ("training_health", "max_grad_norm", 0.0, "max_grad_norm must be > 0"), + ("validation", "every_steps", 0, "validation.every_steps"), ("guidance", "rescale", 1.1, "guidance.rescale must be <= 1"), - ("optim", "betas", [0.9, 1.0], "optim.betas values"), - ("optim", "eps", 0.0, "optim.eps must be > 0"), + ("optimizer", "betas", [0.9, 1.0], "optim.optimizer.betas values"), + ("optimizer", "eps", 0.0, "optim.optimizer.eps must be > 0"), ], ) def test_training_config_gates_fail_during_resolution( tmp_path, section, name, value, message ) -> None: raw = _raw_config(tmp_path) - raw.setdefault(section, {})[name] = value + target = ( + raw["optim"].setdefault("optimizer", {}) + if section == "optimizer" + else raw.setdefault(section, {}) + ) + target[name] = value with pytest.raises(ValueError, match=message): resolve_pdd_recipe_config(raw) diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index 207a50b0ab1..a945dd2e7fc 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -25,9 +25,11 @@ import pathlib import shutil import sys +from types import SimpleNamespace import pytest import torch +import torch.distributed as dist _REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] _FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" @@ -44,7 +46,7 @@ build_pdd_checkpoint_identity, resolve_pdd_training_checkpoint, ) -from pdd.recipe import initialize_pdd_distributed +from pdd.recipe import PDDDiffusionRecipe, initialize_pdd_distributed from pdd.training import prepare_qwen_pdd_batch from pdd.verify_readonly_automodel import snapshot_installed_distribution from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 @@ -119,14 +121,47 @@ def _identity(lifecycle, scheduler, sample_ids): ) +class _StepSchedulerStub: + def __init__(self, trainer, sampler) -> None: + self.trainer = trainer + self.sampler = sampler + self.loaded_state = None + + def state_dict(self): + return {"step": self.trainer.completed_steps, "epoch": self.sampler.epoch} + + def load_state_dict(self, state): + self.loaded_state = dict(state) + + +class _RecordingCheckpointManager(PDDCheckpointManager): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.save_calls = [] + + def save(self): + self.save_calls.append( + { + "live_step": self.step_scheduler.step, + "serialized_step": self.step_scheduler.state_dict()["step"], + "trainer_step": self.trainer.completed_steps, + "live_epoch": self.step_scheduler.epoch, + "sampler_epoch": self.sampler.epoch, + } + ) + return super().save() + + def _manager(root, lifecycle, sampler, rng): checkpointer = _checkpointer(lifecycle, root) + step_scheduler = _StepSchedulerStub(lifecycle.trainer, sampler) manager = PDDCheckpointManager( root=root, checkpointer=checkpointer, model=lifecycle.student, optimizer=lifecycle.optimizer, scheduler=lifecycle.scheduler, + step_scheduler=step_scheduler, trainer=lifecycle.trainer, sampler=sampler, rng=rng, @@ -183,6 +218,244 @@ def test_replayable_sampler_commits_consumed_batches_not_prefetch() -> None: _released_sampler(sample_ids).load_state_dict(bad_ids) +def test_automodel_step_scheduler_serializes_the_completed_yielded_step() -> None: + scheduler_module = pytest.importorskip("nemo_automodel.components.training.step_scheduler") + scheduler = scheduler_module.StepScheduler( + global_batch_size=1, + local_batch_size=1, + dp_size=1, + ckpt_every_steps=2, + save_checkpoint_every_epoch=False, + dataloader=[{"sample": 0}, {"sample": 1}], + val_every_steps=None, + start_step=0, + start_epoch=0, + num_epochs=1, + max_steps=2, + ) + + iterator = iter(scheduler) + assert next(iterator) == [{"sample": 0}] + assert scheduler.step == 0 + assert scheduler.state_dict() == {"step": 1, "epoch": 0} + assert next(iterator) == [{"sample": 1}] + assert scheduler.step == 1 + assert scheduler.is_last_step + assert scheduler.state_dict() == {"step": 2, "epoch": 0} + with pytest.raises(StopIteration): + next(iterator) + + +def _build_recipe_loop( + root, + sample_ids, + *, + max_steps, + num_epochs, + ckpt_every_steps, + restore_from=None, +): + if not dist.is_initialized(): + initialize_pdd_distributed(backend="gloo", timeout_minutes=1) + scheduler_module = pytest.importorskip("nemo_automodel.components.training.step_scheduler") + rng_module = pytest.importorskip("nemo_automodel.components.training.rng") + + lifecycle = build_toy_lifecycle() + sampler = _released_sampler(sample_ids) + rng = rng_module.StatefulRNG(1234, ranked=True) + step_scheduler = scheduler_module.StepScheduler( + global_batch_size=1, + local_batch_size=1, + dp_size=1, + ckpt_every_steps=ckpt_every_steps, + save_checkpoint_every_epoch=False, + dataloader=[None] * len(sampler), + val_every_steps=None, + start_step=0, + start_epoch=0, + num_epochs=num_epochs, + max_steps=max_steps, + ) + checkpointer = _checkpointer(lifecycle, root) + manager = _RecordingCheckpointManager( + root=root, + checkpointer=checkpointer, + model=lifecycle.student, + optimizer=lifecycle.optimizer, + scheduler=lifecycle.scheduler, + step_scheduler=step_scheduler, + trainer=lifecycle.trainer, + sampler=sampler, + rng=rng, + identity=_identity(lifecycle, lifecycle.scheduler, sample_ids), + ) + resume = manager.load(restore_from) + events = SimpleNamespace(first_ids=[], diagnostics=[], validation_steps=[]) + + recipe = object.__new__(PDDDiffusionRecipe) + recipe.config = SimpleNamespace( + step_scheduler=SimpleNamespace(local_batch_size=1, log_every=1), + validation=SimpleNamespace(every_steps=10_000), + checkpoint=SimpleNamespace(enabled=True), + device=torch.device("cpu"), + ) + recipe.training = SimpleNamespace( + pipeline=lifecycle.pipeline, + trainer=lifecycle.trainer, + scheduler=lifecycle.scheduler, + rng=rng, + ) + recipe.setup_artifacts = SimpleNamespace(checkpointer=checkpointer) + recipe.step_scheduler = step_scheduler + recipe.checkpoint_manager = manager + recipe.sampler = sampler + recipe.resume = resume + recipe.resume_pending = resume is not None + recipe.rank = 0 + recipe.world_size = 1 + + def prepared_batches(): + # Bind this iterator to the current sampler plan. The production loader iterator also + # exhausts after that plan even though the recipe commits the sampler into its next epoch. + for _ in range(sampler.remaining_batches): + expected_ids = sampler.expected_next_sample_ids() + if recipe.resume_pending: + assert recipe.resume is not None + recipe.resume.verify_first_batch(expected_ids) + events.first_ids.append(expected_ids) + offset = sum(ord(character) for character in expected_ids[0]) / 10_000 + yield make_batch(expected_ids, offset=offset), expected_ids + + recipe.__dict__["_prepared_training_batches"] = prepared_batches + recipe.__dict__["_run_validation"] = events.validation_steps.append + recipe.__dict__["_log_step"] = lambda diagnostics, _data_wait, _step_time: ( + events.diagnostics.append(diagnostics) + ) + return SimpleNamespace( + recipe=recipe, + lifecycle=lifecycle, + sampler=sampler, + step_scheduler=step_scheduler, + manager=manager, + resume=resume, + events=events, + ) + + +def test_recipe_loop_saves_periodic_and_max_step_checkpoints_once(tmp_path) -> None: + run = _build_recipe_loop( + tmp_path / "periodic", + tuple(f"sample-{index}" for index in range(8)), + max_steps=3, + num_epochs=4, + ckpt_every_steps=2, + ) + run.recipe.run_train_validation_loop() + + assert [call["trainer_step"] for call in run.manager.save_calls] == [2, 3] + assert [call["live_step"] for call in run.manager.save_calls] == [1, 2] + assert [call["serialized_step"] for call in run.manager.save_calls] == [2, 3] + assert sorted(path.name for path in (tmp_path / "periodic").glob("step_*")) == [ + "step_00000002", + "step_00000003", + ] + assert run.events.validation_steps == [3] + + +def test_recipe_loop_epoch_final_save_normalizes_and_restores_epoch(tmp_path) -> None: + root = tmp_path / "epoch" + sample_ids = ("sample-0", "sample-1") + source = _build_recipe_loop( + root, + sample_ids, + max_steps=100, + num_epochs=1, + ckpt_every_steps=100, + ) + source.recipe.run_train_validation_loop() + + assert len(source.manager.save_calls) == 1 + assert source.manager.save_calls[0] == { + "live_step": 1, + "serialized_step": 2, + "trainer_step": 2, + "live_epoch": 0, + "sampler_epoch": 1, + } + manifest = json.loads((root / "step_00000002" / "manifest.json").read_text()) + assert manifest["step_scheduler"] == {"step": 2, "epoch": 1} + + resumed = _build_recipe_loop( + root, + sample_ids, + max_steps=3, + num_epochs=2, + ckpt_every_steps=100, + restore_from="step_00000002", + ) + assert resumed.step_scheduler.step == 2 + assert resumed.step_scheduler.epoch == resumed.sampler.epoch == 1 + assert resumed.resume is not None + expected_ids = resumed.resume.expected_next_sample_ids + resumed.recipe.run_train_validation_loop() + assert resumed.events.first_ids[0] == expected_ids + assert (root / "step_00000003" / "COMPLETE").is_file() + + +def test_recipe_loop_resume_matches_uninterrupted_next_update(tmp_path) -> None: + sample_ids = tuple(f"sample-{index}" for index in range(4)) + control = _build_recipe_loop( + tmp_path / "control", + sample_ids, + max_steps=2, + num_epochs=2, + ckpt_every_steps=100, + ) + control.recipe.run_train_validation_loop() + + staged_root = tmp_path / "staged" + first = _build_recipe_loop( + staged_root, + sample_ids, + max_steps=1, + num_epochs=2, + ckpt_every_steps=100, + ) + first.recipe.run_train_validation_loop() + resumed = _build_recipe_loop( + staged_root, + sample_ids, + max_steps=2, + num_epochs=2, + ckpt_every_steps=100, + restore_from="step_00000001", + ) + assert resumed.resume is not None + expected_next_ids = resumed.resume.expected_next_sample_ids + resumed.recipe.run_train_validation_loop() + + assert resumed.events.first_ids == [expected_next_ids] + assert resumed.events.first_ids[0] == control.events.first_ids[1] + assert resumed.events.diagnostics == [control.events.diagnostics[1]] + for name, tensor in resumed.lifecycle.student.state_dict().items(): + torch.testing.assert_close( + tensor, + control.lifecycle.student.state_dict()[name], + rtol=0, + atol=0, + ) + actual_optimizer = _optimizer_state_by_name(resumed.lifecycle) + expected_optimizer = _optimizer_state_by_name(control.lifecycle) + assert actual_optimizer.keys() == expected_optimizer.keys() + for name in actual_optimizer: + for key, actual in actual_optimizer[name].items(): + expected = expected_optimizer[name][key] + if isinstance(actual, torch.Tensor): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + assert actual == expected + + def test_qwen_batch_preparation_preserves_ids_masks_and_negative_condition() -> None: batch = { "image_latents": torch.ones(2, 3, 4, 4), @@ -496,10 +769,12 @@ def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) step_manifest_path = step_mismatch / "manifest.json" step_manifest = json.loads(step_manifest_path.read_text()) step_manifest["completed_steps"] = 4 + step_manifest["step_scheduler"]["step"] = 4 step_manifest_path.write_text(json.dumps(step_manifest, indent=2, sort_keys=True) + "\n") trainer_state_path = step_mismatch / "trainer_state.json" trainer_state = json.loads(trainer_state_path.read_text()) trainer_state["completed_steps"] = 4 + trainer_state["step_scheduler"]["step"] = 4 trainer_state_path.write_text(json.dumps(trainer_state, indent=2, sort_keys=True) + "\n") _refresh_complete_marker(step_mismatch) with pytest.raises(RuntimeError, match="trainer step"): From 94d6639a349cb8ac4e0b8ea81b4804bb96e3ab7b Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 19:35:41 -0700 Subject: [PATCH 23/45] Support hashless DMD2 caches for PDD resume Signed-off-by: Meng Xin --- .../fastgen/fastgen_data/collate_fns.py | 12 ++- .../fastgen_data/text_to_image_dataset.py | 17 +++- examples/diffusers/fastgen/pdd/README.md | 11 ++- .../fastgen/pdd/configs/qwen_image.yaml | 2 + examples/diffusers/fastgen/pdd/data.py | 14 ++-- examples/diffusers/fastgen/pdd/recipe.py | 9 +- .../diffusers/fastgen/test_dataset_paths.py | 82 ++++++++++++++++++- .../fastgen/test_pdd_recipe_setup.py | 8 ++ 8 files changed, 139 insertions(+), 16 deletions(-) diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index ebf0365cea1..57fb4b18e50 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -198,6 +198,7 @@ def build_text_to_image_multiresolution_dataloader( validation_count: int | None = None, split_seed: int = 2026, exact_resume: bool = False, + verify_payload_hashes: bool | None = None, sampler_seed: int = 42, loader_seed: int | None = None, ) -> tuple[StatefulDataLoader, SequentialBucketSampler | ReplayableBatchSampler]: @@ -225,6 +226,10 @@ def build_text_to_image_multiresolution_dataloader( split_seed: Local seed used to construct deterministic split membership. exact_resume: Wrap the deterministic sampler with a committed cursor that is independent of worker prefetch. Required by the PDD lifecycle. + verify_payload_hashes: Require and authenticate each cached tensor against its + ``cache_sha256`` metadata before loading. ``None`` preserves the historical + builder behavior by following ``exact_resume``; PDD sets this explicitly so + replayable cursor state does not require payload hashes. sampler_seed: Seed for the released deterministic bucket sampler. loader_seed: Optional dedicated seed for DataLoader worker/base-seed generation. PDD supplies this so recreating an iterator cannot consume its restored training RNG. @@ -232,6 +237,11 @@ def build_text_to_image_multiresolution_dataloader( Returns: ``(StatefulDataLoader, SequentialBucketSampler)``. """ + if verify_payload_hashes is None: + verify_payload_hashes = exact_resume + elif type(verify_payload_hashes) is not bool: + raise TypeError("verify_payload_hashes must be bool or None.") + dataset = TextToImageDataset( cache_dir=cache_dir, train_text_encoder=train_text_encoder, @@ -239,7 +249,7 @@ def build_text_to_image_multiresolution_dataloader( split=split, validation_count=validation_count, split_seed=split_seed, - verify_payload_hashes=exact_resume, + verify_payload_hashes=verify_payload_hashes, ) effective_root = dataset.cache_root 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 b97b05daf47..2fb1d410fdd 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -49,7 +49,9 @@ def __init__( split: Optional deterministic ``"train"`` or ``"validation"`` selection. validation_count: Number of validation samples when ``split`` is set. split_seed: Local seed used to construct deterministic split membership. - verify_payload_hashes: Require and authenticate cached tensor content on every load. + verify_payload_hashes: Require and authenticate cached tensor content on every + load. When false, payloads load directly and the cache must remain immutable + for deterministic resume. """ if selected_indices is not None and split is not None: raise ValueError("selected_indices and split are mutually exclusive") @@ -57,6 +59,8 @@ def __init__( raise ValueError("split must be null, 'train', or 'validation'") if split is not None and validation_count is None: raise ValueError("validation_count is required when split is set") + if type(verify_payload_hashes) is not bool: + raise TypeError("verify_payload_hashes must be bool.") self.train_text_encoder = train_text_encoder self.cache_root = resolve_cache_root(cache_dir) self._selected_indices = selected_indices @@ -69,6 +73,11 @@ def __init__( self.dataset_snapshot_sha256: str | None = None super().__init__(str(self.cache_root), quantization=64) + @property + def verify_payload_hashes(self) -> bool: + """Whether sample payload bytes are authenticated before deserialization.""" + return self._verify_payload_hashes + 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") @@ -118,7 +127,7 @@ def _load_metadata(self) -> list[dict]: if self._verify_payload_hashes and cache_sha256 is None: raise ValueError( f"metadata shard {shard_path} item {shard_item_index} has no " - "cache_sha256 required for exact resume" + "cache_sha256 required when verify_payload_hashes=true" ) complete_metadata.append(dict(item)) @@ -126,7 +135,9 @@ def _load_metadata(self) -> list[dict]: raise ValueError(f"No samples found in {metadata_file}") self.total_num_samples = len(complete_metadata) self.metadata_sha256 = digest.hexdigest() - self.payload_hashes_complete = all("cache_sha256" in item for item in complete_metadata) + self.payload_hashes_complete = all( + item.get("cache_sha256") is not None for item in complete_metadata + ) if self._split is None: self.sample_ids = self._validate_selected_indices(self.total_num_samples) else: diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index 34d24969876..8d63efe230e 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -30,9 +30,14 @@ torchrun --standalone --nproc-per-node=8 \ The cache must contain `metadata.json`, its declared shards, cached tensors, and `negative_prompt_embedding.pt`. The environment variable overrides the configured cache root. All metadata, tensor, and negative-embedding paths must still resolve inside that effective root. -For exact resume, every shard item must also contain the `cache_sha256` written by the shared -preprocessor. PDD verifies each tensor's bytes before loading it and binds those expected hashes -plus the negative-prompt embedding hash into the checkpoint dataset identity. +PDD's committed sampler makes the next batch, sample IDs, RNG, optimizer, and scheduler state +exactly replayable without requiring per-payload hashes. Hashless caches must therefore remain +immutable for the duration of a run: metadata and the negative embedding are bound into the +checkpoint identity, but an in-place tensor payload change cannot be detected. + +Set `data.dataloader.verify_payload_hashes: true` to require the `cache_sha256` field written by +the shared preprocessor and authenticate every tensor's bytes before loading it. This stricter mode +has additional read and SHA-256 cost and fails immediately when a hash is missing or mismatched. Training deterministically derives disjoint train and validation membership from metadata ordinals; it does not rewrite the cache or require separate split manifests. The default recipe uses 2,000 diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 03de092a073..6e44d662f4b 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -89,6 +89,8 @@ data: drop_last: true shuffle: true dynamic_batch_size: false + # Cursor-exact resume does not require hashes. Set true to authenticate every payload. + verify_payload_hashes: false negative_prompt_embedding_path: negative_prompt_embedding.pt checkpoint: diff --git a/examples/diffusers/fastgen/pdd/data.py b/examples/diffusers/fastgen/pdd/data.py index f2468cb74d3..fce8aaa9e30 100644 --- a/examples/diffusers/fastgen/pdd/data.py +++ b/examples/diffusers/fastgen/pdd/data.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Authenticated data and collective batch handling for the Qwen-Image PDD recipe.""" +"""Replayable data and collective batch handling for the Qwen-Image PDD recipe.""" from __future__ import annotations @@ -43,6 +43,7 @@ def _dataloader_options(raw: Mapping[str, Any]) -> dict[str, Any]: raise ValueError(f"PDD data.dataloader._target_ must be {expected_target!r}.") if "base_resolution" in options: options["base_resolution"] = tuple(options["base_resolution"]) + options.setdefault("verify_payload_hashes", False) return options @@ -130,11 +131,10 @@ def _validate_dataset_contract( raise RuntimeError("training and validation datasets disagree on total sample count.") if train_dataset.metadata_sha256 != validation_dataset.metadata_sha256: raise RuntimeError("training and validation datasets disagree on metadata content.") - if ( - not train_dataset.payload_hashes_complete - or not validation_dataset.payload_hashes_complete - ): - raise RuntimeError("PDD exact resume requires a cache_sha256 for every tensor payload.") + if train_dataset.verify_payload_hashes != validation_dataset.verify_payload_hashes: + raise RuntimeError("training and validation datasets use different payload policies.") + if train_dataset.payload_hashes_complete != validation_dataset.payload_hashes_complete: + raise RuntimeError("training and validation datasets disagree on payload hash coverage.") if train_dataset.dataset_snapshot_sha256 != validation_dataset.dataset_snapshot_sha256: raise RuntimeError("training and validation datasets disagree on dataset content.") if not isinstance(train_dataset.dataset_snapshot_sha256, str): @@ -148,6 +148,8 @@ def _validate_dataset_contract( "train_samples": len(train_ids), "validation_samples": len(validation_ids), "split_seed": config.validation.split_seed, + "verify_payload_hashes": train_dataset.verify_payload_hashes, + "payload_hashes_complete": train_dataset.payload_hashes_complete, } local_status: dict[str, Any] = { "ok": True, diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index e8c39710f2d..c9d6e816865 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -311,6 +311,10 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: ): raise ValueError("PDD requires cached text embeddings; train_text_encoder must be false.") _require_bool(dataloader.get("shuffle", True), name="data.dataloader.shuffle") + _require_bool( + dataloader.get("verify_payload_hashes", False), + name="data.dataloader.verify_payload_hashes", + ) if "metadata_index" in dataloader: raise ValueError( "PDD uses deterministic ordinal splits from metadata.json; " @@ -1181,11 +1185,14 @@ def _log_setup(self) -> None: ) logging.info( "PDD dataset verified: snapshot_sha256=%s metadata_sha256=%s " - "train=%d validation=%d root=%s", + "train=%d validation=%d payload_hash_verification=%s " + "payload_hashes_complete=%s root=%s", self.snapshot_report["dataset_snapshot_sha256"], self.snapshot_report["metadata_sha256"], self.snapshot_report["train_samples"], self.snapshot_report["validation_samples"], + self.snapshot_report["verify_payload_hashes"], + self.snapshot_report["payload_hashes_complete"], self.snapshot_report["cache_root"], ) logging.info( diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index 33c08c0ed51..f5395be3ef5 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -36,6 +36,7 @@ sys.path.insert(0, str(_FASTGEN_DIR)) from fastgen_data import ( + ReplayableBatchSampler, TextToImageDataset, build_text_to_image_multiresolution_dataloader, resolve_cache_root, @@ -125,7 +126,9 @@ def test_dataset_accepts_absolute_payload_beneath_root(make_fastgen_cache, tmp_p assert dataset[0]["sample_id"] == 0 -def test_exact_resume_requires_and_verifies_payload_hashes(make_fastgen_cache, tmp_path): +def test_payload_hash_authentication_rejects_missing_or_modified_payloads( + make_fastgen_cache, tmp_path +): cache = make_fastgen_cache(tmp_path / "cache") dataset = TextToImageDataset(cache, verify_payload_hashes=True) assert dataset[0]["sample_id"] == 0 @@ -139,10 +142,85 @@ def test_exact_resume_requires_and_verifies_payload_hashes(make_fastgen_cache, t shard[0].pop("cache_sha256") shard_path.write_text(json.dumps(shard)) - with pytest.raises(ValueError, match="cache_sha256 required for exact resume"): + with pytest.raises(ValueError, match="verify_payload_hashes=true"): TextToImageDataset(cache, verify_payload_hashes=True) +def test_builder_preserves_legacy_exact_resume_hash_requirement(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + shard_path = cache / "metadata_shard_0.json" + shard = json.loads(shard_path.read_text()) + shard[0].pop("cache_sha256") + shard_path.write_text(json.dumps(shard)) + + with pytest.raises(ValueError, match="verify_payload_hashes=true"): + build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), num_workers=0, exact_resume=True + ) + + +def test_null_payload_hash_is_incomplete_and_strict_mode_rejects_it( + make_fastgen_cache, tmp_path +): + cache = make_fastgen_cache(tmp_path / "cache") + shard_path = cache / "metadata_shard_0.json" + shard = json.loads(shard_path.read_text()) + shard[0]["cache_sha256"] = None + shard_path.write_text(json.dumps(shard)) + + assert TextToImageDataset(cache).payload_hashes_complete is False + with pytest.raises(ValueError, match="verify_payload_hashes=true"): + TextToImageDataset(cache, verify_payload_hashes=True) + + +def test_hashless_cache_keeps_replayable_exact_cursor(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + shard_path = cache / "metadata_shard_0.json" + shard = json.loads(shard_path.read_text()) + for item in shard: + item.pop("cache_sha256") + shard_path.write_text(json.dumps(shard)) + + options = { + "cache_dir": str(cache), + "batch_size": 1, + "num_workers": 0, + "shuffle": True, + "exact_resume": True, + "verify_payload_hashes": False, + } + loader, sampler = build_text_to_image_multiresolution_dataloader(**options) + assert isinstance(sampler, ReplayableBatchSampler) + assert loader.dataset.verify_payload_hashes is False + assert loader.dataset.payload_hashes_complete is False + + first_batch = next(iter(loader)) + consumed = first_batch["metadata"]["logical_sample_ids"] + sampler.commit(consumed) + state = sampler.state_dict() + expected_next = sampler.expected_next_sample_ids() + + _, restored_sampler = build_text_to_image_multiresolution_dataloader(**options) + restored_sampler.load_state_dict(state) + assert restored_sampler.expected_next_sample_ids() == expected_next + + +def test_hashless_cache_uses_direct_load_by_default(make_fastgen_cache, tmp_path): + cache = make_fastgen_cache(tmp_path / "cache") + shard_path = cache / "metadata_shard_0.json" + shard = json.loads(shard_path.read_text()) + for item in shard: + item.pop("cache_sha256") + shard_path.write_text(json.dumps(shard)) + + loader, sampler = build_text_to_image_multiresolution_dataloader( + cache_dir=str(cache), batch_size=1, num_workers=0 + ) + assert not isinstance(sampler, ReplayableBatchSampler) + assert loader.dataset.verify_payload_hashes is False + assert next(iter(loader))["metadata"]["logical_sample_ids"] + + def test_dataset_snapshot_binds_negative_prompt_embedding(make_fastgen_cache, tmp_path): cache = make_fastgen_cache(tmp_path / "cache") loader, _ = build_text_to_image_multiresolution_dataloader( diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 8b709391d28..2a25b35e4d8 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -458,6 +458,14 @@ def test_training_dataloader_modes_are_gated_during_resolution( resolve_pdd_recipe_config(raw) +def test_payload_hash_verification_mode_must_be_bool(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["data"]["dataloader"]["verify_payload_hashes"] = "false" + + with pytest.raises(TypeError, match="data.dataloader.verify_payload_hashes must be bool"): + resolve_pdd_recipe_config(raw) + + def test_frozen_automodel_distribution_snapshot_is_stable() -> None: _require_exact_automodel() From 8582b790b68d2fce442b1ddff3a6490d1843dd5e Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 20:06:49 -0700 Subject: [PATCH 24/45] Match PDD math to FastGen reference Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/README.md | 10 +++- .../fastgen/pdd/configs/qwen_image.yaml | 6 +- examples/diffusers/fastgen/pdd/recipe.py | 44 ++++++++++++--- examples/diffusers/fastgen/pdd/training.py | 4 ++ modelopt/torch/fastgen/methods/pdd.py | 12 +++- .../torch/fastgen/plugins/qwen_image_pdd.py | 33 ++++++++--- .../fastgen/test_pdd_recipe_setup.py | 49 ++++++++++++++++- .../fastgen/test_pdd_training_lifecycle.py | 9 ++- tests/unit/torch/fastgen/test_pdd_pipeline.py | 15 ++++- .../fastgen/test_qwen_image_pdd_plugin.py | 55 +++++++++++++++++-- 10 files changed, 207 insertions(+), 30 deletions(-) diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index 8d63efe230e..734d3677de3 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -2,8 +2,8 @@ Parallel Decoding Distillation (PDD) trains one Qwen-Image student call to predict several consecutive rectified-flow updates. The student keeps the original transformer backbone and widens -only its output projection to 128 velocity heads. During training it samples different aligned -block lengths up to 64, so the same checkpoint can use different supported block schedules at +only its output projection to 128 velocity heads. During training it samples aligned block starts +and target spans from 1 through 64 intervals, so the same checkpoint can use different supported block schedules at inference. The provided schedules use the 128-interval grid as follows: @@ -41,7 +41,11 @@ has additional read and SHA-256 cost and fails immediately when a hash is missin Training deterministically derives disjoint train and validation membership from metadata ordinals; it does not rewrite the cache or require separate split manifests. The default recipe uses 2,000 -validation samples, learning rate `2e-5`, 128 heads, and sampled block lengths from 4 through 64. +validation samples, learning rate `2e-5`, per-rank batch size 4, 128 heads, start indices aligned by +4, and target spans from 1 through 64 intervals. The learning rate is the cached-Qwen project +treatment; the MR210 reference arm uses `5e-5` with 1,000 warmup steps. On 16 four-GPU nodes the +default per-rank batch gives global batch size 256; other GPU topologies must set per-rank batch to +`256 / world_size` because this recipe does not use gradient accumulation. Checkpoints include the student, optimizer, scheduler, RNG, trainer, and exact replayable sampler state needed to resume the next committed batch. diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 6e44d662f4b..756464081f7 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -57,7 +57,9 @@ step_scheduler: num_epochs: 200 log_every: 10 ckpt_every_steps: 1000 - local_batch_size: 1 + # Matches the original Qwen PDD per-rank batch. On 16 four-GPU nodes this + # gives the project target global batch size of 256 without accumulation. + local_batch_size: 4 save_checkpoint_every_epoch: false # Freeze to 256 only after the production-topology/data gate is approved. global_batch_size: @@ -85,7 +87,7 @@ data: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: data/qwen_image_cache base_resolution: [1024, 1024] - batch_size: 1 + batch_size: 4 drop_last: true shuffle: true dynamic_batch_size: false diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index c9d6e816865..01f5143095e 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -77,7 +77,7 @@ class PDDStepSchedulerConfig: num_epochs: int = 200 log_every: int = 10 ckpt_every_steps: int = 1_000 - local_batch_size: int = 1 + local_batch_size: int = 4 global_batch_size: int | None = None save_checkpoint_every_epoch: bool = False @@ -696,16 +696,18 @@ def _stage_and_shard_training_models( manager: Any, *, device: torch.device, - dtype: torch.dtype, fuse_qkv_projections: bool, ) -> tuple[nn.Module, nn.Module]: - """Move and shard one dense model at a time to bound setup-time GPU memory.""" + """Stage FP32 masters and shard one dense model at a time.""" if fuse_qkv_projections and ( not hasattr(student, "fuse_qkv_projections") or not hasattr(teacher, "fuse_qkv_projections") ): raise AttributeError("QKV fusion requires both Qwen transformers to expose the object API.") - student.to(device=device, dtype=dtype) + # Match the FastGen Qwen PDD recipe: FP32 parameter/optimizer storage, + # with the FSDP policy below casting gathered parameters to the configured + # model dtype for forward/backward compute. + student.to(device=device, dtype=torch.float32) if fuse_qkv_projections: student.fuse_qkv_projections() if not any(getattr(module, "fused_projections", False) for module in student.modules()): @@ -719,7 +721,7 @@ def _stage_and_shard_training_models( # The student is already sharded before the dense teacher reaches the GPU, so multi-rank # setup never holds both complete Qwen transformers on one device. - teacher.to(device=device, dtype=dtype) + teacher.to(device=device, dtype=torch.float32) if fuse_qkv_projections: teacher.fuse_qkv_projections() teacher = manager.parallelize(teacher) @@ -757,6 +759,23 @@ def _materialize_zero_step_adamw_state(optimizer: torch.optim.AdamW) -> None: optimizer.zero_grad(set_to_none=True) +def _require_fp32_optimizer_storage(optimizer: torch.optim.AdamW) -> None: + """Require the FP32 master-parameter and Adam-state contract used by MR210.""" + for group in optimizer.param_groups: + for parameter in group["params"]: + if parameter.dtype != torch.float32: + raise RuntimeError( + f"PDD trainable master parameters must be FP32, got {parameter.dtype}." + ) + state = optimizer.state.get(parameter) + if state is None: + raise RuntimeError("PDD AdamW state is missing after eager materialization.") + for name in ("exp_avg", "exp_avg_sq"): + value = state.get(name) + if not isinstance(value, torch.Tensor) or value.dtype != torch.float32: + raise RuntimeError(f"PDD AdamW {name} state must be FP32.") + + def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: """Compose released AutoModel APIs without editing or patching external packages.""" if not isinstance(config, PDDRecipeConfig): @@ -803,7 +822,17 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: raise ValueError( f"Pure-DP PDD requires fsdp.dp_size ({dp_size}) to equal world size ({world_size})." ) - strategy = FSDP2Config(activation_checkpointing=config.parallel.activation_checkpointing) + from torch.distributed.fsdp import MixedPrecisionPolicy + + strategy = FSDP2Config( + activation_checkpointing=config.parallel.activation_checkpointing, + mp_policy=MixedPrecisionPolicy( + param_dtype=config.dtype, + reduce_dtype=torch.float32, + output_dtype=None, + cast_forward_inputs=False, + ), + ) distributed_setup = DistributedSetup.build( strategy=strategy, parallelism_sizes=ParallelismSizes(dp_size=dp_size), @@ -823,7 +852,6 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: identity, manager, device=config.device, - dtype=config.dtype, fuse_qkv_projections=config.fuse_qkv_projections, ) pipe.transformer = student @@ -850,6 +878,7 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: maximize=False, ) _materialize_zero_step_adamw_state(optimizer) + _require_fp32_optimizer_storage(optimizer) optimizer_parameters = [ parameter for group in optimizer.param_groups for parameter in group["params"] ] @@ -1017,6 +1046,7 @@ def build_pdd_training_artifacts( config.pdd, guidance_rescale=config.guidance.rescale, guidance_eps=config.guidance.eps, + compute_dtype=config.dtype, ) pipeline = PDDPipeline(setup.student, setup.teacher, config.pdd, adapter) scheduler = torch.optim.lr_scheduler.LambdaLR(setup.optimizer, lr_lambda=lambda _: 1.0) diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py index c8780b9f63e..4b9e1480fc5 100644 --- a/examples/diffusers/fastgen/pdd/training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -633,6 +633,10 @@ def train_step( raise RuntimeError( "PDD optimizer produced a zero nominal update from a nonzero gradient." ) + if projection_ratio == 0.0 and grad_norm > 0.0: + raise RuntimeError( + "PDD optimizer produced a zero actual projection update from a nonzero gradient." + ) self.completed_steps += 1 student_velocity_rms = _global_sample_mean(metrics["student_velocity_rms"]) teacher_velocity_rms = _global_sample_mean(metrics["teacher_velocity_rms"]) diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index 6d3324099bf..429e240ed06 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -970,6 +970,16 @@ def sample( kwargs = self._model_kwargs(model_kwargs) resolved_blocks = self._validate_blocks(blocks) grid = self.time_grid(noise.device) + # MR210 derives fused projection coefficients from its float64 schedule, + # then casts the coefficients to FP32. Keep state/time integration on the + # canonical FP32 decoding grid while matching that coefficient path. + 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: @@ -981,7 +991,7 @@ def sample( time, start=start, end=end, - grid=grid, + grid=fusion_grid, condition=condition, **kwargs, ) diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index 39ef65cd352..1a54787fcbb 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -113,6 +113,7 @@ def __init__( *, guidance_rescale: float = 1.0, guidance_eps: float = 1e-5, + compute_dtype: torch.dtype | None = None, ) -> None: """Validate the fixed Qwen continuous-time and packed-CFG contract.""" _validate_qwen_pdd_config(config) @@ -126,6 +127,10 @@ def __init__( raise ValueError("guidance_eps must be finite and > 0.") 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 = ( @@ -133,6 +138,7 @@ def __init__( ) self.guidance_rescale = float(guidance_rescale) self.guidance_eps = float(guidance_eps) + self.compute_dtype = compute_dtype @staticmethod def _validate_state_and_time(state: torch.Tensor, time: torch.Tensor) -> None: @@ -188,8 +194,9 @@ def _parse_condition( raise ValueError(f"{name} tensors must be on {state.device}.") return encoder_hidden_states, attention_mask - @staticmethod - def _model_dtype(model: nn.Module, fallback: torch.dtype) -> torch.dtype: + 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 @@ -447,17 +454,27 @@ def teacher_velocity( f"{tuple(conditional.shape)} and {tuple(unconditional.shape)}." ) - conditional_fp32 = conditional.to(torch.float32) - guided = conditional_fp32 + (float(guidance_scale) - 1.0) * ( - conditional_fp32 - unconditional.to(torch.float32) + # FastGen applies CFG in the model-output dtype, including its BF16 + # rounding, before cfg_rescale promotes the result for norm math. + guided_model_dtype = conditional + (float(guidance_scale) - 1.0) * ( + conditional - unconditional ) + conditional_fp32 = conditional.to(torch.float32) + guided = guided_model_dtype.to(torch.float32) + # MR210 applies CFG after unpacking Qwen output to NCHW and leaves + # ``rescale_dims`` unset, so the norm spans every non-batch element. + # Reducing packed [P, F] here is algebraically identical because + # pack/unpack only reshapes and permutes those elements. + norm_dims = tuple(range(1, conditional_fp32.ndim)) conditional_norm = torch.linalg.vector_norm( conditional_fp32, - dim=-1, + dim=norm_dims, keepdim=True, ) - guided_norm = torch.linalg.vector_norm(guided, dim=-1, keepdim=True) + guided_norm = torch.linalg.vector_norm(guided, dim=norm_dims, keepdim=True) factor = self.guidance_rescale * conditional_norm / guided_norm.clamp_min( self.guidance_eps ) + (1.0 - self.guidance_rescale) - return self._unpack_single(guided * factor, state) + # FastGen's cfg_rescale returns to the conditional model-output dtype + # before PDD promotes the teacher target for FP32 loss math. + return self._unpack_single((guided * factor).to(conditional.dtype), state) diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 2a25b35e4d8..3b8d76cc1dd 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -40,6 +40,7 @@ from pdd.recipe import ( _materialize_zero_step_adamw_state, _projection_identity, + _require_fp32_optimizer_storage, _stage_and_shard_training_models, build_pdd_export_setup, build_pdd_setup, @@ -97,6 +98,23 @@ def test_zero_step_adamw_state_preserves_first_lazy_update() -> None: assert lazy_model.bias not in lazy_optimizer.state +def test_fp32_optimizer_storage_rejects_low_precision_masters_and_state() -> None: + model = nn.Linear(2, 2, dtype=torch.bfloat16) + optimizer = torch.optim.AdamW(model.parameters(), foreach=False, fused=False) + _materialize_zero_step_adamw_state(optimizer) + with pytest.raises(RuntimeError, match="master parameters must be FP32"): + _require_fp32_optimizer_storage(optimizer) + + fp32_model = nn.Linear(2, 2) + fp32_optimizer = torch.optim.AdamW(fp32_model.parameters(), foreach=False, fused=False) + _materialize_zero_step_adamw_state(fp32_optimizer) + fp32_optimizer.state[fp32_model.weight]["exp_avg"] = torch.zeros_like( + fp32_model.weight, dtype=torch.bfloat16 + ) + with pytest.raises(RuntimeError, match="exp_avg state must be FP32"): + _require_fp32_optimizer_storage(fp32_optimizer) + + def _require_exact_automodel() -> None: try: version = importlib.metadata.version("nemo_automodel") @@ -139,12 +157,12 @@ def parallelize(self, model): _projection_identity(projection), TrackedManager(), device=torch.device("cpu"), - dtype=torch.float32, fuse_qkv_projections=False, ) assert staged_student is student assert staged_teacher is teacher + assert {parameter.dtype for parameter in staged_student.parameters()} == {torch.float32} assert events == [ "student.to", "student.parallelize", @@ -153,6 +171,35 @@ def parallelize(self, model): ] +def test_training_setup_upcasts_bf16_models_to_fp32_masters() -> None: + class IdentityManager: + @staticmethod + def parallelize(model): + return model + + student = nn.Module() + student.proj_out = nn.Linear(2, 2, dtype=torch.bfloat16) + teacher = nn.Linear(2, 2, dtype=torch.bfloat16) + projection = convert_to_pdd_output_projection( + student, + PDDLayerSpec("proj_out", "channel_major"), + grid_size=4, + ) + + staged_student, staged_teacher = _stage_and_shard_training_models( + student, + teacher, + projection, + _projection_identity(projection), + IdentityManager(), + device=torch.device("cpu"), + fuse_qkv_projections=False, + ) + + assert {parameter.dtype for parameter in staged_student.parameters()} == {torch.float32} + assert {parameter.dtype for parameter in staged_teacher.parameters()} == {torch.float32} + + def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: return { "model": { diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index a945dd2e7fc..dfc12588808 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -578,7 +578,9 @@ def test_two_direct_updates_have_finite_gradients_updates_and_targeted_coverage( lifecycle.trainer.coverage.require_pairs([(0, 1), (1, 3), (2, 2), (3, 3)]) -def test_training_hard_aborts_for_teacher_gradient_zero_gradient_and_missing_coverage() -> None: +def test_training_hard_aborts_for_teacher_gradient_zero_gradient_and_missing_coverage( + monkeypatch, +) -> None: teacher_gradient = build_toy_lifecycle() teacher_gradient.teacher.scale.grad = torch.ones_like(teacher_gradient.teacher.scale) with pytest.raises(RuntimeError, match="teacher received a gradient"): @@ -593,6 +595,11 @@ def test_training_hard_aborts_for_teacher_gradient_zero_gradient_and_missing_cov with pytest.raises(RuntimeError, match="did not cover"): zero_gradient.trainer.coverage.require_pairs([(3, 3)]) + zero_actual_update = build_toy_lifecycle() + monkeypatch.setattr(zero_actual_update.trainer, "_projection_update_ratio", lambda before: 0.0) + with pytest.raises(RuntimeError, match="zero actual projection update"): + zero_actual_update.trainer.train_step(make_batch(("zero-actual-update",))) + nonfinite = build_toy_lifecycle() bad_batch = make_batch(("nan",)) bad_batch.data.fill_(float("nan")) diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index 6d7301f517b..cc0a76e41ec 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -24,7 +24,7 @@ from torch import nn from modelopt.torch.fastgen import PDDConfig, PDDPipeline -from modelopt.torch.fastgen.flow_matching import fusion_coefficients +from modelopt.torch.fastgen.flow_matching import fusion_coefficients, make_shifted_flow_grid class _HeadModel(nn.Module): @@ -109,7 +109,9 @@ def student_fused_block( "kwargs": model_kwargs, } ) - coefficients = fusion_coefficients(grid, start, end) + # 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 @@ -425,6 +427,12 @@ def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: 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: @@ -442,7 +450,8 @@ def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: 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"], grid) + 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 diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index ee70dfad260..87f4d730be9 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -213,7 +213,7 @@ def test_fused_student_matches_explicit_packed_head_weighting() -> None: assert student.proj_out(state.new_zeros(1, 5)).shape[-1] == 16 -def test_teacher_cfg_and_packed_token_norm_rescale_match_direct_reference() -> None: +def test_teacher_cfg_and_global_norm_rescale_match_mr210_reference() -> None: teacher = _TinyQwenTransformer() config = _config(guidance_scale=4.0) adapter = QwenImagePDDAdapter(config, guidance_rescale=1.0, guidance_eps=1e-5) @@ -233,10 +233,10 @@ def test_teacher_cfg_and_packed_token_norm_rescale_match_direct_reference() -> N guided = conditional + 3.0 * (conditional - unconditional) factor = torch.linalg.vector_norm( conditional, - dim=-1, + dim=(1, 2), keepdim=True, - ) / torch.linalg.vector_norm(guided, dim=-1, keepdim=True).clamp_min(1e-5) - expected = unpack_latents(guided * factor, 4, 4) + ) / torch.linalg.vector_norm(guided, dim=(1, 2), keepdim=True).clamp_min(1e-5) + expected = unpack_latents((guided * factor).to(teacher.calls[0]["output"].dtype), 4, 4) assert actual.dtype == torch.float32 torch.testing.assert_close(actual, expected) @@ -246,6 +246,51 @@ def test_teacher_cfg_and_packed_token_norm_rescale_match_direct_reference() -> N assert all(call["guidance"] is None for call in teacher.calls) +def test_teacher_cfg_returns_to_low_precision_model_output_dtype() -> None: + class LowPrecisionTeacher(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(guidance_embeds=False) + self.anchor = nn.Parameter(torch.zeros((), dtype=torch.bfloat16), requires_grad=False) + self.outputs: list[torch.Tensor] = [] + + def forward(self, *, hidden_states, encoder_hidden_states, **kwargs): + value = encoder_hidden_states.mean(dim=(1, 2), keepdim=True).to(torch.bfloat16) + output = hidden_states.to(torch.bfloat16) + value + self.outputs.append(output.detach().clone()) + return (output,) + + teacher = LowPrecisionTeacher() + state, time, condition, negative_condition = _inputs() + actual = QwenImagePDDAdapter(_config(guidance_scale=4.0)).teacher_velocity( + teacher, + state, + time, + condition=condition, + negative_condition=negative_condition, + ) + + assert actual.dtype == torch.bfloat16 + conditional, unconditional = teacher.outputs + guided_bf16 = conditional + 3.0 * (conditional - unconditional) + conditional_fp32 = conditional.float() + guided_fp32 = guided_bf16.float() + 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) + expected = unpack_latents((guided_fp32 * factor).to(torch.bfloat16), 4, 4) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + unrounded_guided = conditional_fp32 + 3.0 * (conditional_fp32 - unconditional.float()) + unrounded_factor = torch.linalg.vector_norm( + conditional_fp32, dim=(1, 2), keepdim=True + ) / torch.linalg.vector_norm(unrounded_guided, dim=(1, 2), keepdim=True).clamp_min(1e-5) + unrounded = unpack_latents( + (unrounded_guided * unrounded_factor).to(torch.bfloat16), 4, 4 + ) + assert not torch.equal(actual, unrounded) + + def test_guidance_disabled_teacher_is_one_conditional_call_without_negative_condition() -> None: teacher = _TinyQwenTransformer() config = _config(guidance_scale=None) @@ -286,6 +331,8 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N QwenImagePDDAdapter(_config(), guidance_rescale=1.1) with pytest.raises(ValueError, match="guidance_eps"): QwenImagePDDAdapter(_config(), guidance_eps=0.0) + with pytest.raises(TypeError, match="compute_dtype"): + QwenImagePDDAdapter(_config(), compute_dtype=torch.long) transformer = _TinyQwenTransformer() transformer.config.guidance_embeds = True From 95fbcfddb73c92b5ce7713f8f8920c050b5c7af3 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Wed, 15 Jul 2026 21:46:21 -0700 Subject: [PATCH 25/45] Match Qwen PDD execution to FastGen MR210 Signed-off-by: Meng Xin --- .../fastgen/fastgen_data/collate_fns.py | 18 +- examples/diffusers/fastgen/pdd/README.md | 14 +- examples/diffusers/fastgen/pdd/checkpoint.py | 5 +- examples/diffusers/fastgen/pdd/export.py | 17 +- .../fastgen/pdd/export_qwen_image.py | 7 + .../fastgen/pdd/inference_qwen_image.py | 19 +- examples/diffusers/fastgen/pdd/recipe.py | 34 +- .../torch/fastgen/plugins/qwen_image_pdd.py | 309 +++++++++++++++++- .../pdd_checkpoint_failure_distributed.py | 4 +- .../fastgen/pdd_export_distributed.py | 7 +- .../fastgen/test_pdd_inference_checkpoint.py | 51 ++- .../fastgen/test_pdd_recipe_setup.py | 13 +- .../fastgen/test_pdd_training_lifecycle.py | 3 + .../fastgen/test_vendored_migration.py | 31 ++ .../fastgen/test_qwen_image_pdd_plugin.py | 291 +++++++++++++++-- 15 files changed, 761 insertions(+), 62 deletions(-) diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index 57fb4b18e50..76368cc908a 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -40,6 +40,7 @@ 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 @@ -89,7 +90,11 @@ def collate_fn_text_to_image( 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], @@ -108,9 +113,14 @@ def collate_fn_text_to_image( 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] + 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: diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index 734d3677de3..b056fa49a9c 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -19,6 +19,15 @@ The provided schedules use the 128-interval grid as follows: Install the shared requirements from the repository root, then launch with released AutoModel APIs. No AutoModel, Diffusers, or Qwen source changes are required. +The reproducibility arm executes the Qwen forward semantics from FastGen MR210 commit +`c8100b1347b278511336dccfc074a461457216ec`: BF16 image/text compute, FP32 normalized time through +the timestep projection, and the MR attention behavior for zero-padded text. ModelOpt adopts the +loaded Diffusers model into a compatible root without editing or monkeypatching external classes. +The exact FastGen, Diffusers 0.38.0 transformer, and timestep-embedding source identities are +authenticated in every checkpoint and export and are required again by resume and inference. +Guidance-embedded/`zero_cond_t` models, additional time conditioning, ControlNet, PEFT, and QKV +fusion are deliberately rejected because they are outside that executed algorithm. + ```bash pip install -r examples/diffusers/fastgen/requirements.txt export MODELOPT_FASTGEN_DATASET_CACHE_DIR=/absolute/path/to/qwen_image_cache @@ -47,7 +56,10 @@ treatment; the MR210 reference arm uses `5e-5` with 1,000 warmup steps. On 16 fo default per-rank batch gives global batch size 256; other GPU topologies must set per-rank batch to `256 / world_size` because this recipe does not use gradient accumulation. Checkpoints include the student, optimizer, scheduler, RNG, trainer, and exact replayable sampler -state needed to resume the next committed batch. +state needed to resume the next committed batch. FP32 master parameters and Adam state are sharded +while forward/backward compute remains BF16 and gradient reduction remains FP32. The FSDP policy +does not recursively cast forward inputs; ModelOpt casts image/text tensors to BF16 itself so the +FP32 timestep cannot be rounded before the source-locked forward receives it. Start with a one-node smoke and scale only after it passes; project training runs are capped at 16 nodes. diff --git a/examples/diffusers/fastgen/pdd/checkpoint.py b/examples/diffusers/fastgen/pdd/checkpoint.py index 102af9fe37d..0cccaee09c0 100644 --- a/examples/diffusers/fastgen/pdd/checkpoint.py +++ b/examples/diffusers/fastgen/pdd/checkpoint.py @@ -31,11 +31,12 @@ import torch.distributed as dist from modelopt.torch.fastgen import PDDMetadata +from modelopt.torch.fastgen.plugins.qwen_image_pdd import require_qwen_image_pdd_forward_substrate if TYPE_CHECKING: from collections.abc import Sequence -_CHECKPOINT_SCHEMA_VERSION = 2 +_CHECKPOINT_SCHEMA_VERSION = 3 _COMPLETE_SCHEMA_VERSION = 1 _FORBIDDEN_ARTIFACT_TOKENS = ("fake_score", "discriminator", "ema", "r1", "gan") @@ -159,6 +160,7 @@ def _read_json(path: Path) -> dict[str, Any]: def build_pdd_checkpoint_identity( *, metadata: PDDMetadata, + forward_substrate: Mapping[str, Any], model_id: str, model_revision: str | None, guidance_scale: float | None, @@ -228,6 +230,7 @@ def build_pdd_checkpoint_identity( raise ValueError(f"AutoModel snapshot is missing identity keys: {missing}.") return { "schema_version": _CHECKPOINT_SCHEMA_VERSION, + "forward_substrate": require_qwen_image_pdd_forward_substrate(forward_substrate), "model": {"id": model_id, "revision": model_revision, "dtype": dtype}, "pdd_metadata": metadata.to_dict(), "guidance": { diff --git a/examples/diffusers/fastgen/pdd/export.py b/examples/diffusers/fastgen/pdd/export.py index 690cc8a9020..61239b96c30 100644 --- a/examples/diffusers/fastgen/pdd/export.py +++ b/examples/diffusers/fastgen/pdd/export.py @@ -31,7 +31,10 @@ from safetensors.torch import save_file from modelopt.torch.fastgen import PDDConfig, PDDMetadata -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QWEN_IMAGE_PDD_LAYER_SPEC, + require_qwen_image_pdd_forward_substrate, +) from .artifacts import ( load_canonical_json, @@ -41,7 +44,7 @@ write_canonical_json, ) -_EXPORT_SCHEMA_VERSION = 1 +_EXPORT_SCHEMA_VERSION = 2 _COMPLETE_SCHEMA_VERSION = 1 _EXPORT_FORMAT = "modelopt-pdd-safetensors" _CONFIG_FILE = "config.json" @@ -167,12 +170,20 @@ def _bounded_shard_groups( def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> dict[str, Any]: if not isinstance(identity, Mapping): raise TypeError("PDD export identity must be a mapping.") - required = {"automodel", "guidance", "model", "pdd_metadata", "topology"} + required = { + "automodel", + "forward_substrate", + "guidance", + "model", + "pdd_metadata", + "topology", + } missing = sorted(required.difference(identity)) if missing: raise ValueError(f"PDD export identity is missing keys: {missing}.") if identity["pdd_metadata"] != metadata.to_dict(): raise ValueError("PDD export metadata does not match the checkpoint identity.") + require_qwen_image_pdd_forward_substrate(identity["forward_substrate"]) model = _require_exact_mapping( identity["model"], {"id", "revision", "dtype"}, name="identity.model" ) diff --git a/examples/diffusers/fastgen/pdd/export_qwen_image.py b/examples/diffusers/fastgen/pdd/export_qwen_image.py index 8c1fcdf2f0b..67976b87b7c 100644 --- a/examples/diffusers/fastgen/pdd/export_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/export_qwen_image.py @@ -188,6 +188,9 @@ def _collective_publication_preflight(output_dir: Path) -> Mapping[str, Any]: def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: from modelopt.torch.fastgen import PDDMetadata + from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + require_qwen_image_pdd_forward_substrate, + ) identity = manifest.get("identity") if not isinstance(identity, Mapping): @@ -197,6 +200,7 @@ def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, raise RuntimeError("PDD checkpoint has no PDD metadata mapping.") if PDDMetadata.from_dict(pdd_metadata) != setup.metadata: raise RuntimeError("PDD checkpoint metadata does not match the configured student.") + require_qwen_image_pdd_forward_substrate(identity.get("forward_substrate")) if identity.get("model") != { "id": config.model_id, "revision": config.model_revision, @@ -234,7 +238,10 @@ def _collective_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[s def _checkpoint_selector_identity(config: Any, setup: Any) -> dict[str, Any]: + from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_FORWARD_SUBSTRATE + return { + "forward_substrate": dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE), "model": { "id": config.model_id, "revision": config.model_revision, diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index 3cbf39c2991..3de0df28a5c 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -117,19 +117,29 @@ def build_pdd_student(export_dir: str | Path) -> tuple[nn.Module, Any, torch.dty """Reconstruct and strictly load the converted Qwen student on CPU.""" from diffusers import QwenImageTransformer2DModel - from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd + from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + adopt_qwen_image_mr210_forward, + convert_qwen_image_to_pdd, + require_qwen_image_pdd_forward_substrate, + ) from pdd.export import inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata descriptor = inspect_pdd_export(export_dir) model_identity = _model_identity(descriptor) + require_qwen_image_pdd_forward_substrate( + descriptor.manifest["identity"].get("forward_substrate") + ) dtype = _dtype_from_name(model_identity["dtype"]) - student = QwenImageTransformer2DModel.from_config(dict(descriptor.transformer_config)) + loaded_transformer = QwenImageTransformer2DModel.from_config( + dict(descriptor.transformer_config) + ) + student = adopt_qwen_image_mr210_forward(loaded_transformer) metadata = descriptor.metadata _validate_qwen_projection(student, metadata) config = pdd_config_from_metadata(metadata, blocks=metadata.inference_blocks) convert_qwen_image_to_pdd(student, config) - student.to(dtype=dtype) descriptor = load_pdd_export_into_model(export_dir, student) + student.to(dtype=dtype) return student, descriptor, dtype @@ -256,7 +266,7 @@ def main() -> None: ) pipe.to(device) config = pdd_config_from_metadata(descriptor.metadata, schedule=args.schedule) - adapter = QwenImagePDDAdapter(config) + adapter = QwenImagePDDAdapter(config, compute_dtype=dtype) sampler = PDDPipeline(student, nn.Identity(), config, adapter) prompt_embeds, prompt_mask = pipe.encode_prompt( prompt=args.prompt, @@ -318,6 +328,7 @@ def count_invocation( "blocks": list(config.inference_blocks), "height": args.height, "width": args.width, + "forward_substrate_id": descriptor.manifest["identity"]["forward_substrate"]["id"], "export_manifest_sha256": sha256_file(descriptor.root / "manifest.json"), "output": {"path": output_reference, "sha256": sha256_file(output)}, "scheduler_steps": expected_invocations, diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 01f5143095e..4591e865c52 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -32,7 +32,9 @@ from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, QwenImagePDDAdapter, + adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, ) @@ -469,6 +471,8 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: model.get("fuse_qkv_projections", False), name="model.fuse_qkv_projections", ) + if fuse_qkv_projections: + raise ValueError("authenticated Qwen MR210 PDD does not support QKV fusion.") seed = _require_int_at_least(raw.get("seed", 42), name="seed", minimum=0) max_steps = _require_int_at_least( @@ -566,6 +570,9 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: ) if guidance_eps == 0.0: raise ValueError("guidance.eps must be > 0.") + dtype = _resolve_dtype(model.get("torch_dtype", "bfloat16")) + if dtype != torch.bfloat16: + raise ValueError("authenticated Qwen MR210 PDD requires model.torch_dtype='bfloat16'.") return PDDRecipeConfig( model_id=model_id, @@ -608,7 +615,7 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: adam_betas=adam_betas, adam_eps=adam_eps, device=torch.device(model.get("device", "cuda" if torch.cuda.is_available() else "cpu")), - dtype=_resolve_dtype(model.get("torch_dtype", "bfloat16")), + dtype=dtype, fuse_qkv_projections=fuse_qkv_projections, ) @@ -796,7 +803,9 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: ) from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager - pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + student = adopt_qwen_image_mr210_forward(loaded_transformer) + pipe.transformer = student teacher = copy.deepcopy(student).eval().requires_grad_(False) lifecycle.append("load/select") @@ -829,7 +838,7 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: mp_policy=MixedPrecisionPolicy( param_dtype=config.dtype, reduce_dtype=torch.float32, - output_dtype=None, + output_dtype=config.dtype, cast_forward_inputs=False, ), ) @@ -943,7 +952,9 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager lifecycle = ["load/select"] - pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + student = adopt_qwen_image_mr210_forward(loaded_transformer) + pipe.transformer = student raw_transformer_config = getattr(student, "config", None) to_dict = getattr(raw_transformer_config, "to_dict", None) if callable(to_dict): @@ -959,7 +970,7 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: identity = _projection_identity(projection) metadata = PDDMetadata.from_config(config.pdd, projection) lifecycle.append("pdd_conversion") - student.to(device=config.device, dtype=config.dtype) + student.to(device=config.device, dtype=torch.float32) _require_projection_identity(student, projection, identity, stage="device placement") lifecycle.append("device") @@ -977,7 +988,17 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: f"Pure-DP PDD export requires fsdp.dp_size ({dp_size}) to equal world size " f"({world_size})." ) - strategy = FSDP2Config(activation_checkpointing=False) + from torch.distributed.fsdp import MixedPrecisionPolicy + + strategy = FSDP2Config( + activation_checkpointing=False, + mp_policy=MixedPrecisionPolicy( + param_dtype=config.dtype, + reduce_dtype=torch.float32, + output_dtype=config.dtype, + cast_forward_inputs=False, + ), + ) distributed_setup = DistributedSetup.build( strategy=strategy, parallelism_sizes=ParallelismSizes(dp_size=dp_size), @@ -1148,6 +1169,7 @@ def setup(self) -> None: identity = build_pdd_checkpoint_identity( metadata=self.setup_artifacts.metadata, + forward_substrate=QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, model_id=config.model_id, model_revision=config.model_revision, guidance_scale=config.pdd.guidance_scale, diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index 1a54787fcbb..bf81fd0953a 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -17,6 +17,8 @@ from __future__ import annotations +import hashlib +import inspect import math from collections.abc import Mapping from typing import Any @@ -29,27 +31,323 @@ from .qwen_image import build_img_shapes, pack_latents, unpack_latents __all__ = [ + "QWEN_IMAGE_PDD_FORWARD_SUBSTRATE", + "QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID", "QWEN_IMAGE_PDD_LAYER_SPEC", "QwenImagePDDAdapter", + "adopt_qwen_image_mr210_forward", "convert_qwen_image_to_pdd", + "require_qwen_image_pdd_forward_substrate", ] +QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID = ( + "pdd_qwen_mr210_c8100b1347b278511336dccfc074a461457216ec_" + "qwen_33706683487ba16d133b99b73be27b21164c53335441d77b1dcabbfca970f70e" +) +QWEN_IMAGE_PDD_FORWARD_SUBSTRATE = { + "id": QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID, + "fastgen_commit": "c8100b1347b278511336dccfc074a461457216ec", + "fastgen_qwen_source_sha256": ( + "33706683487ba16d133b99b73be27b21164c53335441d77b1dcabbfca970f70e" + ), + "diffusers_version": "0.38.0", + "diffusers_qwen_source_sha256": ( + "34c864b0b066a4a9eb84e40e1bb77b7df303c165e7910600b402a0f5f8d8f94e" + ), + "diffusers_embeddings_source_sha256": ( + "d7a90ef799569e3f0fab41cadde1ecba023abd053af956c022bbfc097662a302" + ), +} + 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", } +_QWEN_IMAGE_ROOT_CHILDREN = ( + "pos_embed", + "time_text_embed", + "txt_norm", + "img_in", + "txt_in", + "transformer_blocks", + "norm_out", + "proj_out", +) +_ADOPTED_QWEN_TYPES: dict[type[nn.Module], type[nn.Module]] = {} + + +def require_qwen_image_pdd_forward_substrate(value: Any) -> dict[str, str]: + """Return the authenticated MR210 Qwen substrate or reject it exactly.""" + if not isinstance(value, Mapping) or dict(value) != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE: + raise ValueError( + "Qwen-Image PDD requires the authenticated MR210 forward substrate " + f"{QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID!r}." + ) + return dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE) + + +def _sha256_source(owner: type[Any]) -> str: + source = inspect.getsourcefile(owner) + if source is None: + raise RuntimeError(f"cannot locate source for {owner.__module__}.{owner.__qualname__}.") + with open(source, "rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + + +def _require_qwen_source_identity(transformer: nn.Module) -> None: + transformer_type = type(transformer) + if ( + transformer_type.__module__ != "diffusers.models.transformers.transformer_qwenimage" + or transformer_type.__name__ != "QwenImageTransformer2DModel" + ): + raise TypeError("MR210 adoption requires the pinned Diffusers QwenImageTransformer2DModel.") + if ( + _sha256_source(transformer_type) + != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_qwen_source_sha256"] + ): + raise RuntimeError("Diffusers Qwen transformer source does not match the MR210 substrate.") + embedding_types = ( + type(transformer.time_text_embed.time_proj), + type(transformer.time_text_embed.timestep_embedder), + ) + if any( + _sha256_source(embedding_type) + != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_embeddings_source_sha256"] + for embedding_type in embedding_types + ): + raise RuntimeError( + "Diffusers timestep embedding source does not match the MR210 substrate." + ) + try: + import diffusers # Optional dependency required only by the Qwen adoption path. + except ImportError as error: # pragma: no cover - the transformer itself requires Diffusers + raise RuntimeError("Diffusers is required for Qwen-Image PDD adoption.") from error + if diffusers.__version__ != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_version"]: + raise RuntimeError( + "Diffusers version does not match the authenticated Qwen MR210 substrate." + ) + + +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 _require_binary_prefix_mask( + encoder_hidden_states: torch.Tensor, + mask: torch.Tensor, +) -> None: + if mask.ndim != 2 or tuple(mask.shape) != tuple(encoder_hidden_states.shape[:2]): + raise ValueError("Qwen MR210 mask must match the text batch and sequence dimensions.") + if mask.dtype.is_floating_point or mask.dtype.is_complex: + raise TypeError("Qwen MR210 mask must use an integer or boolean dtype.") + if mask.device != encoder_hidden_states.device: + raise ValueError("Qwen MR210 mask and text embeddings must share a device.") + if mask.shape[0] == 0 or mask.shape[1] == 0: + raise ValueError("Qwen MR210 masks must have nonempty batch and sequence dimensions.") + mask_int = mask.to(torch.int64) + mask_bool = mask.bool() + binary = torch.all((mask_int == 0) | (mask_int == 1)) + prefix = torch.all(mask_int[:, 1:] <= mask_int[:, :-1]) + lengths = mask_int.sum(dim=1) + valid_lengths = torch.all(lengths > 0) & (lengths.max() == mask.shape[1]) + zero_padding = torch.all(encoder_hidden_states[~mask_bool] == 0) + if not bool((binary & prefix & valid_lengths & zero_padding).item()): + raise ValueError( + "Qwen MR210 requires nonempty binary prefix masks, a longest unpadded row, " + "and zero padding." + ) + + +class _QwenImageMR210ForwardMixin: + """Execute FastGen MR210's Qwen forward without altering Diffusers classes. + + Source contract: ``fastgen/networks/QwenImage/network.py`` at + ``c8100b1347b278511336dccfc074a461457216ec``. + """ + + def forward( + self, + 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, + txt_seq_lens: list[int] | None = None, + guidance: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + controlnet_block_samples: Any = None, + additional_t_cond: Any = None, + return_dict: bool = True, + *, + max_txt_seq_len: int | None = None, + ) -> Any: + """Run the source-locked MR210 regular-output path.""" + if hidden_states.ndim != 3 or hidden_states.dtype != torch.bfloat16: + raise TypeError("Qwen MR210 hidden_states must be packed BF16 [B, P, C].") + if ( + not isinstance(encoder_hidden_states, torch.Tensor) + or encoder_hidden_states.ndim != 3 + or encoder_hidden_states.dtype != torch.bfloat16 + ): + raise TypeError("Qwen MR210 encoder_hidden_states must be BF16 [B, S, D].") + if not isinstance(encoder_hidden_states_mask, torch.Tensor): + raise TypeError("Qwen MR210 requires encoder_hidden_states_mask.") + if not isinstance(timestep, torch.Tensor) or timestep.dtype != torch.float32: + raise TypeError("Qwen MR210 timestep must remain FP32 at transformer entry.") + if timestep.shape != (hidden_states.shape[0],): + raise ValueError("Qwen MR210 timestep must contain one value per batch item.") + if encoder_hidden_states.shape[0] != hidden_states.shape[0]: + raise ValueError("Qwen MR210 image and text batch sizes must match.") + if img_shapes is None or len(img_shapes) != hidden_states.shape[0]: + raise ValueError("Qwen MR210 img_shapes must contain one entry per batch item.") + if txt_seq_lens is not None: + raise ValueError("Qwen MR210 does not support txt_seq_lens.") + if guidance is not None: + raise ValueError("Qwen MR210 does not support transformer guidance embeddings.") + if attention_kwargs: + raise ValueError("Qwen MR210 does not support nonempty attention_kwargs.") + if controlnet_block_samples is not None: + raise ValueError("Qwen MR210 does not support ControlNet residuals.") + if additional_t_cond is not None: + raise ValueError("Qwen MR210 does not support additional time conditioning.") + if type(return_dict) is not bool: + raise TypeError("return_dict must be bool.") + _require_binary_prefix_mask(encoder_hidden_states, encoder_hidden_states_mask) + sequence_length = encoder_hidden_states.shape[1] + if max_txt_seq_len is not None and max_txt_seq_len != sequence_length: + raise ValueError("Qwen MR210 max_txt_seq_len must equal the padded text length.") + + hidden_states = self.img_in(hidden_states) + encoder_hidden_states = self.txt_norm(encoder_hidden_states) + encoder_hidden_states = self.txt_in(encoder_hidden_states) + if timestep.dtype != torch.float32: + raise RuntimeError("Qwen MR210 timestep was rounded before time_text_embed.") + temb = self.time_text_embed(timestep, hidden_states) + image_rotary_emb = self.pos_embed( + img_shapes, + max_txt_seq_len=sequence_length, + 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, + joint_attention_kwargs=attention_kwargs, + ) + + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + if not return_dict: + return (output,) + try: + from diffusers.models.modeling_outputs import ( # Optional Qwen runtime dependency. + Transformer2DModelOutput, + ) + except ImportError as error: # pragma: no cover - adoption already requires Diffusers + raise RuntimeError("Diffusers output types are unavailable.") from error + return Transformer2DModelOutput(sample=output) + + +def _adopted_qwen_type(base: type[nn.Module]) -> type[nn.Module]: + adopted = _ADOPTED_QWEN_TYPES.get(base) + if adopted is None: + adopted = type( + f"ModelOptMR210{base.__name__}", + (_QwenImageMR210ForwardMixin, base), + {"__module__": __name__}, + ) + _ADOPTED_QWEN_TYPES[base] = adopted + return adopted + + +def adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: + """Adopt loaded Qwen children into a Diffusers-compatible MR210 forward root.""" + if isinstance(transformer, _QwenImageMR210ForwardMixin): + return transformer + if not isinstance(transformer, nn.Module): + raise TypeError(f"transformer must be nn.Module, got {type(transformer).__name__}.") + _require_qwen_source_identity(transformer) + if tuple(transformer._modules) != _QWEN_IMAGE_ROOT_CHILDREN: + raise RuntimeError("Qwen root child layout does not match the authenticated substrate.") + if transformer._parameters or transformer._buffers: + raise RuntimeError("Qwen root unexpectedly registers direct parameters or buffers.") + if _config_guidance_embeds(transformer): + raise ValueError("Qwen MR210 does not support transformer guidance embeddings.") + if getattr(transformer, "peft_config", None): + raise ValueError("Qwen MR210 does not support active PEFT adapters.") + if any(getattr(module, "fused_projections", False) for module in transformer.modules()): + raise ValueError("Qwen MR210 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 MR210 requires {name}=False.") + hook_names = ( + "_backward_hooks", + "_backward_pre_hooks", + "_forward_hooks", + "_forward_pre_hooks", + "_load_state_dict_post_hooks", + "_load_state_dict_pre_hooks", + "_state_dict_hooks", + "_state_dict_pre_hooks", + ) + if any(getattr(transformer, name, None) for name in hook_names): + raise RuntimeError("Qwen root hooks must be empty before MR210 adoption.") + + adopted_type = _adopted_qwen_type(type(transformer)) + adopted = adopted_type.__new__(adopted_type) + nn.Module.__init__(adopted) + adopted._internal_dict = transformer._internal_dict + for name in ("out_channels", "inner_dim", "gradient_checkpointing", "zero_cond_t"): + setattr(adopted, name, getattr(transformer, name)) + if hasattr(transformer, "_gradient_checkpointing_func"): + adopted._gradient_checkpointing_func = transformer._gradient_checkpointing_func + for name, child in transformer._modules.items(): + adopted.add_module(name, child) + adopted.train(transformer.training) + if tuple(adopted.state_dict()) != tuple(transformer.state_dict()): + raise RuntimeError("Qwen state keys changed during MR210 adoption.") + if any( + adopted_parameter is not source_parameter + for adopted_parameter, source_parameter in zip( + adopted.parameters(), transformer.parameters(), strict=True + ) + ): + raise RuntimeError("Qwen parameter identity changed during MR210 adoption.") + return adopted + def _config_guidance_embeds(transformer: nn.Module) -> bool: config = getattr(transformer, "config", None) @@ -171,7 +469,7 @@ def _parse_condition( attention_mask, torch.Tensor ): raise TypeError(f"{name} entries must be tensors.") - if encoder_hidden_states.ndim < 2 or attention_mask.ndim != 2: + 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)}." @@ -192,6 +490,7 @@ def _parse_condition( raise ValueError(f"{name} batch size must match state batch size {batch_size}.") if encoder_hidden_states.device != state.device or attention_mask.device != state.device: raise ValueError(f"{name} tensors must be on {state.device}.") + _require_binary_prefix_mask(encoder_hidden_states, attention_mask) return encoder_hidden_states, attention_mask def _model_dtype(self, model: nn.Module, fallback: torch.dtype) -> torch.dtype: @@ -268,7 +567,13 @@ def _call_packed( ) batch_size, _, height, width = state.shape - packed_state = pack_latents(state).to(self._model_dtype(model, state.dtype)) + model_dtype = self._model_dtype(model, state.dtype) + if model_dtype != torch.bfloat16: + raise TypeError("authenticated Qwen MR210 execution requires BF16 compute.") + if time.dtype != torch.float32: + raise TypeError("authenticated Qwen MR210 execution requires FP32 time.") + packed_state = pack_latents(state).to(model_dtype) + encoder_hidden_states = encoder_hidden_states.to(model_dtype) output = model( hidden_states=packed_state, timestep=time, diff --git a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py index bdf976f7b6b..bc13f6c0771 100644 --- a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py @@ -141,7 +141,7 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 2, "topology": {"world_size": 2}}, + identity={"schema_version": 3, "topology": {"world_size": 2}}, ) initial.save() trainer.completed_steps = 2 @@ -156,7 +156,7 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 2, "topology": {"world_size": 2}}, + identity={"schema_version": 3, "topology": {"world_size": 2}}, ) message = None try: diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py index 6fb7a8c6c00..20d42417b91 100644 --- a/tests/examples/diffusers/fastgen/pdd_export_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_export_distributed.py @@ -38,12 +38,14 @@ from pdd.inference_qwen_image import build_pdd_student from pdd.recipe import build_pdd_export_setup, resolve_pdd_recipe_config +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_FORWARD_SUBSTRATE + def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: return { "model": { "pretrained_model_name_or_path": str(model_dir), - "torch_dtype": "float32", + "torch_dtype": "bfloat16", "device": "cpu", "transformer_engine_linear": False, "peft": None, @@ -134,10 +136,11 @@ def main() -> None: ) identity = { "schema_version": 1, + "forward_substrate": dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE), "model": { "id": "Qwen/Qwen-Image", "revision": "3" * 40, - "dtype": "float32", + "dtype": "bfloat16", }, "pdd_metadata": destination.metadata.to_dict(), "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py index 69ecdf20070..7c8e92819f3 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -43,6 +43,7 @@ from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, QwenImagePDDAdapter, convert_qwen_image_to_pdd, ) @@ -52,8 +53,8 @@ class _TinyQwen(nn.Module): def __init__(self) -> None: super().__init__() self.config = SimpleNamespace(guidance_embeds=False, in_channels=4) - self.backbone = nn.Linear(4, 5) - self.proj_out = nn.Linear(5, 4) + self.backbone = nn.Linear(4, 5, dtype=torch.bfloat16) + self.proj_out = nn.Linear(5, 4, dtype=torch.bfloat16) self.calls = 0 def forward( @@ -69,10 +70,14 @@ def forward( ): del img_shapes, guidance, return_dict condition = encoder_hidden_states.mean(dim=(1, 2), keepdim=True) - condition += encoder_hidden_states_mask.sum(dim=1)[:, None, None] / 100 + condition += (encoder_hidden_states_mask.sum(dim=1)[:, None, None] / 100).to( + condition.dtype + ) hidden = torch.tanh(self.backbone(hidden_states)) self.calls += 1 - return (self.proj_out(hidden + condition + timestep[:, None, None] / 10),) + hidden = hidden + condition + hidden = hidden + (timestep[:, None, None] / 10).to(hidden.dtype) + return (self.proj_out(hidden),) def _config(blocks=(32, 32, 32, 32)) -> PDDConfig: @@ -104,7 +109,8 @@ def _converted(seed: int = 17): def _identity(metadata: PDDMetadata) -> dict: return { "schema_version": 1, - "model": {"id": "synthetic-qwen", "revision": "f" * 40, "dtype": "float32"}, + "forward_substrate": dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE), + "model": {"id": "synthetic-qwen", "revision": "f" * 40, "dtype": "bfloat16"}, "pdd_metadata": metadata.to_dict(), "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, "automodel": { @@ -132,17 +138,24 @@ def _write(tmp_path: pathlib.Path): "completed_steps": 10, }, modelopt_source={"commit": "4" * 40, "dirty": False}, - max_shard_bytes=12_000, + max_shard_bytes=5_800, ) return output, model, config, metadata def _condition(): - return torch.tensor([[[0.2, -0.3], [0.1, 0.4]]]), torch.ones(1, 2, dtype=torch.long) + return torch.tensor([[[0.2, -0.3], [0.1, 0.4]]], dtype=torch.bfloat16), torch.ones( + 1, 2, dtype=torch.long + ) def _sample(model: nn.Module, config: PDDConfig, noise: torch.Tensor) -> torch.Tensor: - pipeline = PDDPipeline(model, nn.Identity(), config, QwenImagePDDAdapter(config)) + pipeline = PDDPipeline( + model, + nn.Identity(), + config, + QwenImagePDDAdapter(config, compute_dtype=torch.bfloat16), + ) return pipeline.sample(noise.clone(), condition=_condition()) @@ -267,6 +280,28 @@ def test_export_rejects_unpinned_local_model_identity(tmp_path) -> None: ) +def test_export_rejects_missing_or_mismatched_forward_substrate(tmp_path) -> None: + model, _config_value, metadata = _converted() + for substrate in (None, {**QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, "id": "canonical"}): + identity = _identity(metadata) + identity["forward_substrate"] = substrate + with pytest.raises(ValueError, match="authenticated MR210"): + write_pdd_export( + tmp_path / f"bad-substrate-{substrate is None}", + model.state_dict(), + metadata=metadata, + transformer_config={"in_channels": 4}, + identity=identity, + source_checkpoint={ + "name": "step_00000010", + "manifest_sha256": "3" * 64, + "completed_steps": 10, + }, + modelopt_source={"commit": "4" * 40, "dirty": False}, + max_shard_bytes=12_000, + ) + + def test_export_rejects_nonfinite_and_existing_destination(tmp_path) -> None: output, model, _config_value, metadata = _write(tmp_path) with pytest.raises(FileExistsError): diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 3b8d76cc1dd..d3abf821d9d 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -204,7 +204,7 @@ def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: return { "model": { "pretrained_model_name_or_path": str(model_dir), - "torch_dtype": "float32", + "torch_dtype": "bfloat16", "device": "cpu", "transformer_engine_linear": False, "peft": None, @@ -422,6 +422,8 @@ def test_pdd_finetune_namespace_module_help() -> None: ("model", "guidance_embeds", True, "guidance embeddings"), ("model", "device_map", "auto", "device_map"), ("model", "quantization_config", {"bits": 8}, "quantization_config"), + ("model", "fuse_qkv_projections", True, "QKV fusion"), + ("model", "torch_dtype", "float32", "requires model.torch_dtype='bfloat16'"), ], ) def test_incompatible_modes_fail_during_config_resolution( @@ -509,7 +511,7 @@ def test_payload_hash_verification_mode_must_be_bool(tmp_path) -> None: raw = _raw_config(tmp_path) raw["data"]["dataloader"]["verify_payload_hashes"] = "false" - with pytest.raises(TypeError, match="data.dataloader.verify_payload_hashes must be bool"): + with pytest.raises(TypeError, match=r"data\.dataloader\.verify_payload_hashes must be bool"): resolve_pdd_recipe_config(raw) @@ -575,7 +577,7 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: before = snapshot_installed_distribution() model_dir = create_tiny_qwen_image_pipeline_dir(tmp_path) initialize_pdd_distributed(backend="gloo", timeout_minutes=1) - config = resolve_pdd_recipe_config(_raw_config(model_dir, qkv=True)) + config = resolve_pdd_recipe_config(_raw_config(model_dir)) source = build_pdd_setup(config) @@ -597,6 +599,11 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert source.projection.out_features == source.projection.base_out_features * 4 assert "proj_out.weight" in source.checkpoint_keys assert source.student.state_dict()["proj_out.weight"].shape[0] == source.projection.out_features + policy = source.distributed_setup.strategy_config.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 not any(parameter.requires_grad for parameter in source.teacher.parameters()) optimizer_parameters = [ parameter for group in source.optimizer.param_groups for parameter in group["params"] diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index dfc12588808..d7972d52c8e 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -51,6 +51,8 @@ from pdd.verify_readonly_automodel import snapshot_installed_distribution from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_FORWARD_SUBSTRATE + def _released_sampler(sample_ids: tuple[str, ...]) -> ReplayableBatchSampler: sampler_module = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") @@ -98,6 +100,7 @@ def _checkpointer(lifecycle, checkpoint_dir): def _identity(lifecycle, scheduler, sample_ids): return build_pdd_checkpoint_identity( metadata=lifecycle.metadata, + forward_substrate=QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, model_id="synthetic-pdd-toy", model_revision=None, guidance_scale=None, diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index 8e546110b0c..11965cd7f1a 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -201,6 +201,37 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): assert out["metadata"]["sample_ids"].tolist() == [7, 7] +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), + "sample_id": sample_id, + } + + 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") diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 87f4d730be9..37a7f7627e9 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -30,8 +30,12 @@ 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_FORWARD_SUBSTRATE, + QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID, QWEN_IMAGE_PDD_LAYER_SPEC, + adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, + require_qwen_image_pdd_forward_substrate, ) @@ -41,8 +45,8 @@ class _TinyQwenTransformer(nn.Module): def __init__(self, *, packed_channels: int = 4, hidden_width: int = 5) -> None: super().__init__() self.config = SimpleNamespace(guidance_embeds=False) - self.backbone = nn.Linear(packed_channels, hidden_width) - self.proj_out = nn.Linear(hidden_width, packed_channels) + 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( @@ -62,7 +66,8 @@ def forward( dim=1, keepdim=True ).unsqueeze(-1) hidden = torch.tanh(self.backbone(hidden_states)) - hidden = hidden + condition_value + 0.1 * timestep[:, None, None] + hidden = hidden + condition_value.to(hidden.dtype) + hidden = hidden + (0.1 * timestep[:, None, None]).to(hidden.dtype) output = self.proj_out(hidden) self.calls.append( { @@ -72,6 +77,7 @@ def forward( "encoder_hidden_states_mask": encoder_hidden_states_mask.detach().clone(), "img_shapes": img_shapes, "guidance": guidance, + "projection_input": hidden.detach().clone(), "return_dict": return_dict, "kwargs": kwargs, "output": output.detach().clone(), @@ -98,10 +104,12 @@ 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) - mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long)[:batch_size] - negative_embeddings = torch.randn(batch_size, 3, 2) + 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) @@ -113,9 +121,9 @@ def _call_base_packed( ) -> torch.Tensor: embeddings, mask = condition return model( - hidden_states=pack_latents(state), + hidden_states=pack_latents(state).to(torch.bfloat16), timestep=time, - encoder_hidden_states=embeddings, + encoder_hidden_states=embeddings.to(torch.bfloat16), encoder_hidden_states_mask=mask, img_shapes=build_img_shapes(state.shape[0], state.shape[2], state.shape[3]), guidance=None, @@ -156,7 +164,7 @@ def test_unfused_channel_major_output_maps_each_packed_head_in_order() -> None: projection = convert_qwen_image_to_pdd(student, config) with torch.no_grad(): projection.weight.zero_() - head_bias = torch.arange(16, dtype=torch.float32).reshape(4, 4) / 5 + head_bias = (torch.arange(16, dtype=torch.float32).reshape(4, 4) / 5).to(torch.bfloat16) projection.bias.copy_(head_bias.reshape(-1)) state, time, condition, _ = _inputs(batch_size=1) @@ -181,7 +189,7 @@ def test_unfused_channel_major_output_maps_each_packed_head_in_order() -> None: torch.testing.assert_close(actual, expected) -def test_fused_student_matches_explicit_packed_head_weighting() -> None: +def test_fused_student_matches_explicit_packed_weight_fusion() -> None: student = _TinyQwenTransformer() config = _config() projection = convert_qwen_image_to_pdd(student, config) @@ -192,8 +200,6 @@ def test_fused_student_matches_explicit_packed_head_weighting() -> None: adapter = QwenImagePDDAdapter(config) state, time, condition, _ = _inputs() grid = torch.tensor([1.0, 0.85, 0.55, 0.2, 0.0]) - all_heads = adapter.student_all_heads(student, state, time, condition=condition) - student.calls.clear() actual = adapter.student_fused_block( student, @@ -204,13 +210,22 @@ def test_fused_student_matches_explicit_packed_head_weighting() -> None: grid=grid, condition=condition, ) - coefficients = fusion_coefficients(grid, 1, 4) - expected = torch.einsum("n,bnchw->bchw", coefficients, all_heads[:, 1:4]) + 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(state.new_zeros(1, 5)).shape[-1] == 16 + assert student.proj_out(projection.weight.new_zeros(1, 5)).shape[-1] == 16 def test_teacher_cfg_and_global_norm_rescale_match_mr210_reference() -> None: @@ -228,9 +243,11 @@ def test_teacher_cfg_and_global_norm_rescale_match_mr210_reference() -> None: ) assert len(teacher.calls) == 2 - conditional = teacher.calls[0]["output"].float() - unconditional = teacher.calls[1]["output"].float() - guided = conditional + 3.0 * (conditional - unconditional) + conditional_bf16 = teacher.calls[0]["output"] + unconditional_bf16 = teacher.calls[1]["output"] + guided_bf16 = conditional_bf16 + 3.0 * (conditional_bf16 - unconditional_bf16) + conditional = conditional_bf16.float() + guided = guided_bf16.float() factor = torch.linalg.vector_norm( conditional, dim=(1, 2), @@ -238,7 +255,7 @@ def test_teacher_cfg_and_global_norm_rescale_match_mr210_reference() -> None: ) / torch.linalg.vector_norm(guided, dim=(1, 2), keepdim=True).clamp_min(1e-5) expected = unpack_latents((guided * factor).to(teacher.calls[0]["output"].dtype), 4, 4) - assert actual.dtype == torch.float32 + assert actual.dtype == torch.bfloat16 torch.testing.assert_close(actual, expected) torch.testing.assert_close(teacher.calls[0]["encoder_hidden_states"], condition[0]) torch.testing.assert_close(teacher.calls[1]["encoder_hidden_states"], negative_condition[0]) @@ -285,9 +302,7 @@ def forward(self, *, hidden_states, encoder_hidden_states, **kwargs): unrounded_factor = torch.linalg.vector_norm( conditional_fp32, dim=(1, 2), keepdim=True ) / torch.linalg.vector_norm(unrounded_guided, dim=(1, 2), keepdim=True).clamp_min(1e-5) - unrounded = unpack_latents( - (unrounded_guided * unrounded_factor).to(torch.bfloat16), 4, 4 - ) + unrounded = unpack_latents((unrounded_guided * unrounded_factor).to(torch.bfloat16), 4, 4) assert not torch.equal(actual, unrounded) @@ -324,6 +339,222 @@ def test_conversion_preserves_requires_grad_mode_and_rejects_conflicts() -> None convert_qwen_image_to_pdd(transformer, _config(grid_size=2)) +def test_forward_substrate_identity_is_exact() -> None: + assert QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID == ( + "pdd_qwen_mr210_c8100b1347b278511336dccfc074a461457216ec_" + "qwen_33706683487ba16d133b99b73be27b21164c53335441d77b1dcabbfca970f70e" + ) + assert require_qwen_image_pdd_forward_substrate(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE) == dict( + QWEN_IMAGE_PDD_FORWARD_SUBSTRATE + ) + mismatched = dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE) + mismatched["id"] = "canonical-diffusers" + with pytest.raises(ValueError, match="authenticated MR210"): + require_qwen_image_pdd_forward_substrate(mismatched) + with pytest.raises(ValueError, match="authenticated MR210"): + require_qwen_image_pdd_forward_substrate(None) + + +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_adoption_preserves_diffusers_interface_state_keys_and_parameter_identity( + monkeypatch, +) -> None: + qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") + source = _tiny_diffusers_qwen() + source.eval() + source_type = type(source) + source_keys = tuple(source.state_dict()) + source_parameters = tuple(source.parameters()) + source_config = dict(source.config) + monkeypatch.setattr(qwen_pdd, "_require_qwen_source_identity", lambda _model: None) + + adopted = adopt_qwen_image_mr210_forward(source) + + assert adopted is not source + assert isinstance(adopted, source_type) + assert type(source) is source_type + assert tuple(adopted.state_dict()) == source_keys + assert all( + actual is expected + for actual, expected in zip(adopted.parameters(), source_parameters, strict=True) + ) + assert dict(adopted.config) == source_config + assert adopted.device == source.device + assert adopted.dtype == source.dtype + assert adopted.training is False + assert adopt_qwen_image_mr210_forward(adopted) is adopted + + +def test_adoption_rejects_unpinned_qwen_source(monkeypatch) -> None: + qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") + source = _tiny_diffusers_qwen() + monkeypatch.setattr(qwen_pdd, "_sha256_source", lambda _owner: "0" * 64) + + with pytest.raises(RuntimeError, match="transformer source"): + adopt_qwen_image_mr210_forward(source) + + +def test_adoption_rejects_unpinned_timestep_embedding_and_diffusers_version(monkeypatch) -> None: + diffusers = pytest.importorskip("diffusers") + qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") + source = _tiny_diffusers_qwen() + + def mismatched_embedding_hash(owner): + if owner is type(source): + return QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_qwen_source_sha256"] + return "0" * 64 + + monkeypatch.setattr(qwen_pdd, "_sha256_source", mismatched_embedding_hash) + with pytest.raises(RuntimeError, match="timestep embedding source"): + adopt_qwen_image_mr210_forward(source) + + def pinned_source_hash(owner): + if owner is type(source): + return QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_qwen_source_sha256"] + return QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_embeddings_source_sha256"] + + monkeypatch.setattr(qwen_pdd, "_sha256_source", pinned_source_hash) + monkeypatch.setattr(diffusers, "__version__", "0.0.0") + with pytest.raises(RuntimeError, match="Diffusers version"): + adopt_qwen_image_mr210_forward(source) + + +def test_adoption_rejects_every_material_root_invariant(monkeypatch) -> None: + qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") + baseline = _tiny_diffusers_qwen() + monkeypatch.setattr(qwen_pdd, "_require_qwen_source_identity", lambda _model: None) + + def set_config_flag(source, name): + source._internal_dict = type(source._internal_dict)({**dict(source.config), name: True}) + + def check(mutator, match): + source = copy.deepcopy(baseline) + mutator(source) + with pytest.raises((RuntimeError, ValueError), match=match): + adopt_qwen_image_mr210_forward(source) + + check(lambda source: source.add_module("unexpected", nn.Identity()), "root child layout") + check( + lambda source: setattr( + source, + "_modules", + dict(reversed(tuple(source._modules.items()))), + ), + "root child layout", + ) + check( + lambda source: source.register_parameter("root_parameter", nn.Parameter(torch.zeros(1))), + "direct parameters or buffers", + ) + check( + lambda source: source.register_buffer("root_buffer", torch.zeros(1)), + "direct parameters or buffers", + ) + check(lambda source: set_config_flag(source, "guidance_embeds"), "guidance embeddings") + check(lambda source: setattr(source, "peft_config", {"active": True}), "PEFT") + check( + lambda source: setattr(source.transformer_blocks[0], "fused_projections", True), + "fused QKV", + ) + for name in ("zero_cond_t", "use_additional_t_cond", "use_layer3d_rope"): + check(lambda source, name=name: set_config_flag(source, name), rf"{name}=False") + check(lambda source: source.register_forward_hook(lambda *_args: None), "hooks must be empty") + + +def test_adopted_forward_rejects_every_unsupported_input_contract(monkeypatch) -> None: + qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") + source = _tiny_diffusers_qwen().eval().to(dtype=torch.bfloat16) + monkeypatch.setattr(qwen_pdd, "_require_qwen_source_identity", lambda _model: None) + adopted = adopt_qwen_image_mr210_forward(source) + + generator = torch.Generator().manual_seed(20260715) + hidden_states = torch.randn(2, 4, 8, generator=generator).to(torch.bfloat16) + encoder_hidden_states = torch.randn(2, 3, 12, generator=generator).to(torch.bfloat16) + mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long) + encoder_hidden_states[~mask.bool()] = 0 + base = { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "encoder_hidden_states_mask": mask, + "timestep": torch.tensor([0.875, 0.25], dtype=torch.float32), + "img_shapes": [[(1, 2, 2)], [(1, 2, 2)]], + "return_dict": False, + } + + def condition_for(candidate_mask): + candidate_embeddings = encoder_hidden_states.clone() + candidate_embeddings[~candidate_mask.bool()] = 0 + return candidate_embeddings + + nonbinary_mask = torch.tensor([[1, 1, 1], [1, 2, 0]], dtype=torch.long) + nonprefix_mask = torch.tensor([[1, 1, 1], [1, 0, 1]], dtype=torch.long) + empty_mask = torch.tensor([[1, 1, 1], [0, 0, 0]], dtype=torch.long) + no_full_row_mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long) + nonzero_padding = encoder_hidden_states.clone() + nonzero_padding[1, 1] = 1 + cases = ( + ({"hidden_states": hidden_states.float()}, TypeError, "hidden_states"), + ({"encoder_hidden_states": encoder_hidden_states.float()}, TypeError, "encoder_hidden"), + ({"timestep": base["timestep"].to(torch.bfloat16)}, TypeError, "timestep"), + ( + { + "encoder_hidden_states": condition_for(nonbinary_mask), + "encoder_hidden_states_mask": nonbinary_mask, + }, + ValueError, + "binary prefix masks", + ), + ( + { + "encoder_hidden_states": condition_for(nonprefix_mask), + "encoder_hidden_states_mask": nonprefix_mask, + }, + ValueError, + "binary prefix masks", + ), + ( + { + "encoder_hidden_states": condition_for(empty_mask), + "encoder_hidden_states_mask": empty_mask, + }, + ValueError, + "binary prefix masks", + ), + ( + { + "encoder_hidden_states": condition_for(no_full_row_mask), + "encoder_hidden_states_mask": no_full_row_mask, + }, + ValueError, + "binary prefix masks", + ), + ({"encoder_hidden_states": nonzero_padding}, ValueError, "zero padding"), + ({"max_txt_seq_len": 2}, ValueError, "max_txt_seq_len"), + ({"txt_seq_lens": [3, 1]}, ValueError, "txt_seq_lens"), + ({"guidance": torch.ones(2)}, ValueError, "guidance embeddings"), + ({"attention_kwargs": {"scale": 1.0}}, ValueError, "attention_kwargs"), + ({"controlnet_block_samples": ()}, ValueError, "ControlNet"), + ({"additional_t_cond": torch.ones(2)}, ValueError, "additional time"), + ) + for override, error_type, match in cases: + with pytest.raises(error_type, match=match): + adopted(**(base | override)) + + def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> None: with pytest.raises(ValueError, match="num_train_timesteps=None"): QwenImagePDDAdapter(_config().model_copy(update={"num_train_timesteps": 1000})) @@ -360,6 +591,13 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N convert_qwen_image_to_pdd(transformer, config) with pytest.raises(TypeError, match="tuple"): adapter.student_all_heads(transformer, state, time, condition=condition[0]) + with pytest.raises(ValueError, match="requires batched embeddings"): + adapter.student_all_heads( + transformer, + state, + time, + condition=(condition[0][..., 0], condition[1]), + ) with pytest.raises(ValueError, match="controlled keys"): adapter.student_all_heads( transformer, @@ -377,11 +615,12 @@ def test_raw_head_reference_uses_independent_linear_outputs() -> None: projection = convert_qwen_image_to_pdd(student, config) state, time, condition, _ = _inputs(batch_size=1) embeddings, mask = condition - packed = pack_latents(state) + packed = pack_latents(state).to(torch.bfloat16) hidden = torch.tanh(student.backbone(packed)) - hidden = hidden + embeddings.mean(dim=(1, 2), keepdim=True) - hidden = hidden + 0.01 * mask.sum(dim=1, keepdim=True).unsqueeze(-1) - hidden = hidden + 0.1 * time[:, None, None] + condition_value = embeddings.mean(dim=(1, 2), keepdim=True) + condition_value = condition_value + 0.01 * mask.sum(dim=1, keepdim=True).unsqueeze(-1) + hidden = hidden + condition_value.to(hidden.dtype) + hidden = hidden + (0.1 * time[:, None, None]).to(hidden.dtype) head_weights = projection.weight.reshape(4, 4, 5) head_bias = projection.bias.reshape(4, 4) expected_packed = torch.stack( From 0ac415e268f969af33b31e2512e198a07be42427 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 16 Jul 2026 00:03:37 -0700 Subject: [PATCH 26/45] Use canonical Qwen execution for PDD Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/README.md | 33 +- .../fastgen/pdd/automodel_dependency.json | 14 - examples/diffusers/fastgen/pdd/checkpoint.py | 53 +-- .../fastgen/pdd/configs/qwen_image.yaml | 4 - examples/diffusers/fastgen/pdd/export.py | 55 +-- .../fastgen/pdd/export_qwen_image.py | 76 +--- .../fastgen/pdd/inference_qwen_image.py | 33 +- examples/diffusers/fastgen/pdd/recipe.py | 115 ++---- examples/diffusers/fastgen/pdd/training.py | 2 +- .../fastgen/pdd/verify_readonly_automodel.py | 209 ---------- examples/diffusers/fastgen/requirements.txt | 2 +- modelopt/torch/fastgen/methods/pdd.py | 5 +- .../torch/fastgen/plugins/qwen_image_pdd.py | 340 +--------------- .../pdd_checkpoint_failure_distributed.py | 4 +- .../fastgen/pdd_export_distributed.py | 18 +- .../pdd_validation_oracle_distributed.py | 2 +- .../examples/diffusers/fastgen/test_layout.py | 2 - .../fastgen/test_pdd_inference_checkpoint.py | 75 ++-- .../fastgen/test_pdd_recipe_setup.py | 149 ++++--- .../fastgen/test_pdd_training_lifecycle.py | 29 +- .../fastgen/test_pdd_validation_oracle.py | 2 +- .../fastgen/test_qwen_image_pdd_plugin.py | 384 ++++++++---------- 22 files changed, 385 insertions(+), 1221 deletions(-) delete mode 100644 examples/diffusers/fastgen/pdd/automodel_dependency.json delete mode 100644 examples/diffusers/fastgen/pdd/verify_readonly_automodel.py diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index b056fa49a9c..2c7d6d33727 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -3,8 +3,8 @@ Parallel Decoding Distillation (PDD) trains one Qwen-Image student call to predict several consecutive rectified-flow updates. The student keeps the original transformer backbone and widens only its output projection to 128 velocity heads. During training it samples aligned block starts -and target spans from 1 through 64 intervals, so the same checkpoint can use different supported block schedules at -inference. +and target spans from 1 through 64 intervals, so the same checkpoint can use different supported +block schedules at inference. The provided schedules use the 128-interval grid as follows: @@ -19,14 +19,12 @@ The provided schedules use the 128-interval grid as follows: Install the shared requirements from the repository root, then launch with released AutoModel APIs. No AutoModel, Diffusers, or Qwen source changes are required. -The reproducibility arm executes the Qwen forward semantics from FastGen MR210 commit -`c8100b1347b278511336dccfc074a461457216ec`: BF16 image/text compute, FP32 normalized time through -the timestep projection, and the MR attention behavior for zero-padded text. ModelOpt adopts the -loaded Diffusers model into a compatible root without editing or monkeypatching external classes. -The exact FastGen, Diffusers 0.38.0 transformer, and timestep-embedding source identities are -authenticated in every checkpoint and export and are required again by resume and inference. -Guidance-embedded/`zero_cond_t` models, additional time conditioning, ControlNet, PEFT, and QKV -fusion are deliberately rejected because they are outside that executed algorithm. +The example uses the ordinary Diffusers Qwen transformer without replacing or monkeypatching its +forward. Diffusers owns Qwen timestep conversion and converts each text mask into the joint +text/image attention mask, so padded text tokens do not participate as attention keys. ModelOpt +owns only PDD projection conversion, latent packing, condition validation, and packed per-token +classifier-free guidance using the Qwen pipeline formula. Guidance-embedded models and PEFT remain +outside this first example. ```bash pip install -r examples/diffusers/fastgen/requirements.txt @@ -57,15 +55,22 @@ default per-rank batch gives global batch size 256; other GPU topologies must se `256 / world_size` because this recipe does not use gradient accumulation. Checkpoints include the student, optimizer, scheduler, RNG, trainer, and exact replayable sampler state needed to resume the next committed batch. FP32 master parameters and Adam state are sharded -while forward/backward compute remains BF16 and gradient reduction remains FP32. The FSDP policy -does not recursively cast forward inputs; ModelOpt casts image/text tensors to BF16 itself so the -FP32 timestep cannot be rounded before the source-locked forward receives it. +while forward/backward uses the configured model dtype and gradient reduction remains FP32. The +adapter casts packed image/text inputs to that compute dtype; the ordinary Diffusers Qwen forward +owns timestep conversion. Start with a one-node smoke and scale only after it passes; project training runs are capped at 16 -nodes. +nodes. Checkpointed training, export, and inference require the remote model ID and exact lowercase +40-character Hugging Face commit in the provided config; local model directories are limited to +low-level hermetic setup tests because the frozen teacher is rebuilt rather than checkpointed. ## Export and inference +The v1 export restores the sharded checkpoint with the same total process count that created it; +cross-world-size DCP resharding is not yet part of this example. The command below therefore +applies to a checkpoint trained with eight ranks. For a 64-rank checkpoint, use the cluster +launcher with 64 export ranks before running single-process inference. + ```bash torchrun --standalone --nproc-per-node=8 \ examples/diffusers/fastgen/pdd/export_qwen_image.py \ diff --git a/examples/diffusers/fastgen/pdd/automodel_dependency.json b/examples/diffusers/fastgen/pdd/automodel_dependency.json deleted file mode 100644 index 4cb0535dbf9..00000000000 --- a/examples/diffusers/fastgen/pdd/automodel_dependency.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "distribution": "nemo_automodel", - "import_name": "nemo_automodel", - "package_file_count": 490, - "package_tree_sha256": "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528", - "release_commit": "d02f49cb314554715aabb97e8dba6599c9f6e9e0", - "release_tag": "v0.5.0", - "runtime_versions": { - "diffusers": "0.38.0" - }, - "version": "0.5.0", - "wheel": "nemo_automodel-0.5.0-py3-none-any.whl", - "wheel_sha256": "881aebafc5145752842afbbfe0a42e1c33d06847c3e418ad3d6f154ddc8e0f45" -} diff --git a/examples/diffusers/fastgen/pdd/checkpoint.py b/examples/diffusers/fastgen/pdd/checkpoint.py index 0cccaee09c0..eb04fbdf235 100644 --- a/examples/diffusers/fastgen/pdd/checkpoint.py +++ b/examples/diffusers/fastgen/pdd/checkpoint.py @@ -31,14 +31,12 @@ import torch.distributed as dist from modelopt.torch.fastgen import PDDMetadata -from modelopt.torch.fastgen.plugins.qwen_image_pdd import require_qwen_image_pdd_forward_substrate if TYPE_CHECKING: from collections.abc import Sequence -_CHECKPOINT_SCHEMA_VERSION = 3 +_CHECKPOINT_SCHEMA_VERSION = 4 _COMPLETE_SCHEMA_VERSION = 1 -_FORBIDDEN_ARTIFACT_TOKENS = ("fake_score", "discriminator", "ema", "r1", "gan") def _sha256(path: Path) -> str: @@ -160,13 +158,9 @@ def _read_json(path: Path) -> dict[str, Any]: def build_pdd_checkpoint_identity( *, metadata: PDDMetadata, - forward_substrate: Mapping[str, Any], model_id: str, model_revision: str | None, guidance_scale: float | None, - guidance_rescale: float, - guidance_eps: float, - automodel_snapshot: Mapping[str, Any], ordered_train_id_sha256: str, ordered_heldout_id_sha256: str, dataset_snapshot_sha256: str, @@ -187,11 +181,12 @@ def build_pdd_checkpoint_identity( raise TypeError("metadata must be PDDMetadata.") if not isinstance(model_id, str) or not model_id: raise ValueError("model_id must be a non-empty string.") - if model_revision is not None and not isinstance(model_revision, str): - raise TypeError("model_revision must be a string or null.") - for name, value in (("guidance_rescale", guidance_rescale), ("guidance_eps", guidance_eps)): - if isinstance(value, bool) or not isinstance(value, int | float): - raise TypeError(f"{name} must be a real number.") + if ( + not isinstance(model_revision, str) + or len(model_revision) != 40 + or any(character not in "0123456789abcdef" for character in model_revision) + ): + raise ValueError("model_revision must be an exact lowercase 40-character commit.") if guidance_scale is not None and ( isinstance(guidance_scale, bool) or not isinstance(guidance_scale, int | float) ): @@ -218,39 +213,11 @@ def build_pdd_checkpoint_identity( raise ValueError("dtype must be a non-empty string.") if type(optimizer).__module__ != "torch.optim.adamw" or type(optimizer).__name__ != "AdamW": raise TypeError("PDD checkpoint identity requires the stock torch.optim.AdamW optimizer.") - required_snapshot = { - "distribution", - "package_tree_sha256", - "runtime_versions", - "version", - "wheel_sha256", - } - missing = sorted(required_snapshot.difference(automodel_snapshot)) - if missing: - raise ValueError(f"AutoModel snapshot is missing identity keys: {missing}.") return { "schema_version": _CHECKPOINT_SCHEMA_VERSION, - "forward_substrate": require_qwen_image_pdd_forward_substrate(forward_substrate), "model": {"id": model_id, "revision": model_revision, "dtype": dtype}, "pdd_metadata": metadata.to_dict(), - "guidance": { - "scale": None if guidance_scale is None else float(guidance_scale), - "rescale": float(guidance_rescale), - "eps": float(guidance_eps), - }, - "automodel": { - "distribution": automodel_snapshot["distribution"], - "version": automodel_snapshot["version"], - "package_tree_sha256": _require_sha256( - automodel_snapshot["package_tree_sha256"], - name="automodel.package_tree_sha256", - ), - "wheel_sha256": _require_sha256( - automodel_snapshot["wheel_sha256"], - name="automodel.wheel_sha256", - ), - "runtime_versions": dict(automodel_snapshot["runtime_versions"]), - }, + "guidance": {"scale": None if guidance_scale is None else float(guidance_scale)}, "data": { "ordered_train_id_sha256": _require_sha256( ordered_train_id_sha256, @@ -454,10 +421,6 @@ def validate_pdd_training_checkpoint( raise RuntimeError(f"PDD checkpoint sidecar is missing: {relative}.") if _sha256(path) != manifest["sidecar_sha256"][relative]: raise RuntimeError(f"PDD checkpoint sidecar hash mismatch: {relative}.") - for candidate in checkpoint.rglob("*"): - lowered = candidate.name.lower() - if any(token in lowered for token in _FORBIDDEN_ARTIFACT_TOKENS): - raise RuntimeError(f"PDD checkpoint contains a forbidden DMD artifact: {candidate}.") return manifest diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 756464081f7..27ddf2fb2d1 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -48,10 +48,6 @@ lr_scheduler: lr_warmup_steps: 0 min_lr: 2.0e-5 -guidance: - rescale: 1.0 - eps: 1.0e-5 - step_scheduler: max_steps: 10000 num_epochs: 200 diff --git a/examples/diffusers/fastgen/pdd/export.py b/examples/diffusers/fastgen/pdd/export.py index 61239b96c30..22bb6522b9f 100644 --- a/examples/diffusers/fastgen/pdd/export.py +++ b/examples/diffusers/fastgen/pdd/export.py @@ -31,10 +31,7 @@ from safetensors.torch import save_file from modelopt.torch.fastgen import PDDConfig, PDDMetadata -from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - QWEN_IMAGE_PDD_LAYER_SPEC, - require_qwen_image_pdd_forward_substrate, -) +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC from .artifacts import ( load_canonical_json, @@ -44,7 +41,7 @@ write_canonical_json, ) -_EXPORT_SCHEMA_VERSION = 2 +_EXPORT_SCHEMA_VERSION = 3 _COMPLETE_SCHEMA_VERSION = 1 _EXPORT_FORMAT = "modelopt-pdd-safetensors" _CONFIG_FILE = "config.json" @@ -171,8 +168,6 @@ def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> di if not isinstance(identity, Mapping): raise TypeError("PDD export identity must be a mapping.") required = { - "automodel", - "forward_substrate", "guidance", "model", "pdd_metadata", @@ -183,36 +178,19 @@ def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> di raise ValueError(f"PDD export identity is missing keys: {missing}.") if identity["pdd_metadata"] != metadata.to_dict(): raise ValueError("PDD export metadata does not match the checkpoint identity.") - require_qwen_image_pdd_forward_substrate(identity["forward_substrate"]) model = _require_exact_mapping( identity["model"], {"id", "revision", "dtype"}, name="identity.model" ) if not isinstance(model["id"], str) or not model["id"] or not isinstance(model["dtype"], str): raise ValueError("PDD export checkpoint identity has an invalid model ID or dtype.") revision = model["revision"] - if not isinstance(revision, str) or len(revision) != 40: - raise ValueError("PDD export requires a pinned 40-character model revision.") - try: - int(revision, 16) - except ValueError as error: - raise ValueError("PDD export model revision must be hexadecimal.") from error - automodel = _require_exact_mapping( - identity["automodel"], - {"distribution", "version", "package_tree_sha256", "wheel_sha256", "runtime_versions"}, - name="identity.automodel", - ) if ( - not isinstance(automodel["distribution"], str) - or not automodel["distribution"] - or not isinstance(automodel["version"], str) - or not isinstance(automodel["runtime_versions"], Mapping) + not isinstance(revision, str) + or len(revision) != 40 + or any(character not in "0123456789abcdef" for character in revision) ): - raise ValueError("PDD export AutoModel identity is malformed.") - require_sha256(automodel["package_tree_sha256"], name="AutoModel package tree SHA-256") - require_sha256(automodel["wheel_sha256"], name="AutoModel wheel SHA-256") - guidance = _require_exact_mapping( - identity["guidance"], {"scale", "rescale", "eps"}, name="identity.guidance" - ) + raise ValueError("PDD export model revision must be an exact lowercase commit.") + guidance = _require_exact_mapping(identity["guidance"], {"scale"}, name="identity.guidance") for name, value in guidance.items(): if value is not None and ( isinstance(value, bool) @@ -231,20 +209,6 @@ def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> di return dict(identity) -def _validate_modelopt_source(source: Mapping[str, Any]) -> dict[str, Any]: - source = _require_exact_mapping(source, {"commit", "dirty"}, name="modelopt_source") - commit = source["commit"] - if not isinstance(commit, str) or len(commit) != 40: - raise ValueError("modelopt_source.commit must be a 40-character Git commit.") - try: - int(commit, 16) - except ValueError as error: - raise ValueError("modelopt_source.commit must be hexadecimal.") from error - if source["dirty"] is not False: - raise ValueError("modelopt_source.dirty must be false.") - return dict(source) - - def write_pdd_export( output_dir: str | Path, state_dict: Mapping[str, Any], @@ -253,7 +217,6 @@ def write_pdd_export( transformer_config: Mapping[str, Any], identity: Mapping[str, Any], source_checkpoint: Mapping[str, Any], - modelopt_source: Mapping[str, Any], max_shard_bytes: int, ) -> Path: """Publish a complete PDD export into a previously absent directory.""" @@ -278,7 +241,6 @@ def write_pdd_export( or source_checkpoint["completed_steps"] < 1 ): raise ValueError("source_checkpoint.completed_steps must be an integer >= 1.") - resolved_modelopt_source = _validate_modelopt_source(modelopt_source) resolved_identity = _validate_identity(identity, metadata) tensors = _validate_state_dict(state_dict) @@ -335,7 +297,6 @@ def write_pdd_export( "format": _EXPORT_FORMAT, "identity": resolved_identity, "source_checkpoint": dict(source_checkpoint), - "modelopt_source": resolved_modelopt_source, "max_shard_bytes": max_shard_bytes, "total_tensor_bytes": total_tensor_bytes, "tensors": tensor_specs, @@ -401,7 +362,6 @@ def inspect_pdd_export(export_dir: str | Path) -> PDDExportDescriptor: "format", "identity", "source_checkpoint", - "modelopt_source", "max_shard_bytes", "total_tensor_bytes", "tensors", @@ -411,7 +371,6 @@ def inspect_pdd_export(export_dir: str | Path) -> PDDExportDescriptor: ) if manifest["schema_version"] != _EXPORT_SCHEMA_VERSION or manifest["format"] != _EXPORT_FORMAT: raise ValueError("PDD export manifest schema or format is unsupported.") - _validate_modelopt_source(manifest["modelopt_source"]) if type(manifest["max_shard_bytes"]) is not int or manifest["max_shard_bytes"] <= 0: raise ValueError("PDD export max_shard_bytes is invalid.") if type(manifest["total_tensor_bytes"]) is not int or manifest["total_tensor_bytes"] <= 0: diff --git a/examples/diffusers/fastgen/pdd/export_qwen_image.py b/examples/diffusers/fastgen/pdd/export_qwen_image.py index 67976b87b7c..43ec954be34 100644 --- a/examples/diffusers/fastgen/pdd/export_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/export_qwen_image.py @@ -20,7 +20,6 @@ import argparse import json import math -import subprocess import sys from collections.abc import Mapping from pathlib import Path @@ -142,35 +141,13 @@ def collective_export_memory_preflight( return full_state_bytes, largest_tensor_bytes -def _git_source_identity() -> dict[str, Any]: - commit = subprocess.run( - ["git", "rev-parse", "HEAD"], - cwd=_REPO_ROOT, - check=True, - capture_output=True, - text=True, - ).stdout.strip() - dirty = bool( - subprocess.run( - ["git", "status", "--porcelain", "--untracked-files=normal"], - cwd=_REPO_ROOT, - check=True, - capture_output=True, - text=True, - ).stdout - ) - if dirty: - raise RuntimeError("PDD export requires a clean ModelOpt source checkout.") - return {"commit": commit, "dirty": False} - - -def _collective_publication_preflight(output_dir: Path) -> Mapping[str, Any]: +def _collective_publication_preflight(output_dir: Path) -> None: status = None if dist.get_rank() == 0: try: if output_dir.is_symlink() or output_dir.resolve().exists(): raise FileExistsError(f"PDD export output already exists: {output_dir}.") - status = {"ok": True, "modelopt_source": _git_source_identity()} + status = {"ok": True} except BaseException as error: status = {"ok": False, "error": f"{type(error).__name__}: {error}"} payload = [status] @@ -180,17 +157,10 @@ def _collective_publication_preflight(output_dir: Path) -> Mapping[str, Any]: raise RuntimeError("rank 0 broadcast malformed PDD publication preflight status.") if not status["ok"]: raise RuntimeError(f"PDD publication preflight failed: {status.get('error')}.") - modelopt_source = status.get("modelopt_source") - if not isinstance(modelopt_source, Mapping): - raise RuntimeError("rank 0 broadcast malformed ModelOpt source identity.") - return modelopt_source def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: from modelopt.torch.fastgen import PDDMetadata - from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - require_qwen_image_pdd_forward_substrate, - ) identity = manifest.get("identity") if not isinstance(identity, Mapping): @@ -200,25 +170,12 @@ def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, raise RuntimeError("PDD checkpoint has no PDD metadata mapping.") if PDDMetadata.from_dict(pdd_metadata) != setup.metadata: raise RuntimeError("PDD checkpoint metadata does not match the configured student.") - require_qwen_image_pdd_forward_substrate(identity.get("forward_substrate")) if identity.get("model") != { "id": config.model_id, "revision": config.model_revision, "dtype": str(config.dtype).removeprefix("torch."), }: raise RuntimeError("PDD checkpoint model identity does not match the export config.") - checkpoint_automodel = identity.get("automodel") - if not isinstance(checkpoint_automodel, Mapping): - raise RuntimeError("PDD checkpoint has no AutoModel identity.") - for key in ( - "distribution", - "version", - "package_tree_sha256", - "wheel_sha256", - "runtime_versions", - ): - if checkpoint_automodel.get(key) != setup.automodel_snapshot.get(key): - raise RuntimeError(f"PDD checkpoint AutoModel identity mismatch for {key}.") topology = identity.get("topology") if not isinstance(topology, Mapping) or topology.get("world_size") != dist.get_world_size(): raise RuntimeError("PDD checkpoint topology does not match the export process group.") @@ -238,31 +195,14 @@ def _collective_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[s def _checkpoint_selector_identity(config: Any, setup: Any) -> dict[str, Any]: - from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_FORWARD_SUBSTRATE - return { - "forward_substrate": dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE), "model": { "id": config.model_id, "revision": config.model_revision, "dtype": str(config.dtype).removeprefix("torch."), }, "pdd_metadata": setup.metadata.to_dict(), - "guidance": { - "scale": config.pdd.guidance_scale, - "rescale": config.guidance.rescale, - "eps": config.guidance.eps, - }, - "automodel": { - key: setup.automodel_snapshot[key] - for key in ( - "distribution", - "version", - "package_tree_sha256", - "wheel_sha256", - "runtime_versions", - ) - }, + "guidance": {"scale": config.pdd.guidance_scale}, "topology": {"world_size": dist.get_world_size(), "pure_data_parallel": True}, } @@ -301,6 +241,7 @@ def main() -> None: from pdd.artifacts import sha256_file from pdd.export import write_pdd_export from pdd.recipe import ( + _require_immutable_model_source, build_pdd_export_setup, initialize_pdd_distributed, resolve_pdd_recipe_config, @@ -308,11 +249,7 @@ def main() -> None: raw = yaml.safe_load(args.config.read_text()) config = resolve_pdd_recipe_config(raw) - if Path(config.model_id).is_dir() or config.model_revision is None: - raise ValueError( - "PDD inference export requires a remote model ID and pinned 40-character revision; " - "mutable local model directories are training-only inputs." - ) + _require_immutable_model_source(config, context="PDD export") if not math.isfinite(args.max_shard_size_gib) or args.max_shard_size_gib <= 0: raise ValueError("max_shard_size_gib must be finite and > 0.") if not math.isfinite(args.memory_headroom) or args.memory_headroom < 1.0: @@ -322,7 +259,7 @@ def main() -> None: backend="nccl" if config.device.type == "cuda" else "gloo", timeout_minutes=60, ) - modelopt_source = _collective_publication_preflight(args.output_dir) + _collective_publication_preflight(args.output_dir) restore_from = args.checkpoint or config.checkpoint.restore_from if not restore_from: raise ValueError("PDD export requires --checkpoint or checkpoint.restore_from.") @@ -384,7 +321,6 @@ def main() -> None: "manifest_sha256": sha256_file(checkpoint / "manifest.json"), "completed_steps": checkpoint_manifest["completed_steps"], }, - modelopt_source=modelopt_source, max_shard_bytes=max_shard_bytes, ) publication = {"ok": True, "output": str(output)} diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index 3de0df28a5c..891b76a60e4 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Run conditional-only PDD inference from an authenticated Qwen-Image export.""" +"""Run conditional-only PDD inference from a complete Qwen-Image PDD export.""" from __future__ import annotations @@ -81,12 +81,12 @@ def _model_identity(descriptor: Any) -> Mapping[str, Any]: if not isinstance(model["id"], str) or not model["id"]: raise RuntimeError("PDD export model ID is invalid.") revision = model["revision"] - if not isinstance(revision, str) or len(revision) != 40: - raise RuntimeError("PDD export requires a pinned 40-character model revision.") - try: - int(revision, 16) - except ValueError as error: - raise RuntimeError("PDD export model revision must be hexadecimal.") from error + if ( + not isinstance(revision, str) + or len(revision) != 40 + or any(character not in "0123456789abcdef" for character in revision) + ): + raise RuntimeError("PDD export model revision must be an exact lowercase commit.") return model @@ -105,7 +105,7 @@ def _validate_qwen_projection(student: nn.Module, metadata: Any) -> nn.Linear: or base_projection.out_features != metadata.projection_out_features or (base_projection.bias is not None) != metadata.projection_bias ): - raise RuntimeError("reconstructed Qwen proj_out does not match authenticated metadata.") + raise RuntimeError("reconstructed Qwen proj_out does not match the export metadata.") if base_projection.out_features != in_channels: raise RuntimeError( "Qwen proj_out width must equal transformer in_channels for 2x2 latent packing." @@ -117,23 +117,13 @@ def build_pdd_student(export_dir: str | Path) -> tuple[nn.Module, Any, torch.dty """Reconstruct and strictly load the converted Qwen student on CPU.""" from diffusers import QwenImageTransformer2DModel - from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - adopt_qwen_image_mr210_forward, - convert_qwen_image_to_pdd, - require_qwen_image_pdd_forward_substrate, - ) + from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd from pdd.export import inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata descriptor = inspect_pdd_export(export_dir) model_identity = _model_identity(descriptor) - require_qwen_image_pdd_forward_substrate( - descriptor.manifest["identity"].get("forward_substrate") - ) dtype = _dtype_from_name(model_identity["dtype"]) - loaded_transformer = QwenImageTransformer2DModel.from_config( - dict(descriptor.transformer_config) - ) - student = adopt_qwen_image_mr210_forward(loaded_transformer) + student = QwenImageTransformer2DModel.from_config(dict(descriptor.transformer_config)) metadata = descriptor.metadata _validate_qwen_projection(student, metadata) config = pdd_config_from_metadata(metadata, blocks=metadata.inference_blocks) @@ -318,7 +308,7 @@ def count_invocation( result_json.parent.mkdir(parents=True, exist_ok=True) result = { - "schema_version": 1, + "schema_version": 2, "record_type": "pdd_inference", "condition": args.schedule.replace("-", "_"), "prompt_id": args.prompt_id, @@ -328,7 +318,6 @@ def count_invocation( "blocks": list(config.inference_blocks), "height": args.height, "width": args.width, - "forward_substrate_id": descriptor.manifest["identity"]["forward_substrate"]["id"], "export_manifest_sha256": sha256_file(descriptor.root / "manifest.json"), "output": {"path": output_reference, "sha256": sha256_file(output)}, "scheduler_steps": expected_invocations, diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 4591e865c52..7266618e5e9 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -20,6 +20,7 @@ import copy import logging import math +import re import time from collections.abc import Mapping from dataclasses import dataclass @@ -32,9 +33,7 @@ from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, QwenImagePDDAdapter, - adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, ) @@ -49,7 +48,8 @@ _iter_validation_batches, _validate_dataset_contract, ) -from .verify_readonly_automodel import snapshot_installed_distribution + +_HF_COMMIT_PATTERN = re.compile(r"[0-9a-f]{40}") @dataclass(frozen=True) @@ -102,14 +102,6 @@ class PDDValidationConfig: every_steps: int = 1_000 -@dataclass(frozen=True) -class PDDGuidanceConfig: - """Resolved Qwen packed-CFG norm-rescaling settings.""" - - rescale: float = 1.0 - eps: float = 1e-5 - - @dataclass(frozen=True) class PDDRecipeConfig: """Resolved setup inputs; incompatible mutation modes have already been rejected.""" @@ -122,7 +114,6 @@ class PDDRecipeConfig: step_scheduler: PDDStepSchedulerConfig training_health: PDDTrainingHealthConfig validation: PDDValidationConfig - guidance: PDDGuidanceConfig seed: int learning_rate: float weight_decay: float @@ -148,7 +139,6 @@ class PDDSetupArtifacts: metadata: PDDMetadata checkpoint_keys: tuple[str, ...] lifecycle: tuple[str, ...] - automodel_snapshot: Mapping[str, Any] @dataclass(frozen=True) @@ -175,7 +165,6 @@ class PDDExportSetupArtifacts: checkpoint_keys: tuple[str, ...] transformer_config: Mapping[str, Any] lifecycle: tuple[str, ...] - automodel_snapshot: Mapping[str, Any] def _as_mapping(value: Any, *, name: str) -> Mapping[str, Any]: @@ -246,6 +235,18 @@ def _resolve_dtype(value: Any) -> torch.dtype: ) from error +def _is_exact_hf_commit(value: Any) -> bool: + return isinstance(value, str) and _HF_COMMIT_PATTERN.fullmatch(value) is not None + + +def _require_immutable_model_source(config: PDDRecipeConfig, *, context: str) -> None: + if Path(config.model_id).is_dir() or not _is_exact_hf_commit(config.model_revision): + raise ValueError( + f"{context} requires a Hugging Face model ID and exact lowercase 40-character " + "commit revision." + ) + + def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: """Resolve one canonical DMD2-shaped PDD configuration.""" raw = _config_to_mapping(raw) @@ -279,7 +280,6 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: training_health = _as_mapping(raw.get("training_health", {}), name="training_health") validation = _as_mapping(raw.get("validation", {}), name="validation") checkpoint = _as_mapping(raw.get("checkpoint", {}), name="checkpoint") - guidance = _as_mapping(raw.get("guidance", {}), name="guidance") data = _as_mapping(raw.get("data", {}), name="data") dataloader = _as_mapping(data.get("dataloader", {}), name="data.dataloader") @@ -328,6 +328,7 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: _reject_enabled(model.get("peft_cfg"), name="PEFT/LoRA") _reject_enabled(raw.get("peft"), name="PEFT/LoRA") _reject_enabled(raw.get("peft_cfg"), name="PEFT/LoRA") + _reject_enabled(raw.get("guidance"), name="guidance overrides") _reject_enabled(model.get("guidance_embeds"), name="Qwen guidance embeddings") _reject_enabled(model.get("guidance_embeddings"), name="Qwen guidance embeddings") for option in ( @@ -344,14 +345,13 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: if not isinstance(model_id, str) or not model_id: raise ValueError("model.pretrained_model_name_or_path must be a non-empty string.") model_revision = model.get("revision") - if model_revision is not None and ( - not isinstance(model_revision, str) - or len(model_revision) != 40 - or any(character not in "0123456789abcdefABCDEF" for character in model_revision) - ): - raise ValueError("model.revision must be null or a full 40-character commit hash.") - if not Path(model_id).is_dir() and model_revision is None: - raise ValueError("Remote PDD models require an exact model.revision commit hash.") + if Path(model_id).is_dir(): + if model_revision is not None: + raise ValueError("Local PDD model directories require model.revision=null.") + elif not _is_exact_hf_commit(model_revision): + raise ValueError( + "Remote PDD models require an exact lowercase 40-character model.revision commit." + ) learning_rate = _require_finite_real( optim.get("learning_rate", 2.0e-5), @@ -471,8 +471,6 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: model.get("fuse_qkv_projections", False), name="model.fuse_qkv_projections", ) - if fuse_qkv_projections: - raise ValueError("authenticated Qwen MR210 PDD does not support QKV fusion.") seed = _require_int_at_least(raw.get("seed", 42), name="seed", minimum=0) max_steps = _require_int_at_least( @@ -556,23 +554,7 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: minimum=1, ) - guidance_rescale = _require_finite_real( - guidance.get("rescale", 1.0), - name="guidance.rescale", - minimum=0.0, - ) - if guidance_rescale > 1.0: - raise ValueError("guidance.rescale must be <= 1.") - guidance_eps = _require_finite_real( - guidance.get("eps", 1e-5), - name="guidance.eps", - minimum=0.0, - ) - if guidance_eps == 0.0: - raise ValueError("guidance.eps must be > 0.") dtype = _resolve_dtype(model.get("torch_dtype", "bfloat16")) - if dtype != torch.bfloat16: - raise ValueError("authenticated Qwen MR210 PDD requires model.torch_dtype='bfloat16'.") return PDDRecipeConfig( model_id=model_id, @@ -608,7 +590,6 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: split_seed=split_seed, every_steps=validation_every_steps, ), - guidance=PDDGuidanceConfig(rescale=guidance_rescale, eps=guidance_eps), seed=seed, learning_rate=float(learning_rate), weight_decay=float(weight_decay), @@ -653,17 +634,17 @@ def _require_projection_module( def _resolve_model_source(config: PDDRecipeConfig) -> str: if Path(config.model_id).is_dir(): - return config.model_id + return str(Path(config.model_id).resolve()) from huggingface_hub import snapshot_download - if config.model_revision is None: - raise ValueError("Remote PDD models require a pinned model revision.") - model_source = snapshot_download(config.model_id, revision=config.model_revision) - if Path(model_source).resolve().name != config.model_revision: + model_source = Path( + snapshot_download(config.model_id, revision=config.model_revision) + ).resolve() + if model_source.parent.name != "snapshots" or model_source.name != config.model_revision: raise RuntimeError( - "Hugging Face resolved a model snapshot that does not match the pinned revision." + "Hugging Face resolved a model snapshot that does not match model.revision." ) - return model_source + return str(model_source) def _load_unwrapped_transformer( @@ -711,9 +692,8 @@ def _stage_and_shard_training_models( ): raise AttributeError("QKV fusion requires both Qwen transformers to expose the object API.") - # Match the FastGen Qwen PDD recipe: FP32 parameter/optimizer storage, - # with the FSDP policy below casting gathered parameters to the configured - # model dtype for forward/backward compute. + # Keep parameter and optimizer storage in FP32. The FSDP policy below casts + # gathered parameters to the configured dtype for forward/backward compute. student.to(device=device, dtype=torch.float32) if fuse_qkv_projections: student.fuse_qkv_projections() @@ -767,7 +747,7 @@ def _materialize_zero_step_adamw_state(optimizer: torch.optim.AdamW) -> None: def _require_fp32_optimizer_storage(optimizer: torch.optim.AdamW) -> None: - """Require the FP32 master-parameter and Adam-state contract used by MR210.""" + """Require FP32 master parameters and Adam state for stable small updates.""" for group in optimizer.param_groups: for parameter in group["params"]: if parameter.dtype != torch.float32: @@ -790,10 +770,8 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: if not dist.is_available() or not dist.is_initialized(): raise RuntimeError("Initialize torch.distributed before building the PDD FSDP2 setup.") - automodel_snapshot = snapshot_installed_distribution() lifecycle: list[str] = [] - # Imports are intentionally delayed until the exact installed wheel has passed verification. from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline from nemo_automodel.components.checkpoint.config import CheckpointingConfig from nemo_automodel.components.distributed import ( @@ -803,9 +781,7 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: ) from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager - pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) - student = adopt_qwen_image_mr210_forward(loaded_transformer) - pipe.transformer = student + pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) teacher = copy.deepcopy(student).eval().requires_grad_(False) lifecycle.append("load/select") @@ -839,7 +815,6 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: param_dtype=config.dtype, reduce_dtype=torch.float32, output_dtype=config.dtype, - cast_forward_inputs=False, ), ) distributed_setup = DistributedSetup.build( @@ -930,7 +905,6 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: metadata=metadata, checkpoint_keys=checkpoint_keys, lifecycle=tuple(lifecycle), - automodel_snapshot=automodel_snapshot, ) @@ -940,8 +914,6 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: raise TypeError(f"config must be PDDRecipeConfig, got {type(config).__name__}.") if not dist.is_available() or not dist.is_initialized(): raise RuntimeError("Initialize torch.distributed before building PDD export setup.") - automodel_snapshot = snapshot_installed_distribution() - from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline from nemo_automodel.components.checkpoint.config import CheckpointingConfig from nemo_automodel.components.distributed import ( @@ -952,9 +924,7 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager lifecycle = ["load/select"] - pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) - student = adopt_qwen_image_mr210_forward(loaded_transformer) - pipe.transformer = student + pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) raw_transformer_config = getattr(student, "config", None) to_dict = getattr(raw_transformer_config, "to_dict", None) if callable(to_dict): @@ -996,7 +966,6 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: param_dtype=config.dtype, reduce_dtype=torch.float32, output_dtype=config.dtype, - cast_forward_inputs=False, ), ) distributed_setup = DistributedSetup.build( @@ -1046,7 +1015,6 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: checkpoint_keys=checkpoint_keys, transformer_config=transformer_config, lifecycle=tuple(lifecycle), - automodel_snapshot=automodel_snapshot, ) @@ -1065,8 +1033,6 @@ def build_pdd_training_artifacts( adapter = QwenImagePDDAdapter( config.pdd, - guidance_rescale=config.guidance.rescale, - guidance_eps=config.guidance.eps, compute_dtype=config.dtype, ) pipeline = PDDPipeline(setup.student, setup.teacher, config.pdd, adapter) @@ -1088,8 +1054,7 @@ def build_pdd_training_artifacts( def initialize_pdd_distributed(*, backend: str, timeout_minutes: int = 60) -> Any: - """Verify the wheel, then initialize through AutoModel's released public API.""" - snapshot_installed_distribution() + """Initialize through AutoModel's released public API.""" from nemo_automodel.components.distributed import initialize_distributed return initialize_distributed(backend=backend, timeout_minutes=timeout_minutes) @@ -1106,6 +1071,7 @@ def __init__(self, cfg: Any) -> None: def setup(self) -> None: """Build data, converted models, AutoModel scheduling, and strict resume state.""" config = self.config + _require_immutable_model_source(config, context="Checkpointed PDD training") self.dist_env = initialize_pdd_distributed( backend="nccl" if config.device.type == "cuda" else "gloo", timeout_minutes=60, @@ -1169,13 +1135,9 @@ def setup(self) -> None: identity = build_pdd_checkpoint_identity( metadata=self.setup_artifacts.metadata, - forward_substrate=QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, model_id=config.model_id, model_revision=config.model_revision, guidance_scale=config.pdd.guidance_scale, - guidance_rescale=config.guidance.rescale, - guidance_eps=config.guidance.eps, - automodel_snapshot=self.setup_artifacts.automodel_snapshot, ordered_train_id_sha256=train_ordered_id_sha256, ordered_heldout_id_sha256=heldout_ordered_id_sha256, dataset_snapshot_sha256=self.snapshot_report["dataset_snapshot_sha256"], @@ -1248,10 +1210,9 @@ def _log_setup(self) -> None: self.snapshot_report["cache_root"], ) logging.info( - "PDD setup complete: lifecycle=%s student_keys=%d AutoModel=%s", + "PDD setup complete: lifecycle=%s student_keys=%d", self.setup_artifacts.lifecycle, len(self.setup_artifacts.checkpoint_keys), - self.setup_artifacts.automodel_snapshot["version"], ) def _prepared_training_batches(self): diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py index 4b9e1480fc5..b4ddf7d7d97 100644 --- a/examples/diffusers/fastgen/pdd/training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Direct PDD updates and a logical-ID-stable held-out validation oracle.""" +"""Direct PDD updates and logical-ID-stable held-out validation.""" from __future__ import annotations diff --git a/examples/diffusers/fastgen/pdd/verify_readonly_automodel.py b/examples/diffusers/fastgen/pdd/verify_readonly_automodel.py deleted file mode 100644 index 1ab6154b881..00000000000 --- a/examples/diffusers/fastgen/pdd/verify_readonly_automodel.py +++ /dev/null @@ -1,209 +0,0 @@ -# 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. - -"""Verify and snapshot the exact released AutoModel distribution used by PDD.""" - -from __future__ import annotations - -import argparse -import hashlib -import importlib.metadata -import importlib.util -import json -import os -from pathlib import Path -from typing import Any - -_MANIFEST_PATH = Path(__file__).with_name("automodel_dependency.json") -_GENERATED_NAMES = {"INSTALLER", "RECORD", "REQUESTED"} - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def load_dependency_manifest(path: Path = _MANIFEST_PATH) -> dict[str, Any]: - """Load the immutable AutoModel dependency declaration.""" - data = json.loads(path.read_text()) - required = { - "distribution", - "import_name", - "package_file_count", - "package_tree_sha256", - "release_commit", - "release_tag", - "runtime_versions", - "version", - "wheel", - "wheel_sha256", - } - if set(data) != required: - raise RuntimeError( - f"AutoModel dependency manifest keys mismatch: expected={sorted(required)}, " - f"actual={sorted(data)}." - ) - return data - - -def _distribution_root( - distribution: importlib.metadata.Distribution, manifest: dict[str, Any] -) -> Path: - root = Path(str(distribution.locate_file(""))).resolve() - dist_info = root / f"{manifest['distribution']}-{manifest['version']}.dist-info" - if not dist_info.is_dir() or not dist_info.name.endswith(".dist-info"): - raise RuntimeError(f"AutoModel has no regular wheel dist-info directory: {dist_info}.") - return root - - -def _package_files(root: Path, manifest: dict[str, Any]) -> list[Path]: - package_root = root / manifest["import_name"] - dist_info_root = root / f"{manifest['distribution']}-{manifest['version']}.dist-info" - if not package_root.is_dir() or not dist_info_root.is_dir(): - raise RuntimeError( - "AutoModel package or exact-version dist-info directory is missing from the " - f"installed distribution root {root}." - ) - - files: list[Path] = [] - for base in (package_root, dist_info_root): - for candidate in base.rglob("*"): - if candidate.is_symlink(): - raise RuntimeError(f"AutoModel distribution contains a symlink: {candidate}.") - if not candidate.is_file(): - continue - if "__pycache__" in candidate.parts or candidate.suffix == ".pyc": - continue - if candidate.name in _GENERATED_NAMES: - continue - files.append(candidate) - return sorted(files, key=lambda path: path.relative_to(root).as_posix()) - - -def snapshot_installed_distribution() -> dict[str, Any]: - """Return a deterministic content snapshot after enforcing the frozen wheel tree.""" - manifest = load_dependency_manifest() - distribution = importlib.metadata.distribution(manifest["distribution"]) - if distribution.version != manifest["version"]: - raise RuntimeError( - f"PDD requires {manifest['distribution']}=={manifest['version']}, " - f"found {distribution.version}." - ) - - runtime_versions = { - name: importlib.metadata.version(name) for name in manifest["runtime_versions"] - } - if runtime_versions != manifest["runtime_versions"]: - raise RuntimeError( - "PDD runtime dependency versions mismatch: " - f"expected {manifest['runtime_versions']}, found {runtime_versions}." - ) - - root = _distribution_root(distribution, manifest) - direct_url_text = distribution.read_text("direct_url.json") - if direct_url_text is not None: - direct_url = json.loads(direct_url_text) - if direct_url.get("dir_info", {}).get("editable", False): - raise RuntimeError("PDD rejects editable AutoModel installations.") - - spec = importlib.util.find_spec(manifest["import_name"]) - if spec is None or spec.origin is None: - raise RuntimeError(f"Cannot resolve import {manifest['import_name']!r}.") - import_origin = Path(spec.origin).resolve() - try: - import_origin.relative_to(root) - except ValueError as error: - raise RuntimeError( - f"AutoModel import {import_origin} is shadowing distribution root {root}." - ) from error - files = _package_files(root, manifest) - file_records: list[dict[str, Any]] = [] - tree_digest = hashlib.sha256() - for path in files: - relative = path.relative_to(root).as_posix() - digest = _sha256(path) - size = path.stat().st_size - tree_digest.update(relative.encode()) - tree_digest.update(b"\0") - tree_digest.update(digest.encode()) - tree_digest.update(b"\0") - tree_digest.update(str(size).encode()) - tree_digest.update(b"\n") - file_records.append({"path": relative, "sha256": digest, "size": size}) - - actual_tree_digest = tree_digest.hexdigest() - if len(file_records) != manifest["package_file_count"]: - raise RuntimeError( - "AutoModel package file count does not match the frozen wheel: " - f"expected {manifest['package_file_count']}, found {len(file_records)}." - ) - if actual_tree_digest != manifest["package_tree_sha256"]: - raise RuntimeError( - "AutoModel package tree does not match the frozen official wheel: " - f"expected {manifest['package_tree_sha256']}, found {actual_tree_digest}." - ) - - return { - "distribution": manifest["distribution"], - "files": file_records, - "import_origin": str(import_origin), - "package_file_count": len(file_records), - "package_tree_sha256": actual_tree_digest, - "release_commit": manifest["release_commit"], - "release_tag": manifest["release_tag"], - "root": str(root), - "runtime_versions": runtime_versions, - "version": distribution.version, - "wheel": manifest["wheel"], - "wheel_sha256": manifest["wheel_sha256"], - } - - -def write_snapshot(output: Path) -> None: - """Atomically write a distribution snapshot outside the installation.""" - snapshot = snapshot_installed_distribution() - output = output.resolve() - root = Path(snapshot["root"]) - try: - output.relative_to(root) - except ValueError: - pass - else: - raise ValueError("Snapshot output must be outside the AutoModel distribution.") - output.parent.mkdir(parents=True, exist_ok=True) - temporary = output.with_name(f".{output.name}.{os.getpid()}.tmp") - temporary.write_text(json.dumps(snapshot, indent=2, sort_keys=True) + "\n") - os.replace(temporary, output) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - subparsers = parser.add_subparsers(dest="command", required=True) - snapshot = subparsers.add_parser("snapshot", help="verify and write a content snapshot") - snapshot.add_argument("--output", type=Path, required=True) - return parser.parse_args() - - -def main() -> None: - args = _parse_args() - if args.command == "snapshot": - write_snapshot(args.output) - - -if __name__ == "__main__": - main() diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index ab54870c91c..99049431e9a 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -5,7 +5,7 @@ diffusers==0.38.0 # NeMo AutoModel supplies the parent recipe, FSDP2 wrapping, checkpointer, and the unmodified # upstream diffusion helpers used by the shared data/preprocessing code. These versions are the -# public APIs tested by both DMD2 and PDD; PDD also verifies the exact AutoModel wheel at runtime. +# public APIs tested by both DMD2 and PDD. nemo_automodel[diffusion]==0.5.0 # Optional but recommended for training logs. diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index 429e240ed06..1165fbfc22f 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -970,9 +970,8 @@ def sample( kwargs = self._model_kwargs(model_kwargs) resolved_blocks = self._validate_blocks(blocks) grid = self.time_grid(noise.device) - # MR210 derives fused projection coefficients from its float64 schedule, - # then casts the coefficients to FP32. Keep state/time integration on the - # canonical FP32 decoding grid while matching that coefficient path. + # 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, diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index bf81fd0953a..a9464c76adc 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -17,8 +17,6 @@ from __future__ import annotations -import hashlib -import inspect import math from collections.abc import Mapping from typing import Any @@ -31,34 +29,11 @@ from .qwen_image import build_img_shapes, pack_latents, unpack_latents __all__ = [ - "QWEN_IMAGE_PDD_FORWARD_SUBSTRATE", - "QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID", "QWEN_IMAGE_PDD_LAYER_SPEC", "QwenImagePDDAdapter", - "adopt_qwen_image_mr210_forward", "convert_qwen_image_to_pdd", - "require_qwen_image_pdd_forward_substrate", ] -QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID = ( - "pdd_qwen_mr210_c8100b1347b278511336dccfc074a461457216ec_" - "qwen_33706683487ba16d133b99b73be27b21164c53335441d77b1dcabbfca970f70e" -) -QWEN_IMAGE_PDD_FORWARD_SUBSTRATE = { - "id": QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID, - "fastgen_commit": "c8100b1347b278511336dccfc074a461457216ec", - "fastgen_qwen_source_sha256": ( - "33706683487ba16d133b99b73be27b21164c53335441d77b1dcabbfca970f70e" - ), - "diffusers_version": "0.38.0", - "diffusers_qwen_source_sha256": ( - "34c864b0b066a4a9eb84e40e1bb77b7df303c165e7910600b402a0f5f8d8f94e" - ), - "diffusers_embeddings_source_sha256": ( - "d7a90ef799569e3f0fab41cadde1ecba023abd053af956c022bbfc097662a302" - ), -} - QWEN_IMAGE_PDD_LAYER_SPEC = PDDLayerSpec( projection_path="transformer.proj_out", head_layout="channel_major", @@ -79,274 +54,11 @@ "txt_seq_lens", } -_QWEN_IMAGE_ROOT_CHILDREN = ( - "pos_embed", - "time_text_embed", - "txt_norm", - "img_in", - "txt_in", - "transformer_blocks", - "norm_out", - "proj_out", -) -_ADOPTED_QWEN_TYPES: dict[type[nn.Module], type[nn.Module]] = {} - - -def require_qwen_image_pdd_forward_substrate(value: Any) -> dict[str, str]: - """Return the authenticated MR210 Qwen substrate or reject it exactly.""" - if not isinstance(value, Mapping) or dict(value) != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE: - raise ValueError( - "Qwen-Image PDD requires the authenticated MR210 forward substrate " - f"{QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID!r}." - ) - return dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE) - - -def _sha256_source(owner: type[Any]) -> str: - source = inspect.getsourcefile(owner) - if source is None: - raise RuntimeError(f"cannot locate source for {owner.__module__}.{owner.__qualname__}.") - with open(source, "rb") as stream: - return hashlib.file_digest(stream, "sha256").hexdigest() - - -def _require_qwen_source_identity(transformer: nn.Module) -> None: - transformer_type = type(transformer) - if ( - transformer_type.__module__ != "diffusers.models.transformers.transformer_qwenimage" - or transformer_type.__name__ != "QwenImageTransformer2DModel" - ): - raise TypeError("MR210 adoption requires the pinned Diffusers QwenImageTransformer2DModel.") - if ( - _sha256_source(transformer_type) - != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_qwen_source_sha256"] - ): - raise RuntimeError("Diffusers Qwen transformer source does not match the MR210 substrate.") - embedding_types = ( - type(transformer.time_text_embed.time_proj), - type(transformer.time_text_embed.timestep_embedder), - ) - if any( - _sha256_source(embedding_type) - != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_embeddings_source_sha256"] - for embedding_type in embedding_types - ): - raise RuntimeError( - "Diffusers timestep embedding source does not match the MR210 substrate." - ) - try: - import diffusers # Optional dependency required only by the Qwen adoption path. - except ImportError as error: # pragma: no cover - the transformer itself requires Diffusers - raise RuntimeError("Diffusers is required for Qwen-Image PDD adoption.") from error - if diffusers.__version__ != QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_version"]: - raise RuntimeError( - "Diffusers version does not match the authenticated Qwen MR210 substrate." - ) - -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 _require_binary_prefix_mask( - encoder_hidden_states: torch.Tensor, - mask: torch.Tensor, -) -> None: - if mask.ndim != 2 or tuple(mask.shape) != tuple(encoder_hidden_states.shape[:2]): - raise ValueError("Qwen MR210 mask must match the text batch and sequence dimensions.") - if mask.dtype.is_floating_point or mask.dtype.is_complex: - raise TypeError("Qwen MR210 mask must use an integer or boolean dtype.") - if mask.device != encoder_hidden_states.device: - raise ValueError("Qwen MR210 mask and text embeddings must share a device.") - if mask.shape[0] == 0 or mask.shape[1] == 0: - raise ValueError("Qwen MR210 masks must have nonempty batch and sequence dimensions.") +def _require_binary_mask(mask: torch.Tensor, *, name: str) -> None: mask_int = mask.to(torch.int64) - mask_bool = mask.bool() - binary = torch.all((mask_int == 0) | (mask_int == 1)) - prefix = torch.all(mask_int[:, 1:] <= mask_int[:, :-1]) - lengths = mask_int.sum(dim=1) - valid_lengths = torch.all(lengths > 0) & (lengths.max() == mask.shape[1]) - zero_padding = torch.all(encoder_hidden_states[~mask_bool] == 0) - if not bool((binary & prefix & valid_lengths & zero_padding).item()): - raise ValueError( - "Qwen MR210 requires nonempty binary prefix masks, a longest unpadded row, " - "and zero padding." - ) - - -class _QwenImageMR210ForwardMixin: - """Execute FastGen MR210's Qwen forward without altering Diffusers classes. - - Source contract: ``fastgen/networks/QwenImage/network.py`` at - ``c8100b1347b278511336dccfc074a461457216ec``. - """ - - def forward( - self, - 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, - txt_seq_lens: list[int] | None = None, - guidance: torch.Tensor | None = None, - attention_kwargs: dict[str, Any] | None = None, - controlnet_block_samples: Any = None, - additional_t_cond: Any = None, - return_dict: bool = True, - *, - max_txt_seq_len: int | None = None, - ) -> Any: - """Run the source-locked MR210 regular-output path.""" - if hidden_states.ndim != 3 or hidden_states.dtype != torch.bfloat16: - raise TypeError("Qwen MR210 hidden_states must be packed BF16 [B, P, C].") - if ( - not isinstance(encoder_hidden_states, torch.Tensor) - or encoder_hidden_states.ndim != 3 - or encoder_hidden_states.dtype != torch.bfloat16 - ): - raise TypeError("Qwen MR210 encoder_hidden_states must be BF16 [B, S, D].") - if not isinstance(encoder_hidden_states_mask, torch.Tensor): - raise TypeError("Qwen MR210 requires encoder_hidden_states_mask.") - if not isinstance(timestep, torch.Tensor) or timestep.dtype != torch.float32: - raise TypeError("Qwen MR210 timestep must remain FP32 at transformer entry.") - if timestep.shape != (hidden_states.shape[0],): - raise ValueError("Qwen MR210 timestep must contain one value per batch item.") - if encoder_hidden_states.shape[0] != hidden_states.shape[0]: - raise ValueError("Qwen MR210 image and text batch sizes must match.") - if img_shapes is None or len(img_shapes) != hidden_states.shape[0]: - raise ValueError("Qwen MR210 img_shapes must contain one entry per batch item.") - if txt_seq_lens is not None: - raise ValueError("Qwen MR210 does not support txt_seq_lens.") - if guidance is not None: - raise ValueError("Qwen MR210 does not support transformer guidance embeddings.") - if attention_kwargs: - raise ValueError("Qwen MR210 does not support nonempty attention_kwargs.") - if controlnet_block_samples is not None: - raise ValueError("Qwen MR210 does not support ControlNet residuals.") - if additional_t_cond is not None: - raise ValueError("Qwen MR210 does not support additional time conditioning.") - if type(return_dict) is not bool: - raise TypeError("return_dict must be bool.") - _require_binary_prefix_mask(encoder_hidden_states, encoder_hidden_states_mask) - sequence_length = encoder_hidden_states.shape[1] - if max_txt_seq_len is not None and max_txt_seq_len != sequence_length: - raise ValueError("Qwen MR210 max_txt_seq_len must equal the padded text length.") - - hidden_states = self.img_in(hidden_states) - encoder_hidden_states = self.txt_norm(encoder_hidden_states) - encoder_hidden_states = self.txt_in(encoder_hidden_states) - if timestep.dtype != torch.float32: - raise RuntimeError("Qwen MR210 timestep was rounded before time_text_embed.") - temb = self.time_text_embed(timestep, hidden_states) - image_rotary_emb = self.pos_embed( - img_shapes, - max_txt_seq_len=sequence_length, - 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, - joint_attention_kwargs=attention_kwargs, - ) - - hidden_states = self.norm_out(hidden_states, temb) - output = self.proj_out(hidden_states) - if not return_dict: - return (output,) - try: - from diffusers.models.modeling_outputs import ( # Optional Qwen runtime dependency. - Transformer2DModelOutput, - ) - except ImportError as error: # pragma: no cover - adoption already requires Diffusers - raise RuntimeError("Diffusers output types are unavailable.") from error - return Transformer2DModelOutput(sample=output) - - -def _adopted_qwen_type(base: type[nn.Module]) -> type[nn.Module]: - adopted = _ADOPTED_QWEN_TYPES.get(base) - if adopted is None: - adopted = type( - f"ModelOptMR210{base.__name__}", - (_QwenImageMR210ForwardMixin, base), - {"__module__": __name__}, - ) - _ADOPTED_QWEN_TYPES[base] = adopted - return adopted - - -def adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: - """Adopt loaded Qwen children into a Diffusers-compatible MR210 forward root.""" - if isinstance(transformer, _QwenImageMR210ForwardMixin): - return transformer - if not isinstance(transformer, nn.Module): - raise TypeError(f"transformer must be nn.Module, got {type(transformer).__name__}.") - _require_qwen_source_identity(transformer) - if tuple(transformer._modules) != _QWEN_IMAGE_ROOT_CHILDREN: - raise RuntimeError("Qwen root child layout does not match the authenticated substrate.") - if transformer._parameters or transformer._buffers: - raise RuntimeError("Qwen root unexpectedly registers direct parameters or buffers.") - if _config_guidance_embeds(transformer): - raise ValueError("Qwen MR210 does not support transformer guidance embeddings.") - if getattr(transformer, "peft_config", None): - raise ValueError("Qwen MR210 does not support active PEFT adapters.") - if any(getattr(module, "fused_projections", False) for module in transformer.modules()): - raise ValueError("Qwen MR210 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 MR210 requires {name}=False.") - hook_names = ( - "_backward_hooks", - "_backward_pre_hooks", - "_forward_hooks", - "_forward_pre_hooks", - "_load_state_dict_post_hooks", - "_load_state_dict_pre_hooks", - "_state_dict_hooks", - "_state_dict_pre_hooks", - ) - if any(getattr(transformer, name, None) for name in hook_names): - raise RuntimeError("Qwen root hooks must be empty before MR210 adoption.") - - adopted_type = _adopted_qwen_type(type(transformer)) - adopted = adopted_type.__new__(adopted_type) - nn.Module.__init__(adopted) - adopted._internal_dict = transformer._internal_dict - for name in ("out_channels", "inner_dim", "gradient_checkpointing", "zero_cond_t"): - setattr(adopted, name, getattr(transformer, name)) - if hasattr(transformer, "_gradient_checkpointing_func"): - adopted._gradient_checkpointing_func = transformer._gradient_checkpointing_func - for name, child in transformer._modules.items(): - adopted.add_module(name, child) - adopted.train(transformer.training) - if tuple(adopted.state_dict()) != tuple(transformer.state_dict()): - raise RuntimeError("Qwen state keys changed during MR210 adoption.") - if any( - adopted_parameter is not source_parameter - for adopted_parameter, source_parameter in zip( - adopted.parameters(), transformer.parameters(), strict=True - ) - ): - raise RuntimeError("Qwen parameter identity changed during MR210 adoption.") - return adopted + if not bool(torch.all((mask_int == 0) | (mask_int == 1)).item()): + raise ValueError(f"{name} mask must contain only zero and one values.") def _config_guidance_embeds(transformer: nn.Module) -> bool: @@ -409,20 +121,10 @@ def __init__( self, config: PDDConfig, *, - guidance_rescale: float = 1.0, - guidance_eps: float = 1e-5, compute_dtype: torch.dtype | None = None, ) -> None: """Validate the fixed Qwen continuous-time and packed-CFG contract.""" _validate_qwen_pdd_config(config) - if isinstance(guidance_rescale, bool) or not isinstance(guidance_rescale, int | float): - raise TypeError("guidance_rescale must be a real number.") - if not math.isfinite(guidance_rescale) or not 0.0 <= guidance_rescale <= 1.0: - raise ValueError("guidance_rescale must be finite and in [0, 1].") - if isinstance(guidance_eps, bool) or not isinstance(guidance_eps, int | float): - raise TypeError("guidance_eps must be a real number.") - if not math.isfinite(guidance_eps) or guidance_eps <= 0.0: - raise ValueError("guidance_eps must be finite and > 0.") 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 ( @@ -434,8 +136,6 @@ def __init__( self.guidance_scale = ( None if config.guidance_scale is None else float(config.guidance_scale) ) - self.guidance_rescale = float(guidance_rescale) - self.guidance_eps = float(guidance_eps) self.compute_dtype = compute_dtype @staticmethod @@ -490,7 +190,7 @@ def _parse_condition( raise ValueError(f"{name} batch size must match state batch size {batch_size}.") if encoder_hidden_states.device != state.device or attention_mask.device != state.device: raise ValueError(f"{name} tensors must be on {state.device}.") - _require_binary_prefix_mask(encoder_hidden_states, attention_mask) + _require_binary_mask(attention_mask, name=name) return encoder_hidden_states, attention_mask def _model_dtype(self, model: nn.Module, fallback: torch.dtype) -> torch.dtype: @@ -568,10 +268,6 @@ def _call_packed( batch_size, _, height, width = state.shape model_dtype = self._model_dtype(model, state.dtype) - if model_dtype != torch.bfloat16: - raise TypeError("authenticated Qwen MR210 execution requires BF16 compute.") - if time.dtype != torch.float32: - raise TypeError("authenticated Qwen MR210 execution requires FP32 time.") packed_state = pack_latents(state).to(model_dtype) encoder_hidden_states = encoder_hidden_states.to(model_dtype) output = model( @@ -759,27 +455,7 @@ def teacher_velocity( f"{tuple(conditional.shape)} and {tuple(unconditional.shape)}." ) - # FastGen applies CFG in the model-output dtype, including its BF16 - # rounding, before cfg_rescale promotes the result for norm math. - guided_model_dtype = conditional + (float(guidance_scale) - 1.0) * ( - conditional - unconditional - ) - conditional_fp32 = conditional.to(torch.float32) - guided = guided_model_dtype.to(torch.float32) - # MR210 applies CFG after unpacking Qwen output to NCHW and leaves - # ``rescale_dims`` unset, so the norm spans every non-batch element. - # Reducing packed [P, F] here is algebraically identical because - # pack/unpack only reshapes and permutes those elements. - norm_dims = tuple(range(1, conditional_fp32.ndim)) - conditional_norm = torch.linalg.vector_norm( - conditional_fp32, - dim=norm_dims, - keepdim=True, - ) - guided_norm = torch.linalg.vector_norm(guided, dim=norm_dims, keepdim=True) - factor = self.guidance_rescale * conditional_norm / guided_norm.clamp_min( - self.guidance_eps - ) + (1.0 - self.guidance_rescale) - # FastGen's cfg_rescale returns to the conditional model-output dtype - # before PDD promotes the teacher target for FP32 loss math. - return self._unpack_single((guided * factor).to(conditional.dtype), state) + guided = unconditional + float(guidance_scale) * (conditional - unconditional) + conditional_norm = torch.linalg.vector_norm(conditional, dim=-1, keepdim=True) + guided_norm = torch.linalg.vector_norm(guided, dim=-1, keepdim=True) + return self._unpack_single(guided * (conditional_norm / guided_norm), state) diff --git a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py index bc13f6c0771..8c80fa52685 100644 --- a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py @@ -141,7 +141,7 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 3, "topology": {"world_size": 2}}, + identity={"schema_version": 4, "topology": {"world_size": 2}}, ) initial.save() trainer.completed_steps = 2 @@ -156,7 +156,7 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 3, "topology": {"world_size": 2}}, + identity={"schema_version": 4, "topology": {"world_size": 2}}, ) message = None try: diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py index 20d42417b91..79cd5a072e5 100644 --- a/tests/examples/diffusers/fastgen/pdd_export_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_export_distributed.py @@ -38,8 +38,6 @@ from pdd.inference_qwen_image import build_pdd_student from pdd.recipe import build_pdd_export_setup, resolve_pdd_recipe_config -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_FORWARD_SUBSTRATE - def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: return { @@ -135,25 +133,14 @@ def main() -> None: tensor.numel() * tensor.element_size() for tensor in actual.values() ) identity = { - "schema_version": 1, - "forward_substrate": dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE), + "schema_version": 4, "model": { "id": "Qwen/Qwen-Image", "revision": "3" * 40, "dtype": "bfloat16", }, "pdd_metadata": destination.metadata.to_dict(), - "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, - "automodel": { - key: destination.automodel_snapshot[key] - for key in ( - "distribution", - "version", - "package_tree_sha256", - "wheel_sha256", - "runtime_versions", - ) - }, + "guidance": {"scale": 4.0}, "topology": {"world_size": 2, "pure_data_parallel": True}, } output = write_pdd_export( @@ -167,7 +154,6 @@ def main() -> None: "manifest_sha256": "1" * 64, "completed_steps": 1, }, - modelopt_source={"commit": "2" * 40, "dirty": False}, max_shard_bytes=4 * 1024 * 1024, ) descriptor = inspect_pdd_export(output) diff --git a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py index 1f53db80e8f..6c37eed91ba 100644 --- a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Two-rank CPU/Gloo equivalence harness for the deterministic PDD validation oracle.""" +"""Two-rank CPU/Gloo equivalence harness for deterministic PDD validation.""" from __future__ import annotations diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index d347b067e1d..6f96be4643f 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -52,7 +52,6 @@ "README.md", "__init__.py", "artifacts.py", - "automodel_dependency.json", "checkpoint.py", "configs", "data.py", @@ -62,7 +61,6 @@ "inference_qwen_image.py", "recipe.py", "training.py", - "verify_readonly_automodel.py", } _TEXT_SUFFIXES = {".json", ".md", ".py", ".rst", ".sh", ".toml", ".txt", ".yaml", ".yml"} diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py index 7c8e92819f3..9113d18ea5b 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -39,11 +39,14 @@ pdd_config_from_metadata, write_pdd_export, ) -from pdd.inference_qwen_image import _normalize_prompt_condition, _validate_qwen_projection +from pdd.inference_qwen_image import ( + _model_identity, + _normalize_prompt_condition, + _validate_qwen_projection, +) from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, QwenImagePDDAdapter, convert_qwen_image_to_pdd, ) @@ -108,18 +111,10 @@ def _converted(seed: int = 17): def _identity(metadata: PDDMetadata) -> dict: return { - "schema_version": 1, - "forward_substrate": dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE), + "schema_version": 4, "model": {"id": "synthetic-qwen", "revision": "f" * 40, "dtype": "bfloat16"}, "pdd_metadata": metadata.to_dict(), - "guidance": {"scale": 4.0, "rescale": 1.0, "eps": 1e-5}, - "automodel": { - "distribution": "nemo_automodel", - "version": "0.5.0", - "package_tree_sha256": "1" * 64, - "wheel_sha256": "2" * 64, - "runtime_versions": {"diffusers": "0.38.0"}, - }, + "guidance": {"scale": 4.0}, "topology": {"world_size": 1, "pure_data_parallel": True}, } @@ -137,7 +132,6 @@ def _write(tmp_path: pathlib.Path): "manifest_sha256": "3" * 64, "completed_steps": 10, }, - modelopt_source={"commit": "4" * 40, "dirty": False}, max_shard_bytes=5_800, ) return output, model, config, metadata @@ -259,13 +253,33 @@ def checked_rename(path, target): assert observed -def test_export_rejects_unpinned_local_model_identity(tmp_path) -> None: +def test_export_accepts_an_immutable_model_revision(tmp_path) -> None: model, _config_value, metadata = _converted() identity = _identity(metadata) - identity["model"]["revision"] = None - with pytest.raises(ValueError, match="pinned 40-character"): + output = write_pdd_export( + tmp_path / "immutable-revision", + model.state_dict(), + metadata=metadata, + transformer_config={"in_channels": 4}, + identity=identity, + source_checkpoint={ + "name": "step_00000010", + "manifest_sha256": "3" * 64, + "completed_steps": 10, + }, + max_shard_bytes=12_000, + ) + assert inspect_pdd_export(output).manifest["identity"]["model"]["revision"] == "f" * 40 + + +@pytest.mark.parametrize("revision", [None, "main", "F" * 40]) +def test_export_and_inference_reject_mutable_model_revisions(tmp_path, revision) -> None: + model, _config_value, metadata = _converted() + identity = _identity(metadata) + identity["model"]["revision"] = revision + with pytest.raises(ValueError, match="exact lowercase commit"): write_pdd_export( - tmp_path / "local-model-export", + tmp_path / f"bad-revision-{str(revision)[:8]}", model.state_dict(), metadata=metadata, transformer_config={"in_channels": 4}, @@ -275,31 +289,12 @@ def test_export_rejects_unpinned_local_model_identity(tmp_path) -> None: "manifest_sha256": "3" * 64, "completed_steps": 10, }, - modelopt_source={"commit": "4" * 40, "dirty": False}, max_shard_bytes=12_000, ) - -def test_export_rejects_missing_or_mismatched_forward_substrate(tmp_path) -> None: - model, _config_value, metadata = _converted() - for substrate in (None, {**QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, "id": "canonical"}): - identity = _identity(metadata) - identity["forward_substrate"] = substrate - with pytest.raises(ValueError, match="authenticated MR210"): - write_pdd_export( - tmp_path / f"bad-substrate-{substrate is None}", - model.state_dict(), - metadata=metadata, - transformer_config={"in_channels": 4}, - identity=identity, - source_checkpoint={ - "name": "step_00000010", - "manifest_sha256": "3" * 64, - "completed_steps": 10, - }, - modelopt_source={"commit": "4" * 40, "dirty": False}, - max_shard_bytes=12_000, - ) + descriptor = SimpleNamespace(manifest={"identity": identity}) + with pytest.raises(RuntimeError, match="exact lowercase commit"): + _model_identity(descriptor) def test_export_rejects_nonfinite_and_existing_destination(tmp_path) -> None: @@ -316,7 +311,6 @@ def test_export_rejects_nonfinite_and_existing_destination(tmp_path) -> None: "manifest_sha256": "3" * 64, "completed_steps": 10, }, - modelopt_source={"commit": "4" * 40, "dirty": False}, max_shard_bytes=12_000, ) @@ -334,7 +328,6 @@ def test_export_rejects_nonfinite_and_existing_destination(tmp_path) -> None: "manifest_sha256": "3" * 64, "completed_steps": 10, }, - modelopt_source={"commit": "4" * 40, "dirty": False}, max_shard_bytes=12_000, ) diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index d3abf821d9d..97bb29217aa 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -18,11 +18,8 @@ from __future__ import annotations import copy -import importlib.metadata -import json import os import pathlib -import shutil import subprocess import sys @@ -41,13 +38,14 @@ _materialize_zero_step_adamw_state, _projection_identity, _require_fp32_optimizer_storage, + _require_immutable_model_source, + _resolve_model_source, _stage_and_shard_training_models, build_pdd_export_setup, build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config, ) -from pdd.verify_readonly_automodel import snapshot_installed_distribution from modelopt.torch.fastgen import PDDLayerSpec, convert_to_pdd_output_projection @@ -115,15 +113,6 @@ def test_fp32_optimizer_storage_rejects_low_precision_masters_and_state() -> Non _require_fp32_optimizer_storage(fp32_optimizer) -def _require_exact_automodel() -> None: - try: - version = importlib.metadata.version("nemo_automodel") - except importlib.metadata.PackageNotFoundError: - pytest.skip("nemo_automodel is not installed") - if version != "0.5.0": - pytest.skip(f"requires the official nemo_automodel==0.5.0 wheel, found {version}") - - def test_training_setup_shards_student_before_staging_teacher() -> None: events: list[str] = [] @@ -422,8 +411,6 @@ def test_pdd_finetune_namespace_module_help() -> None: ("model", "guidance_embeds", True, "guidance embeddings"), ("model", "device_map", "auto", "device_map"), ("model", "quantization_config", {"bits": 8}, "quantization_config"), - ("model", "fuse_qkv_projections", True, "QKV fusion"), - ("model", "torch_dtype", "float32", "requires model.torch_dtype='bfloat16'"), ], ) def test_incompatible_modes_fail_during_config_resolution( @@ -437,13 +424,71 @@ def test_incompatible_modes_fail_during_config_resolution( resolve_pdd_recipe_config(raw) -def test_remote_model_requires_full_revision_and_non_dp_parallelism_is_rejected(tmp_path) -> None: +def test_model_revision_and_compute_dtype_follow_loader_contract(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" + raw["model"]["revision"] = "a" * 40 + raw["model"]["torch_dtype"] = "float32" + raw["model"]["fuse_qkv_projections"] = True + + config = resolve_pdd_recipe_config(raw) + + assert config.model_revision == "a" * 40 + assert config.dtype == torch.float32 + assert config.fuse_qkv_projections is True + + +@pytest.mark.parametrize("revision", [None, "main", "A" * 40, "a" * 39]) +def test_remote_model_requires_exact_lowercase_commit(tmp_path, revision) -> None: raw = _raw_config(tmp_path) raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" - with pytest.raises(ValueError, match=r"exact model\.revision"): + raw["model"]["revision"] = revision + + with pytest.raises(ValueError, match="exact lowercase 40-character"): resolve_pdd_recipe_config(raw) - raw["model"]["revision"] = "a" * 40 + +def test_model_source_resolution_requires_the_requested_snapshot(tmp_path, monkeypatch) -> None: + commit = "a" * 40 + raw = _raw_config(tmp_path) + raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" + raw["model"]["revision"] = commit + config = resolve_pdd_recipe_config(raw) + snapshot = tmp_path / "hub" / "snapshots" / commit + snapshot.mkdir(parents=True) + calls = [] + + def matching_snapshot(model_id, *, revision): + calls.append((model_id, revision)) + return str(snapshot) + + monkeypatch.setattr("huggingface_hub.snapshot_download", matching_snapshot) + assert _resolve_model_source(config) == str(snapshot.resolve()) + assert calls == [("Qwen/Qwen-Image", commit)] + _require_immutable_model_source(config, context="test") + + wrong = snapshot.with_name("b" * 40) + wrong.mkdir() + monkeypatch.setattr("huggingface_hub.snapshot_download", lambda *_args, **_kwargs: str(wrong)) + with pytest.raises(RuntimeError, match=r"does not match model\.revision"): + _resolve_model_source(config) + + +def test_local_model_source_is_limited_to_low_level_setup(tmp_path, monkeypatch) -> None: + config = resolve_pdd_recipe_config(_raw_config(tmp_path)) + monkeypatch.setattr( + "huggingface_hub.snapshot_download", + lambda *_args, **_kwargs: pytest.fail("local model resolution must not access the Hub"), + ) + + assert _resolve_model_source(config) == str(tmp_path.resolve()) + with pytest.raises(ValueError, match="Checkpointed PDD training requires"): + _require_immutable_model_source(config, context="Checkpointed PDD training") + + +def test_non_dp_parallelism_is_rejected(tmp_path) -> None: + raw = _raw_config(tmp_path) + raw["fsdp"]["tp_size"] = 2 with pytest.raises(ValueError, match="tp_size must be 1"): resolve_pdd_recipe_config(raw) @@ -460,7 +505,7 @@ def test_remote_model_requires_full_revision_and_non_dp_parallelism_is_rejected( ), ("training_health", "max_grad_norm", 0.0, "max_grad_norm must be > 0"), ("validation", "every_steps", 0, "validation.every_steps"), - ("guidance", "rescale", 1.1, "guidance.rescale must be <= 1"), + ("guidance", "rescale", 1.1, "does not support guidance overrides"), ("optimizer", "betas", [0.9, 1.0], "optim.optimizer.betas values"), ("optimizer", "eps", 0.0, "optim.optimizer.eps must be > 0"), ], @@ -515,66 +560,7 @@ def test_payload_hash_verification_mode_must_be_bool(tmp_path) -> None: resolve_pdd_recipe_config(raw) -def test_frozen_automodel_distribution_snapshot_is_stable() -> None: - _require_exact_automodel() - - before = snapshot_installed_distribution() - after = snapshot_installed_distribution() - - assert before == after - assert before["version"] == "0.5.0" - assert before["release_commit"] == "d02f49cb314554715aabb97e8dba6599c9f6e9e0" - assert before["runtime_versions"] == {"diffusers": "0.38.0"} - assert before["package_file_count"] == 490 - assert before["package_tree_sha256"] == ( - "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528" - ) - - -def test_exact_wheel_install_below_git_checkout_is_accepted(tmp_path) -> None: - _require_exact_automodel() - try: - distribution = importlib.metadata.distribution("nemo_automodel") - except importlib.metadata.PackageNotFoundError: - pytest.skip("nemo_automodel is not installed") - - checkout = tmp_path / "checkout" - (checkout / ".git").mkdir(parents=True) - site_packages = checkout / ".venv" / "lib" / "python" / "site-packages" - site_packages.mkdir(parents=True) - installed_root = pathlib.Path(distribution.locate_file("")).resolve() - shutil.copytree(installed_root / "nemo_automodel", site_packages / "nemo_automodel") - dist_info_name = "nemo_automodel-0.5.0.dist-info" - shutil.copytree(installed_root / dist_info_name, site_packages / dist_info_name) - - output = tmp_path / "snapshot.json" - environment = os.environ.copy() - environment["PYTHONPATH"] = os.pathsep.join( - filter(None, (str(site_packages), environment.get("PYTHONPATH"))) - ) - subprocess.run( - [ - sys.executable, - str(_FASTGEN_DIR / "pdd" / "verify_readonly_automodel.py"), - "snapshot", - "--output", - str(output), - ], - check=True, - env=environment, - ) - - snapshot = json.loads(output.read_text()) - assert pathlib.Path(snapshot["root"]) == site_packages.resolve() - assert pathlib.Path(snapshot["import_origin"]).is_relative_to(site_packages.resolve()) - assert snapshot["package_tree_sha256"] == ( - "b43cb34e04992c66d1888abc0529b760b5b69fc121ff4268b42ecb4a89b1e528" - ) - - def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: - _require_exact_automodel() - before = snapshot_installed_distribution() model_dir = create_tiny_qwen_image_pipeline_dir(tmp_path) initialize_pdd_distributed(backend="gloo", timeout_minutes=1) config = resolve_pdd_recipe_config(_raw_config(model_dir)) @@ -603,7 +589,11 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: 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 policy.cast_forward_inputs is True + assert all( + base.__module__ != "modelopt.torch.fastgen.plugins.qwen_image_pdd" + for base in type(source.student).__mro__ + ) assert not any(parameter.requires_grad for parameter in source.teacher.parameters()) optimizer_parameters = [ parameter for group in source.optimizer.param_groups for parameter in group["params"] @@ -684,7 +674,6 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert destination.optimizer.state[destination_unused]["step"].item() == 0 assert not destination.optimizer.state[destination_unused]["exp_avg"].count_nonzero() assert not destination.optimizer.state[destination_unused]["exp_avg_sq"].count_nonzero() - assert snapshot_installed_distribution() == before source.checkpointer.close() destination.checkpointer.close() export_setup.checkpointer.close() diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index d7972d52c8e..cb8d4f6facd 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -19,7 +19,6 @@ import copy import hashlib -import importlib.metadata import json import math import pathlib @@ -48,11 +47,8 @@ ) from pdd.recipe import PDDDiffusionRecipe, initialize_pdd_distributed from pdd.training import prepare_qwen_pdd_batch -from pdd.verify_readonly_automodel import snapshot_installed_distribution from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_FORWARD_SUBSTRATE - def _released_sampler(sample_ids: tuple[str, ...]) -> ReplayableBatchSampler: sampler_module = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") @@ -100,13 +96,9 @@ def _checkpointer(lifecycle, checkpoint_dir): def _identity(lifecycle, scheduler, sample_ids): return build_pdd_checkpoint_identity( metadata=lifecycle.metadata, - forward_substrate=QWEN_IMAGE_PDD_FORWARD_SUBSTRATE, model_id="synthetic-pdd-toy", - model_revision=None, + model_revision="a" * 40, guidance_scale=None, - guidance_rescale=1.0, - guidance_eps=1e-5, - automodel_snapshot=snapshot_installed_distribution(), ordered_train_id_sha256=ordered_id_sha256(sample_ids), ordered_heldout_id_sha256="1" * 64, dataset_snapshot_sha256="2" * 64, @@ -636,9 +628,6 @@ def test_training_hard_aborts_for_teacher_gradient_zero_gradient_and_missing_cov def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) -> None: pytest.importorskip("nemo_automodel") - version = importlib.metadata.version("nemo_automodel") - if version != "0.5.0": - pytest.skip(f"requires the official nemo_automodel==0.5.0 wheel, found {version}") rng_module = pytest.importorskip("nemo_automodel.components.training.rng") if not torch.distributed.is_initialized(): initialize_pdd_distributed(backend="gloo", timeout_minutes=1) @@ -716,13 +705,6 @@ def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) assert second_resume.completed_steps == 3 assert second_resume.expected_next_sample_ids == destination_sampler.expected_next_sample_ids() - inventory = {path.name.lower() for path in resumed_checkpoint.rglob("*")} - assert not any( - token in name - for name in inventory - for token in ("fake_score", "discriminator", "ema", "r1", "gan") - ) - incomplete = tmp_path / "checkpoints" / "step_99999998" incomplete.mkdir() (tmp_path / "checkpoints" / "LATEST").write_text(incomplete.name + "\n") @@ -748,6 +730,15 @@ def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) ) assert selected == resumed_checkpoint.resolve() assert selected_manifest["identity"] == third_manager.identity + moved_model_identity = copy.deepcopy(third_manager.identity) + moved_model_identity["model"]["revision"] = "b" * 40 + with pytest.raises(RuntimeError, match="identity"): + resolve_pdd_training_checkpoint( + tmp_path / "checkpoints", + resumed_checkpoint.name, + expected_world_size=1, + expected_identity=moved_model_identity, + ) with pytest.raises(RuntimeError, match="identity"): third_manager.resolve(mismatched.name) diff --git a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py index 83b7ce81480..d7b1f0afc80 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Deterministic logical-ID PDD held-out oracle tests.""" +"""Deterministic logical-ID PDD held-out validation tests.""" from __future__ import annotations diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 37a7f7627e9..6df51b54352 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -30,12 +30,8 @@ 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_FORWARD_SUBSTRATE, - QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID, QWEN_IMAGE_PDD_LAYER_SPEC, - adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, - require_qwen_image_pdd_forward_substrate, ) @@ -228,10 +224,10 @@ def test_fused_student_matches_explicit_packed_weight_fusion() -> None: assert student.proj_out(projection.weight.new_zeros(1, 5)).shape[-1] == 16 -def test_teacher_cfg_and_global_norm_rescale_match_mr210_reference() -> None: +def test_teacher_cfg_uses_canonical_packed_per_token_norm_rescale() -> None: teacher = _TinyQwenTransformer() config = _config(guidance_scale=4.0) - adapter = QwenImagePDDAdapter(config, guidance_rescale=1.0, guidance_eps=1e-5) + adapter = QwenImagePDDAdapter(config) state, time, condition, negative_condition = _inputs() actual = adapter.teacher_velocity( @@ -243,17 +239,15 @@ def test_teacher_cfg_and_global_norm_rescale_match_mr210_reference() -> None: ) assert len(teacher.calls) == 2 - conditional_bf16 = teacher.calls[0]["output"] - unconditional_bf16 = teacher.calls[1]["output"] - guided_bf16 = conditional_bf16 + 3.0 * (conditional_bf16 - unconditional_bf16) - conditional = conditional_bf16.float() - guided = guided_bf16.float() + conditional = teacher.calls[0]["output"] + unconditional = teacher.calls[1]["output"] + guided = unconditional + 4.0 * (conditional - unconditional) factor = torch.linalg.vector_norm( conditional, - dim=(1, 2), + dim=-1, keepdim=True, - ) / torch.linalg.vector_norm(guided, dim=(1, 2), keepdim=True).clamp_min(1e-5) - expected = unpack_latents((guided * factor).to(teacher.calls[0]["output"].dtype), 4, 4) + ) / torch.linalg.vector_norm(guided, dim=-1, keepdim=True) + expected = unpack_latents(guided * factor, 4, 4) assert actual.dtype == torch.bfloat16 torch.testing.assert_close(actual, expected) @@ -263,7 +257,7 @@ def test_teacher_cfg_and_global_norm_rescale_match_mr210_reference() -> None: assert all(call["guidance"] is None for call in teacher.calls) -def test_teacher_cfg_returns_to_low_precision_model_output_dtype() -> None: +def test_teacher_cfg_stays_in_model_output_dtype() -> None: class LowPrecisionTeacher(nn.Module): def __init__(self) -> None: super().__init__() @@ -289,22 +283,13 @@ def forward(self, *, hidden_states, encoder_hidden_states, **kwargs): assert actual.dtype == torch.bfloat16 conditional, unconditional = teacher.outputs - guided_bf16 = conditional + 3.0 * (conditional - unconditional) - conditional_fp32 = conditional.float() - guided_fp32 = guided_bf16.float() - 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) - expected = unpack_latents((guided_fp32 * factor).to(torch.bfloat16), 4, 4) + guided = unconditional + 4.0 * (conditional - unconditional) + factor = torch.linalg.vector_norm(conditional, dim=-1, keepdim=True) / torch.linalg.vector_norm( + guided, dim=-1, keepdim=True + ) + expected = unpack_latents(guided * factor, 4, 4) torch.testing.assert_close(actual, expected, rtol=0, atol=0) - unrounded_guided = conditional_fp32 + 3.0 * (conditional_fp32 - unconditional.float()) - unrounded_factor = torch.linalg.vector_norm( - conditional_fp32, dim=(1, 2), keepdim=True - ) / torch.linalg.vector_norm(unrounded_guided, dim=(1, 2), keepdim=True).clamp_min(1e-5) - unrounded = unpack_latents((unrounded_guided * unrounded_factor).to(torch.bfloat16), 4, 4) - assert not torch.equal(actual, unrounded) - def test_guidance_disabled_teacher_is_one_conditional_call_without_negative_condition() -> None: teacher = _TinyQwenTransformer() @@ -339,22 +324,6 @@ def test_conversion_preserves_requires_grad_mode_and_rejects_conflicts() -> None convert_qwen_image_to_pdd(transformer, _config(grid_size=2)) -def test_forward_substrate_identity_is_exact() -> None: - assert QWEN_IMAGE_PDD_FORWARD_SUBSTRATE_ID == ( - "pdd_qwen_mr210_c8100b1347b278511336dccfc074a461457216ec_" - "qwen_33706683487ba16d133b99b73be27b21164c53335441d77b1dcabbfca970f70e" - ) - assert require_qwen_image_pdd_forward_substrate(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE) == dict( - QWEN_IMAGE_PDD_FORWARD_SUBSTRATE - ) - mismatched = dict(QWEN_IMAGE_PDD_FORWARD_SUBSTRATE) - mismatched["id"] = "canonical-diffusers" - with pytest.raises(ValueError, match="authenticated MR210"): - require_qwen_image_pdd_forward_substrate(mismatched) - with pytest.raises(ValueError, match="authenticated MR210"): - require_qwen_image_pdd_forward_substrate(None) - - def _tiny_diffusers_qwen(): diffusers = pytest.importorskip("diffusers") return diffusers.QwenImageTransformer2DModel( @@ -370,198 +339,161 @@ def _tiny_diffusers_qwen(): ) -def test_adoption_preserves_diffusers_interface_state_keys_and_parameter_identity( - monkeypatch, -) -> None: - qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") - source = _tiny_diffusers_qwen() - source.eval() - source_type = type(source) - source_keys = tuple(source.state_dict()) - source_parameters = tuple(source.parameters()) - source_config = dict(source.config) - monkeypatch.setattr(qwen_pdd, "_require_qwen_source_identity", lambda _model: None) - - adopted = adopt_qwen_image_mr210_forward(source) - - assert adopted is not source - assert isinstance(adopted, source_type) - assert type(source) is source_type - assert tuple(adopted.state_dict()) == source_keys - assert all( - actual is expected - for actual, expected in zip(adopted.parameters(), source_parameters, strict=True) - ) - assert dict(adopted.config) == source_config - assert adopted.device == source.device - assert adopted.dtype == source.dtype - assert adopted.training is False - assert adopt_qwen_image_mr210_forward(adopted) is adopted +def test_conversion_preserves_the_ordinary_diffusers_qwen_root() -> None: + student = _tiny_diffusers_qwen().eval() + root_type = type(student) + config = dict(student.config) + convert_qwen_image_to_pdd(student, _config()) -def test_adoption_rejects_unpinned_qwen_source(monkeypatch) -> None: - qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") - source = _tiny_diffusers_qwen() - monkeypatch.setattr(qwen_pdd, "_sha256_source", lambda _owner: "0" * 64) + assert type(student) is root_type + assert isinstance(student.proj_out, PDDOutputProjection) + assert dict(student.config) == config - with pytest.raises(RuntimeError, match="transformer source"): - adopt_qwen_image_mr210_forward(source) +def test_canonical_qwen_conversion_preserves_every_initialized_head() -> None: + base = _tiny_diffusers_qwen().eval() + student = copy.deepcopy(base) + 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) + mask = torch.tensor([[1, 1, 1], [1, 0, 1]], dtype=torch.long) + model_kwargs = { + "hidden_states": pack_latents(state), + "timestep": time, + "encoder_hidden_states": embeddings, + "encoder_hidden_states_mask": mask, + "img_shapes": build_img_shapes(2, 4, 4), + "guidance": None, + "return_dict": False, + } -def test_adoption_rejects_unpinned_timestep_embedding_and_diffusers_version(monkeypatch) -> None: - diffusers = pytest.importorskip("diffusers") - qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") - source = _tiny_diffusers_qwen() - - def mismatched_embedding_hash(owner): - if owner is type(source): - return QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_qwen_source_sha256"] - return "0" * 64 - - monkeypatch.setattr(qwen_pdd, "_sha256_source", mismatched_embedding_hash) - with pytest.raises(RuntimeError, match="timestep embedding source"): - adopt_qwen_image_mr210_forward(source) - - def pinned_source_hash(owner): - if owner is type(source): - return QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_qwen_source_sha256"] - return QWEN_IMAGE_PDD_FORWARD_SUBSTRATE["diffusers_embeddings_source_sha256"] - - monkeypatch.setattr(qwen_pdd, "_sha256_source", pinned_source_hash) - monkeypatch.setattr(diffusers, "__version__", "0.0.0") - with pytest.raises(RuntimeError, match="Diffusers version"): - adopt_qwen_image_mr210_forward(source) - - -def test_adoption_rejects_every_material_root_invariant(monkeypatch) -> None: - qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") - baseline = _tiny_diffusers_qwen() - monkeypatch.setattr(qwen_pdd, "_require_qwen_source_identity", lambda _model: None) - - def set_config_flag(source, name): - source._internal_dict = type(source._internal_dict)({**dict(source.config), name: True}) - - def check(mutator, match): - source = copy.deepcopy(baseline) - mutator(source) - with pytest.raises((RuntimeError, ValueError), match=match): - adopt_qwen_image_mr210_forward(source) - - check(lambda source: source.add_module("unexpected", nn.Identity()), "root child layout") - check( - lambda source: setattr( - source, - "_modules", - dict(reversed(tuple(source._modules.items()))), - ), - "root child layout", - ) - check( - lambda source: source.register_parameter("root_parameter", nn.Parameter(torch.zeros(1))), - "direct parameters or buffers", + with torch.no_grad(): + expected = unpack_latents(base(**model_kwargs)[0], 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_canonical_qwen_mask_makes_masked_padding_numerically_inert() -> None: + student = _tiny_diffusers_qwen().eval() + 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, 3, 12, generator=generator) + mask = torch.tensor([[1, 1, 1], [1, 0, 1]], dtype=torch.long) + poisoned = encoder_hidden_states.clone() + poisoned[~mask.bool()] = ( + torch.randn( + poisoned[~mask.bool()].shape, + generator=generator, + ) + * 100 ) - check( - lambda source: source.register_buffer("root_buffer", torch.zeros(1)), - "direct parameters or buffers", + with torch.no_grad(): + baseline = adapter.student_all_heads( + student, + state, + time, + condition=(encoder_hidden_states, mask), + ) + actual = adapter.student_all_heads( + student, + state, + time, + condition=(poisoned, mask), + ) + + torch.testing.assert_close(actual, baseline, rtol=0, atol=0) + + +def test_canonical_qwen_teacher_cfg_matches_the_pipeline_formula() -> None: + teacher = _tiny_diffusers_qwen().eval() + 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), + torch.tensor([[1, 1, 0], [1, 0, 1]], dtype=torch.long), ) - check(lambda source: set_config_flag(source, "guidance_embeds"), "guidance embeddings") - check(lambda source: setattr(source, "peft_config", {"active": True}), "PEFT") - check( - lambda source: setattr(source.transformer_blocks[0], "fused_projections", True), - "fused QKV", + negative_condition = ( + torch.randn(2, 2, 12, generator=generator), + torch.tensor([[1, 0], [1, 1]], dtype=torch.long), ) - for name in ("zero_cond_t", "use_additional_t_cond", "use_layer3d_rope"): - check(lambda source, name=name: set_config_flag(source, name), rf"{name}=False") - check(lambda source: source.register_forward_hook(lambda *_args: None), "hooks must be empty") + def direct_packed(current_condition): + embeddings, mask = current_condition + return teacher( + hidden_states=pack_latents(state), + timestep=time, + encoder_hidden_states=embeddings, + encoder_hidden_states_mask=mask, + img_shapes=build_img_shapes(2, 4, 4), + guidance=None, + return_dict=False, + )[0] -def test_adopted_forward_rejects_every_unsupported_input_contract(monkeypatch) -> None: - qwen_pdd = pytest.importorskip("modelopt.torch.fastgen.plugins.qwen_image_pdd") - source = _tiny_diffusers_qwen().eval().to(dtype=torch.bfloat16) - monkeypatch.setattr(qwen_pdd, "_require_qwen_source_identity", lambda _model: None) - adopted = adopt_qwen_image_mr210_forward(source) + with torch.no_grad(): + conditional = direct_packed(condition) + unconditional = direct_packed(negative_condition) + guided = unconditional + 4.0 * (conditional - unconditional) + expected = unpack_latents( + guided + * ( + torch.linalg.vector_norm(conditional, dim=-1, keepdim=True) + / torch.linalg.vector_norm(guided, dim=-1, keepdim=True) + ), + 4, + 4, + ) + actual = adapter.teacher_velocity( + teacher, + state, + time, + condition=condition, + negative_condition=negative_condition, + ) - generator = torch.Generator().manual_seed(20260715) - hidden_states = torch.randn(2, 4, 8, generator=generator).to(torch.bfloat16) - encoder_hidden_states = torch.randn(2, 3, 12, generator=generator).to(torch.bfloat16) - mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long) - encoder_hidden_states[~mask.bool()] = 0 - base = { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "encoder_hidden_states_mask": mask, - "timestep": torch.tensor([0.875, 0.25], dtype=torch.float32), - "img_shapes": [[(1, 2, 2)], [(1, 2, 2)]], - "return_dict": False, - } + torch.testing.assert_close(actual, expected, rtol=0, atol=0) - def condition_for(candidate_mask): - candidate_embeddings = encoder_hidden_states.clone() - candidate_embeddings[~candidate_mask.bool()] = 0 - return candidate_embeddings - - nonbinary_mask = torch.tensor([[1, 1, 1], [1, 2, 0]], dtype=torch.long) - nonprefix_mask = torch.tensor([[1, 1, 1], [1, 0, 1]], dtype=torch.long) - empty_mask = torch.tensor([[1, 1, 1], [0, 0, 0]], dtype=torch.long) - no_full_row_mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long) - nonzero_padding = encoder_hidden_states.clone() - nonzero_padding[1, 1] = 1 - cases = ( - ({"hidden_states": hidden_states.float()}, TypeError, "hidden_states"), - ({"encoder_hidden_states": encoder_hidden_states.float()}, TypeError, "encoder_hidden"), - ({"timestep": base["timestep"].to(torch.bfloat16)}, TypeError, "timestep"), - ( - { - "encoder_hidden_states": condition_for(nonbinary_mask), - "encoder_hidden_states_mask": nonbinary_mask, - }, - ValueError, - "binary prefix masks", - ), - ( - { - "encoder_hidden_states": condition_for(nonprefix_mask), - "encoder_hidden_states_mask": nonprefix_mask, - }, - ValueError, - "binary prefix masks", - ), - ( - { - "encoder_hidden_states": condition_for(empty_mask), - "encoder_hidden_states_mask": empty_mask, - }, - ValueError, - "binary prefix masks", - ), - ( - { - "encoder_hidden_states": condition_for(no_full_row_mask), - "encoder_hidden_states_mask": no_full_row_mask, - }, - ValueError, - "binary prefix masks", - ), - ({"encoder_hidden_states": nonzero_padding}, ValueError, "zero padding"), - ({"max_txt_seq_len": 2}, ValueError, "max_txt_seq_len"), - ({"txt_seq_lens": [3, 1]}, ValueError, "txt_seq_lens"), - ({"guidance": torch.ones(2)}, ValueError, "guidance embeddings"), - ({"attention_kwargs": {"scale": 1.0}}, ValueError, "attention_kwargs"), - ({"controlnet_block_samples": ()}, ValueError, "ControlNet"), - ({"additional_t_cond": torch.ones(2)}, ValueError, "additional time"), + +def test_adapter_accepts_arbitrary_binary_masks_nonzero_padding_and_floating_time() -> None: + student = _TinyQwenTransformer() + config = _config() + convert_qwen_image_to_pdd(student, config) + state, time, condition, _ = _inputs() + embeddings = condition[0].clone() + mask = torch.tensor([[1, 0, 1], [0, 0, 0]], dtype=torch.long) + embeddings[~mask.bool()] = 17 + + actual = QwenImagePDDAdapter(config).student_all_heads( + student, + state, + time.to(torch.bfloat16), + condition=(embeddings, mask), ) - for override, error_type, match in cases: - with pytest.raises(error_type, match=match): - adopted(**(base | override)) + + assert actual.shape == (2, 4, 1, 4, 4) + torch.testing.assert_close(student.calls[0]["encoder_hidden_states_mask"], mask) + assert student.calls[0]["timestep"].dtype == torch.bfloat16 def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> None: with pytest.raises(ValueError, match="num_train_timesteps=None"): QwenImagePDDAdapter(_config().model_copy(update={"num_train_timesteps": 1000})) - with pytest.raises(ValueError, match="guidance_rescale"): - QwenImagePDDAdapter(_config(), guidance_rescale=1.1) - with pytest.raises(ValueError, match="guidance_eps"): - QwenImagePDDAdapter(_config(), guidance_eps=0.0) with pytest.raises(TypeError, match="compute_dtype"): QwenImagePDDAdapter(_config(), compute_dtype=torch.long) @@ -606,6 +538,20 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N condition=condition, guidance=torch.ones(state.shape[0]), ) + with pytest.raises(ValueError, match="zero and one"): + adapter.student_all_heads( + transformer, + state, + time, + condition=(condition[0], torch.tensor([[1, 2, 0], [1, 0, 0]])), + ) + with pytest.raises(ValueError, match="integer/bool"): + adapter.student_all_heads( + transformer, + state, + time, + condition=(condition[0], condition[1].float()), + ) def test_raw_head_reference_uses_independent_linear_outputs() -> None: From 64182de6fd4f10fc15c8a161d5fdd030addaf9d1 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 16 Jul 2026 01:43:11 -0700 Subject: [PATCH 27/45] Match Qwen FSDP checkpointing to FastGen Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/recipe.py | 87 +++++++++- .../fastgen/test_pdd_recipe_setup.py | 162 +++++++++++++++++- 2 files changed, 242 insertions(+), 7 deletions(-) diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 7266618e5e9..cd6db9728e7 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -676,6 +676,70 @@ def _load_unwrapped_transformer( return pipe, student +def _apply_qwen_image_activation_checkpointing(model: nn.Module, *, enabled: bool) -> int: + """Apply FastGen-compatible non-reentrant checkpointing to Qwen-Image blocks.""" + if type(enabled) is not bool: + raise TypeError("enabled must be bool.") + + # Diffusers is an optional example dependency; defer imports until this adapter is used. + from diffusers import QwenImageTransformer2DModel + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + CheckpointWrapper, + checkpoint_wrapper, + ) + + if type(model) is not QwenImageTransformer2DModel: + raise TypeError( + "PDD activation checkpointing requires an ordinary " + f"QwenImageTransformer2DModel, got {type(model).__name__}." + ) + blocks = model.transformer_blocks + if not isinstance(blocks, nn.ModuleList) or not blocks: + raise TypeError("Qwen-Image transformer_blocks must be a non-empty nn.ModuleList.") + wrapped_modules = [ + (index, name or "") + for index, block in enumerate(blocks) + for name, module in block.named_modules() + if isinstance(module, CheckpointWrapper) or hasattr(module, "_checkpoint_wrapped_module") + ] + if wrapped_modules: + raise RuntimeError( + "Qwen-Image has already checkpoint-wrapped modules in transformer blocks: " + f"{wrapped_modules}." + ) + unexpected = [ + (index, type(block).__name__) + for index, block in enumerate(blocks) + if type(block) is not QwenImageTransformerBlock + ] + if unexpected: + raise TypeError(f"Qwen-Image transformer block types changed: {unexpected}.") + + model.disable_gradient_checkpointing() + if model.gradient_checkpointing: + raise RuntimeError("Qwen-Image native gradient checkpointing remained enabled.") + if not enabled: + return 0 + + for index, block in enumerate(tuple(blocks)): + blocks[index] = checkpoint_wrapper( + block, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + ) + for index, block in enumerate(blocks): + if not isinstance(block, CheckpointWrapper): + raise RuntimeError(f"Qwen-Image block {index} was not checkpoint-wrapped.") + if block.checkpoint_impl is not CheckpointImpl.NO_REENTRANT: + raise RuntimeError( + f"Qwen-Image block {index} uses the wrong checkpoint implementation." + ) + if type(block._checkpoint_wrapped_module) is not QwenImageTransformerBlock: + raise RuntimeError(f"Qwen-Image block {index} wrapped an unexpected module type.") + return len(blocks) + + def _stage_and_shard_training_models( student: nn.Module, teacher: nn.Module, @@ -685,6 +749,7 @@ def _stage_and_shard_training_models( *, device: torch.device, fuse_qkv_projections: bool, + activation_checkpointing: bool, ) -> tuple[nn.Module, nn.Module]: """Stage FP32 masters and shard one dense model at a time.""" if fuse_qkv_projections and ( @@ -702,6 +767,11 @@ def _stage_and_shard_training_models( "Qwen fuse_qkv_projections() was accepted but produced no fused attention " "modules in the pinned Diffusers release." ) + student_checkpoint_blocks = _apply_qwen_image_activation_checkpointing( + student, + enabled=activation_checkpointing, + ) + logging.info("Qwen student checkpoint-wrapped blocks: %d", student_checkpoint_blocks) _require_projection_identity(student, projection, projection_identity, stage="student staging") student = manager.parallelize(student) _require_projection_module(student, projection, stage="student FSDP2 parallelization") @@ -711,6 +781,11 @@ def _stage_and_shard_training_models( teacher.to(device=device, dtype=torch.float32) if fuse_qkv_projections: teacher.fuse_qkv_projections() + teacher_checkpoint_blocks = _apply_qwen_image_activation_checkpointing( + teacher, + enabled=activation_checkpointing, + ) + logging.info("Qwen teacher checkpoint-wrapped blocks: %d", teacher_checkpoint_blocks) teacher = manager.parallelize(teacher) return student, teacher @@ -810,7 +885,14 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: from torch.distributed.fsdp import MixedPrecisionPolicy strategy = FSDP2Config( - activation_checkpointing=config.parallel.activation_checkpointing, + # Qwen block checkpointing is applied explicitly before FSDP2 because + # AutoModel's generic language-model attribute policy does not match + # Qwen-Image blocks. + activation_checkpointing=False, + # Match FastGen's per-block PyTorch default instead of AutoModel's + # all-but-last ModuleList optimization. AutoModel still keeps the root + # unresharded after forward. + reshard_after_forward=True, mp_policy=MixedPrecisionPolicy( param_dtype=config.dtype, reduce_dtype=torch.float32, @@ -820,7 +902,7 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: distributed_setup = DistributedSetup.build( strategy=strategy, parallelism_sizes=ParallelismSizes(dp_size=dp_size), - activation_checkpointing=config.parallel.activation_checkpointing, + activation_checkpointing=False, world_size=world_size, ) mesh_context = distributed_setup.mesh_context @@ -837,6 +919,7 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: manager, device=config.device, fuse_qkv_projections=config.fuse_qkv_projections, + activation_checkpointing=config.parallel.activation_checkpointing, ) pipe.transformer = student # Keep the public lifecycle summary stable even though placement, optional QKV fusion, and diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 97bb29217aa..96562560d16 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -26,8 +26,18 @@ import pytest import torch import yaml -from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir +from _test_utils.torch.diffusers_models import ( + create_tiny_qwen_image_pipeline_dir, + get_tiny_qwen_image_transformer, +) +from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock from torch import nn +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + CheckpointWrapper, + checkpoint_wrapper, +) +from torch.distributed.checkpoint import FileSystemReader _REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] _FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" @@ -35,6 +45,7 @@ sys.path.insert(0, str(_FASTGEN_DIR)) from pdd.recipe import ( + _apply_qwen_image_activation_checkpointing, _materialize_zero_step_adamw_state, _projection_identity, _require_fp32_optimizer_storage, @@ -50,6 +61,72 @@ from modelopt.torch.fastgen import PDDLayerSpec, convert_to_pdd_output_projection +def _canonical_parameter_names(model: nn.Module) -> tuple[set[str], set[str]]: + raw_names = {name for name, _ in model.named_parameters()} + canonical_names: set[str] = set() + for raw_name in raw_names: + parts = raw_name.split(".") + if parts[0] == "transformer_blocks": + assert len(parts) > 3 + assert parts[1].isdigit() + assert parts[2] == "_checkpoint_wrapped_module" + assert parts.count("_checkpoint_wrapped_module") == 1 + canonical_name = ".".join((*parts[:2], *parts[3:])) + else: + assert "_checkpoint_wrapped_module" not in parts + canonical_name = raw_name + assert canonical_name not in canonical_names + canonical_names.add(canonical_name) + return raw_names, canonical_names + + +def test_qwen_activation_checkpointing_wraps_every_block_once() -> None: + model = get_tiny_qwen_image_transformer(num_layers=2) + model.enable_gradient_checkpointing() + original_parameter_names = {name for name, _ in model.named_parameters()} + original_state_names = set(model.state_dict()) + + wrapped = _apply_qwen_image_activation_checkpointing(model, enabled=True) + + assert wrapped == 2 + assert model.gradient_checkpointing is False + assert len(model.transformer_blocks) == 2 + for block in model.transformer_blocks: + assert isinstance(block, CheckpointWrapper) + assert block.checkpoint_impl is CheckpointImpl.NO_REENTRANT + assert type(block._checkpoint_wrapped_module) is QwenImageTransformerBlock + raw_parameter_names, canonical_parameter_names = _canonical_parameter_names(model) + assert canonical_parameter_names == original_parameter_names + assert set(model.state_dict()) == original_state_names + assert "proj_out.weight" in raw_parameter_names + assert "proj_out.bias" in raw_parameter_names + + with pytest.raises(RuntimeError, match="already checkpoint-wrapped"): + _apply_qwen_image_activation_checkpointing(model, enabled=True) + + +def test_qwen_activation_checkpointing_disabled_keeps_exact_blocks() -> None: + model = get_tiny_qwen_image_transformer(num_layers=2) + model.enable_gradient_checkpointing() + + wrapped = _apply_qwen_image_activation_checkpointing(model, enabled=False) + + assert wrapped == 0 + assert model.gradient_checkpointing is False + assert all(type(block) is QwenImageTransformerBlock for block in model.transformer_blocks) + + +def test_qwen_activation_checkpointing_rejects_pre_wrapped_input() -> None: + model = get_tiny_qwen_image_transformer(num_layers=2) + model.transformer_blocks[0].attn = checkpoint_wrapper( + model.transformer_blocks[0].attn, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + ) + + with pytest.raises(RuntimeError, match="already checkpoint-wrapped"): + _apply_qwen_image_activation_checkpointing(model, enabled=True) + + def test_zero_step_adamw_state_preserves_first_lazy_update() -> None: torch.manual_seed(7) eager_model = nn.Linear(4, 3) @@ -113,7 +190,7 @@ def test_fp32_optimizer_storage_rejects_low_precision_masters_and_state() -> Non _require_fp32_optimizer_storage(fp32_optimizer) -def test_training_setup_shards_student_before_staging_teacher() -> None: +def test_training_setup_shards_student_before_staging_teacher(monkeypatch) -> None: events: list[str] = [] class TrackedModel(nn.Module): @@ -126,11 +203,25 @@ def to(self, *args, **kwargs): events.append(f"{self.label}.to") return super().to(*args, **kwargs) + def fuse_qkv_projections(self): + events.append(f"{self.label}.fuse") + self.fused_projections = True + class TrackedManager: def parallelize(self, model): events.append(f"{model.label}.parallelize") return model + def apply_checkpointing(model, *, enabled): + assert enabled is True + events.append(f"{model.label}.checkpoint") + return 1 + + monkeypatch.setattr( + "pdd.recipe._apply_qwen_image_activation_checkpointing", + apply_checkpointing, + ) + student = TrackedModel("student", projection=True) teacher = TrackedModel("teacher") projection = convert_to_pdd_output_projection( @@ -146,7 +237,8 @@ def parallelize(self, model): _projection_identity(projection), TrackedManager(), device=torch.device("cpu"), - fuse_qkv_projections=False, + fuse_qkv_projections=True, + activation_checkpointing=True, ) assert staged_student is student @@ -154,18 +246,27 @@ def parallelize(self, model): assert {parameter.dtype for parameter in staged_student.parameters()} == {torch.float32} assert events == [ "student.to", + "student.fuse", + "student.checkpoint", "student.parallelize", "teacher.to", + "teacher.fuse", + "teacher.checkpoint", "teacher.parallelize", ] -def test_training_setup_upcasts_bf16_models_to_fp32_masters() -> None: +def test_training_setup_upcasts_bf16_models_to_fp32_masters(monkeypatch) -> None: class IdentityManager: @staticmethod def parallelize(model): return model + monkeypatch.setattr( + "pdd.recipe._apply_qwen_image_activation_checkpointing", + lambda _model, *, enabled: 0, + ) + student = nn.Module() student.proj_out = nn.Linear(2, 2, dtype=torch.bfloat16) teacher = nn.Linear(2, 2, dtype=torch.bfloat16) @@ -183,6 +284,7 @@ def parallelize(model): IdentityManager(), device=torch.device("cpu"), fuse_qkv_projections=False, + activation_checkpointing=False, ) assert {parameter.dtype for parameter in staged_student.parameters()} == {torch.float32} @@ -563,7 +665,9 @@ def test_payload_hash_verification_mode_must_be_bool(tmp_path) -> None: def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: model_dir = create_tiny_qwen_image_pipeline_dir(tmp_path) initialize_pdd_distributed(backend="gloo", timeout_minutes=1) - config = resolve_pdd_recipe_config(_raw_config(model_dir)) + raw_config = _raw_config(model_dir) + raw_config["fsdp"]["activation_checkpointing"] = True + config = resolve_pdd_recipe_config(raw_config) source = build_pdd_setup(config) @@ -590,6 +694,23 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert policy.reduce_dtype == torch.float32 assert policy.output_dtype == torch.bfloat16 assert policy.cast_forward_inputs is True + assert source.distributed_setup.strategy_config.activation_checkpointing is False + assert source.distributed_setup.strategy_config.reshard_after_forward is True + assert config.parallel.activation_checkpointing is True + for model in (source.student, source.teacher): + assert model.gradient_checkpointing is False + assert len(model.transformer_blocks) == 6 + assert all(isinstance(block, CheckpointWrapper) for block in model.transformer_blocks) + assert all( + block.checkpoint_impl is CheckpointImpl.NO_REENTRANT + for block in model.transformer_blocks + ) + assert all( + type(block._checkpoint_wrapped_module) is QwenImageTransformerBlock + for block in model.transformer_blocks + ) + _, canonical_parameter_names = _canonical_parameter_names(model) + assert canonical_parameter_names == set(model.state_dict()) assert all( base.__module__ != "modelopt.torch.fastgen.plugins.qwen_image_pdd" for base in type(source.student).__mro__ @@ -623,10 +744,25 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: source.optimizer.step() expected_weight = source.projection.weight.detach().clone() expected_exp_avg = source.optimizer.state[source.projection.weight]["exp_avg"].clone() + expected_student_state = { + name: value.detach().clone() for name, value in source.student.state_dict().items() + } + expected_optimizer_state = copy.deepcopy(source.optimizer.state_dict()) assert source.optimizer.state[unused_parameter]["step"].item() == 0 checkpoint_root = tmp_path / "checkpoint" source.checkpointer.save_model(source.student, str(checkpoint_root)) source.checkpointer.save_optimizer(source.optimizer, source.student, str(checkpoint_root)) + model_metadata_keys = set( + FileSystemReader(str(checkpoint_root / "model")).read_metadata().state_dict_metadata + ) + optimizer_metadata_keys = set( + FileSystemReader(str(checkpoint_root / "optim")).read_metadata().state_dict_metadata + ) + assert model_metadata_keys + assert optimizer_metadata_keys + assert all("_checkpoint_wrapped_module" not in key for key in model_metadata_keys) + assert all("_checkpoint_wrapped_module" not in key for key in optimizer_metadata_keys) + assert any(key.endswith("proj_out.weight") for key in model_metadata_keys) export_setup = build_pdd_export_setup(config) assert export_setup.lifecycle == ( @@ -648,6 +784,10 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: export_setup.student.state_dict()["proj_out.weight"], expected_weight, ) + export_state = export_setup.student.state_dict() + assert export_state.keys() == expected_student_state.keys() + for name, expected in expected_student_state.items(): + torch.testing.assert_close(export_state[name], expected, rtol=0, atol=0) destination = build_pdd_setup(config) assert destination.metadata == source.metadata @@ -667,6 +807,18 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert destination.student.get_submodule("proj_out") is destination_projection assert id(destination_projection.weight) == destination_weight_id torch.testing.assert_close(destination_projection.weight, expected_weight) + destination_state = destination.student.state_dict() + assert destination_state.keys() == expected_student_state.keys() + for name, expected in expected_student_state.items(): + torch.testing.assert_close(destination_state[name], expected, rtol=0, atol=0) + destination_optimizer_state = destination.optimizer.state_dict() + assert destination_optimizer_state["param_groups"] == expected_optimizer_state["param_groups"] + assert destination_optimizer_state["state"].keys() == expected_optimizer_state["state"].keys() + for parameter_index, expected_state in expected_optimizer_state["state"].items(): + actual_state = destination_optimizer_state["state"][parameter_index] + assert actual_state.keys() == expected_state.keys() + for name, expected in expected_state.items(): + torch.testing.assert_close(actual_state[name], expected, rtol=0, atol=0) torch.testing.assert_close( destination.optimizer.state[destination_projection.weight]["exp_avg"], expected_exp_avg, From d347ef2325596c52328c57f6a5fa8dc5ea5d4568 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 16 Jul 2026 08:46:10 -0700 Subject: [PATCH 28/45] Fix ragged distributed PDD validation Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/training.py | 23 ++----------------- .../pdd_validation_oracle_distributed.py | 21 +++++++++++++++++ 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py index b4ddf7d7d97..493061707cf 100644 --- a/examples/diffusers/fastgen/pdd/training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -852,7 +852,6 @@ def run_pdd_validation( local_error: BaseException | None = None selected: list[PDDValidationAssignment] = [] valid_mask: tuple[bool, ...] = () - signature: Any = None try: if not isinstance(batch, PreparedPDDBatch): raise TypeError("validation batches must contain PreparedPDDBatch values.") @@ -886,30 +885,12 @@ def run_pdd_validation( template.k, ) ) - signature = ( - tuple(batch.data.shape), - tuple(batch.condition[0].shape), - tuple(batch.condition[1].shape), - None - if batch.negative_condition is None - else ( - tuple(batch.negative_condition[0].shape), - tuple(batch.negative_condition[1].shape), - ), - ) except BaseException as error: local_error = error _raise_collective_validation_error(local_error, context="batch preflight") assert isinstance(batch, PreparedPDDBatch) - - if distributed: - signatures: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(signatures, signature) - if any(item != signatures[0] for item in signatures[1:]): - raise RuntimeError( - "distributed PDD validation requires the same padded batch shape " - "on every rank." - ) + # Inputs remain rank-local. Equal batch counts above preserve collective + # ordering, while prompt sequence padding may legitimately differ by rank. noise = torch.stack( [ pdd_validation_noise( diff --git a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py index 6c37eed91ba..e197901f775 100644 --- a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py @@ -91,6 +91,27 @@ def main() -> None: assert abs(distributed.mean_loss - baseline.mean_loss) <= 1e-12 assert distributed.ordered_id_sha256 == baseline.ordered_id_sha256 + # Text conditions are padded only to each rank-local batch maximum in the + # real Qwen loader. Data-parallel validation must therefore accept + # different sequence lengths while preserving the same logical result. + ragged_conditions = [ + dataclasses.replace( + batch, + condition=( + torch.zeros((len(batch.sample_ids), rank + 1, 1)), + torch.ones((len(batch.sample_ids), rank + 1), dtype=torch.long), + ), + ) + for batch in local_batches + ] + ragged = run_pdd_validation( + lifecycle.pipeline, + ragged_conditions, + assignments, + validation_seed=91, + ) + assert ragged == distributed + invalid_mask = list(local_batches) if rank == 0: invalid_mask[0] = dataclasses.replace(invalid_mask[0], valid_mask=()) From d0eb0fc61f25c0e2af8f45416c5620e2fc5423b1 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 16 Jul 2026 11:32:32 -0700 Subject: [PATCH 29/45] Match Qwen PDD execution strictly to FastGen MR210 Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/README.md | 19 +- examples/diffusers/fastgen/pdd/checkpoint.py | 16 +- examples/diffusers/fastgen/pdd/export.py | 11 +- .../fastgen/pdd/export_qwen_image.py | 6 + .../fastgen/pdd/inference_qwen_image.py | 12 +- examples/diffusers/fastgen/pdd/recipe.py | 26 +- .../torch/fastgen/plugins/qwen_image_pdd.py | 234 +++++- .../pdd_checkpoint_failure_distributed.py | 14 +- .../fastgen/pdd_export_distributed.py | 20 +- .../fastgen/pdd_mr210_fsdp_distributed.py | 402 +++++++++++ .../fastgen/test_pdd_inference_checkpoint.py | 56 +- .../fastgen/test_pdd_recipe_setup.py | 31 +- .../fastgen/test_pdd_training_lifecycle.py | 42 +- .../fastgen/test_qwen_image_pdd_plugin.py | 682 ++++++++++++++++-- 14 files changed, 1468 insertions(+), 103 deletions(-) create mode 100644 tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index 2c7d6d33727..da496082a47 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -19,12 +19,13 @@ The provided schedules use the 128-interval grid as follows: Install the shared requirements from the repository root, then launch with released AutoModel APIs. No AutoModel, Diffusers, or Qwen source changes are required. -The example uses the ordinary Diffusers Qwen transformer without replacing or monkeypatching its -forward. Diffusers owns Qwen timestep conversion and converts each text mask into the joint -text/image attention mask, so padded text tokens do not participate as attention keys. ModelOpt -owns only PDD projection conversion, latent packing, condition validation, and packed per-token -classifier-free guidance using the Qwen pipeline formula. Guidance-embedded models and PEFT remain -outside this first example. +The example loads the ordinary Diffusers Qwen transformer, then binds a ModelOpt-owned forward on +that same model instance to reproduce FastGen MR210's executed Qwen path. Normalized time remains +FP32 through `time_text_embed`; the text mask is passed to Qwen blocks as MR210 does instead of being +converted into Diffusers' canonical joint mask. Teacher classifier-free guidance is rounded in the +BF16 model-output dtype, globally norm-rescaled in FP32 with a `1e-5` denominator floor, and cast +back to BF16 before the PDD loss. No AutoModel, Diffusers, or Qwen source is edited. +Guidance-embedded models, PEFT, and QKV fusion remain outside this first example. ```bash pip install -r examples/diffusers/fastgen/requirements.txt @@ -55,9 +56,9 @@ default per-rank batch gives global batch size 256; other GPU topologies must se `256 / world_size` because this recipe does not use gradient accumulation. Checkpoints include the student, optimizer, scheduler, RNG, trainer, and exact replayable sampler state needed to resume the next committed batch. FP32 master parameters and Adam state are sharded -while forward/backward uses the configured model dtype and gradient reduction remains FP32. The -adapter casts packed image/text inputs to that compute dtype; the ordinary Diffusers Qwen forward -owns timestep conversion. +while forward/backward uses BF16 model parameters and outputs and gradient reduction remains FP32. +The adapter casts packed image/text inputs to BF16 while preserving FP32 normalized time at the +FSDP root and Qwen time embedder. Start with a one-node smoke and scale only after it passes; project training runs are capped at 16 nodes. Checkpointed training, export, and inference require the remote model ID and exact lowercase diff --git a/examples/diffusers/fastgen/pdd/checkpoint.py b/examples/diffusers/fastgen/pdd/checkpoint.py index eb04fbdf235..4fbf151d693 100644 --- a/examples/diffusers/fastgen/pdd/checkpoint.py +++ b/examples/diffusers/fastgen/pdd/checkpoint.py @@ -31,11 +31,12 @@ import torch.distributed as dist from modelopt.torch.fastgen import PDDMetadata +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION if TYPE_CHECKING: from collections.abc import Sequence -_CHECKPOINT_SCHEMA_VERSION = 4 +_CHECKPOINT_SCHEMA_VERSION = 5 _COMPLETE_SCHEMA_VERSION = 1 @@ -57,6 +58,13 @@ def _require_sha256(value: Any, *, name: str) -> str: return value.lower() +def _require_qwen_image_execution(identity: Any) -> None: + if not isinstance(identity, Mapping) or identity.get("qwen_image") != { + "execution": QWEN_IMAGE_PDD_EXECUTION + }: + raise RuntimeError("PDD checkpoint has an incompatible Qwen execution identity.") + + def _rank() -> int: return dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 @@ -157,6 +165,7 @@ def _read_json(path: Path) -> dict[str, Any]: def build_pdd_checkpoint_identity( *, + qwen_image_execution: str, metadata: PDDMetadata, model_id: str, model_revision: str | None, @@ -177,6 +186,8 @@ def build_pdd_checkpoint_identity( scheduler: Any, ) -> dict[str, Any]: """Build the strict, path-independent compatibility identity for PDD resume.""" + if qwen_image_execution != QWEN_IMAGE_PDD_EXECUTION: + raise ValueError("qwen_image_execution must identify the bound FastGen MR210 forward.") if not isinstance(metadata, PDDMetadata): raise TypeError("metadata must be PDDMetadata.") if not isinstance(model_id, str) or not model_id: @@ -215,6 +226,7 @@ def build_pdd_checkpoint_identity( raise TypeError("PDD checkpoint identity requires the stock torch.optim.AdamW optimizer.") return { "schema_version": _CHECKPOINT_SCHEMA_VERSION, + "qwen_image": {"execution": qwen_image_execution}, "model": {"id": model_id, "revision": model_revision, "dtype": dtype}, "pdd_metadata": metadata.to_dict(), "guidance": {"scale": None if guidance_scale is None else float(guidance_scale)}, @@ -370,6 +382,7 @@ def validate_pdd_training_checkpoint( if expected_identity is not None and manifest["identity"] != expected_identity: raise RuntimeError("PDD checkpoint identity does not match the current run.") identity = manifest["identity"] + _require_qwen_image_execution(identity) topology = identity.get("topology") if isinstance(identity, Mapping) else None world_size = topology.get("world_size") if isinstance(topology, Mapping) else None if type(world_size) is not int or world_size < 1: @@ -542,6 +555,7 @@ def __init__( self.identity = json.loads(json.dumps(identity, sort_keys=True)) if self.identity.get("schema_version") != _CHECKPOINT_SCHEMA_VERSION: raise ValueError("PDD checkpoint identity has an unsupported schema version.") + _require_qwen_image_execution(self.identity) topology = self.identity.get("topology") if not isinstance(topology, dict) or topology.get("world_size") != _world_size(): raise ValueError("PDD checkpoint identity world size does not match the process group.") diff --git a/examples/diffusers/fastgen/pdd/export.py b/examples/diffusers/fastgen/pdd/export.py index 22bb6522b9f..c41c0c68648 100644 --- a/examples/diffusers/fastgen/pdd/export.py +++ b/examples/diffusers/fastgen/pdd/export.py @@ -31,7 +31,10 @@ from safetensors.torch import save_file from modelopt.torch.fastgen import PDDConfig, PDDMetadata -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_LAYER_SPEC +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QWEN_IMAGE_PDD_EXECUTION, + QWEN_IMAGE_PDD_LAYER_SPEC, +) from .artifacts import ( load_canonical_json, @@ -41,7 +44,7 @@ write_canonical_json, ) -_EXPORT_SCHEMA_VERSION = 3 +_EXPORT_SCHEMA_VERSION = 4 _COMPLETE_SCHEMA_VERSION = 1 _EXPORT_FORMAT = "modelopt-pdd-safetensors" _CONFIG_FILE = "config.json" @@ -171,6 +174,7 @@ def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> di "guidance", "model", "pdd_metadata", + "qwen_image", "topology", } missing = sorted(required.difference(identity)) @@ -178,6 +182,9 @@ def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> di raise ValueError(f"PDD export identity is missing keys: {missing}.") if identity["pdd_metadata"] != metadata.to_dict(): raise ValueError("PDD export metadata does not match the checkpoint identity.") + _require_exact_mapping(identity["qwen_image"], {"execution"}, name="identity.qwen_image") + if identity["qwen_image"]["execution"] != QWEN_IMAGE_PDD_EXECUTION: + raise ValueError("PDD export has an incompatible Qwen execution identity.") model = _require_exact_mapping( identity["model"], {"id", "revision", "dtype"}, name="identity.model" ) diff --git a/examples/diffusers/fastgen/pdd/export_qwen_image.py b/examples/diffusers/fastgen/pdd/export_qwen_image.py index 43ec954be34..e77bc0331e8 100644 --- a/examples/diffusers/fastgen/pdd/export_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/export_qwen_image.py @@ -161,6 +161,7 @@ def _collective_publication_preflight(output_dir: Path) -> None: def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: from modelopt.torch.fastgen import PDDMetadata + from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION identity = manifest.get("identity") if not isinstance(identity, Mapping): @@ -170,6 +171,8 @@ def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, raise RuntimeError("PDD checkpoint has no PDD metadata mapping.") if PDDMetadata.from_dict(pdd_metadata) != setup.metadata: raise RuntimeError("PDD checkpoint metadata does not match the configured student.") + if identity.get("qwen_image") != {"execution": QWEN_IMAGE_PDD_EXECUTION}: + raise RuntimeError("PDD checkpoint Qwen execution identity does not match MR210.") if identity.get("model") != { "id": config.model_id, "revision": config.model_revision, @@ -195,7 +198,10 @@ def _collective_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[s def _checkpoint_selector_identity(config: Any, setup: Any) -> dict[str, Any]: + from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION + return { + "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, "model": { "id": config.model_id, "revision": config.model_revision, diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index 891b76a60e4..ea9dca44ea9 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -72,9 +72,13 @@ def _dtype_from_name(name: Any) -> torch.dtype: def _model_identity(descriptor: Any) -> Mapping[str, Any]: + from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION + identity = descriptor.manifest.get("identity") if not isinstance(identity, Mapping): raise RuntimeError("PDD export has no identity mapping.") + if identity.get("qwen_image") != {"execution": QWEN_IMAGE_PDD_EXECUTION}: + raise RuntimeError("PDD export has an incompatible Qwen execution identity.") model = identity.get("model") if not isinstance(model, Mapping) or set(model) != {"id", "revision", "dtype"}: raise RuntimeError("PDD export model identity is malformed.") @@ -117,7 +121,10 @@ def build_pdd_student(export_dir: str | Path) -> tuple[nn.Module, Any, torch.dty """Reconstruct and strictly load the converted Qwen student on CPU.""" from diffusers import QwenImageTransformer2DModel - from modelopt.torch.fastgen.plugins.qwen_image_pdd import convert_qwen_image_to_pdd + from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + adopt_qwen_image_mr210_forward, + convert_qwen_image_to_pdd, + ) from pdd.export import inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata descriptor = inspect_pdd_export(export_dir) @@ -126,6 +133,7 @@ def build_pdd_student(export_dir: str | Path) -> tuple[nn.Module, Any, torch.dty student = QwenImageTransformer2DModel.from_config(dict(descriptor.transformer_config)) metadata = descriptor.metadata _validate_qwen_projection(student, metadata) + student = adopt_qwen_image_mr210_forward(student) config = pdd_config_from_metadata(metadata, blocks=metadata.inference_blocks) convert_qwen_image_to_pdd(student, config) descriptor = load_pdd_export_into_model(export_dir, student) @@ -254,6 +262,8 @@ def main() -> None: torch_dtype=dtype, use_safetensors=True, ) + if pipe.transformer is not student: + raise RuntimeError("Qwen pipeline did not retain the adopted PDD transformer.") pipe.to(device) config = pdd_config_from_metadata(descriptor.metadata, schedule=args.schedule) adapter = QwenImagePDDAdapter(config, compute_dtype=dtype) diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index cd6db9728e7..f9cefe6d67d 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -34,7 +34,9 @@ from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( QwenImagePDDAdapter, + adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, + require_qwen_image_mr210_forward, ) from .checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity @@ -471,6 +473,8 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: model.get("fuse_qkv_projections", False), name="model.fuse_qkv_projections", ) + if fuse_qkv_projections: + raise ValueError("Qwen MR210 PDD does not support QKV fusion.") seed = _require_int_at_least(raw.get("seed", 42), name="seed", minimum=0) max_steps = _require_int_at_least( @@ -555,6 +559,8 @@ def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: ) dtype = _resolve_dtype(model.get("torch_dtype", "bfloat16")) + if dtype != torch.bfloat16: + raise ValueError("Qwen MR210 PDD requires model.torch_dtype='bfloat16'.") return PDDRecipeConfig( model_id=model_id, @@ -856,7 +862,9 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: ) from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager - pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + student = adopt_qwen_image_mr210_forward(loaded_transformer) + pipe.transformer = student teacher = copy.deepcopy(student).eval().requires_grad_(False) lifecycle.append("load/select") @@ -894,9 +902,10 @@ def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: # unresharded after forward. reshard_after_forward=True, mp_policy=MixedPrecisionPolicy( - param_dtype=config.dtype, + param_dtype=torch.bfloat16, reduce_dtype=torch.float32, - output_dtype=config.dtype, + output_dtype=torch.bfloat16, + cast_forward_inputs=False, ), ) distributed_setup = DistributedSetup.build( @@ -1007,7 +1016,9 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager lifecycle = ["load/select"] - pipe, student = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) + student = adopt_qwen_image_mr210_forward(loaded_transformer) + pipe.transformer = student raw_transformer_config = getattr(student, "config", None) to_dict = getattr(raw_transformer_config, "to_dict", None) if callable(to_dict): @@ -1046,9 +1057,10 @@ def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: strategy = FSDP2Config( activation_checkpointing=False, mp_policy=MixedPrecisionPolicy( - param_dtype=config.dtype, + param_dtype=torch.bfloat16, reduce_dtype=torch.float32, - output_dtype=config.dtype, + output_dtype=torch.bfloat16, + cast_forward_inputs=False, ), ) distributed_setup = DistributedSetup.build( @@ -1192,6 +1204,7 @@ def setup(self) -> None: ) self.setup_artifacts = build_pdd_setup(config) + qwen_image_execution = require_qwen_image_mr210_forward(self.setup_artifacts.student) self.expected_latent_channels, self.expected_condition_features = ( self._resolve_transformer_dimensions(self.setup_artifacts.student) ) @@ -1217,6 +1230,7 @@ def setup(self) -> None: raise ValueError("PDD v1 requires exactly one microbatch per optimizer update.") identity = build_pdd_checkpoint_identity( + qwen_image_execution=qwen_image_execution, metadata=self.setup_artifacts.metadata, model_id=config.model_id, model_revision=config.model_revision, diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index a9464c76adc..7b9caf59b93 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -18,6 +18,7 @@ from __future__ import annotations import math +import types from collections.abc import Mapping from typing import Any @@ -29,11 +30,16 @@ from .qwen_image import build_img_shapes, pack_latents, unpack_latents __all__ = [ + "QWEN_IMAGE_PDD_EXECUTION", "QWEN_IMAGE_PDD_LAYER_SPEC", "QwenImagePDDAdapter", + "adopt_qwen_image_mr210_forward", "convert_qwen_image_to_pdd", + "require_qwen_image_mr210_forward", ] +QWEN_IMAGE_PDD_EXECUTION = "fastgen_mr210" + QWEN_IMAGE_PDD_LAYER_SPEC = PDDLayerSpec( projection_path="transformer.proj_out", head_layout="channel_major", @@ -54,6 +60,18 @@ "txt_seq_lens", } +_QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE = "_modelopt_qwen_image_pdd_execution" +_QWEN_IMAGE_MR210_CHILDREN = ( + "pos_embed", + "time_text_embed", + "txt_norm", + "img_in", + "txt_in", + "transformer_blocks", + "norm_out", + "proj_out", +) + def _require_binary_mask(mask: torch.Tensor, *, name: str) -> None: mask_int = mask.to(torch.int64) @@ -68,6 +86,186 @@ def _config_guidance_embeds(transformer: nn.Module) -> bool: 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_mr210_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, + txt_seq_lens: list[int] | None = None, + guidance: torch.Tensor | None = None, + max_txt_seq_len: int | 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, +) -> Any: + """Run the regular-output Qwen forward used by FastGen MR210.""" + if hidden_states.ndim != 3 or hidden_states.dtype != torch.bfloat16: + raise TypeError("Qwen MR210 hidden_states must be packed BF16 [B, P, C].") + if ( + not isinstance(encoder_hidden_states, torch.Tensor) + or encoder_hidden_states.ndim != 3 + or encoder_hidden_states.dtype != torch.bfloat16 + ): + raise TypeError("Qwen MR210 encoder_hidden_states must be BF16 [B, S, D].") + if not isinstance(encoder_hidden_states_mask, torch.Tensor): + raise TypeError("Qwen MR210 requires encoder_hidden_states_mask.") + if not isinstance(timestep, torch.Tensor) or timestep.dtype != torch.float32: + raise TypeError("Qwen MR210 timestep must remain FP32 at transformer entry.") + batch_size = hidden_states.shape[0] + if encoder_hidden_states.shape[0] != batch_size: + raise ValueError("Qwen MR210 image and text batch sizes must match.") + if timestep.shape != (batch_size,): + raise ValueError("Qwen MR210 timestep must contain one value per batch item.") + if img_shapes is None or len(img_shapes) != batch_size: + raise ValueError("Qwen MR210 img_shapes must contain one entry per batch item.") + if attention_kwargs: + raise ValueError("Qwen MR210 PDD does not support nonempty attention_kwargs.") + if guidance is not None: + raise ValueError("Qwen MR210 PDD does not support transformer guidance embeddings.") + if controlnet_block_samples is not None: + raise ValueError("Qwen MR210 PDD does not support ControlNet block samples.") + if additional_t_cond is not None: + raise ValueError("Qwen MR210 PDD does not support additional timestep conditioning.") + if type(return_dict) is not bool: + raise TypeError("Qwen MR210 return_dict must be a bool.") + if encoder_hidden_states_mask.ndim != 2 or tuple(encoder_hidden_states_mask.shape) != tuple( + encoder_hidden_states.shape[:2] + ): + raise ValueError("Qwen MR210 mask must match the text batch and sequence dimensions.") + if encoder_hidden_states_mask.device != encoder_hidden_states.device: + raise ValueError("Qwen MR210 mask and text embeddings must share a device.") + if ( + encoder_hidden_states_mask.dtype.is_floating_point + or encoder_hidden_states_mask.dtype.is_complex + ): + raise TypeError("Qwen MR210 mask must use an integer or boolean dtype.") + _require_binary_mask(encoder_hidden_states_mask, name="Qwen MR210") + expected_max_txt_seq_len = int( + encoder_hidden_states_mask.sum(dim=1).max().to(torch.int32).item() + ) + if txt_seq_lens is not None: + expected_txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).to(torch.int32).tolist() + if txt_seq_lens != expected_txt_seq_lens: + raise ValueError("Qwen MR210 txt_seq_lens must equal the valid mask lengths.") + if max_txt_seq_len is None: + max_txt_seq_len = expected_max_txt_seq_len + elif max_txt_seq_len != expected_max_txt_seq_len: + raise ValueError("Qwen MR210 max_txt_seq_len must equal the maximum valid mask length.") + + hidden_states = self.img_in(hidden_states) + encoder_hidden_states = self.txt_in(self.txt_norm(encoder_hidden_states)) + if timestep.dtype != torch.float32: + raise RuntimeError("Qwen MR210 timestep was rounded before time_text_embed.") + 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, + joint_attention_kwargs=attention_kwargs, + ) + + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + if not return_dict: + return (output,) + + from diffusers.models.modeling_outputs import Transformer2DModelOutput + + return Transformer2DModelOutput(sample=output) + + +def _is_qwen_image_mr210_forward(model: nn.Module) -> bool: + forward = model.__dict__.get("forward") + return ( + isinstance(forward, types.MethodType) + and forward.__func__ is _qwen_image_mr210_forward + and forward.__self__ is model + and getattr(model, _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE, None) == QWEN_IMAGE_PDD_EXECUTION + ) + + +def require_qwen_image_mr210_forward(model: nn.Module) -> str: + """Require and return the semantic label for the exact bound MR210 forward.""" + if not isinstance(model, nn.Module) or not _is_qwen_image_mr210_forward(model): + raise RuntimeError("Qwen-Image PDD requires the bound FastGen MR210 forward execution.") + return QWEN_IMAGE_PDD_EXECUTION + + +def adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: + """Bind FastGen MR210's regular Qwen forward to the loaded root in place.""" + if not isinstance(transformer, nn.Module): + raise TypeError(f"transformer must be nn.Module, got {type(transformer).__name__}.") + if _is_qwen_image_mr210_forward(transformer): + return transformer + + # Diffusers is an optional dependency used only by the Qwen example. + from diffusers import QwenImageTransformer2DModel + + if type(transformer) is not QwenImageTransformer2DModel: + raise TypeError( + "MR210 forward adoption requires the supported 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.") + missing = [ + name + for name in _QWEN_IMAGE_MR210_CHILDREN + if not isinstance(getattr(transformer, name, None), nn.Module) + ] + if missing: + raise RuntimeError(f"Qwen root is missing required MR210 modules: {missing}.") + if ( + not isinstance(transformer.transformer_blocks, nn.ModuleList) + or not transformer.transformer_blocks + ): + raise RuntimeError("Qwen MR210 requires a nonempty transformer_blocks ModuleList.") + if _config_guidance_embeds(transformer): + raise ValueError("Qwen MR210 PDD does not support transformer guidance embeddings.") + if getattr(transformer, "peft_config", None): + raise ValueError("Qwen MR210 PDD does not support active PEFT adapters.") + if any(getattr(module, "fused_projections", False) for module in transformer.modules()): + raise ValueError("Qwen MR210 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 MR210 PDD requires {name}=False.") + + transformer.forward = types.MethodType(_qwen_image_mr210_forward, transformer) + setattr(transformer, _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE, QWEN_IMAGE_PDD_EXECUTION) + require_qwen_image_mr210_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__}.") @@ -235,6 +433,7 @@ def _prepare_call( condition_name: str, ) -> tuple[torch.Tensor, torch.Tensor]: self._validate_state_and_time(state, time) + require_qwen_image_mr210_forward(model) if _config_guidance_embeds(model): raise ValueError("Qwen-Image PDD does not support transformer guidance embeddings.") encoder_hidden_states, attention_mask = self._parse_condition( @@ -268,15 +467,20 @@ def _call_packed( batch_size, _, height, width = state.shape model_dtype = self._model_dtype(model, state.dtype) + if model_dtype != torch.bfloat16: + raise TypeError("Qwen MR210 PDD execution requires BF16 compute.") + if time.dtype != torch.float32: + raise TypeError("Qwen MR210 PDD execution requires FP32 time.") packed_state = pack_latents(state).to(model_dtype) encoder_hidden_states = encoder_hidden_states.to(model_dtype) + max_txt_seq_len = int(attention_mask.sum(dim=1).max().to(torch.int32).item()) output = model( hidden_states=packed_state, timestep=time, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=attention_mask, img_shapes=build_img_shapes(batch_size, height, width), - guidance=None, + max_txt_seq_len=max_txt_seq_len, return_dict=False, **model_kwargs, ) @@ -357,7 +561,7 @@ def student_all_heads( condition: Any = None, **model_kwargs: Any, ) -> torch.Tensor: - """Return unpacked canonical interval velocities from one Qwen call.""" + """Return unpacked PDD interval velocities from one Qwen call.""" self._projection(model, self.config.grid_size) packed = self._call_packed( model, @@ -455,7 +659,25 @@ def teacher_velocity( f"{tuple(conditional.shape)} and {tuple(unconditional.shape)}." ) - guided = unconditional + float(guidance_scale) * (conditional - unconditional) - conditional_norm = torch.linalg.vector_norm(conditional, dim=-1, keepdim=True) - guided_norm = torch.linalg.vector_norm(guided, dim=-1, keepdim=True) - return self._unpack_single(guided * (conditional_norm / guided_norm), state) + # MR210 unpacks each Qwen prediction before guidance. Keep that + # operation order: the FP32 global reduction order over NCHW is not + # guaranteed to match an algebraically equivalent packed reduction. + conditional_unpacked = self._unpack_single(conditional, state) + unconditional_unpacked = self._unpack_single(unconditional, state) + guided_low_precision = conditional_unpacked + (float(guidance_scale) - 1.0) * ( + conditional_unpacked - unconditional_unpacked + ) + conditional_fp32 = conditional_unpacked.to(torch.float32) + guided_fp32 = guided_low_precision.to(torch.float32) + norm_dims = tuple(range(1, conditional_fp32.ndim)) + conditional_norm = torch.linalg.vector_norm( + conditional_fp32, + dim=norm_dims, + keepdim=True, + ) + guided_norm = torch.linalg.vector_norm( + guided_fp32, + dim=norm_dims, + keepdim=True, + ).clamp_min(1e-5) + return (guided_fp32 * (conditional_norm / guided_norm)).to(conditional.dtype) diff --git a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py index 8c80fa52685..42d5ad394b3 100644 --- a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py @@ -36,6 +36,8 @@ import pdd.checkpoint as pdd_checkpoint_module from pdd.checkpoint import PDDCheckpointManager +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION + class _State: def state_dict(self): @@ -141,7 +143,11 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 4, "topology": {"world_size": 2}}, + identity={ + "schema_version": 5, + "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, + "topology": {"world_size": 2}, + }, ) initial.save() trainer.completed_steps = 2 @@ -156,7 +162,11 @@ def _run_failure(root: pathlib.Path, stage: str) -> None: trainer=trainer, sampler=_Sampler(), rng=_State(), - identity={"schema_version": 4, "topology": {"world_size": 2}}, + identity={ + "schema_version": 5, + "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, + "topology": {"world_size": 2}, + }, ) message = None try: diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py index 79cd5a072e5..f1547c0491d 100644 --- a/tests/examples/diffusers/fastgen/pdd_export_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_export_distributed.py @@ -38,6 +38,11 @@ from pdd.inference_qwen_image import build_pdd_student from pdd.recipe import build_pdd_export_setup, resolve_pdd_recipe_config +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QWEN_IMAGE_PDD_EXECUTION, + require_qwen_image_mr210_forward, +) + def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: return { @@ -133,7 +138,8 @@ def main() -> None: tensor.numel() * tensor.element_size() for tensor in actual.values() ) identity = { - "schema_version": 4, + "schema_version": 5, + "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, "model": { "id": "Qwen/Qwen-Image", "revision": "3" * 40, @@ -158,12 +164,20 @@ def main() -> None: ) descriptor = inspect_pdd_export(output) assert descriptor.metadata == destination.metadata - restored, restored_descriptor, _dtype = build_pdd_student(output) + restored, restored_descriptor, dtype = build_pdd_student(output) + require_qwen_image_mr210_forward(restored) + assert dtype == torch.bfloat16 assert restored_descriptor.metadata == destination.metadata restored_state = restored.state_dict() assert restored_state.keys() == actual.keys() for key in actual: - torch.testing.assert_close(restored_state[key], actual[key], rtol=0, atol=0) + assert restored_state[key].dtype == torch.bfloat16 + torch.testing.assert_close( + restored_state[key], + actual[key].to(torch.bfloat16), + rtol=0, + atol=0, + ) status = {"ok": True} except BaseException as error: status = {"ok": False, "error": f"{type(error).__name__}: {error}"} diff --git a/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py b/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py new file mode 100644 index 00000000000..e96ff52a70a --- /dev/null +++ b/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Two-rank differential for the Qwen MR210 PDD recipe under FSDP2.""" + +from __future__ import annotations + +import argparse +import copy +import json +import os +import pathlib +import shutil +import sys +import tempfile + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor +from torch.func import functional_call + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +for path in (_REPO_ROOT, _REPO_ROOT / "tests", _FASTGEN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir +from diffusers import QwenImageTransformer2DModel +from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock +from pdd.recipe import build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + CheckpointWrapper, +) + +from modelopt.torch.fastgen import PDDPipeline +from modelopt.torch.fastgen.plugins.qwen_image import build_img_shapes, pack_latents +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QwenImagePDDAdapter, + adopt_qwen_image_mr210_forward, + convert_qwen_image_to_pdd, +) + + +class _FSDPBoundaryReferenceAdapter(QwenImagePDDAdapter): + """Emulate FSDP's BF16 parameter views backed by FP32 masters.""" + + def _call_packed( + self, + model, + state, + time, + condition, + model_kwargs, + *, + condition_name, + ): + encoder_hidden_states, attention_mask = self._prepare_call( + model, + state, + time, + condition, + model_kwargs, + condition_name=condition_name, + ) + batch_size, _, height, width = state.shape + max_txt_seq_len = int(attention_mask.sum(dim=1).max().to(torch.int32).item()) + parameters = { + name: parameter.to(torch.bfloat16) if parameter.dtype.is_floating_point else parameter + for name, parameter in model.named_parameters() + } + output = functional_call( + model, + parameters, + (), + { + "hidden_states": pack_latents(state).to(torch.bfloat16), + "timestep": time, + "encoder_hidden_states": encoder_hidden_states.to(torch.bfloat16), + "encoder_hidden_states_mask": attention_mask, + "img_shapes": build_img_shapes(batch_size, height, width), + "max_txt_seq_len": max_txt_seq_len, + "return_dict": False, + **model_kwargs, + }, + strict=False, + ) + return self._extract_packed_output(output) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--activation-checkpointing", action="store_true") + return parser.parse_args() + + +def _raw_config( + model_dir: pathlib.Path, + checkpoint_dir: pathlib.Path, + *, + activation_checkpointing: bool, +) -> dict: + return { + "model": { + "pretrained_model_name_or_path": str(model_dir), + "torch_dtype": "bfloat16", + "device": "cuda", + "transformer_engine_linear": False, + "peft": None, + "guidance_embeds": False, + "fuse_qkv_projections": False, + }, + "pdd": { + "pred_type": "flow", + "num_train_timesteps": None, + "guidance_scale": 4.0, + "student_sample_steps": 2, + "student_sample_type": "ode", + "grid_size": 4, + "grid_max_t": 0.999, + "flow_shift": 5.0, + "block_size_min": 1, + "block_size_max": 4, + "teacher_integrator": "euler", + "inference_blocks": [2, 2], + "data_free": False, + }, + "seed": 42, + "optim": { + "learning_rate": 2.0e-5, + "optimizer": {"_target_": "torch.optim.AdamW", "weight_decay": 0.0}, + }, + "lr_scheduler": { + "lr_decay_style": "constant", + "lr_warmup_steps": 0, + "min_lr": 2.0e-5, + }, + "step_scheduler": { + "max_steps": 1, + "num_epochs": 1, + "log_every": 1, + "ckpt_every_steps": 1, + "local_batch_size": 1, + "global_batch_size": 2, + "save_checkpoint_every_epoch": False, + }, + "training_health": {"max_grad_norm": 1.0, "zero_grad_warmup_steps": 0}, + "validation": {"count": 1, "seed": 11, "split_seed": 7, "every_steps": 1}, + "data": { + "dataloader": { + "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", + "batch_size": 1, + "drop_last": True, + "shuffle": True, + "dynamic_batch_size": False, + } + }, + "fsdp": { + "dp_size": 2, + "tp_size": 1, + "cp_size": 1, + "pp_size": 1, + "ep_size": 1, + "activation_checkpointing": activation_checkpointing, + }, + "checkpoint": { + "enabled": True, + "checkpoint_dir": str(checkpoint_dir), + "model_save_format": "torch_save", + "save_consolidated": False, + }, + } + + +def _full_fp32_gradient(parameter: torch.nn.Parameter) -> torch.Tensor: + gradient = parameter.grad + if gradient is None: + raise RuntimeError("expected a materialized gradient") + if gradient.dtype != torch.float32: + raise RuntimeError(f"expected an FP32 gradient, got {gradient.dtype}") + if isinstance(gradient, DTensor): + gradient = gradient.full_tensor() + if gradient.dtype != torch.float32: + raise RuntimeError(f"expected a gathered FP32 gradient, got {gradient.dtype}") + return gradient.detach() + + +def _assert_checkpointing(model: torch.nn.Module, *, enabled: bool) -> None: + if model.gradient_checkpointing: + raise RuntimeError("native Qwen gradient checkpointing remained enabled") + for index, block in enumerate(model.transformer_blocks): + if enabled: + if not isinstance(block, CheckpointWrapper): + raise RuntimeError(f"Qwen block {index} was not checkpoint-wrapped") + if block.checkpoint_impl is not CheckpointImpl.NO_REENTRANT: + raise RuntimeError(f"Qwen block {index} uses the wrong checkpoint implementation") + if not isinstance(block._checkpoint_wrapped_module, QwenImageTransformerBlock): + raise RuntimeError(f"Qwen block {index} wrapped an unexpected module") + elif not isinstance(block, QwenImageTransformerBlock): + raise RuntimeError(f"Qwen block {index} changed while checkpointing was disabled") + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("Qwen MR210 FSDP2 regression requires CUDA") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + initialize_pdd_distributed(backend="nccl", timeout_minutes=5) + rank = dist.get_rank() + device = torch.device("cuda", torch.cuda.current_device()) + payload = [tempfile.mkdtemp(prefix="modelopt-pdd-mr210-fsdp-") if rank == 0 else None] + dist.broadcast_object_list(payload, src=0, device=device) + root = pathlib.Path(payload[0]) + model_root = root / "model" + model_dir = model_root / "tiny_qwen_image" + completed = False + try: + if rank == 0: + assert create_tiny_qwen_image_pipeline_dir(model_root) == model_dir + dist.barrier() + + config = resolve_pdd_recipe_config( + _raw_config( + model_dir, + root / "checkpoints", + activation_checkpointing=args.activation_checkpointing, + ) + ) + setup = build_pdd_setup(config) + _assert_checkpointing(setup.student, enabled=args.activation_checkpointing) + _assert_checkpointing(setup.teacher, enabled=args.activation_checkpointing) + actual_pipeline = PDDPipeline( + setup.student, + setup.teacher, + config.pdd, + QwenImagePDDAdapter(config.pdd, compute_dtype=torch.bfloat16), + ) + + reference_student = QwenImageTransformer2DModel.from_pretrained( + model_dir, + subfolder="transformer", + torch_dtype=torch.bfloat16, + ) + reference_student = adopt_qwen_image_mr210_forward(reference_student) + reference_teacher = copy.deepcopy(reference_student).eval().requires_grad_(False) + reference_projection = convert_qwen_image_to_pdd(reference_student, config.pdd) + reference_student.to(device=device, dtype=torch.float32) + reference_teacher.to(device=device, dtype=torch.float32) + reference_pipeline = PDDPipeline( + reference_student, + reference_teacher, + config.pdd, + _FSDPBoundaryReferenceAdapter(config.pdd, compute_dtype=torch.bfloat16), + ) + + root_calls = {"student": 0, "teacher": 0} + student_times: list[torch.Tensor] = [] + teacher_times: list[torch.Tensor] = [] + inner_student_calls = 0 + + def student_root_hook(_module, _args, _kwargs): + root_calls["student"] += 1 + + def teacher_root_hook(_module, _args, _kwargs): + root_calls["teacher"] += 1 + + def student_time_hook(_module, args_for_module): + student_times.append(args_for_module[0].detach().clone()) + + def teacher_time_hook(_module, args_for_module): + teacher_times.append(args_for_module[0].detach().clone()) + + def inner_student_hook(_module, _args, _kwargs): + nonlocal inner_student_calls + inner_student_calls += 1 + + student_block = setup.student.transformer_blocks[0] + if isinstance(student_block, CheckpointWrapper): + student_block = student_block._checkpoint_wrapped_module + hooks = [ + setup.student.register_forward_pre_hook(student_root_hook, with_kwargs=True), + setup.teacher.register_forward_pre_hook(teacher_root_hook, with_kwargs=True), + setup.student.time_text_embed.register_forward_pre_hook(student_time_hook), + setup.teacher.time_text_embed.register_forward_pre_hook(teacher_time_hook), + student_block.register_forward_pre_hook(inner_student_hook, with_kwargs=True), + ] + + generator = torch.Generator().manual_seed(20260716) + data = torch.randn(1, 4, 4, 4, generator=generator).to(device) + noise = torch.randn(1, 4, 4, 4, generator=generator).to(device) + condition = ( + torch.randn(1, 3, 16, generator=generator).to( + device=device, + dtype=torch.bfloat16, + ), + torch.tensor([[1, 1, 1]], device=device, dtype=torch.long), + ) + negative_condition = ( + torch.randn(1, 2, 16, generator=generator).to( + device=device, + dtype=torch.bfloat16, + ), + torch.tensor([[1, 1]], device=device, dtype=torch.long), + ) + n = torch.tensor([0], device=device, dtype=torch.long) + k = torch.tensor([2], device=device, dtype=torch.long) + + actual_loss, _ = actual_pipeline.compute_loss( + data, + noise=noise, + condition=condition, + negative_condition=negative_condition, + n=n, + k=k, + ) + actual_loss.backward() + reference_loss, _ = reference_pipeline.compute_loss( + data, + noise=noise, + condition=condition, + negative_condition=negative_condition, + n=n, + k=k, + ) + reference_loss.backward() + for hook in hooks: + hook.remove() + + if actual_loss.dtype != torch.float32 or reference_loss.dtype != torch.float32: + raise RuntimeError( + f"PDD losses must be FP32, got {actual_loss.dtype} and {reference_loss.dtype}" + ) + if root_calls != {"student": 1, "teacher": 2}: + raise RuntimeError(f"unexpected adopted-root call counts: {root_calls}") + if len(student_times) != 1 or len(teacher_times) != 2: + raise RuntimeError( + f"unexpected time capture counts: {len(student_times)}, {len(teacher_times)}" + ) + if any(value.dtype != torch.float32 for value in (*student_times, *teacher_times)): + raise RuntimeError("FSDP rounded an MR210 timestep before time_text_embed") + expected_time = actual_pipeline.time_grid(device)[n] + if not torch.equal(student_times[0], expected_time): + raise RuntimeError("student time does not equal the exact first grid value") + if student_times[0].item() == student_times[0].to(torch.bfloat16).float().item(): + raise RuntimeError("the 0.999 discriminator did not distinguish BF16 rounding") + expected_inner_calls = 2 if args.activation_checkpointing else 1 + if inner_student_calls != expected_inner_calls: + raise RuntimeError( + "unexpected student block call count: " + f"expected {expected_inner_calls}, got {inner_student_calls}" + ) + + torch.testing.assert_close(actual_loss, reference_loss, rtol=2e-3, atol=2e-4) + gradient_pairs = ( + (setup.projection.weight, reference_projection.weight), + (setup.student.img_in.weight, reference_student.img_in.weight), + ) + for actual_parameter, reference_parameter in gradient_pairs: + actual_gradient = _full_fp32_gradient(actual_parameter) + reference_gradient = _full_fp32_gradient(reference_parameter) + if ( + not torch.isfinite(actual_gradient).all() + or not torch.isfinite(reference_gradient).all() + ): + raise FloatingPointError("MR210 FSDP gradient comparison is non-finite") + torch.testing.assert_close( + actual_gradient, + reference_gradient, + rtol=1e-2, + atol=2e-3, + ) + + gathered_losses = [torch.zeros_like(actual_loss) for _ in range(dist.get_world_size())] + dist.all_gather(gathered_losses, actual_loss.detach()) + for value in gathered_losses[1:]: + torch.testing.assert_close(value, gathered_losses[0], rtol=0, atol=0) + if rank == 0: + print( + json.dumps( + { + "activation_checkpointing": args.activation_checkpointing, + "actual_loss": actual_loss.item(), + "reference_loss": reference_loss.item(), + "student_block_calls": inner_student_calls, + "time_0": student_times[0].item(), + "world_size": dist.get_world_size(), + }, + sort_keys=True, + ) + ) + dist.barrier() + completed = True + finally: + if completed and rank == 0: + shutil.rmtree(root, ignore_errors=True) + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py index 9113d18ea5b..b90a257d59d 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -45,8 +45,10 @@ _validate_qwen_projection, ) +import modelopt.torch.fastgen.plugins.qwen_image_pdd as qwen_image_pdd_plugin from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + QWEN_IMAGE_PDD_EXECUTION, QwenImagePDDAdapter, convert_qwen_image_to_pdd, ) @@ -56,6 +58,7 @@ class _TinyQwen(nn.Module): def __init__(self) -> None: super().__init__() self.config = SimpleNamespace(guidance_embeds=False, in_channels=4) + self._modelopt_qwen_image_pdd_execution = QWEN_IMAGE_PDD_EXECUTION self.backbone = nn.Linear(4, 5, dtype=torch.bfloat16) self.proj_out = nn.Linear(5, 4, dtype=torch.bfloat16) self.calls = 0 @@ -68,10 +71,11 @@ def forward( encoder_hidden_states, encoder_hidden_states_mask, img_shapes, - guidance, + max_txt_seq_len, return_dict, ): - del img_shapes, guidance, return_dict + del img_shapes, max_txt_seq_len + assert return_dict is False condition = encoder_hidden_states.mean(dim=(1, 2), keepdim=True) condition += (encoder_hidden_states_mask.sum(dim=1)[:, None, None] / 100).to( condition.dtype @@ -83,6 +87,22 @@ def forward( return (self.proj_out(hidden),) +@pytest.fixture(autouse=True) +def _allow_tiny_qwen_protocol_double(monkeypatch): + require_production_forward = qwen_image_pdd_plugin.require_qwen_image_mr210_forward + + def require_forward(model: nn.Module) -> str: + if type(model) is _TinyQwen: + if model._modelopt_qwen_image_pdd_execution != QWEN_IMAGE_PDD_EXECUTION: + raise RuntimeError( + "Qwen-Image PDD requires the bound FastGen MR210 forward execution." + ) + return QWEN_IMAGE_PDD_EXECUTION + return require_production_forward(model) + + monkeypatch.setattr(qwen_image_pdd_plugin, "require_qwen_image_mr210_forward", require_forward) + + def _config(blocks=(32, 32, 32, 32)) -> PDDConfig: return PDDConfig( grid_size=128, @@ -111,7 +131,8 @@ def _converted(seed: int = 17): def _identity(metadata: PDDMetadata) -> dict: return { - "schema_version": 4, + "schema_version": 5, + "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, "model": {"id": "synthetic-qwen", "revision": "f" * 40, "dtype": "bfloat16"}, "pdd_metadata": metadata.to_dict(), "guidance": {"scale": 4.0}, @@ -272,6 +293,35 @@ def test_export_accepts_an_immutable_model_revision(tmp_path) -> None: assert inspect_pdd_export(output).manifest["identity"]["model"]["revision"] == "f" * 40 +@pytest.mark.parametrize("execution", [None, "canonical_diffusers"]) +def test_export_and_inference_reject_incompatible_qwen_execution(tmp_path, execution) -> None: + model, _config_value, metadata = _converted() + identity = _identity(metadata) + if execution is None: + identity.pop("qwen_image") + else: + identity["qwen_image"] = {"execution": execution} + + with pytest.raises(ValueError, match=r"Qwen execution identity|missing keys"): + write_pdd_export( + tmp_path / f"bad-execution-{execution}", + model.state_dict(), + metadata=metadata, + transformer_config={"in_channels": 4}, + identity=identity, + source_checkpoint={ + "name": "step_00000010", + "manifest_sha256": "3" * 64, + "completed_steps": 10, + }, + max_shard_bytes=12_000, + ) + + descriptor = SimpleNamespace(manifest={"identity": identity}) + with pytest.raises(RuntimeError, match="Qwen execution identity"): + _model_identity(descriptor) + + @pytest.mark.parametrize("revision", [None, "main", "F" * 40]) def test_export_and_inference_reject_mutable_model_revisions(tmp_path, revision) -> None: model, _config_value, metadata = _converted() diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index 96562560d16..a514eaab38d 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -59,6 +59,7 @@ ) from modelopt.torch.fastgen import PDDLayerSpec, convert_to_pdd_output_projection +from modelopt.torch.fastgen.plugins.qwen_image_pdd import require_qwen_image_mr210_forward def _canonical_parameter_names(model: nn.Module) -> tuple[set[str], set[str]]: @@ -526,18 +527,25 @@ def test_incompatible_modes_fail_during_config_resolution( resolve_pdd_recipe_config(raw) -def test_model_revision_and_compute_dtype_follow_loader_contract(tmp_path) -> None: +def test_model_revision_and_strict_compute_contract(tmp_path) -> None: raw = _raw_config(tmp_path) raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" raw["model"]["revision"] = "a" * 40 - raw["model"]["torch_dtype"] = "float32" - raw["model"]["fuse_qkv_projections"] = True config = resolve_pdd_recipe_config(raw) assert config.model_revision == "a" * 40 - assert config.dtype == torch.float32 - assert config.fuse_qkv_projections is True + assert config.dtype == torch.bfloat16 + assert config.fuse_qkv_projections is False + + raw["model"]["torch_dtype"] = "float32" + with pytest.raises(ValueError, match="torch_dtype='bfloat16'"): + resolve_pdd_recipe_config(raw) + + raw["model"]["torch_dtype"] = "bfloat16" + raw["model"]["fuse_qkv_projections"] = True + with pytest.raises(ValueError, match="does not support QKV fusion"): + resolve_pdd_recipe_config(raw) @pytest.mark.parametrize("revision", [None, "main", "A" * 40, "a" * 39]) @@ -685,6 +693,9 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert source.pipe.tokenizer is None assert source.pipe.vae is None assert source.pipe.transformer is source.student + assert source.student.forward.__self__ is source.student + assert source.teacher.forward.__self__ is source.teacher + assert source.teacher.forward.__func__ is source.student.forward.__func__ assert source.student.get_submodule("proj_out") is source.projection assert source.projection.out_features == source.projection.base_out_features * 4 assert "proj_out.weight" in source.checkpoint_keys @@ -693,11 +704,12 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert policy.param_dtype == torch.bfloat16 assert policy.reduce_dtype == torch.float32 assert policy.output_dtype == torch.bfloat16 - assert policy.cast_forward_inputs is True + assert policy.cast_forward_inputs is False assert source.distributed_setup.strategy_config.activation_checkpointing is False assert source.distributed_setup.strategy_config.reshard_after_forward is True assert config.parallel.activation_checkpointing is True for model in (source.student, source.teacher): + require_qwen_image_mr210_forward(model) assert model.gradient_checkpointing is False assert len(model.transformer_blocks) == 6 assert all(isinstance(block, CheckpointWrapper) for block in model.transformer_blocks) @@ -776,6 +788,13 @@ def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: assert export_setup.metadata == source.metadata assert export_setup.checkpoint_keys == source.checkpoint_keys assert not hasattr(export_setup, "optimizer") + require_qwen_image_mr210_forward(export_setup.student) + export_policy = export_setup.distributed_setup.strategy_config.mp_policy + assert export_policy.param_dtype == torch.bfloat16 + assert export_policy.reduce_dtype == torch.float32 + assert export_policy.output_dtype == torch.bfloat16 + assert export_policy.cast_forward_inputs is False + assert export_setup.student.forward.__self__ is export_setup.student export_setup.checkpointer.load_model( export_setup.student, str(checkpoint_root / "model"), diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index cb8d4f6facd..122000e2a25 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -49,6 +49,8 @@ from pdd.training import prepare_qwen_pdd_batch from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 +from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION + def _released_sampler(sample_ids: tuple[str, ...]) -> ReplayableBatchSampler: sampler_module = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") @@ -93,8 +95,15 @@ def _checkpointer(lifecycle, checkpoint_dir): return config.build(dp_rank=0, tp_rank=0, pp_rank=0, moe_mesh=None) -def _identity(lifecycle, scheduler, sample_ids): +def _identity( + lifecycle, + scheduler, + sample_ids, + *, + qwen_image_execution=QWEN_IMAGE_PDD_EXECUTION, +): return build_pdd_checkpoint_identity( + qwen_image_execution=qwen_image_execution, metadata=lifecycle.metadata, model_id="synthetic-pdd-toy", model_revision="a" * 40, @@ -116,6 +125,17 @@ def _identity(lifecycle, scheduler, sample_ids): ) +def test_checkpoint_identity_rejects_unbound_qwen_execution() -> None: + lifecycle = build_toy_lifecycle() + with pytest.raises(ValueError, match="qwen_image_execution"): + _identity( + lifecycle, + lifecycle.scheduler, + ("sample-0",), + qwen_image_execution="canonical_diffusers", + ) + + class _StepSchedulerStub: def __init__(self, trainer, sampler) -> None: self.trainer = trainer @@ -644,6 +664,7 @@ def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) checkpoint = source_manager.save() assert checkpoint.name == "step_00000002" assert source_manager.identity["data"]["dataset_snapshot_sha256"] == "2" * 64 + assert source_manager.identity["qwen_image"] == {"execution": "fastgen_mr210"} assert source_manager.identity["training"] == { "seed": 1234, "validation_seed": 2026, @@ -742,6 +763,25 @@ def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) with pytest.raises(RuntimeError, match="identity"): third_manager.resolve(mismatched.name) + for step, execution in ((99999994, None), (99999995, "canonical_diffusers")): + incompatible = tmp_path / "checkpoints" / f"step_{step:08d}" + shutil.copytree(resumed_checkpoint, incompatible) + incompatible_manifest_path = incompatible / "manifest.json" + incompatible_manifest = json.loads(incompatible_manifest_path.read_text()) + incompatible_manifest["completed_steps"] = step + if execution is None: + incompatible_manifest["identity"].pop("qwen_image") + else: + incompatible_manifest["identity"]["qwen_image"] = {"execution": execution} + incompatible_manifest_path.write_text( + json.dumps(incompatible_manifest, indent=2, sort_keys=True) + "\n" + ) + _refresh_complete_marker(incompatible) + (tmp_path / "checkpoints" / "LATEST").write_text(incompatible.name + "\n") + assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() + with pytest.raises(RuntimeError, match="identity"): + third_manager.resolve(incompatible.name) + missing_model = tmp_path / "checkpoints" / "step_99999996" shutil.copytree(resumed_checkpoint, missing_model) model_payload = next( diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 6df51b54352..1cf105b7146 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -18,29 +18,57 @@ from __future__ import annotations import copy -from types import SimpleNamespace +from types import MethodType, SimpleNamespace import pytest import torch import torch.nn.functional as F from torch import nn -from modelopt.torch.fastgen import PDDConfig, PDDOutputProjection +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, QWEN_IMAGE_PDD_LAYER_SPEC, + adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, + require_qwen_image_mr210_forward, ) -class _TinyQwenTransformer(nn.Module): +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_mr210_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 the bound FastGen MR210 forward execution." + ) + return QWEN_IMAGE_PDD_EXECUTION + return require_production_forward(model) + + monkeypatch.setattr(qwen_image_pdd_plugin, "require_qwen_image_mr210_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]] = [] @@ -53,8 +81,7 @@ def forward( encoder_hidden_states, encoder_hidden_states_mask, img_shapes, - guidance, - return_dict, + max_txt_seq_len, **kwargs, ): condition_value = encoder_hidden_states.mean(dim=(1, 2), keepdim=True) @@ -72,9 +99,8 @@ def forward( "encoder_hidden_states": encoder_hidden_states.detach().clone(), "encoder_hidden_states_mask": encoder_hidden_states_mask.detach().clone(), "img_shapes": img_shapes, - "guidance": guidance, + "max_txt_seq_len": max_txt_seq_len, "projection_input": hidden.detach().clone(), - "return_dict": return_dict, "kwargs": kwargs, "output": output.detach().clone(), } @@ -122,11 +148,58 @@ def _call_base_packed( encoder_hidden_states=embeddings.to(torch.bfloat16), encoder_hidden_states_mask=mask, img_shapes=build_img_shapes(state.shape[0], state.shape[2], state.shape[3]), - guidance=None, - return_dict=False, + max_txt_seq_len=int(mask.sum(dim=1).max().item()), )[0] +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 _tiny_qwen_oracle( + model: _TinyQwenTransformer, + state: torch.Tensor, + time: torch.Tensor, + condition: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + embeddings, mask = condition + packed = _pack_oracle(state).to(torch.bfloat16) + hidden = torch.tanh(model.backbone(packed)) + condition_value = embeddings.mean(dim=(1, 2), keepdim=True) + condition_value = condition_value + 0.01 * mask.sum(dim=1, keepdim=True).unsqueeze(-1) + hidden = hidden + condition_value.to(hidden.dtype) + hidden = hidden + (0.1 * time[:, None, None]).to(hidden.dtype) + return F.linear(hidden, model.proj_out.weight, model.proj_out.bias) + + def test_conversion_is_idempotent_and_every_initialized_head_matches_base() -> None: base = _TinyQwenTransformer() student = copy.deepcopy(base) @@ -150,8 +223,8 @@ def test_conversion_is_idempotent_and_every_initialized_head_matches_base() -> N assert len(student.calls) == 1 torch.testing.assert_close(student.calls[0]["timestep"], time) assert student.calls[0]["img_shapes"] == [[(1, 2, 2)], [(1, 2, 2)]] + assert student.calls[0]["max_txt_seq_len"] == 3 assert "txt_seq_lens" not in student.calls[0]["kwargs"] - assert student.calls[0]["guidance"] is None def test_unfused_channel_major_output_maps_each_packed_head_in_order() -> None: @@ -224,7 +297,7 @@ def test_fused_student_matches_explicit_packed_weight_fusion() -> None: assert student.proj_out(projection.weight.new_zeros(1, 5)).shape[-1] == 16 -def test_teacher_cfg_uses_canonical_packed_per_token_norm_rescale() -> None: +def test_teacher_cfg_uses_mr210_global_fp32_norm_rescale() -> None: teacher = _TinyQwenTransformer() config = _config(guidance_scale=4.0) adapter = QwenImagePDDAdapter(config) @@ -239,29 +312,31 @@ def test_teacher_cfg_uses_canonical_packed_per_token_norm_rescale() -> None: ) assert len(teacher.calls) == 2 - conditional = teacher.calls[0]["output"] - unconditional = teacher.calls[1]["output"] - guided = unconditional + 4.0 * (conditional - unconditional) + conditional = unpack_latents(teacher.calls[0]["output"], 4, 4) + unconditional = unpack_latents(teacher.calls[1]["output"], 4, 4) + guided_low_precision = conditional + 3.0 * (conditional - unconditional) + conditional_fp32 = conditional.float() + guided_fp32 = guided_low_precision.float() factor = torch.linalg.vector_norm( - conditional, - dim=-1, + conditional_fp32, + dim=(1, 2, 3), keepdim=True, - ) / torch.linalg.vector_norm(guided, dim=-1, keepdim=True) - expected = unpack_latents(guided * factor, 4, 4) + ) / torch.linalg.vector_norm(guided_fp32, dim=(1, 2, 3), keepdim=True).clamp_min(1e-5) + expected = (guided_fp32 * factor).to(conditional.dtype) assert actual.dtype == torch.bfloat16 torch.testing.assert_close(actual, expected) torch.testing.assert_close(teacher.calls[0]["encoder_hidden_states"], condition[0]) torch.testing.assert_close(teacher.calls[1]["encoder_hidden_states"], negative_condition[0]) assert all("txt_seq_lens" not in call["kwargs"] for call in teacher.calls) - assert all(call["guidance"] is None for call in teacher.calls) def test_teacher_cfg_stays_in_model_output_dtype() -> None: - class LowPrecisionTeacher(nn.Module): + class LowPrecisionTeacher(_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.outputs: list[torch.Tensor] = [] @@ -282,15 +357,167 @@ def forward(self, *, hidden_states, encoder_hidden_states, **kwargs): ) assert actual.dtype == torch.bfloat16 - conditional, unconditional = teacher.outputs - guided = unconditional + 4.0 * (conditional - unconditional) - factor = torch.linalg.vector_norm(conditional, dim=-1, keepdim=True) / torch.linalg.vector_norm( - guided, dim=-1, keepdim=True - ) - expected = unpack_latents(guided * factor, 4, 4) + conditional, unconditional = (unpack_latents(output, 4, 4) for output in teacher.outputs) + 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, 2, 3), keepdim=True + ) / torch.linalg.vector_norm(guided_fp32, dim=(1, 2, 3), keepdim=True).clamp_min(1e-5) + expected = (guided_fp32 * factor).to(torch.bfloat16) torch.testing.assert_close(actual, expected, rtol=0, atol=0) +def test_teacher_cfg_zero_guided_norm_uses_mr210_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_mr210_qwen_loss_and_backward_match_independent_equations() -> None: + class CapturingAdapter(QwenImagePDDAdapter): + 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 = _TinyQwenTransformer() + actual_student = copy.deepcopy(base) + actual_teacher = copy.deepcopy(base) + oracle_student = copy.deepcopy(base) + oracle_teacher = copy.deepcopy(base) + 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) + pipeline = PDDPipeline(actual_student, actual_teacher, config, adapter) + + generator = torch.Generator().manual_seed(47) + data = torch.randn(1, 1, 4, 4, generator=generator) + noise = torch.randn(1, 1, 4, 4, generator=generator) + condition = ( + torch.randn(1, 3, 2, generator=generator).to(torch.bfloat16), + torch.tensor([[1, 1, 1]], dtype=torch.long), + ) + negative_condition = ( + torch.randn(1, 2, 2, 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, _metrics = 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() + + packed_heads = _tiny_qwen_oracle(oracle_student, x_n, time_n, condition) + batch, patches, _features = packed_heads.shape + packed_heads = packed_heads.reshape(batch, patches, 4, 4).permute(0, 2, 1, 3) + oracle_heads = _unpack_oracle(packed_heads.reshape(4, patches, 4), 4, 4).reshape(1, 4, 1, 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_packed = _tiny_qwen_oracle(oracle_teacher, x_bar_k, time_k, condition) + unconditional_packed = _tiny_qwen_oracle( + oracle_teacher, + x_bar_k, + time_k, + negative_condition, + ) + conditional = _unpack_oracle(conditional_packed, 4, 4) + unconditional = _unpack_oracle(unconditional_packed, 4, 4) + guided_low_precision = conditional + 3.0 * (conditional - unconditional) + conditional_fp32 = conditional.float() + guided_fp32 = guided_low_precision.float() + norm_dims = (1, 2, 3) + teacher_target_low_precision = ( + guided_fp32 + * ( + torch.linalg.vector_norm(conditional_fp32, dim=norm_dims, keepdim=True) + / torch.linalg.vector_norm(guided_fp32, dim=norm_dims, keepdim=True).clamp_min(1e-5) + ) + ).to(torch.bfloat16) + 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_state, + x_bar_k, + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close( + adapter.captured_teacher, + teacher_target_low_precision, + rtol=0, + atol=0, + ) + torch.testing.assert_close(actual_loss, oracle_loss, rtol=1e-6, atol=1e-7) + for name in ("backbone.weight", "proj_out.weight"): + actual_gradient = dict(actual_student.named_parameters())[name].grad + oracle_gradient = dict(oracle_student.named_parameters())[name].grad + torch.testing.assert_close( + actual_gradient, + oracle_gradient, + rtol=1e-6, + atol=1e-7, + ) + + def test_guidance_disabled_teacher_is_one_conditional_call_without_negative_condition() -> None: teacher = _TinyQwenTransformer() config = _config(guidance_scale=None) @@ -339,6 +566,175 @@ def _tiny_diffusers_qwen(): ) +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.""" + 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, + joint_attention_kwargs=None, + ) + 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 = adopt_qwen_image_mr210_forward(copy.deepcopy(base)) + 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 = _unpack_oracle( + oracle_forward(oracle_teacher, x_bar_k, time_k, condition), + 4, + 4, + ) + unconditional = _unpack_oracle( + oracle_forward(oracle_teacher, x_bar_k, time_k, negative_condition), + 4, + 4, + ) + guided_low_precision = conditional + 3.0 * (conditional - unconditional) + conditional_fp32 = conditional.float() + guided_fp32 = guided_low_precision.float() + norm_dims = (1, 2, 3) + teacher_target_low_precision = ( + guided_fp32 + * ( + torch.linalg.vector_norm(conditional_fp32, dim=norm_dims, keepdim=True) + / torch.linalg.vector_norm(guided_fp32, dim=norm_dims, keepdim=True).clamp_min(1e-5) + ) + ).to(torch.bfloat16) + 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_conversion_preserves_the_ordinary_diffusers_qwen_root() -> None: student = _tiny_diffusers_qwen().eval() root_type = type(student) @@ -351,27 +747,90 @@ def test_conversion_preserves_the_ordinary_diffusers_qwen_root() -> None: assert dict(student.config) == config -def test_canonical_qwen_conversion_preserves_every_initialized_head() -> None: - base = _tiny_diffusers_qwen().eval() +def test_mr210_adoption_preserves_root_state_and_deepcopy_binding() -> None: + source = _tiny_diffusers_qwen().eval() + source_type = type(source) + source_state = {name: value.detach().clone() for name, value in source.state_dict().items()} + source_state_keys = tuple(source_state) + source_parameters = tuple(source.parameters()) + source_buffers = tuple(source.buffers()) + custom_attribute = object() + source.custom_attribute = custom_attribute + hook = source.register_forward_pre_hook(lambda *_args: None) + + adopted = adopt_qwen_image_mr210_forward(source) + + assert adopted is source + assert type(adopted) is source_type + assert tuple(adopted.state_dict()) == source_state_keys + assert all( + actual is expected + for actual, expected in zip(adopted.parameters(), source_parameters, strict=True) + ) + assert all( + actual is expected + for actual, expected in zip(adopted.buffers(), source_buffers, strict=True) + ) + assert hook.id in adopted._forward_pre_hooks + assert adopted.custom_attribute is custom_attribute + assert adopt_qwen_image_mr210_forward(adopted) is adopted + + round_trip = copy.deepcopy(adopted) + with torch.no_grad(): + next(round_trip.parameters()).zero_() + round_trip.load_state_dict(source_state) + for name, value in round_trip.state_dict().items(): + torch.testing.assert_close(value, source_state[name], rtol=0, atol=0) + assert round_trip.forward.__self__ is round_trip + + teacher = copy.deepcopy(adopted) + assert teacher is not adopted + assert teacher.forward.__func__ is adopted.forward.__func__ + assert teacher.forward.__self__ is teacher + assert teacher.forward.__self__ is not adopted + require_qwen_image_mr210_forward(teacher) + + tampered = copy.deepcopy(adopted) + tampered.forward = MethodType(lambda self, **_kwargs: self, tampered) + with pytest.raises(RuntimeError, match="MR210 forward execution"): + require_qwen_image_mr210_forward(tampered) + + conflicting = _tiny_diffusers_qwen() + conflicting.forward = MethodType(lambda self, **_kwargs: self, conflicting) + with pytest.raises(RuntimeError, match="instance-level forward override"): + adopt_qwen_image_mr210_forward(conflicting) + + forged = _tiny_diffusers_qwen() + forged._modelopt_qwen_image_pdd_execution = QWEN_IMAGE_PDD_EXECUTION + with pytest.raises(RuntimeError, match="bound FastGen MR210 forward"): + qwen_image_pdd_plugin.require_qwen_image_mr210_forward(forged) + + +def test_mr210_qwen_conversion_preserves_every_initialized_head() -> None: + base = _tiny_diffusers_qwen().eval().to(torch.bfloat16) student = copy.deepcopy(base) + student = adopt_qwen_image_mr210_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) - mask = torch.tensor([[1, 1, 1], [1, 0, 1]], dtype=torch.long) + 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), + "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), - "guidance": None, - "return_dict": False, + "max_txt_seq_len": 3, } with torch.no_grad(): - expected = unpack_latents(base(**model_kwargs)[0], 4, 4) + 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, @@ -383,16 +842,18 @@ def test_canonical_qwen_conversion_preserves_every_initialized_head() -> None: torch.testing.assert_close(actual, expected[:, None].expand_as(actual), rtol=0, atol=0) -def test_canonical_qwen_mask_makes_masked_padding_numerically_inert() -> None: - student = _tiny_diffusers_qwen().eval() +def test_mr210_mask_flow_differs_from_canonical_joint_mask() -> None: + canonical = _tiny_diffusers_qwen().eval().to(torch.bfloat16) + student = copy.deepcopy(canonical) + student = adopt_qwen_image_mr210_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, 3, 12, generator=generator) - mask = torch.tensor([[1, 1, 1], [1, 0, 1]], dtype=torch.long) + encoder_hidden_states = torch.randn(2, 3, 12, generator=generator).to(torch.bfloat16) + mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long) poisoned = encoder_hidden_states.clone() poisoned[~mask.bool()] = ( torch.randn( @@ -400,65 +861,146 @@ def test_canonical_qwen_mask_makes_masked_padding_numerically_inert() -> None: 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 "attention_mask" not in kwargs + 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(): - baseline = adapter.student_all_heads( + 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), ) - actual = adapter.student_all_heads( + 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) + assert not torch.equal(strict_poisoned[1], strict_baseline[1]) + assert len(captured_masks) == 2 + assert all(torch.equal(captured, mask) for captured in captured_masks) + + +def test_mr210_preserves_diffusers_output_and_harmless_call_contract() -> None: + student = adopt_qwen_image_mr210_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), + "txt_seq_lens": [3, 1], + "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) - torch.testing.assert_close(actual, baseline, rtol=0, atol=0) +def test_mr210_time_embed_receives_fp32_grid_value() -> None: + student = adopt_qwen_image_mr210_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()) -def test_canonical_qwen_teacher_cfg_matches_the_pipeline_formula() -> None: - teacher = _tiny_diffusers_qwen().eval() + 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_global_reference() -> None: + teacher = adopt_qwen_image_mr210_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), - torch.tensor([[1, 1, 0], [1, 0, 1]], dtype=torch.long), + 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), + 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), + 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), - guidance=None, + max_txt_seq_len=int(mask.sum(dim=1).max().item()), return_dict=False, )[0] with torch.no_grad(): - conditional = direct_packed(condition) - unconditional = direct_packed(negative_condition) - guided = unconditional + 4.0 * (conditional - unconditional) - expected = unpack_latents( - guided - * ( - torch.linalg.vector_norm(conditional, dim=-1, keepdim=True) - / torch.linalg.vector_norm(guided, dim=-1, keepdim=True) - ), - 4, - 4, - ) + conditional = unpack_latents(direct_packed(condition), 4, 4) + unconditional = unpack_latents(direct_packed(negative_condition), 4, 4) + 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, 2, 3), keepdim=True + ) / torch.linalg.vector_norm(guided_fp32, dim=(1, 2, 3), keepdim=True).clamp_min(1e-5) + expected = (guided_fp32 * factor).to(torch.bfloat16) actual = adapter.teacher_velocity( teacher, state, @@ -470,7 +1012,7 @@ def direct_packed(current_condition): torch.testing.assert_close(actual, expected, rtol=0, atol=0) -def test_adapter_accepts_arbitrary_binary_masks_nonzero_padding_and_floating_time() -> None: +def test_adapter_accepts_binary_masks_nonzero_padding_and_fp32_time() -> None: student = _TinyQwenTransformer() config = _config() convert_qwen_image_to_pdd(student, config) @@ -482,13 +1024,14 @@ def test_adapter_accepts_arbitrary_binary_masks_nonzero_padding_and_floating_tim actual = QwenImagePDDAdapter(config).student_all_heads( student, state, - time.to(torch.bfloat16), + time, condition=(embeddings, mask), ) assert actual.shape == (2, 4, 1, 4, 4) torch.testing.assert_close(student.calls[0]["encoder_hidden_states_mask"], mask) - assert student.calls[0]["timestep"].dtype == torch.bfloat16 + assert student.calls[0]["timestep"].dtype == torch.float32 + assert student.calls[0]["max_txt_seq_len"] == 2 def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> None: @@ -521,6 +1064,13 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N adapter.student_all_heads(transformer, state, time, condition=condition) 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, + ) with pytest.raises(TypeError, match="tuple"): adapter.student_all_heads(transformer, state, time, condition=condition[0]) with pytest.raises(ValueError, match="requires batched embeddings"): @@ -553,6 +1103,12 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N condition=(condition[0], condition[1].float()), ) + unmarked = _TinyQwenTransformer() + delattr(unmarked, "_modelopt_qwen_image_pdd_execution") + convert_qwen_image_to_pdd(unmarked, config) + with pytest.raises(RuntimeError, match="MR210 forward execution"): + adapter.student_all_heads(unmarked, state, time, condition=condition) + def test_raw_head_reference_uses_independent_linear_outputs() -> None: """Pin the widened storage order without calling adapter reshape helpers.""" From 7d3e035964233f52c4e68b0d06c1424c07ff93ac Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 16 Jul 2026 16:58:33 -0700 Subject: [PATCH 30/45] Add repeatable Qwen PDD evaluation Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/README.md | 58 ++ .../fastgen/pdd/evaluate_qwen_image.py | 456 +++++++++++++ .../fastgen/pdd/inference_qwen_image.py | 309 +++------ .../fastgen/pdd/inference_runtime.py | 354 ++++++++++ .../examples/diffusers/fastgen/test_layout.py | 2 + .../fastgen/test_pdd_evaluation_runner.py | 629 ++++++++++++++++++ 6 files changed, 1576 insertions(+), 232 deletions(-) create mode 100644 examples/diffusers/fastgen/pdd/evaluate_qwen_image.py create mode 100644 examples/diffusers/fastgen/pdd/inference_runtime.py create mode 100644 tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index da496082a47..d48fbc8f016 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -84,3 +84,61 @@ python examples/diffusers/fastgen/pdd/inference_qwen_image.py \ --seed 42 --height 1024 --width 1024 \ --output /path/to/pdd4.png --result-json /path/to/pdd4.json ``` + +## Repeatable evaluation records + +`evaluate_qwen_image.py` loads the authenticated export once and evaluates every ordered +prompt/seed pair for one of the source-owned `pdd-2`, `pdd-4`, or `pdd-8` schedules. The prompt +file must be canonical JSON, including its trailing newline. For example: + +```json +{"prompts":[{"prompt":"a small red cube on a white table","prompt_id":"red-cube-0001","seeds":[42]}],"schema_version":1} +``` + +```bash +python examples/diffusers/fastgen/pdd/evaluate_qwen_image.py \ + --export-dir /path/to/pdd-export \ + --prompts /path/to/prompts.json --schedule pdd-4 \ + --output-dir /path/to/evaluation-pdd4 \ + --result-json /path/to/evaluation-pdd4/result.json \ + --warmup-runs 1 --measured-runs 5 \ + --height 1024 --width 1024 --max-sequence-length 512 +``` + +The runner publishes the output directory atomically. Its canonical result records exact export, +prompt, raw-noise, initial-state, and grid identities; requested logical blocks; observed scheduler +calls; actual transformer calls; synchronized end-to-end and transformer timings; throughput; and +peak CUDA allocation. Warmups execute and validate the same path but are excluded from measured +arrays. CPU memory entries are `null`. The record deliberately contains no image-quality score or +effectiveness conclusion. + +For a functional standard-teacher baseline, the public Diffusers pipeline can be pinned to the +same immutable Qwen revision and run for 50 steps: + +```python +import torch +from diffusers import QwenImagePipeline + +revision = "75e0b4be04f60ec59a75f475837eced720f823b6" +pipe = QwenImagePipeline.from_pretrained( + "Qwen/Qwen-Image", + revision=revision, + torch_dtype=torch.bfloat16, + use_safetensors=True, +).to("cuda") +generator = torch.Generator(device="cuda").manual_seed(42) +image = pipe( + prompt="a small red cube on a white table", + height=1024, + width=1024, + num_inference_steps=50, + generator=generator, +).images[0] +image.save("teacher-50.png") +``` + +That command is a functional baseline, not matched scientific evidence. A result-bearing study +must separately freeze the prompt set, seeds, controls, quality metrics, thresholds, and exact +teacher-50, undistilled Euler-2/4/8, and PDD-2/4/8 trajectory/counter schemas. Those controls and +conclusions belong to the reviewed external `scripts/pdd_qwen_effectiveness_v1` experiment package, +not this general ModelOpt example. diff --git a/examples/diffusers/fastgen/pdd/evaluate_qwen_image.py b/examples/diffusers/fastgen/pdd/evaluate_qwen_image.py new file mode 100644 index 00000000000..af00cf79bb2 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/evaluate_qwen_image.py @@ -0,0 +1,456 @@ +# 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. + +"""Evaluate an authenticated Qwen-Image PDD export over a prompt manifest.""" + +from __future__ import annotations + +import argparse +import hashlib +import math +import os +import re +import shutil +import sys +import time +import uuid +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +sys.dont_write_bytecode = True + +_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 pdd.artifacts import load_canonical_json, sha256_file, write_canonical_json # noqa: E402 +from pdd.export import PDD_INFERENCE_SCHEDULES # noqa: E402 +from pdd.inference_runtime import ( # noqa: E402 + QwenPDDInferenceRuntime, + load_qwen_pdd_runtime, + save_png, +) + +_PROMPT_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*\Z") + + +@dataclass(frozen=True) +class PromptPair: + prompt_id: str + prompt: str + seed: int + + +@dataclass(frozen=True) +class Observation: + scheduler_calls: int + transformer_calls: int + transformer_seconds: float + end_to_end_seconds: float + peak_device_memory_bytes: int | None + image: Any + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--export-dir", type=Path, required=True) + parser.add_argument("--prompts", type=Path, required=True) + parser.add_argument("--schedule", choices=tuple(PDD_INFERENCE_SCHEDULES), required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--result-json", type=Path, required=True) + parser.add_argument("--warmup-runs", type=int, default=1) + parser.add_argument("--measured-runs", type=int, default=5) + 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") + return parser.parse_args() + + +def _prompt_pairs(value: Any) -> list[PromptPair]: + if not isinstance(value, dict) or set(value) != {"schema_version", "prompts"}: + raise ValueError("prompt manifest must contain exactly schema_version and prompts.") + if type(value["schema_version"]) is not int or value["schema_version"] != 1: + raise ValueError("prompt manifest schema_version must be 1.") + prompts = value["prompts"] + if not isinstance(prompts, list) or not prompts: + raise ValueError("prompt manifest prompts must be a non-empty list.") + pairs: list[PromptPair] = [] + previous_id: str | None = None + for index, item in enumerate(prompts): + if not isinstance(item, dict) or set(item) != {"prompt_id", "prompt", "seeds"}: + raise ValueError(f"prompts[{index}] has incompatible fields.") + prompt_id = item["prompt_id"] + prompt = item["prompt"] + seeds = item["seeds"] + if not isinstance(prompt_id, str) or _PROMPT_ID.fullmatch(prompt_id) is None: + raise ValueError(f"prompts[{index}].prompt_id is not a safe path component.") + if previous_id is not None and prompt_id <= previous_id: + raise ValueError("prompt IDs must be unique and lexicographically ordered.") + if not isinstance(prompt, str) or not prompt: + raise ValueError(f"prompts[{index}].prompt must be non-empty text.") + if not isinstance(seeds, list) or not seeds: + raise ValueError(f"prompts[{index}].seeds must be non-empty.") + if any(type(seed) is not int or seed < 0 or seed >= 2**63 for seed in seeds): + raise ValueError(f"prompts[{index}].seeds contains an invalid seed.") + if seeds != sorted(set(seeds)): + raise ValueError(f"prompts[{index}].seeds must be sorted and unique.") + pairs.extend(PromptPair(prompt_id, prompt, seed) for seed in seeds) + previous_id = prompt_id + return pairs + + +def _reject_symlink_components(path: Path) -> Path: + if ".." in path.parts: + raise ValueError("evaluation paths cannot contain parent traversal components.") + absolute = path.absolute() + current = Path(absolute.anchor) + for part in absolute.parts[1:]: + current /= part + if current.is_symlink(): + raise ValueError(f"evaluation path traverses a symlink: {current}.") + return absolute.resolve(strict=False) + + +def _resolve_output_paths(output_value: Path, result_value: Path) -> tuple[Path, Path, Path]: + output = _reject_symlink_components(output_value) + result = _reject_symlink_components(result_value) + if output.exists() or output.is_symlink(): + raise FileExistsError(f"evaluation output already exists: {output}.") + if output.name in {"", ".", ".."}: + raise ValueError("evaluation output directory name is invalid.") + if not output.parent.is_dir() or output.parent.is_symlink(): + raise ValueError("evaluation output parent must be an existing regular directory.") + try: + relative_result = result.relative_to(output) + except ValueError as error: + raise ValueError("result JSON must be strictly beneath output_dir.") from error + if not relative_result.parts or relative_result == Path("."): + raise ValueError("result JSON must be strictly beneath output_dir.") + if result.suffix.lower() != ".json": + raise ValueError("result JSON must use a .json suffix.") + staging = output.with_name(f".{output.name}.{uuid.uuid4().hex}.staging") + os.mkdir(staging, mode=0o770) + os.chmod(staging, 0o770, follow_symlinks=False) + return output, relative_result, staging + + +def _restore_instance_method( + owner: Any, name: str, *, had_instance_value: bool, instance_value: Any +) -> None: + if had_instance_value: + setattr(owner, name, instance_value) + elif name in vars(owner): + delattr(owner, name) + + +def _run_repetition( + runtime: QwenPDDInferenceRuntime, + prompt: str, + raw_noise: torch.Tensor, + max_sequence_length: int, +) -> Observation: + scheduler_calls = 0 + transformer_calls = 0 + cpu_forward_started: list[float] = [] + cpu_forward_seconds = 0.0 + cuda_event_pairs: list[tuple[Any, Any]] = [] + scheduler = runtime.scheduler + scheduler_values = vars(scheduler) + scheduler_had_step = "step" in scheduler_values + scheduler_instance_step = scheduler_values.get("step") + original_step = scheduler.step + + def counted_step(*args: Any, **kwargs: Any) -> Any: + nonlocal scheduler_calls + scheduler_calls += 1 + return original_step(*args, **kwargs) + + def before_forward(_module: Any, _args: Any, _kwargs: Any) -> None: + nonlocal transformer_calls + transformer_calls += 1 + if runtime.device.type == "cuda": + started = torch.cuda.Event(enable_timing=True) + ended = torch.cuda.Event(enable_timing=True) + started.record() + cuda_event_pairs.append((started, ended)) + else: + cpu_forward_started.append(time.perf_counter()) + + def after_forward(_module: Any, _args: Any, _kwargs: Any, _output: Any) -> None: + nonlocal cpu_forward_seconds + if runtime.device.type == "cuda": + cuda_event_pairs[-1][1].record() + else: + if not cpu_forward_started: + raise RuntimeError("transformer post-hook ran without its pre-hook.") + cpu_forward_seconds += time.perf_counter() - cpu_forward_started.pop() + + with ExitStack() as cleanup: + cleanup.callback( + _restore_instance_method, + scheduler, + "step", + had_instance_value=scheduler_had_step, + instance_value=scheduler_instance_step, + ) + setattr(scheduler, "step", counted_step) + pre_hook = runtime.student.register_forward_pre_hook(before_forward, with_kwargs=True) + cleanup.callback(pre_hook.remove) + post_hook = runtime.student.register_forward_hook(after_forward, with_kwargs=True) + cleanup.callback(post_hook.remove) + if runtime.device.type == "cuda": + torch.cuda.synchronize(runtime.device) + torch.cuda.reset_peak_memory_stats(runtime.device) + started = time.perf_counter() + condition = runtime.encode_prompt(prompt, max_sequence_length) + images = runtime.sample_decode(condition, raw_noise) + if runtime.device.type == "cuda": + torch.cuda.synchronize(runtime.device) + end_to_end_seconds = time.perf_counter() - started + if runtime.device.type == "cuda": + transformer_seconds = sum( + start.elapsed_time(end) / 1000.0 for start, end in cuda_event_pairs + ) + peak_memory = torch.cuda.max_memory_allocated(runtime.device) + else: + transformer_seconds = cpu_forward_seconds + peak_memory = None + + expected = len(runtime.config.inference_blocks) + if not isinstance(images, list) or len(images) != 1: + raise RuntimeError("one evaluation repetition must return exactly one image.") + if scheduler_calls != 0: + raise RuntimeError(f"PDD evaluation observed {scheduler_calls} scheduler.step calls.") + if transformer_calls != expected: + raise RuntimeError( + f"PDD evaluation observed {transformer_calls} transformer calls; expected {expected}." + ) + if cpu_forward_started: + raise RuntimeError("transformer forward hooks are unbalanced.") + if ( + not math.isfinite(end_to_end_seconds) + or not math.isfinite(transformer_seconds) + or end_to_end_seconds <= 0 + or transformer_seconds <= 0 + or transformer_seconds > end_to_end_seconds + 1e-6 + ): + raise RuntimeError("evaluation timing invariants failed.") + if peak_memory is not None and (type(peak_memory) is not int or peak_memory < 0): + raise RuntimeError("evaluation peak device memory is invalid.") + return Observation( + scheduler_calls=scheduler_calls, + transformer_calls=transformer_calls, + transformer_seconds=transformer_seconds, + end_to_end_seconds=end_to_end_seconds, + peak_device_memory_bytes=peak_memory, + image=images[0], + ) + + +def _summary(values: list[float | int]) -> dict[str, float | int]: + if not values or any(type(value) not in {int, float} for value in values): + raise ValueError("summary values must be a non-empty numeric list.") + ordered = sorted(values) + middle = len(ordered) // 2 + median = ordered[middle] if len(ordered) % 2 else (ordered[middle - 1] + ordered[middle]) / 2 + p95 = ordered[math.ceil(0.95 * len(ordered)) - 1] + return {"median": median, "p95": p95} + + +def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _publish_staging(staging: Path, output: Path) -> None: + """Rename a complete staging tree and roll back a failed parent fsync.""" + staging.rename(output) + try: + _fsync_directory(output.parent) + except BaseException: + output.rename(staging) + _fsync_directory(output.parent) + raise + + +def _record_for_pair( + runtime: QwenPDDInferenceRuntime, + pair: PromptPair, + *, + trajectory: dict[str, Any], + observations: list[Observation], + image_reference: str, + image_sha256: str, +) -> dict[str, Any]: + scheduler = [item.scheduler_calls for item in observations] + actual = [item.transformer_calls for item in observations] + transformer = [item.transformer_seconds for item in observations] + end_to_end = [item.end_to_end_seconds for item in observations] + throughput = [1.0 / value for value in end_to_end] + memory = [item.peak_device_memory_bytes for item in observations] + expected = len(runtime.config.inference_blocks) + if scheduler != [0] * len(observations) or actual != [expected] * len(observations): + raise RuntimeError("measured repetition counters are inconsistent.") + numeric_memory = [value for value in memory if value is not None] + if numeric_memory and len(numeric_memory) != len(memory): + raise RuntimeError("peak-memory observations mix CPU and CUDA domains.") + return { + "prompt_id": pair.prompt_id, + "prompt_sha256": hashlib.sha256(pair.prompt.encode("utf-8")).hexdigest(), + "seed": pair.seed, + "raw_noise_sha256": trajectory["raw_noise_sha256"], + "initial_state_sha256": trajectory["initial_state_sha256"], + "requested_scheduler_steps": expected, + "logical_pdd_blocks": list(runtime.config.inference_blocks), + "logical_pdd_block_count": expected, + "observed_scheduler_step_calls": scheduler, + "actual_transformer_invocations": actual, + "batch_normalized_transformer_evaluations": list(actual), + "transformer_latency_seconds": transformer, + "end_to_end_latency_seconds": end_to_end, + "throughput_images_per_second": throughput, + "peak_device_memory_bytes": memory, + "summaries": { + "transformer_latency_seconds": _summary(transformer), + "end_to_end_latency_seconds": _summary(end_to_end), + "throughput_images_per_second": _summary(throughput), + "peak_device_memory_bytes": _summary(numeric_memory) if numeric_memory else None, + }, + "output": {"path": image_reference, "sha256": image_sha256}, + } + + +@torch.no_grad() +def main() -> None: + args = _parse_args() + if type(args.warmup_runs) is not int or args.warmup_runs < 1: + raise ValueError("warmup_runs must be a positive integer.") + if type(args.measured_runs) is not int or args.measured_runs < 1: + raise ValueError("measured_runs must be a positive integer.") + if args.height < 1 or args.width < 1 or args.max_sequence_length < 1: + raise ValueError("height, width, and max_sequence_length must be positive.") + prompt_manifest = load_canonical_json(args.prompts) + pairs = _prompt_pairs(prompt_manifest) + prompt_manifest_sha256 = sha256_file(args.prompts) + output, result_reference, staging = _resolve_output_paths(args.output_dir, args.result_json) + try: + runtime = load_qwen_pdd_runtime(args.export_dir, args.schedule, args.device) + records: list[dict[str, Any]] = [] + grid_identity: dict[str, Any] | None = None + for pair in pairs: + raw_noise = runtime.make_raw_noise(seed=pair.seed, height=args.height, width=args.width) + trajectory = runtime.trajectory_identity(raw_noise) + if grid_identity is None: + grid_identity = trajectory + elif any( + trajectory[key] != grid_identity[key] + for key in ( + "full_time_nodes", + "full_time_nodes_sha256", + "boundary_indices", + "boundary_time_nodes", + "boundary_time_nodes_sha256", + "first_sigma", + ) + ): + raise RuntimeError("PDD trajectory grid changed between prompt/seed pairs.") + for _ in range(args.warmup_runs): + _run_repetition(runtime, pair.prompt, raw_noise, args.max_sequence_length) + observations = [ + _run_repetition(runtime, pair.prompt, raw_noise, args.max_sequence_length) + for _ in range(args.measured_runs) + ] + image_reference = f"images/{pair.prompt_id}/{pair.seed}.png" + image_path = staging / image_reference + save_png(image_path, observations[0].image) + records.append( + _record_for_pair( + runtime, + pair, + trajectory=trajectory, + observations=observations, + image_reference=image_reference, + image_sha256=sha256_file(image_path), + ) + ) + if grid_identity is None: + raise RuntimeError("evaluation produced no trajectory identity.") + dtype_name = str(runtime.dtype).removeprefix("torch.") + result = { + "schema_version": 1, + "record_type": "pdd_qwen_evaluation", + "identity": { + "export_manifest_sha256": sha256_file(runtime.descriptor.root / "manifest.json"), + "prompt_manifest_sha256": prompt_manifest_sha256, + "model": dict(runtime.model_identity), + "schedule": args.schedule, + "grid": { + "grid_size": runtime.config.grid_size, + "grid_max_t": runtime.config.grid_max_t, + "flow_shift": runtime.config.flow_shift, + "full_time_nodes": grid_identity["full_time_nodes"], + "full_time_nodes_sha256": grid_identity["full_time_nodes_sha256"], + "boundary_indices": grid_identity["boundary_indices"], + "boundary_time_nodes": grid_identity["boundary_time_nodes"], + "boundary_time_nodes_sha256": grid_identity["boundary_time_nodes_sha256"], + "first_sigma": grid_identity["first_sigma"], + }, + }, + "protocol": { + "height": args.height, + "width": args.width, + "max_sequence_length": args.max_sequence_length, + "batch_size": 1, + "warmup_runs": args.warmup_runs, + "measured_runs": args.measured_runs, + "device": str(runtime.device), + "dtype": dtype_name, + "end_to_end_scope": "prompt_encode_through_vae_postprocess", + "transformer_scope": "sum_of_root_student_forward_calls", + "cuda_synchronize": runtime.device.type == "cuda", + "tensor_hash_schema": "pdd_tensor_sha256_v1", + }, + "records": records, + } + staged_result = staging / result_reference + staged_result.parent.mkdir(parents=True, exist_ok=True) + write_canonical_json(staged_result, result) + for directory in sorted( + (path for path in staging.rglob("*") if path.is_dir()), + key=lambda path: len(path.parts), + reverse=True, + ): + _fsync_directory(directory) + _fsync_directory(staging) + _publish_staging(staging, output) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + print(output / result_reference) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index ea9dca44ea9..cabbf9a86fe 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -20,17 +20,18 @@ import argparse import hashlib import math -import os import sys import time -import uuid -from collections.abc import Mapping +from contextlib import ExitStack from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import torch from torch import nn +if TYPE_CHECKING: + from collections.abc import Mapping + sys.dont_write_bytecode = True _THIS_DIR = Path(__file__).resolve().parent @@ -40,6 +41,22 @@ if str(path) not in sys.path: sys.path.insert(0, str(path)) +from pdd.inference_runtime import ( # noqa: E402 + _model_identity, + _normalize_prompt_condition, + _validate_qwen_projection, + build_pdd_student, + load_qwen_pdd_runtime, + save_png, +) + +__all__ = [ + "_model_identity", + "_normalize_prompt_condition", + "_validate_qwen_projection", + "build_pdd_student", +] + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) @@ -57,234 +74,40 @@ def _parse_args() -> argparse.Namespace: return parser.parse_args() -def _dtype_from_name(name: Any) -> torch.dtype: - if not isinstance(name, str): - raise ValueError("PDD export model dtype must be a string.") - dtypes = { - "bfloat16": torch.bfloat16, - "float16": torch.float16, - "float32": torch.float32, - } - try: - return dtypes[name] - except KeyError as error: - raise ValueError(f"PDD inference does not support model dtype {name!r}.") from error - - -def _model_identity(descriptor: Any) -> Mapping[str, Any]: - from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION - - identity = descriptor.manifest.get("identity") - if not isinstance(identity, Mapping): - raise RuntimeError("PDD export has no identity mapping.") - if identity.get("qwen_image") != {"execution": QWEN_IMAGE_PDD_EXECUTION}: - raise RuntimeError("PDD export has an incompatible Qwen execution identity.") - model = identity.get("model") - if not isinstance(model, Mapping) or set(model) != {"id", "revision", "dtype"}: - raise RuntimeError("PDD export model identity is malformed.") - if not isinstance(model["id"], str) or not model["id"]: - raise RuntimeError("PDD export model ID is invalid.") - revision = model["revision"] - if ( - not isinstance(revision, str) - or len(revision) != 40 - or any(character not in "0123456789abcdef" for character in revision) - ): - raise RuntimeError("PDD export model revision must be an exact lowercase commit.") - return model - - -def _validate_qwen_projection(student: nn.Module, metadata: Any) -> nn.Linear: - """Validate the ordinary Qwen projection before widening it for PDD.""" - try: - base_projection = student.get_submodule("proj_out") - except AttributeError as error: - raise RuntimeError("reconstructed Qwen student has no proj_out linear layer.") from error - in_channels = getattr(getattr(student, "config", None), "in_channels", None) - if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: - raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") - if ( - not isinstance(base_projection, nn.Linear) - or base_projection.in_features != metadata.projection_in_features - or base_projection.out_features != metadata.projection_out_features - or (base_projection.bias is not None) != metadata.projection_bias - ): - raise RuntimeError("reconstructed Qwen proj_out does not match the export metadata.") - if base_projection.out_features != in_channels: - raise RuntimeError( - "Qwen proj_out width must equal transformer in_channels for 2x2 latent packing." - ) - return base_projection - - -def build_pdd_student(export_dir: str | Path) -> tuple[nn.Module, Any, torch.dtype]: - """Reconstruct and strictly load the converted Qwen student on CPU.""" - from diffusers import QwenImageTransformer2DModel - - from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - adopt_qwen_image_mr210_forward, - convert_qwen_image_to_pdd, - ) - from pdd.export import inspect_pdd_export, load_pdd_export_into_model, pdd_config_from_metadata - - descriptor = inspect_pdd_export(export_dir) - model_identity = _model_identity(descriptor) - dtype = _dtype_from_name(model_identity["dtype"]) - student = QwenImageTransformer2DModel.from_config(dict(descriptor.transformer_config)) - metadata = descriptor.metadata - _validate_qwen_projection(student, metadata) - student = adopt_qwen_image_mr210_forward(student) - config = pdd_config_from_metadata(metadata, blocks=metadata.inference_blocks) - convert_qwen_image_to_pdd(student, config) - descriptor = load_pdd_export_into_model(export_dir, student) - student.to(dtype=dtype) - return student, descriptor, dtype - - -def _normalize_prompt_condition( - prompt_embeds: Any, - prompt_mask: Any, - *, - device: torch.device, - dtype: torch.dtype, -) -> tuple[torch.Tensor, torch.Tensor]: - """Normalize the pinned Diffusers Qwen prompt-encoding contract for PDD.""" - if not isinstance(prompt_embeds, torch.Tensor) or prompt_embeds.ndim != 3: - raise RuntimeError("Qwen prompt embeddings must have shape [B, S, D].") - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - expected_shape = prompt_embeds.shape[:2] - if prompt_mask is None: - prompt_mask = torch.ones(expected_shape, device=device, dtype=torch.long) - elif not isinstance(prompt_mask, torch.Tensor) or prompt_mask.ndim != 2: - raise RuntimeError("Qwen prompt mask must have shape [B, S] or be None.") - elif tuple(prompt_mask.shape) != tuple(expected_shape): - raise RuntimeError("Qwen prompt mask shape does not match prompt embeddings.") - elif prompt_mask.dtype.is_floating_point or prompt_mask.dtype.is_complex: - raise RuntimeError("Qwen prompt mask must use an integer or boolean dtype.") - else: - prompt_mask = prompt_mask.to(device=device, dtype=torch.long) - return prompt_embeds, prompt_mask - - -def _latent_shape(pipe: Any, *, height: int, width: int) -> tuple[int, int, int, int]: - if type(height) is not int or type(width) is not int or height <= 0 or width <= 0: - raise ValueError("height and width must be positive integers.") - quantum = int(pipe.vae_scale_factor) * 2 - if height % quantum or width % quantum: - raise ValueError(f"height and width must be divisible by {quantum}.") - in_channels = getattr(pipe.transformer.config, "in_channels", None) - if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: - raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") - latent_height = 2 * (height // quantum) - latent_width = 2 * (width // quantum) - return 1, in_channels // 4, latent_height, latent_width - - -def _decode_qwen_latents(pipe: Any, latents: torch.Tensor) -> list[Any]: - if latents.ndim != 4: - raise ValueError("PDD Qwen latents must have shape [B, C, H, W].") - vae = pipe.vae - mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype) - std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype) - if mean.numel() != latents.shape[1] or std.numel() != latents.shape[1]: - raise RuntimeError("Qwen VAE latent statistics do not match the student channels.") - decoded_input = latents.unsqueeze(2) * std.view(1, -1, 1, 1, 1) - decoded_input = decoded_input + mean.view(1, -1, 1, 1, 1) - decoded = vae.decode(decoded_input, return_dict=False)[0] - if decoded.ndim != 5 or decoded.shape[2] != 1: - raise RuntimeError("Qwen VAE must return one-frame 5D image tensors.") - return pipe.image_processor.postprocess(decoded[:, :, 0], output_type="pil") - - -def _save_png(path: Path, image: Any) -> None: - if path.is_symlink(): - raise ValueError("PDD inference output cannot be a symlink.") - path = path.resolve() - if path.suffix.lower() != ".png": - raise ValueError("PDD inference output must use a .png suffix.") - path.parent.mkdir(parents=True, exist_ok=True) - if path.exists() or path.is_symlink(): - raise FileExistsError(f"PDD inference output already exists: {path}.") - staging = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") - try: - image.save(staging, format="PNG") - with staging.open("rb") as stream: - os.fsync(stream.fileno()) - staging.rename(path) - descriptor = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - finally: - staging.unlink(missing_ok=True) - - -@torch.no_grad() -def main() -> None: - args = _parse_args() - from diffusers import QwenImagePipeline - - from modelopt.torch.fastgen import PDDPipeline - from modelopt.torch.fastgen.plugins.qwen_image_pdd import QwenImagePDDAdapter - from pdd.artifacts import sha256_file, write_canonical_json - from pdd.export import pdd_config_from_metadata - - if args.output.is_symlink() or args.result_json.is_symlink(): +def _resolve_outputs(output_value: Path, result_value: Path) -> tuple[Path, Path, str]: + if output_value.is_symlink() or result_value.is_symlink(): raise ValueError("PDD output and result JSON cannot be symlinks.") - output = args.output.resolve() - result_json = args.result_json.resolve() + output = output_value.resolve() + result_json = result_value.resolve() if output.exists() or output.is_symlink(): raise FileExistsError(f"PDD inference output already exists: {output}.") if result_json.exists() or result_json.is_symlink(): raise FileExistsError(f"PDD result JSON already exists: {result_json}.") try: - output_reference = output.relative_to(result_json.parent).as_posix() + reference = output.relative_to(result_json.parent).as_posix() except ValueError as error: raise ValueError("PDD output must be beneath the result JSON directory.") from error + return output, result_json, reference + + +@torch.no_grad() +def main() -> None: + args = _parse_args() + from pdd.artifacts import sha256_file, write_canonical_json + + output, result_json, output_reference = _resolve_outputs(args.output, args.result_json) if not isinstance(args.prompt_id, str) or not args.prompt_id.strip(): raise ValueError("prompt_id must be non-empty.") if args.seed < 0 or args.seed >= 2**63: raise ValueError("seed must be in [0, 2**63).") if args.max_sequence_length < 1: raise ValueError("max_sequence_length must be positive.") - student, descriptor, dtype = build_pdd_student(args.export_dir) - model_identity = _model_identity(descriptor) - device = torch.device(args.device) - if device.type == "cuda" and not torch.cuda.is_available(): - raise RuntimeError("CUDA was requested but is unavailable.") - - student.to(device=device) - pipe = QwenImagePipeline.from_pretrained( - model_identity["id"], - revision=model_identity["revision"], - transformer=student, - torch_dtype=dtype, - use_safetensors=True, - ) - if pipe.transformer is not student: - raise RuntimeError("Qwen pipeline did not retain the adopted PDD transformer.") - pipe.to(device) - config = pdd_config_from_metadata(descriptor.metadata, schedule=args.schedule) - adapter = QwenImagePDDAdapter(config, compute_dtype=dtype) - sampler = PDDPipeline(student, nn.Identity(), config, adapter) - prompt_embeds, prompt_mask = pipe.encode_prompt( - prompt=args.prompt, - device=device, - num_images_per_prompt=1, - max_sequence_length=args.max_sequence_length, - ) - condition = _normalize_prompt_condition( - prompt_embeds, - prompt_mask, - device=device, - dtype=dtype, - ) - generator = torch.Generator(device=device).manual_seed(args.seed) - shape = _latent_shape(pipe, height=args.height, width=args.width) - noise = torch.randn(shape, generator=generator, device=device, dtype=torch.float32) + runtime = load_qwen_pdd_runtime(args.export_dir, args.schedule, args.device) + condition = runtime.encode_prompt(args.prompt, args.max_sequence_length) + noise = runtime.make_raw_noise(seed=args.seed, height=args.height, width=args.width) transformer_invocations = 0 + scheduler_step_calls = 0 def count_invocation( _module: nn.Module, _args: tuple[Any, ...], _kwargs: Mapping[str, Any] @@ -292,29 +115,50 @@ def count_invocation( nonlocal transformer_invocations transformer_invocations += 1 - hook = student.register_forward_pre_hook(count_invocation, with_kwargs=True) - if device.type == "cuda": - torch.cuda.synchronize(device) - started = time.perf_counter() - try: - sampled = sampler.sample(noise, condition=condition) - finally: - hook.remove() - images = _decode_qwen_latents(pipe, sampled.to(dtype)) - if device.type == "cuda": - torch.cuda.synchronize(device) - latency = time.perf_counter() - started - expected_invocations = len(config.inference_blocks) + scheduler = runtime.scheduler + scheduler_state = vars(scheduler).get("step") + scheduler_had_instance_step = "step" in vars(scheduler) + original_scheduler_step = scheduler.step + + def counted_scheduler_step(*call_args: Any, **call_kwargs: Any) -> Any: + nonlocal scheduler_step_calls + scheduler_step_calls += 1 + return original_scheduler_step(*call_args, **call_kwargs) + + def restore_scheduler_step() -> None: + if scheduler_had_instance_step: + setattr(scheduler, "step", scheduler_state) + elif "step" in vars(scheduler): + delattr(scheduler, "step") + + with ExitStack() as cleanup: + cleanup.callback(restore_scheduler_step) + setattr(scheduler, "step", counted_scheduler_step) + hook = runtime.student.register_forward_pre_hook(count_invocation, with_kwargs=True) + cleanup.callback(hook.remove) + if runtime.device.type == "cuda": + torch.cuda.synchronize(runtime.device) + started = time.perf_counter() + images = runtime.sample_decode(condition, noise) + if runtime.device.type == "cuda": + torch.cuda.synchronize(runtime.device) + latency = time.perf_counter() - started + + expected_invocations = len(runtime.config.inference_blocks) if transformer_invocations != expected_invocations: raise RuntimeError( f"PDD sampler made {transformer_invocations} transformer calls; " f"expected {expected_invocations}." ) + if scheduler_step_calls != 0: + raise RuntimeError( + f"PDD sampler unexpectedly called scheduler.step {scheduler_step_calls} times." + ) if len(images) != 1: raise RuntimeError(f"PDD single-prompt inference returned {len(images)} images.") if not math.isfinite(latency) or latency <= 0: raise RuntimeError("PDD inference latency measurement is invalid.") - _save_png(output, images[0]) + save_png(output, images[0]) result_json.parent.mkdir(parents=True, exist_ok=True) result = { @@ -325,12 +169,13 @@ def count_invocation( "prompt_sha256": hashlib.sha256(args.prompt.encode("utf-8")).hexdigest(), "seed": args.seed, "schedule": args.schedule, - "blocks": list(config.inference_blocks), + "blocks": list(runtime.config.inference_blocks), "height": args.height, "width": args.width, - "export_manifest_sha256": sha256_file(descriptor.root / "manifest.json"), + "export_manifest_sha256": sha256_file(runtime.descriptor.root / "manifest.json"), "output": {"path": output_reference, "sha256": sha256_file(output)}, "scheduler_steps": expected_invocations, + "observed_scheduler_step_calls": scheduler_step_calls, "actual_transformer_invocations": transformer_invocations, "batch_normalized_transformer_evaluations": transformer_invocations, "latency_seconds": latency, diff --git a/examples/diffusers/fastgen/pdd/inference_runtime.py b/examples/diffusers/fastgen/pdd/inference_runtime.py new file mode 100644 index 00000000000..da578e5b547 --- /dev/null +++ b/examples/diffusers/fastgen/pdd/inference_runtime.py @@ -0,0 +1,354 @@ +# 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. + +"""Shared authenticated Qwen-Image PDD inference runtime.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from torch import nn + +if TYPE_CHECKING: + from pathlib import Path + +from .artifacts import canonical_json_bytes +from .export import PDD_INFERENCE_SCHEDULES, pdd_config_from_metadata + +_TENSOR_HASH_DOMAINS = { + "raw_noise", + "initial_state", + "full_time_nodes", + "boundary_time_nodes", +} + + +def _dtype_from_name(name: Any) -> torch.dtype: + if not isinstance(name, str): + raise ValueError("PDD export model dtype must be a string.") + dtypes = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + } + try: + return dtypes[name] + except KeyError as error: + raise ValueError(f"PDD inference does not support model dtype {name!r}.") from error + + +def _model_identity(descriptor: Any) -> Mapping[str, Any]: + from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION + + identity = descriptor.manifest.get("identity") + if not isinstance(identity, Mapping): + raise RuntimeError("PDD export has no identity mapping.") + if identity.get("qwen_image") != {"execution": QWEN_IMAGE_PDD_EXECUTION}: + raise RuntimeError("PDD export has an incompatible Qwen execution identity.") + model = identity.get("model") + if not isinstance(model, Mapping) or set(model) != {"id", "revision", "dtype"}: + raise RuntimeError("PDD export model identity is malformed.") + if not isinstance(model["id"], str) or not model["id"]: + raise RuntimeError("PDD export model ID is invalid.") + revision = model["revision"] + if ( + not isinstance(revision, str) + or len(revision) != 40 + or any(character not in "0123456789abcdef" for character in revision) + ): + raise RuntimeError("PDD export model revision must be an exact lowercase commit.") + return model + + +def _validate_qwen_projection(student: nn.Module, metadata: Any) -> nn.Linear: + """Validate the ordinary Qwen projection before widening it for PDD.""" + try: + base_projection = student.get_submodule("proj_out") + except AttributeError as error: + raise RuntimeError("reconstructed Qwen student has no proj_out linear layer.") from error + in_channels = getattr(getattr(student, "config", None), "in_channels", None) + if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: + raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") + if ( + not isinstance(base_projection, nn.Linear) + or base_projection.in_features != metadata.projection_in_features + or base_projection.out_features != metadata.projection_out_features + or (base_projection.bias is not None) != metadata.projection_bias + ): + raise RuntimeError("reconstructed Qwen proj_out does not match the export metadata.") + if base_projection.out_features != in_channels: + raise RuntimeError( + "Qwen proj_out width must equal transformer in_channels for 2x2 latent packing." + ) + return base_projection + + +def build_pdd_student( + export_dir: str | Path, *, schedule: str = "pdd-4" +) -> tuple[nn.Module, Any, torch.dtype]: + """Reconstruct and strictly load a converted Qwen student on CPU.""" + from diffusers import QwenImageTransformer2DModel + + from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( + adopt_qwen_image_mr210_forward, + convert_qwen_image_to_pdd, + ) + + from .export import inspect_pdd_export, load_pdd_export_into_model + + if schedule not in PDD_INFERENCE_SCHEDULES: + raise ValueError( + f"Unknown PDD schedule {schedule!r}; expected {sorted(PDD_INFERENCE_SCHEDULES)}." + ) + descriptor = inspect_pdd_export(export_dir) + model_identity = _model_identity(descriptor) + dtype = _dtype_from_name(model_identity["dtype"]) + student = QwenImageTransformer2DModel.from_config(dict(descriptor.transformer_config)) + _validate_qwen_projection(student, descriptor.metadata) + student = adopt_qwen_image_mr210_forward(student) + config = pdd_config_from_metadata(descriptor.metadata, schedule=schedule) + convert_qwen_image_to_pdd(student, config) + descriptor = load_pdd_export_into_model(export_dir, student) + student.to(dtype=dtype) + return student, descriptor, dtype + + +def _normalize_prompt_condition( + prompt_embeds: Any, + prompt_mask: Any, + *, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance(prompt_embeds, torch.Tensor) or prompt_embeds.ndim != 3: + raise RuntimeError("Qwen prompt embeddings must have shape [B, S, D].") + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + expected_shape = prompt_embeds.shape[:2] + if prompt_mask is None: + prompt_mask = torch.ones(expected_shape, device=device, dtype=torch.long) + elif not isinstance(prompt_mask, torch.Tensor) or prompt_mask.ndim != 2: + raise RuntimeError("Qwen prompt mask must have shape [B, S] or be None.") + elif tuple(prompt_mask.shape) != tuple(expected_shape): + raise RuntimeError("Qwen prompt mask shape does not match prompt embeddings.") + elif prompt_mask.dtype.is_floating_point or prompt_mask.dtype.is_complex: + raise RuntimeError("Qwen prompt mask must use an integer or boolean dtype.") + else: + prompt_mask = prompt_mask.to(device=device, dtype=torch.long) + return prompt_embeds, prompt_mask + + +def _latent_shape(pipe: Any, *, height: int, width: int) -> tuple[int, int, int, int]: + if type(height) is not int or type(width) is not int or height <= 0 or width <= 0: + raise ValueError("height and width must be positive integers.") + quantum = int(pipe.vae_scale_factor) * 2 + if height % quantum or width % quantum: + raise ValueError(f"height and width must be divisible by {quantum}.") + in_channels = getattr(pipe.transformer.config, "in_channels", None) + if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: + raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") + return 1, in_channels // 4, 2 * (height // quantum), 2 * (width // quantum) + + +def _decode_qwen_latents(pipe: Any, latents: torch.Tensor) -> list[Any]: + if latents.ndim != 4: + raise ValueError("PDD Qwen latents must have shape [B, C, H, W].") + vae = pipe.vae + mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype) + std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype) + if mean.numel() != latents.shape[1] or std.numel() != latents.shape[1]: + raise RuntimeError("Qwen VAE latent statistics do not match the student channels.") + decoded_input = latents.unsqueeze(2) * std.view(1, -1, 1, 1, 1) + decoded_input = decoded_input + mean.view(1, -1, 1, 1, 1) + decoded = vae.decode(decoded_input, return_dict=False)[0] + if decoded.ndim != 5 or decoded.shape[2] != 1: + raise RuntimeError("Qwen VAE must return one-frame 5D image tensors.") + return pipe.image_processor.postprocess(decoded[:, :, 0], output_type="pil") + + +def pdd_tensor_sha256(tensor: torch.Tensor, domain: str) -> str: + """Hash one exact FP32 tensor using the evaluation protocol.""" + if domain not in _TENSOR_HASH_DOMAINS: + raise ValueError(f"unknown PDD tensor hash domain {domain!r}.") + if not isinstance(tensor, torch.Tensor) or tensor.dtype != torch.float32: + raise TypeError("PDD tensor hashing requires a float32 tensor.") + if not torch.isfinite(tensor).all().item(): + raise FloatingPointError("PDD tensor hashing rejects non-finite values.") + array = np.ascontiguousarray(tensor.detach().cpu().numpy(), dtype=" None: + """Publish one PNG exclusively and durably.""" + if path.is_symlink(): + raise ValueError("PDD inference output cannot be a symlink.") + path = path.resolve() + if path.suffix.lower() != ".png": + raise ValueError("PDD inference output must use a .png suffix.") + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists() or path.is_symlink(): + raise FileExistsError(f"PDD inference output already exists: {path}.") + staging = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + try: + image.save(staging, format="PNG") + with staging.open("rb") as stream: + os.fsync(stream.fileno()) + staging.rename(path) + descriptor = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + finally: + staging.unlink(missing_ok=True) + + +@dataclass(frozen=True) +class QwenPDDInferenceRuntime: + """One loaded Qwen pipeline and one authenticated PDD sampler.""" + + student: nn.Module + scheduler: Any + descriptor: Any + model_identity: Mapping[str, Any] + dtype: torch.dtype + device: torch.device + config: Any + pipe: Any + sampler: Any + + def encode_prompt(self, prompt: str, max_sequence_length: int) -> Any: + if not isinstance(prompt, str) or not prompt: + raise ValueError("prompt must be a non-empty string.") + if type(max_sequence_length) is not int or max_sequence_length < 1: + raise ValueError("max_sequence_length must be positive.") + prompt_embeds, prompt_mask = self.pipe.encode_prompt( + prompt=prompt, + device=self.device, + num_images_per_prompt=1, + max_sequence_length=max_sequence_length, + ) + return _normalize_prompt_condition( + prompt_embeds, + prompt_mask, + device=self.device, + dtype=self.dtype, + ) + + def make_raw_noise(self, *, seed: int, height: int, width: int) -> torch.Tensor: + if type(seed) is not int or seed < 0 or seed >= 2**63: + raise ValueError("seed must be in [0, 2**63).") + generator = torch.Generator(device=self.device).manual_seed(seed) + shape = _latent_shape(self.pipe, height=height, width=width) + return torch.randn( + shape, + generator=generator, + device=self.device, + dtype=torch.float32, + ) + + def sample_decode(self, condition: Any, raw_noise: torch.Tensor) -> list[Any]: + sampled = self.sampler.sample(raw_noise, condition=condition) + return _decode_qwen_latents(self.pipe, sampled.to(self.dtype)) + + def trajectory_identity(self, raw_noise: torch.Tensor) -> dict[str, Any]: + full = self.sampler.time_grid(raw_noise.device).to(device="cpu", dtype=torch.float32) + boundaries = [0] + for block in self.config.inference_blocks: + boundaries.append(boundaries[-1] + block) + boundary = full[boundaries] + initial = (raw_noise.to(torch.float64) * self.config.grid_max_t).to(torch.float32) + return { + "raw_noise_sha256": pdd_tensor_sha256(raw_noise, "raw_noise"), + "initial_state_sha256": pdd_tensor_sha256(initial, "initial_state"), + "full_time_nodes": full.tolist(), + "full_time_nodes_sha256": pdd_tensor_sha256(full, "full_time_nodes"), + "boundary_indices": boundaries, + "boundary_time_nodes": boundary.tolist(), + "boundary_time_nodes_sha256": pdd_tensor_sha256(boundary, "boundary_time_nodes"), + "first_sigma": float(full[0].item()), + } + + +def load_qwen_pdd_runtime( + export_dir: str | Path, schedule: str, device: str | torch.device +) -> QwenPDDInferenceRuntime: + """Load one authenticated Qwen PDD runtime for a source-owned schedule.""" + from diffusers import QwenImagePipeline + + from modelopt.torch.fastgen import PDDPipeline + from modelopt.torch.fastgen.plugins.qwen_image_pdd import QwenImagePDDAdapter + + if schedule not in PDD_INFERENCE_SCHEDULES: + raise ValueError( + f"Unknown PDD schedule {schedule!r}; expected {sorted(PDD_INFERENCE_SCHEDULES)}." + ) + resolved_device = torch.device(device) + if resolved_device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable.") + student, descriptor, dtype = build_pdd_student(export_dir, schedule=schedule) + model_identity = _model_identity(descriptor) + student.to(device=resolved_device) + pipe = QwenImagePipeline.from_pretrained( + model_identity["id"], + revision=model_identity["revision"], + transformer=student, + torch_dtype=dtype, + use_safetensors=True, + ) + if pipe.transformer is not student: + raise RuntimeError("Qwen pipeline did not retain the adopted PDD transformer.") + pipe.to(resolved_device) + scheduler = getattr(pipe, "scheduler", None) + if scheduler is None or not callable(getattr(scheduler, "step", None)): + raise RuntimeError("Qwen pipeline scheduler does not expose a callable step method.") + config = pdd_config_from_metadata(descriptor.metadata, schedule=schedule) + if tuple(config.inference_blocks) != PDD_INFERENCE_SCHEDULES[schedule]: + raise RuntimeError("authenticated PDD schedule changed after validation.") + sampler = PDDPipeline( + student, + nn.Identity(), + config, + QwenImagePDDAdapter(config, compute_dtype=dtype), + ) + return QwenPDDInferenceRuntime( + student=student, + scheduler=scheduler, + descriptor=descriptor, + model_identity=model_identity, + dtype=dtype, + device=resolved_device, + config=config, + pipe=pipe, + sampler=sampler, + ) diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index 6f96be4643f..4724604c5d9 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -57,7 +57,9 @@ "data.py", "export.py", "export_qwen_image.py", + "evaluate_qwen_image.py", "finetune.py", + "inference_runtime.py", "inference_qwen_image.py", "recipe.py", "training.py", diff --git a/tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py b/tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py new file mode 100644 index 00000000000..a9fa4660847 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py @@ -0,0 +1,629 @@ +# 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. + +"""Tests for the Qwen-Image PDD evaluation runner.""" + +from __future__ import annotations + +import hashlib +import json +import pathlib +import sys +import time +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image +from torch import nn + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +for path in (_REPO_ROOT, _FASTGEN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from pdd.artifacts import canonical_json_bytes, load_canonical_json, write_canonical_json +from pdd.evaluate_qwen_image import ( + _prompt_pairs, + _publish_staging, + _resolve_output_paths, + _run_repetition, + _summary, + main, +) +from pdd.export import PDD_INFERENCE_SCHEDULES +from pdd.inference_qwen_image import main as inference_main +from pdd.inference_runtime import QwenPDDInferenceRuntime, pdd_tensor_sha256 + + +class _Scheduler: + def step(self, value=0): + return value + + +class _Student(nn.Module): + def forward(self, value): + time.sleep(0.0001) + return value + 1 + + +class _FakeRuntime: + def __init__(self, export_root: pathlib.Path, blocks=(1, 1)) -> None: + self.student = _Student() + self.scheduler = _Scheduler() + self.device = torch.device("cpu") + self.dtype = torch.bfloat16 + self.config = SimpleNamespace( + inference_blocks=list(blocks), + grid_size=sum(blocks), + grid_max_t=0.999, + flow_shift=5.0, + ) + self.descriptor = SimpleNamespace(root=export_root) + self.model_identity = { + "id": "Qwen/Qwen-Image", + "revision": "f" * 40, + "dtype": "bfloat16", + } + self.encode_calls = 0 + self.sample_calls = 0 + + def encode_prompt(self, prompt, max_sequence_length): + assert prompt and max_sequence_length > 0 + self.encode_calls += 1 + return torch.tensor(0) + + def make_raw_noise(self, *, seed, height, width): + return torch.randn((1, 1, height, width), generator=torch.Generator().manual_seed(seed)) + + def sample_decode(self, condition, raw_noise): + del condition, raw_noise + self.sample_calls += 1 + value = torch.tensor(0) + for _ in self.config.inference_blocks: + value = self.student(value) + return [Image.new("RGB", (2, 2), color=(int(value), 0, 0))] + + def trajectory_identity(self, raw_noise): + full = torch.linspace(0.999, 0.0, self.config.grid_size + 1, dtype=torch.float32) + boundaries = [0] + for block in self.config.inference_blocks: + boundaries.append(boundaries[-1] + block) + boundary = full[boundaries] + initial = (raw_noise.to(torch.float64) * self.config.grid_max_t).to(torch.float32) + return { + "raw_noise_sha256": pdd_tensor_sha256(raw_noise, "raw_noise"), + "initial_state_sha256": pdd_tensor_sha256(initial, "initial_state"), + "full_time_nodes": full.tolist(), + "full_time_nodes_sha256": pdd_tensor_sha256(full, "full_time_nodes"), + "boundary_indices": boundaries, + "boundary_time_nodes": boundary.tolist(), + "boundary_time_nodes_sha256": pdd_tensor_sha256(boundary, "boundary_time_nodes"), + "first_sigma": float(full[0]), + } + + +def test_prompt_manifest_expands_exact_order_and_rejects_unsafe_data() -> None: + value = { + "schema_version": 1, + "prompts": [ + {"prompt_id": "a", "prompt": "alpha", "seeds": [1, 2]}, + {"prompt_id": "b-2", "prompt": "beta", "seeds": [3]}, + ], + } + assert [(pair.prompt_id, pair.seed) for pair in _prompt_pairs(value)] == [ + ("a", 1), + ("a", 2), + ("b-2", 3), + ] + value["prompts"][1]["prompt_id"] = "../escape" + with pytest.raises(ValueError, match="safe path component"): + _prompt_pairs(value) + + +@pytest.mark.parametrize( + ("value", "message"), + [ + ({"schema_version": 1, "prompts": [], "extra": 1}, "exactly"), + ({"schema_version": 2, "prompts": []}, "schema_version"), + ({"schema_version": True, "prompts": []}, "schema_version"), + ( + { + "schema_version": 1, + "prompts": [ + {"prompt_id": "b", "prompt": "one", "seeds": [1]}, + {"prompt_id": "a", "prompt": "two", "seeds": [2]}, + ], + }, + "lexicographically", + ), + ( + { + "schema_version": 1, + "prompts": [{"prompt_id": "a", "prompt": "one", "seeds": [2, 1]}], + }, + "sorted and unique", + ), + ( + { + "schema_version": 1, + "prompts": [{"prompt_id": "a", "prompt": "one", "seeds": [True]}], + }, + "invalid seed", + ), + ], +) +def test_prompt_manifest_expected_red_matrix(value, message) -> None: + with pytest.raises(ValueError, match=message): + _prompt_pairs(value) + + +def test_tensor_hash_uses_exact_header_little_endian_payload_and_domain() -> None: + tensor = torch.tensor([[1.0, -2.5]], dtype=torch.float32) + header = { + "schema_version": 1, + "domain": "raw_noise", + "dtype": "float32", + "shape": [1, 2], + "byte_order": "little", + "order": "C", + } + payload = np.ascontiguousarray(tensor.numpy(), dtype=" None: + requested_devices = [] + + class Sampler: + def time_grid(self, device): + requested_devices.append(device) + return torch.tensor([0.999, 0.5, 0.0], dtype=torch.float32) + + class RawNoise: + device = torch.device("cuda:7") + + def to(self, dtype): + return torch.ones((1, 1), dtype=dtype) + + runtime = QwenPDDInferenceRuntime( + student=None, + scheduler=None, + descriptor=None, + model_identity={}, + dtype=torch.bfloat16, + device=torch.device("cuda:7"), + config=SimpleNamespace(inference_blocks=(1, 1), grid_max_t=0.999), + pipe=None, + sampler=Sampler(), + ) + monkeypatch.setattr("pdd.inference_runtime.pdd_tensor_sha256", lambda _tensor, domain: domain) + identity = runtime.trajectory_identity(RawNoise()) + assert requested_devices == [torch.device("cuda:7")] + assert identity["full_time_nodes"] == pytest.approx([0.999, 0.5, 0.0]) + + +def test_source_owned_schedules_and_summary_contract() -> None: + assert PDD_INFERENCE_SCHEDULES == { + "pdd-2": (64, 64), + "pdd-4": (32, 32, 32, 32), + "pdd-8": (16, 16, 16, 16, 16, 16, 16, 16), + } + assert _summary([4.0, 1.0, 3.0, 2.0]) == {"median": 2.5, "p95": 4.0} + + +def test_repetition_counts_calls_times_cpu_and_restores_scheduler(tmp_path) -> None: + runtime = _FakeRuntime(tmp_path) + original = runtime.scheduler.step + observation = _run_repetition(runtime, "prompt", torch.zeros(1), 8) + assert observation.scheduler_calls == 0 + assert observation.transformer_calls == 2 + assert observation.peak_device_memory_bytes is None + assert observation.transformer_seconds > 0 + assert observation.end_to_end_seconds >= observation.transformer_seconds + assert runtime.scheduler.step == original + assert not runtime.student._forward_pre_hooks + assert not runtime.student._forward_hooks + + +@pytest.mark.parametrize("blocks", [(64, 64), (32, 32, 32, 32), (16,) * 8]) +def test_repetition_counts_each_supported_schedule(tmp_path, blocks) -> None: + runtime = _FakeRuntime(tmp_path, blocks=blocks) + observation = _run_repetition(runtime, "prompt", torch.zeros(1), 8) + assert observation.transformer_calls == len(blocks) + + +def test_repetition_rejects_scheduler_and_transformer_count_collapse(tmp_path, monkeypatch) -> None: + runtime = _FakeRuntime(tmp_path) + + def scheduler_call(_condition, _noise): + runtime.scheduler.step() + value = runtime.student(torch.tensor(0)) + value = runtime.student(value) + return [Image.new("RGB", (2, 2))] + + monkeypatch.setattr(runtime, "sample_decode", scheduler_call) + with pytest.raises(RuntimeError, match=r"scheduler\.step"): + _run_repetition(runtime, "prompt", torch.zeros(1), 8) + assert "step" not in vars(runtime.scheduler) + + def missing_transformer(_condition, _noise): + return [Image.new("RGB", (2, 2))] + + monkeypatch.setattr(runtime, "sample_decode", missing_transformer) + with pytest.raises(RuntimeError, match="transformer calls"): + _run_repetition(runtime, "prompt", torch.zeros(1), 8) + + +def test_mocked_cuda_instrumentation_orders_sync_reset_events_and_memory( + tmp_path, monkeypatch +) -> None: + runtime = _FakeRuntime(tmp_path) + runtime.device = torch.device("cuda") + actions = [] + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing is True + + def record(self): + actions.append("event") + + def elapsed_time(self, _other): + actions.append("elapsed") + return 0.01 + + monkeypatch.setattr(torch.cuda, "Event", Event) + monkeypatch.setattr(torch.cuda, "synchronize", lambda _device: actions.append("sync")) + monkeypatch.setattr( + torch.cuda, "reset_peak_memory_stats", lambda _device: actions.append("reset") + ) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _device: 123) + observation = _run_repetition(runtime, "prompt", torch.zeros(1), 8) + assert actions[:2] == ["sync", "reset"] + assert actions[-3:] == ["sync", "elapsed", "elapsed"] + assert observation.peak_device_memory_bytes == 123 + + +def test_repetition_restores_instrumentation_on_failure(tmp_path, monkeypatch) -> None: + runtime = _FakeRuntime(tmp_path) + original = runtime.scheduler.step + + def fail(_condition, _noise): + runtime.scheduler.step() + raise RuntimeError("boom") + + monkeypatch.setattr(runtime, "sample_decode", fail) + with pytest.raises(RuntimeError, match="boom"): + _run_repetition(runtime, "prompt", torch.zeros(1), 8) + assert runtime.scheduler.step == original + assert not runtime.student._forward_pre_hooks + assert not runtime.student._forward_hooks + + +def test_repetition_restores_after_initial_sync_failure(tmp_path, monkeypatch) -> None: + runtime = _FakeRuntime(tmp_path) + runtime.device = torch.device("cuda") + original = runtime.scheduler.step + + def fail_sync(_device): + raise RuntimeError("sync failed") + + monkeypatch.setattr(torch.cuda, "synchronize", fail_sync) + with pytest.raises(RuntimeError, match="sync failed"): + _run_repetition(runtime, "prompt", torch.zeros(1), 8) + assert runtime.scheduler.step == original + assert not runtime.student._forward_pre_hooks + assert not runtime.student._forward_hooks + + +def test_repetition_restores_after_partial_hook_install_failure(tmp_path, monkeypatch) -> None: + runtime = _FakeRuntime(tmp_path) + original = runtime.scheduler.step + + def fail_post_hook(*_args, **_kwargs): + raise RuntimeError("post-hook failed") + + monkeypatch.setattr(runtime.student, "register_forward_hook", fail_post_hook) + with pytest.raises(RuntimeError, match="post-hook failed"): + _run_repetition(runtime, "prompt", torch.zeros(1), 8) + assert runtime.scheduler.step == original + assert not runtime.student._forward_pre_hooks + assert not runtime.student._forward_hooks + + +def test_atomic_publish_rolls_back_failed_parent_fsync(tmp_path, monkeypatch) -> None: + staging = tmp_path / ".result.staging" + staging.mkdir() + (staging / "complete").write_text("complete") + output = tmp_path / "result" + module = sys.modules["pdd.evaluate_qwen_image"] + original = module._fsync_directory + failed = False + + def fail_once(path): + nonlocal failed + if path == tmp_path and output.exists() and not failed: + failed = True + raise OSError("fsync failed") + original(path) + + monkeypatch.setattr(module, "_fsync_directory", fail_once) + with pytest.raises(OSError, match="fsync failed"): + _publish_staging(staging, output) + assert failed + assert not output.exists() + assert (staging / "complete").is_file() + + +def test_output_transaction_rejects_escape_collision_and_symlink(tmp_path) -> None: + with pytest.raises(ValueError, match="strictly beneath"): + _resolve_output_paths(tmp_path / "output", tmp_path / "outside.json") + existing = tmp_path / "existing" + existing.mkdir() + with pytest.raises(FileExistsError, match="already exists"): + _resolve_output_paths(existing, existing / "result.json") + real = tmp_path / "real" + real.mkdir() + link = tmp_path / "link" + link.symlink_to(real, target_is_directory=True) + with pytest.raises(ValueError, match="symlink"): + _resolve_output_paths(link / "output", link / "output" / "result.json") + output = real / "output" + unresolved_then_link = tmp_path / "missing" / ".." / "link" / "output" + with pytest.raises(ValueError, match="parent traversal"): + _resolve_output_paths(unresolved_then_link, output / "result.json") + unresolved_result = unresolved_then_link / "result.json" + with pytest.raises(ValueError, match="parent traversal"): + _resolve_output_paths(output, unresolved_result) + + +def test_main_publishes_complete_atomic_cpu_result(tmp_path, monkeypatch) -> None: + export = tmp_path / "export" + export.mkdir() + (export / "manifest.json").write_bytes(b"manifest") + prompts = tmp_path / "prompts.json" + write_canonical_json( + prompts, + { + "schema_version": 1, + "prompts": [{"prompt_id": "sample", "prompt": "text", "seeds": [7]}], + }, + ) + runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) + monkeypatch.setattr( + "pdd.evaluate_qwen_image.load_qwen_pdd_runtime", + lambda _export, _schedule, _device: runtime, + ) + output = tmp_path / "evaluation" + result = output / "result.json" + monkeypatch.setattr( + sys, + "argv", + [ + "evaluate_qwen_image.py", + "--export-dir", + str(export), + "--prompts", + str(prompts), + "--schedule", + "pdd-4", + "--output-dir", + str(output), + "--result-json", + str(result), + "--warmup-runs", + "1", + "--measured-runs", + "2", + "--height", + "2", + "--width", + "2", + "--device", + "cpu", + ], + ) + main() + value = load_canonical_json(result) + assert value["record_type"] == "pdd_qwen_evaluation" + assert value["identity"]["schedule"] == "pdd-4" + record = value["records"][0] + assert record["observed_scheduler_step_calls"] == [0, 0] + assert record["actual_transformer_invocations"] == [4, 4] + assert record["peak_device_memory_bytes"] == [None, None] + assert record["summaries"]["peak_device_memory_bytes"] is None + assert runtime.encode_calls == runtime.sample_calls == 3 + assert (output / record["output"]["path"]).is_file() + assert not list(tmp_path.glob(".evaluation.*.staging")) + assert json.loads(result.read_text())["schema_version"] == 1 + + +def test_main_second_prompt_failure_publishes_nothing(tmp_path, monkeypatch) -> None: + export = tmp_path / "export" + export.mkdir() + (export / "manifest.json").write_bytes(b"manifest") + prompts = tmp_path / "prompts.json" + write_canonical_json( + prompts, + { + "schema_version": 1, + "prompts": [ + {"prompt_id": "a", "prompt": "first", "seeds": [1]}, + {"prompt_id": "b", "prompt": "second", "seeds": [2]}, + ], + }, + ) + runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) + original_encode = runtime.encode_prompt + + def fail_second(prompt, max_sequence_length): + if prompt == "second": + raise RuntimeError("second prompt failed") + return original_encode(prompt, max_sequence_length) + + monkeypatch.setattr(runtime, "encode_prompt", fail_second) + monkeypatch.setattr( + "pdd.evaluate_qwen_image.load_qwen_pdd_runtime", + lambda _export, _schedule, _device: runtime, + ) + output = tmp_path / "evaluation" + monkeypatch.setattr( + sys, + "argv", + [ + "evaluate_qwen_image.py", + "--export-dir", + str(export), + "--prompts", + str(prompts), + "--schedule", + "pdd-4", + "--output-dir", + str(output), + "--result-json", + str(output / "result.json"), + "--warmup-runs", + "1", + "--measured-runs", + "1", + "--height", + "2", + "--width", + "2", + "--device", + "cpu", + ], + ) + with pytest.raises(RuntimeError, match="second prompt failed"): + main() + assert not output.exists() + assert not list(tmp_path.glob(".evaluation.*.staging")) + + +def test_legacy_inference_preserves_scope_and_adds_observed_scheduler_count( + tmp_path, monkeypatch +) -> None: + export = tmp_path / "export" + export.mkdir() + (export / "manifest.json").write_bytes(b"manifest") + runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) + monkeypatch.setattr( + "pdd.inference_qwen_image.load_qwen_pdd_runtime", + lambda _export, _schedule, _device: runtime, + ) + output = tmp_path / "image.png" + result = tmp_path / "result.json" + monkeypatch.setattr( + sys, + "argv", + [ + "inference_qwen_image.py", + "--export-dir", + str(export), + "--prompt", + "text", + "--prompt-id", + "sample", + "--schedule", + "pdd-4", + "--seed", + "7", + "--height", + "2", + "--width", + "2", + "--device", + "cpu", + "--output", + str(output), + "--result-json", + str(result), + ], + ) + inference_main() + value = load_canonical_json(result) + assert value["schema_version"] == 2 + assert value["scheduler_steps"] == 4 + assert value["observed_scheduler_step_calls"] == 0 + assert value["actual_transformer_invocations"] == 4 + assert value["latency_seconds"] > 0 + assert runtime.encode_calls == runtime.sample_calls == 1 + assert output.is_file() + + +def test_legacy_inference_restores_instrumentation_on_initial_sync_failure( + tmp_path, monkeypatch +) -> None: + export = tmp_path / "export" + export.mkdir() + (export / "manifest.json").write_bytes(b"manifest") + runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) + runtime.device = torch.device("cuda") + original = runtime.scheduler.step + monkeypatch.setattr( + "pdd.inference_qwen_image.load_qwen_pdd_runtime", + lambda _export, _schedule, _device: runtime, + ) + + def fail_sync(_device): + raise RuntimeError("sync failed") + + monkeypatch.setattr(torch.cuda, "synchronize", fail_sync) + output = tmp_path / "image.png" + result = tmp_path / "result.json" + monkeypatch.setattr( + sys, + "argv", + [ + "inference_qwen_image.py", + "--export-dir", + str(export), + "--prompt", + "text", + "--prompt-id", + "sample", + "--schedule", + "pdd-4", + "--seed", + "7", + "--height", + "2", + "--width", + "2", + "--device", + "cuda", + "--output", + str(output), + "--result-json", + str(result), + ], + ) + with pytest.raises(RuntimeError, match="sync failed"): + inference_main() + assert runtime.scheduler.step == original + assert not runtime.student._forward_pre_hooks + assert not output.exists() + assert not result.exists() From 929bee68f6d0850eabf993563fa5202d5d5a9f39 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Sat, 18 Jul 2026 03:18:50 -0700 Subject: [PATCH 31/45] Fix FSDP2 reshard after PDD validation Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/training.py | 10 +++++ .../fastgen/pdd_mr210_fsdp_distributed.py | 40 ++++++++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py index 493061707cf..51241bc9295 100644 --- a/examples/diffusers/fastgen/pdd/training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -786,6 +786,14 @@ def _raise_collective_validation_error(error: BaseException | None, *, context: raise RuntimeError(f"distributed PDD validation {context} failed; " + "; ".join(failures)) +def _reshard_fsdp2_modules(model: torch.nn.Module) -> None: + from torch.distributed.fsdp import FSDPModule + + for module in model.modules(): + if isinstance(module, FSDPModule): + module.reshard() + + def run_pdd_validation( pipeline: PDDPipeline, batches: Iterable[PreparedPDDBatch], @@ -950,6 +958,8 @@ def run_pdd_validation( finally: pipeline.student.train(student_was_training) pipeline.teacher.train(teacher_was_training) + _reshard_fsdp2_modules(pipeline.student) + _reshard_fsdp2_modules(pipeline.teacher) gathered: list[list[PDDValidationRecord]] if distributed: diff --git a/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py b/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py index e96ff52a70a..d048482108d 100644 --- a/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py +++ b/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py @@ -28,7 +28,13 @@ from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir from diffusers import QwenImageTransformer2DModel from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock -from pdd.recipe import build_pdd_setup, initialize_pdd_distributed, resolve_pdd_recipe_config +from pdd.recipe import ( + build_pdd_setup, + build_pdd_training_artifacts, + initialize_pdd_distributed, + resolve_pdd_recipe_config, +) +from pdd.training import PDDValidationAssignment, PreparedPDDBatch, run_pdd_validation from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( CheckpointImpl, CheckpointWrapper, @@ -375,6 +381,35 @@ def inner_student_hook(_module, _args, _kwargs): dist.all_gather(gathered_losses, actual_loss.detach()) for value in gathered_losses[1:]: torch.testing.assert_close(value, gathered_losses[0], rtol=0, atol=0) + + setup.optimizer.zero_grad(set_to_none=True) + training = build_pdd_training_artifacts(setup, config) + validation_batch = PreparedPDDBatch( + data, + condition, + negative_condition, + (f"post-validation-update-{rank}",), + ) + run_pdd_validation( + training.pipeline, + [validation_batch], + [ + PDDValidationAssignment(index, f"post-validation-update-{index}", 0, 2) + for index in range(dist.get_world_size()) + ], + validation_seed=11, + ) + post_validation = training.trainer.train_step( + validation_batch, + noise=noise, + n=n, + k=k, + ) + if ( + post_validation.pdd_projection_update_ratio is None + or post_validation.pdd_projection_update_ratio <= 0 + ): + raise RuntimeError("post-validation FSDP projection update was not measured") if rank == 0: print( json.dumps( @@ -382,6 +417,9 @@ def inner_student_hook(_module, _args, _kwargs): "activation_checkpointing": args.activation_checkpointing, "actual_loss": actual_loss.item(), "reference_loss": reference_loss.item(), + "post_validation_projection_update_ratio": ( + post_validation.pdd_projection_update_ratio + ), "student_block_calls": inner_student_calls, "time_0": student_times[0].item(), "world_size": dist.get_world_size(), From 2aa92f43acc30d465a5355e41343f790b42a724d Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Sat, 18 Jul 2026 08:03:01 -0700 Subject: [PATCH 32/45] Match PDD inference partitions to FastGen Signed-off-by: Meng Xin --- modelopt/torch/fastgen/config.py | 23 ++----------------- modelopt/torch/fastgen/methods/pdd.py | 18 +++------------ .../fastgen/test_pdd_inference_checkpoint.py | 5 ++-- tests/unit/torch/fastgen/test_pdd_config.py | 8 +++++-- tests/unit/torch/fastgen/test_pdd_pipeline.py | 4 ++-- 5 files changed, 16 insertions(+), 42 deletions(-) diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py index 11394eeb062..d4542e1d269 100644 --- a/modelopt/torch/fastgen/config.py +++ b/modelopt/torch/fastgen/config.py @@ -261,12 +261,12 @@ class PDDConfig(DistillationConfig): block_size_min: int = ModeloptField( default=4, title="Minimum training block alignment", - description="Alignment of sampled training start indices and inference blocks.", + description="Alignment of sampled training start indices.", ) block_size_max: int = ModeloptField( default=64, title="Maximum trained block size", - description="Largest target span and inference block supported by this training run.", + description="Largest target span sampled during training.", ) teacher_integrator: Literal["euler", "midpoint"] = ModeloptField( default="euler", @@ -330,16 +330,6 @@ def _check_pdd(self) -> PDDConfig: for index, block in enumerate(self.inference_blocks): if block <= 0: raise ValueError(f"inference_blocks[{index}] must be > 0, got {block}.") - if block % self.block_size_min != 0: - raise ValueError( - f"inference_blocks[{index}]={block} must be aligned to " - f"block_size_min={self.block_size_min}." - ) - if block > self.block_size_max: - raise ValueError( - f"inference_blocks[{index}]={block} exceeds " - f"block_size_max={self.block_size_max}." - ) if sum(self.inference_blocks) != self.grid_size: raise ValueError( f"inference_blocks must sum to grid_size={self.grid_size}, got " @@ -351,15 +341,6 @@ def _check_pdd(self) -> PDDConfig: f"{self.student_sample_steps} and {len(self.inference_blocks)}." ) - start = 0 - for index, block in enumerate(self.inference_blocks): - if start % self.block_size_min != 0 or start > self.grid_size - self.block_size_min: - raise ValueError( - f"inference block {index} starts at {start}, which is not a valid " - f"training start aligned to block_size_min={self.block_size_min}." - ) - start += block - default_sample_t_cfg = SampleTimestepConfig() if self.sample_t_cfg.model_dump() != default_sample_t_cfg.model_dump(): raise ValueError( diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index 1165fbfc22f..93b73db4528 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -936,24 +936,12 @@ def _validate_blocks(self, blocks: Sequence[int] | None) -> tuple[int, ...]: resolved = tuple(blocks) if not resolved: raise ValueError("blocks must contain at least one interval count.") - start = 0 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}.") - if block % self.config.block_size_min != 0: - raise ValueError( - f"blocks[{index}]={block} must be aligned to " - f"block_size_min={self.config.block_size_min}." - ) - if block > self.config.block_size_max: - raise ValueError( - f"blocks[{index}]={block} exceeds block_size_max={self.config.block_size_max}." - ) - if start > self.config.grid_size - self.config.block_size_min: - raise ValueError(f"block {index} starts outside the trained support at {start}.") - start += block - if start != self.config.grid_size: - raise ValueError(f"blocks must sum to grid_size={self.config.grid_size}, got {start}.") + 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() diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py index b90a257d59d..92629838596 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py @@ -205,8 +205,9 @@ def test_bounded_safe_export_round_trip_and_seeded_schedules(tmp_path, monkeypat assert source.calls == restored.calls == len(blocks) torch.testing.assert_close(_sample(restored, config, noise), actual, rtol=0, atol=0) - with pytest.raises(ValueError, match="block_size_max"): - pdd_config_from_metadata(metadata, blocks=[128], guidance_scale=4.0) + arbitrary = pdd_config_from_metadata(metadata, blocks=[1, 127], guidance_scale=4.0) + assert arbitrary.inference_blocks == [1, 127] + assert arbitrary.student_sample_steps == 2 def test_inference_config_preserves_authenticated_nondefault_grid_max_t() -> None: diff --git a/tests/unit/torch/fastgen/test_pdd_config.py b/tests/unit/torch/fastgen/test_pdd_config.py index f70bc6468f9..ffb484216a4 100644 --- a/tests/unit/torch/fastgen/test_pdd_config.py +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -122,8 +122,6 @@ def test_pdd_config_accepts_explicit_grid_max_t_upper_boundary(): ({"block_size_max": 129}, "block_size_max <= grid_size"), ({"grid_size": 130}, "must be divisible"), ({"inference_blocks": []}, "at least one block"), - ({"inference_blocks": [30, 34, 32, 32]}, "must be aligned"), - ({"inference_blocks": [128], "student_sample_steps": 1}, "exceeds block_size_max"), ({"inference_blocks": [32, 32, 32]}, "must sum to grid_size"), ( {"inference_blocks": [64, 64], "student_sample_steps": 4}, @@ -136,6 +134,12 @@ def test_pdd_config_rejects_invalid_grid_and_block_boundaries(overrides, message PDDConfig(**overrides) +def test_pdd_config_accepts_inference_partition_outside_training_block_support(): + config = PDDConfig(inference_blocks=[1, 127], student_sample_steps=2) + + assert config.inference_blocks == [1, 127] + + @pytest.mark.parametrize( "overrides", [ diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index cc0a76e41ec..a1cd94338c8 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -413,7 +413,7 @@ def test_sampled_indices_stay_on_exact_uniform_support() -> None: assert observed == set(range(n_value, min(n_value + 4, 8))) -@pytest.mark.parametrize("blocks", [None, [2, 2, 2, 2]]) +@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) @@ -511,7 +511,7 @@ def test_pipeline_rejects_invalid_shapes_dtypes_and_blocks() -> None: ) 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 ([3, 5], [2, 2], [6, 2], [], [2, 2, 2, 4]): + 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) From 30a16df4efcdf25f5a9aa6ef3f71eb8b96ffe0a9 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Sat, 18 Jul 2026 08:16:54 -0700 Subject: [PATCH 33/45] Clean up PDD example scope Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/README.md | 74 +-- .../fastgen/pdd/configs/qwen_image.yaml | 9 +- .../fastgen/pdd/evaluate_qwen_image.py | 456 ------------- .../fastgen/pdd/inference_runtime.py | 53 -- .../examples/diffusers/fastgen/test_layout.py | 1 - .../fastgen/test_pdd_evaluation_runner.py | 629 ------------------ 6 files changed, 10 insertions(+), 1212 deletions(-) delete mode 100644 examples/diffusers/fastgen/pdd/evaluate_qwen_image.py delete mode 100644 tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index d48fbc8f016..6332264c9b6 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -49,21 +49,19 @@ has additional read and SHA-256 cost and fails immediately when a hash is missin Training deterministically derives disjoint train and validation membership from metadata ordinals; it does not rewrite the cache or require separate split manifests. The default recipe uses 2,000 -validation samples, learning rate `2e-5`, per-rank batch size 4, 128 heads, start indices aligned by -4, and target spans from 1 through 64 intervals. The learning rate is the cached-Qwen project -treatment; the MR210 reference arm uses `5e-5` with 1,000 warmup steps. On 16 four-GPU nodes the -default per-rank batch gives global batch size 256; other GPU topologies must set per-rank batch to -`256 / world_size` because this recipe does not use gradient accumulation. +validation samples, constant learning rate `5e-5`, per-rank batch size 4, 128 heads, start indices +aligned by 4, and target spans from 1 through 64 intervals. Other GPU topologies can set per-rank +batch size to obtain the desired global batch size because this recipe does not use gradient +accumulation. Checkpoints include the student, optimizer, scheduler, RNG, trainer, and exact replayable sampler state needed to resume the next committed batch. FP32 master parameters and Adam state are sharded while forward/backward uses BF16 model parameters and outputs and gradient reduction remains FP32. The adapter casts packed image/text inputs to BF16 while preserving FP32 normalized time at the FSDP root and Qwen time embedder. -Start with a one-node smoke and scale only after it passes; project training runs are capped at 16 -nodes. Checkpointed training, export, and inference require the remote model ID and exact lowercase -40-character Hugging Face commit in the provided config; local model directories are limited to -low-level hermetic setup tests because the frozen teacher is rebuilt rather than checkpointed. +Checkpointed training, export, and inference require the remote model ID and exact lowercase +40-character Hugging Face commit in the provided config because the frozen teacher is rebuilt +rather than checkpointed. ## Export and inference @@ -84,61 +82,3 @@ python examples/diffusers/fastgen/pdd/inference_qwen_image.py \ --seed 42 --height 1024 --width 1024 \ --output /path/to/pdd4.png --result-json /path/to/pdd4.json ``` - -## Repeatable evaluation records - -`evaluate_qwen_image.py` loads the authenticated export once and evaluates every ordered -prompt/seed pair for one of the source-owned `pdd-2`, `pdd-4`, or `pdd-8` schedules. The prompt -file must be canonical JSON, including its trailing newline. For example: - -```json -{"prompts":[{"prompt":"a small red cube on a white table","prompt_id":"red-cube-0001","seeds":[42]}],"schema_version":1} -``` - -```bash -python examples/diffusers/fastgen/pdd/evaluate_qwen_image.py \ - --export-dir /path/to/pdd-export \ - --prompts /path/to/prompts.json --schedule pdd-4 \ - --output-dir /path/to/evaluation-pdd4 \ - --result-json /path/to/evaluation-pdd4/result.json \ - --warmup-runs 1 --measured-runs 5 \ - --height 1024 --width 1024 --max-sequence-length 512 -``` - -The runner publishes the output directory atomically. Its canonical result records exact export, -prompt, raw-noise, initial-state, and grid identities; requested logical blocks; observed scheduler -calls; actual transformer calls; synchronized end-to-end and transformer timings; throughput; and -peak CUDA allocation. Warmups execute and validate the same path but are excluded from measured -arrays. CPU memory entries are `null`. The record deliberately contains no image-quality score or -effectiveness conclusion. - -For a functional standard-teacher baseline, the public Diffusers pipeline can be pinned to the -same immutable Qwen revision and run for 50 steps: - -```python -import torch -from diffusers import QwenImagePipeline - -revision = "75e0b4be04f60ec59a75f475837eced720f823b6" -pipe = QwenImagePipeline.from_pretrained( - "Qwen/Qwen-Image", - revision=revision, - torch_dtype=torch.bfloat16, - use_safetensors=True, -).to("cuda") -generator = torch.Generator(device="cuda").manual_seed(42) -image = pipe( - prompt="a small red cube on a white table", - height=1024, - width=1024, - num_inference_steps=50, - generator=generator, -).images[0] -image.save("teacher-50.png") -``` - -That command is a functional baseline, not matched scientific evidence. A result-bearing study -must separately freeze the prompt set, seeds, controls, quality metrics, thresholds, and exact -teacher-50, undistilled Euler-2/4/8, and PDD-2/4/8 trajectory/counter schemas. Those controls and -conclusions belong to the reviewed external `scripts/pdd_qwen_effectiveness_v1` experiment package, -not this general ModelOpt example. diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 27ddf2fb2d1..bc5287b88d7 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -1,4 +1,4 @@ -# Qwen-Image PDD training recipe. Result-bearing paths/topology remain externally gated. +# Qwen-Image PDD training recipe. seed: 42 @@ -30,7 +30,7 @@ pdd: data_free: false optim: - learning_rate: 2.0e-5 + learning_rate: 5.0e-5 optimizer: _target_: torch.optim.AdamW weight_decay: 0.01 @@ -46,18 +46,15 @@ optim: lr_scheduler: lr_decay_style: constant lr_warmup_steps: 0 - min_lr: 2.0e-5 + min_lr: 5.0e-5 step_scheduler: max_steps: 10000 num_epochs: 200 log_every: 10 ckpt_every_steps: 1000 - # Matches the original Qwen PDD per-rank batch. On 16 four-GPU nodes this - # gives the project target global batch size of 256 without accumulation. local_batch_size: 4 save_checkpoint_every_epoch: false - # Freeze to 256 only after the production-topology/data gate is approved. global_batch_size: training_health: diff --git a/examples/diffusers/fastgen/pdd/evaluate_qwen_image.py b/examples/diffusers/fastgen/pdd/evaluate_qwen_image.py deleted file mode 100644 index af00cf79bb2..00000000000 --- a/examples/diffusers/fastgen/pdd/evaluate_qwen_image.py +++ /dev/null @@ -1,456 +0,0 @@ -# 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. - -"""Evaluate an authenticated Qwen-Image PDD export over a prompt manifest.""" - -from __future__ import annotations - -import argparse -import hashlib -import math -import os -import re -import shutil -import sys -import time -import uuid -from contextlib import ExitStack -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import torch - -sys.dont_write_bytecode = True - -_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 pdd.artifacts import load_canonical_json, sha256_file, write_canonical_json # noqa: E402 -from pdd.export import PDD_INFERENCE_SCHEDULES # noqa: E402 -from pdd.inference_runtime import ( # noqa: E402 - QwenPDDInferenceRuntime, - load_qwen_pdd_runtime, - save_png, -) - -_PROMPT_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*\Z") - - -@dataclass(frozen=True) -class PromptPair: - prompt_id: str - prompt: str - seed: int - - -@dataclass(frozen=True) -class Observation: - scheduler_calls: int - transformer_calls: int - transformer_seconds: float - end_to_end_seconds: float - peak_device_memory_bytes: int | None - image: Any - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--export-dir", type=Path, required=True) - parser.add_argument("--prompts", type=Path, required=True) - parser.add_argument("--schedule", choices=tuple(PDD_INFERENCE_SCHEDULES), required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--result-json", type=Path, required=True) - parser.add_argument("--warmup-runs", type=int, default=1) - parser.add_argument("--measured-runs", type=int, default=5) - 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") - return parser.parse_args() - - -def _prompt_pairs(value: Any) -> list[PromptPair]: - if not isinstance(value, dict) or set(value) != {"schema_version", "prompts"}: - raise ValueError("prompt manifest must contain exactly schema_version and prompts.") - if type(value["schema_version"]) is not int or value["schema_version"] != 1: - raise ValueError("prompt manifest schema_version must be 1.") - prompts = value["prompts"] - if not isinstance(prompts, list) or not prompts: - raise ValueError("prompt manifest prompts must be a non-empty list.") - pairs: list[PromptPair] = [] - previous_id: str | None = None - for index, item in enumerate(prompts): - if not isinstance(item, dict) or set(item) != {"prompt_id", "prompt", "seeds"}: - raise ValueError(f"prompts[{index}] has incompatible fields.") - prompt_id = item["prompt_id"] - prompt = item["prompt"] - seeds = item["seeds"] - if not isinstance(prompt_id, str) or _PROMPT_ID.fullmatch(prompt_id) is None: - raise ValueError(f"prompts[{index}].prompt_id is not a safe path component.") - if previous_id is not None and prompt_id <= previous_id: - raise ValueError("prompt IDs must be unique and lexicographically ordered.") - if not isinstance(prompt, str) or not prompt: - raise ValueError(f"prompts[{index}].prompt must be non-empty text.") - if not isinstance(seeds, list) or not seeds: - raise ValueError(f"prompts[{index}].seeds must be non-empty.") - if any(type(seed) is not int or seed < 0 or seed >= 2**63 for seed in seeds): - raise ValueError(f"prompts[{index}].seeds contains an invalid seed.") - if seeds != sorted(set(seeds)): - raise ValueError(f"prompts[{index}].seeds must be sorted and unique.") - pairs.extend(PromptPair(prompt_id, prompt, seed) for seed in seeds) - previous_id = prompt_id - return pairs - - -def _reject_symlink_components(path: Path) -> Path: - if ".." in path.parts: - raise ValueError("evaluation paths cannot contain parent traversal components.") - absolute = path.absolute() - current = Path(absolute.anchor) - for part in absolute.parts[1:]: - current /= part - if current.is_symlink(): - raise ValueError(f"evaluation path traverses a symlink: {current}.") - return absolute.resolve(strict=False) - - -def _resolve_output_paths(output_value: Path, result_value: Path) -> tuple[Path, Path, Path]: - output = _reject_symlink_components(output_value) - result = _reject_symlink_components(result_value) - if output.exists() or output.is_symlink(): - raise FileExistsError(f"evaluation output already exists: {output}.") - if output.name in {"", ".", ".."}: - raise ValueError("evaluation output directory name is invalid.") - if not output.parent.is_dir() or output.parent.is_symlink(): - raise ValueError("evaluation output parent must be an existing regular directory.") - try: - relative_result = result.relative_to(output) - except ValueError as error: - raise ValueError("result JSON must be strictly beneath output_dir.") from error - if not relative_result.parts or relative_result == Path("."): - raise ValueError("result JSON must be strictly beneath output_dir.") - if result.suffix.lower() != ".json": - raise ValueError("result JSON must use a .json suffix.") - staging = output.with_name(f".{output.name}.{uuid.uuid4().hex}.staging") - os.mkdir(staging, mode=0o770) - os.chmod(staging, 0o770, follow_symlinks=False) - return output, relative_result, staging - - -def _restore_instance_method( - owner: Any, name: str, *, had_instance_value: bool, instance_value: Any -) -> None: - if had_instance_value: - setattr(owner, name, instance_value) - elif name in vars(owner): - delattr(owner, name) - - -def _run_repetition( - runtime: QwenPDDInferenceRuntime, - prompt: str, - raw_noise: torch.Tensor, - max_sequence_length: int, -) -> Observation: - scheduler_calls = 0 - transformer_calls = 0 - cpu_forward_started: list[float] = [] - cpu_forward_seconds = 0.0 - cuda_event_pairs: list[tuple[Any, Any]] = [] - scheduler = runtime.scheduler - scheduler_values = vars(scheduler) - scheduler_had_step = "step" in scheduler_values - scheduler_instance_step = scheduler_values.get("step") - original_step = scheduler.step - - def counted_step(*args: Any, **kwargs: Any) -> Any: - nonlocal scheduler_calls - scheduler_calls += 1 - return original_step(*args, **kwargs) - - def before_forward(_module: Any, _args: Any, _kwargs: Any) -> None: - nonlocal transformer_calls - transformer_calls += 1 - if runtime.device.type == "cuda": - started = torch.cuda.Event(enable_timing=True) - ended = torch.cuda.Event(enable_timing=True) - started.record() - cuda_event_pairs.append((started, ended)) - else: - cpu_forward_started.append(time.perf_counter()) - - def after_forward(_module: Any, _args: Any, _kwargs: Any, _output: Any) -> None: - nonlocal cpu_forward_seconds - if runtime.device.type == "cuda": - cuda_event_pairs[-1][1].record() - else: - if not cpu_forward_started: - raise RuntimeError("transformer post-hook ran without its pre-hook.") - cpu_forward_seconds += time.perf_counter() - cpu_forward_started.pop() - - with ExitStack() as cleanup: - cleanup.callback( - _restore_instance_method, - scheduler, - "step", - had_instance_value=scheduler_had_step, - instance_value=scheduler_instance_step, - ) - setattr(scheduler, "step", counted_step) - pre_hook = runtime.student.register_forward_pre_hook(before_forward, with_kwargs=True) - cleanup.callback(pre_hook.remove) - post_hook = runtime.student.register_forward_hook(after_forward, with_kwargs=True) - cleanup.callback(post_hook.remove) - if runtime.device.type == "cuda": - torch.cuda.synchronize(runtime.device) - torch.cuda.reset_peak_memory_stats(runtime.device) - started = time.perf_counter() - condition = runtime.encode_prompt(prompt, max_sequence_length) - images = runtime.sample_decode(condition, raw_noise) - if runtime.device.type == "cuda": - torch.cuda.synchronize(runtime.device) - end_to_end_seconds = time.perf_counter() - started - if runtime.device.type == "cuda": - transformer_seconds = sum( - start.elapsed_time(end) / 1000.0 for start, end in cuda_event_pairs - ) - peak_memory = torch.cuda.max_memory_allocated(runtime.device) - else: - transformer_seconds = cpu_forward_seconds - peak_memory = None - - expected = len(runtime.config.inference_blocks) - if not isinstance(images, list) or len(images) != 1: - raise RuntimeError("one evaluation repetition must return exactly one image.") - if scheduler_calls != 0: - raise RuntimeError(f"PDD evaluation observed {scheduler_calls} scheduler.step calls.") - if transformer_calls != expected: - raise RuntimeError( - f"PDD evaluation observed {transformer_calls} transformer calls; expected {expected}." - ) - if cpu_forward_started: - raise RuntimeError("transformer forward hooks are unbalanced.") - if ( - not math.isfinite(end_to_end_seconds) - or not math.isfinite(transformer_seconds) - or end_to_end_seconds <= 0 - or transformer_seconds <= 0 - or transformer_seconds > end_to_end_seconds + 1e-6 - ): - raise RuntimeError("evaluation timing invariants failed.") - if peak_memory is not None and (type(peak_memory) is not int or peak_memory < 0): - raise RuntimeError("evaluation peak device memory is invalid.") - return Observation( - scheduler_calls=scheduler_calls, - transformer_calls=transformer_calls, - transformer_seconds=transformer_seconds, - end_to_end_seconds=end_to_end_seconds, - peak_device_memory_bytes=peak_memory, - image=images[0], - ) - - -def _summary(values: list[float | int]) -> dict[str, float | int]: - if not values or any(type(value) not in {int, float} for value in values): - raise ValueError("summary values must be a non-empty numeric list.") - ordered = sorted(values) - middle = len(ordered) // 2 - median = ordered[middle] if len(ordered) % 2 else (ordered[middle - 1] + ordered[middle]) / 2 - p95 = ordered[math.ceil(0.95 * len(ordered)) - 1] - return {"median": median, "p95": p95} - - -def _fsync_directory(path: Path) -> None: - descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - - -def _publish_staging(staging: Path, output: Path) -> None: - """Rename a complete staging tree and roll back a failed parent fsync.""" - staging.rename(output) - try: - _fsync_directory(output.parent) - except BaseException: - output.rename(staging) - _fsync_directory(output.parent) - raise - - -def _record_for_pair( - runtime: QwenPDDInferenceRuntime, - pair: PromptPair, - *, - trajectory: dict[str, Any], - observations: list[Observation], - image_reference: str, - image_sha256: str, -) -> dict[str, Any]: - scheduler = [item.scheduler_calls for item in observations] - actual = [item.transformer_calls for item in observations] - transformer = [item.transformer_seconds for item in observations] - end_to_end = [item.end_to_end_seconds for item in observations] - throughput = [1.0 / value for value in end_to_end] - memory = [item.peak_device_memory_bytes for item in observations] - expected = len(runtime.config.inference_blocks) - if scheduler != [0] * len(observations) or actual != [expected] * len(observations): - raise RuntimeError("measured repetition counters are inconsistent.") - numeric_memory = [value for value in memory if value is not None] - if numeric_memory and len(numeric_memory) != len(memory): - raise RuntimeError("peak-memory observations mix CPU and CUDA domains.") - return { - "prompt_id": pair.prompt_id, - "prompt_sha256": hashlib.sha256(pair.prompt.encode("utf-8")).hexdigest(), - "seed": pair.seed, - "raw_noise_sha256": trajectory["raw_noise_sha256"], - "initial_state_sha256": trajectory["initial_state_sha256"], - "requested_scheduler_steps": expected, - "logical_pdd_blocks": list(runtime.config.inference_blocks), - "logical_pdd_block_count": expected, - "observed_scheduler_step_calls": scheduler, - "actual_transformer_invocations": actual, - "batch_normalized_transformer_evaluations": list(actual), - "transformer_latency_seconds": transformer, - "end_to_end_latency_seconds": end_to_end, - "throughput_images_per_second": throughput, - "peak_device_memory_bytes": memory, - "summaries": { - "transformer_latency_seconds": _summary(transformer), - "end_to_end_latency_seconds": _summary(end_to_end), - "throughput_images_per_second": _summary(throughput), - "peak_device_memory_bytes": _summary(numeric_memory) if numeric_memory else None, - }, - "output": {"path": image_reference, "sha256": image_sha256}, - } - - -@torch.no_grad() -def main() -> None: - args = _parse_args() - if type(args.warmup_runs) is not int or args.warmup_runs < 1: - raise ValueError("warmup_runs must be a positive integer.") - if type(args.measured_runs) is not int or args.measured_runs < 1: - raise ValueError("measured_runs must be a positive integer.") - if args.height < 1 or args.width < 1 or args.max_sequence_length < 1: - raise ValueError("height, width, and max_sequence_length must be positive.") - prompt_manifest = load_canonical_json(args.prompts) - pairs = _prompt_pairs(prompt_manifest) - prompt_manifest_sha256 = sha256_file(args.prompts) - output, result_reference, staging = _resolve_output_paths(args.output_dir, args.result_json) - try: - runtime = load_qwen_pdd_runtime(args.export_dir, args.schedule, args.device) - records: list[dict[str, Any]] = [] - grid_identity: dict[str, Any] | None = None - for pair in pairs: - raw_noise = runtime.make_raw_noise(seed=pair.seed, height=args.height, width=args.width) - trajectory = runtime.trajectory_identity(raw_noise) - if grid_identity is None: - grid_identity = trajectory - elif any( - trajectory[key] != grid_identity[key] - for key in ( - "full_time_nodes", - "full_time_nodes_sha256", - "boundary_indices", - "boundary_time_nodes", - "boundary_time_nodes_sha256", - "first_sigma", - ) - ): - raise RuntimeError("PDD trajectory grid changed between prompt/seed pairs.") - for _ in range(args.warmup_runs): - _run_repetition(runtime, pair.prompt, raw_noise, args.max_sequence_length) - observations = [ - _run_repetition(runtime, pair.prompt, raw_noise, args.max_sequence_length) - for _ in range(args.measured_runs) - ] - image_reference = f"images/{pair.prompt_id}/{pair.seed}.png" - image_path = staging / image_reference - save_png(image_path, observations[0].image) - records.append( - _record_for_pair( - runtime, - pair, - trajectory=trajectory, - observations=observations, - image_reference=image_reference, - image_sha256=sha256_file(image_path), - ) - ) - if grid_identity is None: - raise RuntimeError("evaluation produced no trajectory identity.") - dtype_name = str(runtime.dtype).removeprefix("torch.") - result = { - "schema_version": 1, - "record_type": "pdd_qwen_evaluation", - "identity": { - "export_manifest_sha256": sha256_file(runtime.descriptor.root / "manifest.json"), - "prompt_manifest_sha256": prompt_manifest_sha256, - "model": dict(runtime.model_identity), - "schedule": args.schedule, - "grid": { - "grid_size": runtime.config.grid_size, - "grid_max_t": runtime.config.grid_max_t, - "flow_shift": runtime.config.flow_shift, - "full_time_nodes": grid_identity["full_time_nodes"], - "full_time_nodes_sha256": grid_identity["full_time_nodes_sha256"], - "boundary_indices": grid_identity["boundary_indices"], - "boundary_time_nodes": grid_identity["boundary_time_nodes"], - "boundary_time_nodes_sha256": grid_identity["boundary_time_nodes_sha256"], - "first_sigma": grid_identity["first_sigma"], - }, - }, - "protocol": { - "height": args.height, - "width": args.width, - "max_sequence_length": args.max_sequence_length, - "batch_size": 1, - "warmup_runs": args.warmup_runs, - "measured_runs": args.measured_runs, - "device": str(runtime.device), - "dtype": dtype_name, - "end_to_end_scope": "prompt_encode_through_vae_postprocess", - "transformer_scope": "sum_of_root_student_forward_calls", - "cuda_synchronize": runtime.device.type == "cuda", - "tensor_hash_schema": "pdd_tensor_sha256_v1", - }, - "records": records, - } - staged_result = staging / result_reference - staged_result.parent.mkdir(parents=True, exist_ok=True) - write_canonical_json(staged_result, result) - for directory in sorted( - (path for path in staging.rglob("*") if path.is_dir()), - key=lambda path: len(path.parts), - reverse=True, - ): - _fsync_directory(directory) - _fsync_directory(staging) - _publish_staging(staging, output) - except BaseException: - shutil.rmtree(staging, ignore_errors=True) - raise - print(output / result_reference) - - -if __name__ == "__main__": - main() diff --git a/examples/diffusers/fastgen/pdd/inference_runtime.py b/examples/diffusers/fastgen/pdd/inference_runtime.py index da578e5b547..3398547eaa1 100644 --- a/examples/diffusers/fastgen/pdd/inference_runtime.py +++ b/examples/diffusers/fastgen/pdd/inference_runtime.py @@ -17,30 +17,20 @@ from __future__ import annotations -import hashlib import os import uuid from collections.abc import Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any -import numpy as np import torch from torch import nn if TYPE_CHECKING: from pathlib import Path -from .artifacts import canonical_json_bytes from .export import PDD_INFERENCE_SCHEDULES, pdd_config_from_metadata -_TENSOR_HASH_DOMAINS = { - "raw_noise", - "initial_state", - "full_time_nodes", - "boundary_time_nodes", -} - def _dtype_from_name(name: Any) -> torch.dtype: if not isinstance(name, str): @@ -184,30 +174,6 @@ def _decode_qwen_latents(pipe: Any, latents: torch.Tensor) -> list[Any]: return pipe.image_processor.postprocess(decoded[:, :, 0], output_type="pil") -def pdd_tensor_sha256(tensor: torch.Tensor, domain: str) -> str: - """Hash one exact FP32 tensor using the evaluation protocol.""" - if domain not in _TENSOR_HASH_DOMAINS: - raise ValueError(f"unknown PDD tensor hash domain {domain!r}.") - if not isinstance(tensor, torch.Tensor) or tensor.dtype != torch.float32: - raise TypeError("PDD tensor hashing requires a float32 tensor.") - if not torch.isfinite(tensor).all().item(): - raise FloatingPointError("PDD tensor hashing rejects non-finite values.") - array = np.ascontiguousarray(tensor.detach().cpu().numpy(), dtype=" None: """Publish one PNG exclusively and durably.""" if path.is_symlink(): @@ -281,25 +247,6 @@ def sample_decode(self, condition: Any, raw_noise: torch.Tensor) -> list[Any]: sampled = self.sampler.sample(raw_noise, condition=condition) return _decode_qwen_latents(self.pipe, sampled.to(self.dtype)) - def trajectory_identity(self, raw_noise: torch.Tensor) -> dict[str, Any]: - full = self.sampler.time_grid(raw_noise.device).to(device="cpu", dtype=torch.float32) - boundaries = [0] - for block in self.config.inference_blocks: - boundaries.append(boundaries[-1] + block) - boundary = full[boundaries] - initial = (raw_noise.to(torch.float64) * self.config.grid_max_t).to(torch.float32) - return { - "raw_noise_sha256": pdd_tensor_sha256(raw_noise, "raw_noise"), - "initial_state_sha256": pdd_tensor_sha256(initial, "initial_state"), - "full_time_nodes": full.tolist(), - "full_time_nodes_sha256": pdd_tensor_sha256(full, "full_time_nodes"), - "boundary_indices": boundaries, - "boundary_time_nodes": boundary.tolist(), - "boundary_time_nodes_sha256": pdd_tensor_sha256(boundary, "boundary_time_nodes"), - "first_sigma": float(full[0].item()), - } - - def load_qwen_pdd_runtime( export_dir: str | Path, schedule: str, device: str | torch.device ) -> QwenPDDInferenceRuntime: diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index 4724604c5d9..5f8c1215c5e 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -57,7 +57,6 @@ "data.py", "export.py", "export_qwen_image.py", - "evaluate_qwen_image.py", "finetune.py", "inference_runtime.py", "inference_qwen_image.py", diff --git a/tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py b/tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py deleted file mode 100644 index a9fa4660847..00000000000 --- a/tests/examples/diffusers/fastgen/test_pdd_evaluation_runner.py +++ /dev/null @@ -1,629 +0,0 @@ -# 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. - -"""Tests for the Qwen-Image PDD evaluation runner.""" - -from __future__ import annotations - -import hashlib -import json -import pathlib -import sys -import time -from types import SimpleNamespace - -import numpy as np -import pytest -import torch -from PIL import Image -from torch import nn - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -for path in (_REPO_ROOT, _FASTGEN_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - -from pdd.artifacts import canonical_json_bytes, load_canonical_json, write_canonical_json -from pdd.evaluate_qwen_image import ( - _prompt_pairs, - _publish_staging, - _resolve_output_paths, - _run_repetition, - _summary, - main, -) -from pdd.export import PDD_INFERENCE_SCHEDULES -from pdd.inference_qwen_image import main as inference_main -from pdd.inference_runtime import QwenPDDInferenceRuntime, pdd_tensor_sha256 - - -class _Scheduler: - def step(self, value=0): - return value - - -class _Student(nn.Module): - def forward(self, value): - time.sleep(0.0001) - return value + 1 - - -class _FakeRuntime: - def __init__(self, export_root: pathlib.Path, blocks=(1, 1)) -> None: - self.student = _Student() - self.scheduler = _Scheduler() - self.device = torch.device("cpu") - self.dtype = torch.bfloat16 - self.config = SimpleNamespace( - inference_blocks=list(blocks), - grid_size=sum(blocks), - grid_max_t=0.999, - flow_shift=5.0, - ) - self.descriptor = SimpleNamespace(root=export_root) - self.model_identity = { - "id": "Qwen/Qwen-Image", - "revision": "f" * 40, - "dtype": "bfloat16", - } - self.encode_calls = 0 - self.sample_calls = 0 - - def encode_prompt(self, prompt, max_sequence_length): - assert prompt and max_sequence_length > 0 - self.encode_calls += 1 - return torch.tensor(0) - - def make_raw_noise(self, *, seed, height, width): - return torch.randn((1, 1, height, width), generator=torch.Generator().manual_seed(seed)) - - def sample_decode(self, condition, raw_noise): - del condition, raw_noise - self.sample_calls += 1 - value = torch.tensor(0) - for _ in self.config.inference_blocks: - value = self.student(value) - return [Image.new("RGB", (2, 2), color=(int(value), 0, 0))] - - def trajectory_identity(self, raw_noise): - full = torch.linspace(0.999, 0.0, self.config.grid_size + 1, dtype=torch.float32) - boundaries = [0] - for block in self.config.inference_blocks: - boundaries.append(boundaries[-1] + block) - boundary = full[boundaries] - initial = (raw_noise.to(torch.float64) * self.config.grid_max_t).to(torch.float32) - return { - "raw_noise_sha256": pdd_tensor_sha256(raw_noise, "raw_noise"), - "initial_state_sha256": pdd_tensor_sha256(initial, "initial_state"), - "full_time_nodes": full.tolist(), - "full_time_nodes_sha256": pdd_tensor_sha256(full, "full_time_nodes"), - "boundary_indices": boundaries, - "boundary_time_nodes": boundary.tolist(), - "boundary_time_nodes_sha256": pdd_tensor_sha256(boundary, "boundary_time_nodes"), - "first_sigma": float(full[0]), - } - - -def test_prompt_manifest_expands_exact_order_and_rejects_unsafe_data() -> None: - value = { - "schema_version": 1, - "prompts": [ - {"prompt_id": "a", "prompt": "alpha", "seeds": [1, 2]}, - {"prompt_id": "b-2", "prompt": "beta", "seeds": [3]}, - ], - } - assert [(pair.prompt_id, pair.seed) for pair in _prompt_pairs(value)] == [ - ("a", 1), - ("a", 2), - ("b-2", 3), - ] - value["prompts"][1]["prompt_id"] = "../escape" - with pytest.raises(ValueError, match="safe path component"): - _prompt_pairs(value) - - -@pytest.mark.parametrize( - ("value", "message"), - [ - ({"schema_version": 1, "prompts": [], "extra": 1}, "exactly"), - ({"schema_version": 2, "prompts": []}, "schema_version"), - ({"schema_version": True, "prompts": []}, "schema_version"), - ( - { - "schema_version": 1, - "prompts": [ - {"prompt_id": "b", "prompt": "one", "seeds": [1]}, - {"prompt_id": "a", "prompt": "two", "seeds": [2]}, - ], - }, - "lexicographically", - ), - ( - { - "schema_version": 1, - "prompts": [{"prompt_id": "a", "prompt": "one", "seeds": [2, 1]}], - }, - "sorted and unique", - ), - ( - { - "schema_version": 1, - "prompts": [{"prompt_id": "a", "prompt": "one", "seeds": [True]}], - }, - "invalid seed", - ), - ], -) -def test_prompt_manifest_expected_red_matrix(value, message) -> None: - with pytest.raises(ValueError, match=message): - _prompt_pairs(value) - - -def test_tensor_hash_uses_exact_header_little_endian_payload_and_domain() -> None: - tensor = torch.tensor([[1.0, -2.5]], dtype=torch.float32) - header = { - "schema_version": 1, - "domain": "raw_noise", - "dtype": "float32", - "shape": [1, 2], - "byte_order": "little", - "order": "C", - } - payload = np.ascontiguousarray(tensor.numpy(), dtype=" None: - requested_devices = [] - - class Sampler: - def time_grid(self, device): - requested_devices.append(device) - return torch.tensor([0.999, 0.5, 0.0], dtype=torch.float32) - - class RawNoise: - device = torch.device("cuda:7") - - def to(self, dtype): - return torch.ones((1, 1), dtype=dtype) - - runtime = QwenPDDInferenceRuntime( - student=None, - scheduler=None, - descriptor=None, - model_identity={}, - dtype=torch.bfloat16, - device=torch.device("cuda:7"), - config=SimpleNamespace(inference_blocks=(1, 1), grid_max_t=0.999), - pipe=None, - sampler=Sampler(), - ) - monkeypatch.setattr("pdd.inference_runtime.pdd_tensor_sha256", lambda _tensor, domain: domain) - identity = runtime.trajectory_identity(RawNoise()) - assert requested_devices == [torch.device("cuda:7")] - assert identity["full_time_nodes"] == pytest.approx([0.999, 0.5, 0.0]) - - -def test_source_owned_schedules_and_summary_contract() -> None: - assert PDD_INFERENCE_SCHEDULES == { - "pdd-2": (64, 64), - "pdd-4": (32, 32, 32, 32), - "pdd-8": (16, 16, 16, 16, 16, 16, 16, 16), - } - assert _summary([4.0, 1.0, 3.0, 2.0]) == {"median": 2.5, "p95": 4.0} - - -def test_repetition_counts_calls_times_cpu_and_restores_scheduler(tmp_path) -> None: - runtime = _FakeRuntime(tmp_path) - original = runtime.scheduler.step - observation = _run_repetition(runtime, "prompt", torch.zeros(1), 8) - assert observation.scheduler_calls == 0 - assert observation.transformer_calls == 2 - assert observation.peak_device_memory_bytes is None - assert observation.transformer_seconds > 0 - assert observation.end_to_end_seconds >= observation.transformer_seconds - assert runtime.scheduler.step == original - assert not runtime.student._forward_pre_hooks - assert not runtime.student._forward_hooks - - -@pytest.mark.parametrize("blocks", [(64, 64), (32, 32, 32, 32), (16,) * 8]) -def test_repetition_counts_each_supported_schedule(tmp_path, blocks) -> None: - runtime = _FakeRuntime(tmp_path, blocks=blocks) - observation = _run_repetition(runtime, "prompt", torch.zeros(1), 8) - assert observation.transformer_calls == len(blocks) - - -def test_repetition_rejects_scheduler_and_transformer_count_collapse(tmp_path, monkeypatch) -> None: - runtime = _FakeRuntime(tmp_path) - - def scheduler_call(_condition, _noise): - runtime.scheduler.step() - value = runtime.student(torch.tensor(0)) - value = runtime.student(value) - return [Image.new("RGB", (2, 2))] - - monkeypatch.setattr(runtime, "sample_decode", scheduler_call) - with pytest.raises(RuntimeError, match=r"scheduler\.step"): - _run_repetition(runtime, "prompt", torch.zeros(1), 8) - assert "step" not in vars(runtime.scheduler) - - def missing_transformer(_condition, _noise): - return [Image.new("RGB", (2, 2))] - - monkeypatch.setattr(runtime, "sample_decode", missing_transformer) - with pytest.raises(RuntimeError, match="transformer calls"): - _run_repetition(runtime, "prompt", torch.zeros(1), 8) - - -def test_mocked_cuda_instrumentation_orders_sync_reset_events_and_memory( - tmp_path, monkeypatch -) -> None: - runtime = _FakeRuntime(tmp_path) - runtime.device = torch.device("cuda") - actions = [] - - class Event: - def __init__(self, *, enable_timing): - assert enable_timing is True - - def record(self): - actions.append("event") - - def elapsed_time(self, _other): - actions.append("elapsed") - return 0.01 - - monkeypatch.setattr(torch.cuda, "Event", Event) - monkeypatch.setattr(torch.cuda, "synchronize", lambda _device: actions.append("sync")) - monkeypatch.setattr( - torch.cuda, "reset_peak_memory_stats", lambda _device: actions.append("reset") - ) - monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _device: 123) - observation = _run_repetition(runtime, "prompt", torch.zeros(1), 8) - assert actions[:2] == ["sync", "reset"] - assert actions[-3:] == ["sync", "elapsed", "elapsed"] - assert observation.peak_device_memory_bytes == 123 - - -def test_repetition_restores_instrumentation_on_failure(tmp_path, monkeypatch) -> None: - runtime = _FakeRuntime(tmp_path) - original = runtime.scheduler.step - - def fail(_condition, _noise): - runtime.scheduler.step() - raise RuntimeError("boom") - - monkeypatch.setattr(runtime, "sample_decode", fail) - with pytest.raises(RuntimeError, match="boom"): - _run_repetition(runtime, "prompt", torch.zeros(1), 8) - assert runtime.scheduler.step == original - assert not runtime.student._forward_pre_hooks - assert not runtime.student._forward_hooks - - -def test_repetition_restores_after_initial_sync_failure(tmp_path, monkeypatch) -> None: - runtime = _FakeRuntime(tmp_path) - runtime.device = torch.device("cuda") - original = runtime.scheduler.step - - def fail_sync(_device): - raise RuntimeError("sync failed") - - monkeypatch.setattr(torch.cuda, "synchronize", fail_sync) - with pytest.raises(RuntimeError, match="sync failed"): - _run_repetition(runtime, "prompt", torch.zeros(1), 8) - assert runtime.scheduler.step == original - assert not runtime.student._forward_pre_hooks - assert not runtime.student._forward_hooks - - -def test_repetition_restores_after_partial_hook_install_failure(tmp_path, monkeypatch) -> None: - runtime = _FakeRuntime(tmp_path) - original = runtime.scheduler.step - - def fail_post_hook(*_args, **_kwargs): - raise RuntimeError("post-hook failed") - - monkeypatch.setattr(runtime.student, "register_forward_hook", fail_post_hook) - with pytest.raises(RuntimeError, match="post-hook failed"): - _run_repetition(runtime, "prompt", torch.zeros(1), 8) - assert runtime.scheduler.step == original - assert not runtime.student._forward_pre_hooks - assert not runtime.student._forward_hooks - - -def test_atomic_publish_rolls_back_failed_parent_fsync(tmp_path, monkeypatch) -> None: - staging = tmp_path / ".result.staging" - staging.mkdir() - (staging / "complete").write_text("complete") - output = tmp_path / "result" - module = sys.modules["pdd.evaluate_qwen_image"] - original = module._fsync_directory - failed = False - - def fail_once(path): - nonlocal failed - if path == tmp_path and output.exists() and not failed: - failed = True - raise OSError("fsync failed") - original(path) - - monkeypatch.setattr(module, "_fsync_directory", fail_once) - with pytest.raises(OSError, match="fsync failed"): - _publish_staging(staging, output) - assert failed - assert not output.exists() - assert (staging / "complete").is_file() - - -def test_output_transaction_rejects_escape_collision_and_symlink(tmp_path) -> None: - with pytest.raises(ValueError, match="strictly beneath"): - _resolve_output_paths(tmp_path / "output", tmp_path / "outside.json") - existing = tmp_path / "existing" - existing.mkdir() - with pytest.raises(FileExistsError, match="already exists"): - _resolve_output_paths(existing, existing / "result.json") - real = tmp_path / "real" - real.mkdir() - link = tmp_path / "link" - link.symlink_to(real, target_is_directory=True) - with pytest.raises(ValueError, match="symlink"): - _resolve_output_paths(link / "output", link / "output" / "result.json") - output = real / "output" - unresolved_then_link = tmp_path / "missing" / ".." / "link" / "output" - with pytest.raises(ValueError, match="parent traversal"): - _resolve_output_paths(unresolved_then_link, output / "result.json") - unresolved_result = unresolved_then_link / "result.json" - with pytest.raises(ValueError, match="parent traversal"): - _resolve_output_paths(output, unresolved_result) - - -def test_main_publishes_complete_atomic_cpu_result(tmp_path, monkeypatch) -> None: - export = tmp_path / "export" - export.mkdir() - (export / "manifest.json").write_bytes(b"manifest") - prompts = tmp_path / "prompts.json" - write_canonical_json( - prompts, - { - "schema_version": 1, - "prompts": [{"prompt_id": "sample", "prompt": "text", "seeds": [7]}], - }, - ) - runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) - monkeypatch.setattr( - "pdd.evaluate_qwen_image.load_qwen_pdd_runtime", - lambda _export, _schedule, _device: runtime, - ) - output = tmp_path / "evaluation" - result = output / "result.json" - monkeypatch.setattr( - sys, - "argv", - [ - "evaluate_qwen_image.py", - "--export-dir", - str(export), - "--prompts", - str(prompts), - "--schedule", - "pdd-4", - "--output-dir", - str(output), - "--result-json", - str(result), - "--warmup-runs", - "1", - "--measured-runs", - "2", - "--height", - "2", - "--width", - "2", - "--device", - "cpu", - ], - ) - main() - value = load_canonical_json(result) - assert value["record_type"] == "pdd_qwen_evaluation" - assert value["identity"]["schedule"] == "pdd-4" - record = value["records"][0] - assert record["observed_scheduler_step_calls"] == [0, 0] - assert record["actual_transformer_invocations"] == [4, 4] - assert record["peak_device_memory_bytes"] == [None, None] - assert record["summaries"]["peak_device_memory_bytes"] is None - assert runtime.encode_calls == runtime.sample_calls == 3 - assert (output / record["output"]["path"]).is_file() - assert not list(tmp_path.glob(".evaluation.*.staging")) - assert json.loads(result.read_text())["schema_version"] == 1 - - -def test_main_second_prompt_failure_publishes_nothing(tmp_path, monkeypatch) -> None: - export = tmp_path / "export" - export.mkdir() - (export / "manifest.json").write_bytes(b"manifest") - prompts = tmp_path / "prompts.json" - write_canonical_json( - prompts, - { - "schema_version": 1, - "prompts": [ - {"prompt_id": "a", "prompt": "first", "seeds": [1]}, - {"prompt_id": "b", "prompt": "second", "seeds": [2]}, - ], - }, - ) - runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) - original_encode = runtime.encode_prompt - - def fail_second(prompt, max_sequence_length): - if prompt == "second": - raise RuntimeError("second prompt failed") - return original_encode(prompt, max_sequence_length) - - monkeypatch.setattr(runtime, "encode_prompt", fail_second) - monkeypatch.setattr( - "pdd.evaluate_qwen_image.load_qwen_pdd_runtime", - lambda _export, _schedule, _device: runtime, - ) - output = tmp_path / "evaluation" - monkeypatch.setattr( - sys, - "argv", - [ - "evaluate_qwen_image.py", - "--export-dir", - str(export), - "--prompts", - str(prompts), - "--schedule", - "pdd-4", - "--output-dir", - str(output), - "--result-json", - str(output / "result.json"), - "--warmup-runs", - "1", - "--measured-runs", - "1", - "--height", - "2", - "--width", - "2", - "--device", - "cpu", - ], - ) - with pytest.raises(RuntimeError, match="second prompt failed"): - main() - assert not output.exists() - assert not list(tmp_path.glob(".evaluation.*.staging")) - - -def test_legacy_inference_preserves_scope_and_adds_observed_scheduler_count( - tmp_path, monkeypatch -) -> None: - export = tmp_path / "export" - export.mkdir() - (export / "manifest.json").write_bytes(b"manifest") - runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) - monkeypatch.setattr( - "pdd.inference_qwen_image.load_qwen_pdd_runtime", - lambda _export, _schedule, _device: runtime, - ) - output = tmp_path / "image.png" - result = tmp_path / "result.json" - monkeypatch.setattr( - sys, - "argv", - [ - "inference_qwen_image.py", - "--export-dir", - str(export), - "--prompt", - "text", - "--prompt-id", - "sample", - "--schedule", - "pdd-4", - "--seed", - "7", - "--height", - "2", - "--width", - "2", - "--device", - "cpu", - "--output", - str(output), - "--result-json", - str(result), - ], - ) - inference_main() - value = load_canonical_json(result) - assert value["schema_version"] == 2 - assert value["scheduler_steps"] == 4 - assert value["observed_scheduler_step_calls"] == 0 - assert value["actual_transformer_invocations"] == 4 - assert value["latency_seconds"] > 0 - assert runtime.encode_calls == runtime.sample_calls == 1 - assert output.is_file() - - -def test_legacy_inference_restores_instrumentation_on_initial_sync_failure( - tmp_path, monkeypatch -) -> None: - export = tmp_path / "export" - export.mkdir() - (export / "manifest.json").write_bytes(b"manifest") - runtime = _FakeRuntime(export, blocks=(1, 1, 1, 1)) - runtime.device = torch.device("cuda") - original = runtime.scheduler.step - monkeypatch.setattr( - "pdd.inference_qwen_image.load_qwen_pdd_runtime", - lambda _export, _schedule, _device: runtime, - ) - - def fail_sync(_device): - raise RuntimeError("sync failed") - - monkeypatch.setattr(torch.cuda, "synchronize", fail_sync) - output = tmp_path / "image.png" - result = tmp_path / "result.json" - monkeypatch.setattr( - sys, - "argv", - [ - "inference_qwen_image.py", - "--export-dir", - str(export), - "--prompt", - "text", - "--prompt-id", - "sample", - "--schedule", - "pdd-4", - "--seed", - "7", - "--height", - "2", - "--width", - "2", - "--device", - "cuda", - "--output", - str(output), - "--result-json", - str(result), - ], - ) - with pytest.raises(RuntimeError, match="sync failed"): - inference_main() - assert runtime.scheduler.step == original - assert not runtime.student._forward_pre_hooks - assert not output.exists() - assert not result.exists() From 2fe280630185e4972cb9cb2efdf544181c9ec75d Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Mon, 20 Jul 2026 17:33:27 -0700 Subject: [PATCH 34/45] Avoid revalidating PDD parent during saves Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/checkpoint.py | 5 ++- .../fastgen/test_pdd_training_lifecycle.py | 44 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/examples/diffusers/fastgen/pdd/checkpoint.py b/examples/diffusers/fastgen/pdd/checkpoint.py index 4fbf151d693..b8841ed8bf3 100644 --- a/examples/diffusers/fastgen/pdd/checkpoint.py +++ b/examples/diffusers/fastgen/pdd/checkpoint.py @@ -552,6 +552,7 @@ def __init__( self.trainer = trainer self.sampler = sampler self.rng = rng + self._last_checkpoint: Path | None = None self.identity = json.loads(json.dumps(identity, sort_keys=True)) if self.identity.get("schema_version") != _CHECKPOINT_SCHEMA_VERSION: raise ValueError("PDD checkpoint identity has an unsupported schema version.") @@ -778,7 +779,7 @@ def save(self) -> Path: if any(state != step_scheduler_state for state in rank_scheduler_states): raise RuntimeError("PDD ranks disagree on StepScheduler checkpoint state.") final = self.root / f"step_{completed_steps:08d}" - parent = self._collective_resolve("LATEST") + parent = self._last_checkpoint prepare_status = None if _rank() == 0: @@ -852,6 +853,7 @@ def save(self) -> Path: raise RuntimeError( f"rank-0 checkpoint publication failed: {publish_status.get('error')}." ) + self._last_checkpoint = final return final def load(self, restore_from: str | Path | None) -> PDDResumeState | None: @@ -904,6 +906,7 @@ def load(self, restore_from: str | Path | None) -> PDDResumeState | None: if current_lrs != manifest["learning_rates"]: raise RuntimeError("PDD restored learning rate does not match the manifest.") self.checkpointer.load_on_dp_ranks(self.rng, "rng", str(checkpoint)) + self._last_checkpoint = checkpoint return PDDResumeState( checkpoint_path=checkpoint, completed_steps=manifest["completed_steps"], diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py index 122000e2a25..3e70b216f45 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py @@ -854,3 +854,47 @@ def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) source_checkpointer.close() destination_checkpointer.close() third_checkpointer.close() + + +def test_save_chains_validated_parent_without_resolving_latest(tmp_path, monkeypatch) -> None: + pytest.importorskip("nemo_automodel") + rng_module = pytest.importorskip("nemo_automodel.components.training.rng") + if not torch.distributed.is_initialized(): + initialize_pdd_distributed(backend="gloo", timeout_minutes=1) + sample_ids = tuple(f"sample-{index}" for index in range(8)) + + source = build_toy_lifecycle() + source_sampler = _released_sampler(sample_ids) + source_manager, source_checkpointer = _manager( + tmp_path / "checkpoints", + source, + source_sampler, + rng_module.StatefulRNG(1234, ranked=True), + ) + _run_next(source, source_sampler) + first = source_manager.save() + assert json.loads((first / "manifest.json").read_text())["parent_checkpoint"] is None + + resumed = build_toy_lifecycle() + resumed_sampler = _released_sampler(sample_ids) + resumed_manager, resumed_checkpointer = _manager( + tmp_path / "checkpoints", + resumed, + resumed_sampler, + rng_module.StatefulRNG(9999, ranked=True), + ) + assert resumed_manager.load("LATEST") is not None + + def fail_resolution(restore_from): + raise AssertionError(f"save re-resolved {restore_from}") + + monkeypatch.setattr(resumed_manager, "_collective_resolve", fail_resolution) + _run_next(resumed, resumed_sampler) + second = resumed_manager.save() + _run_next(resumed, resumed_sampler) + third = resumed_manager.save() + + assert json.loads((second / "manifest.json").read_text())["parent_checkpoint"] == first.name + assert json.loads((third / "manifest.json").read_text())["parent_checkpoint"] == second.name + source_checkpointer.close() + resumed_checkpointer.close() From ba14577a8283875b15f3fe1f4e6e4717136d23b3 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Mon, 20 Jul 2026 19:11:20 -0700 Subject: [PATCH 35/45] Use native AutoModel lifecycle for PDD Signed-off-by: Meng Xin --- .../fastgen/fastgen_data/__init__.py | 3 - .../fastgen/fastgen_data/collate_fns.py | 83 +- .../fastgen_data/replayable_sampler.py | 227 --- .../fastgen_data/text_to_image_dataset.py | 84 +- examples/diffusers/fastgen/pdd/README.md | 122 +- examples/diffusers/fastgen/pdd/artifacts.py | 165 -- examples/diffusers/fastgen/pdd/checkpoint.py | 916 ---------- .../fastgen/pdd/configs/qwen_image.yaml | 53 +- examples/diffusers/fastgen/pdd/data.py | 401 ----- examples/diffusers/fastgen/pdd/export.py | 598 ------- .../fastgen/pdd/export_qwen_image.py | 349 ---- .../fastgen/pdd/inference_qwen_image.py | 248 ++- .../fastgen/pdd/inference_runtime.py | 301 ---- .../fastgen/pdd/prepare_qwen_image.py | 94 + examples/diffusers/fastgen/pdd/recipe.py | 1529 ++--------------- examples/diffusers/fastgen/pdd/training.py | 1009 +---------- .../preprocess/preprocessing_multiprocess.py | 9 - modelopt/torch/fastgen/methods/pdd.py | 272 +-- .../torch/fastgen/plugins/qwen_image_pdd.py | 129 +- tests/examples/diffusers/fastgen/conftest.py | 2 - .../pdd_checkpoint_failure_distributed.py | 215 --- .../fastgen/pdd_export_distributed.py | 200 --- .../fastgen/pdd_mr210_fsdp_distributed.py | 440 ----- .../diffusers/fastgen/pdd_test_utils.py | 235 --- .../pdd_validation_oracle_distributed.py | 167 -- .../diffusers/fastgen/test_dataset_paths.py | 121 +- .../diffusers/fastgen/test_dataset_splits.py | 37 +- .../examples/diffusers/fastgen/test_layout.py | 7 +- .../diffusers/fastgen/test_pdd_inference.py | 85 + .../fastgen/test_pdd_inference_checkpoint.py | 413 ----- .../fastgen/test_pdd_recipe_setup.py | 860 +-------- .../fastgen/test_pdd_training_lifecycle.py | 900 ---------- .../fastgen/test_pdd_validation_oracle.py | 159 -- .../fastgen/test_vendored_migration.py | 7 +- tests/unit/torch/fastgen/test_pdd_metadata.py | 195 --- tests/unit/torch/fastgen/test_pdd_pipeline.py | 15 + .../unit/torch/fastgen/test_pdd_projection.py | 74 - .../unit/torch/fastgen/test_pdd_public_api.py | 6 +- .../fastgen/test_qwen_image_pdd_plugin.py | 62 +- 39 files changed, 810 insertions(+), 9982 deletions(-) delete mode 100644 examples/diffusers/fastgen/fastgen_data/replayable_sampler.py delete mode 100644 examples/diffusers/fastgen/pdd/artifacts.py delete mode 100644 examples/diffusers/fastgen/pdd/checkpoint.py delete mode 100644 examples/diffusers/fastgen/pdd/data.py delete mode 100644 examples/diffusers/fastgen/pdd/export.py delete mode 100644 examples/diffusers/fastgen/pdd/export_qwen_image.py delete mode 100644 examples/diffusers/fastgen/pdd/inference_runtime.py create mode 100644 examples/diffusers/fastgen/pdd/prepare_qwen_image.py delete mode 100644 tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py delete mode 100644 tests/examples/diffusers/fastgen/pdd_export_distributed.py delete mode 100644 tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py delete mode 100644 tests/examples/diffusers/fastgen/pdd_test_utils.py delete mode 100644 tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py create mode 100644 tests/examples/diffusers/fastgen/test_pdd_inference.py delete mode 100644 tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py delete mode 100644 tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py delete mode 100644 tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py delete mode 100644 tests/unit/torch/fastgen/test_pdd_metadata.py diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index 75a88b4e1eb..d06c001c7c5 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -41,13 +41,11 @@ try: from . import collate_fns as _collate_fns from . import paths as _paths - from . import replayable_sampler as _replayable_sampler from . import resume as _resume from . import splits as _splits from . import text_to_image_dataset as _text_to_image_dataset from .collate_fns import * from .paths import * - from .replayable_sampler import * from .resume import * from .splits import * from .text_to_image_dataset import * @@ -65,7 +63,6 @@ for _module in ( _collate_fns, _paths, - _replayable_sampler, _resume, _splits, _text_to_image_dataset, diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index 76368cc908a..0f25cc854b6 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -13,11 +13,11 @@ # 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 +* :func:`collate_fn_text_to_image` builds the Qwen-Image 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 @@ -32,11 +32,7 @@ """ import functools -import hashlib -import io import logging -from collections.abc import Sequence -from pathlib import Path import torch from nemo_automodel.components.datasets.diffusion.sampler import SequentialBucketSampler @@ -44,7 +40,6 @@ from torchdata.stateful_dataloader import StatefulDataLoader from .paths import resolve_under_root -from .replayable_sampler import ReplayableBatchSampler from .text_to_image_dataset import TextToImageDataset __all__ = [ @@ -60,14 +55,13 @@ def collate_fn_text_to_image( 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). + """Build a 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`` (shape ``[B, seq, dim]``) for CFG-based objectives. negative_text_embeddings_mask: Optional mask for the negative embedding. Returns: @@ -83,7 +77,7 @@ def collate_fn_text_to_image( 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 + # Stack only the keys the FastGen pipelines consume, 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.5.0 unconditionally stacks ``clip_tokens`` / ``t5_tokens``, which # the Qwen-Image cache omits. @@ -103,8 +97,6 @@ def collate_fn_text_to_image( "crop_resolution": torch.stack([item["crop_resolution"] for item in batch]), "original_resolution": torch.stack([item["original_resolution"] for item in batch]), "crop_offset": torch.stack([item["crop_offset"] for item in batch]), - "sample_ids": torch.tensor([item["sample_id"] for item in batch], dtype=torch.long), - "logical_sample_ids": tuple(str(item["sample_id"]) for item in batch), }, } # Optional model-specific embedding fields, when a dataset provides them. @@ -112,7 +104,7 @@ def collate_fn_text_to_image( 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``. + # 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.") @@ -143,16 +135,14 @@ def collate_fn_text_to_image( return image_batch -def _load_negative_prompt_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor, str]: +def _load_negative_prompt_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor]: """Load ``(embed, mask)`` from a negative-prompt-embedding file. Accepts a dict with an ``embed`` tensor (and an optional ``mask`` / ``prompt_embeds_mask`` / ``text_mask``) or a bare embedding tensor; a missing mask defaults to all-ones. """ - payload_bytes = Path(path).read_bytes() - payload_sha256 = hashlib.sha256(payload_bytes).hexdigest() - payload = torch.load(io.BytesIO(payload_bytes), map_location="cpu", weights_only=True) + payload = torch.load(path, map_location="cpu", weights_only=True) neg_embed = payload["embed"] if isinstance(payload, dict) else payload if not torch.is_tensor(neg_embed): raise TypeError( @@ -173,19 +163,7 @@ def _load_negative_prompt_embedding(path: str) -> tuple[torch.Tensor, torch.Tens ) if neg_mask is None: neg_mask = torch.ones(neg_embed.shape[:-1], dtype=torch.long) - return neg_embed, neg_mask, payload_sha256 - - -def _dataset_snapshot_sha256(metadata_sha256: str, negative_sha256: str | None) -> str: - """Bind expected sample content and the static negative condition into one identity.""" - digest = hashlib.sha256(b"modelopt-fastgen-dataset-snapshot-v1\0") - digest.update(bytes.fromhex(metadata_sha256)) - if negative_sha256 is None: - digest.update(b"\0no-negative-prompt") - else: - digest.update(b"\0negative-prompt\0") - digest.update(bytes.fromhex(negative_sha256)) - return digest.hexdigest() + return neg_embed, neg_mask def build_text_to_image_multiresolution_dataloader( @@ -203,16 +181,13 @@ def build_text_to_image_multiresolution_dataloader( pin_memory: bool = True, prefetch_factor: int = 2, negative_prompt_embedding_path: str | None = None, - selected_indices: Sequence[int] | None = None, split: str | None = None, validation_count: int | None = None, split_seed: int = 2026, - exact_resume: bool = False, - verify_payload_hashes: bool | None = None, sampler_seed: int = 42, loader_seed: int | None = None, -) -> tuple[StatefulDataLoader, SequentialBucketSampler | ReplayableBatchSampler]: - """Build the DMD2 text-to-image multiresolution dataloader for ``TrainDiffusionRecipe``. +) -> tuple[StatefulDataLoader, SequentialBucketSampler]: + """Build the shared multiresolution dataloader for ``TrainDiffusionRecipe``. Args: cache_dir: Directory with the preprocessed cache (metadata.json, shards, resolution @@ -229,41 +204,26 @@ 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). - selected_indices: Optional ordered original metadata ordinals to expose. + embedding, bound into the collate and broadcast to every batch. split: Optional deterministic ``"train"`` or ``"validation"`` selection. validation_count: Number of validation samples when ``split`` is set. split_seed: Local seed used to construct deterministic split membership. - exact_resume: Wrap the deterministic sampler with a committed cursor that is - independent of worker prefetch. Required by the PDD lifecycle. - verify_payload_hashes: Require and authenticate each cached tensor against its - ``cache_sha256`` metadata before loading. ``None`` preserves the historical - builder behavior by following ``exact_resume``; PDD sets this explicitly so - replayable cursor state does not require payload hashes. sampler_seed: Seed for the released deterministic bucket sampler. - loader_seed: Optional dedicated seed for DataLoader worker/base-seed generation. PDD - supplies this so recreating an iterator cannot consume its restored training RNG. + loader_seed: Optional dedicated seed for DataLoader worker/base-seed generation. Returns: ``(StatefulDataLoader, SequentialBucketSampler)``. """ - if verify_payload_hashes is None: - verify_payload_hashes = exact_resume - elif type(verify_payload_hashes) is not bool: - raise TypeError("verify_payload_hashes must be bool or None.") - dataset = TextToImageDataset( cache_dir=cache_dir, train_text_encoder=train_text_encoder, - selected_indices=selected_indices, split=split, validation_count=validation_count, split_seed=split_seed, - verify_payload_hashes=verify_payload_hashes, ) effective_root = dataset.cache_root - # Optional negative-prompt embedding for DMD2 CFG: load once, bind into the collate. + # Load the optional negative-prompt embedding once and bind it into the collate. collate_fn = collate_fn_text_to_image if negative_prompt_embedding_path is not None: negative_path = resolve_under_root( @@ -271,8 +231,7 @@ def build_text_to_image_multiresolution_dataloader( negative_prompt_embedding_path, "negative prompt embedding", ) - neg_embed, neg_mask, negative_sha256 = _load_negative_prompt_embedding(str(negative_path)) - dataset.negative_prompt_embedding_sha256 = negative_sha256 + 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", @@ -286,11 +245,6 @@ def build_text_to_image_multiresolution_dataloader( negative_text_embeddings=neg_embed, negative_text_embeddings_mask=neg_mask, ) - dataset.dataset_snapshot_sha256 = _dataset_snapshot_sha256( - dataset.metadata_sha256, - dataset.negative_prompt_embedding_sha256, - ) - sampler = SequentialBucketSampler( dataset, base_batch_size=batch_size, @@ -303,14 +257,13 @@ def build_text_to_image_multiresolution_dataloader( num_replicas=dp_world_size, rank=dp_rank, ) - batch_sampler = ReplayableBatchSampler(sampler) if exact_resume else sampler 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=batch_sampler, + batch_sampler=sampler, collate_fn=collate_fn, num_workers=num_workers, pin_memory=pin_memory, @@ -326,9 +279,9 @@ def build_text_to_image_multiresolution_dataloader( effective_root, len(dataset), dataset.total_num_samples, - len(batch_sampler), + len(sampler), batch_size, dp_rank, dp_world_size, ) - return dataloader, batch_sampler + return dataloader, sampler diff --git a/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py b/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py deleted file mode 100644 index 720f608f0e4..00000000000 --- a/examples/diffusers/fastgen/fastgen_data/replayable_sampler.py +++ /dev/null @@ -1,227 +0,0 @@ -# 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. - -"""Committed-cursor wrapper for deterministic, prefetched batch samplers.""" - -from __future__ import annotations - -import hashlib -import struct -from collections.abc import Iterator, Mapping, Sequence -from typing import Any - -import torch -from torch.utils.data import Sampler - -__all__ = ["ReplayableBatchSampler"] - -_STATE_VERSION = 1 - - -class ReplayableBatchSampler(Sampler[list[int]]): - """Separate actually consumed batches from batches yielded to worker prefetch. - - The wrapped sampler remains the authority for deterministic rank/epoch batch plans. - This wrapper materializes that plan compactly and advances its committed cursor only - after the training loop confirms the logical sample IDs it consumed. - """ - - def __init__(self, sampler: Sampler[list[int]]) -> None: - if not isinstance(sampler, Sampler): - raise TypeError(f"sampler must be a torch Sampler, got {type(sampler).__name__}.") - dataset = getattr(sampler, "dataset", None) - metadata = getattr(dataset, "metadata", None) - if not isinstance(metadata, Sequence): - raise TypeError("sampler.dataset.metadata must be a sequence.") - - self.sampler = sampler - self.dataset = dataset - self.epoch = int(getattr(sampler, "epoch", 0)) - self.committed_batches = 0 - self.sample_slots_consumed = 0 - self._yielded_batches = 0 - self._flat_indices = torch.empty(0, dtype=torch.int64) - self._offsets = torch.zeros(1, dtype=torch.int64) - self._plan_sha256 = "" - self._build_plan() - - def _build_plan(self) -> None: - self.sampler.set_epoch(self.epoch) - self.sampler.load_state_dict({"epoch": self.epoch, "batches_yielded": 0}) - batches = [tuple(int(index) for index in batch) for batch in self.sampler] - if not batches: - raise ValueError("replayable batch plan must contain at least one batch.") - if any(not batch for batch in batches): - raise ValueError("replayable batch plan cannot contain an empty batch.") - - flat = [index for batch in batches for index in batch] - offsets = [0] - for batch in batches: - offsets.append(offsets[-1] + len(batch)) - self._flat_indices = torch.tensor(flat, dtype=torch.int64) - self._offsets = torch.tensor(offsets, dtype=torch.int64) - - digest = hashlib.sha256() - digest.update(b"modelopt-pdd-batch-plan-v1\0") - digest.update(struct.pack(">q", self.epoch)) - digest.update(struct.pack(">q", int(getattr(self.sampler, "rank", 0)))) - digest.update(struct.pack(">q", int(getattr(self.sampler, "num_replicas", 1)))) - for batch in batches: - digest.update(struct.pack(">q", len(batch))) - for index in batch: - digest.update(struct.pack(">q", index)) - self._plan_sha256 = digest.hexdigest() - self._yielded_batches = self.committed_batches - - @property - def plan_sha256(self) -> str: - return self._plan_sha256 - - @property - def remaining_batches(self) -> int: - return len(self) - self.committed_batches - - def _batch_indices(self, batch_index: int) -> list[int]: - if not 0 <= batch_index < len(self): - raise IndexError(f"batch_index={batch_index} is outside [0, {len(self)}).") - start = int(self._offsets[batch_index]) - end = int(self._offsets[batch_index + 1]) - return self._flat_indices[start:end].tolist() - - def _sample_ids(self, batch_index: int) -> tuple[str, ...]: - logical_ids = getattr(self.dataset, "logical_sample_ids", None) - if logical_ids is None: - dataset_sample_ids = getattr(self.dataset, "sample_ids", None) - if not isinstance(dataset_sample_ids, Sequence): - raise TypeError("sampler.dataset must expose logical_sample_ids or sample_ids.") - logical_ids = tuple(str(sample_id) for sample_id in dataset_sample_ids) - if not isinstance(logical_ids, Sequence): - raise TypeError("sampler.dataset.logical_sample_ids must be a sequence.") - batch_sample_ids: list[str] = [] - for index in self._batch_indices(batch_index): - sample_id = logical_ids[index] - if not isinstance(sample_id, str) or not sample_id: - raise ValueError(f"dataset.logical_sample_ids[{index}] must be a nonempty string.") - batch_sample_ids.append(sample_id) - return tuple(batch_sample_ids) - - def expected_next_sample_ids(self) -> tuple[str, ...]: - """Return the next committed batch's logical IDs without consuming it.""" - if self.committed_batches == len(self): - return () - return self._sample_ids(self.committed_batches) - - def commit(self, sample_ids: Sequence[str]) -> None: - """Advance the durable cursor after verifying the collated logical IDs.""" - if isinstance(sample_ids, str) or not isinstance(sample_ids, Sequence): - raise TypeError("sample_ids must be a sequence of strings.") - actual = tuple(sample_ids) - if any(not isinstance(sample_id, str) for sample_id in actual): - raise TypeError("sample_ids must contain only strings.") - expected = self.expected_next_sample_ids() - if not expected: - raise RuntimeError("cannot commit beyond the end of the batch plan.") - if actual != expected: - raise RuntimeError( - f"consumed sample IDs do not match the committed cursor: " - f"expected={expected}, actual={actual}." - ) - self.committed_batches += 1 - self.sample_slots_consumed += len(actual) - - def set_epoch(self, epoch: int) -> None: - if type(epoch) is not int or epoch < 0: - raise ValueError("epoch must be an integer >= 0.") - if epoch == self.epoch: - return - if self.committed_batches != len(self): - raise RuntimeError("cannot change epoch before every planned batch is committed.") - self.epoch = epoch - self.committed_batches = 0 - self._build_plan() - - def state_dict(self) -> dict[str, Any]: - return { - "schema_version": _STATE_VERSION, - "epoch": self.epoch, - "committed_batches": self.committed_batches, - "sample_slots_consumed": self.sample_slots_consumed, - "plan_sha256": self.plan_sha256, - "next_sample_ids": list(self.expected_next_sample_ids()), - } - - def load_state_dict(self, state: Mapping[str, Any]) -> None: - if not isinstance(state, Mapping): - raise TypeError("replayable sampler state must be a mapping.") - expected_keys = { - "schema_version", - "epoch", - "committed_batches", - "sample_slots_consumed", - "plan_sha256", - "next_sample_ids", - } - if set(state) != expected_keys: - raise ValueError( - f"replayable sampler state keys mismatch: expected={sorted(expected_keys)}, " - f"actual={sorted(state)}." - ) - if state["schema_version"] != _STATE_VERSION: - raise ValueError(f"unsupported replayable sampler schema {state['schema_version']!r}.") - epoch = state["epoch"] - committed = state["committed_batches"] - consumed = state["sample_slots_consumed"] - if type(epoch) is not int or epoch < 0: - raise ValueError("saved epoch must be an integer >= 0.") - if type(committed) is not int or committed < 0: - raise ValueError("saved committed_batches must be an integer >= 0.") - if type(consumed) is not int or consumed < 0: - raise ValueError("saved sample_slots_consumed must be an integer >= 0.") - if not isinstance(state["plan_sha256"], str): - raise TypeError("saved plan_sha256 must be a string.") - if not isinstance(state["next_sample_ids"], list) or any( - not isinstance(sample_id, str) for sample_id in state["next_sample_ids"] - ): - raise TypeError("saved next_sample_ids must be a list of strings.") - - self.epoch = epoch - self.committed_batches = 0 - self._build_plan() - if committed > len(self): - raise ValueError( - f"saved committed_batches={committed} exceeds plan length {len(self)}." - ) - if self.plan_sha256 != state["plan_sha256"]: - raise RuntimeError("reconstructed batch plan does not match the saved plan hash.") - self.committed_batches = committed - self.sample_slots_consumed = consumed - self._yielded_batches = committed - if list(self.expected_next_sample_ids()) != state["next_sample_ids"]: - raise RuntimeError("reconstructed next sample IDs do not match the checkpoint.") - - def __iter__(self) -> Iterator[list[int]]: - start = self.committed_batches - flat_indices = self._flat_indices - offsets = self._offsets - total_batches = int(offsets.numel() - 1) - self._yielded_batches = start - for batch_index in range(start, total_batches): - batch_start = int(offsets[batch_index]) - batch_end = int(offsets[batch_index + 1]) - self._yielded_batches = batch_index + 1 - yield flat_indices[batch_start:batch_end].tolist() - - def __len__(self) -> int: - return int(self._offsets.numel() - 1) 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 2fb1d410fdd..c028f04941f 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -13,10 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import hashlib -import io import json -from collections.abc import Sequence from pathlib import Path import torch @@ -35,55 +32,34 @@ def __init__( self, cache_dir: str | Path, train_text_encoder: bool = False, - selected_indices: Sequence[int] | None = None, split: str | None = None, validation_count: int | None = None, split_seed: int = 2026, - verify_payload_hashes: bool = False, ): """ Args: cache_dir: Directory containing preprocessed cache train_text_encoder: If True, returns tokens instead of embeddings - selected_indices: Optional ordered original metadata ordinals to expose. split: Optional deterministic ``"train"`` or ``"validation"`` selection. validation_count: Number of validation samples when ``split`` is set. split_seed: Local seed used to construct deterministic split membership. - verify_payload_hashes: Require and authenticate cached tensor content on every - load. When false, payloads load directly and the cache must remain immutable - for deterministic resume. """ - if selected_indices is not None and split is not None: - raise ValueError("selected_indices and split are mutually exclusive") if split not in (None, "train", "validation"): raise ValueError("split must be null, 'train', or 'validation'") if split is not None and validation_count is None: raise ValueError("validation_count is required when split is set") - if type(verify_payload_hashes) is not bool: - raise TypeError("verify_payload_hashes must be bool.") self.train_text_encoder = train_text_encoder self.cache_root = resolve_cache_root(cache_dir) - self._selected_indices = selected_indices self._split = split self._validation_count = validation_count self._split_seed = split_seed - self._verify_payload_hashes = verify_payload_hashes self._resolved_cache_files: dict[int, Path] = {} - self.negative_prompt_embedding_sha256: str | None = None - self.dataset_snapshot_sha256: str | None = None super().__init__(str(self.cache_root), quantization=64) - @property - def verify_payload_hashes(self) -> bool: - """Whether sample payload bytes are authenticated before deserialization.""" - return self._verify_payload_hashes - 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") - digest = hashlib.sha256(b"modelopt-fastgen-metadata-v1\0") index_bytes = metadata_file.read_bytes() - digest.update(index_bytes) index = json.loads(index_bytes) if not isinstance(index, dict) or not isinstance(index.get("shards"), list): raise ValueError( @@ -98,9 +74,6 @@ def _load_metadata(self) -> list[dict]: self.cache_root, shard_name, f"metadata shard {shard_index}" ) shard_bytes = shard_path.read_bytes() - digest.update(shard_name.encode()) - digest.update(b"\0") - digest.update(shard_bytes) shard = json.loads(shard_bytes) if not isinstance(shard, list): raise ValueError(f"metadata shard {shard_path} must contain a list") @@ -114,32 +87,13 @@ def _load_metadata(self) -> list[dict]: raise TypeError( f"metadata shard {shard_path} item {shard_item_index} has invalid cache_file" ) - cache_sha256 = item.get("cache_sha256") - if cache_sha256 is not None and ( - not isinstance(cache_sha256, str) - or len(cache_sha256) != 64 - or any(character not in "0123456789abcdef" for character in cache_sha256) - ): - raise ValueError( - f"metadata shard {shard_path} item {shard_item_index} has invalid " - "cache_sha256" - ) - if self._verify_payload_hashes and cache_sha256 is None: - raise ValueError( - f"metadata shard {shard_path} item {shard_item_index} has no " - "cache_sha256 required when verify_payload_hashes=true" - ) 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.metadata_sha256 = digest.hexdigest() - self.payload_hashes_complete = all( - item.get("cache_sha256") is not None for item in complete_metadata - ) if self._split is None: - self.sample_ids = self._validate_selected_indices(self.total_num_samples) + self.sample_ids = list(range(self.total_num_samples)) else: if self._validation_count is None: raise RuntimeError("validation_count was not resolved for the requested split") @@ -149,28 +103,8 @@ def _load_metadata(self) -> list[dict]: self._split_seed, ) self.sample_ids = train if self._split == "train" else validation - self.logical_sample_ids = [str(sample_id) for sample_id in self.sample_ids] return [complete_metadata[index] for index in self.sample_ids] - def _validate_selected_indices(self, num_samples: int) -> list[int]: - if self._selected_indices is None: - return list(range(num_samples)) - if isinstance(self._selected_indices, str | bytes) or not isinstance( - self._selected_indices, Sequence - ): - raise TypeError("selected_indices must be a sequence of integers") - selected = list(self._selected_indices) - if not selected: - raise ValueError("selected_indices must not be empty") - for index in selected: - if type(index) is not int: - raise TypeError("selected_indices must contain only non-bool integers") - if not 0 <= index < num_samples: - raise ValueError(f"selected index {index} is outside [0, {num_samples})") - if len(set(selected)) != len(selected): - raise ValueError("selected_indices must be unique") - return selected - def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: """Load a single sample.""" item = self.metadata[idx] @@ -184,20 +118,7 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: ) self._resolved_cache_files[idx] = cache_file - # Exact-resume mode authenticates the same bytes passed to torch.load, avoiding a - # hash-then-reopen race if a shared cache changes during training. - if self._verify_payload_hashes: - payload = cache_file.read_bytes() - actual_sha256 = hashlib.sha256(payload).hexdigest() - expected_sha256 = item["cache_sha256"] - if actual_sha256 != expected_sha256: - raise RuntimeError( - f"sample cache file {sample_id} SHA-256 mismatch: " - f"expected {expected_sha256}, found {actual_sha256}" - ) - data = torch.load(io.BytesIO(payload), map_location="cpu", weights_only=True) - else: - data = torch.load(cache_file, map_location="cpu", weights_only=True) + 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 = { @@ -209,7 +130,6 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: "image_path": data["image_path"], "bucket_id": item["bucket_id"], "aspect_ratio": item.get("aspect_ratio", 1.0), - "sample_id": sample_id, } if self.train_text_encoder: output["clip_tokens"] = data["clip_tokens"].squeeze(0) diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index 6332264c9b6..bd69c01c354 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -1,84 +1,78 @@ # PDD for Qwen-Image -Parallel Decoding Distillation (PDD) trains one Qwen-Image student call to predict several -consecutive rectified-flow updates. The student keeps the original transformer backbone and widens -only its output projection to 128 velocity heads. During training it samples aligned block starts -and target spans from 1 through 64 intervals, so the same checkpoint can use different supported -block schedules at inference. +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 provided schedules use the 128-interval grid as follows: +The frozen Qwen-Image teacher constructs the PDD target. The complete student transformer, +including the widened output projection, is finetuned at a constant learning rate of `5e-5`; this +is not a heads-only run. Training samples aligned target spans from 4 through 64 intervals, so the +same checkpoint supports multiple inference schedules. -| Schedule | Block sizes | Transformer calls | -|---|---|---:| -| `pdd-2` | `[64, 64]` | 2 | -| `pdd-4` | `[32, 32, 32, 32]` | 4 | -| `pdd-8` | `[16, 16, 16, 16, 16, 16, 16, 16]` | 8 | +## Ownership -## Training +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 AutoModel, Diffusers, or Qwen source. -Install the shared requirements from the repository root, then launch with released AutoModel -APIs. No AutoModel, Diffusers, or Qwen source changes are required. +## Prepare the student -The example loads the ordinary Diffusers Qwen transformer, then binds a ModelOpt-owned forward on -that same model instance to reproduce FastGen MR210's executed Qwen path. Normalized time remains -FP32 through `time_text_embed`; the text mask is passed to Qwen blocks as MR210 does instead of being -converted into Diffusers' canonical joint mask. Teacher classifier-free guidance is rounded in the -BF16 model-output dtype, globally norm-rescaled in FP32 with a `1e-5` denominator floor, and cast -back to BF16 before the PDD loss. No AutoModel, Diffusers, or Qwen source is edited. -Guidance-embedded models, PEFT, and QKV fusion remain outside this first example. +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/requirements.txt export MODELOPT_FASTGEN_DATASET_CACHE_DIR=/absolute/path/to/qwen_image_cache + torchrun --standalone --nproc-per-node=8 \ examples/diffusers/fastgen/pdd/finetune.py \ - --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml + --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml \ + --fsdp.dp_size=8 ``` -The cache must contain `metadata.json`, its declared shards, cached tensors, and -`negative_prompt_embedding.pt`. The environment variable overrides the configured cache root. All -metadata, tensor, and negative-embedding paths must still resolve inside that effective root. -PDD's committed sampler makes the next batch, sample IDs, RNG, optimizer, and scheduler state -exactly replayable without requiring per-payload hashes. Hashless caches must therefore remain -immutable for the duration of a run: metadata and the negative embedding are bound into the -checkpoint identity, but an in-place tensor payload change cannot be detected. - -Set `data.dataloader.verify_payload_hashes: true` to require the `cache_sha256` field written by -the shared preprocessor and authenticate every tensor's bytes before loading it. This stricter mode -has additional read and SHA-256 cost and fails immediately when a hash is missing or mismatched. - -Training deterministically derives disjoint train and validation membership from metadata ordinals; -it does not rewrite the cache or require separate split manifests. The default recipe uses 2,000 -validation samples, constant learning rate `5e-5`, per-rank batch size 4, 128 heads, start indices -aligned by 4, and target spans from 1 through 64 intervals. Other GPU topologies can set per-rank -batch size to obtain the desired global batch size because this recipe does not use gradient -accumulation. -Checkpoints include the student, optimizer, scheduler, RNG, trainer, and exact replayable sampler -state needed to resume the next committed batch. FP32 master parameters and Adam state are sharded -while forward/backward uses BF16 model parameters and outputs and gradient reduction remains FP32. -The adapter casts packed image/text inputs to BF16 while preserving FP32 normalized time at the -FSDP root and Qwen time embedder. - -Checkpointed training, export, and inference require the remote model ID and exact lowercase -40-character Hugging Face commit in the provided config because the frozen teacher is rebuilt -rather than checkpointed. - -## Export and inference - -The v1 export restores the sharded checkpoint with the same total process count that created it; -cross-world-size DCP resharding is not yet part of this example. The command below therefore -applies to a checkpoint trained with eight ranks. For a 64-rank checkpoint, use the cluster -launcher with 64 export ranks before running single-process inference. +The cache must contain `metadata.json`, its declared tensor shards, and +`negative_prompt_embedding.pt`. `MODELOPT_FASTGEN_DATASET_CACHE_DIR` overrides the configured +cache root; paths declared by the dataset remain confined to that root. -```bash -torchrun --standalone --nproc-per-node=8 \ - examples/diffusers/fastgen/pdd/export_qwen_image.py \ - --config examples/diffusers/fastgen/pdd/configs/qwen_image.yaml \ - --checkpoint LATEST --output-dir /path/to/pdd-export +The checked-in recipe targets 50,000 optimizer steps with global batch size 256, local batch size +4, and constant `5e-5` learning rate. `checkpoint.restore_from: LATEST` lets AutoModel resume the +latest native checkpoint. 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 \ - --export-dir /path/to/pdd-export --schedule pdd-4 \ - --prompt-id red-cube-0001 --prompt "a small red cube on a white table" \ + --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 /path/to/pdd4.png --result-json /path/to/pdd4.json + --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/artifacts.py b/examples/diffusers/fastgen/pdd/artifacts.py deleted file mode 100644 index 709b6a03ad9..00000000000 --- a/examples/diffusers/fastgen/pdd/artifacts.py +++ /dev/null @@ -1,165 +0,0 @@ -# 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. - -"""Strict canonical-JSON and relative-artifact helpers for the PDD example.""" - -from __future__ import annotations - -import hashlib -import json -import math -import os -from collections.abc import Mapping, Sequence -from pathlib import Path, PurePosixPath -from typing import Any - - -def sha256_file(path: Path) -> str: - """Hash one regular file without following a symlink.""" - if not path.is_file() or path.is_symlink(): - raise RuntimeError(f"PDD artifact is not a regular file: {path}.") - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def require_sha256(value: Any, *, name: str) -> str: - """Validate and normalize a hexadecimal SHA-256 value.""" - if not isinstance(value, str) or len(value) != 64: - raise ValueError(f"{name} must be a 64-character SHA-256 digest.") - try: - int(value, 16) - except ValueError as error: - raise ValueError(f"{name} must be hexadecimal.") from error - return value.lower() - - -def _validate_json_value(value: Any, *, name: str = "JSON") -> None: - if value is None or isinstance(value, str | bool | int): - return - if isinstance(value, float): - if not math.isfinite(value): - raise ValueError(f"{name} contains a non-finite number.") - return - if isinstance(value, Mapping): - if any(not isinstance(key, str) for key in value): - raise TypeError(f"{name} object keys must be strings.") - for key, item in value.items(): - _validate_json_value(item, name=f"{name}.{key}") - return - if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): - for index, item in enumerate(value): - _validate_json_value(item, name=f"{name}[{index}]") - return - raise TypeError(f"{name} contains unsupported type {type(value).__name__}.") - - -def canonical_json_bytes(value: Any) -> bytes: - """Serialize finite JSON deterministically with a trailing newline.""" - _validate_json_value(value) - return ( - json.dumps( - value, - allow_nan=False, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True, - ) - + "\n" - ).encode("utf-8") - - -def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: - value: dict[str, Any] = {} - for key, item in pairs: - if key in value: - raise ValueError(f"canonical JSON contains duplicate key {key!r}.") - value[key] = item - return value - - -def _reject_json_constant(token: str) -> None: - raise ValueError(f"canonical JSON contains {token}.") - - -def load_canonical_json(path: Path) -> Any: - """Load canonical JSON, rejecting duplicates, NaN/Inf, and noncanonical bytes.""" - if not path.is_file() or path.is_symlink(): - raise RuntimeError(f"PDD JSON artifact is not a regular file: {path}.") - raw = path.read_bytes() - try: - value = json.loads( - raw, - object_pairs_hook=_unique_object, - parse_constant=_reject_json_constant, - ) - except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as error: - raise RuntimeError(f"cannot parse canonical PDD JSON {path}.") from error - try: - expected = canonical_json_bytes(value) - except (TypeError, ValueError) as error: - raise RuntimeError(f"invalid canonical PDD JSON {path}.") from error - if raw != expected: - raise RuntimeError(f"PDD JSON is not in canonical form: {path}.") - return value - - -def write_canonical_json(path: Path, value: Any) -> None: - """Create one canonical JSON file and fsync its contents.""" - data = canonical_json_bytes(value) - with path.open("xb") as stream: - stream.write(data) - stream.flush() - os.fsync(stream.fileno()) - - -def resolve_relative_artifact(root: Path, reference: str) -> Path: - """Resolve a normalized POSIX reference beneath root with no symlink component.""" - if not isinstance(reference, str) or not reference: - raise ValueError("artifact reference must be a non-empty string.") - if "\\" in reference: - raise ValueError(f"artifact reference must use POSIX separators: {reference!r}.") - pure = PurePosixPath(reference) - if pure.is_absolute() or any(part in ("", ".", "..") for part in pure.parts): - raise ValueError(f"artifact reference must be normalized and relative: {reference!r}.") - root = root.resolve() - if not root.is_dir() or root.is_symlink(): - raise RuntimeError(f"PDD artifact root is not a regular directory: {root}.") - candidate = root - for part in pure.parts: - candidate = candidate / part - if candidate.is_symlink(): - raise RuntimeError(f"PDD artifact reference traverses a symlink: {reference!r}.") - resolved = candidate.resolve() - try: - resolved.relative_to(root) - except ValueError as error: - raise ValueError(f"artifact reference escapes its root: {reference!r}.") from error - if not resolved.is_file(): - raise FileNotFoundError(f"PDD artifact is missing: {reference!r}.") - return resolved - - -def validate_artifact_reference(root: Path, value: Any, *, name: str) -> Path: - """Validate an exact path/hash reference and return the verified regular file.""" - if not isinstance(value, Mapping) or set(value) != {"path", "sha256"}: - raise ValueError(f"{name} must contain exactly path and sha256.") - path = resolve_relative_artifact(root, value["path"]) - expected = require_sha256(value["sha256"], name=f"{name}.sha256") - if sha256_file(path) != expected: - raise RuntimeError(f"{name} SHA-256 does not match {value['path']!r}.") - return path diff --git a/examples/diffusers/fastgen/pdd/checkpoint.py b/examples/diffusers/fastgen/pdd/checkpoint.py deleted file mode 100644 index b8841ed8bf3..00000000000 --- a/examples/diffusers/fastgen/pdd/checkpoint.py +++ /dev/null @@ -1,916 +0,0 @@ -# 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. - -"""Atomic, strict PDD checkpoint publication around the stock AutoModel Checkpointer.""" - -from __future__ import annotations - -import hashlib -import json -import math -import os -import shutil -import uuid -from collections.abc import Mapping -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING, Any - -import torch.distributed as dist - -from modelopt.torch.fastgen import PDDMetadata -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION - -if TYPE_CHECKING: - from collections.abc import Sequence - -_CHECKPOINT_SCHEMA_VERSION = 5 -_COMPLETE_SCHEMA_VERSION = 1 - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _require_sha256(value: Any, *, name: str) -> str: - if not isinstance(value, str) or len(value) != 64: - raise ValueError(f"{name} must be a 64-character SHA-256 digest.") - try: - int(value, 16) - except ValueError as error: - raise ValueError(f"{name} must be hexadecimal.") from error - return value.lower() - - -def _require_qwen_image_execution(identity: Any) -> None: - if not isinstance(identity, Mapping) or identity.get("qwen_image") != { - "execution": QWEN_IMAGE_PDD_EXECUTION - }: - raise RuntimeError("PDD checkpoint has an incompatible Qwen execution identity.") - - -def _rank() -> int: - return dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 - - -def _world_size() -> int: - return dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 - - -def _barrier() -> None: - if dist.is_available() and dist.is_initialized(): - dist.barrier() - - -def _broadcast_rank0_payload(value: Any) -> Any: - payload = [value] - if dist.is_available() and dist.is_initialized(): - dist.broadcast_object_list(payload, src=0) - return payload[0] - - -def _gather_objects(value: Any) -> list[Any]: - if not dist.is_available() or not dist.is_initialized(): - return [value] - values: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(values, value) - return values - - -def _fsync_file(path: Path) -> None: - with path.open("rb") as stream: - os.fsync(stream.fileno()) - - -def _fsync_directory(path: Path) -> None: - descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - - -def _fsync_tree(root: Path) -> None: - for path in sorted(root.rglob("*"), key=lambda candidate: len(candidate.parts), reverse=True): - if path.is_symlink(): - raise RuntimeError(f"PDD checkpoint staging contains a symlink: {path}.") - if path.is_file(): - _fsync_file(path) - elif path.is_dir(): - _fsync_directory(path) - _fsync_directory(root) - - -def _dcp_payload_hashes(checkpoint: Path) -> dict[str, str]: - hashes: dict[str, str] = {} - for component in ("model", "optim"): - root = checkpoint / component - if not root.is_dir() or root.is_symlink(): - raise RuntimeError(f"PDD checkpoint is missing the {component} DCP directory.") - files: list[Path] = [] - for path in root.rglob("*"): - if path.is_symlink(): - raise RuntimeError(f"PDD {component} DCP tree contains a symlink: {path}.") - if path.is_file(): - files.append(path) - relative_files = {path.relative_to(checkpoint).as_posix() for path in files} - if f"{component}/.metadata" not in relative_files: - raise RuntimeError(f"PDD checkpoint is missing strict {component} DCP metadata.") - if not any(relative != f"{component}/.metadata" for relative in relative_files): - raise RuntimeError(f"PDD checkpoint is missing {component} DCP payload shards.") - for path in files: - hashes[path.relative_to(checkpoint).as_posix()] = _sha256(path) - return hashes - - -def _atomic_text(path: Path, text: str) -> None: - temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") - with temporary.open("w") as stream: - stream.write(text) - stream.flush() - os.fsync(stream.fileno()) - os.replace(temporary, path) - _fsync_directory(path.parent) - - -def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None: - _atomic_text(path, json.dumps(payload, indent=2, sort_keys=True) + "\n") - - -def _read_json(path: Path) -> dict[str, Any]: - try: - value = json.loads(path.read_text()) - except (OSError, json.JSONDecodeError) as error: - raise RuntimeError(f"cannot read PDD checkpoint JSON {path}.") from error - if not isinstance(value, dict): - raise RuntimeError(f"PDD checkpoint JSON {path} must contain an object.") - return value - - -def build_pdd_checkpoint_identity( - *, - qwen_image_execution: str, - metadata: PDDMetadata, - model_id: str, - model_revision: str | None, - guidance_scale: float | None, - ordered_train_id_sha256: str, - ordered_heldout_id_sha256: str, - dataset_snapshot_sha256: str, - local_batch_size: int, - grad_accumulation_steps: int, - training_seed: int, - validation_seed: int, - validation_every_steps: int, - max_grad_norm: float, - zero_grad_warmup_steps: int, - activation_checkpointing: bool, - dtype: str, - optimizer: Any, - scheduler: Any, -) -> dict[str, Any]: - """Build the strict, path-independent compatibility identity for PDD resume.""" - if qwen_image_execution != QWEN_IMAGE_PDD_EXECUTION: - raise ValueError("qwen_image_execution must identify the bound FastGen MR210 forward.") - if not isinstance(metadata, PDDMetadata): - raise TypeError("metadata must be PDDMetadata.") - if not isinstance(model_id, str) or not model_id: - raise ValueError("model_id must be a non-empty string.") - if ( - not isinstance(model_revision, str) - or len(model_revision) != 40 - or any(character not in "0123456789abcdef" for character in model_revision) - ): - raise ValueError("model_revision must be an exact lowercase 40-character commit.") - if guidance_scale is not None and ( - isinstance(guidance_scale, bool) or not isinstance(guidance_scale, int | float) - ): - raise TypeError("guidance_scale must be a real number or null.") - if type(local_batch_size) is not int or local_batch_size < 1: - raise ValueError("local_batch_size must be an integer >= 1.") - if grad_accumulation_steps != 1: - raise ValueError("PDD v1 exact resume requires grad_accumulation_steps=1.") - for name, value, minimum in ( - ("training_seed", training_seed, 0), - ("validation_seed", validation_seed, 0), - ("validation_every_steps", validation_every_steps, 1), - ("zero_grad_warmup_steps", zero_grad_warmup_steps, 0), - ): - if type(value) is not int or value < minimum: - raise ValueError(f"{name} must be an integer >= {minimum}.") - if isinstance(max_grad_norm, bool) or not isinstance(max_grad_norm, int | float): - raise TypeError("max_grad_norm must be a real number.") - if not math.isfinite(max_grad_norm) or max_grad_norm <= 0: - raise ValueError("max_grad_norm must be finite and > 0.") - if type(activation_checkpointing) is not bool: - raise TypeError("activation_checkpointing must be bool.") - if not isinstance(dtype, str) or not dtype: - raise ValueError("dtype must be a non-empty string.") - if type(optimizer).__module__ != "torch.optim.adamw" or type(optimizer).__name__ != "AdamW": - raise TypeError("PDD checkpoint identity requires the stock torch.optim.AdamW optimizer.") - return { - "schema_version": _CHECKPOINT_SCHEMA_VERSION, - "qwen_image": {"execution": qwen_image_execution}, - "model": {"id": model_id, "revision": model_revision, "dtype": dtype}, - "pdd_metadata": metadata.to_dict(), - "guidance": {"scale": None if guidance_scale is None else float(guidance_scale)}, - "data": { - "ordered_train_id_sha256": _require_sha256( - ordered_train_id_sha256, - name="ordered_train_id_sha256", - ), - "ordered_heldout_id_sha256": _require_sha256( - ordered_heldout_id_sha256, - name="ordered_heldout_id_sha256", - ), - "dataset_snapshot_sha256": _require_sha256( - dataset_snapshot_sha256, - name="dataset_snapshot_sha256", - ), - "local_batch_size": local_batch_size, - "grad_accumulation_steps": grad_accumulation_steps, - }, - "topology": {"world_size": _world_size(), "pure_data_parallel": True}, - "training": { - "seed": training_seed, - "validation_seed": validation_seed, - "validation_every_steps": validation_every_steps, - "max_grad_norm": float(max_grad_norm), - "zero_grad_warmup_steps": zero_grad_warmup_steps, - "activation_checkpointing": activation_checkpointing, - }, - "optimizer": { - "class": "torch.optim.AdamW", - "param_groups": [ - { - "lr": float(group["lr"]), - "betas": [float(beta) for beta in group["betas"]], - "eps": float(group["eps"]), - "weight_decay": float(group["weight_decay"]), - "amsgrad": bool(group["amsgrad"]), - "maximize": bool(group["maximize"]), - "capturable": bool(group["capturable"]), - "differentiable": bool(group["differentiable"]), - "foreach": bool(group["foreach"]), - "fused": bool(group["fused"]), - } - for group in optimizer.param_groups - ], - }, - "scheduler": { - "class": f"{type(scheduler).__module__}.{type(scheduler).__qualname__}", - "base_lrs": [float(value) for value in scheduler.base_lrs], - "policy": "constant", - }, - } - - -@dataclass(frozen=True) -class PDDResumeState: - """Restored progress plus the first logical IDs that must be served next.""" - - checkpoint_path: Path - completed_steps: int - sample_slots_consumed: int - expected_next_sample_ids: tuple[str, ...] - parent_checkpoint: str | None - - def verify_first_batch(self, sample_ids: Sequence[str]) -> None: - if tuple(sample_ids) != self.expected_next_sample_ids: - raise RuntimeError( - "first resumed sample IDs do not match the checkpoint: " - f"expected={self.expected_next_sample_ids}, actual={tuple(sample_ids)}." - ) - - -class _StepSchedulerCheckpointState: - """Rank-local carrier for the normalized next-data StepScheduler cursor.""" - - def __init__(self) -> None: - self.state: dict[str, int] = {"step": 0, "epoch": 0} - - def state_dict(self) -> dict[str, int]: - return dict(self.state) - - def load_state_dict(self, state: Mapping[str, Any]) -> None: - if set(state) != {"step", "epoch"}: - raise ValueError("PDD StepScheduler state must contain step and epoch.") - step = state["step"] - epoch = state["epoch"] - if type(step) is not int or step < 0 or type(epoch) is not int or epoch < 0: - raise ValueError("PDD StepScheduler step and epoch must be nonnegative integers.") - self.state = {"step": step, "epoch": epoch} - - -def _checkpoint_sidecar_paths(checkpoint: Path, world_size: int) -> list[Path]: - paths: list[Path] = [] - for rank in range(world_size): - paths.extend( - ( - checkpoint / "rng" / f"rng_dp_rank_{rank}.pt", - checkpoint / "sampler" / f"sampler_dp_rank_{rank}.pt", - checkpoint / "step_scheduler" / f"step_scheduler_dp_rank_{rank}.pt", - checkpoint / "trainer" / f"trainer_dp_rank_{rank}.pt", - ) - ) - return paths - - -def validate_pdd_training_checkpoint( - checkpoint: str | Path, - *, - expected_identity: Mapping[str, Any] | None = None, - expected_world_size: int | None = None, -) -> dict[str, Any]: - """Validate a complete training checkpoint without deserializing pickle sidecars.""" - unresolved_checkpoint = Path(checkpoint) - if unresolved_checkpoint.is_symlink(): - raise RuntimeError(f"PDD checkpoint cannot be a symlink: {unresolved_checkpoint}.") - checkpoint = unresolved_checkpoint.resolve() - if not checkpoint.is_dir(): - raise RuntimeError(f"PDD checkpoint is not a regular directory: {checkpoint}.") - symlinks = [path for path in checkpoint.rglob("*") if path.is_symlink()] - if symlinks: - raise RuntimeError(f"PDD checkpoint contains a symlink: {symlinks[0]}.") - marker_path = checkpoint / "COMPLETE" - manifest_path = checkpoint / "manifest.json" - if ( - not marker_path.is_file() - or marker_path.is_symlink() - or not manifest_path.is_file() - or manifest_path.is_symlink() - ): - raise RuntimeError(f"PDD checkpoint is incomplete: {checkpoint}.") - marker = _read_json(marker_path) - if ( - set(marker) != {"schema_version", "manifest_sha256"} - or marker.get("schema_version") != _COMPLETE_SCHEMA_VERSION - ): - raise RuntimeError("PDD COMPLETE marker is incompatible.") - if marker["manifest_sha256"] != _sha256(manifest_path): - raise RuntimeError("PDD COMPLETE marker does not match manifest content.") - manifest = _read_json(manifest_path) - manifest_keys = { - "schema_version", - "identity", - "completed_steps", - "learning_rates", - "step_scheduler", - "parent_checkpoint", - "rank_progress", - "dcp_sha256", - "sidecar_sha256", - } - if set(manifest) != manifest_keys or manifest["schema_version"] != _CHECKPOINT_SCHEMA_VERSION: - raise RuntimeError("PDD checkpoint manifest schema is incompatible.") - if expected_identity is not None and manifest["identity"] != expected_identity: - raise RuntimeError("PDD checkpoint identity does not match the current run.") - identity = manifest["identity"] - _require_qwen_image_execution(identity) - topology = identity.get("topology") if isinstance(identity, Mapping) else None - world_size = topology.get("world_size") if isinstance(topology, Mapping) else None - if type(world_size) is not int or world_size < 1: - raise RuntimeError("PDD checkpoint identity has an invalid world size.") - if expected_world_size is not None and world_size != expected_world_size: - raise RuntimeError( - f"PDD checkpoint world size {world_size} does not match {expected_world_size}." - ) - if _read_json(checkpoint / "pdd_config.json") != identity: - raise RuntimeError("PDD checkpoint config sidecar does not match the manifest.") - trainer_state = _read_json(checkpoint / "trainer_state.json") - if trainer_state != { - "completed_steps": manifest["completed_steps"], - "learning_rates": manifest["learning_rates"], - "step_scheduler": manifest["step_scheduler"], - "parent_checkpoint": manifest["parent_checkpoint"], - }: - raise RuntimeError("PDD trainer-state sidecar does not match the manifest.") - step_scheduler_state = manifest["step_scheduler"] - if ( - not isinstance(step_scheduler_state, dict) - or set(step_scheduler_state) != {"step", "epoch"} - or step_scheduler_state.get("step") != manifest["completed_steps"] - or type(step_scheduler_state.get("epoch")) is not int - or step_scheduler_state["epoch"] < 0 - ): - raise RuntimeError("PDD checkpoint StepScheduler state is invalid.") - rank_progress = manifest["rank_progress"] - if not isinstance(rank_progress, list) or len(rank_progress) != world_size: - raise RuntimeError("PDD checkpoint rank progress does not match its topology.") - expected_dcp = manifest["dcp_sha256"] - if not isinstance(expected_dcp, dict) or any( - not isinstance(path, str) or not isinstance(digest, str) - for path, digest in expected_dcp.items() - ): - raise RuntimeError("PDD checkpoint DCP hash inventory is malformed.") - if _dcp_payload_hashes(checkpoint) != expected_dcp: - raise RuntimeError("PDD checkpoint DCP payload inventory or hash does not match.") - expected_sidecars = _checkpoint_sidecar_paths(checkpoint, world_size) - expected_relative = {path.relative_to(checkpoint).as_posix() for path in expected_sidecars} - if ( - not isinstance(manifest["sidecar_sha256"], dict) - or set(manifest["sidecar_sha256"]) != expected_relative - ): - raise RuntimeError("PDD checkpoint sidecar inventory does not match the topology.") - for path in expected_sidecars: - relative = path.relative_to(checkpoint).as_posix() - if not path.is_file() or path.is_symlink(): - raise RuntimeError(f"PDD checkpoint sidecar is missing: {relative}.") - if _sha256(path) != manifest["sidecar_sha256"][relative]: - raise RuntimeError(f"PDD checkpoint sidecar hash mismatch: {relative}.") - return manifest - - -def resolve_pdd_training_checkpoint( - root: str | Path, - restore_from: str | Path, - *, - expected_world_size: int, - expected_identity: Mapping[str, Any] | None = None, -) -> tuple[Path, dict[str, Any]]: - """Resolve an explicit checkpoint or the newest compatible complete LATEST candidate.""" - unresolved_root = Path(root) - if unresolved_root.is_symlink(): - raise ValueError("PDD checkpoint_dir cannot be a symlink.") - root = unresolved_root.resolve() - if str(restore_from).upper() != "LATEST": - candidate = Path(restore_from) - if not candidate.is_absolute(): - candidate = root / candidate - if candidate.is_symlink(): - raise ValueError("explicit PDD checkpoint cannot be a symlink.") - candidate = candidate.resolve() - try: - candidate.relative_to(root) - except ValueError as error: - raise ValueError("explicit PDD checkpoint must be beneath checkpoint_dir.") from error - manifest = validate_pdd_training_checkpoint( - candidate, - expected_world_size=expected_world_size, - ) - if expected_identity is not None and not _identity_contains( - manifest.get("identity"), expected_identity - ): - raise RuntimeError("explicit PDD checkpoint identity does not match the selector.") - return candidate, manifest - - if not isinstance(expected_identity, Mapping) or not expected_identity: - raise ValueError("LATEST resolution requires a non-empty expected_identity selector.") - - pointed: tuple[int, Path, dict[str, Any]] | None = None - pointer = root / "LATEST" - if pointer.is_file() and not pointer.is_symlink(): - candidate = (root / pointer.read_text().strip()).resolve() - try: - candidate.relative_to(root) - manifest = validate_pdd_training_checkpoint( - candidate, - expected_world_size=expected_world_size, - ) - completed = manifest["completed_steps"] - if type(completed) is not int or completed < 0: - raise RuntimeError("PDD checkpoint completed_steps is invalid.") - if not _identity_contains(manifest.get("identity"), expected_identity): - raise RuntimeError("pointed PDD checkpoint identity does not match the selector.") - pointed = (completed, candidate, manifest) - except (ValueError, RuntimeError): - pass - - candidates: list[tuple[int, Path, dict[str, Any]]] = [] - if root.is_dir(): - for path in root.iterdir(): - suffix = path.name.removeprefix("step_") - if not path.is_dir() or not path.name.startswith("step_") or not suffix.isdigit(): - continue - if pointed is not None and int(suffix) <= pointed[0]: - continue - try: - manifest = validate_pdd_training_checkpoint( - path, - expected_world_size=expected_world_size, - ) - except RuntimeError: - continue - if not _identity_contains(manifest.get("identity"), expected_identity): - continue - completed = manifest["completed_steps"] - if type(completed) is int and completed >= 0: - candidates.append((completed, path.resolve(), manifest)) - if pointed is not None: - candidates.append(pointed) - if not candidates: - raise FileNotFoundError(f"no complete compatible PDD checkpoint exists beneath {root}.") - return max(candidates, key=lambda item: (item[0], item[1].name))[1:] - - -def _identity_contains(actual: Any, expected: Mapping[str, Any]) -> bool: - if not isinstance(actual, Mapping): - return False - return all(key in actual and actual[key] == value for key, value in expected.items()) - - -class PDDCheckpointManager: - """Publish and restore complete, metadata-compatible PDD checkpoints only.""" - - def __init__( - self, - *, - root: str | Path, - checkpointer: Any, - model: Any, - optimizer: Any, - scheduler: Any, - step_scheduler: Any, - trainer: Any, - sampler: Any, - rng: Any, - identity: Mapping[str, Any], - ) -> None: - self.root = Path(root).resolve() - self.checkpointer = checkpointer - self.model = model - self.optimizer = optimizer - self.scheduler = scheduler - self.step_scheduler = step_scheduler - self._step_scheduler_checkpoint_state = _StepSchedulerCheckpointState() - self.trainer = trainer - self.sampler = sampler - self.rng = rng - self._last_checkpoint: Path | None = None - self.identity = json.loads(json.dumps(identity, sort_keys=True)) - if self.identity.get("schema_version") != _CHECKPOINT_SCHEMA_VERSION: - raise ValueError("PDD checkpoint identity has an unsupported schema version.") - _require_qwen_image_execution(self.identity) - topology = self.identity.get("topology") - if not isinstance(topology, dict) or topology.get("world_size") != _world_size(): - raise ValueError("PDD checkpoint identity world size does not match the process group.") - if bool(getattr(checkpointer.config, "is_async", False)): - raise ValueError("PDD v1 atomic publication requires synchronous checkpoint saves.") - - def _rank_summary(self) -> dict[str, Any]: - sampler_state = self.sampler.state_dict() - return { - "rank": _rank(), - "epoch": sampler_state["epoch"], - "committed_batches": sampler_state["committed_batches"], - "sample_slots_consumed": sampler_state["sample_slots_consumed"], - "plan_sha256": sampler_state["plan_sha256"], - "next_sample_ids": list(sampler_state["next_sample_ids"]), - } - - def _sidecar_paths(self, checkpoint: Path) -> list[Path]: - return _checkpoint_sidecar_paths(checkpoint, _world_size()) - - def _manifest(self, checkpoint: Path) -> dict[str, Any]: - manifest = _read_json(checkpoint / "manifest.json") - expected = { - "schema_version", - "identity", - "completed_steps", - "learning_rates", - "step_scheduler", - "parent_checkpoint", - "rank_progress", - "dcp_sha256", - "sidecar_sha256", - } - if set(manifest) != expected: - raise RuntimeError("PDD checkpoint manifest has incompatible keys.") - if manifest["schema_version"] != _CHECKPOINT_SCHEMA_VERSION: - raise RuntimeError("PDD checkpoint manifest schema is unsupported.") - return manifest - - def _validate_checkpoint(self, checkpoint: Path, *, require_identity: bool) -> dict[str, Any]: - return validate_pdd_training_checkpoint( - checkpoint, - expected_identity=self.identity if require_identity else None, - expected_world_size=_world_size(), - ) - - def _compatible_candidates( - self, - *, - after_completed_steps: int | None = None, - ) -> list[tuple[int, Path]]: - candidates: list[tuple[int, Path]] = [] - if not self.root.is_dir(): - return candidates - for path in self.root.iterdir(): - if not path.is_dir() or path.name.startswith("."): - continue - if after_completed_steps is not None: - prefix = "step_" - suffix = path.name.removeprefix(prefix) - if not path.name.startswith(prefix) or not suffix.isdigit(): - continue - if int(suffix) <= after_completed_steps: - continue - try: - manifest = self._validate_checkpoint(path, require_identity=True) - except RuntimeError: - continue - completed = manifest["completed_steps"] - if type(completed) is int and completed >= 0: - candidates.append((completed, path.resolve())) - return sorted(candidates, key=lambda item: (item[0], item[1].name), reverse=True) - - def resolve(self, restore_from: str | Path | None) -> Path | None: - """Resolve LATEST by scanning only complete, identity-compatible checkpoints.""" - if restore_from is None: - return None - if str(restore_from).upper() == "LATEST": - pointer = self.root / "LATEST" - pointed: tuple[int, Path] | None = None - if pointer.is_file() and not pointer.is_symlink(): - name = pointer.read_text().strip() - candidate = (self.root / name).resolve() - try: - candidate.relative_to(self.root) - manifest = self._validate_checkpoint(candidate, require_identity=True) - completed = manifest["completed_steps"] - if type(completed) is not int or completed < 0: - raise RuntimeError("PDD checkpoint completed_steps is invalid.") - pointed = (completed, candidate) - except (ValueError, RuntimeError): - pass - candidates = self._compatible_candidates( - after_completed_steps=None if pointed is None else pointed[0] - ) - if pointed is not None: - candidates.append(pointed) - candidates.sort(key=lambda item: (item[0], item[1].name), reverse=True) - return candidates[0][1] if candidates else None - - candidate = Path(restore_from) - if not candidate.is_absolute(): - candidate = self.root / candidate - candidate = candidate.resolve() - try: - candidate.relative_to(self.root) - except ValueError as error: - raise ValueError("explicit PDD checkpoint must be beneath checkpoint_dir.") from error - self._validate_checkpoint(candidate, require_identity=True) - return candidate - - def _collective_resolve(self, restore_from: str | Path | None) -> Path | None: - if _world_size() == 1: - return self.resolve(restore_from) - status: dict[str, Any] | None = None - if _rank() == 0: - try: - resolved = self.resolve(restore_from) - status = { - "ok": True, - "path": None if resolved is None else str(resolved), - } - except BaseException as error: - status = { - "ok": False, - "error": f"{type(error).__name__}: {error}", - } - status = _broadcast_rank0_payload(status) - if not isinstance(status, dict) or type(status.get("ok")) is not bool: - raise RuntimeError("rank 0 broadcast a malformed checkpoint resolution status.") - if not status["ok"]: - raise RuntimeError(f"rank-0 checkpoint resolution failed: {status.get('error')}.") - resolved_path = status.get("path") - if resolved_path is None: - return None - if not isinstance(resolved_path, str): - raise RuntimeError("rank 0 broadcast a malformed checkpoint path.") - return Path(resolved_path) - - def _prepare_staging(self, final: Path) -> str: - self.root.mkdir(parents=True, exist_ok=True) - if final.exists(): - try: - self._validate_checkpoint(final, require_identity=False) - except RuntimeError: - shutil.rmtree(final) - else: - raise FileExistsError(f"complete PDD checkpoint already exists: {final}.") - staging_name = f".{final.name}.{uuid.uuid4().hex}.staging" - (self.root / staging_name).mkdir() - return staging_name - - def _publish_staging( - self, - *, - staging: Path, - final: Path, - completed_steps: int, - learning_rates: list[float], - step_scheduler_state: Mapping[str, int], - parent: Path | None, - rank_summaries: list[dict[str, Any]], - ) -> None: - sidecars = self._sidecar_paths(staging) - sidecar_sha256 = {path.relative_to(staging).as_posix(): _sha256(path) for path in sidecars} - manifest = { - "schema_version": _CHECKPOINT_SCHEMA_VERSION, - "identity": self.identity, - "completed_steps": completed_steps, - "learning_rates": learning_rates, - "step_scheduler": dict(step_scheduler_state), - "parent_checkpoint": None if parent is None else parent.name, - "rank_progress": sorted(rank_summaries, key=lambda summary: summary["rank"]), - "dcp_sha256": _dcp_payload_hashes(staging), - "sidecar_sha256": sidecar_sha256, - } - _atomic_json(staging / "pdd_config.json", self.identity) - _atomic_json( - staging / "trainer_state.json", - { - "completed_steps": completed_steps, - "learning_rates": learning_rates, - "step_scheduler": dict(step_scheduler_state), - "parent_checkpoint": manifest["parent_checkpoint"], - }, - ) - _atomic_json(staging / "manifest.json", manifest) - _fsync_tree(staging) - staging.rename(final) - _fsync_directory(self.root) - _atomic_json( - final / "COMPLETE", - { - "schema_version": _COMPLETE_SCHEMA_VERSION, - "manifest_sha256": _sha256(final / "manifest.json"), - }, - ) - self._validate_checkpoint(final, require_identity=True) - _atomic_text(self.root / "LATEST", final.name + "\n") - - def save(self) -> Path: - """Save into staging, publish atomically, mark complete, then update LATEST.""" - completed_steps = self.trainer.completed_steps - if type(completed_steps) is not int or completed_steps <= 0: - raise ValueError("PDD checkpoint requires at least one completed optimizer step.") - rank_summaries = _gather_objects(self._rank_summary()) - if len({summary["sample_slots_consumed"] for summary in rank_summaries}) != 1: - raise RuntimeError("PDD ranks disagree on consumed sample slots.") - learning_rates = [float(group["lr"]) for group in self.optimizer.param_groups] - live_step_scheduler_state = self.step_scheduler.state_dict() - if live_step_scheduler_state.get("step") != completed_steps: - raise RuntimeError("PDD StepScheduler state does not match the completed update.") - sampler_epoch = self.sampler.state_dict()["epoch"] - live_epoch = live_step_scheduler_state.get("epoch") - if sampler_epoch not in {live_epoch, live_epoch + 1}: - raise RuntimeError("PDD sampler epoch is incompatible with the StepScheduler epoch.") - step_scheduler_state = {"step": completed_steps, "epoch": sampler_epoch} - self._step_scheduler_checkpoint_state.load_state_dict(step_scheduler_state) - rank_scheduler_states = _gather_objects(step_scheduler_state) - if any(state != step_scheduler_state for state in rank_scheduler_states): - raise RuntimeError("PDD ranks disagree on StepScheduler checkpoint state.") - final = self.root / f"step_{completed_steps:08d}" - parent = self._last_checkpoint - - prepare_status = None - if _rank() == 0: - try: - prepare_status = {"ok": True, "staging": self._prepare_staging(final)} - except BaseException as error: - prepare_status = { - "ok": False, - "error": f"{type(error).__name__}: {error}", - } - prepare_status = _broadcast_rank0_payload(prepare_status) - if not isinstance(prepare_status, dict) or type(prepare_status.get("ok")) is not bool: - raise RuntimeError("rank 0 broadcast a malformed checkpoint preparation status.") - if not prepare_status["ok"]: - raise RuntimeError( - f"rank-0 checkpoint preparation failed: {prepare_status.get('error')}." - ) - staging = self.root / prepare_status["staging"] - - self.checkpointer.save_model(self.model, str(staging)) - self.checkpointer.save_optimizer( - self.optimizer, - self.model, - str(staging), - self.scheduler, - ) - sidecar_error = None - try: - self.checkpointer.save_on_dp_ranks(self.rng, "rng", str(staging)) - self.checkpointer.save_on_dp_ranks(self.sampler, "sampler", str(staging)) - self.checkpointer.save_on_dp_ranks( - self._step_scheduler_checkpoint_state, - "step_scheduler", - str(staging), - ) - self.checkpointer.save_on_dp_ranks(self.trainer, "trainer", str(staging)) - except BaseException as error: - sidecar_error = f"{type(error).__name__}: {error}" - sidecar_errors = _gather_objects(sidecar_error) - sidecar_failures = [ - f"rank {rank}: {message}" - for rank, message in enumerate(sidecar_errors) - if message is not None - ] - if sidecar_failures: - raise RuntimeError("PDD checkpoint sidecar save failed; " + "; ".join(sidecar_failures)) - _barrier() - - publish_status: dict[str, Any] | None = None - if _rank() == 0: - try: - self._publish_staging( - staging=staging, - final=final, - completed_steps=completed_steps, - learning_rates=learning_rates, - step_scheduler_state=step_scheduler_state, - parent=parent, - rank_summaries=rank_summaries, - ) - publish_status = {"ok": True} - except BaseException as error: - publish_status = { - "ok": False, - "error": f"{type(error).__name__}: {error}", - } - publish_status = _broadcast_rank0_payload(publish_status) - if not isinstance(publish_status, dict) or type(publish_status.get("ok")) is not bool: - raise RuntimeError("rank 0 broadcast a malformed checkpoint publication status.") - if not publish_status["ok"]: - raise RuntimeError( - f"rank-0 checkpoint publication failed: {publish_status.get('error')}." - ) - self._last_checkpoint = final - return final - - def load(self, restore_from: str | Path | None) -> PDDResumeState | None: - """Strictly restore model, optimizer/scheduler, cursor/trainer, and RNG last.""" - checkpoint = self._collective_resolve(restore_from) - if checkpoint is None: - return None - manifest = self._manifest(checkpoint) - self.checkpointer.load_model(self.model, str(checkpoint / "model")) - self.checkpointer.load_optimizer( - self.optimizer, - self.model, - str(checkpoint), - self.scheduler, - ) - self.checkpointer.load_on_dp_ranks(self.trainer, "trainer", str(checkpoint)) - self.checkpointer.load_on_dp_ranks(self.sampler, "sampler", str(checkpoint)) - self.checkpointer.load_on_dp_ranks( - self._step_scheduler_checkpoint_state, - "step_scheduler", - str(checkpoint), - ) - rank_progress = manifest["rank_progress"] - if not isinstance(rank_progress, list) or len(rank_progress) != _world_size(): - raise RuntimeError("PDD checkpoint rank progress does not match world size.") - progress = rank_progress[_rank()] - if progress.get("rank") != _rank(): - raise RuntimeError("PDD checkpoint rank progress is not ordered by rank.") - sampler_state = self.sampler.state_dict() - for key in ( - "epoch", - "committed_batches", - "sample_slots_consumed", - "plan_sha256", - "next_sample_ids", - ): - if sampler_state[key] != progress[key]: - raise RuntimeError(f"PDD restored sampler {key} does not match the manifest.") - if self.trainer.completed_steps != manifest["completed_steps"]: - raise RuntimeError("PDD restored trainer step does not match the manifest.") - step_scheduler_state = self._step_scheduler_checkpoint_state.state_dict() - if step_scheduler_state != manifest["step_scheduler"]: - raise RuntimeError("PDD restored StepScheduler state does not match the manifest.") - if step_scheduler_state["step"] != self.trainer.completed_steps: - raise RuntimeError("PDD restored StepScheduler step does not match the trainer.") - if step_scheduler_state["epoch"] != sampler_state["epoch"]: - raise RuntimeError("PDD restored StepScheduler epoch does not match the sampler.") - self.step_scheduler.load_state_dict(step_scheduler_state) - current_lrs = [float(group["lr"]) for group in self.optimizer.param_groups] - if current_lrs != manifest["learning_rates"]: - raise RuntimeError("PDD restored learning rate does not match the manifest.") - self.checkpointer.load_on_dp_ranks(self.rng, "rng", str(checkpoint)) - self._last_checkpoint = checkpoint - return PDDResumeState( - checkpoint_path=checkpoint, - completed_steps=manifest["completed_steps"], - sample_slots_consumed=progress["sample_slots_consumed"], - expected_next_sample_ids=tuple(progress["next_sample_ids"]), - parent_checkpoint=manifest["parent_checkpoint"], - ) diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index bc5287b88d7..15427d6eb1d 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -3,15 +3,15 @@ seed: 42 model: - pretrained_model_name_or_path: Qwen/Qwen-Image - revision: 75e0b4be04f60ec59a75f475837eced720f823b6 - torch_dtype: bfloat16 - device: cuda + # 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 - peft: - guidance_embeds: false - # Accepted by the composition path, but disabled until the pinned Diffusers Qwen API - # performs an effective fusion rather than its current no-op. fuse_qkv_projections: false pdd: @@ -31,6 +31,7 @@ pdd: optim: learning_rate: 5.0e-5 + clip_grad: 1.0 optimizer: _target_: torch.optim.AdamW weight_decay: 0.01 @@ -43,29 +44,18 @@ optim: fused: false maximize: false -lr_scheduler: - lr_decay_style: constant - lr_warmup_steps: 0 - min_lr: 5.0e-5 +# PDD supplies the loss; AutoModel owns the ordinary diffusion training lifecycle. +flow_matching: + adapter_type: qwen_image step_scheduler: - max_steps: 10000 + max_steps: 50000 num_epochs: 200 log_every: 10 ckpt_every_steps: 1000 local_batch_size: 4 save_checkpoint_every_epoch: false - global_batch_size: - -training_health: - max_grad_norm: 1.0 - zero_grad_warmup_steps: 0 - -validation: - count: 2000 - seed: 2026 - split_seed: 2026 - every_steps: 1000 + global_batch_size: 256 fsdp: dp_size: @@ -73,7 +63,8 @@ fsdp: cp_size: 1 pp_size: 1 ep_size: 1 - activation_checkpointing: true + # The PDD recipe enables Qwen's native block checkpointing after binding the MR210 forward. + activation_checkpointing: false data: dataloader: @@ -84,13 +75,17 @@ data: drop_last: true shuffle: true dynamic_batch_size: false - # Cursor-exact resume does not require hashes. Set true to authenticate every payload. - verify_payload_hashes: false + split: train + validation_count: 2000 + split_seed: 2026 + sampler_seed: 42 + loader_seed: 42 negative_prompt_embedding_path: negative_prompt_embedding.pt checkpoint: enabled: true checkpoint_dir: checkpoints/pdd_qwen_image - model_save_format: torch_save - save_consolidated: false + model_save_format: safetensors + save_consolidated: final + diffusers_compatible: true restore_from: LATEST diff --git a/examples/diffusers/fastgen/pdd/data.py b/examples/diffusers/fastgen/pdd/data.py deleted file mode 100644 index fce8aaa9e30..00000000000 --- a/examples/diffusers/fastgen/pdd/data.py +++ /dev/null @@ -1,401 +0,0 @@ -# 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. - -"""Replayable data and collective batch handling for the Qwen-Image PDD recipe.""" - -from __future__ import annotations - -import dataclasses -import hashlib -import json -from collections.abc import Mapping -from typing import Any - - -def _ordered_id_sha256(sample_ids: tuple[str, ...]) -> str: - digest = hashlib.sha256(b"modelopt-pdd-ordered-sample-ids-v1\0") - for sample_id in sample_ids: - digest.update(sample_id.encode()) - digest.update(b"\n") - return digest.hexdigest() - - -def _dataloader_options(raw: Mapping[str, Any]) -> dict[str, Any]: - data = raw.get("data") - if not isinstance(data, Mapping) or not isinstance(data.get("dataloader"), Mapping): - raise TypeError("PDD config requires a data.dataloader mapping.") - options = dict(data["dataloader"]) - target = options.pop("_target_", None) - expected_target = "fastgen_data.build_text_to_image_multiresolution_dataloader" - if target != expected_target: - raise ValueError(f"PDD data.dataloader._target_ must be {expected_target!r}.") - if "base_resolution" in options: - options["base_resolution"] = tuple(options["base_resolution"]) - options.setdefault("verify_payload_hashes", False) - return options - - -def _build_training_dataloader( - raw: Mapping[str, Any], - config: Any, - *, - dp_rank: int, - dp_world_size: int, -) -> tuple[Any, Any]: - from fastgen_data import ReplayableBatchSampler - from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader - - options = _dataloader_options(raw) - if options.get("drop_last", True) is not True: - raise ValueError("PDD exact sample accounting requires data.dataloader.drop_last=true.") - if options.get("dynamic_batch_size", False) is not False: - raise ValueError("PDD v1 requires data.dataloader.dynamic_batch_size=false.") - options.update( - split="train", - validation_count=config.validation.count, - split_seed=config.validation.split_seed, - dp_rank=dp_rank, - dp_world_size=dp_world_size, - exact_resume=True, - sampler_seed=config.seed, - loader_seed=config.seed, - ) - dataloader, sampler = build_text_to_image_multiresolution_dataloader(**options) - if not isinstance(sampler, ReplayableBatchSampler): - raise RuntimeError("PDD training requires the committed replayable batch sampler.") - if options.get("batch_size", 1) != config.step_scheduler.local_batch_size: - raise RuntimeError("resolved local batch size does not match the built dataloader.") - return dataloader, sampler - - -def _build_validation_dataloader( - raw: Mapping[str, Any], - config: Any, - *, - dp_rank: int, - dp_world_size: int, -) -> tuple[Any, Any]: - from fastgen_data.collate_fns import build_text_to_image_multiresolution_dataloader - - options = _dataloader_options(raw) - options.update( - split="validation", - validation_count=config.validation.count, - split_seed=config.validation.split_seed, - dp_rank=dp_rank, - dp_world_size=dp_world_size, - drop_last=False, - shuffle=False, - dynamic_batch_size=False, - exact_resume=False, - sampler_seed=config.validation.seed, - loader_seed=config.validation.seed, - ) - return build_text_to_image_multiresolution_dataloader(**options) - - -def _validate_dataset_contract( - train_dataset: Any, - validation_dataset: Any, - config: Any, -) -> tuple[Mapping[str, Any], str, str]: - """Collectively verify deterministic splits and the authenticated dataset snapshot.""" - import torch.distributed as dist - - try: - train_ids = tuple(str(value) for value in train_dataset.sample_ids) - validation_ids = tuple(str(value) for value in validation_dataset.sample_ids) - if len(validation_ids) != config.validation.count: - raise RuntimeError( - f"validation split has {len(validation_ids)} samples; " - f"expected {config.validation.count}." - ) - if set(train_ids).intersection(validation_ids): - raise RuntimeError("training and validation splits overlap.") - expected = {str(index) for index in range(train_dataset.total_num_samples)} - if set(train_ids).union(validation_ids) != expected: - raise RuntimeError("training and validation splits do not cover metadata.json.") - if train_dataset.total_num_samples != validation_dataset.total_num_samples: - raise RuntimeError("training and validation datasets disagree on total sample count.") - if train_dataset.metadata_sha256 != validation_dataset.metadata_sha256: - raise RuntimeError("training and validation datasets disagree on metadata content.") - if train_dataset.verify_payload_hashes != validation_dataset.verify_payload_hashes: - raise RuntimeError("training and validation datasets use different payload policies.") - if train_dataset.payload_hashes_complete != validation_dataset.payload_hashes_complete: - raise RuntimeError("training and validation datasets disagree on payload hash coverage.") - if train_dataset.dataset_snapshot_sha256 != validation_dataset.dataset_snapshot_sha256: - raise RuntimeError("training and validation datasets disagree on dataset content.") - if not isinstance(train_dataset.dataset_snapshot_sha256, str): - raise RuntimeError("PDD dataloader did not construct a dataset snapshot identity.") - report = { - "cache_root": str(train_dataset.cache_root), - "metadata_sha256": train_dataset.metadata_sha256, - "negative_prompt_embedding_sha256": (train_dataset.negative_prompt_embedding_sha256), - "dataset_snapshot_sha256": train_dataset.dataset_snapshot_sha256, - "total_samples": train_dataset.total_num_samples, - "train_samples": len(train_ids), - "validation_samples": len(validation_ids), - "split_seed": config.validation.split_seed, - "verify_payload_hashes": train_dataset.verify_payload_hashes, - "payload_hashes_complete": train_dataset.payload_hashes_complete, - } - local_status: dict[str, Any] = { - "ok": True, - "report": report, - "train_hash": _ordered_id_sha256(train_ids), - "validation_hash": _ordered_id_sha256(validation_ids), - } - except BaseException as error: - local_status = {"ok": False, "error": f"{type(error).__name__}: {error}"} - - statuses: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(statuses, local_status) - failures: list[str] = [] - successes: list[Mapping[str, Any]] = [] - for rank, status in enumerate(statuses): - if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: - failures.append(f"rank {rank}: malformed loader authentication status") - continue - if not status["ok"]: - failures.append(f"rank {rank}: {status.get('error')}") - continue - successes.append(status) - if failures: - raise RuntimeError("PDD dataset validation failed: " + "; ".join(failures)) - canonical = {json.dumps(status, sort_keys=True) for status in successes} - if len(canonical) != 1: - raise RuntimeError( - "PDD ranks resolved different dataset roots, metadata, or split membership." - ) - return report, local_status["train_hash"], local_status["validation_hash"] - - -def _build_validation_plan(sampler: Any, config: Any) -> tuple[Any, tuple[tuple[bool, ...], ...]]: - import torch.distributed as dist - - from pdd.training import build_pdd_validation_assignments - - heldout_ids = tuple(str(value) for value in sampler.dataset.sample_ids) - assignments = build_pdd_validation_assignments( - heldout_ids, - config.pdd, - validation_seed=config.validation.seed, - ) - sampler.set_epoch(0) - sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) - local_plan = [ - tuple(str(sampler.dataset.sample_ids[index]) for index in batch) for batch in sampler - ] - sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) - plans: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(plans, local_plan) - batch_counts = {len(plan) for plan in plans} - if len(batch_counts) != 1: - raise RuntimeError("PDD validation sampler produced different batch counts across ranks.") - - masks = [[([False] * len(batch)) for batch in plan] for plan in plans] - seen: set[str] = set() - for batch_index in range(len(local_plan)): - for rank, plan in enumerate(plans): - for position, sample_id in enumerate(plan[batch_index]): - if sample_id not in seen: - masks[rank][batch_index][position] = True - seen.add(sample_id) - if seen != set(heldout_ids): - missing = sorted(set(heldout_ids) - seen) - extra = sorted(seen - set(heldout_ids)) - raise RuntimeError( - f"PDD validation sampler does not cover the held-out split: " - f"missing={missing[:5]}, extra={extra[:5]}." - ) - local_masks = tuple(tuple(batch) for batch in masks[dist.get_rank()]) - return assignments, local_masks - - -def _iter_validation_batches( - dataloader: Any, - masks: tuple[tuple[bool, ...], ...], - config: Any, - expected_latent_channels: int, - expected_condition_features: int, -): - from pdd.training import prepare_qwen_pdd_batch - - count = 0 - for count, (raw_batch, valid_mask) in enumerate(zip(dataloader, masks, strict=True), start=1): - prepared = prepare_qwen_pdd_batch( - raw_batch, - device=config.device, - dtype=config.dtype, - require_negative_condition=config.pdd.guidance_scale is not None, - expected_latent_channels=expected_latent_channels, - expected_condition_features=expected_condition_features, - ) - yield dataclasses.replace(prepared, valid_mask=valid_mask) - if count != len(masks): - raise RuntimeError( - f"PDD validation loader produced {count} batches for a {len(masks)}-batch plan." - ) - - -def _coverage_axis(counts: Any, loss_sums: Any) -> dict[int, dict[str, float | int]]: - return { - index: {"count": int(count), "mean_loss": float(loss_sums[index] / count)} - for index, count in enumerate(counts.tolist()) - if count - } - - -def _collective_training_iterator(dataloader: Any, sampler: Any) -> Any: - """Advance epochs and construct rank-local iterators under a collective error gate.""" - import torch.distributed as dist - - iterator = None - error_message = None - try: - if sampler.remaining_batches == 0: - sampler.set_epoch(sampler.epoch + 1) - iterator = iter(dataloader) - if iterator is None: - raise RuntimeError("PDD dataloader returned no iterator.") - except BaseException as error: - error_message = f"{type(error).__name__}: {error}" - errors: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(errors, error_message) - failures = [f"rank {rank}: {message}" for rank, message in enumerate(errors) if message] - if failures: - raise RuntimeError( - "distributed PDD training iterator construction failed; " + "; ".join(failures) - ) - if iterator is None: - raise RuntimeError("local PDD iterator construction succeeded without an iterator.") - return iterator - - -def _collective_training_batch( - iterator: Any, - *, - sampler: Any, - resume: Any, - resume_pending: bool, - device: Any, - dtype: Any, - require_negative_condition: bool, - expected_batch_size: int, - expected_latent_channels: int, - expected_condition_features: int, -) -> tuple[Any, tuple[str, ...]] | None: - """Prepare one rank-local batch, then agree on success before any model call.""" - import torch.distributed as dist - - from pdd.training import prepare_qwen_pdd_batch - - prepared = None - sample_ids: tuple[str, ...] = () - status: dict[str, Any] - try: - raw_batch = next(iterator) - except StopIteration: - if resume_pending: - status = { - "state": "error", - "error": "RuntimeError: resumed dataloader ended before its first batch", - "resume_pending": True, - } - else: - status = {"state": "end", "resume_pending": False} - except BaseException as error: - status = { - "state": "error", - "error": f"{type(error).__name__}: {error}", - "resume_pending": resume_pending, - } - else: - try: - metadata = raw_batch["metadata"] - raw_ids = metadata.get("logical_sample_ids", metadata.get("sample_ids")) - if hasattr(raw_ids, "tolist"): - raw_ids = raw_ids.tolist() - sample_ids = tuple(str(value) for value in raw_ids) - expected_ids = sampler.expected_next_sample_ids() - if sample_ids != expected_ids: - raise RuntimeError( - "prefetched PDD batch does not match committed cursor: " - f"expected={expected_ids}, actual={sample_ids}." - ) - if resume_pending: - if resume is None: - raise RuntimeError("resume_pending is true without a PDD resume state.") - resume.verify_first_batch(sample_ids) - prepared = prepare_qwen_pdd_batch( - raw_batch, - device=device, - dtype=dtype, - require_negative_condition=require_negative_condition, - expected_latent_channels=expected_latent_channels, - expected_condition_features=expected_condition_features, - ) - if prepared is None: - raise RuntimeError("PDD batch preparation returned no prepared batch.") - if len(sample_ids) != expected_batch_size: - raise RuntimeError( - f"PDD training batch has {len(sample_ids)} samples; " - f"expected {expected_batch_size}." - ) - status = { - "state": "batch", - "batch_size": len(sample_ids), - "resume_pending": resume_pending, - } - except BaseException as error: - status = { - "state": "error", - "error": f"{type(error).__name__}: {error}", - "resume_pending": resume_pending, - } - - statuses: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(statuses, status) - malformed = [rank for rank, item in enumerate(statuses) if not isinstance(item, Mapping)] - if malformed: - raise RuntimeError(f"PDD training ranks returned malformed statuses: {malformed}.") - failures = [ - f"rank {rank}: {item.get('error')}" - for rank, item in enumerate(statuses) - if item.get("state") == "error" - ] - if failures: - raise RuntimeError( - "distributed PDD training batch preflight failed; " + "; ".join(failures) - ) - states = {item.get("state") for item in statuses} - if states == {"end"}: - return None - if states != {"batch"}: - raise RuntimeError( - "distributed PDD training ranks produced different dataloader lengths: " - f"{[item.get('state') for item in statuses]}." - ) - pending = {item.get("resume_pending") for item in statuses} - if len(pending) != 1: - raise RuntimeError("distributed PDD training ranks disagree on resume verification state.") - batch_sizes = {item.get("batch_size") for item in statuses} - if batch_sizes != {expected_batch_size}: - raise RuntimeError( - f"distributed PDD training ranks disagree on batch size: {sorted(batch_sizes)}." - ) - if prepared is None: - raise RuntimeError("local PDD batch preparation succeeded without a prepared batch.") - return prepared, sample_ids diff --git a/examples/diffusers/fastgen/pdd/export.py b/examples/diffusers/fastgen/pdd/export.py deleted file mode 100644 index c41c0c68648..00000000000 --- a/examples/diffusers/fastgen/pdd/export.py +++ /dev/null @@ -1,598 +0,0 @@ -# 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. - -"""Authenticated, bounded safetensors export and strict PDD reconstruction helpers.""" - -from __future__ import annotations - -import os -import re -import shutil -import uuid -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import torch -from safetensors import safe_open -from safetensors.torch import save_file - -from modelopt.torch.fastgen import PDDConfig, PDDMetadata -from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - QWEN_IMAGE_PDD_EXECUTION, - QWEN_IMAGE_PDD_LAYER_SPEC, -) - -from .artifacts import ( - load_canonical_json, - require_sha256, - resolve_relative_artifact, - sha256_file, - write_canonical_json, -) - -_EXPORT_SCHEMA_VERSION = 4 -_COMPLETE_SCHEMA_VERSION = 1 -_EXPORT_FORMAT = "modelopt-pdd-safetensors" -_CONFIG_FILE = "config.json" -_METADATA_FILE = "pdd_metadata.json" -_INDEX_FILE = "diffusion_pytorch_model.safetensors.index.json" -_MANIFEST_FILE = "manifest.json" -_COMPLETE_FILE = "COMPLETE" -_SHARD_PATTERN = "diffusion_pytorch_model-{index:05d}-of-{count:05d}.safetensors" - -PDD_INFERENCE_SCHEDULES: Mapping[str, tuple[int, ...]] = { - "pdd-2": (64, 64), - "pdd-4": (32, 32, 32, 32), - "pdd-8": (16, 16, 16, 16, 16, 16, 16, 16), -} - - -@dataclass(frozen=True) -class PDDExportDescriptor: - """Validated non-tensor export metadata.""" - - root: Path - manifest: Mapping[str, Any] - metadata: PDDMetadata - transformer_config: Mapping[str, Any] - - -def _fsync_directory(path: Path) -> None: - descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - - -def _fsync_tree(root: Path) -> None: - for path in sorted(root.rglob("*"), key=lambda item: len(item.parts), reverse=True): - if path.is_symlink(): - raise RuntimeError(f"PDD export staging contains a symlink: {path}.") - if path.is_file(): - with path.open("rb") as stream: - os.fsync(stream.fileno()) - elif path.is_dir(): - _fsync_directory(path) - _fsync_directory(root) - - -def _tensor_nbytes(tensor: torch.Tensor) -> int: - return tensor.numel() * tensor.element_size() - - -def _validate_state_dict(state_dict: Mapping[str, Any]) -> dict[str, torch.Tensor]: - if not isinstance(state_dict, Mapping) or not state_dict: - raise ValueError("PDD export state_dict must be a non-empty mapping.") - tensors: dict[str, torch.Tensor] = {} - for key in sorted(state_dict): - tensor = state_dict[key] - if not isinstance(key, str) or not key: - raise ValueError("PDD export tensor keys must be non-empty strings.") - if not isinstance(tensor, torch.Tensor): - raise TypeError(f"PDD export value {key!r} is not a tensor.") - if tensor.device.type != "cpu" or tensor.is_meta: - raise ValueError(f"PDD export tensor {key!r} must be a materialized CPU tensor.") - if tensor.layout != torch.strided or tensor.is_quantized: - raise TypeError(f"PDD export tensor {key!r} must be a dense, unquantized tensor.") - if tensor.dtype.is_floating_point and not torch.isfinite(tensor).all().item(): - raise FloatingPointError(f"PDD export tensor {key!r} is non-finite.") - tensors[key] = tensor.detach().contiguous() - return tensors - - -def _save_probe(staging: Path, keys: Sequence[str], tensors: Mapping[str, torch.Tensor]) -> int: - probe = staging / f".probe-{uuid.uuid4().hex}.safetensors" - payload = {key: tensors[key].clone() for key in keys} - try: - save_file(payload, str(probe), metadata={"format": "pt"}) - return probe.stat().st_size - finally: - probe.unlink(missing_ok=True) - - -def _bounded_shard_groups( - staging: Path, - tensors: Mapping[str, torch.Tensor], - max_shard_bytes: int, -) -> list[tuple[str, ...]]: - if type(max_shard_bytes) is not int or max_shard_bytes <= 0: - raise ValueError("max_shard_bytes must be a positive integer.") - initial: list[tuple[str, ...]] = [] - current: list[str] = [] - current_bytes = 0 - for key, tensor in tensors.items(): - size = _tensor_nbytes(tensor) - if size >= max_shard_bytes: - raise ValueError( - f"tensor {key!r} has {size} bytes and cannot fit beneath the physical " - f"shard bound {max_shard_bytes}." - ) - if current and current_bytes + size >= max_shard_bytes: - initial.append(tuple(current)) - current = [] - current_bytes = 0 - current.append(key) - current_bytes += size - if current: - initial.append(tuple(current)) - - bounded: list[tuple[str, ...]] = [] - pending = list(initial) - while pending: - keys = pending.pop(0) - if _save_probe(staging, keys, tensors) <= max_shard_bytes: - bounded.append(keys) - continue - if len(keys) == 1: - raise ValueError( - f"tensor {keys[0]!r} plus safetensors metadata exceeds max_shard_bytes." - ) - midpoint = len(keys) // 2 - pending[0:0] = [keys[:midpoint], keys[midpoint:]] - return bounded - - -def _validate_identity(identity: Mapping[str, Any], metadata: PDDMetadata) -> dict[str, Any]: - if not isinstance(identity, Mapping): - raise TypeError("PDD export identity must be a mapping.") - required = { - "guidance", - "model", - "pdd_metadata", - "qwen_image", - "topology", - } - missing = sorted(required.difference(identity)) - if missing: - raise ValueError(f"PDD export identity is missing keys: {missing}.") - if identity["pdd_metadata"] != metadata.to_dict(): - raise ValueError("PDD export metadata does not match the checkpoint identity.") - _require_exact_mapping(identity["qwen_image"], {"execution"}, name="identity.qwen_image") - if identity["qwen_image"]["execution"] != QWEN_IMAGE_PDD_EXECUTION: - raise ValueError("PDD export has an incompatible Qwen execution identity.") - model = _require_exact_mapping( - identity["model"], {"id", "revision", "dtype"}, name="identity.model" - ) - if not isinstance(model["id"], str) or not model["id"] or not isinstance(model["dtype"], str): - raise ValueError("PDD export checkpoint identity has an invalid model ID or dtype.") - revision = model["revision"] - if ( - not isinstance(revision, str) - or len(revision) != 40 - or any(character not in "0123456789abcdef" for character in revision) - ): - raise ValueError("PDD export model revision must be an exact lowercase commit.") - guidance = _require_exact_mapping(identity["guidance"], {"scale"}, name="identity.guidance") - for name, value in guidance.items(): - if value is not None and ( - isinstance(value, bool) - or not isinstance(value, int | float) - or not torch.isfinite(torch.tensor(float(value))).item() - ): - raise ValueError(f"PDD export guidance {name} must be finite or null.") - topology = identity["topology"] - if ( - not isinstance(topology, Mapping) - or type(topology.get("world_size")) is not int - or topology["world_size"] < 1 - or topology.get("pure_data_parallel") is not True - ): - raise ValueError("PDD export checkpoint identity has invalid pure-DP topology.") - return dict(identity) - - -def write_pdd_export( - output_dir: str | Path, - state_dict: Mapping[str, Any], - *, - metadata: PDDMetadata, - transformer_config: Mapping[str, Any], - identity: Mapping[str, Any], - source_checkpoint: Mapping[str, Any], - max_shard_bytes: int, -) -> Path: - """Publish a complete PDD export into a previously absent directory.""" - if not isinstance(metadata, PDDMetadata): - raise TypeError("metadata must be PDDMetadata.") - if metadata.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: - raise ValueError("PDD export supports only the fixed Qwen-Image layer specification.") - if not isinstance(transformer_config, Mapping): - raise TypeError("transformer_config must be a mapping.") - checkpoint_keys = {"name", "manifest_sha256", "completed_steps"} - if not isinstance(source_checkpoint, Mapping) or set(source_checkpoint) != checkpoint_keys: - raise ValueError(f"source_checkpoint must contain exactly {sorted(checkpoint_keys)}.") - if ( - not isinstance(source_checkpoint["name"], str) - or not source_checkpoint["name"] - or Path(source_checkpoint["name"]).name != source_checkpoint["name"] - ): - raise ValueError("source_checkpoint.name must be a basename.") - require_sha256(source_checkpoint["manifest_sha256"], name="source manifest SHA-256") - if ( - type(source_checkpoint["completed_steps"]) is not int - or source_checkpoint["completed_steps"] < 1 - ): - raise ValueError("source_checkpoint.completed_steps must be an integer >= 1.") - resolved_identity = _validate_identity(identity, metadata) - tensors = _validate_state_dict(state_dict) - - unresolved_output = Path(output_dir) - if unresolved_output.is_symlink(): - raise ValueError("PDD export output cannot be a symlink.") - output = unresolved_output.resolve() - output.parent.mkdir(parents=True, exist_ok=True) - if output.exists() or output.is_symlink(): - raise FileExistsError(f"PDD export output already exists: {output}.") - staging = output.with_name(f".{output.name}.{uuid.uuid4().hex}.staging") - staging.mkdir() - published = False - try: - groups = _bounded_shard_groups(staging, tensors, max_shard_bytes) - shard_names = [ - _SHARD_PATTERN.format(index=index, count=len(groups)) - for index in range(1, len(groups) + 1) - ] - weight_map: dict[str, str] = {} - for name, keys in zip(shard_names, groups): - payload = {key: tensors[key].clone() for key in keys} - save_file(payload, str(staging / name), metadata={"format": "pt"}) - if (staging / name).stat().st_size > max_shard_bytes: - raise RuntimeError(f"PDD safetensors shard exceeds its physical bound: {name}.") - weight_map.update(dict.fromkeys(keys, name)) - - total_tensor_bytes = sum(_tensor_nbytes(tensor) for tensor in tensors.values()) - write_canonical_json(staging / _CONFIG_FILE, dict(transformer_config)) - write_canonical_json(staging / _METADATA_FILE, metadata.to_dict()) - write_canonical_json( - staging / _INDEX_FILE, - {"metadata": {"total_size": total_tensor_bytes}, "weight_map": weight_map}, - ) - file_names = [_CONFIG_FILE, _METADATA_FILE, _INDEX_FILE, *shard_names] - files = { - name: { - "sha256": sha256_file(staging / name), - "size": (staging / name).stat().st_size, - } - for name in file_names - } - tensor_specs = { - key: { - "dtype": str(tensor.dtype).removeprefix("torch."), - "nbytes": _tensor_nbytes(tensor), - "shape": list(tensor.shape), - "shard": weight_map[key], - } - for key, tensor in tensors.items() - } - manifest = { - "schema_version": _EXPORT_SCHEMA_VERSION, - "format": _EXPORT_FORMAT, - "identity": resolved_identity, - "source_checkpoint": dict(source_checkpoint), - "max_shard_bytes": max_shard_bytes, - "total_tensor_bytes": total_tensor_bytes, - "tensors": tensor_specs, - "files": files, - } - write_canonical_json(staging / _MANIFEST_FILE, manifest) - write_canonical_json( - staging / _COMPLETE_FILE, - { - "schema_version": _COMPLETE_SCHEMA_VERSION, - "manifest_sha256": sha256_file(staging / _MANIFEST_FILE), - }, - ) - _fsync_tree(staging) - inspect_pdd_export(staging) - staging.rename(output) - published = True - _fsync_directory(output.parent) - return output - except BaseException: - target = output if published else staging - if target.exists() and not target.is_symlink(): - shutil.rmtree(target) - raise - - -def _require_exact_mapping(value: Any, keys: set[str], *, name: str) -> Mapping[str, Any]: - if not isinstance(value, Mapping) or set(value) != keys: - actual = sorted(value) if isinstance(value, Mapping) else type(value).__name__ - raise ValueError(f"{name} keys mismatch: expected={sorted(keys)}, actual={actual}.") - return value - - -def inspect_pdd_export(export_dir: str | Path) -> PDDExportDescriptor: - """Authenticate a PDD export without loading its tensor payloads.""" - unresolved_root = Path(export_dir) - if unresolved_root.is_symlink(): - raise RuntimeError(f"PDD export cannot be a symlink: {unresolved_root}.") - root = unresolved_root.resolve() - if not root.is_dir(): - raise RuntimeError(f"PDD export is not a regular directory: {root}.") - for path in root.rglob("*"): - if path.is_symlink(): - raise RuntimeError(f"PDD export contains a symlink: {path}.") - if path.is_dir(): - raise RuntimeError(f"PDD export must be flat, found directory: {path}.") - - complete = _require_exact_mapping( - load_canonical_json(root / _COMPLETE_FILE), - {"schema_version", "manifest_sha256"}, - name="PDD COMPLETE", - ) - if complete["schema_version"] != _COMPLETE_SCHEMA_VERSION: - raise ValueError("PDD COMPLETE schema version is unsupported.") - if require_sha256( - complete["manifest_sha256"], name="PDD COMPLETE manifest SHA-256" - ) != sha256_file(root / _MANIFEST_FILE): - raise RuntimeError("PDD COMPLETE does not match the export manifest.") - manifest = _require_exact_mapping( - load_canonical_json(root / _MANIFEST_FILE), - { - "schema_version", - "format", - "identity", - "source_checkpoint", - "max_shard_bytes", - "total_tensor_bytes", - "tensors", - "files", - }, - name="PDD export manifest", - ) - if manifest["schema_version"] != _EXPORT_SCHEMA_VERSION or manifest["format"] != _EXPORT_FORMAT: - raise ValueError("PDD export manifest schema or format is unsupported.") - if type(manifest["max_shard_bytes"]) is not int or manifest["max_shard_bytes"] <= 0: - raise ValueError("PDD export max_shard_bytes is invalid.") - if type(manifest["total_tensor_bytes"]) is not int or manifest["total_tensor_bytes"] <= 0: - raise ValueError("PDD export total_tensor_bytes is invalid.") - source = _require_exact_mapping( - manifest["source_checkpoint"], - {"name", "manifest_sha256", "completed_steps"}, - name="source_checkpoint", - ) - if ( - not isinstance(source["name"], str) - or not source["name"] - or Path(source["name"]).name != source["name"] - ): - raise ValueError("source_checkpoint.name must be a basename.") - require_sha256(source["manifest_sha256"], name="source checkpoint manifest SHA-256") - if type(source["completed_steps"]) is not int or source["completed_steps"] < 1: - raise ValueError("source_checkpoint.completed_steps is invalid.") - - files = manifest["files"] - if not isinstance(files, Mapping) or not files: - raise ValueError("PDD export file inventory must be a non-empty mapping.") - expected_names = set(files) | {_MANIFEST_FILE, _COMPLETE_FILE} - actual_names = {path.name for path in root.iterdir() if path.is_file()} - if actual_names != expected_names: - raise RuntimeError( - f"PDD export file inventory mismatch: expected={sorted(expected_names)}, " - f"actual={sorted(actual_names)}." - ) - for name, record in files.items(): - if Path(name).name != name: - raise ValueError(f"PDD export file name must be a basename: {name!r}.") - record = _require_exact_mapping(record, {"sha256", "size"}, name=f"files[{name!r}]") - path = resolve_relative_artifact(root, name) - if type(record["size"]) is not int or record["size"] < 0: - raise ValueError(f"PDD export file size is invalid for {name!r}.") - if path.stat().st_size != record["size"]: - raise RuntimeError(f"PDD export file size mismatch for {name!r}.") - if sha256_file(path) != require_sha256(record["sha256"], name=f"files[{name!r}].sha256"): - raise RuntimeError(f"PDD export file SHA-256 mismatch for {name!r}.") - if name.endswith(".safetensors") and record["size"] > manifest["max_shard_bytes"]: - raise RuntimeError(f"PDD export shard exceeds max_shard_bytes: {name!r}.") - mandatory = {_CONFIG_FILE, _METADATA_FILE, _INDEX_FILE} - if not mandatory.issubset(files): - raise RuntimeError( - f"PDD export is missing mandatory files: {sorted(mandatory - set(files))}." - ) - - metadata_data = load_canonical_json(root / _METADATA_FILE) - metadata = PDDMetadata.from_dict(metadata_data) - if metadata.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: - raise ValueError("PDD export carries a non-Qwen layer specification.") - identity = manifest["identity"] - _validate_identity(identity, metadata) - transformer_config = load_canonical_json(root / _CONFIG_FILE) - if not isinstance(transformer_config, Mapping): - raise ValueError("PDD transformer config must contain an object.") - - tensors = manifest["tensors"] - if not isinstance(tensors, Mapping) or not tensors: - raise ValueError("PDD export tensor inventory must be a non-empty mapping.") - index = _require_exact_mapping( - load_canonical_json(root / _INDEX_FILE), - {"metadata", "weight_map"}, - name="safetensors index", - ) - index_metadata = _require_exact_mapping( - index["metadata"], {"total_size"}, name="index metadata" - ) - if index_metadata["total_size"] != manifest["total_tensor_bytes"]: - raise RuntimeError("PDD safetensors index total size does not match the manifest.") - weight_map = index["weight_map"] - if not isinstance(weight_map, Mapping) or set(weight_map) != set(tensors): - raise RuntimeError("PDD safetensors index keys do not match the tensor inventory.") - shard_names = {name for name in files if name.endswith(".safetensors")} - shard_pattern = re.compile(r"diffusion_pytorch_model-(\d{5})-of-(\d{5})\.safetensors") - shard_numbers = [] - for name in shard_names: - match = shard_pattern.fullmatch(name) - if match is None: - raise ValueError(f"PDD export has a noncanonical shard name: {name!r}.") - shard_numbers.append((int(match.group(1)), int(match.group(2)))) - if not shard_numbers: - raise RuntimeError("PDD export has no safetensors shards.") - shard_count = len(shard_numbers) - if sorted(shard_numbers) != [(index, shard_count) for index in range(1, shard_count + 1)]: - raise RuntimeError("PDD export safetensors shard numbering is inconsistent.") - if set(weight_map.values()) != shard_names: - raise RuntimeError("PDD safetensors index shard inventory does not match export files.") - - total = 0 - for key, spec in tensors.items(): - if not isinstance(key, str) or not key: - raise ValueError("PDD tensor inventory keys must be non-empty strings.") - spec = _require_exact_mapping( - spec, - {"dtype", "nbytes", "shape", "shard"}, - name=f"tensors[{key!r}]", - ) - if ( - not isinstance(spec["dtype"], str) - or type(spec["nbytes"]) is not int - or spec["nbytes"] <= 0 - or not isinstance(spec["shape"], list) - or any(type(size) is not int or size < 0 for size in spec["shape"]) - or spec["shard"] != weight_map[key] - ): - raise ValueError(f"PDD tensor specification is malformed for {key!r}.") - total += spec["nbytes"] - if total != manifest["total_tensor_bytes"]: - raise RuntimeError("PDD tensor byte inventory does not match total_tensor_bytes.") - return PDDExportDescriptor(root, manifest, metadata, transformer_config) - - -def pdd_config_from_metadata( - metadata: PDDMetadata, - *, - blocks: Sequence[int] | None = None, - schedule: str | None = None, - guidance_scale: float | None = None, -) -> PDDConfig: - """Build a fresh validated inference config from authenticated metadata.""" - if (blocks is None) == (schedule is None): - raise ValueError("Specify exactly one of blocks or schedule.") - if schedule is not None: - try: - resolved = PDD_INFERENCE_SCHEDULES[schedule] - except KeyError as error: - raise ValueError( - f"Unknown PDD schedule {schedule!r}; expected {sorted(PDD_INFERENCE_SCHEDULES)}." - ) from error - else: - if isinstance(blocks, str | bytes) or not isinstance(blocks, Sequence): - raise TypeError("blocks must be a sequence of integers.") - resolved = tuple(blocks) - return PDDConfig( - grid_size=metadata.grid_size, - grid_max_t=metadata.grid_max_t, - flow_shift=metadata.flow_shift, - block_size_min=metadata.block_size_min, - block_size_max=metadata.block_size_max, - inference_blocks=list(resolved), - student_sample_steps=len(resolved), - teacher_integrator=metadata.teacher_integrator, - guidance_scale=guidance_scale, - num_train_timesteps=None, - ) - - -def _load_shard( - descriptor: PDDExportDescriptor, - shard_name: str, -) -> dict[str, torch.Tensor]: - expected = { - key: spec - for key, spec in descriptor.manifest["tensors"].items() - if spec["shard"] == shard_name - } - path = descriptor.root / shard_name - loaded: dict[str, torch.Tensor] = {} - with safe_open(str(path), framework="pt", device="cpu") as stream: - keys = list(stream.keys()) - if len(keys) != len(set(keys)) or set(keys) != set(expected): - raise RuntimeError(f"PDD safetensors keys do not match the index for {shard_name!r}.") - for key in keys: - tensor = stream.get_tensor(key) - spec = expected[key] - if list(tensor.shape) != spec["shape"]: - raise RuntimeError(f"PDD tensor shape mismatch for {key!r}.") - if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: - raise RuntimeError(f"PDD tensor dtype mismatch for {key!r}.") - if _tensor_nbytes(tensor) != spec["nbytes"]: - raise RuntimeError(f"PDD tensor byte-size mismatch for {key!r}.") - if tensor.dtype.is_floating_point and not torch.isfinite(tensor).all().item(): - raise FloatingPointError(f"PDD tensor {key!r} is non-finite.") - loaded[key] = tensor - return loaded - - -def load_pdd_export_into_model( - export_dir: str | Path, - model: torch.nn.Module, -) -> PDDExportDescriptor: - """Strictly stream authenticated safetensors shards into a converted CPU model.""" - if not isinstance(model, torch.nn.Module): - raise TypeError("model must be an nn.Module.") - descriptor = inspect_pdd_export(export_dir) - expected_state = model.state_dict() - specs = descriptor.manifest["tensors"] - if set(expected_state) != set(specs): - missing = sorted(set(expected_state) - set(specs)) - extra = sorted(set(specs) - set(expected_state)) - raise RuntimeError(f"PDD model/export keys mismatch: missing={missing}, extra={extra}.") - for key, expected in expected_state.items(): - spec = specs[key] - if list(expected.shape) != spec["shape"]: - raise RuntimeError(f"PDD model/export shape mismatch for {key!r}.") - if str(expected.dtype).removeprefix("torch.") != spec["dtype"]: - raise RuntimeError(f"PDD model/export dtype mismatch for {key!r}.") - - shard_names = sorted({spec["shard"] for spec in specs.values()}) - loaded_keys: set[str] = set() - for shard_name in shard_names: - shard = _load_shard(descriptor, shard_name) - if loaded_keys.intersection(shard): - raise RuntimeError("PDD tensor appears in more than one safetensors shard.") - incompatible = model.load_state_dict(shard, strict=False) - if incompatible.unexpected_keys: - raise RuntimeError(f"PDD shard has unexpected keys: {incompatible.unexpected_keys}.") - loaded_keys.update(shard) - if loaded_keys != set(expected_state): - raise RuntimeError("PDD safe load did not assign every expected tensor.") - if any(parameter.is_meta for parameter in model.parameters()) or any( - buffer.is_meta for buffer in model.buffers() - ): - raise RuntimeError("PDD safe load left meta tensors in the reconstructed model.") - model.eval().requires_grad_(False) - return descriptor diff --git a/examples/diffusers/fastgen/pdd/export_qwen_image.py b/examples/diffusers/fastgen/pdd/export_qwen_image.py deleted file mode 100644 index e77bc0331e8..00000000000 --- a/examples/diffusers/fastgen/pdd/export_qwen_image.py +++ /dev/null @@ -1,349 +0,0 @@ -# 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. - -"""Collectively restore a PDD checkpoint and publish a safe Qwen-Image export.""" - -from __future__ import annotations - -import argparse -import json -import math -import sys -from collections.abc import Mapping -from pathlib import Path -from typing import Any - -import torch -import torch.distributed as dist -import yaml - -sys.dont_write_bytecode = True - -_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)) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--config", - type=Path, - default=_THIS_DIR / "configs" / "qwen_image.yaml", - ) - parser.add_argument( - "--checkpoint", - help="Checkpoint basename/path beneath checkpoint_dir, or LATEST; defaults to config.", - ) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--max-shard-size-gib", type=float, default=5.0) - parser.add_argument("--memory-headroom", type=float, default=1.25) - return parser.parse_args() - - -def _read_integer(path: Path) -> int | None: - try: - value = path.read_text().strip() - except OSError: - return None - if value == "max": - return None - try: - return int(value) - except ValueError: - return None - - -def host_available_bytes() -> int: - """Return the strictest visible host/cgroup memory availability estimate.""" - candidates: list[int] = [] - try: - for line in Path("/proc/meminfo").read_text().splitlines(): - if line.startswith("MemAvailable:"): - candidates.append(int(line.split()[1]) * 1024) - break - except (OSError, ValueError, IndexError): - pass - for limit_path, used_path in ( - (Path("/sys/fs/cgroup/memory.max"), Path("/sys/fs/cgroup/memory.current")), - ( - Path("/sys/fs/cgroup/memory/memory.limit_in_bytes"), - Path("/sys/fs/cgroup/memory/memory.usage_in_bytes"), - ), - ): - limit = _read_integer(limit_path) - used = _read_integer(used_path) - if limit is not None and used is not None and 0 < limit < (1 << 62): - candidates.append(max(0, limit - used)) - if not candidates: - raise RuntimeError("cannot determine host memory availability.") - return min(candidates) - - -def _state_sizes(model: torch.nn.Module) -> tuple[int, int]: - sizes = [value.numel() * value.element_size() for value in model.state_dict().values()] - if not sizes: - raise RuntimeError("PDD export model has an empty state dictionary.") - return sum(sizes), max(sizes) - - -def collective_export_memory_preflight( - model: torch.nn.Module, - *, - max_shard_bytes: int, - headroom: float, - device: torch.device, -) -> tuple[int, int]: - """Abort collectively before full-state gathering when host/GPU headroom is insufficient.""" - if not math.isfinite(headroom) or headroom < 1.0: - raise ValueError("memory_headroom must be finite and >= 1.") - full_state_bytes, largest_tensor_bytes = _state_sizes(model) - local_error = None - try: - required_gpu = math.ceil(largest_tensor_bytes * headroom) - if device.type == "cuda": - free_gpu, _total_gpu = torch.cuda.mem_get_info(device) - if free_gpu < required_gpu: - raise MemoryError( - f"rank {dist.get_rank()} has {free_gpu} free GPU bytes; " - f"full-state gather requires at least {required_gpu}." - ) - if dist.get_rank() == 0: - required_host = math.ceil((full_state_bytes + max_shard_bytes) * headroom) - available_host = host_available_bytes() - if available_host < required_host: - raise MemoryError( - f"rank 0 has {available_host} available host bytes; export requires at " - f"least {required_host}." - ) - except BaseException as error: - local_error = f"{type(error).__name__}: {error}" - errors: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(errors, local_error) - failures = [f"rank {rank}: {error}" for rank, error in enumerate(errors) if error] - if failures: - raise RuntimeError("PDD export memory preflight failed; " + "; ".join(failures)) - return full_state_bytes, largest_tensor_bytes - - -def _collective_publication_preflight(output_dir: Path) -> None: - status = None - if dist.get_rank() == 0: - try: - if output_dir.is_symlink() or output_dir.resolve().exists(): - raise FileExistsError(f"PDD export output already exists: {output_dir}.") - status = {"ok": True} - except BaseException as error: - status = {"ok": False, "error": f"{type(error).__name__}: {error}"} - payload = [status] - dist.broadcast_object_list(payload, src=0) - status = payload[0] - if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: - raise RuntimeError("rank 0 broadcast malformed PDD publication preflight status.") - if not status["ok"]: - raise RuntimeError(f"PDD publication preflight failed: {status.get('error')}.") - - -def _require_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: - from modelopt.torch.fastgen import PDDMetadata - from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION - - identity = manifest.get("identity") - if not isinstance(identity, Mapping): - raise RuntimeError("PDD checkpoint has no identity mapping.") - pdd_metadata = identity.get("pdd_metadata") - if not isinstance(pdd_metadata, Mapping): - raise RuntimeError("PDD checkpoint has no PDD metadata mapping.") - if PDDMetadata.from_dict(pdd_metadata) != setup.metadata: - raise RuntimeError("PDD checkpoint metadata does not match the configured student.") - if identity.get("qwen_image") != {"execution": QWEN_IMAGE_PDD_EXECUTION}: - raise RuntimeError("PDD checkpoint Qwen execution identity does not match MR210.") - if identity.get("model") != { - "id": config.model_id, - "revision": config.model_revision, - "dtype": str(config.dtype).removeprefix("torch."), - }: - raise RuntimeError("PDD checkpoint model identity does not match the export config.") - topology = identity.get("topology") - if not isinstance(topology, Mapping) or topology.get("world_size") != dist.get_world_size(): - raise RuntimeError("PDD checkpoint topology does not match the export process group.") - - -def _collective_checkpoint_identity(config: Any, setup: Any, manifest: Mapping[str, Any]) -> None: - local_error = None - try: - _require_checkpoint_identity(config, setup, manifest) - except BaseException as error: - local_error = f"{type(error).__name__}: {error}" - errors: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(errors, local_error) - failures = [f"rank {rank}: {error}" for rank, error in enumerate(errors) if error] - if failures: - raise RuntimeError("PDD checkpoint identity validation failed; " + "; ".join(failures)) - - -def _checkpoint_selector_identity(config: Any, setup: Any) -> dict[str, Any]: - from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION - - return { - "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, - "model": { - "id": config.model_id, - "revision": config.model_revision, - "dtype": str(config.dtype).removeprefix("torch."), - }, - "pdd_metadata": setup.metadata.to_dict(), - "guidance": {"scale": config.pdd.guidance_scale}, - "topology": {"world_size": dist.get_world_size(), "pure_data_parallel": True}, - } - - -def _collective_checkpoint_resolution( - config: Any, setup: Any, restore_from: str -) -> tuple[Path, Mapping[str, Any]]: - from pdd.checkpoint import resolve_pdd_training_checkpoint - - status = None - if dist.get_rank() == 0: - try: - checkpoint, manifest = resolve_pdd_training_checkpoint( - config.checkpoint.checkpoint_dir, - restore_from, - expected_world_size=dist.get_world_size(), - expected_identity=_checkpoint_selector_identity(config, setup), - ) - status = {"ok": True, "checkpoint": str(checkpoint), "manifest": manifest} - except BaseException as error: - status = {"ok": False, "error": f"{type(error).__name__}: {error}"} - payload = [status] - dist.broadcast_object_list(payload, src=0) - status = payload[0] - if not isinstance(status, Mapping) or type(status.get("ok")) is not bool: - raise RuntimeError("rank 0 broadcast malformed PDD checkpoint resolution status.") - if not status["ok"]: - raise RuntimeError(f"PDD checkpoint resolution failed: {status.get('error')}.") - return Path(status["checkpoint"]), status["manifest"] - - -def main() -> None: - args = _parse_args() - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict - - from pdd.artifacts import sha256_file - from pdd.export import write_pdd_export - from pdd.recipe import ( - _require_immutable_model_source, - build_pdd_export_setup, - initialize_pdd_distributed, - resolve_pdd_recipe_config, - ) - - raw = yaml.safe_load(args.config.read_text()) - config = resolve_pdd_recipe_config(raw) - _require_immutable_model_source(config, context="PDD export") - if not math.isfinite(args.max_shard_size_gib) or args.max_shard_size_gib <= 0: - raise ValueError("max_shard_size_gib must be finite and > 0.") - if not math.isfinite(args.memory_headroom) or args.memory_headroom < 1.0: - raise ValueError("memory_headroom must be finite and >= 1.") - max_shard_bytes = int(args.max_shard_size_gib * (1 << 30)) - initialize_pdd_distributed( - backend="nccl" if config.device.type == "cuda" else "gloo", - timeout_minutes=60, - ) - _collective_publication_preflight(args.output_dir) - restore_from = args.checkpoint or config.checkpoint.restore_from - if not restore_from: - raise ValueError("PDD export requires --checkpoint or checkpoint.restore_from.") - setup = build_pdd_export_setup(config) - try: - checkpoint, checkpoint_manifest = _collective_checkpoint_resolution( - config, setup, restore_from - ) - _collective_checkpoint_identity(config, setup, checkpoint_manifest) - setup.checkpointer.load_model(setup.student, str(checkpoint / "model")) - full_state_bytes, largest_tensor_bytes = collective_export_memory_preflight( - setup.student, - max_shard_bytes=max_shard_bytes, - headroom=args.memory_headroom, - device=config.device, - ) - state_dict = get_model_state_dict( - setup.student, - options=StateDictOptions(full_state_dict=True, cpu_offload=True), - ) - local_error = None - try: - if dist.get_rank() == 0: - if set(state_dict) != set(setup.checkpoint_keys): - raise RuntimeError("gathered PDD full-state keys do not match the model.") - if ( - sum(tensor.numel() * tensor.element_size() for tensor in state_dict.values()) - != full_state_bytes - ): - raise RuntimeError( - "gathered PDD full-state byte count changed after preflight." - ) - gathered_largest = max( - tensor.numel() * tensor.element_size() for tensor in state_dict.values() - ) - if gathered_largest != largest_tensor_bytes: - raise RuntimeError("gathered PDD largest tensor changed after preflight.") - elif state_dict: - raise RuntimeError("nonzero rank received a full CPU state dictionary.") - except BaseException as error: - local_error = f"{type(error).__name__}: {error}" - gather_errors: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(gather_errors, local_error) - failures = [f"rank {rank}: {error}" for rank, error in enumerate(gather_errors) if error] - if failures: - raise RuntimeError("PDD full-state gather validation failed; " + "; ".join(failures)) - - publication = None - if dist.get_rank() == 0: - try: - output = write_pdd_export( - args.output_dir, - state_dict, - metadata=setup.metadata, - transformer_config=setup.transformer_config, - identity=checkpoint_manifest["identity"], - source_checkpoint={ - "name": checkpoint.name, - "manifest_sha256": sha256_file(checkpoint / "manifest.json"), - "completed_steps": checkpoint_manifest["completed_steps"], - }, - max_shard_bytes=max_shard_bytes, - ) - publication = {"ok": True, "output": str(output)} - except BaseException as error: - publication = {"ok": False, "error": f"{type(error).__name__}: {error}"} - payload = [publication] - dist.broadcast_object_list(payload, src=0) - publication = payload[0] - if not isinstance(publication, Mapping) or type(publication.get("ok")) is not bool: - raise RuntimeError("rank 0 broadcast malformed PDD publication status.") - if not publication["ok"]: - raise RuntimeError(f"PDD export publication failed: {publication.get('error')}.") - if dist.get_rank() == 0: - print(json.dumps(publication, indent=2, sort_keys=True)) - finally: - setup.checkpointer.close() - - -if __name__ == "__main__": - main() diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index cabbf9a86fe..b77c518cebb 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -13,27 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Run conditional-only PDD inference from a complete Qwen-Image PDD export.""" +"""Generate an image with a trained Qwen-Image PDD transformer.""" from __future__ import annotations import argparse -import hashlib -import math import sys import time -from contextlib import ExitStack from pathlib import Path -from typing import TYPE_CHECKING, Any import torch +import yaml +from diffusers import QwenImagePipeline, QwenImageTransformer2DModel from torch import nn -if TYPE_CHECKING: - from collections.abc import Mapping - -sys.dont_write_bytecode = True - _THIS_DIR = Path(__file__).resolve().parent _FASTGEN_DIR = _THIS_DIR.parent _REPO_ROOT = _FASTGEN_DIR.parents[2] @@ -41,147 +34,142 @@ if str(path) not in sys.path: sys.path.insert(0, str(path)) -from pdd.inference_runtime import ( # noqa: E402 - _model_identity, - _normalize_prompt_condition, - _validate_qwen_projection, - build_pdd_student, - load_qwen_pdd_runtime, - save_png, +from modelopt.torch.fastgen import PDDConfig, PDDPipeline # noqa: E402 +from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( # noqa: E402 + QwenImagePDDAdapter, + adopt_qwen_image_mr210_forward, + restore_qwen_image_pdd_projection, ) -__all__ = [ - "_model_identity", - "_normalize_prompt_condition", - "_validate_qwen_projection", - "build_pdd_student", -] - def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--export-dir", type=Path, required=True) + 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("--prompt-id", required=True) - parser.add_argument("--schedule", choices=("pdd-2", "pdd-4", "pdd-8"), default="pdd-4") + 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) - parser.add_argument("--result-json", type=Path, required=True) return parser.parse_args() -def _resolve_outputs(output_value: Path, result_value: Path) -> tuple[Path, Path, str]: - if output_value.is_symlink() or result_value.is_symlink(): - raise ValueError("PDD output and result JSON cannot be symlinks.") - output = output_value.resolve() - result_json = result_value.resolve() - if output.exists() or output.is_symlink(): - raise FileExistsError(f"PDD inference output already exists: {output}.") - if result_json.exists() or result_json.is_symlink(): - raise FileExistsError(f"PDD result JSON already exists: {result_json}.") +def _parse_blocks(value: str) -> list[int]: try: - reference = output.relative_to(result_json.parent).as_posix() + blocks = [int(part.strip()) for part in value.split(",")] except ValueError as error: - raise ValueError("PDD output must be beneath the result JSON directory.") from error - return output, result_json, reference + 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.update(inference_blocks=blocks, student_sample_steps=len(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.no_grad() +@torch.inference_mode() def main() -> None: args = _parse_args() - from pdd.artifacts import sha256_file, write_canonical_json - - output, result_json, output_reference = _resolve_outputs(args.output, args.result_json) - if not isinstance(args.prompt_id, str) or not args.prompt_id.strip(): - raise ValueError("prompt_id must be non-empty.") - if args.seed < 0 or args.seed >= 2**63: - raise ValueError("seed must be in [0, 2**63).") - if args.max_sequence_length < 1: - raise ValueError("max_sequence_length must be positive.") - - runtime = load_qwen_pdd_runtime(args.export_dir, args.schedule, args.device) - condition = runtime.encode_prompt(args.prompt, args.max_sequence_length) - noise = runtime.make_raw_noise(seed=args.seed, height=args.height, width=args.width) - transformer_invocations = 0 - scheduler_step_calls = 0 - - def count_invocation( - _module: nn.Module, _args: tuple[Any, ...], _kwargs: Mapping[str, Any] - ) -> None: - nonlocal transformer_invocations - transformer_invocations += 1 - - scheduler = runtime.scheduler - scheduler_state = vars(scheduler).get("step") - scheduler_had_instance_step = "step" in vars(scheduler) - original_scheduler_step = scheduler.step - - def counted_scheduler_step(*call_args: Any, **call_kwargs: Any) -> Any: - nonlocal scheduler_step_calls - scheduler_step_calls += 1 - return original_scheduler_step(*call_args, **call_kwargs) - - def restore_scheduler_step() -> None: - if scheduler_had_instance_step: - setattr(scheduler, "step", scheduler_state) - elif "step" in vars(scheduler): - delattr(scheduler, "step") - - with ExitStack() as cleanup: - cleanup.callback(restore_scheduler_step) - setattr(scheduler, "step", counted_scheduler_step) - hook = runtime.student.register_forward_pre_hook(count_invocation, with_kwargs=True) - cleanup.callback(hook.remove) - if runtime.device.type == "cuda": - torch.cuda.synchronize(runtime.device) - started = time.perf_counter() - images = runtime.sample_decode(condition, noise) - if runtime.device.type == "cuda": - torch.cuda.synchronize(runtime.device) - latency = time.perf_counter() - started - - expected_invocations = len(runtime.config.inference_blocks) - if transformer_invocations != expected_invocations: - raise RuntimeError( - f"PDD sampler made {transformer_invocations} transformer calls; " - f"expected {expected_invocations}." - ) - if scheduler_step_calls != 0: - raise RuntimeError( - f"PDD sampler unexpectedly called scheduler.step {scheduler_step_calls} times." - ) - if len(images) != 1: - raise RuntimeError(f"PDD single-prompt inference returned {len(images)} images.") - if not math.isfinite(latency) or latency <= 0: - raise RuntimeError("PDD inference latency measurement is invalid.") - save_png(output, images[0]) - - result_json.parent.mkdir(parents=True, exist_ok=True) - result = { - "schema_version": 2, - "record_type": "pdd_inference", - "condition": args.schedule.replace("-", "_"), - "prompt_id": args.prompt_id, - "prompt_sha256": hashlib.sha256(args.prompt.encode("utf-8")).hexdigest(), - "seed": args.seed, - "schedule": args.schedule, - "blocks": list(runtime.config.inference_blocks), - "height": args.height, - "width": args.width, - "export_manifest_sha256": sha256_file(runtime.descriptor.root / "manifest.json"), - "output": {"path": output_reference, "sha256": sha256_file(output)}, - "scheduler_steps": expected_invocations, - "observed_scheduler_step_calls": scheduler_step_calls, - "actual_transformer_invocations": transformer_invocations, - "batch_normalized_transformer_evaluations": transformer_invocations, - "latency_seconds": latency, - } - write_canonical_json(result_json, result) - print(result_json) + 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) + adopt_qwen_image_mr210_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, + nn.Identity(), + 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__": diff --git a/examples/diffusers/fastgen/pdd/inference_runtime.py b/examples/diffusers/fastgen/pdd/inference_runtime.py deleted file mode 100644 index 3398547eaa1..00000000000 --- a/examples/diffusers/fastgen/pdd/inference_runtime.py +++ /dev/null @@ -1,301 +0,0 @@ -# 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. - -"""Shared authenticated Qwen-Image PDD inference runtime.""" - -from __future__ import annotations - -import os -import uuid -from collections.abc import Mapping -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -import torch -from torch import nn - -if TYPE_CHECKING: - from pathlib import Path - -from .export import PDD_INFERENCE_SCHEDULES, pdd_config_from_metadata - - -def _dtype_from_name(name: Any) -> torch.dtype: - if not isinstance(name, str): - raise ValueError("PDD export model dtype must be a string.") - dtypes = { - "bfloat16": torch.bfloat16, - "float16": torch.float16, - "float32": torch.float32, - } - try: - return dtypes[name] - except KeyError as error: - raise ValueError(f"PDD inference does not support model dtype {name!r}.") from error - - -def _model_identity(descriptor: Any) -> Mapping[str, Any]: - from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION - - identity = descriptor.manifest.get("identity") - if not isinstance(identity, Mapping): - raise RuntimeError("PDD export has no identity mapping.") - if identity.get("qwen_image") != {"execution": QWEN_IMAGE_PDD_EXECUTION}: - raise RuntimeError("PDD export has an incompatible Qwen execution identity.") - model = identity.get("model") - if not isinstance(model, Mapping) or set(model) != {"id", "revision", "dtype"}: - raise RuntimeError("PDD export model identity is malformed.") - if not isinstance(model["id"], str) or not model["id"]: - raise RuntimeError("PDD export model ID is invalid.") - revision = model["revision"] - if ( - not isinstance(revision, str) - or len(revision) != 40 - or any(character not in "0123456789abcdef" for character in revision) - ): - raise RuntimeError("PDD export model revision must be an exact lowercase commit.") - return model - - -def _validate_qwen_projection(student: nn.Module, metadata: Any) -> nn.Linear: - """Validate the ordinary Qwen projection before widening it for PDD.""" - try: - base_projection = student.get_submodule("proj_out") - except AttributeError as error: - raise RuntimeError("reconstructed Qwen student has no proj_out linear layer.") from error - in_channels = getattr(getattr(student, "config", None), "in_channels", None) - if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: - raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") - if ( - not isinstance(base_projection, nn.Linear) - or base_projection.in_features != metadata.projection_in_features - or base_projection.out_features != metadata.projection_out_features - or (base_projection.bias is not None) != metadata.projection_bias - ): - raise RuntimeError("reconstructed Qwen proj_out does not match the export metadata.") - if base_projection.out_features != in_channels: - raise RuntimeError( - "Qwen proj_out width must equal transformer in_channels for 2x2 latent packing." - ) - return base_projection - - -def build_pdd_student( - export_dir: str | Path, *, schedule: str = "pdd-4" -) -> tuple[nn.Module, Any, torch.dtype]: - """Reconstruct and strictly load a converted Qwen student on CPU.""" - from diffusers import QwenImageTransformer2DModel - - from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - adopt_qwen_image_mr210_forward, - convert_qwen_image_to_pdd, - ) - - from .export import inspect_pdd_export, load_pdd_export_into_model - - if schedule not in PDD_INFERENCE_SCHEDULES: - raise ValueError( - f"Unknown PDD schedule {schedule!r}; expected {sorted(PDD_INFERENCE_SCHEDULES)}." - ) - descriptor = inspect_pdd_export(export_dir) - model_identity = _model_identity(descriptor) - dtype = _dtype_from_name(model_identity["dtype"]) - student = QwenImageTransformer2DModel.from_config(dict(descriptor.transformer_config)) - _validate_qwen_projection(student, descriptor.metadata) - student = adopt_qwen_image_mr210_forward(student) - config = pdd_config_from_metadata(descriptor.metadata, schedule=schedule) - convert_qwen_image_to_pdd(student, config) - descriptor = load_pdd_export_into_model(export_dir, student) - student.to(dtype=dtype) - return student, descriptor, dtype - - -def _normalize_prompt_condition( - prompt_embeds: Any, - prompt_mask: Any, - *, - device: torch.device, - dtype: torch.dtype, -) -> tuple[torch.Tensor, torch.Tensor]: - if not isinstance(prompt_embeds, torch.Tensor) or prompt_embeds.ndim != 3: - raise RuntimeError("Qwen prompt embeddings must have shape [B, S, D].") - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - expected_shape = prompt_embeds.shape[:2] - if prompt_mask is None: - prompt_mask = torch.ones(expected_shape, device=device, dtype=torch.long) - elif not isinstance(prompt_mask, torch.Tensor) or prompt_mask.ndim != 2: - raise RuntimeError("Qwen prompt mask must have shape [B, S] or be None.") - elif tuple(prompt_mask.shape) != tuple(expected_shape): - raise RuntimeError("Qwen prompt mask shape does not match prompt embeddings.") - elif prompt_mask.dtype.is_floating_point or prompt_mask.dtype.is_complex: - raise RuntimeError("Qwen prompt mask must use an integer or boolean dtype.") - else: - prompt_mask = prompt_mask.to(device=device, dtype=torch.long) - return prompt_embeds, prompt_mask - - -def _latent_shape(pipe: Any, *, height: int, width: int) -> tuple[int, int, int, int]: - if type(height) is not int or type(width) is not int or height <= 0 or width <= 0: - raise ValueError("height and width must be positive integers.") - quantum = int(pipe.vae_scale_factor) * 2 - if height % quantum or width % quantum: - raise ValueError(f"height and width must be divisible by {quantum}.") - in_channels = getattr(pipe.transformer.config, "in_channels", None) - if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: - raise RuntimeError("Qwen transformer in_channels must be a positive multiple of four.") - return 1, in_channels // 4, 2 * (height // quantum), 2 * (width // quantum) - - -def _decode_qwen_latents(pipe: Any, latents: torch.Tensor) -> list[Any]: - if latents.ndim != 4: - raise ValueError("PDD Qwen latents must have shape [B, C, H, W].") - vae = pipe.vae - mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype) - std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype) - if mean.numel() != latents.shape[1] or std.numel() != latents.shape[1]: - raise RuntimeError("Qwen VAE latent statistics do not match the student channels.") - decoded_input = latents.unsqueeze(2) * std.view(1, -1, 1, 1, 1) - decoded_input = decoded_input + mean.view(1, -1, 1, 1, 1) - decoded = vae.decode(decoded_input, return_dict=False)[0] - if decoded.ndim != 5 or decoded.shape[2] != 1: - raise RuntimeError("Qwen VAE must return one-frame 5D image tensors.") - return pipe.image_processor.postprocess(decoded[:, :, 0], output_type="pil") - - -def save_png(path: Path, image: Any) -> None: - """Publish one PNG exclusively and durably.""" - if path.is_symlink(): - raise ValueError("PDD inference output cannot be a symlink.") - path = path.resolve() - if path.suffix.lower() != ".png": - raise ValueError("PDD inference output must use a .png suffix.") - path.parent.mkdir(parents=True, exist_ok=True) - if path.exists() or path.is_symlink(): - raise FileExistsError(f"PDD inference output already exists: {path}.") - staging = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") - try: - image.save(staging, format="PNG") - with staging.open("rb") as stream: - os.fsync(stream.fileno()) - staging.rename(path) - descriptor = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - finally: - staging.unlink(missing_ok=True) - - -@dataclass(frozen=True) -class QwenPDDInferenceRuntime: - """One loaded Qwen pipeline and one authenticated PDD sampler.""" - - student: nn.Module - scheduler: Any - descriptor: Any - model_identity: Mapping[str, Any] - dtype: torch.dtype - device: torch.device - config: Any - pipe: Any - sampler: Any - - def encode_prompt(self, prompt: str, max_sequence_length: int) -> Any: - if not isinstance(prompt, str) or not prompt: - raise ValueError("prompt must be a non-empty string.") - if type(max_sequence_length) is not int or max_sequence_length < 1: - raise ValueError("max_sequence_length must be positive.") - prompt_embeds, prompt_mask = self.pipe.encode_prompt( - prompt=prompt, - device=self.device, - num_images_per_prompt=1, - max_sequence_length=max_sequence_length, - ) - return _normalize_prompt_condition( - prompt_embeds, - prompt_mask, - device=self.device, - dtype=self.dtype, - ) - - def make_raw_noise(self, *, seed: int, height: int, width: int) -> torch.Tensor: - if type(seed) is not int or seed < 0 or seed >= 2**63: - raise ValueError("seed must be in [0, 2**63).") - generator = torch.Generator(device=self.device).manual_seed(seed) - shape = _latent_shape(self.pipe, height=height, width=width) - return torch.randn( - shape, - generator=generator, - device=self.device, - dtype=torch.float32, - ) - - def sample_decode(self, condition: Any, raw_noise: torch.Tensor) -> list[Any]: - sampled = self.sampler.sample(raw_noise, condition=condition) - return _decode_qwen_latents(self.pipe, sampled.to(self.dtype)) - -def load_qwen_pdd_runtime( - export_dir: str | Path, schedule: str, device: str | torch.device -) -> QwenPDDInferenceRuntime: - """Load one authenticated Qwen PDD runtime for a source-owned schedule.""" - from diffusers import QwenImagePipeline - - from modelopt.torch.fastgen import PDDPipeline - from modelopt.torch.fastgen.plugins.qwen_image_pdd import QwenImagePDDAdapter - - if schedule not in PDD_INFERENCE_SCHEDULES: - raise ValueError( - f"Unknown PDD schedule {schedule!r}; expected {sorted(PDD_INFERENCE_SCHEDULES)}." - ) - resolved_device = torch.device(device) - if resolved_device.type == "cuda" and not torch.cuda.is_available(): - raise RuntimeError("CUDA was requested but is unavailable.") - student, descriptor, dtype = build_pdd_student(export_dir, schedule=schedule) - model_identity = _model_identity(descriptor) - student.to(device=resolved_device) - pipe = QwenImagePipeline.from_pretrained( - model_identity["id"], - revision=model_identity["revision"], - transformer=student, - torch_dtype=dtype, - use_safetensors=True, - ) - if pipe.transformer is not student: - raise RuntimeError("Qwen pipeline did not retain the adopted PDD transformer.") - pipe.to(resolved_device) - scheduler = getattr(pipe, "scheduler", None) - if scheduler is None or not callable(getattr(scheduler, "step", None)): - raise RuntimeError("Qwen pipeline scheduler does not expose a callable step method.") - config = pdd_config_from_metadata(descriptor.metadata, schedule=schedule) - if tuple(config.inference_blocks) != PDD_INFERENCE_SCHEDULES[schedule]: - raise RuntimeError("authenticated PDD schedule changed after validation.") - sampler = PDDPipeline( - student, - nn.Identity(), - config, - QwenImagePDDAdapter(config, compute_dtype=dtype), - ) - return QwenPDDInferenceRuntime( - student=student, - scheduler=scheduler, - descriptor=descriptor, - model_identity=model_identity, - dtype=dtype, - device=resolved_device, - config=config, - pipe=pipe, - sampler=sampler, - ) 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 index f9cefe6d67d..2ae0d3c227e 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -13,1461 +13,132 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ModelOpt-owned construction of a Qwen-Image PDD student and frozen teacher.""" +"""Thin PDD objective integration for AutoModel's diffusion recipe.""" from __future__ import annotations -import copy import logging -import math -import re -import time from collections.abc import Mapping -from dataclasses import dataclass from pathlib import Path from typing import Any -import torch import torch.distributed as dist +from huggingface_hub import snapshot_download from torch import nn -from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDOutputProjection, PDDPipeline +try: + from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.training.rng import ScopedRNG + from nemo_automodel.recipes.diffusion.train import ( + TrainDiffusionRecipe, + _build_diffusion_parallel_manager_args, + ) +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, adopt_qwen_image_mr210_forward, - convert_qwen_image_to_pdd, - require_qwen_image_mr210_forward, -) - -from .checkpoint import PDDCheckpointManager, build_pdd_checkpoint_identity -from .data import ( - _build_training_dataloader, - _build_validation_dataloader, - _build_validation_plan, - _collective_training_batch, - _collective_training_iterator, - _coverage_axis, - _iter_validation_batches, - _validate_dataset_contract, ) -_HF_COMMIT_PATTERN = re.compile(r"[0-9a-f]{40}") - - -@dataclass(frozen=True) -class PDDParallelConfig: - """Pure-data-parallel FSDP2 settings for the first Qwen-Image example.""" - - dp_size: int | None = None - activation_checkpointing: bool = False - - -@dataclass(frozen=True) -class PDDCheckpointConfig: - """AutoModel Checkpointer settings needed by the PDD lifecycle.""" - - checkpoint_dir: str = "checkpoints/pdd_qwen_image" - enabled: bool = True - model_save_format: str = "torch_save" - restore_from: str | None = None - save_consolidated: bool = False - - -@dataclass(frozen=True) -class PDDStepSchedulerConfig: - """AutoModel-compatible batch, cadence, and termination settings.""" - - max_steps: int = 10_000 - num_epochs: int = 200 - log_every: int = 10 - ckpt_every_steps: int = 1_000 - local_batch_size: int = 4 - global_batch_size: int | None = None - save_checkpoint_every_epoch: bool = False - - -@dataclass(frozen=True) -class PDDTrainingHealthConfig: - """PDD-only update health settings.""" - - max_grad_norm: float = 1.0 - zero_grad_warmup_steps: int = 0 - - -@dataclass(frozen=True) -class PDDValidationConfig: - """Deterministic held-out validation settings.""" - - count: int = 2_000 - seed: int = 2026 - split_seed: int = 2026 - every_steps: int = 1_000 - - -@dataclass(frozen=True) -class PDDRecipeConfig: - """Resolved setup inputs; incompatible mutation modes have already been rejected.""" - - model_id: str - model_revision: str | None - pdd: PDDConfig - parallel: PDDParallelConfig - checkpoint: PDDCheckpointConfig - step_scheduler: PDDStepSchedulerConfig - training_health: PDDTrainingHealthConfig - validation: PDDValidationConfig - seed: int - learning_rate: float - weight_decay: float - adam_betas: tuple[float, float] - adam_eps: float - device: torch.device - dtype: torch.dtype - fuse_qkv_projections: bool - - -@dataclass(frozen=True) -class PDDSetupArtifacts: - """Objects produced in the required load-to-checkpoint construction order.""" - - pipe: Any - student: nn.Module - teacher: nn.Module - projection: PDDOutputProjection - optimizer: torch.optim.Optimizer - distributed_setup: Any - fsdp_manager: Any - checkpointer: Any - metadata: PDDMetadata - checkpoint_keys: tuple[str, ...] - lifecycle: tuple[str, ...] - - -@dataclass(frozen=True) -class PDDTrainingArtifacts: - """Direct-update objects layered on the already-constructed Task-7 setup.""" - - pipeline: PDDPipeline - trainer: Any - scheduler: torch.optim.lr_scheduler.LRScheduler - rng: Any - +from .training import PDDFlowMatchingStepAdapter -@dataclass(frozen=True) -class PDDExportSetupArtifacts: - """Student-only FSDP2 objects needed for collective DCP export.""" - pipe: Any - student: nn.Module - projection: PDDOutputProjection - distributed_setup: Any - fsdp_manager: Any - checkpointer: Any - metadata: PDDMetadata - checkpoint_keys: tuple[str, ...] - transformer_config: Mapping[str, Any] - lifecycle: tuple[str, ...] +def _config_mapping(value: Any) -> dict[str, Any]: + if hasattr(value, "to_dict"): + return value.to_dict() + return dict(value) -def _as_mapping(value: Any, *, name: str) -> Mapping[str, Any]: - if not isinstance(value, Mapping): - raise TypeError(f"{name} must be a mapping, got {type(value).__name__}.") - return value - - -def _config_to_mapping(value: Any) -> Mapping[str, Any]: - """Materialize AutoModel config targets as their original YAML dotted paths.""" - if isinstance(value, Mapping): - return value - to_yaml_dict = getattr(value, "to_yaml_dict", None) - if callable(to_yaml_dict): - return _as_mapping( - to_yaml_dict(resolve_env=True, use_orig_values=True), - name="ConfigNode.to_yaml_dict() result", - ) - to_dict = getattr(value, "to_dict", None) - if callable(to_dict): - return _as_mapping(to_dict(), name="config.to_dict() result") - raise TypeError(f"config must be a mapping or ConfigNode, got {type(value).__name__}.") - - -def _reject_enabled(value: Any, *, name: str) -> None: - if value is None or value is False or value == {}: - return - raise ValueError(f"PDD does not support {name}; disable it before model loading.") - - -def _require_bool(value: Any, *, name: str) -> bool: - if type(value) is not bool: - raise TypeError(f"{name} must be bool.") - return value - - -def _require_int_at_least(value: Any, *, name: str, minimum: int) -> int: - if type(value) is not int or value < minimum: - raise ValueError(f"{name} must be an integer >= {minimum}.") - return value - - -def _require_finite_real(value: Any, *, name: str, minimum: float | None = None) -> float: - if isinstance(value, bool) or not isinstance(value, int | float): - raise TypeError(f"{name} must be a real number.") - resolved = float(value) - if not math.isfinite(resolved) or (minimum is not None and resolved < minimum): - qualifier = "finite" if minimum is None else f"finite and >= {minimum}" - raise ValueError(f"{name} must be {qualifier}.") - return resolved - - -def _resolve_dtype(value: Any) -> torch.dtype: - if isinstance(value, torch.dtype): - return value - if not isinstance(value, str): - raise TypeError("model.torch_dtype must be a torch dtype name.") - dtypes = { - "bfloat16": torch.bfloat16, - "float16": torch.float16, - "float32": torch.float32, - } +def _validate_prepared_student(model: nn.Module, config: PDDConfig) -> None: + """Require the widened PDD projection to exist before AutoModel setup.""" try: - return dtypes[value] - except KeyError as error: - raise ValueError( - f"Unsupported model.torch_dtype={value!r}; expected {sorted(dtypes)}." - ) from error - - -def _is_exact_hf_commit(value: Any) -> bool: - return isinstance(value, str) and _HF_COMMIT_PATTERN.fullmatch(value) is not None - - -def _require_immutable_model_source(config: PDDRecipeConfig, *, context: str) -> None: - if Path(config.model_id).is_dir() or not _is_exact_hf_commit(config.model_revision): - raise ValueError( - f"{context} requires a Hugging Face model ID and exact lowercase 40-character " - "commit revision." - ) - - -def resolve_pdd_recipe_config(raw: Any) -> PDDRecipeConfig: - """Resolve one canonical DMD2-shaped PDD configuration.""" - raw = _config_to_mapping(raw) - - legacy_training = _as_mapping(raw.get("training", {}), name="training") - legacy_replacements = { - "seed": "seed", - "max_steps": "step_scheduler.max_steps", - "log_every_steps": "step_scheduler.log_every", - "checkpoint_every_steps": "step_scheduler.ckpt_every_steps", - "validation_every_steps": "validation.every_steps", - "local_batch_size": "step_scheduler.local_batch_size", - "global_batch_size": "step_scheduler.global_batch_size", - "grad_accumulation_steps": "step_scheduler.global_batch_size", - "validation_seed": "validation.seed", - "max_grad_norm": "training_health.max_grad_norm", - "zero_grad_warmup_steps": "training_health.zero_grad_warmup_steps", - } - if legacy_training: - key = next(iter(legacy_training)) - replacement_key = legacy_replacements.get(key, "the canonical PDD schema") - raise ValueError(f"training.{key} is unsupported; use {replacement_key}.") - - model = _as_mapping(raw.get("model"), name="model") - pdd_raw = _as_mapping(raw.get("pdd"), name="pdd") - fsdp = _as_mapping(raw.get("fsdp", {}), name="fsdp") - optim = _as_mapping(raw.get("optim", {}), name="optim") - optimizer_cfg = _as_mapping(optim.get("optimizer", {}), name="optim.optimizer") - lr_scheduler = _as_mapping(raw.get("lr_scheduler", {}), name="lr_scheduler") - step_scheduler = _as_mapping(raw.get("step_scheduler", {}), name="step_scheduler") - training_health = _as_mapping(raw.get("training_health", {}), name="training_health") - validation = _as_mapping(raw.get("validation", {}), name="validation") - checkpoint = _as_mapping(raw.get("checkpoint", {}), name="checkpoint") - data = _as_mapping(raw.get("data", {}), name="data") - dataloader = _as_mapping(data.get("dataloader", {}), name="data.dataloader") - - for legacy_key in ("weight_decay", "betas", "eps"): - if legacy_key in optim: - raise ValueError( - f"optim.{legacy_key} is unsupported; use optim.optimizer.{legacy_key}." - ) - - for legacy_key, replacement_key in ( - ("validation_count", "validation.count"), - ("split_seed", "validation.split_seed"), - ): - if legacy_key in data: - raise ValueError(f"data.{legacy_key} is unsupported; use {replacement_key}.") - - target = dataloader.get("_target_") - expected_target = "fastgen_data.build_text_to_image_multiresolution_dataloader" - if target is not None and target != expected_target: - raise ValueError(f"data.dataloader._target_ must be {expected_target!r}.") - if _require_bool(dataloader.get("drop_last", True), name="data.dataloader.drop_last") is False: - raise ValueError("PDD exact sample accounting requires data.dataloader.drop_last=true.") - if _require_bool( - dataloader.get("dynamic_batch_size", False), - name="data.dataloader.dynamic_batch_size", - ): - raise ValueError("PDD v1 requires data.dataloader.dynamic_batch_size=false.") - if _require_bool( - dataloader.get("train_text_encoder", False), - name="data.dataloader.train_text_encoder", - ): - raise ValueError("PDD requires cached text embeddings; train_text_encoder must be false.") - _require_bool(dataloader.get("shuffle", True), name="data.dataloader.shuffle") - _require_bool( - dataloader.get("verify_payload_hashes", False), - name="data.dataloader.verify_payload_hashes", - ) - if "metadata_index" in dataloader: + 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( - "PDD uses deterministic ordinal splits from metadata.json; " - "data.dataloader.metadata_index is unsupported." - ) - - _reject_enabled(model.get("transformer_engine_linear"), name="global TE-linear conversion") - _reject_enabled(model.get("peft"), name="PEFT/LoRA") - _reject_enabled(model.get("peft_cfg"), name="PEFT/LoRA") - _reject_enabled(raw.get("peft"), name="PEFT/LoRA") - _reject_enabled(raw.get("peft_cfg"), name="PEFT/LoRA") - _reject_enabled(raw.get("guidance"), name="guidance overrides") - _reject_enabled(model.get("guidance_embeds"), name="Qwen guidance embeddings") - _reject_enabled(model.get("guidance_embeddings"), name="Qwen guidance embeddings") - for option in ( - "device_map", - "load_in_4bit", - "load_in_8bit", - "offload_folder", - "offload_state_dict", - "quantization_config", - ): - _reject_enabled(model.get(option), name=f"model loader option {option!r}") - - model_id = model.get("pretrained_model_name_or_path") - if not isinstance(model_id, str) or not model_id: - raise ValueError("model.pretrained_model_name_or_path must be a non-empty string.") - model_revision = model.get("revision") - if Path(model_id).is_dir(): - if model_revision is not None: - raise ValueError("Local PDD model directories require model.revision=null.") - elif not _is_exact_hf_commit(model_revision): - raise ValueError( - "Remote PDD models require an exact lowercase 40-character model.revision commit." - ) - - learning_rate = _require_finite_real( - optim.get("learning_rate", 2.0e-5), - name="optim.learning_rate", - minimum=0.0, - ) - if learning_rate == 0.0: - raise ValueError("optim.learning_rate must be > 0.") - optimizer_target = optimizer_cfg.get("_target_", "torch.optim.AdamW") - if optimizer_target != "torch.optim.AdamW": - raise ValueError("PDD v1 requires optim.optimizer._target_='torch.optim.AdamW'.") - allowed_optimizer_keys = { - "_target_", - "weight_decay", - "betas", - "eps", - "amsgrad", - "capturable", - "differentiable", - "foreach", - "fused", - "maximize", - } - unsupported_optimizer_keys = sorted(set(optimizer_cfg) - allowed_optimizer_keys) - if unsupported_optimizer_keys: - raise ValueError(f"unsupported PDD optimizer keys: {unsupported_optimizer_keys}.") - for flag in ("amsgrad", "capturable", "differentiable", "foreach", "fused", "maximize"): - if _require_bool(optimizer_cfg.get(flag, False), name=f"optim.optimizer.{flag}"): - raise ValueError(f"PDD v1 requires optim.optimizer.{flag}=false.") - weight_decay = _require_finite_real( - optimizer_cfg.get("weight_decay", 0.0), - name="optim.optimizer.weight_decay", - minimum=0.0, - ) - adam_betas_raw = optimizer_cfg.get("betas", [0.9, 0.999]) - if ( - not isinstance(adam_betas_raw, list | tuple) - or len(adam_betas_raw) != 2 - or any( - isinstance(beta, bool) or not isinstance(beta, int | float) for beta in adam_betas_raw - ) - ): - raise TypeError("optim.optimizer.betas must contain two real numbers.") - adam_betas = (float(adam_betas_raw[0]), float(adam_betas_raw[1])) - if any(not math.isfinite(beta) or not 0.0 <= beta < 1.0 for beta in adam_betas): - raise ValueError("optim.optimizer.betas values must be finite and in [0, 1).") - adam_eps = _require_finite_real( - optimizer_cfg.get("eps", 1e-8), - name="optim.optimizer.eps", - minimum=0.0, - ) - if adam_eps == 0.0: - raise ValueError("optim.optimizer.eps must be > 0.") - - lr_decay_style = lr_scheduler.get("lr_decay_style", "constant") - if lr_decay_style != "constant": - raise ValueError("PDD v1 requires lr_scheduler.lr_decay_style='constant'.") - lr_warmup_steps = _require_int_at_least( - lr_scheduler.get("lr_warmup_steps", 0), - name="lr_scheduler.lr_warmup_steps", - minimum=0, - ) - if lr_warmup_steps != 0: - raise ValueError("PDD v1 requires lr_scheduler.lr_warmup_steps=0.") - min_lr = _require_finite_real( - lr_scheduler.get("min_lr", learning_rate), - name="lr_scheduler.min_lr", - minimum=0.0, - ) - if min_lr != learning_rate: - raise ValueError("lr_scheduler.min_lr must equal optim.learning_rate for constant PDD LR.") - if "max_lr" in lr_scheduler: - max_lr = _require_finite_real( - lr_scheduler["max_lr"], - name="lr_scheduler.max_lr", - minimum=0.0, + "Prepare the Qwen PDD student before training: expected proj_out.out_features=" + f"{expected_out_features}, got {projection.out_features}." ) - if max_lr != learning_rate: - raise ValueError( - "lr_scheduler.max_lr must equal optim.learning_rate for constant PDD LR." - ) - - dp_size = fsdp.get("dp_size") - if dp_size is not None and (type(dp_size) is not int or dp_size < 1): - raise ValueError("fsdp.dp_size must be null or an integer >= 1.") - activation_checkpointing = fsdp.get("activation_checkpointing", False) - if type(activation_checkpointing) is not bool: - raise TypeError("fsdp.activation_checkpointing must be bool.") - for dimension in ("tp_size", "cp_size", "pp_size", "ep_size"): - value = fsdp.get(dimension, 1) - if type(value) is not int or value != 1: - raise ValueError(f"PDD v1 supports pure data parallelism; fsdp.{dimension} must be 1.") - - pdd = PDDConfig(**dict(pdd_raw)) - if pdd.num_train_timesteps is not None: - raise ValueError("Qwen-Image PDD requires pdd.num_train_timesteps=null.") - - checkpoint_enabled = _require_bool(checkpoint.get("enabled", True), name="checkpoint.enabled") - save_consolidated = _require_bool( - checkpoint.get("save_consolidated", False), - name="checkpoint.save_consolidated", - ) - if save_consolidated: - raise ValueError("PDD training checkpoints require checkpoint.save_consolidated=false.") - checkpoint_dir = checkpoint.get("checkpoint_dir", "checkpoints/pdd_qwen_image") - if not isinstance(checkpoint_dir, str) or not checkpoint_dir: - raise ValueError("checkpoint.checkpoint_dir must be a non-empty string.") - model_save_format = checkpoint.get("model_save_format", "torch_save") - if model_save_format != "torch_save": - raise ValueError("PDD training checkpoints require model_save_format='torch_save'.") - restore_from = checkpoint.get("restore_from") - if restore_from is not None and (not isinstance(restore_from, str) or not restore_from): - raise ValueError("checkpoint.restore_from must be null or a non-empty string.") - if not checkpoint_enabled and restore_from is not None: - raise ValueError("checkpoint.restore_from requires checkpoint.enabled=true.") - fuse_qkv_projections = _require_bool( - model.get("fuse_qkv_projections", False), - name="model.fuse_qkv_projections", - ) - if fuse_qkv_projections: - raise ValueError("Qwen MR210 PDD does not support QKV fusion.") - seed = _require_int_at_least(raw.get("seed", 42), name="seed", minimum=0) - max_steps = _require_int_at_least( - step_scheduler.get("max_steps", 10_000), - name="step_scheduler.max_steps", - minimum=1, - ) - num_epochs = _require_int_at_least( - step_scheduler.get("num_epochs", 200), - name="step_scheduler.num_epochs", - minimum=1, - ) - log_every = _require_int_at_least( - step_scheduler.get("log_every", 10), - name="step_scheduler.log_every", - minimum=1, - ) - ckpt_every_steps = _require_int_at_least( - step_scheduler.get("ckpt_every_steps", 1_000), - name="step_scheduler.ckpt_every_steps", - minimum=1, - ) - save_checkpoint_every_epoch = _require_bool( - step_scheduler.get("save_checkpoint_every_epoch", False), - name="step_scheduler.save_checkpoint_every_epoch", - ) - if save_checkpoint_every_epoch: - raise ValueError( - "PDD exact resume requires step_scheduler.save_checkpoint_every_epoch=false." - ) - local_batch_size = _require_int_at_least( - step_scheduler.get("local_batch_size", 1), - name="step_scheduler.local_batch_size", - minimum=1, - ) - data_batch_size = _require_int_at_least( - dataloader.get("batch_size", local_batch_size), - name="data.dataloader.batch_size", - minimum=1, - ) - if data_batch_size != local_batch_size: - raise ValueError("data.dataloader.batch_size must equal step_scheduler.local_batch_size.") - global_batch_size = step_scheduler.get("global_batch_size") - if global_batch_size is not None: - global_batch_size = _require_int_at_least( - global_batch_size, - name="step_scheduler.global_batch_size", - minimum=1, - ) - - max_grad_norm = _require_finite_real( - training_health.get("max_grad_norm", 1.0), - name="training_health.max_grad_norm", - minimum=0.0, - ) - if max_grad_norm == 0.0: - raise ValueError("training_health.max_grad_norm must be > 0.") - zero_grad_warmup_steps = _require_int_at_least( - training_health.get("zero_grad_warmup_steps", 0), - name="training_health.zero_grad_warmup_steps", - minimum=0, - ) - validation_count = _require_int_at_least( - validation.get("count", 2_000), - name="validation.count", - minimum=1, - ) - validation_seed = _require_int_at_least( - validation.get("seed", 2026), - name="validation.seed", - minimum=0, - ) - split_seed = _require_int_at_least( - validation.get("split_seed", 2026), - name="validation.split_seed", - minimum=0, - ) - validation_every_steps = _require_int_at_least( - validation.get("every_steps", 1_000), - name="validation.every_steps", - minimum=1, - ) - - dtype = _resolve_dtype(model.get("torch_dtype", "bfloat16")) - if dtype != torch.bfloat16: - raise ValueError("Qwen MR210 PDD requires model.torch_dtype='bfloat16'.") - - return PDDRecipeConfig( - model_id=model_id, - model_revision=model_revision, - pdd=pdd, - parallel=PDDParallelConfig( - dp_size=dp_size, - activation_checkpointing=activation_checkpointing, - ), - checkpoint=PDDCheckpointConfig( - checkpoint_dir=checkpoint_dir, - enabled=checkpoint_enabled, - model_save_format=model_save_format, - restore_from=restore_from, - save_consolidated=save_consolidated, - ), - step_scheduler=PDDStepSchedulerConfig( - max_steps=max_steps, - num_epochs=num_epochs, - log_every=log_every, - ckpt_every_steps=ckpt_every_steps, - local_batch_size=local_batch_size, - global_batch_size=global_batch_size, - save_checkpoint_every_epoch=save_checkpoint_every_epoch, - ), - training_health=PDDTrainingHealthConfig( - max_grad_norm=max_grad_norm, - zero_grad_warmup_steps=zero_grad_warmup_steps, - ), - validation=PDDValidationConfig( - count=validation_count, - seed=validation_seed, - split_seed=split_seed, - every_steps=validation_every_steps, - ), - seed=seed, - learning_rate=float(learning_rate), - weight_decay=float(weight_decay), - adam_betas=adam_betas, - adam_eps=adam_eps, - device=torch.device(model.get("device", "cuda" if torch.cuda.is_available() else "cpu")), - dtype=dtype, - fuse_qkv_projections=fuse_qkv_projections, - ) - - -def _projection_identity(projection: PDDOutputProjection) -> tuple[int, int, int | None]: - return ( - id(projection), - id(projection.weight), - None if projection.bias is None else id(projection.bias), - ) - - -def _require_projection_identity( - student: nn.Module, - projection: PDDOutputProjection, - expected: tuple[int, int, int | None], - *, - stage: str, -) -> None: - if student.get_submodule("proj_out") is not projection: - raise RuntimeError(f"PDD projection was replaced during {stage}.") - if _projection_identity(projection) != expected: - raise RuntimeError(f"PDD projection parameter identity changed during {stage}.") - - -def _require_projection_module( - student: nn.Module, - projection: PDDOutputProjection, - *, - stage: str, -) -> None: - if student.get_submodule("proj_out") is not projection: - raise RuntimeError(f"PDD projection module was replaced during {stage}.") - - -def _resolve_model_source(config: PDDRecipeConfig) -> str: - if Path(config.model_id).is_dir(): - return str(Path(config.model_id).resolve()) - from huggingface_hub import snapshot_download - - model_source = Path( - snapshot_download(config.model_id, revision=config.model_revision) - ).resolve() - if model_source.parent.name != "snapshots" or model_source.name != config.model_revision: - raise RuntimeError( - "Hugging Face resolved a model snapshot that does not match model.revision." - ) - return str(model_source) - - -def _load_unwrapped_transformer( - config: PDDRecipeConfig, pipeline_type: Any -) -> tuple[Any, nn.Module]: - pipe, loader_managers = pipeline_type.from_pretrained( - _resolve_model_source(config), - parallel_scheme=None, - device=None, - torch_dtype=config.dtype, - move_to_device=False, - load_for_training=True, - components_to_load=["transformer"], - peft_cfg=None, - active_transformer="transformer", - transformer_engine_linear=False, - fuse_qkv_projections=False, - compact_fused_qkv_projections=False, - low_cpu_mem_usage=True, - text_encoder=None, - tokenizer=None, - vae=None, - ) - if loader_managers: - raise RuntimeError("Unwrapped AutoModel load unexpectedly created parallel managers.") - student = pipe.transformer - if not isinstance(student, nn.Module): - raise TypeError("AutoModel pipeline did not return an nn.Module transformer.") - return pipe, student - -def _apply_qwen_image_activation_checkpointing(model: nn.Module, *, enabled: bool) -> int: - """Apply FastGen-compatible non-reentrant checkpointing to Qwen-Image blocks.""" - if type(enabled) is not bool: - raise TypeError("enabled must be bool.") - - # Diffusers is an optional example dependency; defer imports until this adapter is used. - from diffusers import QwenImageTransformer2DModel - from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock - from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( - CheckpointImpl, - CheckpointWrapper, - checkpoint_wrapper, - ) - - if type(model) is not QwenImageTransformer2DModel: - raise TypeError( - "PDD activation checkpointing requires an ordinary " - f"QwenImageTransformer2DModel, got {type(model).__name__}." - ) - blocks = model.transformer_blocks - if not isinstance(blocks, nn.ModuleList) or not blocks: - raise TypeError("Qwen-Image transformer_blocks must be a non-empty nn.ModuleList.") - wrapped_modules = [ - (index, name or "") - for index, block in enumerate(blocks) - for name, module in block.named_modules() - if isinstance(module, CheckpointWrapper) or hasattr(module, "_checkpoint_wrapped_module") - ] - if wrapped_modules: - raise RuntimeError( - "Qwen-Image has already checkpoint-wrapped modules in transformer blocks: " - f"{wrapped_modules}." - ) - unexpected = [ - (index, type(block).__name__) - for index, block in enumerate(blocks) - if type(block) is not QwenImageTransformerBlock - ] - if unexpected: - raise TypeError(f"Qwen-Image transformer block types changed: {unexpected}.") - - model.disable_gradient_checkpointing() - if model.gradient_checkpointing: - raise RuntimeError("Qwen-Image native gradient checkpointing remained enabled.") - if not enabled: - return 0 - - for index, block in enumerate(tuple(blocks)): - blocks[index] = checkpoint_wrapper( - block, - checkpoint_impl=CheckpointImpl.NO_REENTRANT, - ) - for index, block in enumerate(blocks): - if not isinstance(block, CheckpointWrapper): - raise RuntimeError(f"Qwen-Image block {index} was not checkpoint-wrapped.") - if block.checkpoint_impl is not CheckpointImpl.NO_REENTRANT: - raise RuntimeError( - f"Qwen-Image block {index} uses the wrong checkpoint implementation." - ) - if type(block._checkpoint_wrapped_module) is not QwenImageTransformerBlock: - raise RuntimeError(f"Qwen-Image block {index} wrapped an unexpected module type.") - return len(blocks) - - -def _stage_and_shard_training_models( - student: nn.Module, - teacher: nn.Module, - projection: PDDOutputProjection, - projection_identity: tuple[int, int, int | None], - manager: Any, - *, - device: torch.device, - fuse_qkv_projections: bool, - activation_checkpointing: bool, -) -> tuple[nn.Module, nn.Module]: - """Stage FP32 masters and shard one dense model at a time.""" - if fuse_qkv_projections and ( - not hasattr(student, "fuse_qkv_projections") or not hasattr(teacher, "fuse_qkv_projections") - ): - raise AttributeError("QKV fusion requires both Qwen transformers to expose the object API.") - - # Keep parameter and optimizer storage in FP32. The FSDP policy below casts - # gathered parameters to the configured dtype for forward/backward compute. - student.to(device=device, dtype=torch.float32) - if fuse_qkv_projections: - student.fuse_qkv_projections() - if not any(getattr(module, "fused_projections", False) for module in student.modules()): - logging.warning( - "Qwen fuse_qkv_projections() was accepted but produced no fused attention " - "modules in the pinned Diffusers release." - ) - student_checkpoint_blocks = _apply_qwen_image_activation_checkpointing( - student, - enabled=activation_checkpointing, - ) - logging.info("Qwen student checkpoint-wrapped blocks: %d", student_checkpoint_blocks) - _require_projection_identity(student, projection, projection_identity, stage="student staging") - student = manager.parallelize(student) - _require_projection_module(student, projection, stage="student FSDP2 parallelization") - - # The student is already sharded before the dense teacher reaches the GPU, so multi-rank - # setup never holds both complete Qwen transformers on one device. - teacher.to(device=device, dtype=torch.float32) - if fuse_qkv_projections: - teacher.fuse_qkv_projections() - teacher_checkpoint_blocks = _apply_qwen_image_activation_checkpointing( - teacher, - enabled=activation_checkpointing, - ) - logging.info("Qwen teacher checkpoint-wrapped blocks: %d", teacher_checkpoint_blocks) - teacher = manager.parallelize(teacher) - return student, teacher - - -def _materialize_zero_step_adamw_state(optimizer: torch.optim.AdamW) -> None: - """Create complete strict-DCP state without changing parameters or update numbering.""" - if type(optimizer) is not torch.optim.AdamW: - raise TypeError("PDD state materialization requires the stock torch.optim.AdamW optimizer.") - if optimizer.state: - raise RuntimeError("PDD AdamW state must be empty before materialization.") - parameters = [parameter for group in optimizer.param_groups for parameter in group["params"]] - if any(parameter.grad is not None for parameter in parameters): - raise RuntimeError("PDD AdamW parameters must not have gradients before materialization.") - - learning_rates = [group["lr"] for group in optimizer.param_groups] - try: - for group in optimizer.param_groups: - group["lr"] = 0.0 - for parameter in parameters: - parameter.grad = torch.zeros_like(parameter) - optimizer.step() - for parameter in parameters: - state = optimizer.state.get(parameter) - if state is None or set(state) != {"step", "exp_avg", "exp_avg_sq"}: - raise RuntimeError("PDD AdamW did not create complete checkpoint state.") - step = state["step"] - if not isinstance(step, torch.Tensor) or step.numel() != 1 or step.item() != 1: - raise RuntimeError("PDD AdamW created an unexpected initial step.") - step.zero_() - finally: - for group, learning_rate in zip(optimizer.param_groups, learning_rates, strict=True): - group["lr"] = learning_rate - optimizer.zero_grad(set_to_none=True) - - -def _require_fp32_optimizer_storage(optimizer: torch.optim.AdamW) -> None: - """Require FP32 master parameters and Adam state for stable small updates.""" - for group in optimizer.param_groups: - for parameter in group["params"]: - if parameter.dtype != torch.float32: - raise RuntimeError( - f"PDD trainable master parameters must be FP32, got {parameter.dtype}." - ) - state = optimizer.state.get(parameter) - if state is None: - raise RuntimeError("PDD AdamW state is missing after eager materialization.") - for name in ("exp_avg", "exp_avg_sq"): - value = state.get(name) - if not isinstance(value, torch.Tensor) or value.dtype != torch.float32: - raise RuntimeError(f"PDD AdamW {name} state must be FP32.") - - -def build_pdd_setup(config: PDDRecipeConfig) -> PDDSetupArtifacts: - """Compose released AutoModel APIs without editing or patching external packages.""" - if not isinstance(config, PDDRecipeConfig): - raise TypeError(f"config must be PDDRecipeConfig, got {type(config).__name__}.") - if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("Initialize torch.distributed before building the PDD FSDP2 setup.") - - lifecycle: list[str] = [] - - from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline - from nemo_automodel.components.checkpoint.config import CheckpointingConfig - from nemo_automodel.components.distributed import ( - DistributedSetup, - FSDP2Config, - ParallelismSizes, - ) - from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager - - pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) - student = adopt_qwen_image_mr210_forward(loaded_transformer) - pipe.transformer = student - teacher = copy.deepcopy(student).eval().requires_grad_(False) - lifecycle.append("load/select") - - projection = convert_qwen_image_to_pdd(student, config.pdd) - identity = _projection_identity(projection) - metadata = PDDMetadata.from_config(config.pdd, projection) - lifecycle.append("pdd_conversion") - - world_size = dist.get_world_size() - if config.step_scheduler.global_batch_size is not None: - effective_global_batch = config.step_scheduler.local_batch_size * world_size - if effective_global_batch != config.step_scheduler.global_batch_size: - raise ValueError( - "PDD global batch mismatch: " - f"local_batch_size={config.step_scheduler.local_batch_size} * " - f"world_size={world_size} = {effective_global_batch}, configured " - "step_scheduler.global_batch_size=" - f"{config.step_scheduler.global_batch_size}. PDD v1 requires one microbatch per " - "optimizer update." - ) - dp_size = config.parallel.dp_size or world_size - if dp_size != world_size: - raise ValueError( - f"Pure-DP PDD requires fsdp.dp_size ({dp_size}) to equal world size ({world_size})." - ) - from torch.distributed.fsdp import MixedPrecisionPolicy - - strategy = FSDP2Config( - # Qwen block checkpointing is applied explicitly before FSDP2 because - # AutoModel's generic language-model attribute policy does not match - # Qwen-Image blocks. - activation_checkpointing=False, - # Match FastGen's per-block PyTorch default instead of AutoModel's - # all-but-last ModuleList optimization. AutoModel still keeps the root - # unresharded after forward. - reshard_after_forward=True, - mp_policy=MixedPrecisionPolicy( - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - output_dtype=torch.bfloat16, - cast_forward_inputs=False, - ), - ) - distributed_setup = DistributedSetup.build( - strategy=strategy, - parallelism_sizes=ParallelismSizes(dp_size=dp_size), - activation_checkpointing=False, - world_size=world_size, - ) - mesh_context = distributed_setup.mesh_context - manager = FSDP2Manager( - distributed_setup.strategy_config, - device_mesh=mesh_context.device_mesh, - moe_mesh=mesh_context.moe_mesh, - ) - student, teacher = _stage_and_shard_training_models( - student, - teacher, - projection, - identity, - manager, - device=config.device, - fuse_qkv_projections=config.fuse_qkv_projections, - activation_checkpointing=config.parallel.activation_checkpointing, - ) - pipe.transformer = student - # Keep the public lifecycle summary stable even though placement, optional QKV fusion, and - # FSDP2 are deliberately interleaved per model to cap peak device memory. - lifecycle.extend(("device", "qkv", "parallelize")) - - trainable = [parameter for parameter in student.parameters() if parameter.requires_grad] - if not trainable: - raise RuntimeError("PDD student has no trainable parameters after FSDP2 setup.") - if any(parameter.requires_grad for parameter in teacher.parameters()): - raise RuntimeError("PDD teacher became trainable during setup.") - optimizer = torch.optim.AdamW( - trainable, - lr=config.learning_rate, - weight_decay=config.weight_decay, - betas=config.adam_betas, - eps=config.adam_eps, - amsgrad=False, - capturable=False, - differentiable=False, - foreach=False, - fused=False, - maximize=False, - ) - _materialize_zero_step_adamw_state(optimizer) - _require_fp32_optimizer_storage(optimizer) - optimizer_parameters = [ - parameter for group in optimizer.param_groups for parameter in group["params"] - ] - if not any(parameter is projection.weight for parameter in optimizer_parameters): - raise RuntimeError("PDD projection parameters are missing from the optimizer.") - lifecycle.append("optimizer") - - checkpoint_keys = tuple(student.state_dict().keys()) - projection_key = "proj_out.weight" - if projection_key not in checkpoint_keys: - raise RuntimeError( - f"PDD projection key {projection_key!r} is missing from checkpoint state." - ) - checkpoint_config = CheckpointingConfig( - enabled=config.checkpoint.enabled, - checkpoint_dir=config.checkpoint.checkpoint_dir, - model_save_format=config.checkpoint.model_save_format, - model_repo_id=config.model_id, - save_consolidated=config.checkpoint.save_consolidated, - is_peft=False, - model_state_dict_keys=list(checkpoint_keys), - ) - checkpointer = checkpoint_config.build( - dp_rank=dist.get_rank(), - tp_rank=0, - pp_rank=0, - moe_mesh=None, - ) - lifecycle.append("checkpoint") - - return PDDSetupArtifacts( - pipe=pipe, - student=student, - teacher=teacher, - projection=projection, - optimizer=optimizer, - distributed_setup=distributed_setup, - fsdp_manager=manager, - checkpointer=checkpointer, - metadata=metadata, - checkpoint_keys=checkpoint_keys, - lifecycle=tuple(lifecycle), - ) - - -def build_pdd_export_setup(config: PDDRecipeConfig) -> PDDExportSetupArtifacts: - """Build only the converted/sharded student needed for collective export.""" - if not isinstance(config, PDDRecipeConfig): - raise TypeError(f"config must be PDDRecipeConfig, got {type(config).__name__}.") - if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("Initialize torch.distributed before building PDD export setup.") - from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline - from nemo_automodel.components.checkpoint.config import CheckpointingConfig - from nemo_automodel.components.distributed import ( - DistributedSetup, - FSDP2Config, - ParallelismSizes, - ) - from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager - - lifecycle = ["load/select"] - pipe, loaded_transformer = _load_unwrapped_transformer(config, NeMoAutoDiffusionPipeline) - student = adopt_qwen_image_mr210_forward(loaded_transformer) - pipe.transformer = student - raw_transformer_config = getattr(student, "config", None) - to_dict = getattr(raw_transformer_config, "to_dict", None) - if callable(to_dict): - transformer_config = to_dict() - elif isinstance(raw_transformer_config, Mapping): - transformer_config = dict(raw_transformer_config) - else: - raise TypeError("Qwen transformer config must expose to_dict() or Mapping.") - if not isinstance(transformer_config, Mapping): - raise TypeError("Qwen transformer to_dict() must return a mapping.") - - projection = convert_qwen_image_to_pdd(student, config.pdd) - identity = _projection_identity(projection) - metadata = PDDMetadata.from_config(config.pdd, projection) - lifecycle.append("pdd_conversion") - student.to(device=config.device, dtype=torch.float32) - _require_projection_identity(student, projection, identity, stage="device placement") - lifecycle.append("device") - - if config.fuse_qkv_projections: - if not hasattr(student, "fuse_qkv_projections"): - raise AttributeError("QKV fusion requires Qwen to expose the object API.") - student.fuse_qkv_projections() - _require_projection_identity(student, projection, identity, stage="QKV fusion") - lifecycle.append("qkv") - - world_size = dist.get_world_size() - dp_size = config.parallel.dp_size or world_size - if dp_size != world_size: - raise ValueError( - f"Pure-DP PDD export requires fsdp.dp_size ({dp_size}) to equal world size " - f"({world_size})." - ) - from torch.distributed.fsdp import MixedPrecisionPolicy - - strategy = FSDP2Config( - activation_checkpointing=False, - mp_policy=MixedPrecisionPolicy( - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - output_dtype=torch.bfloat16, - cast_forward_inputs=False, - ), - ) - distributed_setup = DistributedSetup.build( - strategy=strategy, - parallelism_sizes=ParallelismSizes(dp_size=dp_size), - activation_checkpointing=False, - world_size=world_size, - ) - mesh_context = distributed_setup.mesh_context - manager = FSDP2Manager( - distributed_setup.strategy_config, - device_mesh=mesh_context.device_mesh, - moe_mesh=mesh_context.moe_mesh, - ) - student = manager.parallelize(student) - pipe.transformer = student - _require_projection_module(student, projection, stage="FSDP2 export parallelization") - lifecycle.append("parallelize") - - checkpoint_keys = tuple(student.state_dict()) - if "proj_out.weight" not in checkpoint_keys: - raise RuntimeError("PDD projection is missing from the export checkpoint key inventory.") - checkpoint_config = CheckpointingConfig( - enabled=True, - checkpoint_dir=config.checkpoint.checkpoint_dir, - model_save_format="torch_save", - model_repo_id=config.model_id, - save_consolidated=False, - is_peft=False, - model_state_dict_keys=list(checkpoint_keys), - ) - checkpointer = checkpoint_config.build( - dp_rank=dist.get_rank(), - tp_rank=0, - pp_rank=0, - moe_mesh=None, - ) - lifecycle.append("checkpoint") - return PDDExportSetupArtifacts( - pipe=pipe, - student=student, - projection=projection, - distributed_setup=distributed_setup, - fsdp_manager=manager, - checkpointer=checkpointer, - metadata=metadata, - checkpoint_keys=checkpoint_keys, - transformer_config=transformer_config, - lifecycle=tuple(lifecycle), - ) - - -def build_pdd_training_artifacts( - setup: PDDSetupArtifacts, - config: PDDRecipeConfig, -) -> PDDTrainingArtifacts: - """Layer the direct-update pipeline, constant-LR scheduler, and ranked RNG on setup.""" - if not isinstance(setup, PDDSetupArtifacts): - raise TypeError("setup must be PDDSetupArtifacts.") - if not isinstance(config, PDDRecipeConfig): - raise TypeError("config must be PDDRecipeConfig.") - from nemo_automodel.components.training.rng import StatefulRNG - - from .training import PDDTrainer - - adapter = QwenImagePDDAdapter( - config.pdd, - compute_dtype=config.dtype, - ) - pipeline = PDDPipeline(setup.student, setup.teacher, config.pdd, adapter) - scheduler = torch.optim.lr_scheduler.LambdaLR(setup.optimizer, lr_lambda=lambda _: 1.0) - rng = StatefulRNG(config.seed, ranked=True) - trainer = PDDTrainer( - pipeline, - setup.optimizer, - projection=setup.projection, - max_grad_norm=config.training_health.max_grad_norm, - warmup_steps=config.training_health.zero_grad_warmup_steps, - ) - return PDDTrainingArtifacts( - pipeline=pipeline, - trainer=trainer, - scheduler=scheduler, - rng=rng, - ) - - -def initialize_pdd_distributed(*, backend: str, timeout_minutes: int = 60) -> Any: - """Initialize through AutoModel's released public API.""" - from nemo_automodel.components.distributed import initialize_distributed - - return initialize_distributed(backend=backend, timeout_minutes=timeout_minutes) - - -class PDDDiffusionRecipe: - """Compose released AutoModel components around the PDD-specific update.""" - - def __init__(self, cfg: Any) -> None: - self.cfg = cfg - self.raw_config = _config_to_mapping(cfg) - self.config = resolve_pdd_recipe_config(self.raw_config) +class PDDDiffusionRecipe(TrainDiffusionRecipe): + """Use AutoModel's native lifecycle with a PDD loss and frozen teacher.""" def setup(self) -> None: - """Build data, converted models, AutoModel scheduling, and strict resume state.""" - config = self.config - _require_immutable_model_source(config, context="Checkpointed PDD training") - self.dist_env = initialize_pdd_distributed( - backend="nccl" if config.device.type == "cuda" else "gloo", - timeout_minutes=60, - ) - from nemo_automodel.components.loggers.log_utils import setup_logging - from nemo_automodel.components.training.step_scheduler import StepScheduler - - setup_logging() - self.rank = dist.get_rank() - self.world_size = dist.get_world_size() - self.dataloader, self.sampler = _build_training_dataloader( - self.raw_config, - config, - dp_rank=self.rank, - dp_world_size=self.world_size, - ) - self.validation_dataloader, self.validation_sampler = _build_validation_dataloader( - self.raw_config, - config, - dp_rank=self.rank, - dp_world_size=self.world_size, - ) - ( - self.snapshot_report, - train_ordered_id_sha256, - heldout_ordered_id_sha256, - ) = _validate_dataset_contract( - self.sampler.dataset, - self.validation_sampler.dataset, - config, - ) - self.validation_assignments, self.validation_masks = _build_validation_plan( - self.validation_sampler, - config, - ) - - self.setup_artifacts = build_pdd_setup(config) - qwen_image_execution = require_qwen_image_mr210_forward(self.setup_artifacts.student) - self.expected_latent_channels, self.expected_condition_features = ( - self._resolve_transformer_dimensions(self.setup_artifacts.student) - ) - self.training = build_pdd_training_artifacts(self.setup_artifacts, config) - global_batch_size = ( - config.step_scheduler.global_batch_size - or config.step_scheduler.local_batch_size * self.world_size - ) - self.step_scheduler = StepScheduler( - global_batch_size=global_batch_size, - local_batch_size=config.step_scheduler.local_batch_size, - dp_size=self.world_size, - ckpt_every_steps=config.step_scheduler.ckpt_every_steps, - save_checkpoint_every_epoch=False, - dataloader=self.dataloader, - val_every_steps=None, - start_step=0, - start_epoch=0, - num_epochs=config.step_scheduler.num_epochs, - max_steps=config.step_scheduler.max_steps, - ) - if self.step_scheduler.grad_acc_steps != 1: - raise ValueError("PDD v1 requires exactly one microbatch per optimizer update.") - - identity = build_pdd_checkpoint_identity( - qwen_image_execution=qwen_image_execution, - metadata=self.setup_artifacts.metadata, - model_id=config.model_id, - model_revision=config.model_revision, - guidance_scale=config.pdd.guidance_scale, - ordered_train_id_sha256=train_ordered_id_sha256, - ordered_heldout_id_sha256=heldout_ordered_id_sha256, - dataset_snapshot_sha256=self.snapshot_report["dataset_snapshot_sha256"], - local_batch_size=config.step_scheduler.local_batch_size, - grad_accumulation_steps=1, - training_seed=config.seed, - validation_seed=config.validation.seed, - validation_every_steps=config.validation.every_steps, - max_grad_norm=config.training_health.max_grad_norm, - zero_grad_warmup_steps=config.training_health.zero_grad_warmup_steps, - activation_checkpointing=config.parallel.activation_checkpointing, - dtype=str(config.dtype).removeprefix("torch."), - optimizer=self.setup_artifacts.optimizer, - scheduler=self.training.scheduler, - ) - self.checkpoint_manager = PDDCheckpointManager( - root=config.checkpoint.checkpoint_dir, - checkpointer=self.setup_artifacts.checkpointer, - model=self.setup_artifacts.student, - optimizer=self.setup_artifacts.optimizer, - scheduler=self.training.scheduler, - step_scheduler=self.step_scheduler, - trainer=self.training.trainer, - sampler=self.sampler, - rng=self.training.rng, - identity=identity, - ) - self.resume = self.checkpoint_manager.load(config.checkpoint.restore_from) - self.resume_pending = self.resume is not None - self._log_setup() - - @staticmethod - def _resolve_transformer_dimensions(student: nn.Module) -> tuple[int, int]: - transformer_config = getattr(student, "config", None) - if isinstance(transformer_config, Mapping): - in_channels = transformer_config.get("in_channels") - condition_features = transformer_config.get("joint_attention_dim") - else: - in_channels = getattr(transformer_config, "in_channels", None) - condition_features = getattr(transformer_config, "joint_attention_dim", None) - if type(in_channels) is not int or in_channels <= 0 or in_channels % 4: - raise RuntimeError("constructed Qwen transformer has invalid packed in_channels.") - if type(condition_features) is not int or condition_features <= 0: - raise RuntimeError("constructed Qwen transformer has invalid joint_attention_dim.") - return in_channels // 4, condition_features - - def _log_setup(self) -> None: - if self.rank != 0: - return - if self.resume is not None: - logging.info( - "PDD resume selected: checkpoint=%s parent=%s step=%d sample_slots=%d " - "expected_first_sample_ids=%s", - self.resume.checkpoint_path, - self.resume.parent_checkpoint, - self.resume.completed_steps, - self.resume.sample_slots_consumed, - self.resume.expected_next_sample_ids, + 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 MR210 forward here changes behavior only; it creates no parameters. + adopt_qwen_image_mr210_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), + ) + self.flow_matching_pipeline = PDDFlowMatchingStepAdapter(pdd_pipeline) + logging.info("[PDD] AutoModel lifecycle enabled; grid_size=%d", self.pdd_config.grid_size) + + 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 = _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), ) - logging.info( - "PDD dataset verified: snapshot_sha256=%s metadata_sha256=%s " - "train=%d validation=%d payload_hash_verification=%s " - "payload_hashes_complete=%s root=%s", - self.snapshot_report["dataset_snapshot_sha256"], - self.snapshot_report["metadata_sha256"], - self.snapshot_report["train_samples"], - self.snapshot_report["validation_samples"], - self.snapshot_report["verify_payload_hashes"], - self.snapshot_report["payload_hashes_complete"], - self.snapshot_report["cache_root"], - ) - logging.info( - "PDD setup complete: lifecycle=%s student_keys=%d", - self.setup_artifacts.lifecycle, - len(self.setup_artifacts.checkpoint_keys), - ) - - def _prepared_training_batches(self): - iterator = _collective_training_iterator(self.dataloader, self.sampler) - while True: - next_batch = _collective_training_batch( - iterator, - sampler=self.sampler, - resume=self.resume, - resume_pending=self.resume_pending, - device=self.config.device, - dtype=self.config.dtype, - require_negative_condition=self.config.pdd.guidance_scale is not None, - expected_batch_size=self.config.step_scheduler.local_batch_size, - expected_latent_channels=self.expected_latent_channels, - expected_condition_features=self.expected_condition_features, + 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, ) - if next_batch is None: - return - yield next_batch - - def _run_validation(self, completed_step: int) -> None: - from .training import run_pdd_validation - - self.validation_sampler.set_epoch(0) - self.validation_sampler.load_state_dict({"epoch": 0, "batches_yielded": 0}) - result = run_pdd_validation( - self.training.pipeline, - _iter_validation_batches( - self.validation_dataloader, - self.validation_masks, - self.config, - self.expected_latent_channels, - self.expected_condition_features, - ), - self.validation_assignments, - validation_seed=self.config.validation.seed, - ) - if self.rank == 0: - logging.info( - "PDD validation step=%d loss=%.12g pairs=%d starts=%d heads=%d " - "ordered_id_sha256=%s records=%d", - completed_step, - result.mean_loss, - result.pair_count, - result.start_count, - result.head_count, - result.ordered_id_sha256, - len(result.records), - ) - - def _log_step(self, diagnostics: Any, data_wait_seconds: float, step_seconds: float) -> None: - timing = torch.tensor( - [data_wait_seconds, step_seconds], - dtype=torch.float64, - device=self.config.device, - ) - dist.all_reduce(timing, op=dist.ReduceOp.MAX) - peak_memory = ( - torch.cuda.max_memory_allocated(self.config.device) - if self.config.device.type == "cuda" - else 0 - ) - memory = torch.tensor(peak_memory, dtype=torch.int64, device=self.config.device) - dist.all_reduce(memory, op=dist.ReduceOp.MAX) - global_samples = self.config.step_scheduler.local_batch_size * self.world_size - throughput = global_samples / max(float(timing[1].item()), 1e-12) - coverage = self.training.trainer.coverage - bin_loss = [ - None if count == 0 else float(loss_sum / count) - for loss_sum, count in zip( - coverage.bin_loss_sums.tolist(), - coverage.bin_counts.tolist(), - ) - ] - if self.rank == 0: - logging.info( - "PDD step=%d loss=%.6g grad_norm=%.6g nominal_update_ratio=%.6g " - "projection_update_ratio=%s lr=%.6g student_rms=%.6g " - "teacher_rms=%.6g student_teacher_rms_ratio=%.6g " - "reconstruction_rms=%.6g pairs=%d n_coverage=%s k_coverage=%s " - "bins=%s bin_loss=%s samples_per_second=%.3f " - "data_wait_seconds=%.4f peak_memory_bytes=%d", - diagnostics.completed_step, - diagnostics.loss, - diagnostics.grad_norm, - diagnostics.student_adamw_nominal_update_ratio, - diagnostics.pdd_projection_update_ratio, - diagnostics.learning_rate, - diagnostics.student_velocity_rms, - diagnostics.teacher_velocity_rms, - diagnostics.student_teacher_velocity_rms_ratio, - diagnostics.reconstructed_state_rms, - int((coverage.pair_counts > 0).sum()), - _coverage_axis(coverage.n_counts, coverage.n_loss_sums), - _coverage_axis(coverage.k_counts, coverage.k_loss_sums), - coverage.bin_counts.tolist(), - bin_loss, - throughput, - float(timing[0].item()), - int(memory.item()), - ) - if self.config.device.type == "cuda": - torch.cuda.reset_peak_memory_stats(self.config.device) - - def run_train_validation_loop(self) -> None: - """Train through AutoModel StepScheduler without weakening PDD resume semantics.""" - data_wait_started = time.perf_counter() - try: - for _epoch in self.step_scheduler.epochs: - self.step_scheduler.dataloader = self._prepared_training_batches() - for batch_group in self.step_scheduler: - if len(batch_group) != 1: - raise RuntimeError("PDD v1 requires one microbatch per optimizer update.") - if self.step_scheduler.step != self.training.trainer.completed_steps: - raise RuntimeError( - "PDD trainer and AutoModel StepScheduler disagree before the update." - ) - (batch, sample_ids) = batch_group[0] - data_wait_seconds = time.perf_counter() - data_wait_started - if self.resume_pending: - if self.resume is None: - raise RuntimeError("PDD resume is pending without restored state.") - if self.rank == 0: - logging.info( - "PDD resume first batch verified: checkpoint=%s sample_ids=%s", - self.resume.checkpoint_path, - sample_ids, - ) - self.resume_pending = False - - step_started = time.perf_counter() - next_step = self.training.trainer.completed_steps + 1 - diagnostics = self.training.trainer.train_step( - batch, - measure_updates=(next_step % self.config.step_scheduler.log_every == 0), - ) - self.training.scheduler.step() - self.sampler.commit(sample_ids) - if self.sampler.remaining_batches == 0: - self.sampler.set_epoch(self.sampler.epoch + 1) - if self.training.trainer.completed_steps != self.step_scheduler.step + 1: - raise RuntimeError( - "PDD trainer and AutoModel StepScheduler disagree after the update." - ) - serialized_step = self.step_scheduler.state_dict()["step"] - if serialized_step != self.training.trainer.completed_steps: - raise RuntimeError("AutoModel StepScheduler serialized the wrong PDD step.") - step_seconds = time.perf_counter() - step_started - - completed_step = diagnostics.completed_step - is_final_step = self.step_scheduler.is_last_step - if completed_step % self.config.step_scheduler.log_every == 0: - self._log_step(diagnostics, data_wait_seconds, step_seconds) - if completed_step % self.config.validation.every_steps == 0 or is_final_step: - self._run_validation(completed_step) - if self.config.checkpoint.enabled and self.step_scheduler.is_ckpt_step: - self.checkpoint_manager.save() - data_wait_started = time.perf_counter() - finally: - self.setup_artifacts.checkpointer.close() + teacher = pipe.transformer + del pipe + adopt_qwen_image_mr210_forward(teacher) + teacher.eval().requires_grad_(False) + return teacher diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py index 51241bc9295..3ebf3b1cf7d 100644 --- a/examples/diffusers/fastgen/pdd/training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -13,985 +13,62 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Direct PDD updates and logical-ID-stable held-out validation.""" +"""PDD objective adapter for AutoModel's diffusion training loop.""" from __future__ import annotations -import hashlib -import math -from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any import torch -import torch.distributed as dist +from torch import nn -from modelopt.torch.fastgen import PDDConfig, PDDOutputProjection, PDDPipeline +if TYPE_CHECKING: + from modelopt.torch.fastgen import PDDPipeline -_TRAINER_STATE_VERSION = 1 -_VALIDATION_SCHEMA_VERSION = 1 -_VALIDATION_ORDER_DOMAIN = b"modelopt-pdd-validation-order-v1\0" -_VALIDATION_PAIR_DOMAIN = b"modelopt-pdd-validation-pair-v1\0" -_VALIDATION_NOISE_DOMAIN = b"modelopt-pdd-validation-noise-v1\0" +class PDDFlowMatchingStepAdapter: + """Expose the PDD objective through AutoModel's flow-matching ``step`` API.""" -@dataclass(frozen=True) -class PreparedPDDBatch: - """Canonical PDD inputs extracted from one Qwen cache batch.""" - - data: torch.Tensor - condition: tuple[torch.Tensor, torch.Tensor] - negative_condition: tuple[torch.Tensor, torch.Tensor] | None - sample_ids: tuple[str, ...] - valid_mask: tuple[bool, ...] | None = None - - -@dataclass(frozen=True) -class PDDStepDiagnostics: - """Host-side diagnostics for one completed direct student update.""" - - completed_step: int - loss: float - grad_norm: float - student_adamw_nominal_update_ratio: float | None - pdd_projection_update_ratio: float | None - learning_rate: float - n: tuple[int, ...] - k: tuple[int, ...] - student_velocity_rms: float - teacher_velocity_rms: float - student_teacher_velocity_rms_ratio: float - reconstructed_state_rms: float - - -@dataclass(frozen=True) -class PDDValidationAssignment: - """One canonical held-out logical ID and its explicit PDD target indices.""" - - ordinal: int - sample_id: str - n: int - k: int - - -@dataclass(frozen=True) -class PDDValidationRecord: - """Per-sample deterministic held-out result.""" - - ordinal: int - sample_id: str - n: int - k: int - loss: float - - -@dataclass(frozen=True) -class PDDValidationResult: - """Rank-invariant held-out records and canonical float64 aggregate.""" - - records: tuple[PDDValidationRecord, ...] - mean_loss: float - ordered_id_sha256: str - pair_count: int - start_count: int - head_count: int - schema_version: int = _VALIDATION_SCHEMA_VERSION - - -class PDDCoverage: - """Exact host-side n/k/pair coverage with deterministic coarse head bins.""" - - def __init__(self, config: PDDConfig, *, bins: int = 8) -> None: - if type(bins) is not int or bins <= 0: - raise ValueError("bins must be a positive integer.") - self.grid_size = config.grid_size - self.block_size_min = config.block_size_min - self.block_size_max = config.block_size_max - self.bins = min(bins, config.grid_size) - self.n_counts = torch.zeros(config.grid_size, dtype=torch.int64) - self.k_counts = torch.zeros(config.grid_size, dtype=torch.int64) - self.pair_counts = torch.zeros((config.grid_size, config.grid_size), dtype=torch.int64) - self.bin_counts = torch.zeros(self.bins, dtype=torch.int64) - self.n_loss_sums = torch.zeros(config.grid_size, dtype=torch.float64) - self.k_loss_sums = torch.zeros(config.grid_size, dtype=torch.float64) - self.pair_loss_sums = torch.zeros((config.grid_size, config.grid_size), dtype=torch.float64) - self.bin_loss_sums = torch.zeros(self.bins, dtype=torch.float64) - - def update(self, n: torch.Tensor, k: torch.Tensor, losses: torch.Tensor) -> None: - n_cpu = n.detach().to(device="cpu", dtype=torch.int64).reshape(-1) - k_cpu = k.detach().to(device="cpu", dtype=torch.int64).reshape(-1) - losses_cpu = losses.detach().to(device="cpu", dtype=torch.float64).reshape(-1) - if n_cpu.shape != k_cpu.shape or n_cpu.shape != losses_cpu.shape: - raise ValueError("n, k, and loss coverage tensors must have identical shapes.") - upper = torch.minimum( - n_cpu + self.block_size_max, - torch.full_like(n_cpu, self.grid_size), - ) - valid = ( - (n_cpu >= 0) - & (n_cpu < self.grid_size) - & (n_cpu.remainder(self.block_size_min) == 0) - & (k_cpu >= n_cpu) - & (k_cpu < upper) - ) - if not bool(valid.all()): - raise RuntimeError("observed n/k lies outside the exact configured support.") - device = n.device - n_device = n.detach().to(device=device, dtype=torch.int64).reshape(-1) - k_device = k.detach().to(device=device, dtype=torch.int64).reshape(-1) - losses_device = losses.detach().to(device=device, dtype=torch.float64).reshape(-1) - bins_device = torch.minimum( - k_device * self.bins // self.grid_size, - torch.full_like(k_device, self.bins - 1), - ) - pair_indices = n_device * self.grid_size + k_device - counts = ( - torch.bincount(n_device, minlength=self.grid_size), - torch.bincount(k_device, minlength=self.grid_size), - torch.bincount(pair_indices, minlength=self.grid_size * self.grid_size), - torch.bincount(bins_device, minlength=self.bins), - ) - loss_sums = [] - for indices, size in ( - (n_device, self.grid_size), - (k_device, self.grid_size), - (pair_indices, self.grid_size * self.grid_size), - (bins_device, self.bins), - ): - sums = torch.zeros(size, dtype=torch.float64, device=device) - sums.index_add_(0, indices, losses_device) - loss_sums.append(sums) - if dist.is_available() and dist.is_initialized(): - for tensor in (*counts, *loss_sums): - dist.all_reduce(tensor, op=dist.ReduceOp.SUM) - self.n_counts += counts[0].cpu() - self.k_counts += counts[1].cpu() - self.pair_counts += counts[2].reshape(self.grid_size, self.grid_size).cpu() - self.bin_counts += counts[3].cpu() - self.n_loss_sums += loss_sums[0].cpu() - self.k_loss_sums += loss_sums[1].cpu() - self.pair_loss_sums += loss_sums[2].reshape(self.grid_size, self.grid_size).cpu() - self.bin_loss_sums += loss_sums[3].cpu() - - def require_pairs(self, expected: Sequence[tuple[int, int]]) -> None: - missing = [pair for pair in expected if int(self.pair_counts[pair]) == 0] - if missing: - raise RuntimeError(f"targeted PDD smoke did not cover required n/k pairs: {missing}.") - - def state_dict(self) -> dict[str, Any]: - return { - "grid_size": self.grid_size, - "block_size_min": self.block_size_min, - "block_size_max": self.block_size_max, - "bins": self.bins, - "n_counts": self.n_counts.clone(), - "k_counts": self.k_counts.clone(), - "pair_counts": self.pair_counts.clone(), - "bin_counts": self.bin_counts.clone(), - "n_loss_sums": self.n_loss_sums.clone(), - "k_loss_sums": self.k_loss_sums.clone(), - "pair_loss_sums": self.pair_loss_sums.clone(), - "bin_loss_sums": self.bin_loss_sums.clone(), - } - - def load_state_dict(self, state: Mapping[str, Any]) -> None: - expected = { - "grid_size", - "block_size_min", - "block_size_max", - "bins", - "n_counts", - "k_counts", - "pair_counts", - "bin_counts", - "n_loss_sums", - "k_loss_sums", - "pair_loss_sums", - "bin_loss_sums", - } - if not isinstance(state, Mapping) or set(state) != expected: - raise ValueError("PDD coverage state has incompatible keys.") - identity = ( - state["grid_size"], - state["block_size_min"], - state["block_size_max"], - state["bins"], - ) - if identity != ( - self.grid_size, - self.block_size_min, - self.block_size_max, - self.bins, - ): - raise ValueError("PDD coverage state does not match the current configuration.") - for name in ( - "n_counts", - "k_counts", - "pair_counts", - "bin_counts", - "n_loss_sums", - "k_loss_sums", - "pair_loss_sums", - "bin_loss_sums", - ): - saved = state[name] - current = getattr(self, name) - if not isinstance(saved, torch.Tensor) or saved.shape != current.shape: - raise ValueError(f"PDD coverage {name} has an incompatible tensor shape.") - current.copy_(saved.to(device="cpu", dtype=current.dtype)) - - -def prepare_qwen_pdd_batch( - batch: Mapping[str, Any], - *, - device: torch.device, - dtype: torch.dtype, - require_negative_condition: bool, - expected_latent_channels: int, - expected_condition_features: int, -) -> PreparedPDDBatch: - """Move a portable Qwen cache batch into the PDD adapter contract.""" - if type(expected_latent_channels) is not int or expected_latent_channels <= 0: - raise ValueError("expected_latent_channels must be a positive integer.") - if type(expected_condition_features) is not int or expected_condition_features <= 0: - raise ValueError("expected_condition_features must be a positive integer.") - if not dtype.is_floating_point: - raise TypeError("Qwen PDD model dtype must be floating point.") - if not isinstance(batch, Mapping): - raise TypeError(f"batch must be a mapping, got {type(batch).__name__}.") - required = {"image_latents", "text_embeddings", "text_embeddings_mask", "metadata"} - missing = sorted(required.difference(batch)) - if missing: - raise KeyError(f"Qwen PDD batch is missing required keys: {missing}.") - data = batch["image_latents"] - text = batch["text_embeddings"] - mask = batch["text_embeddings_mask"] - metadata = batch["metadata"] - if not all(isinstance(value, torch.Tensor) for value in (data, text, mask)): - raise TypeError("Qwen PDD latent, text embedding, and mask values must be tensors.") - if data.ndim != 4: - raise ValueError(f"Qwen PDD image_latents must be 4D, got {tuple(data.shape)}.") - if not data.dtype.is_floating_point: - raise TypeError("Qwen PDD image_latents must use a floating-point dtype.") - if data.shape[0] <= 0 or data.shape[1] != expected_latent_channels: - raise ValueError( - "Qwen PDD image_latents must have a non-empty batch and exactly " - f"{expected_latent_channels} channels, got {tuple(data.shape)}." - ) - if data.shape[2] <= 0 or data.shape[3] <= 0 or data.shape[2] % 2 or data.shape[3] % 2: - raise ValueError( - "Qwen PDD image_latents must have positive even spatial dimensions, got " - f"{tuple(data.shape[2:])}." - ) - if not isinstance(metadata, Mapping): - raise TypeError("Qwen PDD batch metadata must be a mapping.") - sample_ids = metadata.get("logical_sample_ids", metadata.get("sample_ids")) - if isinstance(sample_ids, torch.Tensor): - sample_ids = tuple(str(value) for value in sample_ids.tolist()) - if isinstance(sample_ids, str) or not isinstance(sample_ids, Sequence): - raise TypeError("Qwen PDD metadata.sample_ids must be a sequence of strings.") - sample_ids = tuple(sample_ids) - if len(sample_ids) != data.shape[0] or any( - not isinstance(sample_id, str) or not sample_id for sample_id in sample_ids - ): - raise ValueError("Qwen PDD sample_ids must be non-empty strings matching batch size.") - - def prepare_condition( - embeddings: torch.Tensor, - attention_mask: torch.Tensor, - *, - name: str, - ) -> tuple[torch.Tensor, torch.Tensor]: - if not embeddings.dtype.is_floating_point: - raise TypeError(f"{name} embeddings must use a 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 boolean dtype.") - if embeddings.ndim not in (2, 3): - raise ValueError(f"{name} embeddings must be 2D or 3D, got {embeddings.ndim}D.") - if attention_mask.ndim not in (1, 2): - raise ValueError(f"{name} mask must be 1D or 2D, got {attention_mask.ndim}D.") - if embeddings.shape[-2] <= 0 or embeddings.shape[-1] <= 0: - raise ValueError(f"{name} embeddings must have non-empty sequence and feature axes.") - if embeddings.shape[-1] != expected_condition_features: - raise ValueError( - f"{name} embeddings must have exactly {expected_condition_features} features." - ) - if attention_mask.shape[-1] != embeddings.shape[-2]: - raise ValueError(f"{name} mask sequence length must match its embeddings.") - if embeddings.ndim == 3 and embeddings.shape[0] != data.shape[0]: - raise ValueError(f"{name} embedding batch size must match image_latents.") - if attention_mask.ndim == 2 and attention_mask.shape[0] != data.shape[0]: - raise ValueError(f"{name} mask batch size must match image_latents.") - - embeddings = embeddings.to(device=device, dtype=dtype) - attention_mask = attention_mask.to(device=device) - if embeddings.ndim == 2: - embeddings = embeddings.unsqueeze(0).expand(data.shape[0], -1, -1).contiguous() - if attention_mask.ndim == 1: - attention_mask = attention_mask.unsqueeze(0).expand(data.shape[0], -1).contiguous() - return embeddings, attention_mask - - data = data.to(device=device, dtype=dtype) - condition = prepare_condition(text, mask, name="Qwen PDD condition") - - negative: tuple[torch.Tensor, torch.Tensor] | None = None - negative_text = batch.get("negative_text_embeddings") - negative_mask = batch.get("negative_text_embeddings_mask") - if negative_text is not None or negative_mask is not None: - if not isinstance(negative_text, torch.Tensor) or not isinstance( - negative_mask, torch.Tensor - ): - raise TypeError("negative Qwen conditioning requires embedding and mask tensors.") - negative = prepare_condition( - negative_text, - negative_mask, - name="negative Qwen PDD condition", - ) - if require_negative_condition and negative is None: - raise ValueError("guided Qwen PDD training requires negative prompt conditioning.") - return PreparedPDDBatch(data, condition, negative, sample_ids, (True,) * len(sample_ids)) - - -def _local_tensor(value: torch.Tensor) -> torch.Tensor: - to_local = getattr(value, "to_local", None) - return to_local() if callable(to_local) else value - - -def _replication_factor(value: torch.Tensor) -> int: - placements = getattr(value, "placements", ()) - mesh = getattr(value, "device_mesh", None) - factor = 1 - if mesh is not None: - for dimension, placement in enumerate(placements): - if type(placement).__name__ == "Replicate": - factor *= int(mesh.size(dimension)) - return factor - - -def _global_squared_sum(values: Sequence[torch.Tensor], *, device: torch.device) -> torch.Tensor: - total = torch.zeros((), dtype=torch.float64, device=device) - for value in values: - local = _local_tensor(value.detach()) - contribution = local.float().square().sum(dtype=torch.float64) - total += contribution / _replication_factor(value) - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(total, op=dist.ReduceOp.SUM) - return total - - -def _global_any(flag: bool, *, device: torch.device) -> bool: - tensor = torch.tensor(int(flag), dtype=torch.int64, device=device) - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(tensor, op=dist.ReduceOp.MAX) - return bool(tensor.item()) - - -def _all_parameters_finite(parameters: Sequence[torch.Tensor], *, device: torch.device) -> bool: - local_finite = torch.ones((), dtype=torch.bool, device=device) - for parameter in parameters: - local_finite.logical_and_(torch.isfinite(_local_tensor(parameter.detach())).all()) - return not _global_any(not bool(local_finite.item()), device=device) - - -def _global_sample_mean(value: torch.Tensor) -> float: - value = value.detach().reshape(-1) - totals = torch.stack( - ( - value.double().sum(), - torch.tensor(float(value.numel()), dtype=torch.float64, device=value.device), - ) - ) - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(totals, op=dist.ReduceOp.SUM) - if totals[1] <= 0: - raise RuntimeError("cannot aggregate an empty PDD sample metric.") - return float((totals[0] / totals[1]).item()) - - -def _require_supported_adamw(optimizer: torch.optim.Optimizer) -> None: - if type(optimizer) is not torch.optim.AdamW: - raise TypeError("PDD v1 diagnostics require the stock torch.optim.AdamW optimizer.") - rejected_truthy = ("amsgrad", "maximize", "capturable", "differentiable", "foreach", "fused") - for group_index, group in enumerate(optimizer.param_groups): - for name in rejected_truthy: - if group.get(name) is not False: - raise ValueError( - f"PDD v1 requires AdamW param_groups[{group_index}][{name!r}]=False." - ) - lr = group.get("lr") - betas = group.get("betas") - eps = group.get("eps") - weight_decay = group.get("weight_decay") - if not isinstance(lr, float) or not isinstance(weight_decay, float): - raise TypeError("PDD v1 AdamW learning rate and weight decay must be scalar floats.") - if ( - not isinstance(betas, tuple) - or len(betas) != 2 - or any(not isinstance(beta, float) for beta in betas) - or not isinstance(eps, float) - ): - raise TypeError("PDD v1 AdamW betas and epsilon must be scalar floats.") - if not math.isfinite(lr) or lr <= 0 or not math.isfinite(weight_decay) or weight_decay < 0: - raise ValueError("PDD v1 AdamW learning rate/weight decay is invalid.") - if any(not math.isfinite(beta) or not 0.0 <= beta < 1.0 for beta in betas): - raise ValueError("PDD v1 AdamW betas must be finite and in [0, 1).") - if not math.isfinite(eps) or eps <= 0: - raise ValueError("PDD v1 AdamW epsilon must be finite and > 0.") - - -def _adamw_nominal_update_ratio( - optimizer: torch.optim.Optimizer, - *, - device: torch.device, -) -> float: - """Stream the public AdamW equation over local shards without cloning the model.""" - update_squared = torch.zeros((), dtype=torch.float64, device=device) - parameter_squared = torch.zeros((), dtype=torch.float64, device=device) - for group in optimizer.param_groups: - lr = group["lr"] - beta1, beta2 = group["betas"] - eps = group["eps"] - decay = 1.0 - lr * group["weight_decay"] - if decay <= 0.0: - raise RuntimeError("AdamW decoupled weight decay factor must remain positive.") - for parameter in group["params"]: - if parameter.grad is None: - continue - state = optimizer.state.get(parameter) - if not state or "exp_avg" not in state or "exp_avg_sq" not in state: - continue - step_value = state.get("step") - if isinstance(step_value, torch.Tensor): - step = float(step_value.item()) - else: - step = float(step_value) - if step <= 0: - raise RuntimeError("AdamW state has a non-positive step after optimizer.step().") - parameter_local = _local_tensor(parameter.detach()).float() - exp_avg = _local_tensor(state["exp_avg"].detach()).float() - exp_avg_sq = _local_tensor(state["exp_avg_sq"].detach()).float() - bias_correction1 = 1.0 - beta1**step - bias_correction2_sqrt = math.sqrt(1.0 - beta2**step) - denominator = exp_avg_sq.sqrt().div_(bias_correction2_sqrt).add_(eps) - direction = exp_avg.div(denominator).div_(bias_correction1) - parameter_before = (parameter_local + lr * direction) / decay - nominal_delta = parameter_local - parameter_before - factor = _replication_factor(parameter) - update_squared += nominal_delta.square().sum(dtype=torch.float64) / factor - parameter_squared += parameter_before.square().sum(dtype=torch.float64) / factor - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(update_squared, op=dist.ReduceOp.SUM) - dist.all_reduce(parameter_squared, op=dist.ReduceOp.SUM) - if not bool(torch.isfinite(update_squared) & torch.isfinite(parameter_squared)): - raise RuntimeError("PDD AdamW nominal update diagnostics became non-finite.") - return float((update_squared.sqrt() / parameter_squared.sqrt().clamp_min(1e-30)).item()) - - -class PDDTrainer: - """Own direct student updates while leaving algorithm and checkpoint state separate.""" - - def __init__( - self, - pipeline: PDDPipeline, - optimizer: torch.optim.Optimizer, - *, - projection: PDDOutputProjection, - max_grad_norm: float, - warmup_steps: int = 0, - ) -> None: - if not isinstance(pipeline, PDDPipeline): - raise TypeError(f"pipeline must be PDDPipeline, got {type(pipeline).__name__}.") - if not isinstance(projection, PDDOutputProjection): - raise TypeError("projection must be PDDOutputProjection.") - if isinstance(max_grad_norm, bool) or not isinstance(max_grad_norm, int | float): - raise TypeError("max_grad_norm must be a real number.") - if not math.isfinite(max_grad_norm) or max_grad_norm <= 0: - raise ValueError("max_grad_norm must be finite and > 0.") - if type(warmup_steps) is not int or warmup_steps < 0: - raise ValueError("warmup_steps must be an integer >= 0.") - _require_supported_adamw(optimizer) + def __init__(self, pipeline: PDDPipeline) -> None: self.pipeline = pipeline - self.optimizer = optimizer - self.projection = projection - self.max_grad_norm = float(max_grad_norm) - self.warmup_steps = warmup_steps - self.completed_steps = 0 - self.consecutive_zero_grad_steps = 0 - self.coverage = PDDCoverage(pipeline.config) - def _projection_snapshot(self) -> list[torch.Tensor]: - parameters = [self.projection.weight] - if self.projection.bias is not None: - parameters.append(self.projection.bias) - return [_local_tensor(parameter.detach()).clone() for parameter in parameters] - - def _projection_update_ratio(self, before: Sequence[torch.Tensor]) -> float: - parameters = [self.projection.weight] - if self.projection.bias is not None: - parameters.append(self.projection.bias) - update_squared = torch.zeros((), dtype=torch.float64, device=self.pipeline.device) - parameter_squared = torch.zeros((), dtype=torch.float64, device=self.pipeline.device) - for parameter, saved in zip(parameters, before): - local = _local_tensor(parameter.detach()).float() - factor = _replication_factor(parameter) - update_squared += (local - saved.float()).square().sum(dtype=torch.float64) / factor - parameter_squared += local.square().sum(dtype=torch.float64) / factor - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(update_squared, op=dist.ReduceOp.SUM) - dist.all_reduce(parameter_squared, op=dist.ReduceOp.SUM) - if not bool(torch.isfinite(update_squared) & torch.isfinite(parameter_squared)): - raise RuntimeError("PDD projection update diagnostics became non-finite.") - return float((update_squared.sqrt() / parameter_squared.sqrt().clamp_min(1e-30)).item()) - - def train_step( + def step( self, - batch: PreparedPDDBatch, - *, - noise: torch.Tensor | None = None, - n: torch.Tensor | None = None, - k: torch.Tensor | None = None, - generator: torch.Generator | None = None, - measure_updates: bool = True, - ) -> PDDStepDiagnostics: - """Run one direct PDD student update and enforce all immediate hard aborts.""" - if not isinstance(batch, PreparedPDDBatch): - raise TypeError("batch must be PreparedPDDBatch.") - self.pipeline.student.train() - self.pipeline.teacher.eval() - self.optimizer.zero_grad(set_to_none=True) - before_projection = self._projection_snapshot() if measure_updates else None - loss, metrics = self.pipeline.compute_loss( - batch.data, - noise=noise, - condition=batch.condition, - negative_condition=batch.negative_condition, - n=n, - k=k, - generator=generator, - ) - finite_metrics = ( - "all_student_heads_finite", - "student_target_finite", - "teacher_target_finite", - "reconstructed_state_finite", - "loss_finite", - ) - local_nonfinite = not bool(torch.isfinite(loss)) or any( - not bool(metrics[name].all()) for name in finite_metrics - ) - if _global_any(local_nonfinite, device=batch.data.device): - raise FloatingPointError( - "PDD loss, prediction, target, or reconstruction is non-finite." - ) - self.coverage.update(metrics["n"], metrics["k"], metrics["student_target_mse"]) - loss.backward() - - teacher_gradient = any( - parameter.grad is not None for parameter in self.pipeline.teacher.parameters() - ) - if _global_any(teacher_gradient, device=batch.data.device): - raise RuntimeError("PDD frozen teacher received a gradient.") - trainable = [ - parameter for parameter in self.pipeline.student.parameters() if parameter.requires_grad - ] - gradients = [parameter.grad for parameter in trainable if parameter.grad is not None] - if not gradients: - grad_norm = 0.0 - else: - grad_squared = _global_squared_sum(gradients, device=batch.data.device) - if not bool(torch.isfinite(grad_squared)): - raise FloatingPointError("PDD student gradient became non-finite.") - grad_norm = float(grad_squared.sqrt().item()) - if self.completed_steps >= self.warmup_steps and grad_norm == 0.0: - self.consecutive_zero_grad_steps += 1 - if self.consecutive_zero_grad_steps >= 2: - raise RuntimeError("PDD student gradient was zero for two consecutive updates.") - else: - self.consecutive_zero_grad_steps = 0 - clip_coefficient = min(1.0, self.max_grad_norm / (grad_norm + 1e-6)) - if clip_coefficient < 1.0: - for gradient in gradients: - gradient.mul_(clip_coefficient) - - self.optimizer.step() - if not _all_parameters_finite(trainable, device=batch.data.device): - raise FloatingPointError("PDD student parameter update became non-finite.") - nominal_ratio = ( - _adamw_nominal_update_ratio( - self.optimizer, - device=batch.data.device, - ) - if measure_updates - else None - ) - projection_ratio = ( - self._projection_update_ratio(before_projection) - if before_projection is not None - else None - ) - if nominal_ratio == 0.0 and grad_norm > 0.0: - raise RuntimeError( - "PDD optimizer produced a zero nominal update from a nonzero gradient." - ) - if projection_ratio == 0.0 and grad_norm > 0.0: - raise RuntimeError( - "PDD optimizer produced a zero actual projection update from a nonzero gradient." + 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, global_step + + data = batch["image_latents"].to(device=device, dtype=dtype, non_blocking=True) + condition = ( + batch["text_embeddings"].to(device=device, dtype=dtype, non_blocking=True), + batch["text_embeddings_mask"].to(device=device, non_blocking=True), + ) + negative_condition = None + if self.pipeline.config.guidance_scale is not None: + negative_condition = ( + batch["negative_text_embeddings"].to( + device=device, + dtype=dtype, + non_blocking=True, + ), + batch["negative_text_embeddings_mask"].to(device=device, non_blocking=True), ) - self.completed_steps += 1 - student_velocity_rms = _global_sample_mean(metrics["student_velocity_rms"]) - teacher_velocity_rms = _global_sample_mean(metrics["teacher_velocity_rms"]) - return PDDStepDiagnostics( - completed_step=self.completed_steps, - loss=_global_sample_mean(metrics["student_target_mse"]), - grad_norm=grad_norm, - student_adamw_nominal_update_ratio=nominal_ratio, - pdd_projection_update_ratio=projection_ratio, - learning_rate=float(self.optimizer.param_groups[0]["lr"]), - n=tuple(int(value) for value in metrics["n"].detach().cpu().tolist()), - k=tuple(int(value) for value in metrics["k"].detach().cpu().tolist()), - student_velocity_rms=student_velocity_rms, - teacher_velocity_rms=teacher_velocity_rms, - student_teacher_velocity_rms_ratio=student_velocity_rms - / max(teacher_velocity_rms, 1e-30), - reconstructed_state_rms=_global_sample_mean(metrics["reconstructed_state_rms"]), - ) - - def state_dict(self) -> dict[str, Any]: - return { - "schema_version": _TRAINER_STATE_VERSION, - "completed_steps": self.completed_steps, - "consecutive_zero_grad_steps": self.consecutive_zero_grad_steps, - "coverage": self.coverage.state_dict(), - } - - def load_state_dict(self, state: Mapping[str, Any]) -> None: - expected = { - "schema_version", - "completed_steps", - "consecutive_zero_grad_steps", - "coverage", - } - if not isinstance(state, Mapping) or set(state) != expected: - raise ValueError("PDD trainer state has incompatible keys.") - if state["schema_version"] != _TRAINER_STATE_VERSION: - raise ValueError(f"unsupported PDD trainer schema {state['schema_version']!r}.") - completed = state["completed_steps"] - zero_steps = state["consecutive_zero_grad_steps"] - if type(completed) is not int or completed < 0: - raise ValueError("saved completed_steps must be an integer >= 0.") - if type(zero_steps) is not int or zero_steps < 0: - raise ValueError("saved consecutive_zero_grad_steps must be an integer >= 0.") - self.completed_steps = completed - self.consecutive_zero_grad_steps = zero_steps - self.coverage.load_state_dict(state["coverage"]) - - -def _stable_digest(domain: bytes, validation_seed: int, payload: str) -> bytes: - digest = hashlib.sha256() - digest.update(domain) - digest.update(str(validation_seed).encode()) - digest.update(b"\0") - digest.update(payload.encode()) - return digest.digest() - - -def pdd_validation_support(config: PDDConfig) -> tuple[tuple[int, int], ...]: - """Return the exact lexicographic support of explicit PDD validation pairs.""" - return tuple( - (n, k) - for n in range(0, config.grid_size, config.block_size_min) - for k in range(n, min(n + config.block_size_max, config.grid_size)) - ) - - -def build_pdd_validation_assignments( - sample_ids: Sequence[str], - config: PDDConfig, - *, - validation_seed: int, - require_full_coverage: bool = True, -) -> tuple[PDDValidationAssignment, ...]: - """Assign every logical ID a rank/batch-order-independent explicit n/k pair.""" - if isinstance(sample_ids, str) or not isinstance(sample_ids, Sequence): - raise TypeError("sample_ids must be a sequence of strings.") - if any(not isinstance(sample_id, str) or not sample_id for sample_id in sample_ids): - raise ValueError("sample_ids must contain non-empty strings.") - if not sample_ids: - raise ValueError("sample_ids must contain at least one logical ID.") - if len(set(sample_ids)) != len(sample_ids): - raise ValueError("held-out validation sample_ids must be unique.") - if type(validation_seed) is not int or validation_seed < 0: - raise ValueError("validation_seed must be an integer >= 0.") - support = pdd_validation_support(config) - if require_full_coverage and len(sample_ids) < len(support): - raise ValueError( - f"full PDD validation coverage requires at least {len(support)} logical IDs, " - f"found {len(sample_ids)}." - ) - ordered_ids = sorted( - sample_ids, - key=lambda sample_id: ( - _stable_digest(_VALIDATION_ORDER_DOMAIN, validation_seed, sample_id), - sample_id, - ), - ) - permuted_support = sorted( - support, - key=lambda pair: ( - _stable_digest(_VALIDATION_PAIR_DOMAIN, validation_seed, f"{pair[0]}:{pair[1]}"), - pair, - ), - ) - return tuple( - PDDValidationAssignment(ordinal, sample_id, *permuted_support[ordinal % len(support)]) - for ordinal, sample_id in enumerate(ordered_ids) - ) - - -def pdd_validation_noise( - sample_id: str, - shape: Sequence[int], - *, - validation_seed: int, - device: torch.device, -) -> torch.Tensor: - """Generate per-ID CPU float32 noise without touching the global RNG.""" - digest = _stable_digest(_VALIDATION_NOISE_DOMAIN, validation_seed, sample_id) - seed = int.from_bytes(digest[:8], "big") & ((1 << 63) - 1) - generator = torch.Generator(device="cpu") - generator.manual_seed(seed) - return torch.randn(tuple(shape), generator=generator, dtype=torch.float32).to(device) - - -def _ordered_id_digest(records: Sequence[PDDValidationRecord]) -> str: - digest = hashlib.sha256() - digest.update(b"modelopt-pdd-ordered-validation-ids-v1\0") - for record in records: - digest.update(record.sample_id.encode()) - digest.update(b"\n") - return digest.hexdigest() - - -def _raise_collective_validation_error(error: BaseException | None, *, context: str) -> None: - if not dist.is_available() or not dist.is_initialized(): - if error is not None: - raise error - return - local = None if error is None else f"{type(error).__name__}: {error}" - errors: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(errors, local) - failures = [f"rank {rank}: {message}" for rank, message in enumerate(errors) if message] - if failures: - raise RuntimeError(f"distributed PDD validation {context} failed; " + "; ".join(failures)) - - -def _reshard_fsdp2_modules(model: torch.nn.Module) -> None: - from torch.distributed.fsdp import FSDPModule - - for module in model.modules(): - if isinstance(module, FSDPModule): - module.reshard() - - -def run_pdd_validation( - pipeline: PDDPipeline, - batches: Iterable[PreparedPDDBatch], - assignments: Sequence[PDDValidationAssignment], - *, - validation_seed: int, -) -> PDDValidationResult: - """Evaluate explicit per-ID targets and aggregate identically across rank partitions.""" - if not isinstance(pipeline, PDDPipeline): - raise TypeError("pipeline must be PDDPipeline.") - distributed = dist.is_available() and dist.is_initialized() - assignment_by_id = {assignment.sample_id: assignment for assignment in assignments} - assignment_error: BaseException | None = None - if len(assignment_by_id) != len(assignments): - assignment_error = ValueError("validation assignments contain duplicate logical IDs.") - elif not assignments: - assignment_error = ValueError("validation assignments cannot be empty.") - _raise_collective_validation_error(assignment_error, context="assignment preflight") - if distributed: - assignment_identity = tuple( - (item.ordinal, item.sample_id, item.n, item.k) for item in assignments + loss, metrics = self.pipeline.compute_loss( + data, + condition=condition, + negative_condition=negative_condition, + collect_metrics=collect_metrics, ) - assignment_identities: list[Any] = [None] * dist.get_world_size() - dist.all_gather_object(assignment_identities, assignment_identity) - if any(identity != assignment_identities[0] for identity in assignment_identities[1:]): - raise RuntimeError("distributed PDD validation assignments differ across ranks.") - student_was_training = pipeline.student.training - teacher_was_training = pipeline.teacher.training - pipeline.student.eval() - pipeline.teacher.eval() - local_records: list[PDDValidationRecord] = [] - try: - with torch.no_grad(): - iterator = iter(batches) - batch_index = 0 - while True: - batch = None - next_error: BaseException | None = None - exhausted = False - try: - batch = next(iterator) - except StopIteration: - exhausted = True - except BaseException as error: - next_error = error - if distributed: - status = "error" if next_error is not None else "end" if exhausted else "batch" - statuses: list[str] = [""] * dist.get_world_size() - dist.all_gather_object(statuses, status) - if "error" in statuses: - _raise_collective_validation_error(next_error, context="iteration") - if all(item == "end" for item in statuses): - break - if any(item != "batch" for item in statuses): - raise RuntimeError( - "distributed PDD validation ranks produced different batch counts." - ) - else: - if next_error is not None: - raise next_error - if exhausted: - break + if check_loss and not bool(torch.isfinite(loss)): + raise FloatingPointError("PDD loss is non-finite.") - local_error: BaseException | None = None - selected: list[PDDValidationAssignment] = [] - valid_mask: tuple[bool, ...] = () - try: - if not isinstance(batch, PreparedPDDBatch): - raise TypeError("validation batches must contain PreparedPDDBatch values.") - valid_mask = ( - (True,) * len(batch.sample_ids) - if batch.valid_mask is None - else batch.valid_mask - ) - if len(valid_mask) != len(batch.sample_ids) or any( - type(valid) is not bool for valid in valid_mask - ): - raise ValueError( - "validation valid_mask must contain one bool per sample ID." - ) - for position, (sample_id, valid) in enumerate( - zip(batch.sample_ids, valid_mask) - ): - if valid and sample_id not in assignment_by_id: - raise ValueError( - f"validation batch contains unassigned sample ID {sample_id!r}." - ) - if valid: - selected.append(assignment_by_id[sample_id]) - else: - template = assignments[(batch_index + position) % len(assignments)] - selected.append( - PDDValidationAssignment( - template.ordinal, - f"__pdd_dummy__:{batch_index}:{position}", - template.n, - template.k, - ) - ) - except BaseException as error: - local_error = error - _raise_collective_validation_error(local_error, context="batch preflight") - assert isinstance(batch, PreparedPDDBatch) - # Inputs remain rank-local. Equal batch counts above preserve collective - # ordering, while prompt sequence padding may legitimately differ by rank. - noise = torch.stack( - [ - pdd_validation_noise( - assignment.sample_id, - batch.data.shape[1:], - validation_seed=validation_seed, - device=batch.data.device, - ) - for assignment in selected - ] - ) - n = torch.tensor( - [assignment.n for assignment in selected], - dtype=torch.long, - device=batch.data.device, - ) - k = torch.tensor( - [assignment.k for assignment in selected], - dtype=torch.long, - device=batch.data.device, - ) - _, metrics = pipeline.compute_loss( - batch.data, - noise=noise, - condition=batch.condition, - negative_condition=batch.negative_condition, - n=n, - k=k, - ) - finite_metrics = ( - "all_student_heads_finite", - "student_target_finite", - "teacher_target_finite", - "reconstructed_state_finite", - "loss_finite", - ) - local_nonfinite = any(not bool(metrics[name].all()) for name in finite_metrics) - losses = metrics["student_target_mse"].double().cpu().tolist() - local_nonfinite = local_nonfinite or any(not math.isfinite(loss) for loss in losses) - if _global_any(local_nonfinite, device=batch.data.device): - raise FloatingPointError( - "deterministic PDD validation produced a non-finite prediction, target, " - "or loss." - ) - local_records.extend( - PDDValidationRecord( - assignment.ordinal, - assignment.sample_id, - assignment.n, - assignment.k, - float(loss), - ) - for assignment, loss, valid in zip(selected, losses, valid_mask) - if valid - ) - batch_index += 1 - finally: - pipeline.student.train(student_was_training) - pipeline.teacher.train(teacher_was_training) - _reshard_fsdp2_modules(pipeline.student) - _reshard_fsdp2_modules(pipeline.teacher) - - gathered: list[list[PDDValidationRecord]] - if distributed: - gathered = [[] for _ in range(dist.get_world_size())] - dist.all_gather_object(gathered, local_records) - else: - gathered = [local_records] - records = sorted( - (record for rank_records in gathered for record in rank_records), - key=lambda record: record.ordinal, - ) - if len(records) != len(assignments): - raise RuntimeError( - f"validation contributed {len(records)} records for {len(assignments)} assignments." - ) - if [record.ordinal for record in records] != list(range(len(assignments))): - raise RuntimeError("validation ordinals are duplicated or missing across ranks.") - for record, assignment in zip(records, assignments): - if (record.sample_id, record.n, record.k) != ( - assignment.sample_id, - assignment.n, - assignment.k, - ): - raise RuntimeError("validation record does not match its canonical assignment.") - pairs = {(record.n, record.k) for record in records} - starts = {record.n for record in records} - heads = {record.k for record in records} - return PDDValidationResult( - records=tuple(records), - mean_loss=math.fsum(record.loss for record in records) / len(records), - ordered_id_sha256=_ordered_id_digest(records), - pair_count=len(pairs), - start_count=len(starts), - head_count=len(heads), - ) + 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/preprocessing_multiprocess.py b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py index bb933479fee..c08f731cbe2 100644 --- a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py +++ b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py @@ -95,14 +95,6 @@ def _get_media_files(media_dir: Path, extensions: set) -> list[Path]: return sorted(media_files) -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - def _save_metadata_shards( all_metadata: list[dict], output_dir: Path, @@ -136,7 +128,6 @@ def _save_metadata_shards( { **item, "cache_file": str(cache_file), - "cache_sha256": _sha256_file(cache_file), } ) diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index 93b73db4528..fab7d09b206 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -43,7 +43,6 @@ __all__ = [ "PDDLayerSpec", - "PDDMetadata", "PDDModelAdapter", "PDDOutputProjection", "PDDPipeline", @@ -54,17 +53,6 @@ PDDHeadLayout = Literal["channel_major", "patch_major"] _HEAD_LAYOUTS = ("channel_major", "patch_major") -_METADATA_SCHEMA_VERSION = 1 - - -def _require_exact_keys(mapping: Mapping[str, Any], expected: set[str], *, name: str) -> None: - if any(not isinstance(key, str) for key in mapping): - raise ValueError(f"{name} keys must all be strings.") - actual = set(mapping) - if actual != expected: - missing = sorted(expected - actual) - extra = sorted(actual - expected) - raise ValueError(f"{name} keys mismatch: missing={missing}, extra={extra}.") def _require_int(value: Any, *, name: str, minimum: int = 1) -> int: @@ -103,222 +91,6 @@ def __post_init__(self) -> None: else: _require_int(self.output_channels, name="output_channels") - def to_dict(self) -> dict[str, Any]: - """Serialize to a strict primitive mapping.""" - return { - "projection_path": self.projection_path, - "head_layout": self.head_layout, - "output_channels": self.output_channels, - } - - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> PDDLayerSpec: - """Deserialize a strict primitive mapping.""" - if not isinstance(data, Mapping): - raise TypeError(f"layer_spec must be a mapping, got {type(data).__name__}.") - _require_exact_keys( - data, - {"projection_path", "head_layout", "output_channels"}, - name="layer_spec", - ) - if not isinstance(data["projection_path"], str): - raise ValueError("layer_spec.projection_path must be a string.") - raw_head_layout = data["head_layout"] - if raw_head_layout == "channel_major": - head_layout: PDDHeadLayout = "channel_major" - elif raw_head_layout == "patch_major": - head_layout = "patch_major" - else: - raise ValueError(f"layer_spec.head_layout must be one of {_HEAD_LAYOUTS}.") - output_channels = data["output_channels"] - if output_channels is not None and type(output_channels) is not int: - raise ValueError("layer_spec.output_channels must be an integer or null.") - return cls( - projection_path=data["projection_path"], - head_layout=head_layout, - output_channels=output_channels, - ) - - -@dataclass(frozen=True) -class PDDMetadata: - """Versioned minimum metadata required to reconstruct a PDD projection.""" - - grid_size: int - grid_max_t: float - flow_shift: float - block_size_min: int - block_size_max: int - inference_blocks: tuple[int, ...] - teacher_integrator: Literal["euler", "midpoint"] - layer_spec: PDDLayerSpec - projection_in_features: int - projection_out_features: int - projection_bias: bool - schema_version: int = _METADATA_SCHEMA_VERSION - - def __post_init__(self) -> None: - _require_int(self.schema_version, name="schema_version") - if self.schema_version != _METADATA_SCHEMA_VERSION: - raise ValueError( - f"unsupported PDD metadata schema_version={self.schema_version}; " - f"expected {_METADATA_SCHEMA_VERSION}." - ) - _require_int(self.grid_size, name="grid_size") - if type(self.grid_max_t) is not float: - raise ValueError(f"grid_max_t must be a float, got {self.grid_max_t!r}.") - if type(self.flow_shift) is not float: - raise ValueError(f"flow_shift must be a float, got {self.flow_shift!r}.") - _require_int(self.block_size_min, name="block_size_min") - _require_int(self.block_size_max, name="block_size_max") - if not isinstance(self.inference_blocks, tuple) or any( - type(block) is not int for block in self.inference_blocks - ): - raise ValueError("inference_blocks must be a tuple of integers.") - if self.teacher_integrator not in ("euler", "midpoint"): - raise ValueError( - "teacher_integrator must be either 'euler' or 'midpoint', got " - f"{self.teacher_integrator!r}." - ) - _require_int(self.projection_in_features, name="projection_in_features") - _require_int(self.projection_out_features, name="projection_out_features") - if type(self.projection_bias) is not bool: - raise ValueError(f"projection_bias must be bool, got {self.projection_bias!r}.") - if not isinstance(self.layer_spec, PDDLayerSpec): - raise TypeError( - f"layer_spec must be PDDLayerSpec, got {type(self.layer_spec).__name__}." - ) - if self.layer_spec.head_layout == "patch_major": - output_channels = self.layer_spec.output_channels - if output_channels is None or self.projection_out_features % output_channels != 0: - raise ValueError( - f"projection_out_features={self.projection_out_features} must be divisible by " - f"output_channels={output_channels} for patch_major layout." - ) - - PDDConfig( - grid_size=self.grid_size, - grid_max_t=self.grid_max_t, - flow_shift=self.flow_shift, - block_size_min=self.block_size_min, - block_size_max=self.block_size_max, - inference_blocks=list(self.inference_blocks), - student_sample_steps=len(self.inference_blocks), - teacher_integrator=self.teacher_integrator, - ) - - @classmethod - def from_config(cls, config: PDDConfig, projection: PDDOutputProjection) -> PDDMetadata: - """Build reconstruction metadata from a validated config and projection.""" - if not isinstance(config, PDDConfig): - raise TypeError(f"config must be PDDConfig, got {type(config).__name__}.") - if not isinstance(projection, PDDOutputProjection): - raise TypeError( - f"projection must be PDDOutputProjection, got {type(projection).__name__}." - ) - if config.grid_size != projection.grid_size: - raise ValueError( - f"config grid_size={config.grid_size} does not match projection " - f"grid_size={projection.grid_size}." - ) - return cls( - grid_size=config.grid_size, - grid_max_t=config.grid_max_t, - flow_shift=config.flow_shift, - block_size_min=config.block_size_min, - block_size_max=config.block_size_max, - inference_blocks=tuple(config.inference_blocks), - teacher_integrator=config.teacher_integrator, - layer_spec=projection.layer_spec, - projection_in_features=projection.in_features, - projection_out_features=projection.base_out_features, - projection_bias=projection.bias is not None, - ) - - def to_dict(self) -> dict[str, Any]: - """Serialize to a strict, JSON/YAML-safe mapping.""" - return { - "schema_version": self.schema_version, - "grid_size": self.grid_size, - "grid_max_t": self.grid_max_t, - "flow_shift": self.flow_shift, - "block_size_min": self.block_size_min, - "block_size_max": self.block_size_max, - "inference_blocks": list(self.inference_blocks), - "teacher_integrator": self.teacher_integrator, - "layer_spec": self.layer_spec.to_dict(), - "base_projection": { - "in_features": self.projection_in_features, - "out_features": self.projection_out_features, - "bias": self.projection_bias, - }, - } - - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> PDDMetadata: - """Deserialize a strict, versioned metadata mapping.""" - if not isinstance(data, Mapping): - raise TypeError(f"PDD metadata must be a mapping, got {type(data).__name__}.") - _require_exact_keys( - data, - { - "schema_version", - "grid_size", - "grid_max_t", - "flow_shift", - "block_size_min", - "block_size_max", - "inference_blocks", - "teacher_integrator", - "layer_spec", - "base_projection", - }, - name="PDD metadata", - ) - schema_version = _require_int(data["schema_version"], name="schema_version") - if type(data["grid_max_t"]) is not float: - raise ValueError(f"grid_max_t must be a float, got {data['grid_max_t']!r}.") - if type(data["flow_shift"]) is not float: - raise ValueError(f"flow_shift must be a float, got {data['flow_shift']!r}.") - if not isinstance(data["inference_blocks"], list) or any( - type(block) is not int for block in data["inference_blocks"] - ): - raise ValueError("inference_blocks must be a list of integers.") - if data["teacher_integrator"] not in ("euler", "midpoint"): - raise ValueError( - "teacher_integrator must be either 'euler' or 'midpoint', got " - f"{data['teacher_integrator']!r}." - ) - base_projection = data["base_projection"] - if not isinstance(base_projection, Mapping): - raise TypeError("base_projection must be a mapping.") - _require_exact_keys( - base_projection, - {"in_features", "out_features", "bias"}, - name="base_projection", - ) - if type(base_projection["bias"]) is not bool: - raise ValueError("base_projection.bias must be bool.") - - return cls( - schema_version=schema_version, - grid_size=_require_int(data["grid_size"], name="grid_size"), - grid_max_t=data["grid_max_t"], - flow_shift=data["flow_shift"], - block_size_min=_require_int(data["block_size_min"], name="block_size_min"), - block_size_max=_require_int(data["block_size_max"], name="block_size_max"), - inference_blocks=tuple(data["inference_blocks"]), - teacher_integrator=data["teacher_integrator"], - layer_spec=PDDLayerSpec.from_dict(data["layer_spec"]), - projection_in_features=_require_int( - base_projection["in_features"], name="base_projection.in_features" - ), - projection_out_features=_require_int( - base_projection["out_features"], name="base_projection.out_features" - ), - projection_bias=base_projection["bias"], - ) - @dataclass(frozen=True) class _FusionRequest: @@ -806,8 +578,11 @@ def compute_loss( 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( @@ -903,28 +678,33 @@ def compute_loss( squared_error = (student_target - teacher_target).square() loss = squared_error.mean() metric_dims = tuple(range(1, squared_error.ndim)) - all_head_dims = tuple(range(1, student_heads.ndim)) with torch.no_grad(): metrics = { - "n": n.detach(), - "k": k.detach(), - "target_span": (k - n).detach(), "student_target_mse": squared_error.mean(dim=metric_dims).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(), } + if collect_metrics: + all_head_dims = tuple(range(1, student_heads.ndim)) + metrics.update( + n=n.detach(), + k=k.detach(), + target_span=(k - n).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 def _validate_blocks(self, blocks: Sequence[int] | None) -> tuple[int, ...]: diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index 7b9caf59b93..8365b976943 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -36,6 +36,7 @@ "adopt_qwen_image_mr210_forward", "convert_qwen_image_to_pdd", "require_qwen_image_mr210_forward", + "restore_qwen_image_pdd_projection", ] QWEN_IMAGE_PDD_EXECUTION = "fastgen_mr210" @@ -172,6 +173,16 @@ def _qwen_image_mr210_forward( max_txt_seq_len=max_txt_seq_len, device=hidden_states.device, ) + image_mask = torch.ones( + (batch_size, hidden_states.shape[1]), + dtype=torch.bool, + device=hidden_states.device, + ) + joint_attention_mask = torch.cat( + (encoder_hidden_states_mask.to(torch.bool), image_mask), + dim=1, + )[:, None, None, :] + block_attention_kwargs = {"attention_mask": joint_attention_mask} for block in self.transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: @@ -179,18 +190,19 @@ def _qwen_image_mr210_forward( block, hidden_states, encoder_hidden_states, - encoder_hidden_states_mask, + None, temb, image_rotary_emb, + block_attention_kwargs, ) 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, + encoder_hidden_states_mask=None, temb=temb, image_rotary_emb=image_rotary_emb, - joint_attention_kwargs=attention_kwargs, + joint_attention_kwargs=block_attention_kwargs, ) hidden_states = self.norm_out(hidden_states, temb) @@ -230,7 +242,7 @@ def adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: # Diffusers is an optional dependency used only by the Qwen example. from diffusers import QwenImageTransformer2DModel - if type(transformer) is not QwenImageTransformer2DModel: + if not isinstance(transformer, QwenImageTransformer2DModel): raise TypeError( "MR210 forward adoption requires the supported QwenImageTransformer2DModel, " f"got {type(transformer).__name__}." @@ -312,6 +324,59 @@ def convert_qwen_image_to_pdd( 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.") + try: + current = transformer.get_submodule("proj_out") + except AttributeError as error: + raise ValueError( + "Qwen-Image transformer must register an nn.Linear at 'proj_out'." + ) from error + if not isinstance(current, nn.Linear): + raise TypeError(f"Qwen-Image proj_out must be nn.Linear, got {type(current).__name__}.") + 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 + if transformer.get_submodule("proj_out") is not projection: + raise RuntimeError("Qwen-Image proj_out replacement did not remain registered.") + return projection + + class QwenImagePDDAdapter: """Adapt raw Qwen packed-token calls to the framework-neutral PDD protocol.""" @@ -532,24 +597,52 @@ def _unpack_single(self, packed: torch.Tensor, state: torch.Tensor) -> torch.Ten return unpack_latents(packed, state.shape[2], state.shape[3]) @staticmethod - def _projection(model: nn.Module, grid_size: int) -> PDDOutputProjection: + def _all_heads_projection( + model: nn.Module, + grid_size: int, + base_out_features: int, + ) -> nn.Linear: try: projection = model.get_submodule("proj_out") except AttributeError as error: raise ValueError( - "Qwen student must register a PDD projection at 'proj_out'." + "Qwen student must register an output projection at 'proj_out'." ) from error - if not isinstance(projection, PDDOutputProjection): - raise TypeError( - "Qwen student proj_out must be converted to PDDOutputProjection before use." + if not isinstance(projection, nn.Linear): + raise TypeError("Qwen student proj_out must be a widened nn.Linear for PDD training.") + expected_out_features = grid_size * base_out_features + if projection.out_features != expected_out_features: + raise ValueError( + f"Qwen PDD proj_out has {projection.out_features} outputs; expected " + f"{expected_out_features} ({grid_size} heads x {base_out_features})." ) - if projection.grid_size != grid_size: + if isinstance(projection, PDDOutputProjection): + if projection.grid_size != grid_size: + raise ValueError( + f"Qwen PDD projection grid_size={projection.grid_size} does not match " + f"config grid_size={grid_size}." + ) + if projection.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: + raise ValueError("Qwen PDD projection carries an incompatible layer specification.") + return projection + + @classmethod + def _fused_projection(cls, model: nn.Module, grid_size: int) -> PDDOutputProjection: + try: + projection = model.get_submodule("proj_out") + except AttributeError as error: raise ValueError( - f"Qwen PDD projection grid_size={projection.grid_size} does not match " - f"config grid_size={grid_size}." + "Qwen student must register an output projection at 'proj_out'." + ) from error + if not isinstance(projection, PDDOutputProjection): + raise TypeError( + "Qwen fused PDD inference requires proj_out to be a PDDOutputProjection." ) - if projection.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: - raise ValueError("Qwen PDD projection carries an incompatible layer specification.") + cls._all_heads_projection( + model, + grid_size, + base_out_features=projection.base_out_features, + ) return projection def student_all_heads( @@ -562,7 +655,11 @@ def student_all_heads( **model_kwargs: Any, ) -> torch.Tensor: """Return unpacked PDD interval velocities from one Qwen call.""" - self._projection(model, self.config.grid_size) + self._all_heads_projection( + model, + self.config.grid_size, + base_out_features=state.shape[1] * 4, + ) packed = self._call_packed( model, state, @@ -586,7 +683,7 @@ def student_fused_block( **model_kwargs: Any, ) -> torch.Tensor: """Run one conditional Qwen call with its final projection fused for a block.""" - projection = self._projection(model, self.config.grid_size) + projection = self._fused_projection(model, self.config.grid_size) with projection.fuse_block(start, end, grid): packed = self._call_packed( model, diff --git a/tests/examples/diffusers/fastgen/conftest.py b/tests/examples/diffusers/fastgen/conftest.py index 3e052b6c7c9..e4e7d7d04c2 100644 --- a/tests/examples/diffusers/fastgen/conftest.py +++ b/tests/examples/diffusers/fastgen/conftest.py @@ -15,7 +15,6 @@ from __future__ import annotations -import hashlib import json from typing import TYPE_CHECKING @@ -60,7 +59,6 @@ def _make( metadata.append( { "cache_file": str(cache_file), - "cache_sha256": hashlib.sha256(payload_path.read_bytes()).hexdigest(), "bucket_resolution": resolution, "original_resolution": resolution, "bucket_id": sample_id % 2, diff --git a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py b/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py deleted file mode 100644 index 42d5ad394b3..00000000000 --- a/tests/examples/diffusers/fastgen/pdd_checkpoint_failure_distributed.py +++ /dev/null @@ -1,215 +0,0 @@ -# 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. - -"""Two-rank proof that rank-0 checkpoint failures propagate instead of deadlocking.""" - -from __future__ import annotations - -import pathlib -import shutil -import sys -import tempfile -from types import SimpleNamespace - -import torch -import torch.distributed as dist - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) -if str(_FASTGEN_DIR) not in sys.path: - sys.path.insert(0, str(_FASTGEN_DIR)) - -import pdd.checkpoint as pdd_checkpoint_module -from pdd.checkpoint import PDDCheckpointManager - -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION - - -class _State: - def state_dict(self): - return {"value": 1} - - -class _Sampler(_State): - def state_dict(self): - return { - "epoch": 0, - "committed_batches": 1, - "sample_slots_consumed": 1, - "plan_sha256": "0" * 64, - "next_sample_ids": ["next"], - } - - -class _Trainer(_State): - def __init__(self, completed_steps: int = 1) -> None: - self.completed_steps = completed_steps - - -class _StepScheduler: - def __init__(self, trainer: _Trainer) -> None: - self.trainer = trainer - - def state_dict(self): - return {"step": self.trainer.completed_steps, "epoch": 0} - - -class _Checkpointer: - def __init__(self, rank: int, *, fail_sidecar: bool = False) -> None: - self.config = SimpleNamespace(is_async=False) - self.rank = rank - self.fail_sidecar = fail_sidecar - - def save_model(self, model, path: str) -> None: - del model - if self.rank == 0: - root = pathlib.Path(path) / "model" - root.mkdir(parents=True) - (root / ".metadata").write_bytes(b"metadata") - (root / "__0_0.distcp").write_bytes(b"model") - - def save_optimizer(self, optimizer, model, path: str, scheduler) -> None: - del optimizer, model, scheduler - if self.rank == 0: - root = pathlib.Path(path) / "optim" - root.mkdir(parents=True) - (root / ".metadata").write_bytes(b"metadata") - (root / "__0_0.distcp").write_bytes(b"optim") - - def save_on_dp_ranks(self, state, state_name: str, path: str) -> None: - if self.fail_sidecar and self.rank == 1 and state_name == "sampler": - raise OSError("injected rank-1 sidecar failure") - root = pathlib.Path(path) / state_name - root.mkdir(parents=True, exist_ok=True) - torch.save(state.state_dict(), root / f"{state_name}_dp_rank_{self.rank}.pt") - - -class _FailingManager(PDDCheckpointManager): - def __init__(self, *, failure_stage: str, **kwargs) -> None: - super().__init__(**kwargs) - self.failure_stage = failure_stage - - def _prepare_staging(self, final: pathlib.Path) -> str: - if self.failure_stage == "prepare": - raise OSError("injected preparation failure") - return super()._prepare_staging(final) - - def _publish_staging(self, **kwargs) -> None: - if self.failure_stage == "publish": - raise OSError("injected publication failure") - if self.failure_stage != "latest": - super()._publish_staging(**kwargs) - return - original = pdd_checkpoint_module._atomic_text - - def fail_latest(path: pathlib.Path, text: str) -> None: - if path.name == "LATEST": - raise OSError("injected LATEST update failure") - original(path, text) - - pdd_checkpoint_module._atomic_text = fail_latest - try: - super()._publish_staging(**kwargs) - finally: - pdd_checkpoint_module._atomic_text = original - - -def _run_failure(root: pathlib.Path, stage: str) -> None: - rank = dist.get_rank() - trainer = _Trainer() - if stage == "latest": - initial = _FailingManager( - failure_stage="none", - root=root / stage, - checkpointer=_Checkpointer(rank), - model=object(), - optimizer=SimpleNamespace(param_groups=[{"lr": 2.0e-5}]), - scheduler=object(), - step_scheduler=_StepScheduler(trainer), - trainer=trainer, - sampler=_Sampler(), - rng=_State(), - identity={ - "schema_version": 5, - "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, - "topology": {"world_size": 2}, - }, - ) - initial.save() - trainer.completed_steps = 2 - manager = _FailingManager( - failure_stage=stage, - root=root / stage, - checkpointer=_Checkpointer(rank, fail_sidecar=stage == "sidecar"), - model=object(), - optimizer=SimpleNamespace(param_groups=[{"lr": 2.0e-5}]), - scheduler=object(), - step_scheduler=_StepScheduler(trainer), - trainer=trainer, - sampler=_Sampler(), - rng=_State(), - identity={ - "schema_version": 5, - "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, - "topology": {"world_size": 2}, - }, - ) - message = None - try: - manager.save() - except RuntimeError as error: - message = str(error) - messages: list[str | None] = [None] * dist.get_world_size() - dist.all_gather_object(messages, message) - if stage == "sidecar": - assert all( - item is not None and "checkpoint sidecar save failed" in item for item in messages - ) - else: - expected = "preparation" if stage == "prepare" else "publication" - assert all( - item is not None and f"rank-0 checkpoint {expected}" in item for item in messages - ) - if stage == "latest": - assert (root / stage / "LATEST").read_text().strip() == "step_00000001" - assert manager.resolve("LATEST").name == "step_00000002" - - -def main() -> None: - dist.init_process_group("gloo") - root_payload = [ - tempfile.mkdtemp(prefix="modelopt-pdd-rank0-failure-") if dist.get_rank() == 0 else None - ] - dist.broadcast_object_list(root_payload, src=0) - root = pathlib.Path(root_payload[0]) - try: - _run_failure(root, "prepare") - dist.barrier() - _run_failure(root, "publish") - dist.barrier() - _run_failure(root, "sidecar") - dist.barrier() - _run_failure(root, "latest") - finally: - dist.barrier() - if dist.get_rank() == 0: - shutil.rmtree(root) - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/examples/diffusers/fastgen/pdd_export_distributed.py b/tests/examples/diffusers/fastgen/pdd_export_distributed.py deleted file mode 100644 index f1547c0491d..00000000000 --- a/tests/examples/diffusers/fastgen/pdd_export_distributed.py +++ /dev/null @@ -1,200 +0,0 @@ -# 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. - -"""Two-rank released-AutoModel DCP-to-full-state export proof for the PDD example.""" - -from __future__ import annotations - -import pathlib -import shutil -import sys -import tempfile - -import torch -import torch.distributed as dist -from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -for path in (_REPO_ROOT, _REPO_ROOT / "tests", _FASTGEN_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - -from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir -from pdd.export import inspect_pdd_export, write_pdd_export -from pdd.export_qwen_image import collective_export_memory_preflight -from pdd.inference_qwen_image import build_pdd_student -from pdd.recipe import build_pdd_export_setup, resolve_pdd_recipe_config - -from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - QWEN_IMAGE_PDD_EXECUTION, - require_qwen_image_mr210_forward, -) - - -def _raw_config(model_dir: pathlib.Path, checkpoint_dir: pathlib.Path) -> dict: - return { - "model": { - "pretrained_model_name_or_path": str(model_dir), - "torch_dtype": "bfloat16", - "device": "cpu", - "transformer_engine_linear": False, - "peft": None, - "guidance_embeds": False, - "fuse_qkv_projections": False, - }, - "pdd": { - "pred_type": "flow", - "num_train_timesteps": None, - "guidance_scale": 4.0, - "student_sample_steps": 2, - "student_sample_type": "ode", - "grid_size": 4, - "grid_max_t": 0.999, - "flow_shift": 5.0, - "block_size_min": 1, - "block_size_max": 4, - "teacher_integrator": "euler", - "inference_blocks": [2, 2], - "data_free": False, - }, - "optim": { - "learning_rate": 2.0e-5, - "optimizer": {"weight_decay": 0.01}, - }, - "fsdp": { - "dp_size": 2, - "tp_size": 1, - "cp_size": 1, - "pp_size": 1, - "ep_size": 1, - "activation_checkpointing": False, - }, - "checkpoint": { - "enabled": True, - "checkpoint_dir": str(checkpoint_dir), - "model_save_format": "torch_save", - "save_consolidated": False, - }, - } - - -def _full_state(model: torch.nn.Module) -> dict[str, torch.Tensor]: - return get_model_state_dict( - model, - options=StateDictOptions(full_state_dict=True, cpu_offload=True), - ) - - -def main() -> None: - dist.init_process_group("gloo") - payload = [tempfile.mkdtemp(prefix="modelopt-pdd-export-") if dist.get_rank() == 0 else None] - dist.broadcast_object_list(payload, src=0) - root = pathlib.Path(payload[0]) - model_root = root / "model" - model_dir = model_root / "tiny_qwen_image" - try: - if dist.get_rank() == 0: - assert create_tiny_qwen_image_pipeline_dir(model_root) == model_dir - dist.barrier() - config = resolve_pdd_recipe_config(_raw_config(model_dir, root / "checkpoints")) - source = build_pdd_export_setup(config) - expected = _full_state(source.student) - source.checkpointer.save_model(source.student, str(root / "dcp")) - source.checkpointer.close() - - destination = build_pdd_export_setup(config) - destination.checkpointer.load_model(destination.student, str(root / "dcp" / "model")) - full_state_bytes, largest_tensor_bytes = collective_export_memory_preflight( - destination.student, - max_shard_bytes=4 * 1024 * 1024, - headroom=1.0, - device=torch.device("cpu"), - ) - actual = _full_state(destination.student) - status = None - if dist.get_rank() == 0: - try: - assert expected and actual - assert expected.keys() == actual.keys() - for key in expected: - torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) - assert full_state_bytes == sum( - tensor.numel() * tensor.element_size() for tensor in actual.values() - ) - assert largest_tensor_bytes == max( - tensor.numel() * tensor.element_size() for tensor in actual.values() - ) - identity = { - "schema_version": 5, - "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, - "model": { - "id": "Qwen/Qwen-Image", - "revision": "3" * 40, - "dtype": "bfloat16", - }, - "pdd_metadata": destination.metadata.to_dict(), - "guidance": {"scale": 4.0}, - "topology": {"world_size": 2, "pure_data_parallel": True}, - } - output = write_pdd_export( - root / "export", - actual, - metadata=destination.metadata, - transformer_config=destination.transformer_config, - identity=identity, - source_checkpoint={ - "name": "step_00000001", - "manifest_sha256": "1" * 64, - "completed_steps": 1, - }, - max_shard_bytes=4 * 1024 * 1024, - ) - descriptor = inspect_pdd_export(output) - assert descriptor.metadata == destination.metadata - restored, restored_descriptor, dtype = build_pdd_student(output) - require_qwen_image_mr210_forward(restored) - assert dtype == torch.bfloat16 - assert restored_descriptor.metadata == destination.metadata - restored_state = restored.state_dict() - assert restored_state.keys() == actual.keys() - for key in actual: - assert restored_state[key].dtype == torch.bfloat16 - torch.testing.assert_close( - restored_state[key], - actual[key].to(torch.bfloat16), - rtol=0, - atol=0, - ) - status = {"ok": True} - except BaseException as error: - status = {"ok": False, "error": f"{type(error).__name__}: {error}"} - else: - assert not expected and not actual - payload = [status] - dist.broadcast_object_list(payload, src=0) - if not payload[0]["ok"]: - raise RuntimeError(payload[0]["error"]) - destination.checkpointer.close() - dist.barrier() - finally: - dist.barrier() - if dist.get_rank() == 0: - shutil.rmtree(root) - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py b/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py deleted file mode 100644 index d048482108d..00000000000 --- a/tests/examples/diffusers/fastgen/pdd_mr210_fsdp_distributed.py +++ /dev/null @@ -1,440 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Two-rank differential for the Qwen MR210 PDD recipe under FSDP2.""" - -from __future__ import annotations - -import argparse -import copy -import json -import os -import pathlib -import shutil -import sys -import tempfile - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor -from torch.func import functional_call - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -for path in (_REPO_ROOT, _REPO_ROOT / "tests", _FASTGEN_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - -from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir -from diffusers import QwenImageTransformer2DModel -from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock -from pdd.recipe import ( - build_pdd_setup, - build_pdd_training_artifacts, - initialize_pdd_distributed, - resolve_pdd_recipe_config, -) -from pdd.training import PDDValidationAssignment, PreparedPDDBatch, run_pdd_validation -from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( - CheckpointImpl, - CheckpointWrapper, -) - -from modelopt.torch.fastgen import PDDPipeline -from modelopt.torch.fastgen.plugins.qwen_image import build_img_shapes, pack_latents -from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - QwenImagePDDAdapter, - adopt_qwen_image_mr210_forward, - convert_qwen_image_to_pdd, -) - - -class _FSDPBoundaryReferenceAdapter(QwenImagePDDAdapter): - """Emulate FSDP's BF16 parameter views backed by FP32 masters.""" - - def _call_packed( - self, - model, - state, - time, - condition, - model_kwargs, - *, - condition_name, - ): - encoder_hidden_states, attention_mask = self._prepare_call( - model, - state, - time, - condition, - model_kwargs, - condition_name=condition_name, - ) - batch_size, _, height, width = state.shape - max_txt_seq_len = int(attention_mask.sum(dim=1).max().to(torch.int32).item()) - parameters = { - name: parameter.to(torch.bfloat16) if parameter.dtype.is_floating_point else parameter - for name, parameter in model.named_parameters() - } - output = functional_call( - model, - parameters, - (), - { - "hidden_states": pack_latents(state).to(torch.bfloat16), - "timestep": time, - "encoder_hidden_states": encoder_hidden_states.to(torch.bfloat16), - "encoder_hidden_states_mask": attention_mask, - "img_shapes": build_img_shapes(batch_size, height, width), - "max_txt_seq_len": max_txt_seq_len, - "return_dict": False, - **model_kwargs, - }, - strict=False, - ) - return self._extract_packed_output(output) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--activation-checkpointing", action="store_true") - return parser.parse_args() - - -def _raw_config( - model_dir: pathlib.Path, - checkpoint_dir: pathlib.Path, - *, - activation_checkpointing: bool, -) -> dict: - return { - "model": { - "pretrained_model_name_or_path": str(model_dir), - "torch_dtype": "bfloat16", - "device": "cuda", - "transformer_engine_linear": False, - "peft": None, - "guidance_embeds": False, - "fuse_qkv_projections": False, - }, - "pdd": { - "pred_type": "flow", - "num_train_timesteps": None, - "guidance_scale": 4.0, - "student_sample_steps": 2, - "student_sample_type": "ode", - "grid_size": 4, - "grid_max_t": 0.999, - "flow_shift": 5.0, - "block_size_min": 1, - "block_size_max": 4, - "teacher_integrator": "euler", - "inference_blocks": [2, 2], - "data_free": False, - }, - "seed": 42, - "optim": { - "learning_rate": 2.0e-5, - "optimizer": {"_target_": "torch.optim.AdamW", "weight_decay": 0.0}, - }, - "lr_scheduler": { - "lr_decay_style": "constant", - "lr_warmup_steps": 0, - "min_lr": 2.0e-5, - }, - "step_scheduler": { - "max_steps": 1, - "num_epochs": 1, - "log_every": 1, - "ckpt_every_steps": 1, - "local_batch_size": 1, - "global_batch_size": 2, - "save_checkpoint_every_epoch": False, - }, - "training_health": {"max_grad_norm": 1.0, "zero_grad_warmup_steps": 0}, - "validation": {"count": 1, "seed": 11, "split_seed": 7, "every_steps": 1}, - "data": { - "dataloader": { - "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", - "batch_size": 1, - "drop_last": True, - "shuffle": True, - "dynamic_batch_size": False, - } - }, - "fsdp": { - "dp_size": 2, - "tp_size": 1, - "cp_size": 1, - "pp_size": 1, - "ep_size": 1, - "activation_checkpointing": activation_checkpointing, - }, - "checkpoint": { - "enabled": True, - "checkpoint_dir": str(checkpoint_dir), - "model_save_format": "torch_save", - "save_consolidated": False, - }, - } - - -def _full_fp32_gradient(parameter: torch.nn.Parameter) -> torch.Tensor: - gradient = parameter.grad - if gradient is None: - raise RuntimeError("expected a materialized gradient") - if gradient.dtype != torch.float32: - raise RuntimeError(f"expected an FP32 gradient, got {gradient.dtype}") - if isinstance(gradient, DTensor): - gradient = gradient.full_tensor() - if gradient.dtype != torch.float32: - raise RuntimeError(f"expected a gathered FP32 gradient, got {gradient.dtype}") - return gradient.detach() - - -def _assert_checkpointing(model: torch.nn.Module, *, enabled: bool) -> None: - if model.gradient_checkpointing: - raise RuntimeError("native Qwen gradient checkpointing remained enabled") - for index, block in enumerate(model.transformer_blocks): - if enabled: - if not isinstance(block, CheckpointWrapper): - raise RuntimeError(f"Qwen block {index} was not checkpoint-wrapped") - if block.checkpoint_impl is not CheckpointImpl.NO_REENTRANT: - raise RuntimeError(f"Qwen block {index} uses the wrong checkpoint implementation") - if not isinstance(block._checkpoint_wrapped_module, QwenImageTransformerBlock): - raise RuntimeError(f"Qwen block {index} wrapped an unexpected module") - elif not isinstance(block, QwenImageTransformerBlock): - raise RuntimeError(f"Qwen block {index} changed while checkpointing was disabled") - - -def main() -> None: - args = _parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("Qwen MR210 FSDP2 regression requires CUDA") - torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) - initialize_pdd_distributed(backend="nccl", timeout_minutes=5) - rank = dist.get_rank() - device = torch.device("cuda", torch.cuda.current_device()) - payload = [tempfile.mkdtemp(prefix="modelopt-pdd-mr210-fsdp-") if rank == 0 else None] - dist.broadcast_object_list(payload, src=0, device=device) - root = pathlib.Path(payload[0]) - model_root = root / "model" - model_dir = model_root / "tiny_qwen_image" - completed = False - try: - if rank == 0: - assert create_tiny_qwen_image_pipeline_dir(model_root) == model_dir - dist.barrier() - - config = resolve_pdd_recipe_config( - _raw_config( - model_dir, - root / "checkpoints", - activation_checkpointing=args.activation_checkpointing, - ) - ) - setup = build_pdd_setup(config) - _assert_checkpointing(setup.student, enabled=args.activation_checkpointing) - _assert_checkpointing(setup.teacher, enabled=args.activation_checkpointing) - actual_pipeline = PDDPipeline( - setup.student, - setup.teacher, - config.pdd, - QwenImagePDDAdapter(config.pdd, compute_dtype=torch.bfloat16), - ) - - reference_student = QwenImageTransformer2DModel.from_pretrained( - model_dir, - subfolder="transformer", - torch_dtype=torch.bfloat16, - ) - reference_student = adopt_qwen_image_mr210_forward(reference_student) - reference_teacher = copy.deepcopy(reference_student).eval().requires_grad_(False) - reference_projection = convert_qwen_image_to_pdd(reference_student, config.pdd) - reference_student.to(device=device, dtype=torch.float32) - reference_teacher.to(device=device, dtype=torch.float32) - reference_pipeline = PDDPipeline( - reference_student, - reference_teacher, - config.pdd, - _FSDPBoundaryReferenceAdapter(config.pdd, compute_dtype=torch.bfloat16), - ) - - root_calls = {"student": 0, "teacher": 0} - student_times: list[torch.Tensor] = [] - teacher_times: list[torch.Tensor] = [] - inner_student_calls = 0 - - def student_root_hook(_module, _args, _kwargs): - root_calls["student"] += 1 - - def teacher_root_hook(_module, _args, _kwargs): - root_calls["teacher"] += 1 - - def student_time_hook(_module, args_for_module): - student_times.append(args_for_module[0].detach().clone()) - - def teacher_time_hook(_module, args_for_module): - teacher_times.append(args_for_module[0].detach().clone()) - - def inner_student_hook(_module, _args, _kwargs): - nonlocal inner_student_calls - inner_student_calls += 1 - - student_block = setup.student.transformer_blocks[0] - if isinstance(student_block, CheckpointWrapper): - student_block = student_block._checkpoint_wrapped_module - hooks = [ - setup.student.register_forward_pre_hook(student_root_hook, with_kwargs=True), - setup.teacher.register_forward_pre_hook(teacher_root_hook, with_kwargs=True), - setup.student.time_text_embed.register_forward_pre_hook(student_time_hook), - setup.teacher.time_text_embed.register_forward_pre_hook(teacher_time_hook), - student_block.register_forward_pre_hook(inner_student_hook, with_kwargs=True), - ] - - generator = torch.Generator().manual_seed(20260716) - data = torch.randn(1, 4, 4, 4, generator=generator).to(device) - noise = torch.randn(1, 4, 4, 4, generator=generator).to(device) - condition = ( - torch.randn(1, 3, 16, generator=generator).to( - device=device, - dtype=torch.bfloat16, - ), - torch.tensor([[1, 1, 1]], device=device, dtype=torch.long), - ) - negative_condition = ( - torch.randn(1, 2, 16, generator=generator).to( - device=device, - dtype=torch.bfloat16, - ), - torch.tensor([[1, 1]], device=device, dtype=torch.long), - ) - n = torch.tensor([0], device=device, dtype=torch.long) - k = torch.tensor([2], device=device, dtype=torch.long) - - actual_loss, _ = actual_pipeline.compute_loss( - data, - noise=noise, - condition=condition, - negative_condition=negative_condition, - n=n, - k=k, - ) - actual_loss.backward() - reference_loss, _ = reference_pipeline.compute_loss( - data, - noise=noise, - condition=condition, - negative_condition=negative_condition, - n=n, - k=k, - ) - reference_loss.backward() - for hook in hooks: - hook.remove() - - if actual_loss.dtype != torch.float32 or reference_loss.dtype != torch.float32: - raise RuntimeError( - f"PDD losses must be FP32, got {actual_loss.dtype} and {reference_loss.dtype}" - ) - if root_calls != {"student": 1, "teacher": 2}: - raise RuntimeError(f"unexpected adopted-root call counts: {root_calls}") - if len(student_times) != 1 or len(teacher_times) != 2: - raise RuntimeError( - f"unexpected time capture counts: {len(student_times)}, {len(teacher_times)}" - ) - if any(value.dtype != torch.float32 for value in (*student_times, *teacher_times)): - raise RuntimeError("FSDP rounded an MR210 timestep before time_text_embed") - expected_time = actual_pipeline.time_grid(device)[n] - if not torch.equal(student_times[0], expected_time): - raise RuntimeError("student time does not equal the exact first grid value") - if student_times[0].item() == student_times[0].to(torch.bfloat16).float().item(): - raise RuntimeError("the 0.999 discriminator did not distinguish BF16 rounding") - expected_inner_calls = 2 if args.activation_checkpointing else 1 - if inner_student_calls != expected_inner_calls: - raise RuntimeError( - "unexpected student block call count: " - f"expected {expected_inner_calls}, got {inner_student_calls}" - ) - - torch.testing.assert_close(actual_loss, reference_loss, rtol=2e-3, atol=2e-4) - gradient_pairs = ( - (setup.projection.weight, reference_projection.weight), - (setup.student.img_in.weight, reference_student.img_in.weight), - ) - for actual_parameter, reference_parameter in gradient_pairs: - actual_gradient = _full_fp32_gradient(actual_parameter) - reference_gradient = _full_fp32_gradient(reference_parameter) - if ( - not torch.isfinite(actual_gradient).all() - or not torch.isfinite(reference_gradient).all() - ): - raise FloatingPointError("MR210 FSDP gradient comparison is non-finite") - torch.testing.assert_close( - actual_gradient, - reference_gradient, - rtol=1e-2, - atol=2e-3, - ) - - gathered_losses = [torch.zeros_like(actual_loss) for _ in range(dist.get_world_size())] - dist.all_gather(gathered_losses, actual_loss.detach()) - for value in gathered_losses[1:]: - torch.testing.assert_close(value, gathered_losses[0], rtol=0, atol=0) - - setup.optimizer.zero_grad(set_to_none=True) - training = build_pdd_training_artifacts(setup, config) - validation_batch = PreparedPDDBatch( - data, - condition, - negative_condition, - (f"post-validation-update-{rank}",), - ) - run_pdd_validation( - training.pipeline, - [validation_batch], - [ - PDDValidationAssignment(index, f"post-validation-update-{index}", 0, 2) - for index in range(dist.get_world_size()) - ], - validation_seed=11, - ) - post_validation = training.trainer.train_step( - validation_batch, - noise=noise, - n=n, - k=k, - ) - if ( - post_validation.pdd_projection_update_ratio is None - or post_validation.pdd_projection_update_ratio <= 0 - ): - raise RuntimeError("post-validation FSDP projection update was not measured") - if rank == 0: - print( - json.dumps( - { - "activation_checkpointing": args.activation_checkpointing, - "actual_loss": actual_loss.item(), - "reference_loss": reference_loss.item(), - "post_validation_projection_update_ratio": ( - post_validation.pdd_projection_update_ratio - ), - "student_block_calls": inner_student_calls, - "time_0": student_times[0].item(), - "world_size": dist.get_world_size(), - }, - sort_keys=True, - ) - ) - dist.barrier() - completed = True - finally: - if completed and rank == 0: - shutil.rmtree(root, ignore_errors=True) - if dist.is_initialized(): - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/examples/diffusers/fastgen/pdd_test_utils.py b/tests/examples/diffusers/fastgen/pdd_test_utils.py deleted file mode 100644 index 668f282bcb7..00000000000 --- a/tests/examples/diffusers/fastgen/pdd_test_utils.py +++ /dev/null @@ -1,235 +0,0 @@ -# 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. - -"""Small plain-torch objects shared by PDD example lifecycle tests.""" - -from __future__ import annotations - -import hashlib -from dataclasses import dataclass -from typing import Any - -import torch -from pdd.training import PDDTrainer, PreparedPDDBatch -from torch import nn - -from modelopt.torch.fastgen import ( - PDDConfig, - PDDLayerSpec, - PDDMetadata, - PDDOutputProjection, - PDDPipeline, - convert_to_pdd_output_projection, -) - - -def ordered_sample_ids_sha256(sample_ids: tuple[str, ...] | list[str]) -> str: - digest = hashlib.sha256(b"modelopt-pdd-ordered-sample-ids-v1\0") - for sample_id in sample_ids: - digest.update(sample_id.encode()) - digest.update(b"\n") - return digest.hexdigest() - - -class ToyStudent(nn.Module): - def __init__(self, width: int = 3) -> None: - super().__init__() - self.backbone = nn.Linear(width, width) - self.projection = nn.Linear(width, width) - - def forward(self, state: torch.Tensor) -> torch.Tensor: - return self.projection(torch.tanh(self.backbone(state))) - - -class ToyTeacher(nn.Module): - def __init__(self) -> None: - super().__init__() - self.scale = nn.Parameter(torch.tensor(-0.25)) - self.bias = nn.Parameter(torch.tensor(0.125)) - - def forward(self, state: torch.Tensor, time: torch.Tensor) -> torch.Tensor: - return self.scale * state + self.bias + 0.1 * time[:, None] - - -class ToyAdapter: - def __init__(self, grid_size: int, *, zero_student_gradient: bool = False) -> None: - self.grid_size = grid_size - self.zero_student_gradient = zero_student_gradient - - def student_all_heads( - self, - model: ToyStudent, - state: torch.Tensor, - time: torch.Tensor, - *, - condition: Any = None, - **model_kwargs: Any, - ) -> torch.Tensor: - del time, condition, model_kwargs - raw = model(state) - if self.zero_student_gradient: - raw = raw * 0.0 - return raw.reshape(state.shape[0], self.grid_size, state.shape[1]) - - def student_fused_block( - self, - model: ToyStudent, - 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) - with projection.fuse_block(start, end, grid): - return model(state) - - def teacher_velocity( - self, - model: ToyTeacher, - state: torch.Tensor, - time: torch.Tensor, - *, - condition: Any = None, - negative_condition: Any = None, - **model_kwargs: Any, - ) -> torch.Tensor: - del condition, negative_condition, model_kwargs - return model(state, time) - - -@dataclass -class ToyLifecycle: - config: PDDConfig - student: ToyStudent - teacher: ToyTeacher - projection: PDDOutputProjection - pipeline: PDDPipeline - optimizer: torch.optim.AdamW - scheduler: torch.optim.lr_scheduler.LambdaLR - trainer: PDDTrainer - metadata: PDDMetadata - - -def build_toy_lifecycle( - *, - seed: int = 17, - zero_student_gradient: bool = False, - weight_decay: float = 0.01, -) -> ToyLifecycle: - torch.manual_seed(seed) - config = PDDConfig( - grid_size=4, - grid_max_t=0.999, - flow_shift=5.0, - block_size_min=1, - block_size_max=4, - inference_blocks=[2, 2], - student_sample_steps=2, - guidance_scale=None, - ) - student = ToyStudent() - projection = convert_to_pdd_output_projection( - student, - PDDLayerSpec("projection", "channel_major"), - config.grid_size, - ) - teacher = ToyTeacher() - pipeline = PDDPipeline( - student, - teacher, - config, - ToyAdapter(config.grid_size, zero_student_gradient=zero_student_gradient), - ) - optimizer = torch.optim.AdamW( - student.parameters(), - lr=2e-3, - weight_decay=weight_decay, - amsgrad=False, - capturable=False, - differentiable=False, - foreach=False, - fused=False, - maximize=False, - ) - scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda _: 1.0) - trainer = PDDTrainer( - pipeline, - optimizer, - projection=projection, - max_grad_norm=0.5, - ) - return ToyLifecycle( - config, - student, - teacher, - projection, - pipeline, - optimizer, - scheduler, - trainer, - PDDMetadata.from_config(config, projection), - ) - - -def make_batch(sample_ids: tuple[str, ...], *, offset: float = 0.0) -> PreparedPDDBatch: - data = torch.stack( - [ - torch.tensor([0.5 + offset + index / 10, -1.0, 0.25], dtype=torch.float32) - for index in range(len(sample_ids)) - ] - ) - condition = ( - torch.zeros((len(sample_ids), 1, 1), dtype=torch.float32), - torch.ones((len(sample_ids), 1), dtype=torch.long), - ) - return PreparedPDDBatch(data, condition, None, sample_ids) - - -class SamplerDataset: - def __init__(self, sample_ids: tuple[str, ...]) -> None: - self.logical_sample_ids = sample_ids - self.metadata = [ - { - "sample_id": sample_id, - "bucket_id": "64x64", - "bucket_resolution": [64, 64], - } - for sample_id in sample_ids - ] - self.bucket_groups = { - (64, 64): { - "indices": list(range(len(sample_ids))), - "resolution": (64, 64), - "aspect_name": "square", - } - } - self.sorted_bucket_keys = [(64, 64)] - self.calculator = None - - def __len__(self) -> int: - return len(self.metadata) - - def __getitem__(self, index: int) -> int: - return index - - -def ordered_id_sha256(sample_ids: tuple[str, ...]) -> str: - return ordered_sample_ids_sha256(sample_ids) diff --git a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py b/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py deleted file mode 100644 index e197901f775..00000000000 --- a/tests/examples/diffusers/fastgen/pdd_validation_oracle_distributed.py +++ /dev/null @@ -1,167 +0,0 @@ -# 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. - -"""Two-rank CPU/Gloo equivalence harness for deterministic PDD validation.""" - -from __future__ import annotations - -import dataclasses -import pathlib -import sys - -import torch -import torch.distributed as dist - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) -if str(_FASTGEN_DIR) not in sys.path: - sys.path.insert(0, str(_FASTGEN_DIR)) -if str(pathlib.Path(__file__).parent) not in sys.path: - sys.path.insert(0, str(pathlib.Path(__file__).parent)) - -from pdd.training import build_pdd_validation_assignments, run_pdd_validation -from pdd_test_utils import build_toy_lifecycle, make_batch - - -def _expect_failure(error_type, callback) -> None: - try: - callback() - except error_type: - return - raise AssertionError(f"expected {error_type.__name__}") - - -def main() -> None: - lifecycle = build_toy_lifecycle() - sample_ids = tuple(f"distributed-validation-{index:02d}" for index in range(13)) - assignments = build_pdd_validation_assignments( - sample_ids, - lifecycle.config, - validation_seed=91, - require_full_coverage=False, - ) - all_batches = [ - make_batch((assignment.sample_id,), offset=assignment.ordinal / 100) - for assignment in assignments - ] - baseline = run_pdd_validation( - lifecycle.pipeline, - all_batches, - assignments, - validation_seed=91, - ) - - dist.init_process_group(backend="gloo") - try: - rank = dist.get_rank() - world_size = dist.get_world_size() - assert world_size == 2 - local_batches = all_batches[rank::world_size] - padded_batch_count = max( - len(all_batches[candidate_rank::world_size]) for candidate_rank in range(world_size) - ) - while len(local_batches) < padded_batch_count: - local_batches.append( - dataclasses.replace( - make_batch((f"dummy-rank-{rank}",)), - valid_mask=(False,), - ) - ) - distributed = run_pdd_validation( - lifecycle.pipeline, - local_batches, - assignments, - validation_seed=91, - ) - assert distributed.records == baseline.records - assert abs(distributed.mean_loss - baseline.mean_loss) <= 1e-12 - assert distributed.ordered_id_sha256 == baseline.ordered_id_sha256 - - # Text conditions are padded only to each rank-local batch maximum in the - # real Qwen loader. Data-parallel validation must therefore accept - # different sequence lengths while preserving the same logical result. - ragged_conditions = [ - dataclasses.replace( - batch, - condition=( - torch.zeros((len(batch.sample_ids), rank + 1, 1)), - torch.ones((len(batch.sample_ids), rank + 1), dtype=torch.long), - ), - ) - for batch in local_batches - ] - ragged = run_pdd_validation( - lifecycle.pipeline, - ragged_conditions, - assignments, - validation_seed=91, - ) - assert ragged == distributed - - invalid_mask = list(local_batches) - if rank == 0: - invalid_mask[0] = dataclasses.replace(invalid_mask[0], valid_mask=()) - _expect_failure( - RuntimeError, - lambda: run_pdd_validation( - lifecycle.pipeline, - invalid_mask, - assignments, - validation_seed=91, - ), - ) - dist.barrier() - - unassigned = list(local_batches) - if rank == 0: - unassigned[0] = dataclasses.replace( - unassigned[0], - sample_ids=("not-in-heldout-assignments",), - valid_mask=(True,), - ) - _expect_failure( - RuntimeError, - lambda: run_pdd_validation( - lifecycle.pipeline, - unassigned, - assignments, - validation_seed=91, - ), - ) - dist.barrier() - - nonfinite = list(local_batches) - if rank == 0: - nonfinite[0] = dataclasses.replace( - nonfinite[0], - data=torch.full_like(nonfinite[0].data, float("nan")), - ) - _expect_failure( - FloatingPointError, - lambda: run_pdd_validation( - lifecycle.pipeline, - nonfinite, - assignments, - validation_seed=91, - ), - ) - finally: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index f5395be3ef5..89fc3ef1f6c 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -17,7 +17,6 @@ from __future__ import annotations -import hashlib import json import logging import pathlib @@ -36,7 +35,6 @@ sys.path.insert(0, str(_FASTGEN_DIR)) from fastgen_data import ( - ReplayableBatchSampler, TextToImageDataset, build_text_to_image_multiresolution_dataloader, resolve_cache_root, @@ -123,123 +121,7 @@ def test_dataset_accepts_absolute_payload_beneath_root(make_fastgen_cache, tmp_p cache = make_fastgen_cache(tmp_path / "cache", absolute_payloads=True) dataset = TextToImageDataset(cache) - assert dataset[0]["sample_id"] == 0 - - -def test_payload_hash_authentication_rejects_missing_or_modified_payloads( - make_fastgen_cache, tmp_path -): - cache = make_fastgen_cache(tmp_path / "cache") - dataset = TextToImageDataset(cache, verify_payload_hashes=True) - assert dataset[0]["sample_id"] == 0 - - shard_path = cache / "metadata_shard_0.json" - shard = json.loads(shard_path.read_text()) - payload_path = cache / shard[0]["cache_file"] - torch.save({"latent": torch.full((4, 2, 2), 999.0)}, payload_path) - with pytest.raises(RuntimeError, match="SHA-256 mismatch"): - dataset[0] - - shard[0].pop("cache_sha256") - shard_path.write_text(json.dumps(shard)) - with pytest.raises(ValueError, match="verify_payload_hashes=true"): - TextToImageDataset(cache, verify_payload_hashes=True) - - -def test_builder_preserves_legacy_exact_resume_hash_requirement(make_fastgen_cache, tmp_path): - cache = make_fastgen_cache(tmp_path / "cache") - shard_path = cache / "metadata_shard_0.json" - shard = json.loads(shard_path.read_text()) - shard[0].pop("cache_sha256") - shard_path.write_text(json.dumps(shard)) - - with pytest.raises(ValueError, match="verify_payload_hashes=true"): - build_text_to_image_multiresolution_dataloader( - cache_dir=str(cache), num_workers=0, exact_resume=True - ) - - -def test_null_payload_hash_is_incomplete_and_strict_mode_rejects_it( - make_fastgen_cache, tmp_path -): - cache = make_fastgen_cache(tmp_path / "cache") - shard_path = cache / "metadata_shard_0.json" - shard = json.loads(shard_path.read_text()) - shard[0]["cache_sha256"] = None - shard_path.write_text(json.dumps(shard)) - - assert TextToImageDataset(cache).payload_hashes_complete is False - with pytest.raises(ValueError, match="verify_payload_hashes=true"): - TextToImageDataset(cache, verify_payload_hashes=True) - - -def test_hashless_cache_keeps_replayable_exact_cursor(make_fastgen_cache, tmp_path): - cache = make_fastgen_cache(tmp_path / "cache") - shard_path = cache / "metadata_shard_0.json" - shard = json.loads(shard_path.read_text()) - for item in shard: - item.pop("cache_sha256") - shard_path.write_text(json.dumps(shard)) - - options = { - "cache_dir": str(cache), - "batch_size": 1, - "num_workers": 0, - "shuffle": True, - "exact_resume": True, - "verify_payload_hashes": False, - } - loader, sampler = build_text_to_image_multiresolution_dataloader(**options) - assert isinstance(sampler, ReplayableBatchSampler) - assert loader.dataset.verify_payload_hashes is False - assert loader.dataset.payload_hashes_complete is False - - first_batch = next(iter(loader)) - consumed = first_batch["metadata"]["logical_sample_ids"] - sampler.commit(consumed) - state = sampler.state_dict() - expected_next = sampler.expected_next_sample_ids() - - _, restored_sampler = build_text_to_image_multiresolution_dataloader(**options) - restored_sampler.load_state_dict(state) - assert restored_sampler.expected_next_sample_ids() == expected_next - - -def test_hashless_cache_uses_direct_load_by_default(make_fastgen_cache, tmp_path): - cache = make_fastgen_cache(tmp_path / "cache") - shard_path = cache / "metadata_shard_0.json" - shard = json.loads(shard_path.read_text()) - for item in shard: - item.pop("cache_sha256") - shard_path.write_text(json.dumps(shard)) - - loader, sampler = build_text_to_image_multiresolution_dataloader( - cache_dir=str(cache), batch_size=1, num_workers=0 - ) - assert not isinstance(sampler, ReplayableBatchSampler) - assert loader.dataset.verify_payload_hashes is False - assert next(iter(loader))["metadata"]["logical_sample_ids"] - - -def test_dataset_snapshot_binds_negative_prompt_embedding(make_fastgen_cache, tmp_path): - cache = make_fastgen_cache(tmp_path / "cache") - loader, _ = build_text_to_image_multiresolution_dataloader( - cache_dir=str(cache), - num_workers=0, - negative_prompt_embedding_path="negative_prompt_embedding.pt", - exact_resume=True, - ) - first_snapshot = loader.dataset.dataset_snapshot_sha256 - - torch.save(torch.full((2, 3), 7.0), cache / "negative_prompt_embedding.pt") - rebuilt, _ = build_text_to_image_multiresolution_dataloader( - cache_dir=str(cache), - num_workers=0, - negative_prompt_embedding_path="negative_prompt_embedding.pt", - exact_resume=True, - ) - - assert rebuilt.dataset.dataset_snapshot_sha256 != first_snapshot + assert "latent" in dataset[0] def test_environment_redirects_samples_and_relative_negative_embedding( @@ -324,4 +206,3 @@ def test_preprocessing_publishes_absolute_paths_for_relative_output(monkeypatch, assert published.is_absolute() assert published == payload.resolve() published.relative_to(output.resolve()) - assert shard[0]["cache_sha256"] == hashlib.sha256(payload.read_bytes()).hexdigest() diff --git a/tests/examples/diffusers/fastgen/test_dataset_splits.py b/tests/examples/diffusers/fastgen/test_dataset_splits.py index 639b5af0be3..f458549cafd 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_splits.py +++ b/tests/examples/diffusers/fastgen/test_dataset_splits.py @@ -32,7 +32,6 @@ sys.path.insert(0, str(_FASTGEN_DIR)) from fastgen_data import ( - TextToImageDataset, build_text_to_image_multiresolution_dataloader, make_train_validation_indices, ) @@ -77,25 +76,6 @@ def test_split_rejects_invalid_inputs(num_samples, validation_count, seed, error make_train_validation_indices(num_samples, validation_count, seed) -@pytest.mark.parametrize( - "selected_indices", - [[], [0, 0], [-1], [6], [True], [1.5], ["1"]], -) -def test_dataset_rejects_invalid_selected_indices(make_fastgen_cache, selected_indices, tmp_path): - cache = make_fastgen_cache(tmp_path / "cache") - with pytest.raises((TypeError, ValueError)): - TextToImageDataset(cache, selected_indices=selected_indices) - - -def test_dataset_preserves_selected_original_ordinals(make_fastgen_cache, tmp_path): - cache = make_fastgen_cache(tmp_path / "cache") - dataset = TextToImageDataset(cache, selected_indices=[5, 1, 3]) - - assert dataset.total_num_samples == 6 - assert dataset.sample_ids == [5, 1, 3] - assert [dataset[index]["sample_id"] for index in range(len(dataset))] == [5, 1, 3] - - def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( make_fastgen_cache, tmp_path ): @@ -124,14 +104,8 @@ def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( drop_last=False, ) - train_seen = torch.cat([batch["metadata"]["sample_ids"] for batch in train_loader]).tolist() - validation_seen = torch.cat( - [batch["metadata"]["sample_ids"] for batch in validation_loader] - ).tolist() - - assert sorted(train_seen) == train_ids - assert sorted(validation_seen) == validation_ids - assert set(train_seen).isdisjoint(validation_seen) + assert train_loader.dataset.sample_ids == train_ids + assert validation_loader.dataset.sample_ids == validation_ids assert ( train_loader.dataset.cache_root == validation_loader.dataset.cache_root == cache.resolve() ) @@ -139,11 +113,4 @@ def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( assert not validation_sampler.shuffle_buckets assert not validation_sampler.shuffle_within_bucket assert validation_sampler.drop_last is False - assert train_loader.dataset.logical_sample_ids == [str(value) for value in train_ids] - assert validation_loader.dataset.logical_sample_ids == [str(value) for value in validation_ids] - assert all( - batch["metadata"]["sample_ids"].dtype == torch.long - and batch["metadata"]["sample_ids"].device.type == "cpu" - for batch in validation_loader - ) assert _snapshot(cache) == before diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index 5f8c1215c5e..c4aaee9d616 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -51,15 +51,10 @@ _EXPECTED_PDD_FILES = { "README.md", "__init__.py", - "artifacts.py", - "checkpoint.py", "configs", - "data.py", - "export.py", - "export_qwen_image.py", "finetune.py", - "inference_runtime.py", "inference_qwen_image.py", + "prepare_qwen_image.py", "recipe.py", "training.py", } 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..c1d4abf1fd9 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_pdd_inference.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. + +"""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], + student_sample_steps=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" + " student_sample_steps: 2\n" + ) + + blocks = _parse_blocks("2, 2,4") + assert blocks == [2, 2, 4] + assert _load_config(config_path, blocks).inference_blocks == blocks diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py b/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py deleted file mode 100644 index 92629838596..00000000000 --- a/tests/examples/diffusers/fastgen/test_pdd_inference_checkpoint.py +++ /dev/null @@ -1,413 +0,0 @@ -# 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. - -"""Tests for authenticated PDD export, reconstruction, and schedules.""" - -from __future__ import annotations - -import copy -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" -for path in (_REPO_ROOT, _FASTGEN_DIR): - if str(path) not in sys.path: - sys.path.insert(0, str(path)) - -from pdd.export import ( - PDD_INFERENCE_SCHEDULES, - inspect_pdd_export, - load_pdd_export_into_model, - pdd_config_from_metadata, - write_pdd_export, -) -from pdd.inference_qwen_image import ( - _model_identity, - _normalize_prompt_condition, - _validate_qwen_projection, -) - -import modelopt.torch.fastgen.plugins.qwen_image_pdd as qwen_image_pdd_plugin -from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline -from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( - QWEN_IMAGE_PDD_EXECUTION, - QwenImagePDDAdapter, - convert_qwen_image_to_pdd, -) - - -class _TinyQwen(nn.Module): - def __init__(self) -> None: - super().__init__() - self.config = SimpleNamespace(guidance_embeds=False, in_channels=4) - self._modelopt_qwen_image_pdd_execution = QWEN_IMAGE_PDD_EXECUTION - self.backbone = nn.Linear(4, 5, dtype=torch.bfloat16) - self.proj_out = nn.Linear(5, 4, dtype=torch.bfloat16) - self.calls = 0 - - def forward( - self, - *, - hidden_states, - timestep, - encoder_hidden_states, - encoder_hidden_states_mask, - img_shapes, - max_txt_seq_len, - return_dict, - ): - del img_shapes, max_txt_seq_len - assert return_dict is False - condition = encoder_hidden_states.mean(dim=(1, 2), keepdim=True) - condition += (encoder_hidden_states_mask.sum(dim=1)[:, None, None] / 100).to( - condition.dtype - ) - hidden = torch.tanh(self.backbone(hidden_states)) - self.calls += 1 - hidden = hidden + condition - hidden = hidden + (timestep[:, None, None] / 10).to(hidden.dtype) - return (self.proj_out(hidden),) - - -@pytest.fixture(autouse=True) -def _allow_tiny_qwen_protocol_double(monkeypatch): - require_production_forward = qwen_image_pdd_plugin.require_qwen_image_mr210_forward - - def require_forward(model: nn.Module) -> str: - if type(model) is _TinyQwen: - if model._modelopt_qwen_image_pdd_execution != QWEN_IMAGE_PDD_EXECUTION: - raise RuntimeError( - "Qwen-Image PDD requires the bound FastGen MR210 forward execution." - ) - return QWEN_IMAGE_PDD_EXECUTION - return require_production_forward(model) - - monkeypatch.setattr(qwen_image_pdd_plugin, "require_qwen_image_mr210_forward", require_forward) - - -def _config(blocks=(32, 32, 32, 32)) -> PDDConfig: - return PDDConfig( - grid_size=128, - grid_max_t=0.999, - flow_shift=5.0, - block_size_min=4, - block_size_max=64, - inference_blocks=list(blocks), - student_sample_steps=len(blocks), - guidance_scale=4.0, - num_train_timesteps=None, - ) - - -def _converted(seed: int = 17): - torch.manual_seed(seed) - model = _TinyQwen() - config = _config() - projection = convert_qwen_image_to_pdd(model, config) - generator = torch.Generator().manual_seed(seed + 1) - with torch.no_grad(): - for parameter in model.parameters(): - parameter.copy_(torch.randn(parameter.shape, generator=generator) / 10) - return model, config, PDDMetadata.from_config(config, projection) - - -def _identity(metadata: PDDMetadata) -> dict: - return { - "schema_version": 5, - "qwen_image": {"execution": QWEN_IMAGE_PDD_EXECUTION}, - "model": {"id": "synthetic-qwen", "revision": "f" * 40, "dtype": "bfloat16"}, - "pdd_metadata": metadata.to_dict(), - "guidance": {"scale": 4.0}, - "topology": {"world_size": 1, "pure_data_parallel": True}, - } - - -def _write(tmp_path: pathlib.Path): - model, config, metadata = _converted() - output = write_pdd_export( - tmp_path / "export", - model.state_dict(), - metadata=metadata, - transformer_config={"_class_name": "SyntheticQwen", "in_channels": 4}, - identity=_identity(metadata), - source_checkpoint={ - "name": "step_00000010", - "manifest_sha256": "3" * 64, - "completed_steps": 10, - }, - max_shard_bytes=5_800, - ) - return output, model, config, metadata - - -def _condition(): - return torch.tensor([[[0.2, -0.3], [0.1, 0.4]]], dtype=torch.bfloat16), torch.ones( - 1, 2, dtype=torch.long - ) - - -def _sample(model: nn.Module, config: PDDConfig, noise: torch.Tensor) -> torch.Tensor: - pipeline = PDDPipeline( - model, - nn.Identity(), - config, - QwenImagePDDAdapter(config, compute_dtype=torch.bfloat16), - ) - return pipeline.sample(noise.clone(), condition=_condition()) - - -def test_bounded_safe_export_round_trip_and_seeded_schedules(tmp_path, monkeypatch) -> None: - output, source, _source_config, metadata = _write(tmp_path) - descriptor = inspect_pdd_export(output) - - shards = sorted(output.glob("*.safetensors")) - assert len(shards) >= 2 - assert all(path.stat().st_size <= descriptor.manifest["max_shard_bytes"] for path in shards) - assert descriptor.metadata == metadata - - restored = _TinyQwen() - convert_qwen_image_to_pdd(restored, _config()) - monkeypatch.setattr(torch, "load", lambda *args, **kwargs: pytest.fail("unsafe torch.load")) - load_pdd_export_into_model(output, restored) - for key, tensor in source.state_dict().items(): - torch.testing.assert_close(restored.state_dict()[key], tensor, rtol=0, atol=0) - - noise = torch.randn((1, 1, 4, 4), generator=torch.Generator().manual_seed(91)) - for schedule, blocks in PDD_INFERENCE_SCHEDULES.items(): - config = pdd_config_from_metadata( - metadata, - schedule=schedule, - guidance_scale=4.0, - ) - source.calls = 0 - restored.calls = 0 - expected = _sample(source, config, noise) - actual = _sample(restored, config, noise) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - assert source.calls == restored.calls == len(blocks) - torch.testing.assert_close(_sample(restored, config, noise), actual, rtol=0, atol=0) - - arbitrary = pdd_config_from_metadata(metadata, blocks=[1, 127], guidance_scale=4.0) - assert arbitrary.inference_blocks == [1, 127] - assert arbitrary.student_sample_steps == 2 - - -def test_inference_config_preserves_authenticated_nondefault_grid_max_t() -> None: - _model, _config_value, metadata = _converted() - payload = metadata.to_dict() - payload["grid_max_t"] = 1.0 - boundary_metadata = PDDMetadata.from_dict(payload) - - config = pdd_config_from_metadata( - boundary_metadata, - schedule="pdd-4", - guidance_scale=4.0, - ) - - assert config.grid_max_t == 1.0 - assert ( - PDDPipeline( - _TinyQwen(), - nn.Identity(), - config, - QwenImagePDDAdapter(config), - ).time_grid()[0] - == 1.0 - ) - - -def test_pinned_qwen_none_prompt_mask_is_normalized_for_pdd() -> None: - """Diffusers 0.38 returns None when the single-prompt mask is all ones.""" - embeddings = torch.randn(1, 3, 5, dtype=torch.float32) - resolved_embeddings, resolved_mask = _normalize_prompt_condition( - embeddings, - None, - device=torch.device("cpu"), - dtype=torch.bfloat16, - ) - assert resolved_embeddings.dtype == torch.bfloat16 - assert resolved_mask.dtype == torch.long - assert resolved_mask.shape == embeddings.shape[:2] - assert torch.equal(resolved_mask, torch.ones_like(resolved_mask)) - - -def test_qwen_projection_rejects_inconsistent_packed_width() -> None: - _model, _config_value, metadata = _converted() - base = _TinyQwen() - assert _validate_qwen_projection(base, metadata) is base.proj_out - base.config.in_channels = 8 - with pytest.raises(RuntimeError, match="proj_out width"): - _validate_qwen_projection(base, metadata) - - -def test_export_is_complete_before_atomic_rename(tmp_path, monkeypatch) -> None: - original_rename = pathlib.Path.rename - observed = False - - def checked_rename(path, target): - nonlocal observed - if path.name.endswith(".staging"): - inspect_pdd_export(path) - assert (path / "COMPLETE").is_file() - observed = True - return original_rename(path, target) - - monkeypatch.setattr(pathlib.Path, "rename", checked_rename) - _write(tmp_path) - assert observed - - -def test_export_accepts_an_immutable_model_revision(tmp_path) -> None: - model, _config_value, metadata = _converted() - identity = _identity(metadata) - output = write_pdd_export( - tmp_path / "immutable-revision", - model.state_dict(), - metadata=metadata, - transformer_config={"in_channels": 4}, - identity=identity, - source_checkpoint={ - "name": "step_00000010", - "manifest_sha256": "3" * 64, - "completed_steps": 10, - }, - max_shard_bytes=12_000, - ) - assert inspect_pdd_export(output).manifest["identity"]["model"]["revision"] == "f" * 40 - - -@pytest.mark.parametrize("execution", [None, "canonical_diffusers"]) -def test_export_and_inference_reject_incompatible_qwen_execution(tmp_path, execution) -> None: - model, _config_value, metadata = _converted() - identity = _identity(metadata) - if execution is None: - identity.pop("qwen_image") - else: - identity["qwen_image"] = {"execution": execution} - - with pytest.raises(ValueError, match=r"Qwen execution identity|missing keys"): - write_pdd_export( - tmp_path / f"bad-execution-{execution}", - model.state_dict(), - metadata=metadata, - transformer_config={"in_channels": 4}, - identity=identity, - source_checkpoint={ - "name": "step_00000010", - "manifest_sha256": "3" * 64, - "completed_steps": 10, - }, - max_shard_bytes=12_000, - ) - - descriptor = SimpleNamespace(manifest={"identity": identity}) - with pytest.raises(RuntimeError, match="Qwen execution identity"): - _model_identity(descriptor) - - -@pytest.mark.parametrize("revision", [None, "main", "F" * 40]) -def test_export_and_inference_reject_mutable_model_revisions(tmp_path, revision) -> None: - model, _config_value, metadata = _converted() - identity = _identity(metadata) - identity["model"]["revision"] = revision - with pytest.raises(ValueError, match="exact lowercase commit"): - write_pdd_export( - tmp_path / f"bad-revision-{str(revision)[:8]}", - model.state_dict(), - metadata=metadata, - transformer_config={"in_channels": 4}, - identity=identity, - source_checkpoint={ - "name": "step_00000010", - "manifest_sha256": "3" * 64, - "completed_steps": 10, - }, - max_shard_bytes=12_000, - ) - - descriptor = SimpleNamespace(manifest={"identity": identity}) - with pytest.raises(RuntimeError, match="exact lowercase commit"): - _model_identity(descriptor) - - -def test_export_rejects_nonfinite_and_existing_destination(tmp_path) -> None: - output, model, _config_value, metadata = _write(tmp_path) - with pytest.raises(FileExistsError): - write_pdd_export( - output, - model.state_dict(), - metadata=metadata, - transformer_config={"in_channels": 4}, - identity=_identity(metadata), - source_checkpoint={ - "name": "step_00000010", - "manifest_sha256": "3" * 64, - "completed_steps": 10, - }, - max_shard_bytes=12_000, - ) - - bad = copy.deepcopy(model.state_dict()) - bad["backbone.weight"][0, 0] = float("nan") - with pytest.raises(FloatingPointError, match="non-finite"): - write_pdd_export( - tmp_path / "bad", - bad, - metadata=metadata, - transformer_config={"in_channels": 4}, - identity=_identity(metadata), - source_checkpoint={ - "name": "step_00000010", - "manifest_sha256": "3" * 64, - "completed_steps": 10, - }, - max_shard_bytes=12_000, - ) - - -@pytest.mark.parametrize("corruption", ["complete", "shard", "extra", "symlink"]) -def test_export_authentication_rejects_corruption(tmp_path, corruption) -> None: - output, _model, _config_value, _metadata = _write(tmp_path) - if corruption == "complete": - (output / "COMPLETE").unlink() - elif corruption == "shard": - shard = next(output.glob("*.safetensors")) - with shard.open("ab") as stream: - stream.write(b"corrupt") - elif corruption == "extra": - (output / "undeclared.bin").write_bytes(b"extra") - else: - (output / "linked").symlink_to(output / "config.json") - - with pytest.raises((FileNotFoundError, RuntimeError)): - inspect_pdd_export(output) - - -def test_safe_load_rejects_wrong_model_inventory(tmp_path) -> None: - output, _model, _config_value, _metadata = _write(tmp_path) - unconverted = _TinyQwen() - with pytest.raises(RuntimeError, match="shape mismatch"): - load_pdd_export_into_model(output, unconverted) - - wrong_dtype = _TinyQwen().to(torch.float64) - convert_qwen_image_to_pdd(wrong_dtype, _config()) - with pytest.raises(RuntimeError, match="dtype mismatch"): - load_pdd_export_into_model(output, wrong_dtype) diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index a514eaab38d..e6c1c74a132 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -13,845 +13,89 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Released-AutoModel seam tests for the ModelOpt-owned PDD setup.""" +"""Focused tests for the PDD loss seam used by AutoModel's diffusion recipe.""" from __future__ import annotations -import copy -import os import pathlib -import subprocess import sys +from types import SimpleNamespace import pytest import torch -import yaml -from _test_utils.torch.diffusers_models import ( - create_tiny_qwen_image_pipeline_dir, - get_tiny_qwen_image_transformer, -) -from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock from torch import nn -from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( - CheckpointImpl, - CheckpointWrapper, - checkpoint_wrapper, -) -from torch.distributed.checkpoint import FileSystemReader _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.recipe import ( - _apply_qwen_image_activation_checkpointing, - _materialize_zero_step_adamw_state, - _projection_identity, - _require_fp32_optimizer_storage, - _require_immutable_model_source, - _resolve_model_source, - _stage_and_shard_training_models, - build_pdd_export_setup, - build_pdd_setup, - initialize_pdd_distributed, - resolve_pdd_recipe_config, -) +from pdd.recipe import _validate_prepared_student +from pdd.training import PDDFlowMatchingStepAdapter -from modelopt.torch.fastgen import PDDLayerSpec, convert_to_pdd_output_projection -from modelopt.torch.fastgen.plugins.qwen_image_pdd import require_qwen_image_mr210_forward +from modelopt.torch.fastgen import PDDConfig -def _canonical_parameter_names(model: nn.Module) -> tuple[set[str], set[str]]: - raw_names = {name for name, _ in model.named_parameters()} - canonical_names: set[str] = set() - for raw_name in raw_names: - parts = raw_name.split(".") - if parts[0] == "transformer_blocks": - assert len(parts) > 3 - assert parts[1].isdigit() - assert parts[2] == "_checkpoint_wrapped_module" - assert parts.count("_checkpoint_wrapped_module") == 1 - canonical_name = ".".join((*parts[:2], *parts[3:])) - else: - assert "_checkpoint_wrapped_module" not in parts - canonical_name = raw_name - assert canonical_name not in canonical_names - canonical_names.add(canonical_name) - return raw_names, canonical_names +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) -def test_qwen_activation_checkpointing_wraps_every_block_once() -> None: - model = get_tiny_qwen_image_transformer(num_layers=2) - model.enable_gradient_checkpointing() - original_parameter_names = {name for name, _ in model.named_parameters()} - original_state_names = set(model.state_dict()) +class _LossPipeline: + def __init__(self, *, guidance_scale: float | None = None) -> None: + self.config = SimpleNamespace(guidance_scale=guidance_scale) + self.scale = nn.Parameter(torch.tensor(2.0)) + self.last_call = None - wrapped = _apply_qwen_image_activation_checkpointing(model, enabled=True) + 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} - assert wrapped == 2 - assert model.gradient_checkpointing is False - assert len(model.transformer_blocks) == 2 - for block in model.transformer_blocks: - assert isinstance(block, CheckpointWrapper) - assert block.checkpoint_impl is CheckpointImpl.NO_REENTRANT - assert type(block._checkpoint_wrapped_module) is QwenImageTransformerBlock - raw_parameter_names, canonical_parameter_names = _canonical_parameter_names(model) - assert canonical_parameter_names == original_parameter_names - assert set(model.state_dict()) == original_state_names - assert "proj_out.weight" in raw_parameter_names - assert "proj_out.bias" in raw_parameter_names - with pytest.raises(RuntimeError, match="already checkpoint-wrapped"): - _apply_qwen_image_activation_checkpointing(model, enabled=True) - - -def test_qwen_activation_checkpointing_disabled_keeps_exact_blocks() -> None: - model = get_tiny_qwen_image_transformer(num_layers=2) - model.enable_gradient_checkpointing() - - wrapped = _apply_qwen_image_activation_checkpointing(model, enabled=False) - - assert wrapped == 0 - assert model.gradient_checkpointing is False - assert all(type(block) is QwenImageTransformerBlock for block in model.transformer_blocks) - - -def test_qwen_activation_checkpointing_rejects_pre_wrapped_input() -> None: - model = get_tiny_qwen_image_transformer(num_layers=2) - model.transformer_blocks[0].attn = checkpoint_wrapper( - model.transformer_blocks[0].attn, - checkpoint_impl=CheckpointImpl.NO_REENTRANT, - ) - - with pytest.raises(RuntimeError, match="already checkpoint-wrapped"): - _apply_qwen_image_activation_checkpointing(model, enabled=True) - - -def test_zero_step_adamw_state_preserves_first_lazy_update() -> None: - torch.manual_seed(7) - eager_model = nn.Linear(4, 3) - lazy_model = copy.deepcopy(eager_model) - optimizer_options = { - "lr": 2.0e-5, - "weight_decay": 0.01, - "foreach": False, - "fused": False, - } - eager_optimizer = torch.optim.AdamW(eager_model.parameters(), **optimizer_options) - lazy_optimizer = torch.optim.AdamW(lazy_model.parameters(), **optimizer_options) - parameters_before = { - name: parameter.detach().clone() for name, parameter in eager_model.named_parameters() - } - - _materialize_zero_step_adamw_state(eager_optimizer) - - for name, parameter in eager_model.named_parameters(): - torch.testing.assert_close(parameter, parameters_before[name], rtol=0, atol=0) - assert parameter.grad is None - state = eager_optimizer.state[parameter] - assert state["step"].item() == 0 - assert not state["exp_avg"].count_nonzero() - assert not state["exp_avg_sq"].count_nonzero() - - eager_model.weight.square().mean().backward() - lazy_model.weight.square().mean().backward() - eager_optimizer.step() - lazy_optimizer.step() - - for eager_parameter, lazy_parameter in zip( - eager_model.parameters(), lazy_model.parameters(), strict=True - ): - torch.testing.assert_close(eager_parameter, lazy_parameter, rtol=0, atol=0) - for key in ("step", "exp_avg", "exp_avg_sq"): - torch.testing.assert_close( - eager_optimizer.state[eager_model.weight][key], - lazy_optimizer.state[lazy_model.weight][key], - rtol=0, - atol=0, - ) - assert eager_optimizer.state[eager_model.bias]["step"].item() == 0 - assert lazy_model.bias not in lazy_optimizer.state - - -def test_fp32_optimizer_storage_rejects_low_precision_masters_and_state() -> None: - model = nn.Linear(2, 2, dtype=torch.bfloat16) - optimizer = torch.optim.AdamW(model.parameters(), foreach=False, fused=False) - _materialize_zero_step_adamw_state(optimizer) - with pytest.raises(RuntimeError, match="master parameters must be FP32"): - _require_fp32_optimizer_storage(optimizer) - - fp32_model = nn.Linear(2, 2) - fp32_optimizer = torch.optim.AdamW(fp32_model.parameters(), foreach=False, fused=False) - _materialize_zero_step_adamw_state(fp32_optimizer) - fp32_optimizer.state[fp32_model.weight]["exp_avg"] = torch.zeros_like( - fp32_model.weight, dtype=torch.bfloat16 - ) - with pytest.raises(RuntimeError, match="exp_avg state must be FP32"): - _require_fp32_optimizer_storage(fp32_optimizer) - - -def test_training_setup_shards_student_before_staging_teacher(monkeypatch) -> None: - events: list[str] = [] - - class TrackedModel(nn.Module): - def __init__(self, label: str, *, projection: bool = False) -> None: - super().__init__() - self.label = label - self.proj_out = nn.Linear(2, 2) if projection else nn.Identity() - - def to(self, *args, **kwargs): - events.append(f"{self.label}.to") - return super().to(*args, **kwargs) - - def fuse_qkv_projections(self): - events.append(f"{self.label}.fuse") - self.fused_projections = True - - class TrackedManager: - def parallelize(self, model): - events.append(f"{model.label}.parallelize") - return model - - def apply_checkpointing(model, *, enabled): - assert enabled is True - events.append(f"{model.label}.checkpoint") - return 1 - - monkeypatch.setattr( - "pdd.recipe._apply_qwen_image_activation_checkpointing", - apply_checkpointing, - ) - - student = TrackedModel("student", projection=True) - teacher = TrackedModel("teacher") - projection = convert_to_pdd_output_projection( - student, - PDDLayerSpec("proj_out", "channel_major"), - grid_size=4, - ) - - staged_student, staged_teacher = _stage_and_shard_training_models( - student, - teacher, - projection, - _projection_identity(projection), - TrackedManager(), - device=torch.device("cpu"), - fuse_qkv_projections=True, - activation_checkpointing=True, - ) - - assert staged_student is student - assert staged_teacher is teacher - assert {parameter.dtype for parameter in staged_student.parameters()} == {torch.float32} - assert events == [ - "student.to", - "student.fuse", - "student.checkpoint", - "student.parallelize", - "teacher.to", - "teacher.fuse", - "teacher.checkpoint", - "teacher.parallelize", - ] - - -def test_training_setup_upcasts_bf16_models_to_fp32_masters(monkeypatch) -> None: - class IdentityManager: - @staticmethod - def parallelize(model): - return model - - monkeypatch.setattr( - "pdd.recipe._apply_qwen_image_activation_checkpointing", - lambda _model, *, enabled: 0, - ) - - student = nn.Module() - student.proj_out = nn.Linear(2, 2, dtype=torch.bfloat16) - teacher = nn.Linear(2, 2, dtype=torch.bfloat16) - projection = convert_to_pdd_output_projection( - student, - PDDLayerSpec("proj_out", "channel_major"), - grid_size=4, - ) - - staged_student, staged_teacher = _stage_and_shard_training_models( - student, - teacher, - projection, - _projection_identity(projection), - IdentityManager(), - device=torch.device("cpu"), - fuse_qkv_projections=False, - activation_checkpointing=False, - ) - - assert {parameter.dtype for parameter in staged_student.parameters()} == {torch.float32} - assert {parameter.dtype for parameter in staged_teacher.parameters()} == {torch.float32} - - -def _raw_config(model_dir: pathlib.Path, *, qkv: bool = False) -> dict: +def _batch() -> dict[str, torch.Tensor]: return { - "model": { - "pretrained_model_name_or_path": str(model_dir), - "torch_dtype": "bfloat16", - "device": "cpu", - "transformer_engine_linear": False, - "peft": None, - "guidance_embeds": False, - "fuse_qkv_projections": qkv, - }, - "pdd": { - "pred_type": "flow", - "num_train_timesteps": None, - "guidance_scale": 4.0, - "student_sample_steps": 2, - "student_sample_type": "ode", - "grid_size": 4, - "grid_max_t": 0.999, - "flow_shift": 5.0, - "block_size_min": 1, - "block_size_max": 4, - "teacher_integrator": "euler", - "inference_blocks": [2, 2], - "data_free": False, - }, - "seed": 42, - "optim": { - "learning_rate": 2.0e-5, - "optimizer": { - "_target_": "torch.optim.AdamW", - "weight_decay": 0.01, - }, - }, - "lr_scheduler": { - "lr_decay_style": "constant", - "lr_warmup_steps": 0, - "min_lr": 2.0e-5, - }, - "step_scheduler": { - "max_steps": 10, - "num_epochs": 2, - "log_every": 1, - "ckpt_every_steps": 5, - "local_batch_size": 1, - "global_batch_size": 1, - "save_checkpoint_every_epoch": False, - }, - "training_health": {"max_grad_norm": 1.0, "zero_grad_warmup_steps": 0}, - "validation": {"count": 3, "seed": 11, "split_seed": 7, "every_steps": 5}, - "data": { - "dataloader": { - "_target_": "fastgen_data.build_text_to_image_multiresolution_dataloader", - "batch_size": 1, - "drop_last": True, - "shuffle": True, - "dynamic_batch_size": False, - } - }, - "fsdp": { - "dp_size": 1, - "tp_size": 1, - "cp_size": 1, - "pp_size": 1, - "ep_size": 1, - "activation_checkpointing": False, - }, - "checkpoint": { - "enabled": True, - "checkpoint_dir": "checkpoints/test", - "model_save_format": "torch_save", - "save_consolidated": False, - }, + "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_example_recipe_explicitly_pins_grid_max_t() -> None: - raw = yaml.safe_load((_FASTGEN_DIR / "pdd" / "configs" / "qwen_image.yaml").read_text()) - assert type(raw["pdd"]["grid_max_t"]) is float - assert raw["pdd"]["grid_max_t"] == 0.999 - assert raw["validation"]["count"] == 2000 - assert raw["validation"]["split_seed"] == 2026 - - -def test_split_config_fields_are_strict(tmp_path) -> None: - raw = _raw_config(tmp_path) - raw["validation"] = {"count": 3, "seed": 11, "split_seed": 7, "every_steps": 5} - config = resolve_pdd_recipe_config(raw) - assert config.validation.count == 3 - assert config.validation.split_seed == 7 - - for invalid_count in (0, -1, True, 1.5): - raw["validation"]["count"] = invalid_count - with pytest.raises((TypeError, ValueError), match=r"validation\.count"): - resolve_pdd_recipe_config(raw) - - raw["validation"] = {"count": 3, "seed": 11, "split_seed": -1, "every_steps": 5} - with pytest.raises(ValueError, match="split_seed"): - resolve_pdd_recipe_config(raw) - - -def test_config_node_and_canonical_dotted_values_are_consumed(tmp_path) -> None: - raw = _raw_config(tmp_path) - raw["step_scheduler"].update( - max_steps=50_000, - ckpt_every_steps=1_000, - global_batch_size=1, - ) - raw["optim"]["learning_rate"] = 3.0e-5 - raw["lr_scheduler"]["min_lr"] = 3.0e-5 - - class ConfigNodeLike: - def to_dict(self): - return copy.deepcopy(raw) - - config = resolve_pdd_recipe_config(ConfigNodeLike()) - assert config.step_scheduler.max_steps == 50_000 - assert config.step_scheduler.ckpt_every_steps == 1_000 - assert config.step_scheduler.global_batch_size == 1 - assert config.learning_rate == 3.0e-5 - - -def test_automodel_parser_dotted_overrides_reach_the_pdd_resolver(tmp_path, monkeypatch) -> None: - parser_module = pytest.importorskip("nemo_automodel.components.config._arg_parser") - config_path = tmp_path / "pdd.yaml" - config_path.write_text(yaml.safe_dump(_raw_config(tmp_path))) - monkeypatch.setattr( - sys, - "argv", - [ - "finetune.py", - "--config", - str(config_path), - "--step_scheduler.max_steps=50000", - "--step_scheduler.ckpt_every_steps=1000", - "--step_scheduler.global_batch_size=1", - "--optim.learning_rate=3e-5", - "--lr_scheduler.min_lr=3e-5", - ], +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], + student_sample_steps=2, ) + _validate_prepared_student(_PreparedStudent(out_features=32), config) - parsed = parser_module.parse_args_and_load_config(str(config_path)) - resolved = resolve_pdd_recipe_config(parsed) - assert resolved.step_scheduler.max_steps == 50_000 - assert resolved.step_scheduler.ckpt_every_steps == 1_000 - assert resolved.step_scheduler.global_batch_size == 1 - assert resolved.learning_rate == 3.0e-5 - - -@pytest.mark.parametrize( - ("legacy_key", "replacement"), - [ - ("max_steps", "step_scheduler.max_steps"), - ("global_batch_size", "step_scheduler.global_batch_size"), - ("checkpoint_every_steps", "step_scheduler.ckpt_every_steps"), - ("log_every_steps", "step_scheduler.log_every"), - ], -) -def test_legacy_training_lifecycle_keys_are_rejected(tmp_path, legacy_key, replacement) -> None: - raw = _raw_config(tmp_path) - raw["training"] = {legacy_key: 2} - with pytest.raises(ValueError, match=replacement.replace(".", r"\.")): - resolve_pdd_recipe_config(raw) - - -@pytest.mark.parametrize( - ("field", "value", "message"), - [ - ("lr_decay_style", "cosine", "lr_decay_style='constant'"), - ("lr_warmup_steps", 1, "lr_warmup_steps=0"), - ("min_lr", 1.0e-5, "min_lr must equal optim.learning_rate"), - ], -) -def test_nonconstant_lr_declarations_are_rejected(tmp_path, field, value, message) -> None: - raw = _raw_config(tmp_path) - raw["lr_scheduler"][field] = value - with pytest.raises(ValueError, match=message): - resolve_pdd_recipe_config(raw) - - -@pytest.mark.parametrize("field", ["weight_decay", "betas", "eps"]) -def test_legacy_optimizer_fields_are_rejected(tmp_path, field) -> None: - raw = _raw_config(tmp_path) - raw["optim"][field] = { - "weight_decay": 0.01, - "betas": [0.9, 0.999], - "eps": 1.0e-8, - }[field] - with pytest.raises(ValueError, match=rf"optim\.optimizer\.{field}"): - resolve_pdd_recipe_config(raw) + with pytest.raises(ValueError, match="Prepare the Qwen PDD student"): + _validate_prepared_student(_PreparedStudent(out_features=4), config) -def test_pdd_rejects_external_split_manifest(tmp_path) -> None: - raw = _raw_config(tmp_path) - raw["data"] = {"dataloader": {"metadata_index": "metadata_train.json"}} - with pytest.raises(ValueError, match="metadata_index is unsupported"): - resolve_pdd_recipe_config(raw) +@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) - -def test_pdd_finetune_namespace_module_help() -> None: - environment = os.environ.copy() - environment["PYTHONDONTWRITEBYTECODE"] = "1" - result = subprocess.run( - [sys.executable, "-m", "examples.diffusers.fastgen.pdd.finetune", "--help"], - cwd=_REPO_ROOT, - env=environment, - check=True, - capture_output=True, - text=True, - ) - assert "Qwen-Image PDD training" in result.stdout - - -@pytest.mark.parametrize( - ("scope", "name", "value", "message"), - [ - ("model", "transformer_engine_linear", True, "TE-linear"), - ("model", "peft", {"rank": 8}, "PEFT/LoRA"), - ("root", "peft_cfg", {"rank": 8}, "PEFT/LoRA"), - ("model", "guidance_embeds", True, "guidance embeddings"), - ("model", "device_map", "auto", "device_map"), - ("model", "quantization_config", {"bits": 8}, "quantization_config"), - ], -) -def test_incompatible_modes_fail_during_config_resolution( - tmp_path, scope, name, value, message -) -> None: - raw = _raw_config(tmp_path) - target = raw if scope == "root" else raw[scope] - target[name] = value - - with pytest.raises(ValueError, match=message): - resolve_pdd_recipe_config(raw) - - -def test_model_revision_and_strict_compute_contract(tmp_path) -> None: - raw = _raw_config(tmp_path) - raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" - raw["model"]["revision"] = "a" * 40 - - config = resolve_pdd_recipe_config(raw) - - assert config.model_revision == "a" * 40 - assert config.dtype == torch.bfloat16 - assert config.fuse_qkv_projections is False - - raw["model"]["torch_dtype"] = "float32" - with pytest.raises(ValueError, match="torch_dtype='bfloat16'"): - resolve_pdd_recipe_config(raw) - - raw["model"]["torch_dtype"] = "bfloat16" - raw["model"]["fuse_qkv_projections"] = True - with pytest.raises(ValueError, match="does not support QKV fusion"): - resolve_pdd_recipe_config(raw) - - -@pytest.mark.parametrize("revision", [None, "main", "A" * 40, "a" * 39]) -def test_remote_model_requires_exact_lowercase_commit(tmp_path, revision) -> None: - raw = _raw_config(tmp_path) - raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" - raw["model"]["revision"] = revision - - with pytest.raises(ValueError, match="exact lowercase 40-character"): - resolve_pdd_recipe_config(raw) - - -def test_model_source_resolution_requires_the_requested_snapshot(tmp_path, monkeypatch) -> None: - commit = "a" * 40 - raw = _raw_config(tmp_path) - raw["model"]["pretrained_model_name_or_path"] = "Qwen/Qwen-Image" - raw["model"]["revision"] = commit - config = resolve_pdd_recipe_config(raw) - snapshot = tmp_path / "hub" / "snapshots" / commit - snapshot.mkdir(parents=True) - calls = [] - - def matching_snapshot(model_id, *, revision): - calls.append((model_id, revision)) - return str(snapshot) - - monkeypatch.setattr("huggingface_hub.snapshot_download", matching_snapshot) - assert _resolve_model_source(config) == str(snapshot.resolve()) - assert calls == [("Qwen/Qwen-Image", commit)] - _require_immutable_model_source(config, context="test") - - wrong = snapshot.with_name("b" * 40) - wrong.mkdir() - monkeypatch.setattr("huggingface_hub.snapshot_download", lambda *_args, **_kwargs: str(wrong)) - with pytest.raises(RuntimeError, match=r"does not match model\.revision"): - _resolve_model_source(config) - - -def test_local_model_source_is_limited_to_low_level_setup(tmp_path, monkeypatch) -> None: - config = resolve_pdd_recipe_config(_raw_config(tmp_path)) - monkeypatch.setattr( - "huggingface_hub.snapshot_download", - lambda *_args, **_kwargs: pytest.fail("local model resolution must not access the Hub"), - ) - - assert _resolve_model_source(config) == str(tmp_path.resolve()) - with pytest.raises(ValueError, match="Checkpointed PDD training requires"): - _require_immutable_model_source(config, context="Checkpointed PDD training") - - -def test_non_dp_parallelism_is_rejected(tmp_path) -> None: - raw = _raw_config(tmp_path) - - raw["fsdp"]["tp_size"] = 2 - with pytest.raises(ValueError, match="tp_size must be 1"): - resolve_pdd_recipe_config(raw) - - -@pytest.mark.parametrize( - ("section", "name", "value", "message"), - [ - ( - "step_scheduler", - "save_checkpoint_every_epoch", - True, - "save_checkpoint_every_epoch=false", - ), - ("training_health", "max_grad_norm", 0.0, "max_grad_norm must be > 0"), - ("validation", "every_steps", 0, "validation.every_steps"), - ("guidance", "rescale", 1.1, "does not support guidance overrides"), - ("optimizer", "betas", [0.9, 1.0], "optim.optimizer.betas values"), - ("optimizer", "eps", 0.0, "optim.optimizer.eps must be > 0"), - ], -) -def test_training_config_gates_fail_during_resolution( - tmp_path, section, name, value, message -) -> None: - raw = _raw_config(tmp_path) - target = ( - raw["optim"].setdefault("optimizer", {}) - if section == "optimizer" - else raw.setdefault(section, {}) - ) - target[name] = value - - with pytest.raises(ValueError, match=message): - resolve_pdd_recipe_config(raw) - - -def test_restore_requires_enabled_checkpointing(tmp_path) -> None: - raw = _raw_config(tmp_path) - raw["checkpoint"]["enabled"] = False - raw["checkpoint"]["restore_from"] = "LATEST" - - with pytest.raises(ValueError, match=r"restore_from requires checkpoint\.enabled=true"): - resolve_pdd_recipe_config(raw) - - -@pytest.mark.parametrize( - ("name", "value", "message"), - [ - ("drop_last", False, "drop_last=true"), - ("dynamic_batch_size", True, "dynamic_batch_size=false"), - ("train_text_encoder", True, "cached text embeddings"), - ], -) -def test_training_dataloader_modes_are_gated_during_resolution( - tmp_path, name, value, message -) -> None: - raw = _raw_config(tmp_path) - raw["data"] = {"dataloader": {name: value}} - - with pytest.raises(ValueError, match=message): - resolve_pdd_recipe_config(raw) - - -def test_payload_hash_verification_mode_must_be_bool(tmp_path) -> None: - raw = _raw_config(tmp_path) - raw["data"]["dataloader"]["verify_payload_hashes"] = "false" - - with pytest.raises(TypeError, match=r"data\.dataloader\.verify_payload_hashes must be bool"): - resolve_pdd_recipe_config(raw) - - -def test_real_loader_manager_optimizer_and_checkpoint_restore(tmp_path) -> None: - model_dir = create_tiny_qwen_image_pipeline_dir(tmp_path) - initialize_pdd_distributed(backend="gloo", timeout_minutes=1) - raw_config = _raw_config(model_dir) - raw_config["fsdp"]["activation_checkpointing"] = True - config = resolve_pdd_recipe_config(raw_config) - - source = build_pdd_setup(config) - - assert source.lifecycle == ( - "load/select", - "pdd_conversion", - "device", - "qkv", - "parallelize", - "optimizer", - "checkpoint", - ) - assert type(source.pipe).__name__ == "QwenImagePipeline" - assert source.pipe.text_encoder is None - assert source.pipe.tokenizer is None - assert source.pipe.vae is None - assert source.pipe.transformer is source.student - assert source.student.forward.__self__ is source.student - assert source.teacher.forward.__self__ is source.teacher - assert source.teacher.forward.__func__ is source.student.forward.__func__ - assert source.student.get_submodule("proj_out") is source.projection - assert source.projection.out_features == source.projection.base_out_features * 4 - assert "proj_out.weight" in source.checkpoint_keys - assert source.student.state_dict()["proj_out.weight"].shape[0] == source.projection.out_features - policy = source.distributed_setup.strategy_config.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 source.distributed_setup.strategy_config.activation_checkpointing is False - assert source.distributed_setup.strategy_config.reshard_after_forward is True - assert config.parallel.activation_checkpointing is True - for model in (source.student, source.teacher): - require_qwen_image_mr210_forward(model) - assert model.gradient_checkpointing is False - assert len(model.transformer_blocks) == 6 - assert all(isinstance(block, CheckpointWrapper) for block in model.transformer_blocks) - assert all( - block.checkpoint_impl is CheckpointImpl.NO_REENTRANT - for block in model.transformer_blocks - ) - assert all( - type(block._checkpoint_wrapped_module) is QwenImageTransformerBlock - for block in model.transformer_blocks - ) - _, canonical_parameter_names = _canonical_parameter_names(model) - assert canonical_parameter_names == set(model.state_dict()) - assert all( - base.__module__ != "modelopt.torch.fastgen.plugins.qwen_image_pdd" - for base in type(source.student).__mro__ - ) - assert not any(parameter.requires_grad for parameter in source.teacher.parameters()) - optimizer_parameters = [ - parameter for group in source.optimizer.param_groups for parameter in group["params"] - ] - assert any(parameter is source.projection.weight for parameter in optimizer_parameters) - assert set(source.optimizer.state) == set(optimizer_parameters) - assert all( - source.optimizer.state[parameter]["step"].item() == 0 for parameter in optimizer_parameters - ) - unused_parameter = next( - parameter for parameter in optimizer_parameters if parameter is not source.projection.weight - ) - unused_name = next( - name - for name, parameter in source.student.named_parameters() - if parameter is unused_parameter - ) - # Diffusers 0.38 accepts the Qwen object API but currently performs no effective fusion. - assert not any( - getattr(module, "fused_projections", False) for module in source.student.modules() - ) - - source.optimizer.zero_grad(set_to_none=True) - # Exercise strict stock-DCP restore after a partial-gradient update. Eager step-zero state must - # retain exact lazy-Adam semantics for untouched parameters while keeping every DCP key present. - source.projection.weight.float().square().mean().backward() - source.optimizer.step() - expected_weight = source.projection.weight.detach().clone() - expected_exp_avg = source.optimizer.state[source.projection.weight]["exp_avg"].clone() - expected_student_state = { - name: value.detach().clone() for name, value in source.student.state_dict().items() - } - expected_optimizer_state = copy.deepcopy(source.optimizer.state_dict()) - assert source.optimizer.state[unused_parameter]["step"].item() == 0 - checkpoint_root = tmp_path / "checkpoint" - source.checkpointer.save_model(source.student, str(checkpoint_root)) - source.checkpointer.save_optimizer(source.optimizer, source.student, str(checkpoint_root)) - model_metadata_keys = set( - FileSystemReader(str(checkpoint_root / "model")).read_metadata().state_dict_metadata - ) - optimizer_metadata_keys = set( - FileSystemReader(str(checkpoint_root / "optim")).read_metadata().state_dict_metadata - ) - assert model_metadata_keys - assert optimizer_metadata_keys - assert all("_checkpoint_wrapped_module" not in key for key in model_metadata_keys) - assert all("_checkpoint_wrapped_module" not in key for key in optimizer_metadata_keys) - assert any(key.endswith("proj_out.weight") for key in model_metadata_keys) - - export_setup = build_pdd_export_setup(config) - assert export_setup.lifecycle == ( - "load/select", - "pdd_conversion", - "device", - "qkv", - "parallelize", - "checkpoint", - ) - assert export_setup.metadata == source.metadata - assert export_setup.checkpoint_keys == source.checkpoint_keys - assert not hasattr(export_setup, "optimizer") - require_qwen_image_mr210_forward(export_setup.student) - export_policy = export_setup.distributed_setup.strategy_config.mp_policy - assert export_policy.param_dtype == torch.bfloat16 - assert export_policy.reduce_dtype == torch.float32 - assert export_policy.output_dtype == torch.bfloat16 - assert export_policy.cast_forward_inputs is False - assert export_setup.student.forward.__self__ is export_setup.student - export_setup.checkpointer.load_model( - export_setup.student, - str(checkpoint_root / "model"), - ) - torch.testing.assert_close( - export_setup.student.state_dict()["proj_out.weight"], - expected_weight, - ) - export_state = export_setup.student.state_dict() - assert export_state.keys() == expected_student_state.keys() - for name, expected in expected_student_state.items(): - torch.testing.assert_close(export_state[name], expected, rtol=0, atol=0) - - destination = build_pdd_setup(config) - assert destination.metadata == source.metadata - destination_projection = destination.projection - destination_unused = dict(destination.student.named_parameters())[unused_name] - destination_weight_id = id(destination_projection.weight) - destination.checkpointer.load_model( - destination.student, - str(checkpoint_root / "model"), - ) - destination.checkpointer.load_optimizer( - destination.optimizer, - destination.student, - str(checkpoint_root), - ) - - assert destination.student.get_submodule("proj_out") is destination_projection - assert id(destination_projection.weight) == destination_weight_id - torch.testing.assert_close(destination_projection.weight, expected_weight) - destination_state = destination.student.state_dict() - assert destination_state.keys() == expected_student_state.keys() - for name, expected in expected_student_state.items(): - torch.testing.assert_close(destination_state[name], expected, rtol=0, atol=0) - destination_optimizer_state = destination.optimizer.state_dict() - assert destination_optimizer_state["param_groups"] == expected_optimizer_state["param_groups"] - assert destination_optimizer_state["state"].keys() == expected_optimizer_state["state"].keys() - for parameter_index, expected_state in expected_optimizer_state["state"].items(): - actual_state = destination_optimizer_state["state"][parameter_index] - assert actual_state.keys() == expected_state.keys() - for name, expected in expected_state.items(): - torch.testing.assert_close(actual_state[name], expected, rtol=0, atol=0) - torch.testing.assert_close( - destination.optimizer.state[destination_projection.weight]["exp_avg"], - expected_exp_avg, + per_sample, loss, prediction, metrics = adapter.step( + model=nn.Identity(), + batch=_batch(), + device=torch.device("cpu"), + dtype=torch.float32, ) - assert destination.optimizer.state[destination_unused]["step"].item() == 0 - assert not destination.optimizer.state[destination_unused]["exp_avg"].count_nonzero() - assert not destination.optimizer.state[destination_unused]["exp_avg_sq"].count_nonzero() - source.checkpointer.close() - destination.checkpointer.close() - export_setup.checkpointer.close() - + loss.backward() -def test_qwen_pdd_adapter_has_no_automodel_import() -> None: - source = ( - _REPO_ROOT / "modelopt" / "torch" / "fastgen" / "plugins" / "qwen_image_pdd.py" - ).read_text() - assert "nemo_automodel" not in source + 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 diff --git a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py b/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py deleted file mode 100644 index 3e70b216f45..00000000000 --- a/tests/examples/diffusers/fastgen/test_pdd_training_lifecycle.py +++ /dev/null @@ -1,900 +0,0 @@ -# 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. - -"""Tests for direct updates, committed cursors, and strict PDD resume.""" - -from __future__ import annotations - -import copy -import hashlib -import json -import math -import pathlib -import shutil -import sys -from types import SimpleNamespace - -import pytest -import torch -import torch.distributed as dist - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) -if str(_FASTGEN_DIR) not in sys.path: - sys.path.insert(0, str(_FASTGEN_DIR)) -if str(pathlib.Path(__file__).parent) not in sys.path: - sys.path.insert(0, str(pathlib.Path(__file__).parent)) - -from fastgen_data.replayable_sampler import ReplayableBatchSampler -from pdd.checkpoint import ( - PDDCheckpointManager, - build_pdd_checkpoint_identity, - resolve_pdd_training_checkpoint, -) -from pdd.recipe import PDDDiffusionRecipe, initialize_pdd_distributed -from pdd.training import prepare_qwen_pdd_batch -from pdd_test_utils import SamplerDataset, build_toy_lifecycle, make_batch, ordered_id_sha256 - -from modelopt.torch.fastgen.plugins.qwen_image_pdd import QWEN_IMAGE_PDD_EXECUTION - - -def _released_sampler(sample_ids: tuple[str, ...]) -> ReplayableBatchSampler: - sampler_module = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") - dataset = SamplerDataset(sample_ids) - sampler = sampler_module.SequentialBucketSampler( - dataset, - base_batch_size=1, - base_resolution=(64, 64), - drop_last=True, - shuffle_buckets=True, - shuffle_within_bucket=True, - dynamic_batch_size=False, - seed=31, - num_replicas=1, - rank=0, - ) - return ReplayableBatchSampler(sampler) - - -def _run_next(lifecycle, sampler): - sample_ids = sampler.expected_next_sample_ids() - assert sample_ids - diagnostics = lifecycle.trainer.train_step(make_batch(sample_ids)) - lifecycle.scheduler.step() - sampler.commit(sample_ids) - if sampler.remaining_batches == 0: - sampler.set_epoch(sampler.epoch + 1) - return sample_ids, diagnostics - - -def _checkpointer(lifecycle, checkpoint_dir): - checkpoint_module = pytest.importorskip("nemo_automodel.components.checkpoint.config") - config = checkpoint_module.CheckpointingConfig( - enabled=True, - checkpoint_dir=str(checkpoint_dir), - model_save_format="torch_save", - model_repo_id="synthetic-pdd-toy", - save_consolidated=False, - is_peft=False, - model_state_dict_keys=list(lifecycle.student.state_dict()), - ) - return config.build(dp_rank=0, tp_rank=0, pp_rank=0, moe_mesh=None) - - -def _identity( - lifecycle, - scheduler, - sample_ids, - *, - qwen_image_execution=QWEN_IMAGE_PDD_EXECUTION, -): - return build_pdd_checkpoint_identity( - qwen_image_execution=qwen_image_execution, - metadata=lifecycle.metadata, - model_id="synthetic-pdd-toy", - model_revision="a" * 40, - guidance_scale=None, - ordered_train_id_sha256=ordered_id_sha256(sample_ids), - ordered_heldout_id_sha256="1" * 64, - dataset_snapshot_sha256="2" * 64, - local_batch_size=1, - grad_accumulation_steps=1, - training_seed=1234, - validation_seed=2026, - validation_every_steps=100, - max_grad_norm=0.5, - zero_grad_warmup_steps=0, - activation_checkpointing=False, - dtype="float32", - optimizer=lifecycle.optimizer, - scheduler=scheduler, - ) - - -def test_checkpoint_identity_rejects_unbound_qwen_execution() -> None: - lifecycle = build_toy_lifecycle() - with pytest.raises(ValueError, match="qwen_image_execution"): - _identity( - lifecycle, - lifecycle.scheduler, - ("sample-0",), - qwen_image_execution="canonical_diffusers", - ) - - -class _StepSchedulerStub: - def __init__(self, trainer, sampler) -> None: - self.trainer = trainer - self.sampler = sampler - self.loaded_state = None - - def state_dict(self): - return {"step": self.trainer.completed_steps, "epoch": self.sampler.epoch} - - def load_state_dict(self, state): - self.loaded_state = dict(state) - - -class _RecordingCheckpointManager(PDDCheckpointManager): - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - self.save_calls = [] - - def save(self): - self.save_calls.append( - { - "live_step": self.step_scheduler.step, - "serialized_step": self.step_scheduler.state_dict()["step"], - "trainer_step": self.trainer.completed_steps, - "live_epoch": self.step_scheduler.epoch, - "sampler_epoch": self.sampler.epoch, - } - ) - return super().save() - - -def _manager(root, lifecycle, sampler, rng): - checkpointer = _checkpointer(lifecycle, root) - step_scheduler = _StepSchedulerStub(lifecycle.trainer, sampler) - manager = PDDCheckpointManager( - root=root, - checkpointer=checkpointer, - model=lifecycle.student, - optimizer=lifecycle.optimizer, - scheduler=lifecycle.scheduler, - step_scheduler=step_scheduler, - trainer=lifecycle.trainer, - sampler=sampler, - rng=rng, - identity=_identity(lifecycle, lifecycle.scheduler, tuple(f"sample-{i}" for i in range(8))), - ) - return manager, checkpointer - - -def _optimizer_state_by_name(lifecycle): - names = {parameter: name for name, parameter in lifecycle.student.named_parameters()} - return { - names[parameter]: { - key: value.detach().clone() if isinstance(value, torch.Tensor) else value - for key, value in state.items() - } - for parameter, state in lifecycle.optimizer.state.items() - } - - -def _refresh_complete_marker(checkpoint: pathlib.Path) -> None: - manifest_path = checkpoint / "manifest.json" - marker = { - "schema_version": 1, - "manifest_sha256": hashlib.sha256(manifest_path.read_bytes()).hexdigest(), - } - (checkpoint / "COMPLETE").write_text(json.dumps(marker, indent=2, sort_keys=True) + "\n") - - -def test_replayable_sampler_commits_consumed_batches_not_prefetch() -> None: - sample_ids = tuple(f"sample-{index}" for index in range(8)) - sampler = _released_sampler(sample_ids) - first_expected = sampler.expected_next_sample_ids() - iterator = iter(sampler) - next(iterator) - next(iterator) - - assert sampler.committed_batches == 0 - assert sampler.expected_next_sample_ids() == first_expected - sampler.commit(first_expected) - state = sampler.state_dict() - - restored = _released_sampler(sample_ids) - restored.load_state_dict(state) - assert restored.state_dict() == state - next_indices = next(iter(restored)) - assert tuple(restored.dataset.metadata[index]["sample_id"] for index in next_indices) == tuple( - state["next_sample_ids"] - ) - bad_hash = dict(state, plan_sha256="0" * 64) - with pytest.raises(RuntimeError, match="plan hash"): - _released_sampler(sample_ids).load_state_dict(bad_hash) - bad_ids = dict(state, next_sample_ids=["wrong-id"]) - with pytest.raises(RuntimeError, match="next sample IDs"): - _released_sampler(sample_ids).load_state_dict(bad_ids) - - -def test_automodel_step_scheduler_serializes_the_completed_yielded_step() -> None: - scheduler_module = pytest.importorskip("nemo_automodel.components.training.step_scheduler") - scheduler = scheduler_module.StepScheduler( - global_batch_size=1, - local_batch_size=1, - dp_size=1, - ckpt_every_steps=2, - save_checkpoint_every_epoch=False, - dataloader=[{"sample": 0}, {"sample": 1}], - val_every_steps=None, - start_step=0, - start_epoch=0, - num_epochs=1, - max_steps=2, - ) - - iterator = iter(scheduler) - assert next(iterator) == [{"sample": 0}] - assert scheduler.step == 0 - assert scheduler.state_dict() == {"step": 1, "epoch": 0} - assert next(iterator) == [{"sample": 1}] - assert scheduler.step == 1 - assert scheduler.is_last_step - assert scheduler.state_dict() == {"step": 2, "epoch": 0} - with pytest.raises(StopIteration): - next(iterator) - - -def _build_recipe_loop( - root, - sample_ids, - *, - max_steps, - num_epochs, - ckpt_every_steps, - restore_from=None, -): - if not dist.is_initialized(): - initialize_pdd_distributed(backend="gloo", timeout_minutes=1) - scheduler_module = pytest.importorskip("nemo_automodel.components.training.step_scheduler") - rng_module = pytest.importorskip("nemo_automodel.components.training.rng") - - lifecycle = build_toy_lifecycle() - sampler = _released_sampler(sample_ids) - rng = rng_module.StatefulRNG(1234, ranked=True) - step_scheduler = scheduler_module.StepScheduler( - global_batch_size=1, - local_batch_size=1, - dp_size=1, - ckpt_every_steps=ckpt_every_steps, - save_checkpoint_every_epoch=False, - dataloader=[None] * len(sampler), - val_every_steps=None, - start_step=0, - start_epoch=0, - num_epochs=num_epochs, - max_steps=max_steps, - ) - checkpointer = _checkpointer(lifecycle, root) - manager = _RecordingCheckpointManager( - root=root, - checkpointer=checkpointer, - model=lifecycle.student, - optimizer=lifecycle.optimizer, - scheduler=lifecycle.scheduler, - step_scheduler=step_scheduler, - trainer=lifecycle.trainer, - sampler=sampler, - rng=rng, - identity=_identity(lifecycle, lifecycle.scheduler, sample_ids), - ) - resume = manager.load(restore_from) - events = SimpleNamespace(first_ids=[], diagnostics=[], validation_steps=[]) - - recipe = object.__new__(PDDDiffusionRecipe) - recipe.config = SimpleNamespace( - step_scheduler=SimpleNamespace(local_batch_size=1, log_every=1), - validation=SimpleNamespace(every_steps=10_000), - checkpoint=SimpleNamespace(enabled=True), - device=torch.device("cpu"), - ) - recipe.training = SimpleNamespace( - pipeline=lifecycle.pipeline, - trainer=lifecycle.trainer, - scheduler=lifecycle.scheduler, - rng=rng, - ) - recipe.setup_artifacts = SimpleNamespace(checkpointer=checkpointer) - recipe.step_scheduler = step_scheduler - recipe.checkpoint_manager = manager - recipe.sampler = sampler - recipe.resume = resume - recipe.resume_pending = resume is not None - recipe.rank = 0 - recipe.world_size = 1 - - def prepared_batches(): - # Bind this iterator to the current sampler plan. The production loader iterator also - # exhausts after that plan even though the recipe commits the sampler into its next epoch. - for _ in range(sampler.remaining_batches): - expected_ids = sampler.expected_next_sample_ids() - if recipe.resume_pending: - assert recipe.resume is not None - recipe.resume.verify_first_batch(expected_ids) - events.first_ids.append(expected_ids) - offset = sum(ord(character) for character in expected_ids[0]) / 10_000 - yield make_batch(expected_ids, offset=offset), expected_ids - - recipe.__dict__["_prepared_training_batches"] = prepared_batches - recipe.__dict__["_run_validation"] = events.validation_steps.append - recipe.__dict__["_log_step"] = lambda diagnostics, _data_wait, _step_time: ( - events.diagnostics.append(diagnostics) - ) - return SimpleNamespace( - recipe=recipe, - lifecycle=lifecycle, - sampler=sampler, - step_scheduler=step_scheduler, - manager=manager, - resume=resume, - events=events, - ) - - -def test_recipe_loop_saves_periodic_and_max_step_checkpoints_once(tmp_path) -> None: - run = _build_recipe_loop( - tmp_path / "periodic", - tuple(f"sample-{index}" for index in range(8)), - max_steps=3, - num_epochs=4, - ckpt_every_steps=2, - ) - run.recipe.run_train_validation_loop() - - assert [call["trainer_step"] for call in run.manager.save_calls] == [2, 3] - assert [call["live_step"] for call in run.manager.save_calls] == [1, 2] - assert [call["serialized_step"] for call in run.manager.save_calls] == [2, 3] - assert sorted(path.name for path in (tmp_path / "periodic").glob("step_*")) == [ - "step_00000002", - "step_00000003", - ] - assert run.events.validation_steps == [3] - - -def test_recipe_loop_epoch_final_save_normalizes_and_restores_epoch(tmp_path) -> None: - root = tmp_path / "epoch" - sample_ids = ("sample-0", "sample-1") - source = _build_recipe_loop( - root, - sample_ids, - max_steps=100, - num_epochs=1, - ckpt_every_steps=100, - ) - source.recipe.run_train_validation_loop() - - assert len(source.manager.save_calls) == 1 - assert source.manager.save_calls[0] == { - "live_step": 1, - "serialized_step": 2, - "trainer_step": 2, - "live_epoch": 0, - "sampler_epoch": 1, - } - manifest = json.loads((root / "step_00000002" / "manifest.json").read_text()) - assert manifest["step_scheduler"] == {"step": 2, "epoch": 1} - - resumed = _build_recipe_loop( - root, - sample_ids, - max_steps=3, - num_epochs=2, - ckpt_every_steps=100, - restore_from="step_00000002", - ) - assert resumed.step_scheduler.step == 2 - assert resumed.step_scheduler.epoch == resumed.sampler.epoch == 1 - assert resumed.resume is not None - expected_ids = resumed.resume.expected_next_sample_ids - resumed.recipe.run_train_validation_loop() - assert resumed.events.first_ids[0] == expected_ids - assert (root / "step_00000003" / "COMPLETE").is_file() - - -def test_recipe_loop_resume_matches_uninterrupted_next_update(tmp_path) -> None: - sample_ids = tuple(f"sample-{index}" for index in range(4)) - control = _build_recipe_loop( - tmp_path / "control", - sample_ids, - max_steps=2, - num_epochs=2, - ckpt_every_steps=100, - ) - control.recipe.run_train_validation_loop() - - staged_root = tmp_path / "staged" - first = _build_recipe_loop( - staged_root, - sample_ids, - max_steps=1, - num_epochs=2, - ckpt_every_steps=100, - ) - first.recipe.run_train_validation_loop() - resumed = _build_recipe_loop( - staged_root, - sample_ids, - max_steps=2, - num_epochs=2, - ckpt_every_steps=100, - restore_from="step_00000001", - ) - assert resumed.resume is not None - expected_next_ids = resumed.resume.expected_next_sample_ids - resumed.recipe.run_train_validation_loop() - - assert resumed.events.first_ids == [expected_next_ids] - assert resumed.events.first_ids[0] == control.events.first_ids[1] - assert resumed.events.diagnostics == [control.events.diagnostics[1]] - for name, tensor in resumed.lifecycle.student.state_dict().items(): - torch.testing.assert_close( - tensor, - control.lifecycle.student.state_dict()[name], - rtol=0, - atol=0, - ) - actual_optimizer = _optimizer_state_by_name(resumed.lifecycle) - expected_optimizer = _optimizer_state_by_name(control.lifecycle) - assert actual_optimizer.keys() == expected_optimizer.keys() - for name in actual_optimizer: - for key, actual in actual_optimizer[name].items(): - expected = expected_optimizer[name][key] - if isinstance(actual, torch.Tensor): - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - else: - assert actual == expected - - -def test_qwen_batch_preparation_preserves_ids_masks_and_negative_condition() -> None: - batch = { - "image_latents": torch.ones(2, 3, 4, 4), - "text_embeddings": torch.ones(2, 5, 6), - "text_embeddings_mask": torch.ones(2, 5, dtype=torch.bool), - "negative_text_embeddings": torch.zeros(5, 6), - "negative_text_embeddings_mask": torch.ones(5, dtype=torch.bool), - "metadata": {"sample_ids": ["qwen-a", "qwen-b"]}, - } - - prepared = prepare_qwen_pdd_batch( - batch, - device=torch.device("cpu"), - dtype=torch.float32, - require_negative_condition=True, - expected_latent_channels=3, - expected_condition_features=6, - ) - - assert prepared.sample_ids == ("qwen-a", "qwen-b") - assert prepared.valid_mask == (True, True) - assert prepared.data.shape == (2, 3, 4, 4) - assert prepared.condition[0].shape == (2, 5, 6) - assert prepared.negative_condition is not None - assert prepared.negative_condition[0].shape == (2, 5, 6) - - without_negative = dict(batch) - without_negative.pop("negative_text_embeddings") - without_negative.pop("negative_text_embeddings_mask") - with pytest.raises(ValueError, match="requires negative prompt conditioning"): - prepare_qwen_pdd_batch( - without_negative, - device=torch.device("cpu"), - dtype=torch.float32, - require_negative_condition=True, - expected_latent_channels=3, - expected_condition_features=6, - ) - - -def test_qwen_batch_preparation_rejects_every_pre_model_shape_and_dtype_mismatch() -> None: - base = { - "image_latents": torch.ones(2, 3, 4, 4), - "text_embeddings": torch.ones(2, 5, 6), - "text_embeddings_mask": torch.ones(2, 5, dtype=torch.long), - "negative_text_embeddings": torch.zeros(2, 5, 6), - "negative_text_embeddings_mask": torch.ones(2, 5, dtype=torch.bool), - "metadata": {"sample_ids": ["qwen-a", "qwen-b"]}, - } - cases = ( - ("image_latents", torch.ones(2, 3, 4, 4, dtype=torch.long), "floating-point dtype"), - ("image_latents", torch.ones(2, 4, 4, 4), "exactly 3 channels"), - ("image_latents", torch.ones(2, 3, 3, 4), "positive even spatial dimensions"), - ("text_embeddings", torch.ones(2, 5, 6, dtype=torch.long), "floating-point dtype"), - ("text_embeddings", torch.ones(2, 5, 7), "exactly 6 features"), - ("text_embeddings", torch.ones(2, 5, 6, 1), "must be 2D or 3D"), - ("text_embeddings_mask", torch.ones(2, 5), "integer or boolean dtype"), - ("text_embeddings_mask", torch.ones(2, 4, dtype=torch.long), "sequence length"), - ( - "negative_text_embeddings_mask", - torch.ones(2, 5), - "integer or boolean dtype", - ), - ) - - for field, value, message in cases: - batch = dict(base) - batch[field] = value - with pytest.raises((TypeError, ValueError), match=message): - prepare_qwen_pdd_batch( - batch, - device=torch.device("cpu"), - dtype=torch.float32, - require_negative_condition=True, - expected_latent_channels=3, - expected_condition_features=6, - ) - - -def test_two_direct_updates_have_finite_gradients_updates_and_targeted_coverage() -> None: - lifecycle = build_toy_lifecycle(weight_decay=0.02) - first_batch = make_batch(("a", "b")) - before = [parameter.detach().clone() for parameter in lifecycle.student.parameters()] - first = lifecycle.trainer.train_step( - first_batch, - noise=torch.tensor([[0.25, -0.5, 1.0], [-0.25, 0.5, -1.0]]), - n=torch.tensor([0, 1]), - k=torch.tensor([1, 3]), - ) - lifecycle.scheduler.step() - actual_update = math.sqrt( - sum( - (parameter.detach() - saved).double().square().sum().item() - for parameter, saved in zip(lifecycle.student.parameters(), before) - ) - ) - before_norm = math.sqrt(sum(saved.double().square().sum().item() for saved in before)) - - assert math.isfinite(first.loss) and first.loss > 0 - assert math.isfinite(first.grad_norm) and first.grad_norm > 0 - assert first.pdd_projection_update_ratio is not None - assert first.pdd_projection_update_ratio > 0 - assert math.isfinite(first.student_teacher_velocity_rms_ratio) - assert first.student_teacher_velocity_rms_ratio >= 0 - assert first.student_adamw_nominal_update_ratio == pytest.approx( - actual_update / before_norm, - rel=2e-5, - abs=1e-8, - ) - assert all(parameter.grad is None for parameter in lifecycle.teacher.parameters()) - - second = lifecycle.trainer.train_step( - make_batch(("c", "d"), offset=0.25), - noise=torch.tensor([[0.75, 0.0, -0.5], [-0.75, 0.0, 0.5]]), - n=torch.tensor([2, 3]), - k=torch.tensor([2, 3]), - ) - lifecycle.scheduler.step() - assert second.completed_step == 2 - lifecycle.trainer.coverage.require_pairs([(0, 1), (1, 3), (2, 2), (3, 3)]) - - -def test_training_hard_aborts_for_teacher_gradient_zero_gradient_and_missing_coverage( - monkeypatch, -) -> None: - teacher_gradient = build_toy_lifecycle() - teacher_gradient.teacher.scale.grad = torch.ones_like(teacher_gradient.teacher.scale) - with pytest.raises(RuntimeError, match="teacher received a gradient"): - teacher_gradient.trainer.train_step(make_batch(("teacher-grad",))) - - zero_gradient = build_toy_lifecycle(zero_student_gradient=True, weight_decay=0.0) - first = zero_gradient.trainer.train_step(make_batch(("zero-1",))) - assert first.grad_norm == 0.0 - with pytest.raises(RuntimeError, match="zero for two consecutive"): - zero_gradient.trainer.train_step(make_batch(("zero-2",))) - - with pytest.raises(RuntimeError, match="did not cover"): - zero_gradient.trainer.coverage.require_pairs([(3, 3)]) - - zero_actual_update = build_toy_lifecycle() - monkeypatch.setattr(zero_actual_update.trainer, "_projection_update_ratio", lambda before: 0.0) - with pytest.raises(RuntimeError, match="zero actual projection update"): - zero_actual_update.trainer.train_step(make_batch(("zero-actual-update",))) - - nonfinite = build_toy_lifecycle() - bad_batch = make_batch(("nan",)) - bad_batch.data.fill_(float("nan")) - with pytest.raises(FloatingPointError, match="non-finite"): - nonfinite.trainer.train_step(bad_batch) - - nonfinite_gradient = build_toy_lifecycle() - handle = nonfinite_gradient.projection.weight.register_hook( - lambda gradient: torch.full_like(gradient, float("inf")) - ) - with pytest.raises(FloatingPointError, match="gradient became non-finite"): - nonfinite_gradient.trainer.train_step(make_batch(("inf-gradient",))) - handle.remove() - - invalid_support = build_toy_lifecycle() - with pytest.raises(RuntimeError, match="k must satisfy"): - invalid_support.trainer.train_step( - make_batch(("invalid-support",)), - n=torch.tensor([2]), - k=torch.tensor([1]), - ) - - nonfinite_update = build_toy_lifecycle() - nonfinite_update.optimizer.param_groups[0]["eps"] = float("nan") - with pytest.raises(FloatingPointError, match="parameter update became non-finite"): - nonfinite_update.trainer.train_step( - make_batch(("inf-update",)), - measure_updates=False, - ) - - -def test_stock_dcp_resume_recovers_rng_scheduler_cursor_and_next_loss(tmp_path) -> None: - pytest.importorskip("nemo_automodel") - rng_module = pytest.importorskip("nemo_automodel.components.training.rng") - if not torch.distributed.is_initialized(): - initialize_pdd_distributed(backend="gloo", timeout_minutes=1) - sample_ids = tuple(f"sample-{index}" for index in range(8)) - - source = build_toy_lifecycle() - source_sampler = _released_sampler(sample_ids) - source_rng = rng_module.StatefulRNG(1234, ranked=True) - source_manager, source_checkpointer = _manager( - tmp_path / "checkpoints", source, source_sampler, source_rng - ) - _run_next(source, source_sampler) - _run_next(source, source_sampler) - checkpoint = source_manager.save() - assert checkpoint.name == "step_00000002" - assert source_manager.identity["data"]["dataset_snapshot_sha256"] == "2" * 64 - assert source_manager.identity["qwen_image"] == {"execution": "fastgen_mr210"} - assert source_manager.identity["training"] == { - "seed": 1234, - "validation_seed": 2026, - "validation_every_steps": 100, - "max_grad_norm": 0.5, - "zero_grad_warmup_steps": 0, - "activation_checkpointing": False, - } - expected_next_ids, expected_next = _run_next(source, source_sampler) - expected_model = copy.deepcopy(source.student.state_dict()) - expected_optimizer = _optimizer_state_by_name(source) - - destination = build_toy_lifecycle() - destination_sampler = _released_sampler(sample_ids) - destination_rng = rng_module.StatefulRNG(9999, ranked=True) - destination_manager, destination_checkpointer = _manager( - tmp_path / "checkpoints", - destination, - destination_sampler, - destination_rng, - ) - resume = destination_manager.load("LATEST") - assert resume is not None - assert resume.completed_steps == 2 - assert resume.sample_slots_consumed == 2 - assert resume.expected_next_sample_ids == expected_next_ids - resume.verify_first_batch(destination_sampler.expected_next_sample_ids()) - with pytest.raises(RuntimeError, match="first resumed sample IDs"): - resume.verify_first_batch(("wrong-id",)) - actual_ids, actual_next = _run_next(destination, destination_sampler) - - assert actual_ids == expected_next_ids - assert actual_next.n == expected_next.n - assert actual_next.k == expected_next.k - assert actual_next.loss == expected_next.loss - assert actual_next.learning_rate == expected_next.learning_rate - for name, tensor in destination.student.state_dict().items(): - torch.testing.assert_close(tensor, expected_model[name], rtol=0, atol=0) - actual_optimizer = _optimizer_state_by_name(destination) - assert actual_optimizer.keys() == expected_optimizer.keys() - for name in actual_optimizer: - for key, actual in actual_optimizer[name].items(): - expected = expected_optimizer[name][key] - if isinstance(actual, torch.Tensor): - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - else: - assert actual == expected - - resumed_checkpoint = destination_manager.save() - assert resumed_checkpoint.name == "step_00000003" - third = build_toy_lifecycle() - third_sampler = _released_sampler(sample_ids) - third_rng = rng_module.StatefulRNG(7, ranked=True) - third_manager, third_checkpointer = _manager( - tmp_path / "checkpoints", third, third_sampler, third_rng - ) - second_resume = third_manager.load("LATEST") - assert second_resume is not None - assert second_resume.completed_steps == 3 - assert second_resume.expected_next_sample_ids == destination_sampler.expected_next_sample_ids() - - incomplete = tmp_path / "checkpoints" / "step_99999998" - incomplete.mkdir() - (tmp_path / "checkpoints" / "LATEST").write_text(incomplete.name + "\n") - assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() - with pytest.raises(RuntimeError, match="incomplete"): - third_manager.resolve(incomplete.name) - - mismatched = tmp_path / "checkpoints" / "step_99999999" - shutil.copytree(resumed_checkpoint, mismatched) - manifest_path = mismatched / "manifest.json" - manifest = json.loads(manifest_path.read_text()) - manifest["identity"]["model"]["id"] = "different-model" - manifest["completed_steps"] = 99999999 - manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") - _refresh_complete_marker(mismatched) - (tmp_path / "checkpoints" / "LATEST").write_text(mismatched.name + "\n") - assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() - selected, selected_manifest = resolve_pdd_training_checkpoint( - tmp_path / "checkpoints", - "LATEST", - expected_world_size=1, - expected_identity=third_manager.identity, - ) - assert selected == resumed_checkpoint.resolve() - assert selected_manifest["identity"] == third_manager.identity - moved_model_identity = copy.deepcopy(third_manager.identity) - moved_model_identity["model"]["revision"] = "b" * 40 - with pytest.raises(RuntimeError, match="identity"): - resolve_pdd_training_checkpoint( - tmp_path / "checkpoints", - resumed_checkpoint.name, - expected_world_size=1, - expected_identity=moved_model_identity, - ) - with pytest.raises(RuntimeError, match="identity"): - third_manager.resolve(mismatched.name) - - for step, execution in ((99999994, None), (99999995, "canonical_diffusers")): - incompatible = tmp_path / "checkpoints" / f"step_{step:08d}" - shutil.copytree(resumed_checkpoint, incompatible) - incompatible_manifest_path = incompatible / "manifest.json" - incompatible_manifest = json.loads(incompatible_manifest_path.read_text()) - incompatible_manifest["completed_steps"] = step - if execution is None: - incompatible_manifest["identity"].pop("qwen_image") - else: - incompatible_manifest["identity"]["qwen_image"] = {"execution": execution} - incompatible_manifest_path.write_text( - json.dumps(incompatible_manifest, indent=2, sort_keys=True) + "\n" - ) - _refresh_complete_marker(incompatible) - (tmp_path / "checkpoints" / "LATEST").write_text(incompatible.name + "\n") - assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() - with pytest.raises(RuntimeError, match="identity"): - third_manager.resolve(incompatible.name) - - missing_model = tmp_path / "checkpoints" / "step_99999996" - shutil.copytree(resumed_checkpoint, missing_model) - model_payload = next( - path for path in (missing_model / "model").iterdir() if path.name != ".metadata" - ) - model_payload.unlink() - (tmp_path / "checkpoints" / "LATEST").write_text(missing_model.name + "\n") - assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() - with pytest.raises(RuntimeError, match="DCP"): - third_manager.resolve(missing_model.name) - - corrupt_optimizer = tmp_path / "checkpoints" / "step_99999997" - shutil.copytree(resumed_checkpoint, corrupt_optimizer) - optimizer_payload = next( - path for path in (corrupt_optimizer / "optim").iterdir() if path.name != ".metadata" - ) - with optimizer_payload.open("ab") as stream: - stream.write(b"corrupt") - (tmp_path / "checkpoints" / "LATEST").write_text(corrupt_optimizer.name + "\n") - assert third_manager.resolve("LATEST") == resumed_checkpoint.resolve() - with pytest.raises(RuntimeError, match="DCP"): - third_manager.resolve(corrupt_optimizer.name) - - step_mismatch = tmp_path / "checkpoints" / "step_00000004" - shutil.copytree(resumed_checkpoint, step_mismatch) - step_manifest_path = step_mismatch / "manifest.json" - step_manifest = json.loads(step_manifest_path.read_text()) - step_manifest["completed_steps"] = 4 - step_manifest["step_scheduler"]["step"] = 4 - step_manifest_path.write_text(json.dumps(step_manifest, indent=2, sort_keys=True) + "\n") - trainer_state_path = step_mismatch / "trainer_state.json" - trainer_state = json.loads(trainer_state_path.read_text()) - trainer_state["completed_steps"] = 4 - trainer_state["step_scheduler"]["step"] = 4 - trainer_state_path.write_text(json.dumps(trainer_state, indent=2, sort_keys=True) + "\n") - _refresh_complete_marker(step_mismatch) - with pytest.raises(RuntimeError, match="trainer step"): - third_manager.load(step_mismatch.name) - - lr_mismatch = tmp_path / "checkpoints" / "step_00000005" - shutil.copytree(resumed_checkpoint, lr_mismatch) - lr_manifest_path = lr_mismatch / "manifest.json" - lr_manifest = json.loads(lr_manifest_path.read_text()) - lr_manifest["learning_rates"] = [0.123] - lr_manifest_path.write_text(json.dumps(lr_manifest, indent=2, sort_keys=True) + "\n") - lr_trainer_path = lr_mismatch / "trainer_state.json" - lr_trainer = json.loads(lr_trainer_path.read_text()) - lr_trainer["learning_rates"] = [0.123] - lr_trainer_path.write_text(json.dumps(lr_trainer, indent=2, sort_keys=True) + "\n") - _refresh_complete_marker(lr_mismatch) - with pytest.raises(RuntimeError, match="learning rate"): - third_manager.load(lr_mismatch.name) - - cursor_mismatch = tmp_path / "checkpoints" / "step_00000006" - shutil.copytree(resumed_checkpoint, cursor_mismatch) - sampler_path = cursor_mismatch / "sampler" / "sampler_dp_rank_0.pt" - sampler_state = torch.load(sampler_path, weights_only=False) - sampler_state["plan_sha256"] = "0" * 64 - torch.save(sampler_state, sampler_path) - cursor_manifest_path = cursor_mismatch / "manifest.json" - cursor_manifest = json.loads(cursor_manifest_path.read_text()) - cursor_manifest["sidecar_sha256"]["sampler/sampler_dp_rank_0.pt"] = hashlib.sha256( - sampler_path.read_bytes() - ).hexdigest() - cursor_manifest_path.write_text(json.dumps(cursor_manifest, indent=2, sort_keys=True) + "\n") - _refresh_complete_marker(cursor_mismatch) - with pytest.raises(RuntimeError, match="plan hash"): - third_manager.load(cursor_mismatch.name) - - source_checkpointer.close() - destination_checkpointer.close() - third_checkpointer.close() - - -def test_save_chains_validated_parent_without_resolving_latest(tmp_path, monkeypatch) -> None: - pytest.importorskip("nemo_automodel") - rng_module = pytest.importorskip("nemo_automodel.components.training.rng") - if not torch.distributed.is_initialized(): - initialize_pdd_distributed(backend="gloo", timeout_minutes=1) - sample_ids = tuple(f"sample-{index}" for index in range(8)) - - source = build_toy_lifecycle() - source_sampler = _released_sampler(sample_ids) - source_manager, source_checkpointer = _manager( - tmp_path / "checkpoints", - source, - source_sampler, - rng_module.StatefulRNG(1234, ranked=True), - ) - _run_next(source, source_sampler) - first = source_manager.save() - assert json.loads((first / "manifest.json").read_text())["parent_checkpoint"] is None - - resumed = build_toy_lifecycle() - resumed_sampler = _released_sampler(sample_ids) - resumed_manager, resumed_checkpointer = _manager( - tmp_path / "checkpoints", - resumed, - resumed_sampler, - rng_module.StatefulRNG(9999, ranked=True), - ) - assert resumed_manager.load("LATEST") is not None - - def fail_resolution(restore_from): - raise AssertionError(f"save re-resolved {restore_from}") - - monkeypatch.setattr(resumed_manager, "_collective_resolve", fail_resolution) - _run_next(resumed, resumed_sampler) - second = resumed_manager.save() - _run_next(resumed, resumed_sampler) - third = resumed_manager.save() - - assert json.loads((second / "manifest.json").read_text())["parent_checkpoint"] == first.name - assert json.loads((third / "manifest.json").read_text())["parent_checkpoint"] == second.name - source_checkpointer.close() - resumed_checkpointer.close() diff --git a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py b/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py deleted file mode 100644 index d7b1f0afc80..00000000000 --- a/tests/examples/diffusers/fastgen/test_pdd_validation_oracle.py +++ /dev/null @@ -1,159 +0,0 @@ -# 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. - -"""Deterministic logical-ID PDD held-out validation tests.""" - -from __future__ import annotations - -import pathlib -import sys - -import pytest -import torch - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) -if str(_FASTGEN_DIR) not in sys.path: - sys.path.insert(0, str(_FASTGEN_DIR)) -if str(pathlib.Path(__file__).parent) not in sys.path: - sys.path.insert(0, str(pathlib.Path(__file__).parent)) - -from pdd.training import ( - build_pdd_validation_assignments, - pdd_validation_noise, - pdd_validation_support, - run_pdd_validation, -) -from pdd_test_utils import build_toy_lifecycle, make_batch - -from modelopt.torch.fastgen import PDDConfig - - -def test_canonical_2k_assignment_covers_all_1568_pairs_32_starts_and_128_heads() -> None: - config = PDDConfig( - grid_size=128, - grid_max_t=0.999, - flow_shift=5.0, - block_size_min=4, - block_size_max=64, - inference_blocks=[32, 32, 32, 32], - student_sample_steps=4, - ) - sample_ids = [f"heldout-{index:04d}" for index in range(2000)] - assignments = build_pdd_validation_assignments( - list(reversed(sample_ids)), - config, - validation_seed=2026, - ) - - assert len(pdd_validation_support(config)) == 1568 - assert len(assignments) == 2000 - assert len({(assignment.n, assignment.k) for assignment in assignments}) == 1568 - assert len({assignment.n for assignment in assignments}) == 32 - assert len({assignment.k for assignment in assignments}) == 128 - assert assignments == build_pdd_validation_assignments( - sample_ids, - config, - validation_seed=2026, - ) - - -def test_assignment_rejects_duplicate_missing_coverage_and_noise_is_per_id_stable() -> None: - lifecycle = build_toy_lifecycle() - with pytest.raises(ValueError, match="unique"): - build_pdd_validation_assignments( - ["duplicate", "duplicate"], - lifecycle.config, - validation_seed=1, - require_full_coverage=False, - ) - with pytest.raises(ValueError, match="at least"): - build_pdd_validation_assignments( - ["too-small"], - lifecycle.config, - validation_seed=1, - ) - - before = torch.get_rng_state().clone() - first = pdd_validation_noise( - "logical-id", - (3,), - validation_seed=7, - device=torch.device("cpu"), - ) - second = pdd_validation_noise( - "logical-id", - (3,), - validation_seed=7, - device=torch.device("cpu"), - ) - different = pdd_validation_noise( - "different-id", - (3,), - validation_seed=7, - device=torch.device("cpu"), - ) - assert torch.equal(torch.get_rng_state(), before) - assert torch.equal(first, second) - assert not torch.equal(first, different) - - -def test_repeated_validation_is_exact_and_does_not_change_training_state_or_rng() -> None: - lifecycle = build_toy_lifecycle() - sample_ids = tuple(f"validation-{index:02d}" for index in range(12)) - assignments = build_pdd_validation_assignments( - sample_ids, - lifecycle.config, - validation_seed=44, - require_full_coverage=False, - ) - batches = [ - make_batch((assignment.sample_id,), offset=assignment.ordinal / 100) - for assignment in assignments - ] - parameter_before = { - name: parameter.detach().clone() for name, parameter in lifecycle.student.named_parameters() - } - optimizer_before = lifecycle.optimizer.state_dict() - scheduler_before = lifecycle.scheduler.state_dict() - rng_before = torch.get_rng_state().clone() - student_mode_before = lifecycle.student.training - teacher_mode_before = lifecycle.teacher.training - - first = run_pdd_validation( - lifecycle.pipeline, - (batch for batch in batches), - assignments, - validation_seed=44, - ) - second = run_pdd_validation( - lifecycle.pipeline, - list(reversed(batches)), - assignments, - validation_seed=44, - ) - - assert first == second - assert first.mean_loss == second.mean_loss - assert first.pair_count == len({(item.n, item.k) for item in assignments}) - assert torch.equal(torch.get_rng_state(), rng_before) - assert lifecycle.student.training is student_mode_before - assert lifecycle.teacher.training is teacher_mode_before - assert lifecycle.optimizer.state_dict() == optimizer_before - assert lifecycle.scheduler.state_dict() == scheduler_before - for name, parameter in lifecycle.student.named_parameters(): - torch.testing.assert_close(parameter, parameter_before[name], rtol=0, atol=0) diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index 11965cd7f1a..070a8358e7a 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -149,7 +149,7 @@ def test_formerly_vendored_files_use_standard_nvidia_header(): def test_data_builders_importable_and_accept_shared_cache_options(): - """The real-data builder exposes the negative embedding and stable-ID selection seams.""" + """The real-data builder exposes the negative embedding and deterministic split seams.""" pytest.importorskip("nemo_automodel") pytest.importorskip("torch") @@ -158,7 +158,7 @@ def test_data_builders_importable_and_accept_shared_cache_options(): assert callable(fastgen_data.build_text_to_image_multiresolution_dataloader) sig = inspect.signature(fastgen_data.build_text_to_image_multiresolution_dataloader) assert "negative_prompt_embedding_path" in sig.parameters - assert "selected_indices" in sig.parameters + assert "split" in sig.parameters # Default None => CFG-less construction works without the negative embedding (it is optional). assert sig.parameters["negative_prompt_embedding_path"].default is None @@ -185,7 +185,6 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): "aspect_ratio": 1.0, "prompt_embeds": torch.randn(seq, dim), "prompt_embeds_mask": torch.ones(seq, dtype=torch.long), - "sample_id": 7, } batch = [dict(sample), dict(sample)] neg = torch.randn(seq, dim) @@ -198,7 +197,6 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): assert out["negative_text_embeddings"].shape[0] == len( batch ) # broadcast [seq,dim]->[B,seq,dim] - assert out["metadata"]["sample_ids"].tolist() == [7, 7] def test_collate_zero_pads_variable_length_qwen_embeddings_and_masks(): @@ -219,7 +217,6 @@ def sample(sample_id, sequence_length): "aspect_ratio": 1.0, "prompt_embeds": torch.full((sequence_length, 3), float(sample_id)), "prompt_embeds_mask": torch.ones(sequence_length, dtype=torch.long), - "sample_id": sample_id, } result = collate_fn_text_to_image([sample(1, 5), sample(2, 3)]) diff --git a/tests/unit/torch/fastgen/test_pdd_metadata.py b/tests/unit/torch/fastgen/test_pdd_metadata.py deleted file mode 100644 index 1c3c65f6f2d..00000000000 --- a/tests/unit/torch/fastgen/test_pdd_metadata.py +++ /dev/null @@ -1,195 +0,0 @@ -# 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. - -"""Plain-torch PDD lifecycle and serialized-metadata reconstruction evidence.""" - -from __future__ import annotations - -import copy -import json -from typing import Any - -import pytest -import torch -from torch import nn - -from modelopt.torch.fastgen import ( - PDDConfig, - PDDLayerSpec, - PDDMetadata, - PDDOutputProjection, - PDDPipeline, - convert_to_pdd_output_projection, -) - - -class _ToyStudent(nn.Module): - def __init__(self, width: int = 3) -> None: - super().__init__() - self.backbone = nn.Linear(width, width) - self.projection = nn.Linear(width, width) - - def forward(self, state: torch.Tensor) -> torch.Tensor: - return self.projection(torch.tanh(self.backbone(state))) - - -class _ToyTeacher(nn.Module): - def __init__(self) -> None: - super().__init__() - self.scale = nn.Parameter(torch.tensor(-0.25)) - - def forward(self, state: torch.Tensor, time: torch.Tensor) -> torch.Tensor: - return self.scale * state + 0.1 * time[:, None] - - -class _ToyAdapter: - def __init__(self, grid_size: int) -> None: - self.grid_size = grid_size - - def _projection(self, model: _ToyStudent) -> PDDOutputProjection: - projection = model.projection - assert isinstance(projection, PDDOutputProjection) - return projection - - def student_all_heads( - self, - model: _ToyStudent, - state: torch.Tensor, - time: torch.Tensor, - *, - condition: Any = None, - **model_kwargs: Any, - ) -> torch.Tensor: - del time, condition, model_kwargs - raw = model(state) - return raw.reshape(state.shape[0], self.grid_size, state.shape[1]) - - def student_fused_block( - self, - model: _ToyStudent, - 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 - with self._projection(model).fuse_block(start, end, grid): - return model(state) - - def teacher_velocity( - self, - model: _ToyTeacher, - state: torch.Tensor, - time: torch.Tensor, - *, - condition: Any = None, - negative_condition: Any = None, - **model_kwargs: Any, - ) -> torch.Tensor: - del condition, negative_condition, model_kwargs - return model(state, time) - - -def _config() -> PDDConfig: - return PDDConfig( - grid_size=4, - grid_max_t=0.999, - flow_shift=5.0, - block_size_min=1, - block_size_max=4, - inference_blocks=[2, 2], - student_sample_steps=2, - ) - - -def test_plain_torch_training_sampling_and_strict_metadata_reconstruction(tmp_path) -> None: - torch.manual_seed(19) - config = _config() - layer_spec = PDDLayerSpec("projection", "channel_major") - student = _ToyStudent() - projection = convert_to_pdd_output_projection(student, layer_spec, config.grid_size) - pipeline = PDDPipeline(student, _ToyTeacher(), config, _ToyAdapter(config.grid_size)) - optimizer = torch.optim.SGD(student.parameters(), lr=0.05) - data = torch.tensor([[0.5, -1.0, 0.25], [-0.75, 0.5, 1.25]]) - noise = torch.tensor([[-0.25, 0.75, 1.0], [0.5, -1.25, 0.0]]) - projection_before = projection.weight.detach().clone() - - loss, _ = pipeline.compute_loss( - data, - noise=noise, - n=torch.tensor([0, 1]), - k=torch.tensor([2, 3]), - ) - loss.backward() - optimizer.step() - - assert torch.isfinite(loss) - assert not torch.equal(projection.weight, projection_before) - inference_noise = torch.tensor([[1.0, -0.5, 0.25], [-0.25, 0.75, 1.5]]) - expected_sample = pipeline.sample(inference_noise) - - metadata = PDDMetadata.from_config(config, projection) - metadata_path = tmp_path / "pdd_metadata.json" - metadata_path.write_text(json.dumps(metadata.to_dict(), sort_keys=True), encoding="utf-8") - restored_metadata = PDDMetadata.from_dict(json.loads(metadata_path.read_text(encoding="utf-8"))) - assert restored_metadata == metadata - - restored_config = PDDConfig( - grid_size=restored_metadata.grid_size, - grid_max_t=restored_metadata.grid_max_t, - flow_shift=restored_metadata.flow_shift, - block_size_min=restored_metadata.block_size_min, - block_size_max=restored_metadata.block_size_max, - inference_blocks=list(restored_metadata.inference_blocks), - student_sample_steps=len(restored_metadata.inference_blocks), - teacher_integrator=restored_metadata.teacher_integrator, - ) - - restored_student = _ToyStudent(width=restored_metadata.projection_in_features) - restored_projection = convert_to_pdd_output_projection( - restored_student, - restored_metadata.layer_spec, - restored_metadata.grid_size, - ) - load_result = restored_student.load_state_dict(copy.deepcopy(student.state_dict()), strict=True) - assert load_result.missing_keys == [] - assert load_result.unexpected_keys == [] - assert restored_projection.base_out_features == restored_metadata.projection_out_features - assert (restored_projection.bias is not None) is restored_metadata.projection_bias - - restored_pipeline = PDDPipeline( - restored_student, - copy.deepcopy(pipeline.teacher), - restored_config, - _ToyAdapter(restored_metadata.grid_size), - ) - torch.testing.assert_close(restored_pipeline.sample(inference_noise), expected_sample) - - -def test_strict_restore_rejects_checkpoint_with_different_projection_grid() -> None: - config = _config() - layer_spec = PDDLayerSpec("projection", "channel_major") - student = _ToyStudent() - convert_to_pdd_output_projection(student, layer_spec, config.grid_size) - checkpoint = copy.deepcopy(student.state_dict()) - incompatible = _ToyStudent() - convert_to_pdd_output_projection(incompatible, layer_spec, grid_size=2) - - with pytest.raises(RuntimeError, match="size mismatch"): - incompatible.load_state_dict(checkpoint, strict=True) diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index a1cd94338c8..75716a0c005 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -329,6 +329,21 @@ def test_selected_head_low_precision_outputs_use_float32_mse() -> None: torch.testing.assert_close(loss, expected) +def test_loss_can_skip_optional_diagnostics() -> None: + pipeline, _ = _pipeline() + data = torch.ones(2, 3) + + _, metrics = pipeline.compute_loss( + data, + noise=torch.zeros_like(data), + n=torch.tensor([0, 2]), + k=torch.tensor([1, 3]), + collect_metrics=False, + ) + + assert set(metrics) == {"student_target_mse"} + + def test_small_grid_accepts_exactly_the_trained_index_support() -> None: pipeline, _ = _pipeline() data = torch.ones(1, 3) diff --git a/tests/unit/torch/fastgen/test_pdd_projection.py b/tests/unit/torch/fastgen/test_pdd_projection.py index f859e96660f..398e685ed81 100644 --- a/tests/unit/torch/fastgen/test_pdd_projection.py +++ b/tests/unit/torch/fastgen/test_pdd_projection.py @@ -20,7 +20,6 @@ import copy import threading from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait -from dataclasses import replace import pytest import torch @@ -28,9 +27,7 @@ from torch import nn from modelopt.torch.fastgen import ( - PDDConfig, PDDLayerSpec, - PDDMetadata, PDDOutputProjection, convert_to_pdd_output_projection, get_module_by_path, @@ -352,74 +349,3 @@ def test_base_and_widened_checkpoint_load_order_is_strict_and_nonmutating(): with pytest.raises(RuntimeError, match="size mismatch"): _NestedModel().load_state_dict(widened_state, strict=True) _assert_state_unchanged(widened_state, widened_original) - - -@pytest.mark.parametrize("layout", ["channel_major", "patch_major"]) -def test_metadata_round_trips_exactly_without_mutating_mapping(layout): - projection = PDDOutputProjection.from_linear(_base_linear(), 128, _spec(layout)) - metadata = PDDMetadata.from_config(PDDConfig(), projection) - payload = metadata.to_dict() - original = copy.deepcopy(payload) - - restored = PDDMetadata.from_dict(payload) - - assert restored == metadata - assert restored.to_dict() == payload - assert payload == original - - -@pytest.mark.parametrize( - "changes", - [ - {"schema_version": True}, - {"grid_max_t": 1}, - {"flow_shift": 5}, - {"inference_blocks": [32, 32, 32, 32]}, - {"projection_bias": 1}, - ], -) -def test_direct_metadata_construction_is_as_strict_as_deserialization(changes): - projection = PDDOutputProjection.from_linear(_base_linear(), 128, _spec("channel_major")) - metadata = PDDMetadata.from_config(PDDConfig(), projection) - with pytest.raises(ValueError): - replace(metadata, **changes) - - -def _valid_patch_metadata_payload(): - projection = PDDOutputProjection.from_linear(_base_linear(), 128, _spec("patch_major")) - return PDDMetadata.from_config(PDDConfig(), projection).to_dict() - - -@pytest.mark.parametrize( - ("mutate", "message"), - [ - (lambda data: data.update(schema_version=2), "unsupported.*schema_version"), - (lambda data: data.update(extra=True), "keys mismatch"), - (lambda data: data.update({1: "bad"}), "keys must all be strings"), - (lambda data: data.pop("grid_max_t"), "keys mismatch"), - (lambda data: data.update(grid_max_t=1), "grid_max_t must be a float"), - (lambda data: data.update(grid_max_t=0.0), "0 < grid_max_t <= 1"), - (lambda data: data.update(flow_shift=5), "flow_shift must be a float"), - (lambda data: data.update(grid_size=124), "inference_blocks must sum"), - (lambda data: data.update(inference_blocks=(32, 32, 32, 32)), "list of integers"), - ( - lambda data: data["layer_spec"].update(head_layout="unsupported"), - "head_layout must be one of", - ), - ( - lambda data: data["base_projection"].update(out_features=5), - "must be divisible", - ), - ], -) -def test_metadata_rejects_unsupported_schema_schedule_shape_and_layout(mutate, message): - payload = _valid_patch_metadata_payload() - mutate(payload) - with pytest.raises(ValueError, match=message): - PDDMetadata.from_dict(payload) - - -def test_metadata_rejects_projection_config_grid_mismatch(): - projection = PDDOutputProjection.from_linear(_base_linear(), 64, _spec("channel_major")) - with pytest.raises(ValueError, match="does not match projection"): - PDDMetadata.from_config(PDDConfig(), projection) diff --git a/tests/unit/torch/fastgen/test_pdd_public_api.py b/tests/unit/torch/fastgen/test_pdd_public_api.py index 617b6a9b46c..9aae47cf634 100644 --- a/tests/unit/torch/fastgen/test_pdd_public_api.py +++ b/tests/unit/torch/fastgen/test_pdd_public_api.py @@ -32,7 +32,6 @@ from modelopt.torch.fastgen.loader import load_pdd_config from modelopt.torch.fastgen.methods.pdd import ( PDDLayerSpec, - PDDMetadata, PDDModelAdapter, PDDOutputProjection, PDDPipeline, @@ -89,7 +88,6 @@ def test_core_pdd_symbols_are_exported_from_the_public_package() -> None: expected = { "PDDConfig": PDDConfig, "PDDLayerSpec": PDDLayerSpec, - "PDDMetadata": PDDMetadata, "PDDModelAdapter": PDDModelAdapter, "PDDOutputProjection": PDDOutputProjection, "PDDPipeline": PDDPipeline, @@ -132,9 +130,9 @@ def find_spec(self, fullname, path=None, target=None): baseline = set(sys.modules) blocker.attempts.clear() -from modelopt.torch.fastgen import PDDConfig, PDDMetadata, PDDPipeline +from modelopt.torch.fastgen import PDDConfig, PDDPipeline -assert all(symbol is not None for symbol in (PDDConfig, PDDMetadata, PDDPipeline)) +assert all(symbol is not None for symbol in (PDDConfig, PDDPipeline)) assert blocker.attempts == [] loaded = set(sys.modules) - baseline for prefix in ( diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 1cf105b7146..c27842b71c8 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -258,6 +258,32 @@ def test_unfused_channel_major_output_maps_each_packed_head_in_order() -> None: torch.testing.assert_close(actual, expected) +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() @@ -584,14 +610,23 @@ def _mr210_qwen_forward_oracle( max_txt_seq_len=max_txt_seq_len, device=hidden_states.device, ) + image_mask = torch.ones( + (hidden_states.shape[0], hidden_states.shape[1]), + dtype=torch.bool, + device=hidden_states.device, + ) + joint_attention_mask = torch.cat( + (encoder_hidden_states_mask.to(torch.bool), image_mask), + dim=1, + )[:, None, None, :] 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, + encoder_hidden_states_mask=None, temb=temb, image_rotary_emb=image_rotary_emb, - joint_attention_kwargs=None, + joint_attention_kwargs={"attention_mask": joint_attention_mask}, ) return model.proj_out(model.norm_out(hidden_states, temb)) @@ -616,6 +651,7 @@ def teacher_velocity(self, model, state, time, **kwargs): torch.manual_seed(20260716) base = _tiny_diffusers_qwen().eval() actual_student = adopt_qwen_image_mr210_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) @@ -747,6 +783,14 @@ def test_conversion_preserves_the_ordinary_diffusers_qwen_root() -> None: assert dict(student.config) == config +def test_mr210_adoption_accepts_a_dynamic_qwen_subclass() -> None: + student = _tiny_diffusers_qwen().eval() + student.__class__ = type("FSDPQwenImageTransformer2DModel", (type(student),), {}) + + assert adopt_qwen_image_mr210_forward(student) is student + assert require_qwen_image_mr210_forward(student) == QWEN_IMAGE_PDD_EXECUTION + + def test_mr210_adoption_preserves_root_state_and_deepcopy_binding() -> None: source = _tiny_diffusers_qwen().eval() source_type = type(source) @@ -842,7 +886,7 @@ def test_mr210_qwen_conversion_preserves_every_initialized_head() -> None: torch.testing.assert_close(actual, expected[:, None].expand_as(actual), rtol=0, atol=0) -def test_mr210_mask_flow_differs_from_canonical_joint_mask() -> None: +def test_mr210_joint_mask_ignores_padded_token_values() -> None: canonical = _tiny_diffusers_qwen().eval().to(torch.bfloat16) student = copy.deepcopy(canonical) student = adopt_qwen_image_mr210_forward(student) @@ -873,8 +917,8 @@ def test_mr210_mask_flow_differs_from_canonical_joint_mask() -> None: captured_masks: list[torch.Tensor] = [] def capture_block_mask(_module, _args, kwargs): - assert "attention_mask" not in kwargs - captured_masks.append(kwargs["encoder_hidden_states_mask"].detach().clone()) + assert kwargs["encoder_hidden_states_mask"] is None + captured_masks.append(kwargs["joint_attention_kwargs"]["attention_mask"].detach().clone()) hook = student.transformer_blocks[0].register_forward_pre_hook( capture_block_mask, @@ -904,9 +948,11 @@ def capture_block_mask(_module, _args, kwargs): hook.remove() torch.testing.assert_close(canonical_poisoned, canonical_baseline, rtol=0, atol=0) - assert not torch.equal(strict_poisoned[1], strict_baseline[1]) + torch.testing.assert_close(strict_poisoned, strict_baseline, rtol=0, atol=0) assert len(captured_masks) == 2 - assert all(torch.equal(captured, mask) for captured in captured_masks) + expected_mask = torch.cat((mask.bool(), torch.ones(2, 4, dtype=torch.bool)), dim=1) + expected_mask = expected_mask[:, None, None, :] + assert all(torch.equal(captured, expected_mask) for captured in captured_masks) def test_mr210_preserves_diffusers_output_and_harmless_call_contract() -> None: @@ -1060,7 +1106,7 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N negative_condition=condition[0], ) assert transformer.calls == [] - with pytest.raises(TypeError, match="converted to PDDOutputProjection"): + with pytest.raises(ValueError, match="has 4 outputs; expected 16"): adapter.student_all_heads(transformer, state, time, condition=condition) convert_qwen_image_to_pdd(transformer, config) From fdea352f28b5887f8220e75ed43f97163bb697f7 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 21 Jul 2026 01:28:39 -0700 Subject: [PATCH 36/45] Preserve FP32 PDD timesteps under FSDP Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/recipe.py | 67 ++++++++++++++----- .../fastgen/test_pdd_recipe_setup.py | 31 ++++++++- 2 files changed, 80 insertions(+), 18 deletions(-) diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 2ae0d3c227e..768c3f14ce8 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -18,21 +18,22 @@ from __future__ import annotations import logging -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from pathlib import Path from typing import Any +import torch import torch.distributed as dist from huggingface_hub import snapshot_download from torch import nn +from torch.distributed.fsdp import MixedPrecisionPolicy try: + import nemo_automodel.recipes.diffusion.train as automodel_diffusion_train from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline from nemo_automodel.components.training.rng import ScopedRNG - from nemo_automodel.recipes.diffusion.train import ( - TrainDiffusionRecipe, - _build_diffusion_parallel_manager_args, - ) + from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe except ImportError as exc: raise ImportError( "The PDD example requires nemo_automodel. Install " @@ -48,6 +49,37 @@ from .training import PDDFlowMatchingStepAdapter +@contextmanager +def _preserve_fp32_timestep_inputs() -> Iterator[None]: + """Keep Qwen's continuous timestep in FP32 while FSDP computes in BF16.""" + original_builder = automodel_diffusion_train._build_diffusion_parallel_manager_args + + 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 + + automodel_diffusion_train._build_diffusion_parallel_manager_args = build_manager_args + try: + yield + finally: + automodel_diffusion_train._build_diffusion_parallel_manager_args = original_builder + + def _config_mapping(value: Any) -> dict[str, Any]: if hasattr(value, "to_dict"): return value.to_dict() @@ -83,20 +115,21 @@ class PDDDiffusionRecipe(TrainDiffusionRecipe): """Use AutoModel's native lifecycle with a PDD loss and frozen teacher.""" def setup(self) -> None: - super().setup() + with _preserve_fp32_timestep_inputs(): + super().setup() - raw_pdd = _config_mapping(self.cfg.get("pdd", {})) - self.pdd_config = PDDConfig.model_validate(raw_pdd) + 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 MR210 forward here changes behavior only; it creates no parameters. - adopt_qwen_image_mr210_forward(self.model) - _validate_prepared_student(self.model, self.pdd_config) - self.model.enable_gradient_checkpointing() + # The student artifact is widened before AutoModel creates FSDP and optimizer state. + # Binding the MR210 forward here changes behavior only; it creates no parameters. + adopt_qwen_image_mr210_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() + # ``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, @@ -110,7 +143,7 @@ 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 = _build_diffusion_parallel_manager_args( + manager_args = automodel_diffusion_train._build_diffusion_parallel_manager_args( fsdp_cfg=fsdp_cfg, ddp_cfg=ddp_cfg, world_size=self.world_size, diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index e6c1c74a132..a8d3d66a8e2 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -30,7 +30,8 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -from pdd.recipe import _validate_prepared_student +from pdd import recipe as pdd_recipe +from pdd.recipe import _preserve_fp32_timestep_inputs, _validate_prepared_student from pdd.training import PDDFlowMatchingStepAdapter from modelopt.torch.fastgen import PDDConfig @@ -79,6 +80,34 @@ def test_prepared_student_width_is_validated_before_training() -> None: _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_recipe.automodel_diffusion_train, + "_build_diffusion_parallel_manager_args", + build_manager_args, + ) + with _preserve_fp32_timestep_inputs(): + manager_args = pdd_recipe.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_recipe.automodel_diffusion_train._build_diffusion_parallel_manager_args + is build_manager_args + ) + + @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) From 300fde453740a10519b5735f87a006af9260799d Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 11 Aug 2026 02:16:28 -0700 Subject: [PATCH 37/45] Add data-free PDD training for Qwen-Image Signed-off-by: Meng Xin --- .../fastgen/fastgen_data/__init__.py | 10 +- .../fastgen/fastgen_data/collate_fns.py | 55 ++--- .../fastgen_data/text_to_image_dataset.py | 6 +- examples/diffusers/fastgen/pdd/README.md | 29 ++- .../fastgen/pdd/configs/qwen_image.yaml | 31 ++- .../fastgen/pdd/inference_qwen_image.py | 3 +- examples/diffusers/fastgen/pdd/recipe.py | 50 ++++- examples/diffusers/fastgen/pdd/training.py | 145 +++++++++++-- modelopt/torch/fastgen/config.py | 8 +- modelopt/torch/fastgen/methods/pdd.py | 200 +++++++++++++----- .../torch/fastgen/plugins/qwen_image_pdd.py | 179 +++++++++++++--- .../diffusers/fastgen/test_dataset_paths.py | 20 ++ .../fastgen/test_pdd_recipe_setup.py | 168 ++++++++++++++- .../fastgen/test_vendored_migration.py | 27 +++ tests/unit/torch/fastgen/test_pdd_config.py | 3 +- tests/unit/torch/fastgen/test_pdd_pipeline.py | 53 ++++- .../fastgen/test_qwen_image_pdd_plugin.py | 151 +++++++++---- 17 files changed, 924 insertions(+), 214 deletions(-) diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index d06c001c7c5..c1b1c7f6a2b 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -19,11 +19,11 @@ 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`` diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index 0f25cc854b6..d15b3e28faf 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -17,10 +17,10 @@ Self-contained on **stock** ``nemo_automodel`` (no AutoModel patch required): -* :func:`collate_fn_text_to_image` builds the Qwen-Image 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 +* :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 @@ -44,45 +44,28 @@ __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 a 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]``) for CFG-based objectives. - 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 FastGen pipelines consume, 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.5.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": pad_sequence( [item["prompt_embeds"] for item in batch], @@ -117,7 +100,7 @@ def collate_fn_text_to_image( 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() @@ -135,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. @@ -170,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, @@ -193,6 +192,7 @@ def build_text_to_image_multiresolution_dataloader( 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. @@ -217,6 +217,7 @@ def build_text_to_image_multiresolution_dataloader( dataset = TextToImageDataset( cache_dir=cache_dir, train_text_encoder=train_text_encoder, + prompt_only=prompt_only, split=split, validation_count=validation_count, split_seed=split_seed, @@ -224,7 +225,7 @@ def build_text_to_image_multiresolution_dataloader( effective_root = dataset.cache_root # Load the optional negative-prompt embedding once and bind it into the collate. - collate_fn = collate_fn_text_to_image + collate_fn = collate_fn_text_prompts if prompt_only else collate_fn_text_to_image if negative_prompt_embedding_path is not None: negative_path = resolve_under_root( effective_root, @@ -241,7 +242,7 @@ def build_text_to_image_multiresolution_dataloader( 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, ) 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 c028f04941f..29b04f86797 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -32,6 +32,7 @@ def __init__( self, cache_dir: str | Path, train_text_encoder: bool = False, + prompt_only: bool = False, split: str | None = None, validation_count: int | None = None, split_seed: int = 2026, @@ -40,6 +41,7 @@ def __init__( 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. split: Optional deterministic ``"train"`` or ``"validation"`` selection. validation_count: Number of validation samples when ``split`` is set. split_seed: Local seed used to construct deterministic split membership. @@ -49,6 +51,7 @@ def __init__( if split is not None and validation_count is None: raise ValueError("validation_count is required when split is set") self.train_text_encoder = train_text_encoder + self.prompt_only = prompt_only self.cache_root = resolve_cache_root(cache_dir) self._split = split self._validation_count = validation_count @@ -122,7 +125,6 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: # 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"]), @@ -131,6 +133,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 index bd69c01c354..b3bfd17f462 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -6,10 +6,20 @@ Qwen-Image student with 128 output heads. A block schedule such as `[32, 32, 32, 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. The complete student transformer, -including the widened output projection, is finetuned at a constant learning rate of `5e-5`; this -is not a heads-only run. Training samples aligned target spans from 4 through 64 intervals, so the -same checkpoint supports multiple inference schedules. +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 @@ -45,13 +55,16 @@ torchrun --standalone --nproc-per-node=8 \ --fsdp.dp_size=8 ``` -The cache must contain `metadata.json`, its declared tensor shards, and +The cache must contain `metadata.json`, its declared prompt-embedding shards, and `negative_prompt_embedding.pt`. `MODELOPT_FASTGEN_DATASET_CACHE_DIR` overrides the configured cache root; paths declared by the dataset remain confined to that root. -The checked-in recipe targets 50,000 optimizer steps with global batch size 256, local batch size -4, and constant `5e-5` learning rate. `checkpoint.restore_from: LATEST` lets AutoModel resume the -latest native checkpoint. For wall-time-limited Slurm jobs, request an early signal such as +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. diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 15427d6eb1d..0023b42ada9 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -13,6 +13,8 @@ model: mode: finetune transformer_engine_linear: false fuse_qkv_projections: false + # Native Qwen-Image latent shape at 1024x1024. + latent_shape: [16, 128, 128] pdd: pred_type: flow @@ -23,18 +25,18 @@ pdd: grid_size: 128 grid_max_t: 0.999 flow_shift: 5.0 - block_size_min: 4 + block_size_min: 16 block_size_max: 64 - teacher_integrator: euler + teacher_integrator: midpoint inference_blocks: [32, 32, 32, 32] - data_free: false + data_free: true optim: - learning_rate: 5.0e-5 + learning_rate: 1.0e-5 clip_grad: 1.0 optimizer: _target_: torch.optim.AdamW - weight_decay: 0.01 + weight_decay: 0.0 betas: [0.9, 0.999] eps: 1.0e-8 amsgrad: false @@ -44,18 +46,26 @@ optim: 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: 50000 + max_steps: 3000 num_epochs: 200 log_every: 10 - ckpt_every_steps: 1000 + ckpt_every_steps: 250 local_batch_size: 4 save_checkpoint_every_epoch: false - global_batch_size: 256 + global_batch_size: 2048 fsdp: dp_size: @@ -70,6 +80,7 @@ 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 @@ -84,8 +95,8 @@ data: checkpoint: enabled: true - checkpoint_dir: checkpoints/pdd_qwen_image + checkpoint_dir: checkpoints/pdd_qwen_image_data_free_midpoint_3k model_save_format: safetensors save_consolidated: final diffusers_compatible: true - restore_from: LATEST + restore_from: diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index b77c518cebb..4298744447a 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -25,7 +25,6 @@ import torch import yaml from diffusers import QwenImagePipeline, QwenImageTransformer2DModel -from torch import nn _THIS_DIR = Path(__file__).resolve().parent _FASTGEN_DIR = _THIS_DIR.parent @@ -153,7 +152,7 @@ def main() -> None: ) sampler = PDDPipeline( transformer, - nn.Identity(), + None, config, QwenImagePDDAdapter(config, compute_dtype=torch.bfloat16), ) diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 768c3f14ce8..18643f2effd 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -31,7 +31,7 @@ try: import nemo_automodel.recipes.diffusion.train as automodel_diffusion_train - from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + 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: @@ -44,6 +44,7 @@ from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( QwenImagePDDAdapter, adopt_qwen_image_mr210_forward, + freeze_qwen_image_mr210_unused_parameters, ) from .training import PDDFlowMatchingStepAdapter @@ -111,11 +112,39 @@ def _validate_prepared_student(model: nn.Module, config: PDDConfig) -> None: ) +@contextmanager +def _freeze_unused_qwen_parameters_before_optimizer() -> Iterator[None]: + """Freeze unused Qwen outputs before AutoModel collects optimizer parameters.""" + pipeline_cls = automodel_diffusion_train.NeMoAutoDiffusionPipeline + original_descriptor = pipeline_cls.__dict__["from_pretrained"] + original_from_pretrained = pipeline_cls.from_pretrained + + def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any: + del cls + pipe, managers = original_from_pretrained(*args, **kwargs) + if not kwargs.get("load_for_training", False): + return pipe, managers + model = pipe.transformer + frozen_names = freeze_qwen_image_mr210_unused_parameters(model) + logging.info( + "[PDD] Full training excludes %d unused final-block text-output tensors: %s", + len(frozen_names), + ", ".join(frozen_names), + ) + return pipe, managers + + pipeline_cls.from_pretrained = classmethod(from_pretrained) + try: + yield + finally: + pipeline_cls.from_pretrained = original_descriptor + + class PDDDiffusionRecipe(TrainDiffusionRecipe): """Use AutoModel's native lifecycle with a PDD loss and frozen teacher.""" def setup(self) -> None: - with _preserve_fp32_timestep_inputs(): + with _preserve_fp32_timestep_inputs(), _freeze_unused_qwen_parameters_before_optimizer(): super().setup() raw_pdd = _config_mapping(self.cfg.get("pdd", {})) @@ -136,8 +165,21 @@ def setup(self) -> None: self.pdd_config, QwenImagePDDAdapter(self.pdd_config, compute_dtype=self.compute_dtype), ) - self.flow_matching_pipeline = PDDFlowMatchingStepAdapter(pdd_pipeline) - logging.info("[PDD] AutoModel lifecycle enabled; grid_size=%d", self.pdd_config.grid_size) + 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.""" diff --git a/examples/diffusers/fastgen/pdd/training.py b/examples/diffusers/fastgen/pdd/training.py index 3ebf3b1cf7d..dd42f882f8b 100644 --- a/examples/diffusers/fastgen/pdd/training.py +++ b/examples/diffusers/fastgen/pdd/training.py @@ -17,20 +17,101 @@ 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) -> None: + 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, @@ -43,30 +124,54 @@ def step( 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, global_step + del model - data = batch["image_latents"].to(device=device, dtype=dtype, non_blocking=True) - condition = ( - batch["text_embeddings"].to(device=device, dtype=dtype, non_blocking=True), - batch["text_embeddings_mask"].to(device=device, non_blocking=True), - ) + incoming_condition = self._condition(batch, device=device, dtype=dtype) negative_condition = None if self.pipeline.config.guidance_scale is not None: - negative_condition = ( - batch["negative_text_embeddings"].to( - device=device, - dtype=dtype, - non_blocking=True, - ), - batch["negative_text_embeddings_mask"].to(device=device, non_blocking=True), + negative_condition = self._condition( + batch, + device=device, + dtype=dtype, + prefix="negative_", ) - loss, metrics = self.pipeline.compute_loss( - data, - condition=condition, - negative_condition=negative_condition, - collect_metrics=collect_metrics, - ) + 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.") diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py index d4542e1d269..bf5f8c3a673 100644 --- a/modelopt/torch/fastgen/config.py +++ b/modelopt/torch/fastgen/config.py @@ -221,7 +221,7 @@ class DistillationConfig(ModeloptBaseConfig): class PDDConfig(DistillationConfig): - """Hyperparameters for data-dependent Parallel Decoding Distillation (PDD). + """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 @@ -278,10 +278,12 @@ class PDDConfig(DistillationConfig): title="Fused inference block schedule", description="Contiguous interval counts that partition the complete PDD grid.", ) - data_free: Literal[False] = ModeloptField( + data_free: bool = ModeloptField( default=False, title="Data-free training", - description="Data-free PDD is unsupported; training uses noised real latents.", + description=( + "Carry student-generated trajectories from fresh noise instead of noising real latents." + ), ) def __setattr__(self, name: str, value: object) -> None: diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index fab7d09b206..bcd96f6877f 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -15,9 +15,10 @@ """Framework-neutral projection, training, and sampling primitives for PDD. -This module owns projection layout and fusion plus the data-dependent objective -and block sampler. Model calls and architecture-specific packing remain behind -the adapter protocol and belong in ``modelopt.torch.fastgen.plugins``. +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 @@ -424,19 +425,24 @@ def teacher_velocity( class PDDPipeline(DistillationPipeline): - """Data-dependent PDD loss and fused sampler over a single core-owned grid.""" + """PDD losses and fused sampler over a single core-owned grid.""" def __init__( self, student: nn.Module, - teacher: nn.Module, + teacher: nn.Module | None, config: PDDConfig, adapter: PDDModelAdapter, ) -> None: - """Store the models/config/adapter and freeze the teacher.""" + """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__}.") - super().__init__(student, teacher, config) + if teacher is None: + self.student = student + self.teacher = None + self.config = config + else: + super().__init__(student, teacher, config) self.adapter = adapter def time_grid(self, device: torch.device | str | None = None) -> torch.Tensor: @@ -567,73 +573,46 @@ 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( + def _compute_loss_from_state( self, - data: torch.Tensor, + state: 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, + n: torch.Tensor, + k: torch.Tensor, 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) - + 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) - data_fp32 = data.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_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,) * (data.ndim - 1) - x_n = add_noise(data_fp32, noise_fp32, time_n) + broadcast_shape = (batch_size,) + (1,) * (state.ndim - 1) student_heads = self.adapter.student_all_heads( self.student, - x_n, + state_fp32, time_n, condition=condition, **kwargs, ) - expected_head_shape = torch.Size((batch_size, self.config.grid_size, *data.shape[1:])) + 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=data.device, + device=state.device, name="student_all_heads", ) with torch.no_grad(): - x_bar_k = integrate_interval_velocities(x_n, student_heads, grid, n, k) + x_bar_k = integrate_interval_velocities(state_fp32, student_heads, grid, n, k) - batch_ids = torch.arange(batch_size, device=data.device) + batch_ids = torch.arange(batch_size, device=state.device) student_target = student_heads[batch_ids, k] time_k = grid[k] with torch.no_grad(): @@ -648,8 +627,8 @@ def compute_loss( ) teacher_first = self._normalize_velocity( teacher_first, - expected_shape=data.shape, - device=data.device, + expected_shape=state.shape, + device=state.device, name="teacher_velocity", ) if self.config.teacher_integrator == "euler": @@ -670,11 +649,26 @@ def compute_loss( ) teacher_target = self._normalize_velocity( teacher_target, - expected_shape=data.shape, - device=data.device, + 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)) @@ -687,7 +681,7 @@ def compute_loss( metrics.update( n=n.detach(), k=k.detach(), - target_span=(k - n).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(), @@ -705,8 +699,100 @@ def compute_loss( .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) diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index 8365b976943..fa8b95be2ed 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -1,3 +1,6 @@ +# Adapted from the Qwen-Image implementation in Diffusers: +# https://github.com/huggingface/diffusers/blob/275869dcae4ebcfee6a80253fdabc56033335020/src/diffusers/models/transformers/transformer_qwenimage.py +# SPDX-FileCopyrightText: Copyright (c) 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved. # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -23,6 +26,7 @@ from typing import Any import torch +import torch.distributed as dist from torch import nn from ..config import PDDConfig @@ -35,6 +39,7 @@ "QwenImagePDDAdapter", "adopt_qwen_image_mr210_forward", "convert_qwen_image_to_pdd", + "freeze_qwen_image_mr210_unused_parameters", "require_qwen_image_mr210_forward", "restore_qwen_image_pdd_projection", ] @@ -232,6 +237,34 @@ def require_qwen_image_mr210_forward(model: nn.Module) -> str: return QWEN_IMAGE_PDD_EXECUTION +def freeze_qwen_image_mr210_unused_parameters(transformer: nn.Module) -> tuple[str, ...]: + """Freeze final-block text outputs that the MR210 forward does not consume.""" + blocks = getattr(transformer, "transformer_blocks", None) + if not isinstance(blocks, nn.ModuleList) or not blocks: + raise RuntimeError("Qwen MR210 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 MR210 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 adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: """Bind FastGen MR210's regular Qwen forward to the loaded root in place.""" if not isinstance(transformer, nn.Module): @@ -520,15 +553,19 @@ def _call_packed( model_kwargs: Mapping[str, Any], *, condition_name: str, + prepared_condition: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: - encoder_hidden_states, attention_mask = self._prepare_call( - model, - state, - time, - condition, - model_kwargs, - condition_name=condition_name, - ) + if prepared_condition is None: + encoder_hidden_states, attention_mask = self._prepare_call_collectively( + model, + state, + time, + condition, + model_kwargs, + condition_name=condition_name, + ) + else: + encoder_hidden_states, attention_mask = prepared_condition batch_size, _, height, width = state.shape model_dtype = self._model_dtype(model, state.dtype) @@ -551,6 +588,92 @@ def _call_packed( ) return self._extract_packed_output(output) + @staticmethod + def _raise_collective_preflight_error( + local_error: Exception | None, + *, + state: torch.Tensor, + ) -> None: + if dist.is_available() and dist.is_initialized(): + failed = torch.tensor(local_error is not None, dtype=torch.int32, device=state.device) + dist.all_reduce(failed, op=dist.ReduceOp.MAX) + if bool(failed): + if local_error is not None: + raise local_error + raise RuntimeError("Qwen PDD preflight failed on another rank.") + elif local_error is not None: + raise local_error + + def _prepare_call_collectively( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + condition: Any, + model_kwargs: Mapping[str, Any], + *, + condition_name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + prepared: tuple[torch.Tensor, torch.Tensor] | None = None + local_error: Exception | None = None + try: + prepared = self._prepare_call( + model, + state, + time, + condition, + model_kwargs, + condition_name=condition_name, + ) + except Exception as error: + local_error = error + + self._raise_collective_preflight_error(local_error, state=state) + if prepared is None: + raise RuntimeError("Qwen PDD preflight did not prepare a model call.") + return prepared + + def _prepare_teacher_cfg_calls( + self, + model: nn.Module, + state: torch.Tensor, + time: torch.Tensor, + condition: Any, + negative_condition: Any, + model_kwargs: Mapping[str, Any], + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + """Make rank-local CFG validation fail collectively before either teacher call.""" + prepared_condition: tuple[torch.Tensor, torch.Tensor] | None = None + prepared_negative_condition: tuple[torch.Tensor, torch.Tensor] | None = None + local_error: Exception | None = None + try: + prepared_condition = self._prepare_call( + model, + state, + time, + condition, + model_kwargs, + condition_name="condition", + ) + prepared_negative_condition = self._prepare_call( + model, + state, + time, + negative_condition, + model_kwargs, + condition_name="negative_condition", + ) + except Exception as error: + local_error = error + + self._raise_collective_preflight_error(local_error, state=state) + if prepared_condition is None or prepared_negative_condition is None: + raise RuntimeError("Qwen teacher CFG preflight did not prepare both model calls.") + return prepared_condition, prepared_negative_condition + @staticmethod def _expected_packed_shape(state: torch.Tensor, *, output_features: int) -> torch.Size: return torch.Size( @@ -708,27 +831,18 @@ def teacher_velocity( ) -> torch.Tensor: """Return conditional or fixed two-pass packed-CFG Qwen teacher velocity.""" guidance_scale = self.guidance_scale - if guidance_scale is not None and negative_condition is None: - raise ValueError("negative_condition is required when Qwen teacher CFG is enabled.") if guidance_scale is not None: - # Validate both collective-participating calls before either model - # call so malformed rank-local conditioning cannot split call counts. - self._prepare_call( + prepared_condition, prepared_negative_condition = self._prepare_teacher_cfg_calls( model, state, time, condition, - model_kwargs, - condition_name="condition", - ) - self._prepare_call( - model, - state, - time, negative_condition, model_kwargs, - condition_name="negative_condition", ) + else: + prepared_condition = None + prepared_negative_condition = None conditional = self._call_packed( model, @@ -737,6 +851,7 @@ def teacher_velocity( condition, model_kwargs, condition_name="condition", + prepared_condition=prepared_condition, ) if guidance_scale is None: return self._unpack_single(conditional, state) @@ -748,6 +863,7 @@ def teacher_velocity( negative_condition, model_kwargs, condition_name="negative_condition", + prepared_condition=prepared_negative_condition, ) expected = self._expected_packed_shape(state, output_features=state.shape[1] * 4) if conditional.shape != expected or unconditional.shape != expected: @@ -756,25 +872,22 @@ def teacher_velocity( f"{tuple(conditional.shape)} and {tuple(unconditional.shape)}." ) - # MR210 unpacks each Qwen prediction before guidance. Keep that - # operation order: the FP32 global reduction order over NCHW is not - # guaranteed to match an algebraically equivalent packed reduction. - conditional_unpacked = self._unpack_single(conditional, state) - unconditional_unpacked = self._unpack_single(unconditional, state) - guided_low_precision = conditional_unpacked + (float(guidance_scale) - 1.0) * ( - conditional_unpacked - unconditional_unpacked + # 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_unpacked.to(torch.float32) + conditional_fp32 = conditional.to(torch.float32) guided_fp32 = guided_low_precision.to(torch.float32) - norm_dims = tuple(range(1, conditional_fp32.ndim)) conditional_norm = torch.linalg.vector_norm( conditional_fp32, - dim=norm_dims, + dim=-1, keepdim=True, ) guided_norm = torch.linalg.vector_norm( guided_fp32, - dim=norm_dims, + dim=-1, keepdim=True, ).clamp_min(1e-5) - return (guided_fp32 * (conditional_norm / guided_norm)).to(conditional.dtype) + guided = (guided_fp32 * (conditional_norm / guided_norm)).to(conditional.dtype) + return self._unpack_single(guided, state) diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index 89fc3ef1f6c..d19e8c90e88 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -124,6 +124,26 @@ def test_dataset_accepts_absolute_payload_beneath_root(make_fastgen_cache, tmp_p 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_environment_redirects_samples_and_relative_negative_embedding( make_fastgen_cache, monkeypatch, tmp_path ): diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index a8d3d66a8e2..ab61a300b79 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -31,7 +31,11 @@ sys.path.insert(0, str(_FASTGEN_DIR)) from pdd import recipe as pdd_recipe -from pdd.recipe import _preserve_fp32_timestep_inputs, _validate_prepared_student +from pdd.recipe import ( + _freeze_unused_qwen_parameters_before_optimizer, + _preserve_fp32_timestep_inputs, + _validate_prepared_student, +) from pdd.training import PDDFlowMatchingStepAdapter from modelopt.torch.fastgen import PDDConfig @@ -44,9 +48,27 @@ def __init__(self, out_features: int) -> None: 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) + self.config = SimpleNamespace(guidance_scale=guidance_scale, data_free=False) self.scale = nn.Parameter(torch.tensor(2.0)) self.last_call = None @@ -56,6 +78,44 @@ def compute_loss(self, data, *, condition, negative_condition, collect_metrics): 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), @@ -108,6 +168,41 @@ def build_manager_args(**kwargs): ) +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_recipe.automodel_diffusion_train, + "NeMoAutoDiffusionPipeline", + _Pipeline, + ) + + original_descriptor = _Pipeline.__dict__["from_pretrained"] + with _freeze_unused_qwen_parameters_before_optimizer(): + student, _ = _Pipeline.from_pretrained("student", load_for_training=True) + teacher, _ = _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 + + @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) @@ -128,3 +223,72 @@ def test_step_adapter_returns_native_tuple_and_preserves_pdd_gradient(guidance_s _, _, 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_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index 070a8358e7a..890bd4bc5f2 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -199,6 +199,33 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): ) # broadcast [seq,dim]->[B,seq,dim] +def test_prompt_only_collate_omits_image_latents_and_keeps_cfg_conditioning(): + pytest.importorskip("nemo_automodel") + torch = pytest.importorskip("torch") + + from fastgen_data import collate_fn_text_prompts + + sample = { + "crop_resolution": torch.tensor([1024, 1024]), + "original_resolution": torch.tensor([1024, 1024]), + "crop_offset": torch.tensor([0, 0]), + "prompt": "a test prompt", + "image_path": "/unused/source.png", + "bucket_id": 0, + "aspect_ratio": 1.0, + "prompt_embeds": torch.randn(5, 16), + "prompt_embeds_mask": torch.ones(5, dtype=torch.long), + } + negative = torch.randn(5, 16) + + result = collate_fn_text_prompts([sample, sample], negative_text_embeddings=negative) + + assert "image_latents" not in result + assert result["text_embeddings"].shape == (2, 5, 16) + assert result["text_embeddings_mask"].shape == (2, 5) + assert result["negative_text_embeddings"].shape == (2, 5, 16) + + def test_collate_zero_pads_variable_length_qwen_embeddings_and_masks(): pytest.importorskip("nemo_automodel") torch = pytest.importorskip("torch") diff --git a/tests/unit/torch/fastgen/test_pdd_config.py b/tests/unit/torch/fastgen/test_pdd_config.py index ffb484216a4..186c757d0b3 100644 --- a/tests/unit/torch/fastgen/test_pdd_config.py +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -62,11 +62,13 @@ def test_pdd_config_accepts_supported_schedule_and_adapter_time_scale(): student_sample_steps=2, teacher_integrator="midpoint", num_train_timesteps=1000, + data_free=True, ) assert config.inference_blocks == [64, 64] assert config.teacher_integrator == "midpoint" assert config.num_train_timesteps == 1000 + assert config.data_free is True def test_rejected_attribute_assignment_leaves_pdd_config_unchanged(): @@ -147,7 +149,6 @@ def test_pdd_config_accepts_inference_partition_outside_training_block_support() {"student_sample_type": "sde"}, {"teacher_integrator": "heun"}, {"teacher_integrator": "rk4"}, - {"data_free": True}, ], ) def test_pdd_config_locks_algorithm_modes(overrides): diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index 75716a0c005..24013526f7c 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -273,7 +273,7 @@ def test_euler_loss_matches_analytic_empty_and_tail_reconstruction() -> None: 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([0, 1])) + 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"]) @@ -309,6 +309,46 @@ def test_midpoint_target_uses_exact_final_interval_midpoint() -> None: 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 @@ -472,6 +512,17 @@ def test_fused_sampler_matches_explicit_block_updates(blocks) -> None: 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( diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index c27842b71c8..235c8216914 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -35,6 +35,7 @@ QWEN_IMAGE_PDD_LAYER_SPEC, adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, + freeze_qwen_image_mr210_unused_parameters, require_qwen_image_mr210_forward, ) @@ -323,7 +324,7 @@ def test_fused_student_matches_explicit_packed_weight_fusion() -> None: assert student.proj_out(projection.weight.new_zeros(1, 5)).shape[-1] == 16 -def test_teacher_cfg_uses_mr210_global_fp32_norm_rescale() -> None: +def test_teacher_cfg_uses_qwen_per_token_fp32_norm_rescale() -> None: teacher = _TinyQwenTransformer() config = _config(guidance_scale=4.0) adapter = QwenImagePDDAdapter(config) @@ -338,17 +339,17 @@ def test_teacher_cfg_uses_mr210_global_fp32_norm_rescale() -> None: ) assert len(teacher.calls) == 2 - conditional = unpack_latents(teacher.calls[0]["output"], 4, 4) - unconditional = unpack_latents(teacher.calls[1]["output"], 4, 4) + conditional = teacher.calls[0]["output"] + unconditional = teacher.calls[1]["output"] 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, 2, 3), + dim=-1, keepdim=True, - ) / torch.linalg.vector_norm(guided_fp32, dim=(1, 2, 3), keepdim=True).clamp_min(1e-5) - expected = (guided_fp32 * factor).to(conditional.dtype) + ) / torch.linalg.vector_norm(guided_fp32, dim=-1, keepdim=True).clamp_min(1e-5) + expected = unpack_latents((guided_fp32 * factor).to(conditional.dtype), 4, 4) assert actual.dtype == torch.bfloat16 torch.testing.assert_close(actual, expected) @@ -357,6 +358,59 @@ def test_teacher_cfg_uses_mr210_global_fp32_norm_rescale() -> None: assert all("txt_seq_lens" not in call["kwargs"] for call in teacher.calls) +def test_teacher_cfg_remote_preflight_failure_stops_both_model_calls(monkeypatch) -> None: + teacher = _TinyQwenTransformer() + adapter = QwenImagePDDAdapter(_config(guidance_scale=4.0)) + state, time, condition, negative_condition = _inputs() + + monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_available", lambda: True) + monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_initialized", lambda: True) + + def report_remote_failure(failed, *, op): + assert not bool(failed) + assert op is torch.distributed.ReduceOp.MAX + failed.fill_(1) + + monkeypatch.setattr(qwen_image_pdd_plugin.dist, "all_reduce", report_remote_failure) + + with pytest.raises(RuntimeError, match="preflight failed on another rank"): + adapter.teacher_velocity( + teacher, + state, + time, + condition=condition, + negative_condition=negative_condition, + ) + + assert teacher.calls == [] + + +def test_teacher_cfg_local_missing_negative_condition_fails_collectively(monkeypatch) -> None: + teacher = _TinyQwenTransformer() + adapter = QwenImagePDDAdapter(_config(guidance_scale=4.0)) + state, time, condition, _ = _inputs() + + monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_available", lambda: True) + monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_initialized", lambda: True) + + def preserve_local_failure(failed, *, op): + assert bool(failed) + assert op is torch.distributed.ReduceOp.MAX + + monkeypatch.setattr(qwen_image_pdd_plugin.dist, "all_reduce", preserve_local_failure) + + with pytest.raises(TypeError, match="negative_condition must be a tuple"): + adapter.teacher_velocity( + teacher, + state, + time, + condition=condition, + negative_condition=None, + ) + + assert teacher.calls == [] + + def test_teacher_cfg_stays_in_model_output_dtype() -> None: class LowPrecisionTeacher(_QwenImageTestDouble): def __init__(self) -> None: @@ -383,18 +437,18 @@ def forward(self, *, hidden_states, encoder_hidden_states, **kwargs): ) assert actual.dtype == torch.bfloat16 - conditional, unconditional = (unpack_latents(output, 4, 4) for output in teacher.outputs) + conditional, unconditional = teacher.outputs 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, 2, 3), keepdim=True - ) / torch.linalg.vector_norm(guided_fp32, dim=(1, 2, 3), keepdim=True).clamp_min(1e-5) - expected = (guided_fp32 * factor).to(torch.bfloat16) + 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) torch.testing.assert_close(actual, expected, rtol=0, atol=0) -def test_teacher_cfg_zero_guided_norm_uses_mr210_clamp() -> None: +def test_teacher_cfg_zero_guided_norm_uses_qwen_clamp() -> None: class ZeroGuidedTeacher(_QwenImageTestDouble): def __init__(self) -> None: super().__init__() @@ -497,19 +551,17 @@ def teacher_velocity(self, _model, state, time, **kwargs): time_k, negative_condition, ) - conditional = _unpack_oracle(conditional_packed, 4, 4) - unconditional = _unpack_oracle(unconditional_packed, 4, 4) - guided_low_precision = conditional + 3.0 * (conditional - unconditional) - conditional_fp32 = conditional.float() + guided_low_precision = conditional_packed + 3.0 * (conditional_packed - unconditional_packed) + conditional_fp32 = conditional_packed.float() guided_fp32 = guided_low_precision.float() - norm_dims = (1, 2, 3) - teacher_target_low_precision = ( + teacher_target_packed = ( guided_fp32 * ( - torch.linalg.vector_norm(conditional_fp32, dim=norm_dims, keepdim=True) - / torch.linalg.vector_norm(guided_fp32, dim=norm_dims, keepdim=True).clamp_min(1e-5) + 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() @@ -592,6 +644,33 @@ def _tiny_diffusers_qwen(): ) +def test_mr210_freezes_exactly_the_structurally_unused_parameters() -> None: + student = _tiny_diffusers_qwen() + + frozen_names = freeze_qwen_image_mr210_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, @@ -720,27 +799,19 @@ def oracle_forward(model, state, time, current_condition): 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 = _unpack_oracle( - oracle_forward(oracle_teacher, x_bar_k, time_k, condition), - 4, - 4, - ) - unconditional = _unpack_oracle( - oracle_forward(oracle_teacher, x_bar_k, time_k, negative_condition), - 4, - 4, - ) + 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() - norm_dims = (1, 2, 3) - teacher_target_low_precision = ( + teacher_target_packed = ( guided_fp32 * ( - torch.linalg.vector_norm(conditional_fp32, dim=norm_dims, keepdim=True) - / torch.linalg.vector_norm(guided_fp32, dim=norm_dims, keepdim=True).clamp_min(1e-5) + 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() @@ -1009,7 +1080,7 @@ def capture_time(_module, args): assert captured[0].item() != time.to(torch.bfloat16).float().item() -def test_mr210_qwen_teacher_cfg_matches_global_reference() -> None: +def test_mr210_qwen_teacher_cfg_matches_per_token_reference() -> None: teacher = adopt_qwen_image_mr210_forward(_tiny_diffusers_qwen().eval().to(torch.bfloat16)) config = _config(guidance_scale=4.0) adapter = QwenImagePDDAdapter(config) @@ -1038,15 +1109,15 @@ def direct_packed(current_condition): )[0] with torch.no_grad(): - conditional = unpack_latents(direct_packed(condition), 4, 4) - unconditional = unpack_latents(direct_packed(negative_condition), 4, 4) + 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, 2, 3), keepdim=True - ) / torch.linalg.vector_norm(guided_fp32, dim=(1, 2, 3), keepdim=True).clamp_min(1e-5) - expected = (guided_fp32 * factor).to(torch.bfloat16) + 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) actual = adapter.teacher_velocity( teacher, state, @@ -1095,7 +1166,7 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N config = _config() adapter = QwenImagePDDAdapter(config) state, time, condition, _ = _inputs() - with pytest.raises(ValueError, match="negative_condition is required"): + with pytest.raises(TypeError, match="negative_condition must be a tuple"): adapter.teacher_velocity(transformer, state, time, condition=condition) with pytest.raises(TypeError, match="negative_condition must be a tuple"): adapter.teacher_velocity( From 0c18015c0295b9606267beb1d956ff62644239de Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Tue, 11 Aug 2026 19:16:54 -0700 Subject: [PATCH 38/45] Align PDD integration with ModelOpt conventions Signed-off-by: Meng Xin --- examples/diffusers/fastgen/pdd/README.md | 5 + examples/diffusers/fastgen/pdd/__init__.py | 2 + examples/diffusers/fastgen/pdd/compat.py | 134 ++++++++++++++++++ .../fastgen/pdd/configs/qwen_image.yaml | 2 +- .../fastgen/pdd/inference_qwen_image.py | 4 +- examples/diffusers/fastgen/pdd/recipe.py | 76 +--------- modelopt/torch/fastgen/__init__.py | 15 +- modelopt/torch/fastgen/methods/pdd.py | 14 +- .../torch/fastgen/plugins/qwen_image_pdd.py | 105 +++++++------- .../examples/diffusers/fastgen/test_layout.py | 119 ++-------------- .../fastgen/test_pdd_recipe_setup.py | 36 +++-- .../unit/torch/fastgen/test_pdd_projection.py | 21 +-- .../unit/torch/fastgen/test_pdd_public_api.py | 7 +- .../fastgen/test_qwen_image_pdd_plugin.py | 54 ++++--- 14 files changed, 274 insertions(+), 320 deletions(-) create mode 100644 examples/diffusers/fastgen/pdd/compat.py diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index b3bfd17f462..0252cd60a66 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -29,6 +29,11 @@ clipping, optimizer and learning-rate state, step scheduling, SIGTERM handling, `LATEST`, and resume. The example does not define a custom training loop or checkpoint manager and does not modify AutoModel, Diffusers, or Qwen source. +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()`. + ## Prepare the student Widen the Qwen output projection before AutoModel constructs FSDP and the optimizer: diff --git a/examples/diffusers/fastgen/pdd/__init__.py b/examples/diffusers/fastgen/pdd/__init__.py index 3d2b6c232ef..029ab41833a 100644 --- a/examples/diffusers/fastgen/pdd/__init__.py +++ b/examples/diffusers/fastgen/pdd/__init__.py @@ -14,3 +14,5 @@ # 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..045373d3094 --- /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: + @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 index 0023b42ada9..72669dc4b51 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -73,7 +73,7 @@ fsdp: cp_size: 1 pp_size: 1 ep_size: 1 - # The PDD recipe enables Qwen's native block checkpointing after binding the MR210 forward. + # The PDD recipe enables Qwen's native block checkpointing after binding the PDD forward. activation_checkpointing: false data: diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index 4298744447a..a12c87e54ec 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -36,7 +36,7 @@ from modelopt.torch.fastgen import PDDConfig, PDDPipeline # noqa: E402 from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( # noqa: E402 QwenImagePDDAdapter, - adopt_qwen_image_mr210_forward, + enable_qwen_image_pdd_forward, restore_qwen_image_pdd_projection, ) @@ -124,7 +124,7 @@ def main() -> None: low_cpu_mem_usage=True, ) restore_qwen_image_pdd_projection(transformer, config) - adopt_qwen_image_mr210_forward(transformer) + enable_qwen_image_pdd_forward(transformer) transformer.eval() pipe = QwenImagePipeline.from_pretrained( diff --git a/examples/diffusers/fastgen/pdd/recipe.py b/examples/diffusers/fastgen/pdd/recipe.py index 18643f2effd..e589aff2924 100644 --- a/examples/diffusers/fastgen/pdd/recipe.py +++ b/examples/diffusers/fastgen/pdd/recipe.py @@ -18,16 +18,13 @@ from __future__ import annotations import logging -from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Mapping from pathlib import Path from typing import Any -import torch import torch.distributed as dist from huggingface_hub import snapshot_download from torch import nn -from torch.distributed.fsdp import MixedPrecisionPolicy try: import nemo_automodel.recipes.diffusion.train as automodel_diffusion_train @@ -43,44 +40,13 @@ from modelopt.torch.fastgen import PDDConfig, PDDPipeline from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( QwenImagePDDAdapter, - adopt_qwen_image_mr210_forward, - freeze_qwen_image_mr210_unused_parameters, + enable_qwen_image_pdd_forward, ) +from .compat import automodel_pdd_setup from .training import PDDFlowMatchingStepAdapter -@contextmanager -def _preserve_fp32_timestep_inputs() -> Iterator[None]: - """Keep Qwen's continuous timestep in FP32 while FSDP computes in BF16.""" - original_builder = automodel_diffusion_train._build_diffusion_parallel_manager_args - - 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 - - automodel_diffusion_train._build_diffusion_parallel_manager_args = build_manager_args - try: - yield - finally: - automodel_diffusion_train._build_diffusion_parallel_manager_args = original_builder - - def _config_mapping(value: Any) -> dict[str, Any]: if hasattr(value, "to_dict"): return value.to_dict() @@ -112,47 +78,19 @@ def _validate_prepared_student(model: nn.Module, config: PDDConfig) -> None: ) -@contextmanager -def _freeze_unused_qwen_parameters_before_optimizer() -> Iterator[None]: - """Freeze unused Qwen outputs before AutoModel collects optimizer parameters.""" - pipeline_cls = automodel_diffusion_train.NeMoAutoDiffusionPipeline - original_descriptor = pipeline_cls.__dict__["from_pretrained"] - original_from_pretrained = pipeline_cls.from_pretrained - - def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any: - del cls - pipe, managers = original_from_pretrained(*args, **kwargs) - if not kwargs.get("load_for_training", False): - return pipe, managers - model = pipe.transformer - frozen_names = freeze_qwen_image_mr210_unused_parameters(model) - logging.info( - "[PDD] Full training excludes %d unused final-block text-output tensors: %s", - len(frozen_names), - ", ".join(frozen_names), - ) - return pipe, managers - - pipeline_cls.from_pretrained = classmethod(from_pretrained) - try: - yield - finally: - pipeline_cls.from_pretrained = original_descriptor - - class PDDDiffusionRecipe(TrainDiffusionRecipe): """Use AutoModel's native lifecycle with a PDD loss and frozen teacher.""" def setup(self) -> None: - with _preserve_fp32_timestep_inputs(), _freeze_unused_qwen_parameters_before_optimizer(): + 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 MR210 forward here changes behavior only; it creates no parameters. - adopt_qwen_image_mr210_forward(self.model) + # 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() @@ -214,6 +152,6 @@ def _load_teacher(self) -> nn.Module: ) teacher = pipe.transformer del pipe - adopt_qwen_image_mr210_forward(teacher) + enable_qwen_image_pdd_forward(teacher) teacher.eval().requires_grad_(False) return teacher diff --git a/modelopt/torch/fastgen/__init__.py b/modelopt/torch/fastgen/__init__.py index da576540145..961c39c04e6 100644 --- a/modelopt/torch/fastgen/__init__.py +++ b/modelopt/torch/fastgen/__init__.py @@ -67,14 +67,7 @@ from .methods.pdd import * from .pipeline import * - -def __getattr__(name: str): - """Load optional model plugins only when the public namespace is requested.""" - if name != "plugins": - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - import importlib - - loaded = importlib.import_module(f"{__name__}.plugins") - globals()[name] = loaded - return loaded +# isort: off +# 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/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index bcd96f6877f..611a63bcc93 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -48,8 +48,6 @@ "PDDOutputProjection", "PDDPipeline", "convert_to_pdd_output_projection", - "get_module_by_path", - "replace_module_by_path", ] PDDHeadLayout = Literal["channel_major", "patch_major"] @@ -335,7 +333,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return F.linear(input, fused_weight, fused_bias) -def get_module_by_path(model: nn.Module, path: str) -> nn.Module: +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__}.") @@ -349,13 +347,13 @@ def get_module_by_path(model: nn.Module, path: str) -> nn.Module: ) from error -def replace_module_by_path(model: nn.Module, path: str, replacement: nn.Module) -> nn.Module: +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) + 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 + parent = _get_module_by_path(model, parent_path) if parent_path else model setattr(parent, name, replacement) return previous @@ -366,7 +364,7 @@ def convert_to_pdd_output_projection( grid_size: int, ) -> PDDOutputProjection: """Explicitly replace ``layer_spec.projection_path`` with a PDD projection.""" - current = get_module_by_path(model, layer_spec.projection_path) + 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, " @@ -374,7 +372,7 @@ def convert_to_pdd_output_projection( ) 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) + 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 diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index fa8b95be2ed..bff9dbd402d 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -1,5 +1,6 @@ # Adapted from the Qwen-Image implementation in Diffusers: # https://github.com/huggingface/diffusers/blob/275869dcae4ebcfee6a80253fdabc56033335020/src/diffusers/models/transformers/transformer_qwenimage.py +# The masked joint-attention execution follows FastGen merge request 210. # SPDX-FileCopyrightText: Copyright (c) 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved. # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 @@ -37,14 +38,14 @@ "QWEN_IMAGE_PDD_EXECUTION", "QWEN_IMAGE_PDD_LAYER_SPEC", "QwenImagePDDAdapter", - "adopt_qwen_image_mr210_forward", "convert_qwen_image_to_pdd", - "freeze_qwen_image_mr210_unused_parameters", - "require_qwen_image_mr210_forward", + "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 = "fastgen_mr210" +QWEN_IMAGE_PDD_EXECUTION = "qwen_image_pdd_masked_joint_attention_v1" QWEN_IMAGE_PDD_LAYER_SPEC = PDDLayerSpec( projection_path="transformer.proj_out", @@ -67,7 +68,7 @@ } _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE = "_modelopt_qwen_image_pdd_execution" -_QWEN_IMAGE_MR210_CHILDREN = ( +_QWEN_IMAGE_PDD_CHILDREN = ( "pos_embed", "time_text_embed", "txt_norm", @@ -99,7 +100,7 @@ def _config_value(transformer: nn.Module, name: str, default: Any = None) -> Any return getattr(config, name, default) -def _qwen_image_mr210_forward( +def _qwen_image_pdd_forward( self: nn.Module, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, @@ -114,64 +115,64 @@ def _qwen_image_mr210_forward( additional_t_cond: torch.Tensor | None = None, return_dict: bool = True, ) -> Any: - """Run the regular-output Qwen forward used by FastGen MR210.""" + """Run Qwen-Image with the masked joint-attention contract required by PDD.""" if hidden_states.ndim != 3 or hidden_states.dtype != torch.bfloat16: - raise TypeError("Qwen MR210 hidden_states must be packed BF16 [B, P, C].") + raise TypeError("Qwen PDD hidden_states must be packed BF16 [B, P, C].") if ( not isinstance(encoder_hidden_states, torch.Tensor) or encoder_hidden_states.ndim != 3 or encoder_hidden_states.dtype != torch.bfloat16 ): - raise TypeError("Qwen MR210 encoder_hidden_states must be BF16 [B, S, D].") + raise TypeError("Qwen PDD encoder_hidden_states must be BF16 [B, S, D].") if not isinstance(encoder_hidden_states_mask, torch.Tensor): - raise TypeError("Qwen MR210 requires encoder_hidden_states_mask.") + raise TypeError("Qwen PDD requires encoder_hidden_states_mask.") if not isinstance(timestep, torch.Tensor) or timestep.dtype != torch.float32: - raise TypeError("Qwen MR210 timestep must remain FP32 at transformer entry.") + raise TypeError("Qwen PDD timestep must remain FP32 at transformer entry.") batch_size = hidden_states.shape[0] if encoder_hidden_states.shape[0] != batch_size: - raise ValueError("Qwen MR210 image and text batch sizes must match.") + raise ValueError("Qwen PDD image and text batch sizes must match.") if timestep.shape != (batch_size,): - raise ValueError("Qwen MR210 timestep must contain one value per batch item.") + raise ValueError("Qwen PDD timestep must contain one value per batch item.") if img_shapes is None or len(img_shapes) != batch_size: - raise ValueError("Qwen MR210 img_shapes must contain one entry per batch item.") + raise ValueError("Qwen PDD img_shapes must contain one entry per batch item.") if attention_kwargs: - raise ValueError("Qwen MR210 PDD does not support nonempty attention_kwargs.") + raise ValueError("Qwen PDD does not support nonempty attention_kwargs.") if guidance is not None: - raise ValueError("Qwen MR210 PDD does not support transformer guidance embeddings.") + raise ValueError("Qwen PDD does not support transformer guidance embeddings.") if controlnet_block_samples is not None: - raise ValueError("Qwen MR210 PDD does not support ControlNet block samples.") + raise ValueError("Qwen PDD does not support ControlNet block samples.") if additional_t_cond is not None: - raise ValueError("Qwen MR210 PDD does not support additional timestep conditioning.") + raise ValueError("Qwen PDD does not support additional timestep conditioning.") if type(return_dict) is not bool: - raise TypeError("Qwen MR210 return_dict must be a bool.") + raise TypeError("Qwen PDD return_dict must be a bool.") if encoder_hidden_states_mask.ndim != 2 or tuple(encoder_hidden_states_mask.shape) != tuple( encoder_hidden_states.shape[:2] ): - raise ValueError("Qwen MR210 mask must match the text batch and sequence dimensions.") + raise ValueError("Qwen PDD mask must match the text batch and sequence dimensions.") if encoder_hidden_states_mask.device != encoder_hidden_states.device: - raise ValueError("Qwen MR210 mask and text embeddings must share a device.") + raise ValueError("Qwen PDD mask and text embeddings must share a device.") if ( encoder_hidden_states_mask.dtype.is_floating_point or encoder_hidden_states_mask.dtype.is_complex ): - raise TypeError("Qwen MR210 mask must use an integer or boolean dtype.") - _require_binary_mask(encoder_hidden_states_mask, name="Qwen MR210") + raise TypeError("Qwen PDD mask must use an integer or boolean dtype.") + _require_binary_mask(encoder_hidden_states_mask, name="Qwen PDD") expected_max_txt_seq_len = int( encoder_hidden_states_mask.sum(dim=1).max().to(torch.int32).item() ) if txt_seq_lens is not None: expected_txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).to(torch.int32).tolist() if txt_seq_lens != expected_txt_seq_lens: - raise ValueError("Qwen MR210 txt_seq_lens must equal the valid mask lengths.") + raise ValueError("Qwen PDD txt_seq_lens must equal the valid mask lengths.") if max_txt_seq_len is None: max_txt_seq_len = expected_max_txt_seq_len elif max_txt_seq_len != expected_max_txt_seq_len: - raise ValueError("Qwen MR210 max_txt_seq_len must equal the maximum valid mask length.") + raise ValueError("Qwen PDD max_txt_seq_len must equal the maximum valid mask length.") hidden_states = self.img_in(hidden_states) encoder_hidden_states = self.txt_in(self.txt_norm(encoder_hidden_states)) if timestep.dtype != torch.float32: - raise RuntimeError("Qwen MR210 timestep was rounded before time_text_embed.") + raise RuntimeError("Qwen PDD timestep was rounded before time_text_embed.") temb = self.time_text_embed(timestep, hidden_states) image_rotary_emb = self.pos_embed( img_shapes, @@ -220,28 +221,28 @@ def _qwen_image_mr210_forward( return Transformer2DModelOutput(sample=output) -def _is_qwen_image_mr210_forward(model: nn.Module) -> bool: +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_mr210_forward + 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_mr210_forward(model: nn.Module) -> str: - """Require and return the semantic label for the exact bound MR210 forward.""" - if not isinstance(model, nn.Module) or not _is_qwen_image_mr210_forward(model): - raise RuntimeError("Qwen-Image PDD requires the bound FastGen MR210 forward 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_mr210_unused_parameters(transformer: nn.Module) -> tuple[str, ...]: - """Freeze final-block text outputs that the MR210 forward does not consume.""" +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 MR210 requires a nonempty transformer_blocks ModuleList.") + raise RuntimeError("Qwen PDD requires a nonempty transformer_blocks ModuleList.") final_block = blocks[-1] local_names = ( @@ -258,18 +259,18 @@ def freeze_qwen_image_mr210_unused_parameters(transformer: nn.Module) -> tuple[s parameter = final_block.get_parameter(local_name) except AttributeError as error: raise RuntimeError( - f"Qwen MR210 final block is missing required parameter {local_name!r}." + 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 adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: - """Bind FastGen MR210's regular Qwen forward to the loaded root in place.""" +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_mr210_forward(transformer): + if _is_qwen_image_pdd_forward(transformer): return transformer # Diffusers is an optional dependency used only by the Qwen example. @@ -277,7 +278,7 @@ def adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: if not isinstance(transformer, QwenImageTransformer2DModel): raise TypeError( - "MR210 forward adoption requires the supported QwenImageTransformer2DModel, " + "Qwen PDD forward binding requires QwenImageTransformer2DModel, " f"got {type(transformer).__name__}." ) existing_forward = transformer.__dict__.get("forward") @@ -285,29 +286,29 @@ def adopt_qwen_image_mr210_forward(transformer: nn.Module) -> nn.Module: raise RuntimeError("Qwen root already has a different instance-level forward override.") missing = [ name - for name in _QWEN_IMAGE_MR210_CHILDREN + for name in _QWEN_IMAGE_PDD_CHILDREN if not isinstance(getattr(transformer, name, None), nn.Module) ] if missing: - raise RuntimeError(f"Qwen root is missing required MR210 modules: {missing}.") + raise RuntimeError(f"Qwen root is missing required PDD forward modules: {missing}.") if ( not isinstance(transformer.transformer_blocks, nn.ModuleList) or not transformer.transformer_blocks ): - raise RuntimeError("Qwen MR210 requires a nonempty transformer_blocks ModuleList.") + raise RuntimeError("Qwen PDD requires a nonempty transformer_blocks ModuleList.") if _config_guidance_embeds(transformer): - raise ValueError("Qwen MR210 PDD does not support transformer guidance embeddings.") + raise ValueError("Qwen PDD does not support transformer guidance embeddings.") if getattr(transformer, "peft_config", None): - raise ValueError("Qwen MR210 PDD does not support active PEFT adapters.") + 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 MR210 PDD does not support fused QKV projections.") + 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 MR210 PDD requires {name}=False.") + raise ValueError(f"Qwen PDD requires {name}=False.") - transformer.forward = types.MethodType(_qwen_image_mr210_forward, transformer) + transformer.forward = types.MethodType(_qwen_image_pdd_forward, transformer) setattr(transformer, _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE, QWEN_IMAGE_PDD_EXECUTION) - require_qwen_image_mr210_forward(transformer) + require_qwen_image_pdd_forward(transformer) return transformer @@ -531,7 +532,7 @@ def _prepare_call( condition_name: str, ) -> tuple[torch.Tensor, torch.Tensor]: self._validate_state_and_time(state, time) - require_qwen_image_mr210_forward(model) + require_qwen_image_pdd_forward(model) if _config_guidance_embeds(model): raise ValueError("Qwen-Image PDD does not support transformer guidance embeddings.") encoder_hidden_states, attention_mask = self._parse_condition( @@ -570,9 +571,9 @@ def _call_packed( batch_size, _, height, width = state.shape model_dtype = self._model_dtype(model, state.dtype) if model_dtype != torch.bfloat16: - raise TypeError("Qwen MR210 PDD execution requires BF16 compute.") + raise TypeError("Qwen PDD execution requires BF16 compute.") if time.dtype != torch.float32: - raise TypeError("Qwen MR210 PDD execution requires FP32 time.") + raise TypeError("Qwen PDD execution requires FP32 time.") packed_state = pack_latents(state).to(model_dtype) encoder_hidden_states = encoder_hidden_states.to(model_dtype) max_txt_seq_len = int(attention_mask.sum(dim=1).max().to(torch.int32).item()) diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py index c4aaee9d616..db937491ec3 100644 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ b/tests/examples/diffusers/fastgen/test_layout.py @@ -20,7 +20,6 @@ import json import os import pathlib -import re import subprocess import sys @@ -28,96 +27,17 @@ _REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] _FASTGEN_ROOT = _REPO_ROOT / "examples" / "diffusers" / "fastgen" -_EXPECTED_ROOT_ENTRIES = { - "README.md", - "dmd2", - "fastgen_data", - "make_negative_prompt_embedding.py", - "preprocess", - "preprocess_qwen_image.py", - "pdd", - "requirements.txt", -} -_EXPECTED_DMD2_FILES = { - "README.md", - "__init__.py", - "checkpoint.py", - "configs", - "export_qwen_image.py", - "finetune.py", - "inference_qwen_image.py", - "recipe.py", -} -_EXPECTED_PDD_FILES = { - "README.md", - "__init__.py", - "configs", - "finetune.py", - "inference_qwen_image.py", - "prepare_qwen_image.py", - "recipe.py", - "training.py", -} -_TEXT_SUFFIXES = {".json", ".md", ".py", ".rst", ".sh", ".toml", ".txt", ".yaml", ".yml"} -def _old_name(prefix: str, suffix: str) -> str: - return f"{prefix}_{suffix}" +def test_fastgen_algorithm_entrypoints_are_namespaced() -> None: + for algorithm in ("dmd2", "pdd"): + package = _FASTGEN_ROOT / algorithm + assert (package / "__init__.py").is_file() + assert (package / "finetune.py").is_file() + assert (package / "configs" / "qwen_image.yaml").is_file() - -def _old_modules() -> tuple[str, ...]: - return ( - _old_name("dmd2", "finetune"), - _old_name("dmd2", "recipe"), - _old_name("fastgen", "checkpoint"), - _old_name("export", "diffusers_qwen_image"), - _old_name("inference", "dmd2_qwen_image"), - _old_name("pdd", "artifacts"), - _old_name("pdd", "checkpoint"), - _old_name("pdd", "export"), - _old_name("pdd", "finetune"), - _old_name("pdd", "recipe"), - _old_name("pdd", "training"), - _old_name("export", "pdd_qwen_image"), - _old_name("inference", "pdd_qwen_image"), - ) - - -def _source_text_files() -> list[pathlib.Path]: - completed = subprocess.run( - ["git", "ls-files", "-co", "--exclude-standard", "-z"], - cwd=_REPO_ROOT, - check=True, - capture_output=True, - ) - return [ - _REPO_ROOT / relative - for relative in completed.stdout.decode().split("\0") - if relative - and pathlib.Path(relative).suffix in _TEXT_SUFFIXES - and (_REPO_ROOT / relative).is_file() - ] - - -def test_fastgen_root_has_closed_shared_and_algorithm_ownership() -> None: - root_entries = {path.name for path in _FASTGEN_ROOT.iterdir() if path.name != "__pycache__"} - assert root_entries == _EXPECTED_ROOT_ENTRIES - assert not (_FASTGEN_ROOT / "configs").exists() - - dmd2_entries = { - path.name for path in (_FASTGEN_ROOT / "dmd2").iterdir() if path.name != "__pycache__" - } - assert dmd2_entries == _EXPECTED_DMD2_FILES - assert {path.name for path in (_FASTGEN_ROOT / "dmd2" / "configs").iterdir()} == { - "qwen_image.yaml" - } - pdd_entries = { - path.name for path in (_FASTGEN_ROOT / "pdd").iterdir() if path.name != "__pycache__" - } - assert pdd_entries == _EXPECTED_PDD_FILES - assert {path.name for path in (_FASTGEN_ROOT / "pdd" / "configs").iterdir()} == { - "qwen_image.yaml" - } + assert not (_FASTGEN_ROOT / "dmd2_finetune.py").exists() + assert not (_FASTGEN_ROOT / "pdd_finetune.py").exists() def test_dmd2_config_retains_accepted_semantics() -> None: @@ -133,29 +53,6 @@ def test_dmd2_config_retains_accepted_semantics() -> None: assert "metadata_index" not in value["data"]["dataloader"] -def test_repository_sources_have_no_flat_algorithm_paths() -> None: - old_modules = _old_modules() - stale = ( - *(f"{module}.py" for module in old_modules), - "configs/" + _old_name("dmd2", "qwen_image") + ".yaml", - ) - failures: list[str] = [] - for path in _source_text_files(): - if path == pathlib.Path(__file__): - continue - text = path.read_text(errors="strict") - relative = path.relative_to(_REPO_ROOT).as_posix() - failures.extend(f"{relative}: {token}" for token in stale if token in text) - if path.suffix == ".py": - failures.extend( - f"{relative}: stale import {module}" - for module in old_modules - if re.search(rf"(?m)^\s*(?:from|import)\s+{re.escape(module)}(?:\s|\.|$)", text) - or re.search(rf"['\"]{re.escape(module)}['\"]", text) - ) - assert not failures, "\n".join(failures) - - def test_dmd2_package_is_import_light() -> None: probe = f""" import importlib, json, sys diff --git a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index ab61a300b79..968069ad4e2 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -30,12 +30,8 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -from pdd import recipe as pdd_recipe -from pdd.recipe import ( - _freeze_unused_qwen_parameters_before_optimizer, - _preserve_fp32_timestep_inputs, - _validate_prepared_student, -) +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 @@ -146,12 +142,12 @@ def build_manager_args(**kwargs): return {"_manager_type": "fsdp2"} monkeypatch.setattr( - pdd_recipe.automodel_diffusion_train, + pdd_compat.automodel_diffusion_train, "_build_diffusion_parallel_manager_args", build_manager_args, ) - with _preserve_fp32_timestep_inputs(): - manager_args = pdd_recipe.automodel_diffusion_train._build_diffusion_parallel_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, @@ -163,7 +159,7 @@ def build_manager_args(**kwargs): assert policy.output_dtype == torch.bfloat16 assert policy.cast_forward_inputs is False assert ( - pdd_recipe.automodel_diffusion_train._build_diffusion_parallel_manager_args + pdd_compat.automodel_diffusion_train._build_diffusion_parallel_manager_args is build_manager_args ) @@ -176,15 +172,16 @@ def from_pretrained(cls, *args, **kwargs): return SimpleNamespace(transformer=_PreparedQwenStudent(out_features=32)), {} monkeypatch.setattr( - pdd_recipe.automodel_diffusion_train, + pdd_compat.automodel_diffusion_train, "NeMoAutoDiffusionPipeline", _Pipeline, ) original_descriptor = _Pipeline.__dict__["from_pretrained"] - with _freeze_unused_qwen_parameters_before_optimizer(): - student, _ = _Pipeline.from_pretrained("student", load_for_training=True) - teacher, _ = _Pipeline.from_pretrained("teacher", load_for_training=False) + with pdd_compat.automodel_pdd_setup(): + setup_pipeline = pdd_compat.automodel_diffusion_train.NeMoAutoDiffusionPipeline + student, _ = setup_pipeline.from_pretrained("student", load_for_training=True) + teacher, _ = setup_pipeline.from_pretrained("teacher", load_for_training=False) frozen_names = { name @@ -201,6 +198,17 @@ def from_pretrained(cls, *args, **kwargs): } 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]) diff --git a/tests/unit/torch/fastgen/test_pdd_projection.py b/tests/unit/torch/fastgen/test_pdd_projection.py index 398e685ed81..8884f8d079b 100644 --- a/tests/unit/torch/fastgen/test_pdd_projection.py +++ b/tests/unit/torch/fastgen/test_pdd_projection.py @@ -30,8 +30,6 @@ PDDLayerSpec, PDDOutputProjection, convert_to_pdd_output_projection, - get_module_by_path, - replace_module_by_path, ) @@ -149,32 +147,19 @@ def test_layer_spec_rejects_unsupported_layout_metadata(kwargs): def test_nested_conversion_is_explicit_idempotent_and_conflict_safe(): model = _NestedModel() spec = _spec("channel_major") - original = get_module_by_path(model, spec.projection_path) + 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 get_module_by_path(model, spec.projection_path) is projection + 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 get_module_by_path(model, spec.projection_path) is projection - - -def test_nested_module_helpers_require_existing_registered_modules(): - model = _NestedModel() - replacement = nn.Linear(2, 6) - previous = replace_module_by_path(model, "transformer.blocks.1", replacement) - - assert isinstance(previous, nn.Linear) - assert get_module_by_path(model, "transformer.blocks.1") is replacement - with pytest.raises(ValueError, match="does not resolve"): - get_module_by_path(model, "transformer.missing") - with pytest.raises(ValueError, match="non-empty dotted"): - get_module_by_path(model, "") + assert model.get_submodule(spec.projection_path) is projection @pytest.mark.parametrize("layout", ["channel_major", "patch_major"]) diff --git a/tests/unit/torch/fastgen/test_pdd_public_api.py b/tests/unit/torch/fastgen/test_pdd_public_api.py index 9aae47cf634..e871098fc0d 100644 --- a/tests/unit/torch/fastgen/test_pdd_public_api.py +++ b/tests/unit/torch/fastgen/test_pdd_public_api.py @@ -36,8 +36,6 @@ PDDOutputProjection, PDDPipeline, convert_to_pdd_output_projection, - get_module_by_path, - replace_module_by_path, ) _CORE_SOURCES = ( @@ -93,17 +91,15 @@ def test_core_pdd_symbols_are_exported_from_the_public_package() -> None: "PDDPipeline": PDDPipeline, "convert_to_pdd_output_projection": convert_to_pdd_output_projection, "fusion_coefficients": fusion_coefficients, - "get_module_by_path": get_module_by_path, "integrate_interval_velocities": integrate_interval_velocities, "load_pdd_config": load_pdd_config, "make_shifted_flow_grid": make_shifted_flow_grid, - "replace_module_by_path": replace_module_by_path, } assert {name: getattr(fastgen, name) for name in expected} == expected -def test_fresh_core_import_does_not_load_model_plugins_or_frameworks() -> None: +def test_fresh_core_import_does_not_require_external_frameworks() -> None: repository = Path(__file__).resolve().parents[4] script = r""" import importlib.abc @@ -139,7 +135,6 @@ def find_spec(self, fullname, path=None, target=None): "diffusers", "fastgen", "nemo_automodel", - "modelopt.torch.fastgen.plugins.qwen_image", "transformers", ): assert not any(name == prefix or name.startswith(prefix + ".") for name in loaded), ( diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 235c8216914..30c52c72edc 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -33,10 +33,10 @@ from modelopt.torch.fastgen.plugins.qwen_image_pdd import ( QWEN_IMAGE_PDD_EXECUTION, QWEN_IMAGE_PDD_LAYER_SPEC, - adopt_qwen_image_mr210_forward, convert_qwen_image_to_pdd, - freeze_qwen_image_mr210_unused_parameters, - require_qwen_image_mr210_forward, + enable_qwen_image_pdd_forward, + freeze_qwen_image_pdd_unused_parameters, + require_qwen_image_pdd_forward, ) @@ -46,7 +46,7 @@ class _QwenImageTestDouble(nn.Module): @pytest.fixture(autouse=True) def _allow_qwen_image_test_doubles(monkeypatch): - require_production_forward = qwen_image_pdd_plugin.require_qwen_image_mr210_forward + require_production_forward = qwen_image_pdd_plugin.require_qwen_image_pdd_forward def require_forward(model: nn.Module) -> str: if isinstance(model, _QwenImageTestDouble): @@ -54,13 +54,11 @@ def require_forward(model: nn.Module) -> str: getattr(model, "_modelopt_qwen_image_pdd_execution", None) != QWEN_IMAGE_PDD_EXECUTION ): - raise RuntimeError( - "Qwen-Image PDD requires the bound FastGen MR210 forward 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_mr210_forward", require_forward) + monkeypatch.setattr(qwen_image_pdd_plugin, "require_qwen_image_pdd_forward", require_forward) class _TinyQwenTransformer(_QwenImageTestDouble): @@ -647,7 +645,7 @@ def _tiny_diffusers_qwen(): def test_mr210_freezes_exactly_the_structurally_unused_parameters() -> None: student = _tiny_diffusers_qwen() - frozen_names = freeze_qwen_image_mr210_unused_parameters(student) + frozen_names = freeze_qwen_image_pdd_unused_parameters(student) assert set(frozen_names) == { "transformer_blocks.0.attn.to_add_out.weight", @@ -729,7 +727,7 @@ def teacher_velocity(self, model, state, time, **kwargs): torch.manual_seed(20260716) base = _tiny_diffusers_qwen().eval() - actual_student = adopt_qwen_image_mr210_forward(copy.deepcopy(base)) + 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) @@ -854,15 +852,15 @@ def test_conversion_preserves_the_ordinary_diffusers_qwen_root() -> None: assert dict(student.config) == config -def test_mr210_adoption_accepts_a_dynamic_qwen_subclass() -> None: +def test_qwen_pdd_forward_binding_accepts_a_dynamic_qwen_subclass() -> None: student = _tiny_diffusers_qwen().eval() student.__class__ = type("FSDPQwenImageTransformer2DModel", (type(student),), {}) - assert adopt_qwen_image_mr210_forward(student) is student - assert require_qwen_image_mr210_forward(student) == QWEN_IMAGE_PDD_EXECUTION + assert enable_qwen_image_pdd_forward(student) is student + assert require_qwen_image_pdd_forward(student) == QWEN_IMAGE_PDD_EXECUTION -def test_mr210_adoption_preserves_root_state_and_deepcopy_binding() -> None: +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()} @@ -873,7 +871,7 @@ def test_mr210_adoption_preserves_root_state_and_deepcopy_binding() -> None: source.custom_attribute = custom_attribute hook = source.register_forward_pre_hook(lambda *_args: None) - adopted = adopt_qwen_image_mr210_forward(source) + adopted = enable_qwen_image_pdd_forward(source) assert adopted is source assert type(adopted) is source_type @@ -888,7 +886,7 @@ def test_mr210_adoption_preserves_root_state_and_deepcopy_binding() -> None: ) assert hook.id in adopted._forward_pre_hooks assert adopted.custom_attribute is custom_attribute - assert adopt_qwen_image_mr210_forward(adopted) is adopted + assert enable_qwen_image_pdd_forward(adopted) is adopted round_trip = copy.deepcopy(adopted) with torch.no_grad(): @@ -903,28 +901,28 @@ def test_mr210_adoption_preserves_root_state_and_deepcopy_binding() -> None: assert teacher.forward.__func__ is adopted.forward.__func__ assert teacher.forward.__self__ is teacher assert teacher.forward.__self__ is not adopted - require_qwen_image_mr210_forward(teacher) + require_qwen_image_pdd_forward(teacher) tampered = copy.deepcopy(adopted) tampered.forward = MethodType(lambda self, **_kwargs: self, tampered) - with pytest.raises(RuntimeError, match="MR210 forward execution"): - require_qwen_image_mr210_forward(tampered) + with pytest.raises(RuntimeError, match="masked joint-attention forward"): + require_qwen_image_pdd_forward(tampered) conflicting = _tiny_diffusers_qwen() conflicting.forward = MethodType(lambda self, **_kwargs: self, conflicting) with pytest.raises(RuntimeError, match="instance-level forward override"): - adopt_qwen_image_mr210_forward(conflicting) + enable_qwen_image_pdd_forward(conflicting) forged = _tiny_diffusers_qwen() forged._modelopt_qwen_image_pdd_execution = QWEN_IMAGE_PDD_EXECUTION - with pytest.raises(RuntimeError, match="bound FastGen MR210 forward"): - qwen_image_pdd_plugin.require_qwen_image_mr210_forward(forged) + with pytest.raises(RuntimeError, match="masked joint-attention forward"): + qwen_image_pdd_plugin.require_qwen_image_pdd_forward(forged) def test_mr210_qwen_conversion_preserves_every_initialized_head() -> None: base = _tiny_diffusers_qwen().eval().to(torch.bfloat16) student = copy.deepcopy(base) - student = adopt_qwen_image_mr210_forward(student) + student = enable_qwen_image_pdd_forward(student) config = _config() generator = torch.Generator().manual_seed(20260715) state = torch.randn(2, 2, 4, 4, generator=generator) @@ -960,7 +958,7 @@ def test_mr210_qwen_conversion_preserves_every_initialized_head() -> None: def test_mr210_joint_mask_ignores_padded_token_values() -> None: canonical = _tiny_diffusers_qwen().eval().to(torch.bfloat16) student = copy.deepcopy(canonical) - student = adopt_qwen_image_mr210_forward(student) + student = enable_qwen_image_pdd_forward(student) config = _config() convert_qwen_image_to_pdd(student, config) adapter = QwenImagePDDAdapter(config) @@ -1027,7 +1025,7 @@ def capture_block_mask(_module, _args, kwargs): def test_mr210_preserves_diffusers_output_and_harmless_call_contract() -> None: - student = adopt_qwen_image_mr210_forward(_tiny_diffusers_qwen().eval().to(torch.bfloat16)) + 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( @@ -1051,7 +1049,7 @@ def test_mr210_preserves_diffusers_output_and_harmless_call_contract() -> None: def test_mr210_time_embed_receives_fp32_grid_value() -> None: - student = adopt_qwen_image_mr210_forward(_tiny_diffusers_qwen().eval().to(torch.bfloat16)) + 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] = [] @@ -1081,7 +1079,7 @@ def capture_time(_module, args): def test_mr210_qwen_teacher_cfg_matches_per_token_reference() -> None: - teacher = adopt_qwen_image_mr210_forward(_tiny_diffusers_qwen().eval().to(torch.bfloat16)) + 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) @@ -1223,7 +1221,7 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N unmarked = _TinyQwenTransformer() delattr(unmarked, "_modelopt_qwen_image_pdd_execution") convert_qwen_image_to_pdd(unmarked, config) - with pytest.raises(RuntimeError, match="MR210 forward execution"): + with pytest.raises(RuntimeError, match="masked joint-attention forward"): adapter.student_all_heads(unmarked, state, time, condition=condition) From e00ac00b93b213d84371f2e97e47283cc8d1be52 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 3 Sep 2026 21:25:16 -0700 Subject: [PATCH 39/45] tests: streamline PDD coverage Signed-off-by: Meng Xin --- .../examples/diffusers/fastgen/test_layout.py | 105 ---- .../fastgen/test_resume_dataloader.py | 50 +- .../fastgen/test_vendored_migration.py | 27 - tests/gpu/torch/fastgen/test_pdd_toy.py | 84 +-- tests/unit/recipe/test_loader.py | 13 - tests/unit/torch/fastgen/test_pdd_config.py | 33 +- .../fastgen/test_pdd_gradient_routing.py | 129 ----- tests/unit/torch/fastgen/test_pdd_pipeline.py | 151 +----- .../unit/torch/fastgen/test_pdd_projection.py | 75 --- .../unit/torch/fastgen/test_pdd_public_api.py | 160 ------ .../torch/fastgen/test_pdd_reference_math.py | 201 +------ .../fastgen/test_qwen_image_pdd_plugin.py | 500 +----------------- 12 files changed, 63 insertions(+), 1465 deletions(-) delete mode 100644 tests/examples/diffusers/fastgen/test_layout.py delete mode 100644 tests/unit/torch/fastgen/test_pdd_gradient_routing.py delete mode 100644 tests/unit/torch/fastgen/test_pdd_public_api.py diff --git a/tests/examples/diffusers/fastgen/test_layout.py b/tests/examples/diffusers/fastgen/test_layout.py deleted file mode 100644 index db937491ec3..00000000000 --- a/tests/examples/diffusers/fastgen/test_layout.py +++ /dev/null @@ -1,105 +0,0 @@ -# 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. - -"""Closed layout contract for the FastGen Diffusers example.""" - -from __future__ import annotations - -import json -import os -import pathlib -import subprocess -import sys - -import yaml - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] -_FASTGEN_ROOT = _REPO_ROOT / "examples" / "diffusers" / "fastgen" - - -def test_fastgen_algorithm_entrypoints_are_namespaced() -> None: - for algorithm in ("dmd2", "pdd"): - package = _FASTGEN_ROOT / algorithm - assert (package / "__init__.py").is_file() - assert (package / "finetune.py").is_file() - assert (package / "configs" / "qwen_image.yaml").is_file() - - assert not (_FASTGEN_ROOT / "dmd2_finetune.py").exists() - assert not (_FASTGEN_ROOT / "pdd_finetune.py").exists() - - -def test_dmd2_config_retains_accepted_semantics() -> None: - config_path = _FASTGEN_ROOT / "dmd2" / "configs" / "qwen_image.yaml" - value = yaml.safe_load(config_path.read_text()) - assert ( - value["data"]["dataloader"]["_target_"] - == "fastgen_data.build_text_to_image_multiresolution_dataloader" - ) - assert value["data"]["dataloader"]["negative_prompt_embedding_path"] == ( - "negative_prompt_embedding.pt" - ) - assert "metadata_index" not in value["data"]["dataloader"] - - -def test_dmd2_package_is_import_light() -> None: - probe = f""" -import importlib, json, sys -sys.path.insert(0, {str(_FASTGEN_ROOT)!r}) -before = set(sys.modules) -module = importlib.import_module('dmd2') -forbidden = ('torch', 'modelopt', 'nemo_automodel', 'diffusers', 'transformers') -loaded = sorted(name for name in set(sys.modules) - before if name.split('.')[0] in forbidden) -print(json.dumps({{'loaded': loaded, 'all': getattr(module, '__all__', None)}})) -""" - completed = subprocess.run( - [sys.executable, "-c", probe], - cwd=_REPO_ROOT, - check=True, - capture_output=True, - text=True, - env=dict(os.environ, PYTHONDONTWRITEBYTECODE="1"), - ) - assert json.loads(completed.stdout) == {"loaded": [], "all": []} - - -def test_dmd2_help_works_from_repo_root_without_training_imports() -> None: - script = _FASTGEN_ROOT / "dmd2" / "finetune.py" - probe = f""" -import json, runpy, sys -sys.argv = [{str(script)!r}, '--help'] -before = set(sys.modules) -try: - runpy.run_path({str(script)!r}, run_name='__main__') -except SystemExit as error: - if error.code != 0: - raise -forbidden = ('torch', 'modelopt', 'nemo_automodel', 'diffusers', 'transformers') -loaded = sorted(name for name in set(sys.modules) - before if name.split('.')[0] in forbidden) -print('__FASTGEN_LOADED__=' + json.dumps(loaded)) -""" - completed = subprocess.run( - [sys.executable, "-c", probe], - cwd=_REPO_ROOT, - check=True, - capture_output=True, - text=True, - env=dict(os.environ, PYTHONDONTWRITEBYTECODE="1"), - ) - assert "DMD2 Qwen-Image training" in completed.stdout - assert "--config" in completed.stdout - loaded_line = next( - line for line in completed.stdout.splitlines() if line.startswith("__FASTGEN_LOADED__=") - ) - assert json.loads(loaded_line.partition("=")[2]) == [] diff --git a/tests/examples/diffusers/fastgen/test_resume_dataloader.py b/tests/examples/diffusers/fastgen/test_resume_dataloader.py index 18fd9514635..2610affdfaf 100644 --- a/tests/examples/diffusers/fastgen/test_resume_dataloader.py +++ b/tests/examples/diffusers/fastgen/test_resume_dataloader.py @@ -34,7 +34,6 @@ from __future__ import annotations -import inspect import pathlib import sys from types import SimpleNamespace @@ -54,7 +53,6 @@ _sampler_mod = pytest.importorskip("nemo_automodel.components.datasets.diffusion.sampler") _stateful_dataloader_mod = pytest.importorskip("torchdata.stateful_dataloader") dmd2_recipe = pytest.importorskip("dmd2.recipe") -from fastgen_data import rebuild_stateful_dataloader SequentialBucketSampler = _sampler_mod.SequentialBucketSampler StatefulDataLoader = _stateful_dataloader_mod.StatefulDataLoader @@ -84,7 +82,7 @@ def load_state_dict(self, state_dict): super().load_state_dict(state_dict) -def _build(n, sampler_cls, loader_cls, **loader_kwargs): +def _build(n, sampler_cls, loader_cls): """A real sampler + StatefulDataLoader over one shared synthetic dataset.""" ds = _Dataset(n) sampler = sampler_cls( @@ -104,7 +102,6 @@ def _build(n, sampler_cls, loader_cls, **loader_kwargs): batch_sampler=sampler, collate_fn=lambda b: b, num_workers=0, - **loader_kwargs, ) return sampler, loader @@ -194,48 +191,3 @@ def test_resume_reset_is_noop_on_fresh_start(monkeypatch): assert recipe.dataloader is loader, "fresh start must not rebuild the dataloader" assert recipe.sampler.state_dict() == {"epoch": 0, "batches_yielded": 0} assert recipe.step_scheduler.epoch == 0 - - -def test_resume_helper_preserves_public_loader_options(): - generator = pytest.importorskip("torch").Generator().manual_seed(11) - sampler, loader = _build( - _N, - _RecordingSampler, - StatefulDataLoader, - timeout=0, - worker_init_fn=None, - generator=generator, - pin_memory=False, - in_order=True, - snapshot_every_n_steps=7, - ) - scheduler = SimpleNamespace(epoch_len=_N, grad_acc_steps=1, epoch=0, dataloader=loader) - - rebuilt = rebuild_stateful_dataloader(loader, sampler, scheduler, global_step=3) - - assert rebuilt is scheduler.dataloader - for name in ( - "collate_fn", - "num_workers", - "pin_memory", - "timeout", - "worker_init_fn", - "multiprocessing_context", - "generator", - "prefetch_factor", - "persistent_workers", - "pin_memory_device", - "in_order", - "snapshot_every_n_steps", - ): - assert getattr(rebuilt, name) is getattr(loader, name) or getattr(rebuilt, name) == getattr( - loader, name - ) - - -def test_recipe_resume_path_has_no_private_loader_or_sampler_access(): - source = inspect.getsource(dmd2_recipe.DMD2DiffusionRecipe._rebuild_dataloader_for_resume) - loop_source = inspect.getsource(dmd2_recipe.DMD2DiffusionRecipe.run_train_validation_loop) - - assert "_batches_to_skip" not in source + loop_source - assert '__dict__["dataloader"]' not in source diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index 890bd4bc5f2..070a8358e7a 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -199,33 +199,6 @@ def test_collate_emits_contract_keys_and_broadcasts_negative_prompt(): ) # broadcast [seq,dim]->[B,seq,dim] -def test_prompt_only_collate_omits_image_latents_and_keeps_cfg_conditioning(): - pytest.importorskip("nemo_automodel") - torch = pytest.importorskip("torch") - - from fastgen_data import collate_fn_text_prompts - - sample = { - "crop_resolution": torch.tensor([1024, 1024]), - "original_resolution": torch.tensor([1024, 1024]), - "crop_offset": torch.tensor([0, 0]), - "prompt": "a test prompt", - "image_path": "/unused/source.png", - "bucket_id": 0, - "aspect_ratio": 1.0, - "prompt_embeds": torch.randn(5, 16), - "prompt_embeds_mask": torch.ones(5, dtype=torch.long), - } - negative = torch.randn(5, 16) - - result = collate_fn_text_prompts([sample, sample], negative_text_embeddings=negative) - - assert "image_latents" not in result - assert result["text_embeddings"].shape == (2, 5, 16) - assert result["text_embeddings_mask"].shape == (2, 5) - assert result["negative_text_embeddings"].shape == (2, 5, 16) - - def test_collate_zero_pads_variable_length_qwen_embeddings_and_masks(): pytest.importorskip("nemo_automodel") torch = pytest.importorskip("torch") diff --git a/tests/gpu/torch/fastgen/test_pdd_toy.py b/tests/gpu/torch/fastgen/test_pdd_toy.py index fc137d2686d..10f480ad2d8 100644 --- a/tests/gpu/torch/fastgen/test_pdd_toy.py +++ b/tests/gpu/torch/fastgen/test_pdd_toy.py @@ -17,9 +17,7 @@ from __future__ import annotations -import importlib.util -import sys -from typing import TYPE_CHECKING, Any +from typing import Any import torch from torch import nn @@ -32,10 +30,6 @@ convert_to_pdd_output_projection, ) -if TYPE_CHECKING: - from pathlib import Path - -_FORBIDDEN_MODULES = ("diffusers", "fastgen", "nemo_automodel") _WIDTH = 8 _GRID_SIZE = 4 @@ -154,18 +148,10 @@ def _build( return student, teacher, projection, pipeline, adapter -def _assert_optional_frameworks_absent() -> None: - resolvable = sorted(name for name in _FORBIDDEN_MODULES if importlib.util.find_spec(name)) - assert not resolvable, f"plain PDD GPU environment resolves optional frameworks: {resolvable}" - imported = sorted(name for name in _FORBIDDEN_MODULES if name in sys.modules) - assert not imported, f"plain PDD GPU test imported optional frameworks: {imported}" - - -def test_bf16_loss_gradient_update_reload_and_fused_sample(tmp_path: Path) -> None: - _assert_optional_frameworks_absent() - assert torch.cuda.is_available(), "Task-10 BF16 gate requires a real CUDA device" +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 gate requires Ampere or newer" + 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} @@ -175,15 +161,6 @@ def test_bf16_loss_gradient_update_reload_and_fused_sample(tmp_path: Path) -> No 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) - optimizer = torch.optim.AdamW( - student.parameters(), - lr=2.0e-3, - weight_decay=0.0, - foreach=False, - fused=False, - ) - - optimizer.zero_grad(set_to_none=True) loss, metrics = pipeline.compute_loss(data, noise=noise, n=n, k=k) assert loss.dtype == torch.float32 assert torch.isfinite(loss) @@ -199,55 +176,12 @@ def test_bf16_loss_gradient_update_reload_and_fused_sample(tmp_path: Path) -> No assert all(parameter.grad is None for parameter in teacher.parameters()) assert projection.weight.grad is not None - weight_grad = projection.weight.grad.reshape(_GRID_SIZE, _WIDTH, _WIDTH) - bias_grad = projection.bias.grad.reshape(_GRID_SIZE, _WIDTH) - assert torch.count_nonzero(weight_grad[2]) > 0 - assert torch.count_nonzero(bias_grad[2]) > 0 - assert torch.count_nonzero(weight_grad[[0, 1, 3]]) == 0 - assert torch.count_nonzero(bias_grad[[0, 1, 3]]) == 0 assert student.backbone.weight.grad is not None - assert torch.count_nonzero(student.backbone.weight.grad) > 0 - - gradients = [ - parameter.grad.float().square().sum() - for parameter in student.parameters() - if parameter.grad is not None - ] - grad_norm = torch.stack(gradients).sum().sqrt() - assert torch.isfinite(grad_norm) and grad_norm > 0 - before = {name: parameter.detach().clone() for name, parameter in student.named_parameters()} - optimizer.step() - update_norm = ( - torch.stack( - [ - (parameter.detach() - before[name]).float().square().sum() - for name, parameter in student.named_parameters() - ] - ) - .sum() - .sqrt() - ) - assert torch.isfinite(update_norm) and update_norm > 0 - - checkpoint = tmp_path / "pdd_bf16_state.pt" - torch.save(student.state_dict(), checkpoint) - saved = torch.load(checkpoint, map_location=device, weights_only=True) - assert saved.keys() == student.state_dict().keys() - assert all(value.dtype == torch.bfloat16 for value in saved.values()) - - restored, _teacher, restored_projection, restored_pipeline, restored_adapter = _build(device) - incompatible = restored.load_state_dict(saved, strict=True) - assert incompatible.missing_keys == [] - assert incompatible.unexpected_keys == [] - assert restored_projection.weight.shape == projection.weight.shape - time = pipeline.time_grid(device)[n] - expected = adapter.student_all_heads(student, data.float(), time) - actual = restored_adapter.student_all_heads(restored, data.float(), time) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - sampled = restored_pipeline.sample(noise, blocks=[2, 2]) + 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 restored_adapter.fused_calls == 2 + assert adapter.fused_calls == 2 torch.cuda.synchronize(device) - _assert_optional_frameworks_absent() diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 7275d91481c..28661c0562a 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -86,19 +86,6 @@ def test_load_pdd_config_builtin_recipe(): assert config.inference_blocks == [32, 32, 32, 32] -def test_load_pdd_config_filesystem_precedes_same_named_builtin(tmp_path, monkeypatch): - """An explicit same-named filesystem recipe takes precedence over the built-in.""" - relative_path = tmp_path / "general" / "distillation" / "pdd_qwen_image.yaml" - relative_path.parent.mkdir(parents=True) - relative_path.write_text("guidance_scale: 7.0\n", encoding="utf-8") - monkeypatch.chdir(tmp_path) - - config = load_pdd_config("general/distillation/pdd_qwen_image") - - assert config.guidance_scale == 7.0 - 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 index 186c757d0b3..aff9c2e15a1 100644 --- a/tests/unit/torch/fastgen/test_pdd_config.py +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -19,14 +19,7 @@ import pytest -from modelopt.torch.fastgen import ( - PDDConfig, - SampleTimestepConfig, - fusion_coefficients, - integrate_interval_velocities, - load_pdd_config, - make_shifted_flow_grid, -) +from modelopt.torch.fastgen import PDDConfig, SampleTimestepConfig, load_pdd_config def test_default_pdd_config_is_canonical_and_lists_are_independent(): @@ -50,12 +43,6 @@ def test_default_pdd_config_is_canonical_and_lists_are_independent(): assert second.inference_blocks == [32, 32, 32, 32] -def test_pdd_public_surface_exports_stateless_math_helpers(): - assert callable(make_shifted_flow_grid) - assert callable(integrate_interval_velocities) - assert callable(fusion_coefficients) - - def test_pdd_config_accepts_supported_schedule_and_adapter_time_scale(): config = PDDConfig( inference_blocks=[64, 64], @@ -102,25 +89,15 @@ def test_pdd_config_rejects_non_float_grid_max_t_before_coercion(value): assert config.grid_max_t == 0.999 -def test_pdd_config_accepts_explicit_grid_max_t_upper_boundary(): - config = PDDConfig(grid_max_t=1.0) - assert config.grid_max_t == 1.0 - - @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": -0.1}, "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"), - ({"grid_max_t": float("inf")}, "0 < grid_max_t <= 1"), - ({"grid_max_t": float("-inf")}, "0 < grid_max_t <= 1"), ({"flow_shift": 0.5}, "flow_shift must be finite and >= 1"), - ({"flow_shift": float("inf")}, "flow_shift must be finite and >= 1"), ({"block_size_min": 0}, "0 < block_size_min"), - ({"block_size_min": 65}, "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"), @@ -160,14 +137,6 @@ def test_pdd_config_rejects_nondefault_sample_timestep_config(): with pytest.raises(ValueError, match="sample_t_cfg is unused by PDD"): PDDConfig(sample_t_cfg=SampleTimestepConfig(shift=6.0)) - with pytest.raises(ValueError, match="sample_t_cfg is unused by PDD"): - PDDConfig(sample_t_cfg=SampleTimestepConfig(t_list=[1.0, 0.0])) - - -def test_pdd_config_accepts_explicit_default_sample_timestep_config(): - config = PDDConfig(sample_t_cfg=SampleTimestepConfig()) - assert config.sample_t_cfg == SampleTimestepConfig() - def test_pdd_config_loads_filesystem_yaml_with_optional_suffix(tmp_path): config_path = tmp_path / "pdd.yaml" diff --git a/tests/unit/torch/fastgen/test_pdd_gradient_routing.py b/tests/unit/torch/fastgen/test_pdd_gradient_routing.py deleted file mode 100644 index 36657709cc6..00000000000 --- a/tests/unit/torch/fastgen/test_pdd_gradient_routing.py +++ /dev/null @@ -1,129 +0,0 @@ -# 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. - -"""Gradient-routing contract for data-dependent PDD targets.""" - -from __future__ import annotations - -from typing import Any - -import torch -from torch import nn - -from modelopt.torch.fastgen import PDDConfig, PDDPipeline - - -class _GradientStudent(nn.Module): - def __init__(self) -> None: - super().__init__() - self.shared = nn.Parameter(torch.tensor(0.5)) - self.heads = nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(8, 2) / 11) - - -class _GradientTeacher(nn.Module): - def __init__(self) -> None: - super().__init__() - self.scale = nn.Parameter(torch.tensor(-0.25)) - - -class _GradientAdapter: - def __init__(self) -> None: - self.raw_heads: torch.Tensor | None = None - self.teacher_query: torch.Tensor | None = None - - def student_all_heads( - self, - model: _GradientStudent, - state: torch.Tensor, - time: torch.Tensor, - *, - condition: Any = None, - **model_kwargs: Any, - ) -> torch.Tensor: - del time, condition, model_kwargs - output = model.shared * state[:, None] + model.heads[None] - output.retain_grad() - self.raw_heads = output - return output - - def student_fused_block(self, *args: Any, **kwargs: Any) -> torch.Tensor: - raise AssertionError("sampling is not part of this gradient test") - - def teacher_velocity( - self, - model: _GradientTeacher, - state: torch.Tensor, - time: torch.Tensor, - *, - condition: Any = None, - negative_condition: Any = None, - **model_kwargs: Any, - ) -> torch.Tensor: - del condition, negative_condition, model_kwargs - self.teacher_query = state - return model.scale * state + time[:, None] - - -def test_only_selected_head_and_shared_backbone_receive_gradients() -> None: - student = _GradientStudent() - teacher = _GradientTeacher() - adapter = _GradientAdapter() - config = 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], - student_sample_steps=2, - ) - pipeline = PDDPipeline(student, teacher, config, adapter) - data = torch.tensor([[1.0, -0.5]]) - noise = torch.tensor([[-0.25, 2.0]]) - original_heads = student.heads.detach().clone() - optimizer = torch.optim.SGD(student.parameters(), lr=0.1) - - loss, metrics = pipeline.compute_loss( - data, - noise=noise, - n=torch.tensor([0]), - k=torch.tensor([3]), - ) - loss.backward() - - assert student.shared.grad is not None - assert not torch.equal(student.shared.grad, torch.zeros_like(student.shared.grad)) - assert torch.all(torch.isfinite(student.shared.grad)) - assert student.heads.grad is not None - assert torch.count_nonzero(student.heads.grad[3]) > 0 - assert torch.count_nonzero(student.heads.grad[:3]) == 0 - assert torch.count_nonzero(student.heads.grad[4:]) == 0 - assert torch.all(torch.isfinite(student.heads.grad)) - assert adapter.raw_heads is not None - assert adapter.raw_heads.grad is not None - assert torch.count_nonzero(adapter.raw_heads.grad[:, 3]) > 0 - assert torch.count_nonzero(adapter.raw_heads.grad[:, :3]) == 0 - assert torch.count_nonzero(adapter.raw_heads.grad[:, 4:]) == 0 - assert adapter.teacher_query is not None - assert adapter.teacher_query.requires_grad is False - assert teacher.scale.requires_grad is False - assert teacher.scale.grad is None - assert loss.requires_grad is True - assert all(not value.requires_grad for value in metrics.values()) - - optimizer.step() - assert torch.equal(student.heads[:3], original_heads[:3]) - assert not torch.equal(student.heads[3], original_heads[3]) - assert torch.equal(student.heads[4:], original_heads[4:]) diff --git a/tests/unit/torch/fastgen/test_pdd_pipeline.py b/tests/unit/torch/fastgen/test_pdd_pipeline.py index 24013526f7c..03d9af3dc89 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -194,48 +194,6 @@ def _reference_rf_forward_process( return (data_64 * (1.0 - time_64) + noise_64 * time_64).to(torch.float32) -def test_student_input_matches_fastgen_float64_forward_process() -> None: - pipeline, adapter = _pipeline() - data = torch.tensor( - [[-0.2654421329498291, 0.5161616802215576, -0.7285917401313782]], - dtype=torch.float32, - ) - noise = torch.tensor( - [[0.3856363296508789, -0.34849217534065247, -0.11881951987743378]], - dtype=torch.float32, - ) - n = torch.tensor([0]) - - pipeline.compute_loss(data, noise=noise, n=n, k=torch.tensor([0])) - - time = pipeline.time_grid()[n] - expected = _reference_rf_forward_process(data, noise, time) - stale_direct = (1.0 - time[:, None]) * data + time[:, None] * noise - assert torch.equal( - expected, - torch.tensor([[0.3849852681159973, -0.34762752056121826, -0.11942928284406662]]), - ) - assert not torch.equal(stale_direct, expected) - assert torch.equal(adapter.student_calls[0]["state"], expected) - assert torch.equal(adapter.student_calls[0]["time"], time) - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) -def test_student_input_normalizes_supported_floating_dtypes_to_float32(dtype) -> None: - pipeline, adapter = _pipeline() - data = torch.tensor([[0.75, -1.25, 0.125]], dtype=dtype) - noise = torch.tensor([[-0.5, 1.5, 2.25]], dtype=dtype) - n = torch.tensor([2]) - - pipeline.compute_loss(data, noise=noise, n=n, k=torch.tensor([3])) - - time = pipeline.time_grid()[n] - expected = _reference_rf_forward_process(data, noise, time) - assert adapter.student_calls[0]["state"].dtype == torch.float32 - assert torch.equal(adapter.student_calls[0]["state"], expected) - assert torch.equal(adapter.student_calls[0]["time"], time) - - 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]]) @@ -369,95 +327,16 @@ def test_selected_head_low_precision_outputs_use_float32_mse() -> None: torch.testing.assert_close(loss, expected) -def test_loss_can_skip_optional_diagnostics() -> None: - pipeline, _ = _pipeline() - data = torch.ones(2, 3) - - _, metrics = pipeline.compute_loss( - data, - noise=torch.zeros_like(data), - n=torch.tensor([0, 2]), - k=torch.tensor([1, 3]), - collect_metrics=False, - ) - - assert set(metrics) == {"student_target_mse"} - - -def test_small_grid_accepts_exactly_the_trained_index_support() -> None: - pipeline, _ = _pipeline() - data = torch.ones(1, 3) - noise = torch.zeros_like(data) - expected = { - (0, 0), - (0, 1), - (0, 2), - (0, 3), - (2, 2), - (2, 3), - (2, 4), - (2, 5), - (4, 4), - (4, 5), - (4, 6), - (4, 7), - (6, 6), - (6, 7), - } - accepted = set() - - for n_value in range(-1, 9): - for k_value in range(-1, 9): - pair = (n_value, k_value) - if pair not in expected: - with pytest.raises(RuntimeError): - pipeline.compute_loss( - data, - noise=noise, - n=torch.tensor([n_value]), - k=torch.tensor([k_value]), - ) - continue - pipeline.compute_loss( - data, - noise=noise, - n=torch.tensor([n_value]), - k=torch.tensor([k_value]), - ) - accepted.add(pair) - - assert accepted == expected - - -def test_explicit_k_requires_explicit_n_but_explicit_n_can_sample_k() -> None: - pipeline, _ = _pipeline() - data = torch.ones(16, 3) - noise = torch.zeros_like(data) - - with pytest.raises(ValueError, match="explicit k requires explicit n"): - pipeline.compute_loss(data, noise=noise, k=torch.zeros(16, dtype=torch.long)) - - _, metrics = pipeline.compute_loss( - data, - noise=noise, - n=torch.full((16,), 6, dtype=torch.long), - generator=torch.Generator().manual_seed(7), - ) - assert torch.equal(metrics["n"], torch.full((16,), 6, dtype=torch.long)) - assert set(metrics["k"].tolist()) == {6, 7} - - def test_sampled_indices_stay_on_exact_uniform_support() -> None: pipeline, _ = _pipeline() generator = torch.Generator().manual_seed(1234) - n, k = pipeline._resolve_indices( - batch_size=4096, - device=torch.device("cpu"), - n=None, - k=None, + _, 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) @@ -600,3 +479,25 @@ def test_pipeline_freezes_teacher_but_not_student() -> None: 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 index 8884f8d079b..270a4fc443a 100644 --- a/tests/unit/torch/fastgen/test_pdd_projection.py +++ b/tests/unit/torch/fastgen/test_pdd_projection.py @@ -17,10 +17,6 @@ from __future__ import annotations -import copy -import threading -from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait - import pytest import torch import torch.nn.functional as F @@ -204,66 +200,6 @@ def test_fused_forward_matches_independent_weighted_head_sum(layout, bias): assert torch.equal(grid, original_grid) -def test_fusion_contexts_nest_and_restore_in_order(): - 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]) - normal = projection(inputs) - - with projection.fuse_block(0, 2, grid): - outer_before = projection(inputs) - with projection.fuse_block(1, 3, grid): - inner = projection(inputs) - outer_after = projection(inputs) - - torch.testing.assert_close(outer_before, outer_after) - assert outer_before.shape == inner.shape == (1, 6) - assert projection(inputs).shape == normal.shape == (1, 18) - torch.testing.assert_close(projection(inputs), normal) - - -def test_active_fusion_rejects_forward_from_another_thread(): - 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]) - - with projection.fuse_block(0, 2, grid), ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(projection, inputs) - with pytest.raises(RuntimeError, match="non-owning thread"): - future.result() - - -def test_simultaneous_fusion_entry_admits_exactly_one_thread(): - projection = PDDOutputProjection.from_linear(_base_linear(), 3, _spec("channel_major")) - grid = torch.tensor([1.0, 0.7, 0.2, 0.0]) - start_barrier = threading.Barrier(3) - release_owner = threading.Event() - - def _enter(start, end): - start_barrier.wait() - try: - with projection.fuse_block(start, end, grid): - release_owner.wait(timeout=5) - return "admitted" - except RuntimeError as error: - return f"rejected: {error}" - - with ThreadPoolExecutor(max_workers=2) as executor: - futures = [executor.submit(_enter, 0, 2), executor.submit(_enter, 1, 3)] - start_barrier.wait() - done, _ = wait(futures, timeout=5, return_when=FIRST_COMPLETED) - release_owner.set() - results = [future.result(timeout=5) for future in futures] - - assert len(done) == 1 - assert results.count("admitted") == 1 - rejected = [result for result in results if result != "admitted"] - assert len(rejected) == 1 - assert "another thread" in rejected[0] - assert projection._fusion_stack == [] - assert projection._fusion_owner_thread is None - - def test_fusion_context_exception_cleans_up_and_allows_reuse(): projection = PDDOutputProjection.from_linear(_base_linear(), 3, _spec("channel_major")) inputs = torch.tensor([[1.0, -0.5]]) @@ -277,17 +213,6 @@ def test_fusion_context_exception_cleans_up_and_allows_reuse(): assert projection(inputs).shape == (1, 18) -def test_projection_deepcopy_recreates_inactive_fusion_lock(): - projection = PDDOutputProjection.from_linear(_base_linear(), 3, _spec("channel_major")) - copied = copy.deepcopy(projection) - inputs = torch.tensor([[1.0, -0.5]]) - grid = torch.tensor([1.0, 0.7, 0.2, 0.0]) - - with copied.fuse_block(0, 2, grid): - assert copied(inputs).shape == (1, 6) - assert copied._fusion_lock is not projection._fusion_lock - - @pytest.mark.parametrize( ("start", "end", "message"), [(-1, 2, "0 <= start"), (1, 1, "0 <= start"), (1, 4, "0 <= start")], diff --git a/tests/unit/torch/fastgen/test_pdd_public_api.py b/tests/unit/torch/fastgen/test_pdd_public_api.py deleted file mode 100644 index e871098fc0d..00000000000 --- a/tests/unit/torch/fastgen/test_pdd_public_api.py +++ /dev/null @@ -1,160 +0,0 @@ -# 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-surface and optional-dependency isolation checks for core PDD.""" - -from __future__ import annotations - -import ast -import subprocess -import sys -from pathlib import Path - -import modelopt.torch.fastgen as fastgen -from modelopt.torch.fastgen.config import PDDConfig -from modelopt.torch.fastgen.flow_matching import ( - fusion_coefficients, - integrate_interval_velocities, - make_shifted_flow_grid, -) -from modelopt.torch.fastgen.loader import load_pdd_config -from modelopt.torch.fastgen.methods.pdd import ( - PDDLayerSpec, - PDDModelAdapter, - PDDOutputProjection, - PDDPipeline, - convert_to_pdd_output_projection, -) - -_CORE_SOURCES = ( - "modelopt/torch/fastgen/config.py", - "modelopt/torch/fastgen/flow_matching.py", - "modelopt/torch/fastgen/loader.py", - "modelopt/torch/fastgen/methods/pdd.py", -) -_FORBIDDEN_IMPORT_ROOTS = {"diffusers", "fastgen", "nemo_automodel", "transformers"} -_FORBIDDEN_RELATIVE_COMPONENTS = {"plugins", "qwen_image", "qwen_image_pdd"} - - -def test_core_pdd_sources_do_not_import_model_or_framework_packages() -> None: - repository = Path(__file__).resolve().parents[4] - violations = [] - - for relative_path in _CORE_SOURCES: - source_path = repository / relative_path - tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) - for node in ast.walk(tree): - if isinstance(node, ast.Import): - imported = [alias.name for alias in node.names] - forbidden = [ - name for name in imported if name.split(".", 1)[0] in _FORBIDDEN_IMPORT_ROOTS - ] - elif isinstance(node, ast.ImportFrom) and node.module is not None: - imported = [node.module] - if node.level == 0: - forbidden = [ - name - for name in imported - if name.split(".", 1)[0] in _FORBIDDEN_IMPORT_ROOTS - ] - else: - forbidden = [ - name - for name in imported - if set(name.split(".")).intersection(_FORBIDDEN_RELATIVE_COMPONENTS) - ] - else: - continue - violations.extend((relative_path, name) for name in forbidden) - - assert violations == [] - - -def test_core_pdd_symbols_are_exported_from_the_public_package() -> None: - expected = { - "PDDConfig": PDDConfig, - "PDDLayerSpec": PDDLayerSpec, - "PDDModelAdapter": PDDModelAdapter, - "PDDOutputProjection": PDDOutputProjection, - "PDDPipeline": PDDPipeline, - "convert_to_pdd_output_projection": convert_to_pdd_output_projection, - "fusion_coefficients": fusion_coefficients, - "integrate_interval_velocities": integrate_interval_velocities, - "load_pdd_config": load_pdd_config, - "make_shifted_flow_grid": make_shifted_flow_grid, - } - - assert {name: getattr(fastgen, name) for name in expected} == expected - - -def test_fresh_core_import_does_not_require_external_frameworks() -> None: - repository = Path(__file__).resolve().parents[4] - script = r""" -import importlib.abc -import sys - - -class _UnavailableOptionalFrameworks(importlib.abc.MetaPathFinder): - prefixes = ("diffusers", "fastgen", "nemo_automodel", "transformers") - - def __init__(self): - self.attempts = [] - - def find_spec(self, fullname, path=None, target=None): - del path, target - if any(fullname == prefix or fullname.startswith(prefix + ".") for prefix in self.prefixes): - self.attempts.append(fullname) - raise ModuleNotFoundError(f"unavailable optional framework {fullname}", name=fullname) - return None - - -blocker = _UnavailableOptionalFrameworks() -sys.meta_path.insert(0, blocker) -import modelopt.torch - -baseline = set(sys.modules) -blocker.attempts.clear() -from modelopt.torch.fastgen import PDDConfig, PDDPipeline - -assert all(symbol is not None for symbol in (PDDConfig, PDDPipeline)) -assert blocker.attempts == [] -loaded = set(sys.modules) - baseline -for prefix in ( - "diffusers", - "fastgen", - "nemo_automodel", - "transformers", -): - assert not any(name == prefix or name.startswith(prefix + ".") for name in loaded), ( - prefix, - sorted(name for name in loaded if name == prefix or name.startswith(prefix + ".")), - ) -""" - result = subprocess.run( - [sys.executable, "-c", script], - cwd=repository, - check=False, - capture_output=True, - close_fds=True, - start_new_session=True, - stdin=subprocess.DEVNULL, - text=True, - ) - - assert result.returncode == 0, result.stdout + result.stderr - - -def test_optional_plugins_remain_available_through_explicit_public_access() -> None: - assert fastgen.plugins.__name__ == "modelopt.torch.fastgen.plugins" diff --git a/tests/unit/torch/fastgen/test_pdd_reference_math.py b/tests/unit/torch/fastgen/test_pdd_reference_math.py index 449ce471ebb..de91bd14e97 100644 --- a/tests/unit/torch/fastgen/test_pdd_reference_math.py +++ b/tests/unit/torch/fastgen/test_pdd_reference_math.py @@ -23,7 +23,6 @@ import pytest import torch -import torch.nn.functional as F from modelopt.torch.fastgen.flow_matching import ( add_noise, @@ -71,124 +70,6 @@ def _reference_integrate( return result -def _reference_fused_parameters( - weight: torch.Tensor, - bias: torch.Tensor | None, - grid: torch.Tensor, - start: int, - end: int, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Fuse per-interval linear parameters for the block ``[start, end)``.""" - denominator = grid[end] - grid[start] - fused_weight = torch.zeros_like(weight[0], dtype=torch.float64) - fused_bias = None if bias is None else torch.zeros_like(bias[0], dtype=torch.float64) - - for index in range(start, end): - coefficient = (grid[index + 1] - grid[index]) / denominator - fused_weight += coefficient * weight[index].to(torch.float64) - if fused_bias is not None: - fused_bias += coefficient * bias[index].to(torch.float64) - - return fused_weight, fused_bias - - -def test_shifted_grid_matches_hand_calculated_values_and_preserves_float32_intervals(): - grid = _reference_shifted_grid(grid_size=4, shift=5.0, max_t=1.0) - - expected = torch.tensor([1.0, 0.9375, 5.0 / 6.0, 0.625, 0.0], dtype=torch.float64) - torch.testing.assert_close(grid, expected, rtol=0.0, atol=0.0) - - canonical_grid = _reference_shifted_grid(grid_size=128, shift=5.0, max_t=0.999) - assert torch.all(torch.diff(canonical_grid.to(torch.float32)) < 0) - assert canonical_grid.to(torch.bfloat16)[1] == 1.0 - - for start in range(0, 128, 32): - assert canonical_grid[start + 32] - canonical_grid[start] != 0 - - -def test_half_open_integration_uses_only_selected_interval_heads(): - grid = _reference_shifted_grid(grid_size=4, shift=5.0, max_t=1.0) - state = torch.tensor([3.0, -2.0]) - velocities = torch.tensor( - [ - [1000.0, 1000.0], - [2.0, -1.0], - [-3.0, 4.0], - [-1000.0, -1000.0], - ] - ) - - result = _reference_integrate(state, velocities, grid, start=1, end=3) - expected = ( - state.to(torch.float64) - + (grid[2] - grid[1]) * velocities[1].to(torch.float64) - + (grid[3] - grid[2]) * velocities[2].to(torch.float64) - ) - - torch.testing.assert_close(result, expected, rtol=0.0, atol=1e-15) - torch.testing.assert_close( - _reference_integrate(state, velocities, grid, start=2, end=2), - state.to(torch.float64), - rtol=0.0, - atol=0.0, - ) - - -def test_final_half_open_block_advances_exactly_four_intervals(): - grid = _reference_shifted_grid(grid_size=8, shift=5.0, max_t=1.0) - state = torch.tensor([1.25]) - velocities = torch.ones(8, 1) - - result = _reference_integrate(state, velocities, grid, start=4, end=8) - expected = state.to(torch.float64) + grid[8] - grid[4] - - torch.testing.assert_close(result, expected, rtol=0.0, atol=1e-15) - - -def test_fused_projection_matches_weighted_sum_and_explicit_block_update(): - grid = _reference_shifted_grid(grid_size=4, shift=5.0, max_t=1.0) - inputs = torch.tensor([[2.0, -1.0], [-0.5, 3.0]], dtype=torch.float64) - weight = torch.tensor( - [ - [[1.0, 0.0], [0.0, 1.0]], - [[0.0, 2.0], [1.0, -1.0]], - [[-1.0, 1.0], [2.0, 0.5]], - [[3.0, -2.0], [-0.5, 1.5]], - ], - dtype=torch.float64, - ) - bias = torch.tensor( - [[0.0, 0.5], [1.0, -1.0], [2.0, 0.25], [-1.0, 3.0]], - dtype=torch.float64, - ) - original = (inputs.clone(), weight.clone(), bias.clone(), grid.clone()) - start, end = 1, 4 - - fused_weight, fused_bias = _reference_fused_parameters(weight, bias, grid, start, end) - fused_output = F.linear(inputs, fused_weight, fused_bias) - - head_outputs = torch.stack( - [F.linear(inputs, weight[index], bias[index]) for index in range(weight.shape[0])] - ) - coefficients = torch.tensor([1.0 / 9.0, 2.0 / 9.0, 2.0 / 3.0], dtype=torch.float64) - explicit_output = torch.einsum("i,ibo->bo", coefficients, head_outputs[start:end]) - explicit_update = ( - (-5.0 / 48.0) * head_outputs[1] - + (-5.0 / 24.0) * head_outputs[2] - + (-5.0 / 8.0) * head_outputs[3] - ) - - torch.testing.assert_close(fused_output, explicit_output, rtol=1e-14, atol=1e-14) - torch.testing.assert_close( - (grid[end] - grid[start]) * fused_output, - explicit_update, - rtol=1e-14, - atol=1e-14, - ) - for value, unchanged in zip((inputs, weight, bias, grid), original): - assert torch.equal(value, unchanged) - - 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) @@ -199,6 +80,9 @@ def test_production_shifted_grid_matches_independent_oracle(): 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) @@ -242,50 +126,19 @@ def test_production_rf_forward_process_matches_float64_intermediate_oracle(): assert torch.equal(value, unchanged) -def test_production_grid_promotes_low_precision_requests(): - grid = make_shifted_flow_grid(128, 5.0, max_t=0.999, dtype=torch.bfloat16) - - assert grid.dtype == torch.float32 - assert torch.all(torch.diff(grid) < 0) - - @pytest.mark.parametrize( - ("grid_size", "shift", "message"), + ("args", "kwargs", "error"), [ - (0, 5.0, "positive integer"), - (128, 0.5, "finite and >= 1"), - (128, float("nan"), "finite and >= 1"), + ((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_boundaries(grid_size, shift, message): - with pytest.raises(ValueError, match=message): - make_shifted_flow_grid(grid_size, shift, max_t=0.999) - - -@pytest.mark.parametrize("max_t", [True, 1, 0]) -def test_production_grid_rejects_non_float_max_t(max_t): - with pytest.raises(TypeError, match="max_t must be a float"): - make_shifted_flow_grid(4, 5.0, max_t=max_t) - - -@pytest.mark.parametrize("max_t", [float("nan"), float("inf"), float("-inf"), 0.0, -0.1, 1.0001]) -def test_production_grid_rejects_invalid_max_t(max_t): - with pytest.raises(ValueError, match="0 < max_t <= 1"): - make_shifted_flow_grid(4, 5.0, max_t=max_t) - - -def test_production_grid_requires_explicit_max_t_and_accepts_one(): - with pytest.raises(TypeError, match="max_t"): - make_shifted_flow_grid(4, 5.0) - - grid = make_shifted_flow_grid(4, 5.0, max_t=1.0) - assert grid[0] == 1.0 - assert grid[-1] == 0.0 - - -def test_production_grid_rejects_non_floating_dtype(): - with pytest.raises(TypeError, match="floating-point dtype"): - make_shifted_flow_grid(128, 5.0, max_t=0.999, dtype=torch.int64) +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(): @@ -339,13 +192,15 @@ def test_production_integration_promotes_bfloat16_math_to_float32(): torch.testing.assert_close(result, torch.full((1, 2), -0.999)) -def test_production_integration_rejects_out_of_range_half_open_block(): +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(): @@ -362,29 +217,3 @@ def test_production_fusion_coefficients_match_independent_oracle(): 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) - - -def test_production_fusion_coefficients_reject_empty_block(): - grid = make_shifted_flow_grid(4, 5.0, max_t=0.999) - with pytest.raises(ValueError, match="0 <= start < end <= 4"): - fusion_coefficients(grid, start=2, end=2) - - -def test_production_helpers_do_not_extract_meta_tensor_scalars(): - grid = make_shifted_flow_grid(4, 5.0, max_t=0.999, device="meta") - state = torch.empty(2, 3, device="meta") - velocities = torch.empty(2, 4, 3, device="meta") - - integrated = integrate_interval_velocities( - state, - velocities, - grid, - start=torch.tensor([0, 2], device="meta"), - end=torch.tensor([2, 4], device="meta"), - ) - coefficients = fusion_coefficients(grid, start=1, end=4) - - assert integrated.device.type == "meta" - assert integrated.shape == state.shape - assert coefficients.device.type == "meta" - assert coefficients.shape == (3,) diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 30c52c72edc..a1dd2858dee 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -26,13 +26,12 @@ 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 import PDDConfig, 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, - QWEN_IMAGE_PDD_LAYER_SPEC, convert_qwen_image_to_pdd, enable_qwen_image_pdd_forward, freeze_qwen_image_pdd_unused_parameters, @@ -134,23 +133,6 @@ def _inputs(batch_size: int = 2): return state, time, (embeddings, mask), (negative_embeddings, negative_mask) -def _call_base_packed( - model: _TinyQwenTransformer, - state: torch.Tensor, - time: torch.Tensor, - condition: tuple[torch.Tensor, torch.Tensor], -) -> torch.Tensor: - embeddings, mask = condition - return model( - hidden_states=pack_latents(state).to(torch.bfloat16), - timestep=time, - encoder_hidden_states=embeddings.to(torch.bfloat16), - encoder_hidden_states_mask=mask, - img_shapes=build_img_shapes(state.shape[0], state.shape[2], state.shape[3]), - max_txt_seq_len=int(mask.sum(dim=1).max().item()), - )[0] - - def _pack_oracle(latents: torch.Tensor) -> torch.Tensor: batch, channels, height, width = latents.shape return ( @@ -183,80 +165,6 @@ def _mr210_rollout_oracle( return state.float() + torch.einsum("bn,bn...->b...", weighted_intervals, heads.float()) -def _tiny_qwen_oracle( - model: _TinyQwenTransformer, - state: torch.Tensor, - time: torch.Tensor, - condition: tuple[torch.Tensor, torch.Tensor], -) -> torch.Tensor: - embeddings, mask = condition - packed = _pack_oracle(state).to(torch.bfloat16) - hidden = torch.tanh(model.backbone(packed)) - condition_value = embeddings.mean(dim=(1, 2), keepdim=True) - condition_value = condition_value + 0.01 * mask.sum(dim=1, keepdim=True).unsqueeze(-1) - hidden = hidden + condition_value.to(hidden.dtype) - hidden = hidden + (0.1 * time[:, None, None]).to(hidden.dtype) - return F.linear(hidden, model.proj_out.weight, model.proj_out.bias) - - -def test_conversion_is_idempotent_and_every_initialized_head_matches_base() -> None: - base = _TinyQwenTransformer() - student = copy.deepcopy(base) - state, time, condition, _ = _inputs() - base_packed = _call_base_packed(base, state, time, condition) - base_velocity = unpack_latents(base_packed, 4, 4) - config = _config() - - projection = convert_qwen_image_to_pdd(student, config) - repeated = convert_qwen_image_to_pdd(student, config) - adapter = QwenImagePDDAdapter(config) - actual = adapter.student_all_heads(student, state, time, condition=condition) - - assert projection is repeated - assert student.proj_out is projection - assert isinstance(projection, PDDOutputProjection) - assert projection.layer_spec == QWEN_IMAGE_PDD_LAYER_SPEC - assert projection.layer_spec.projection_path == "transformer.proj_out" - assert actual.shape == (2, 4, 1, 4, 4) - torch.testing.assert_close(actual, base_velocity[:, None].expand_as(actual)) - assert len(student.calls) == 1 - torch.testing.assert_close(student.calls[0]["timestep"], time) - assert student.calls[0]["img_shapes"] == [[(1, 2, 2)], [(1, 2, 2)]] - assert student.calls[0]["max_txt_seq_len"] == 3 - assert "txt_seq_lens" not in student.calls[0]["kwargs"] - - -def test_unfused_channel_major_output_maps_each_packed_head_in_order() -> None: - student = _TinyQwenTransformer() - config = _config() - projection = convert_qwen_image_to_pdd(student, config) - with torch.no_grad(): - projection.weight.zero_() - head_bias = (torch.arange(16, dtype=torch.float32).reshape(4, 4) / 5).to(torch.bfloat16) - projection.bias.copy_(head_bias.reshape(-1)) - state, time, condition, _ = _inputs(batch_size=1) - - actual = QwenImagePDDAdapter(config).student_all_heads( - student, - state, - time, - condition=condition, - ) - expected = torch.stack( - [ - unpack_latents( - head_bias[index].reshape(1, 1, 4).expand(1, 4, 4), - 4, - 4, - ) - for index in range(4) - ], - dim=1, - ) - - torch.testing.assert_close(actual, expected) - - def test_all_head_training_accepts_a_serialized_widened_linear() -> None: student = _TinyQwenTransformer() config = _config() @@ -322,40 +230,6 @@ def test_fused_student_matches_explicit_packed_weight_fusion() -> None: assert student.proj_out(projection.weight.new_zeros(1, 5)).shape[-1] == 16 -def test_teacher_cfg_uses_qwen_per_token_fp32_norm_rescale() -> None: - teacher = _TinyQwenTransformer() - config = _config(guidance_scale=4.0) - adapter = QwenImagePDDAdapter(config) - state, time, condition, negative_condition = _inputs() - - actual = adapter.teacher_velocity( - teacher, - state, - time, - condition=condition, - negative_condition=negative_condition, - ) - - assert len(teacher.calls) == 2 - conditional = teacher.calls[0]["output"] - unconditional = teacher.calls[1]["output"] - 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(conditional.dtype), 4, 4) - - assert actual.dtype == torch.bfloat16 - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(teacher.calls[0]["encoder_hidden_states"], condition[0]) - torch.testing.assert_close(teacher.calls[1]["encoder_hidden_states"], negative_condition[0]) - assert all("txt_seq_lens" not in call["kwargs"] for call in teacher.calls) - - def test_teacher_cfg_remote_preflight_failure_stops_both_model_calls(monkeypatch) -> None: teacher = _TinyQwenTransformer() adapter = QwenImagePDDAdapter(_config(guidance_scale=4.0)) @@ -383,69 +257,6 @@ def report_remote_failure(failed, *, op): assert teacher.calls == [] -def test_teacher_cfg_local_missing_negative_condition_fails_collectively(monkeypatch) -> None: - teacher = _TinyQwenTransformer() - adapter = QwenImagePDDAdapter(_config(guidance_scale=4.0)) - state, time, condition, _ = _inputs() - - monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_available", lambda: True) - monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_initialized", lambda: True) - - def preserve_local_failure(failed, *, op): - assert bool(failed) - assert op is torch.distributed.ReduceOp.MAX - - monkeypatch.setattr(qwen_image_pdd_plugin.dist, "all_reduce", preserve_local_failure) - - with pytest.raises(TypeError, match="negative_condition must be a tuple"): - adapter.teacher_velocity( - teacher, - state, - time, - condition=condition, - negative_condition=None, - ) - - assert teacher.calls == [] - - -def test_teacher_cfg_stays_in_model_output_dtype() -> None: - class LowPrecisionTeacher(_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.outputs: list[torch.Tensor] = [] - - def forward(self, *, hidden_states, encoder_hidden_states, **kwargs): - value = encoder_hidden_states.mean(dim=(1, 2), keepdim=True).to(torch.bfloat16) - output = hidden_states.to(torch.bfloat16) + value - self.outputs.append(output.detach().clone()) - return (output,) - - teacher = LowPrecisionTeacher() - state, time, condition, negative_condition = _inputs() - actual = QwenImagePDDAdapter(_config(guidance_scale=4.0)).teacher_velocity( - teacher, - state, - time, - condition=condition, - negative_condition=negative_condition, - ) - - assert actual.dtype == torch.bfloat16 - conditional, unconditional = teacher.outputs - 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) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - def test_teacher_cfg_zero_guided_norm_uses_qwen_clamp() -> None: class ZeroGuidedTeacher(_QwenImageTestDouble): def __init__(self) -> None: @@ -474,126 +285,6 @@ def forward(self, *, hidden_states, **_kwargs): torch.testing.assert_close(actual, torch.zeros_like(actual), rtol=0, atol=0) -def test_mr210_qwen_loss_and_backward_match_independent_equations() -> None: - class CapturingAdapter(QwenImagePDDAdapter): - 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 = _TinyQwenTransformer() - actual_student = copy.deepcopy(base) - actual_teacher = copy.deepcopy(base) - oracle_student = copy.deepcopy(base) - oracle_teacher = copy.deepcopy(base) - 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) - pipeline = PDDPipeline(actual_student, actual_teacher, config, adapter) - - generator = torch.Generator().manual_seed(47) - data = torch.randn(1, 1, 4, 4, generator=generator) - noise = torch.randn(1, 1, 4, 4, generator=generator) - condition = ( - torch.randn(1, 3, 2, generator=generator).to(torch.bfloat16), - torch.tensor([[1, 1, 1]], dtype=torch.long), - ) - negative_condition = ( - torch.randn(1, 2, 2, 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, _metrics = 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() - - packed_heads = _tiny_qwen_oracle(oracle_student, x_n, time_n, condition) - batch, patches, _features = packed_heads.shape - packed_heads = packed_heads.reshape(batch, patches, 4, 4).permute(0, 2, 1, 3) - oracle_heads = _unpack_oracle(packed_heads.reshape(4, patches, 4), 4, 4).reshape(1, 4, 1, 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_packed = _tiny_qwen_oracle(oracle_teacher, x_bar_k, time_k, condition) - unconditional_packed = _tiny_qwen_oracle( - oracle_teacher, - x_bar_k, - time_k, - negative_condition, - ) - guided_low_precision = conditional_packed + 3.0 * (conditional_packed - unconditional_packed) - conditional_fp32 = conditional_packed.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_state, - x_bar_k, - rtol=1e-6, - atol=1e-7, - ) - torch.testing.assert_close( - adapter.captured_teacher, - teacher_target_low_precision, - rtol=0, - atol=0, - ) - torch.testing.assert_close(actual_loss, oracle_loss, rtol=1e-6, atol=1e-7) - for name in ("backbone.weight", "proj_out.weight"): - actual_gradient = dict(actual_student.named_parameters())[name].grad - oracle_gradient = dict(oracle_student.named_parameters())[name].grad - torch.testing.assert_close( - actual_gradient, - oracle_gradient, - rtol=1e-6, - atol=1e-7, - ) - - def test_guidance_disabled_teacher_is_one_conditional_call_without_negative_condition() -> None: teacher = _TinyQwenTransformer() config = _config(guidance_scale=None) @@ -612,21 +303,6 @@ def test_guidance_disabled_teacher_is_one_conditional_call_without_negative_cond torch.testing.assert_close(actual, expected) -def test_conversion_preserves_requires_grad_mode_and_rejects_conflicts() -> None: - transformer = _TinyQwenTransformer() - transformer.eval() - transformer.proj_out.weight.requires_grad_(False) - transformer.proj_out.bias.requires_grad_(False) - projection = convert_qwen_image_to_pdd(transformer, _config()) - - assert transformer.training is False - assert projection.training is False - assert projection.weight.requires_grad is False - assert projection.bias.requires_grad is False - with pytest.raises(ValueError, match="incompatible"): - convert_qwen_image_to_pdd(transformer, _config(grid_size=2)) - - def _tiny_diffusers_qwen(): diffusers = pytest.importorskip("diffusers") return diffusers.QwenImageTransformer2DModel( @@ -840,84 +516,29 @@ def oracle_forward(model, state, time, current_condition): ) -def test_conversion_preserves_the_ordinary_diffusers_qwen_root() -> None: - student = _tiny_diffusers_qwen().eval() - root_type = type(student) - config = dict(student.config) - - convert_qwen_image_to_pdd(student, _config()) - - assert type(student) is root_type - assert isinstance(student.proj_out, PDDOutputProjection) - assert dict(student.config) == config - - -def test_qwen_pdd_forward_binding_accepts_a_dynamic_qwen_subclass() -> None: - student = _tiny_diffusers_qwen().eval() - student.__class__ = type("FSDPQwenImageTransformer2DModel", (type(student),), {}) - - assert enable_qwen_image_pdd_forward(student) is student - assert require_qwen_image_pdd_forward(student) == QWEN_IMAGE_PDD_EXECUTION - - 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()} - source_state_keys = tuple(source_state) - source_parameters = tuple(source.parameters()) - source_buffers = tuple(source.buffers()) - custom_attribute = object() - source.custom_attribute = custom_attribute - hook = source.register_forward_pre_hook(lambda *_args: None) adopted = enable_qwen_image_pdd_forward(source) assert adopted is source assert type(adopted) is source_type - assert tuple(adopted.state_dict()) == source_state_keys - assert all( - actual is expected - for actual, expected in zip(adopted.parameters(), source_parameters, strict=True) - ) - assert all( - actual is expected - for actual, expected in zip(adopted.buffers(), source_buffers, strict=True) - ) - assert hook.id in adopted._forward_pre_hooks - assert adopted.custom_attribute is custom_attribute 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) - with torch.no_grad(): - next(round_trip.parameters()).zero_() - round_trip.load_state_dict(source_state) - for name, value in round_trip.state_dict().items(): - torch.testing.assert_close(value, source_state[name], rtol=0, atol=0) assert round_trip.forward.__self__ is round_trip - - teacher = copy.deepcopy(adopted) - assert teacher is not adopted - assert teacher.forward.__func__ is adopted.forward.__func__ - assert teacher.forward.__self__ is teacher - assert teacher.forward.__self__ is not adopted - require_qwen_image_pdd_forward(teacher) - - tampered = copy.deepcopy(adopted) - tampered.forward = MethodType(lambda self, **_kwargs: self, tampered) - with pytest.raises(RuntimeError, match="masked joint-attention forward"): - require_qwen_image_pdd_forward(tampered) + 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) - forged = _tiny_diffusers_qwen() - forged._modelopt_qwen_image_pdd_execution = QWEN_IMAGE_PDD_EXECUTION - with pytest.raises(RuntimeError, match="masked joint-attention forward"): - qwen_image_pdd_plugin.require_qwen_image_pdd_forward(forged) - def test_mr210_qwen_conversion_preserves_every_initialized_head() -> None: base = _tiny_diffusers_qwen().eval().to(torch.bfloat16) @@ -1116,6 +737,10 @@ def direct_packed(current_condition): 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, @@ -1124,59 +749,21 @@ def direct_packed(current_condition): negative_condition=negative_condition, ) + assert not torch.equal(expected, global_expected) torch.testing.assert_close(actual, expected, rtol=0, atol=0) -def test_adapter_accepts_binary_masks_nonzero_padding_and_fp32_time() -> None: - student = _TinyQwenTransformer() - config = _config() - convert_qwen_image_to_pdd(student, config) - state, time, condition, _ = _inputs() - embeddings = condition[0].clone() - mask = torch.tensor([[1, 0, 1], [0, 0, 0]], dtype=torch.long) - embeddings[~mask.bool()] = 17 - - actual = QwenImagePDDAdapter(config).student_all_heads( - student, - state, - time, - condition=(embeddings, mask), - ) - - assert actual.shape == (2, 4, 1, 4, 4) - torch.testing.assert_close(student.calls[0]["encoder_hidden_states_mask"], mask) - assert student.calls[0]["timestep"].dtype == torch.float32 - assert student.calls[0]["max_txt_seq_len"] == 2 - - -def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> None: +def test_qwen_pdd_rejects_unsupported_configuration_and_inputs() -> None: with pytest.raises(ValueError, match="num_train_timesteps=None"): QwenImagePDDAdapter(_config().model_copy(update={"num_train_timesteps": 1000})) - with pytest.raises(TypeError, match="compute_dtype"): - QwenImagePDDAdapter(_config(), compute_dtype=torch.long) transformer = _TinyQwenTransformer() - transformer.config.guidance_embeds = True - with pytest.raises(ValueError, match="guidance embeddings"): - convert_qwen_image_to_pdd(transformer, _config()) - - transformer.config.guidance_embeds = False 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) - with pytest.raises(TypeError, match="negative_condition must be a tuple"): - adapter.teacher_velocity( - transformer, - state, - time, - condition=condition, - negative_condition=condition[0], - ) assert transformer.calls == [] - with pytest.raises(ValueError, match="has 4 outputs; expected 16"): - adapter.student_all_heads(transformer, state, time, condition=condition) convert_qwen_image_to_pdd(transformer, config) with pytest.raises(TypeError, match="FP32 time"): @@ -1186,74 +773,9 @@ def test_qwen_pdd_rejects_unsupported_config_condition_and_call_contracts() -> N time.to(torch.bfloat16), condition=condition, ) - with pytest.raises(TypeError, match="tuple"): - adapter.student_all_heads(transformer, state, time, condition=condition[0]) - with pytest.raises(ValueError, match="requires batched embeddings"): - adapter.student_all_heads( - transformer, - state, - time, - condition=(condition[0][..., 0], condition[1]), - ) - with pytest.raises(ValueError, match="controlled keys"): - adapter.student_all_heads( - transformer, - state, - time, - condition=condition, - guidance=torch.ones(state.shape[0]), - ) - with pytest.raises(ValueError, match="zero and one"): - adapter.student_all_heads( - transformer, - state, - time, - condition=(condition[0], torch.tensor([[1, 2, 0], [1, 0, 0]])), - ) - with pytest.raises(ValueError, match="integer/bool"): - adapter.student_all_heads( - transformer, - state, - time, - condition=(condition[0], condition[1].float()), - ) 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) - - -def test_raw_head_reference_uses_independent_linear_outputs() -> None: - """Pin the widened storage order without calling adapter reshape helpers.""" - student = _TinyQwenTransformer() - config = _config() - projection = convert_qwen_image_to_pdd(student, config) - state, time, condition, _ = _inputs(batch_size=1) - embeddings, mask = condition - packed = pack_latents(state).to(torch.bfloat16) - hidden = torch.tanh(student.backbone(packed)) - condition_value = embeddings.mean(dim=(1, 2), keepdim=True) - condition_value = condition_value + 0.01 * mask.sum(dim=1, keepdim=True).unsqueeze(-1) - hidden = hidden + condition_value.to(hidden.dtype) - hidden = hidden + (0.1 * time[:, None, None]).to(hidden.dtype) - head_weights = projection.weight.reshape(4, 4, 5) - head_bias = projection.bias.reshape(4, 4) - expected_packed = torch.stack( - [F.linear(hidden, head_weights[index], head_bias[index]) for index in range(4)], - dim=1, - ) - expected = torch.stack( - [unpack_latents(expected_packed[:, index], 4, 4) for index in range(4)], - dim=1, - ) - - actual = QwenImagePDDAdapter(config).student_all_heads( - student, - state, - time, - condition=condition, - ) - - torch.testing.assert_close(actual, expected) From 6550a921679508413f19d61a209a8eff907b2ecb Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 3 Sep 2026 23:06:58 -0700 Subject: [PATCH 40/45] refactor: simplify PDD integration Signed-off-by: Meng Xin --- .../fastgen/fastgen_data/__init__.py | 31 +-- .../fastgen/fastgen_data/collate_fns.py | 9 - .../diffusers/fastgen/fastgen_data/splits.py | 52 ---- .../fastgen_data/text_to_image_dataset.py | 28 +-- examples/diffusers/fastgen/pdd/README.md | 7 +- .../fastgen/pdd/configs/qwen_image.yaml | 7 - .../fastgen/pdd/inference_qwen_image.py | 2 +- modelopt/torch/fastgen/config.py | 53 +---- modelopt/torch/fastgen/methods/pdd.py | 114 +++------ .../torch/fastgen/plugins/qwen_image_pdd.py | 224 ++++-------------- .../general/distillation/pdd_qwen_image.yaml | 7 - .../diffusers/fastgen/test_dataset_paths.py | 9 +- .../diffusers/fastgen/test_dataset_splits.py | 116 --------- .../diffusers/fastgen/test_pdd_inference.py | 4 +- .../fastgen/test_pdd_recipe_setup.py | 1 - .../fastgen/test_vendored_migration.py | 3 +- tests/gpu/torch/fastgen/test_pdd_toy.py | 13 +- tests/unit/recipe/test_loader.py | 2 +- tests/unit/torch/fastgen/test_pdd_config.py | 102 +++----- tests/unit/torch/fastgen/test_pdd_pipeline.py | 1 - .../unit/torch/fastgen/test_pdd_projection.py | 17 +- .../fastgen/test_qwen_image_pdd_plugin.py | 70 ++---- 22 files changed, 181 insertions(+), 691 deletions(-) delete mode 100644 examples/diffusers/fastgen/fastgen_data/splits.py delete mode 100644 tests/examples/diffusers/fastgen/test_dataset_splits.py diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index c1b1c7f6a2b..8338f44e517 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -39,16 +39,12 @@ # (``nemo_automodel.components.datasets.diffusion.{sampler,base_dataset}``). # Convert a missing-helper ImportError into an actionable message naming the supported release. try: - from . import collate_fns as _collate_fns - from . import paths as _paths - from . import resume as _resume - from . import splits as _splits - from . import text_to_image_dataset as _text_to_image_dataset - from .collate_fns import * - from .paths import * - from .resume import * - from .splits import * - from .text_to_image_dataset import * + from .collate_fns import ( + build_text_to_image_multiresolution_dataloader, + collate_fn_text_prompts, + collate_fn_text_to_image, + ) + 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 " @@ -59,15 +55,12 @@ f"Underlying import error: {exc!r}" ) from exc -__all__: list[str] = [] -for _module in ( - _collate_fns, - _paths, - _resume, - _splits, - _text_to_image_dataset, -): - __all__.extend(_module.__all__) +__all__ = [ + "build_text_to_image_multiresolution_dataloader", + "collate_fn_text_prompts", + "collate_fn_text_to_image", + "rebuild_stateful_dataloader", +] def _warn_if_unsupported_upstream() -> None: diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index d15b3e28faf..9a20f2b9bfa 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -180,9 +180,6 @@ def build_text_to_image_multiresolution_dataloader( pin_memory: bool = True, prefetch_factor: int = 2, negative_prompt_embedding_path: str | None = None, - split: str | None = None, - validation_count: int | None = None, - split_seed: int = 2026, sampler_seed: int = 42, loader_seed: int | None = None, ) -> tuple[StatefulDataLoader, SequentialBucketSampler]: @@ -205,9 +202,6 @@ def build_text_to_image_multiresolution_dataloader( 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. - split: Optional deterministic ``"train"`` or ``"validation"`` selection. - validation_count: Number of validation samples when ``split`` is set. - split_seed: Local seed used to construct deterministic split membership. sampler_seed: Seed for the released deterministic bucket sampler. loader_seed: Optional dedicated seed for DataLoader worker/base-seed generation. @@ -218,9 +212,6 @@ def build_text_to_image_multiresolution_dataloader( cache_dir=cache_dir, train_text_encoder=train_text_encoder, prompt_only=prompt_only, - split=split, - validation_count=validation_count, - split_seed=split_seed, ) effective_root = dataset.cache_root diff --git a/examples/diffusers/fastgen/fastgen_data/splits.py b/examples/diffusers/fastgen/fastgen_data/splits.py deleted file mode 100644 index 465f4ea6329..00000000000 --- a/examples/diffusers/fastgen/fastgen_data/splits.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. - -"""Deterministic train/validation membership for FastGen cache ordinals.""" - -from __future__ import annotations - -import torch - -__all__ = ["make_train_validation_indices"] - - -def _require_integer(name: str, value: int) -> int: - if type(value) is not int: - raise TypeError(f"{name} must be an integer; got {type(value).__name__}") - return value - - -def make_train_validation_indices( - num_samples: int, - validation_count: int, - seed: int, -) -> tuple[list[int], list[int]]: - """Return disjoint ordered metadata ordinals using a local CPU generator.""" - num_samples = _require_integer("num_samples", num_samples) - validation_count = _require_integer("validation_count", validation_count) - seed = _require_integer("seed", seed) - if num_samples <= 0: - raise ValueError("num_samples must be positive") - if not 1 <= validation_count < num_samples: - raise ValueError("validation_count must be in [1, num_samples)") - if seed < 0: - raise ValueError("seed must be nonnegative") - - generator = torch.Generator(device="cpu") - generator.manual_seed(seed) - permutation = torch.randperm(num_samples, generator=generator, device="cpu").tolist() - validation = sorted(permutation[:validation_count]) - train = sorted(permutation[validation_count:]) - return train, validation 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 29b04f86797..5272443bca7 100644 --- a/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py +++ b/examples/diffusers/fastgen/fastgen_data/text_to_image_dataset.py @@ -20,7 +20,6 @@ from nemo_automodel.components.datasets.diffusion.base_dataset import BaseMultiresolutionDataset from .paths import resolve_cache_root, resolve_under_root -from .splits import make_train_validation_indices __all__ = ["TextToImageDataset"] @@ -33,29 +32,16 @@ def __init__( cache_dir: str | Path, train_text_encoder: bool = False, prompt_only: bool = False, - split: str | None = None, - validation_count: int | None = None, - split_seed: int = 2026, ): """ 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. - split: Optional deterministic ``"train"`` or ``"validation"`` selection. - validation_count: Number of validation samples when ``split`` is set. - split_seed: Local seed used to construct deterministic split membership. """ - if split not in (None, "train", "validation"): - raise ValueError("split must be null, 'train', or 'validation'") - if split is not None and validation_count is None: - raise ValueError("validation_count is required when split is set") self.train_text_encoder = train_text_encoder self.prompt_only = prompt_only self.cache_root = resolve_cache_root(cache_dir) - self._split = split - self._validation_count = validation_count - self._split_seed = split_seed self._resolved_cache_files: dict[int, Path] = {} super().__init__(str(self.cache_root), quantization=64) @@ -95,18 +81,8 @@ def _load_metadata(self) -> list[dict]: if not complete_metadata: raise ValueError(f"No samples found in {metadata_file}") self.total_num_samples = len(complete_metadata) - if self._split is None: - self.sample_ids = list(range(self.total_num_samples)) - else: - if self._validation_count is None: - raise RuntimeError("validation_count was not resolved for the requested split") - train, validation = make_train_validation_indices( - self.total_num_samples, - self._validation_count, - self._split_seed, - ) - self.sample_ids = train if self._split == "train" else validation - return [complete_metadata[index] for index in self.sample_ids] + 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.""" diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index 0252cd60a66..f958c7743c7 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -27,13 +27,18 @@ ModelOpt owns only the PDD model transformation, Qwen execution adapter, loss, a 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 AutoModel, Diffusers, or Qwen source. +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: diff --git a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml index 72669dc4b51..596e49d2017 100644 --- a/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/pdd/configs/qwen_image.yaml @@ -17,11 +17,7 @@ model: latent_shape: [16, 128, 128] pdd: - pred_type: flow - num_train_timesteps: guidance_scale: 4.0 - student_sample_steps: 4 - student_sample_type: ode grid_size: 128 grid_max_t: 0.999 flow_shift: 5.0 @@ -86,9 +82,6 @@ data: drop_last: true shuffle: true dynamic_batch_size: false - split: train - validation_count: 2000 - split_seed: 2026 sampler_seed: 42 loader_seed: 42 negative_prompt_embedding_path: negative_prompt_embedding.pt diff --git a/examples/diffusers/fastgen/pdd/inference_qwen_image.py b/examples/diffusers/fastgen/pdd/inference_qwen_image.py index a12c87e54ec..e541cdc092f 100644 --- a/examples/diffusers/fastgen/pdd/inference_qwen_image.py +++ b/examples/diffusers/fastgen/pdd/inference_qwen_image.py @@ -87,7 +87,7 @@ def _parse_blocks(value: str) -> list[int]: def _load_config(path: Path, blocks: list[int]) -> PDDConfig: raw = yaml.safe_load(path.read_text()) values = dict(raw["pdd"]) - values.update(inference_blocks=blocks, student_sample_steps=len(blocks)) + values["inference_blocks"] = blocks return PDDConfig.model_validate(values) diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py index bf5f8c3a673..06d706bfbb5 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``. @@ -172,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( @@ -220,7 +218,7 @@ class DistillationConfig(ModeloptBaseConfig): ) -class PDDConfig(DistillationConfig): +class PDDConfig(ModeloptBaseConfig): """Hyperparameters for Parallel Decoding Distillation (PDD). PDD trains one velocity head per interval on a fixed shifted rectified-flow @@ -228,20 +226,10 @@ class PDDConfig(DistillationConfig): not define a second timestep schedule. """ - pred_type: Literal["flow"] = ModeloptField( - default="flow", - title="Network prediction parameterization", - description="PDD is defined for rectified-flow velocity prediction.", - ) - student_sample_type: Literal["ode"] = ModeloptField( - default="ode", - title="Student sampling mode", - description="PDD fused inference follows the fixed rectified-flow ODE grid.", - ) - student_sample_steps: int = ModeloptField( - default=4, - title="Student inference steps", - description="Number of contiguous blocks in ``inference_blocks``.", + 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, @@ -273,8 +261,8 @@ class PDDConfig(DistillationConfig): title="Teacher target integrator", description="Integrator used to estimate the teacher mean velocity for an interval.", ) - inference_blocks: list[int] = Field( - default_factory=lambda: [32, 32, 32, 32], + 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.", ) @@ -287,12 +275,7 @@ class PDDConfig(DistillationConfig): ) def __setattr__(self, name: str, value: object) -> None: - """Validate a complete candidate config before changing an initialized field. - - Pydantic's after-model validators otherwise run after assignment and leave the - rejected value stored. PDD has cross-field schedule invariants, so attribute - and mutable-mapping updates must be transactional. - """ + """Validate cross-field invariants before changing an initialized field.""" if name in type(self).model_fields and name in self.__dict__: candidate = self.model_dump() candidate[name] = value @@ -337,18 +320,6 @@ def _check_pdd(self) -> PDDConfig: f"inference_blocks must sum to grid_size={self.grid_size}, got " f"{sum(self.inference_blocks)}." ) - if self.student_sample_steps != len(self.inference_blocks): - raise ValueError( - "student_sample_steps must equal len(inference_blocks), got " - f"{self.student_sample_steps} and {len(self.inference_blocks)}." - ) - - default_sample_t_cfg = SampleTimestepConfig() - if self.sample_t_cfg.model_dump() != default_sample_t_cfg.model_dump(): - raise ValueError( - "sample_t_cfg is unused by PDD and cannot be overridden; PDD samples " - "discrete interval indices from its fixed shifted grid." - ) return self @classmethod diff --git a/modelopt/torch/fastgen/methods/pdd.py b/modelopt/torch/fastgen/methods/pdd.py index 611a63bcc93..bd173c785ae 100644 --- a/modelopt/torch/fastgen/methods/pdd.py +++ b/modelopt/torch/fastgen/methods/pdd.py @@ -23,9 +23,7 @@ from __future__ import annotations -import contextlib -import threading -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any, Literal, Protocol @@ -40,7 +38,6 @@ integrate_interval_velocities, make_shifted_flow_grid, ) -from ..pipeline import DistillationPipeline __all__ = [ "PDDLayerSpec", @@ -91,21 +88,12 @@ def __post_init__(self) -> None: _require_int(self.output_channels, name="output_channels") -@dataclass(frozen=True) -class _FusionRequest: - start: int - end: int - grid: torch.Tensor - - class PDDOutputProjection(nn.Linear): """A widened linear projection with one output head per PDD interval. - Outside :meth:`fuse_block`, ``forward`` returns the full widened output. - Inside the context, ``forward`` computes the selected block's weighted - projection in float32 and returns the base-sized output. Fusion state is - synchronous and thread-owned; nested contexts in one thread are supported, - while concurrent access from another thread is rejected. + ``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__( @@ -145,23 +133,6 @@ def __init__( self.base_out_features = base_out_features self.grid_size = grid_size self.layer_spec = layer_spec - self._fusion_stack: list[_FusionRequest] = [] - self._fusion_owner_thread: int | None = None - self._fusion_lock = threading.Lock() - - def __getstate__(self) -> dict[str, Any]: - """Exclude the process-local lock while preserving ordinary module deepcopy.""" - with self._fusion_lock: - if self._fusion_stack: - raise RuntimeError("cannot copy or serialize an active PDD fusion context.") - state = super().__getstate__() - state.pop("_fusion_lock", None) - return state - - def __setstate__(self, state: dict[str, Any]) -> None: - """Restore module state with a fresh process-local fusion lock.""" - super().__setstate__(state) - self._fusion_lock = threading.Lock() @property def patch_factor(self) -> int: @@ -262,14 +233,13 @@ def _tensor_by_head(self, tensor: torch.Tensor) -> torch.Tensor: .reshape(self.grid_size, self.base_out_features, *trailing_shape) ) - @contextlib.contextmanager - def fuse_block(self, start: int, end: int, grid: torch.Tensor) -> Iterator[PDDOutputProjection]: - """Temporarily return the fused base-sized projection for ``[start, end)``. - - Contexts may nest synchronously in one thread. Because fusion selection is - stored on the module, a second thread may neither enter a context nor call - ``forward`` until the owning context exits. - """ + 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): @@ -278,58 +248,33 @@ def fuse_block(self, start: int, end: int, grid: torch.Tensor) -> Iterator[PDDOu raise ValueError( f"grid must contain {self.grid_size + 1} nodes, got shape {tuple(grid.shape)}." ) - fusion_coefficients(grid, start, end) - - thread_id = threading.get_ident() - request = _FusionRequest(start=start, end=end, grid=grid) - with self._fusion_lock: - if self._fusion_stack and self._fusion_owner_thread != thread_id: - raise RuntimeError("PDD fusion context is already active in another thread.") - if not self._fusion_stack: - self._fusion_owner_thread = thread_id - self._fusion_stack.append(request) - try: - yield self - finally: - with self._fusion_lock: - if not self._fusion_stack or self._fusion_stack[-1] is not request: - raise RuntimeError("PDD fusion contexts exited out of order.") - self._fusion_stack.pop() - if not self._fusion_stack: - self._fusion_owner_thread = None - - def _fused_parameters( - self, request: _FusionRequest - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Compute block-fused parameters from the registered widened parameter.""" - coefficients = fusion_coefficients(request.grid, request.start, request.end).to( + coefficients = fusion_coefficients(grid, start, end).to( device=self.weight.device, dtype=torch.float32, ) - head_weights = self._tensor_by_head(self.weight)[request.start : request.end] + 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)[request.start : request.end] + 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) -> torch.Tensor: - """Apply the widened or currently scoped fused projection.""" - with self._fusion_lock: - if not self._fusion_stack: - request = None - else: - if self._fusion_owner_thread != threading.get_ident(): - raise RuntimeError("PDD fused forward was called from a non-owning thread.") - request = self._fusion_stack[-1] - if request is None: + 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) - fused_weight, fused_bias = self._fused_parameters(request) + start, end, grid = fusion + fused_weight, fused_bias = self._fused_parameters(start, end, grid) return F.linear(input, fused_weight, fused_bias) @@ -422,7 +367,7 @@ def teacher_velocity( ... -class PDDPipeline(DistillationPipeline): +class PDDPipeline: """PDD losses and fused sampler over a single core-owned grid.""" def __init__( @@ -435,12 +380,9 @@ def __init__( """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__}.") - if teacher is None: - self.student = student - self.teacher = None - self.config = config - else: - super().__init__(student, teacher, config) + 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: diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index bff9dbd402d..dc93e03770b 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -27,7 +27,6 @@ from typing import Any import torch -import torch.distributed as dist from torch import nn from ..config import PDDConfig @@ -65,19 +64,10 @@ "return_dict", "timestep", "txt_seq_lens", + "_pdd_fusion", } _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE = "_modelopt_qwen_image_pdd_execution" -_QWEN_IMAGE_PDD_CHILDREN = ( - "pos_embed", - "time_text_embed", - "txt_norm", - "img_in", - "txt_in", - "transformer_blocks", - "norm_out", - "proj_out", -) def _require_binary_mask(mask: torch.Tensor, *, name: str) -> None: @@ -114,27 +104,15 @@ def _qwen_image_pdd_forward( 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 hidden_states.ndim != 3 or hidden_states.dtype != torch.bfloat16: - raise TypeError("Qwen PDD hidden_states must be packed BF16 [B, P, C].") - if ( - not isinstance(encoder_hidden_states, torch.Tensor) - or encoder_hidden_states.ndim != 3 - or encoder_hidden_states.dtype != torch.bfloat16 - ): - raise TypeError("Qwen PDD encoder_hidden_states must be BF16 [B, S, D].") - if not isinstance(encoder_hidden_states_mask, torch.Tensor): - raise TypeError("Qwen PDD requires encoder_hidden_states_mask.") if not isinstance(timestep, torch.Tensor) or timestep.dtype != torch.float32: raise TypeError("Qwen PDD timestep must remain FP32 at transformer entry.") - batch_size = hidden_states.shape[0] - if encoder_hidden_states.shape[0] != batch_size: - raise ValueError("Qwen PDD image and text batch sizes must match.") - if timestep.shape != (batch_size,): - raise ValueError("Qwen PDD timestep must contain one value per batch item.") - if img_shapes is None or len(img_shapes) != batch_size: - raise ValueError("Qwen PDD img_shapes must contain one entry per batch item.") + 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: @@ -143,76 +121,44 @@ def _qwen_image_pdd_forward( 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.") - if type(return_dict) is not bool: - raise TypeError("Qwen PDD return_dict must be a bool.") - if encoder_hidden_states_mask.ndim != 2 or tuple(encoder_hidden_states_mask.shape) != tuple( - encoder_hidden_states.shape[:2] - ): - raise ValueError("Qwen PDD mask must match the text batch and sequence dimensions.") - if encoder_hidden_states_mask.device != encoder_hidden_states.device: - raise ValueError("Qwen PDD mask and text embeddings must share a device.") - if ( - encoder_hidden_states_mask.dtype.is_floating_point - or encoder_hidden_states_mask.dtype.is_complex - ): - raise TypeError("Qwen PDD mask must use an integer or boolean dtype.") - _require_binary_mask(encoder_hidden_states_mask, name="Qwen PDD") - expected_max_txt_seq_len = int( - encoder_hidden_states_mask.sum(dim=1).max().to(torch.int32).item() - ) - if txt_seq_lens is not None: - expected_txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).to(torch.int32).tolist() - if txt_seq_lens != expected_txt_seq_lens: - raise ValueError("Qwen PDD txt_seq_lens must equal the valid mask lengths.") - if max_txt_seq_len is None: - max_txt_seq_len = expected_max_txt_seq_len - elif max_txt_seq_len != expected_max_txt_seq_len: - raise ValueError("Qwen PDD max_txt_seq_len must equal the maximum valid mask length.") + del txt_seq_lens + 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)) - if timestep.dtype != torch.float32: - raise RuntimeError("Qwen PDD timestep was rounded before time_text_embed.") 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, ) - image_mask = torch.ones( - (batch_size, hidden_states.shape[1]), - dtype=torch.bool, - device=hidden_states.device, - ) - joint_attention_mask = torch.cat( - (encoder_hidden_states_mask.to(torch.bool), image_mask), - dim=1, - )[:, None, None, :] - block_attention_kwargs = {"attention_mask": joint_attention_mask} - 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, - None, + encoder_hidden_states_mask, temb, image_rotary_emb, - block_attention_kwargs, ) else: encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, - encoder_hidden_states_mask=None, + encoder_hidden_states_mask=encoder_hidden_states_mask, temb=temb, image_rotary_emb=image_rotary_emb, - joint_attention_kwargs=block_attention_kwargs, ) hidden_states = self.norm_out(hidden_states, temb) - output = self.proj_out(hidden_states) + 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,) @@ -284,18 +230,6 @@ def enable_qwen_image_pdd_forward(transformer: nn.Module) -> nn.Module: existing_forward = transformer.__dict__.get("forward") if existing_forward is not None: raise RuntimeError("Qwen root already has a different instance-level forward override.") - missing = [ - name - for name in _QWEN_IMAGE_PDD_CHILDREN - if not isinstance(getattr(transformer, name, None), nn.Module) - ] - if missing: - raise RuntimeError(f"Qwen root is missing required PDD forward modules: {missing}.") - if ( - not isinstance(transformer.transformer_blocks, nn.ModuleList) - or not transformer.transformer_blocks - ): - raise RuntimeError("Qwen PDD requires a nonempty transformer_blocks ModuleList.") if _config_guidance_embeds(transformer): raise ValueError("Qwen PDD does not support transformer guidance embeddings.") if getattr(transformer, "peft_config", None): @@ -315,11 +249,6 @@ def enable_qwen_image_pdd_forward(transformer: nn.Module) -> nn.Module: def _validate_qwen_pdd_config(config: PDDConfig) -> None: if not isinstance(config, PDDConfig): raise TypeError(f"config must be PDDConfig, got {type(config).__name__}.") - if config.num_train_timesteps is not None: - raise ValueError( - "Qwen-Image PDD requires num_train_timesteps=None because the adapter " - "forwards normalized continuous grid time." - ) def convert_qwen_image_to_pdd( @@ -555,9 +484,10 @@ def _call_packed( *, condition_name: str, prepared_condition: tuple[torch.Tensor, torch.Tensor] | None = None, + fusion: tuple[int, int, torch.Tensor] | None = None, ) -> torch.Tensor: if prepared_condition is None: - encoder_hidden_states, attention_mask = self._prepare_call_collectively( + encoder_hidden_states, attention_mask = self._prepare_call( model, state, time, @@ -576,64 +506,20 @@ def _call_packed( raise TypeError("Qwen PDD execution requires FP32 time.") packed_state = pack_latents(state).to(model_dtype) encoder_hidden_states = encoder_hidden_states.to(model_dtype) - max_txt_seq_len = int(attention_mask.sum(dim=1).max().to(torch.int32).item()) + 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=build_img_shapes(batch_size, height, width), - max_txt_seq_len=max_txt_seq_len, return_dict=False, - **model_kwargs, + **call_kwargs, ) return self._extract_packed_output(output) - @staticmethod - def _raise_collective_preflight_error( - local_error: Exception | None, - *, - state: torch.Tensor, - ) -> None: - if dist.is_available() and dist.is_initialized(): - failed = torch.tensor(local_error is not None, dtype=torch.int32, device=state.device) - dist.all_reduce(failed, op=dist.ReduceOp.MAX) - if bool(failed): - if local_error is not None: - raise local_error - raise RuntimeError("Qwen PDD preflight failed on another rank.") - elif local_error is not None: - raise local_error - - def _prepare_call_collectively( - self, - model: nn.Module, - state: torch.Tensor, - time: torch.Tensor, - condition: Any, - model_kwargs: Mapping[str, Any], - *, - condition_name: str, - ) -> tuple[torch.Tensor, torch.Tensor]: - prepared: tuple[torch.Tensor, torch.Tensor] | None = None - local_error: Exception | None = None - try: - prepared = self._prepare_call( - model, - state, - time, - condition, - model_kwargs, - condition_name=condition_name, - ) - except Exception as error: - local_error = error - - self._raise_collective_preflight_error(local_error, state=state) - if prepared is None: - raise RuntimeError("Qwen PDD preflight did not prepare a model call.") - return prepared - def _prepare_teacher_cfg_calls( self, model: nn.Module, @@ -646,33 +532,23 @@ def _prepare_teacher_cfg_calls( tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor], ]: - """Make rank-local CFG validation fail collectively before either teacher call.""" - prepared_condition: tuple[torch.Tensor, torch.Tensor] | None = None - prepared_negative_condition: tuple[torch.Tensor, torch.Tensor] | None = None - local_error: Exception | None = None - try: - prepared_condition = self._prepare_call( - model, - state, - time, - condition, - model_kwargs, - condition_name="condition", - ) - prepared_negative_condition = self._prepare_call( - model, - state, - time, - negative_condition, - model_kwargs, - condition_name="negative_condition", - ) - except Exception as error: - local_error = error - - self._raise_collective_preflight_error(local_error, state=state) - if prepared_condition is None or prepared_negative_condition is None: - raise RuntimeError("Qwen teacher CFG preflight did not prepare both model calls.") + """Validate both local CFG calls before either model execution begins.""" + prepared_condition = self._prepare_call( + model, + state, + time, + condition, + model_kwargs, + condition_name="condition", + ) + prepared_negative_condition = self._prepare_call( + model, + state, + time, + negative_condition, + model_kwargs, + condition_name="negative_condition", + ) return prepared_condition, prepared_negative_condition @staticmethod @@ -807,16 +683,16 @@ def student_fused_block( **model_kwargs: Any, ) -> torch.Tensor: """Run one conditional Qwen call with its final projection fused for a block.""" - projection = self._fused_projection(model, self.config.grid_size) - with projection.fuse_block(start, end, grid): - packed = self._call_packed( - model, - state, - time, - condition, - model_kwargs, - condition_name="condition", - ) + self._fused_projection(model, self.config.grid_size) + packed = self._call_packed( + model, + state, + time, + condition, + model_kwargs, + condition_name="condition", + fusion=(start, end, grid), + ) return self._unpack_single(packed, state) @torch.no_grad() diff --git a/modelopt_recipes/general/distillation/pdd_qwen_image.yaml b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml index b44de0ae14e..ecc4c2b9e4e 100644 --- a/modelopt_recipes/general/distillation/pdd_qwen_image.yaml +++ b/modelopt_recipes/general/distillation/pdd_qwen_image.yaml @@ -4,13 +4,6 @@ # model-call defaults only; data roots, training topology, checkpoint paths, and # cluster settings belong to the framework example and run manifest. -pred_type: flow -student_sample_type: ode -student_sample_steps: 4 - -# Qwen consumes the normalized continuous time supplied by the PDD grid. -num_train_timesteps: - # 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 diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index d19e8c90e88..8ddfb8d408b 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -34,12 +34,9 @@ if str(_FASTGEN_DIR) not in sys.path: sys.path.insert(0, str(_FASTGEN_DIR)) -from fastgen_data import ( - TextToImageDataset, - build_text_to_image_multiresolution_dataloader, - resolve_cache_root, - resolve_under_root, -) +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_uses_unset_or_empty_fallback(make_fastgen_cache, monkeypatch, tmp_path): diff --git a/tests/examples/diffusers/fastgen/test_dataset_splits.py b/tests/examples/diffusers/fastgen/test_dataset_splits.py deleted file mode 100644 index f458549cafd..00000000000 --- a/tests/examples/diffusers/fastgen/test_dataset_splits.py +++ /dev/null @@ -1,116 +0,0 @@ -# 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. - -"""Deterministic stable-ID split contract for the shared FastGen cache.""" - -from __future__ import annotations - -import pathlib -import sys - -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, - make_train_validation_indices, -) - - -def _snapshot(root: pathlib.Path) -> dict[str, bytes]: - return { - path.relative_to(root).as_posix(): path.read_bytes() - for path in root.rglob("*") - if path.is_file() - } - - -def test_split_has_frozen_membership_and_does_not_change_global_rng(): - torch.manual_seed(1234) - rng_before = torch.random.get_rng_state().clone() - - train, validation = make_train_validation_indices(10, validation_count=3, seed=17) - - assert validation == [0, 7, 9] - assert train == [1, 2, 3, 4, 5, 6, 8] - assert set(train).isdisjoint(validation) - assert sorted(train + validation) == list(range(10)) - assert torch.equal(torch.random.get_rng_state(), rng_before) - assert make_train_validation_indices(10, 3, 17) == (train, validation) - - -@pytest.mark.parametrize( - ("num_samples", "validation_count", "seed", "error"), - [ - (True, 1, 0, TypeError), - (1, 1, 0, ValueError), - (4, False, 0, TypeError), - (4, 0, 0, ValueError), - (4, 4, 0, ValueError), - (4, 1, True, TypeError), - (4, 1, -1, ValueError), - ], -) -def test_split_rejects_invalid_inputs(num_samples, validation_count, seed, error): - with pytest.raises(error): - make_train_validation_indices(num_samples, validation_count, seed) - - -def test_train_and_validation_loaders_are_disjoint_stable_and_read_only( - make_fastgen_cache, tmp_path -): - cache = make_fastgen_cache(tmp_path / "cache") - before = _snapshot(cache) - train_ids, validation_ids = make_train_validation_indices(6, validation_count=2, seed=17) - - train_loader, train_sampler = build_text_to_image_multiresolution_dataloader( - cache_dir=str(cache), - split="train", - validation_count=2, - split_seed=17, - batch_size=1, - num_workers=0, - shuffle=True, - drop_last=True, - ) - validation_loader, validation_sampler = build_text_to_image_multiresolution_dataloader( - cache_dir=str(cache), - split="validation", - validation_count=2, - split_seed=17, - batch_size=1, - num_workers=0, - shuffle=False, - drop_last=False, - ) - - assert train_loader.dataset.sample_ids == train_ids - assert validation_loader.dataset.sample_ids == validation_ids - assert ( - train_loader.dataset.cache_root == validation_loader.dataset.cache_root == cache.resolve() - ) - assert train_sampler.shuffle_buckets and train_sampler.shuffle_within_bucket - assert not validation_sampler.shuffle_buckets - assert not validation_sampler.shuffle_within_bucket - assert validation_sampler.drop_last is False - assert _snapshot(cache) == before diff --git a/tests/examples/diffusers/fastgen/test_pdd_inference.py b/tests/examples/diffusers/fastgen/test_pdd_inference.py index c1d4abf1fd9..cab6b8142c9 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_inference.py +++ b/tests/examples/diffusers/fastgen/test_pdd_inference.py @@ -44,7 +44,6 @@ def test_widened_diffusers_projection_restores_pdd_fusion_metadata(tmp_path) -> block_size_min=1, block_size_max=4, inference_blocks=[2, 2], - student_sample_steps=2, ) transformer = get_tiny_qwen_image_transformer(num_layers=1) base_out_channels = transformer.out_channels @@ -77,9 +76,8 @@ def test_inference_block_override_is_validated(tmp_path) -> None: " block_size_min: 1\n" " block_size_max: 8\n" " inference_blocks: [4, 4]\n" - " student_sample_steps: 2\n" ) blocks = _parse_blocks("2, 2,4") assert blocks == [2, 2, 4] - assert _load_config(config_path, blocks).inference_blocks == blocks + 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 index 968069ad4e2..d0f6bfcf99b 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -128,7 +128,6 @@ def test_prepared_student_width_is_validated_before_training() -> None: block_size_min=1, block_size_max=8, inference_blocks=[4, 4], - student_sample_steps=2, ) _validate_prepared_student(_PreparedStudent(out_features=32), config) diff --git a/tests/examples/diffusers/fastgen/test_vendored_migration.py b/tests/examples/diffusers/fastgen/test_vendored_migration.py index 070a8358e7a..a869552154a 100644 --- a/tests/examples/diffusers/fastgen/test_vendored_migration.py +++ b/tests/examples/diffusers/fastgen/test_vendored_migration.py @@ -149,7 +149,7 @@ def test_formerly_vendored_files_use_standard_nvidia_header(): def test_data_builders_importable_and_accept_shared_cache_options(): - """The real-data builder exposes the negative embedding and deterministic split seams.""" + """The real-data builder exposes the shared negative-embedding seam.""" pytest.importorskip("nemo_automodel") pytest.importorskip("torch") @@ -158,7 +158,6 @@ def test_data_builders_importable_and_accept_shared_cache_options(): assert callable(fastgen_data.build_text_to_image_multiresolution_dataloader) sig = inspect.signature(fastgen_data.build_text_to_image_multiresolution_dataloader) assert "negative_prompt_embedding_path" in sig.parameters - assert "split" in sig.parameters # Default None => CFG-less construction works without the negative embedding (it is optional). assert sig.parameters["negative_prompt_embedding_path"].default is None diff --git a/tests/gpu/torch/fastgen/test_pdd_toy.py b/tests/gpu/torch/fastgen/test_pdd_toy.py index 10f480ad2d8..bf19a1f28e3 100644 --- a/tests/gpu/torch/fastgen/test_pdd_toy.py +++ b/tests/gpu/torch/fastgen/test_pdd_toy.py @@ -40,8 +40,13 @@ def __init__(self) -> None: self.backbone = nn.Linear(_WIDTH, _WIDTH) self.projection = nn.Linear(_WIDTH, _WIDTH) - def forward(self, state: torch.Tensor) -> torch.Tensor: - return self.projection(torch.tanh(self.backbone(state))) + 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): @@ -90,8 +95,7 @@ def student_fused_block( projection = model.projection assert isinstance(projection, PDDOutputProjection) self.fused_calls += 1 - with projection.fuse_block(start, end, grid): - return model(state.to(self._model_dtype(model))) + return model(state.to(self._model_dtype(model)), fusion=(start, end, grid)) def teacher_velocity( self, @@ -131,7 +135,6 @@ def _build( block_size_min=1, block_size_max=_GRID_SIZE, inference_blocks=[2, 2], - student_sample_steps=2, guidance_scale=None, ) student = _Student().to(device=device, dtype=torch.bfloat16) diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 28661c0562a..7c4c8574b3b 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -83,7 +83,7 @@ def test_load_pdd_config_builtin_recipe(): 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] + assert config.inference_blocks == (32, 32, 32, 32) QUANTIZER_ATTRIBUTE_SCHEMA = ( diff --git a/tests/unit/torch/fastgen/test_pdd_config.py b/tests/unit/torch/fastgen/test_pdd_config.py index aff9c2e15a1..ddb1a9caba3 100644 --- a/tests/unit/torch/fastgen/test_pdd_config.py +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -19,75 +19,40 @@ import pytest -from modelopt.torch.fastgen import PDDConfig, SampleTimestepConfig, load_pdd_config +from modelopt.torch.fastgen import PDDConfig, load_pdd_config -def test_default_pdd_config_is_canonical_and_lists_are_independent(): - first = PDDConfig() - second = PDDConfig() - - assert first.pred_type == "flow" - assert first.student_sample_type == "ode" - assert first.student_sample_steps == 4 - assert first.grid_size == 128 - assert first.grid_max_t == 0.999 - assert first.flow_shift == 5.0 - assert first.block_size_min == 4 - assert first.block_size_max == 64 - assert first.teacher_integrator == "euler" - assert first.inference_blocks == [32, 32, 32, 32] - assert first.data_free is False - assert first.inference_blocks is not second.inference_blocks +def test_default_pdd_config_is_canonical(): + config = PDDConfig() - first.inference_blocks[0] = 16 - assert second.inference_blocks == [32, 32, 32, 32] + 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_and_adapter_time_scale(): +def test_pdd_config_accepts_supported_schedule(): config = PDDConfig( inference_blocks=[64, 64], - student_sample_steps=2, teacher_integrator="midpoint", - num_train_timesteps=1000, data_free=True, ) - assert config.inference_blocks == [64, 64] + assert config.inference_blocks == (64, 64) assert config.teacher_integrator == "midpoint" - assert config.num_train_timesteps == 1000 assert config.data_free is True -def test_rejected_attribute_assignment_leaves_pdd_config_unchanged(): - config = PDDConfig() - - with pytest.raises(ValueError, match="student_sample_steps must equal"): - config.student_sample_steps = 3 - - assert config.student_sample_steps == 4 - assert config.inference_blocks == [32, 32, 32, 32] - - -def test_rejected_mapping_assignment_leaves_pdd_config_unchanged(): - config = PDDConfig() - - with pytest.raises(ValueError, match="student_sample_steps must equal"): - config["inference_blocks"] = [64, 64] - - assert config.student_sample_steps == 4 - assert config.inference_blocks == [32, 32, 32, 32] - - @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) - config = PDDConfig() - with pytest.raises(ValueError, match="grid_max_t must be a float"): - config.grid_max_t = value - assert config.grid_max_t == 0.999 - @pytest.mark.parametrize( ("overrides", "message"), @@ -102,10 +67,6 @@ def test_pdd_config_rejects_non_float_grid_max_t_before_coercion(value): ({"grid_size": 130}, "must be divisible"), ({"inference_blocks": []}, "at least one block"), ({"inference_blocks": [32, 32, 32]}, "must sum to grid_size"), - ( - {"inference_blocks": [64, 64], "student_sample_steps": 4}, - "student_sample_steps must equal", - ), ], ) def test_pdd_config_rejects_invalid_grid_and_block_boundaries(overrides, message): @@ -114,34 +75,37 @@ def test_pdd_config_rejects_invalid_grid_and_block_boundaries(overrides, message def test_pdd_config_accepts_inference_partition_outside_training_block_support(): - config = PDDConfig(inference_blocks=[1, 127], student_sample_steps=2) + config = PDDConfig(inference_blocks=[1, 127]) - assert config.inference_blocks == [1, 127] + assert config.inference_blocks == (1, 127) + + +@pytest.mark.parametrize("mapping_assignment", [False, True]) +def test_rejected_assignment_leaves_pdd_config_unchanged(mapping_assignment): + config = PDDConfig() + + with pytest.raises(ValueError, match="must sum to grid_size"): + if mapping_assignment: + config["grid_size"] = 64 + else: + config.grid_size = 64 + + assert config.grid_size == 128 + assert config.inference_blocks == (32, 32, 32, 32) @pytest.mark.parametrize( - "overrides", - [ - {"pred_type": "x0"}, - {"student_sample_type": "sde"}, - {"teacher_integrator": "heun"}, - {"teacher_integrator": "rk4"}, - ], + "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_rejects_nondefault_sample_timestep_config(): - with pytest.raises(ValueError, match="sample_t_cfg is unused by PDD"): - PDDConfig(sample_t_cfg=SampleTimestepConfig(shift=6.0)) - - 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]\nstudent_sample_steps: 2\nteacher_integrator: midpoint\n", + "inference_blocks: [64, 64]\nteacher_integrator: midpoint\n", encoding="utf-8", ) @@ -149,5 +113,5 @@ def test_pdd_config_loads_filesystem_yaml_with_optional_suffix(tmp_path): from_class = PDDConfig.from_yaml(config_path) assert loaded == from_class - assert loaded.inference_blocks == [64, 64] + 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 index 03d9af3dc89..fb5c98d2c72 100644 --- a/tests/unit/torch/fastgen/test_pdd_pipeline.py +++ b/tests/unit/torch/fastgen/test_pdd_pipeline.py @@ -148,7 +148,6 @@ def _config(*, teacher_integrator: str = "euler") -> PDDConfig: block_size_min=2, block_size_max=4, inference_blocks=[4, 4], - student_sample_steps=2, teacher_integrator=teacher_integrator, ) diff --git a/tests/unit/torch/fastgen/test_pdd_projection.py b/tests/unit/torch/fastgen/test_pdd_projection.py index 270a4fc443a..d6c9d233005 100644 --- a/tests/unit/torch/fastgen/test_pdd_projection.py +++ b/tests/unit/torch/fastgen/test_pdd_projection.py @@ -186,8 +186,7 @@ def test_fused_forward_matches_independent_weighted_head_sum(layout, bias): coefficients = torch.tensor([0.5 / 0.8, 0.3 / 0.8]) expected = torch.einsum("n,bno->bo", coefficients, explicit_heads[:, 1:3]) - with projection.fuse_block(1, 3, grid): - actual = projection(inputs) + actual = projection(inputs, fusion=(1, 3, grid)) torch.testing.assert_close(actual, expected) torch.testing.assert_close( @@ -200,16 +199,12 @@ def test_fused_forward_matches_independent_weighted_head_sum(layout, bias): assert torch.equal(grid, original_grid) -def test_fusion_context_exception_cleans_up_and_allows_reuse(): +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]) - with pytest.raises(RuntimeError, match="body failed"), projection.fuse_block(0, 2, grid): - raise RuntimeError("body failed") - - with projection.fuse_block(1, 3, grid): - assert projection(inputs).shape == (1, 6) + assert projection(inputs, fusion=(1, 3, grid)).shape == (1, 6) assert projection(inputs).shape == (1, 18) @@ -217,11 +212,11 @@ def test_fusion_context_exception_cleans_up_and_allows_reuse(): ("start", "end", "message"), [(-1, 2, "0 <= start"), (1, 1, "0 <= start"), (1, 4, "0 <= start")], ) -def test_fusion_context_rejects_invalid_blocks(start, end, message): +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.fuse_block(start, end, grid): - pass + with pytest.raises(ValueError, match=message): + projection(torch.zeros(1, 2), fusion=(start, end, grid)) def _state_clone(state): diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index a1dd2858dee..9e474621ed3 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -26,7 +26,7 @@ from torch import nn import modelopt.torch.fastgen.plugins.qwen_image_pdd as qwen_image_pdd_plugin -from modelopt.torch.fastgen import PDDConfig, PDDPipeline +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 @@ -79,9 +79,12 @@ def forward( encoder_hidden_states, encoder_hidden_states_mask, img_shapes, - max_txt_seq_len, + 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 @@ -89,7 +92,11 @@ def forward( hidden = torch.tanh(self.backbone(hidden_states)) hidden = hidden + condition_value.to(hidden.dtype) hidden = hidden + (0.1 * timestep[:, None, None]).to(hidden.dtype) - output = self.proj_out(hidden) + 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(), @@ -114,9 +121,7 @@ def _config(*, guidance_scale: float | None = 4.0, grid_size: int = 4) -> PDDCon block_size_min=1, block_size_max=grid_size, inference_blocks=[2, 2] if grid_size == 4 else [grid_size], - student_sample_steps=2 if grid_size == 4 else 1, guidance_scale=guidance_scale, - num_train_timesteps=None, ) @@ -230,33 +235,6 @@ def test_fused_student_matches_explicit_packed_weight_fusion() -> None: assert student.proj_out(projection.weight.new_zeros(1, 5)).shape[-1] == 16 -def test_teacher_cfg_remote_preflight_failure_stops_both_model_calls(monkeypatch) -> None: - teacher = _TinyQwenTransformer() - adapter = QwenImagePDDAdapter(_config(guidance_scale=4.0)) - state, time, condition, negative_condition = _inputs() - - monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_available", lambda: True) - monkeypatch.setattr(qwen_image_pdd_plugin.dist, "is_initialized", lambda: True) - - def report_remote_failure(failed, *, op): - assert not bool(failed) - assert op is torch.distributed.ReduceOp.MAX - failed.fill_(1) - - monkeypatch.setattr(qwen_image_pdd_plugin.dist, "all_reduce", report_remote_failure) - - with pytest.raises(RuntimeError, match="preflight failed on another rank"): - adapter.teacher_velocity( - teacher, - state, - time, - condition=condition, - negative_condition=negative_condition, - ) - - assert teacher.calls == [] - - def test_teacher_cfg_zero_guided_norm_uses_qwen_clamp() -> None: class ZeroGuidedTeacher(_QwenImageTestDouble): def __init__(self) -> None: @@ -355,6 +333,7 @@ def _mr210_qwen_forward_oracle( 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) @@ -363,23 +342,13 @@ def _mr210_qwen_forward_oracle( max_txt_seq_len=max_txt_seq_len, device=hidden_states.device, ) - image_mask = torch.ones( - (hidden_states.shape[0], hidden_states.shape[1]), - dtype=torch.bool, - device=hidden_states.device, - ) - joint_attention_mask = torch.cat( - (encoder_hidden_states_mask.to(torch.bool), image_mask), - dim=1, - )[:, None, None, :] 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=None, + encoder_hidden_states_mask=encoder_hidden_states_mask, temb=temb, image_rotary_emb=image_rotary_emb, - joint_attention_kwargs={"attention_mask": joint_attention_mask}, ) return model.proj_out(model.norm_out(hidden_states, temb)) @@ -586,8 +555,8 @@ def test_mr210_joint_mask_ignores_padded_token_values() -> None: 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, 3, 12, generator=generator).to(torch.bfloat16) - mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long) + 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( @@ -607,8 +576,8 @@ def test_mr210_joint_mask_ignores_padded_token_values() -> None: captured_masks: list[torch.Tensor] = [] def capture_block_mask(_module, _args, kwargs): - assert kwargs["encoder_hidden_states_mask"] is None - captured_masks.append(kwargs["joint_attention_kwargs"]["attention_mask"].detach().clone()) + 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, @@ -640,9 +609,7 @@ def capture_block_mask(_module, _args, kwargs): 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 - expected_mask = torch.cat((mask.bool(), torch.ones(2, 4, dtype=torch.bool)), dim=1) - expected_mask = expected_mask[:, None, None, :] - assert all(torch.equal(captured, expected_mask) for captured in captured_masks) + assert all(torch.equal(captured, mask.bool()) for captured in captured_masks) def test_mr210_preserves_diffusers_output_and_harmless_call_contract() -> None: @@ -754,9 +721,6 @@ def direct_packed(current_condition): def test_qwen_pdd_rejects_unsupported_configuration_and_inputs() -> None: - with pytest.raises(ValueError, match="num_train_timesteps=None"): - QwenImagePDDAdapter(_config().model_copy(update={"num_train_timesteps": 1000})) - transformer = _TinyQwenTransformer() config = _config() adapter = QwenImagePDDAdapter(config) From 1ac442aef075cd841fb91ebfed67acdbd1ee9e0e Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 3 Sep 2026 23:18:25 -0700 Subject: [PATCH 41/45] docs: move PDD note to current release Signed-off-by: Meng Xin --- CHANGELOG.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2d29e2e1ca6..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** @@ -211,7 +213,6 @@ Changelog - Add offline DFlash speculative decoding training. Train the draft module from pre-computed base-model hidden states dumped by ``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py``; base-model transformer layers are deleted after conversion to save memory. Controlled by the auto-derived ``dflash_offline`` flag on ``DFlashConfig`` (derived from ``data_args.offline_data_path``). The dump scripts now share ``collect_hidden_states/common.py`` for aux-layer selection (``--aux-layers eagle|dflash|``) and optional assistant-token ``loss_mask`` for answer-only-loss training. - Add ``mtsa.config.SKIP_SOFTMAX_TRITON_CALIB`` for skip-softmax attention-sparsity calibration through the fused Triton ``attention_calibrate`` kernel (HF ``modelopt_triton`` backend), measuring multi-threshold tile-skip statistics the way the Triton inference kernel actually skips tiles for both prefill and decode. Exposed as ``--sparse_attn_cfg skip_softmax_triton_calib`` in ``examples/llm_sparsity/attention_sparsity/hf_sa.py`` (with a new ``--calib_data_dir`` flag for RULER calibration data). - Add DMD2 distillation for few-step diffusion models in ``examples/diffusers/fastgen/``: distill Qwen-Image into a 4/8-step student via Distribution Matching Distillation. See `examples/diffusers/fastgen/README.md `_ for details. -- 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. - Make ``.agents/skills/`` the canonical location for agent skills; agent-specific directories (``.claude/skills/``, etc.) are now relative symlinks into ``.agents/``, so one skill suite serves multiple coding agents (Claude Code, Codex). See ``.agents/README.md``. - Extend Claude Code agent skills for PTQ, deployment, evaluation, monitoring, and baseline-vs-quantized result comparison. Adds evaluation task references for additional benchmarks, stronger PTQ checkpoint validation gates, and session-scoped workspace/job tracking. - Add SLURM Quality of Service (QoS) support to the ModelOpt launcher. Users can set QoS via ``slurm_config.qos`` or ``SLURM_QOS`` and the value is forwarded to ``nemo_run.SlurmExecutor``. From b8eeb30255671c038f78889508bd11c3f2f6b01d Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 3 Sep 2026 23:44:54 -0700 Subject: [PATCH 42/45] chore: document Qwen Image source attribution Signed-off-by: Meng Xin --- .pre-commit-config.yaml | 1 + LICENSE | 1 + .../torch/fastgen/plugins/qwen_image_pdd.py | 25 +++++++++++++++---- 3 files changed, 22 insertions(+), 5 deletions(-) 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/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/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index dc93e03770b..63e84334a4a 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -1,7 +1,19 @@ -# Adapted from the Qwen-Image implementation in Diffusers: -# https://github.com/huggingface/diffusers/blob/275869dcae4ebcfee6a80253fdabc56033335020/src/diffusers/models/transformers/transformer_qwenimage.py -# The masked joint-attention execution follows FastGen merge request 210. -# SPDX-FileCopyrightText: Copyright (c) 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved. +# 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 # @@ -17,7 +29,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen-Image adapter and explicit output-projection conversion for PDD.""" +"""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 From 519264faf93f262a3fba92c78b3957df9d484bf8 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Fri, 4 Sep 2026 00:13:25 -0700 Subject: [PATCH 43/45] refactor: simplify Qwen Image PDD adapter Signed-off-by: Meng Xin --- .../torch/fastgen/plugins/qwen_image_pdd.py | 319 ++++++------------ .../fastgen/test_qwen_image_pdd_plugin.py | 2 +- 2 files changed, 111 insertions(+), 210 deletions(-) diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index 63e84334a4a..f24dc241e11 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -85,12 +85,6 @@ _QWEN_IMAGE_PDD_EXECUTION_ATTRIBUTE = "_modelopt_qwen_image_pdd_execution" -def _require_binary_mask(mask: torch.Tensor, *, name: str) -> None: - mask_int = mask.to(torch.int64) - if not bool(torch.all((mask_int == 0) | (mask_int == 1)).item()): - raise ValueError(f"{name} mask must contain only zero and one values.") - - def _config_guidance_embeds(transformer: nn.Module) -> bool: config = getattr(transformer, "config", None) if isinstance(config, Mapping): @@ -266,6 +260,37 @@ def _validate_qwen_pdd_config(config: PDDConfig) -> None: 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, @@ -281,14 +306,7 @@ def convert_qwen_image_to_pdd( 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.") - try: - current = transformer.get_submodule("proj_out") - except AttributeError as error: - raise ValueError( - "Qwen-Image transformer must register an nn.Linear at 'proj_out'." - ) from error - if not isinstance(current, nn.Linear): - raise TypeError(f"Qwen-Image proj_out must be nn.Linear, got {type(current).__name__}.") + current = _output_projection(transformer) projection = PDDOutputProjection.from_linear( current, @@ -297,8 +315,6 @@ def convert_qwen_image_to_pdd( ) if projection is not current: transformer.proj_out = projection - if transformer.get_submodule("proj_out") is not projection: - raise RuntimeError("Qwen-Image proj_out replacement did not remain registered.") return projection @@ -312,14 +328,7 @@ def restore_qwen_image_pdd_projection( 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.") - try: - current = transformer.get_submodule("proj_out") - except AttributeError as error: - raise ValueError( - "Qwen-Image transformer must register an nn.Linear at 'proj_out'." - ) from error - if not isinstance(current, nn.Linear): - raise TypeError(f"Qwen-Image proj_out must be nn.Linear, got {type(current).__name__}.") + current = _output_projection(transformer) if isinstance(current, PDDOutputProjection): return PDDOutputProjection.from_linear( current, @@ -350,8 +359,6 @@ def restore_qwen_image_pdd_projection( projection.bias = current.bias projection.train(current.training) transformer.proj_out = projection - if transformer.get_submodule("proj_out") is not projection: - raise RuntimeError("Qwen-Image proj_out replacement did not remain registered.") return projection @@ -379,23 +386,6 @@ def __init__( ) self.compute_dtype = compute_dtype - @staticmethod - def _validate_state_and_time(state: torch.Tensor, time: torch.Tensor) -> None: - if state.ndim != 4: - raise ValueError( - f"Qwen-Image PDD state must have shape [B, C, H, W], got {tuple(state.shape)}." - ) - if state.shape[2] % 2 or state.shape[3] % 2: - raise ValueError("Qwen-Image PDD requires even latent height and width.") - if time.shape != (state.shape[0],): - raise ValueError( - f"Qwen-Image PDD time must have shape ({state.shape[0]},), got {tuple(time.shape)}." - ) - if time.device != state.device: - raise ValueError(f"time must be on {state.device}, got {time.device}.") - if not time.dtype.is_floating_point: - raise TypeError(f"time must use a real floating-point dtype, got {time.dtype}.") - @staticmethod def _parse_condition( condition: Any, @@ -417,21 +407,14 @@ def _parse_condition( ) 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 - or attention_mask.shape[1] != encoder_hidden_states.shape[1] - ): - raise ValueError( - f"{name} mask must be an integer/bool tensor matching the embedding " - "sequence length." - ) - batch_size = state.shape[0] - if encoder_hidden_states.shape[0] != batch_size or attention_mask.shape[0] != batch_size: - raise ValueError(f"{name} batch size must match state batch size {batch_size}.") + 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}.") - _require_binary_mask(attention_mask, name=name) return encoder_hidden_states, attention_mask def _model_dtype(self, model: nn.Module, fallback: torch.dtype) -> torch.dtype: @@ -444,82 +427,60 @@ def _model_dtype(self, model: nn.Module, fallback: torch.dtype) -> torch.dtype: @staticmethod def _extract_packed_output(output: Any) -> torch.Tensor: - if isinstance(output, tuple): - if not output: - raise TypeError("Qwen-Image model returned an empty tuple.") - packed = output[0] - elif isinstance(output, torch.Tensor): - packed = output - elif hasattr(output, "sample"): - packed = output.sample - else: - raise TypeError( - "Qwen-Image PDD could not extract a tensor from model output of type " - f"{type(output).__name__}." - ) - if not isinstance(packed, torch.Tensor): - raise TypeError("Qwen-Image model output payload must be a 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_call( + def _prepare_model_call( self, model: nn.Module, state: torch.Tensor, time: torch.Tensor, - condition: Any, model_kwargs: Mapping[str, Any], - *, - condition_name: str, - ) -> tuple[torch.Tensor, torch.Tensor]: - self._validate_state_and_time(state, time) + ) -> 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.") - encoder_hidden_states, attention_mask = self._parse_condition( - condition, - state=state, - name=condition_name, - ) conflicts = sorted(_CONTROLLED_MODEL_KWARGS.intersection(model_kwargs)) if conflicts: raise ValueError(f"Qwen-Image PDD model_kwargs contains controlled keys: {conflicts}.") - return encoder_hidden_states, attention_mask + 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, - state: torch.Tensor, + packed_state: torch.Tensor, time: torch.Tensor, - condition: Any, model_kwargs: Mapping[str, Any], *, - condition_name: str, - prepared_condition: tuple[torch.Tensor, torch.Tensor] | None = None, + 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: - if prepared_condition is None: - encoder_hidden_states, attention_mask = self._prepare_call( - model, - state, - time, - condition, - model_kwargs, - condition_name=condition_name, - ) - else: - encoder_hidden_states, attention_mask = prepared_condition + encoder_hidden_states, attention_mask = prepared_condition - batch_size, _, height, width = state.shape - 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) encoder_hidden_states = encoder_hidden_states.to(model_dtype) call_kwargs = dict(model_kwargs) if fusion is not None: @@ -529,43 +490,12 @@ def _call_packed( timestep=time, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=attention_mask, - img_shapes=build_img_shapes(batch_size, height, width), + img_shapes=img_shapes, return_dict=False, **call_kwargs, ) return self._extract_packed_output(output) - def _prepare_teacher_cfg_calls( - self, - model: nn.Module, - state: torch.Tensor, - time: torch.Tensor, - condition: Any, - negative_condition: Any, - model_kwargs: Mapping[str, Any], - ) -> tuple[ - tuple[torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor], - ]: - """Validate both local CFG calls before either model execution begins.""" - prepared_condition = self._prepare_call( - model, - state, - time, - condition, - model_kwargs, - condition_name="condition", - ) - prepared_negative_condition = self._prepare_call( - model, - state, - time, - negative_condition, - model_kwargs, - condition_name="negative_condition", - ) - return prepared_condition, prepared_negative_condition - @staticmethod def _expected_packed_shape(state: torch.Tensor, *, output_features: int) -> torch.Size: return torch.Size( @@ -611,55 +541,6 @@ def _unpack_single(self, packed: torch.Tensor, state: torch.Tensor) -> torch.Ten ) return unpack_latents(packed, state.shape[2], state.shape[3]) - @staticmethod - def _all_heads_projection( - model: nn.Module, - grid_size: int, - base_out_features: int, - ) -> nn.Linear: - try: - projection = model.get_submodule("proj_out") - except AttributeError as error: - raise ValueError( - "Qwen student must register an output projection at 'proj_out'." - ) from error - if not isinstance(projection, nn.Linear): - raise TypeError("Qwen student proj_out must be a widened nn.Linear for PDD training.") - expected_out_features = grid_size * base_out_features - if projection.out_features != expected_out_features: - raise ValueError( - f"Qwen PDD proj_out has {projection.out_features} outputs; expected " - f"{expected_out_features} ({grid_size} heads x {base_out_features})." - ) - if isinstance(projection, PDDOutputProjection): - if projection.grid_size != grid_size: - raise ValueError( - f"Qwen PDD projection grid_size={projection.grid_size} does not match " - f"config grid_size={grid_size}." - ) - if projection.layer_spec != QWEN_IMAGE_PDD_LAYER_SPEC: - raise ValueError("Qwen PDD projection carries an incompatible layer specification.") - return projection - - @classmethod - def _fused_projection(cls, model: nn.Module, grid_size: int) -> PDDOutputProjection: - try: - projection = model.get_submodule("proj_out") - except AttributeError as error: - raise ValueError( - "Qwen student must register an output projection at 'proj_out'." - ) from error - if not isinstance(projection, PDDOutputProjection): - raise TypeError( - "Qwen fused PDD inference requires proj_out to be a PDDOutputProjection." - ) - cls._all_heads_projection( - model, - grid_size, - base_out_features=projection.base_out_features, - ) - return projection - def student_all_heads( self, model: nn.Module, @@ -670,18 +551,23 @@ def student_all_heads( **model_kwargs: Any, ) -> torch.Tensor: """Return unpacked PDD interval velocities from one Qwen call.""" - self._all_heads_projection( + _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, - state, + packed_state, time, - condition, model_kwargs, - condition_name="condition", + prepared_condition=prepared_condition, + model_dtype=model_dtype, + img_shapes=img_shapes, ) return self._unpack_all_heads(packed, state) @@ -698,14 +584,30 @@ def student_fused_block( **model_kwargs: Any, ) -> torch.Tensor: """Run one conditional Qwen call with its final projection fused for a block.""" - self._fused_projection(model, self.config.grid_size) + 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, - state, + packed_state, time, - condition, model_kwargs, - condition_name="condition", + prepared_condition=prepared_condition, + model_dtype=model_dtype, + img_shapes=img_shapes, fusion=(start, end, grid), ) return self._unpack_single(packed, state) @@ -722,40 +624,39 @@ def teacher_velocity( **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_condition, prepared_negative_condition = self._prepare_teacher_cfg_calls( - model, - state, - time, - condition, + prepared_negative_condition = self._parse_condition( negative_condition, - model_kwargs, + state=state, + name="negative_condition", ) - else: - prepared_condition = None - prepared_negative_condition = None conditional = self._call_packed( model, - state, + packed_state, time, - condition, model_kwargs, - condition_name="condition", 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, - state, + packed_state, time, - negative_condition, model_kwargs, - condition_name="negative_condition", 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: diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index 9e474621ed3..e69f089ac7d 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -247,7 +247,7 @@ def __init__(self) -> None: 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) + return (torch.full_like(hidden_states, value, dtype=torch.bfloat16),) teacher = ZeroGuidedTeacher() state, time, condition, negative_condition = _inputs(batch_size=1) From dac6128107313e76e07aa963602054682eac661d Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Fri, 4 Sep 2026 01:55:05 -0700 Subject: [PATCH 44/45] fix: address PDD review feedback Signed-off-by: Meng Xin --- examples/diffusers/fastgen/dmd2/README.md | 2 +- .../diffusers/fastgen/fastgen_data/__init__.py | 18 +++++++++--------- examples/diffusers/fastgen/pdd/README.md | 8 +++++--- examples/diffusers/fastgen/pdd/compat.py | 2 +- .../diffusers/fastgen/pdd/requirements.txt | 4 ++++ .../preprocess/preprocessing_multiprocess.py | 4 ++-- examples/diffusers/fastgen/requirements.txt | 9 ++++----- modelopt/torch/fastgen/config.py | 8 -------- modelopt/torch/fastgen/flow_matching.py | 13 ++++--------- .../torch/fastgen/plugins/qwen_image_pdd.py | 3 --- .../diffusers/fastgen/test_dataset_paths.py | 11 +++++++---- .../diffusers/fastgen/test_pdd_recipe_setup.py | 1 + tests/unit/torch/fastgen/test_pdd_config.py | 14 -------------- .../torch/fastgen/test_pdd_reference_math.py | 15 --------------- .../fastgen/test_qwen_image_pdd_plugin.py | 4 +--- 15 files changed, 39 insertions(+), 77 deletions(-) create mode 100644 examples/diffusers/fastgen/pdd/requirements.txt diff --git a/examples/diffusers/fastgen/dmd2/README.md b/examples/diffusers/fastgen/dmd2/README.md index a67adc51d2f..e523bc86327 100644 --- a/examples/diffusers/fastgen/dmd2/README.md +++ b/examples/diffusers/fastgen/dmd2/README.md @@ -13,7 +13,7 @@ output distribution. Built on `modelopt.torch.fastgen` and NeMo AutoModel's ## Requirements & self-contained data path -This example runs against **stock upstream `nemo_automodel==0.5.0`** (see +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: diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index 8338f44e517..77bc814b841 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -15,9 +15,9 @@ """Self-contained shared dataloaders for the FastGen diffusion examples. -The data path builds on stock ``nemo_automodel==0.5.0`` where it is model-agnostic and implements -the example-owned batch contract locally, so the published example does not depend on AutoModel -source modifications: +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 functions + dataloader builder. It reuses the upstream ``SequentialBucketSampler`` and emits either the ordinary latent-conditioned batch or a @@ -37,7 +37,7 @@ # 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 release. +# Convert a missing-helper ImportError into an actionable message naming the supported range. try: from .collate_fns import ( build_text_to_image_multiresolution_dataloader, @@ -48,7 +48,7 @@ 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.5.0 install (it imports the unmodified 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" @@ -78,12 +78,12 @@ def _warn_if_unsupported_upstream() -> None: raw = str(getattr(nemo_automodel, "__version__", "") or "") match = re.match(r"^(\d+)\.(\d+)\.(\d+)", raw) version = tuple(int(part) for part in match.groups()) if match else () - if version != (0, 5, 0): + if not ((0, 4, 0) <= version < (0, 6, 0)): logging.getLogger(__name__).warning( - "fastgen_data: installed nemo_automodel %s does not match the tested release " - "(==0.5.0). The vendored data/preprocessing code imports unmodified upstream " + "fastgen_data: installed nemo_automodel %s is outside the tested range " + "(>=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 release.", + "fail or behavior drifts, pin nemo_automodel to the supported range.", raw or "", ) except Exception: # pragma: no cover - never block import on a version probe diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index f958c7743c7..3bb4929075c 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -56,7 +56,7 @@ The output is a full Diffusers pipeline overlay with a widened transformer. Poin ## Train and resume ```bash -pip install -r examples/diffusers/fastgen/requirements.txt +pip install -r examples/diffusers/fastgen/pdd/requirements.txt export MODELOPT_FASTGEN_DATASET_CACHE_DIR=/absolute/path/to/qwen_image_cache torchrun --standalone --nproc-per-node=8 \ @@ -66,8 +66,10 @@ torchrun --standalone --nproc-per-node=8 \ ``` The cache must contain `metadata.json`, its declared prompt-embedding shards, and -`negative_prompt_embedding.pt`. `MODELOPT_FASTGEN_DATASET_CACHE_DIR` overrides the configured -cache root; paths declared by the dataset remain confined to that root. +`negative_prompt_embedding.pt`. `MODELOPT_FASTGEN_DATASET_CACHE_DIR` lets a cluster launcher +redirect an unchanged recipe to its mounted cache; it overrides the configured cache root and +moves sample payloads and the negative embedding together. Paths declared by the dataset remain +confined to that root. 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. diff --git a/examples/diffusers/fastgen/pdd/compat.py b/examples/diffusers/fastgen/pdd/compat.py index 045373d3094..d168562abb2 100644 --- a/examples/diffusers/fastgen/pdd/compat.py +++ b/examples/diffusers/fastgen/pdd/compat.py @@ -110,7 +110,7 @@ def build_manager_args(**kwargs: Any) -> dict[str, Any]: ) return manager_args - class PDDSetupPipeline: + class PDDSetupPipeline(automodel_diffusion_train.NeMoAutoDiffusionPipeline): @classmethod def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any: del cls 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/preprocess/preprocessing_multiprocess.py b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py index c08f731cbe2..fdd2e9fe20a 100644 --- a/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py +++ b/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py @@ -118,7 +118,7 @@ def _save_metadata_shards( for item_index, item in enumerate(all_metadata): cache_file = Path(item["cache_file"]).resolve(strict=True) try: - cache_file.relative_to(output_root) + 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 " @@ -127,7 +127,7 @@ def _save_metadata_shards( normalized_metadata.append( { **item, - "cache_file": str(cache_file), + "cache_file": relative_cache_file.as_posix(), } ) diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index 90159ecb257..59c6c3c9fec 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -1,11 +1,10 @@ -# Runtime requirements for the Qwen-Image AutoModel examples. +# 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. -diffusers==0.38.0 # NeMo AutoModel supplies the parent recipe, FSDP2 wrapping, checkpointer, and upstream diffusion -# data helpers. PDD's guarded setup compatibility seam is validated against release 0.5.0. -nemo_automodel[diffusion]==0.5.0 +# 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 and training logs. +# Optional but recommended for training logs. wandb diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py index 06d706bfbb5..83a55b81fac 100644 --- a/modelopt/torch/fastgen/config.py +++ b/modelopt/torch/fastgen/config.py @@ -274,14 +274,6 @@ class PDDConfig(ModeloptBaseConfig): ), ) - def __setattr__(self, name: str, value: object) -> None: - """Validate cross-field invariants before changing an initialized field.""" - if name in type(self).model_fields and name in self.__dict__: - candidate = self.model_dump() - candidate[name] = value - type(self).model_validate(candidate) - super().__setattr__(name, value) - @field_validator("grid_max_t", mode="before") @classmethod def _check_grid_max_t_type(cls, value: object) -> object: diff --git a/modelopt/torch/fastgen/flow_matching.py b/modelopt/torch/fastgen/flow_matching.py index 3696572e042..735dd3f9398 100644 --- a/modelopt/torch/fastgen/flow_matching.py +++ b/modelopt/torch/fastgen/flow_matching.py @@ -188,16 +188,11 @@ def integrate_interval_velocities( interval_ids[None] < end_tensor[:, None] ) widths = torch.diff(grid.to(device=state.device, dtype=result_dtype)) - velocity_mask = mask.reshape(mask.shape + (1,) * (velocities.ndim - 2)) - selected_velocities = torch.where( - velocity_mask, - velocities.to(device=state.device, dtype=result_dtype), - torch.zeros((), device=state.device, dtype=result_dtype), - ) + weights = mask.to(result_dtype) * widths[None] update = torch.einsum( - "n,bn...->b...", - widths, - selected_velocities, + "bn,bn...->b...", + weights, + velocities.to(device=state.device, dtype=result_dtype), ) return state.to(dtype=result_dtype) + update diff --git a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py index f24dc241e11..056f6e46943 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image_pdd.py +++ b/modelopt/torch/fastgen/plugins/qwen_image_pdd.py @@ -106,9 +106,7 @@ def _qwen_image_pdd_forward( encoder_hidden_states_mask: torch.Tensor | None = None, timestep: torch.Tensor | None = None, img_shapes: list[Any] | None = None, - txt_seq_lens: list[int] | None = None, guidance: torch.Tensor | None = None, - max_txt_seq_len: int | None = None, attention_kwargs: dict[str, Any] | None = None, controlnet_block_samples: Any = None, additional_t_cond: torch.Tensor | None = None, @@ -130,7 +128,6 @@ def _qwen_image_pdd_forward( 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.") - del txt_seq_lens max_txt_seq_len = encoder_hidden_states.shape[1] encoder_hidden_states_mask = encoder_hidden_states_mask.to(torch.bool) diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index 8ddfb8d408b..d00d56bbb51 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -196,7 +196,7 @@ def test_builder_logs_effective_root_once_on_rank_zero(make_fastgen_cache, caplo assert "selected=6/6" in messages[0] -def test_preprocessing_publishes_absolute_paths_for_relative_output(monkeypatch, tmp_path): +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")) @@ -220,6 +220,9 @@ def test_preprocessing_publishes_absolute_paths_for_relative_output(monkeypatch, shard = json.loads((output / "metadata_shard_s0000.json").read_text()) published = pathlib.Path(shard[0]["cache_file"]) - assert published.is_absolute() - assert published == payload.resolve() - published.relative_to(output.resolve()) + 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_recipe_setup.py b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py index d0f6bfcf99b..ebc6ae3e1bd 100644 --- a/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py +++ b/tests/examples/diffusers/fastgen/test_pdd_recipe_setup.py @@ -179,6 +179,7 @@ def from_pretrained(cls, *args, **kwargs): 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) diff --git a/tests/unit/torch/fastgen/test_pdd_config.py b/tests/unit/torch/fastgen/test_pdd_config.py index ddb1a9caba3..ebbfaae9fd6 100644 --- a/tests/unit/torch/fastgen/test_pdd_config.py +++ b/tests/unit/torch/fastgen/test_pdd_config.py @@ -80,20 +80,6 @@ def test_pdd_config_accepts_inference_partition_outside_training_block_support() assert config.inference_blocks == (1, 127) -@pytest.mark.parametrize("mapping_assignment", [False, True]) -def test_rejected_assignment_leaves_pdd_config_unchanged(mapping_assignment): - config = PDDConfig() - - with pytest.raises(ValueError, match="must sum to grid_size"): - if mapping_assignment: - config["grid_size"] = 64 - else: - config.grid_size = 64 - - assert config.grid_size == 128 - assert config.inference_blocks == (32, 32, 32, 32) - - @pytest.mark.parametrize( "overrides", [{"teacher_integrator": "heun"}, {"teacher_integrator": "rk4"}] ) diff --git a/tests/unit/torch/fastgen/test_pdd_reference_math.py b/tests/unit/torch/fastgen/test_pdd_reference_math.py index de91bd14e97..e467fe35555 100644 --- a/tests/unit/torch/fastgen/test_pdd_reference_math.py +++ b/tests/unit/torch/fastgen/test_pdd_reference_math.py @@ -166,21 +166,6 @@ def test_production_half_open_integration_matches_independent_oracle_per_sample( torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) -def test_production_integration_does_not_consume_excluded_nonfinite_heads(): - grid = make_shifted_flow_grid(4, 5.0, max_t=0.999, dtype=torch.float64) - state = torch.tensor([[3.0, -2.0]], dtype=torch.float64) - velocities = torch.tensor( - [[[torch.nan, torch.nan], [2.0, -1.0], [-3.0, 4.0], [torch.inf, -torch.inf]]], - dtype=torch.float64, - ) - - actual = integrate_interval_velocities(state, velocities, grid, start=1, end=3) - expected = _reference_integrate(state[0], velocities[0], grid, start=1, end=3)[None] - - assert torch.all(torch.isfinite(actual)) - 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) diff --git a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py index e69f089ac7d..625816cac53 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_pdd_plugin.py @@ -612,7 +612,7 @@ def capture_block_mask(_module, _args, kwargs): assert all(torch.equal(captured, mask.bool()) for captured in captured_masks) -def test_mr210_preserves_diffusers_output_and_harmless_call_contract() -> None: +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 = { @@ -623,7 +623,6 @@ def test_mr210_preserves_diffusers_output_and_harmless_call_contract() -> None: "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), - "txt_seq_lens": [3, 1], "guidance": None, } @@ -690,7 +689,6 @@ def direct_packed(current_condition): encoder_hidden_states=embeddings, encoder_hidden_states_mask=mask, img_shapes=build_img_shapes(2, 4, 4), - max_txt_seq_len=int(mask.sum(dim=1).max().item()), return_dict=False, )[0] From 58cc97c73805ca2399a7e27f260829498eafc670 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Fri, 4 Sep 2026 02:11:44 -0700 Subject: [PATCH 45/45] refactor: use config for FastGen cache path Signed-off-by: Meng Xin --- examples/diffusers/fastgen/dmd2/README.md | 18 +++---- .../fastgen/dmd2/configs/qwen_image.yaml | 3 +- .../diffusers/fastgen/fastgen_data/paths.py | 16 +------ examples/diffusers/fastgen/pdd/README.md | 9 ++-- .../diffusers/fastgen/test_dataset_paths.py | 48 ++----------------- 5 files changed, 19 insertions(+), 75 deletions(-) diff --git a/examples/diffusers/fastgen/dmd2/README.md b/examples/diffusers/fastgen/dmd2/README.md index e523bc86327..cc3adff434f 100644 --- a/examples/diffusers/fastgen/dmd2/README.md +++ b/examples/diffusers/fastgen/dmd2/README.md @@ -52,16 +52,16 @@ 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 override it without -editing YAML: +Then point the config's `data.dataloader.cache_dir` at ``, or append this dotted +override to the training command: ```bash -export MODELOPT_FASTGEN_DATASET_CACHE_DIR=/absolute/path/to/cache +--data.dataloader.cache_dir=/absolute/path/to/cache ``` -The environment override must name an absolute existing directory. The config keeps -`negative_prompt_embedding_path: negative_prompt_embedding.pt`, so both samples and the negative -embedding are resolved from that same selected root. Paths that escape it are rejected. +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 @@ -107,8 +107,8 @@ this example subclasses. `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`, or export - `MODELOPT_FASTGEN_DATASET_CACHE_DIR` to an absolute existing cache directory. +- **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`. @@ -188,7 +188,7 @@ student). | `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` | Environment-overridable latent cache root + optional root-relative CFG embedding. | +| `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 diff --git a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml index d08ba727438..11d91a916ce 100644 --- a/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml +++ b/examples/diffusers/fastgen/dmd2/configs/qwen_image.yaml @@ -149,8 +149,7 @@ data: dataloader: _target_: fastgen_data.build_text_to_image_multiresolution_dataloader cache_dir: /path/to/preprocessed/qwen_image_1024p - # The environment override changes the effective root; referenced files must stay beneath it. - # MODELOPT_FASTGEN_DATASET_CACHE_DIR overrides cache_dir when set to a non-empty value. + # Referenced files must stay beneath this root. base_resolution: [1024, 1024] batch_size: 1 drop_last: false diff --git a/examples/diffusers/fastgen/fastgen_data/paths.py b/examples/diffusers/fastgen/fastgen_data/paths.py index ebcefb31b87..1c63b43ccf1 100644 --- a/examples/diffusers/fastgen/fastgen_data/paths.py +++ b/examples/diffusers/fastgen/fastgen_data/paths.py @@ -17,13 +17,10 @@ from __future__ import annotations -import os from pathlib import Path __all__ = ["resolve_cache_root", "resolve_under_root"] -_CACHE_ROOT_ENV = "MODELOPT_FASTGEN_DATASET_CACHE_DIR" - def _existing_directory(path: str | Path, label: str) -> Path: resolved = Path(path).resolve(strict=True) @@ -33,18 +30,7 @@ def _existing_directory(path: str | Path, label: str) -> Path: def resolve_cache_root(configured_root: str | Path) -> Path: - """Return the effective cache root selected by config and environment. - - An unset or exactly empty ``MODELOPT_FASTGEN_DATASET_CACHE_DIR`` falls back to - ``configured_root``. A nonempty override must already be an absolute path to an existing - directory. - """ - override = os.environ.get(_CACHE_ROOT_ENV) - if override: - override_path = Path(override) - if not override_path.is_absolute(): - raise ValueError(f"{_CACHE_ROOT_ENV} must be an absolute path; got {override!r}") - return _existing_directory(override_path, _CACHE_ROOT_ENV) + """Resolve the configured cache root to an existing directory.""" return _existing_directory(configured_root, "configured cache root") diff --git a/examples/diffusers/fastgen/pdd/README.md b/examples/diffusers/fastgen/pdd/README.md index 3bb4929075c..f7ef06fa7db 100644 --- a/examples/diffusers/fastgen/pdd/README.md +++ b/examples/diffusers/fastgen/pdd/README.md @@ -57,19 +57,18 @@ The output is a full Diffusers pipeline overlay with a widened transformer. Poin ```bash pip install -r examples/diffusers/fastgen/pdd/requirements.txt -export MODELOPT_FASTGEN_DATASET_CACHE_DIR=/absolute/path/to/qwen_image_cache 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`. `MODELOPT_FASTGEN_DATASET_CACHE_DIR` lets a cluster launcher -redirect an unchanged recipe to its mounted cache; it overrides the configured cache root and -moves sample payloads and the negative embedding together. Paths declared by the dataset remain -confined to that root. +`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. diff --git a/tests/examples/diffusers/fastgen/test_dataset_paths.py b/tests/examples/diffusers/fastgen/test_dataset_paths.py index d00d56bbb51..86cada06dff 100644 --- a/tests/examples/diffusers/fastgen/test_dataset_paths.py +++ b/tests/examples/diffusers/fastgen/test_dataset_paths.py @@ -39,39 +39,20 @@ from fastgen_data.text_to_image_dataset import TextToImageDataset -def test_cache_root_uses_unset_or_empty_fallback(make_fastgen_cache, monkeypatch, tmp_path): +def test_cache_root_resolves_configured_directory(make_fastgen_cache, tmp_path): cache = make_fastgen_cache(tmp_path / "cache") - monkeypatch.delenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", raising=False) assert resolve_cache_root(cache) == cache.resolve() - monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", "") - assert resolve_cache_root(cache) == cache.resolve() - - -@pytest.mark.parametrize("override", ["relative/cache", "~/cache", " "]) -def test_cache_root_rejects_nonempty_relative_override(monkeypatch, override, tmp_path): - fallback = tmp_path / "fallback" - fallback.mkdir() - monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", override) - - with pytest.raises(ValueError, match="absolute"): - resolve_cache_root(fallback) - -def test_cache_root_rejects_missing_or_non_directory_override(monkeypatch, tmp_path): - fallback = tmp_path / "fallback" - fallback.mkdir() - - monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", str(tmp_path / "missing")) +def test_cache_root_rejects_missing_or_non_directory(tmp_path): with pytest.raises(FileNotFoundError): - resolve_cache_root(fallback) + resolve_cache_root(tmp_path / "missing") regular_file = tmp_path / "file" regular_file.write_text("not a directory") - monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", str(regular_file)) with pytest.raises(NotADirectoryError): - resolve_cache_root(fallback) + resolve_cache_root(regular_file) def test_resolve_under_root_rejects_traversal_absolute_and_symlink_escape(tmp_path): @@ -141,27 +122,6 @@ def test_prompt_only_dataset_and_loader_do_not_emit_image_latents(make_fastgen_c assert "negative_text_embeddings" in batch -def test_environment_redirects_samples_and_relative_negative_embedding( - make_fastgen_cache, monkeypatch, tmp_path -): - fallback = make_fastgen_cache(tmp_path / "fallback", marker=1.0) - override = make_fastgen_cache(tmp_path / "override", marker=9.0) - monkeypatch.setenv("MODELOPT_FASTGEN_DATASET_CACHE_DIR", str(override.resolve())) - - loader, _ = build_text_to_image_multiresolution_dataloader( - cache_dir=str(fallback), - batch_size=1, - num_workers=0, - shuffle=False, - negative_prompt_embedding_path="negative_prompt_embedding.pt", - ) - batch = next(iter(loader)) - - assert loader.dataset.cache_root == override.resolve() - assert batch["metadata"]["prompts"][0].startswith("prompt-9.0-") - assert torch.equal(batch["negative_text_embeddings"], torch.full((1, 2, 3), 9.0)) - - def test_builder_rejects_negative_embedding_escape(make_fastgen_cache, tmp_path): cache = make_fastgen_cache(tmp_path / "cache") outside = tmp_path / "negative.pt"