diff --git a/examples/diffusers/fastgen/qad/README.md b/examples/diffusers/fastgen/qad/README.md new file mode 100644 index 00000000000..d887d637444 --- /dev/null +++ b/examples/diffusers/fastgen/qad/README.md @@ -0,0 +1,225 @@ +# 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 bundle validation contracts. In +both cases, `model.pretrained_model_name_or_path` is the only student artifact +path: it points to a complete, calibrated Diffusers pipeline written by +`quantize.py --output-bundle`. QAD restores the pipeline's weights and +component-local ModelOpt state together before FSDP; it does not accept a second +quantizer-state or transformer-checkpoint path. + +### Regular NVFP4 + +Set `qad.student.mode=nvfp4` and point +`model.pretrained_model_name_or_path` at a regular NVFP4 training bundle. The +bundle includes the calibrated weights and ModelOpt quantizer topology/state. +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 standalone weight-free NVFP4 quantizer-state file is not a QAD student bundle. +This is especially important for SVDQuant: calibration subtracts the low-rank +branch from the original weight, so both the resulting residual weight and the +PEFT factors are required. Deployment artifacts are not training bundles 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. + +## Prepare a student bundle + +Patch Diffusers ModelMixin support and save the complete pipeline through the +quantization entry point. `quantize.py` does this automatically before model +load and calls `pipe.save_pretrained(output_bundle)` after calibration. For +example: + +```bash +# Regular NVFP4 +python examples/diffusers/quantization/quantize.py \ + --model qwen-image \ + --override-model-path /path/to/Qwen-Image \ + --model-dtype BFloat16 \ + --format fp4 \ + --quant-algo max \ + --block-size 16 \ + --batch-size 1 \ + --calib-size 32 \ + --n-steps 50 \ + --extra-param true_cfg_scale=4.0 \ + --extra-param "negative_prompt= " \ + --output-bundle /path/to/Qwen-Image-NVFP4-Calib32 + +# NVFP4 SVDQuant, rank 32 +python examples/diffusers/quantization/quantize.py \ + --model qwen-image \ + --override-model-path /path/to/Qwen-Image \ + --model-dtype BFloat16 \ + --format fp4 \ + --quant-algo svdquant \ + --lowrank 32 \ + --block-size 16 \ + --batch-size 1 \ + --calib-size 32 \ + --n-steps 50 \ + --extra-param true_cfg_scale=4.0 \ + --extra-param "negative_prompt= " \ + --output-bundle /path/to/Qwen-Image-NVFP4-SVDQuant-Calib32 +``` + +The saved root must contain `model_index.json`; the converted transformer must +contain `transformer/modelopt_state.pth`. For Qwen-Image-Flash, use `--n-steps 4` +and `--extra-param true_cfg_scale=1.0`; omit `negative_prompt`. Standard output +includes ModelOpt's full quantizer summary; capture it with `tee` and retain that +log with the bundle. + +## 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, quantization 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..92c977a6c50 --- /dev/null +++ b/examples/diffusers/fastgen/qad/artifacts.py @@ -0,0 +1,404 @@ +# 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: + +* validate the ModelOpt topology restored by a native Diffusers training bundle + before FSDP; +* after FSDP, optionally freeze everything except ModelOpt SVDQuant's HF PEFT A/B + parameters and rebuild AdamW from the live sharded parameters. + +Quantization itself is never calibrated here. The complete topology, weights, +and quantizer buffers must already be present in the student bundle. +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import inspect +import logging +import re +from typing import TYPE_CHECKING, Any + +import modelopt.torch.opt as mto +from modelopt.torch.quantization.nn import TensorQuantizer + +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" + + 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" and self.train_scope != "all": + raise ValueError("Regular NVFP4 supports only qad.student.train_scope=all.") + + +@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 _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 _modelopt_mode_states(model: nn.Module) -> dict[str, dict[str, Any]]: + if not mto.ModeloptStateManager.is_converted(model): + return {} + return dict(mto.modelopt_state(model)["modelopt_state_dict"]) + + +def _reject_non_training_modes(mode_states: dict[str, dict[str, Any]]) -> None: + if "real_quantize" in mode_states: + raise RuntimeError( + "QAD cannot train a compressed real-quantized bundle. Recalibrate without " + "quantize.py --compress and provide the resulting fake-quantized training bundle." + ) + + +def _validate_regular_bundle(model: nn.Module) -> int: + mode_states = _modelopt_mode_states(model) + if not mode_states: + raise RuntimeError( + "qad.student.mode=nvfp4 requires a ModelOpt-aware Diffusers training bundle. " + "Calibrate it with quantize.py --output-bundle before starting QAD." + ) + _reject_non_training_modes(mode_states) + if "svdquant_calibrate" in mode_states: + raise RuntimeError( + "qad.student.mode=nvfp4 received an SVDQuant bundle; use mode=nvfp4_svdquant." + ) + quantizers = [module for module in model.modules() if isinstance(module, TensorQuantizer)] + if not quantizers: + raise RuntimeError("The regular NVFP4 student bundle restored no TensorQuantizers.") + _validate_nvfp4_quantizers(model, artifact_name="The regular NVFP4 student bundle") + logging.info( + "[QAD] validated regular ModelOpt NVFP4 bundle before FSDP: %d quantizers", + len(quantizers), + ) + return len(quantizers) + + +def _validate_svdquant_bundle(model: nn.Module) -> tuple[str, ...]: + mode_states = _modelopt_mode_states(model) + _reject_non_training_modes(mode_states) + mode_state = mode_states.get("svdquant_calibrate") + if mode_state is None: + raise RuntimeError( + "qad.student.mode=nvfp4_svdquant requires a bundle containing the " + "svdquant_calibrate ModelOpt mode." + ) + metadata = mode_state.get("metadata", {}).get("svdquant_peft") + if not metadata: + raise RuntimeError( + "The SVDQuant bundle is malformed or predates the HF PEFT contract: its " + "svdquant_calibrate mode has no svdquant_peft metadata. Recalibrate it " + "with quantize.py --output-bundle." + ) + + 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": + state.quantizer_count = _validate_regular_bundle(transformer) + else: + state.svdquant_parameter_names = _validate_svdquant_bundle(transformer) + state.quantizer_count = sum( + isinstance(module, TensorQuantizer) for module in transformer.modules() + ) + 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..2cc67afd74c --- /dev/null +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml @@ -0,0 +1,106 @@ +# Qwen-Image QAD with a regular ModelOpt NVFP4 student. +# +# model.pretrained_model_name_or_path must be a complete ModelOpt-aware +# Diffusers training bundle produced by quantize.py --output-bundle. + +seed: 42 + +wandb: + project: fastgen-qad-qwen-image + mode: online + name: qwen_image_qad_nvfp4 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + # Complete calibrated NVFP4 bundle; AutoModel records this field in checkpoints. + pretrained_model_name_or_path: /path/to/qwen-image-nvfp4-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: + # 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 + 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..12ecc1a3535 --- /dev/null +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml @@ -0,0 +1,111 @@ +# Qwen-Image QAD with a ModelOpt NVFP4 SVDQuant + Hugging Face PEFT student. +# +# model.pretrained_model_name_or_path must be a complete ModelOpt-aware +# Diffusers training bundle produced by quantize.py --output-bundle. + +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 calibrated NVFP4 SVDQuant 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..05136cd3d7a --- /dev/null +++ b/examples/diffusers/fastgen/qad/recipe.py @@ -0,0 +1,616 @@ +# 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 comes 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 both + # regular NVFP4 and SVDQuant bundles rebuild their ModelOpt topology and + # load component-local modelopt_state.pth before AutoModel applies FSDP. + 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")) + secondary_artifact_fields = sorted( + field + for field in ("quant_state_path", "modelopt_state_path") + if student_cfg.get(field) is not None + ) + if secondary_artifact_fields: + raise ValueError( + "QAD accepts one complete Diffusers student bundle through " + "model.pretrained_model_name_or_path; remove unsupported secondary " + "artifact field(s): " + + ", ".join(f"qad.student.{field}" for field in secondary_artifact_fields) + ) + 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), + 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, + "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")) + secondary_artifact_fields = sorted( + field + for field in ("quant_state_path", "modelopt_state_path") + if student_cfg.get(field) is not None + ) + if secondary_artifact_fields: + raise RuntimeError( + "The saved QAD checkpoint uses unsupported secondary student artifact " + "field(s): " + + ", ".join(f"qad.student.{field}" for field in secondary_artifact_fields) + ) + 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", "")), + 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 bundle, quantization mode, 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") diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index bebc61970a3..4c9b8fd1f9a 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -82,6 +82,15 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: """ self.logger.info(f"Starting calibration with {self.config.num_batches} batches") extra_args = MODEL_DEFAULTS.get(self.model_type, {}).get("inference_extra_args", {}) + if self.model_type == ModelType.QWEN_IMAGE: + extra_params = self.pipeline_manager.config.extra_params + self.logger.info( + "Qwen-Image calibration path: steps=%d true_cfg_scale=%s " + "negative_prompt=%s output_type=latent", + self.config.n_steps, + extra_params.get("true_cfg_scale", "pipeline default"), + "provided" if "negative_prompt" in extra_params else "omitted", + ) with tqdm(total=self.config.num_batches, desc="Calibration", unit="batch") as pbar: for i, prompt_batch in enumerate(batched_prompts): @@ -99,6 +108,8 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: elif self.model_type == ModelType.QWEN_IMAGE_DMD2: # DMD2 students use a custom few-step sampler, not the standard loop. self._run_qwen_image_dmd2_calibration(prompt_batch) + elif self.model_type == ModelType.QWEN_IMAGE: + self._run_qwen_image_calibration(prompt_batch, extra_args) else: common_args = { "prompt": prompt_batch, @@ -109,6 +120,31 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: self.logger.debug(f"Completed calibration batch {i + 1}/{self.config.num_batches}") self.logger.info("Calibration completed successfully") + def _run_qwen_image_calibration( + self, prompt_batch: list[str], extra_args: dict[str, Any] + ) -> None: + """Run Qwen-Image's standard denoising loop without the unused VAE decode. + + ``true_cfg_scale`` alone does not enable true CFG in QwenImagePipeline; a + negative prompt is also required. Keep both values explicit launch-time + parameters so the 50-step base model and four-step Flash model can share + this model type while calibrating against their actual inference paths. + """ + extra_params = self.pipeline_manager.config.extra_params + kwargs = { + "height": extra_params.get("height", extra_args.get("height", 1024)), + "width": extra_params.get("width", extra_args.get("width", 1024)), + "num_inference_steps": self.config.n_steps, + "output_type": "latent", + } + for name in ("negative_prompt", "true_cfg_scale", "guidance_scale"): + if name in extra_params: + value = extra_params[name] + if name == "negative_prompt" and isinstance(value, str): + value = [value] * len(prompt_batch) + kwargs[name] = value + self.pipe(prompt=prompt_batch, **kwargs).images + def _run_qwen_image_dmd2_calibration(self, prompt_batch: list[str]) -> None: """Calibrate a DMD2 Qwen-Image student via its few-step sampler. diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 26aac8ae07f..5fbaa7c9ccb 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -52,6 +52,7 @@ ) from utils import check_conv_and_mha, check_lora, restore_quantizer_state, save_quantizer_state +import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint @@ -316,6 +317,29 @@ def save_checkpoint( self.logger.info("Checkpoint saved successfully") + def save_training_bundle(self, pipe: DiffusionPipeline) -> None: + """Save a calibrated pipeline for QAD through native Diffusers checkpointing.""" + if not self.config.output_bundle: + return + + output_bundle = self.config.output_bundle + self.logger.info("Saving ModelOpt training bundle to %s", output_bundle) + pipe.save_pretrained(output_bundle) + + model_index_path = output_bundle / "model_index.json" + if not model_index_path.is_file(): + raise RuntimeError(f"Training bundle is missing {model_index_path}.") + if self.pipeline_manager is None: + raise RuntimeError("Pipeline manager is required to validate the training bundle.") + for backbone_name, backbone in self.pipeline_manager.iter_backbones(): + if mto.ModeloptStateManager.is_converted(backbone): + state_path = output_bundle / backbone_name / "modelopt_state.pth" + if not state_path.is_file(): + raise RuntimeError( + f"ModelOpt state was not saved for {backbone_name}: {state_path}" + ) + self.logger.info("ModelOpt training bundle saved successfully") + def export_onnx( self, pipe: DiffusionPipeline, @@ -583,6 +607,11 @@ def create_argument_parser() -> argparse.ArgumentParser: type=str, help="Path to save quantized PyTorch checkpoint", ) + export_group.add_argument( + "--output-bundle", + type=str, + help="Directory for a native ModelOpt-aware Diffusers training bundle used by QAD", + ) export_group.add_argument("--onnx-dir", type=str, help="Directory for ONNX export") export_group.add_argument( "--hf-ckpt-dir", @@ -613,6 +642,10 @@ def main() -> None: parser = create_argument_parser() args, unknown_args = parser.parse_known_args() + if args.output_bundle: + # Install ModelOpt's Diffusers save/load hooks before pipeline construction. + mto.enable_huggingface_checkpointing() + model_type = ModelType(args.model) if args.backbone is None: args.backbone = [MODEL_DEFAULTS[model_type]["backbone"]] @@ -672,6 +705,7 @@ def main() -> None: quantized_torch_ckpt_path=Path(args.quantized_torch_ckpt_save_path) if args.quantized_torch_ckpt_save_path else None, + output_bundle=Path(args.output_bundle) if args.output_bundle else None, onnx_dir=Path(args.onnx_dir) if args.onnx_dir else None, hf_ckpt_dir=Path(args.hf_ckpt_dir) if args.hf_ckpt_dir else None, restore_from=Path(args.restore_from) if args.restore_from else None, @@ -730,6 +764,7 @@ def forward_loop(mod): export_manager.save_checkpoint(backbone, backbone_name) pipeline_manager.print_quant_summary() + export_manager.save_training_bundle(pipe) for backbone_name, backbone in pipeline_manager.iter_backbones(): export_manager.export_onnx( diff --git a/examples/diffusers/quantization/quantize_config.py b/examples/diffusers/quantization/quantize_config.py index a92dd4e8147..7bbb724ee59 100644 --- a/examples/diffusers/quantization/quantize_config.py +++ b/examples/diffusers/quantization/quantize_config.py @@ -141,6 +141,7 @@ class ExportConfig: """Configuration for model export.""" quantized_torch_ckpt_path: Path | None = None + output_bundle: Path | None = None onnx_dir: Path | None = None hf_ckpt_dir: Path | None = None restore_from: Path | None = None @@ -155,6 +156,19 @@ def validate(self) -> None: if not parent_dir.exists(): parent_dir.mkdir(parents=True, exist_ok=True) + if self.output_bundle: + parent_dir = self.output_bundle.parent + if not parent_dir.exists(): + parent_dir.mkdir(parents=True, exist_ok=True) + if self.output_bundle.exists() and not self.output_bundle.is_dir(): + raise FileExistsError( + f"Output training bundle path is not a directory: {self.output_bundle}" + ) + if self.output_bundle.exists() and any(self.output_bundle.iterdir()): + raise FileExistsError( + f"Output training bundle directory is not empty: {self.output_bundle}" + ) + if self.onnx_dir and not self.onnx_dir.exists(): self.onnx_dir.mkdir(parents=True, exist_ok=True)