From 602879b28070f9d6c67d56eb66999fe5b68d2125 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:03:58 +0000 Subject: [PATCH 1/6] refactor(export): split the diffusers exporter into its own module The transformers and diffusers export paths shared a file but almost no code: they meet only at the dispatch in export_hf_checkpoint. Move the diffusers half -- _export_diffusers_checkpoint, _postprocess_safetensors, _fuse_qkv_linears_diffusion and four helpers, 498 lines -- to unified_export_diffusers.py. unified_export_hf.py goes 1685 -> 1187. The diffusers-only imports (generate_diffusion_dummy_forward_fn, get_diffusion_components, merge_diffusion_checkpoint and the rest) leave with it; only is_diffusers_object, is_qkv_projection and get_qkv_group_key stay, for the dispatch check and the shared QKV fusion. The dispatch imports _export_diffusers_checkpoint lazily for now, because the diffusers module still imports the module-walking helpers back from here. The following commits move those out and the lazy import goes away. Two test files imported _postprocess_safetensors from the old location and are updated rather than shimmed; test_export_diffusers.py's monkeypatches move to the new module, since a `from X import Y` binding is not affected by patching Y on X. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../torch/export/unified_export_diffusers.py | 577 ++++++++++++++++++ modelopt/torch/export/unified_export_hf.py | 536 +--------------- .../torch/export/test_export_diffusers.py | 17 +- tests/unit/torch/export/test_nvfp4_utils.py | 2 +- 4 files changed, 593 insertions(+), 539 deletions(-) create mode 100644 modelopt/torch/export/unified_export_diffusers.py diff --git a/modelopt/torch/export/unified_export_diffusers.py b/modelopt/torch/export/unified_export_diffusers.py new file mode 100644 index 00000000000..77400d10216 --- /dev/null +++ b/modelopt/torch/export/unified_export_diffusers.py @@ -0,0 +1,577 @@ +# 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. + +"""Unified HF checkpoint export for diffusers models. + +Split out of :mod:`unified_export_hf`, which it shares almost nothing with: the +transformers and diffusers paths meet only at the dispatch in ``export_hf_checkpoint``. +That dispatch imports :func:`_export_diffusers_checkpoint` lazily, since this module +imports the shared module-walking helpers back from ``unified_export_hf``. +""" + +import json +import warnings +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from safetensors import safe_open +from safetensors.torch import save_file + +from .convert_hf_config import convert_hf_quant_config_format +from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales +from .layer_utils import is_quantlinear +from .model_config import QUANTIZATION_NONE +from .quant_utils import get_quant_config, get_quantization_format, has_quantized_modules +from .unified_export_hf import ( + _fuse_shared_input_modules, + _process_quantized_modules, + collect_shared_input_modules, +) + +try: + import diffusers + + from .diffusers_utils import ( + generate_diffusion_dummy_forward_fn, + get_diffusion_components, + get_diffusion_model_type, + hide_quantizers_from_state_dict, + infer_dtype_from_model, + merge_diffusion_checkpoint, + ) + + HAS_DIFFUSERS = True +except ImportError: + HAS_DIFFUSERS = False + +try: + from modelopt.torch.sparsity.attention_sparsity.conversion import export_sparse_attention_config +except ImportError: + export_sparse_attention_config = None + + +def _save_component_state_dict_safetensors( + component: nn.Module, + component_export_dir: Path, +) -> None: + """Save component state dict as a plain safetensors file. + + Args: + component: The nn.Module to save. + component_export_dir: Directory to save model.safetensors and config.json. + """ + cpu_state_dict = {k: v.detach().contiguous().cpu() for k, v in component.state_dict().items()} + metadata = { + "_export_format": "safetensors_state_dict", + "_class_name": type(component).__name__, + } + + save_file( + cpu_state_dict, + str(component_export_dir / "model.safetensors"), + metadata=metadata, + ) + + with open(component_export_dir / "config.json", "w") as f: + json.dump(metadata, f, indent=4) + + +def _postprocess_safetensors( + export_dir: Path, + pipe: Any | None = None, + hf_quant_config: dict | None = None, + **kwargs, +) -> None: + """Post-process saved safetensors files for deployment compatibility. + + Loads each ``.safetensors`` file in *export_dir* and applies all requested + transformations in order, then re-saves in-place with updated metadata: + + 1. **Merge** with base checkpoint — combines quantized transformer weights with + non-transformer components (VAE, vocoder, text encoders) from a base + ``.safetensors`` file to produce a single-file checkpoint (e.g., for ComfyUI). + 2. **Pad** NVFP4 weight/scale tensors — ensures dimensions are multiples of 16 + for hardware alignment requirements. + 3. **Swizzle** NVFP4 block scales — rearranges from flat layout to cuBLAS 2-D + block-scaling-factors tiled layout for optimized inference. + 4. **Inject metadata** — embeds ``quantization_config`` and per-layer + ``_quantization_metadata`` so inference runtimes can detect and handle + quantized layers. + + All of these target single-file deployment runtimes (e.g. ComfyUI) and are + opt-in; ModelOpt itself reads the quant config from ``config.json`` on reload. If + the caller passes none of ``merged_base_safetensor_path``, ``padding_strategy``, + ``enable_swizzle_layout``, or ``enable_layerwise_quant_metadata``, this function + does nothing and leaves the standard exported checkpoint untouched. + + Args: + export_dir: Directory containing the saved ``.safetensors`` file(s). + pipe: The diffusion pipeline / model. Used to infer the model type + (via :func:`get_diffusion_model_type`) when + ``merged_base_safetensor_path`` is set. + hf_quant_config: Quantization config dict to embed in metadata. + **kwargs: Runtime-specific keyword arguments: + merged_base_safetensor_path (str, optional): When provided, merges + the exported transformer weights with non-transformer components + (VAE, vocoder, text encoders, etc.) from this base safetensors + file to produce a single-file checkpoint compatible with ComfyUI. + Value should be the path to a full base model ``.safetensors`` + file (e.g. ``"path/to/ltx-2-19b-dev.safetensors"``). + enable_layerwise_quant_metadata (bool, optional): When True, embeds + ``quantization_config`` and per-layer ``_quantization_metadata`` in the + safetensors header so single-file runtimes (e.g., ComfyUI) can identify + which layers are quantized and in what format. Defaults to False (no + header metadata; this alone leaves the export untouched). + enable_swizzle_layout (bool, optional): When True, rearranges NVFP4 + block scales from ModelOpt's flat layout to cuBLAS 2-D tiled + layout. Required for runtimes that consume cuBLAS block-scaled + GEMM (e.g., comfy_kitchen). Defaults to False. + padding_strategy (str | None, optional): Padding strategy for NVFP4 + weight and scale tensors. ``"row"`` pads rows to multiples of + 16 (columns assumed already aligned). ``"row_col"`` pads both + dimensions. ``None`` (default) disables padding. Independent of + ``enable_swizzle_layout``. + + """ + merged_base_safetensor_path: str | None = kwargs.get("merged_base_safetensor_path") + enable_layerwise_quant_metadata: bool = kwargs.get("enable_layerwise_quant_metadata", False) + enable_swizzle_layout: bool = kwargs.get("enable_swizzle_layout", False) + padding_strategy: str | None = kwargs.get("padding_strategy") + + # This post-processing only produces single-file deployment checkpoints (e.g. + # ComfyUI): merging with a base checkpoint, NVFP4 padding/swizzling, and embedding + # quant metadata in the safetensors header. None of it is read back by ModelOpt + # (the diffusers reload uses ``config.json``), so if the user has not opted into any + # of these options there is nothing to do — leave the exported checkpoint untouched. + if not ( + merged_base_safetensor_path is not None + or padding_strategy is not None + or enable_swizzle_layout + or enable_layerwise_quant_metadata + ): + return + + safetensor_files = sorted(export_dir.glob("*.safetensors")) + if not safetensor_files: + return + + if list(export_dir.glob("*.safetensors.index.json")) and ( + merged_base_safetensor_path is not None or enable_layerwise_quant_metadata + ): + raise NotImplementedError( + "Post-processing sharded safetensors is not supported. " + "Export with a larger max_shard_size or disable merge/metadata options." + ) + + model_type: str | None = None + if merged_base_safetensor_path is not None: + if pipe is None: + raise ValueError("`pipe` must be provided when `merged_base_safetensor_path` is set.") + model_type = get_diffusion_model_type(pipe) + + for sf_path in safetensor_files: + with safe_open(str(sf_path), framework="pt") as f: + metadata = dict(f.metadata() or {}) + sd = {k: f.get_tensor(k).clone() for k in f.keys()} # noqa: SIM118 + + if merged_base_safetensor_path is not None and model_type is not None: + sd, base_metadata = merge_diffusion_checkpoint( + sd, merged_base_safetensor_path, model_type, hf_quant_config=None + ) + base_metadata.update(metadata) + metadata = base_metadata + + if padding_strategy is not None: + sd = pad_nvfp4_weights(sd, padding_strategy) + + if enable_swizzle_layout: + sd = swizzle_nvfp4_scales(sd) + + if hf_quant_config is not None: + metadata["quantization_config"] = json.dumps(hf_quant_config) + if enable_layerwise_quant_metadata: + metadata["_quantization_metadata"] = build_layerwise_quant_metadata( + sd, hf_quant_config + ) + + save_file(sd, str(sf_path), metadata=metadata) + + +def _fuse_qkv_linears_diffusion( + model: nn.Module, + dummy_forward_fn: Callable[[], None] | None = None, + strict: bool = False, +) -> None: + """Fuse QKV linear layers that share the same input for diffusion models. + + This function uses forward hooks to dynamically identify linear modules that + share the same input tensor (e.g., q_proj, k_proj, v_proj in attention). + For these modules, it unifies their input and weight amax values. + + Note: This is a simplified version for diffusion models that: + - Handles QKV fusion (shared input detection) + - Filters to only fuse actual QKV projection layers (not AdaLN, FFN, etc.) + - Skips pre_quant_scale *fusion* (the export path promotes pre_quant_scale to + module-level keys separately; see _promote_quantizer_tensors_to_module) + - Skips FFN fusion with layernorm (TODO for future) + + Args: + model: The diffusion model component (e.g., transformer, unet). + dummy_forward_fn: Optional callable to run a dummy forward pass. Use this + for diffusion-like models whose forward signature is not compatible + with `generate_diffusion_dummy_inputs`. + """ + quantization_format = get_quantization_format(model) + + if quantization_format == QUANTIZATION_NONE: + return + + if dummy_forward_fn is None: + dummy_forward_fn = generate_diffusion_dummy_forward_fn(model) + + # Collect modules sharing the same input + try: + input_to_linear, _ = collect_shared_input_modules( + model, dummy_forward_fn, collect_layernorms=False + ) + except Exception as e: + if strict: + raise RuntimeError( + f"QKV fusion dummy forward failed for {type(model).__name__}; a working " + f"dummy forward is required to export this model correctly. Original error: {e}" + ) from e + print(f"Warning: Failed to run dummy forward for QKV fusion: {e}") + print("Skipping QKV fusion. Quantization may still work but amax values won't be unified.") + return + + if not input_to_linear: + print("No quantized linear modules found for QKV fusion.") + return + + # Fuse the collected modules (QKV only for diffusion) + _fuse_shared_input_modules( + model, + input_to_linear, + output_to_layernorm=None, + qkv_only=True, + fuse_layernorms=False, + quantization_format=quantization_format, + ) + + +def _detect_svdquant_rank(component: nn.Module) -> int | None: + """Return the single SVDQuant low-rank dimension shared by the SVDQuant linears. + + ``svdquant_lora_a`` has shape ``(rank, in_features)``, so its first dimension is + the low-rank size. A single global ``lora_rank`` is written to the checkpoint + config, so all SVDQuant linears are expected to share one rank; an inconsistency + is raised rather than silently recording one module's rank for all. Returns + ``None`` when no SVDQuant LoRA factors are present. + """ + ranks: set[int] = set() + for _, sub_module in component.named_modules(): + weight_quantizer = getattr(sub_module, "weight_quantizer", None) + lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) + if lora_a is not None: + ranks.add(int(lora_a.shape[0])) + if not ranks: + return None + if len(ranks) > 1: + raise ValueError(f"Inconsistent SVDQuant ranks across modules: {sorted(ranks)}") + return next(iter(ranks)) + + +def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: + """Promote quantizer-owned export tensors onto their parent linear module. + + The diffusers export path saves via ``save_pretrained`` inside + :func:`hide_quantizers_from_state_dict` (which deletes the ``weight_quantizer`` + / ``input_quantizer`` submodules) and -- unlike the transformers path -- does + NOT run :func:`postprocess_state_dict`. Without this step the AWQ smoothing + scale and the SVDQuant low-rank factors would be dropped from the exported + checkpoint. We register them as module buffers under clean, AWQ-aligned keys + so they are embedded in the component's main safetensors: + + - ``input_quantizer._pre_quant_scale`` -> ``.pre_quant_scale`` + (the same key the transformers/AWQ path produces via postprocess_state_dict) + - ``weight_quantizer.svdquant_lora_a`` -> ``.svdquant_lora_a`` + - ``weight_quantizer.svdquant_lora_b`` -> ``.svdquant_lora_b`` + + This runs after :func:`_process_quantized_modules` (which leaves these + quantizer buffers in place) and before ``save_pretrained``. + """ + for _, sub_module in component.named_modules(): + if not is_quantlinear(sub_module): + continue + + # register_buffer overwrites an existing buffer of the same name, so a + # repeated export refreshes (rather than keeps stale) promoted tensors. + input_quantizer = getattr(sub_module, "input_quantizer", None) + pre_quant_scale = getattr(input_quantizer, "_pre_quant_scale", None) + if pre_quant_scale is not None: + sub_module.register_buffer("pre_quant_scale", pre_quant_scale.detach().clone()) + + weight_quantizer = getattr(sub_module, "weight_quantizer", None) + lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) + lora_b = getattr(weight_quantizer, "svdquant_lora_b", None) + if lora_a is not None and lora_b is not None: + sub_module.register_buffer("svdquant_lora_a", lora_a.detach().clone()) + sub_module.register_buffer("svdquant_lora_b", lora_b.detach().clone()) + + +def _remove_promoted_quantizer_tensors(component: nn.Module) -> None: + """Undo :func:`_promote_quantizer_tensors_to_module`. + + Removes the temporary module-level export buffers (``svdquant_lora_a/b`` and + ``pre_quant_scale``) so the live module is unchanged after export, keeping + repeated export / post-export module reuse correct. The quantizer-owned tensors + (``weight_quantizer.svdquant_lora_a/b``, ``input_quantizer._pre_quant_scale``) + are left untouched. + """ + for _, sub_module in component.named_modules(): + for buffer_name in ("svdquant_lora_a", "svdquant_lora_b", "pre_quant_scale"): + if buffer_name in getattr(sub_module, "_buffers", {}): + del sub_module._buffers[buffer_name] + + +def _export_diffusers_checkpoint( + pipe: Any, + dtype: torch.dtype | None, + export_dir: Path, + components: list[str] | None, + max_shard_size: int | str = "10GB", + **kwargs, +) -> None: + """Internal: Export diffusion(-like) model/pipeline checkpoint. + + This function handles the export of: + - diffusers models: DiffusionPipeline and individual ModelMixin components. + - LTX-2 pipelines (duck-typed): exports stage-1 transformer only. + + Args: + pipe: The model or pipeline to export. + dtype: The data type for weight conversion. If None, will be inferred from model. + export_dir: The directory to save the exported checkpoint. + components: Optional list of component names to export. Only used for pipelines. + If None, all components are exported. + max_shard_size: Maximum size of each shard file. If the model exceeds this size, + it will be sharded into multiple files and a .safetensors.index.json will be + created. Use smaller values like "5GB" or "2GB" to force sharding. + **kwargs: Runtime-specific post-processing options forwarded to + :func:`_postprocess_safetensors`. See its docstring for details. + """ + export_dir = Path(export_dir) + + # Get all pipeline components (nn.Module, tokenizers, schedulers, etc.) + all_components = get_diffusion_components(pipe, components) + + if not all_components: + warnings.warn("No exportable components found in the model.") + return + + # Separate nn.Module components for quantization-aware export + module_components = { + name: comp for name, comp in all_components.items() if isinstance(comp, nn.Module) + } + + # Best-effort diffusers pipeline check (kept for folder layout + model_index.json behavior) + is_diffusers_pipe = False + if HAS_DIFFUSERS: + try: + from diffusers import DiffusionPipeline as _DiffusionPipeline + + is_diffusers_pipe = isinstance(pipe, _DiffusionPipeline) + except Exception: + is_diffusers_pipe = False + + # Export each nn.Module component with quantization handling + for component_name, component in module_components.items(): + is_quantized = has_quantized_modules(component) + status = "quantized" if is_quantized else "non-quantized" + print(f"Exporting component: {component_name} ({status})") + + # Determine component export directory + # For pipelines, each component goes in a subfolder + if is_diffusers_pipe: + component_export_dir = export_dir / component_name + else: + component_export_dir = export_dir + + component_export_dir.mkdir(parents=True, exist_ok=True) + + # Infer dtype if not provided + component_dtype = dtype if dtype is not None else infer_dtype_from_model(component) + + if is_quantized: + # Fuse QKV linears that share the same input (unify amax values) + # This is similar to requantize_resmooth_fused_llm_layers but simplified for diffusion + # TODO: Add FFN fusion for AWQ-style quantization (pre_quant_scale is + # promoted to module keys at export by _promote_quantizer_tensors_to_module below) + print(f" Running QKV fusion for {component_name}...") + # Qwen-Image's packed-latent forward signature is non-standard; if the + # dummy forward fails for it, fail loudly rather than silently skipping + # fusion (which would export un-unified amax values). + is_qwen_component = "qwen" in type(component).__name__.lower() + _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) + + # Process quantized modules (convert weights, register scales) + _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) + + # Promote quantizer-owned tensors (AWQ pre_quant_scale and SVDQuant + # LoRA factors) onto the module so they survive + # hide_quantizers_from_state_dict and are embedded in the component's + # main safetensors under clean, AWQ-aligned keys. + _promote_quantizer_tensors_to_module(component) + + # Build the quantization config + save inside try/finally so the temporary + # promoted buffers are always removed, even if save / post-process / config + # update raises (keeps the live module reusable for a repeated export). + try: + quant_config = get_quant_config(component, is_modelopt_qlora=False) + if quant_config: + quantization_details = quant_config.get("quantization", {}) + # Record the SVDQuant low-rank size so consumers know the LoRA shape. + if quantization_details.get("quant_algo") == "NVFP4_SVD": + svdquant_rank = _detect_svdquant_rank(component) + if svdquant_rank is not None: + quantization_details["lora_rank"] = svdquant_rank + hf_quant_config = ( + convert_hf_quant_config_format(quant_config) if quant_config else None + ) + + # Save the component + # - diffusers ModelMixin.save_pretrained does NOT accept state_dict parameter + # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save + if hasattr(component, "save_pretrained"): + with hide_quantizers_from_state_dict(component): + component.save_pretrained( + component_export_dir, max_shard_size=max_shard_size + ) + else: + with hide_quantizers_from_state_dict(component): + _save_component_state_dict_safetensors(component, component_export_dir) + + # Post-process — merge, metadata, padding, swizzle + _postprocess_safetensors( + component_export_dir, + pipe, + hf_quant_config=hf_quant_config, + **kwargs, + ) + + # Update config.json with quantization info + if hf_quant_config is not None: + config_path = component_export_dir / "config.json" + if config_path.exists(): + with open(config_path) as file: + config_data = json.load(file) + config_data["quantization_config"] = hf_quant_config + with open(config_path, "w") as file: + json.dump(config_data, file, indent=4) + finally: + # Drop the temporary promoted export buffers so the live module is + # unchanged after export (supports repeated export / module reuse). + _remove_promoted_quantizer_tensors(component) + # Non-quantized component: just save as-is + elif hasattr(component, "save_pretrained"): + component.save_pretrained(component_export_dir, max_shard_size=max_shard_size) + else: + _save_component_state_dict_safetensors(component, component_export_dir) + + # Update config.json with sparse attention info (both quantized and non-quantized) + if export_sparse_attention_config is not None: + sparse_attn_config = export_sparse_attention_config(component) + if sparse_attn_config is not None: + config_path = component_export_dir / "config.json" + if config_path.exists(): + with open(config_path) as file: + config_data = json.load(file) + config_data["sparse_attention_config"] = sparse_attn_config + with open(config_path, "w") as file: + json.dump(config_data, file, indent=4) + print(f" Added sparse_attention_config to {config_path.name}") + + print(f" Saved to: {component_export_dir}") + + # Export non-nn.Module components (tokenizers, schedulers, feature extractors, etc.) + if is_diffusers_pipe: + for component_name, component in all_components.items(): + # Skip nn.Module components (already handled above) + if isinstance(component, nn.Module): + continue + + component_export_dir = export_dir / component_name + component_export_dir.mkdir(parents=True, exist_ok=True) + + print(f"Exporting component: {component_name} ({type(component).__name__})") + + # Handle different component types + if hasattr(component, "save_pretrained"): + # Tokenizers, feature extractors, image processors + component.save_pretrained(component_export_dir) + elif hasattr(component, "save_config"): + # Schedulers + component.save_config(component_export_dir) + else: + warnings.warn( + f"Component '{component_name}' of type {type(component).__name__} " + "does not have save_pretrained or save_config method. Skipping." + ) + continue + + print(f" Saved to: {component_export_dir}") + + # For pipelines, also save model_index.json + if is_diffusers_pipe: + model_index_path = export_dir / "model_index.json" + is_partial_export = components is not None + + # For full export, preserve original model_index.json when possible. + # For partial export, skip this to avoid listing non-exported components. + if not is_partial_export: + source_path = getattr(pipe, "name_or_path", None) or getattr( + getattr(pipe, "config", None), "_name_or_path", None + ) + if source_path: + candidate_model_index = Path(source_path) / "model_index.json" + if candidate_model_index.exists(): + with open(candidate_model_index) as file: + model_index = json.load(file) + with open(model_index_path, "w") as file: + json.dump(model_index, file, indent=4) + + # Full-export fallback to Diffusers-native config serialization. + # Partial export skips this for the same reason as above. + if not is_partial_export and not model_index_path.exists() and hasattr(pipe, "save_config"): + pipe.save_config(export_dir) + + # Last resort: synthesize a minimal model_index.json from exported components. + if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None: + model_index = { + "_class_name": type(pipe).__name__, + "_diffusers_version": diffusers.__version__, + } + for name, comp in all_components.items(): + module = type(comp).__module__ + library = module.split(".")[0] + model_index[name] = [library, type(comp).__name__] + + with open(model_index_path, "w") as file: + json.dump(model_index, file, indent=4) + + print(f"Export complete. Saved to: {export_dir}") diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 2a605ed6d9e..d8d9386d379 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -27,25 +27,9 @@ import torch import torch.nn as nn -from safetensors import safe_open -from safetensors.torch import save_file - -from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales try: - import diffusers - - from .diffusers_utils import ( - generate_diffusion_dummy_forward_fn, - get_diffusion_components, - get_diffusion_model_type, - get_qkv_group_key, - hide_quantizers_from_state_dict, - infer_dtype_from_model, - is_diffusers_object, - is_qkv_projection, - merge_diffusion_checkpoint, - ) + from .diffusers_utils import get_qkv_group_key, is_diffusers_object, is_qkv_projection HAS_DIFFUSERS = True except ImportError: @@ -108,7 +92,6 @@ get_weight_block_size, get_weight_scaling_factor, get_weight_scaling_factor_2, - has_quantized_modules, maybe_transpose_expert_weight_dimensions, postprocess_state_dict, preprocess_linear_fusion, @@ -130,153 +113,6 @@ def _is_enabled_quantizer(quantizer): return False -def _save_component_state_dict_safetensors( - component: nn.Module, - component_export_dir: Path, -) -> None: - """Save component state dict as a plain safetensors file. - - Args: - component: The nn.Module to save. - component_export_dir: Directory to save model.safetensors and config.json. - """ - cpu_state_dict = {k: v.detach().contiguous().cpu() for k, v in component.state_dict().items()} - metadata = { - "_export_format": "safetensors_state_dict", - "_class_name": type(component).__name__, - } - - save_file( - cpu_state_dict, - str(component_export_dir / "model.safetensors"), - metadata=metadata, - ) - - with open(component_export_dir / "config.json", "w") as f: - json.dump(metadata, f, indent=4) - - -def _postprocess_safetensors( - export_dir: Path, - pipe: Any | None = None, - hf_quant_config: dict | None = None, - **kwargs, -) -> None: - """Post-process saved safetensors files for deployment compatibility. - - Loads each ``.safetensors`` file in *export_dir* and applies all requested - transformations in order, then re-saves in-place with updated metadata: - - 1. **Merge** with base checkpoint — combines quantized transformer weights with - non-transformer components (VAE, vocoder, text encoders) from a base - ``.safetensors`` file to produce a single-file checkpoint (e.g., for ComfyUI). - 2. **Pad** NVFP4 weight/scale tensors — ensures dimensions are multiples of 16 - for hardware alignment requirements. - 3. **Swizzle** NVFP4 block scales — rearranges from flat layout to cuBLAS 2-D - block-scaling-factors tiled layout for optimized inference. - 4. **Inject metadata** — embeds ``quantization_config`` and per-layer - ``_quantization_metadata`` so inference runtimes can detect and handle - quantized layers. - - All of these target single-file deployment runtimes (e.g. ComfyUI) and are - opt-in; ModelOpt itself reads the quant config from ``config.json`` on reload. If - the caller passes none of ``merged_base_safetensor_path``, ``padding_strategy``, - ``enable_swizzle_layout``, or ``enable_layerwise_quant_metadata``, this function - does nothing and leaves the standard exported checkpoint untouched. - - Args: - export_dir: Directory containing the saved ``.safetensors`` file(s). - pipe: The diffusion pipeline / model. Used to infer the model type - (via :func:`get_diffusion_model_type`) when - ``merged_base_safetensor_path`` is set. - hf_quant_config: Quantization config dict to embed in metadata. - **kwargs: Runtime-specific keyword arguments: - merged_base_safetensor_path (str, optional): When provided, merges - the exported transformer weights with non-transformer components - (VAE, vocoder, text encoders, etc.) from this base safetensors - file to produce a single-file checkpoint compatible with ComfyUI. - Value should be the path to a full base model ``.safetensors`` - file (e.g. ``"path/to/ltx-2-19b-dev.safetensors"``). - enable_layerwise_quant_metadata (bool, optional): When True, embeds - ``quantization_config`` and per-layer ``_quantization_metadata`` in the - safetensors header so single-file runtimes (e.g., ComfyUI) can identify - which layers are quantized and in what format. Defaults to False (no - header metadata; this alone leaves the export untouched). - enable_swizzle_layout (bool, optional): When True, rearranges NVFP4 - block scales from ModelOpt's flat layout to cuBLAS 2-D tiled - layout. Required for runtimes that consume cuBLAS block-scaled - GEMM (e.g., comfy_kitchen). Defaults to False. - padding_strategy (str | None, optional): Padding strategy for NVFP4 - weight and scale tensors. ``"row"`` pads rows to multiples of - 16 (columns assumed already aligned). ``"row_col"`` pads both - dimensions. ``None`` (default) disables padding. Independent of - ``enable_swizzle_layout``. - - """ - merged_base_safetensor_path: str | None = kwargs.get("merged_base_safetensor_path") - enable_layerwise_quant_metadata: bool = kwargs.get("enable_layerwise_quant_metadata", False) - enable_swizzle_layout: bool = kwargs.get("enable_swizzle_layout", False) - padding_strategy: str | None = kwargs.get("padding_strategy") - - # This post-processing only produces single-file deployment checkpoints (e.g. - # ComfyUI): merging with a base checkpoint, NVFP4 padding/swizzling, and embedding - # quant metadata in the safetensors header. None of it is read back by ModelOpt - # (the diffusers reload uses ``config.json``), so if the user has not opted into any - # of these options there is nothing to do — leave the exported checkpoint untouched. - if not ( - merged_base_safetensor_path is not None - or padding_strategy is not None - or enable_swizzle_layout - or enable_layerwise_quant_metadata - ): - return - - safetensor_files = sorted(export_dir.glob("*.safetensors")) - if not safetensor_files: - return - - if list(export_dir.glob("*.safetensors.index.json")) and ( - merged_base_safetensor_path is not None or enable_layerwise_quant_metadata - ): - raise NotImplementedError( - "Post-processing sharded safetensors is not supported. " - "Export with a larger max_shard_size or disable merge/metadata options." - ) - - model_type: str | None = None - if merged_base_safetensor_path is not None: - if pipe is None: - raise ValueError("`pipe` must be provided when `merged_base_safetensor_path` is set.") - model_type = get_diffusion_model_type(pipe) - - for sf_path in safetensor_files: - with safe_open(str(sf_path), framework="pt") as f: - metadata = dict(f.metadata() or {}) - sd = {k: f.get_tensor(k).clone() for k in f.keys()} # noqa: SIM118 - - if merged_base_safetensor_path is not None and model_type is not None: - sd, base_metadata = merge_diffusion_checkpoint( - sd, merged_base_safetensor_path, model_type, hf_quant_config=None - ) - base_metadata.update(metadata) - metadata = base_metadata - - if padding_strategy is not None: - sd = pad_nvfp4_weights(sd, padding_strategy) - - if enable_swizzle_layout: - sd = swizzle_nvfp4_scales(sd) - - if hf_quant_config is not None: - metadata["quantization_config"] = json.dumps(hf_quant_config) - if enable_layerwise_quant_metadata: - metadata["_quantization_metadata"] = build_layerwise_quant_metadata( - sd, hf_quant_config - ) - - save_file(sd, str(sf_path), metadata=metadata) - - def collect_shared_input_modules( model: nn.Module, dummy_forward_fn: Callable[[], None], @@ -1059,371 +895,6 @@ def _export_transformers_checkpoint( return quantized_state_dict, quant_config -def _fuse_qkv_linears_diffusion( - model: nn.Module, - dummy_forward_fn: Callable[[], None] | None = None, - strict: bool = False, -) -> None: - """Fuse QKV linear layers that share the same input for diffusion models. - - This function uses forward hooks to dynamically identify linear modules that - share the same input tensor (e.g., q_proj, k_proj, v_proj in attention). - For these modules, it unifies their input and weight amax values. - - Note: This is a simplified version for diffusion models that: - - Handles QKV fusion (shared input detection) - - Filters to only fuse actual QKV projection layers (not AdaLN, FFN, etc.) - - Skips pre_quant_scale *fusion* (the export path promotes pre_quant_scale to - module-level keys separately; see _promote_quantizer_tensors_to_module) - - Skips FFN fusion with layernorm (TODO for future) - - Args: - model: The diffusion model component (e.g., transformer, unet). - dummy_forward_fn: Optional callable to run a dummy forward pass. Use this - for diffusion-like models whose forward signature is not compatible - with `generate_diffusion_dummy_inputs`. - """ - quantization_format = get_quantization_format(model) - - if quantization_format == QUANTIZATION_NONE: - return - - if dummy_forward_fn is None: - dummy_forward_fn = generate_diffusion_dummy_forward_fn(model) - - # Collect modules sharing the same input - try: - input_to_linear, _ = collect_shared_input_modules( - model, dummy_forward_fn, collect_layernorms=False - ) - except Exception as e: - if strict: - raise RuntimeError( - f"QKV fusion dummy forward failed for {type(model).__name__}; a working " - f"dummy forward is required to export this model correctly. Original error: {e}" - ) from e - print(f"Warning: Failed to run dummy forward for QKV fusion: {e}") - print("Skipping QKV fusion. Quantization may still work but amax values won't be unified.") - return - - if not input_to_linear: - print("No quantized linear modules found for QKV fusion.") - return - - # Fuse the collected modules (QKV only for diffusion) - _fuse_shared_input_modules( - model, - input_to_linear, - output_to_layernorm=None, - qkv_only=True, - fuse_layernorms=False, - quantization_format=quantization_format, - ) - - -def _detect_svdquant_rank(component: nn.Module) -> int | None: - """Return the single SVDQuant low-rank dimension shared by the SVDQuant linears. - - ``svdquant_lora_a`` has shape ``(rank, in_features)``, so its first dimension is - the low-rank size. A single global ``lora_rank`` is written to the checkpoint - config, so all SVDQuant linears are expected to share one rank; an inconsistency - is raised rather than silently recording one module's rank for all. Returns - ``None`` when no SVDQuant LoRA factors are present. - """ - ranks: set[int] = set() - for _, sub_module in component.named_modules(): - weight_quantizer = getattr(sub_module, "weight_quantizer", None) - lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) - if lora_a is not None: - ranks.add(int(lora_a.shape[0])) - if not ranks: - return None - if len(ranks) > 1: - raise ValueError(f"Inconsistent SVDQuant ranks across modules: {sorted(ranks)}") - return next(iter(ranks)) - - -def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: - """Promote quantizer-owned export tensors onto their parent linear module. - - The diffusers export path saves via ``save_pretrained`` inside - :func:`hide_quantizers_from_state_dict` (which deletes the ``weight_quantizer`` - / ``input_quantizer`` submodules) and -- unlike the transformers path -- does - NOT run :func:`postprocess_state_dict`. Without this step the AWQ smoothing - scale and the SVDQuant low-rank factors would be dropped from the exported - checkpoint. We register them as module buffers under clean, AWQ-aligned keys - so they are embedded in the component's main safetensors: - - - ``input_quantizer._pre_quant_scale`` -> ``.pre_quant_scale`` - (the same key the transformers/AWQ path produces via postprocess_state_dict) - - ``weight_quantizer.svdquant_lora_a`` -> ``.svdquant_lora_a`` - - ``weight_quantizer.svdquant_lora_b`` -> ``.svdquant_lora_b`` - - This runs after :func:`_process_quantized_modules` (which leaves these - quantizer buffers in place) and before ``save_pretrained``. - """ - for _, sub_module in component.named_modules(): - if not is_quantlinear(sub_module): - continue - - # register_buffer overwrites an existing buffer of the same name, so a - # repeated export refreshes (rather than keeps stale) promoted tensors. - input_quantizer = getattr(sub_module, "input_quantizer", None) - pre_quant_scale = getattr(input_quantizer, "_pre_quant_scale", None) - if pre_quant_scale is not None: - sub_module.register_buffer("pre_quant_scale", pre_quant_scale.detach().clone()) - - weight_quantizer = getattr(sub_module, "weight_quantizer", None) - lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) - lora_b = getattr(weight_quantizer, "svdquant_lora_b", None) - if lora_a is not None and lora_b is not None: - sub_module.register_buffer("svdquant_lora_a", lora_a.detach().clone()) - sub_module.register_buffer("svdquant_lora_b", lora_b.detach().clone()) - - -def _remove_promoted_quantizer_tensors(component: nn.Module) -> None: - """Undo :func:`_promote_quantizer_tensors_to_module`. - - Removes the temporary module-level export buffers (``svdquant_lora_a/b`` and - ``pre_quant_scale``) so the live module is unchanged after export, keeping - repeated export / post-export module reuse correct. The quantizer-owned tensors - (``weight_quantizer.svdquant_lora_a/b``, ``input_quantizer._pre_quant_scale``) - are left untouched. - """ - for _, sub_module in component.named_modules(): - for buffer_name in ("svdquant_lora_a", "svdquant_lora_b", "pre_quant_scale"): - if buffer_name in getattr(sub_module, "_buffers", {}): - del sub_module._buffers[buffer_name] - - -def _export_diffusers_checkpoint( - pipe: Any, - dtype: torch.dtype | None, - export_dir: Path, - components: list[str] | None, - max_shard_size: int | str = "10GB", - **kwargs, -) -> None: - """Internal: Export diffusion(-like) model/pipeline checkpoint. - - This function handles the export of: - - diffusers models: DiffusionPipeline and individual ModelMixin components. - - LTX-2 pipelines (duck-typed): exports stage-1 transformer only. - - Args: - pipe: The model or pipeline to export. - dtype: The data type for weight conversion. If None, will be inferred from model. - export_dir: The directory to save the exported checkpoint. - components: Optional list of component names to export. Only used for pipelines. - If None, all components are exported. - max_shard_size: Maximum size of each shard file. If the model exceeds this size, - it will be sharded into multiple files and a .safetensors.index.json will be - created. Use smaller values like "5GB" or "2GB" to force sharding. - **kwargs: Runtime-specific post-processing options forwarded to - :func:`_postprocess_safetensors`. See its docstring for details. - """ - export_dir = Path(export_dir) - - # Get all pipeline components (nn.Module, tokenizers, schedulers, etc.) - all_components = get_diffusion_components(pipe, components) - - if not all_components: - warnings.warn("No exportable components found in the model.") - return - - # Separate nn.Module components for quantization-aware export - module_components = { - name: comp for name, comp in all_components.items() if isinstance(comp, nn.Module) - } - - # Best-effort diffusers pipeline check (kept for folder layout + model_index.json behavior) - is_diffusers_pipe = False - if HAS_DIFFUSERS: - try: - from diffusers import DiffusionPipeline as _DiffusionPipeline - - is_diffusers_pipe = isinstance(pipe, _DiffusionPipeline) - except Exception: - is_diffusers_pipe = False - - # Export each nn.Module component with quantization handling - for component_name, component in module_components.items(): - is_quantized = has_quantized_modules(component) - status = "quantized" if is_quantized else "non-quantized" - print(f"Exporting component: {component_name} ({status})") - - # Determine component export directory - # For pipelines, each component goes in a subfolder - if is_diffusers_pipe: - component_export_dir = export_dir / component_name - else: - component_export_dir = export_dir - - component_export_dir.mkdir(parents=True, exist_ok=True) - - # Infer dtype if not provided - component_dtype = dtype if dtype is not None else infer_dtype_from_model(component) - - if is_quantized: - # Fuse QKV linears that share the same input (unify amax values) - # This is similar to requantize_resmooth_fused_llm_layers but simplified for diffusion - # TODO: Add FFN fusion for AWQ-style quantization (pre_quant_scale is - # promoted to module keys at export by _promote_quantizer_tensors_to_module below) - print(f" Running QKV fusion for {component_name}...") - # Qwen-Image's packed-latent forward signature is non-standard; if the - # dummy forward fails for it, fail loudly rather than silently skipping - # fusion (which would export un-unified amax values). - is_qwen_component = "qwen" in type(component).__name__.lower() - _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) - - # Process quantized modules (convert weights, register scales) - _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) - - # Promote quantizer-owned tensors (AWQ pre_quant_scale and SVDQuant - # LoRA factors) onto the module so they survive - # hide_quantizers_from_state_dict and are embedded in the component's - # main safetensors under clean, AWQ-aligned keys. - _promote_quantizer_tensors_to_module(component) - - # Build the quantization config + save inside try/finally so the temporary - # promoted buffers are always removed, even if save / post-process / config - # update raises (keeps the live module reusable for a repeated export). - try: - quant_config = get_quant_config(component, is_modelopt_qlora=False) - if quant_config: - quantization_details = quant_config.get("quantization", {}) - # Record the SVDQuant low-rank size so consumers know the LoRA shape. - if quantization_details.get("quant_algo") == "NVFP4_SVD": - svdquant_rank = _detect_svdquant_rank(component) - if svdquant_rank is not None: - quantization_details["lora_rank"] = svdquant_rank - hf_quant_config = ( - convert_hf_quant_config_format(quant_config) if quant_config else None - ) - - # Save the component - # - diffusers ModelMixin.save_pretrained does NOT accept state_dict parameter - # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save - if hasattr(component, "save_pretrained"): - with hide_quantizers_from_state_dict(component): - component.save_pretrained( - component_export_dir, max_shard_size=max_shard_size - ) - else: - with hide_quantizers_from_state_dict(component): - _save_component_state_dict_safetensors(component, component_export_dir) - - # Post-process — merge, metadata, padding, swizzle - _postprocess_safetensors( - component_export_dir, - pipe, - hf_quant_config=hf_quant_config, - **kwargs, - ) - - # Update config.json with quantization info - if hf_quant_config is not None: - config_path = component_export_dir / "config.json" - if config_path.exists(): - with open(config_path) as file: - config_data = json.load(file) - config_data["quantization_config"] = hf_quant_config - with open(config_path, "w") as file: - json.dump(config_data, file, indent=4) - finally: - # Drop the temporary promoted export buffers so the live module is - # unchanged after export (supports repeated export / module reuse). - _remove_promoted_quantizer_tensors(component) - # Non-quantized component: just save as-is - elif hasattr(component, "save_pretrained"): - component.save_pretrained(component_export_dir, max_shard_size=max_shard_size) - else: - _save_component_state_dict_safetensors(component, component_export_dir) - - # Update config.json with sparse attention info (both quantized and non-quantized) - if export_sparse_attention_config is not None: - sparse_attn_config = export_sparse_attention_config(component) - if sparse_attn_config is not None: - config_path = component_export_dir / "config.json" - if config_path.exists(): - with open(config_path) as file: - config_data = json.load(file) - config_data["sparse_attention_config"] = sparse_attn_config - with open(config_path, "w") as file: - json.dump(config_data, file, indent=4) - print(f" Added sparse_attention_config to {config_path.name}") - - print(f" Saved to: {component_export_dir}") - - # Export non-nn.Module components (tokenizers, schedulers, feature extractors, etc.) - if is_diffusers_pipe: - for component_name, component in all_components.items(): - # Skip nn.Module components (already handled above) - if isinstance(component, nn.Module): - continue - - component_export_dir = export_dir / component_name - component_export_dir.mkdir(parents=True, exist_ok=True) - - print(f"Exporting component: {component_name} ({type(component).__name__})") - - # Handle different component types - if hasattr(component, "save_pretrained"): - # Tokenizers, feature extractors, image processors - component.save_pretrained(component_export_dir) - elif hasattr(component, "save_config"): - # Schedulers - component.save_config(component_export_dir) - else: - warnings.warn( - f"Component '{component_name}' of type {type(component).__name__} " - "does not have save_pretrained or save_config method. Skipping." - ) - continue - - print(f" Saved to: {component_export_dir}") - - # For pipelines, also save model_index.json - if is_diffusers_pipe: - model_index_path = export_dir / "model_index.json" - is_partial_export = components is not None - - # For full export, preserve original model_index.json when possible. - # For partial export, skip this to avoid listing non-exported components. - if not is_partial_export: - source_path = getattr(pipe, "name_or_path", None) or getattr( - getattr(pipe, "config", None), "_name_or_path", None - ) - if source_path: - candidate_model_index = Path(source_path) / "model_index.json" - if candidate_model_index.exists(): - with open(candidate_model_index) as file: - model_index = json.load(file) - with open(model_index_path, "w") as file: - json.dump(model_index, file, indent=4) - - # Full-export fallback to Diffusers-native config serialization. - # Partial export skips this for the same reason as above. - if not is_partial_export and not model_index_path.exists() and hasattr(pipe, "save_config"): - pipe.save_config(export_dir) - - # Last resort: synthesize a minimal model_index.json from exported components. - if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None: - model_index = { - "_class_name": type(pipe).__name__, - "_diffusers_version": diffusers.__version__, - } - for name, comp in all_components.items(): - module = type(comp).__module__ - library = module.split(".")[0] - model_index[name] = [library, type(comp).__name__] - - with open(model_index_path, "w") as file: - json.dump(model_index, file, indent=4) - - print(f"Export complete. Saved to: {export_dir}") - - # TODO: Remove this workaround once HuggingFace fixes revert_weight_conversion to handle # scalar (0-d) tensors. transformers' Chunk.convert() calls torch.chunk() on quantization # scale buffers that are 0-d scalars, raising RuntimeError ("chunk expects at least a @@ -1566,6 +1037,11 @@ def export_hf_checkpoint( if HAS_DIFFUSERS: is_diffusers_obj = is_diffusers_object(model) if is_diffusers_obj: + # Imported here rather than at module scope: the diffusers exporter imports the + # shared module-walking helpers from this module, so a top-level import would be + # circular. The cycle goes away once those helpers move to their own modules. + from .unified_export_diffusers import _export_diffusers_checkpoint + _export_diffusers_checkpoint( model, dtype, diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 753c81a4b0e..371b5a77e7b 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -31,14 +31,15 @@ from safetensors import safe_open from safetensors.torch import save_file -import modelopt.torch.export.unified_export_hf as unified_export_hf +import modelopt.torch.export.unified_export_diffusers as unified_export_diffusers import modelopt.torch.quantization as mtq from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.diffusers_utils import ( generate_diffusion_dummy_inputs, hide_quantizers_from_state_dict, ) -from modelopt.torch.export.unified_export_hf import _postprocess_safetensors, export_hf_checkpoint +from modelopt.torch.export.unified_export_diffusers import _postprocess_safetensors +from modelopt.torch.export.unified_export_hf import export_hf_checkpoint def _load_config(config_path): @@ -92,7 +93,7 @@ def test_export_diffusers_unet_quantized_matches_llm_config(tmp_path, monkeypatc model = get_tiny_unet() export_dir = tmp_path / "export_unet_quant" - monkeypatch.setattr(unified_export_hf, "has_quantized_modules", lambda *_: True) + monkeypatch.setattr(unified_export_diffusers, "has_quantized_modules", lambda *_: True) fuse_calls = {"count": 0} process_calls = {"count": 0} @@ -103,15 +104,15 @@ def _fuse_stub(*_args, **_kwargs): def _process_stub(*_args, **_kwargs): process_calls["count"] += 1 - monkeypatch.setattr(unified_export_hf, "_fuse_qkv_linears_diffusion", _fuse_stub) - monkeypatch.setattr(unified_export_hf, "_process_quantized_modules", _process_stub) + monkeypatch.setattr(unified_export_diffusers, "_fuse_qkv_linears_diffusion", _fuse_stub) + monkeypatch.setattr(unified_export_diffusers, "_process_quantized_modules", _process_stub) dummy_quant_config = { "quantization": {"quant_algo": "FP8", "kv_cache_quant_algo": "FP8"}, "producer": {"name": "modelopt", "version": "0.0"}, } monkeypatch.setattr( - unified_export_hf, "get_quant_config", lambda *_args, **_kwargs: dummy_quant_config + unified_export_diffusers, "get_quant_config", lambda *_args, **_kwargs: dummy_quant_config ) export_hf_checkpoint(model, export_dir=export_dir) @@ -178,7 +179,7 @@ def test_svdquant_diffusers_export_promotes_clean_keys(): assert getattr(linear.input_quantizer, "_pre_quant_scale", None) is not None # Export promotes them to clean module-level keys and hides the quantizers. - unified_export_hf._promote_quantizer_tensors_to_module(model) + unified_export_diffusers._promote_quantizer_tensors_to_module(model) with hide_quantizers_from_state_dict(model): keys = set(model.state_dict().keys()) @@ -190,7 +191,7 @@ def test_svdquant_diffusers_export_promotes_clean_keys(): ) # The promotion is undone after export, leaving the live module unchanged. - unified_export_hf._remove_promoted_quantizer_tensors(model) + unified_export_diffusers._remove_promoted_quantizer_tensors(model) keys_after = set(model.state_dict().keys()) assert not any( k.endswith((".svdquant_lora_a", ".svdquant_lora_b", ".pre_quant_scale")) for k in keys_after diff --git a/tests/unit/torch/export/test_nvfp4_utils.py b/tests/unit/torch/export/test_nvfp4_utils.py index 7aed23f0b7a..b14d43237a8 100644 --- a/tests/unit/torch/export/test_nvfp4_utils.py +++ b/tests/unit/torch/export/test_nvfp4_utils.py @@ -27,7 +27,7 @@ pad_nvfp4_weights, swizzle_nvfp4_scales, ) -from modelopt.torch.export.unified_export_hf import _postprocess_safetensors +from modelopt.torch.export.unified_export_diffusers import _postprocess_safetensors def _make_nvfp4_state_dict(rows=32, cols=64): From b3037f891068bb4a5c6031e859a201ea773dda66 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:12:29 +0000 Subject: [PATCH 2/6] refactor(export): extract model-level preparation into hf_export_prep Every unified HF exporter runs the same preparation before packing a single weight: resolve the dtype, prepare MoE input quantizers, resmooth and fuse shared-input modules, adjust the quant config, and patch transformers while artifacts are written. That code sat in unified_export_hf.py, so the exporters had to import it back from the module that dispatches to them -- which is the only reason the lazy imports exist. Move those 13 symbols (364 lines) to hf_export_prep.py. It imports nothing else from the export package, so it sits at the bottom of the graph and the three exporters can depend on it without a cycle. unified_export_hf.py goes 1187 -> 823. The QKV fusion helpers travel with _fuse_shared_input_modules, so only is_diffusers_object remains of the diffusers imports here. External importers are repointed rather than shimmed: plugins/vllm_fakequant_hf.py for collect_shared_input_modules, and tests/gpu/.../test_fsdp2_export.py for requantize_resmooth_fused_llm_layers. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/hf_export_prep.py | 455 ++++++++++++++++++ .../torch/export/plugins/vllm_fakequant_hf.py | 2 +- .../torch/export/unified_export_diffusers.py | 7 +- modelopt/torch/export/unified_export_hf.py | 420 +--------------- .../export/unified_export_hf_streaming.py | 10 +- tests/gpu/torch/export/test_fsdp2_export.py | 6 +- 6 files changed, 478 insertions(+), 422 deletions(-) create mode 100644 modelopt/torch/export/hf_export_prep.py diff --git a/modelopt/torch/export/hf_export_prep.py b/modelopt/torch/export/hf_export_prep.py new file mode 100644 index 00000000000..b5e59450738 --- /dev/null +++ b/modelopt/torch/export/hf_export_prep.py @@ -0,0 +1,455 @@ +# 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-level preparation shared by every unified HF exporter. + +Everything here runs on the whole model before any weight is packed: dtype resolution, +MoE input-quantizer preparation, resmoothing and shared-input fusion, quant-config +adjustments, and the transformers patches needed while writing artifacts. + +This module depends on nothing else in the export package, which is what lets the +three exporters import it without a cycle. +""" + +import re +import warnings +from collections import defaultdict +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn as nn + +from modelopt.torch.quantization import set_quantizer_by_cfg_context +from modelopt.torch.quantization.nn import SequentialQuantizer +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update +from modelopt.torch.utils.dataset_utils import _disable_use_cache + +from .layer_utils import ( + get_experts_list, + is_layernorm, + is_moe, + is_quantlinear, + sync_moe_gate_up_amax, +) +from .model_config import ( + QUANTIZATION_FP8, + QUANTIZATION_FP8_PB_REAL, + QUANTIZATION_NONE, + QUANTIZATION_NVFP4_SVDQUANT, +) +from .model_utils import get_language_model_from_vl, is_multimodal_model +from .quant_utils import ( + fuse_prequant_layernorm, + fuse_prequant_to_linear, + get_quantization_format, + preprocess_linear_fusion, +) +from .registry import ExportContext, PrepareMoEInputsRegistry + +try: + from .diffusers_utils import get_qkv_group_key, is_qkv_projection +except ImportError: # diffusers not installed; QKV fusion is diffusers-only + get_qkv_group_key = is_qkv_projection = None + + +def _is_enabled_quantizer(quantizer): + if hasattr(quantizer, "is_enabled") and quantizer.is_enabled: + return True + + if isinstance(quantizer, SequentialQuantizer): + return any(q.is_enabled for q in quantizer) + + return False + + +def collect_shared_input_modules( + model: nn.Module, + dummy_forward_fn: Callable[[], None], + collect_layernorms: bool = False, +) -> tuple[dict, dict | None]: + """Collect modules that share the same input using forward hooks. + + This is a common helper for both LLM and diffusion model fusion. + + Args: + model: The model to analyze. + dummy_forward_fn: A callable that runs a dummy forward pass on the model. + Should be a function that takes no arguments. + collect_layernorms: If True, also collect layernorm output mappings (for AWQ). + + Returns: + A tuple of (input_to_linear, output_to_layernorm). + input_to_linear: Dict mapping input tensor to list of modules sharing that input. + output_to_layernorm: Dict mapping layernorm output to the layernorm module (or None). + """ + input_to_linear: dict = defaultdict(list) + output_to_layernorm: dict | None = defaultdict(lambda: None) if collect_layernorms else None + + def _input_hook(module, input, output): + """Update dictionary with list of all modules that share the same input.""" + if len(input) > 0 and isinstance(input[0], torch.Tensor): + # TODO: Handle DBRX MoE case + input_to_linear[input[0]].append(module) + + def _output_hook(module, input, output): + """Update dictionary with mapping of layernorms and their outputs.""" + if output_to_layernorm is not None and isinstance(output, torch.Tensor): + output_to_layernorm[output] = module + + handles = [] + + # Register hooks on all quantized linear modules (and optionally layernorms) + for name, module in model.named_modules(): + if collect_layernorms and is_layernorm(module): + module.name = name + handle = module.register_forward_hook(_output_hook) + handles.append(handle) + elif is_quantlinear(module) and ( + _is_enabled_quantizer(module.input_quantizer) + or _is_enabled_quantizer(module.weight_quantizer) + ): + module.name = name + handle = module.register_forward_hook(_input_hook) + handles.append(handle) + + if not handles: + return input_to_linear, output_to_layernorm + + # Run dummy forward pass to collect modules sharing same input. + # `_disable_use_cache` keeps the probe forward working on configs that don't + # set `use_cache` (e.g., stepfun-ai/Step-3.5-Flash's Step3p5Config). + try: + with ( + torch.no_grad(), + set_quantizer_by_cfg_context(model, [{"quantizer_name": "*", "enable": False}]), + _disable_use_cache(model), + ): + dummy_forward_fn() + finally: + # Always remove hooks + for handle in handles: + handle.remove() + + return input_to_linear, output_to_layernorm + + +def _fuse_shared_input_modules( + model: nn.Module, + input_to_linear: dict, + output_to_layernorm: dict | None = None, + qkv_only: bool = False, + fuse_layernorms: bool = False, + quantization_format: str | None = None, +) -> dict[str, list[str]]: + """Fuse modules that share the same input. + + This is a common helper for both LLM and diffusion model fusion. + + Args: + model: The model being processed (for FSDP-aware updates). + input_to_linear: Dict mapping input tensor to list of modules sharing that input. + output_to_layernorm: Dict mapping layernorm output to the layernorm module (optional). + qkv_only: If True, only fuse QKV projection layers (for diffusion models). + fuse_layernorms: If True, also fuse layernorms with pre_quant_scale (for AWQ). + quantization_format: The quantization format of the model. + + Returns: + Dict mapping first module name to list of all fused module names. + """ + fused_linears = {} + fused_count = 0 + + for tensor, modules in input_to_linear.items(): + # Get quantization format for this group of modules + # (must be re-evaluated per group as different modules may have different formats) + group_quant_format = get_quantization_format(modules[0]) if modules else quantization_format + + if len(modules) > 1 and group_quant_format not in [ + QUANTIZATION_FP8, + QUANTIZATION_NONE, + QUANTIZATION_FP8_PB_REAL, + ]: + if qkv_only: + # Filter to only include QKV projection layers (diffusion models) + qkv_modules = [m for m in modules if is_qkv_projection(getattr(m, "name", ""))] + + if len(qkv_modules) > 1: + # Group QKV modules by their parent attention block + qkv_groups: dict[str, list[nn.Module]] = defaultdict(list) + for m in qkv_modules: + group_key = get_qkv_group_key(getattr(m, "name", "")) + qkv_groups[group_key].append(m) + + # Fuse each group separately + for group_key, group_modules in qkv_groups.items(): + if len(group_modules) >= 2: + preprocess_linear_fusion(group_modules, resmooth_only=False) + fused_count += 1 + module_names = [getattr(m, "name", "unknown") for m in group_modules] + print(f" Fused QKV group: {module_names}") + else: + # Fuse all modules that have the same input (LLM models) + with fsdp2_aware_weight_update(model, modules): + preprocess_linear_fusion(modules) + fused_linears[modules[0].name] = [module.name for module in modules] + fused_count += 1 + + # Fuse layernorms (for AWQ) + if ( + fuse_layernorms + and output_to_layernorm is not None + and group_quant_format is not None + and group_quant_format != QUANTIZATION_NONE + and "awq" in group_quant_format + and tensor in output_to_layernorm + ): + with fsdp2_aware_weight_update(model, output_to_layernorm[tensor]): + fuse_prequant_layernorm(output_to_layernorm[tensor], modules) + + if qkv_only: + if fused_count > 0: + print(f"Fused {fused_count} QKV group(s) for unified amax values.") + else: + print("No QKV groups found to fuse.") + + return fused_linears + + +def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): + """Group modules that take the same input and register shared parameters in module.""" + # TODO: Handle DBRX MoE + quantization_format = get_quantization_format(model) + model_type = type(model).__name__.lower() + module_names = set() + + # NVFP4 SVDQuant does not need pre-quant scale fusion (either into previous linear or layernorm) because + # 1) its kernel handles pre-quant scale. + # 2) fusing into previous linear will need to change the lora_up in up_proj which may cause issue in + # the later gate up fusion. + # Fuse pre_quant_scale to the linear weights if possible + if quantization_format is not None and "nvfp4_awq" in quantization_format.lower(): + fuse_prequant_to_linear(model) + + # Pre-process MoE experts + for name, module in model.named_modules(): + module_names.add(name) + + # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts + if is_moe(module) and ( + quantization_format is not QUANTIZATION_NONE + and ("awq" in quantization_format or quantization_format == QUANTIZATION_NVFP4_SVDQUANT) + ): + # update_experts_avg_prequant_scale(module) + grouped_experts = get_experts_list(module, model_type) + for modules in grouped_experts: + with fsdp2_aware_weight_update(model, modules): + preprocess_linear_fusion(modules, resmooth_only=True) + + # Define the dummy forward function for LLM + def llm_dummy_forward(): + fake_input = torch.ones([1, 2], dtype=torch.long).to(model.device) + decoder_fake_input = fake_input + + # Check if this is a VL model that needs special input handling + is_vl_model = is_multimodal_model(model) + + if model_type.startswith("whisper"): + # For Whisper models, we need to pass a fake input with the specific sequence length + from transformers import AutoFeatureExtractor + + feature_extractor = AutoFeatureExtractor.from_pretrained(model.name_or_path) + fake_input = torch.ones( + [1, model.config.num_mel_bins, feature_extractor.nb_max_frames], dtype=model.dtype + ).to(model.device) + + if is_vl_model and "nemotron" in model_type: + # For Nemotron VL models, run optimization on just the language model/decoder. + # This avoids needing pixel_values for the vision encoder. + language_model_lineage = get_language_model_from_vl(model) + + if language_model_lineage is not None: + language_model = language_model_lineage[-1] + print( + f"Running optimization on language model with fake_input shape: {fake_input.shape}" + ) + # Pass use_cache=False to avoid KV cache issues in encoder-decoder models + language_model(fake_input, use_cache=False) + else: + raise ValueError( + f"Cannot extract language_model from Nemotron VL model (type: {model_type}). " + "This is required for requantization/resmoothing optimization. " + "Please ensure the model architecture is supported or file an issue." + ) + elif getattr(model.config, "is_encoder_decoder", False): + # For other encoder-decoder models (non-VL), pass both encoder and decoder input ids + model(fake_input, decoder_input_ids=decoder_fake_input) + elif hasattr(model, "get_dummy_inputs"): + # For speculative decoding models (EAGLE, etc.), use model-provided dummy inputs + model(**model.get_dummy_inputs()) + else: + model(fake_input) + + input_to_linear, output_to_layernorm = collect_shared_input_modules( + model, llm_dummy_forward, collect_layernorms=True + ) + + fused_linears = _fuse_shared_input_modules( + model, + input_to_linear, + output_to_layernorm, + qkv_only=False, + fuse_layernorms=True, + quantization_format=quantization_format, + ) + + # The dummy forward may not be able to activate all the experts. + # Process experts by naming rules like experts.0, experts.1, etc. + for name, modules_fused in fused_linears.items(): + if re.search(r"experts?\.\d+", name): + expert_id = 0 + while True: + new_expert_name = re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", name, count=1) + if new_expert_name in fused_linears: + expert_id += 1 + continue + if new_expert_name not in module_names: + break + + new_expert_modules = [] + for name_fused in modules_fused: + new_expert_name = re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", name_fused) + assert new_expert_name in module_names + new_expert_modules.append(model.get_submodule(new_expert_name)) + + with fsdp2_aware_weight_update(model, new_expert_modules): + preprocess_linear_fusion(new_expert_modules) + + expert_id += 1 + + +def _resolve_export_dtype(model: nn.Module, dtype: torch.dtype | None) -> torch.dtype: + """Return the export dtype, defaulting to the model's own and warning on a mismatch.""" + if dtype is None: + return model.config.torch_dtype + if dtype != model.config.torch_dtype: + warnings.warn( + f"Model's original dtype ({model.config.torch_dtype}) differs from target dtype " + f"({dtype}), which may lead to numerical errors." + ) + return dtype + + +def _prepare_moe_inputs(model: nn.Module, dtype: torch.dtype, is_modelopt_qlora: bool) -> None: + """Handle input quantizers of experts that are not calibrated. + + Each MoE block is dispatched by its experts container to the matching preparation + handler. + """ + prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + for name, sub_module in model.named_modules(): + if is_moe(sub_module) and hasattr(sub_module, "experts"): + handler = PrepareMoEInputsRegistry.match(sub_module.experts) + if handler is None: + # Unsupported MoE model structure + raise NotImplementedError( + f"MoE model with experts type '{type(sub_module.experts).__name__}' is not supported in export." + f"Please file an issue or add support for this model architecture." + ) + handler(name, sub_module, prepare_ctx) + + +def _add_mtp_exclusions(model: nn.Module, quant_config: dict) -> None: + """Add MTP layer prefixes to exclude_modules if they were excluded from quantization. + + This ensures they appear in ``quantization_config["ignore"]`` in ``config.json``. + """ + mtp_layer_prefixes = getattr(model, "_mtp_layer_prefixes", None) + if mtp_layer_prefixes: + exclude_modules = quant_config["quantization"].setdefault("exclude_modules", []) + for prefix in mtp_layer_prefixes: + # Add wildcard pattern to exclude all submodules under this MTP layer + pattern = f"{prefix}*" + if pattern not in exclude_modules: + exclude_modules.append(pattern) + print(f"Adding MTP layer to quantization_config ignore: {pattern}") + + +def _warn_on_unsynced_moe_gate_up(model: nn.Module) -> None: + """Safety net for gate/up weight quantizer amaxes that resmoothing did not reach. + + ``requantize_resmooth_fused_llm_layers`` can miss experts that the dummy forward + never activated, or that use non-standard expert naming. + """ + synced = sync_moe_gate_up_amax(model) + if synced: + warnings.warn( + f"Found {synced} MoE expert gate/up projection pair(s) with mismatched " + f"weight_scale_2 after requantize_resmooth_fused_llm_layers. " + f"This typically means the dummy forward did not activate these experts. " + f"Taking element-wise max of amaxes for serving-engine fusion." + ) + + +def _revert_weight_conversion_noop(model: Any, state_dict: dict) -> dict: + """No-op replacement for transformers' revert_weight_conversion.""" + return state_dict + + +def _try_patch_module(mod_path: str) -> tuple[Any, Any] | None: + """Try to patch revert_weight_conversion in a single module.""" + import importlib + + try: + mod = importlib.import_module(mod_path) + if hasattr(mod, "revert_weight_conversion"): + original = getattr(mod, "revert_weight_conversion") + setattr(mod, "revert_weight_conversion", _revert_weight_conversion_noop) + return (mod, original) + except (ImportError, AttributeError): + pass + return None + + +def _patch_revert_weight_conversion() -> list[tuple[Any, Any]]: + """Patch revert_weight_conversion in transformers to avoid RuntimeError on scalar tensors.""" + patches: list[tuple[Any, Any]] = [] + for mod_path in [ + "transformers.core_model_loading", + "transformers.modeling_utils", + ]: + result = _try_patch_module(mod_path) + if result is not None: + patches.append(result) + return patches + + +def _unpatch_revert_weight_conversion(patches: list[tuple[Any, Any]]) -> None: + """Restore the original revert_weight_conversion functions.""" + for mod, original in patches: + mod.revert_weight_conversion = original + + +def _sanitize_generation_config_for_save(model: torch.nn.Module) -> None: + """Force ``do_sample=True`` when generation_config has ``top_k``/``top_p`` set. + + Newer transformers reject ``do_sample=False`` mixed with sampling attrs in + ``save_pretrained``'s strict validate. + """ + gc = getattr(model, "generation_config", None) + if gc is None: + return + if getattr(gc, "top_k", None) is not None or getattr(gc, "top_p", None) is not None: + gc.do_sample = True diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index acb1968e070..4c1a1beb94c 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -35,9 +35,9 @@ from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import get_unwrapped_name, safe_save +from ..hf_export_prep import collect_shared_input_modules from ..layer_utils import get_experts_list, is_moe from ..quant_utils import get_quantization_format -from ..unified_export_hf import collect_shared_input_modules __all__ = [ "export_hf_vllm_fq_checkpoint", diff --git a/modelopt/torch/export/unified_export_diffusers.py b/modelopt/torch/export/unified_export_diffusers.py index 77400d10216..fab216dc433 100644 --- a/modelopt/torch/export/unified_export_diffusers.py +++ b/modelopt/torch/export/unified_export_diffusers.py @@ -34,14 +34,11 @@ from .convert_hf_config import convert_hf_quant_config_format from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales +from .hf_export_prep import _fuse_shared_input_modules, collect_shared_input_modules from .layer_utils import is_quantlinear from .model_config import QUANTIZATION_NONE from .quant_utils import get_quant_config, get_quantization_format, has_quantized_modules -from .unified_export_hf import ( - _fuse_shared_input_modules, - _process_quantized_modules, - collect_shared_input_modules, -) +from .unified_export_hf import _process_quantized_modules try: import diffusers diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index d8d9386d379..82bbab04818 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -16,12 +16,8 @@ """Code that export quantized Hugging Face models for deployment.""" import json -import re import tempfile import warnings -from builtins import ValueError -from collections import defaultdict -from collections.abc import Callable from pathlib import Path from typing import Any @@ -29,7 +25,7 @@ import torch.nn as nn try: - from .diffusers_utils import get_qkv_group_key, is_diffusers_object, is_qkv_projection + from .diffusers_utils import is_diffusers_object HAS_DIFFUSERS = True except ImportError: @@ -38,14 +34,12 @@ from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from torch.distributed.fsdp import FSDPModule -from modelopt.torch.quantization import set_quantizer_by_cfg_context from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 -from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names +from modelopt.torch.quantization.utils import quantizer_attr_names from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload -from modelopt.torch.utils.dataset_utils import _disable_use_cache from modelopt.torch.utils.distributed import is_fsdp2_model try: @@ -56,12 +50,15 @@ # Importing the built-in handlers installs their entries in the two registries. from . import hf_export_handlers as _hf_export_handlers # noqa: F401 from .convert_hf_config import convert_hf_quant_config_format -from .layer_utils import ( - get_experts_list, - is_layernorm, - is_moe, - is_quantlinear, - sync_moe_gate_up_amax, +from .hf_export_prep import ( + _add_mtp_exclusions, + _patch_revert_weight_conversion, + _prepare_moe_inputs, + _resolve_export_dtype, + _sanitize_generation_config_for_save, + _unpatch_revert_weight_conversion, + _warn_on_unsynced_moe_gate_up, + requantize_resmooth_fused_llm_layers, ) from .model_config import ( QUANTIZATION_FP8, @@ -76,7 +73,7 @@ QUANTIZATION_W4A8_NVFP4_FP8, QUANTIZATION_W4A16_NVFP4, ) -from .model_utils import _reorder_canonical_first, get_language_model_from_vl, is_multimodal_model +from .model_utils import _reorder_canonical_first from .plugins import SpeculativeDecodingExporter, has_spec_opt, sanitize_hf_config_for_deployment from .quant_aware_conversion import ( build_reverse_name_mapper, @@ -84,8 +81,6 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import ( - fuse_prequant_layernorm, - fuse_prequant_to_linear, get_activation_scaling_factor, get_quant_config, get_quantization_format, @@ -94,290 +89,14 @@ get_weight_scaling_factor_2, maybe_transpose_expert_weight_dimensions, postprocess_state_dict, - preprocess_linear_fusion, sync_tied_input_amax, to_quantized_weight, ) -from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry +from .registry import ExportContext, ExportModuleRegistry __all__ = ["export_hf_checkpoint", "export_speculative_decoding"] -def _is_enabled_quantizer(quantizer): - if hasattr(quantizer, "is_enabled") and quantizer.is_enabled: - return True - - if isinstance(quantizer, SequentialQuantizer): - return any(q.is_enabled for q in quantizer) - - return False - - -def collect_shared_input_modules( - model: nn.Module, - dummy_forward_fn: Callable[[], None], - collect_layernorms: bool = False, -) -> tuple[dict, dict | None]: - """Collect modules that share the same input using forward hooks. - - This is a common helper for both LLM and diffusion model fusion. - - Args: - model: The model to analyze. - dummy_forward_fn: A callable that runs a dummy forward pass on the model. - Should be a function that takes no arguments. - collect_layernorms: If True, also collect layernorm output mappings (for AWQ). - - Returns: - A tuple of (input_to_linear, output_to_layernorm). - input_to_linear: Dict mapping input tensor to list of modules sharing that input. - output_to_layernorm: Dict mapping layernorm output to the layernorm module (or None). - """ - input_to_linear: dict = defaultdict(list) - output_to_layernorm: dict | None = defaultdict(lambda: None) if collect_layernorms else None - - def _input_hook(module, input, output): - """Update dictionary with list of all modules that share the same input.""" - if len(input) > 0 and isinstance(input[0], torch.Tensor): - # TODO: Handle DBRX MoE case - input_to_linear[input[0]].append(module) - - def _output_hook(module, input, output): - """Update dictionary with mapping of layernorms and their outputs.""" - if output_to_layernorm is not None and isinstance(output, torch.Tensor): - output_to_layernorm[output] = module - - handles = [] - - # Register hooks on all quantized linear modules (and optionally layernorms) - for name, module in model.named_modules(): - if collect_layernorms and is_layernorm(module): - module.name = name - handle = module.register_forward_hook(_output_hook) - handles.append(handle) - elif is_quantlinear(module) and ( - _is_enabled_quantizer(module.input_quantizer) - or _is_enabled_quantizer(module.weight_quantizer) - ): - module.name = name - handle = module.register_forward_hook(_input_hook) - handles.append(handle) - - if not handles: - return input_to_linear, output_to_layernorm - - # Run dummy forward pass to collect modules sharing same input. - # `_disable_use_cache` keeps the probe forward working on configs that don't - # set `use_cache` (e.g., stepfun-ai/Step-3.5-Flash's Step3p5Config). - try: - with ( - torch.no_grad(), - set_quantizer_by_cfg_context(model, [{"quantizer_name": "*", "enable": False}]), - _disable_use_cache(model), - ): - dummy_forward_fn() - finally: - # Always remove hooks - for handle in handles: - handle.remove() - - return input_to_linear, output_to_layernorm - - -def _fuse_shared_input_modules( - model: nn.Module, - input_to_linear: dict, - output_to_layernorm: dict | None = None, - qkv_only: bool = False, - fuse_layernorms: bool = False, - quantization_format: str | None = None, -) -> dict[str, list[str]]: - """Fuse modules that share the same input. - - This is a common helper for both LLM and diffusion model fusion. - - Args: - model: The model being processed (for FSDP-aware updates). - input_to_linear: Dict mapping input tensor to list of modules sharing that input. - output_to_layernorm: Dict mapping layernorm output to the layernorm module (optional). - qkv_only: If True, only fuse QKV projection layers (for diffusion models). - fuse_layernorms: If True, also fuse layernorms with pre_quant_scale (for AWQ). - quantization_format: The quantization format of the model. - - Returns: - Dict mapping first module name to list of all fused module names. - """ - fused_linears = {} - fused_count = 0 - - for tensor, modules in input_to_linear.items(): - # Get quantization format for this group of modules - # (must be re-evaluated per group as different modules may have different formats) - group_quant_format = get_quantization_format(modules[0]) if modules else quantization_format - - if len(modules) > 1 and group_quant_format not in [ - QUANTIZATION_FP8, - QUANTIZATION_NONE, - QUANTIZATION_FP8_PB_REAL, - ]: - if qkv_only: - # Filter to only include QKV projection layers (diffusion models) - qkv_modules = [m for m in modules if is_qkv_projection(getattr(m, "name", ""))] - - if len(qkv_modules) > 1: - # Group QKV modules by their parent attention block - qkv_groups: dict[str, list[nn.Module]] = defaultdict(list) - for m in qkv_modules: - group_key = get_qkv_group_key(getattr(m, "name", "")) - qkv_groups[group_key].append(m) - - # Fuse each group separately - for group_key, group_modules in qkv_groups.items(): - if len(group_modules) >= 2: - preprocess_linear_fusion(group_modules, resmooth_only=False) - fused_count += 1 - module_names = [getattr(m, "name", "unknown") for m in group_modules] - print(f" Fused QKV group: {module_names}") - else: - # Fuse all modules that have the same input (LLM models) - with fsdp2_aware_weight_update(model, modules): - preprocess_linear_fusion(modules) - fused_linears[modules[0].name] = [module.name for module in modules] - fused_count += 1 - - # Fuse layernorms (for AWQ) - if ( - fuse_layernorms - and output_to_layernorm is not None - and group_quant_format is not None - and group_quant_format != QUANTIZATION_NONE - and "awq" in group_quant_format - and tensor in output_to_layernorm - ): - with fsdp2_aware_weight_update(model, output_to_layernorm[tensor]): - fuse_prequant_layernorm(output_to_layernorm[tensor], modules) - - if qkv_only: - if fused_count > 0: - print(f"Fused {fused_count} QKV group(s) for unified amax values.") - else: - print("No QKV groups found to fuse.") - - return fused_linears - - -def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): - """Group modules that take the same input and register shared parameters in module.""" - # TODO: Handle DBRX MoE - quantization_format = get_quantization_format(model) - model_type = type(model).__name__.lower() - module_names = set() - - # NVFP4 SVDQuant does not need pre-quant scale fusion (either into previous linear or layernorm) because - # 1) its kernel handles pre-quant scale. - # 2) fusing into previous linear will need to change the lora_up in up_proj which may cause issue in - # the later gate up fusion. - # Fuse pre_quant_scale to the linear weights if possible - if quantization_format is not None and "nvfp4_awq" in quantization_format.lower(): - fuse_prequant_to_linear(model) - - # Pre-process MoE experts - for name, module in model.named_modules(): - module_names.add(name) - - # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts - if is_moe(module) and ( - quantization_format is not QUANTIZATION_NONE - and ("awq" in quantization_format or quantization_format == QUANTIZATION_NVFP4_SVDQUANT) - ): - # update_experts_avg_prequant_scale(module) - grouped_experts = get_experts_list(module, model_type) - for modules in grouped_experts: - with fsdp2_aware_weight_update(model, modules): - preprocess_linear_fusion(modules, resmooth_only=True) - - # Define the dummy forward function for LLM - def llm_dummy_forward(): - fake_input = torch.ones([1, 2], dtype=torch.long).to(model.device) - decoder_fake_input = fake_input - - # Check if this is a VL model that needs special input handling - is_vl_model = is_multimodal_model(model) - - if model_type.startswith("whisper"): - # For Whisper models, we need to pass a fake input with the specific sequence length - from transformers import AutoFeatureExtractor - - feature_extractor = AutoFeatureExtractor.from_pretrained(model.name_or_path) - fake_input = torch.ones( - [1, model.config.num_mel_bins, feature_extractor.nb_max_frames], dtype=model.dtype - ).to(model.device) - - if is_vl_model and "nemotron" in model_type: - # For Nemotron VL models, run optimization on just the language model/decoder. - # This avoids needing pixel_values for the vision encoder. - language_model_lineage = get_language_model_from_vl(model) - - if language_model_lineage is not None: - language_model = language_model_lineage[-1] - print( - f"Running optimization on language model with fake_input shape: {fake_input.shape}" - ) - # Pass use_cache=False to avoid KV cache issues in encoder-decoder models - language_model(fake_input, use_cache=False) - else: - raise ValueError( - f"Cannot extract language_model from Nemotron VL model (type: {model_type}). " - "This is required for requantization/resmoothing optimization. " - "Please ensure the model architecture is supported or file an issue." - ) - elif getattr(model.config, "is_encoder_decoder", False): - # For other encoder-decoder models (non-VL), pass both encoder and decoder input ids - model(fake_input, decoder_input_ids=decoder_fake_input) - elif hasattr(model, "get_dummy_inputs"): - # For speculative decoding models (EAGLE, etc.), use model-provided dummy inputs - model(**model.get_dummy_inputs()) - else: - model(fake_input) - - input_to_linear, output_to_layernorm = collect_shared_input_modules( - model, llm_dummy_forward, collect_layernorms=True - ) - - fused_linears = _fuse_shared_input_modules( - model, - input_to_linear, - output_to_layernorm, - qkv_only=False, - fuse_layernorms=True, - quantization_format=quantization_format, - ) - - # The dummy forward may not be able to activate all the experts. - # Process experts by naming rules like experts.0, experts.1, etc. - for name, modules_fused in fused_linears.items(): - if re.search(r"experts?\.\d+", name): - expert_id = 0 - while True: - new_expert_name = re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", name, count=1) - if new_expert_name in fused_linears: - expert_id += 1 - continue - if new_expert_name not in module_names: - break - - new_expert_modules = [] - for name_fused in modules_fused: - new_expert_name = re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", name_fused) - assert new_expert_name in module_names - new_expert_modules.append(model.get_submodule(new_expert_name)) - - with fsdp2_aware_weight_update(model, new_expert_modules): - preprocess_linear_fusion(new_expert_modules) - - expert_id += 1 - - def _compressed_per_block_scale( weight_quantizer: TensorQuantizer, weight: QTensorWrapper ) -> torch.Tensor | None: @@ -699,69 +418,6 @@ def _dispatch_export_handler(name: str, sub_module: nn.Module, ctx: ExportContex handler(name, sub_module, ctx) -def _resolve_export_dtype(model: nn.Module, dtype: torch.dtype | None) -> torch.dtype: - """Return the export dtype, defaulting to the model's own and warning on a mismatch.""" - if dtype is None: - return model.config.torch_dtype - if dtype != model.config.torch_dtype: - warnings.warn( - f"Model's original dtype ({model.config.torch_dtype}) differs from target dtype " - f"({dtype}), which may lead to numerical errors." - ) - return dtype - - -def _prepare_moe_inputs(model: nn.Module, dtype: torch.dtype, is_modelopt_qlora: bool) -> None: - """Handle input quantizers of experts that are not calibrated. - - Each MoE block is dispatched by its experts container to the matching preparation - handler. - """ - prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) - for name, sub_module in model.named_modules(): - if is_moe(sub_module) and hasattr(sub_module, "experts"): - handler = PrepareMoEInputsRegistry.match(sub_module.experts) - if handler is None: - # Unsupported MoE model structure - raise NotImplementedError( - f"MoE model with experts type '{type(sub_module.experts).__name__}' is not supported in export." - f"Please file an issue or add support for this model architecture." - ) - handler(name, sub_module, prepare_ctx) - - -def _add_mtp_exclusions(model: nn.Module, quant_config: dict) -> None: - """Add MTP layer prefixes to exclude_modules if they were excluded from quantization. - - This ensures they appear in ``quantization_config["ignore"]`` in ``config.json``. - """ - mtp_layer_prefixes = getattr(model, "_mtp_layer_prefixes", None) - if mtp_layer_prefixes: - exclude_modules = quant_config["quantization"].setdefault("exclude_modules", []) - for prefix in mtp_layer_prefixes: - # Add wildcard pattern to exclude all submodules under this MTP layer - pattern = f"{prefix}*" - if pattern not in exclude_modules: - exclude_modules.append(pattern) - print(f"Adding MTP layer to quantization_config ignore: {pattern}") - - -def _warn_on_unsynced_moe_gate_up(model: nn.Module) -> None: - """Safety net for gate/up weight quantizer amaxes that resmoothing did not reach. - - ``requantize_resmooth_fused_llm_layers`` can miss experts that the dummy forward - never activated, or that use non-standard expert naming. - """ - synced = sync_moe_gate_up_amax(model) - if synced: - warnings.warn( - f"Found {synced} MoE expert gate/up projection pair(s) with mismatched " - f"weight_scale_2 after requantize_resmooth_fused_llm_layers. " - f"This typically means the dummy forward did not activate these experts. " - f"Taking element-wise max of amaxes for serving-engine fusion." - ) - - def _process_quantized_modules( model: nn.Module, dtype: torch.dtype, @@ -900,56 +556,6 @@ def _export_transformers_checkpoint( # scale buffers that are 0-d scalars, raising RuntimeError ("chunk expects at least a # 1-dimensional tensor"). Confirmed in transformers 5.12.0. # See: transformers/core_model_loading.py, Chunk.convert() -def _revert_weight_conversion_noop(model: Any, state_dict: dict) -> dict: - """No-op replacement for transformers' revert_weight_conversion.""" - return state_dict - - -def _try_patch_module(mod_path: str) -> tuple[Any, Any] | None: - """Try to patch revert_weight_conversion in a single module.""" - import importlib - - try: - mod = importlib.import_module(mod_path) - if hasattr(mod, "revert_weight_conversion"): - original = getattr(mod, "revert_weight_conversion") - setattr(mod, "revert_weight_conversion", _revert_weight_conversion_noop) - return (mod, original) - except (ImportError, AttributeError): - pass - return None - - -def _patch_revert_weight_conversion() -> list[tuple[Any, Any]]: - """Patch revert_weight_conversion in transformers to avoid RuntimeError on scalar tensors.""" - patches: list[tuple[Any, Any]] = [] - for mod_path in [ - "transformers.core_model_loading", - "transformers.modeling_utils", - ]: - result = _try_patch_module(mod_path) - if result is not None: - patches.append(result) - return patches - - -def _unpatch_revert_weight_conversion(patches: list[tuple[Any, Any]]) -> None: - """Restore the original revert_weight_conversion functions.""" - for mod, original in patches: - mod.revert_weight_conversion = original - - -def _sanitize_generation_config_for_save(model: torch.nn.Module) -> None: - """Force ``do_sample=True`` when generation_config has ``top_k``/``top_p`` set. - - Newer transformers reject ``do_sample=False`` mixed with sampling attrs in - ``save_pretrained``'s strict validate. - """ - gc = getattr(model, "generation_config", None) - if gc is None: - return - if getattr(gc, "top_k", None) is not None or getattr(gc, "top_p", None) is not None: - gc.do_sample = True def export_speculative_decoding( diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 5e2d682c770..7a7b2ab1056 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -33,12 +33,8 @@ import torch.nn as nn from safetensors.torch import save_file -from .quant_aware_conversion import build_reverse_name_mapper -from .quant_utils import _postprocess_single_tensor, get_quant_config -from .registry import ExportContext -from .unified_export_hf import ( +from .hf_export_prep import ( _add_mtp_exclusions, - _dispatch_export_handler, _patch_revert_weight_conversion, _prepare_moe_inputs, _resolve_export_dtype, @@ -47,6 +43,10 @@ _warn_on_unsynced_moe_gate_up, requantize_resmooth_fused_llm_layers, ) +from .quant_aware_conversion import build_reverse_name_mapper +from .quant_utils import _postprocess_single_tensor, get_quant_config +from .registry import ExportContext +from .unified_export_hf import _dispatch_export_handler __all__ = ["_export_transformers_checkpoint_streaming"] diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 1897f3d6db3..00af6bf483f 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -22,11 +22,9 @@ from torch.distributed._composable.fsdp import fully_shard import modelopt.torch.quantization as mtq +from modelopt.torch.export.hf_export_prep import requantize_resmooth_fused_llm_layers from modelopt.torch.export.layer_utils import is_quantlinear -from modelopt.torch.export.unified_export_hf import ( - _export_quantized_weight, - requantize_resmooth_fused_llm_layers, -) +from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, patch_fsdp_mp_dtypes From 04a7382b109e728879a970be0566eb557fbdc53c Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:28:33 +0000 Subject: [PATCH 3/6] refactor(export): extract per-module weight export into hf_weight_export _export_quantized_weight is the leaf of the pipeline -- it packs one module's weight and registers the scale buffers beside it -- but it lived in the same file as the exporters that call it, so moe_utils.py and hf_export_handlers.py had to reach it through function-local imports to dodge the cycle. Move it, _compressed_per_block_scale, _dispatch_export_handler and _process_quantized_modules (349 lines) to hf_weight_export.py. Like hf_export_prep, it imports nothing else from the export package. unified_export_hf.py goes 823 -> 474. Thirteen files are repointed rather than shimmed: moe_utils.py, hf_export_handlers.py, the two other exporters, and nine test modules. The patch targets in test_fused_experts.py move too, since patching a name on the old module no longer reaches the callers' bindings. The lazy imports in moe_utils.py and hf_export_handlers.py still point at the new module; hoisting them is the next commit. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/hf_export_handlers.py | 2 +- modelopt/torch/export/hf_weight_export.py | 415 ++++++++++++++++++ modelopt/torch/export/moe_utils.py | 2 +- .../torch/export/unified_export_diffusers.py | 2 +- modelopt/torch/export/unified_export_hf.py | 391 +---------------- .../export/unified_export_hf_streaming.py | 2 +- .../gpu/torch/export/test_export_embedding.py | 2 +- .../torch/export/test_export_weight_gpu.py | 2 +- tests/gpu/torch/export/test_fsdp2_export.py | 2 +- tests/gpu/torch/quantization/test_gptq.py | 2 +- .../export/test_export_compressed_nvfp4.py | 2 +- .../unit/torch/export/test_export_registry.py | 2 +- tests/unit/torch/export/test_export_weight.py | 2 +- .../unit/torch/export/test_offload_export.py | 2 +- .../torch/export/test_unified_export_hf.py | 2 +- .../plugins/test_fused_experts.py | 4 +- 16 files changed, 432 insertions(+), 404 deletions(-) create mode 100644 modelopt/torch/export/hf_weight_export.py diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 800b51daca9..adc16a44708 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -43,7 +43,7 @@ def _export_weight( ) -> None: # Imported lazily to avoid a cycle: unified_export_hf imports this module to # install the built-in handlers while retaining this legacy helper's import path. - from .unified_export_hf import _export_quantized_weight + from .hf_weight_export import _export_quantized_weight _export_quantized_weight(module, ctx.dtype, weight_name, _tied_cache=ctx.tied_cache) diff --git a/modelopt/torch/export/hf_weight_export.py b/modelopt/torch/export/hf_weight_export.py new file mode 100644 index 00000000000..aff47529b1b --- /dev/null +++ b/modelopt/torch/export/hf_weight_export.py @@ -0,0 +1,415 @@ +# 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. + +"""Per-module quantized weight export. + +The leaf of the export pipeline: packing one module's weight into its quantized +representation and registering the scale buffers beside it, plus the registry dispatch +and the whole-model walk that drive it. + +Like :mod:`hf_export_prep`, this depends on nothing else in the export package, so the +exporters and the MoE/handler plugins can import it directly instead of lazily. +""" + +import torch +import torch.nn as nn +from torch.distributed.fsdp import FSDPModule + +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor +from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper +from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 +from modelopt.torch.quantization.utils import quantizer_attr_names + +from .model_config import ( + QUANTIZATION_FP8, + QUANTIZATION_FP8_PB_REAL, + QUANTIZATION_FP8_PC_PT, + QUANTIZATION_MXFP8, + QUANTIZATION_NONE, + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_NVFP4_SVDQUANT, + QUANTIZATION_W4A8_AWQ, + QUANTIZATION_W4A8_NVFP4_FP8, + QUANTIZATION_W4A16_NVFP4, +) +from .quant_utils import ( + get_activation_scaling_factor, + get_quantization_format, + get_weight_block_size, + get_weight_scaling_factor, + get_weight_scaling_factor_2, + maybe_transpose_expert_weight_dimensions, + to_quantized_weight, +) +from .registry import ExportContext, ExportModuleRegistry + + +def _compressed_per_block_scale( + weight_quantizer: TensorQuantizer, weight: QTensorWrapper +) -> torch.Tensor | None: + """Per-block scale captured at compression time, in the modelopt E4M3 layout. + + ``NVFP4QTensor.quantize(..., try_tensorrt=True)`` returns a cutlass-swizzled 1-D uint8 scale + when TensorRT-LLM is available on an FP4-capable device, so normalize it the way + ``NVFP4QTensor.dequantize`` does before it is used as an exported ``weight_scale``. + """ + scale = getattr(weight_quantizer, "_scale", None) + if scale is None or not (scale.dtype == torch.uint8 and scale.ndim == 1): + return scale + try: + from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( + cutlass_fp4_scale_to_modelopt_fp4_scale, + ) + except ImportError as e: + raise ImportError( + "This weight was compressed by TensorRT-LLM, so its NVFP4 block scale is " + "cutlass-swizzled, but tensorrt_llm cannot be imported to convert it for export." + ) from e + return cutlass_fp4_scale_to_modelopt_fp4_scale(scale, weight.metadata["shape"][-2:]) + + +def _export_quantized_weight( + sub_module: nn.Module, + dtype: torch.dtype, + weight_name: str = "weight", + _tied_cache: dict[int, nn.Module] | None = None, +): + """For the given weight attr of the sub_module, export the quantization info of it. + + The export includes converting weight tensor to correct quantized values and quantized dtype, + and registering scaling factors. + + Tied-weight dedup is opt-in via ``_tied_cache``: the setattr below replaces + ``.weight`` with a fresh ``nn.Parameter`` wrapping packed bytes, breaking + any HF-level tie. When the caller passes a ``_tied_cache`` dict (keyed by + the pre-pack ``weight.data_ptr()``), the alias step at the end re-points + ``weight`` / ``weight_scale`` / ``weight_scale_2`` at a previously-processed + module sharing the same source memory so the downstream data_ptr dedup can + collapse them. The cache is owned by the caller (typically + ``_export_transformers_checkpoint``) and scoped to one export invocation; + when ``_tied_cache`` is ``None`` (the default) the alias step is skipped + entirely. Uses memory identity only — no ``_tied_weights_keys`` lookup, + no-op for non-tied modules. + """ + quantization_format = get_quantization_format(sub_module) + if quantization_format == QUANTIZATION_NONE: + return + + block_size = get_weight_block_size(sub_module, weight_name) + quantizer_attrs = quantizer_attr_names(weight_name) + weight: nn.Parameter = getattr(sub_module, weight_name) + + if weight.is_meta: + raise RuntimeError( + f"Weight '{weight_name}' of {type(sub_module).__name__} is a meta tensor during " + "export. If the model was loaded with disk/CPU offload, use export_hf_checkpoint() " + "which dispatches to the streaming writer that materialises weights layer-by-layer." + ) + + # Capture source identity BEFORE any tensor-creating operation below. + # For HF-tied weights this matches across all modules sharing the + # underlying Parameter; the cache lookup at the end of this function + # uses it to detect ties whose Python identity is about to be broken + # by the setattr on `weight_name` further down. + _tied_source_data_ptr = weight.data_ptr() + weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr( + sub_module, quantizer_attrs.weight_quantizer + ) + input_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr( + sub_module, quantizer_attrs.input_quantizer, None + ) + output_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr( + sub_module, quantizer_attrs.output_quantizer, None + ) + + # Already real-quantized weights (``mtq.compress`` / ``hf_ptq --low_memory_mode``) hold packed + # nibbles -- half the logical last dim -- so per-block scales cannot be recomputed from them. + # Use the scale the quantizer captured at compression time instead. + uses_compressed_nvfp4_scale = isinstance(weight, QTensorWrapper) and quantization_format in [ + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_NVFP4_SVDQUANT, + QUANTIZATION_W4A16_NVFP4, + ] + compressed_weight_scale = ( + _compressed_per_block_scale(weight_quantizer, weight) + if uses_compressed_nvfp4_scale + else None + ) + compressed_weight_scale_2 = ( + getattr(weight_quantizer, "_double_scale", None) if uses_compressed_nvfp4_scale else None + ) + use_compressed_scale = ( + compressed_weight_scale is not None and compressed_weight_scale_2 is not None + ) + + if quantization_format == QUANTIZATION_FP8: + # Convert amax to float32 + weight_quantizer._amax = weight_quantizer._amax.to(torch.float32) + + if weight_quantizer._amax.dim() == 1: + # Per-tensor amax + weight_scaling_factor = torch.tensor( + weight_quantizer.amax.item() / weight_quantizer.maxbound + ) + else: + # Per-channel amax + weight_scaling_factor = torch.tensor(weight_quantizer.amax / weight_quantizer.maxbound) + + sub_module.register_buffer( + quantizer_attrs.weight_scale, + weight_scaling_factor, + ) + + if hasattr(input_quantizer, "_amax"): + assert input_quantizer is not None + input_quantizer._amax = input_quantizer._amax.to(torch.float32) + + sub_module.register_buffer( + quantizer_attrs.input_scale, + get_activation_scaling_factor( + sub_module, input_quantizer_name=quantizer_attrs.input_quantizer + ).squeeze(), + ) + + if hasattr(output_quantizer, "_amax"): + assert output_quantizer is not None + output_quantizer._amax = output_quantizer._amax.to(torch.float32) + else: + # Register weight_scale and input_scale + if quantization_format == QUANTIZATION_FP8_PB_REAL: + sub_module.register_buffer( + quantizer_attrs.weight_scale, + weight_quantizer._scale.to(torch.float32), + ) + del weight_quantizer._scale + elif quantization_format == QUANTIZATION_MXFP8: + # MXFP8 uses dynamic block quantization with E8M0 scales (uint8) + weight = getattr(sub_module, weight_name) + e8m0_scale = MXFP8QTensor.get_weights_scaling_factor_from_quantizer( + weight, weight_quantizer + ) + sub_module.register_buffer(quantizer_attrs.weight_scale, e8m0_scale) + if hasattr(weight_quantizer, "_scale") and weight_quantizer._scale is not None: + del weight_quantizer._scale + elif not use_compressed_scale: + sub_module.register_buffer( + quantizer_attrs.weight_scale, get_weight_scaling_factor(sub_module, weight_name) + ) + + if ( + input_quantizer is not None + and "disabled" not in repr(input_quantizer) + and input_quantizer.amax is not None + ): + sub_module.register_buffer( + quantizer_attrs.input_scale, + get_activation_scaling_factor( + sub_module, input_quantizer_name=quantizer_attrs.input_quantizer + ).squeeze(), + ) + + if quantization_format in [ + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_NVFP4_SVDQUANT, + QUANTIZATION_NVFP4, + QUANTIZATION_W4A16_NVFP4, + QUANTIZATION_W4A8_AWQ, + QUANTIZATION_W4A8_NVFP4_FP8, + ]: + # Register weight_scale_2 + sub_module.register_buffer( + quantizer_attrs.weight_scale_2, + get_weight_scaling_factor_2(sub_module, weight_name).squeeze(), + ) + + weight_scale: torch.Tensor | None = getattr(sub_module, quantizer_attrs.weight_scale, None) + weight_scale_2: torch.Tensor | None = getattr(sub_module, quantizer_attrs.weight_scale_2, None) + + # Transpose weight for bmm-style expert quantization (llama4, gpt-oss) + # Check if this is a BMM-style expert weight that needs transposition + is_bmm_expert_weight = weight.dim() == 3 and any( + expert_type in type(sub_module).__name__ + for expert_type in ["Llama4TextExperts", "GptOssExperts"] + ) + # NVFP4StaticQuantizer + BMM-style experts: route through the static-aware + # ``_from_quantizer`` helper so the pinned per-block ``_amax`` (e.g. set by + # the MXFP4->NVFP4 cast to ``6 * 2^k_j``) is used to derive the FP8 + # per-block scale. The plain ``get_weights_scaling_factor`` would ignore + # ``_amax`` and recompute per-block max from the BF16 weight, which + # rebuckets nibbles and loses bit-exactness when ``max_nibble < 6``. + + if quantization_format in [ + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_NVFP4_SVDQUANT, + QUANTIZATION_W4A16_NVFP4, + ]: + # Transpose weight from (num_experts, input_dim, output_dim) to (num_experts, output_dim, input_dim) + # for NVFP4 quantization functions that expect input_dim as the last dimension for block quantization + weight, _ = maybe_transpose_expert_weight_dimensions( + weight, is_bmm_expert_weight=is_bmm_expert_weight + ) + + if use_compressed_scale and weight_scale_2 is not None: + # Dequant is ``nibble * weight_scale * weight_scale_2``; the stored per-block scale is + # normalized against the compression-time global scale, so rescale to keep that product. + # The nibbles cannot be re-quantized here (the high-precision weight is gone), so once + # ``preprocess_linear_fusion`` unifies ``weight_scale_2`` over a fused group the ratio + # below is 1 only for the member owning the group max; the others take one extra E4M3 + # rounding (<= half-ULP, 6.25%). Avoiding that needs a shared scale at compress time. + assert compressed_weight_scale is not None and compressed_weight_scale_2 is not None + device = compressed_weight_scale.device + weight_scale = _cast_per_block_scale_to_fp8( + compressed_weight_scale.float() + * compressed_weight_scale_2.float().to(device) + / weight_scale_2.float().to(device) + ) + elif NVFP4QTensor._is_static_quantizer(weight_quantizer): + weight_scale = NVFP4QTensor.get_weights_scaling_factor_from_quantizer( + weight_quantizer, + weight, + weight_scale_2, + )[0] + else: + weight_scale = NVFP4QTensor.get_weights_scaling_factor( + weight, + block_size=block_size, + weights_scaling_factor_2=weight_scale_2, + )[0] + + quantized_weight = to_quantized_weight( + weight.to(dtype), + weight_scale, + quantization_format, + weight_scale_2, + block_size, + ) + + quantized_weight, weight_scale = maybe_transpose_expert_weight_dimensions( + quantized_weight, weight_scale, is_bmm_expert_weight=is_bmm_expert_weight + ) + elif quantization_format == QUANTIZATION_FP8_PC_PT and is_bmm_expert_weight: + # For FP8_PC_PT with BMM-style experts, transpose only the weight (not weight_scale) + weight, _ = maybe_transpose_expert_weight_dimensions( + weight, is_bmm_expert_weight=is_bmm_expert_weight + ) + + quantized_weight = to_quantized_weight( + weight.to(dtype), + weight_scale, + quantization_format, + weight_scale_2, + block_size, + ) + + # Transpose back to original BMM format + quantized_weight, _ = maybe_transpose_expert_weight_dimensions( + quantized_weight, is_bmm_expert_weight=is_bmm_expert_weight + ) + else: + quantized_weight = to_quantized_weight( + weight.to(dtype), + weight_scale, + quantization_format, + weight_scale_2, + block_size, + ) + + setattr(sub_module, weight_name, nn.Parameter(quantized_weight, requires_grad=False)) + + # Register the corrected weight_scale as a buffer + if weight_scale is not None: + sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale) + + # Tied-weight dedup: if a previously-processed module shared the same + # source weight memory, alias the packed weight + scale buffers so the + # downstream data_ptr dedup in postprocess_state_dict can collapse them. + # input_scale is safe to alias because sync_tied_input_amax (earlier in + # this export) already max-merged the per-side amaxes. Gated on the + # caller-owned _tied_cache so the dedup state is scoped to one export. + if _tied_cache is not None: + _prior = _tied_cache.get(_tied_source_data_ptr) + if _prior is not None and _prior is not sub_module: + if hasattr(_prior, weight_name): + setattr(sub_module, weight_name, getattr(_prior, weight_name)) + for _attr in ( + quantizer_attrs.weight_scale, + quantizer_attrs.weight_scale_2, + quantizer_attrs.input_scale, + ): + if not hasattr(_prior, _attr): + continue + if _attr in sub_module._buffers: + del sub_module._buffers[_attr] + elif hasattr(sub_module, _attr): + delattr(sub_module, _attr) + sub_module.register_buffer(_attr, getattr(_prior, _attr)) + else: + _tied_cache[_tied_source_data_ptr] = sub_module + + torch.cuda.empty_cache() + + +def _dispatch_export_handler(name: str, sub_module: nn.Module, ctx: ExportContext) -> None: + """QLoRA skip, unpack-weight preprocessing, and handler dispatch for one module.""" + if ctx.is_modelopt_qlora and hasattr(sub_module, "base_layer"): + return + # Restore unpacked weight so the export path can read the live quantizer state. + if hasattr(sub_module, "weight_packed") or ( + "QuantFP8Linear" in type(sub_module).__name__ and sub_module.weight.element_size() <= 1 + ): + sub_module.unpack_weight() + handler = ExportModuleRegistry.match(sub_module) + if handler is not None: + handler(name, sub_module, ctx) + + +def _process_quantized_modules( + model: nn.Module, + dtype: torch.dtype, + is_modelopt_qlora: bool = False, +) -> None: + """Process all quantized modules in model, export weights in-place. + + This function iterates through all modules in the model and invokes the first matching + handler in :data:`ExportModuleRegistry`. Modules matching no handler are left untouched. + + Args: + model: The model containing quantized modules. + dtype: The data type for weight conversion. + is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. + If True, modules with base_layer attribute are skipped. + """ + # Per-call tied-weight dedup caches inside the context. Created fresh on + # every invocation so cache state is scoped to one export and cannot leak + # into a later call (see ExportContext). + ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + fsdp_module_to_reshard = None + + for name, sub_module in model.named_modules(): + # Optimization to perform resharding only once per decoder layer to avoid extra communication overhead + if isinstance(sub_module, FSDPModule): + # Every time we encounter a new FSDPModule, the previous decoder layer is fully processed. + # We need to reshard the previous FSDPModule to prevent potential OOM. + # This hack reduces the number of unshard reshard operations, to avoid unnecessary communication. + if fsdp_module_to_reshard is not None: + fsdp_module_to_reshard.reshard() + + fsdp_module_to_reshard = sub_module + + _dispatch_export_handler(name, sub_module, ctx) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 787e173959e..2b177436d23 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -114,7 +114,7 @@ def _export_fused_experts( ``_export_transformers_checkpoint``) and scoped to one export invocation; when ``None`` the corresponding alias step is skipped. """ - from modelopt.torch.export.unified_export_hf import _export_quantized_weight + from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim n = module.num_experts diff --git a/modelopt/torch/export/unified_export_diffusers.py b/modelopt/torch/export/unified_export_diffusers.py index fab216dc433..0e01b738649 100644 --- a/modelopt/torch/export/unified_export_diffusers.py +++ b/modelopt/torch/export/unified_export_diffusers.py @@ -35,10 +35,10 @@ from .convert_hf_config import convert_hf_quant_config_format from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales from .hf_export_prep import _fuse_shared_input_modules, collect_shared_input_modules +from .hf_weight_export import _process_quantized_modules from .layer_utils import is_quantlinear from .model_config import QUANTIZATION_NONE from .quant_utils import get_quant_config, get_quantization_format, has_quantized_modules -from .unified_export_hf import _process_quantized_modules try: import diffusers diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 82bbab04818..0de1e476a74 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -32,13 +32,7 @@ HAS_DIFFUSERS = False from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict -from torch.distributed.fsdp import FSDPModule -from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer -from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor -from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper -from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 -from modelopt.torch.quantization.utils import quantizer_attr_names from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload from modelopt.torch.utils.distributed import is_fsdp2_model @@ -60,19 +54,7 @@ _warn_on_unsynced_moe_gate_up, requantize_resmooth_fused_llm_layers, ) -from .model_config import ( - QUANTIZATION_FP8, - QUANTIZATION_FP8_PB_REAL, - QUANTIZATION_FP8_PC_PT, - QUANTIZATION_MXFP8, - QUANTIZATION_NONE, - QUANTIZATION_NVFP4, - QUANTIZATION_NVFP4_AWQ, - QUANTIZATION_NVFP4_SVDQUANT, - QUANTIZATION_W4A8_AWQ, - QUANTIZATION_W4A8_NVFP4_FP8, - QUANTIZATION_W4A16_NVFP4, -) +from .hf_weight_export import _process_quantized_modules from .model_utils import _reorder_canonical_first from .plugins import SpeculativeDecodingExporter, has_spec_opt, sanitize_hf_config_for_deployment from .quant_aware_conversion import ( @@ -80,380 +62,11 @@ revert_quant_config_names, revert_weight_conversion_quant_aware, ) -from .quant_utils import ( - get_activation_scaling_factor, - get_quant_config, - get_quantization_format, - get_weight_block_size, - get_weight_scaling_factor, - get_weight_scaling_factor_2, - maybe_transpose_expert_weight_dimensions, - postprocess_state_dict, - sync_tied_input_amax, - to_quantized_weight, -) -from .registry import ExportContext, ExportModuleRegistry +from .quant_utils import get_quant_config, postprocess_state_dict, sync_tied_input_amax __all__ = ["export_hf_checkpoint", "export_speculative_decoding"] -def _compressed_per_block_scale( - weight_quantizer: TensorQuantizer, weight: QTensorWrapper -) -> torch.Tensor | None: - """Per-block scale captured at compression time, in the modelopt E4M3 layout. - - ``NVFP4QTensor.quantize(..., try_tensorrt=True)`` returns a cutlass-swizzled 1-D uint8 scale - when TensorRT-LLM is available on an FP4-capable device, so normalize it the way - ``NVFP4QTensor.dequantize`` does before it is used as an exported ``weight_scale``. - """ - scale = getattr(weight_quantizer, "_scale", None) - if scale is None or not (scale.dtype == torch.uint8 and scale.ndim == 1): - return scale - try: - from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( - cutlass_fp4_scale_to_modelopt_fp4_scale, - ) - except ImportError as e: - raise ImportError( - "This weight was compressed by TensorRT-LLM, so its NVFP4 block scale is " - "cutlass-swizzled, but tensorrt_llm cannot be imported to convert it for export." - ) from e - return cutlass_fp4_scale_to_modelopt_fp4_scale(scale, weight.metadata["shape"][-2:]) - - -def _export_quantized_weight( - sub_module: nn.Module, - dtype: torch.dtype, - weight_name: str = "weight", - _tied_cache: dict[int, nn.Module] | None = None, -): - """For the given weight attr of the sub_module, export the quantization info of it. - - The export includes converting weight tensor to correct quantized values and quantized dtype, - and registering scaling factors. - - Tied-weight dedup is opt-in via ``_tied_cache``: the setattr below replaces - ``.weight`` with a fresh ``nn.Parameter`` wrapping packed bytes, breaking - any HF-level tie. When the caller passes a ``_tied_cache`` dict (keyed by - the pre-pack ``weight.data_ptr()``), the alias step at the end re-points - ``weight`` / ``weight_scale`` / ``weight_scale_2`` at a previously-processed - module sharing the same source memory so the downstream data_ptr dedup can - collapse them. The cache is owned by the caller (typically - ``_export_transformers_checkpoint``) and scoped to one export invocation; - when ``_tied_cache`` is ``None`` (the default) the alias step is skipped - entirely. Uses memory identity only — no ``_tied_weights_keys`` lookup, - no-op for non-tied modules. - """ - quantization_format = get_quantization_format(sub_module) - if quantization_format == QUANTIZATION_NONE: - return - - block_size = get_weight_block_size(sub_module, weight_name) - quantizer_attrs = quantizer_attr_names(weight_name) - weight: nn.Parameter = getattr(sub_module, weight_name) - - if weight.is_meta: - raise RuntimeError( - f"Weight '{weight_name}' of {type(sub_module).__name__} is a meta tensor during " - "export. If the model was loaded with disk/CPU offload, use export_hf_checkpoint() " - "which dispatches to the streaming writer that materialises weights layer-by-layer." - ) - - # Capture source identity BEFORE any tensor-creating operation below. - # For HF-tied weights this matches across all modules sharing the - # underlying Parameter; the cache lookup at the end of this function - # uses it to detect ties whose Python identity is about to be broken - # by the setattr on `weight_name` further down. - _tied_source_data_ptr = weight.data_ptr() - weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr( - sub_module, quantizer_attrs.weight_quantizer - ) - input_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr( - sub_module, quantizer_attrs.input_quantizer, None - ) - output_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr( - sub_module, quantizer_attrs.output_quantizer, None - ) - - # Already real-quantized weights (``mtq.compress`` / ``hf_ptq --low_memory_mode``) hold packed - # nibbles -- half the logical last dim -- so per-block scales cannot be recomputed from them. - # Use the scale the quantizer captured at compression time instead. - uses_compressed_nvfp4_scale = isinstance(weight, QTensorWrapper) and quantization_format in [ - QUANTIZATION_NVFP4, - QUANTIZATION_NVFP4_AWQ, - QUANTIZATION_NVFP4_SVDQUANT, - QUANTIZATION_W4A16_NVFP4, - ] - compressed_weight_scale = ( - _compressed_per_block_scale(weight_quantizer, weight) - if uses_compressed_nvfp4_scale - else None - ) - compressed_weight_scale_2 = ( - getattr(weight_quantizer, "_double_scale", None) if uses_compressed_nvfp4_scale else None - ) - use_compressed_scale = ( - compressed_weight_scale is not None and compressed_weight_scale_2 is not None - ) - - if quantization_format == QUANTIZATION_FP8: - # Convert amax to float32 - weight_quantizer._amax = weight_quantizer._amax.to(torch.float32) - - if weight_quantizer._amax.dim() == 1: - # Per-tensor amax - weight_scaling_factor = torch.tensor( - weight_quantizer.amax.item() / weight_quantizer.maxbound - ) - else: - # Per-channel amax - weight_scaling_factor = torch.tensor(weight_quantizer.amax / weight_quantizer.maxbound) - - sub_module.register_buffer( - quantizer_attrs.weight_scale, - weight_scaling_factor, - ) - - if hasattr(input_quantizer, "_amax"): - assert input_quantizer is not None - input_quantizer._amax = input_quantizer._amax.to(torch.float32) - - sub_module.register_buffer( - quantizer_attrs.input_scale, - get_activation_scaling_factor( - sub_module, input_quantizer_name=quantizer_attrs.input_quantizer - ).squeeze(), - ) - - if hasattr(output_quantizer, "_amax"): - assert output_quantizer is not None - output_quantizer._amax = output_quantizer._amax.to(torch.float32) - else: - # Register weight_scale and input_scale - if quantization_format == QUANTIZATION_FP8_PB_REAL: - sub_module.register_buffer( - quantizer_attrs.weight_scale, - weight_quantizer._scale.to(torch.float32), - ) - del weight_quantizer._scale - elif quantization_format == QUANTIZATION_MXFP8: - # MXFP8 uses dynamic block quantization with E8M0 scales (uint8) - weight = getattr(sub_module, weight_name) - e8m0_scale = MXFP8QTensor.get_weights_scaling_factor_from_quantizer( - weight, weight_quantizer - ) - sub_module.register_buffer(quantizer_attrs.weight_scale, e8m0_scale) - if hasattr(weight_quantizer, "_scale") and weight_quantizer._scale is not None: - del weight_quantizer._scale - elif not use_compressed_scale: - sub_module.register_buffer( - quantizer_attrs.weight_scale, get_weight_scaling_factor(sub_module, weight_name) - ) - - if ( - input_quantizer is not None - and "disabled" not in repr(input_quantizer) - and input_quantizer.amax is not None - ): - sub_module.register_buffer( - quantizer_attrs.input_scale, - get_activation_scaling_factor( - sub_module, input_quantizer_name=quantizer_attrs.input_quantizer - ).squeeze(), - ) - - if quantization_format in [ - QUANTIZATION_NVFP4_AWQ, - QUANTIZATION_NVFP4_SVDQUANT, - QUANTIZATION_NVFP4, - QUANTIZATION_W4A16_NVFP4, - QUANTIZATION_W4A8_AWQ, - QUANTIZATION_W4A8_NVFP4_FP8, - ]: - # Register weight_scale_2 - sub_module.register_buffer( - quantizer_attrs.weight_scale_2, - get_weight_scaling_factor_2(sub_module, weight_name).squeeze(), - ) - - weight_scale: torch.Tensor | None = getattr(sub_module, quantizer_attrs.weight_scale, None) - weight_scale_2: torch.Tensor | None = getattr(sub_module, quantizer_attrs.weight_scale_2, None) - - # Transpose weight for bmm-style expert quantization (llama4, gpt-oss) - # Check if this is a BMM-style expert weight that needs transposition - is_bmm_expert_weight = weight.dim() == 3 and any( - expert_type in type(sub_module).__name__ - for expert_type in ["Llama4TextExperts", "GptOssExperts"] - ) - # NVFP4StaticQuantizer + BMM-style experts: route through the static-aware - # ``_from_quantizer`` helper so the pinned per-block ``_amax`` (e.g. set by - # the MXFP4->NVFP4 cast to ``6 * 2^k_j``) is used to derive the FP8 - # per-block scale. The plain ``get_weights_scaling_factor`` would ignore - # ``_amax`` and recompute per-block max from the BF16 weight, which - # rebuckets nibbles and loses bit-exactness when ``max_nibble < 6``. - - if quantization_format in [ - QUANTIZATION_NVFP4, - QUANTIZATION_NVFP4_AWQ, - QUANTIZATION_NVFP4_SVDQUANT, - QUANTIZATION_W4A16_NVFP4, - ]: - # Transpose weight from (num_experts, input_dim, output_dim) to (num_experts, output_dim, input_dim) - # for NVFP4 quantization functions that expect input_dim as the last dimension for block quantization - weight, _ = maybe_transpose_expert_weight_dimensions( - weight, is_bmm_expert_weight=is_bmm_expert_weight - ) - - if use_compressed_scale and weight_scale_2 is not None: - # Dequant is ``nibble * weight_scale * weight_scale_2``; the stored per-block scale is - # normalized against the compression-time global scale, so rescale to keep that product. - # The nibbles cannot be re-quantized here (the high-precision weight is gone), so once - # ``preprocess_linear_fusion`` unifies ``weight_scale_2`` over a fused group the ratio - # below is 1 only for the member owning the group max; the others take one extra E4M3 - # rounding (<= half-ULP, 6.25%). Avoiding that needs a shared scale at compress time. - assert compressed_weight_scale is not None and compressed_weight_scale_2 is not None - device = compressed_weight_scale.device - weight_scale = _cast_per_block_scale_to_fp8( - compressed_weight_scale.float() - * compressed_weight_scale_2.float().to(device) - / weight_scale_2.float().to(device) - ) - elif NVFP4QTensor._is_static_quantizer(weight_quantizer): - weight_scale = NVFP4QTensor.get_weights_scaling_factor_from_quantizer( - weight_quantizer, - weight, - weight_scale_2, - )[0] - else: - weight_scale = NVFP4QTensor.get_weights_scaling_factor( - weight, - block_size=block_size, - weights_scaling_factor_2=weight_scale_2, - )[0] - - quantized_weight = to_quantized_weight( - weight.to(dtype), - weight_scale, - quantization_format, - weight_scale_2, - block_size, - ) - - quantized_weight, weight_scale = maybe_transpose_expert_weight_dimensions( - quantized_weight, weight_scale, is_bmm_expert_weight=is_bmm_expert_weight - ) - elif quantization_format == QUANTIZATION_FP8_PC_PT and is_bmm_expert_weight: - # For FP8_PC_PT with BMM-style experts, transpose only the weight (not weight_scale) - weight, _ = maybe_transpose_expert_weight_dimensions( - weight, is_bmm_expert_weight=is_bmm_expert_weight - ) - - quantized_weight = to_quantized_weight( - weight.to(dtype), - weight_scale, - quantization_format, - weight_scale_2, - block_size, - ) - - # Transpose back to original BMM format - quantized_weight, _ = maybe_transpose_expert_weight_dimensions( - quantized_weight, is_bmm_expert_weight=is_bmm_expert_weight - ) - else: - quantized_weight = to_quantized_weight( - weight.to(dtype), - weight_scale, - quantization_format, - weight_scale_2, - block_size, - ) - - setattr(sub_module, weight_name, nn.Parameter(quantized_weight, requires_grad=False)) - - # Register the corrected weight_scale as a buffer - if weight_scale is not None: - sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale) - - # Tied-weight dedup: if a previously-processed module shared the same - # source weight memory, alias the packed weight + scale buffers so the - # downstream data_ptr dedup in postprocess_state_dict can collapse them. - # input_scale is safe to alias because sync_tied_input_amax (earlier in - # this export) already max-merged the per-side amaxes. Gated on the - # caller-owned _tied_cache so the dedup state is scoped to one export. - if _tied_cache is not None: - _prior = _tied_cache.get(_tied_source_data_ptr) - if _prior is not None and _prior is not sub_module: - if hasattr(_prior, weight_name): - setattr(sub_module, weight_name, getattr(_prior, weight_name)) - for _attr in ( - quantizer_attrs.weight_scale, - quantizer_attrs.weight_scale_2, - quantizer_attrs.input_scale, - ): - if not hasattr(_prior, _attr): - continue - if _attr in sub_module._buffers: - del sub_module._buffers[_attr] - elif hasattr(sub_module, _attr): - delattr(sub_module, _attr) - sub_module.register_buffer(_attr, getattr(_prior, _attr)) - else: - _tied_cache[_tied_source_data_ptr] = sub_module - - torch.cuda.empty_cache() - - -def _dispatch_export_handler(name: str, sub_module: nn.Module, ctx: ExportContext) -> None: - """QLoRA skip, unpack-weight preprocessing, and handler dispatch for one module.""" - if ctx.is_modelopt_qlora and hasattr(sub_module, "base_layer"): - return - # Restore unpacked weight so the export path can read the live quantizer state. - if hasattr(sub_module, "weight_packed") or ( - "QuantFP8Linear" in type(sub_module).__name__ and sub_module.weight.element_size() <= 1 - ): - sub_module.unpack_weight() - handler = ExportModuleRegistry.match(sub_module) - if handler is not None: - handler(name, sub_module, ctx) - - -def _process_quantized_modules( - model: nn.Module, - dtype: torch.dtype, - is_modelopt_qlora: bool = False, -) -> None: - """Process all quantized modules in model, export weights in-place. - - This function iterates through all modules in the model and invokes the first matching - handler in :data:`ExportModuleRegistry`. Modules matching no handler are left untouched. - - Args: - model: The model containing quantized modules. - dtype: The data type for weight conversion. - is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. - If True, modules with base_layer attribute are skipped. - """ - # Per-call tied-weight dedup caches inside the context. Created fresh on - # every invocation so cache state is scoped to one export and cannot leak - # into a later call (see ExportContext). - ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) - fsdp_module_to_reshard = None - - for name, sub_module in model.named_modules(): - # Optimization to perform resharding only once per decoder layer to avoid extra communication overhead - if isinstance(sub_module, FSDPModule): - # Every time we encounter a new FSDPModule, the previous decoder layer is fully processed. - # We need to reshard the previous FSDPModule to prevent potential OOM. - # This hack reduces the number of unshard reshard operations, to avoid unnecessary communication. - if fsdp_module_to_reshard is not None: - fsdp_module_to_reshard.reshard() - - fsdp_module_to_reshard = sub_module - - _dispatch_export_handler(name, sub_module, ctx) - - def _export_transformers_checkpoint( model: nn.Module, dtype: torch.dtype | None = None, diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 7a7b2ab1056..1d7a43a0866 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -43,10 +43,10 @@ _warn_on_unsynced_moe_gate_up, requantize_resmooth_fused_llm_layers, ) +from .hf_weight_export import _dispatch_export_handler from .quant_aware_conversion import build_reverse_name_mapper from .quant_utils import _postprocess_single_tensor, get_quant_config from .registry import ExportContext -from .unified_export_hf import _dispatch_export_handler __all__ = ["_export_transformers_checkpoint_streaming"] diff --git a/tests/gpu/torch/export/test_export_embedding.py b/tests/gpu/torch/export/test_export_embedding.py index 1795bd87aaa..6113344e5df 100644 --- a/tests/gpu/torch/export/test_export_embedding.py +++ b/tests/gpu/torch/export/test_export_embedding.py @@ -25,7 +25,7 @@ import torch.nn as nn import modelopt.torch.quantization as mtq -from modelopt.torch.export.unified_export_hf import _process_quantized_modules +from modelopt.torch.export.hf_weight_export import _process_quantized_modules from modelopt.torch.quantization.utils import quantizer_attr_names VOCAB_SIZE = 16 diff --git a/tests/gpu/torch/export/test_export_weight_gpu.py b/tests/gpu/torch/export/test_export_weight_gpu.py index 9db2b51114b..87002e10815 100644 --- a/tests/gpu/torch/export/test_export_weight_gpu.py +++ b/tests/gpu/torch/export/test_export_weight_gpu.py @@ -23,8 +23,8 @@ from torch.nn import init import modelopt.torch.quantization as mtq +from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.export.quant_utils import postprocess_state_dict -from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.nn.modules.quant_module import QuantModule, QuantModuleRegistry from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer from modelopt.torch.quantization.tensor_quant import QUANT_DESC_8BIT_PER_TENSOR diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 00af6bf483f..c2c87077503 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -23,8 +23,8 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.hf_export_prep import requantize_resmooth_fused_llm_layers +from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.export.layer_utils import is_quantlinear -from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, patch_fsdp_mp_dtypes diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 0a1849544b8..5b8937ea3a8 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -22,7 +22,7 @@ from conftest import requires_triton import modelopt.torch.quantization as mtq -from modelopt.torch.export.unified_export_hf import _export_quantized_weight +from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.quantization.model_calib import gptq from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor from modelopt.torch.quantization.utils.calib_utils import ( diff --git a/tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py b/tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py index 878b0320450..0ccc742d2b8 100644 --- a/tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py +++ b/tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py @@ -20,7 +20,7 @@ from _test_utils.torch.export.utils import ToyModel import modelopt.torch.quantization as mtq -from modelopt.torch.export.unified_export_hf import _export_quantized_weight +from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.quantization.backends.utils import fp4_compatible from modelopt.torch.quantization.utils import quantizer_attr_names diff --git a/tests/unit/torch/export/test_export_registry.py b/tests/unit/torch/export/test_export_registry.py index 67647f9c31f..9748907a0ea 100644 --- a/tests/unit/torch/export/test_export_registry.py +++ b/tests/unit/torch/export/test_export_registry.py @@ -33,13 +33,13 @@ _prepare_fused_experts, _prepare_iterable_experts, ) +from modelopt.torch.export.hf_weight_export import _process_quantized_modules from modelopt.torch.export.registry import ( ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry, _ExportHandlerRegistryCls, ) -from modelopt.torch.export.unified_export_hf import _process_quantized_modules class _Experts(nn.Module): diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 6fc17d982e8..5675b2b9774 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -20,7 +20,7 @@ from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config import modelopt.torch.quantization as mtq -from modelopt.torch.export.unified_export_hf import ( +from modelopt.torch.export.hf_weight_export import ( _export_quantized_weight, _process_quantized_modules, ) diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index a8847b2ba32..99eebf03d9f 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -31,10 +31,10 @@ pytest.skip("accelerate not available", allow_module_level=True) import modelopt.torch.quantization as mtq +from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.export.model_config import KV_CACHE_FP8 from modelopt.torch.export.quant_utils import _postprocess_single_tensor from modelopt.torch.export.registry import ExportContext -from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.export.unified_export_hf_streaming import ( _parse_shard_size, _StreamingShardWriter, diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 118331ce3d9..0aa0b92326e 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -24,12 +24,12 @@ ) import modelopt.torch.quantization as mtq +from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.export.model_utils import ( _collect_canonical_tied_patterns, _reorder_canonical_first, ) from modelopt.torch.export.quant_utils import fuse_prequant_layernorm, sync_tied_input_amax -from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.nn import TensorQuantizer diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index c435b3698be..978f4b29e7b 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -512,7 +512,7 @@ def _spy_export(wrapper, dtype, **_kwargs): return monkeypatch.setattr( - "modelopt.torch.export.unified_export_hf._export_quantized_weight", + "modelopt.torch.export.hf_weight_export._export_quantized_weight", _spy_export, ) @@ -601,7 +601,7 @@ def _spy_export(wrapper, dtype, **_kwargs): return monkeypatch.setattr( - "modelopt.torch.export.unified_export_hf._export_quantized_weight", + "modelopt.torch.export.hf_weight_export._export_quantized_weight", _spy_export, ) From 04c59043965288da4671fd8a5aa8f7ac5eab8697 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:46:55 +0000 Subject: [PATCH 4/6] refactor(export): drop the lazy imports the layering made unnecessary With preparation and weight packing in their own modules, the export package is a DAG: hf_export_prep, hf_weight_export -> (nothing in the package) unified_export_hf_streaming -> prep unified_export_diffusers -> prep, weight unified_export_hf -> the three exporters, prep, weight so the four function-local imports that existed only to dodge a cycle become ordinary module-scope ones: - moe_utils.py and hf_export_handlers.py reach _export_quantized_weight directly. These predate this work -- they were dodging the cycle through unified_export_hf. - export_hf_checkpoint imports both the diffusers and streaming exporters at module scope. The streaming one was added in #2008 with a comment saying it could go once the shared helpers moved; this is that. Verified by importing each of the eight modules first, and the package. One test consequence, since hoisting changes name binding: the spies in test_fused_experts.py patched _export_quantized_weight on the module that defines it, which worked while moe_utils imported it lazily. Now that moe_utils holds a module-scope reference, the patch has to target moe_utils._export_quantized_weight instead. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/hf_export_handlers.py | 2 +- modelopt/torch/export/moe_utils.py | 3 ++- modelopt/torch/export/unified_export_diffusers.py | 4 ++-- modelopt/torch/export/unified_export_hf.py | 11 ++--------- modelopt/torch/export/unified_export_hf_streaming.py | 6 +++--- .../torch/quantization/plugins/test_fused_experts.py | 4 ++-- 6 files changed, 12 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index adc16a44708..86c3d5647f4 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -22,6 +22,7 @@ from modelopt.torch.quantization.utils import fsdp2_aware_weight_update +from .hf_weight_export import _export_quantized_weight from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax from .model_config import QUANTIZATION_NONE from .moe_utils import _export_fused_experts @@ -43,7 +44,6 @@ def _export_weight( ) -> None: # Imported lazily to avoid a cycle: unified_export_hf imports this module to # install the built-in handlers while retaining this legacy helper's import path. - from .hf_weight_export import _export_quantized_weight _export_quantized_weight(module, ctx.dtype, weight_name, _tied_cache=ctx.tied_cache) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 2b177436d23..5cebf7c0842 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -22,6 +22,8 @@ import torch import torch.nn as nn +from .hf_weight_export import _export_quantized_weight + def _alias_per_expert_subtree_from_prior(module: nn.Module, prior: nn.Module, n: int) -> None: """Build per-expert subtree on ``module`` by aliasing ``prior``'s packed buffers. @@ -114,7 +116,6 @@ def _export_fused_experts( ``_export_transformers_checkpoint``) and scoped to one export invocation; when ``None`` the corresponding alias step is skipped. """ - from modelopt.torch.export.hf_weight_export import _export_quantized_weight from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim n = module.num_experts diff --git a/modelopt/torch/export/unified_export_diffusers.py b/modelopt/torch/export/unified_export_diffusers.py index 0e01b738649..03fee771eff 100644 --- a/modelopt/torch/export/unified_export_diffusers.py +++ b/modelopt/torch/export/unified_export_diffusers.py @@ -17,8 +17,8 @@ Split out of :mod:`unified_export_hf`, which it shares almost nothing with: the transformers and diffusers paths meet only at the dispatch in ``export_hf_checkpoint``. -That dispatch imports :func:`_export_diffusers_checkpoint` lazily, since this module -imports the shared module-walking helpers back from ``unified_export_hf``. +What they do share -- model preparation and per-module weight packing -- lives in +:mod:`hf_export_prep` and :mod:`hf_weight_export`. """ import json diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 0de1e476a74..4fba1f842a1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -63,6 +63,8 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import get_quant_config, postprocess_state_dict, sync_tied_input_amax +from .unified_export_diffusers import _export_diffusers_checkpoint +from .unified_export_hf_streaming import _export_transformers_checkpoint_streaming __all__ = ["export_hf_checkpoint", "export_speculative_decoding"] @@ -256,11 +258,6 @@ def export_hf_checkpoint( if HAS_DIFFUSERS: is_diffusers_obj = is_diffusers_object(model) if is_diffusers_obj: - # Imported here rather than at module scope: the diffusers exporter imports the - # shared module-walking helpers from this module, so a top-level import would be - # circular. The cycle goes away once those helpers move to their own modules. - from .unified_export_diffusers import _export_diffusers_checkpoint - _export_diffusers_checkpoint( model, dtype, @@ -283,10 +280,6 @@ def export_hf_checkpoint( try: if _offloaded: - # Imported here rather than at module scope: the streaming exporter imports the - # shared prep helpers from this module, so a top-level import would be circular. - from .unified_export_hf_streaming import _export_transformers_checkpoint_streaming - if save_modelopt_state: warnings.warn( "save_modelopt_state=True is not supported in the streaming offload export " diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 1d7a43a0866..6d89065d5b5 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -16,9 +16,9 @@ """Streaming HF checkpoint export for disk/CPU-offloaded models. Kept apart from :mod:`unified_export_hf` so the resident exporter cannot drift back into -being offload-aware: the only edge between them is the dispatch in -``export_hf_checkpoint``, which imports :func:`_export_transformers_checkpoint_streaming` -lazily to keep the dependency acyclic. +being offload-aware. Shared work lives in :mod:`hf_export_prep` and +:mod:`hf_weight_export`, so the only edge between the two exporters is the dispatch in +``export_hf_checkpoint``. """ import contextlib diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 978f4b29e7b..01f6a1eeca8 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -512,7 +512,7 @@ def _spy_export(wrapper, dtype, **_kwargs): return monkeypatch.setattr( - "modelopt.torch.export.hf_weight_export._export_quantized_weight", + "modelopt.torch.export.moe_utils._export_quantized_weight", _spy_export, ) @@ -601,7 +601,7 @@ def _spy_export(wrapper, dtype, **_kwargs): return monkeypatch.setattr( - "modelopt.torch.export.hf_weight_export._export_quantized_weight", + "modelopt.torch.export.moe_utils._export_quantized_weight", _spy_export, ) From ee7f050b888ef7f61dbf342989e5f20a53849328 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:50:08 +0000 Subject: [PATCH 5/6] fix(export): correct the diffusers probe and stale pointers from the split Review findings on the module split. HAS_DIFFUSERS was the real one. Replacing the `import diffusers` probe with an import from .diffusers_utils changed behavior: that module catches its own diffusers ImportError and still imports cleanly, so the except branch could never fire and the flag was unconditionally True. Verified: without diffusers it read True here while unified_export_diffusers read False -- two identically named flags disagreeing. Both now read diffusers_utils._HAS_DIFFUSERS, so there is one probe. unified_export_diffusers keeps a use-site `import diffusers` for the one place it needs __version__. Also from the split: - hf_export_prep wrapped the QKV helpers in an `except ImportError` that could not fire, whose None fallback would have turned a missing dependency into `TypeError: 'NoneType' object is not callable` at the call site. Removed. - Both new module docstrings claimed to depend on nothing else in the export package; each imports several leaf modules. They now state the real invariant: leaf helpers only, never an exporter. - hf_export_handlers kept the comment explaining a lazy import that commit 04c5904396 deleted, describing a cycle that no longer exists. - registry.py, model_utils.py and the ptq skill reference still pointed at unified_export_hf.py for code that moved. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../ptq/references/unsupported-models.md | 2 +- modelopt/torch/export/hf_export_handlers.py | 3 -- modelopt/torch/export/hf_export_prep.py | 11 +++---- modelopt/torch/export/hf_weight_export.py | 5 +-- modelopt/torch/export/model_utils.py | 4 +-- modelopt/torch/export/registry.py | 2 +- .../torch/export/unified_export_diffusers.py | 31 +++++++++---------- modelopt/torch/export/unified_export_hf.py | 13 +++----- 8 files changed, 30 insertions(+), 41 deletions(-) diff --git a/.agents/skills/ptq/references/unsupported-models.md b/.agents/skills/ptq/references/unsupported-models.md index 361669f70c2..d6ef420446a 100644 --- a/.agents/skills/ptq/references/unsupported-models.md +++ b/.agents/skills/ptq/references/unsupported-models.md @@ -224,7 +224,7 @@ quant_cfg["quant_cfg"]["*vision*"] = {"enable": False} quant_cfg["quant_cfg"]["*multi_modal_projector*"] = {"enable": False} ``` -**Known VLM export issue**: The export step (`requantize_resmooth_fused_llm_layers` in `unified_export_hf.py`) may try to run a dummy forward pass on the full VLM instead of the language model backbone. This currently only handles Nemotron VLMs. If hit, patch the export to use `is_multimodal_model()` for the VLM check instead of model-specific string matching. +**Known VLM export issue**: The export step (`requantize_resmooth_fused_llm_layers` in `hf_export_prep.py`) may try to run a dummy forward pass on the full VLM instead of the language model backbone. This currently only handles Nemotron VLMs. If hit, patch the export to use `is_multimodal_model()` for the VLM check instead of model-specific string matching. ## Pattern 5: FP8 Checkpoint Handling diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 86c3d5647f4..e9fd815d22d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -42,9 +42,6 @@ def _export_weight( ctx: ExportContext, weight_name: str = "weight", ) -> None: - # Imported lazily to avoid a cycle: unified_export_hf imports this module to - # install the built-in handlers while retaining this legacy helper's import path. - _export_quantized_weight(module, ctx.dtype, weight_name, _tied_cache=ctx.tied_cache) diff --git a/modelopt/torch/export/hf_export_prep.py b/modelopt/torch/export/hf_export_prep.py index b5e59450738..fabc9d1c6c1 100644 --- a/modelopt/torch/export/hf_export_prep.py +++ b/modelopt/torch/export/hf_export_prep.py @@ -19,8 +19,9 @@ MoE input-quantizer preparation, resmoothing and shared-input fusion, quant-config adjustments, and the transformers patches needed while writing artifacts. -This module depends on nothing else in the export package, which is what lets the -three exporters import it without a cycle. +It imports only leaf helpers (layer_utils, model_config, model_utils, quant_utils, +registry, diffusers_utils) and never an exporter, which is what lets all three exporters +import it without a cycle. """ import re @@ -37,6 +38,7 @@ from modelopt.torch.quantization.utils import fsdp2_aware_weight_update from modelopt.torch.utils.dataset_utils import _disable_use_cache +from .diffusers_utils import get_qkv_group_key, is_qkv_projection from .layer_utils import ( get_experts_list, is_layernorm, @@ -59,11 +61,6 @@ ) from .registry import ExportContext, PrepareMoEInputsRegistry -try: - from .diffusers_utils import get_qkv_group_key, is_qkv_projection -except ImportError: # diffusers not installed; QKV fusion is diffusers-only - get_qkv_group_key = is_qkv_projection = None - def _is_enabled_quantizer(quantizer): if hasattr(quantizer, "is_enabled") and quantizer.is_enabled: diff --git a/modelopt/torch/export/hf_weight_export.py b/modelopt/torch/export/hf_weight_export.py index aff47529b1b..0ce7b1d4ed8 100644 --- a/modelopt/torch/export/hf_weight_export.py +++ b/modelopt/torch/export/hf_weight_export.py @@ -19,8 +19,9 @@ representation and registering the scale buffers beside it, plus the registry dispatch and the whole-model walk that drive it. -Like :mod:`hf_export_prep`, this depends on nothing else in the export package, so the -exporters and the MoE/handler plugins can import it directly instead of lazily. +Like :mod:`hf_export_prep`, this imports only leaf helpers (model_config, quant_utils, +registry) and never an exporter, so the exporters and the MoE/handler plugins can import +it directly instead of lazily. """ import torch diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py index 307ea9aac51..a6e4bc43b34 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -210,8 +210,8 @@ def _reorder_canonical_first(state_dict: dict, model: nn.Module) -> dict: name to scope the reorder to DiffusionGemma; other tied encoder-decoder models that ship dict-style ``_tied_weights_keys`` can be added to the allowlist here. Mirrors the ``model_type`` - dispatch used for the Whisper and Nemotron-VL branches elsewhere - in ``unified_export_hf.py``. + dispatch used for the Whisper and Nemotron-VL branches in + ``hf_export_prep.py``. """ model_type = type(model).__name__.lower() if "diffusiongemma" not in model_type and "diffusion_gemma" not in model_type: diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 6f4d4be0a88..50289077f4d 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -23,7 +23,7 @@ Preparation and export use separate registries because they have independent matching precedence. Registering a handler for a new module type replaces what previously required -editing if/elif chains inside ``unified_export_hf.py``. +editing if/elif chains inside ``hf_weight_export.py``. """ from collections.abc import Callable diff --git a/modelopt/torch/export/unified_export_diffusers.py b/modelopt/torch/export/unified_export_diffusers.py index 03fee771eff..24c65c7e533 100644 --- a/modelopt/torch/export/unified_export_diffusers.py +++ b/modelopt/torch/export/unified_export_diffusers.py @@ -33,29 +33,24 @@ from safetensors.torch import save_file from .convert_hf_config import convert_hf_quant_config_format -from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales +from .diffusers_utils import _HAS_DIFFUSERS as HAS_DIFFUSERS +from .diffusers_utils import ( + build_layerwise_quant_metadata, + generate_diffusion_dummy_forward_fn, + get_diffusion_components, + get_diffusion_model_type, + hide_quantizers_from_state_dict, + infer_dtype_from_model, + merge_diffusion_checkpoint, + pad_nvfp4_weights, + swizzle_nvfp4_scales, +) from .hf_export_prep import _fuse_shared_input_modules, collect_shared_input_modules from .hf_weight_export import _process_quantized_modules from .layer_utils import is_quantlinear from .model_config import QUANTIZATION_NONE from .quant_utils import get_quant_config, get_quantization_format, has_quantized_modules -try: - import diffusers - - from .diffusers_utils import ( - generate_diffusion_dummy_forward_fn, - get_diffusion_components, - get_diffusion_model_type, - hide_quantizers_from_state_dict, - infer_dtype_from_model, - merge_diffusion_checkpoint, - ) - - HAS_DIFFUSERS = True -except ImportError: - HAS_DIFFUSERS = False - try: from modelopt.torch.sparsity.attention_sparsity.conversion import export_sparse_attention_config except ImportError: @@ -559,6 +554,8 @@ def _export_diffusers_checkpoint( # Last resort: synthesize a minimal model_index.json from exported components. if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None: + import diffusers # only reachable on the diffusers path + model_index = { "_class_name": type(pipe).__name__, "_diffusers_version": diffusers.__version__, diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 4fba1f842a1..1a6ed2d5020 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -23,19 +23,16 @@ import torch import torch.nn as nn - -try: - from .diffusers_utils import is_diffusers_object - - HAS_DIFFUSERS = True -except ImportError: - HAS_DIFFUSERS = False - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload from modelopt.torch.utils.distributed import is_fsdp2_model +# _HAS_DIFFUSERS is diffusers_utils' own probe; re-deriving it here would drift, since +# that module imports cleanly whether or not diffusers is installed. +from .diffusers_utils import _HAS_DIFFUSERS as HAS_DIFFUSERS +from .diffusers_utils import is_diffusers_object + try: from modelopt.torch.sparsity.attention_sparsity.conversion import export_sparse_attention_config except ImportError: From c927f89f86cc1b692a0415d789f137e05a45af3f Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:14:21 +0000 Subject: [PATCH 6/6] docs(export): reattach the revert_weight_conversion TODO to the code it documents Commit b3037f8910 moved _revert_weight_conversion_noop and _patch_revert_weight_conversion to hf_export_prep.py but left the TODO explaining them behind in unified_export_hf.py, where it dangled between _export_transformers_checkpoint and export_speculative_decoding -- two functions it has nothing to do with. That note is the only record of the transformers 5.12.0 0-d-scalar bug and the condition for dropping the workaround, so detached it left the patch helpers with no rationale and pointed anyone revisiting them at the wrong file. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/hf_export_prep.py | 5 +++++ modelopt/torch/export/unified_export_hf.py | 7 ------- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/export/hf_export_prep.py b/modelopt/torch/export/hf_export_prep.py index fabc9d1c6c1..28bc7c1b243 100644 --- a/modelopt/torch/export/hf_export_prep.py +++ b/modelopt/torch/export/hf_export_prep.py @@ -400,6 +400,11 @@ def _warn_on_unsynced_moe_gate_up(model: nn.Module) -> None: ) +# TODO: Remove this workaround once HuggingFace fixes revert_weight_conversion to handle +# scalar (0-d) tensors. transformers' Chunk.convert() calls torch.chunk() on quantization +# scale buffers that are 0-d scalars, raising RuntimeError ("chunk expects at least a +# 1-dimensional tensor"). Confirmed in transformers 5.12.0. +# See: transformers/core_model_loading.py, Chunk.convert() def _revert_weight_conversion_noop(model: Any, state_dict: dict) -> dict: """No-op replacement for transformers' revert_weight_conversion.""" return state_dict diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 1a6ed2d5020..56dc6b7cc98 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -163,13 +163,6 @@ def _export_transformers_checkpoint( return quantized_state_dict, quant_config -# TODO: Remove this workaround once HuggingFace fixes revert_weight_conversion to handle -# scalar (0-d) tensors. transformers' Chunk.convert() calls torch.chunk() on quantization -# scale buffers that are 0-d scalars, raising RuntimeError ("chunk expects at least a -# 1-dimensional tensor"). Confirmed in transformers 5.12.0. -# See: transformers/core_model_loading.py, Chunk.convert() - - def export_speculative_decoding( model: torch.nn.Module, dtype: torch.dtype | None = None,