Add DMD2 distillation for Qwen-Image (fastgen)

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 <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/fastgen>`_ for details.
 
 **Bug Fixes**
 

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)

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=<gpus> \
+#       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