diff --git a/CHANGELOG.rst b/CHANGELOG.rst index da02b315f67..46894b75cef 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -50,6 +50,7 @@ Changelog The legacy FSDP1 accelerate config is removed; ``llm_qat`` now documents FSDP2, DeepSpeed, and DDP backends. - The PTQ example scripts ``examples/llm_ptq/hf_ptq.py``, ``examples/llm_ptq/multinode_ptq.py`` and ``examples/megatron_bridge/quantize.py`` now derive their ``--qformat`` / ``--kv_cache_qformat`` (``--quant_cfg`` / ``--kv_cache_quant`` for Megatron-Bridge) CLI vocabularies by discovering the YAML presets under ``modelopt_recipes/configs/ptq/presets/{model,kv}/`` rather than carrying hardcoded ``QUANT_CFG_CHOICES`` / ``KV_QUANT_CFG_CHOICES`` tables. The discovery helper, alias table and ready-built ``QUANT_CFG_CHOICES`` / ``KV_QUANT_CFG_CHOICES`` mappings now live in ``modelopt.recipe.presets`` and are shared by all three scripts. Presets are loaded eagerly into a plain dict at import. Adding a new preset YAML makes it available on the CLI of all three with no script change — note this means each script now accepts every preset under those directories, not just a previously curated subset. All previously-supported short names (``int8_sq``, ``nvfp4_awq``, ``fp8_pb_wo``, ``nvfp4_mse``, ``w4a8_awq``, ``nvfp4_local_hessian``, ``fp8_pc_pt``, ``int8_wo``) keep working via a small deprecation alias table; new formats should be exposed as preset YAMLs (or, longer term, as full ``--recipe`` recipes). - Add ``configs/ptq/presets/kv/fp8_cast.yaml`` and ``configs/ptq/presets/kv/nvfp4_cast.yaml``, promoting ``fp8_cast`` / ``nvfp4_cast`` to first-class KV presets composed from the existing ``kv_fp8_cast`` / ``kv_nvfp4_cast`` unit fragments. The previous runtime ``use_constant_amax`` post-edit in ``hf_ptq.py`` is removed; ``use_constant_amax: true`` now lives in the YAML and is therefore authoritative. **Custom (out-of-tree) recipes that target a cast KV format must set ``use_constant_amax: true`` themselves on the ``[kv]_bmm_quantizer`` config** — in-tree recipes already do via the ``kv_*_cast`` units. +- 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. **Bug Fixes** diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md new file mode 100644 index 00000000000..b3c3bc780f9 --- /dev/null +++ b/examples/diffusers/fastgen/README.md @@ -0,0 +1,180 @@ +# 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. + +## 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. + +## Quick start — mock data (no dataset needed) + +The smoke config feeds random tensors at Qwen-Image's shapes, so it runs end-to-end with +**no dataset to prepare** — it exercises the full training loop (FSDP2 sharding, phase +alternation, checkpoint save/restore). Use it to validate your environment: + +```bash +torchrun --nproc-per-node=8 \ + examples/diffusers/fastgen/dmd2_finetune.py \ + --config examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml +``` + +Scale `--fsdp.dp_size` to your GPU count. You'll see alternating `phase=student` / +`phase=fake_score` log lines and a checkpoint written at the last step. + +> The mock loop validates wiring only — it does **not** produce meaningful images. For +> that, train on real data (below). + +## 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` | Real latent cache vs. `build_mock_t2i_dataloader`. | +| `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`), +enable `--fsdp.activation_checkpointing=true`, or use the mock smoke for wiring checks. + +**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/configs/dmd2_qwen_image.yaml b/examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml new file mode 100644 index 00000000000..791b0efbaa9 --- /dev/null +++ b/examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml @@ -0,0 +1,173 @@ +# Qwen-Image DMD2 — canonical real-data training config. +# +# Enables the full Qwen-Image DMD2 setup: 4-step student, CFG, and the GAN + R1 +# branch. This is the real-data training config; the mock-data wiring smoke +# (no dataset required) lives in ``dmd2_qwen_image_smoke.yaml``. +# +# 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 \ +# --step_scheduler.max_steps=5000 +# +# The data.* and checkpoint.* paths below are placeholders — point them at your +# own preprocessed latent cache and output directory before launching. + +seed: 42 + +wandb: + project: fastgen-dmd2-qwen-image + mode: online + name: qwen_image_dmd2 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + pretrained_model_name_or_path: Qwen/Qwen-Image + mode: finetune + +step_scheduler: + # Must be divisible by local_batch_size * dp_size. For the canonical + # 32-node GB200 run below, dp_size=128 and local_batch_size=1, so GBS=128 + # gives one micro-batch per rank and no gradient accumulation. + global_batch_size: 128 + local_batch_size: 1 + ckpt_every_steps: 500 + # With 40K cached samples and GBS=128, one epoch is ~313 optimizer steps. + # 16 epochs gives >=5000 optimizer steps; max_steps below stops exactly at 5000. + num_epochs: 16 + log_every: 1 + # Production placeholder; CLI override expected for the long run. + max_steps: 5000 + +# ─── DMD2 block ───────────────────────────────────────────────────────────────── +# ``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 +# consulted for anything set here. Pydantic's ``model_copy(update=...)`` is +# a shallow merge, so when overriding any ``sample_t_cfg`` or ``ema`` +# sub-field we re-list every field in that block to avoid silent drops. +dmd2: + recipe_path: general/distillation/dmd2_qwen_image + pipeline_plugin: qwen_image + qwen_image_guidance: + + # ── DMD2 method core ── + pred_type: flow # Qwen-Image is rectified flow + num_train_timesteps: # continuous t ∈ [0, 1]; QwenImageDMDPipeline forwards t verbatim + guidance_scale: 4.0 # CFG strength on teacher during the student update + student_sample_steps: 4 # 4-step student (Phase 2) + student_sample_type: ode # Euler integration when unrolling the student + # Default keeps FastGen Qwen parity: train each rung from noised real latents. + # Override with ``--dmd2.backward_simulation=true`` to no-grad unroll the + # current student from the first rung before training the selected rung. + backward_simulation: false + student_update_freq: 5 # one student step per 5 fake-score / discriminator steps + fake_score_pred_type: x0 # fake_score regresses x0; teacher/student live in flow space + + # ── GAN branch (Phase 2) ── + gan_loss_weight_gen: 0.03 # generator weight in student loss + gan_use_same_t_noise: true # share (t, noise) between generator + discriminator updates + gan_r1_reg_weight: 0.1 # R1 gradient penalty on real samples + gan_r1_reg_alpha: 0.1 # R1 EMA coefficient + + # ── Optimizer LRs (student LR comes from ``optim.learning_rate`` below) ── + fake_score_lr: 2.0e-6 + discriminator_lr: 2.0e-6 + + # ── GAN discriminator placement & dim ── + gan_feature_indices: [30] # tap transformer_blocks[30] for the feature head + gan_num_blocks: 60 # Qwen-Image has 60 transformer blocks total + gan_inner_dim: 3072 # Qwen-Image hidden_size (matches FastGen reference) + + # ── 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 + # 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% + # train↔inference skew on each non-endpoint timestep. + sample_t_cfg: + # ``time_dist_type`` governs the *perturbation* timestep ``t`` that gets + # sampled on every loss path — VSD perturbation in compute_student_loss + # (dmd.py:417), fake-score DSM perturbation in compute_fake_score_loss + # (dmd.py:529), and GAN/discriminator perturbation in + # compute_discriminator_loss (dmd.py:605). All three call + # ``self.sample_timesteps`` → ``sample_t`` → reads ``time_dist_type``. + # + # It does NOT govern the student's *starting* timestep ``t_student``: + # under student_sample_steps > 1, ``_build_student_input`` calls + # ``sample_from_t_list`` (dmd.py:346) which samples uniformly from + # ``t_list[:-1]`` regardless of ``time_dist_type``. + # + # ``uniform`` matches FastGen's reference Qwen DMD2 config. Earlier formal + # runs used ``logitnormal`` (concentrates ``t`` toward the middle of + # [min_t, max_t]) to follow Automodel's Qwen finetune reference; we now + # default to uniform here to keep launchers parity-aligned without an + # EXTRA_ARGS override. Flip to ``logitnormal`` on the CLI for an ablation. + time_dist_type: uniform + min_t: 0.001 + max_t: 0.999 + p_mean: 0.0 + p_std: 1.0 + t_list: [0.999, 0.74925, 0.4995, 0.24975, 0.0] + + # ── Student EMA ── + ema: + decay: 0.9999 + type: constant + start_iter: 0 + fsdp2: true + mode: full_tensor + +optim: + learning_rate: 2.0e-6 + optimizer: + weight_decay: 0.01 + betas: [0.9, 0.999] + +lr_scheduler: + lr_decay_style: constant + lr_warmup_steps: 0 + min_lr: 2.0e-6 + +fsdp: + tp_size: 1 + cp_size: 1 + pp_size: 1 + dp_replicate_size: 1 + dp_size: 128 # 32 nodes × 4 GPUs/node = 128 GPUs + activation_checkpointing: true + +# Real cached latents (40K items, 1024x1024). Negative-prompt embedding for +# CFG is loaded inside the dataloader via ``negative_prompt_embedding_path``. +data: + dataloader: + _target_: nemo_automodel.components.datasets.diffusion.build_text_to_image_multiresolution_dataloader + cache_dir: /path/to/preprocessed/qwen_image_1024p + base_resolution: [1024, 1024] + batch_size: 1 + drop_last: false + shuffle: true + num_workers: 0 + negative_prompt_embedding_path: /path/to/preprocessed/qwen_image_1024p/negative_prompt_embedding.pt + +# Inference-loadable safetensors saves so checkpoints are usable without a +# secondary export pass. +checkpoint: + enabled: true + checkpoint_dir: /path/to/output/qwen_image_dmd2/checkpoints + model_save_format: safetensors + save_consolidated: true + v4_compatible: true + diffusers_compatible: true + # ``LATEST`` auto-resumes from the most recent checkpoint in checkpoint_dir + # if one exists, and starts fresh otherwise (safe on first launch). To pin a + # specific checkpoint, override on the CLI: + # --checkpoint.restore_from=epoch_0_step_500 + restore_from: LATEST diff --git a/examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml b/examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml new file mode 100644 index 00000000000..2d8f2ad3581 --- /dev/null +++ b/examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml @@ -0,0 +1,109 @@ +# DMD2 on Qwen-Image — mock-data wiring smoke (NOT for real training). +# +# Feeds random tensors at Qwen-Image's shapes/dtypes via +# ``build_mock_t2i_dataloader``. Useful for end-to-end wiring tests (FSDP2, +# phase routing, checkpoint save/restore) but useless for image quality — real +# training uses ``dmd2_qwen_image.yaml`` (real cache + CFG + GAN). + +seed: 42 + +wandb: + project: fastgen-dmd2-qwen-image + mode: online + name: phase1_smoke + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + # Qwen-Image text-to-image checkpoint (HF id, or a local snapshot path to + # avoid hitting HF on every job). + pretrained_model_name_or_path: Qwen/Qwen-Image + mode: finetune + +step_scheduler: + global_batch_size: 8 + local_batch_size: 1 + ckpt_every_steps: 100 + num_epochs: 1 + log_every: 1 + # Hard cap the Phase 1 smoke at 100 optimizer steps. Flip to null for full runs. + max_steps: 100 + +# ─── DMD2-specific block ──────────────────────────────────────────────────────────── +dmd2: + # Built-in fastgen recipe for Qwen-Image: pred_type=flow, num_train_timesteps=null + # (Qwen normalises t internally), logit_normal time sampling, student_update_freq=5, + # fake_score_pred_type=x0, gan_loss_weight_gen=0 (Phase 1), EMA on. + recipe_path: general/distillation/dmd2_qwen_image + + # Phase 1 overrides: + # * GAN disabled — no discriminator shipped in Phase 1. + # * CFG disabled — no negative-prompt embedding precompute yet; guidance_scale=null + # short-circuits the negative-conditioning branch inside compute_student_loss. + gan_loss_weight_gen: 0.0 + guidance_scale: + + # Phase-1-only knobs (NOT on DMDConfig — consumed directly by the recipe): + # LR for the fake-score AdamW; matches the student LR below. + fake_score_lr: 1.0e-5 + + # Explicit pipeline plugin selector. Auto-detect via model_id substring works for a + # local Qwen-Image snapshot path, but spelling it out keeps the choice visible. + pipeline_plugin: qwen_image + + # Optional guidance scalar forwarded to the transformer's ``guidance`` kwarg every + # call. The shipped ``Qwen/Qwen-Image`` checkpoint has guidance_embeds=false, so + # leave this null. Set to e.g. 3.5 only if you've fine-tuned a guidance-embed + # variant. + qwen_image_guidance: + +# Student LR + optimizer. +optim: + learning_rate: 1.0e-5 + optimizer: + weight_decay: 0.01 + betas: [0.9, 0.999] + +# Constant LR for the smoke — flip to cosine/linear once the loop is validated. +lr_scheduler: + lr_decay_style: constant + lr_warmup_steps: 0 + min_lr: 1.0e-5 + +# FSDP2 config. Scale dp_size to match the GPU count of the run. +fsdp: + tp_size: 1 + cp_size: 1 + pp_size: 1 + dp_replicate_size: 1 + dp_size: 8 + activation_checkpointing: true + +# Mock data by default — the debug target is the DMD2 loop, not the data pipeline. +# Swap _target_ to build_text_to_image_multiresolution_dataloader once a real +# preprocessed cache is available. +data: + dataloader: + _target_: nemo_automodel.components.datasets.diffusion.build_mock_t2i_dataloader + # Qwen-Image VAE: 16 latent channels, 8x spatial downsample. + # 256x256 image -> 32x32 latent. Must be even (2x2 patch packing). + num_channels: 16 + spatial_h: 32 + spatial_w: 32 + # Qwen2.5-VL hidden_dim = 3584 (text_encoder/config.json: hidden_size=3584). + text_seq_len: 512 + text_embed_dim: 3584 + length: 256 + num_workers: 0 + shuffle: true + +checkpoint: + enabled: true + checkpoint_dir: /path/to/output/qwen_image_dmd2_smoke/checkpoints + model_save_format: torch_save + save_consolidated: false + diffusers_compatible: false + # Set to LATEST or epoch_0_step_100 (etc.) to resume. Null = start fresh. + restore_from: diff --git a/examples/diffusers/fastgen/dmd2_finetune.py b/examples/diffusers/fastgen/dmd2_finetune.py new file mode 100644 index 00000000000..0f7936ef1bb --- /dev/null +++ b/examples/diffusers/fastgen/dmd2_finetune.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Entrypoint for the DMD2 Qwen-Image AutoModel example. + +Parses the YAML config + CLI overrides with AutoModel's argument parser, then hands +control to :class:`DMD2DiffusionRecipe`. +""" + +from __future__ import annotations + +from dmd2_recipe import DMD2DiffusionRecipe +from nemo_automodel.components.config._arg_parser import parse_args_and_load_config + + +def main( + default_config_path: str = "examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml", +) -> None: + cfg = parse_args_and_load_config(default_config_path) + recipe = DMD2DiffusionRecipe(cfg) + recipe.setup() + recipe.run_train_validation_loop() + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/dmd2_recipe.py b/examples/diffusers/fastgen/dmd2_recipe.py new file mode 100644 index 00000000000..9614d4283a5 --- /dev/null +++ b/examples/diffusers/fastgen/dmd2_recipe.py @@ -0,0 +1,1317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""DMD2 distillation recipe built on NeMo AutoModel. + +This recipe subclasses :class:`nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe` +so it inherits AutoModel's student + optimizer + dataloader + checkpoint plumbing, then +drives ``modelopt.torch.fastgen.DMDPipeline`` (or a plugin subclass) through the +three-phase DMD2 alternation (student update / fake-score update / EMA step). + +Backbone: **Qwen-Image** (``Qwen/Qwen-Image``) — 4D ``image_latents``, +:class:`QwenImageDMDPipeline` handles 2x2 patch packing / img_shapes / +unpacking. Configs: ``configs/dmd2_qwen_image.yaml`` for the canonical +real-data run (4-step + CFG + GAN); ``configs/dmd2_qwen_image_smoke.yaml`` +for the mock-data wiring smoke (no dataset required). + +Launch:: + + # Mock-data wiring smoke (no real cache required). + torchrun --nproc-per-node=8 \\ + examples/diffusers/fastgen/dmd2_finetune.py \\ + --config examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml + # Real-data formal training (canonical). + torchrun --nproc-per-node=8 \\ + examples/diffusers/fastgen/dmd2_finetune.py \\ + --config examples/diffusers/fastgen/configs/dmd2_qwen_image.yaml + +See ``examples/diffusers/fastgen/README.md`` for the three-phase +alternation diagram + troubleshooting notes. +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import shutil +from typing import Any + +import torch +import torch.distributed as dist + +# 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 +# real stack — a previous gate that fell back to ``object`` silently masked missing deps +# and surfaced as a downstream ``TypeError: takes no arguments``. +try: + from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe, is_main_process +except ImportError as exc: + raise ImportError( + "The DMD2 fastgen example requires `nemo_automodel`. Install the example " + "dependencies with:\n" + " pip install -r examples/diffusers/fastgen/requirements.txt" + ) from exc +from torch import nn + +import modelopt.torch.fastgen as mtf +from modelopt.torch.fastgen.config import DMDConfig +from modelopt.torch.fastgen.discriminators import Discriminator_ImageDiT +from modelopt.torch.fastgen.methods.dmd import DMDPipeline +from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin + +# 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 +# ``modelopt_recipes/general/distillation/dmd2_qwen_image.yaml`` file. +_DMD_CONFIG_OVERRIDE_KEYS = frozenset(DMDConfig.model_fields.keys()) + + +def _deep_merge_dicts(base: dict, override: dict) -> dict: + """Recursively merge ``override`` onto ``base`` and return a new dict. + + Nested dicts (e.g. the ``sample_t_cfg`` / ``ema`` sub-configs) are merged key-by-key + rather than replaced wholesale, so a YAML block that overrides a single sub-field + keeps the recipe's other sub-fields instead of silently resetting them to + :class:`DMDConfig` defaults. + """ + merged = dict(base) + for key, value in override.items(): + existing = merged.get(key) + if isinstance(value, dict) and isinstance(existing, dict): + merged[key] = _deep_merge_dicts(existing, value) + else: + merged[key] = value + return merged + + +# Auto-detect substrings (matched case-insensitively against ``model_id``) that map to +# DMDPipeline plugin subclasses. Keep this list small — adding a new entry is only the +# right move when the model has a non-diffusers transformer signature that requires a +# pack/unpack wrapper. Models with the standard ``(hidden_states, timestep, +# encoder_hidden_states)`` signature work with the base :class:`DMDPipeline`. +_PIPELINE_PLUGIN_BY_MODEL_SUBSTR = ( + ("qwen-image", "qwen_image"), + ("qwen_image", "qwen_image"), +) + +_DMD_COMPLETE_MARKER = "dmd2_complete.marker" + + +class DMD2DiffusionRecipe(TrainDiffusionRecipe): + """DMD2 recipe that reuses ``TrainDiffusionRecipe`` for the student path. + + What the superclass gives us (reused unchanged): + + - Student transformer + AdamW optimizer + LR scheduler, loaded via + :class:`NeMoAutoDiffusionPipeline` with FSDP2 sharding. + - ``self.dataloader`` / ``self.sampler`` (swapped to AutoModel's mock dataloader + when ``data.use_mock: true`` — see :meth:`_build_dataloader`). + - ``self.step_scheduler`` (gradient accumulation + checkpoint cadence). + - ``self.checkpointer`` (DCP student weights + optimizer). + - ``self.device`` / ``self.bf16`` / ``self.clip_grad_max_norm`` / etc. + + What this recipe adds: + + - A frozen teacher loaded via a second :meth:`NeMoAutoDiffusionPipeline.from_pretrained` + call with the same ``parallel_scheme`` so it lands with the same FSDP2 sharding + as the student. + - A trainable fake-score transformer loaded the same way (weights identical to the + teacher on step 0). + - A separate AdamW optimizer for the fake-score phase. + - An :class:`mtf.DMDPipeline` driving VSD + DSM + EMA. + - Sidecar checkpoint save / restore for fake-score weights, fake-score optimizer, + EMA shadow, and DMDPipeline iteration counters. + + 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`` and off in the mock-data smoke. See + ``examples/diffusers/fastgen/README.md`` for details. + """ + + # ------------------------------------------------------------------ # + # Setup # + # ------------------------------------------------------------------ # + + def setup(self) -> None: + """Build the student via ``super()``, then add teacher / fake_score / DMDPipeline. + + The extras (``_teacher``, ``_fake_score``, ``_fake_score_optimizer``, + ``_dmd_pipeline``, ``_dmd_config``) are assigned through ``self.__dict__[...]`` + to bypass :meth:`BaseRecipe.__setattr__`'s auto-tracking — otherwise they'd be + added to ``__state_tracked`` and clobber the superclass's single-model / + single-optimizer checkpoint loop. + """ + # 1. Run the parent setup. Builds self.model / self.optimizer / self.lr_scheduler / + # self.dataloader / self.step_scheduler / self.checkpointer / etc. The parent's + # trailing call to self.load_checkpoint(self.restore_from) runs BEFORE our + # extras exist, so it only restores the student — that is intentional and safe. + # + # For the mock-data smoke, ``data.dataloader._target_`` in the YAML points at + # ``nemo_automodel.components.datasets.diffusion.build_mock_dataloader`` so the + # parent wires up the mock dataloader for us — no swap needed. + super().setup() + + # 2. Load the frozen teacher. Same from_pretrained path, same parallel_scheme, but + # ``load_for_training=False`` so the transformer comes back in eval mode with + # requires_grad=False. Bypass __setattr__ to stay invisible to the parent's + # __state_tracked loop. + self.__dict__["_teacher"] = self._load_frozen_teacher() + + # 4. Load the trainable fake-score. Third from_pretrained call — weights start + # identical to the teacher (both come from the same HF checkpoint). + self.__dict__["_fake_score"] = self._load_fake_score() + + # 5. Resolve the DMDConfig: load the fastgen built-in recipe, then apply any + # inline overrides under the YAML ``dmd2:`` block. + self.__dict__["_dmd_config"] = self._resolve_dmd_config() + + # 6. Optimizer for the fake-score phase. LR defaults to student LR when + # ``dmd2.fake_score_lr`` isn't set. + self.__dict__["_fake_score_optimizer"] = self._build_fake_score_optimizer() + + # 7. Optional GAN discriminator. Built when ``gan_loss_weight_gen > 0`` so the + # DMDPipeline constructor's assert is satisfied; otherwise ``discriminator=None`` + # and that assert fires if a YAML enables GAN for an unsupported backbone. + self.__dict__["_discriminator"] = self._build_discriminator() + self.__dict__["_discriminator_optimizer"] = self._build_discriminator_optimizer() + if self._discriminator is not None: + self._attach_gan_feature_capture() + + # 8. DMDPipeline. + # + # Dispatch to a plugin subclass when the backbone needs a non-diffusers call + # signature (e.g. Qwen-Image's packed-latents path). Default = base pipeline. + pipeline_cls = self._resolve_pipeline_cls() + pipeline_kwargs = self._resolve_pipeline_kwargs(pipeline_cls) + self.__dict__["_dmd_pipeline"] = pipeline_cls( + student=self.model, + teacher=self._teacher, + fake_score=self._fake_score, + config=self._dmd_config, + discriminator=self._discriminator, + **pipeline_kwargs, + ) + + # 8. Drop the parent's flow_matching_pipeline — we replace the training loop, + # so keeping it around is pure deadweight. The attribute is not tracked by + # ``__state_tracked`` (FlowMatchingPipeline is a plain class), so ``del`` is + # safe. + if hasattr(self, "flow_matching_pipeline"): + del self.flow_matching_pipeline + + # 9. Extend the student-only restore that super().setup() already ran: also + # restore the fake_score / fake_score_optimizer / EMA / DMD state from the + # same checkpoint directory. + self._restore_dmd_extras(getattr(self, "_dmd2_resolved_restore_from", self.restore_from)) + + if is_main_process(): + logging.info("[DMD2] recipe initialized: %s", self._dmd_config_summary()) + logging.info("[DMD2] full configuration:\n%s", self._dmd_full_config_log()) + + # ------------------------------------------------------------------ # + # Training loop # + # ------------------------------------------------------------------ # + + def run_train_validation_loop(self) -> None: + """Three-phase DMD2 alternation driven by ``step_scheduler``. + + Each outer iteration picks either the student or fake-score phase based on + ``global_step % student_update_freq``. The student phase runs + ``compute_student_loss`` + ``update_ema``. The fake-score phase runs + ``compute_fake_score_loss`` and, when a discriminator is configured + (``gan_loss_weight_gen > 0``), ``compute_discriminator_loss``. + + Mirrors the gating in ``FastGen/fastgen/methods/distribution_matching/dmd2.py`` + (``_student_update_step`` / ``_fake_score_discriminator_update_step``). + """ + dmd = self._dmd_pipeline + cfg = self._dmd_config + + logging.info( + "[DMD2] Starting DMD2 training on %s (pipeline=%s)", + self.model_id, + type(self._dmd_pipeline).__name__, + ) + # Dataloader target (mock vs real cache) is non-obvious from the per-step + # logs; surface it explicitly here so §16's "mock or real dataloader + # target" bullet is checkable from the startup log. + try: + dl_target = type(self.dataloader.dataset).__name__ + logging.info("[DMD2] Dataloader dataset class: %s", dl_target) + except Exception: + pass + logging.info( + "[DMD2] Global batch size: %s; local batch size: %s; DP size: %s", + self.global_batch_size, + self.local_batch_size, + self.dp_size, + ) + logging.info( + "[DMD2] student_update_freq=%d; fake_score_pred_type=%s; guidance_scale=%s;" + " gan_loss_weight_gen=%s", + cfg.student_update_freq, + cfg.fake_score_pred_type, + cfg.guidance_scale, + cfg.gan_loss_weight_gen, + ) + + global_step = int(self.step_scheduler.step) + + for epoch in self.step_scheduler.epochs: + if self.sampler is not None and hasattr(self.sampler, "set_epoch"): + self.sampler.set_epoch(epoch) + + # On resume, the diffusion sampler's load_state_dict primes + # ``_batches_to_skip`` so the next ``__iter__`` skips already-yielded + # batches. Forward that to tqdm's ``initial=`` so the progress bar + # reads e.g. ``187/313`` instead of the misleading ``0/313`` (the + # sampler resets the counter to 0 on the next ``__iter__`` call, + # so reading it here is a one-shot for the resumed epoch only). + tqdm_initial = int(getattr(self.sampler, "_batches_to_skip", 0) or 0) + + if is_main_process(): + from tqdm import tqdm + + self.step_scheduler.dataloader = tqdm( + self.dataloader, + desc=f"Epoch {epoch + 1}/{self.num_epochs}", + initial=tqdm_initial, + ) + else: + self.step_scheduler.dataloader = self.dataloader + + epoch_student_loss = 0.0 + epoch_fake_score_loss = 0.0 + student_steps = 0 + fake_score_steps = 0 + + for batch_group in self.step_scheduler: + is_student_phase = (global_step % cfg.student_update_freq) == 0 + + if is_student_phase: + self.optimizer.zero_grad(set_to_none=True) + else: + self._fake_score_optimizer.zero_grad(set_to_none=True) + + self._set_grad_requirements(is_student_phase) + + micro_losses: list[float] = [] + micro_vsd_losses: list[float] = [] + micro_disc_losses: list[float] = [] + for micro_batch in batch_group: + ( + latents, + noise, + text_embeds, + text_mask, + neg_text_embeds, + neg_text_mask, + ) = self._prepare_micro_batch(micro_batch) + + if is_student_phase: + # ``compute_student_loss`` reads ``guidance_scale`` from the + # DMDConfig when this kwarg is None. We pass the negative + # embedding unconditionally — the function ignores it when + # CFG is disabled, and raises a clear ValueError when CFG + # is enabled but no negative was supplied. + losses = dmd.compute_student_loss( + latents, + noise, + encoder_hidden_states=text_embeds, + encoder_hidden_states_mask=text_mask, + negative_encoder_hidden_states=neg_text_embeds, + negative_encoder_hidden_states_mask=neg_text_mask, + guidance_scale=None, + ) + micro_vsd_losses.append(float(losses["vsd"].item())) + else: + losses = dmd.compute_fake_score_loss( + latents, + noise, + encoder_hidden_states=text_embeds, + encoder_hidden_states_mask=text_mask, + ) + + (losses["total"] / len(batch_group)).backward() + micro_losses.append(float(losses["total"].item())) + + # GAN: in the fake-score phase, also update the discriminator + # on the same batch (FastGen pattern: + # _fake_score_discriminator_update_step). + if ( + not is_student_phase + and self._discriminator is not None + and self._discriminator_optimizer is not None + ): + self._discriminator_optimizer.zero_grad(set_to_none=True) + disc_losses = dmd.compute_discriminator_loss( + latents, + noise, + encoder_hidden_states=text_embeds, + encoder_hidden_states_mask=text_mask, + ) + (disc_losses["total"] / len(batch_group)).backward() + # Manual gradient all-reduce across DP ranks (the + # discriminator is replicated, not FSDP-sharded). + if dist.is_initialized(): + for p in self._discriminator.parameters(): + if p.grad is not None: + dist.all_reduce(p.grad, op=dist.ReduceOp.AVG) + torch.nn.utils.clip_grad_norm_( + self._discriminator.parameters(), + max_norm=self.clip_grad_max_norm, + ) + self._discriminator_optimizer.step() + micro_disc_losses.append(float(disc_losses["total"].item())) + + # Grad clip on whichever module is the active trainable. + if is_student_phase: + grad_norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), max_norm=self.clip_grad_max_norm + ) + else: + grad_norm = torch.nn.utils.clip_grad_norm_( + self._fake_score.parameters(), max_norm=self.clip_grad_max_norm + ) + grad_norm = float(grad_norm) if torch.is_tensor(grad_norm) else grad_norm + + # Step. + if is_student_phase: + self.optimizer.step() + dmd.update_ema() + if self.lr_scheduler is not None: + self.lr_scheduler[0].step(1) + else: + self._fake_score_optimizer.step() + + group_loss_mean = float(sum(micro_losses) / len(micro_losses)) + if is_student_phase: + epoch_student_loss += group_loss_mean + student_steps += 1 + else: + epoch_fake_score_loss += group_loss_mean + fake_score_steps += 1 + + global_step = int(self.step_scheduler.step) + + if ( + self.log_every + and self.log_every > 0 + and is_main_process() + and (global_step % self.log_every == 0) + ): + self._log_step( + global_step=global_step, + is_student_phase=is_student_phase, + group_loss=group_loss_mean, + grad_norm=grad_norm, + vsd_loss=(sum(micro_vsd_losses) / len(micro_vsd_losses)) + if micro_vsd_losses + else None, + disc_loss=(sum(micro_disc_losses) / len(micro_disc_losses)) + if micro_disc_losses + else None, + ) + + if self.step_scheduler.is_ckpt_step: + # Use the group mean of the active phase as the reported train loss. + self.save_checkpoint(epoch, global_step, group_loss_mean) + + # End-of-epoch logging. + if is_main_process(): + avg_student = ( + (epoch_student_loss / student_steps) if student_steps else float("nan") + ) + avg_fake = ( + epoch_fake_score_loss / fake_score_steps if fake_score_steps else float("nan") + ) + logging.info( + "[DMD2] Epoch %d complete. student_avg=%.6f (%d steps) " + "fake_score_avg=%.6f (%d steps)", + epoch + 1, + avg_student, + student_steps, + avg_fake, + fake_score_steps, + ) + + if torch.cuda.is_available(): + peak = torch.cuda.max_memory_allocated() / (1024**3) + reserved = torch.cuda.max_memory_reserved() / (1024**3) + rank = dist.get_rank() if dist.is_initialized() else 0 + logging.info( + "[DMD2] PEAK_MEM rank=%d max_allocated=%.2fGiB max_reserved=%.2fGiB", + rank, + peak, + reserved, + ) + + if is_main_process(): + logging.info("[DMD2] Training complete. Final step: %s", global_step) + + # ------------------------------------------------------------------ # + # Checkpoint save / restore (sidecars next to student DCP) # + # ------------------------------------------------------------------ # + + def load_checkpoint(self, restore_from: str | None = None): + """Load only from checkpoints whose DMD2 sidecars are complete. + + ``TrainDiffusionRecipe.setup()`` calls this before the DMD2-only objects exist, + so this method only resolves the path and delegates the student restore to the + parent. The sidecars are restored later by ``_restore_dmd_extras``. + """ + resolved = self._resolve_complete_dmd_checkpoint(restore_from) + self.__dict__["_dmd2_resolved_restore_from"] = resolved + + if resolved is None: + if ( + restore_from is not None + and str(restore_from).upper() == "LATEST" + and is_main_process() + ): + logging.warning( + "[DMD2] restore_from=LATEST but no complete DMD2 checkpoint was found in %s. " + "Starting fresh.", + self.checkpointer.config.checkpoint_dir, + ) + return + + super().load_checkpoint(resolved) + + def save_checkpoint( + self, + epoch: int, + step: int, + train_loss: float, + val_loss: dict[str, float] | None = None, + best_metric_key: str = "default", + ) -> None: + """Delegate student save to ``super()``, then sidecar the DMD2 extras.""" + # Recover from a partial save from a previous run (e.g. SLURM time + # limit killed the job between super().save_checkpoint() — which writes + # the model + step_scheduler + dataloader — and our DMD2 sidecar + # writes below). The parent's save_checkpoint refuses to overwrite an + # existing directory and raises FileExistsError, so without this we'd + # need a manual cleanup every time a SLURM kill landed mid-save. + path = os.path.join(self.checkpointer.config.checkpoint_dir, f"epoch_{epoch}_step_{step}") + if is_main_process() and self.checkpointer.config.enabled and os.path.exists(path): + if not self._is_dmd_checkpoint_complete(path): + logging.warning( + "[DMD2] cleaning up incomplete checkpoint directory left by a previous run: %s", + path, + ) + shutil.rmtree(path) + if dist.is_initialized(): + dist.barrier() + + previous_complete = None + if self.checkpointer.config.enabled: + previous_complete = self._find_latest_complete_dmd_checkpoint( + self.checkpointer.config.checkpoint_dir + ) + + super().save_checkpoint(epoch, step, train_loss, val_loss, best_metric_key) + + if not self.checkpointer.config.enabled: + return + + # The parent save updates LATEST before DMD2 sidecars exist. Until the marker is + # written below, keep LATEST on the previous complete DMD2 checkpoint. + if is_main_process(): + if previous_complete is not None: + self._update_latest_symlink(previous_complete) + else: + self._remove_checkpoint_pointer("LATEST") + if dist.is_initialized(): + dist.barrier() + + self._save_dmd_extras(path) + + if dist.is_initialized(): + dist.barrier() + + if is_main_process(): + self._write_dmd_complete_marker(path) + self._update_checkpoint_symlink("DMD2_LATEST", path) + self._update_latest_symlink(path) + if dist.is_initialized(): + dist.barrier() + + def _save_dmd_extras(self, path: str) -> None: + """Write fake_score DCP + fake_score_optimizer DCP + ema_shadow.pt + dmd_state.pt.""" + # fake_score weights — DCP sharded save via the same Checkpointer the parent uses + # for the student. Each rank writes its own shard. + fs_weights_dir = os.path.join(path, "fake_score") + os.makedirs(fs_weights_dir, exist_ok=True) + self.checkpointer.save_model( + model=self._fake_score, + weights_path=fs_weights_dir, + peft_config=None, + tokenizer=None, + ) + # fake_score optimizer — also DCP sharded. ``save_optimizer`` takes the optimizer + # and its owning model in order to rebuild the parameter mapping. + fs_opt_dir = os.path.join(path, "fake_score_optimizer") + os.makedirs(fs_opt_dir, exist_ok=True) + self.checkpointer.save_optimizer( + self._fake_score_optimizer, self._fake_score, fs_opt_dir, None + ) + + # EMA shadow + DMD scalar state — rank-0 torch.save. EMA's ``state_dict`` already + # materialises full tensors via ``DTensor.full_tensor()`` under FSDP2 full_tensor + # mode, so this is a single unsharded file. + if is_main_process(): + logging.info("[DMD2] saved fake_score weights -> %s", fs_weights_dir) + logging.info("[DMD2] saved fake_score optimizer -> %s", fs_opt_dir) + if self._dmd_pipeline.ema is not None: + ema_path = os.path.join(path, "ema_shadow.pt") + torch.save(self._dmd_pipeline.ema.state_dict(), ema_path) + logging.info("[DMD2] saved ema_shadow -> %s", ema_path) + state_path = os.path.join(path, "dmd_state.pt") + torch.save({"iteration": self._dmd_pipeline._iteration}, state_path) + logging.info( + "[DMD2] saved dmd_state (iteration=%d) -> %s", + int(self._dmd_pipeline._iteration), + state_path, + ) + # Discriminator + its optimizer — replicated across ranks (no FSDP), + # so rank-0 torch.save of the canonical state_dict suffices. + if self._discriminator is not None: + disc_path = os.path.join(path, "discriminator.pt") + torch.save(self._discriminator.state_dict(), disc_path) + logging.info("[DMD2] saved discriminator -> %s", disc_path) + if self._discriminator_optimizer is not None: + disc_opt_path = os.path.join(path, "discriminator_optimizer.pt") + torch.save(self._discriminator_optimizer.state_dict(), disc_opt_path) + logging.info("[DMD2] saved discriminator optimizer -> %s", disc_opt_path) + + def _write_dmd_complete_marker(self, path: str) -> None: + marker_path = os.path.join(path, _DMD_COMPLETE_MARKER) + payload = { + "checkpoint": os.path.basename(os.path.realpath(path)), + "dmd_iteration": int(self._dmd_pipeline._iteration), + } + with open(marker_path, "w") as f: + json.dump(payload, f) + f.write("\n") + logging.info("[DMD2] marked checkpoint complete -> %s", marker_path) + + def _remove_checkpoint_pointer(self, link_name: str) -> None: + ckpt_root = self.checkpointer.config.checkpoint_dir + for path in ( + os.path.join(ckpt_root, link_name), + os.path.join(ckpt_root, f"{link_name}.txt"), + ): + if os.path.lexists(path): + os.remove(path) + + def _restore_dmd_extras(self, restore_from: str | None) -> None: + """Restore fake_score + fake_score optimizer + EMA + DMD scalar state. + + No-op when no checkpoint is being restored. ``load_checkpoint`` resolves + ``LATEST`` to the latest complete DMD2 checkpoint before this method runs. + """ + if restore_from is None: + return + + ckpt_dir = self._resolve_extras_dir(restore_from) + if ckpt_dir is None or not os.path.isdir(ckpt_dir): + return + + fs_weights_dir = os.path.join(ckpt_dir, "fake_score") + fs_opt_dir = os.path.join(ckpt_dir, "fake_score_optimizer") + ema_path = os.path.join(ckpt_dir, "ema_shadow.pt") + state_path = os.path.join(ckpt_dir, "dmd_state.pt") + + # Checkpointer.save_model writes DCP shards to ``/model/``; + # load_model expects that *inner* ``model/`` dir as ``model_path`` (see + # ``BaseRecipe.load_checkpoint`` which passes ``os.path.join(ckpt_dir, "model")``). + # The kwarg name differs between save (``weights_path``) and load (``model_path``). + fs_weights_model_dir = os.path.join(fs_weights_dir, "model") + if os.path.isdir(fs_weights_model_dir): + self.checkpointer.load_model(model=self._fake_score, model_path=fs_weights_model_dir) + if is_main_process(): + logging.info("[DMD2] restored fake_score weights <- %s", fs_weights_model_dir) + elif is_main_process(): + logging.info( + "[DMD2] WARN: fake_score weights dir missing at %s -- skipping", + fs_weights_model_dir, + ) + # load_optimizer, in contrast, appends ``optim/`` internally — pass the base dir. + if os.path.isdir(os.path.join(fs_opt_dir, "optim")): + self.checkpointer.load_optimizer( + self._fake_score_optimizer, self._fake_score, fs_opt_dir, None + ) + if is_main_process(): + logging.info("[DMD2] restored fake_score optimizer <- %s", fs_opt_dir) + elif is_main_process(): + logging.info( + "[DMD2] WARN: fake_score optimizer dir missing at %s -- skipping", + fs_opt_dir, + ) + + if os.path.isfile(ema_path) and self._dmd_pipeline.ema is not None: + ema_state = torch.load(ema_path, map_location="cpu", weights_only=False) + self._dmd_pipeline.ema.load_state_dict(ema_state) + if is_main_process(): + logging.info("[DMD2] restored ema_shadow <- %s", ema_path) + if os.path.isfile(state_path): + state = torch.load(state_path, map_location="cpu", weights_only=False) + self._dmd_pipeline._iteration = int(state.get("iteration", 0)) + if is_main_process(): + logging.info( + "[DMD2] restored dmd_state (iteration=%d) <- %s", + self._dmd_pipeline._iteration, + state_path, + ) + + # Discriminator + its optimizer. + if self._discriminator is not None: + disc_path = os.path.join(ckpt_dir, "discriminator.pt") + if os.path.isfile(disc_path): + disc_state = torch.load(disc_path, map_location="cpu", weights_only=False) + self._discriminator.load_state_dict(disc_state) + if is_main_process(): + logging.info("[DMD2] restored discriminator <- %s", disc_path) + elif is_main_process(): + logging.info("[DMD2] WARN: discriminator file missing at %s -- skipping", disc_path) + if self._discriminator_optimizer is not None: + disc_opt_path = os.path.join(ckpt_dir, "discriminator_optimizer.pt") + if os.path.isfile(disc_opt_path): + disc_opt_state = torch.load(disc_opt_path, map_location="cpu", weights_only=False) + self._discriminator_optimizer.load_state_dict(disc_opt_state) + if is_main_process(): + logging.info("[DMD2] restored discriminator optimizer <- %s", disc_opt_path) + elif is_main_process(): + logging.info( + "[DMD2] WARN: discriminator optimizer file missing at %s -- skipping", + disc_opt_path, + ) + + def _resolve_extras_dir(self, restore_from: str) -> str | None: + """Best-effort resolve of the checkpoint dir, matching BaseRecipe's convention. + + For explicit paths we pass through; for ``"LATEST"`` we look under + ``checkpointer.config.checkpoint_dir``. This keeps resolution simple and delegates + the hard cases (async symlinks, cross-node shared filesystems) to the user. + """ + if os.path.isabs(restore_from): + return restore_from + # Try the checkpoint_dir-relative form first (matches the parent's symlink + # naming — "LATEST" or an explicit ``epoch_N_step_M`` subdir). + candidate = os.path.join(self.checkpointer.config.checkpoint_dir, restore_from) + if os.path.exists(candidate): + return os.path.realpath(candidate) + return None + + def _resolve_complete_dmd_checkpoint(self, restore_from: str | None) -> str | None: + ckpt_root = self.checkpointer.config.checkpoint_dir + + if restore_from is None or str(restore_from).upper() in {"LATEST", "DMD2_LATEST"}: + return self._find_latest_complete_dmd_checkpoint(ckpt_root) + + if os.path.isabs(restore_from): + candidate = restore_from + else: + candidate = os.path.join(ckpt_root, restore_from) + candidate = os.path.realpath(candidate) + + if not os.path.isdir(candidate): + return candidate + if not self._is_dmd_checkpoint_complete(candidate): + raise RuntimeError( + f"DMD2 checkpoint is incomplete and cannot be restored: {candidate}. " + "Use a complete older checkpoint or remove the partial directory." + ) + return candidate + + def _find_latest_complete_dmd_checkpoint(self, ckpt_root: str) -> str | None: + dmd2_latest = os.path.join(ckpt_root, "DMD2_LATEST") + for pointer in (dmd2_latest, os.path.join(ckpt_root, "LATEST")): + resolved = self._resolve_checkpoint_pointer(pointer) + if resolved is not None and self._is_dmd_checkpoint_complete(resolved): + return resolved + + candidates = [] + if os.path.isdir(ckpt_root): + for name in os.listdir(ckpt_root): + path = os.path.join(ckpt_root, name) + if ( + os.path.isdir(path) + and "_step_" in name + and self._is_dmd_checkpoint_complete(path) + ): + candidates.append(os.path.realpath(path)) + if not candidates: + return None + return max(candidates, key=self._checkpoint_step) + + def _resolve_checkpoint_pointer(self, pointer: str) -> str | None: + resolved = None + if os.path.islink(pointer): + try: + resolved = os.readlink(pointer) + except OSError: + return None + elif os.path.isfile(pointer + ".txt"): + try: + with open(pointer + ".txt") as f: + resolved = f.read().strip() + except OSError: + return None + if not resolved: + return None + if not os.path.isabs(resolved): + resolved = os.path.abspath(os.path.join(os.path.dirname(pointer), resolved)) + return os.path.realpath(resolved) if os.path.isdir(resolved) else None + + def _is_dmd_checkpoint_complete(self, path: str) -> bool: + path = os.path.realpath(path) + if not os.path.isdir(path): + return False + if os.path.isfile(os.path.join(path, _DMD_COMPLETE_MARKER)): + return True + + fs_model_dir = os.path.join(path, "fake_score", "model") + fs_opt_metadata = os.path.join(path, "fake_score_optimizer", "optim", ".metadata") + dmd_state = os.path.join(path, "dmd_state.pt") + complete = ( + self._dir_has_regular_file(fs_model_dir) + and os.path.isfile(fs_opt_metadata) + and os.path.isfile(dmd_state) + ) + if not complete: + return False + + if self._cfg_gan_enabled(): + return os.path.isfile(os.path.join(path, "discriminator.pt")) and os.path.isfile( + os.path.join(path, "discriminator_optimizer.pt") + ) + return True + + def _cfg_gan_enabled(self) -> bool: + cfg = getattr(self, "cfg", None) + if cfg is None: + return False + try: + return float(cfg.get("dmd2.gan_loss_weight_gen", 0.0) or 0.0) > 0 + except (TypeError, ValueError): + return False + + @staticmethod + def _dir_has_regular_file(path: str) -> bool: + if not os.path.isdir(path): + return False + try: + with os.scandir(path) as entries: + return any(entry.is_file() for entry in entries) + except OSError: + return False + + @staticmethod + def _checkpoint_step(path: str) -> int: + match = re.search(r"_step_(\d+)$", os.path.basename(os.path.realpath(path))) + return int(match.group(1)) if match else -1 + + # ------------------------------------------------------------------ # + # Helpers — teacher / fake_score loading, DMDConfig resolution # + # ------------------------------------------------------------------ # + + def _load_frozen_teacher(self) -> nn.Module: + """Load a second copy of the pretrained transformer, frozen + FSDP2-sharded. + + The same pretrained path + ``parallel_scheme`` as the student. Setting + ``load_for_training=False`` walks the parameters once and flips + ``requires_grad=False`` after FSDP2 wrapping; we also call ``.eval()`` on the + returned module just to be defensive. + """ + parallel_scheme = self._build_parallel_scheme_snapshot() + pipe, _ = NeMoAutoDiffusionPipeline.from_pretrained( + self.model_id, + torch_dtype=self.bf16, + device=self.device, + parallel_scheme=parallel_scheme, + components_to_load=["transformer"], + load_for_training=False, + low_cpu_mem_usage=True, + ) + teacher = pipe.transformer + teacher.eval() + for p in teacher.parameters(): + p.requires_grad_(False) + return teacher + + def _build_discriminator(self) -> nn.Module | None: + """Construct the Discriminator_ImageDiT when GAN is enabled. + + Returns ``None`` when ``dmd2.gan_loss_weight_gen`` is zero so the + DMDPipeline runs without a discriminator (any run with the GAN branch disabled). + """ + gan_weight = float(self.cfg.get("dmd2.gan_loss_weight_gen", 0.0) or 0.0) + if gan_weight <= 0.0: + return None + + # GAN-specific knobs read directly from the YAML so callers don't have to + # touch the built-in DMDConfig recipe just to flip feature indices. + feature_indices = self.cfg.get( + "dmd2.gan_feature_indices", [30] + ) # middle of Qwen-Image's 60 blocks + num_blocks = int(self.cfg.get("dmd2.gan_num_blocks", 60)) + inner_dim = int(self.cfg.get("dmd2.gan_inner_dim", 3072)) + + disc = Discriminator_ImageDiT( + feature_indices={int(i) for i in feature_indices}, + num_blocks=num_blocks, + inner_dim=inner_dim, + ) + disc.to(device=self.device, dtype=self.bf16) + disc.train() + for p in disc.parameters(): + p.requires_grad_(True) + if is_main_process(): + logging.info( + "[DMD2] Built discriminator: %s | num_features=%d num_blocks=%d inner_dim=%d " + "params=%d", + type(disc).__name__, + disc.num_features, + num_blocks, + inner_dim, + sum(p.numel() for p in disc.parameters()), + ) + return disc + + def _build_discriminator_optimizer(self) -> torch.optim.Optimizer | None: + """AdamW on the discriminator. No FSDP wrap — manual grad all-reduce keeps it simple.""" + if self._discriminator is None: + return None + lr = float(self.cfg.get("dmd2.discriminator_lr", 1.0e-5) or 1.0e-5) + opt = torch.optim.AdamW( + self._discriminator.parameters(), + lr=lr, + weight_decay=0.01, + betas=(0.9, 0.999), # FastGen Qwen-Image DMD2 inherits BaseOptimizerConfig betas + ) + if is_main_process(): + logging.info("[DMD2] Built discriminator optimizer: AdamW lr=%g betas=(0.9, 0.999)", lr) + return opt + + def _attach_gan_feature_capture(self) -> None: + """Install Qwen-Image feature-capture hooks on the teacher when GAN is enabled. + + Reads the latent resolution from the dataloader so the hook can reshape + ``[B, num_image_patches, 3072]`` into ``[B, 3072, H_lat//2, W_lat//2]``. + Mock dataloader → spatial_h/spatial_w from the YAML. Real dataloader → + base_resolution / vae_scale. + """ + feature_indices = list(self.cfg.get("dmd2.gan_feature_indices", [30])) + + # Resolve h_lat / w_lat. Mock has it in the YAML; real cache uses the + # configured base_resolution divided by the VAE 8x downsample. + # Convert the dataloader subtree to a plain dict — AutoModel's ConfigNode + # doesn't expose deep dotted paths like ``data.dataloader.spatial_h``. + dl_node = self.cfg.get("data.dataloader", None) + if dl_node is not None and hasattr(dl_node, "to_dict"): + dl_dict = dl_node.to_dict() + elif dl_node is not None: + try: + dl_dict = dict(dl_node) + except (TypeError, ValueError): + dl_dict = {} + else: + dl_dict = {} + + spatial_h = dl_dict.get("spatial_h") + spatial_w = dl_dict.get("spatial_w") + base_resolution = dl_dict.get("base_resolution") + if spatial_h is not None and spatial_w is not None: + h_lat = int(spatial_h) + w_lat = int(spatial_w) + elif base_resolution is not None: + h_lat = int(base_resolution[0]) // 8 + w_lat = int(base_resolution[1]) // 8 + else: + # Fallback: hope it's 64x64 (512px image). Smoke tests pin this explicitly. + h_lat, w_lat = 64, 64 + if is_main_process(): + logging.warning( + "[DMD2] Could not infer h_lat/w_lat from data.dataloader; defaulting to 64x64." + ) + + qwen_image_plugin.attach_feature_capture( + self._teacher, + feature_indices=feature_indices, + h_lat=h_lat, + w_lat=w_lat, + ) + if is_main_process(): + logging.info( + "[DMD2] Attached GAN feature capture: indices=%s h_lat=%d w_lat=%d", + feature_indices, + h_lat, + w_lat, + ) + + def _load_fake_score(self) -> nn.Module: + """Load a third copy, trainable. Weights start identical to the teacher.""" + parallel_scheme = self._build_parallel_scheme_snapshot() + pipe, _ = NeMoAutoDiffusionPipeline.from_pretrained( + self.model_id, + torch_dtype=self.bf16, + device=self.device, + parallel_scheme=parallel_scheme, + components_to_load=["transformer"], + load_for_training=True, + low_cpu_mem_usage=True, + ) + fake_score = pipe.transformer + fake_score.train() + for p in fake_score.parameters(): + p.requires_grad_(True) + return fake_score + + def _build_parallel_scheme_snapshot(self) -> dict[str, dict[str, Any]]: + """Reconstruct the FSDP2 manager_args used for the student. + + Mirrors ``build_model_and_optimizer`` in ``nemo_automodel.recipes.diffusion.train``. + We can't capture the student's ``parallel_scheme`` directly (the parent doesn't + stash it), so we rebuild it from the same YAML knobs the parent consumed. + """ + from torch.distributed.fsdp import MixedPrecisionPolicy + + fsdp_cfg = self.cfg.get("fsdp", None) or {} + ddp_cfg = self.cfg.get("ddp", None) + + world_size = dist.get_world_size() if dist.is_initialized() else 1 + + if ddp_cfg is not None: + return { + "transformer": { + "_manager_type": "ddp", + "backend": ddp_cfg.get("backend", "nccl"), + "world_size": world_size, + "activation_checkpointing": ddp_cfg.get("activation_checkpointing", False), + } + } + + dp_size = fsdp_cfg.get("dp_size") + tp_size = fsdp_cfg.get("tp_size", 1) + cp_size = fsdp_cfg.get("cp_size", 1) + pp_size = fsdp_cfg.get("pp_size", 1) + if dp_size is None: + denom = max(1, tp_size * cp_size * pp_size) + dp_size = max(1, world_size // denom) + + return { + "transformer": { + "_manager_type": "fsdp2", + "dp_size": dp_size, + "dp_replicate_size": fsdp_cfg.get("dp_replicate_size", None), + "tp_size": tp_size, + "cp_size": cp_size, + "pp_size": pp_size, + "backend": "nccl", + "world_size": world_size, + "use_hf_tp_plan": fsdp_cfg.get("use_hf_tp_plan", False), + "activation_checkpointing": fsdp_cfg.get("activation_checkpointing", True), + "mp_policy": MixedPrecisionPolicy( + param_dtype=self.bf16, + reduce_dtype=torch.float32, + output_dtype=self.bf16, + ), + } + } + + def _resolve_pipeline_cls(self) -> type[DMDPipeline]: + """Pick the DMDPipeline subclass for the current backbone. + + Resolution order: + + 1. ``dmd2.pipeline_plugin`` in the YAML (explicit override, ``null`` for base). + 2. Substring match on ``model.pretrained_model_name_or_path`` + (e.g. ``Qwen-Image`` -> ``qwen_image`` plugin). + 3. Fall back to :class:`DMDPipeline`. + """ + explicit = self.cfg.get("dmd2.pipeline_plugin", None) + if explicit is None: + model_id_lc = (self.model_id or "").lower() + for needle, plugin_name in _PIPELINE_PLUGIN_BY_MODEL_SUBSTR: + if needle in model_id_lc: + explicit = plugin_name + break + if explicit in (None, "base", "DMDPipeline"): + return DMDPipeline + if explicit == "qwen_image": + # Imported lazily so ``base`` users don't pay the import cost. + from modelopt.torch.fastgen.plugins.qwen_image import QwenImageDMDPipeline + + return QwenImageDMDPipeline + raise ValueError( + f"Unknown dmd2.pipeline_plugin={explicit!r}. Supported: null/'base', 'qwen_image'." + ) + + def _resolve_pipeline_kwargs(self, pipeline_cls: type[DMDPipeline]) -> dict[str, Any]: + """Extra kwargs to forward to the pipeline subclass constructor (plugin-specific).""" + if pipeline_cls.__name__ == "QwenImageDMDPipeline": + # Optional ``guidance`` value passed to the transformer's guidance kwarg every + # call. Independent of DMDConfig.guidance_scale (which drives the negative- + # prompt CFG path on the teacher). Leave ``None`` to skip the embedding when + # the transformer was built with ``guidance_embeds=false`` (default for + # ``Qwen/Qwen-Image``). + return {"guidance": self.cfg.get("dmd2.qwen_image_guidance", None)} + return {} + + def _resolve_dmd_config(self) -> DMDConfig: + """Load the built-in fastgen recipe, then apply inline YAML overrides.""" + dmd_cfg_node = self.cfg.get("dmd2", None) + if dmd_cfg_node is None: + raise ValueError( + "Missing ``dmd2:`` block in the YAML config. Expected at minimum " + "``dmd2.recipe_path`` pointing at a fastgen DMDConfig recipe " + "(e.g. ``general/distillation/dmd2_qwen_image``)." + ) + dmd_dict = ( + dmd_cfg_node.to_dict() if hasattr(dmd_cfg_node, "to_dict") else dict(dmd_cfg_node) + ) + + recipe_path = dmd_dict.pop("recipe_path", None) + if recipe_path is None: + raise ValueError( + "``dmd2.recipe_path`` is required — Phase 1 relies on the built-in " + "``modelopt_recipes`` path resolver to hydrate the full DMDConfig." + ) + base_config = mtf.load_dmd_config(recipe_path) + + # Filter overrides to the subset that actually corresponds to DMDConfig fields. + # Non-matching keys (e.g. ``fake_score_lr``, ``cfg_mode``) are kept as top-level + # recipe knobs and read via ``self.cfg.get("dmd2.")``. + overrides = {k: v for k, v in dmd_dict.items() if k in _DMD_CONFIG_OVERRIDE_KEYS} + if not overrides: + return base_config + # Deep-merge so a YAML block that overrides a single ``sample_t_cfg`` / ``ema`` + # sub-field keeps the recipe's other sub-fields — a shallow ``dict.update`` would + # replace the whole sub-config and silently reset its siblings to defaults. + # Re-validate the merged dict so the nested blocks become their Pydantic config + # objects instead of raw dicts. + merged = _deep_merge_dicts(base_config.model_dump(), overrides) + return DMDConfig.model_validate(merged) + + def _build_fake_score_optimizer(self) -> torch.optim.Optimizer: + """AdamW on fake_score params. LR defaults to student LR; overridable via YAML.""" + fs_lr = self.cfg.get("dmd2.fake_score_lr", None) + if fs_lr is None: + fs_lr = self.learning_rate + optimizer_cfg = self.cfg.get("optim.optimizer", {}) or {} + optimizer_cfg = ( + optimizer_cfg.to_dict() if hasattr(optimizer_cfg, "to_dict") else dict(optimizer_cfg) + ) + weight_decay = optimizer_cfg.get("weight_decay", 0.01) + betas = tuple(optimizer_cfg.get("betas", (0.9, 0.999))) + + trainable_params = [p for p in self._fake_score.parameters() if p.requires_grad] + if not trainable_params: + raise RuntimeError("No trainable parameters found in fake_score.") + return torch.optim.AdamW(trainable_params, lr=fs_lr, weight_decay=weight_decay, betas=betas) + + # ------------------------------------------------------------------ # + # Inner helpers # + # ------------------------------------------------------------------ # + + def _set_grad_requirements(self, is_student_phase: bool) -> None: + """Toggle train/eval + requires_grad across modules for the active phase. + + Mirrors FastGen's ``_setup_grad_requirements`` (``dmd2.py`` lines 67-77), + INCLUDING the discriminator toggle that was previously omitted. + + Why the discriminator toggle matters: ``compute_student_loss`` calls + ``self.discriminator(fake_feat)`` for the ``gan_gen`` term, so the + discriminator is in the student-phase backward graph. With its params + left at ``requires_grad=True``, ``total.backward()`` allocates and + fills ``.grad`` for every discriminator parameter — gradients which + the student optimizer never consumes and which the next discriminator + ``zero_grad(set_to_none=True)`` simply wipes. Freezing the discriminator + during the student phase skips that wasted memory + backward compute + without changing any numerics (the student still receives the GAN + signal through the discriminator's input-side gradient, which doesn't + require the discriminator's own params to be in the autograd graph). + + Called every step; cheap enough that we don't bother caching the last state. + """ + if is_student_phase: + self.model.train() + for p in self.model.parameters(): + p.requires_grad_(True) + self._fake_score.eval() + for p in self._fake_score.parameters(): + p.requires_grad_(False) + if self._discriminator is not None: + self._discriminator.eval() + for p in self._discriminator.parameters(): + p.requires_grad_(False) + else: + self.model.eval() + for p in self.model.parameters(): + p.requires_grad_(False) + self._fake_score.train() + for p in self._fake_score.parameters(): + p.requires_grad_(True) + if self._discriminator is not None: + self._discriminator.train() + for p in self._discriminator.parameters(): + p.requires_grad_(True) + + def _prepare_micro_batch( + self, micro_batch: dict[str, Any] + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + """Extract latents, noise, text conditioning, and optional masks from a batch. + + Accepts both 5D ``video_latents`` and 4D ``image_latents`` + (Qwen-Image / Flux / SD3). Mirrors the key dispatch in + ``nemo_automodel.components.flow_matching.pipeline.FlowMatchingPipeline.step``. + + ``negative_text_embeddings`` is optional — present when the dataloader + supplies it (mock T2I, real cache with precomputed empty-prompt + embedding) and consumed by ``compute_student_loss`` only when CFG is + enabled (``dmd2.guidance_scale is not None``). + """ + if "image_latents" in micro_batch: + latents = micro_batch["image_latents"].to(self.device, dtype=self.bf16) + elif "video_latents" in micro_batch: + latents = micro_batch["video_latents"].to(self.device, dtype=self.bf16) + else: + raise KeyError( + "Batch must contain either 'image_latents' (4D) or 'video_latents' (5D). " + f"Got keys: {sorted(micro_batch.keys())}." + ) + text_embeds = micro_batch["text_embeddings"].to(self.device, dtype=self.bf16) + if text_embeds.ndim == 2: + text_embeds = text_embeds.unsqueeze(0) + text_mask = micro_batch.get("text_embeddings_mask") + if text_mask is not None: + text_mask = text_mask.to(self.device) + if text_mask.ndim == 1: + text_mask = text_mask.unsqueeze(0) + negative_text_embeds = micro_batch.get("negative_text_embeddings") + if negative_text_embeds is not None: + negative_text_embeds = negative_text_embeds.to(self.device, dtype=self.bf16) + if negative_text_embeds.ndim == 2: + negative_text_embeds = negative_text_embeds.unsqueeze(0) + negative_text_mask = micro_batch.get("negative_text_embeddings_mask") + if negative_text_mask is not None: + negative_text_mask = negative_text_mask.to(self.device) + if negative_text_mask.ndim == 1: + negative_text_mask = negative_text_mask.unsqueeze(0) + # Fresh noise per micro-batch — DMD2 samples noise independently at each loss call. + noise = torch.randn_like(latents) + return latents, noise, text_embeds, text_mask, negative_text_embeds, negative_text_mask + + def _log_step( + self, + *, + global_step: int, + is_student_phase: bool, + group_loss: float, + grad_norm: float, + vsd_loss: float | None, + disc_loss: float | None = None, + ) -> None: + """Log a single step. Stdout always; wandb when the parent set it up.""" + phase = "student" if is_student_phase else "fake_score" + + # Stdout + suffix = f" vsd={vsd_loss:.4f}" if vsd_loss is not None else "" + if disc_loss is not None: + suffix += f" disc={disc_loss:.4f}" + logging.info( + "[STEP %d] phase=%s loss=%.4f grad_norm=%.4f%s lr=%.2e", + global_step, + phase, + group_loss, + grad_norm, + suffix, + self.optimizer.param_groups[0]["lr"], + ) + + # wandb + try: + import wandb + + if wandb.run is not None: + log_dict: dict[str, Any] = { + f"{phase}/loss": group_loss, + f"{phase}/grad_norm": grad_norm, + "global_step": global_step, + "lr_student": self.optimizer.param_groups[0]["lr"], + "lr_fake_score": self._fake_score_optimizer.param_groups[0]["lr"], + } + if vsd_loss is not None: + log_dict["student/vsd"] = vsd_loss + if disc_loss is not None: + log_dict["discriminator/loss"] = disc_loss + wandb.log(log_dict, step=global_step) + except Exception: + # wandb not installed or not initialised — silent no-op. + pass + + def _dmd_config_summary(self) -> str: + """Compact one-line summary of the active DMDConfig for startup logging.""" + cfg = self._dmd_config + t_list = cfg.sample_t_cfg.t_list if cfg.sample_t_cfg is not None else None + return ( + f"pred_type={cfg.pred_type} fake_score_pred_type={cfg.fake_score_pred_type} " + f"num_train_timesteps={cfg.num_train_timesteps} " + f"student_update_freq={cfg.student_update_freq} " + f"student_sample_steps={cfg.student_sample_steps} " + f"student_sample_type={cfg.student_sample_type} " + f"backward_simulation={cfg.backward_simulation} " + f"t_list={t_list} " + f"gan_loss_weight_gen={cfg.gan_loss_weight_gen} " + f"guidance_scale={cfg.guidance_scale} ema={'on' if cfg.ema is not None else 'off'}" + ) + + def _dmd_full_config_log(self) -> str: + """Full multi-line dump of every DMD2 parameter for startup tracing. + + Two sections: the resolved DMDConfig (every Pydantic field, including + nested ``sample_t_cfg`` and ``ema`` blocks) and the recipe-side keys + under ``dmd2:`` that aren't DMDConfig fields (e.g. ``fake_score_lr``, + ``gan_feature_indices``, ``pipeline_plugin``). Combined they cover + every knob that ends up driving the DMD2 method at runtime. + """ + cfg = self._dmd_config + dmd_node = self.cfg.get("dmd2", {}) or {} + if hasattr(dmd_node, "to_dict"): + dmd_node = dmd_node.to_dict() + else: + dmd_node = dict(dmd_node) + recipe_extras = { + k: v + for k, v in dmd_node.items() + if k not in _DMD_CONFIG_OVERRIDE_KEYS and k != "recipe_path" + } + combined = { + "DMDConfig_resolved": cfg.model_dump(), + "recipe_side_extras": recipe_extras, + } + return json.dumps(combined, indent=2, default=str) diff --git a/examples/diffusers/fastgen/export_diffusers_qwen_image.py b/examples/diffusers/fastgen/export_diffusers_qwen_image.py new file mode 100644 index 00000000000..65a17c9c3c0 --- /dev/null +++ b/examples/diffusers/fastgen/export_diffusers_qwen_image.py @@ -0,0 +1,202 @@ +# 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. + +"""Export a DMD2 student as a full diffusers-loadable QwenImagePipeline dir. + +The §10 safetensors addendum produces a transformer-only dir: + + epoch_0_step_N/model/consolidated/ + config.json + diffusion_pytorch_model.safetensors.index.json + model-00001-of-00001.safetensors + +That dir is loadable by ``QwenImageTransformer2DModel.from_pretrained`` but NOT +by ``QwenImagePipeline.from_pretrained`` or ``DiffusionPipeline.from_pretrained`` +because it lacks ``model_index.json`` and the sibling component dirs +(``vae/``, ``text_encoder/``, ``tokenizer/``, ``scheduler/``). + +This utility assembles a full pipeline dir by: + + / + model_index.json (copied from base) + transformer/ (symlinked or copied from the consolidated student) + vae/ (symlinked from base Qwen-Image) + text_encoder/ (symlinked from base Qwen-Image) + tokenizer/ (symlinked from base Qwen-Image) + scheduler/ (symlinked from base Qwen-Image) + +After this runs, the dir loads with ``QwenImagePipeline.from_pretrained(output_dir)`` +or ``DiffusionPipeline.from_pretrained(output_dir)``. + +Symlinks are the default — the base Qwen-Image components are huge (text encoder +alone is ~12 GB) and never change between DMD2 students, so copying them per +checkpoint wastes disk. Use ``--copy`` if the output dir must be portable. + +Usage:: + + python export_diffusers_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 \\ + [--copy] + +Smoke test (``--verify``) loads the assembled dir via QwenImagePipeline and +checks the transformer config matches the student's. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import shutil +import sys +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +logger = logging.getLogger(__name__) + +# Components borrowed from the base Qwen-Image checkpoint. The trained +# transformer replaces the corresponding entry. +BASE_COMPONENTS = ("vae", "text_encoder", "tokenizer", "scheduler") + + +def _link_or_copy(src: str, dst: str, copy: bool) -> None: + if os.path.lexists(dst): + if os.path.islink(dst) or os.path.isfile(dst): + os.remove(dst) + else: + shutil.rmtree(dst) + if copy: + if os.path.isdir(src): + shutil.copytree(src, dst, symlinks=True) + else: + shutil.copy2(src, dst) + else: + os.symlink(os.path.abspath(src), dst) + + +def export_diffusers( + student_path: str | Path, + base_pipeline_path: str | Path, + output_dir: str | Path, + copy: bool = False, +) -> None: + student_path = str(student_path) + base_pipeline_path = str(base_pipeline_path) + output_dir = str(output_dir) + + if not os.path.isdir(student_path): + raise FileNotFoundError(f"student_path is not a directory: {student_path}") + if not os.path.isdir(base_pipeline_path): + raise FileNotFoundError(f"base_pipeline_path is not a directory: {base_pipeline_path}") + base_index = os.path.join(base_pipeline_path, "model_index.json") + if not os.path.isfile(base_index): + raise FileNotFoundError(f"base pipeline missing model_index.json: {base_index}") + + os.makedirs(output_dir, exist_ok=True) + logger.info( + "[Diffusers-Export] Output dir: %s (mode=%s)", output_dir, "copy" if copy else "symlink" + ) + + # 1. model_index.json — copy verbatim (the class registry is the same + # whether the transformer weights are live or DMD-distilled). + dst_index = os.path.join(output_dir, "model_index.json") + with open(base_index) as f: + index = json.load(f) + with open(dst_index, "w") as f: + json.dump(index, f, indent=2) + logger.info("[Diffusers-Export] Wrote %s", dst_index) + + # 2. transformer/ — link/copy from the consolidated student. + dst_transformer = os.path.join(output_dir, "transformer") + _link_or_copy(student_path, dst_transformer, copy=copy) + logger.info("[Diffusers-Export] transformer/ <- %s", student_path) + + # 3. vae / text_encoder / tokenizer / scheduler — link/copy from base. + for comp in BASE_COMPONENTS: + src = os.path.join(base_pipeline_path, comp) + dst = os.path.join(output_dir, comp) + if not os.path.isdir(src): + raise FileNotFoundError(f"base pipeline missing component: {src}") + _link_or_copy(src, dst, copy=copy) + logger.info("[Diffusers-Export] %s/ <- %s", comp, src) + + logger.info( + "[Diffusers-Export] Done. Load via QwenImagePipeline.from_pretrained(%r)", output_dir + ) + + +def _verify(output_dir: str) -> dict: + """Load via QwenImagePipeline.from_pretrained and report a small status dict.""" + import torch + from diffusers import QwenImagePipeline + + logger.info("[Diffusers-Export-Verify] Loading %s via QwenImagePipeline", output_dir) + pipe = QwenImagePipeline.from_pretrained(output_dir, torch_dtype=torch.bfloat16) + n_transformer_params = sum(p.numel() for p in pipe.transformer.parameters()) + n_vae_params = sum(p.numel() for p in pipe.vae.parameters()) + n_text_params = sum(p.numel() for p in pipe.text_encoder.parameters()) + stats = { + "loaded_class": type(pipe).__name__, + "transformer_class": type(pipe.transformer).__name__, + "transformer_params": int(n_transformer_params), + "vae_class": type(pipe.vae).__name__, + "vae_params": int(n_vae_params), + "text_encoder_class": type(pipe.text_encoder).__name__, + "text_encoder_params": int(n_text_params), + "tokenizer_class": type(pipe.tokenizer).__name__, + "scheduler_class": type(pipe.scheduler).__name__, + } + return stats + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--student_path", + required=True, + help="Consolidated student dir (e.g. .../epoch_0_step_5/model/consolidated)", + ) + parser.add_argument( + "--base_pipeline_path", + required=True, + help="Base Qwen-Image pipeline dir (vae / text_encoder / tokenizer / scheduler source)", + ) + parser.add_argument("--output_dir", required=True, help="Where to write the full pipeline dir") + parser.add_argument( + "--copy", + action="store_true", + help="Copy components instead of symlinking. Off by default to save disk.", + ) + parser.add_argument( + "--verify", action="store_true", help="Re-load via QwenImagePipeline.from_pretrained" + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + export_diffusers( + student_path=args.student_path, + base_pipeline_path=args.base_pipeline_path, + output_dir=args.output_dir, + copy=args.copy, + ) + if args.verify: + stats = _verify(args.output_dir) + print(json.dumps(stats, indent=2)) + sys.exit(0) diff --git a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py new file mode 100644 index 00000000000..7748eb15ec3 --- /dev/null +++ b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py @@ -0,0 +1,528 @@ +# 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. + +"""Inference pipeline for DMD2-trained Qwen-Image students. + +Loads the consolidated safetensors transformer (from a §10 checkpoint with +``model_save_format=safetensors`` + ``save_consolidated=true``) plus the base +Qwen-Image VAE / text-encoder / tokenizer, and exposes a diffusers-style +``pipe(prompt=...).images[0]`` call that runs the DMD few-step sampler. + +Math is bit-aligned with the training-time ``_build_student_input`` in +``modelopt/torch/fastgen/methods/dmd.py``: + + Single-step (Phase 1 default): + x_T = noise * max_t # initial latent + v = student(x_T, t=max_t, text_emb) # one forward + x_0 = x_T - max_t * v # RF identity + image = vae.decode(x_0) + + Multi-step (Phase 2, ``num_inference_steps > 1``): + x = noise * t_list[0] # initial latent at t_max + for (t_cur, t_next) in zip(t_list[:-1], t_list[1:]): + v = student(x, t=t_cur, text_emb) # flow at t_cur + x_0 = x - t_cur * v # RF identity → x_0 estimate + if t_next > 0: + eps = (x - (1 - t_cur) * x_0) / t_cur # ODE: invert RF forward + x = (1 - t_next) * x_0 + t_next * eps # re-noise to t_next + else: + x = x_0 # final step + image = vae.decode(x) + +Usage:: + + from inference_dmd2_qwen_image import QwenImageDMDInferencePipeline + import torch + + 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 "…/epoch_0_step_5/ema_shadow.pt" + torch_dtype=torch.bfloat16, + ).to("cuda") + + image = pipe( + prompt="a small red cube on a white table", + num_inference_steps=1, + height=512, + width=512, + generator=torch.Generator("cuda").manual_seed(42), + ).images[0] + image.save("dmd2_smoke.png") +""" + +from __future__ import annotations + +import itertools +import logging +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from diffusers import QwenImagePipeline, QwenImageTransformer2DModel +from diffusers.utils.torch_utils import randn_tensor + +if TYPE_CHECKING: + from pathlib import Path + +logger = logging.getLogger(__name__) + + +@dataclass +class QwenImageDMDOutput: + """Container for inference outputs — mirrors diffusers' pipeline outputs.""" + + images: list + + +class QwenImageDMDInferencePipeline: + """Thin inference wrapper around a DMD2-trained Qwen-Image student. + + Wraps a stock ``diffusers.QwenImagePipeline`` whose ``transformer`` field + has been swapped for our trained student. Re-uses the pipeline's VAE, + text encoder, tokenizer, and image-processor for everything *except* the + denoising loop, which is replaced by the DMD few-step sampler. + """ + + def __init__( + self, + base_pipeline: QwenImagePipeline, + max_t: float = 0.999, + ) -> None: + self._pipe = base_pipeline + self.max_t = max_t + + # ------------------------------------------------------------------ # + # Loading # + # ------------------------------------------------------------------ # + + @classmethod + def from_pretrained( + cls, + student_path: str | Path, + base_pipeline_path: str | Path, + ema_path: str | Path | None = None, + torch_dtype: torch.dtype = torch.bfloat16, + max_t: float = 0.999, + ) -> QwenImageDMDInferencePipeline: + """Load the student + base Qwen-Image components. + + Args: + student_path: A consolidated dir from §10's safetensors save — must + contain ``config.json``, ``diffusion_pytorch_model.safetensors.index.json``, + and one or more ``*.safetensors`` shards. Loadable directly via + ``QwenImageTransformer2DModel.from_pretrained``. + base_pipeline_path: The base Qwen-Image checkpoint (e.g. + ``Qwen/Qwen-Image`` or a local snapshot). Used only for the + ``vae`` / ``text_encoder`` / ``tokenizer`` / ``image_processor``; + the transformer is replaced. + ema_path: Optional ``ema_shadow.pt`` produced by ``_save_dmd_extras``. + If provided, the EMA shadow weights are overlaid onto the + student after the safetensors load. EMA usually yields cleaner + samples than the live student. + torch_dtype: dtype for the student + VAE. ``bfloat16`` is what we + trained with. + max_t: Initial timestep for the 1-step sampler. Must match the + recipe's ``sample_t_cfg.max_t`` (default 0.999). + """ + student_path = str(student_path) + base_pipeline_path = str(base_pipeline_path) + if not os.path.isdir(student_path): + raise FileNotFoundError(f"student_path is not a directory: {student_path}") + if not os.path.isdir(base_pipeline_path): + raise FileNotFoundError(f"base_pipeline_path is not a directory: {base_pipeline_path}") + + logger.info("[DMD2-Inference] Loading trained student from %s", student_path) + student = QwenImageTransformer2DModel.from_pretrained(student_path, torch_dtype=torch_dtype) + + if ema_path is not None: + logger.info("[DMD2-Inference] Overlaying EMA shadow from %s", ema_path) + ema_state = torch.load(str(ema_path), map_location="cpu", weights_only=False) + shadow = ( + ema_state.get("shadow", ema_state) if isinstance(ema_state, dict) else ema_state + ) + if not isinstance(shadow, dict): + raise ValueError( + f"ema_shadow.pt content has unexpected type {type(shadow).__name__}; " + "expected dict[str, Tensor]." + ) + missing, unexpected = student.load_state_dict(shadow, strict=False) + if unexpected: + logger.warning( + "[DMD2-Inference] EMA overlay had %d unexpected keys (first: %s)", + len(unexpected), + unexpected[:3], + ) + if missing: + logger.warning( + "[DMD2-Inference] EMA overlay missed %d student keys (first: %s)", + len(missing), + missing[:3], + ) + + student.eval() + + logger.info( + "[DMD2-Inference] Loading base Qwen-Image pipeline from %s (transformer replaced)", + base_pipeline_path, + ) + # Passing transformer= bypasses loading the original transformer from disk; + # the rest (vae, text_encoder, tokenizer, scheduler, image_processor) loads + # normally. + pipe = QwenImagePipeline.from_pretrained( + base_pipeline_path, + transformer=student, + torch_dtype=torch_dtype, + ) + + return cls(base_pipeline=pipe, max_t=max_t) + + def to(self, device: str | torch.device) -> QwenImageDMDInferencePipeline: + self._pipe.to(device) + return self + + @property + def device(self) -> torch.device: + return self._pipe.transformer.device + + @property + def dtype(self) -> torch.dtype: + return next(self._pipe.transformer.parameters()).dtype + + # ------------------------------------------------------------------ # + # Inference # + # ------------------------------------------------------------------ # + + @torch.no_grad() + def __call__( + self, + prompt: str | list, + negative_prompt: str | list | None = None, + num_inference_steps: int = 1, + guidance_scale: float = 1.0, + height: int = 1024, + width: int = 1024, + num_images_per_prompt: int = 1, + generator: torch.Generator | None = None, + max_t: float | None = None, + t_list: list | None = None, + sample_type: str = "ode", + output_type: str = "pil", + max_sequence_length: int = 512, + ) -> QwenImageDMDOutput: + """Generate image(s) from the trained DMD2 student. + + ``num_inference_steps == 1`` runs the canonical Phase 1 single-step + sampler (``x_0 = x_T - max_t * v``). ``num_inference_steps > 1`` runs + the multi-step DMD unroll using ``t_list`` (or + ``linspace(max_t, 0, num_inference_steps + 1)`` if ``t_list`` is None). + + ``sample_type`` selects between deterministic (``"ode"``, recovers eps + from x_0 via RF identity) and stochastic (``"sde"``, fresh noise per + step). Matches FastGen's ``student_sample_type``. + + ``guidance_scale != 1.0`` activates inference-time classifier-free + guidance: the student is called twice (positive + negative prompt) per + step and the two flows are blended as ``v = v_neg + s*(v_pos - v_neg)``. + ``negative_prompt`` defaults to ``""`` (the canonical empty prompt) when + unset and CFG is engaged. + + **Note on DMD2 students trained with CFG**: a student trained with + ``dmd2.guidance_scale=4.0`` has *already* internalised CFG (its + single-pass output is the CFG-augmented teacher target). Pass + ``guidance_scale=1.0`` at inference to avoid double-CFG. Use + ``guidance_scale > 1.0`` only for students that were trained without + CFG (``dmd2.guidance_scale=null``). + """ + if sample_type not in ("ode", "sde"): + raise ValueError(f"sample_type must be 'ode' or 'sde', got {sample_type!r}") + do_cfg = guidance_scale != 1.0 + if do_cfg and negative_prompt is None: + # Default to the canonical empty unconditional prompt. + negative_prompt = "" + + max_t = float(max_t) if max_t is not None else float(self.max_t) + pipe = self._pipe + device = self.device + dtype = self.dtype + + # ---- 1. Resolve t_list ----------------------------------------------- + if num_inference_steps == 1: + schedule = [max_t, 0.0] + elif t_list is not None: + if len(t_list) != num_inference_steps + 1: + raise ValueError( + f"t_list must have num_inference_steps+1 entries " + f"(got {len(t_list)} for num_inference_steps={num_inference_steps})" + ) + schedule = [float(t) for t in t_list] + if abs(schedule[-1]) > 1e-6: + raise ValueError( + f"t_list must end at 0.0 (got {schedule[-1]}); the final step lands on x_0." + ) + else: + # Default: linear schedule from max_t to 0 (matches FastGen's + # torch.linspace(max_t, 0, sample_steps + 1) fallback). + schedule = torch.linspace(max_t, 0.0, num_inference_steps + 1).tolist() + + # ---- 2. Encode prompt(s) --------------------------------------------- + prompt_embeds, prompt_embeds_mask = pipe.encode_prompt( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + neg_prompt_embeds = None + neg_prompt_embeds_mask = None + if do_cfg: + neg_prompt_embeds, neg_prompt_embeds_mask = pipe.encode_prompt( + prompt=negative_prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + txt_seq_lens = ( + prompt_embeds_mask.sum(dim=1).int().tolist() if prompt_embeds_mask is not None else None + ) + neg_txt_seq_lens = ( + neg_prompt_embeds_mask.sum(dim=1).int().tolist() + if neg_prompt_embeds_mask is not None + else None + ) + + # ---- 3. Build initial noisy latents at t = schedule[0] --------------- + if isinstance(prompt, str): + batch_size = 1 + else: + batch_size = len(prompt) + batch_size = batch_size * num_images_per_prompt + + num_channels_latents = pipe.transformer.config.in_channels // 4 # 64 // 4 = 16 + h_lat = 2 * (height // (pipe.vae_scale_factor * 2)) + w_lat = 2 * (width // (pipe.vae_scale_factor * 2)) + latent_shape = (batch_size, 1, num_channels_latents, h_lat, w_lat) + + # DMD initial latents: noise * schedule[0] (RF: sigma(t0) = t0). + noise = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) + latents_5d = noise * schedule[0] + + latents_packed = pipe._pack_latents( + latents_5d, batch_size, num_channels_latents, h_lat, w_lat + ) + img_shapes = [[(1, h_lat // 2, w_lat // 2)]] * batch_size + + # ---- 4. DMD few-step unroll ----------------------------------------- + x_packed = latents_packed + for t_cur, t_next in itertools.pairwise(schedule): + timestep = torch.tensor([t_cur], device=device, dtype=dtype).expand(batch_size) + flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + if do_cfg: + # CFG two-pass: ``v_cfg = v_neg + s*(v_pos - v_neg)``. Equivalent + # to the FastGen training-time formula + # ``teacher_pos + (s-1) * (teacher_pos - teacher_neg)`` once + # expanded. Engaged only when the caller passes + # ``guidance_scale != 1.0`` — DMD2 students trained with a + # non-null ``dmd2.guidance_scale`` already internalise CFG, so + # leave guidance_scale=1.0 for those. + neg_flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=neg_prompt_embeds, + encoder_hidden_states_mask=neg_prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=neg_txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + flow_packed = ( + neg_flow_packed.to(torch.float64) + + float(guidance_scale) + * (flow_packed.to(torch.float64) - neg_flow_packed.to(torch.float64)) + ).to(dtype) + # RF identity: x_0 = x_t - t_cur * v (computed in fp64 for stability). + x0_packed = ( + x_packed.to(torch.float64) - float(t_cur) * flow_packed.to(torch.float64) + ).to(dtype) + + if t_next > 1e-6: + # Re-noise x_0 forward to t_next. + if sample_type == "ode": + # Deterministic: invert the RF forward to recover the implied eps, + # then re-noise. eps = (x_t - (1 - t_cur) * x_0) / t_cur + alpha_cur = 1.0 - float(t_cur) + eps_packed = ( + (x_packed.to(torch.float64) - alpha_cur * x0_packed.to(torch.float64)) + / max(float(t_cur), 1e-6) + ).to(dtype) + else: + # Stochastic: fresh Gaussian noise. + eps_packed = torch.randn( + x_packed.shape, generator=generator, device=device, dtype=dtype + ) + # RF forward: x_{t_next} = (1 - t_next) * x_0 + t_next * eps. + alpha_next = 1.0 - float(t_next) + x_packed = ( + alpha_next * x0_packed.to(torch.float64) + + float(t_next) * eps_packed.to(torch.float64) + ).to(dtype) + else: + # Last step: x_0 is the output. + x_packed = x0_packed + + # Unpack to 5D for VAE. + x0_5d = pipe._unpack_latents(x_packed, height, width, pipe.vae_scale_factor) + + # ---- 5. VAE decode --------------------------------------------------- + # Reverse the VAE-side scaling that diffusers applied at encoding time. + latents_mean = ( + torch.tensor(pipe.vae.config.latents_mean) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device=device, dtype=dtype) + ) + latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view( + 1, pipe.vae.config.z_dim, 1, 1, 1 + ).to(device=device, dtype=dtype) + x0_scaled = x0_5d / latents_std + latents_mean + + # vae.decode returns 5D; the trailing [:, :, 0] drops the temporal dim + # since Qwen-Image treats images as 1-frame videos. + image_5d = pipe.vae.decode(x0_scaled, return_dict=False)[0] + image_4d = image_5d[:, :, 0] # [B, C, H, W] + + images = pipe.image_processor.postprocess(image_4d, output_type=output_type) + return QwenImageDMDOutput(images=images) + + +# ---------------------------------------------------------------------------- # +# Standalone smoke-test driver. Validates end-to-end wiring against the §10 # +# safetensors checkpoint. Mock-data training means the image won't be # +# coherent — pass criterion is just "the pipeline produces a finite image # +# tensor and the file writes successfully". # +# ---------------------------------------------------------------------------- # + + +def _smoke_test( + student_path: str, + base_pipeline_path: str, + output_png: str, + ema_path: str | None = None, + prompt: str = "a small red cube on a white table", + height: int = 512, + width: int = 512, + seed: int = 42, +) -> None: + """Run a one-shot inference and dump a PNG. + + Writes a small JSON sidecar next to the PNG with shape/dtype/range stats + so the success criteria can be machine-checked. + """ + import json + + logging.basicConfig(level=logging.INFO) + pipe = QwenImageDMDInferencePipeline.from_pretrained( + student_path=student_path, + base_pipeline_path=base_pipeline_path, + ema_path=ema_path, + torch_dtype=torch.bfloat16, + ) + device = "cuda" if torch.cuda.is_available() else "cpu" + pipe = pipe.to(device) + gen = torch.Generator(device=device).manual_seed(seed) + + out = pipe( + prompt=prompt, + num_inference_steps=1, + height=height, + width=width, + generator=gen, + ) + image = out.images[0] + + # PIL image; sanity-check shape + range. + import numpy as np + + arr = np.array(image) + stats = { + "prompt": prompt, + "height": height, + "width": width, + "seed": seed, + "ema_overlay": ema_path is not None, + "image_shape": list(arr.shape), + "image_dtype": str(arr.dtype), + "image_min": int(arr.min()), + "image_max": int(arr.max()), + "image_mean": float(arr.mean()), + "image_std": float(arr.std()), + "is_finite": bool(np.isfinite(arr).all()), + "is_not_constant": bool(arr.std() > 0), + } + + os.makedirs(os.path.dirname(output_png), exist_ok=True) + image.save(output_png) + sidecar = output_png.replace(".png", "_stats.json") + with open(sidecar, "w") as f: + json.dump(stats, f, indent=2) + print(json.dumps(stats, indent=2)) + print(f"\nImage saved to: {output_png}") + print(f"Stats sidecar: {sidecar}") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--student_path", + required=True, + help="Path to the consolidated safetensors student checkpoint " + "(e.g. .../epoch_0_step_500/model/consolidated).", + ) + parser.add_argument( + "--base_pipeline_path", + default="Qwen/Qwen-Image", + help="Base Qwen-Image pipeline (HF id or local snapshot) for the VAE / text-encoder / tokenizer.", + ) + parser.add_argument( + "--output_png", + default="./outputs/dmd2_sample.png", + ) + parser.add_argument("--ema_path", default=None) + parser.add_argument("--prompt", default="a small red cube on a white table") + parser.add_argument("--height", type=int, default=512) + parser.add_argument("--width", type=int, default=512) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + _smoke_test( + student_path=args.student_path, + base_pipeline_path=args.base_pipeline_path, + output_png=args.output_png, + ema_path=args.ema_path, + prompt=args.prompt, + height=args.height, + width=args.width, + seed=args.seed, + ) diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt new file mode 100644 index 00000000000..1cba3ce7868 --- /dev/null +++ b/examples/diffusers/fastgen/requirements.txt @@ -0,0 +1,10 @@ +# Runtime requirements for the DMD2 Qwen-Image AutoModel example. +# 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, dataloader, FSDP2 wrapping). +# The diffusion extras install diffusers + accelerate with matching pins. +nemo_automodel[diffusion] + +# Optional but recommended for the smoke logs. +wandb diff --git a/modelopt/torch/fastgen/__init__.py b/modelopt/torch/fastgen/__init__.py new file mode 100644 index 00000000000..507acfddd72 --- /dev/null +++ b/modelopt/torch/fastgen/__init__.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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-agnostic diffusion step-distillation losses (FastGen port). + +``modelopt.torch.fastgen`` is a loss-computation library. It accepts already-built +``nn.Module`` references (student / teacher / fake-score / optional discriminator) and +returns scalar loss tensors. It does **not** load models, manage optimizers, wrap +anything as a ``DynamicModule``, or register itself in any mode registry. + +Typical usage with a YAML-driven config:: + + import modelopt.torch.fastgen as mtf + + student, teacher = build_wan_student_and_teacher(...) + fake_score = mtf.create_fake_score(teacher) + + cfg = mtf.load_dmd_config("general/distillation/dmd2_qwen_image") + + # If GAN is enabled, expose intermediate teacher features to the discriminator. + if cfg.gan_loss_weight_gen > 0: + mtf.plugins.qwen_image.attach_feature_capture(teacher, feature_indices=[30]) + + pipeline = mtf.DMDPipeline(student, teacher, fake_score, cfg, discriminator=disc) + + # Inside the training loop (framework-owned): + if step % cfg.student_update_freq == 0: + losses = pipeline.compute_student_loss( + latents, noise, text_embeds, negative_encoder_hidden_states=neg_embeds + ) + losses["total"].backward() + student_opt.step() + pipeline.update_ema() + else: + f = pipeline.compute_fake_score_loss(latents, noise, text_embeds) + f["total"].backward() + fake_score_opt.step() + if disc is not None: + d = pipeline.compute_discriminator_loss(latents, noise, text_embeds) + d["total"].backward() + disc_opt.step() +""" + +from . import flow_matching, losses, utils +from .config import * +from .ema import * +from .factory import * +from .loader import * +from .methods.dmd 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 diff --git a/modelopt/torch/fastgen/config.py b/modelopt/torch/fastgen/config.py new file mode 100644 index 00000000000..30d78b8720f --- /dev/null +++ b/modelopt/torch/fastgen/config.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""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. + +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``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal + +from pydantic import Field, model_validator + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +if TYPE_CHECKING: + from pathlib import Path + +__all__ = [ + "DMDConfig", + "DistillationConfig", + "EMAConfig", + "SampleTimestepConfig", +] + +PredType = Literal["x0", "eps", "v", "flow"] +TimeDistType = Literal["uniform", "logitnormal", "lognormal", "shifted", "polynomial"] + + +class SampleTimestepConfig(ModeloptBaseConfig): + """Timestep sampling distribution for diffusion training.""" + + time_dist_type: TimeDistType = ModeloptField( + default="shifted", + title="Timestep distribution", + description=( + "Distribution used to sample the training timestep ``t``. Rectified-flow models" + " typically use ``shifted`` (Wan 2.2) or ``logitnormal`` (SD3, Flux)." + ), + ) + min_t: float = ModeloptField( + default=0.001, + title="Minimum t", + description="Lower bound of the sampling range (clamped before use).", + ) + max_t: float = ModeloptField( + default=0.999, + title="Maximum t", + description="Upper bound of the sampling range (clamped before use).", + ) + shift: float = ModeloptField( + default=5.0, + title="Shift factor", + description="Shift factor for ``time_dist_type='shifted'``; must be >= 1.", + ) + p_mean: float = ModeloptField( + default=0.0, + title="Distribution mean (log-space)", + description="Mean of the underlying normal for ``logitnormal`` / ``lognormal``.", + ) + p_std: float = ModeloptField( + default=1.0, + title="Distribution std (log-space)", + description="Standard deviation of the underlying normal for ``logitnormal`` / ``lognormal``.", + ) + t_list: list[float] | None = ModeloptField( + default=None, + title="Multi-step student timesteps", + description=( + "Explicit timestep schedule used when ``DMDConfig.student_sample_steps > 1``." + " The final element must be ``0.0``." + ), + ) + + @model_validator(mode="after") + def _check_bounds(self) -> SampleTimestepConfig: + assert 0.0 <= self.min_t < self.max_t, ( + f"require 0 <= min_t < max_t, got min_t={self.min_t}, max_t={self.max_t}" + ) + assert self.shift >= 1.0, f"shift must be >= 1, got {self.shift}" + if self.t_list is not None: + assert len(self.t_list) >= 2, "t_list must contain at least 2 entries (including t=0)" + assert self.t_list[-1] == 0.0, f"t_list[-1] must be 0.0, got {self.t_list[-1]}" + return self + + +class EMAConfig(ModeloptBaseConfig): + """Exponential moving average (EMA) hyperparameters for the student network.""" + + decay: float = ModeloptField( + default=0.9999, + title="EMA decay", + description="Decay coefficient for ``type='constant'``. Ignored for ``halflife``/``power``.", + ) + type: Literal["constant", "halflife", "power"] = ModeloptField( + default="constant", + title="EMA decay schedule", + description="Schedule used to compute the per-step decay coefficient.", + ) + start_iter: int = ModeloptField( + default=0, + title="EMA start iteration", + description="Iteration at which EMA tracking begins (EMA is initialized from the live weights at this step).", + ) + gamma: float = ModeloptField( + default=16.97, + title="Power schedule gamma", + description="Exponent for ``type='power'`` (``beta = (1 - 1/iter)**(gamma + 1)``).", + ) + halflife_kimg: float = ModeloptField( + default=500.0, + title="Halflife (kimg)", + description="Halflife in thousands of images for ``type='halflife'``.", + ) + rampup_ratio: float | None = ModeloptField( + default=0.05, + title="Halflife rampup ratio", + description="Rampup fraction for ``type='halflife'``; pass ``None`` to disable rampup.", + ) + batch_size: int = ModeloptField( + default=1, + title="Effective batch size", + description="Per-step global batch size used to convert iterations to nimg for the halflife schedule.", + ) + fsdp2: bool = ModeloptField( + default=True, + title="FSDP2 enabled", + description="If True, the EMA uses ``DTensor.full_tensor()`` to gather sharded parameters before updating.", + ) + mode: Literal["full_tensor", "local_shard"] = ModeloptField( + default="full_tensor", + title="FSDP2 gather mode", + description=( + "``full_tensor`` performs an all_gather per parameter (higher memory, exact global EMA)." + " ``local_shard`` updates each rank's local DTensor shard in place (low memory fallback)." + ), + ) + dtype: Literal["float32", "bfloat16", "float16"] | None = ModeloptField( + default="float32", + title="EMA shadow dtype", + description=( + "Precision of the EMA parameter shadows. Defaults to ``float32`` so EMA updates" + " remain numerically meaningful even when the live model is bf16/fp16 (cf. FastGen," + " which instantiates its EMA module in the net's construction dtype — typically" + " fp32). Pass ``None`` to keep param shadows in the live parameter's dtype." + " Buffer shadows always track the live dtype regardless of this setting." + ), + ) + + +class DistillationConfig(ModeloptBaseConfig): + """Shared hyperparameters for diffusion step-distillation methods. + + Concrete methods subclass this config to add method-specific fields + (see :class:`DMDConfig`). + """ + + pred_type: PredType = ModeloptField( + default="flow", + title="Network prediction parameterization", + description="Quantity predicted by the teacher / student network.", + ) + guidance_scale: float | None = ModeloptField( + default=None, + title="CFG scale", + description="Classifier-free guidance scale. If ``None`` CFG is disabled.", + ) + # ``ModeloptField`` hard-asserts on ``default_factory``; use Pydantic's ``Field`` + # directly for this mutable sub-config so each DMDConfig instance gets its own + # SampleTimestepConfig instead of sharing a single mutable default. + sample_t_cfg: SampleTimestepConfig = Field( + default_factory=SampleTimestepConfig, + title="Timestep sampling", + description="Timestep distribution used for both the teacher forward and the VSD / DSM losses.", + ) + student_sample_steps: int = ModeloptField( + default=1, + title="Student inference steps", + description="Number of denoising steps the distilled student performs at inference.", + ) + student_sample_type: Literal["sde", "ode"] = ModeloptField( + default="ode", + title="Student sampling mode", + description=( + "Integrator used when unrolling the student over ``student_sample_steps > 1`` steps." + " Consumed by inference samplers and by DMDPipeline when" + " ``DMDConfig.backward_simulation`` is enabled." + ), + ) + num_train_timesteps: int | None = ModeloptField( + default=None, + title="Training-time discrete timestep count", + description=( + "If set, the pipeline rescales the continuous RF timestep ``t ∈ [0, 1]`` to" + " ``num_train_timesteps * t`` before passing it to the model. Matches the" + " diffusers convention used by Wan 2.2 / SD3 / Flux (``num_train_timesteps = 1000``)." + " Leave ``None`` when the model wrapper already handles the rescaling internally." + ), + ) + + +class DMDConfig(DistillationConfig): + """Hyperparameters for DMD / DMD2 distribution-matching distillation. + + Default values are tuned for Wan 2.2 5B; callers fine-tune them per model. + See ``FastGen/fastgen/configs/experiments/WanT2V/config_dmd2_wan22_5b.py``. + """ + + student_update_freq: int = ModeloptField( + default=5, + title="Student update frequency", + description=( + "One student step for every ``student_update_freq`` fake-score / discriminator steps." + " Matches FastGen's DMD2 alternation. Not read by DMDPipeline; the training loop is" + " expected to enforce the alternation." + ), + ) + fake_score_pred_type: PredType | None = ModeloptField( + default="x0", + title="Fake-score prediction parameterization", + description=( + "Parameterization used when training the fake score. If ``None`` falls back to" + " :attr:`DistillationConfig.pred_type`." + ), + ) + backward_simulation: bool = ModeloptField( + default=False, + title="Backward simulation", + description=( + "When True for multi-step students, build the selected student input by" + " no-grad unrolling the current student from the first schedule rung through" + " earlier rungs, then re-noising the generated x0 at the selected rung." + " When False, use FastGen's Qwen-style noised-real latent path." + ), + ) + gan_loss_weight_gen: float = ModeloptField( + default=0.0, + title="Generator GAN weight", + description="Weight of the GAN generator term in the student loss. ``0`` disables the GAN branch.", + ) + gan_use_same_t_noise: bool = ModeloptField( + default=False, + title="Share t/noise across real and fake", + description="If True, reuse the same ``t`` and ``eps`` for real and fake samples in the discriminator update.", + ) + gan_r1_reg_weight: float = ModeloptField( + default=0.0, + title="R1 regularization weight", + description=( + "Weight of the approximate-R1 regularization term for the discriminator update. ``0`` disables R1." + " Recommended range when enabled: 100-1000." + ), + ) + gan_r1_reg_alpha: float = ModeloptField( + default=0.1, + title="R1 regularization noise scale", + description=( + "Standard deviation of the perturbation applied to real data when computing the" + " approximate R1 term." + ), + ) + ema: EMAConfig | None = ModeloptField( + default=None, + title="Student EMA", + description=( + "If set, an exponential moving average of the student is maintained and updated" + " via ``DMDPipeline.update_ema``." + ), + ) + + @model_validator(mode="after") + def _check_gan(self) -> DMDConfig: + if self.gan_r1_reg_weight > 0 and self.gan_loss_weight_gen <= 0: + raise ValueError( + "gan_r1_reg_weight > 0 requires gan_loss_weight_gen > 0 (the discriminator must be enabled)." + ) + if self.backward_simulation: + if self.student_sample_steps <= 1: + raise ValueError("backward_simulation=True requires student_sample_steps > 1.") + if self.sample_t_cfg.t_list is None: + raise ValueError("backward_simulation=True requires sample_t_cfg.t_list to be set.") + return self + + @classmethod + 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. + """ + # Imported lazily to avoid a circular import between this module and + # ``modelopt.torch.fastgen.loader`` (which imports :class:`DMDConfig`). + from .loader import load_dmd_config + + return load_dmd_config(config_file) diff --git a/modelopt/torch/fastgen/discriminators.py b/modelopt/torch/fastgen/discriminators.py new file mode 100644 index 00000000000..8e55233d052 --- /dev/null +++ b/modelopt/torch/fastgen/discriminators.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Discriminator modules for the DMD2 GAN branch. + +Ports FastGen's image-DiT discriminator from +``source/FastGen/fastgen/networks/discriminators.py`` so that ModelOpt's +:class:`~modelopt.torch.fastgen.methods.dmd.DMDPipeline` can run the GAN branch +without a FastGen dependency. The discriminator is **model-agnostic**: it takes +a list of spatial feature tensors ``[B, C, H, W]`` and returns concatenated +logits ``[B, num_heads]``. The model-specific work of producing those tensors +(installing forward hooks, reshaping packed-token streams into spatial maps) +lives in the per-model plugins (``plugins/qwen_image.py``). +""" + +from __future__ import annotations + +import torch +from torch import nn + +__all__ = ["Discriminator", "Discriminator_ImageDiT"] + + +def _get_optimal_groups(num_channels: int) -> int: + """Return a GroupNorm group count that divides ``num_channels`` evenly. + + Matches the heuristic in FastGen's discriminator: prefer 32 groups when + possible, fall back to the largest divisor below 32, and use + ``num_channels // 4`` for very small channel counts. + """ + if num_channels <= 32: + groups = max(1, num_channels // 4) + else: + groups = 32 + while groups > 1 and num_channels % groups != 0: + groups -= 1 + assert num_channels % groups == 0, f"{num_channels} not divisible by {groups}" + return groups + + +class Discriminator(nn.Module): + """Base class for DMD2 discriminators.""" + + def __init__(self, feature_indices: set[int] | None = None) -> None: + """Store the teacher block indices whose features feed the discriminator.""" + super().__init__() + self.feature_indices = feature_indices + + def forward(self, feats: list[torch.Tensor]) -> torch.Tensor: + """Map captured teacher features to discriminator logits (overridden by subclasses).""" + raise NotImplementedError("Subclasses must implement forward()") + + +# Class name kept verbatim from the FastGen reference implementation. +class Discriminator_ImageDiT(Discriminator): # noqa: N801 + """Image-DiT discriminator with one lightweight conv head per captured block. + + Input: list of feature tensors with shape ``[B, inner_dim, H, W]``, one per + block index in :attr:`feature_indices`. + + Output: concatenated logits ``[B, num_heads]`` (one column per head). The + DMD2 generator/discriminator losses read this as a 2D tensor. + + Per-head parameter count is ~``inner_dim * (inner_dim // 2) * 16 + ...``; + for ``inner_dim=3072`` (Flux / Qwen-Image) that's ~75 M params per head, so + keep ``len(feature_indices)`` small (≤3 heads is typical). + """ + + def __init__( + self, + feature_indices: set[int] | None = None, + num_blocks: int = 57, + inner_dim: int = 3072, + ) -> None: + """Build one lightweight conv classification head per captured block.""" + super().__init__(feature_indices=feature_indices) + + if self.feature_indices is None: + self.feature_indices = {int(num_blocks // 2)} + self.feature_indices = {i for i in self.feature_indices if i < num_blocks} + self.num_features = len(self.feature_indices) + self.inner_dim = inner_dim + + hidden_channels = inner_dim // 2 + self.cls_pred_heads = nn.ModuleList() + for _ in range(self.num_features): + head = nn.Sequential( + nn.Conv2d( + in_channels=inner_dim, + out_channels=hidden_channels, + kernel_size=4, + stride=2, + padding=1, + ), + nn.GroupNorm( + num_groups=_get_optimal_groups(hidden_channels), + num_channels=hidden_channels, + ), + nn.LeakyReLU(0.2), + nn.Conv2d( + in_channels=hidden_channels, + out_channels=1, + kernel_size=1, + stride=1, + padding=0, + ), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + ) + self.cls_pred_heads.append(head) + + def forward(self, feats: list[torch.Tensor]) -> torch.Tensor: + """Run each per-block conv head and concatenate their logits to ``[B, num_heads]``.""" + if not isinstance(feats, list) or len(feats) != self.num_features: + raise ValueError( + f"Expected list of {self.num_features} feature tensors, " + f"got {type(feats).__name__} with length " + f"{len(feats) if isinstance(feats, list) else 'N/A'}" + ) + all_logits = [] + for head, feat in zip(self.cls_pred_heads, feats): + logits = head(feat) + all_logits.append(logits) + return torch.cat(all_logits, dim=1) diff --git a/modelopt/torch/fastgen/ema.py b/modelopt/torch/fastgen/ema.py new file mode 100644 index 00000000000..3088acfaa62 --- /dev/null +++ b/modelopt/torch/fastgen/ema.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Exponential moving average of a student network, FSDP2 DTensor aware. + +Ported from ``FastGen/fastgen/callbacks/ema.py`` (lines 20-169) but exposed as a plain +class rather than a framework-specific callback. The caller decides when to call +:meth:`update` (typically after ``optimizer.step()``), how to persist the shadow state +(via :meth:`state_dict`), and when to publish the EMA weights back to a target module +(via :meth:`copy_to`). +""" + +from __future__ import annotations + +import copy +from typing import TYPE_CHECKING + +import torch +from torch import nn + +if TYPE_CHECKING: + from .config import EMAConfig + +__all__ = ["ExponentialMovingAverage"] + + +_DTYPE_MAP: dict[str, torch.dtype] = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, +} + + +def _resolve_dtype(config_dtype: str | None, fallback: torch.dtype) -> torch.dtype: + """Map an ``EMAConfig.dtype`` string to a ``torch.dtype``. + + ``config_dtype is None`` falls through to ``fallback`` (the live parameter's dtype). + """ + if config_dtype is None: + return fallback + try: + return _DTYPE_MAP[config_dtype] + except KeyError as exc: + raise ValueError( + f"Unsupported EMA dtype {config_dtype!r}; expected one of {sorted(_DTYPE_MAP)} or None." + ) from exc + + +def _is_distributed_tensor(t: torch.Tensor) -> bool: + """Return True if ``t`` is a ``torch.distributed.DTensor`` supporting ``full_tensor()``.""" + return hasattr(t, "full_tensor") and callable(t.full_tensor) + + +def _gather_full(param: torch.Tensor, *, fsdp2: bool) -> torch.Tensor: + """Return a materialised full tensor for ``param``. + + Mirrors the FSDP2 branch in ``FastGen/fastgen/callbacks/ema.py:128-139``: if CPU + offloading is enabled the local shard must be moved to CUDA before ``full_tensor()`` + can perform the all-gather (which requires a CUDA backend). + """ + if fsdp2 and _is_distributed_tensor(param): + if param.device.type == "cpu": + return param.to("cuda").full_tensor() + return param.full_tensor() + return param + + +def _strip_checkpoint_prefix(name: str) -> str: + """Remove the ``_checkpoint_wrapped_module.`` prefix injected by FSDP2 activation checkpointing.""" + return name.replace("_checkpoint_wrapped_module.", "") + + +class ExponentialMovingAverage: + """FSDP2-aware EMA tracker for a PyTorch module. + + The tracker stores a shadow state dict: parameters are promoted per + :attr:`EMAConfig.dtype` (default fp32) while buffers are kept in the live module's + dtype. Buffers are replicated across ranks and stepped via ``copy_`` rather than + ``lerp_``, so the bf16-roundoff argument that motivates parameter promotion + doesn't apply — preserving the live dtype makes the buffer restore exact. + + By default the tracker materialises the full tensor per parameter + (``mode='full_tensor'``) so the EMA represents the globally averaged weights even + when the model is sharded across ranks. A ``mode='local_shard'`` fallback is + available for memory-constrained settings — it does not all-gather and therefore + each rank holds an EMA of its local shard only. + + Example:: + + ema = ExponentialMovingAverage(student, EMAConfig(decay=0.999)) + for step in range(max_steps): + ... # compute loss, backward, optimizer.step() + ema.update(student, iteration=step) + + ema.copy_to(student_for_eval) # publish for inference + """ + + def __init__(self, model: nn.Module, config: EMAConfig) -> None: + """Pre-allocate the shadow state from ``model``'s parameters and buffers.""" + self.config = config + self._shadow: dict[str, torch.Tensor] = {} + self._buffer_shadow: dict[str, torch.Tensor] = {} + self._initialized = False + + # Pre-allocate shadow storage as a deepcopy of the live parameters on their + # current devices. Shadow dtype is promoted to ``EMAConfig.dtype`` (default + # fp32) so EMA updates remain meaningful even when the live model is + # bf16/fp16: the per-step increment ``(live - shadow) * (1 - beta)`` rounds + # to zero in bf16 (unit roundoff ~2^-8 of |shadow|) long before the live + # weights have converged. Pass ``dtype=None`` to fall back to the live + # parameter's dtype. + with torch.no_grad(): + for name, p in model.named_parameters(): + clean = _strip_checkpoint_prefix(name) + full = _gather_full(p.detach(), fsdp2=config.fsdp2) + target_dtype = _resolve_dtype(config.dtype, full.dtype) + self._shadow[clean] = copy.deepcopy(full).to(dtype=target_dtype) + # Buffers are replicated (not averaged) across ranks and stepped via + # ``copy_``, so the bf16-roundoff argument that drives ``EMAConfig.dtype`` + # on parameters doesn't apply — keep buffers in the live dtype for exact + # restore. + for name, b in model.named_buffers(): + clean = _strip_checkpoint_prefix(name) + self._buffer_shadow[clean] = copy.deepcopy(b.detach()) + + # ------------------------------------------------------------------ # + # Decay schedules # + # ------------------------------------------------------------------ # + + def _beta(self, iteration: int) -> float: + cfg = self.config + if cfg.type == "constant": + return cfg.decay + if cfg.type == "power": + # (1 - 1/iter) ** (gamma + 1); iteration must be > 0 for this to be finite. + safe_iter = max(iteration, 1) + return (1.0 - 1.0 / safe_iter) ** (cfg.gamma + 1.0) + if cfg.type == "halflife": + ema_halflife_nimg = cfg.halflife_kimg * 1000.0 + cur_nimg = iteration * cfg.batch_size + if cfg.rampup_ratio is not None: + ema_halflife_nimg = min(ema_halflife_nimg, cur_nimg * cfg.rampup_ratio) + return 0.5 ** (cfg.batch_size / max(ema_halflife_nimg, 1e-8)) + raise ValueError(f"Unsupported EMA type: {cfg.type!r}") + + # ------------------------------------------------------------------ # + # Public API # + # ------------------------------------------------------------------ # + + @torch.no_grad() + def update(self, model: nn.Module, *, iteration: int) -> None: + """Update the shadow state from ``model`` at the given iteration. + + Skips updates before :attr:`EMAConfig.start_iter`. On the iteration that equals + ``start_iter`` the shadow is (re-)initialised from the live weights; after that + it is updated with ``shadow = beta * shadow + (1 - beta) * live``. + """ + if iteration < self.config.start_iter: + return + + # (Re-)initialise the shadow from the live weights. Both arms are intentional: + # ``iteration == start_iter`` inits exactly at start when start_iter > 0 (earlier + # iterations are skipped above), while ``not self._initialized`` covers start_iter + # == 0 — where the auto-incremented counter never passes 0 — plus the first call + # after a resume. + if iteration == self.config.start_iter or not self._initialized: + self._copy_from_model(model) + self._initialized = True + return + + beta = self._beta(iteration) + + for name, p in model.named_parameters(): + clean = _strip_checkpoint_prefix(name) + if clean not in self._shadow: + continue + shadow = self._shadow[clean] + if self.config.mode == "full_tensor": + live = _gather_full(p.detach(), fsdp2=self.config.fsdp2) + else: + live = p.detach().to_local() if _is_distributed_tensor(p) else p.detach() + shadow.lerp_(live.to(device=shadow.device, dtype=shadow.dtype), 1.0 - beta) + + # Buffers are replicated across ranks under FSDP2, so we just copy. + for name, b in model.named_buffers(): + clean = _strip_checkpoint_prefix(name) + if clean in self._buffer_shadow: + shadow = self._buffer_shadow[clean] + shadow.copy_(b.detach().to(device=shadow.device, dtype=shadow.dtype)) + + @torch.no_grad() + def copy_to(self, target: nn.Module) -> None: + """Load the shadow state into ``target`` (which should share the tracked module's structure). + + The target is expected to be an unsharded module (i.e. the caller has unwrapped + any FSDP2 wrappers before calling). For sharded targets, prefer saving the + shadow via :meth:`state_dict` and reloading it through the framework's usual + checkpoint path. + """ + for name, p in target.named_parameters(): + clean = _strip_checkpoint_prefix(name) + if clean in self._shadow: + shadow = self._shadow[clean] + p.data.copy_(shadow.to(device=p.device, dtype=p.dtype)) + for name, b in target.named_buffers(): + clean = _strip_checkpoint_prefix(name) + if clean in self._buffer_shadow: + shadow = self._buffer_shadow[clean] + b.data.copy_(shadow.to(device=b.device, dtype=b.dtype)) + + def state_dict(self) -> dict[str, torch.Tensor]: + """Return the shadow state (parameters + buffers) for checkpointing.""" + merged: dict[str, torch.Tensor] = {} + merged.update(self._shadow) + merged.update(self._buffer_shadow) + return merged + + def load_state_dict(self, state: dict[str, torch.Tensor]) -> None: + """Restore the shadow state from a previously saved dict.""" + for k, v in state.items(): + if k in self._shadow: + shadow = self._shadow[k] + shadow.copy_(v.to(device=shadow.device, dtype=shadow.dtype)) + elif k in self._buffer_shadow: + shadow = self._buffer_shadow[k] + shadow.copy_(v.to(device=shadow.device, dtype=shadow.dtype)) + self._initialized = True + + # ------------------------------------------------------------------ # + # Internals # + # ------------------------------------------------------------------ # + + @torch.no_grad() + def _copy_from_model(self, model: nn.Module) -> None: + for name, p in model.named_parameters(): + clean = _strip_checkpoint_prefix(name) + if clean not in self._shadow: + continue + shadow = self._shadow[clean] + live = _gather_full(p.detach(), fsdp2=self.config.fsdp2) + shadow.copy_(live.to(device=shadow.device, dtype=shadow.dtype)) + for name, b in model.named_buffers(): + clean = _strip_checkpoint_prefix(name) + if clean in self._buffer_shadow: + shadow = self._buffer_shadow[clean] + shadow.copy_(b.detach().to(device=shadow.device, dtype=shadow.dtype)) diff --git a/modelopt/torch/fastgen/factory.py b/modelopt/torch/fastgen/factory.py new file mode 100644 index 00000000000..25ca7f31331 --- /dev/null +++ b/modelopt/torch/fastgen/factory.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Convenience factory helpers for constructing the auxiliary DMD networks. + +These helpers are intentionally tiny — the training framework is free to build the +fake score directly (e.g. under a meta-init context for FSDP2) instead of calling +:func:`create_fake_score`. See the ModelOpt ↔ FastGen design doc (FASTGEN_MODELOPT.md, +section "How the framework can build the fake_score") for both options. +""" + +from __future__ import annotations + +import copy +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from torch import nn + +__all__ = ["create_fake_score"] + + +def _looks_fsdp_wrapped(module: nn.Module) -> bool: + """Best-effort detection of an FSDP-wrapped module. + + Matches two shapes: + + - **FSDP1** (``torch.distributed.fsdp.FullyShardedDataParallel``): child modules + carry an ``_fsdp_wrapped_module`` attribute. + - **FSDP2** (``torch.distributed._composable.fsdp.fully_shard``): parameters are + ``DTensor`` instances exposing a ``full_tensor`` method. + + Probes only the first parameter to avoid iterating a large model. Intended for + the ``create_fake_score`` fast-fail check, not as a general-purpose predicate. + """ + if any(hasattr(m, "_fsdp_wrapped_module") for m in module.modules()): + return True + for p in module.parameters(): + if hasattr(p, "full_tensor"): + return True + break # only probe the first param + return False + + +def create_fake_score(teacher: nn.Module, *, deep_copy: bool = True) -> nn.Module: + """Return a trainable fake-score network initialized from the teacher. + + This is the unit-test / single-script path; frameworks that do meta-init + FSDP2 + wrapping will typically construct the fake score themselves and pass it directly + into :class:`~modelopt.torch.fastgen.methods.dmd.DMDPipeline`. + + Args: + teacher: The already-built teacher module. Must already have its weights loaded. + deep_copy: If True, :func:`copy.deepcopy` the teacher; if False, reuse the same + instance (only sensible if the caller can guarantee it is no longer held + elsewhere as the frozen teacher). + + Returns: + A copy of ``teacher`` in training mode with all parameters requiring gradients. + + FSDP2 caveat + ------------ + ``copy.deepcopy(teacher)`` is **not safe** when the teacher is already FSDP2-wrapped + (DTensor parameters + FSDP pre/post hooks + meta-init bookkeeping). For Stage-2 FSDP2 + training, skip this factory and construct the fake score under meta-init, then + rank-0-load weights and let ``sync_module_states`` broadcast:: + + with meta_init_context(): + fake_score = build_teacher_from_config(teacher_config) + if is_rank0(): + fake_score.load_state_dict(teacher.state_dict(), strict=False) + # Wrap with FSDP2(..., sync_module_states=True) to broadcast from rank 0. + + The pattern mirrors FastGen's + ``methods/distribution_matching/dmd2.py::DMD2Model.build_model``. A dedicated + ``create_fake_score_meta`` factory is planned alongside the Stage-2 training example. + + Raises: + RuntimeError: When ``deep_copy=True`` and the teacher looks FSDP-wrapped + (either FSDP1 via ``_fsdp_wrapped_module`` or FSDP2 via DTensor + parameters). The ``deep_copy=False`` branch skips the check because + reusing the teacher directly is compatible with an FSDP-wrapped input. + """ + if deep_copy and _looks_fsdp_wrapped(teacher): + raise RuntimeError( + "create_fake_score(deep_copy=True) is not safe on an FSDP-wrapped teacher " + "(DTensor parameters + FSDP hooks + meta-init bookkeeping don't survive " + "copy.deepcopy). Construct the fake score under meta-init and rank-0-load " + "weights instead — see the 'FSDP2 caveat' section of this function's docstring." + ) + fake_score = copy.deepcopy(teacher) if deep_copy else teacher + fake_score.train() + fake_score.requires_grad_(True) + return fake_score diff --git a/modelopt/torch/fastgen/flow_matching.py b/modelopt/torch/fastgen/flow_matching.py new file mode 100644 index 00000000000..66925dd00d1 --- /dev/null +++ b/modelopt/torch/fastgen/flow_matching.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Rectified-flow (RF) helpers: forward process, inversions, timestep sampling. + +This module intentionally does **not** define a ``NoiseScheduler`` class. It exposes the +handful of primitives that DMD2 actually needs as plain functions, so callers can plug +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. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +from torch.distributions import Normal + +from .utils import expand_like + +if TYPE_CHECKING: + from .config import SampleTimestepConfig + +__all__ = [ + "add_noise", + "pred_noise_to_pred_x0", + "pred_x0_from_flow", + "rf_alpha", + "rf_sigma", + "sample_from_t_list", + "sample_timesteps", + "x0_to_eps", + "x0_to_flow", +] + + +def rf_alpha(t: torch.Tensor) -> torch.Tensor: + """Rectified-flow data coefficient ``alpha_t = 1 - t``.""" + return 1.0 - t + + +def rf_sigma(t: torch.Tensor) -> torch.Tensor: + """Rectified-flow noise coefficient ``sigma_t = t``.""" + return t + + +def add_noise(x0: torch.Tensor, eps: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Forward process under rectified flow: ``x_t = (1 - t) * x_0 + t * eps``. + + ``t`` is broadcast across the spatial axes of ``x_0`` via :func:`expand_like`. + Computation is performed in ``float64`` for numerical stability and the output is + cast back to ``x_0``'s dtype. + """ + original_dtype = x0.dtype + x0_64 = x0.to(torch.float64) + eps_64 = eps.to(torch.float64) + t_64 = t.to(torch.float64) + alpha = expand_like(rf_alpha(t_64), x0_64) + sigma = expand_like(rf_sigma(t_64), x0_64) + x_t = x0_64 * alpha + eps_64 * sigma + return x_t.to(original_dtype) + + +def pred_noise_to_pred_x0( + pred_noise: torch.Tensor, + noisy_latents: torch.Tensor, + t: torch.Tensor, +) -> torch.Tensor: + """Convert an ``eps``-parameterized prediction to an ``x_0`` prediction under RF. + + Solves ``x_t = (1 - t) * x_0 + t * eps`` for ``x_0``: + ``x_0 = (x_t - t * eps) / (1 - t)``. + """ + original_dtype = noisy_latents.dtype + x_t = noisy_latents.to(torch.float64) + pred_noise_64 = pred_noise.to(torch.float64) + t_64 = t.to(torch.float64) + alpha = expand_like(rf_alpha(t_64), x_t) + sigma = expand_like(rf_sigma(t_64), x_t) + x0 = (x_t - sigma * pred_noise_64) / alpha.clamp_min(1e-6) + return x0.to(original_dtype) + + +def pred_x0_from_flow( + pred_flow: torch.Tensor, + noisy_latents: torch.Tensor, + t: torch.Tensor, +) -> torch.Tensor: + """Convert a flow-parameterized prediction (``v = eps - x_0``) to an ``x_0`` prediction. + + Under RF ``x_t = (1 - t) * x_0 + t * eps`` and ``v = eps - x_0`` combine to + ``x_t = x_0 + t * v``, so ``x_0 = x_t - t * v``. + """ + original_dtype = noisy_latents.dtype + x_t = noisy_latents.to(torch.float64) + v = pred_flow.to(torch.float64) + t_64 = t.to(torch.float64) + sigma = expand_like(rf_sigma(t_64), x_t) + x0 = x_t - sigma * v + return x0.to(original_dtype) + + +def x0_to_eps( + x0: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, +) -> torch.Tensor: + """Invert the RF forward process: ``eps = (x_t - (1 - t) * x_0) / t``. + + Used when unrolling the student in ODE mode — given the current ``x_t`` and the + student's ``x_0`` prediction, we can recover the implied ``eps`` deterministically. + """ + original_dtype = x0.dtype + x0_64 = x0.to(torch.float64) + x_t_64 = x_t.to(torch.float64) + t_64 = t.to(torch.float64) + alpha = expand_like(rf_alpha(t_64), x0_64) + sigma = expand_like(rf_sigma(t_64), x0_64) + eps = (x_t_64 - alpha * x0_64) / sigma.clamp_min(1e-6) + return eps.to(original_dtype) + + +def x0_to_flow( + x0: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, +) -> torch.Tensor: + """Convert an ``x_0`` prediction back into a flow-parameterized prediction under RF. + + Under RF ``x_t = (1 - t) * x_0 + t * eps`` and ``v = eps - x_0``, so + ``x_t - x_0 = t * (eps - x_0) = t * v`` and therefore ``v = (x_t - x_0) / t``. + + Used when the fake score is flow-native but the DSM loss is computed in a + different target parameterization: convert raw flow → x_0 via + :func:`pred_x0_from_flow`, then back to the loss space (which may coincide + with flow, in which case the round-trip is identity up to fp64 round-off). + """ + original_dtype = x0.dtype + x0_64 = x0.to(torch.float64) + x_t_64 = x_t.to(torch.float64) + t_64 = t.to(torch.float64) + sigma = expand_like(rf_sigma(t_64), x0_64) + flow = (x_t_64 - x0_64) / sigma.clamp_min(1e-6) + return flow.to(original_dtype) + + +# ---------------------------------------------------------------------------- # +# Timestep sampling # +# ---------------------------------------------------------------------------- # + + +def _truncated_lognormal( + n: int, + mean: float, + std: float, + *, + min_t: float, + max_t: float, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Sample ``n`` values from a lognormal truncated to ``(min_t, max_t)``. + + Implementation ported from ``FastGen/fastgen/networks/noise_schedule.py`` (EDM + ``_truncated_lognormal_sample``). Uses CDF inversion on the underlying normal for + exact truncation. + """ + min_t = max(min_t, 1e-12) + log_min_t = torch.tensor(math.log(min_t), dtype=torch.float64) + log_max_t = torch.tensor(math.log(max_t), dtype=torch.float64) + normal = Normal( + torch.tensor(mean, dtype=torch.float64), + torch.tensor(std, dtype=torch.float64), + ) + cdf_min = normal.cdf(log_min_t) + cdf_max = normal.cdf(log_max_t) + u = torch.rand(n, dtype=torch.float64) * (cdf_max - cdf_min) + cdf_min + t = normal.icdf(u).exp() + return t.to(device=device, dtype=dtype) + + +def sample_timesteps( + n: int, + cfg: SampleTimestepConfig, + *, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Sample ``n`` training timesteps according to ``cfg``. + + Supports ``uniform``, ``logitnormal``, ``lognormal``, ``shifted``, and ``polynomial`` + distributions. ``polynomial`` for RF degenerates to discrete uniform sampling from a + ``linspace(min_t, max_t, 1000)`` grid (EDM's polynomial-spaced ``_sigmas`` is + EDM-specific and not applicable under RF). + """ + min_t = cfg.min_t + max_t = cfg.max_t + + if cfg.time_dist_type == "uniform": + t = torch.rand(n, device=device, dtype=dtype) * (max_t - min_t) + min_t + elif cfg.time_dist_type == "logitnormal": + t = ( + torch.sigmoid(torch.randn(n, device=device, dtype=dtype) * cfg.p_std + cfg.p_mean) + * (max_t - min_t) + + min_t + ) + elif cfg.time_dist_type == "lognormal": + t = _truncated_lognormal( + n, + cfg.p_mean, + cfg.p_std, + min_t=min_t, + max_t=max_t, + device=device, + dtype=dtype, + ) + elif cfg.time_dist_type == "shifted": + t = torch.rand(n, device=device, dtype=dtype) * (max_t - min_t) + min_t + t = t * cfg.shift / (t * (cfg.shift - 1.0) + 1.0) + elif cfg.time_dist_type == "polynomial": + grid = torch.linspace(min_t, max_t, 1000, device=device, dtype=dtype) + idx = torch.randint(0, grid.numel(), (n,), device=device) + t = grid[idx] + else: + raise ValueError(f"Unsupported time_dist_type={cfg.time_dist_type!r}") + + return t.clamp(min_t, max_t) + + +def sample_from_t_list( + n: int, + t_list: list[float], + *, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Sample ``n`` starting timesteps uniformly from ``t_list[:-1]``. + + Used for multi-step student training: ``t_list`` encodes the inference trajectory + (``t_list[-1]`` must be ``0``), and a random intermediate timestep is sampled so the + student is trained at every rung of the trajectory. + """ + assert len(t_list) >= 2, "t_list must have at least 2 entries (including the final 0)" + assert t_list[-1] == 0.0, f"t_list[-1] must be 0.0, got {t_list[-1]}" + t_tensor = torch.tensor(t_list, device=device, dtype=dtype) + ids = torch.randint(0, t_tensor.numel() - 1, (n,), device=device) + return t_tensor[ids] diff --git a/modelopt/torch/fastgen/loader.py b/modelopt/torch/fastgen/loader.py new file mode 100644 index 00000000000..175ffa23661 --- /dev/null +++ b/modelopt/torch/fastgen/loader.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""YAML-driven configuration loading for fastgen distillation pipelines. + +YAML is the first-class entry point for DMD configurations — the fastgen library +does not expect callers to hand-build Python dicts. Typical usage:: + + from modelopt.torch.fastgen import DMDConfig, load_dmd_config + + # (a) Load a built-in recipe by relative path + cfg = load_dmd_config("general/distillation/dmd2_qwen_image") + + # (b) Load a user-provided file + cfg = load_dmd_config("/path/to/my_dmd.yaml") + + # (c) Equivalent classmethod + cfg = DMDConfig.from_yaml("/path/to/my_dmd.yaml") + +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). + +Suffixes ``.yml`` and ``.yaml`` are both accepted. +""" + +from __future__ import annotations + +import contextlib +from importlib.resources import files +from pathlib import Path +from typing import TYPE_CHECKING, Any + +# ``Traversable`` moved out of ``importlib.abc`` in Python 3.11. We only need it for +# type hints, but suppress ImportError so older runtimes can still import this module. +with contextlib.suppress(ImportError): + from importlib.resources.abc import Traversable + +import yaml + +from .config import DMDConfig + +if TYPE_CHECKING: + from importlib.abc import Traversable + +__all__ = ["load_config", "load_dmd_config"] + + +# Root to all built-in recipes shipped with modelopt. +_BUILTIN_RECIPES_LIB = files("modelopt_recipes") + + +_SUFFIXES = (".yml", ".yaml") + + +def _candidate_paths(config_file: str | Path) -> list[Path | Traversable]: + """Return the ordered list of locations to probe for ``config_file``.""" + candidates: list[Path | Traversable] = [] + + # Normalize to string for suffix probing; keep Path/Traversable behavior otherwise. + if isinstance(config_file, str): + base = config_file + if base.endswith(_SUFFIXES): + candidates.append(Path(base)) + candidates.append(_BUILTIN_RECIPES_LIB.joinpath(base)) + else: + candidates.extend(Path(base + suffix) for suffix in _SUFFIXES) + candidates.extend(_BUILTIN_RECIPES_LIB.joinpath(base + suffix) for suffix in _SUFFIXES) + elif isinstance(config_file, Path): + if config_file.suffix in _SUFFIXES: + candidates.append(config_file) + if not config_file.is_absolute(): + candidates.append(_BUILTIN_RECIPES_LIB.joinpath(str(config_file))) + else: + candidates.extend(Path(str(config_file) + suffix) for suffix in _SUFFIXES) + if not config_file.is_absolute(): + candidates.extend( + _BUILTIN_RECIPES_LIB.joinpath(str(config_file) + suffix) for suffix in _SUFFIXES + ) + else: + raise TypeError( + f"Expected str or Path for config_file, got {type(config_file).__name__!r}." + ) + return candidates + + +def load_config(config_file: str | Path) -> dict[str, Any]: + """Load a YAML file and return the parsed mapping. + + Mirrors :func:`modelopt.recipe._config_loader.load_config` in spirit but without + 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. + + Returns: + The parsed dictionary. An empty file yields ``{}``. + """ + for candidate in _candidate_paths(config_file): + if candidate.is_file(): + data = yaml.safe_load(candidate.read_text(encoding="utf-8")) + if data is None: + return {} + if not isinstance(data, dict): + raise ValueError( + f"Config file {candidate!s} must contain a YAML mapping, got {type(data).__name__}." + ) + return data + raise FileNotFoundError( + f"Cannot locate config file {config_file!r}; searched both the built-in " + f"recipe library and the filesystem." + ) + + +def load_dmd_config(config_file: str | Path) -> DMDConfig: + """Load a YAML file and construct a :class:`DMDConfig`. + + The YAML is validated against :class:`DMDConfig`'s Pydantic schema — unknown keys + raise ``ValidationError``. + + Example YAML:: + + pred_type: flow + guidance_scale: 5.0 + student_sample_steps: 2 + gan_loss_weight_gen: 0.03 + sample_t_cfg: + time_dist_type: shifted + t_list: [0.999, 0.833, 0.0] + ema: + decay: 0.9999 + """ + data = load_config(config_file) + return DMDConfig(**data) diff --git a/modelopt/torch/fastgen/losses.py b/modelopt/torch/fastgen/losses.py new file mode 100644 index 00000000000..8b5b1c9faf2 --- /dev/null +++ b/modelopt/torch/fastgen/losses.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Pure loss functions used by the fastgen distillation pipelines. + +All functions in this module are stateless: they take tensors in and return a scalar +loss tensor. They do not touch any ``nn.Module``. Higher-level orchestration (teacher +forward, CFG, noise scheduling) lives in :mod:`modelopt.torch.fastgen.methods.dmd`. + +Math ported from ``FastGen/fastgen/methods/common_loss.py`` (lines 12-136) and +``FastGen/fastgen/methods/distribution_matching/dmd2.py`` lines 287-317 (R1). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F + +from .utils import expand_like + +if TYPE_CHECKING: + from collections.abc import Callable + +__all__ = [ + "dsm_loss", + "gan_disc_loss", + "gan_gen_loss", + "r1_loss", + "vsd_loss", +] + + +def dsm_loss( + pred_type: str, + net_pred: torch.Tensor, + *, + x0: torch.Tensor | None = None, + eps: torch.Tensor | None = None, + t: torch.Tensor | None = None, + alpha_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, + sigma_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, +) -> torch.Tensor: + """Denoising score-matching loss for ``x0`` / ``eps`` / ``v`` / ``flow`` predictions. + + The forward process is ``x_t = alpha_t * x_0 + sigma_t * eps``. For + ``pred_type='v'`` we need ``alpha_t`` and ``sigma_t``, which are supplied as + callables rather than a full noise-scheduler object so this function stays + scheduler-agnostic. + + Args: + pred_type: One of ``"x0"``, ``"eps"``, ``"v"``, ``"flow"``. + net_pred: The network output; its interpretation is determined by ``pred_type``. + x0: Clean data. Required for all ``pred_type`` except ``"eps"``. + eps: Noise used in the forward process. Required for all ``pred_type`` except ``"x0"``. + t: Timesteps in ``[0, 1]`` (or scheduler convention). Required for ``pred_type='v'``. + alpha_fn: Callable mapping ``t`` -> ``alpha_t``. Required for ``pred_type='v'``. + sigma_fn: Callable mapping ``t`` -> ``sigma_t``. Required for ``pred_type='v'``. + + Returns: + Scalar MSE loss. + """ + if pred_type == "x0": + assert x0 is not None, "x0 is required for pred_type='x0'" + return F.mse_loss(x0, net_pred, reduction="mean") + if pred_type == "eps": + assert eps is not None, "eps is required for pred_type='eps'" + return F.mse_loss(eps, net_pred, reduction="mean") + if pred_type == "v": + assert x0 is not None and eps is not None and t is not None, ( + "x0, eps, and t are required for pred_type='v'" + ) + assert alpha_fn is not None and sigma_fn is not None, ( + "alpha_fn and sigma_fn are required for pred_type='v'" + ) + alpha_t = expand_like(alpha_fn(t), x0).to(device=x0.device, dtype=x0.dtype) + sigma_t = expand_like(sigma_fn(t), x0).to(device=x0.device, dtype=x0.dtype) + v = alpha_t * eps - sigma_t * x0 + return F.mse_loss(v, net_pred, reduction="mean") + if pred_type == "flow": + assert x0 is not None and eps is not None, "x0 and eps are required for pred_type='flow'" + flow_velocity = eps - x0 + return F.mse_loss(flow_velocity, net_pred, reduction="mean") + raise ValueError(f"Unknown pred_type {pred_type!r}; expected one of 'x0', 'eps', 'v', 'flow'.") + + +def vsd_loss( + gen_data: torch.Tensor, + teacher_x0: torch.Tensor, + fake_score_x0: torch.Tensor, + additional_scale: torch.Tensor | None = None, +) -> torch.Tensor: + """Variational score-distillation (VSD) loss used by the DMD student update. + + Implements the FastGen formulation: a per-sample weight + ``w = 1 / (mean_abs(gen_data - teacher_x0) + 1e-6)`` is computed in fp32 for + numerical stability, then the gradient ``(fake_score_x0 - teacher_x0) * w`` is + subtracted from the generated data to form a pseudo-target. The loss is + ``0.5 * MSE(gen_data, pseudo_target)``. + + Args: + gen_data: Student-generated clean data ``x_0``. + teacher_x0: Teacher ``x_0`` prediction (after CFG, if enabled). Detached. + fake_score_x0: Fake-score ``x_0`` prediction. Detached. + additional_scale: Optional per-sample scale applied multiplicatively to the weight. + + Returns: + Scalar VSD loss. + """ + dims = tuple(range(1, teacher_x0.ndim)) + + with torch.no_grad(): + original_dtype = gen_data.dtype + gen_data_fp32 = gen_data.float() + teacher_x0_fp32 = teacher_x0.float() + + diff_abs_mean = (gen_data_fp32 - teacher_x0_fp32).abs().mean(dim=dims, keepdim=True) + w_fp32 = 1.0 / (diff_abs_mean + 1e-6) + + if additional_scale is not None: + w_fp32 = w_fp32 * expand_like(additional_scale.float(), w_fp32) + + w = w_fp32.to(dtype=original_dtype) + vsd_grad = (fake_score_x0 - teacher_x0) * w + pseudo_target = gen_data - vsd_grad + + return 0.5 * F.mse_loss(gen_data, pseudo_target, reduction="mean") + + +def gan_gen_loss(fake_logits: torch.Tensor) -> torch.Tensor: + """Softplus GAN generator loss: ``E[softplus(-fake_logits)]``. + + Args: + fake_logits: Discriminator logits on generated samples. Must be 2D: ``(B, num_heads)``. + """ + assert fake_logits.ndim == 2, f"fake_logits must be 2D, got shape {tuple(fake_logits.shape)}" + return F.softplus(-fake_logits).mean() + + +def gan_disc_loss(real_logits: torch.Tensor, fake_logits: torch.Tensor) -> torch.Tensor: + """Softplus GAN discriminator loss: ``E[softplus(fake_logits)] + E[softplus(-real_logits)]``.""" + assert real_logits.ndim == 2, f"real_logits must be 2D, got shape {tuple(real_logits.shape)}" + assert fake_logits.ndim == 2, f"fake_logits must be 2D, got shape {tuple(fake_logits.shape)}" + return F.softplus(fake_logits).mean() + F.softplus(-real_logits).mean() + + +def r1_loss( + real_logits: torch.Tensor, + perturbed_real_logits: torch.Tensor, +) -> torch.Tensor: + """Approximate R1 regularization (APT formulation). + + Penalizes the discriminator for being sensitive to small noise perturbations of + the real data. The caller is responsible for computing ``perturbed_real_logits`` + by re-running the teacher feature extractor and discriminator on real data that + has been perturbed with ``alpha * randn_like(real)``; this function only applies + the final MSE between the two logit sets. + + See ``FastGen/fastgen/methods/distribution_matching/dmd2.py`` lines 287-317. + """ + assert real_logits.shape == perturbed_real_logits.shape, ( + f"real_logits {tuple(real_logits.shape)} and perturbed_real_logits " + f"{tuple(perturbed_real_logits.shape)} must have matching shapes" + ) + return F.mse_loss(real_logits, perturbed_real_logits, reduction="mean") diff --git a/modelopt/torch/fastgen/methods/__init__.py b/modelopt/torch/fastgen/methods/__init__.py new file mode 100644 index 00000000000..c999ca87f6f --- /dev/null +++ b/modelopt/torch/fastgen/methods/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Concrete distillation method implementations (DMD, future: Self-Forcing, CausVid, ...).""" + +from .dmd import * diff --git a/modelopt/torch/fastgen/methods/dmd.py b/modelopt/torch/fastgen/methods/dmd.py new file mode 100644 index 00000000000..a9d9add2439 --- /dev/null +++ b/modelopt/torch/fastgen/methods/dmd.py @@ -0,0 +1,812 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Distribution Matching Distillation (DMD2) pipeline. + +:class:`DMDPipeline` holds references to the student / teacher / fake-score / (optional) +discriminator and exposes the three loss-computation entry points that a training loop +calls from each update step: + +- :meth:`DMDPipeline.compute_student_loss` — variational score-distillation loss plus an optional + GAN generator term. +- :meth:`DMDPipeline.compute_fake_score_loss` — denoising score matching against the student's + generated samples. +- :meth:`DMDPipeline.compute_discriminator_loss` — GAN discriminator loss plus an optional R1 + regularizer. + +The pipeline does **not** own optimizers, schedulers, gradient toggles, or device placement. +Callers drive the alternation between student / fake-score / discriminator updates, toggle +``requires_grad``, and call the appropriate ``compute_*_loss`` each step. + +Math is a close port of ``FastGen/fastgen/methods/distribution_matching/dmd2.py`` (lines +45-455). See the docstrings on the individual methods for line-level references. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +import torch.distributed as dist +from torch import nn + +from ..ema import ExponentialMovingAverage +from ..flow_matching import ( + add_noise, + pred_noise_to_pred_x0, + pred_x0_from_flow, + rf_alpha, + rf_sigma, + sample_from_t_list, + x0_to_eps, + x0_to_flow, +) +from ..losses import dsm_loss, gan_disc_loss, gan_gen_loss, r1_loss, vsd_loss +from ..pipeline import DistillationPipeline +from ..utils import classifier_free_guidance + +if TYPE_CHECKING: + from ..config import DMDConfig + +__all__ = ["DMDPipeline"] + + +# ---------------------------------------------------------------------------- # +# Feature capture helper (duck-typed so tests can bypass the capture plugin) # +# ---------------------------------------------------------------------------- # + + +def _drain_if_hooked(module: nn.Module) -> list[torch.Tensor] | None: + """Drain the feature-capture buffer on ``module`` if hooks are attached. + + Returns the captured tensors in insertion order (clearing the buffer in-place), or + ``None`` when no hooks are installed. Non-raising by design so pipeline-internal + call sites can drain unconditionally after every teacher forward — this prevents + the buffer from growing across steps when hooks are attached but the GAN branch is + disabled (e.g. an ablation). Callers that need the strict "did you forget to attach + hooks?" failure mode should call :func:`_require_hooked` on the result. + """ + captured = getattr(module, "_fastgen_captured", None) + if captured is None: + return None + out = list(captured) + captured.clear() + return out + + +def _require_hooked( + features: list[torch.Tensor] | None, + *, + which: str, +) -> list[torch.Tensor]: + """Adapter that turns a ``None`` drain result into a clear ``RuntimeError``. + + Use at pipeline sites that *must* consume captured features (i.e. the GAN-enabled + paths in ``compute_student_loss`` / ``compute_discriminator_loss``). Keeps the + non-raising ``_drain_if_hooked`` primitive for the "drain-and-discard" sites. + + The message names the attribute the pipeline looks for + (``teacher._fastgen_captured``) so a debugger can grep straight to the hook + installation site. + """ + if features is None: + raise RuntimeError( + f"Feature-capture hooks are required on the teacher ({which} branch): " + "teacher._fastgen_captured is missing. Call " + "modelopt.torch.fastgen.plugins.qwen_image.attach_feature_capture(teacher, ...) " + "before running this loss." + ) + return features + + +# ---------------------------------------------------------------------------- # +# DMDPipeline # +# ---------------------------------------------------------------------------- # + + +class DMDPipeline(DistillationPipeline): + """DMD2 loss pipeline. + + Args: + student: Trainable student module. Must be callable with ``(hidden_states, timestep, + encoder_hidden_states=..., **kwargs)`` and return either a ``Tensor``, a + ``(Tensor, ...)`` tuple (as diffusers returns with ``return_dict=False``), or an + object with a ``.sample`` attribute. + teacher: Frozen reference module with the same call signature. If ``discriminator`` + is provided, feature-capture hooks must be attached to ``teacher`` before + calling ``compute_*_loss`` — see :func:`modelopt.torch.fastgen.plugins.qwen_image.attach_feature_capture`. + fake_score: Trainable auxiliary module (same signature as teacher/student). Used to + approximate the student's generated distribution for the VSD gradient. + config: :class:`~modelopt.torch.fastgen.config.DMDConfig` with the hyperparameters. + discriminator: Optional discriminator. Required when ``config.gan_loss_weight_gen > 0``. + Must accept ``list[Tensor]`` (the captured teacher features) and return a 2D logit tensor. + """ + + def __init__( + self, + student: nn.Module, + teacher: nn.Module, + fake_score: nn.Module, + config: DMDConfig, + *, + discriminator: nn.Module | None = None, + ) -> None: + """Wire up student / teacher / fake-score / discriminator and create the EMA tracker.""" + super().__init__(student, teacher, config) + self.fake_score = fake_score + self.discriminator = discriminator + self._ema: ExponentialMovingAverage | None = ( + ExponentialMovingAverage(student, config.ema) if config.ema is not None else None + ) + self._iteration = 0 + + if config.gan_loss_weight_gen > 0 and discriminator is None: + raise ValueError( + "gan_loss_weight_gen > 0 requires a discriminator to be provided to DMDPipeline." + ) + + # Re-declare config at the class level so type checkers see ``DMDConfig`` here + # even though the base class stores it as ``DistillationConfig``. At runtime the + # attribute is set by :meth:`DistillationPipeline.__init__`. + config: DMDConfig + + @property + def ema(self) -> ExponentialMovingAverage | None: + """Reference to the student EMA tracker, if configured.""" + return self._ema + + # ================================================================== # + # Model-call helpers # + # ================================================================== # + + def _call_model( + self, + model: nn.Module, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Forward a diffusers-style transformer and return the raw prediction tensor. + + Assumes the target module accepts ``hidden_states`` / ``timestep`` / + ``encoder_hidden_states`` as kwargs and returns one of: + + * a ``torch.Tensor`` (custom modules), + * a ``tuple`` whose first element is the prediction (diffusers ``return_dict=False``), + * an object with a ``.sample`` attribute (diffusers ``return_dict=True``). + + **Timestep convention.** ``timestep`` is passed verbatim to the model by default. + If :attr:`DistillationConfig.num_train_timesteps` is set, the continuous RF time + ``t ∈ [0, 1]`` is rescaled to ``num_train_timesteps * t`` before the call — which + matches the diffusers training convention for Wan 2.2, SD3, Flux. Leave + ``num_train_timesteps=None`` when the upstream model wrapper (e.g. a VaceWan-style + module) already scales the timestep internally. + + Subclass and override this method for modules with non-diffusers signatures + (positional-only args, alternate kwarg names) or bespoke timestep transforms. + """ + call_kwargs: dict[str, Any] = dict(model_kwargs) + call_kwargs["hidden_states"] = hidden_states + if self.config.num_train_timesteps is not None: + # Cast to match the hidden-state dtype, mirroring FastGen's VaceWan + # wrapper (``noise_scheduler.rescale_t(t).to(dtype=x_t.dtype)``). + timestep = (timestep * float(self.config.num_train_timesteps)).to( + dtype=hidden_states.dtype + ) + call_kwargs["timestep"] = timestep + if encoder_hidden_states is not None: + call_kwargs["encoder_hidden_states"] = encoder_hidden_states + + out = model(**call_kwargs) + if isinstance(out, torch.Tensor): + return out + if isinstance(out, tuple): + return out[0] + if hasattr(out, "sample"): + return out.sample + raise TypeError( + f"DMDPipeline._call_model could not extract a tensor from output of type " + f"{type(out).__name__!r}. Override ``_call_model`` in a subclass to handle " + f"custom module signatures." + ) + + @staticmethod + def _raw_to_x0( + raw: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, + *, + native_pred_type: str, + ) -> torch.Tensor: + """Convert a raw model output in ``native_pred_type`` space to an ``x_0`` estimate. + + ``native_pred_type`` is the parameterization the module *actually* predicts + (i.e. its architecture's native output), not the space a downstream loss + wants to operate in. Under RF, ``flow`` and ``v`` are equivalent (both are + ``eps - x_0``). + """ + if native_pred_type == "x0": + return raw + if native_pred_type == "eps": + return pred_noise_to_pred_x0(raw, x_t, t) + if native_pred_type in ("flow", "v"): + return pred_x0_from_flow(raw, x_t, t) + raise ValueError(f"Unsupported native_pred_type={native_pred_type!r}") + + @staticmethod + def _x0_to_raw( + x0: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, + *, + target_pred_type: str, + ) -> torch.Tensor: + """Inverse of :meth:`_raw_to_x0` — project an ``x_0`` estimate into ``target_pred_type`` space.""" + if target_pred_type == "x0": + return x0 + if target_pred_type == "eps": + return x0_to_eps(x0, x_t, t) + if target_pred_type in ("flow", "v"): + return x0_to_flow(x0, x_t, t) + raise ValueError(f"Unsupported target_pred_type={target_pred_type!r}") + + def _convert_pred( + self, + raw: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, + *, + from_pred_type: str, + to_pred_type: str, + ) -> torch.Tensor: + """Project a prediction between parameterizations via the ``x_0`` hub. + + Used by :meth:`compute_fake_score_loss` to land the fake-score's raw + output in the DSM loss's target space. Short-circuits to identity when + both spaces agree. + """ + if from_pred_type == to_pred_type: + return raw + x0 = self._raw_to_x0(raw, x_t, t, native_pred_type=from_pred_type) + if to_pred_type == "x0": + return x0 + return self._x0_to_raw(x0, x_t, t, target_pred_type=to_pred_type) + + def _predict_x0( + self, + model: nn.Module, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + *, + native_pred_type: str | None = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Forward ``model`` and return its ``x_0`` estimate. + + ``native_pred_type`` declares the module's **architectural** output + parameterization — NOT any downstream loss's target space. In the DMD2 + setup the student / teacher / fake_score are arch-twins, so this defaults + to :attr:`DistillationConfig.pred_type`; callers should only override it + when wiring in a model whose architecture genuinely differs. + """ + raw = self._call_model( + model, hidden_states, timestep, encoder_hidden_states, **model_kwargs + ) + native_pred_type = native_pred_type or self.config.pred_type + return self._raw_to_x0(raw, hidden_states, timestep, native_pred_type=native_pred_type) + + # ================================================================== # + # Noise / timestep sampling # + # ================================================================== # + + def _build_backward_simulated_student_input( + self, + noise: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build a multi-step student input by no-grad unrolling the student. + + This mirrors the SDXL DMD2 ``--backward_simulation`` idea in RF space: + choose a schedule rung, generate an x0 distribution by running the current + student through all earlier rungs, then re-noise that generated x0 at the + selected rung. ``student_sample_type`` controls whether intermediate + transitions reuse the implied ODE noise or draw fresh SDE noise. + """ + cfg = self.config + t_list = cfg.sample_t_cfg.t_list + if t_list is None: + raise ValueError( + "backward_simulation=True requires DMDConfig.sample_t_cfg.t_list to be set." + ) + if len(t_list) != cfg.student_sample_steps + 1: + raise ValueError( + "backward_simulation=True expects len(sample_t_cfg.t_list) == " + "student_sample_steps + 1, got " + f"{len(t_list)} vs {cfg.student_sample_steps + 1}." + ) + + batch_size = noise.shape[0] + device = noise.device + dtype = noise.dtype + num_train_rungs = len(t_list) - 1 + selected_idx_tensor = torch.randint( + 0, num_train_rungs, (1,), device=device, dtype=torch.long + ) + if dist.is_available() and dist.is_initialized(): + dist.broadcast(selected_idx_tensor, src=0) + selected_idx = int(selected_idx_tensor.item()) + t_student = torch.full( + (batch_size,), float(t_list[selected_idx]), device=device, dtype=torch.float32 + ) + + # First rung is the initial RF noise state, matching inference's + # ``latents = noise * schedule[0]`` and SDXL's pure-noise special case. + if selected_idx == 0: + input_student = (noise.to(torch.float64) * float(t_list[0])).to(dtype) + return input_student, t_student + + current = (torch.randn_like(noise).to(torch.float64) * float(t_list[0])).to(dtype) + generated_x0: torch.Tensor | None = None + with torch.no_grad(): + for step_idx in range(selected_idx): + t_cur = torch.full( + (batch_size,), float(t_list[step_idx]), device=device, dtype=torch.float32 + ) + generated_x0 = self._predict_x0( + self.student, + current, + t_cur, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + + if step_idx == selected_idx - 1: + break + + t_next = torch.full( + (batch_size,), + float(t_list[step_idx + 1]), + device=device, + dtype=torch.float32, + ) + if cfg.student_sample_type == "ode": + step_noise = x0_to_eps(generated_x0, current, t_cur) + elif cfg.student_sample_type == "sde": + step_noise = torch.randn_like(noise) + else: + raise ValueError( + "student_sample_type must be one of {'ode', 'sde'}, got " + f"{cfg.student_sample_type!r}." + ) + current = add_noise(generated_x0, step_noise, t_next) + + if generated_x0 is None: + raise RuntimeError("backward simulation did not produce a generated x0.") + input_student = add_noise(generated_x0.detach(), noise, t_student) + return input_student, t_student + + def _build_student_input( + self, + latents: torch.Tensor, + noise: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Construct ``(input_student, t_student)`` for the student forward pass. + + - ``student_sample_steps == 1``: Use the maximum training timestep and set the + student's input to ``sigma(max_t) * noise = max_t * noise`` (RF). + - ``student_sample_steps > 1`` and ``backward_simulation=False``: sample a + random intermediate timestep from ``config.sample_t_cfg.t_list`` and + noise the real latents up to that timestep. + - ``student_sample_steps > 1`` and ``backward_simulation=True``: no-grad + unroll the current student to the selected rung and noise that generated + x0, matching the SDXL DMD2 backward-simulation training regime. + """ + cfg = self.config + batch_size = latents.shape[0] + device = latents.device + + if cfg.student_sample_steps == 1: + max_t = cfg.sample_t_cfg.max_t + t_student = torch.full((batch_size,), max_t, device=device, dtype=torch.float32) + # Under RF, ``sigma(max_t) = max_t``. Do the scaling in fp64 and cast back + # to mirror FastGen's ``BaseNoiseSchedule.latents`` — matters for bf16 + # student input at ``max_t ≈ 0.999`` where naive bf16 multiply loses + # ~10 bits of mantissa relative to the fp64 path. + original_dtype = noise.dtype + input_student = (noise.to(torch.float64) * float(max_t)).to(original_dtype) + else: + if cfg.sample_t_cfg.t_list is None: + raise ValueError( + "student_sample_steps > 1 requires DMDConfig.sample_t_cfg.t_list to be set." + ) + if cfg.backward_simulation: + return self._build_backward_simulated_student_input( + noise, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + t_student = sample_from_t_list( + batch_size, + cfg.sample_t_cfg.t_list, + device=device, + dtype=torch.float32, + ) + input_student = add_noise(latents, noise, t_student) + return input_student, t_student + + # ================================================================== # + # Public API # + # ================================================================== # + + def compute_student_loss( + self, + latents: torch.Tensor, + noise: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + *, + negative_encoder_hidden_states: torch.Tensor | None = None, + negative_encoder_hidden_states_mask: torch.Tensor | None = None, + guidance_scale: float | None = None, + **model_kwargs: Any, + ) -> dict[str, torch.Tensor]: + """Compute the student update losses. + + The returned dict always contains ``"vsd"`` and ``"total"``. When the GAN branch + is enabled (``discriminator is not None`` and ``config.gan_loss_weight_gen > 0``), + ``"gan_gen"`` is also present. + + Gradient flow summary: + + - VSD gradient: flows through ``student`` only (``teacher_x0`` is detached, + ``fake_score_x0`` is computed under ``torch.no_grad()``). + - GAN generator gradient: flows through ``student`` via the feature-capture + hooks on the teacher. The teacher forward is therefore **not** wrapped in + ``torch.no_grad()`` when the GAN branch is active. + + Args: + latents: Real clean-data latents ``x_0``. Used only when + ``student_sample_steps > 1`` to construct ``input_student``. + noise: Pure Gaussian noise tensor matching ``latents`` in shape/dtype. + encoder_hidden_states: Positive conditioning passed unchanged to all three + models. + negative_encoder_hidden_states: Negative conditioning used by classifier-free + guidance. Required when ``guidance_scale`` (or :attr:`DMDConfig.guidance_scale`) + is not ``None``. + negative_encoder_hidden_states_mask: Optional negative-conditioning mask. Used + for models such as Qwen-Image whose positional embedding depends on the + real text sequence length. + guidance_scale: Overrides :attr:`DMDConfig.guidance_scale` for this call. + ``None`` keeps the config-level value. + **model_kwargs: Forwarded verbatim to ``student``, ``teacher``, and ``fake_score``. + + Returns: + Dictionary with keys ``"vsd"``, ``"total"``, and optionally ``"gan_gen"``. + """ + cfg = self.config + batch_size = latents.shape[0] + device = latents.device + gan_enabled = self.discriminator is not None and cfg.gan_loss_weight_gen > 0 + + # 1. Student input. + input_student, t_student = self._build_student_input( + latents, + noise, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + + # 2. Student forward -> x0. + gen_data = self._predict_x0( + self.student, + input_student, + t_student, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + + # 3. Sample perturbation timesteps and noise, perturb gen_data. + t = self.sample_timesteps(batch_size, device=device, dtype=torch.float32) + eps = torch.randn_like(latents) + perturbed = add_noise(gen_data, eps, t) + + # 4. Fake score prediction (no grad). + # + # VSD always operates in x_0 space, regardless of ``fake_score_pred_type`` + # (which controls the DSM loss space on the fake-score side — see + # :meth:`compute_fake_score_loss`). The fake_score's architecture matches the + # student's in the DMD2 setup, so its native output parameterization is + # ``cfg.pred_type``; ``_predict_x0`` converts from that to x_0 automatically. + with torch.no_grad(): + fake_score_x0 = self._predict_x0( + self.fake_score, + perturbed, + t, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + + # 5. Teacher forward. + fake_feat: list[torch.Tensor] | None = None + if gan_enabled: + # Grad must flow through the teacher for the GAN generator term, since the + # captured features depend on perturbed -> gen_data -> student weights. + teacher_x0 = self._predict_x0( + self.teacher, + perturbed, + t, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + fake_feat = _require_hooked(_drain_if_hooked(self.teacher), which="student-fake") + else: + with torch.no_grad(): + teacher_x0 = self._predict_x0( + self.teacher, + perturbed, + t, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + # Drain any hooks attached but not consumed (e.g. hooks left over from a + # previous GAN-enabled ablation). No-op when hooks aren't installed. + _ = _drain_if_hooked(self.teacher) + + # 6. Classifier-free guidance (applied to teacher_x0 prior to detach). + effective_scale = guidance_scale if guidance_scale is not None else cfg.guidance_scale + if effective_scale is not None: + if negative_encoder_hidden_states is None: + raise ValueError( + "guidance_scale is set but negative_encoder_hidden_states was not provided." + ) + with torch.no_grad(): + negative_model_kwargs = dict(model_kwargs) + if negative_encoder_hidden_states_mask is not None: + negative_model_kwargs["encoder_hidden_states_mask"] = ( + negative_encoder_hidden_states_mask + ) + else: + negative_model_kwargs.pop("encoder_hidden_states_mask", None) + negative_model_kwargs.pop("txt_seq_lens", None) + teacher_x0_neg = self._predict_x0( + self.teacher, + perturbed, + t, + encoder_hidden_states=negative_encoder_hidden_states, + **negative_model_kwargs, + ) + # Negative-branch features are never used for GAN — drain unconditionally so + # the buffer stays clean for subsequent calls. + _ = _drain_if_hooked(self.teacher) + teacher_x0 = classifier_free_guidance(teacher_x0, teacher_x0_neg, effective_scale) + + teacher_x0 = teacher_x0.detach() + + # 7. Losses. + vsd = vsd_loss(gen_data, teacher_x0, fake_score_x0) + + if gan_enabled: + # ``fake_feat`` is guaranteed non-None by ``_require_hooked`` above; + # ``gan_enabled`` implies a discriminator was provided. + assert self.discriminator is not None + gan_gen = gan_gen_loss(self.discriminator(fake_feat)) + total = vsd + cfg.gan_loss_weight_gen * gan_gen + return {"vsd": vsd, "gan_gen": gan_gen, "total": total} + + return {"vsd": vsd, "total": vsd} + + def compute_fake_score_loss( + self, + latents: torch.Tensor, + noise: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> dict[str, torch.Tensor]: + """Compute the fake-score (auxiliary) update loss. + + The fake score is trained with denoising score matching against the student's + generated samples. The student forward is wrapped in ``torch.no_grad()`` — the + gradient here is w.r.t. ``fake_score`` only. + + Returns a dict with ``"fake_score"`` and ``"total"`` (both equal). + """ + cfg = self.config + batch_size = latents.shape[0] + device = latents.device + + # 1. Build student input. + input_student, t_student = self._build_student_input( + latents, + noise, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + + # 2. Generate data from student (no grad). + with torch.no_grad(): + gen_data = self._predict_x0( + self.student, + input_student, + t_student, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + + # 3. Perturb gen_data. + t = self.sample_timesteps(batch_size, device=device, dtype=torch.float32) + eps = torch.randn_like(latents) + perturbed = add_noise(gen_data, eps, t) + + # 4. Fake-score forward (grad flows here). + # + # The fake_score's architectural output parameterization is ``cfg.pred_type`` + # (same arch as teacher/student in DMD2). ``fake_score_pred_type`` controls + # which space the DSM loss is computed in — it is a loss-side knob, not a + # model-side one. When the two differ (e.g. the Wan 2.2 recipe with + # flow-native models and ``fake_score_pred_type='x0'``), we project the raw + # output through the ``x_0`` hub into the loss space before calling + # ``dsm_loss``. When they agree, ``_convert_pred`` short-circuits to identity. + fake_pred_type = cfg.fake_score_pred_type or cfg.pred_type + raw = self._call_model( + self.fake_score, + perturbed, + t, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + pred_in_loss_space = self._convert_pred( + raw, + perturbed, + t, + from_pred_type=cfg.pred_type, + to_pred_type=fake_pred_type, + ) + + # 5. DSM loss in the chosen parameterization. + loss = dsm_loss( + fake_pred_type, + pred_in_loss_space, + x0=gen_data, + eps=eps, + t=t, + alpha_fn=rf_alpha, + sigma_fn=rf_sigma, + ) + return {"fake_score": loss, "total": loss} + + def compute_discriminator_loss( + self, + latents: torch.Tensor, + noise: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> dict[str, torch.Tensor]: + """Compute the discriminator update loss (GAN + optional R1). + + Teacher and student forwards are wrapped in ``torch.no_grad()``; gradient flows + only through the discriminator. + + Returns a dict with ``"gan_disc"`` and ``"total"``. When + ``config.gan_r1_reg_weight > 0`` the dict also contains ``"r1"``. + """ + cfg = self.config + if self.discriminator is None: + raise RuntimeError( + "compute_discriminator_loss requires a discriminator to be set on DMDPipeline." + ) + batch_size = latents.shape[0] + device = latents.device + + # 1. Build student input and generate gen_data (no grad). + input_student, t_student = self._build_student_input( + latents, + noise, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + with torch.no_grad(): + gen_data = self._predict_x0( + self.student, + input_student, + t_student, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + + # 2. Sample fake-branch timesteps and noise. + t = self.sample_timesteps(batch_size, device=device, dtype=torch.float32) + eps = torch.randn_like(latents) + perturbed_fake = add_noise(gen_data, eps, t) + + # 3. Teacher forward on fake data to capture features. + _ = self._predict_x0( + self.teacher, + perturbed_fake, + t, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + fake_feat = _require_hooked(_drain_if_hooked(self.teacher), which="disc-fake") + + # 4. Real branch: same t/eps or re-sampled. + if cfg.gan_use_same_t_noise: + t_real = t + eps_real = eps + else: + t_real = self.sample_timesteps(batch_size, device=device, dtype=torch.float32) + eps_real = torch.randn_like(latents) + perturbed_real = add_noise(latents, eps_real, t_real) + + _ = self._predict_x0( + self.teacher, + perturbed_real, + t_real, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + real_feat = _require_hooked(_drain_if_hooked(self.teacher), which="disc-real") + + # 5. Discriminator on real / fake (grad required). + real_logits = self.discriminator(real_feat) + fake_logits = self.discriminator(fake_feat) + disc = gan_disc_loss(real_logits, fake_logits) + + result: dict[str, torch.Tensor] = {"gan_disc": disc} + + # 6. Optional R1 regularization. + if cfg.gan_r1_reg_weight > 0: + with torch.no_grad(): + perturbed_real_alpha = latents + cfg.gan_r1_reg_alpha * torch.randn_like(latents) + _ = self._predict_x0( + self.teacher, + perturbed_real_alpha, + t_real, + encoder_hidden_states=encoder_hidden_states, + **model_kwargs, + ) + real_feat_alpha = _require_hooked(_drain_if_hooked(self.teacher), which="disc-r1") + real_logits_alpha = self.discriminator(real_feat_alpha) + r1 = r1_loss(real_logits, real_logits_alpha) + total = disc + cfg.gan_r1_reg_weight * r1 + result["r1"] = r1 + result["total"] = total + else: + result["total"] = disc + return result + + # ================================================================== # + # EMA # + # ================================================================== # + + def update_ema(self, *, iteration: int | None = None) -> None: + """Update the student EMA tracker (no-op if ``config.ema`` is ``None``). + + Typically called after the student optimizer step. If ``iteration`` is not + provided, an internal counter is auto-incremented. + """ + if self._ema is None: + return + if iteration is not None: + self._iteration = iteration + else: + # Counter starts at 0 and pre-increments, so the first auto call passes 1. + # With start_iter=0 the shadow is therefore first initialised via EMA.update's + # ``not self._initialized`` arm, not the ``iteration == start_iter`` one. + self._iteration += 1 + self._ema.update(self.student, iteration=self._iteration) diff --git a/modelopt/torch/fastgen/pipeline.py b/modelopt/torch/fastgen/pipeline.py new file mode 100644 index 00000000000..5c39f70d6ae --- /dev/null +++ b/modelopt/torch/fastgen/pipeline.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Base class for diffusion step-distillation pipelines. + +:class:`DistillationPipeline` is deliberately minimal: it is **not** an ``nn.Module``, +does not wrap the student or teacher, does not manage optimizers or lifecycle state, +and does not register itself in any mode registry. It exists only to hold references +to the student / teacher and to freeze the teacher in a single place. + +Concrete methods — for now :class:`~modelopt.torch.fastgen.methods.dmd.DMDPipeline` — +subclass this and add ``compute_*_loss`` methods. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch import nn + +from .flow_matching import sample_timesteps + +if TYPE_CHECKING: + from .config import DistillationConfig + +__all__ = ["DistillationPipeline"] + + +class DistillationPipeline: + """Hold student/teacher references and expose shared utilities. + + Args: + student: Trainable student module. The pipeline does not wrap it — its lifecycle + (``train()`` / ``eval()``, ``requires_grad_``, sharding, optimizer) remains + owned by the caller. + teacher: Reference module. Frozen here via ``eval()`` + ``requires_grad_(False)``. + config: A :class:`DistillationConfig` (or subclass). + """ + + def __init__( + self, + student: nn.Module, + teacher: nn.Module, + config: DistillationConfig, + ) -> None: + """Store student / teacher references and freeze the teacher.""" + self.student = student + self.teacher = teacher.eval().requires_grad_(False) + self.config = config + + # ------------------------------------------------------------------ # + # Device / dtype inferred from the student # + # ------------------------------------------------------------------ # + + @property + def device(self) -> torch.device: + """Device of the first student parameter (best-effort; falls back to CPU).""" + for p in self.student.parameters(): + return p.device + return torch.device("cpu") + + @property + def dtype(self) -> torch.dtype: + """Dtype of the first student parameter (best-effort; falls back to float32).""" + for p in self.student.parameters(): + return p.dtype + return torch.float32 + + # ------------------------------------------------------------------ # + # Shared helpers # + # ------------------------------------------------------------------ # + + def sample_timesteps( + self, + n: int, + *, + device: torch.device | None = None, + dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + """Sample ``n`` training timesteps according to :attr:`config`.``sample_t_cfg``.""" + return sample_timesteps( + n, + self.config.sample_t_cfg, + device=device or self.device, + dtype=dtype, + ) diff --git a/modelopt/torch/fastgen/plugins/__init__.py b/modelopt/torch/fastgen/plugins/__init__.py new file mode 100644 index 00000000000..8810470b26f --- /dev/null +++ b/modelopt/torch/fastgen/plugins/__init__.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""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. +""" + +from modelopt.torch.utils import import_plugin + +with import_plugin("qwen_image"): + from .qwen_image import * diff --git a/modelopt/torch/fastgen/plugins/qwen_image.py b/modelopt/torch/fastgen/plugins/qwen_image.py new file mode 100644 index 00000000000..08a32b09301 --- /dev/null +++ b/modelopt/torch/fastgen/plugins/qwen_image.py @@ -0,0 +1,381 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 plumbing for the DMD2 pipeline. + +Qwen-Image's ``QwenImageTransformer2DModel`` does not accept the diffusers-standard +``(hidden_states[B,C,H,W], timestep, encoder_hidden_states)`` triple. Instead it expects +*packed* latents ``[B, (H//2)*(W//2), C*4]`` plus three extra kwargs +(``encoder_hidden_states_mask``, ``img_shapes``, ``guidance``) and returns its prediction +in the same packed layout. The packing / unpacking step mirrors +``QwenImagePipeline._pack_latents`` in diffusers. + +DMD2's internal math (noise injection, VSD / DSM losses, EMA updates, fake-score updates) +all operates on the *unpacked* latent ``[B, C, H, W]``, so we keep that as the +:class:`DMDPipeline` external contract and push the pack / call / unpack triple into a +single override of :meth:`DMDPipeline._call_model` on :class:`QwenImageDMDPipeline`. + +Usage from a training recipe:: + + from modelopt.torch.fastgen.plugins.qwen_image import QwenImageDMDPipeline + + pipeline = QwenImageDMDPipeline( + student=student_transformer, + teacher=teacher_transformer, + fake_score=fake_score_transformer, + config=dmd_config, + discriminator=None, + ) + +The companion ``modelopt_recipes/general/distillation/dmd2_qwen_image.yaml`` must keep +``num_train_timesteps: null`` so the continuous RF time ``t ∈ [0, 1]`` is forwarded +verbatim to the transformer (Qwen-Image normalises timesteps to ``[0, 1]`` internally; +the diffusers ``[0, 1000]`` scale used for Wan / SD3 / Flux does NOT apply here). +""" + +from __future__ import annotations + +import contextlib +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +from ..methods.dmd import DMDPipeline + +if TYPE_CHECKING: + from ..config import DMDConfig + +__all__ = [ + "QwenImageDMDPipeline", + "attach_feature_capture", + "build_img_shapes", + "pack_latents", + "remove_feature_capture", + "unpack_latents", +] + + +# ---------------------------------------------------------------------------- # +# Latent pack / unpack helpers (2x2 patch grouping) # +# ---------------------------------------------------------------------------- # + + +def pack_latents(latents: torch.Tensor) -> torch.Tensor: + """Pack ``[B, C, H, W]`` latents into ``[B, (H//2)*(W//2), C*4]`` for Qwen-Image. + + Mirrors ``QwenImagePipeline._pack_latents`` (diffusers): groups every ``2x2`` spatial + block on each channel and lays the four values out along the channel axis so the + transformer's ``in_channels = 4 * out_channels`` patch embedding sees them as a + single token. ``H`` and ``W`` must both be even. + """ + if latents.ndim != 4: + raise ValueError( + f"pack_latents expects [B, C, H, W] (got {latents.ndim}D tensor of shape " + f"{tuple(latents.shape)})." + ) + b, c, h, w = latents.shape + if h % 2 or w % 2: + raise ValueError( + f"pack_latents requires even spatial dims, got H={h}, W={w}. " + "Increase the latent resolution or pad before packing." + ) + x = latents.view(b, c, h // 2, 2, w // 2, 2) + x = x.permute(0, 2, 4, 1, 3, 5) + x = x.reshape(b, (h // 2) * (w // 2), c * 4) + return x + + +def unpack_latents(packed: torch.Tensor, height: int, width: int) -> torch.Tensor: + """Inverse of :func:`pack_latents`. ``height`` / ``width`` are the unpacked latent dims.""" + if packed.ndim != 3: + raise ValueError( + f"unpack_latents expects [B, num_patches, C*4] (got {packed.ndim}D tensor of " + f"shape {tuple(packed.shape)})." + ) + b, num_patches, c4 = packed.shape + if c4 % 4: + raise ValueError(f"unpack_latents expects last dim divisible by 4, got {c4}.") + c = c4 // 4 + if height % 2 or width % 2: + raise ValueError( + f"unpack_latents requires even target spatial dims, got H={height}, W={width}." + ) + h2, w2 = height // 2, width // 2 + if num_patches != h2 * w2: + raise ValueError( + f"num_patches ({num_patches}) does not match H//2 * W//2 ({h2 * w2}) for " + f"target shape H={height}, W={width}." + ) + x = packed.view(b, h2, w2, c, 2, 2) + x = x.permute(0, 3, 1, 4, 2, 5) + x = x.reshape(b, c, height, width) + return x + + +def build_img_shapes(batch_size: int, h_lat: int, w_lat: int) -> list[list[tuple[int, int, int]]]: + """Build the ``img_shapes`` kwarg expected by ``QwenImageTransformer2DModel``. + + Each entry is ``[(1, h_lat // 2, w_lat // 2)]`` — one tuple per sample in the batch. + The leading ``1`` is the temporal dim (single frame for T2I). + """ + if h_lat % 2 or w_lat % 2: + raise ValueError( + f"build_img_shapes requires even latent dims, got h_lat={h_lat}, w_lat={w_lat}." + ) + return [[(1, h_lat // 2, w_lat // 2)]] * batch_size + + +# ---------------------------------------------------------------------------- # +# Pipeline subclass # +# ---------------------------------------------------------------------------- # + + +class QwenImageDMDPipeline(DMDPipeline): + """DMD2 pipeline that targets Qwen-Image's packed transformer signature. + + Drops in for :class:`DMDPipeline` and overrides :meth:`_call_model` only. All other + behaviour (noise schedules, VSD / DSM losses, EMA, GAN paths) is inherited unchanged. + + The student / teacher / fake_score modules must be the raw diffusers + ``QwenImageTransformer2DModel`` (or FSDP-sharded copy thereof). The pipeline handles + pack / call / unpack on every internal forward. + + Args: + guidance: Optional scalar guidance value forwarded to the transformer's + ``guidance`` kwarg every call. Only used when the transformer was built with + ``guidance_embeds=true`` (off by default for ``Qwen/Qwen-Image``). Leave + ``None`` to skip the guidance embedding entirely — this is independent of + :attr:`DMDConfig.guidance_scale`, which controls classifier-free guidance on + the teacher. + """ + + def __init__( + self, + student: nn.Module, + teacher: nn.Module, + fake_score: nn.Module, + config: DMDConfig, + *, + discriminator: nn.Module | None = None, + guidance: float | None = None, + ) -> None: + """Wrap the base DMD pipeline with Qwen-Image patch packing / guidance handling.""" + super().__init__( + student=student, + teacher=teacher, + fake_score=fake_score, + config=config, + discriminator=discriminator, + ) + if config.num_train_timesteps is not None: + raise ValueError( + "QwenImageDMDPipeline requires DMDConfig.num_train_timesteps=None — " + f"Qwen-Image normalises timesteps to [0, 1] internally (got " + f"num_train_timesteps={config.num_train_timesteps})." + ) + self._guidance_value = guidance + + def _call_model( + self, + model: nn.Module, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Pack [B, C, H, W] -> packed -> call transformer -> unpack -> [B, C, H, W].""" + if hidden_states.ndim != 4: + raise ValueError( + f"QwenImageDMDPipeline._call_model expects 4D hidden_states " + f"[B, C, H, W] (got {hidden_states.ndim}D)." + ) + b, _c, h, w = hidden_states.shape + + packed = pack_latents(hidden_states) + img_shapes = build_img_shapes(b, h, w) + + call_kwargs: dict[str, Any] = dict(model_kwargs) + call_kwargs.pop("hidden_states", None) + encoder_hidden_states_mask = call_kwargs.pop("encoder_hidden_states_mask", None) + call_kwargs.pop("img_shapes", None) + call_kwargs.pop("guidance", None) + call_kwargs.pop("return_dict", None) + txt_seq_lens = call_kwargs.pop("txt_seq_lens", None) + if txt_seq_lens is None and encoder_hidden_states_mask is not None: + txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).int().tolist() + + guidance = None + if self._guidance_value is not None: + guidance = torch.full( + (b,), + float(self._guidance_value), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + out = model( + hidden_states=packed, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=encoder_hidden_states_mask, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + guidance=guidance, + return_dict=False, + **call_kwargs, + ) + + if isinstance(out, tuple): + raw_packed = out[0] + elif isinstance(out, torch.Tensor): + raw_packed = out + elif hasattr(out, "sample"): + raw_packed = out.sample + else: + raise TypeError( + "QwenImageDMDPipeline._call_model could not extract a tensor from output " + f"of type {type(out).__name__!r}." + ) + + return unpack_latents(raw_packed, h, w) + + +# ---------------------------------------------------------------------------- # +# GAN feature capture # +# ---------------------------------------------------------------------------- # + +# These attribute names are what the shared +# :func:`~modelopt.torch.fastgen.methods.dmd._drain_if_hooked` / +# :func:`~modelopt.torch.fastgen.methods.dmd._require_hooked` helpers look for. +_CAPTURED_ATTR = "_fastgen_captured" +_HANDLES_ATTR = "_fastgen_capture_handles" +_INDICES_ATTR = "_fastgen_capture_indices" +_SHAPE_ATTR = "_fastgen_capture_shape" + + +def attach_feature_capture( + teacher: nn.Module, + feature_indices: list[int], + h_lat: int, + w_lat: int, + *, + blocks_attr: str = "transformer_blocks", +) -> None: + """Install forward hooks on ``teacher.transformer_blocks[i]`` for each ``i`` in ``feature_indices``. + + Qwen-Image's ``QwenImageTransformerBlock.forward`` returns + ``(encoder_hidden_states, hidden_states)`` where ``hidden_states`` has shape + ``[B, num_image_patches, 3072]`` with ``num_image_patches == (h_lat // 2) * (w_lat // 2)`` + (joint dual-stream attention; the text branch is the first tuple element and + is discarded). The hook unpacks the image branch and reshapes it to + ``[B, 3072, h_lat // 2, w_lat // 2]`` so the discriminator can consume it + as standard NCHW spatial features. + + Captured tensors land in ``teacher._fastgen_captured`` (a list) for the + DMD2 ``_drain_if_hooked`` / ``_require_hooked`` helpers to pop after each + teacher forward. + + Args: + teacher: The teacher transformer module. + feature_indices: Block indices to capture (e.g. ``[30]`` or + ``[15, 30, 45]`` for the 60-block Qwen-Image teacher). + h_lat: Latent height passed to the transformer this step. Must be even. + w_lat: Latent width passed to the transformer this step. Must be even. + blocks_attr: Attribute under which the teacher exposes its block stack. + Default ``"transformer_blocks"`` matches diffusers' + ``QwenImageTransformer2DModel``. + + Raises: + AttributeError: ``teacher`` does not expose ``blocks_attr``. + IndexError: An entry of ``feature_indices`` is out of range. + ValueError: ``h_lat`` or ``w_lat`` is odd. + """ + if h_lat % 2 != 0 or w_lat % 2 != 0: + raise ValueError( + f"attach_feature_capture requires even latent dims, got h_lat={h_lat}, w_lat={w_lat}." + ) + + remove_feature_capture(teacher) + + blocks = getattr(teacher, blocks_attr, None) + if blocks is None: + raise AttributeError( + f"Teacher {type(teacher).__name__!r} does not expose a ``{blocks_attr}`` attribute; " + f"pass blocks_attr=... if the block stack is named differently." + ) + try: + num_blocks = len(blocks) + except TypeError as exc: + raise TypeError( + f"Teacher ``{blocks_attr}`` is not a sequence (got {type(blocks).__name__!r})." + ) from exc + + sorted_indices = sorted(set(feature_indices)) + for idx in sorted_indices: + if not (0 <= idx < num_blocks): + raise IndexError( + f"feature_indices entry {idx} is out of range for teacher with {num_blocks} blocks." + ) + + captured: list[torch.Tensor] = [] + setattr(teacher, _CAPTURED_ATTR, captured) + setattr(teacher, _INDICES_ATTR, list(sorted_indices)) + setattr(teacher, _SHAPE_ATTR, (h_lat // 2, w_lat // 2)) + + handles: list[Any] = [] + h_half = h_lat // 2 + w_half = w_lat // 2 + for idx in sorted_indices: + block = blocks[idx] + + def _hook(_module: nn.Module, _inputs: Any, output: Any) -> None: + # Qwen-Image block.forward returns (encoder_hidden_states, hidden_states). + if isinstance(output, tuple) and len(output) == 2: + hidden = output[1] + elif isinstance(output, torch.Tensor): + hidden = output + else: + raise TypeError( + f"Unexpected QwenImage block output type {type(output).__name__!r}; " + "expected (encoder_hidden_states, hidden_states) tuple or Tensor." + ) + # hidden: [B, num_image_patches, C] -> [B, C, H_half, W_half]. + b, s, c = hidden.shape + expected_s = h_half * w_half + if s != expected_s: + raise RuntimeError( + f"QwenImage feature-capture got hidden_states seq_len={s} but expected " + f"{expected_s} = (h_lat // 2) * (w_lat // 2). Did the input resolution " + f"drift from the attach_feature_capture-time setting?" + ) + feat = hidden.permute(0, 2, 1).reshape(b, c, h_half, w_half) + captured.append(feat) + + handles.append(block.register_forward_hook(_hook)) + + setattr(teacher, _HANDLES_ATTR, handles) + + +def remove_feature_capture(teacher: nn.Module) -> None: + """Remove previously installed feature-capture hooks (no-op if none are installed).""" + handles = getattr(teacher, _HANDLES_ATTR, None) + if handles: + for h in handles: + h.remove() + for attr in (_HANDLES_ATTR, _CAPTURED_ATTR, _INDICES_ATTR, _SHAPE_ATTR): + if hasattr(teacher, attr): + with contextlib.suppress(AttributeError): + delattr(teacher, attr) diff --git a/modelopt/torch/fastgen/utils.py b/modelopt/torch/fastgen/utils.py new file mode 100644 index 00000000000..2e9ce6a7001 --- /dev/null +++ b/modelopt/torch/fastgen/utils.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 tensor helpers shared across the fastgen subpackage.""" + +from __future__ import annotations + +import torch + +__all__ = ["classifier_free_guidance", "expand_like"] + + +def expand_like(x: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Pad ``x`` with trailing singleton dims until it has the same ndim as ``target``. + + Used to broadcast per-sample scalars like ``alpha_t`` / ``sigma_t`` across the + spatial / temporal axes of a video latent. + + Example:: + + x = torch.ones(5) # shape (5,) + target = torch.ones(5, 4, 16, 16) + expand_like(x, target).shape # (5, 1, 1, 1) + """ + x = torch.atleast_1d(x) + while x.ndim < target.ndim: + x = x[..., None] + return x + + +def classifier_free_guidance( + cond_pred: torch.Tensor, + uncond_pred: torch.Tensor, + guidance_scale: float, +) -> torch.Tensor: + """Combine conditional and unconditional predictions via classifier-free guidance. + + Uses the DMD2 convention ``cond + (scale - 1) * (cond - uncond)``, which is + mathematically equivalent to the standard CFG formula + ``uncond + scale * (cond - uncond)``. + """ + return cond_pred + (guidance_scale - 1.0) * (cond_pred - uncond_pred) diff --git a/modelopt_recipes/general/distillation/dmd2_qwen_image.yaml b/modelopt_recipes/general/distillation/dmd2_qwen_image.yaml new file mode 100644 index 00000000000..79f8124e6fe --- /dev/null +++ b/modelopt_recipes/general/distillation/dmd2_qwen_image.yaml @@ -0,0 +1,80 @@ +# DMD2 distillation recipe for Qwen-Image (text-to-image). +# +# Maps to :class:`modelopt.torch.fastgen.DMDConfig`. Load with:: +# +# from modelopt.torch.fastgen import load_dmd_config +# cfg = load_dmd_config("general/distillation/dmd2_qwen_image") +# +# Reference: NeMo AutoModel's Qwen-Image flow-matching config +# (Automodel/examples/diffusion/finetune/qwen_image_t2i_flow.yaml). + +# Qwen-Image is rectified-flow. Under RF, "flow" and "v" are equivalent. +pred_type: flow + +# Qwen-Image normalises timesteps to [0, 1] internally (see +# QwenImageAdapter.prepare_inputs: ``timesteps = context.timesteps / 1000``). +# QwenImageDMDPipeline asserts num_train_timesteps is null so the continuous RF +# time ``t ∈ [0, 1]`` is forwarded verbatim. +num_train_timesteps: + +# Classifier-free guidance strength applied to the teacher during the student update. +# null disables CFG (skips the negative-conditioning teacher forward). +guidance_scale: + +# Phase 2: 4-step student to match FastGen's Qwen-Image DMD2 default +# (config_dmd2.py:52). ``sample_t_cfg.t_list`` below pins the per-step schedule. +student_sample_steps: 4 +student_sample_type: ode +# Default keeps FastGen Qwen parity: train each rung from noised real latents. +# Override with ``--dmd2.backward_simulation=true`` to no-grad unroll the current +# student from the first rung before training the selected rung. +backward_simulation: false + +# One student step per N fake-score / discriminator steps. +student_update_freq: 5 + +# Fake score trains in x0 space while the student/teacher operate in flow space. +fake_score_pred_type: x0 + +# GAN generator weight. 0 disables the discriminator branch (Phase 1). +gan_loss_weight_gen: 0.0 +gan_use_same_t_noise: false +gan_r1_reg_weight: 0.0 +gan_r1_reg_alpha: 0.1 + +sample_t_cfg: + # ``time_dist_type`` governs the *perturbation* timestep ``t`` sampled on + # every loss path — VSD perturbation in compute_student_loss (dmd.py:417), + # fake-score DSM perturbation in compute_fake_score_loss (dmd.py:529), and + # GAN/discriminator perturbation in compute_discriminator_loss + # (dmd.py:605). All three call ``self.sample_timesteps`` → ``sample_t`` → + # reads ``time_dist_type``. + # + # It does NOT govern the student's *starting* timestep ``t_student``: under + # student_sample_steps > 1, ``_build_student_input`` calls + # ``sample_from_t_list`` (dmd.py:346) which samples uniformly from + # ``t_list[:-1]`` regardless of ``time_dist_type``. So ``time_dist_type`` + # is a knob that's active at every loss path; the 4-step setup just makes + # it irrelevant for ``t_student`` specifically. + time_dist_type: logitnormal + min_t: 0.001 + max_t: 0.999 + p_mean: 0.0 + 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 + # 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 + # were really ``linspace(1.0, 0, 5)`` with t=1 shaved to 0.999, leaving a + # silent ~0.3% train↔inference skew on each non-endpoint timestep. + t_list: [0.999, 0.74925, 0.4995, 0.24975, 0.0] + +# Student EMA. Omit this block to disable EMA tracking entirely. +ema: + decay: 0.9999 + type: constant + start_iter: 0 + fsdp2: true + mode: full_tensor diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index 352cc60d793..a42ebdd8ffb 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -42,6 +42,16 @@ except Exception: # pragma: no cover - optional diffusers models AutoencoderKLWan = None +try: + from diffusers.models.transformers import QwenImageTransformer2DModel +except Exception: # pragma: no cover - optional diffusers models + QwenImageTransformer2DModel = None + +try: + from diffusers.models.autoencoders import AutoencoderKLQwenImage +except Exception: # pragma: no cover - optional diffusers models + AutoencoderKLQwenImage = None + import modelopt.torch.opt as mto @@ -250,3 +260,120 @@ def create_tiny_wan22_pipeline_dir(tmp_path: Path) -> Path: save_dir = tmp_path / "tiny_wan22" pipe.save_pretrained(save_dir) return save_dir + + +def get_tiny_qwen_image_transformer(**config_kwargs): + """Create a tiny QwenImageTransformer2DModel for testing. + + Scaled down from the real Qwen-Image config (60 layers, 24 heads, head_dim 128, + joint_attention_dim 3584). Two constraints to keep in mind: + - ``axes_dims_rope`` must sum to ``attention_head_dim``. + - ``joint_attention_dim`` must match the text-embedding dim the model is fed. In + the DMD2 mock-data training path that is the *dataloader's* ``text_embed_dim`` + (the bundled text encoder is bypassed), so pair this with + ``--data.dataloader.text_embed_dim=``. + """ + if QwenImageTransformer2DModel is None: + pytest.skip("QwenImageTransformer2DModel is not available in this diffusers version.") + + kwargs = { + "patch_size": 2, + "in_channels": 64, # vae z_dim (16) * 2x2 patch + "out_channels": 16, # = vae z_dim + "num_layers": 2, + "attention_head_dim": 16, + "num_attention_heads": 2, # hidden_dim = 32 + "joint_attention_dim": 32, + "pooled_projection_dim": 32, + "guidance_embeds": False, + "axes_dims_rope": (8, 4, 4), # sums to attention_head_dim (16) + } + kwargs.update(**config_kwargs) + return QwenImageTransformer2DModel(**kwargs) + + +def get_tiny_qwen_image_vae(**config_kwargs): + """Create a tiny AutoencoderKLQwenImage for testing (z_dim=16 to match the transformer).""" + if AutoencoderKLQwenImage is None: + pytest.skip("AutoencoderKLQwenImage is not available in this diffusers version.") + + kwargs = { + "base_dim": 8, + "z_dim": 16, # = transformer out_channels + "dim_mult": [1, 2], + "num_res_blocks": 1, + "temperal_downsample": [True], # len == len(dim_mult) - 1 + "attn_scales": [], + "latents_mean": [0.0] * 16, # length must == z_dim + "latents_std": [1.0] * 16, + } + kwargs.update(**config_kwargs) + return AutoencoderKLQwenImage(**kwargs) + + +def create_tiny_qwen_image_pipeline_dir(tmp_path: Path) -> Path: + """Create and save a tiny Qwen-Image pipeline to a directory (SKETCH). + + Mirrors ``create_tiny_wan22_pipeline_dir``. Needs in-container validation; the + fragile piece is the Qwen2.5-VL text encoder. This prefers a tiny-random HF model + (as Wan uses ``hf-internal-testing/tiny-random-t5``); if that id drifts or the + config schema differs across transformers versions, copy the text-encoder + construction from diffusers' own QwenImage fast test + (``tests/pipelines/qwenimage/test_qwenimage.py``). + + For the DMD2 mock-data training path the transformer consumes the dataloader's + embeddings rather than the text encoder, so the bundled tiny text encoder only + needs to load; its hidden size is intentionally decoupled from the transformer's + ``joint_attention_dim`` (set the dataloader's ``text_embed_dim`` to match instead). + The saved dir loads with ``QwenImagePipeline.from_pretrained(path)``. + """ + if QwenImageTransformer2DModel is None or AutoencoderKLQwenImage is None: + pytest.skip("QwenImage diffusers classes not available in this diffusers version.") + from diffusers import FlowMatchEulerDiscreteScheduler, QwenImagePipeline + + transformers = pytest.importorskip("transformers") + + # Tiny Qwen2.5-VL text encoder + matching Qwen2 tokenizer (loaded, but bypassed + # during DMD2 mock-data training). + # NOTE (validated 2026-06-06): the hf-internal-testing id below does NOT exist on the + # Hub, so this fixture currently skips. To make the recipe e2e runnable in CI, + # construct the encoder inline from a tiny ``Qwen2_5_VLConfig`` (nested text + vision + # config) — mirror diffusers' ``QwenImagePipelineFastTests.get_dummy_components`` in + # ``tests/pipelines/qwenimage/test_qwenimage.py``. + tiny_id = "hf-internal-testing/tiny-random-Qwen2_5_VLForConditionalGeneration" + try: + text_encoder = transformers.Qwen2_5_VLForConditionalGeneration.from_pretrained(tiny_id) + tokenizer = transformers.Qwen2Tokenizer.from_pretrained(tiny_id) + except Exception as exc: # pragma: no cover - depends on hub availability / version + pytest.skip( + f"tiny Qwen2.5-VL text encoder unavailable ({exc}); " + "copy the fixture from diffusers' QwenImage fast test" + ) + + torch.manual_seed(0) + transformer = get_tiny_qwen_image_transformer() + torch.manual_seed(0) + vae = get_tiny_qwen_image_vae() + + scheduler = FlowMatchEulerDiscreteScheduler( + base_image_seq_len=256, + base_shift=0.5, + max_image_seq_len=8192, + max_shift=0.9, + num_train_timesteps=1000, + shift=1.0, + shift_terminal=0.02, + use_dynamic_shifting=True, + time_shift_type="exponential", + ) + + pipe = QwenImagePipeline( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + ) + save_dir = tmp_path / "tiny_qwen_image" + pipe.save_pretrained(save_dir) + return save_dir diff --git a/tests/examples/diffusers/conftest.py b/tests/examples/diffusers/conftest.py index 8893d188d9e..e704f6d5879 100644 --- a/tests/examples/diffusers/conftest.py +++ b/tests/examples/diffusers/conftest.py @@ -29,3 +29,20 @@ def tiny_wan22_path(tmp_path_factory): tmp_path = tmp_path_factory.mktemp("wan22") return str(create_tiny_wan22_pipeline_dir(tmp_path)) + + +@pytest.fixture(scope="session") +def tiny_qwen_image_path(tmp_path_factory): + """Create a tiny Qwen-Image pipeline and return its path (built once per session). + + SKETCH fixture for the recipe-level DMD2 e2e (``test_fastgen_recipe_e2e.py``). + See ``create_tiny_qwen_image_pipeline_dir`` for caveats — notably the tiny + Qwen2.5-VL text encoder, which needs in-container validation. + """ + try: + from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir + except ImportError: + pytest.skip("Qwen-Image diffusers models not available") + + tmp_path = tmp_path_factory.mktemp("qwen_image") + return str(create_tiny_qwen_image_pipeline_dir(tmp_path)) diff --git a/tests/unit/torch/fastgen/conftest.py b/tests/unit/torch/fastgen/conftest.py new file mode 100644 index 00000000000..26110005bc5 --- /dev/null +++ b/tests/unit/torch/fastgen/conftest.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 fixtures for ``tests/unit/torch/fastgen/``. + +Keeps the duplicated ``_ToyTransformer`` / ``_ToyDiscriminator`` / pipeline-builder +helpers in one place so individual test files can focus on assertions rather than +wiring. +""" + +from __future__ import annotations + +import copy + +import pytest +import torch +from torch import nn + +from modelopt.torch.fastgen import DMDConfig, DMDPipeline + + +class ToyTransformer(nn.Module): + """Minimal diffusers-shaped transformer: output = Linear(hidden_states). + + Accepts ``hidden_states`` / ``timestep`` / ``encoder_hidden_states`` / **kwargs + but ignores all of them except the first. + """ + + def __init__(self, d: int) -> None: + super().__init__() + self.linear = nn.Linear(d, d, bias=False) + + def forward(self, hidden_states, timestep=None, encoder_hidden_states=None, **kwargs): + return self.linear(hidden_states) + + +class ToyDiscriminator(nn.Module): + """Consumes ``list[Tensor]`` and returns 2D logits ``(B, 1)`` by averaging features.""" + + def forward(self, features): + x = features[0] + return x.flatten(start_dim=1).mean(dim=-1, keepdim=True) + + +@pytest.fixture +def toy_transformer_factory(): + """Return a callable ``d -> ToyTransformer(d)`` factory.""" + return ToyTransformer + + +@pytest.fixture +def toy_discriminator_factory(): + """Return a callable ``() -> ToyDiscriminator()`` factory.""" + return ToyDiscriminator + + +@pytest.fixture +def build_pipeline(): + """Factory that constructs a :class:`DMDPipeline` with toy student/teacher/fake_score. + + Usage:: + + pipeline = build_pipeline( + d=4, pred_type="flow", gan_loss_weight_gen=0.03, discriminator=ToyDiscriminator() + ) + """ + + def _build( + d: int, + *, + pred_type: str = "flow", + fake_score_pred_type: str | None = None, + num_train_timesteps: int | None = None, + gan_loss_weight_gen: float = 0.0, + gan_use_same_t_noise: bool = False, + gan_r1_reg_weight: float = 0.0, + ema=None, + discriminator: nn.Module | None = None, + seed: int = 0, + ) -> DMDPipeline: + torch.manual_seed(seed) + student = ToyTransformer(d) + teacher = ToyTransformer(d) + fake_score = copy.deepcopy(teacher) + cfg = DMDConfig( + pred_type=pred_type, + fake_score_pred_type=fake_score_pred_type, + num_train_timesteps=num_train_timesteps, + gan_loss_weight_gen=gan_loss_weight_gen, + gan_use_same_t_noise=gan_use_same_t_noise, + gan_r1_reg_weight=gan_r1_reg_weight, + ema=ema, + ) + return DMDPipeline( + student, + teacher, + fake_score, + cfg, + discriminator=discriminator, + ) + + return _build diff --git a/tests/unit/torch/fastgen/test_dmd_gradient_routing.py b/tests/unit/torch/fastgen/test_dmd_gradient_routing.py new file mode 100644 index 00000000000..5c842bd8201 --- /dev/null +++ b/tests/unit/torch/fastgen/test_dmd_gradient_routing.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 tests on ``DMDPipeline`` with tiny modules. + +Ports the in-process gradient-isolation bullets from checklist §3 (3.1, 3.2): +the student loss must only touch student params, and the fake-score loss must +only touch fake-score params. The source-grep bullets (3.4 / 3.6 / 3.7 / 3.8) +intentionally stay in ``experiments/qwen.3/run_section_3.py`` — they are +recipe-source linting, not unit-testable logic. + +Uses a ``_TinyTransformer`` with a timestep bias (the checklist's §3 module) so +the gradient signal definitely flows through the transformer when the +``compute_*_loss`` paths are exercised. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from modelopt.torch.fastgen.config import DMDConfig, SampleTimestepConfig +from modelopt.torch.fastgen.methods.dmd import DMDPipeline + + +class _TinyTransformer(nn.Module): + """``(hidden_states, timestep, encoder_hidden_states, **kw) -> Tensor`` module. + + Linear projection over the flattened spatial axes plus a timestep-derived + bias. Cheap enough to run on CPU and returns a tensor in flow-space so it + can play either student / teacher / fake_score role. + """ + + def __init__(self, channels: int = 16, dim: int = 8) -> None: + super().__init__() + self.channels = channels + self.dim = dim + flat = channels * dim * dim + self.proj = nn.Linear(flat, flat, bias=True) + self.t_proj = nn.Linear(1, flat, bias=False) + + def forward(self, hidden_states, timestep, encoder_hidden_states=None, **_kw): + b = hidden_states.shape[0] + x = hidden_states.reshape(b, -1) + t = timestep.reshape(b, 1).to(x.dtype) + return (self.proj(x) + self.t_proj(t)).reshape_as(hidden_states) + + +def _make_pipeline() -> tuple[DMDPipeline, _TinyTransformer, _TinyTransformer, _TinyTransformer]: + torch.manual_seed(0) + student = _TinyTransformer() + teacher = _TinyTransformer() + fake_score = _TinyTransformer() + cfg = DMDConfig( + pred_type="flow", + num_train_timesteps=None, + student_sample_steps=1, + student_update_freq=5, + fake_score_pred_type="x0", + gan_loss_weight_gen=0.0, + guidance_scale=None, + sample_t_cfg=SampleTimestepConfig( + time_dist_type="shifted", min_t=0.001, max_t=0.999, shift=1.0 + ), + ema=None, + ) + return ( + DMDPipeline(student=student, teacher=teacher, fake_score=fake_score, config=cfg), + student, + teacher, + fake_score, + ) + + +def _has_grad(module: nn.Module) -> bool: + return any(p.grad is not None and p.grad.abs().sum().item() > 0 for p in module.parameters()) + + +# ---------------------------------------------------------------------------- # +# §3.1 — compute_student_loss: only the student gets grads # +# ---------------------------------------------------------------------------- # + + +def test_compute_student_loss_routes_gradients_to_student_only(): + pipe, student, teacher, fake_score = _make_pipeline() + + # Mirror the recipe's _set_grad_requirements for the student phase. + student.train() + for p in student.parameters(): + p.requires_grad_(True) + fake_score.eval() + for p in fake_score.parameters(): + p.requires_grad_(False) + teacher.eval() + for p in teacher.parameters(): + p.requires_grad_(False) + + torch.manual_seed(1) + latents = torch.randn(2, 16, 8, 8) + noise = torch.randn_like(latents) + text = torch.randn(2, 8, 4) + + losses = pipe.compute_student_loss( + latents, + noise, + encoder_hidden_states=text, + negative_encoder_hidden_states=None, + guidance_scale=None, + ) + losses["total"].backward() + + assert "vsd" in losses + assert "total" in losses + assert _has_grad(student) + assert not _has_grad(teacher) + assert not _has_grad(fake_score) + + +# ---------------------------------------------------------------------------- # +# §3.2 — compute_fake_score_loss: only the fake_score gets grads # +# ---------------------------------------------------------------------------- # + + +def test_compute_fake_score_loss_routes_gradients_to_fake_score_only(): + pipe, student, teacher, fake_score = _make_pipeline() + + # Mirror the fake-score phase grad config. + student.eval() + for p in student.parameters(): + p.requires_grad_(False) + fake_score.train() + for p in fake_score.parameters(): + p.requires_grad_(True) + teacher.eval() + for p in teacher.parameters(): + p.requires_grad_(False) + + torch.manual_seed(2) + latents = torch.randn(2, 16, 8, 8) + noise = torch.randn_like(latents) + text = torch.randn(2, 8, 4) + + losses = pipe.compute_fake_score_loss(latents, noise, encoder_hidden_states=text) + losses["total"].backward() + + assert "fake_score" in losses + assert "total" in losses + assert _has_grad(fake_score) + assert not _has_grad(student) + assert not _has_grad(teacher) diff --git a/tests/unit/torch/fastgen/test_dmd_math.py b/tests/unit/torch/fastgen/test_dmd_math.py new file mode 100644 index 00000000000..91db5a61392 --- /dev/null +++ b/tests/unit/torch/fastgen/test_dmd_math.py @@ -0,0 +1,419 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""DMD math parity tests against the FastGen reference implementation. + +Ports checklist §2 bullets — flow-matching identities, dsm/vsd/gan losses, CFG +formula, and the fake-score flow→x0→DSM conversion chain. The FastGen reference +math is inlined verbatim from +``source/FastGen/fastgen/methods/common_loss.py`` and +``source/FastGen/fastgen/methods/distribution_matching/dmd2.py`` so the test is +hermetic — no FastGen import required. + +Numerical tolerance is ``1e-6`` for floating-point losses; pack/permute paths +(``add_noise``, ``pred_x0_from_flow``, ``x0_to_flow``, CFG) use ``torch.equal`` +because both implementations route through fp64 intermediates with the same +operation order. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from modelopt.torch.fastgen.config import DMDConfig, SampleTimestepConfig +from modelopt.torch.fastgen.flow_matching import ( + add_noise, + pred_x0_from_flow, + rf_alpha, + rf_sigma, + x0_to_flow, +) +from modelopt.torch.fastgen.losses import dsm_loss, gan_disc_loss, gan_gen_loss, vsd_loss +from modelopt.torch.fastgen.methods import dmd as dmd_module +from modelopt.torch.fastgen.methods.dmd import DMDPipeline +from modelopt.torch.fastgen.utils import classifier_free_guidance + +# ---------------------------------------------------------------------------- # +# FastGen reference impls (math only) — inlined verbatim # +# ---------------------------------------------------------------------------- # + + +def _expand_like(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + while t.ndim < target.ndim: + t = t.unsqueeze(-1) + return t + + +def _fastgen_forward_process(x: torch.Tensor, eps: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """``BaseNoiseSchedule.forward_process`` with RF alpha/sigma inlined.""" + original_dtype = x.dtype + t64 = t.to(torch.float64) + x64 = x.to(torch.float64) + eps64 = eps.to(torch.float64) + alpha_t = _expand_like(1.0 - t64, x64) + sigma_t = _expand_like(t64, eps64) + return (x64 * alpha_t + eps64 * sigma_t).to(original_dtype) + + +def _fastgen_dsm(pred_type, net_pred, *, x0=None, eps=None, t=None): + """common_loss.py:12-60. RF scheduler so alpha=1-t, sigma=t for 'v'.""" + if pred_type == "x0": + return F.mse_loss(x0, net_pred, reduction="mean") + if pred_type == "eps": + return F.mse_loss(eps, net_pred, reduction="mean") + if pred_type == "v": + alpha_t = _expand_like((1.0 - t).to(dtype=x0.dtype), x0).to(device=x0.device) + sigma_t = _expand_like(t.to(dtype=x0.dtype), x0).to(device=x0.device) + v = alpha_t * eps - sigma_t * x0 + return F.mse_loss(v, net_pred, reduction="mean") + if pred_type == "flow": + return F.mse_loss(eps - x0, net_pred, reduction="mean") + raise NotImplementedError(pred_type) + + +def _fastgen_vsd(gen_data, teacher_x0, fake_score_x0): + """common_loss.py:63-103.""" + dims = tuple(range(1, teacher_x0.ndim)) + with torch.no_grad(): + original_dtype = gen_data.dtype + gen_fp32 = gen_data.float() + teacher_fp32 = teacher_x0.float() + diff_abs_mean = (gen_fp32 - teacher_fp32).abs().mean(dim=dims, keepdim=True) + w = (1 / (diff_abs_mean + 1e-6)).to(dtype=original_dtype) + pseudo_target = gen_data - (fake_score_x0 - teacher_x0) * w + return 0.5 * F.mse_loss(gen_data, pseudo_target, reduction="mean") + + +def _fastgen_cfg(cond, uncond, scale): + """dmd2.py:184 — ``teacher_x0 + (scale - 1) * (teacher_x0 - teacher_x0_neg)``.""" + return cond + (scale - 1) * (cond - uncond) + + +def _fastgen_gan_gen(fake_logits): + return F.softplus(-fake_logits).mean() + + +def _fastgen_gan_disc(real_logits, fake_logits): + return F.softplus(fake_logits).mean() + F.softplus(-real_logits).mean() + + +# ---------------------------------------------------------------------------- # +# §2.1 — RF forward process # +# ---------------------------------------------------------------------------- # + + +def test_rf_forward_matches_fastgen(): + torch.manual_seed(0) + x0 = torch.randn(2, 16, 8, 8, dtype=torch.float32) + eps = torch.randn_like(x0) + t = torch.tensor([0.1, 0.7], dtype=torch.float64) + assert torch.equal(add_noise(x0, eps, t), _fastgen_forward_process(x0, eps, t)) + + +# ---------------------------------------------------------------------------- # +# §2.2 / §2.3 — student input for single-step and multi-step # +# ---------------------------------------------------------------------------- # + + +def _student_input_pipeline(*, sample_steps: int, t_list=None) -> DMDPipeline: + cfg = DMDConfig( + pred_type="flow", + num_train_timesteps=None, + student_sample_steps=sample_steps, + student_update_freq=5, + sample_t_cfg=SampleTimestepConfig( + time_dist_type="shifted", + min_t=0.001, + max_t=0.999, + shift=5.0, + t_list=t_list, + ), + ) + return DMDPipeline( + student=nn.Identity(), teacher=nn.Identity(), fake_score=nn.Identity(), config=cfg + ) + + +def test_build_student_input_single_step_matches_max_t_noise(): + pipe = _student_input_pipeline(sample_steps=1) + latents = torch.randn(2, 16, 8, 8, dtype=torch.float32) + noise = torch.randn_like(latents) + input_student, t_student = pipe._build_student_input(latents, noise) + max_t = float(pipe.config.sample_t_cfg.max_t) + expected_input = (noise.to(torch.float64) * max_t).to(noise.dtype) + expected_t = torch.full((2,), max_t, dtype=torch.float32) + assert torch.equal(input_student, expected_input) + assert torch.equal(t_student, expected_t) + + +def test_build_student_input_multi_step_uses_t_list_prefix_and_add_noise(): + pipe = _student_input_pipeline(sample_steps=2, t_list=[0.999, 0.5, 0.0]) + torch.manual_seed(0) + latents = torch.randn(8, 16, 4, 4, dtype=torch.float32) + noise = torch.randn_like(latents) + input_student, t_student = pipe._build_student_input(latents, noise) + allowed = list(pipe.config.sample_t_cfg.t_list[:-1]) + actual = t_student.detach().cpu().tolist() + # fp32 round-trip — compare with tolerance, not set membership. + assert all(any(abs(v - a) < 1e-5 for a in allowed) for v in actual) + assert torch.equal(input_student, add_noise(latents, noise, t_student)) + + +class _ZeroFlow(nn.Module): + def forward(self, hidden_states, timestep, encoder_hidden_states=None, **_kwargs): + return torch.zeros_like(hidden_states) + + +def _backward_simulation_pipeline(*, sample_type: str = "ode") -> DMDPipeline: + cfg = DMDConfig( + pred_type="flow", + num_train_timesteps=None, + student_sample_steps=2, + student_sample_type=sample_type, + backward_simulation=True, + sample_t_cfg=SampleTimestepConfig( + time_dist_type="uniform", + min_t=0.001, + max_t=0.9, + t_list=[0.9, 0.5, 0.0], + ), + ) + model = _ZeroFlow() + return DMDPipeline(student=model, teacher=model, fake_score=model, config=cfg) + + +def test_build_student_input_backward_simulation_uses_generated_distribution(monkeypatch): + def _fixed_randint(low, high, size, *, device=None, dtype=None, **_kwargs): + assert low == 0 + assert high == 2 + return torch.ones(size, device=device, dtype=dtype or torch.long) + + def _fixed_randn_like(x, *args, **kwargs): + return torch.full_like(x, 2.0) + + monkeypatch.setattr(torch, "randint", _fixed_randint) + monkeypatch.setattr(torch, "randn_like", _fixed_randn_like) + + pipe = _backward_simulation_pipeline() + latents = torch.zeros(2, 16, 4, 4, dtype=torch.float32) + final_noise = torch.full_like(latents, 3.0) + input_student, t_student = pipe._build_student_input(latents, final_noise) + + expected_t = torch.full((2,), 0.5, dtype=torch.float32) + generated_x0 = torch.full_like(latents, 2.0 * 0.9) + expected_input = add_noise(generated_x0, final_noise, expected_t) + assert torch.equal(t_student, expected_t) + assert torch.equal(input_student, expected_input) + assert not torch.equal(input_student, add_noise(latents, final_noise, expected_t)) + + +def test_backward_simulation_selected_rung_is_broadcast(monkeypatch): + calls = [] + + def _fixed_randint(low, high, size, *, device=None, dtype=None, **_kwargs): + assert low == 0 + assert high == 2 + return torch.ones(size, device=device, dtype=dtype or torch.long) + + def _broadcast_to_first_rung(tensor, src): + assert src == 0 + calls.append(tensor.clone()) + tensor.zero_() + + monkeypatch.setattr(torch, "randint", _fixed_randint) + monkeypatch.setattr(dmd_module.dist, "is_available", lambda: True) + monkeypatch.setattr(dmd_module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(dmd_module.dist, "broadcast", _broadcast_to_first_rung) + + pipe = _backward_simulation_pipeline() + latents = torch.zeros(2, 16, 4, 4, dtype=torch.float32) + final_noise = torch.full_like(latents, 3.0) + input_student, t_student = pipe._build_student_input(latents, final_noise) + + expected_t = torch.full((2,), 0.9, dtype=torch.float32) + expected_input = (final_noise.to(torch.float64) * 0.9).to(final_noise.dtype) + assert len(calls) == 1 + assert torch.equal(calls[0], torch.ones(1, dtype=torch.long)) + assert torch.equal(t_student, expected_t) + assert torch.equal(input_student, expected_input) + + +# ---------------------------------------------------------------------------- # +# §2.4 / §2.5 — flow ↔ x0 identities # +# ---------------------------------------------------------------------------- # + + +def test_pred_x0_from_flow_matches_identity(): + torch.manual_seed(1) + x_t = torch.randn(2, 16, 8, 8, dtype=torch.float32) + flow = torch.randn_like(x_t) + t = torch.tensor([0.3, 0.7], dtype=torch.float32) + mo = pred_x0_from_flow(flow, x_t, t) + t64 = _expand_like(t.to(torch.float64), x_t.to(torch.float64)) + ref = (x_t.to(torch.float64) - t64 * flow.to(torch.float64)).to(x_t.dtype) + assert torch.equal(mo, ref) + + +def test_x0_to_flow_matches_identity(): + torch.manual_seed(2) + x0 = torch.randn(2, 16, 8, 8, dtype=torch.float32) + x_t = torch.randn_like(x0) + t = torch.tensor([0.3, 0.7], dtype=torch.float32) + mo = x0_to_flow(x0, x_t, t) + t64 = _expand_like(t.to(torch.float64), x0.to(torch.float64)) + ref = ((x_t.to(torch.float64) - x0.to(torch.float64)) / t64.clamp_min(1e-6)).to(x0.dtype) + assert torch.equal(mo, ref) + + +# ---------------------------------------------------------------------------- # +# §2.6 — dsm_loss for x0 / eps / flow / v # +# ---------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("pred_type", ["x0", "eps", "flow", "v"]) +def test_dsm_loss_matches_fastgen(pred_type): + torch.manual_seed(3) + x0 = torch.randn(2, 16, 8, 8, dtype=torch.float32) + eps = torch.randn_like(x0) + t = torch.tensor([0.4, 0.6], dtype=torch.float32) + net_pred = torch.randn_like(x0) + kwargs = {"x0": x0, "eps": eps, "t": t} + if pred_type == "v": + kwargs["alpha_fn"] = rf_alpha + kwargs["sigma_fn"] = rf_sigma + mo = dsm_loss(pred_type, net_pred, **kwargs).item() + fg = _fastgen_dsm(pred_type, net_pred, x0=x0, eps=eps, t=t).item() + assert abs(mo - fg) < 1e-6 + + +# ---------------------------------------------------------------------------- # +# §2.7 — vsd_loss # +# ---------------------------------------------------------------------------- # + + +def test_vsd_loss_matches_fastgen(): + torch.manual_seed(4) + gen_data = torch.randn(2, 16, 8, 8, dtype=torch.float32, requires_grad=True) + teacher_x0 = torch.randn_like(gen_data).detach() + fake_score_x0 = torch.randn_like(gen_data).detach() + mo = vsd_loss(gen_data, teacher_x0, fake_score_x0).item() + fg = _fastgen_vsd(gen_data.detach(), teacher_x0, fake_score_x0).item() + assert abs(mo - fg) < 1e-6 + + +# ---------------------------------------------------------------------------- # +# §2.8 — fake-score DSM target: ModelOpt flow→x0→DSM matches FastGen direct DSM('x0') # +# ---------------------------------------------------------------------------- # + + +def test_fake_score_flow_to_x0_dsm_matches_fastgen(): + torch.manual_seed(5) + x0_real = torch.randn(2, 16, 8, 8, dtype=torch.float32) + eps = torch.randn_like(x0_real) + t = torch.tensor([0.3, 0.7], dtype=torch.float32) + x_t = add_noise(x0_real, eps, t) + raw_flow = torch.randn_like(x0_real) + x0_pred_modelopt = DMDPipeline._raw_to_x0(raw_flow, x_t, t, native_pred_type="flow") + loss_modelopt = dsm_loss("x0", x0_pred_modelopt, x0=x0_real).item() + + # FastGen "direct" reference: x0 = x_t - t * flow. + t64 = _expand_like(t.to(torch.float64), x_t.to(torch.float64)) + x0_pred_ref = (x_t.to(torch.float64) - t64 * raw_flow.to(torch.float64)).to(x0_real.dtype) + loss_fastgen = _fastgen_dsm("x0", x0_pred_ref, x0=x0_real).item() + assert abs(loss_modelopt - loss_fastgen) < 1e-6 + + +# ---------------------------------------------------------------------------- # +# §2.9 — classifier-free guidance # +# ---------------------------------------------------------------------------- # + + +def test_classifier_free_guidance_matches_fastgen(): + torch.manual_seed(6) + cond = torch.randn(2, 16, 8, 8, dtype=torch.float32) + uncond = torch.randn_like(cond) + assert torch.equal(classifier_free_guidance(cond, uncond, 4.0), _fastgen_cfg(cond, uncond, 4.0)) + + +class _RecordingFlow(nn.Module): + def __init__(self): + super().__init__() + self.masks: list[torch.Tensor | None] = [] + + def forward(self, hidden_states, timestep, encoder_hidden_states=None, **kwargs): + mask = kwargs.get("encoder_hidden_states_mask") + self.masks.append(mask.detach().clone() if torch.is_tensor(mask) else None) + return torch.zeros_like(hidden_states) + + +def test_compute_student_loss_uses_separate_negative_cfg_mask(): + cfg = DMDConfig( + pred_type="flow", + num_train_timesteps=None, + student_sample_steps=1, + guidance_scale=4.0, + sample_t_cfg=SampleTimestepConfig(time_dist_type="uniform", min_t=0.001, max_t=0.999), + ) + student = _RecordingFlow() + teacher = _RecordingFlow() + fake_score = _RecordingFlow() + pipe = DMDPipeline(student=student, teacher=teacher, fake_score=fake_score, config=cfg) + + torch.manual_seed(7) + latents = torch.randn(2, 16, 4, 4) + noise = torch.randn_like(latents) + text = torch.randn(2, 8, 4) + neg_text = torch.randn(2, 3, 4) + text_mask = torch.ones(2, 8, dtype=torch.long) + neg_mask = torch.ones(2, 3, dtype=torch.long) + + pipe.compute_student_loss( + latents, + noise, + encoder_hidden_states=text, + encoder_hidden_states_mask=text_mask, + negative_encoder_hidden_states=neg_text, + negative_encoder_hidden_states_mask=neg_mask, + ) + + assert torch.equal(teacher.masks[0], text_mask) + assert torch.equal(teacher.masks[1], neg_mask) + + +# ---------------------------------------------------------------------------- # +# §2.10 — GAN gen/disc/R1 losses # +# ---------------------------------------------------------------------------- # + + +def test_gan_gen_loss_matches_fastgen(): + torch.manual_seed(7) + fake_logits = torch.randn(8, 1) + assert abs(gan_gen_loss(fake_logits).item() - _fastgen_gan_gen(fake_logits).item()) < 1e-6 + + +def test_gan_disc_loss_matches_fastgen(): + torch.manual_seed(7) + fake_logits = torch.randn(8, 1) + real_logits = torch.randn(8, 1) + assert ( + abs( + gan_disc_loss(real_logits, fake_logits).item() + - _fastgen_gan_disc(real_logits, fake_logits).item() + ) + < 1e-6 + ) diff --git a/tests/unit/torch/fastgen/test_dmd_pipeline_step.py b/tests/unit/torch/fastgen/test_dmd_pipeline_step.py new file mode 100644 index 00000000000..71d7a93e640 --- /dev/null +++ b/tests/unit/torch/fastgen/test_dmd_pipeline_step.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 DMD2 training-step test (pipeline-level). + +Exercises one full DMD2 optimizer step through the real ``QwenImageDMDPipeline`` +loss path — the student VSD phase and the fake-score DSM phase — on tiny stub +transformers, with no real Qwen weights, NCCL, or FSDP2. A single test +transitively covers the plugin's pack/unpack ``_call_model`` path, the VSD/DSM +loss math, gradient isolation between phases, an optimizer update, and a +checkpoint round-trip. +""" + +from __future__ import annotations + +import copy + +import torch +from torch import nn + +from modelopt.torch.fastgen.config import DMDConfig, SampleTimestepConfig +from modelopt.torch.fastgen.plugins.qwen_image import QwenImageDMDPipeline + + +class _TinyQwenTransformer(nn.Module): + """Grad-capable stub over the packed Qwen hidden states ``[B, num_patches, C*4]``. + + ``QwenImageDMDPipeline._call_model`` packs ``[B, C, H, W] -> [B, P, C*4]`` before + the forward and unpacks after, so a ``Linear`` over the last (``C*4``) dim gives a + real gradient path while matching the expected return shape. + """ + + def __init__(self, packed_dim: int = 64) -> None: + super().__init__() + self.proj = nn.Linear(packed_dim, packed_dim) + + def forward(self, hidden_states, **_kwargs): + return self.proj(hidden_states) + + +def _build_pipeline(): + torch.manual_seed(0) + student = _TinyQwenTransformer() + teacher = _TinyQwenTransformer() + fake_score = _TinyQwenTransformer() + cfg = DMDConfig( + pred_type="flow", + num_train_timesteps=None, # required by QwenImageDMDPipeline + student_sample_steps=1, + student_update_freq=5, + fake_score_pred_type="x0", + gan_loss_weight_gen=0.0, # no GAN branch -> no discriminator / feature hooks needed + guidance_scale=None, + sample_t_cfg=SampleTimestepConfig(time_dist_type="uniform", min_t=0.001, max_t=0.999), + ema=None, + ) + pipe = QwenImageDMDPipeline( + student=student, + teacher=teacher, + fake_score=fake_score, + config=cfg, + discriminator=None, + ) + return pipe, student, teacher, fake_score + + +def _mock_batch(batch_size: int = 2): + latents = torch.randn(batch_size, 16, 8, 8) # even dims -> packs to [B, 16, 64] + noise = torch.randn_like(latents) + text = torch.randn(batch_size, 8, 64) # shape is arbitrary; the stub ignores it + return latents, noise, text + + +def _snapshot(module: nn.Module) -> dict[str, torch.Tensor]: + return {k: v.detach().clone() for k, v in module.state_dict().items()} + + +def _params_changed(before: dict[str, torch.Tensor], module: nn.Module) -> bool: + return any(not torch.equal(before[k], v) for k, v in module.state_dict().items()) + + +def _train_only(active: nn.Module, *frozen: nn.Module) -> None: + active.train() + for p in active.parameters(): + p.requires_grad_(True) + for m in frozen: + m.eval() + for p in m.parameters(): + p.requires_grad_(False) + + +def test_dmd2_student_then_fake_score_step_updates_only_active_module(): + """One student VSD step then one fake-score DSM step: each phase yields a + finite loss, steps only its own module, and leaves the others untouched.""" + pipe, student, teacher, fake_score = _build_pipeline() + latents, noise, text = _mock_batch() + + # ---- student (VSD) phase ---- + _train_only(student, teacher, fake_score) + student_before, teacher_before = _snapshot(student), _snapshot(teacher) + opt_s = torch.optim.Adam(student.parameters(), lr=1e-2) + opt_s.zero_grad() + losses = pipe.compute_student_loss( + latents, + noise, + encoder_hidden_states=text, + negative_encoder_hidden_states=None, + guidance_scale=None, + ) + assert "vsd" in losses and "total" in losses + assert torch.isfinite(losses["total"]) + losses["total"].backward() + opt_s.step() + + assert _params_changed(student_before, student) # student learned + assert not _params_changed(teacher_before, teacher) # teacher stayed frozen + + # ---- fake-score (DSM) phase ---- + _train_only(fake_score, student, teacher) + student_after_student_phase = _snapshot(student) + fake_before = _snapshot(fake_score) + opt_f = torch.optim.Adam(fake_score.parameters(), lr=1e-2) + opt_f.zero_grad() + fs_losses = pipe.compute_fake_score_loss(latents, noise, encoder_hidden_states=text) + assert "fake_score" in fs_losses and "total" in fs_losses + assert torch.isfinite(fs_losses["total"]) + fs_losses["total"].backward() + opt_f.step() + + assert _params_changed(fake_before, fake_score) # fake_score learned + assert not _params_changed(student_after_student_phase, student) # student untouched + + +def test_dmd2_student_state_dict_round_trips(): + """After a training step, the student's state_dict reloads bit-exactly into a + fresh module — the save/restore contract the recipe's checkpointing relies on.""" + pipe, student, _teacher, _fake = _build_pipeline() + latents, noise, text = _mock_batch() + + _train_only(student, _teacher, _fake) + opt = torch.optim.Adam(student.parameters(), lr=1e-2) + opt.zero_grad() + pipe.compute_student_loss(latents, noise, encoder_hidden_states=text, guidance_scale=None)[ + "total" + ].backward() + opt.step() + + saved = copy.deepcopy(student.state_dict()) + reloaded = _TinyQwenTransformer() + reloaded.load_state_dict(saved) + for k, v in student.state_dict().items(): + assert torch.equal(v, reloaded.state_dict()[k]), k diff --git a/tests/unit/torch/fastgen/test_hook_requirements.py b/tests/unit/torch/fastgen/test_hook_requirements.py new file mode 100644 index 00000000000..00914b5d53d --- /dev/null +++ b/tests/unit/torch/fastgen/test_hook_requirements.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 "did you attach hooks?" / "is this FSDP-wrapped?" runtime guards. + +Covers R2.1 (GAN branches must raise a clear ``RuntimeError`` when +``teacher._fastgen_captured`` is missing) and R2.5 (``create_fake_score`` must +reject FSDP-wrapped teachers when ``deep_copy=True``). +""" + +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from modelopt.torch.fastgen import DMDConfig, DMDPipeline, create_fake_score + + +class _ToyTransformer(nn.Module): + """Linear-on-hidden-states transformer that matches the pipeline's call convention.""" + + def __init__(self, d: int) -> None: + super().__init__() + self.linear = nn.Linear(d, d, bias=False) + + def forward(self, hidden_states, timestep=None, encoder_hidden_states=None, **kwargs): + return self.linear(hidden_states) + + +class _ToyDiscriminator(nn.Module): + """Consumes ``list[Tensor]`` and returns 2D logits ``(B, 1)``.""" + + def forward(self, features): + x = features[0] + return x.flatten(start_dim=1).mean(dim=-1, keepdim=True) + + +def _build_gan_pipeline(d: int) -> DMDPipeline: + torch.manual_seed(0) + cfg = DMDConfig(pred_type="flow", gan_loss_weight_gen=0.03) + return DMDPipeline( + _ToyTransformer(d), + _ToyTransformer(d), + _ToyTransformer(d), + cfg, + discriminator=_ToyDiscriminator(), + ) + + +def test_compute_student_loss_raises_when_hooks_missing(): + """GAN-enabled ``compute_student_loss`` without ``attach_feature_capture`` + must raise a ``RuntimeError`` naming the attach helper, not strip under + ``-O`` like the previous ``assert``.""" + d, b = 4, 2 + pipeline = _build_gan_pipeline(d) + latents = torch.randn(b, d) + noise = torch.randn(b, d) + + with pytest.raises(RuntimeError, match="attach_feature_capture"): + pipeline.compute_student_loss(latents, noise) + + +def test_compute_discriminator_loss_raises_when_hooks_missing(): + """``compute_discriminator_loss`` without ``attach_feature_capture`` must + raise ``RuntimeError`` at the first drain, before the discriminator is called.""" + d, b = 4, 2 + pipeline = _build_gan_pipeline(d) + latents = torch.randn(b, d) + noise = torch.randn(b, d) + + with pytest.raises(RuntimeError, match="attach_feature_capture"): + pipeline.compute_discriminator_loss(latents, noise) + + +def test_create_fake_score_raises_on_fsdp2_wrapped(): + """``create_fake_score(deep_copy=True)`` on a module whose first parameter + looks like a DTensor (``full_tensor`` attribute) must raise with a message + pointing at the meta-init recipe in the docstring.""" + m = nn.Linear(4, 4) + # Monkey-patch the first parameter to look like a DTensor. + first = next(m.parameters()) + first.full_tensor = lambda: first # type: ignore[attr-defined] + + with pytest.raises(RuntimeError, match="meta-init"): + create_fake_score(m, deep_copy=True) + + +def test_create_fake_score_no_copy_skips_fsdp_check(): + """``deep_copy=False`` reuses the teacher directly — the FSDP-wrap check + is skipped because there is no ``copy.deepcopy`` to protect against.""" + m = nn.Linear(4, 4) + first = next(m.parameters()) + first.full_tensor = lambda: first # type: ignore[attr-defined] + + fake_score = create_fake_score(m, deep_copy=False) + assert fake_score is m + assert fake_score.training # create_fake_score must set .train() regardless + assert all(p.requires_grad for p in fake_score.parameters()) diff --git a/tests/unit/torch/fastgen/test_pred_type_conversion.py b/tests/unit/torch/fastgen/test_pred_type_conversion.py new file mode 100644 index 00000000000..69346b81838 --- /dev/null +++ b/tests/unit/torch/fastgen/test_pred_type_conversion.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Regression tests for fastgen pred_type conversion, timestep rescaling, and EMA dtype promotion. + +Covers three round-1 fix guards: + +1. ``fake_score_pred_type`` regression — under ``pred_type='flow'`` with + ``fake_score_pred_type='x0'``, the fake-score DSM loss must operate on the + ``x_0`` projection of the raw flow output, not on the raw flow tensor. +2. ``num_train_timesteps`` rescale — the pipeline must scale the RF + ``t ∈ [0, 1]`` to ``num_train_timesteps * t`` before passing it to the + transformer, when the knob is set. +3. EMA shadow dtype promotion — by default the shadow lives in ``float32`` + even when the live model is ``bfloat16``. +""" + +from __future__ import annotations + +import copy +from unittest import mock + +import torch +import torch.nn.functional as F +from torch import nn + +from modelopt.torch.fastgen import DMDConfig, DMDPipeline, EMAConfig, ExponentialMovingAverage +from modelopt.torch.fastgen.flow_matching import add_noise, pred_x0_from_flow +from modelopt.torch.fastgen.methods import dmd as dmd_module + + +class _ToyTransformer(nn.Module): + """Minimal diffusers-shaped transformer. Output is linear in ``hidden_states`` and + ignores ``timestep`` / ``encoder_hidden_states`` — keeps the expected-value + reconstruction analytic.""" + + def __init__(self, d: int) -> None: + super().__init__() + self.linear = nn.Linear(d, d, bias=False) + + def forward(self, hidden_states, timestep=None, encoder_hidden_states=None, **kwargs): + return self.linear(hidden_states) + + +class _TimestepEchoModel(nn.Module): + """Returns ``hidden_states + timestep`` broadcast — used to observe the rescale knob.""" + + def forward(self, hidden_states, timestep, encoder_hidden_states=None, **kwargs): + return hidden_states + timestep.view(-1, 1).to(hidden_states.dtype) + + +def _build_pipeline( + d: int, + *, + pred_type: str = "flow", + fake_score_pred_type: str | None = "x0", + num_train_timesteps: int | None = None, +): + torch.manual_seed(0) + student = _ToyTransformer(d) + teacher = _ToyTransformer(d) + fake_score = copy.deepcopy(teacher) + cfg = DMDConfig( + pred_type=pred_type, + fake_score_pred_type=fake_score_pred_type, + num_train_timesteps=num_train_timesteps, + ) + pipeline = DMDPipeline(student, teacher, fake_score, cfg) + return pipeline, student, teacher, fake_score, cfg + + +def test_fake_score_dsm_matches_manual_flow_to_x0(): + """compute_fake_score_loss under ``(flow, x0)`` must equal the manual + ``F.mse_loss(gen_data, pred_x0_from_flow(raw, x_t, t))`` reconstruction.""" + d, b = 4, 2 + pipeline, student, _teacher, fake_score, cfg = _build_pipeline( + d, pred_type="flow", fake_score_pred_type="x0" + ) + latents = torch.randn(b, d) + noise = torch.randn(b, d) + + fixed_t = torch.full((b,), 0.5) + fixed_eps = torch.randn(b, d) + + with ( + mock.patch.object(pipeline, "sample_timesteps", return_value=fixed_t), + mock.patch.object(dmd_module.torch, "randn_like", return_value=fixed_eps), + ): + actual = pipeline.compute_fake_score_loss(latents, noise)["fake_score"] + + with torch.no_grad(): + max_t = cfg.sample_t_cfg.max_t + t_student = torch.full((b,), max_t, dtype=torch.float32) + input_student = noise * max_t + gen_data = pred_x0_from_flow( + student(hidden_states=input_student, timestep=t_student), + input_student, + t_student, + ) + perturbed = add_noise(gen_data, fixed_eps, fixed_t) + fake_raw = fake_score(hidden_states=perturbed, timestep=fixed_t) + pred_x0 = pred_x0_from_flow(fake_raw, perturbed, fixed_t) + expected = F.mse_loss(gen_data, pred_x0) + + assert torch.allclose(actual, expected, atol=1e-6), ( + f"compute_fake_score_loss={actual.item():.3e}, manual={expected.item():.3e}" + ) + + +def test_student_vsd_sees_x0_not_raw_flow(): + """compute_student_loss must feed ``vsd_loss`` the x0-converted flow output of the + fake score (the prior bug was to forward the raw flow tensor instead).""" + d, b = 4, 2 + pipeline, student, _teacher, fake_score, cfg = _build_pipeline( + d, pred_type="flow", fake_score_pred_type="x0" + ) + latents = torch.randn(b, d) + noise = torch.randn(b, d) + fixed_t = torch.full((b,), 0.5) + fixed_eps = torch.randn(b, d) + + captured: dict[str, torch.Tensor] = {} + orig_vsd_loss = dmd_module.vsd_loss + + def spy(gen_data, teacher_x0, fake_score_x0, additional_scale=None): + captured["fake_score_x0"] = fake_score_x0.detach().clone() + return orig_vsd_loss(gen_data, teacher_x0, fake_score_x0, additional_scale) + + with ( + mock.patch.object(pipeline, "sample_timesteps", return_value=fixed_t), + mock.patch.object(dmd_module.torch, "randn_like", return_value=fixed_eps), + mock.patch.object(dmd_module, "vsd_loss", side_effect=spy), + ): + pipeline.compute_student_loss(latents, noise) + + with torch.no_grad(): + max_t = cfg.sample_t_cfg.max_t + t_student = torch.full((b,), max_t, dtype=torch.float32) + input_student = noise * max_t + gen_data_expected = pred_x0_from_flow( + student(hidden_states=input_student, timestep=t_student), + input_student, + t_student, + ) + perturbed = add_noise(gen_data_expected, fixed_eps, fixed_t) + fake_raw = fake_score(hidden_states=perturbed, timestep=fixed_t) + fake_x0_expected = pred_x0_from_flow(fake_raw, perturbed, fixed_t) + + assert torch.allclose(captured["fake_score_x0"], fake_x0_expected, atol=1e-6) + + +def test_call_model_rescales_timestep_when_num_train_timesteps_set(): + """``num_train_timesteps=1000`` rescales ``t`` by 1000 inside ``_call_model`` + and casts it to ``hidden_states.dtype`` (matching FastGen's VaceWan wrapper); + ``None`` leaves ``t`` untouched.""" + d, b = 3, 2 + model = _TimestepEchoModel() + x = torch.zeros(b, d) + t = torch.tensor([0.1, 0.7]) + + pipe_scaled = DMDPipeline( + _ToyTransformer(d), + _ToyTransformer(d), + _ToyTransformer(d), + DMDConfig(pred_type="x0", num_train_timesteps=1000), + ) + out_scaled = pipe_scaled._call_model(model, x, t) + assert torch.allclose(out_scaled, x + (t * 1000.0).view(-1, 1)) + + pipe_none = DMDPipeline( + _ToyTransformer(d), + _ToyTransformer(d), + _ToyTransformer(d), + DMDConfig(pred_type="x0", num_train_timesteps=None), + ) + out_none = pipe_none._call_model(model, x, t) + assert torch.allclose(out_none, x + t.view(-1, 1)) + + # bf16 hidden_states: rescaled timestep must be cast to bf16 so the + # addition inside the model returns a bf16 tensor (parity with FastGen's + # ``.to(dtype=x_t.dtype)``). fp32 timestep + bf16 hidden_states without the + # cast would either upcast the result or push dtype juggling into the model. + x_bf16 = torch.zeros(b, d, dtype=torch.bfloat16) + out_bf16 = pipe_scaled._call_model(model, x_bf16, t) + assert out_bf16.dtype == torch.bfloat16 + assert torch.allclose( + out_bf16, + x_bf16 + (t * 1000.0).to(torch.bfloat16).view(-1, 1), + ) + + +def test_ema_shadow_dtype_promotion(): + """``EMAConfig.dtype='float32'`` on a bf16 student gives an fp32 shadow; + ``dtype=None`` falls back to the live parameter's dtype.""" + torch.manual_seed(0) + student_bf16 = nn.Linear(4, 4).to(torch.bfloat16) + + cfg_fp32 = EMAConfig(fsdp2=False, dtype="float32") + ema_fp32 = ExponentialMovingAverage(student_bf16, cfg_fp32) + for name, shadow in ema_fp32.state_dict().items(): + assert shadow.dtype == torch.float32, f"{name}: expected float32, got {shadow.dtype}" + + cfg_none = EMAConfig(fsdp2=False, dtype=None) + ema_none = ExponentialMovingAverage(student_bf16, cfg_none) + for name, shadow in ema_none.state_dict().items(): + assert shadow.dtype == torch.bfloat16, f"{name}: expected bfloat16, got {shadow.dtype}" diff --git a/tests/unit/torch/fastgen/test_qwen_image_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_plugin.py new file mode 100644 index 00000000000..498b6ce5f9f --- /dev/null +++ b/tests/unit/torch/fastgen/test_qwen_image_plugin.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Unit tests for ``modelopt.torch.fastgen.plugins.qwen_image``. + +Ports the checklist §1 bullets (pack/unpack + FastGen parity + _call_model wiring) +into pytest form so they live in-repo and run under ``pytest tests/examples/diffusers``. +Adds the §6-specific bullet ``num_train_timesteps != None`` constructor error. + +The parity comparison against the FastGen reference extract is bit-exact +(``torch.equal``) — both are pure permute+reshape operations with no +floating-point arithmetic. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from modelopt.torch.fastgen import DMDConfig +from modelopt.torch.fastgen.plugins.qwen_image import ( + QwenImageDMDPipeline, + build_img_shapes, + pack_latents, + unpack_latents, +) + +# ---------------------------------------------------------------------------- # +# §1.1 — pack / unpack inverse for representative latent sizes # +# ---------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("shape", "dtype"), + [ + ((1, 16, 128, 128), torch.float32), # production size + ((2, 16, 32, 32), torch.bfloat16), # batch>1 + bf16 dtype preserved + ], +) +def test_pack_unpack_inverse(shape, dtype): + """Round-trip is bit-exact; dtype/device/contiguity preserved (incl. bf16).""" + x = torch.randn(*shape, dtype=dtype) + p = pack_latents(x) + y = unpack_latents(p, shape[2], shape[3]) + assert torch.equal(x, y) + assert x.dtype == p.dtype == y.dtype + assert x.device == p.device == y.device + assert p.is_contiguous() + assert y.is_contiguous() + + +# ---------------------------------------------------------------------------- # +# §1.2 — odd spatial dims raise a clear ValueError # +# ---------------------------------------------------------------------------- # + + +def test_pack_rejects_odd_spatial(): + with pytest.raises(ValueError, match="even"): + pack_latents(torch.randn(1, 16, 31, 32)) + + +def test_unpack_rejects_odd_target(): + with pytest.raises(ValueError, match="even"): + unpack_latents(torch.randn(1, 256, 64), 31, 32) + + +# ---------------------------------------------------------------------------- # +# §1.5 — parity vs the FastGen reference # +# # +# FastGen's QwenImage class pulls heavy deps; we inline the two methods # +# verbatim from ``source/FastGen/fastgen/networks/QwenImage/network.py`` so # +# the parity check is hermetic. # +# ---------------------------------------------------------------------------- # + + +def _fastgen_pack(latents: torch.Tensor) -> torch.Tensor: + batch_size, channels, height, width = latents.shape + latents = latents.view(batch_size, channels, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + return latents.reshape(batch_size, (height // 2) * (width // 2), channels * 4) + + +def _fastgen_unpack(latents: torch.Tensor, height: int, width: int) -> torch.Tensor: + batch_size = latents.shape[0] + channels = latents.shape[2] // 4 + latents = latents.reshape(batch_size, height // 2, width // 2, channels, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + return latents.reshape(batch_size, channels, height, width) + + +def test_pack_unpack_parity_vs_fastgen(): + x = torch.randn(2, 16, 32, 32) + fg_p = _fastgen_pack(x) + mo_p = pack_latents(x) + assert torch.equal(fg_p, mo_p) + + fg_u = _fastgen_unpack(fg_p, 32, 32) + mo_u = unpack_latents(mo_p, 32, 32) + assert torch.equal(fg_u, mo_u) + assert torch.equal(fg_u, x) # FastGen unpack round-trips the input + + +# ---------------------------------------------------------------------------- # +# §1.6 — build_img_shapes structural equality # +# ---------------------------------------------------------------------------- # + + +def test_build_img_shapes_structure(): + out = build_img_shapes(batch_size=2, h_lat=32, w_lat=32) + assert out == [[(1, 16, 16)], [(1, 16, 16)]] + + +# ---------------------------------------------------------------------------- # +# §1.7 / §1.8 — _call_model kwarg forwarding + unpack return styles # +# ---------------------------------------------------------------------------- # + + +class _CapturingModel(nn.Module): + """Stub transformer that records the kwargs it was called with and emits + a zero tensor of the requested shape in one of three return styles.""" + + def __init__(self, out_shape: tuple[int, ...], style: str = "tensor") -> None: + super().__init__() + self.out_shape = out_shape + self.style = style + self.last_kwargs: dict[str, object] = {} + + def forward(self, **kwargs): + self.last_kwargs = dict(kwargs) + out = torch.zeros(*self.out_shape, dtype=kwargs["hidden_states"].dtype) + if self.style == "tensor": + return out + if self.style == "tuple": + return (out, "extra") + if self.style == "sample": + return SimpleNamespace(sample=out) + raise ValueError(self.style) + + +def _make_pipeline(student: nn.Module) -> QwenImageDMDPipeline: + return QwenImageDMDPipeline( + student=student, + teacher=nn.Identity(), + fake_score=nn.Identity(), + config=DMDConfig(num_train_timesteps=None), + discriminator=None, + ) + + +def test_call_model_forwards_qwen_kwargs(): + """``_call_model`` must forward the exact Qwen signature (hidden_states packed + to ``[B, num_patches, 64]``, encoder_hidden_states verbatim, + encoder_hidden_states_mask verbatim, txt_seq_lens derived from the mask, + img_shapes as ``[[(1, h//2, w//2)]] * B``, guidance=None, return_dict=False, + timestep verbatim with no /1000 rescale).""" + b, c, h, w = 2, 16, 32, 32 + student = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4), style="tensor") + pipe = _make_pipeline(student) + + hidden = torch.randn(b, c, h, w) + t = torch.full((b,), 0.5, dtype=torch.float32) + enc = torch.randn(b, 512, 3584) + mask = torch.zeros(b, 512, dtype=torch.long) + mask[0, :37] = 1 + mask[1, :42] = 1 + + out = pipe._call_model( + student, + hidden, + t, + encoder_hidden_states=enc, + encoder_hidden_states_mask=mask, + ) + + kw = student.last_kwargs + assert tuple(kw["hidden_states"].shape) == (b, (h // 2) * (w // 2), c * 4) + assert tuple(kw["encoder_hidden_states"].shape) == (b, 512, 3584) + assert torch.equal(kw["encoder_hidden_states_mask"], mask) + assert kw["txt_seq_lens"] == [37, 42] + assert kw["img_shapes"] == [[(1, h // 2, w // 2)]] * b + assert kw["guidance"] is None + assert kw["return_dict"] is False + assert torch.equal(kw["timestep"], t) # no /1000 rescale when num_train_timesteps=None + assert tuple(out.shape) == (b, c, h, w) + + +@pytest.mark.parametrize("style", ["tensor", "tuple", "sample"]) +def test_call_model_unpacks_return_styles(style): + """``_call_model`` must unpack ``tensor`` / ``tuple`` / ``.sample`` return + styles into ``[B, C, H, W]`` of the input's dtype.""" + b, c, h, w = 1, 16, 32, 32 + model = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4), style=style) + pipe = _make_pipeline(model) + hidden = torch.randn(b, c, h, w) + t = torch.full((b,), 0.5, dtype=torch.float32) + enc = torch.randn(b, 512, 3584) + out = pipe._call_model(model, hidden, t, encoder_hidden_states=enc) + assert tuple(out.shape) == (b, c, h, w) + assert out.dtype == hidden.dtype + + +# ---------------------------------------------------------------------------- # +# §6.X — QwenImageDMDPipeline constructor rejects num_train_timesteps != None # +# ---------------------------------------------------------------------------- # + + +def test_constructor_rejects_non_null_num_train_timesteps(): + """The pipeline normalizes ``t ∈ [0, 1]`` internally and forwards continuous + ``t`` to the Qwen transformer. ``num_train_timesteps`` is a discretization + knob that doesn't apply — the constructor must refuse it loudly.""" + cfg = DMDConfig(num_train_timesteps=1000) + with pytest.raises(ValueError, match="num_train_timesteps"): + QwenImageDMDPipeline( + student=nn.Identity(), + teacher=nn.Identity(), + fake_score=nn.Identity(), + config=cfg, + discriminator=None, + )