From 1678e3629bcfc900692d9bef355e3ed9ed6dc202 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Wed, 5 Aug 2026 00:20:09 -0700 Subject: [PATCH] Add FastGen quantization-aware distillation Signed-off-by: Jingyu Xin --- examples/diffusers/fastgen/qad/README.md | 174 +++++ examples/diffusers/fastgen/qad/__init__.py | 16 + examples/diffusers/fastgen/qad/artifacts.py | 428 +++++++++++++ .../fastgen/qad/configs/qwen_image_nvfp4.yaml | 108 ++++ .../configs/qwen_image_svdquant_nvfp4.yaml | 112 ++++ examples/diffusers/fastgen/qad/finetune.py | 43 ++ examples/diffusers/fastgen/qad/modeling.py | 259 ++++++++ examples/diffusers/fastgen/qad/pipeline.py | 74 +++ examples/diffusers/fastgen/qad/recipe.py | 595 ++++++++++++++++++ 9 files changed, 1809 insertions(+) create mode 100644 examples/diffusers/fastgen/qad/README.md create mode 100644 examples/diffusers/fastgen/qad/__init__.py create mode 100644 examples/diffusers/fastgen/qad/artifacts.py create mode 100644 examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml create mode 100644 examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml create mode 100644 examples/diffusers/fastgen/qad/finetune.py create mode 100644 examples/diffusers/fastgen/qad/modeling.py create mode 100644 examples/diffusers/fastgen/qad/pipeline.py create mode 100644 examples/diffusers/fastgen/qad/recipe.py diff --git a/examples/diffusers/fastgen/qad/README.md b/examples/diffusers/fastgen/qad/README.md new file mode 100644 index 00000000000..175f59a2e02 --- /dev/null +++ b/examples/diffusers/fastgen/qad/README.md @@ -0,0 +1,174 @@ +# FastGen Quantization-Aware Distillation + +This example trains a quantized diffusion student against a frozen, BF16 +Diffusers teacher with ModelOpt's distillation API. It is a standalone FastGen +recipe: it does not use DMD2, reduce the sampling schedule, create a fake-score +model, or add a GAN/EMA training phase. + +The initial Qwen-Image recipe uses the official `Qwen/Qwen-Image` Diffusers +checkpoint as the teacher. Set `qad.teacher_model_name_or_path` to +`nvidia/Qwen-Image-Flash` when the four-step DMD2-trained Qwen-Image checkpoint +should be the teacher instead. Both follow the standard Diffusers checkpoint +interface. QAD intentionally does not interpret FastGen/DMD2 intermediate +checkpoint sidecars or standalone transformer safetensors as teacher inputs. + +Every training micro-batch samples one noisy latent and one timestep, then sends +the same latent, timestep, prompt conditioning, and guidance inputs to the +teacher and student. + +## Supported students + +The `qad.student.mode` field selects one of two restore contracts. + +### Regular NVFP4 + +Set `qad.student.mode=nvfp4` and point +`model.pretrained_model_name_or_path` at the unquantized Diffusers model. The recipe +loads its weights first and then restores the weight-free ModelOpt NVFP4 state +from `qad.student.quant_state_path`. The quantizer state must have been calibrated +against exactly the same student weights. This mode trains all student +parameters, so its only valid `train_scope` is `all`. + +Use [`configs/qwen_image_nvfp4.yaml`](configs/qwen_image_nvfp4.yaml) as the +starting configuration. + +### NVFP4 SVDQuant with Hugging Face PEFT + +Set `qad.student.mode=nvfp4_svdquant` and point +`model.pretrained_model_name_or_path` at a user-prepared, ModelOpt-enabled Diffusers +training bundle. The bundle must contain the complete SVDQuant student: + +- a DiffusionPipeline root with `model_index.json` (not only a standalone + transformer `save_pretrained` directory); +- the ModelOpt topology and quantizer state; +- the residual weights produced by SVDQuant calibration; and +- the Hugging Face PEFT A/B factors for the SVDQuant low-rank branch. + +For the standard Diffusers layout, the transformer files and ModelOpt sidecar +are under `transformer/`, including `transformer/modelopt_state.pth`. The path +given to QAD is the parent DiffusionPipeline directory. + +A weight-free NVFP4 quantizer-state file is not a valid SVDQuant bundle. +SVDQuant subtracts the low-rank branch from the original weight, so both the +resulting residual weight and the PEFT factors are required. A unified +deployment export is also not a training bundle and must not be used here. + +The SVDQuant topology is restored before FSDP and before optimizer construction. +`qad.student.train_scope=all` is the default and trains both the residual/base +parameters and the PEFT factors. Set it to `lora_only` to freeze every student +parameter except the SVDQuant PEFT A/B factors. In both scopes, +`pre_quant_scale` remains ModelOpt buffer state and is never placed in the +optimizer. + +Use [`configs/qwen_image_svdquant_nvfp4.yaml`](configs/qwen_image_svdquant_nvfp4.yaml) +as the starting configuration. + +QAD is restore-only in both modes. It does not calibrate a student during +distributed training. + +## Distillation losses + +Output distillation is always MSE. The canonical setting is: + +```yaml +qad: + output_loss: + type: mse + weight: 1.0 + task_loss: + weight: 0.0 +``` + +At `weight: 1.0`, the optimized objective is pure teacher-output MSE and the +ordinary flow-matching target has weight zero because `task_loss.weight` defaults +to `0.0`. All coefficients are independent and additive. For example, setting +both output and task weights to `0.5` produces an equal output-MSE/flow-matching +mixture; adding layerwise terms does not silently renormalize either coefficient. + +Optional layerwise MSE can be added without changing the output loss: + +```yaml +qad: + layerwise: + enabled: true + pairs: + - student_layer: transformer_blocks.29 + teacher_layer: transformer_blocks.29 + selector: hidden_states + weight: 0.05 +``` + +Each pair is an exact module name relative to the student or teacher +transformer. Its weight is additive to the output/task objective. Start with +output-only training: layer hooks retain activations and therefore increase +memory use, especially when activation checkpointing is enabled. + +The recipe logs the flow-matching loss, output MSE, every configured layerwise +MSE, and the final combined loss separately. + +## Configuration and launch + +The entry point is `examples/diffusers/fastgen/qad/finetune.py`. It uses the +same YAML plus dotted-command-line override convention as the other FastGen +recipes: + +```bash +torchrun --nproc-per-node=4 \ + examples/diffusers/fastgen/qad/finetune.py \ + --config examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml \ + --fsdp.dp_size=4 \ + --model.pretrained_model_name_or_path=/path/to/qwen-image-nvfp4-svdquant-training-bundle \ + --data.dataloader.cache_dir=/path/to/qwen_image_1024p \ + --checkpoint.checkpoint_dir=/path/to/qad/checkpoints +``` + +Cluster launchers can keep the established `CONFIG`, `RUN_ID`, and +`EXTRA_ARGS` interface. For example: + +```bash +EXTRA_ARGS="--step_scheduler.max_steps=50000 \ +--step_scheduler.ckpt_every_steps=1000 \ +--step_scheduler.num_epochs=200 \ +--step_scheduler.global_batch_size=64 \ +--optim.learning_rate=2e-6 \ +--lr_scheduler.min_lr=2e-6 \ +--fsdp.dp_size=64 \ +--qad.teacher_model_name_or_path=Qwen/Qwen-Image \ +--qad.output_loss.weight=1.0 \ +--qad.task_loss.weight=0.0 \ +--qad.student.mode=nvfp4_svdquant \ +--model.pretrained_model_name_or_path=/path/to/qwen-image-nvfp4-svdquant-training-bundle \ +--qad.student.train_scope=all \ +--data.dataloader.cache_dir=/path/to/qwen_image_1024p" \ +CONFIG=examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml \ +RUN_ID=qad_qwen_image_svdquant_nvfp4_16n \ +NODES=16 \ +GPUS_PER_NODE=4 \ +TIME=05:00:00 \ +PARTITION=batch \ +bash /path/to/experiments/qad_qwen_image/launch.sh +``` + +The launcher must invoke `examples/diffusers/fastgen/qad/finetune.py`. +Pointing the existing DMD2 launcher at a QAD YAML is not sufficient when that +launcher still hard-codes `dmd2_finetune.py`. + +The launch environment contains no Attention Grill settings. It also contains +no DMD2 timestep, fake-score, discriminator, negative-prompt, GAN, or EMA +settings. + +## Restore and checkpoint invariants + +On a fresh run the recipe restores the complete student first, constructs its +final ModelOpt/PEFT topology, applies FSDP, builds the optimizer from the selected +training scope, and only then creates the frozen teacher and distillation +controller. On resume, the same immutable student source reconstructs the +topology before the QAD checkpoint is loaded. + +The teacher and the transient ModelOpt distillation controller are not training +checkpoint payloads. Checkpoints contain the student state required by the +selected training scope together with optimizer, scheduler, dataloader, RNG, and +global-step state. Resolved dotted CLI overrides are materialized into the saved +`config.yaml`. Resume validates the student bundle, quantizer state, mode, train +scope, teacher, and loss configuration before loading optimizer shards; do not +change them while resuming an existing run. diff --git a/examples/diffusers/fastgen/qad/__init__.py b/examples/diffusers/fastgen/qad/__init__.py new file mode 100644 index 00000000000..d44f8e3a87e --- /dev/null +++ b/examples/diffusers/fastgen/qad/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Quantization-aware distillation example for Diffusers models.""" diff --git a/examples/diffusers/fastgen/qad/artifacts.py b/examples/diffusers/fastgen/qad/artifacts.py new file mode 100644 index 00000000000..2721fb93eeb --- /dev/null +++ b/examples/diffusers/fastgen/qad/artifacts.py @@ -0,0 +1,428 @@ +# 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. + +"""Student artifact restore and train-scope handling for the QAD example. + +The generic AutoModel diffusion builder intentionally owns FSDP and optimizer +construction. QAD only needs two narrowly-scoped hooks around that builder: + +* restore a standalone, weight-free ModelOpt state before FSDP for regular NVFP4; +* after FSDP, optionally freeze everything except ModelOpt SVDQuant's HF PEFT A/B + parameters and rebuild AdamW from the live sharded parameters. + +SVDQuant itself is never calibrated here. Its complete topology and weights must +already be present in a ModelOpt-enabled Diffusers training bundle. +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import inspect +import logging +import os +import re +from typing import TYPE_CHECKING, Any + +import modelopt.torch.opt as mto +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import set_quantizer_state_dict + +if TYPE_CHECKING: + from collections.abc import Iterator + + import torch + from torch import nn + +_SVDQUANT_PARAMETER_RE = re.compile(r"(?:^|\.)lora_[AB]\.modelopt_svdquant\.weight$") +_SUPPORTED_STUDENT_MODES = frozenset({"nvfp4", "nvfp4_svdquant"}) +_SUPPORTED_TRAIN_SCOPES = frozenset({"all", "lora_only"}) + + +@dataclasses.dataclass(frozen=True) +class StudentSettings: + """Resolved ``qad.student`` configuration.""" + + mode: str + model_name_or_path: str + train_scope: str = "all" + quant_state_path: str | None = None + + def validate(self) -> None: + if self.mode not in _SUPPORTED_STUDENT_MODES: + raise ValueError( + f"qad.student.mode must be one of {sorted(_SUPPORTED_STUDENT_MODES)}, " + f"got {self.mode!r}." + ) + if not self.model_name_or_path: + raise ValueError("model.pretrained_model_name_or_path is required for the student.") + if self.train_scope not in _SUPPORTED_TRAIN_SCOPES: + raise ValueError( + f"qad.student.train_scope must be 'all' or 'lora_only', got {self.train_scope!r}." + ) + + if self.mode == "nvfp4": + if not self.quant_state_path: + raise ValueError( + "qad.student.quant_state_path is required for a regular NVFP4 student." + ) + if self.train_scope != "all": + raise ValueError("Regular NVFP4 supports only qad.student.train_scope=all.") + elif self.quant_state_path: + raise ValueError( + "A SVDQuant student must be restored from a complete ModelOpt-enabled " + "Diffusers training bundle; do not set qad.student.quant_state_path." + ) + + +@dataclasses.dataclass +class StudentBuildState: + """Information captured while AutoModel builds the student.""" + + parallel_scheme: dict[str, dict[str, Any]] | None = None + quantizer_count: int = 0 + svdquant_parameter_names: tuple[str, ...] = () + + +def _restore_regular_quant_state(model: nn.Module, path: str, device: torch.device) -> int: + """Restore a regular weight-free ModelOpt quantizer state before FSDP.""" + if not os.path.isfile(path): + raise FileNotFoundError(f"ModelOpt quantizer state does not exist: {path}") + if mto.ModeloptStateManager.is_converted(model): + raise RuntimeError( + "The regular NVFP4 base model already contains ModelOpt topology while " + "qad.student.quant_state_path was also supplied. Use exactly one restore source." + ) + + state = mto.load_modelopt_state(path) + mode_states = state.get("modelopt_state_dict", ()) + if any( + mode_name == "svdquant_calibrate" or mode_state.get("metadata", {}).get("svdquant_peft") + for mode_name, mode_state in mode_states + ): + raise RuntimeError( + "qad.student.mode=nvfp4 cannot restore an SVDQuant state. Use " + "mode=nvfp4_svdquant with a complete ModelOpt-enabled Diffusers bundle." + ) + quantizer_state = state.pop("modelopt_state_weights", None) + if quantizer_state is None: + raise RuntimeError( + f"{path} has no modelopt_state_weights payload. Expected the weight-free " + "quantizer state produced by the Diffusers quantization example." + ) + + mto.restore_from_modelopt_state(model, state) + set_quantizer_state_dict(model, quantizer_state) + + quantizers = [module for module in model.modules() if isinstance(module, TensorQuantizer)] + if not quantizers: + raise RuntimeError(f"No TensorQuantizer modules were restored from {path}.") + for quantizer in quantizers: + quantizer.to(device) + _validate_nvfp4_quantizers(model, artifact_name=path) + + logging.info( + "[QAD] restored regular ModelOpt quantizer state before FSDP: %s (%d quantizers)", + path, + len(quantizers), + ) + return len(quantizers) + + +def _is_block16_nvfp4(quantizer: TensorQuantizer) -> bool: + block_sizes = quantizer.block_sizes or {} + return bool( + (quantizer.is_nvfp4_dynamic or quantizer.is_nvfp4_static) and block_sizes.get(-1) == 16 + ) + + +def _enabled_quantizer_leaves(module: Any) -> tuple[TensorQuantizer, ...]: + if module is None or not hasattr(module, "modules"): + return () + return tuple( + child + for child in module.modules() + if isinstance(child, TensorQuantizer) and child.is_enabled + ) + + +def _validate_nvfp4_quantizers( + model: nn.Module, + *, + artifact_name: str, + required_targets: tuple[str, ...] = (), +) -> None: + """Reject non-NVFP4 artifacts before FSDP obscures their module topology.""" + enabled_by_slot: dict[str, list[tuple[str, TensorQuantizer]]] = { + "weight": [], + "input": [], + } + for name, module in model.named_modules(): + if not isinstance(module, TensorQuantizer) or not module.is_enabled: + continue + path_parts = name.split(".") + for slot in enabled_by_slot: + if f"{slot}_quantizer" in path_parts: + enabled_by_slot[slot].append((name, module)) + + missing_slots = [slot for slot, entries in enabled_by_slot.items() if not entries] + if missing_slots: + raise RuntimeError( + f"{artifact_name} is not an NVFP4 W4A4 training artifact: no enabled " + + "/".join(missing_slots) + + " quantizers were found." + ) + + incompatible = [ + name + for entries in enabled_by_slot.values() + for name, quantizer in entries + if not _is_block16_nvfp4(quantizer) + ] + if incompatible: + raise RuntimeError( + f"{artifact_name} contains enabled GEMM quantizers that are not block-16 NVFP4 " + "(E2M1 values with E4M3 scales): " + ", ".join(incompatible[:5]) + ) + + for target_name in required_targets: + target = model.get_submodule(target_name) + get_base_layer = getattr(target, "get_base_layer", None) + base_layer = get_base_layer() if callable(get_base_layer) else target + for slot in ("weight", "input"): + leaves = _enabled_quantizer_leaves(getattr(base_layer, f"{slot}_quantizer", None)) + if not leaves or any(not _is_block16_nvfp4(quantizer) for quantizer in leaves): + raise RuntimeError( + f"SVDQuant target {target_name!r} does not have an enabled block-16 " + f"NVFP4 {slot}_quantizer." + ) + + logging.info( + "[QAD] validated block-16 NVFP4 W4A4 quantizers before FSDP: %d weight, %d input", + len(enabled_by_slot["weight"]), + len(enabled_by_slot["input"]), + ) + + +def _svdquant_metadata(model: nn.Module) -> dict[str, Any] | None: + if not mto.ModeloptStateManager.is_converted(model): + return None + for mode_name, mode_state in mto.modelopt_state(model)["modelopt_state_dict"]: + if mode_name == "svdquant_calibrate": + return mode_state.get("metadata", {}).get("svdquant_peft") + return None + + +def _validate_svdquant_bundle(model: nn.Module) -> tuple[str, ...]: + metadata = _svdquant_metadata(model) + if not metadata: + raise RuntimeError( + "qad.student.mode=nvfp4_svdquant requires a ModelOpt-enabled Diffusers " + "training bundle with svdquant_peft metadata. The topology must be restored " + "during Diffusers from_pretrained(), before FSDP." + ) + + expected_targets = tuple(metadata.get("target_modules", ())) + names = tuple( + name for name, _ in model.named_parameters() if _SVDQUANT_PARAMETER_RE.search(name) + ) + expected_names = { + f"{target_name}.lora_{factor}.modelopt_svdquant.weight" + for target_name in expected_targets + for factor in ("A", "B") + } + if not expected_targets or set(names) != expected_names: + raise RuntimeError( + "The SVDQuant bundle did not restore a complete pair of " + "lora_A/lora_B.modelopt_svdquant weights for every target module. " + "A weight-free quantizer state or a deployment export is not a valid " + "QAD training bundle." + ) + _validate_nvfp4_quantizers( + model, + artifact_name="The SVDQuant student bundle", + required_targets=expected_targets, + ) + missing_pre_quant_scale_buffers: list[str] = [] + for target_name in expected_targets: + target = model.get_submodule(target_name) + get_base_layer = getattr(target, "get_base_layer", None) + base_layer = get_base_layer() if callable(get_base_layer) else target + input_quantizer = getattr(base_layer, "input_quantizer", None) + pre_quant_scale = getattr(input_quantizer, "_pre_quant_scale", None) + if ( + pre_quant_scale is None + or getattr(input_quantizer, "_buffers", {}).get("_pre_quant_scale") + is not pre_quant_scale + ): + missing_pre_quant_scale_buffers.append(target_name) + if missing_pre_quant_scale_buffers: + raise RuntimeError( + "SVDQuant pre_quant_scale must be restored as frozen TensorQuantizer buffer " + "state for every target; missing or non-buffer targets: " + + ", ".join(missing_pre_quant_scale_buffers[:5]) + ) + logging.info( + "[QAD] validated SVDQuant training bundle before FSDP: %d targets, %d A/B tensors", + len(expected_targets), + len(names), + ) + return names + + +def _apply_train_scope(model: nn.Module, scope: str) -> list[nn.Parameter]: + if scope == "lora_only": + for name, parameter in model.named_parameters(): + parameter.requires_grad_(_SVDQUANT_PARAMETER_RE.search(name) is not None) + + trainable = [parameter for parameter in model.parameters() if parameter.requires_grad] + if not trainable: + raise RuntimeError(f"qad.student.train_scope={scope!r} selected no parameters.") + + if scope == "lora_only": + live_names = tuple( + name for name, parameter in model.named_parameters() if parameter.requires_grad + ) + invalid = [name for name in live_names if not _SVDQUANT_PARAMETER_RE.search(name)] + if invalid: + raise RuntimeError( + "lora_only left non-SVDQuant parameters trainable: " + ", ".join(invalid[:5]) + ) + + parameter_pre_scales = [ + name for name, _ in model.named_parameters() if "pre_quant_scale" in name + ] + if parameter_pre_scales: + raise RuntimeError( + "pre_quant_scale must remain a buffer and must never enter the optimizer: " + + ", ".join(parameter_pre_scales[:5]) + ) + return trainable + + +def _rebuild_optimizer_from_live_parameters( + optimizer: torch.optim.Optimizer, + parameters: list[nn.Parameter], +) -> torch.optim.Optimizer: + """Recreate the just-built optimizer without carrying stale parameter refs.""" + if optimizer.state: + raise RuntimeError("QAD expected a newly-created optimizer with no state.") + if len(optimizer.param_groups) != 1: + raise RuntimeError( + "QAD lora_only currently expects AutoModel to create one optimizer parameter group." + ) + return type(optimizer)(parameters, **dict(optimizer.defaults)) + + +def _validate_optimizer_membership( + model: nn.Module, + optimizer: torch.optim.Optimizer, +) -> None: + expected = {id(parameter) for parameter in model.parameters() if parameter.requires_grad} + actual_list = [parameter for group in optimizer.param_groups for parameter in group["params"]] + actual = {id(parameter) for parameter in actual_list} + if len(actual) != len(actual_list): + raise RuntimeError("The student optimizer contains duplicate parameter references.") + if actual != expected: + raise RuntimeError( + "Student optimizer membership does not exactly match the live post-FSDP " + f"trainable parameters (missing={len(expected - actual)}, extra={len(actual - expected)})." + ) + + +def _guard_automodel_hooks(diffusion_train: Any, auto_pipeline: Any) -> None: + builder_parameters = inspect.signature(diffusion_train.build_model_and_optimizer).parameters + required_builder_parameters = { + "model_id", + "learning_rate", + "device", + "dtype", + "optimizer_cfg", + } + missing = required_builder_parameters - set(builder_parameters) + if missing: + raise RuntimeError( + "Unsupported nemo_automodel diffusion builder; missing parameters: " + + ", ".join(sorted(missing)) + ) + if not hasattr(auto_pipeline, "_apply_parallelization"): + raise RuntimeError( + "Unsupported nemo_automodel: auto_diffusion_pipeline._apply_parallelization is missing." + ) + + +@contextlib.contextmanager +def patch_student_build( + settings: StudentSettings, +) -> Iterator[StudentBuildState]: + """Patch the two example-local seams needed during the parent ``setup`` call. + + Both module globals are restored in ``finally``. The patch is active only while + the one student is being constructed; teacher construction happens afterwards. + """ + from nemo_automodel._diffusers import auto_diffusion_pipeline as auto_pipeline + from nemo_automodel.recipes.diffusion import train as diffusion_train + + _guard_automodel_hooks(diffusion_train, auto_pipeline) + original_apply_parallelization = auto_pipeline._apply_parallelization + original_build_model_and_optimizer = diffusion_train.build_model_and_optimizer + state = StudentBuildState() + apply_calls = 0 + + def apply_parallelization(pipe, parallel_scheme): + nonlocal apply_calls + apply_calls += 1 + if apply_calls != 1: + raise RuntimeError( + "QAD's guarded student build expected exactly one parallelized component load." + ) + state.parallel_scheme = parallel_scheme + transformer = pipe.transformer + if settings.mode == "nvfp4": + quant_state_path = settings.quant_state_path + assert quant_state_path is not None + state.quantizer_count = _restore_regular_quant_state( + transformer, + quant_state_path, + next(transformer.parameters()).device, + ) + else: + state.svdquant_parameter_names = _validate_svdquant_bundle(transformer) + return original_apply_parallelization(pipe, parallel_scheme) + + def build_model_and_optimizer(**kwargs): + pipe, optimizer, device_mesh = original_build_model_and_optimizer(**kwargs) + trainable = _apply_train_scope(pipe.transformer, settings.train_scope) + if settings.train_scope == "lora_only": + optimizer = _rebuild_optimizer_from_live_parameters(optimizer, trainable) + logging.info( + "[QAD] rebuilt optimizer after FSDP for lora_only: %d live A/B tensors", + len(trainable), + ) + _validate_optimizer_membership(pipe.transformer, optimizer) + return pipe, optimizer, device_mesh + + auto_pipeline._apply_parallelization = apply_parallelization + diffusion_train.build_model_and_optimizer = build_model_and_optimizer + try: + yield state + finally: + diffusion_train.build_model_and_optimizer = original_build_model_and_optimizer + auto_pipeline._apply_parallelization = original_apply_parallelization + + if apply_calls != 1 or state.parallel_scheme is None: + raise RuntimeError( + "QAD did not observe the expected pre-FSDP student parallelization point." + ) diff --git a/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml b/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml new file mode 100644 index 00000000000..70bc99f8d09 --- /dev/null +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml @@ -0,0 +1,108 @@ +# Qwen-Image QAD with a regular ModelOpt NVFP4 student. +# +# The student is loaded from its Diffusers checkpoint, then the weight-free +# quantizer state is restored. Paths are placeholders and should be supplied by +# dotted CLI overrides in production launches. + +seed: 42 + +wandb: + project: fastgen-qad-qwen-image + mode: online + name: qwen_image_qad_nvfp4 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + # Canonical student source; AutoModel records this field in checkpoints. + pretrained_model_name_or_path: Qwen/Qwen-Image + mode: finetune + +step_scheduler: + global_batch_size: 64 + local_batch_size: 1 + ckpt_every_steps: 1000 + num_epochs: 200 + log_every: 1 + max_steps: 50000 + +qad: + # The teacher is always an unquantized, frozen Diffusers model. + teacher_model_name_or_path: Qwen/Qwen-Image + + # weight=1.0 is pure teacher-output MSE; the flow-matching task weight is 0. + output_loss: + type: mse + weight: 1.0 + + # Independent additive coefficient for the ordinary flow-matching target. + task_loss: + weight: 0.0 + + # Layer pairs use exact module names relative to each transformer. + layerwise: + enabled: false + pairs: [] + + student: + mode: nvfp4 + quant_state_path: /path/to/qwen-image/quant/transformer.nvfp4.pt + train_scope: all + +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: 1 + activation_checkpointing: true + +flow_matching: + adapter_type: qwen_image + timestep_sampling: logit_normal + logit_mean: 0.0 + logit_std: 1.0 + flow_shift: 3.0 + mix_uniform_ratio: 0.1 + use_sigma_noise: true + sigma_min: 0.0 + sigma_max: 1.0 + num_train_timesteps: 1000 + cfg_dropout_prob: 0.0 + use_loss_weighting: true + loss_weighting_scheme: linear + adapter_kwargs: + guidance_scale: 3.5 + use_guidance_embeds: false + +data: + dataloader: + _target_: fastgen_data.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 + +checkpoint: + enabled: true + checkpoint_dir: /path/to/output/qwen_image_qad_nvfp4/checkpoints + model_save_format: safetensors + # Save full sharded student DCP. Deployment export is a separate operation. + save_consolidated: false + diffusers_compatible: false + restore_from: LATEST diff --git a/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml b/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml new file mode 100644 index 00000000000..99873aeae43 --- /dev/null +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml @@ -0,0 +1,112 @@ +# Qwen-Image QAD with a ModelOpt NVFP4 SVDQuant + Hugging Face PEFT student. +# +# model.pretrained_model_name_or_path must be a complete, user-prepared ModelOpt-enabled +# Diffusers training bundle. A weight-free quantizer-state file or a unified +# deployment export is not sufficient. + +seed: 42 + +wandb: + project: fastgen-qad-qwen-image + mode: online + name: qwen_image_qad_svdquant_nvfp4 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + # Complete user-prepared ModelOpt-enabled Diffusers training bundle. + pretrained_model_name_or_path: /path/to/qwen-image-nvfp4-svdquant-training-bundle + mode: finetune + +step_scheduler: + global_batch_size: 64 + local_batch_size: 1 + ckpt_every_steps: 1000 + num_epochs: 200 + log_every: 1 + max_steps: 50000 + +qad: + # Keep the teacher independent from the quantized student bundle. + teacher_model_name_or_path: Qwen/Qwen-Image + + # weight=1.0 is pure teacher-output MSE; the flow-matching task weight is 0. + output_loss: + type: mse + weight: 1.0 + + # Independent additive coefficient for the ordinary flow-matching target. + task_loss: + weight: 0.0 + + # Example when enabled: + # pairs: + # - student_layer: transformer_blocks.29 + # teacher_layer: transformer_blocks.29 + # weight: 0.05 + layerwise: + enabled: false + pairs: [] + + student: + mode: nvfp4_svdquant + # all is canonical; lora_only trains only the SVDQuant HF PEFT A/B factors. + train_scope: all + +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: 1 + activation_checkpointing: true + +flow_matching: + adapter_type: qwen_image + timestep_sampling: logit_normal + logit_mean: 0.0 + logit_std: 1.0 + flow_shift: 3.0 + mix_uniform_ratio: 0.1 + use_sigma_noise: true + sigma_min: 0.0 + sigma_max: 1.0 + num_train_timesteps: 1000 + cfg_dropout_prob: 0.0 + use_loss_weighting: true + loss_weighting_scheme: linear + adapter_kwargs: + guidance_scale: 3.5 + use_guidance_embeds: false + +data: + dataloader: + _target_: fastgen_data.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 + +checkpoint: + enabled: true + checkpoint_dir: /path/to/output/qwen_image_qad_svdquant_nvfp4/checkpoints + model_save_format: safetensors + # Save full sharded student DCP. Deployment export is a separate operation. + save_consolidated: false + diffusers_compatible: false + restore_from: LATEST diff --git a/examples/diffusers/fastgen/qad/finetune.py b/examples/diffusers/fastgen/qad/finetune.py new file mode 100644 index 00000000000..495c2bef484 --- /dev/null +++ b/examples/diffusers/fastgen/qad/finetune.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Entrypoint for FastGen quantization-aware distillation.""" + +from __future__ import annotations + +import os +import sys + +_QAD_DIR = os.path.dirname(os.path.abspath(__file__)) +_FASTGEN_DIR = os.path.dirname(_QAD_DIR) +if _FASTGEN_DIR not in sys.path: + sys.path.insert(0, _FASTGEN_DIR) + +from nemo_automodel.components.config._arg_parser import parse_args_and_load_config # noqa: E402 + +from qad.recipe import QADDiffusionRecipe # noqa: E402 + + +def main( + default_config_path: str = ("examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml"), +) -> None: + cfg = parse_args_and_load_config(default_config_path) + recipe = QADDiffusionRecipe(cfg) + recipe.setup() + recipe.run_train_validation_loop() + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/qad/modeling.py b/examples/diffusers/fastgen/qad/modeling.py new file mode 100644 index 00000000000..759302795bb --- /dev/null +++ b/examples/diffusers/fastgen/qad/modeling.py @@ -0,0 +1,259 @@ +# 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. + +"""Model and ModelOpt-distillation helpers for QAD.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn.functional as F +from torch import nn + +import modelopt.torch.distill as mtd + +if TYPE_CHECKING: + from collections.abc import Sequence + +try: + from nemo_automodel.components.distributed.parallelizer import ( + PARALLELIZATION_STRATEGIES, + DefaultParallelizationStrategy, + register_parallel_strategy, + ) +except ImportError as exc: + raise ImportError( + "The FastGen QAD example requires nemo_automodel. Install " + "examples/diffusers/fastgen/requirements.txt." + ) from exc + + +class _QwenImageParallelizationStrategy(DefaultParallelizationStrategy): + """Checkpoint complete Qwen transformer blocks before AutoModel applies FSDP.""" + + def parallelize( + self, + model, + device_mesh, + activation_checkpointing: bool = False, + **kwargs, + ): + if activation_checkpointing: + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + checkpoint_wrapper, + ) + + blocks = getattr(model, "transformer_blocks", None) + if blocks is None: + raise AttributeError( + "QwenImageTransformer2DModel does not expose transformer_blocks." + ) + for index, block in enumerate(blocks): + blocks[index] = checkpoint_wrapper( + block, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + ) + logging.info( + "[QAD] Qwen-Image activation checkpointing enabled for %d full blocks", + len(blocks), + ) + + return super().parallelize( + model, + device_mesh, + activation_checkpointing=False, + **kwargs, + ) + + +def register_qwen_image_parallelization_strategy() -> None: + """Register the Qwen strategy unless AutoModel already ships a native strategy.""" + model_class_name = "QwenImageTransformer2DModel" + if model_class_name not in PARALLELIZATION_STRATEGIES: + register_parallel_strategy(name=model_class_name)(_QwenImageParallelizationStrategy) + + +register_qwen_image_parallelization_strategy() + + +def _extract_tensor(output: Any, selector: str) -> torch.Tensor: + """Select a tensor from a Diffusers root or Qwen dual-stream block output.""" + normalized = selector.lower() + if torch.is_tensor(output): + return output + if hasattr(output, "sample") and normalized in {"sample", "output", "tensor"}: + return output.sample + if isinstance(output, dict): + if selector not in output: + raise KeyError(f"Selector {selector!r} is not present in layer output keys.") + selected = output[selector] + if not torch.is_tensor(selected): + raise TypeError(f"Layer output {selector!r} is not a tensor.") + return selected + if isinstance(output, tuple | list): + index_by_name = { + "sample": 0, + "output": 0, + "tensor": 0, + "first": 0, + "encoder_hidden_states": 0, + "text": 0, + "hidden_states": 1, + "image": 1, + "last": -1, + } + if normalized not in index_by_name: + raise ValueError( + f"Unsupported tuple selector {selector!r}; use hidden_states/image, " + "encoder_hidden_states/text, first, last, or sample." + ) + selected = output[index_by_name[normalized]] + if not torch.is_tensor(selected): + raise TypeError(f"Selected {selector!r} output is not a tensor.") + return selected + raise TypeError(f"Cannot select {selector!r} from output type {type(output).__name__}.") + + +class TensorOutputDelegate(nn.Module): + """Forward to a model while exposing only its final tensor to ModelOpt KD. + + The wrapped model is deliberately stored outside ``nn.Module._modules``. This + keeps the controller parameter-free and prevents its state_dict from duplicating + either the FSDP student or teacher. ``get_submodule`` still routes layerwise + criterion paths to the live wrapped transformer. + """ + + def __init__(self, target: nn.Module): + super().__init__() + self.__dict__["_qad_target"] = target + + @property + def target(self) -> nn.Module: + return self.__dict__["_qad_target"] + + def forward(self, *args, **kwargs) -> torch.Tensor: + return _extract_tensor(self.target(*args, **kwargs), "sample") + + def get_submodule(self, target: str) -> nn.Module: + if target == "": + return self + return self.target.get_submodule(target) + + +class SelectedMSELoss(nn.modules.loss._Loss): + """FP32 MSE after selecting a stream from a captured layer output.""" + + def __init__(self, selector: str = "sample"): + super().__init__(reduction="mean") + self.selector = selector + + def forward(self, student_output: Any, teacher_output: Any) -> torch.Tensor: + student = _extract_tensor(student_output, self.selector) + teacher = _extract_tensor(teacher_output, self.selector) + return F.mse_loss(student.float(), teacher.float(), reduction="mean") + + +class AdditiveLossBalancer(mtd.DistillationLossBalancer): + """Apply independent, additive weights to task and KD loss terms.""" + + def __init__(self, *, task_weight: float, kd_weights: Sequence[float]): + super().__init__() + self.task_weight = float(task_weight) + self.kd_weights = tuple(float(weight) for weight in kd_weights) + + def forward(self, losses: dict[str, torch.Tensor]) -> torch.Tensor: + losses = dict(losses) + student_loss = losses.pop("student_loss", None) + if not losses: + raise RuntimeError("QAD received no KD loss terms.") + total = None + if self.task_weight != 0.0: + if student_loss is None: + raise RuntimeError("A nonzero QAD task weight requires student_loss.") + total = student_loss * self.task_weight + + if len(losses) != len(self.kd_weights): + raise RuntimeError( + "ModelOpt returned an unexpected number of KD losses: " + f"expected {len(self.kd_weights)}, got {len(losses)}." + ) + for loss, weight in zip(losses.values(), self.kd_weights): + # Multiplying a disabled NaN/Inf term by zero would still poison the + # objective. Skip disabled terms completely while retaining their + # detached diagnostics in the pipeline. + if weight == 0.0: + continue + weighted_loss = loss * weight + total = weighted_loss if total is None else total + weighted_loss + if total is None: + raise RuntimeError("QAD has no nonzero loss coefficient.") + return total + + +def build_distillation_controller( + *, + student: nn.Module, + teacher: nn.Module, + output_weight: float, + task_weight: float, + layer_pairs: Sequence[dict[str, Any]], +) -> tuple[nn.Module, tuple[str, ...]]: + """Create a parameter-free ModelOpt KD controller around live FSDP models.""" + criterion: dict[tuple[str, str], nn.modules.loss._Loss] = {("", ""): SelectedMSELoss("sample")} + names = ["output_mse"] + weights = [float(output_weight)] + seen_pairs = {("", "")} + + for index, pair in enumerate(layer_pairs): + student_layer = str(pair["student_layer"]) + teacher_layer = str(pair.get("teacher_layer", student_layer)) + selector = str(pair.get("selector", "hidden_states")) + weight = float(pair.get("weight", 1.0)) + key = (student_layer, teacher_layer) + if key in seen_pairs: + raise ValueError(f"Duplicate QAD layer pair: {key!r}") + seen_pairs.add(key) + criterion[key] = SelectedMSELoss(selector) + names.append(f"layer_{index}_{student_layer}_{selector}_mse") + weights.append(weight) + + controller = mtd.convert( + TensorOutputDelegate(student), + mode=[ + ( + "kd_loss", + { + "teacher_model": TensorOutputDelegate(teacher), + "criterion": criterion, + "loss_balancer": AdditiveLossBalancer( + task_weight=task_weight, + kd_weights=weights, + ), + "expose_minimal_state_dict": True, + }, + ) + ], + ) + return controller, tuple(names) + + +def clear_captured_outputs(controller: nn.Module) -> None: + """Release activation references before forwards and after checkpoint recompute.""" + for student_layer, teacher_layer in controller._layers_to_loss: + student_layer._intermediate_output = None + teacher_layer._intermediate_output = None diff --git a/examples/diffusers/fastgen/qad/pipeline.py b/examples/diffusers/fastgen/qad/pipeline.py new file mode 100644 index 00000000000..734a904fd31 --- /dev/null +++ b/examples/diffusers/fastgen/qad/pipeline.py @@ -0,0 +1,74 @@ +# 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. + +"""QAD loss pipeline layered on AutoModel's flow-matching input preparation.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from .modeling import clear_captured_outputs + + +class QADPipeline: + """Run teacher/student on identical inputs and aggregate ModelOpt KD losses.""" + + def __init__(self, flow_matching_pipeline, controller: nn.Module, loss_names: tuple[str, ...]): + self.flow_matching_pipeline = flow_matching_pipeline + self.controller = controller + self.loss_names = loss_names + + def step( + self, + *, + batch: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + global_step: int, + check_loss: bool, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + clear_captured_outputs(self.controller) + _, task_loss, _, _ = self.flow_matching_pipeline.step( + model=self.controller, + batch=batch, + device=device, + dtype=dtype, + global_step=global_step, + collect_metrics=False, + # The flow target is optional in QAD. Validate the actual combined loss below. + check_loss=False, + ) + losses = self.controller.compute_kd_loss( + student_loss=task_loss, + skip_balancer=True, + ) + total = self.controller.loss_balancer(losses) + if check_loss and not bool(torch.isfinite(total.detach()).all()): + raise FloatingPointError(f"Non-finite QAD loss at step {global_step}.") + + kd_values = [value for key, value in losses.items() if key != "student_loss"] + if len(kd_values) != len(self.loss_names): + raise RuntimeError( + "QAD loss-name mapping is out of sync with ModelOpt's returned losses." + ) + metrics = {"task_loss": task_loss.detach(), "total_loss": total.detach()} + metrics.update({name: value.detach() for name, value in zip(self.loss_names, kd_values)}) + return total, metrics + + def clear(self) -> None: + clear_captured_outputs(self.controller) diff --git a/examples/diffusers/fastgen/qad/recipe.py b/examples/diffusers/fastgen/qad/recipe.py new file mode 100644 index 00000000000..276abac32ac --- /dev/null +++ b/examples/diffusers/fastgen/qad/recipe.py @@ -0,0 +1,595 @@ +# 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. + +"""FastGen quantization-aware distillation recipe. + +QAD is deliberately separate from DMD2: one frozen Diffusers teacher and one +quantized student see the same noisy latent, timestep, and conditioning, and +ModelOpt's standard ``kd_loss`` API supplies output and optional representation +MSE losses. +""" + +from __future__ import annotations + +import logging +import math +import os +from typing import Any + +import torch +import wandb +import yaml +from torch import nn +from torchdata.stateful_dataloader import StatefulDataLoader + +import modelopt.torch.distill as mtd +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.nn import TensorQuantizer + +try: + from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.training.utils import ( + clip_grad_norm, + prepare_after_first_microbatch, + prepare_for_final_backward, + prepare_for_grad_accumulation, + ) + from nemo_automodel.recipes.base_recipe import ( + _find_latest_checkpoint, + _resolve_restore_from_to_ckpt_dir, + ) + from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe, is_main_process +except ImportError as exc: + raise ImportError( + "The FastGen QAD example requires nemo_automodel. Install dependencies with:\n" + " pip install -r examples/diffusers/fastgen/requirements.txt" + ) from exc + +from fastgen_checkpoint import make_optimizer_partial_load_tolerant + +from .artifacts import StudentSettings, patch_student_build +from .modeling import build_distillation_controller, clear_captured_outputs +from .pipeline import QADPipeline + + +def _as_dict(value: Any) -> dict[str, Any]: + if value is None: + return {} + if hasattr(value, "to_dict"): + return value.to_dict() + return dict(value) + + +class QADDiffusionRecipe(TrainDiffusionRecipe): + """AutoModel diffusion recipe with a ModelOpt KD controller.""" + + def __init__(self, cfg) -> None: + # AutoModel's dotted CLI setter updates live ConfigNodes but not the + # raw_config later written beside checkpoints. Materialize the resolved + # runtime values so QAD paths, scope, teacher, and losses are reproducible. + if hasattr(cfg, "to_yaml_dict"): + cfg.__dict__["_raw_config"] = cfg.to_yaml_dict( + resolve_env=False, + redact_sensitive=True, + use_orig_values=False, + ) + super().__init__(cfg) + + def setup(self) -> None: + settings, loss_config = self._resolve_qad_config() + self.__dict__["_qad_resume_signature"] = self._resume_signature( + settings, + loss_config, + ) + + if self.cfg.get("peft", None) is not None: + raise ValueError( + "Do not set AutoModel's top-level peft block for QAD. SVDQuant's " + "modelopt_svdquant HF PEFT topology must come from the student bundle." + ) + if str(self.cfg.get("model.mode", "finetune")).lower() != "finetune": + raise ValueError("QAD supports model.mode=finetune only.") + if self.cfg.get("ddp", None) is not None: + raise ValueError("QAD currently supports AutoModel FSDP2, not DDP.") + + # Diffusers' ModelMixin must be patched before from_pretrained so a + # SVDQuant bundle rebuilds quant/PEFT topology before loading its weights. + mto.enable_huggingface_checkpointing() + + with patch_student_build(settings) as build_state: + super().setup() + + # Diffusers loads ModelMixin objects in eval mode. QAD owns the student + # train/eval boundary because the controller delegate intentionally does + # not register the live FSDP module as a child. + self.model.train() + + # Parent checkpoint restore established the exact next-step RNG state. + # Teacher construction/sharding is transient setup and must not perturb + # the first fresh or resumed training sample. + training_rng_state = self.rng.state_dict() + try: + parallel_scheme = build_state.parallel_scheme + if parallel_scheme is None: + raise RuntimeError("QAD failed to capture the student's parallel scheme.") + teacher = self._load_frozen_teacher( + loss_config["teacher_model_name_or_path"], + parallel_scheme, + ) + controller, loss_names = build_distillation_controller( + student=self.model, + teacher=teacher, + output_weight=loss_config["output_weight"], + task_weight=loss_config["task_weight"], + layer_pairs=loss_config["layer_pairs"], + ) + if any(True for _ in controller.parameters()): + raise RuntimeError( + "The QAD controller must remain parameter-free; optimizer/checkpoint " + "ownership belongs exclusively to self.model." + ) + + # BaseRecipe tracks nn.Module assignments. Bypass it for the frozen teacher + # and transient controller so checkpoint selection cannot mistake either for + # the student. + object.__setattr__(self, "_qad_teacher", teacher) + object.__setattr__(self, "_qad_controller", controller) + object.__setattr__( + self, + "_qad_pipeline", + QADPipeline(self.flow_matching_pipeline, controller, loss_names), + ) + object.__setattr__(self, "_qad_student_settings", settings) + object.__setattr__(self, "_qad_loss_config", loss_config) + + tracked = self.__dict__.get("__state_tracked", set()) + forbidden = {"_qad_teacher", "_qad_controller", "_qad_pipeline"} & set(tracked) + if forbidden: + raise RuntimeError( + f"QAD transient objects were accidentally state-tracked: {forbidden}" + ) + self._validate_state_ownership() + finally: + self.rng.load_state_dict(training_rng_state) + + if is_main_process(): + logging.info( + "[QAD] initialized: teacher=%s student=%s mode=%s train_scope=%s " + "task_weight=%g output_weight=%g layer_pairs=%d", + loss_config["teacher_model_name_or_path"], + settings.model_name_or_path, + settings.mode, + settings.train_scope, + loss_config["task_weight"], + loss_config["output_weight"], + len(loss_config["layer_pairs"]), + ) + logging.info("[QAD] student quantizer summary:") + mtq.print_quant_summary(self.model) + + def _resolve_qad_config(self) -> tuple[StudentSettings, dict[str, Any]]: + qad = _as_dict(self.cfg.get("qad", None)) + if not qad: + raise ValueError("Missing required qad configuration block.") + + student_cfg = _as_dict(qad.get("student")) + model_name_or_path = self.cfg.get("model.pretrained_model_name_or_path", None) + if not model_name_or_path: + raise ValueError( + "model.pretrained_model_name_or_path is required and is the canonical " + "student source recorded in checkpoints." + ) + duplicate_student_path = student_cfg.get("model_name_or_path") + if duplicate_student_path is not None and str(duplicate_student_path) != str( + model_name_or_path + ): + raise ValueError( + "qad.student.model_name_or_path conflicts with the canonical " + "model.pretrained_model_name_or_path. Remove the duplicate QAD field." + ) + mode = str(student_cfg.get("mode", "nvfp4")).lower() + # Accept the early design spelling while emitting one canonical name. + if mode == "svdquant_nvfp4": + mode = "nvfp4_svdquant" + settings = StudentSettings( + mode=mode, + model_name_or_path=str(model_name_or_path), + quant_state_path=student_cfg.get("quant_state_path"), + train_scope=str(student_cfg.get("train_scope", "all")).lower(), + ) + settings.validate() + + teacher_model_name_or_path = qad.get("teacher_model_name_or_path") + if not teacher_model_name_or_path: + raise ValueError("qad.teacher_model_name_or_path is required.") + + output_cfg = _as_dict(qad.get("output_loss")) + if str(output_cfg.get("type", "mse")).lower() != "mse": + raise ValueError("QAD currently supports only output_loss.type=mse.") + output_weight = float(output_cfg.get("weight", 1.0)) + + task_cfg = _as_dict(qad.get("task_loss")) + task_weight = float(task_cfg.get("weight", 0.0)) + + layerwise_cfg = _as_dict(qad.get("layerwise")) + layer_pairs = layerwise_cfg.get("pairs", []) if layerwise_cfg.get("enabled", False) else [] + layer_pairs = [_as_dict(pair) for pair in layer_pairs] + + all_weights = [output_weight, task_weight] + [ + float(pair.get("weight", 1.0)) for pair in layer_pairs + ] + if any(not math.isfinite(weight) or weight < 0.0 for weight in all_weights): + raise ValueError("QAD loss weights must be finite and non-negative.") + if not any(weight > 0.0 for weight in all_weights): + raise ValueError("At least one QAD loss weight must be positive.") + for index, pair in enumerate(layer_pairs): + if not pair.get("student_layer"): + raise ValueError(f"qad.layerwise.pairs[{index}].student_layer is required.") + + return settings, { + "teacher_model_name_or_path": str(teacher_model_name_or_path), + "output_weight": output_weight, + "task_weight": task_weight, + "layer_pairs": layer_pairs, + } + + @staticmethod + def _resume_signature( + settings: StudentSettings, + loss_config: dict[str, Any], + ) -> dict[str, Any]: + return { + "student_source": settings.model_name_or_path, + "student_mode": settings.mode, + "quant_state_path": settings.quant_state_path, + "train_scope": settings.train_scope, + "teacher_source": loss_config["teacher_model_name_or_path"], + "output_weight": float(loss_config["output_weight"]), + "task_weight": float(loss_config["task_weight"]), + "layer_pairs": tuple( + ( + str(pair["student_layer"]), + str(pair.get("teacher_layer", pair["student_layer"])), + str(pair.get("selector", "hidden_states")), + float(pair.get("weight", 1.0)), + ) + for pair in loss_config["layer_pairs"] + ), + } + + @classmethod + def _resume_signature_from_saved_config(cls, config: dict[str, Any]) -> dict[str, Any]: + model_cfg = _as_dict(config.get("model")) + qad_cfg = _as_dict(config.get("qad")) + student_cfg = _as_dict(qad_cfg.get("student")) + output_cfg = _as_dict(qad_cfg.get("output_loss")) + task_cfg = _as_dict(qad_cfg.get("task_loss")) + layerwise_cfg = _as_dict(qad_cfg.get("layerwise")) + + mode = str(student_cfg.get("mode", "nvfp4")).lower() + if mode == "svdquant_nvfp4": + mode = "nvfp4_svdquant" + raw_pairs = layerwise_cfg.get("pairs", []) if layerwise_cfg.get("enabled", False) else [] + loss_config = { + "teacher_model_name_or_path": str(qad_cfg.get("teacher_model_name_or_path", "")), + "output_weight": float(output_cfg.get("weight", 1.0)), + "task_weight": float(task_cfg.get("weight", 0.0)), + "layer_pairs": [_as_dict(pair) for pair in raw_pairs], + } + settings = StudentSettings( + mode=mode, + model_name_or_path=str(model_cfg.get("pretrained_model_name_or_path", "")), + quant_state_path=student_cfg.get("quant_state_path"), + train_scope=str(student_cfg.get("train_scope", "all")).lower(), + ) + return cls._resume_signature(settings, loss_config) + + def _resolved_checkpoint_dir(self, restore_from: str | None) -> str | None: + if not self.checkpointer.config.enabled: + return None + if restore_from: + resolved = _resolve_restore_from_to_ckpt_dir( + self.checkpointer.config.checkpoint_dir, + restore_from, + ) + else: + resolved = _find_latest_checkpoint(self.checkpointer.config.checkpoint_dir) + if resolved is None: + return None + return os.fspath(resolved) + + def _validate_qad_checkpoint_signature(self, checkpoint_dir: str) -> None: + config_path = os.path.join(checkpoint_dir, "config.yaml") + if not os.path.isfile(config_path): + raise RuntimeError( + "QAD cannot safely restore optimizer shards from a checkpoint without " + f"config.yaml: {checkpoint_dir}" + ) + with open(config_path) as config_file: + saved_config = yaml.safe_load(config_file) or {} + saved_signature = self._resume_signature_from_saved_config(saved_config) + current_signature = self.__dict__["_qad_resume_signature"] + if saved_signature != current_signature: + changed = [ + key + for key in current_signature + if saved_signature.get(key) != current_signature[key] + ] + raise RuntimeError( + "QAD resume contract changed for " + + ", ".join(changed) + + ". Use the same student artifact, quantization mode/state, train scope, " + "teacher, and loss configuration as the saved run." + ) + + def load_checkpoint(self, restore_from: str | None = None) -> None: + """Validate QAD topology before enabling FSDP2 partial-shard optimizer load.""" + checkpoint_dir = self._resolved_checkpoint_dir(restore_from) + if checkpoint_dir is not None and os.path.isdir(checkpoint_dir): + self._validate_qad_checkpoint_signature(checkpoint_dir) + make_optimizer_partial_load_tolerant(self.checkpointer) + super().load_checkpoint(restore_from) + + def _rebuild_dataloader_for_resume(self, global_step: int) -> None: + """Rebuild the loader and deterministically skip to the restored data position.""" + epoch_len = int(getattr(self.step_scheduler, "epoch_len", 0) or 0) + grad_acc = int(getattr(self.step_scheduler, "grad_acc_steps", 1) or 1) + if epoch_len <= 0 or self.sampler is None or global_step <= 0: + return + + current_epoch = global_step // epoch_len + skip_batches = (global_step % epoch_len) * grad_acc + old_dataloader = self.dataloader + dataloader_kwargs = { + "collate_fn": getattr(old_dataloader, "collate_fn", None), + "num_workers": int(getattr(old_dataloader, "num_workers", 0) or 0), + "pin_memory": bool(getattr(old_dataloader, "pin_memory", False)), + } + if dataloader_kwargs["num_workers"] > 0: + dataloader_kwargs["prefetch_factor"] = getattr( + old_dataloader, + "prefetch_factor", + 2, + ) + dataloader_kwargs["persistent_workers"] = bool( + getattr(old_dataloader, "persistent_workers", False) + ) + + # Keep the parent's existing tracked state key while replacing the + # StatefulDataLoader object whose restored cursor is known to stick. + self.__dict__["dataloader"] = StatefulDataLoader( + old_dataloader.dataset, + batch_sampler=self.sampler, + **dataloader_kwargs, + ) + self.step_scheduler.epoch = current_epoch + self.sampler.set_epoch(current_epoch) + self.sampler._batches_to_skip = skip_batches + if is_main_process(): + logging.info( + "[QAD][resume] rebuilt dataloader at epoch=%d skip_batches=%d " + "(global_step=%d epoch_len=%d grad_acc=%d)", + current_epoch, + skip_batches, + global_step, + epoch_len, + grad_acc, + ) + + def _load_frozen_teacher( + self, + model_name_or_path: str, + parallel_scheme: dict[str, dict[str, Any]], + ) -> nn.Module: + pipe, _ = NeMoAutoDiffusionPipeline.from_pretrained( + model_name_or_path, + 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 + if mto.ModeloptStateManager.is_converted(teacher): + raise RuntimeError( + "QAD teacher must be a plain BF16 Diffusers checkpoint without ModelOpt modes." + ) + if any(isinstance(module, TensorQuantizer) for module in teacher.modules()): + raise RuntimeError("QAD teacher must be an unquantized BF16 Diffusers checkpoint.") + teacher.eval() + teacher.requires_grad_(False) + return teacher + + def _validate_state_ownership(self) -> None: + optimizer_parameters = { + id(parameter) for group in self.optimizer.param_groups for parameter in group["params"] + } + student_parameters = { + id(parameter) for parameter in self.model.parameters() if parameter.requires_grad + } + teacher_parameters = {id(parameter) for parameter in self._qad_teacher.parameters()} + if optimizer_parameters != student_parameters: + raise RuntimeError("QAD optimizer does not exactly own the trainable student state.") + if optimizer_parameters & teacher_parameters: + raise RuntimeError("Frozen teacher parameters leaked into the student optimizer.") + if any(parameter.requires_grad for parameter in self._qad_teacher.parameters()): + raise RuntimeError("QAD teacher must be completely frozen.") + + def run_train_validation_loop(self) -> None: + """Run a conventional optimizer loop using the QAD objective.""" + self.model.train() + logging.info( + "[QAD] starting training: global_batch_size=%s local_batch_size=%s dp_size=%s", + self.global_batch_size, + self.local_batch_size, + self.dp_size, + ) + global_step = int(self.step_scheduler.step) + self._rebuild_dataloader_for_resume(global_step) + + try: + for epoch in self.step_scheduler.epochs: + if self.sampler is not None and hasattr(self.sampler, "set_epoch"): + self.sampler.set_epoch(epoch) + + 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} (global step {global_step})", + initial=tqdm_initial, + ) + else: + self.step_scheduler.dataloader = self.dataloader + + epoch_loss = 0.0 + num_steps = 0 + for batch_group in self.step_scheduler: + # StepScheduler increments only after control returns to its + # generator, so refresh at the top of every yielded group. + global_step = int(self.step_scheduler.step) + self.optimizer.zero_grad(set_to_none=True) + prepare_for_grad_accumulation([self.model], pp_enabled=False) + num_microbatches = len(batch_group) + micro_metrics: list[dict[str, torch.Tensor]] = [] + + for microbatch_index, micro_batch in enumerate(batch_group): + if microbatch_index == num_microbatches - 1: + prepare_for_final_backward([self.model], pp_enabled=False) + try: + total_loss, metrics = self._qad_pipeline.step( + batch=micro_batch, + device=self.device, + dtype=self.bf16, + global_step=global_step, + check_loss=self.check_loss, + ) + (total_loss / num_microbatches).backward() + micro_metrics.append(metrics) + finally: + # Full-block NO_REENTRANT checkpoint wrappers avoid hook + # repopulation during recompute; this final cleanup is also + # safe when activation checkpointing is disabled. + self._qad_pipeline.clear() + + if microbatch_index == 0: + prepare_after_first_microbatch() + + self._validate_first_step_gradients(global_step) + grad_norm = clip_grad_norm( + self.clip_grad_max_norm, + [self.model], + foreach=self.grad_clip_foreach, + ) + grad_norm = float(grad_norm) if torch.is_tensor(grad_norm) else grad_norm + self.optimizer.step() + if self.lr_scheduler is not None: + self.lr_scheduler[0].step(1) + + reduced_metrics = { + name: float( + torch.stack([metrics[name] for metrics in micro_metrics]).mean().item() + ) + for name in micro_metrics[0] + } + group_loss = reduced_metrics["total_loss"] + epoch_loss += group_loss + num_steps += 1 + + if self.log_every and global_step % self.log_every == 0 and is_main_process(): + log_dict = { + "train_loss": group_loss, + "train_avg_loss": epoch_loss / num_steps, + "lr": self.optimizer.param_groups[0]["lr"], + "grad_norm": grad_norm, + "epoch": epoch, + "global_step": global_step, + **{f"qad/{name}": value for name, value in reduced_metrics.items()}, + } + if wandb.run is not None: + wandb.log(log_dict, step=global_step) + component_text = " ".join( + f"{name}={value:.6f}" for name, value in reduced_metrics.items() + ) + logging.info( + "[QAD][TRAIN] step=%d epoch=%d %s lr=%.3e grad_norm=%.3f", + global_step, + epoch, + component_text, + self.optimizer.param_groups[0]["lr"], + grad_norm, + ) + if hasattr(self.step_scheduler.dataloader, "set_postfix"): + self.step_scheduler.dataloader.set_postfix( + { + "loss": f"{group_loss:.4f}", + "lr": f"{self.optimizer.param_groups[0]['lr']:.2e}", + "gn": f"{grad_norm:.2f}", + } + ) + + if self.step_scheduler.is_ckpt_step: + self.save_checkpoint(epoch, global_step, epoch_loss / num_steps) + + if num_steps == 0: + logging.info( + "[QAD] epoch %d skipped (already completed in previous run)", epoch + 1 + ) + continue + logging.info( + "[QAD] epoch %d complete: avg_loss=%.6f", + epoch + 1, + epoch_loss / num_steps, + ) + + if is_main_process() and wandb.run is not None: + wandb.finish() + logging.info("[QAD] training complete at step %d", global_step) + finally: + self._release_distillation_controller() + + def _validate_first_step_gradients(self, global_step: int) -> None: + if global_step != 0: + return + trainable = [parameter for parameter in self.model.parameters() if parameter.requires_grad] + if not any(parameter.grad is not None for parameter in trainable): + raise RuntimeError("QAD produced no gradients for any trainable student parameter.") + if self._qad_student_settings.train_scope == "lora_only": + missing = [ + name + for name, parameter in self.model.named_parameters() + if parameter.requires_grad and parameter.grad is None + ] + if missing: + raise RuntimeError( + "SVDQuant lora_only parameters missing gradients on the first step: " + + ", ".join(missing[:5]) + ) + + def _release_distillation_controller(self) -> None: + controller = getattr(self, "_qad_controller", None) + if controller is None or not hasattr(controller, "_layers_to_loss"): + return + layer_pairs = tuple(controller._layers_to_loss) + clear_captured_outputs(controller) + mtd.export(controller) + for student_layer, teacher_layer in layer_pairs: + for layer in (student_layer, teacher_layer): + if hasattr(layer, "_intermediate_output"): + delattr(layer, "_intermediate_output")