From eb08b940df9ef2c6d754bb535b05f511ea06c952 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:52:03 -0700 Subject: [PATCH 01/21] feat(ptq): single-GPU disk-offload layerwise PTQ (G1/G2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit G1 — offload-aware unified HF export - Add _has_accelerate_offload() to detect CPU/disk-offload hooks - Add _process_quantized_modules_offloaded() that materializes decoder layers one at a time inside enable_weight_access_and_writeback, then collects the full state dict via a second loop for non-decoder offloaded modules (embed, norm, lm_head) - Branch _export_transformers_checkpoint on the offload flag; remove hooks only after the offloaded state dict has been collected - Add meta-tensor guard in _export_quantized_weight to catch accidental standard-path use on offloaded models G2 — disk-offload CLI wiring in hf_ptq - Add --offload_folder, --max_gpu_memory_gb, --max_cpu_memory_gb args - Inject max_memory budget into load_model_from_config; skip seq_device_map when offload folder is set - Add nvfp4_experts_only-kv_fp8_layerwise_offload.yaml recipe Other - Fix _FP8BF16Fallback shim: nested try avoids ambiguous except; remove how-comments from matmul; tighten outer except to Exception only - Fix get_nemotron_h_decoder_layers to check both backbone.layers and model.layers Tests - 7 CPU-only unit tests (test_offload_export.py) - 2 GPU integration tests (tests/gpu/torch/export/test_offload_export.py) Signed-off-by: Fridah-nv Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 129 +++++++++++++--- examples/hf_ptq/hf_ptq.py | 45 ++++++ modelopt/torch/export/unified_export_hf.py | 143 ++++++++++++++---- .../torch/quantization/plugins/huggingface.py | 12 +- ...experts_only-kv_fp8_layerwise_offload.yaml | 49 ++++++ tests/gpu/torch/export/test_offload_export.py | 119 +++++++++++++++ .../unit/torch/export/test_offload_export.py | 116 ++++++++++++++ 7 files changed, 559 insertions(+), 54 deletions(-) create mode 100644 modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml create mode 100644 tests/gpu/torch/export/test_offload_export.py create mode 100644 tests/unit/torch/export/test_offload_export.py diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d74ffb34efb..2263d9908ec 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -30,6 +30,60 @@ import torch import transformers + +# Shim for is_torch_fx_available removed in transformers >=5.x; older model files (e.g. +# DeepSeek-R1 bundled modeling_deepseek.py) import it from transformers.utils.import_utils. +try: + from transformers.utils.import_utils import is_torch_fx_available # noqa: F401 +except ImportError: + import transformers.utils.import_utils as _tui + + _tui.is_torch_fx_available = lambda: False # type: ignore[attr-defined] + +# Shim for broken flash_attn installs (undefined symbol in .so). Probe the actual import; +# if it fails, force transformers' availability checks to return False so bundled remote-code +# model files (e.g. modeling_deepseek.py) skip the flash_attn import block. +# Must patch both transformers.utils.import_utils AND transformers.utils since bundled models +# import from either location. +try: + import flash_attn as _flash_attn_probe # noqa: F401 +except Exception: + import transformers.utils as _tu + import transformers.utils.import_utils as _tui + + for _mod in (_tu, _tui): + _mod.is_flash_attn_2_available = lambda: False # type: ignore[attr-defined] + _mod.is_flash_attn_available = lambda: False # type: ignore[attr-defined] + _mod.is_flash_attn_greater_or_equal_2_10 = lambda: False # type: ignore[attr-defined] + +# On nodes without the `kernels` package, DSR1 block-scaled FP8 matmul fails at import. +# Patch the loader with a BF16 dequant fallback so calibration forward passes succeed +# (amax collection only — not suitable for production inference). +try: + import transformers.integrations.finegrained_fp8 as _ff8 + + try: + _ff8._load_finegrained_fp8_kernel() + except ImportError: + + class _FP8BF16Fallback: + @staticmethod + def matmul(input, weight, weight_scale_inv, block_size, output_dtype=None, activation_scale=None): + out_f, in_f = weight.shape[-2], weight.shape[-1] + nb_out, nb_in = weight_scale_inv.shape[-2], weight_scale_inv.shape[-1] + scale = ( + weight_scale_inv.float() + .repeat_interleave(out_f // nb_out, -2) + .repeat_interleave(in_f // nb_in, -1) + ) + w_bf16 = (weight.float() * scale).to(torch.bfloat16) + out = torch.nn.functional.linear(input.to(torch.bfloat16), w_bf16) + return out if output_dtype is None else out.to(output_dtype) + + _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback # type: ignore[attr-defined] +except Exception: + pass + from accelerate import infer_auto_device_map, init_empty_weights from accelerate.utils import get_max_memory from safetensors import safe_open @@ -659,6 +713,17 @@ def _apply_dtype_to_config(model_kwargs, config_dtype, architecture, apply_confi return model_kwargs +def _fmt_max_memory(max_memory: dict) -> str: + """Format a ``{device: bytes}`` budget dict into a human-readable string.""" + parts = [] + for key in sorted(max_memory.keys(), key=lambda k: (isinstance(k, str), k)): + val = max_memory[key] + label = f"{val / 1024 ** 3:.1f} GiB" if isinstance(val, int) else str(val) + key_str = f"GPU {key}" if isinstance(key, int) else str(key) + parts.append(f" {key_str}: {label}") + return "\n".join(parts) + + def get_model( ckpt_path, device="cuda", @@ -666,9 +731,21 @@ def get_model( trust_remote_code=False, use_seq_device_map=False, attn_implementation=None, + offload_folder=None, + max_cpu_memory_gb=None, + max_gpu_memory_gb=None, ): print(f"Initializing model from {ckpt_path}") + _disk_offload = offload_folder is not None + if _disk_offload and max_cpu_memory_gb is None: + warnings.warn( + "offload_folder is set but max_cpu_memory_gb is not specified. " + "CPU memory usage during model load will be unbounded. " + "Pass max_cpu_memory_gb to cap CPU usage.", + UserWarning, + ) + device_map = "auto" if device == "cpu": device_map = "cpu" @@ -772,12 +849,11 @@ def has_pack_quantized_config(config): raise ValueError(f"Model config at {ckpt_path} has no architectures defined") architecture = hf_config.architectures[0] - if not hasattr(transformers, architecture) or "Deepseek" in architecture: - if not hasattr(transformers, architecture): - warnings.warn( - f"Architecture {architecture} not found in transformers: {transformers.__version__}. " - "Falling back to AutoModelForCausalLM (or AutoModel for non-causal architectures)." - ) + if not hasattr(transformers, architecture): + warnings.warn( + f"Architecture {architecture} not found in transformers: {transformers.__version__}. " + "Falling back to AutoModelForCausalLM (or AutoModel for non-causal architectures)." + ) assert trust_remote_code, ( "Please set trust_remote_code to True if you want to use this architecture" ) @@ -809,24 +885,39 @@ def has_pack_quantized_config(config): model = from_config(config_for_init, **model_kwargs2) max_memory = get_max_memory() - inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) - - on_cpu = "cpu" in inferred_device_map.values() - - if on_cpu: - for _device in max_memory: - if isinstance(_device, int): - max_memory[_device] *= gpu_mem_percentage + if _disk_offload: + if max_gpu_memory_gb is not None: + for _k in max_memory: + if isinstance(_k, int): + max_memory[_k] = int(max_gpu_memory_gb * 1024**3) + if max_cpu_memory_gb is not None: + max_memory["cpu"] = int(max_cpu_memory_gb * 1024**3) + model_kwargs["max_memory"] = max_memory print( - "Model does not fit to the GPU mem. " - f"We apply the following memory limit for calibration: \n{max_memory}\n" - "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " - "reduce the calibration `batch_size` manually." + "Disk-offload mode enabled. " + f"Memory budgets: {_fmt_max_memory(max_memory)}\n" + f"Offload folder: {offload_folder}\n" + "Weights exceeding GPU+CPU budgets will be streamed from disk." ) - model_kwargs["max_memory"] = max_memory + else: + inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) + if "cpu" in inferred_device_map.values(): + for _device in max_memory: + if isinstance(_device, int): + max_memory[_device] *= gpu_mem_percentage + + print( + "Model does not fit to the GPU mem. " + f"We apply the following memory limit for calibration: \n{max_memory}\n" + "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " + "reduce the calibration `batch_size` manually." + ) + model_kwargs["max_memory"] = max_memory model_kwargs2 = _apply_dtype_to_config(model_kwargs, config_dtype, architecture) + if _disk_offload: + model_kwargs2["offload_folder"] = offload_folder model = auto_model_module.from_pretrained( ckpt_path, device_map=device_map, diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6a4bd476984..a57c6e1d84a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -578,6 +578,9 @@ def load_model(args: argparse.Namespace): trust_remote_code=args.trust_remote_code, use_seq_device_map=args.use_seq_device_map, attn_implementation=args.attn_implementation, + offload_folder=args.offload_folder, + max_cpu_memory_gb=args.max_cpu_memory_gb, + max_gpu_memory_gb=args.max_gpu_memory_gb, ) else: assert args.qformat in QUANT_CFG_CHOICES, ( @@ -1621,6 +1624,38 @@ def parse_args() -> argparse.Namespace: "openai/gpt-oss-20b) and the target qformat is NVFP4-family." ), ) + parser.add_argument( + "--offload_folder", + type=str, + default=None, + help=( + "Path to a local folder for disk-offloaded model weights. " + "When set, activates disk-offload mode: model weights that exceed the GPU+CPU " + "budgets are streamed from disk during calibration and export. " + "Pair with --max_cpu_memory_gb to cap CPU RAM usage. " + "Incompatible with --low_memory_mode and --use_seq_device_map." + ), + ) + parser.add_argument( + "--max_cpu_memory_gb", + type=float, + default=None, + help=( + "Maximum CPU RAM budget in GiB for disk-offload model loading. " + "Only effective when --offload_folder is set. " + "Weights beyond this limit are streamed from disk." + ), + ) + parser.add_argument( + "--max_gpu_memory_gb", + type=float, + default=None, + help=( + "Maximum GPU memory budget per device in GiB for disk-offload model loading. " + "Only effective when --offload_folder is set. " + "Defaults to 80%% of available GPU memory when not specified." + ), + ) args = parser.parse_args() if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): @@ -1655,6 +1690,16 @@ def parse_args() -> argparse.Namespace: if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4: parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.") + if args.offload_folder is not None and args.low_memory_mode: + parser.error("--offload_folder (disk-offload) is not compatible with --low_memory_mode.") + + if args.offload_folder is not None and args.use_seq_device_map: + parser.error( + "--offload_folder (disk-offload) is not compatible with --use_seq_device_map; " + "device_map=auto is used for disk-offload to let accelerate place layers across " + "GPU, CPU, and disk." + ) + return args diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f51eea17b1a..45eacbecad2 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -570,6 +570,14 @@ def _export_quantized_weight( 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, export must run inside an " + "enable_weight_access_and_writeback context. Use the offload-aware export path " + "(_process_quantized_modules_offloaded) rather than _process_quantized_modules." + ) + # 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 @@ -780,6 +788,20 @@ def _export_quantized_weight( 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, @@ -813,20 +835,72 @@ def _process_quantized_modules( fsdp_module_to_reshard = sub_module - # We skip QuantLoraLinear module for modelopt QLoRA - if ctx.is_modelopt_qlora and hasattr(sub_module, "base_layer"): - continue + _dispatch_export_handler(name, sub_module, ctx) - # Preprocessing: restore unpacked weight so the export path can read - # the live quantizer state. Falls through to the handler dispatch below. - 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 _has_accelerate_offload(model: nn.Module) -> bool: + """Return True if any module in model has a CPU- or disk-offload accelerate hook.""" + try: + from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook + except ImportError: + return False + for mod in model.modules(): + hook = getattr(mod, "_hf_hook", None) + if hook is not None and _get_offload_hook(hook) is not None: + return True + return False + + +def _process_quantized_modules_offloaded( + model: nn.Module, + dtype: torch.dtype, + is_modelopt_qlora: bool = False, +) -> dict[str, Any]: + """Export quantized decoder-layer weights for an offloaded model, one layer at a time. + + Returns a full-model state dict with no meta tensors. + + Limitation: only decoder layers discovered by LayerActivationCollector are + materialized. Non-decoder quantized modules (e.g. a quantized lm_head) are + collected from model.state_dict() in their current form. Default FP8/NVFP4 + configs exclude lm_head, so this is typically harmless, but custom configs + that quantize non-decoder modules will export those layers without quantization applied. + """ + from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + decoder_layers = LayerActivationCollector.get_decoder_layers(model) + if decoder_layers is None: + raise RuntimeError( + "Disk/CPU-offloaded export requires discoverable decoder layers. " + "The model architecture is not supported by LayerActivationCollector." + ) + decoder_layer_ids = {id(m) for m in decoder_layers} + + ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + layer_tensors: dict[str, torch.Tensor] = {} + + for name, module in model.named_modules(): + if id(module) not in decoder_layer_ids: + continue + with enable_weight_access_and_writeback(module, module, writeback=True): + for sub_name, sub_mod in module.named_modules(): + full_name = f"{name}.{sub_name}" if sub_name else name + _dispatch_export_handler(full_name, sub_mod, ctx) + + # Snapshot inside the context: post-exit, post_forward re-offloads params to meta. + prefix = f"{name}." if name else "" + for key, tensor in module.state_dict().items(): + assert not tensor.is_meta, ( + f"Expected real tensor for '{prefix + key}' inside materialization context" + ) + layer_tensors[prefix + key] = tensor.detach() + + # model.state_dict() gives real tensors for non-offloaded parts (embed, lm_head, norms, …); + # meta placeholders for offloaded decoder layers are overridden by layer_tensors. + full_sd = model.state_dict() + full_sd.update(layer_tensors) + return full_sd def _export_transformers_checkpoint( @@ -873,13 +947,18 @@ def _export_transformers_checkpoint( # TODO: Handle mixed precision requantize_resmooth_fused_llm_layers(model) - # Remove all hooks from the model - try: - from accelerate.hooks import remove_hook_from_module + # Detect accelerate offload before removing hooks; offloaded models need weights + # materialized layer-by-layer during export (hooks must stay alive for that pass). + _offloaded = _has_accelerate_offload(model) - remove_hook_from_module(model, recurse=True) - except ImportError: - warnings.warn("accelerate is not installed, hooks will not be removed") + # Remove all hooks from the model (deferred for offloaded models) + if not _offloaded: + try: + from accelerate.hooks import remove_hook_from_module + + remove_hook_from_module(model, recurse=True) + except ImportError: + warnings.warn("accelerate is not installed, hooks will not be removed") quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) @@ -917,22 +996,24 @@ def _export_transformers_checkpoint( ) # Process all quantized modules and export weights - _process_quantized_modules(model, dtype, is_modelopt_qlora) - - # Reconstruct fused MoELinear: per-expert _QuantLinear weights → original 3D format from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear - _reconstruct_fused_moe_linear(model) - - if is_fsdp2_model(model): - # FSDP2: gather the full (unsharded) state_dict to CPU on rank 0. - quantized_state_dict = get_model_state_dict( - model, - options=StateDictOptions(full_state_dict=True, cpu_offload=True), - ) + if _offloaded: + quantized_state_dict = _process_quantized_modules_offloaded(model, dtype, is_modelopt_qlora) + # Reconstruct fused MoELinear: per-expert _QuantLinear weights → original 3D format + _reconstruct_fused_moe_linear(model) else: - # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict). - quantized_state_dict = model.state_dict() + _reconstruct_fused_moe_linear(model) + + if is_fsdp2_model(model): + # FSDP2: gather the full (unsharded) state_dict to CPU on rank 0. + quantized_state_dict = get_model_state_dict( + model, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + else: + # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict). + quantized_state_dict = model.state_dict() # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization. kv_cache_max_bound = 448 diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 69b8711da78..d26367cfb8e 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1750,10 +1750,14 @@ def get_nemotron_h_decoder_layers(model: nn.Module) -> nn.ModuleList | None: if not _is_supported_hf_model(model): return None - if hasattr(model, "backbone") and hasattr(model.backbone, "layers"): - layers = model.backbone.layers - if len(layers) > 0 and hasattr(layers[0], "block_type"): - return layers + # Custom remote-code checkpoint uses model.backbone.layers; + # native transformers NemotronHModel uses model.model.layers. + for container_attr in ("backbone", "model"): + container = getattr(model, container_attr, None) + if container is not None and hasattr(container, "layers"): + layers = container.layers + if layers and hasattr(layers[0], "block_type"): + return layers return None diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml new file mode 100644 index 00000000000..aa525cb4188 --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml @@ -0,0 +1,49 @@ +# 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. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: > + NVFP4 static weight and dynamic activation for expert layers only (W4A4), FP8 KV cache, + max layerwise calibration with calib_mutates_weights=False for disk-offloaded single-GPU + PTQ. Weights stay as meta tensors between layers; export_hf_checkpoint materializes them. +quantize: + algorithm: + method: max + layerwise: + enable: true + calib_mutates_weights: false + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers diff --git a/tests/gpu/torch/export/test_offload_export.py b/tests/gpu/torch/export/test_offload_export.py new file mode 100644 index 00000000000..6276081da06 --- /dev/null +++ b/tests/gpu/torch/export/test_offload_export.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. + +"""GPU integration tests for offload-aware unified HF export. + +Tests the full round-trip: + tiny LLaMA (CPU-offloaded via accelerate) + → FP8 layerwise calibration (calib_mutates_weights=False) + → export_hf_checkpoint + → assert no meta tensors in output safetensors + → assert hf_quant_config.json present with fp8 format +""" + +import copy +import json + +import pytest +import torch +from _test_utils.torch.transformers_models import create_tiny_llama_dir +from accelerate import init_empty_weights, load_checkpoint_and_dispatch +from safetensors import safe_open +from transformers import AutoConfig, AutoModelForCausalLM + +import modelopt.torch.quantization as mtq +from modelopt.torch.export import export_hf_checkpoint + + +def _make_cpu_offloaded_model(tmp_path, num_hidden_layers=3): + """Tiny LLaMA with first decoder layer offloaded to CPU, rest on GPU.""" + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_hidden_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + + # First layer on CPU to exercise the offload path; lm_head / embed on GPU. + device_map = {} + for n, _m in model.named_modules(): + if "layers" not in n or n.split("layers.")[-1].isdigit(): + device_map[n] = 0 + device_map["model.layers.0"] = "cpu" + + model = load_checkpoint_and_dispatch(model, tiny_llama_dir, device_map=device_map) + return model, config, tiny_llama_dir + + +def _layerwise_fp8_cfg(): + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + algo = cfg.get("algorithm", "max") + method = algo if isinstance(algo, str) else algo.get("method", "max") + # calib_mutates_weights is a field of LayerwiseConfig (nested), not of the algorithm. + cfg["algorithm"] = {"method": method, "layerwise": {"calib_mutates_weights": False}} + return cfg + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, _layerwise_fp8_cfg()]) +def test_export_hf_checkpoint_cpu_offloaded(tmp_path, quant_cfg): + """export_hf_checkpoint must succeed on a CPU-offloaded model and produce valid weights. + + Regression guard against the pre-fix bug where remove_hook_from_module was called + before weight materialization, causing meta tensors to be serialized as empty safetensors. + """ + num_hidden_layers = 3 + model, _config, _llama_dir = _make_cpu_offloaded_model( + tmp_path / "offloaded", num_hidden_layers=num_hidden_layers + ) + model.eval() + + def forward_loop(m): + ids = torch.randint(0, m.config.vocab_size, (1, 32)).cuda() + with torch.no_grad(): + m(ids) + + model = mtq.quantize(model, quant_cfg, forward_loop) + + export_dir = tmp_path / "hf_export" + export_dir.mkdir() + export_hf_checkpoint(model, export_dir=str(export_dir)) + + # --- Assertions --- + + # 1. hf_quant_config.json must exist and declare fp8 + quant_config_path = export_dir / "hf_quant_config.json" + assert quant_config_path.exists(), "hf_quant_config.json not written" + with open(quant_config_path) as f: + quant_config = json.load(f) + assert quant_config["quantization"]["quant_algo"] == "FP8", ( + f"Expected FP8, got {quant_config['quantization'].get('quant_algo')}" + ) + + # 2. All tensors in safetensors shards must be non-empty (no meta serialized as zeros) + safetensor_files = list(export_dir.glob("*.safetensors")) + assert safetensor_files, "No safetensors files written" + + for st_file in safetensor_files: + with safe_open(str(st_file), framework="pt") as st: + for key in st.keys(): + tensor = st.get_tensor(key) + assert tensor.numel() > 0, f"Zero-numel tensor for key '{key}' in {st_file.name}" + assert not tensor.is_meta, f"Meta tensor for key '{key}' in {st_file.name}" + # Weight tensors (not scales) must have non-zero norm — guards against all-zeros + # from meta serialization + if "weight" in key and "scale" not in key and "quantizer" not in key: + assert tensor.float().abs().sum() > 0, ( + f"All-zero weight tensor '{key}' in {st_file.name} — " + "possible meta tensor serialization bug" + ) diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py new file mode 100644 index 00000000000..7898b8ef0d0 --- /dev/null +++ b/tests/unit/torch/export/test_offload_export.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for offload-aware unified HF export helpers (CPU-only, no GPU required).""" + +import pytest +import torch +import torch.nn as nn + +try: + from accelerate.hooks import AlignDevicesHook, add_hook_to_module + from accelerate.utils import set_module_tensor_to_device +except ImportError: + pytest.skip("accelerate not available", allow_module_level=True) + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.unified_export_hf import ( + _export_quantized_weight, + _has_accelerate_offload, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_offloaded_linear(dim: int = 16): + """Return a Linear with a CPU-offload AlignDevicesHook attached and params on meta.""" + linear = nn.Linear(dim, dim, bias=False) + weights_map = {"weight": linear.weight.data.clone().cpu()} + hook = AlignDevicesHook(execution_device="cpu", offload=True, weights_map=weights_map) + add_hook_to_module(linear, hook) + set_module_tensor_to_device(linear, "weight", "meta") + return linear, weights_map + + +# --------------------------------------------------------------------------- +# _has_accelerate_offload +# --------------------------------------------------------------------------- + + +def test_has_accelerate_offload_true(): + linear, _ = _make_offloaded_linear() + assert _has_accelerate_offload(linear) is True + + +def test_has_accelerate_offload_false_no_hooks(): + linear = nn.Linear(16, 16) + assert _has_accelerate_offload(linear) is False + + +def test_has_accelerate_offload_false_non_offload_hook(): + """A hook with offload=False should not be detected as offloaded.""" + linear = nn.Linear(16, 16) + hook = AlignDevicesHook(execution_device="cpu", offload=False) + add_hook_to_module(linear, hook) + assert _has_accelerate_offload(linear) is False + + +def test_has_accelerate_offload_detects_nested_module(): + """Offload hook on a child module should be detected when scanning the parent.""" + + class _Parent(nn.Module): + def __init__(self): + super().__init__() + self.child = nn.Linear(8, 8, bias=False) + + def forward(self, x): + return self.child(x) + + parent = _Parent() + weights_map = {"weight": parent.child.weight.data.clone().cpu()} + hook = AlignDevicesHook(execution_device="cpu", offload=True, weights_map=weights_map) + add_hook_to_module(parent.child, hook) + set_module_tensor_to_device(parent.child, "weight", "meta") + + assert _has_accelerate_offload(parent) is True + + +# --------------------------------------------------------------------------- +# _export_quantized_weight meta guard +# --------------------------------------------------------------------------- + + +def test_meta_guard_raises_on_meta_weight(): + """_export_quantized_weight must raise RuntimeError when weight is a meta tensor.""" + linear = nn.Linear(16, 16, bias=False) + + mtq.quantize(linear, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(1, 16))) + + # Manually set weight to meta to simulate what happens after hooks are removed. + linear.weight = nn.Parameter(torch.empty(16, 16, device="meta")) + + with pytest.raises(RuntimeError, match="meta tensor"): + _export_quantized_weight(linear, torch.float32) + + +def test_meta_guard_not_raised_for_real_weight(): + """No RuntimeError when weight is a real (non-meta) tensor.""" + linear = nn.Linear(32, 32, bias=False) + mtq.quantize(linear, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(1, 32))) + # Should not raise + _export_quantized_weight(linear, torch.float32) From 9f2b6793e080e4ac24009591fe3b1ea63440cca8 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 23 Jul 2026 13:44:44 -0700 Subject: [PATCH 02/21] fix(export): materialize non-decoder disk-offloaded tensors before save embed_tokens, final norms, and lm_head are disk-offloaded alongside decoder layers on single-GPU runs. model.state_dict() returns meta placeholders for them; after revert_weight_conversion_quant_aware renames to hub-original keys, transformers' save_pretrained looks up the tensors by hub name and crashes (e.g. NemotronHForCausalLM has no attribute 'backbone'). Fix: add a second loop in _process_quantized_modules_offloaded that iterates non-decoder modules, skips any without a live offload hook, checks for DIRECT meta parameters/buffers (avoids re-collecting decoder children already captured above), and materializes each via enable_weight_access_and_writeback. Signed-off-by: Fridah-nv Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/unified_export_hf.py | 36 ++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 45eacbecad2..303a582d585 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -896,8 +896,40 @@ def _process_quantized_modules_offloaded( ) layer_tensors[prefix + key] = tensor.detach() - # model.state_dict() gives real tensors for non-offloaded parts (embed, lm_head, norms, …); - # meta placeholders for offloaded decoder layers are overridden by layer_tensors. + # Also collect direct parameters of non-decoder modules that are disk-offloaded. + # model.state_dict() returns meta for ANY disk-offloaded tensor, including + # embed_tokens, final norms, and lm_head. After revert_weight_conversion renames + # these to hub-original names (e.g. backbone.*), transformers' save_pretrained + # looks them up in the model by hub name and crashes if they are still meta. + # Fix: materialize each such module in-place and capture the real tensor. + from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook + + for name, module in model.named_modules(): + if id(module) in decoder_layer_ids: + continue + if not hasattr(module, "_hf_hook"): + continue + if _get_offload_hook(module._hf_hook) is None: + continue + # Only handle modules that have DIRECT meta parameters/buffers. + # Child decoder layers (already quantized above) must not be re-collected. + if not ( + any(p is not None and p.is_meta for p in module._parameters.values()) + or any(b is not None and b.is_meta for b in module._buffers.values()) + ): + continue + with enable_weight_access_and_writeback(module, module, writeback=False): + prefix = f"{name}." if name else "" + for pname, param in module._parameters.items(): + if param is not None and not param.is_meta: + layer_tensors[prefix + pname] = param.data.detach().cpu() + for bname, buf in module._buffers.items(): + if buf is not None and not buf.is_meta: + layer_tensors[prefix + bname] = buf.detach().cpu() + + # model.state_dict() fills in non-offloaded parts (GPU-resident tensors). + # layer_tensors overrides both decoder-layer placeholders and non-decoder + # offloaded placeholders so the returned dict contains no meta tensors. full_sd = model.state_dict() full_sd.update(layer_tensors) return full_sd From ce143d90befab24fb3810a475a372786b95b59e1 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 23 Jul 2026 14:06:15 -0700 Subject: [PATCH 03/21] =?UTF-8?q?fix(export):=20correct=20offloaded=20expo?= =?UTF-8?q?rt=20path=20=E2=80=94=20MoE=20ordering,=20lm=5Fhead=20handler,?= =?UTF-8?q?=20memory?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four correctness / efficiency bugs in _process_quantized_modules_offloaded: 1. MoE reconstruction ordering: _reconstruct_fused_moe_linear was called after _process_quantized_modules_offloaded returned the state dict, so fused-MoE decoder layers shipped with per-expert 2D keys (or meta weights after context exit). Move it per-layer inside the materialization window before the state dict snapshot, mirroring the non-offloaded path. 2. Quantized non-decoder modules missed: lm_head (or any quantized module outside the decoder stack) never had _dispatch_export_handler called, so it exported raw unquantized weights. Add handler dispatch inside the non-decoder materialization context. 3. Decoder snapshots accumulating on GPU: tensor.detach() kept every layer's materialized weights on the GPU, defeating the layer-by-layer memory goal. Change to tensor.detach().cpu(). 4. writeback=True on decoder context: on context exit the quantized weights were written back to the offload store (disk → CPU promotion per layer), wasting CPU memory. Changed to writeback=False since weights are captured in layer_tensors immediately after. Also: apply gpu_mem_percentage (default 80 %) when --max_gpu_memory_gb is not specified in disk-offload mode, matching the documented CLI default. Restore test_non_decoder_offloaded_tensors_are_collected (lost during squash). Signed-off-by: Fridah-nv Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 8 +- modelopt/torch/export/unified_export_hf.py | 54 +++++++------- .../unit/torch/export/test_offload_export.py | 73 +++++++++++++++++++ 3 files changed, 106 insertions(+), 29 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 2263d9908ec..b3b16467d01 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -887,10 +887,12 @@ def has_pack_quantized_config(config): max_memory = get_max_memory() if _disk_offload: - if max_gpu_memory_gb is not None: - for _k in max_memory: - if isinstance(_k, int): + for _k in max_memory: + if isinstance(_k, int): + if max_gpu_memory_gb is not None: max_memory[_k] = int(max_gpu_memory_gb * 1024**3) + else: + max_memory[_k] = int(max_memory[_k] * gpu_mem_percentage) if max_cpu_memory_gb is not None: max_memory["cpu"] = int(max_cpu_memory_gb * 1024**3) model_kwargs["max_memory"] = max_memory diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 303a582d585..9a047a59138 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -856,16 +856,17 @@ def _process_quantized_modules_offloaded( dtype: torch.dtype, is_modelopt_qlora: bool = False, ) -> dict[str, Any]: - """Export quantized decoder-layer weights for an offloaded model, one layer at a time. + """Export quantized weights for a disk/CPU-offloaded model, one layer at a time. - Returns a full-model state dict with no meta tensors. + Decoder layers are processed one at a time via enable_weight_access_and_writeback. + Non-decoder modules that are also disk-offloaded (embed_tokens, norms, lm_head) are + materialized individually; any quantized non-decoder module (e.g. lm_head) has its + export handler invoked in the same context. - Limitation: only decoder layers discovered by LayerActivationCollector are - materialized. Non-decoder quantized modules (e.g. a quantized lm_head) are - collected from model.state_dict() in their current form. Default FP8/NVFP4 - configs exclude lm_head, so this is typically harmless, but custom configs - that quantize non-decoder modules will export those layers without quantization applied. + Returns a full-model state dict with no meta tensors. """ + from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook + from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -883,27 +884,29 @@ def _process_quantized_modules_offloaded( for name, module in model.named_modules(): if id(module) not in decoder_layer_ids: continue - with enable_weight_access_and_writeback(module, module, writeback=True): + # writeback=False: weights are captured in layer_tensors below; no need to promote + # the quantized values back to the offload store on context exit. + with enable_weight_access_and_writeback(module, module, writeback=False): for sub_name, sub_mod in module.named_modules(): full_name = f"{name}.{sub_name}" if sub_name else name _dispatch_export_handler(full_name, sub_mod, ctx) + # Mirror the non-offloaded path: reconstruct fused MoE per-expert weights + # into 3D tensors BEFORE snapshotting, so captured keys match the original + # MoE format (e.g. moe.up_proj.weight [N, out, in]). + _reconstruct_fused_moe_linear(module) + # Snapshot inside the context: post-exit, post_forward re-offloads params to meta. prefix = f"{name}." if name else "" for key, tensor in module.state_dict().items(): assert not tensor.is_meta, ( f"Expected real tensor for '{prefix + key}' inside materialization context" ) - layer_tensors[prefix + key] = tensor.detach() - - # Also collect direct parameters of non-decoder modules that are disk-offloaded. - # model.state_dict() returns meta for ANY disk-offloaded tensor, including - # embed_tokens, final norms, and lm_head. After revert_weight_conversion renames - # these to hub-original names (e.g. backbone.*), transformers' save_pretrained - # looks them up in the model by hub name and crashes if they are still meta. - # Fix: materialize each such module in-place and capture the real tensor. - from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook + layer_tensors[prefix + key] = tensor.detach().cpu() + # Also collect non-decoder modules that are disk-offloaded (embed_tokens, norms, lm_head). + # model.state_dict() returns meta for these; materialize, run any export handlers + # (e.g. a quantized lm_head), then snapshot the real tensors. for name, module in model.named_modules(): if id(module) in decoder_layer_ids: continue @@ -912,20 +915,20 @@ def _process_quantized_modules_offloaded( if _get_offload_hook(module._hf_hook) is None: continue # Only handle modules that have DIRECT meta parameters/buffers. - # Child decoder layers (already quantized above) must not be re-collected. + # Child decoder layers (already captured above) must not be re-collected. if not ( any(p is not None and p.is_meta for p in module._parameters.values()) or any(b is not None and b.is_meta for b in module._buffers.values()) ): continue with enable_weight_access_and_writeback(module, module, writeback=False): + for sub_name, sub_mod in module.named_modules(): + full_name = f"{name}.{sub_name}" if sub_name else name + _dispatch_export_handler(full_name, sub_mod, ctx) prefix = f"{name}." if name else "" - for pname, param in module._parameters.items(): - if param is not None and not param.is_meta: - layer_tensors[prefix + pname] = param.data.detach().cpu() - for bname, buf in module._buffers.items(): - if buf is not None and not buf.is_meta: - layer_tensors[prefix + bname] = buf.detach().cpu() + for key, tensor in module.state_dict().items(): + if not tensor.is_meta: + layer_tensors[prefix + key] = tensor.detach().cpu() # model.state_dict() fills in non-offloaded parts (GPU-resident tensors). # layer_tensors overrides both decoder-layer placeholders and non-decoder @@ -1031,9 +1034,8 @@ def _export_transformers_checkpoint( from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear if _offloaded: + # MoE reconstruction happens per-layer inside _process_quantized_modules_offloaded. quantized_state_dict = _process_quantized_modules_offloaded(model, dtype, is_modelopt_qlora) - # Reconstruct fused MoELinear: per-expert _QuantLinear weights → original 3D format - _reconstruct_fused_moe_linear(model) else: _reconstruct_fused_moe_linear(model) diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 7898b8ef0d0..ef5144ff67f 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -29,6 +29,7 @@ from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _has_accelerate_offload, + _process_quantized_modules_offloaded, ) @@ -114,3 +115,75 @@ def test_meta_guard_not_raised_for_real_weight(): mtq.quantize(linear, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(1, 32))) # Should not raise _export_quantized_weight(linear, torch.float32) + + +# --------------------------------------------------------------------------- +# _process_quantized_modules_offloaded — non-decoder materialization +# --------------------------------------------------------------------------- + + +def test_non_decoder_offloaded_tensors_are_collected(): + """Non-decoder modules with disk-offload hooks must have no meta tensors in the result. + + Reproduces the NemotronH 550B crash: embed_tokens (and norm, lm_head) are + disk-offloaded and return meta from model.state_dict(). After + revert_weight_conversion_quant_aware renames them to hub-original names, transformers' + remove_tied_weights_from_state_dict tries to look them up in the model by that name + and crashes. Fix: _process_quantized_modules_offloaded materialises non-decoder + offloaded modules directly so the returned state dict contains no meta tensors. + + The decoder layer here is NOT disk-offloaded (all weights GPU-resident) so the + decoder-layer loop exercises the null-context path and we focus on the non-decoder + collection pass that was previously missing. + """ + + class _TinyLayer(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(8, 8, bias=False) + + def forward(self, x): + return self.proj(x) + + class _TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.embed = nn.Embedding(16, 8) + self.layers = nn.ModuleList([_TinyLayer()]) + + def forward(self, x): + return self.layers[0](self.embed(x)) + + model = _TinyModel() + + # Install a CPU-offload hook on embed ONLY (non-decoder module). + # The decoder layer is left GPU-resident so enable_weight_access_and_writeback + # returns a no-op nullcontext and the decoder-layer state_dict() returns real tensors. + embed_val = model.embed.weight.data.clone().cpu() + embed_weights_map = {"weight": embed_val} + embed_hook = AlignDevicesHook( + execution_device="cpu", offload=True, weights_map=embed_weights_map + ) + add_hook_to_module(model.embed, embed_hook) + set_module_tensor_to_device(model.embed, "weight", "meta") + + from unittest.mock import patch + + with patch( + "modelopt.torch.quantization.utils.layerwise_calib" + ".LayerActivationCollector.get_decoder_layers", + return_value=list(model.layers), + ): + result = _process_quantized_modules_offloaded(model, torch.float32) + + assert "embed.weight" in result, "embed.weight missing from state dict" + emb = result["embed.weight"] + assert not emb.is_meta, "embed.weight must not be meta in exported state dict" + assert emb.shape == (16, 8) + + assert "layers.0.proj.weight" in result + assert not result["layers.0.proj.weight"].is_meta + + for key, val in result.items(): + if isinstance(val, torch.Tensor): + assert not val.is_meta, f"meta tensor found for key '{key}'" From 841ea5fcfe55eab807765db9049683fe0c4a6545 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 23 Jul 2026 14:40:34 -0700 Subject: [PATCH 04/21] feat(export): streaming shard writer for 80 GB CPU RAM target Replace the accumulate-then-save pattern in the offloaded export path with a stream-then-save pattern. Peak memory drops from ~764 GiB (Ultra 550B full state dict in RAM) to 1 decoder layer + 1 shard buffer (~57 GB). New pieces: - `_postprocess_single_tensor` in quant_utils.py: per-tensor subset of postprocess_state_dict for use in the streaming loop - `_StreamingShardWriter`: buffers tensors up to max_shard_size, flushes to temp part files, renames to canonical shard names at finalize() - `_parse_shard_size`: converts "10GB"/"500MB" strings to bytes - `_export_transformers_checkpoint_streaming`: streams decoder layers one at a time via enable_weight_access_and_writeback, applies per-tensor postprocessing and name reversal inline, handles tied-weight dedup from _tied_weights_keys - `export_hf_checkpoint` dispatch: branches on _has_accelerate_offload to call the streaming path instead of _export_transformers_checkpoint for offloaded models; hf_quant_config.json and config.json update are shared between paths Unit tests: 10 new tests covering _StreamingShardWriter (single-shard, multi-shard, readback) and _postprocess_single_tensor (passthrough, filter, rename, squeeze, real-quant drop, kv scale divide). Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/quant_utils.py | 83 +++++ modelopt/torch/export/unified_export_hf.py | 352 +++++++++++++++++- .../unit/torch/export/test_offload_export.py | 131 ++++++- 3 files changed, 546 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..ac825194bda 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -959,6 +959,89 @@ def from_quantized_weight( raise NotImplementedError(f"quantization format {quantization} not supported") +def _postprocess_single_tensor( + key: str, + value: torch.Tensor, + kv_cache_max_bound: float, + kv_cache_format: str | None, + is_modelopt_qlora: bool = False, +) -> tuple[str | None, torch.Tensor | None]: + """Per-tensor subset of :func:`postprocess_state_dict`, for streaming export. + + Returns ``(new_key, new_value)`` to emit, or ``(None, None)`` to skip. + Tied-weight dedup is NOT performed here; callers should pre-compute alias + keys from ``model._tied_weights_keys`` and filter them at the call site. + """ + replacements = { + "k_bmm_quantizer._amax": "k_proj.k_scale", + "v_bmm_quantizer._amax": "v_proj.v_scale", + "k_bmm_quantizer._bias_value": "k_proj.k_bias", + "v_bmm_quantizer._bias_value": "v_proj.v_bias", + "input_quantizer._pre_quant_scale": "pre_quant_scale", + } + skip_keys = [ + "output_quantizer", + "_amax", + "_bias_value", + "input_quantizer._pre_quant_scale", + "weight_shape", + ] + if is_modelopt_qlora: + replacements.update( + { + "base_layer.weight": "weight", + "base_layer.input_scale": "input_scale", + "base_layer.weight_scale": "weight_scale", + } + ) + skip_keys.append("base_layer") + + # Skip problematic VL model parameters + if key == "vision_model.radio_model.summary_idxs": + return None, None + + # Skip real quant parameters + if any(key.endswith("weight_quantizer." + q) for q in RealQuantLinear.list_of_scale_tensors): + return None, None + + # Skip LoRA adapters for QLoRA models + if is_modelopt_qlora and "lora" in key: + return None, None + + # Keys not related to quantizers: keep as-is + if all(sk not in key for sk in skip_keys): + if "scale" in key and isinstance(value, torch.Tensor) and value.dim() == 3 and value.shape[0] == 1: + value = value.squeeze(0) + return key, value + + # Apply replacements if the key matches any suffix in the replacements dict + for old_suffix, new_suffix in replacements.items(): + if key.endswith(old_suffix): + prefix = key[: -len(old_suffix)] + if "_amax" in key: + assert kv_cache_format in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], ( + "Invalid KV cache quantization format." + ) + assert kv_cache_max_bound > 0, "Maxbound must be greater than zero." + value = value.float() / kv_cache_max_bound + if kv_cache_format == KV_CACHE_FP8 and value.item() > 0.5: + logger.warning( + "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." + ) + new_key = prefix + new_suffix + if ( + "scale" in new_key + and isinstance(value, torch.Tensor) + and value.dim() == 3 + and value.shape[0] == 1 + ): + value = value.squeeze(0) + return new_key, value + + # Key has a skip_key but no replacement matched — drop it + return None, None + + def postprocess_state_dict( state_dict: dict, maxbound: float, diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 9a047a59138..00ea4cb6b7a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -97,6 +97,7 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import ( + _postprocess_single_tensor, fuse_prequant_layernorm, fuse_prequant_to_linear, get_activation_scaling_factor, @@ -938,6 +939,291 @@ def _process_quantized_modules_offloaded( return full_sd +class _StreamingShardWriter: + """Write tensors to safetensors shard files without accumulating the full state dict. + + Buffers tensors up to ``max_shard_size`` bytes, flushes to a numbered temp file, then + at :meth:`finalize` renames temp files to canonical shard names once the total shard + count is known. + + Peak memory = 1 layer (being materialized) + 1 shard buffer, not the full checkpoint. + """ + + def __init__(self, export_dir: Path | str, max_shard_size: int) -> None: + self._export_dir = Path(export_dir) + self._max_shard_size = max_shard_size + self._buffer: dict[str, torch.Tensor] = {} + self._buffer_bytes: int = 0 + self._part_files: list[Path] = [] + self._part_bytes: list[int] = [] + # Maps tensor key → part-file index (recorded at flush time) + self._key_to_part: dict[str, int] = {} + + def _flush(self) -> None: + if not self._buffer: + return + part_idx = len(self._part_files) + part_path = self._export_dir / f"__shard_part_{part_idx:05d}.safetensors" + save_file(self._buffer, str(part_path)) + for key in self._buffer: + self._key_to_part[key] = part_idx + self._part_files.append(part_path) + self._part_bytes.append(self._buffer_bytes) + self._buffer = {} + self._buffer_bytes = 0 + + def add(self, key: str, tensor: torch.Tensor) -> None: + """Buffer a tensor, flushing the current shard to disk when it is full.""" + self._buffer[key] = tensor + self._buffer_bytes += tensor.nbytes + if self._buffer_bytes >= self._max_shard_size: + self._flush() + + def finalize(self) -> dict[str, str]: + """Flush remaining buffer, rename part files, write model.safetensors.index.json. + + Returns the weight_map ``{key: shard_filename}`` written to the index. + Single-shard exports use ``model.safetensors`` without an index file. + """ + self._flush() + n_shards = len(self._part_files) + if n_shards == 0: + return {} + + if n_shards == 1: + final_name = "model.safetensors" + self._part_files[0].rename(self._export_dir / final_name) + return dict.fromkeys(self._key_to_part, final_name) + + for i, part_path in enumerate(self._part_files): + part_path.rename(self._export_dir / f"model-{i + 1:05d}-of-{n_shards:05d}.safetensors") + + weight_map = { + key: f"model-{part_idx + 1:05d}-of-{n_shards:05d}.safetensors" + for key, part_idx in self._key_to_part.items() + } + total_size = sum(self._part_bytes) + index_path = self._export_dir / "model.safetensors.index.json" + with open(index_path, "w") as f: + json.dump({"metadata": {"total_size": total_size}, "weight_map": weight_map}, f) + return weight_map + + +def _parse_shard_size(size: int | str) -> int: + """Convert a shard-size string (e.g. ``"10GB"``, ``"500MB"``) to bytes.""" + try: + from transformers.utils import convert_file_size_to_int + + return convert_file_size_to_int(size) + except (ImportError, Exception): + pass + if isinstance(size, int): + return size + s = size.strip().upper() + if s.endswith("GIB"): + return int(float(s[:-3]) * 1024**3) + if s.endswith("GB"): + return int(float(s[:-2]) * 1024**3) + if s.endswith("MIB"): + return int(float(s[:-3]) * 1024**2) + if s.endswith("MB"): + return int(float(s[:-2]) * 1024**2) + return int(s) + + +def _export_transformers_checkpoint_streaming( + model: nn.Module, + dtype: torch.dtype | None = None, + is_modelopt_qlora: bool = False, + export_dir: Path | str = ".", + max_shard_size: int | str = "10GB", + **kwargs, +) -> tuple[None, dict[str, Any]]: + """Export a disk/CPU-offloaded model by streaming tensors layer-by-layer to shard files. + + Peak memory = 1 decoder layer + 1 shard buffer, rather than the full quantized state + dict accumulated in RAM (which reaches ~764 GiB for Ultra 550B). + + Returns ``(None, quant_config)``; shard files, ``config.json``, and + ``generation_config.json`` are written to ``export_dir`` directly. The caller is + responsible for writing ``hf_quant_config.json`` and updating ``config.json`` with + ``quantization_config``. + """ + from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook + from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + export_dir = Path(export_dir) + + # --- Same model-level setup as _export_transformers_checkpoint --- + if dtype is None: + dtype = model.config.torch_dtype + elif 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." + ) + + 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: + 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) + + requantize_resmooth_fused_llm_layers(model) + + quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) + + 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: + pattern = f"{prefix}*" + if pattern not in exclude_modules: + exclude_modules.append(pattern) + print(f"Adding MTP layer to quantization_config ignore: {pattern}") + + 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." + ) + + synced_input = sync_tied_input_amax(model) + if synced_input: + print( + f"sync_tied_input_amax: max-merged input_quantizer amaxes across " + f"{synced_input} tied module group(s)" + ) + + # --- Per-tensor constants --- + kv_cache_max_bound = 448 + kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] + + # --- Tied alias keys to skip (data_ptr() is unreliable for disk-offloaded weights) --- + raw_tied_keys: set[str] = set(getattr(model, "_tied_weights_keys", None) or []) + + # --- Name mapper for per-tensor key reversal --- + # Tensor names are applied inline; quant config names are handled by the caller. + name_mapper = None + try: + name_mapper = build_reverse_name_mapper(model) + except Exception as exc: + warnings.warn( + f"Reverse name mapper unavailable ({exc}); exported tensor names may not match " + "the original HF hub checkpoint." + ) + + tied_alias_keys: set[str] = ( + {name_mapper(k) for k in raw_tied_keys} if name_mapper is not None else raw_tied_keys + ) + + # --- Decoder layers --- + decoder_layers = LayerActivationCollector.get_decoder_layers(model) + if decoder_layers is None: + raise RuntimeError( + "Streaming export requires discoverable decoder layers. " + "The model architecture is not supported by LayerActivationCollector." + ) + decoder_layer_ids = {id(m) for m in decoder_layers} + + # --- Stream tensors to shard files --- + shard_size_bytes = _parse_shard_size(max_shard_size) + writer = _StreamingShardWriter(export_dir, shard_size_bytes) + ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + seen_keys: set[str] = set() + + def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: + new_key, new_value = _postprocess_single_tensor( + full_key, tensor, kv_cache_max_bound, kv_cache_format, is_modelopt_qlora + ) + if new_key is None: + return + if name_mapper is not None: + new_key = name_mapper(new_key) + if new_key in tied_alias_keys: + return + writer.add(new_key, new_value.detach().cpu()) + + # Decoder layers (offloaded: materialize one at a time) + for layer_name, layer_module in model.named_modules(): + if id(layer_module) not in decoder_layer_ids: + continue + with enable_weight_access_and_writeback(layer_module, layer_module, writeback=False): + for sub_name, sub_mod in layer_module.named_modules(): + full_name = f"{layer_name}.{sub_name}" if sub_name else layer_name + _dispatch_export_handler(full_name, sub_mod, ctx) + _reconstruct_fused_moe_linear(layer_module) + prefix = f"{layer_name}." if layer_name else "" + for key, tensor in layer_module.state_dict().items(): + full_key = prefix + key + if full_key in seen_keys: + continue + seen_keys.add(full_key) + _stream_tensor(full_key, tensor) + + # Non-decoder modules with offload hooks (embed_tokens, norm, lm_head, etc.) + for name, module in model.named_modules(): + if id(module) in decoder_layer_ids: + continue + if not hasattr(module, "_hf_hook"): + continue + if _get_offload_hook(module._hf_hook) is None: + continue + if not ( + any(p is not None and p.is_meta for p in module._parameters.values()) + or any(b is not None and b.is_meta for b in module._buffers.values()) + ): + continue + with enable_weight_access_and_writeback(module, module, writeback=False): + for sub_name, sub_mod in module.named_modules(): + full_name = f"{name}.{sub_name}" if sub_name else name + _dispatch_export_handler(full_name, sub_mod, ctx) + prefix = f"{name}." if name else "" + for key, tensor in module.state_dict().items(): + full_key = prefix + key + if full_key in seen_keys or tensor.is_meta: + continue + seen_keys.add(full_key) + _stream_tensor(full_key, tensor) + + # GPU-resident parameters and buffers (not covered by the above loops) + for name, param in model.named_parameters(): + if name in seen_keys or param is None or param.is_meta: + continue + seen_keys.add(name) + _stream_tensor(name, param) + + for name, buf in model.named_buffers(): + if name in seen_keys or buf is None or buf.is_meta: + continue + seen_keys.add(name) + _stream_tensor(name, buf) + + writer.finalize() + + # Write non-weight artifacts (model.save_pretrained skipped to avoid OOM) + import contextlib + + _sanitize_generation_config_for_save(model) + model.config.save_pretrained(str(export_dir)) + gc = getattr(model, "generation_config", None) + if gc is not None: + with contextlib.suppress(Exception): + gc.save_pretrained(str(export_dir)) + + return None, quant_config + + def _export_transformers_checkpoint( model: nn.Module, dtype: torch.dtype | None = None, @@ -1557,14 +1843,38 @@ def export_hf_checkpoint( and torch.distributed.is_initialized() and is_fsdp2_model(model) ) + # Streaming path writes shard files layer-by-layer without accumulating the full + # state dict in RAM (peak = 1 layer + 1 shard buffer vs. ~764 GiB for Ultra 550B). + _offloaded = _has_accelerate_offload(model) + try: - post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) + if _offloaded: + if save_modelopt_state: + warnings.warn( + "save_modelopt_state=True is not supported in the streaming offload export " + "path and will be ignored." + ) + if extra_state_dict: + warnings.warn( + "extra_state_dict is not supported in the streaming offload export path " + "and will be ignored." + ) + _, hf_quant_config = _export_transformers_checkpoint_streaming( + model, + dtype, + export_dir=export_dir, + max_shard_size=max_shard_size, + **kwargs, + ) + else: + post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) # Remove hf_quantizer from model so post_state_dict can be exported. if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None - export_state_dict = {**post_state_dict, **(extra_state_dict or {})} + if not _offloaded: + export_state_dict = {**post_state_dict, **(extra_state_dict or {})} # transformers may have applied a load-time conversion_mapping (fused gate_up_proj, # renamed MoE leaves, reordered model/language_model prefix), so the in-memory names @@ -1579,7 +1889,10 @@ def export_hf_checkpoint( # weights and config so they stay mutually consistent. try: name_mapper = build_reverse_name_mapper(model) - export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict) + if not _offloaded: + # Streaming path applies per-tensor renaming inline inside + # _export_transformers_checkpoint_streaming; skip full-dict reversal here. + export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict) if name_mapper is not None and hf_quant_config: revert_quant_config_names(hf_quant_config.get("quantization", {}), name_mapper) except Exception as exc: @@ -1612,24 +1925,25 @@ def export_hf_checkpoint( else: hf_quant_config = None - # Keep transformers' own revert_weight_conversion disabled (the quant-aware reverse - # above replaces it): it can't handle quantized state dicts (RuntimeError on 0-d scalar - # scale tensors). Patch both the source and importing module since modeling_utils does - # `from core_model_loading import revert_weight_conversion`. - _patches = _patch_revert_weight_conversion() + if not _offloaded: + # Keep transformers' own revert_weight_conversion disabled (the quant-aware reverse + # above replaces it): it can't handle quantized state dicts (RuntimeError on 0-d scalar + # scale tensors). Patch both the source and importing module since modeling_utils does + # `from core_model_loading import revert_weight_conversion`. + _patches = _patch_revert_weight_conversion() - _sanitize_generation_config_for_save(model) + _sanitize_generation_config_for_save(model) - # TODO: parallelize the disk write across ranks (avoid single-process speed + rank-0 OOM). - try: - model.save_pretrained( - export_dir, - state_dict=export_state_dict, - save_modelopt_state=save_modelopt_state, - max_shard_size=max_shard_size, - ) - finally: - _unpatch_revert_weight_conversion(_patches) + # TODO: parallelize the disk write across ranks (avoid single-process speed + rank-0 OOM). + try: + model.save_pretrained( + export_dir, + state_dict=export_state_dict, + save_modelopt_state=save_modelopt_state, + max_shard_size=max_shard_size, + ) + finally: + _unpatch_revert_weight_conversion(_patches) original_config = f"{export_dir}/config.json" config_data = {} diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index ef5144ff67f..05f9974df6e 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -15,9 +15,14 @@ """Unit tests for offload-aware unified HF export helpers (CPU-only, no GPU required).""" +import json +import tempfile +from pathlib import Path + import pytest import torch import torch.nn as nn +from safetensors import safe_open try: from accelerate.hooks import AlignDevicesHook, add_hook_to_module @@ -26,13 +31,14 @@ pytest.skip("accelerate not available", allow_module_level=True) import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import _postprocess_single_tensor from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _has_accelerate_offload, _process_quantized_modules_offloaded, + _StreamingShardWriter, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -187,3 +193,126 @@ def forward(self, x): for key, val in result.items(): if isinstance(val, torch.Tensor): assert not val.is_meta, f"meta tensor found for key '{key}'" + + +# --------------------------------------------------------------------------- +# _StreamingShardWriter +# --------------------------------------------------------------------------- + + +def test_streaming_shard_writer_single_shard(): + """Small tensors that fit in one shard produce model.safetensors without an index.""" + with tempfile.TemporaryDirectory() as tmpdir: + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("a", torch.ones(4, 4)) + writer.add("b", torch.zeros(2, 2)) + weight_map = writer.finalize() + + single = Path(tmpdir) / "model.safetensors" + index = Path(tmpdir) / "model.safetensors.index.json" + assert single.exists(), "model.safetensors not written" + assert not index.exists(), "index file must not exist for single-shard export" + assert set(weight_map.values()) == {"model.safetensors"} + assert set(weight_map.keys()) == {"a", "b"} + + +def test_streaming_shard_writer_multi_shard(): + """Tensors exceeding max_shard_size produce multiple shards and an index file.""" + with tempfile.TemporaryDirectory() as tmpdir: + # One float32 4x4 tensor = 64 bytes; set limit to 64 so each tensor goes to a new shard + writer = _StreamingShardWriter(tmpdir, max_shard_size=64) + writer.add("x", torch.ones(4, 4)) + writer.add("y", torch.ones(4, 4)) + weight_map = writer.finalize() + + index_path = Path(tmpdir) / "model.safetensors.index.json" + assert index_path.exists(), "model.safetensors.index.json not written" + assert weight_map["x"] != weight_map["y"], "keys must be in different shards" + + with open(index_path) as f: + index = json.load(f) + assert "weight_map" in index + assert "metadata" in index + assert index["metadata"]["total_size"] > 0 + + +def test_streaming_shard_writer_tensors_readable(): + """Tensors written by the shard writer can be read back correctly.""" + with tempfile.TemporaryDirectory() as tmpdir: + t = torch.randn(8, 8) + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("weight", t) + weight_map = writer.finalize() + + shard_file = Path(tmpdir) / weight_map["weight"] + with safe_open(str(shard_file), framework="pt") as f: + recovered = f.get_tensor("weight") + assert torch.allclose(recovered, t), "recovered tensor does not match original" + + +# --------------------------------------------------------------------------- +# _postprocess_single_tensor +# --------------------------------------------------------------------------- + + +def test_postprocess_passthrough_normal_key(): + """Non-quantizer weights pass through unchanged.""" + key, val = _postprocess_single_tensor("model.layers.0.self_attn.q_proj.weight", torch.randn(4, 4), 448.0, None) + assert key == "model.layers.0.self_attn.q_proj.weight" + assert val is not None + assert val.shape == (4, 4) + + +def test_postprocess_amax_dropped(): + """weight_quantizer._amax matches skip_keys but has no replacement — dropped.""" + key, val = _postprocess_single_tensor("model.layers.0.weight_quantizer._amax", torch.tensor(1.0), 448.0, None) + assert key is None + assert val is None + + +def test_postprocess_output_quantizer_dropped(): + """output_quantizer keys are always dropped.""" + key, val = _postprocess_single_tensor( + "model.layers.0.output_quantizer._amax", torch.tensor(0.5), 448.0, None + ) + assert key is None + + +def test_postprocess_kv_scale_renamed_and_divided(): + """k_bmm_quantizer._amax is renamed to k_proj.k_scale and divided by maxbound.""" + from modelopt.torch.export.model_config import KV_CACHE_FP8 + + key, val = _postprocess_single_tensor( + "model.layers.0.self_attn.k_bmm_quantizer._amax", + torch.tensor(224.0), + 448.0, + KV_CACHE_FP8, + ) + assert key == "model.layers.0.self_attn.k_proj.k_scale" + assert abs(val.item() - 0.5) < 1e-5 + + +def test_postprocess_scale_squeezed(): + """3D scale tensors with shape[0]==1 are squeezed.""" + t = torch.ones(1, 4, 4) + key, val = _postprocess_single_tensor("model.weight_scale", t, 448.0, None) + assert key == "model.weight_scale" + assert val.shape == (4, 4), f"expected (4, 4), got {val.shape}" + + +def test_postprocess_real_quant_param_dropped(): + """Keys matching RealQuantLinear scale tensors are dropped.""" + from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear + + for q_key in RealQuantLinear.list_of_scale_tensors: + full_key = f"model.layers.0.weight_quantizer.{q_key}" + key, val = _postprocess_single_tensor(full_key, torch.tensor(1.0), 448.0, None) + assert key is None, f"expected None for real quant key '{full_key}'" + + +def test_postprocess_vision_model_summary_idxs_dropped(): + """The vision model summary_idxs parameter is always skipped.""" + key, val = _postprocess_single_tensor( + "vision_model.radio_model.summary_idxs", torch.tensor([0, 1]), 448.0, None + ) + assert key is None From 0c62df660c0ea892cb3e62c6b21abcf5f89d7505 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 23 Jul 2026 14:54:48 -0700 Subject: [PATCH 05/21] refactor(export): simplify streaming-export additions per /simplify review - Extract _KV_CACHE_REPLACEMENTS / _QLORA_REPLACEMENTS / _BASE_SKIP_KEYS / _QLORA_SKIP_KEYS as module-level constants; eliminate the per-call rebuild in both _postprocess_single_tensor and postprocess_state_dict. - Add _maybe_squeeze_scale helper; remove three identical inline squeeze expressions. - Replace _StreamingShardWriter._part_bytes list (used only for sum()) with a scalar _total_bytes accumulator. - Collapse the two GPU-resident named_parameters / named_buffers loops into a single itertools.chain loop. - Move import contextlib to module level (was deferred inside function body). - Initialize export_state_dict = None before the _offloaded branch to prevent a latent NameError on future edits. - Narrow except (ImportError, Exception) to except ImportError in _parse_shard_size so parser errors propagate instead of silently falling through to the manual fallback. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/quant_utils.py | 105 ++++++++------------- modelopt/torch/export/unified_export_hf.py | 25 ++--- 2 files changed, 47 insertions(+), 83 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ac825194bda..912c0003484 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -959,6 +959,36 @@ def from_quantized_weight( raise NotImplementedError(f"quantization format {quantization} not supported") +_KV_CACHE_REPLACEMENTS: dict[str, str] = { + "k_bmm_quantizer._amax": "k_proj.k_scale", + "v_bmm_quantizer._amax": "v_proj.v_scale", + "k_bmm_quantizer._bias_value": "k_proj.k_bias", + "v_bmm_quantizer._bias_value": "v_proj.v_bias", + "input_quantizer._pre_quant_scale": "pre_quant_scale", +} +_QLORA_REPLACEMENTS: dict[str, str] = { + **_KV_CACHE_REPLACEMENTS, + "base_layer.weight": "weight", + "base_layer.input_scale": "input_scale", + "base_layer.weight_scale": "weight_scale", +} +_BASE_SKIP_KEYS: tuple[str, ...] = ( + "output_quantizer", + "_amax", + "_bias_value", + "input_quantizer._pre_quant_scale", + "weight_shape", +) +_QLORA_SKIP_KEYS: tuple[str, ...] = (*_BASE_SKIP_KEYS, "base_layer") + + +def _maybe_squeeze_scale(key: str, value: Any) -> Any: + """Squeeze a leading dim=1 from 3-D scale tensors of shape (1, n, m).""" + if "scale" in key and isinstance(value, torch.Tensor) and value.dim() == 3 and value.shape[0] == 1: + return value.squeeze(0) + return value + + def _postprocess_single_tensor( key: str, value: torch.Tensor, @@ -972,29 +1002,8 @@ def _postprocess_single_tensor( Tied-weight dedup is NOT performed here; callers should pre-compute alias keys from ``model._tied_weights_keys`` and filter them at the call site. """ - replacements = { - "k_bmm_quantizer._amax": "k_proj.k_scale", - "v_bmm_quantizer._amax": "v_proj.v_scale", - "k_bmm_quantizer._bias_value": "k_proj.k_bias", - "v_bmm_quantizer._bias_value": "v_proj.v_bias", - "input_quantizer._pre_quant_scale": "pre_quant_scale", - } - skip_keys = [ - "output_quantizer", - "_amax", - "_bias_value", - "input_quantizer._pre_quant_scale", - "weight_shape", - ] - if is_modelopt_qlora: - replacements.update( - { - "base_layer.weight": "weight", - "base_layer.input_scale": "input_scale", - "base_layer.weight_scale": "weight_scale", - } - ) - skip_keys.append("base_layer") + replacements = _QLORA_REPLACEMENTS if is_modelopt_qlora else _KV_CACHE_REPLACEMENTS + skip_keys = _QLORA_SKIP_KEYS if is_modelopt_qlora else _BASE_SKIP_KEYS # Skip problematic VL model parameters if key == "vision_model.radio_model.summary_idxs": @@ -1010,9 +1019,7 @@ def _postprocess_single_tensor( # Keys not related to quantizers: keep as-is if all(sk not in key for sk in skip_keys): - if "scale" in key and isinstance(value, torch.Tensor) and value.dim() == 3 and value.shape[0] == 1: - value = value.squeeze(0) - return key, value + return key, _maybe_squeeze_scale(key, value) # Apply replacements if the key matches any suffix in the replacements dict for old_suffix, new_suffix in replacements.items(): @@ -1029,14 +1036,7 @@ def _postprocess_single_tensor( "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." ) new_key = prefix + new_suffix - if ( - "scale" in new_key - and isinstance(value, torch.Tensor) - and value.dim() == 3 - and value.shape[0] == 1 - ): - value = value.squeeze(0) - return new_key, value + return new_key, _maybe_squeeze_scale(new_key, value) # Key has a skip_key but no replacement matched — drop it return None, None @@ -1059,31 +1059,8 @@ def postprocess_state_dict( Returns: The filtered state_dict without unnecessary keys like '_amax' and non KV cache output quantizers. """ - replacements = { - "k_bmm_quantizer._amax": "k_proj.k_scale", - "v_bmm_quantizer._amax": "v_proj.v_scale", - "k_bmm_quantizer._bias_value": "k_proj.k_bias", - "v_bmm_quantizer._bias_value": "v_proj.v_bias", - "input_quantizer._pre_quant_scale": "pre_quant_scale", - } - skip_keys = [ - "output_quantizer", - "_amax", - "_bias_value", - "input_quantizer._pre_quant_scale", - "weight_shape", - ] - - # For modelopt-trained LoRA models, we need to remove the base_layer prefix from the keys for deployment - if is_modelopt_qlora: - replacements.update( - { - "base_layer.weight": "weight", - "base_layer.input_scale": "input_scale", - "base_layer.weight_scale": "weight_scale", - } - ) - skip_keys.append("base_layer") + replacements = _QLORA_REPLACEMENTS if is_modelopt_qlora else _KV_CACHE_REPLACEMENTS + skip_keys = _QLORA_SKIP_KEYS if is_modelopt_qlora else _BASE_SKIP_KEYS post_state_dict = {} @@ -1119,15 +1096,7 @@ def postprocess_state_dict( post_state_dict[prefix + new_suffix] = value break - # Squeeze scales with a leading dimension of 1 - for key, value in post_state_dict.items(): - if ( - "scale" in key - and isinstance(value, torch.Tensor) - and value.dim() == 3 - and value.shape[0] == 1 - ): - post_state_dict[key] = value.squeeze(0) + post_state_dict = {k: _maybe_squeeze_scale(k, v) for k, v in post_state_dict.items()} # remove real quant parameters from the state dict keys_to_delete = [] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 00ea4cb6b7a..8d71d6088d0 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,6 +15,8 @@ """Code that export quantized Hugging Face models for deployment.""" +import contextlib +import itertools import json import re import tempfile @@ -955,7 +957,7 @@ def __init__(self, export_dir: Path | str, max_shard_size: int) -> None: self._buffer: dict[str, torch.Tensor] = {} self._buffer_bytes: int = 0 self._part_files: list[Path] = [] - self._part_bytes: list[int] = [] + self._total_bytes: int = 0 # Maps tensor key → part-file index (recorded at flush time) self._key_to_part: dict[str, int] = {} @@ -968,7 +970,7 @@ def _flush(self) -> None: for key in self._buffer: self._key_to_part[key] = part_idx self._part_files.append(part_path) - self._part_bytes.append(self._buffer_bytes) + self._total_bytes += self._buffer_bytes self._buffer = {} self._buffer_bytes = 0 @@ -1002,7 +1004,7 @@ def finalize(self) -> dict[str, str]: key: f"model-{part_idx + 1:05d}-of-{n_shards:05d}.safetensors" for key, part_idx in self._key_to_part.items() } - total_size = sum(self._part_bytes) + total_size = self._total_bytes index_path = self._export_dir / "model.safetensors.index.json" with open(index_path, "w") as f: json.dump({"metadata": {"total_size": total_size}, "weight_map": weight_map}, f) @@ -1015,7 +1017,7 @@ def _parse_shard_size(size: int | str) -> int: from transformers.utils import convert_file_size_to_int return convert_file_size_to_int(size) - except (ImportError, Exception): + except ImportError: pass if isinstance(size, int): return size @@ -1197,23 +1199,15 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: _stream_tensor(full_key, tensor) # GPU-resident parameters and buffers (not covered by the above loops) - for name, param in model.named_parameters(): - if name in seen_keys or param is None or param.is_meta: + for name, tensor in itertools.chain(model.named_parameters(), model.named_buffers()): + if name in seen_keys or tensor is None or tensor.is_meta: continue seen_keys.add(name) - _stream_tensor(name, param) - - for name, buf in model.named_buffers(): - if name in seen_keys or buf is None or buf.is_meta: - continue - seen_keys.add(name) - _stream_tensor(name, buf) + _stream_tensor(name, tensor) writer.finalize() # Write non-weight artifacts (model.save_pretrained skipped to avoid OOM) - import contextlib - _sanitize_generation_config_for_save(model) model.config.save_pretrained(str(export_dir)) gc = getattr(model, "generation_config", None) @@ -1846,6 +1840,7 @@ def export_hf_checkpoint( # Streaming path writes shard files layer-by-layer without accumulating the full # state dict in RAM (peak = 1 layer + 1 shard buffer vs. ~764 GiB for Ultra 550B). _offloaded = _has_accelerate_offload(model) + export_state_dict = None try: if _offloaded: From 9164813b7ffb37dca64ac5a8222e2bb72c64b17f Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 23 Jul 2026 15:28:22 -0700 Subject: [PATCH 06/21] fix(export): harden streaming export path per code review - Gate tied-weight dedup on model.config.tie_word_embeddings to avoid incorrectly dropping lm_head.weight when embeddings are not tied - Add _is_persistent_buffer helper; filter named_buffers() in the GPU-resident pass to match state_dict() semantics - Add .contiguous() before .cpu() in _stream_tensor to handle non-contiguous views from accelerate writeback - Copy trust_remote_code modeling files via model.save_pretrained( state_dict={}) with single-shard rename protection so custom model class files are not lost - Refactor example_utils.py shims: module-level _FP8BF16Fallback class, _install_transformers_compat_shims() called lazily from get_model(), explicit UserWarning for lossy BF16 fallback path Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 110 +++++++++++---------- modelopt/torch/export/unified_export_hf.py | 59 ++++++++--- 2 files changed, 103 insertions(+), 66 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index b3b16467d01..b23e52043e7 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import copy import glob import hashlib @@ -30,60 +31,6 @@ import torch import transformers - -# Shim for is_torch_fx_available removed in transformers >=5.x; older model files (e.g. -# DeepSeek-R1 bundled modeling_deepseek.py) import it from transformers.utils.import_utils. -try: - from transformers.utils.import_utils import is_torch_fx_available # noqa: F401 -except ImportError: - import transformers.utils.import_utils as _tui - - _tui.is_torch_fx_available = lambda: False # type: ignore[attr-defined] - -# Shim for broken flash_attn installs (undefined symbol in .so). Probe the actual import; -# if it fails, force transformers' availability checks to return False so bundled remote-code -# model files (e.g. modeling_deepseek.py) skip the flash_attn import block. -# Must patch both transformers.utils.import_utils AND transformers.utils since bundled models -# import from either location. -try: - import flash_attn as _flash_attn_probe # noqa: F401 -except Exception: - import transformers.utils as _tu - import transformers.utils.import_utils as _tui - - for _mod in (_tu, _tui): - _mod.is_flash_attn_2_available = lambda: False # type: ignore[attr-defined] - _mod.is_flash_attn_available = lambda: False # type: ignore[attr-defined] - _mod.is_flash_attn_greater_or_equal_2_10 = lambda: False # type: ignore[attr-defined] - -# On nodes without the `kernels` package, DSR1 block-scaled FP8 matmul fails at import. -# Patch the loader with a BF16 dequant fallback so calibration forward passes succeed -# (amax collection only — not suitable for production inference). -try: - import transformers.integrations.finegrained_fp8 as _ff8 - - try: - _ff8._load_finegrained_fp8_kernel() - except ImportError: - - class _FP8BF16Fallback: - @staticmethod - def matmul(input, weight, weight_scale_inv, block_size, output_dtype=None, activation_scale=None): - out_f, in_f = weight.shape[-2], weight.shape[-1] - nb_out, nb_in = weight_scale_inv.shape[-2], weight_scale_inv.shape[-1] - scale = ( - weight_scale_inv.float() - .repeat_interleave(out_f // nb_out, -2) - .repeat_interleave(in_f // nb_in, -1) - ) - w_bf16 = (weight.float() * scale).to(torch.bfloat16) - out = torch.nn.functional.linear(input.to(torch.bfloat16), w_bf16) - return out if output_dtype is None else out.to(output_dtype) - - _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback # type: ignore[attr-defined] -except Exception: - pass - from accelerate import infer_auto_device_map, init_empty_weights from accelerate.utils import get_max_memory from safetensors import safe_open @@ -165,6 +112,60 @@ def validate_fsdp2_supported(args, config): + "\nRemove --use_fsdp2 or use a standard causal-LM checkpoint." ) +class _FP8BF16Fallback: + """BF16 dequant fallback for block-scaled FP8 matmul when the kernels package is absent. + + Calibration amax collection only — not accurate for production inference. + """ + + @staticmethod + def matmul(input, weight, weight_scale_inv, block_size, output_dtype=None, activation_scale=None): + out_f, in_f = weight.shape[-2], weight.shape[-1] + nb_out, nb_in = weight_scale_inv.shape[-2], weight_scale_inv.shape[-1] + scale = ( + weight_scale_inv.float() + .repeat_interleave(out_f // nb_out, -2) + .repeat_interleave(in_f // nb_in, -1) + ) + w_bf16 = (weight.float() * scale).to(torch.bfloat16) + out = torch.nn.functional.linear(input.to(torch.bfloat16), w_bf16) + return out if output_dtype is None else out.to(output_dtype) + + +def _install_transformers_compat_shims() -> None: + """Patch transformers so older remote-code models (e.g. DeepSeek-R1) load on + newer/partial installs. Call once before loading a trust_remote_code checkpoint.""" + import transformers.utils as _tu + import transformers.utils.import_utils as _tui + + # transformers >=5 removed is_torch_fx_available; older bundled model files still import it. + if not hasattr(_tui, "is_torch_fx_available"): + _tui.is_torch_fx_available = lambda: False # type: ignore[attr-defined] + + # Broken flash_attn installs (.so undefined-symbol) crash at import time, not find_spec time. + # Force transformers' availability checks to False so bundled models skip the flash-attn path. + try: + import flash_attn # noqa: F401 + except Exception: + for _mod in (_tu, _tui): + for _fn in ("is_flash_attn_2_available", "is_flash_attn_available", + "is_flash_attn_greater_or_equal_2_10"): + setattr(_mod, _fn, lambda: False) + + # No `kernels` package → block-scaled FP8 matmul fails; swap in lossy BF16 fallback. + with contextlib.suppress(Exception): + import transformers.integrations.finegrained_fp8 as _ff8 + try: + _ff8._load_finegrained_fp8_kernel() + except ImportError: + warnings.warn( + "finegrained-fp8 kernel unavailable; using a lossy BF16 dequant fallback " + "for FP8 matmul. Suitable for calibration amax collection only.", + UserWarning, + stacklevel=2, + ) + _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback # type: ignore[attr-defined] + def run_nemotron_vl_preview( full_model, @@ -735,6 +736,7 @@ def get_model( max_cpu_memory_gb=None, max_gpu_memory_gb=None, ): + _install_transformers_compat_shims() print(f"Initializing model from {ckpt_path}") _disk_offload = offload_folder is not None diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 8d71d6088d0..be845f88726 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,7 +15,6 @@ """Code that export quantized Hugging Face models for deployment.""" -import contextlib import itertools import json import re @@ -1111,8 +1110,15 @@ def _export_transformers_checkpoint_streaming( kv_cache_max_bound = 448 kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] - # --- Tied alias keys to skip (data_ptr() is unreliable for disk-offloaded weights) --- - raw_tied_keys: set[str] = set(getattr(model, "_tied_weights_keys", None) or []) + # --- Tied alias keys to skip --- + # data_ptr() is unreliable for disk-offloaded weights, so we use _tied_weights_keys. + # Only apply when tie_word_embeddings=True: _tied_weights_keys can list keys whose + # weights are not actually shared (e.g. if the model was saved with tie_word_embeddings=False + # but the attribute was never cleared), which would incorrectly drop lm_head.weight. + if getattr(model.config, "tie_word_embeddings", False): + raw_tied_keys: set[str] = set(getattr(model, "_tied_weights_keys", None) or []) + else: + raw_tied_keys: set[str] = set() # --- Name mapper for per-tensor key reversal --- # Tensor names are applied inline; quant config names are handled by the caller. @@ -1138,6 +1144,14 @@ def _export_transformers_checkpoint_streaming( ) decoder_layer_ids = {id(m) for m in decoder_layers} + # --- Persistent-buffer predicate (mirrors state_dict() which excludes non-persistent) --- + def _is_persistent_buffer(name: str) -> bool: + parts = name.split(".") + mod: nn.Module = model + for part in parts[:-1]: + mod = getattr(mod, part, mod) + return parts[-1] not in getattr(mod, "_non_persistent_buffers_set", frozenset()) + # --- Stream tensors to shard files --- shard_size_bytes = _parse_shard_size(max_shard_size) writer = _StreamingShardWriter(export_dir, shard_size_bytes) @@ -1154,7 +1168,7 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: new_key = name_mapper(new_key) if new_key in tied_alias_keys: return - writer.add(new_key, new_value.detach().cpu()) + writer.add(new_key, new_value.detach().contiguous().cpu()) # Decoder layers (offloaded: materialize one at a time) for layer_name, layer_module in model.named_modules(): @@ -1198,8 +1212,12 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: seen_keys.add(full_key) _stream_tensor(full_key, tensor) - # GPU-resident parameters and buffers (not covered by the above loops) - for name, tensor in itertools.chain(model.named_parameters(), model.named_buffers()): + # GPU-resident parameters and persistent buffers (not covered by the above loops). + # named_buffers() includes non-persistent buffers that state_dict() excludes; filter them. + for name, tensor in itertools.chain( + model.named_parameters(), + ((n, b) for n, b in model.named_buffers() if _is_persistent_buffer(n)), + ): if name in seen_keys or tensor is None or tensor.is_meta: continue seen_keys.add(name) @@ -1207,13 +1225,30 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: writer.finalize() - # Write non-weight artifacts (model.save_pretrained skipped to avoid OOM) + # Write non-weight artifacts: config.json, generation_config.json, tokenizer, and + # the custom modeling *.py files that trust_remote_code models (e.g. NemotronH) need. + # model.save_pretrained with an empty state dict is the only reliable way to trigger + # transformers' custom-code copy logic without holding the full checkpoint in RAM. + # Protect any real shard already written by _StreamingShardWriter (single-shard path + # renames its output to model.safetensors, which save_pretrained would overwrite). + _single_shard = export_dir / "model.safetensors" + _protected = export_dir / "__modelopt_protected_model.safetensors" + if _single_shard.exists(): + _single_shard.rename(_protected) + _sanitize_generation_config_for_save(model) - model.config.save_pretrained(str(export_dir)) - gc = getattr(model, "generation_config", None) - if gc is not None: - with contextlib.suppress(Exception): - gc.save_pretrained(str(export_dir)) + _patches = _patch_revert_weight_conversion() + try: + model.save_pretrained(str(export_dir), state_dict={}) + finally: + _unpatch_revert_weight_conversion(_patches) + + # Remove the empty placeholder shard save_pretrained created for state_dict={}. + if _single_shard.exists() and _single_shard.stat().st_size < 512: + _single_shard.unlink() + # Restore the real single-shard if we protected it. + if _protected.exists(): + _protected.rename(_single_shard) return None, quant_config From ffb16e84485246eedccdac3a89853b518e54beed Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:42:03 -0700 Subject: [PATCH 07/21] fix(export): avoid save_pretrained shared-tensor crash on MoE models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit model.save_pretrained(state_dict={}) triggers safetensors' shared-tensor validation on the live model parameters even when no weights are being saved. DSR1 and other MoE models have expert weights that share storage across layers, so this check always fails — crashing the export after all shards are correctly written and leaving hf_quant_config.json unwritten. Replace with targeted saves: - model.config.save_pretrained() for config.json - model.generation_config.save_pretrained() for generation_config.json - shutil.copy2(*.py) for trust_remote_code custom modeling files Also add missing contextlib and shutil stdlib imports. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/unified_export_hf.py | 39 +++++++++++----------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index be845f88726..109e929c24e 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,9 +15,11 @@ """Code that export quantized Hugging Face models for deployment.""" +import contextlib import itertools import json import re +import shutil import tempfile import warnings from builtins import ValueError @@ -1225,30 +1227,29 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: writer.finalize() - # Write non-weight artifacts: config.json, generation_config.json, tokenizer, and - # the custom modeling *.py files that trust_remote_code models (e.g. NemotronH) need. - # model.save_pretrained with an empty state dict is the only reliable way to trigger - # transformers' custom-code copy logic without holding the full checkpoint in RAM. - # Protect any real shard already written by _StreamingShardWriter (single-shard path - # renames its output to model.safetensors, which save_pretrained would overwrite). - _single_shard = export_dir / "model.safetensors" - _protected = export_dir / "__modelopt_protected_model.safetensors" - if _single_shard.exists(): - _single_shard.rename(_protected) - + # Write non-weight artifacts: config.json, generation_config.json, and the custom + # modeling *.py files that trust_remote_code models (e.g. NemotronH) need. + # We avoid model.save_pretrained(state_dict={}) here because MoE models (e.g. DSR1) + # have expert weights that share underlying storage across layers; safetensors' shared- + # tensor check fires even when saving an empty state dict, crashing the export after + # all shards are already written correctly. _sanitize_generation_config_for_save(model) _patches = _patch_revert_weight_conversion() try: - model.save_pretrained(str(export_dir), state_dict={}) + model.config.save_pretrained(str(export_dir)) finally: _unpatch_revert_weight_conversion(_patches) - - # Remove the empty placeholder shard save_pretrained created for state_dict={}. - if _single_shard.exists() and _single_shard.stat().st_size < 512: - _single_shard.unlink() - # Restore the real single-shard if we protected it. - if _protected.exists(): - _protected.rename(_single_shard) + if hasattr(model, "generation_config") and model.generation_config is not None: + with contextlib.suppress(Exception): + model.generation_config.save_pretrained(str(export_dir)) + + # Copy custom modeling *.py files for trust_remote_code checkpoints. + _src_dir = Path(getattr(model.config, "_name_or_path", "") or "") + if _src_dir.is_dir(): + for _py in _src_dir.glob("*.py"): + _dst = export_dir / _py.name + if not _dst.exists(): + shutil.copy2(_py, _dst) return None, quant_config From 1cc0ec8fc732f2c7fd1b947bb56af145565b1851 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 16:00:28 -0700 Subject: [PATCH 08/21] refactor(export): remove dead offloaded branch, early-return dispatch, shard regression test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Delete _process_quantized_modules_offloaded: dead from export_hf_checkpoint (public path dispatches to streaming writer; function accumulated full state dict in RAM, defeating offload) - Replace if _offloaded: branch inside _export_transformers_checkpoint with NotImplementedError — streaming path owns that case via export_hf_checkpoint - Add _write_hf_export_config helper (hf_quant_config.json + config.json patching) - Refactor export_hf_checkpoint: early return after offloaded path eliminates three scattered if not _offloaded: guards; both paths share the helper - Update meta-guard error message to point to export_hf_checkpoint - Tests: remove test_non_decoder_offloaded_tensors_are_collected (tests deleted fn); add test_multi_shard_files_exist_after_finalize (regression: save_pretrained(state_dict={}) triggered transformers cleanup loop that deleted model-NNNNN-of-NNNNN shards) Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/unified_export_hf.py | 231 ++++++------------ .../unit/torch/export/test_offload_export.py | 107 +++----- 2 files changed, 105 insertions(+), 233 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 109e929c24e..8d80ee59516 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -577,9 +577,8 @@ def _export_quantized_weight( 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, export must run inside an " - "enable_weight_access_and_writeback context. Use the offload-aware export path " - "(_process_quantized_modules_offloaded) rather than _process_quantized_modules." + "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. @@ -855,93 +854,6 @@ def _has_accelerate_offload(model: nn.Module) -> bool: return False -def _process_quantized_modules_offloaded( - model: nn.Module, - dtype: torch.dtype, - is_modelopt_qlora: bool = False, -) -> dict[str, Any]: - """Export quantized weights for a disk/CPU-offloaded model, one layer at a time. - - Decoder layers are processed one at a time via enable_weight_access_and_writeback. - Non-decoder modules that are also disk-offloaded (embed_tokens, norms, lm_head) are - materialized individually; any quantized non-decoder module (e.g. lm_head) has its - export handler invoked in the same context. - - Returns a full-model state dict with no meta tensors. - """ - from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear - from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback - from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector - - decoder_layers = LayerActivationCollector.get_decoder_layers(model) - if decoder_layers is None: - raise RuntimeError( - "Disk/CPU-offloaded export requires discoverable decoder layers. " - "The model architecture is not supported by LayerActivationCollector." - ) - decoder_layer_ids = {id(m) for m in decoder_layers} - - ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) - layer_tensors: dict[str, torch.Tensor] = {} - - for name, module in model.named_modules(): - if id(module) not in decoder_layer_ids: - continue - # writeback=False: weights are captured in layer_tensors below; no need to promote - # the quantized values back to the offload store on context exit. - with enable_weight_access_and_writeback(module, module, writeback=False): - for sub_name, sub_mod in module.named_modules(): - full_name = f"{name}.{sub_name}" if sub_name else name - _dispatch_export_handler(full_name, sub_mod, ctx) - - # Mirror the non-offloaded path: reconstruct fused MoE per-expert weights - # into 3D tensors BEFORE snapshotting, so captured keys match the original - # MoE format (e.g. moe.up_proj.weight [N, out, in]). - _reconstruct_fused_moe_linear(module) - - # Snapshot inside the context: post-exit, post_forward re-offloads params to meta. - prefix = f"{name}." if name else "" - for key, tensor in module.state_dict().items(): - assert not tensor.is_meta, ( - f"Expected real tensor for '{prefix + key}' inside materialization context" - ) - layer_tensors[prefix + key] = tensor.detach().cpu() - - # Also collect non-decoder modules that are disk-offloaded (embed_tokens, norms, lm_head). - # model.state_dict() returns meta for these; materialize, run any export handlers - # (e.g. a quantized lm_head), then snapshot the real tensors. - for name, module in model.named_modules(): - if id(module) in decoder_layer_ids: - continue - if not hasattr(module, "_hf_hook"): - continue - if _get_offload_hook(module._hf_hook) is None: - continue - # Only handle modules that have DIRECT meta parameters/buffers. - # Child decoder layers (already captured above) must not be re-collected. - if not ( - any(p is not None and p.is_meta for p in module._parameters.values()) - or any(b is not None and b.is_meta for b in module._buffers.values()) - ): - continue - with enable_weight_access_and_writeback(module, module, writeback=False): - for sub_name, sub_mod in module.named_modules(): - full_name = f"{name}.{sub_name}" if sub_name else name - _dispatch_export_handler(full_name, sub_mod, ctx) - prefix = f"{name}." if name else "" - for key, tensor in module.state_dict().items(): - if not tensor.is_meta: - layer_tensors[prefix + key] = tensor.detach().cpu() - - # model.state_dict() fills in non-offloaded parts (GPU-resident tensors). - # layer_tensors overrides both decoder-layer placeholders and non-decoder - # offloaded placeholders so the returned dict contains no meta tensors. - full_sd = model.state_dict() - full_sd.update(layer_tensors) - return full_sd - - class _StreamingShardWriter: """Write tensors to safetensors shard files without accumulating the full state dict. @@ -1309,7 +1221,7 @@ def _export_transformers_checkpoint( remove_hook_from_module(model, recurse=True) except ImportError: - warnings.warn("accelerate is not installed, hooks will not be removed") + pass # no accelerate installed → no offload hooks exist to remove quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) @@ -1350,8 +1262,10 @@ def _export_transformers_checkpoint( from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear if _offloaded: - # MoE reconstruction happens per-layer inside _process_quantized_modules_offloaded. - quantized_state_dict = _process_quantized_modules_offloaded(model, dtype, is_modelopt_qlora) + raise NotImplementedError( + "_export_transformers_checkpoint does not support disk/CPU-offloaded models. " + "Use export_hf_checkpoint() which dispatches to _export_transformers_checkpoint_streaming." + ) else: _reconstruct_fused_moe_linear(model) @@ -1816,6 +1730,38 @@ def export_speculative_decoding( exporter.export(export_dir, dtype) +def _write_hf_export_config( + model: nn.Module, + hf_quant_config: dict | None, + export_dir: Path, +) -> None: + """Write hf_quant_config.json (if quantized) and embed quantization_config into config.json.""" + quantization_details = (hf_quant_config or {}).get("quantization", {}) + is_quantized_export = ( + quantization_details.get("quant_algo") is not None + or quantization_details.get("kv_cache_quant_algo") is not None + ) + if is_quantized_export: + with open(f"{export_dir}/hf_quant_config.json", "w") as file: + json.dump(hf_quant_config, file, indent=4) + hf_quant_config = convert_hf_quant_config_format(hf_quant_config) + else: + hf_quant_config = None + + original_config = f"{export_dir}/config.json" + with open(original_config) as file: + config_data = json.load(file) + sanitize_hf_config_for_deployment(config_data, model) + if hf_quant_config is not None: + config_data["quantization_config"] = hf_quant_config + if export_sparse_attention_config is not None: + sparse_attn_config = export_sparse_attention_config(model) + if sparse_attn_config is not None: + config_data["sparse_attention_config"] = sparse_attn_config + with open(original_config, "w") as file: + json.dump(config_data, file, indent=4) + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1876,7 +1822,6 @@ def export_hf_checkpoint( # Streaming path writes shard files layer-by-layer without accumulating the full # state dict in RAM (peak = 1 layer + 1 shard buffer vs. ~764 GiB for Ultra 550B). _offloaded = _has_accelerate_offload(model) - export_state_dict = None try: if _offloaded: @@ -1897,15 +1842,27 @@ def export_hf_checkpoint( max_shard_size=max_shard_size, **kwargs, ) - else: - post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) + if getattr(model, "hf_quantizer", None) is not None: + model.hf_quantizer = None + try: + name_mapper = build_reverse_name_mapper(model) + if name_mapper is not None and hf_quant_config: + revert_quant_config_names(hf_quant_config.get("quantization", {}), name_mapper) + except Exception as exc: + warnings.warn( + f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor " + "names may not match the original HF hub checkpoint." + ) + _write_hf_export_config(model, hf_quant_config, export_dir) + return + + post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) # Remove hf_quantizer from model so post_state_dict can be exported. if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None - if not _offloaded: - export_state_dict = {**post_state_dict, **(extra_state_dict or {})} + export_state_dict = {**post_state_dict, **(extra_state_dict or {})} # transformers may have applied a load-time conversion_mapping (fused gate_up_proj, # renamed MoE leaves, reordered model/language_model prefix), so the in-memory names @@ -1920,10 +1877,7 @@ def export_hf_checkpoint( # weights and config so they stay mutually consistent. try: name_mapper = build_reverse_name_mapper(model) - if not _offloaded: - # Streaming path applies per-tensor renaming inline inside - # _export_transformers_checkpoint_streaming; skip full-dict reversal here. - export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict) + export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict) if name_mapper is not None and hf_quant_config: revert_quant_config_names(hf_quant_config.get("quantization", {}), name_mapper) except Exception as exc: @@ -1936,65 +1890,26 @@ def export_hf_checkpoint( if is_distributed and torch.distributed.get_rank() != 0: return - # Only treat the export as quantized when at least one quant_algo field is set. - # get_quant_config always returns a dict (even for sparsity-only or unmodified models), - # so emitting hf_quant_config.json unconditionally produces a file with - # "quant_algo": null that downstream loaders (e.g. TensorRT-LLM) reject as a - # malformed pre-quantized checkpoint. - quantization_details = (hf_quant_config or {}).get("quantization", {}) - is_quantized_export = ( - quantization_details.get("quant_algo") is not None - or quantization_details.get("kv_cache_quant_algo") is not None - ) - - if is_quantized_export: - # Save hf_quant_config.json for backward compatibility - with open(f"{export_dir}/hf_quant_config.json", "w") as file: - json.dump(hf_quant_config, file, indent=4) - - hf_quant_config = convert_hf_quant_config_format(hf_quant_config) - else: - hf_quant_config = None + # Keep transformers' own revert_weight_conversion disabled (the quant-aware reverse + # above replaces it): it can't handle quantized state dicts (RuntimeError on 0-d scalar + # scale tensors). Patch both the source and importing module since modeling_utils does + # `from core_model_loading import revert_weight_conversion`. + _patches = _patch_revert_weight_conversion() - if not _offloaded: - # Keep transformers' own revert_weight_conversion disabled (the quant-aware reverse - # above replaces it): it can't handle quantized state dicts (RuntimeError on 0-d scalar - # scale tensors). Patch both the source and importing module since modeling_utils does - # `from core_model_loading import revert_weight_conversion`. - _patches = _patch_revert_weight_conversion() - - _sanitize_generation_config_for_save(model) - - # TODO: parallelize the disk write across ranks (avoid single-process speed + rank-0 OOM). - try: - model.save_pretrained( - export_dir, - state_dict=export_state_dict, - save_modelopt_state=save_modelopt_state, - max_shard_size=max_shard_size, - ) - finally: - _unpatch_revert_weight_conversion(_patches) + _sanitize_generation_config_for_save(model) - original_config = f"{export_dir}/config.json" - config_data = {} - - with open(original_config) as file: - config_data = json.load(file) - - sanitize_hf_config_for_deployment(config_data, model) - - if hf_quant_config is not None: - config_data["quantization_config"] = hf_quant_config - - # Add sparse attention config if available - if export_sparse_attention_config is not None: - sparse_attn_config = export_sparse_attention_config(model) - if sparse_attn_config is not None: - config_data["sparse_attention_config"] = sparse_attn_config + # TODO: parallelize the disk write across ranks (avoid single-process speed + rank-0 OOM). + try: + model.save_pretrained( + export_dir, + state_dict=export_state_dict, + save_modelopt_state=save_modelopt_state, + max_shard_size=max_shard_size, + ) + finally: + _unpatch_revert_weight_conversion(_patches) - with open(original_config, "w") as file: - json.dump(config_data, file, indent=4) + _write_hf_export_config(model, hf_quant_config, export_dir) except Exception as e: warnings.warn( diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 05f9974df6e..d6f5747ca69 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -35,7 +35,6 @@ from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _has_accelerate_offload, - _process_quantized_modules_offloaded, _StreamingShardWriter, ) @@ -123,78 +122,6 @@ def test_meta_guard_not_raised_for_real_weight(): _export_quantized_weight(linear, torch.float32) -# --------------------------------------------------------------------------- -# _process_quantized_modules_offloaded — non-decoder materialization -# --------------------------------------------------------------------------- - - -def test_non_decoder_offloaded_tensors_are_collected(): - """Non-decoder modules with disk-offload hooks must have no meta tensors in the result. - - Reproduces the NemotronH 550B crash: embed_tokens (and norm, lm_head) are - disk-offloaded and return meta from model.state_dict(). After - revert_weight_conversion_quant_aware renames them to hub-original names, transformers' - remove_tied_weights_from_state_dict tries to look them up in the model by that name - and crashes. Fix: _process_quantized_modules_offloaded materialises non-decoder - offloaded modules directly so the returned state dict contains no meta tensors. - - The decoder layer here is NOT disk-offloaded (all weights GPU-resident) so the - decoder-layer loop exercises the null-context path and we focus on the non-decoder - collection pass that was previously missing. - """ - - class _TinyLayer(nn.Module): - def __init__(self): - super().__init__() - self.proj = nn.Linear(8, 8, bias=False) - - def forward(self, x): - return self.proj(x) - - class _TinyModel(nn.Module): - def __init__(self): - super().__init__() - self.embed = nn.Embedding(16, 8) - self.layers = nn.ModuleList([_TinyLayer()]) - - def forward(self, x): - return self.layers[0](self.embed(x)) - - model = _TinyModel() - - # Install a CPU-offload hook on embed ONLY (non-decoder module). - # The decoder layer is left GPU-resident so enable_weight_access_and_writeback - # returns a no-op nullcontext and the decoder-layer state_dict() returns real tensors. - embed_val = model.embed.weight.data.clone().cpu() - embed_weights_map = {"weight": embed_val} - embed_hook = AlignDevicesHook( - execution_device="cpu", offload=True, weights_map=embed_weights_map - ) - add_hook_to_module(model.embed, embed_hook) - set_module_tensor_to_device(model.embed, "weight", "meta") - - from unittest.mock import patch - - with patch( - "modelopt.torch.quantization.utils.layerwise_calib" - ".LayerActivationCollector.get_decoder_layers", - return_value=list(model.layers), - ): - result = _process_quantized_modules_offloaded(model, torch.float32) - - assert "embed.weight" in result, "embed.weight missing from state dict" - emb = result["embed.weight"] - assert not emb.is_meta, "embed.weight must not be meta in exported state dict" - assert emb.shape == (16, 8) - - assert "layers.0.proj.weight" in result - assert not result["layers.0.proj.weight"].is_meta - - for key, val in result.items(): - if isinstance(val, torch.Tensor): - assert not val.is_meta, f"meta tensor found for key '{key}'" - - # --------------------------------------------------------------------------- # _StreamingShardWriter # --------------------------------------------------------------------------- @@ -236,6 +163,32 @@ def test_streaming_shard_writer_multi_shard(): assert index["metadata"]["total_size"] > 0 +def test_multi_shard_files_exist_after_finalize(): + """All numbered shard files referenced in the index must exist on disk after finalize(). + + Regression guard: an earlier code path called model.save_pretrained(state_dict={}) after + finalize(), triggering transformers' stale-shard cleanup loop which matched and deleted + every model-NNNNN-of-NNNNN.safetensors file because filename_to_tensors was empty. + """ + with tempfile.TemporaryDirectory() as tmpdir: + writer = _StreamingShardWriter(tmpdir, max_shard_size=64) + writer.add("x", torch.ones(4, 4)) + writer.add("y", torch.ones(4, 4)) + writer.finalize() + + index_path = Path(tmpdir) / "model.safetensors.index.json" + assert index_path.exists() + with open(index_path) as f: + index = json.load(f) + + for key, shard_name in index["weight_map"].items(): + shard_path = Path(tmpdir) / shard_name + assert shard_path.exists(), ( + f"Shard '{shard_name}' (for key '{key}') missing from disk after finalize()" + ) + assert shard_path.stat().st_size > 0, f"Shard file {shard_name} is empty" + + def test_streaming_shard_writer_tensors_readable(): """Tensors written by the shard writer can be read back correctly.""" with tempfile.TemporaryDirectory() as tmpdir: @@ -257,7 +210,9 @@ def test_streaming_shard_writer_tensors_readable(): def test_postprocess_passthrough_normal_key(): """Non-quantizer weights pass through unchanged.""" - key, val = _postprocess_single_tensor("model.layers.0.self_attn.q_proj.weight", torch.randn(4, 4), 448.0, None) + key, val = _postprocess_single_tensor( + "model.layers.0.self_attn.q_proj.weight", torch.randn(4, 4), 448.0, None + ) assert key == "model.layers.0.self_attn.q_proj.weight" assert val is not None assert val.shape == (4, 4) @@ -265,7 +220,9 @@ def test_postprocess_passthrough_normal_key(): def test_postprocess_amax_dropped(): """weight_quantizer._amax matches skip_keys but has no replacement — dropped.""" - key, val = _postprocess_single_tensor("model.layers.0.weight_quantizer._amax", torch.tensor(1.0), 448.0, None) + key, val = _postprocess_single_tensor( + "model.layers.0.weight_quantizer._amax", torch.tensor(1.0), 448.0, None + ) assert key is None assert val is None From 5fc72c52e466dbea020577ba3ce284bc7886891e Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:20:04 +0000 Subject: [PATCH 09/21] chore: document new offload recipe and fix pre-commit failures Addresses PR review: recipe-doc parity, ruff, and mypy all fail on this branch. - modelopt_recipes/ptq.md: add the nvfp4_experts_only-kv_fp8_layerwise_offload row and bump the summary count to 21, restoring the parity enforced by tests/unit/recipe/test_recipe_docs.py. - Apply ruff format to example_utils.py and quant_utils.py. - test_offload_export.py: wrap st.keys() in list() to clear SIM118. Ruff's own suggested fix (iterating the handle directly) would break the test -- safetensors safe_open handles are not iterable. - Fix 5 pre-existing mypy errors: duplicate raw_tied_keys annotation, missing None guard on the _postprocess_single_tensor result (its documented contract is to return (None, None) to skip), rebinding hf_quant_config to a different type in _write_hf_export_config, and two stale type: ignore comments. No behavior change. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 18 ++++++++++----- modelopt/torch/export/quant_utils.py | 7 +++++- modelopt/torch/export/unified_export_hf.py | 22 +++++++++---------- modelopt_recipes/ptq.md | 3 ++- tests/gpu/torch/export/test_offload_export.py | 2 +- 5 files changed, 32 insertions(+), 20 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index b23e52043e7..5078e6d9b4f 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -119,7 +119,9 @@ class _FP8BF16Fallback: """ @staticmethod - def matmul(input, weight, weight_scale_inv, block_size, output_dtype=None, activation_scale=None): + def matmul( + input, weight, weight_scale_inv, block_size, output_dtype=None, activation_scale=None + ): out_f, in_f = weight.shape[-2], weight.shape[-1] nb_out, nb_in = weight_scale_inv.shape[-2], weight_scale_inv.shape[-1] scale = ( @@ -140,7 +142,7 @@ def _install_transformers_compat_shims() -> None: # transformers >=5 removed is_torch_fx_available; older bundled model files still import it. if not hasattr(_tui, "is_torch_fx_available"): - _tui.is_torch_fx_available = lambda: False # type: ignore[attr-defined] + _tui.is_torch_fx_available = lambda: False # Broken flash_attn installs (.so undefined-symbol) crash at import time, not find_spec time. # Force transformers' availability checks to False so bundled models skip the flash-attn path. @@ -148,13 +150,17 @@ def _install_transformers_compat_shims() -> None: import flash_attn # noqa: F401 except Exception: for _mod in (_tu, _tui): - for _fn in ("is_flash_attn_2_available", "is_flash_attn_available", - "is_flash_attn_greater_or_equal_2_10"): + for _fn in ( + "is_flash_attn_2_available", + "is_flash_attn_available", + "is_flash_attn_greater_or_equal_2_10", + ): setattr(_mod, _fn, lambda: False) # No `kernels` package → block-scaled FP8 matmul fails; swap in lossy BF16 fallback. with contextlib.suppress(Exception): import transformers.integrations.finegrained_fp8 as _ff8 + try: _ff8._load_finegrained_fp8_kernel() except ImportError: @@ -164,7 +170,7 @@ def _install_transformers_compat_shims() -> None: UserWarning, stacklevel=2, ) - _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback # type: ignore[attr-defined] + _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback def run_nemotron_vl_preview( @@ -719,7 +725,7 @@ def _fmt_max_memory(max_memory: dict) -> str: parts = [] for key in sorted(max_memory.keys(), key=lambda k: (isinstance(k, str), k)): val = max_memory[key] - label = f"{val / 1024 ** 3:.1f} GiB" if isinstance(val, int) else str(val) + label = f"{val / 1024**3:.1f} GiB" if isinstance(val, int) else str(val) key_str = f"GPU {key}" if isinstance(key, int) else str(key) parts.append(f" {key_str}: {label}") return "\n".join(parts) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 912c0003484..d4f7199b7d3 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -984,7 +984,12 @@ def from_quantized_weight( def _maybe_squeeze_scale(key: str, value: Any) -> Any: """Squeeze a leading dim=1 from 3-D scale tensors of shape (1, n, m).""" - if "scale" in key and isinstance(value, torch.Tensor) and value.dim() == 3 and value.shape[0] == 1: + if ( + "scale" in key + and isinstance(value, torch.Tensor) + and value.dim() == 3 + and value.shape[0] == 1 + ): return value.squeeze(0) return value diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 8d80ee59516..cbac80eeff3 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1029,10 +1029,11 @@ def _export_transformers_checkpoint_streaming( # Only apply when tie_word_embeddings=True: _tied_weights_keys can list keys whose # weights are not actually shared (e.g. if the model was saved with tie_word_embeddings=False # but the attribute was never cleared), which would incorrectly drop lm_head.weight. - if getattr(model.config, "tie_word_embeddings", False): - raw_tied_keys: set[str] = set(getattr(model, "_tied_weights_keys", None) or []) - else: - raw_tied_keys: set[str] = set() + raw_tied_keys: set[str] = ( + set(getattr(model, "_tied_weights_keys", None) or []) + if getattr(model.config, "tie_word_embeddings", False) + else set() + ) # --- Name mapper for per-tensor key reversal --- # Tensor names are applied inline; quant config names are handled by the caller. @@ -1076,7 +1077,7 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: new_key, new_value = _postprocess_single_tensor( full_key, tensor, kv_cache_max_bound, kv_cache_format, is_modelopt_qlora ) - if new_key is None: + if new_key is None or new_value is None: return if name_mapper is not None: new_key = name_mapper(new_key) @@ -1741,19 +1742,18 @@ def _write_hf_export_config( quantization_details.get("quant_algo") is not None or quantization_details.get("kv_cache_quant_algo") is not None ) - if is_quantized_export: + quantization_config = None + if hf_quant_config is not None and is_quantized_export: with open(f"{export_dir}/hf_quant_config.json", "w") as file: json.dump(hf_quant_config, file, indent=4) - hf_quant_config = convert_hf_quant_config_format(hf_quant_config) - else: - hf_quant_config = None + quantization_config = convert_hf_quant_config_format(hf_quant_config) original_config = f"{export_dir}/config.json" with open(original_config) as file: config_data = json.load(file) sanitize_hf_config_for_deployment(config_data, model) - if hf_quant_config is not None: - config_data["quantization_config"] = hf_quant_config + if quantization_config is not None: + config_data["quantization_config"] = quantization_config if export_sparse_attention_config is not None: sparse_attn_config = export_sparse_attention_config(model) if sparse_attn_config is not None: diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 255544ffe1e..1761b7e8398 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -29,7 +29,7 @@ supported combinations. ### The shipped recipes
-All 20 general/ptq/ recipes (click to expand) +All 21 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -46,6 +46,7 @@ supported combinations. | `nvfp4_experts_only-kv_fp8` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max | | `nvfp4_experts_only-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | max | | `nvfp4_experts_only-kv_fp8_layerwise` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise | +| `nvfp4_experts_only-kv_fp8_layerwise_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (non-mutating, for disk offload) | | `nvfp4_experts_only_mse-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | MSE + FP8 sweep | | `nvfp4_experts_only_input_scale1-kv_fp8_cast` | NVFP4 W4A4, MoE experts only, expert `input_scale` pinned to 1.0 | FP8 (constant amax) | max (weights); expert activations uncalibrated | | `nvfp4_omlp_only-kv_fp8` | NVFP4 W4A4, o_proj + MLP/MoE | FP8 (calibrated) | max | diff --git a/tests/gpu/torch/export/test_offload_export.py b/tests/gpu/torch/export/test_offload_export.py index 6276081da06..5bfb4d122e4 100644 --- a/tests/gpu/torch/export/test_offload_export.py +++ b/tests/gpu/torch/export/test_offload_export.py @@ -106,7 +106,7 @@ def forward_loop(m): for st_file in safetensor_files: with safe_open(str(st_file), framework="pt") as st: - for key in st.keys(): + for key in list(st.keys()): tensor = st.get_tensor(key) assert tensor.numel() > 0, f"Zero-numel tensor for key '{key}' in {st_file.name}" assert not tensor.is_meta, f"Meta tensor for key '{key}' in {st_file.name}" From 8721fdc7a895457ac0e2e743d0cd82d91de12317 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:30:18 +0000 Subject: [PATCH 10/21] fix(export): guard streaming shards against aliases; make compat shims opt-in Addresses review feedback on the streaming export path and the hf_ptq shims. _StreamingShardWriter.add now resolves shared storage before buffering. safetensors.save_file rejects a dict holding two tensors that share memory, and _stream_tensor's .detach().contiguous().cpu() is a no-op for an already-contiguous CPU tensor, so the name-based _tied_weights_keys filter was the only guard -- it misses ties transformers does not declare. data_ptr() is unreliable across the export but reliable within one buffer, since buffered tensors stay alive until flush. Exact (pointer, shape, dtype) matches are dropped as genuine ties; partial matches are cloned so a distinct view is never silently lost. _install_transformers_compat_shims is now opt-in via --allow_compat_shims. It was running on every get_model() call, applying process-wide monkeypatches -- silently disabling flash attention and substituting a lossy BF16 FP8 matmul -- for users who never asked. Gating on trust_remote_code would be wrong: the FP8 shim patches the native HF path, which is what DeepSeek-R1 now uses. The suppress(Exception) around the FP8 block is narrowed so a missing _load_finegrained_fp8_kernel warns instead of silently no-opping; on transformers 5.7 that symbol is absent, so the fallback was never actually installed. _FP8BF16Fallback.matmul now expands scales by block_size rather than out_f // nb_out. The scale grid is ceil-divided, so the ratio mis-groups the last block when a dimension is not a multiple of block_size (out_f=300, block=128 gave 100) and truncates outright when the ratio does not divide evenly. Verified against a per-block reference for divisible and non-divisible shapes. Also drops the fp32 intermediate (~528 MB per call at DSR1's 7168x18432) by scaling in bf16. Per review, hoist the offloaded refusal in _export_transformers_checkpoint to a guard clause beside the detection, removing the else-after-raise and the if not _offloaded wrapper around hook removal. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 71 +++++++++++++------ examples/hf_ptq/hf_ptq.py | 12 ++++ modelopt/torch/export/unified_export_hf.py | 70 +++++++++++------- .../unit/torch/export/test_offload_export.py | 36 ++++++++++ 4 files changed, 142 insertions(+), 47 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 5078e6d9b4f..5a6f5273d65 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import contextlib import copy import glob import hashlib @@ -122,21 +121,29 @@ class _FP8BF16Fallback: def matmul( input, weight, weight_scale_inv, block_size, output_dtype=None, activation_scale=None ): + # Expand by block_size, not by out_f // nb_out: the scale grid is ceil-divided, so + # deriving the factor from the ratio mis-groups (or truncates) the last block + # whenever a dimension is not an exact multiple of block_size. + block_out, block_in = ( + (block_size, block_size) if isinstance(block_size, int) else block_size + ) out_f, in_f = weight.shape[-2], weight.shape[-1] - nb_out, nb_in = weight_scale_inv.shape[-2], weight_scale_inv.shape[-1] scale = ( - weight_scale_inv.float() - .repeat_interleave(out_f // nb_out, -2) - .repeat_interleave(in_f // nb_in, -1) + weight_scale_inv.to(torch.bfloat16) + .repeat_interleave(block_out, -2) + .repeat_interleave(block_in, -1)[..., :out_f, :in_f] ) - w_bf16 = (weight.float() * scale).to(torch.bfloat16) + w_bf16 = weight.to(torch.bfloat16) * scale out = torch.nn.functional.linear(input.to(torch.bfloat16), w_bf16) return out if output_dtype is None else out.to(output_dtype) def _install_transformers_compat_shims() -> None: - """Patch transformers so older remote-code models (e.g. DeepSeek-R1) load on - newer/partial installs. Call once before loading a trust_remote_code checkpoint.""" + """Patch transformers so large FP8 checkpoints (e.g. DeepSeek-R1) load on partial installs. + + Opt-in only (``--allow_compat_shims``): these are process-wide monkeypatches, and the + FP8 one degrades numerics, so they must never be applied on a user's behalf. + """ import transformers.utils as _tu import transformers.utils.import_utils as _tui @@ -148,7 +155,13 @@ def _install_transformers_compat_shims() -> None: # Force transformers' availability checks to False so bundled models skip the flash-attn path. try: import flash_attn # noqa: F401 - except Exception: + except Exception as exc: + warnings.warn( + f"flash_attn is unavailable ({exc}); forcing transformers' flash-attention " + "availability checks to False for this process.", + UserWarning, + stacklevel=2, + ) for _mod in (_tu, _tui): for _fn in ( "is_flash_attn_2_available", @@ -158,19 +171,33 @@ def _install_transformers_compat_shims() -> None: setattr(_mod, _fn, lambda: False) # No `kernels` package → block-scaled FP8 matmul fails; swap in lossy BF16 fallback. - with contextlib.suppress(Exception): + # Only ImportError is suppressed: a missing _load_finegrained_fp8_kernel means the + # transformers internals moved, which must surface rather than silently no-op. + try: import transformers.integrations.finegrained_fp8 as _ff8 + except ImportError: + return - try: - _ff8._load_finegrained_fp8_kernel() - except ImportError: - warnings.warn( - "finegrained-fp8 kernel unavailable; using a lossy BF16 dequant fallback " - "for FP8 matmul. Suitable for calibration amax collection only.", - UserWarning, - stacklevel=2, - ) - _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback + if not hasattr(_ff8, "_load_finegrained_fp8_kernel"): + warnings.warn( + "transformers.integrations.finegrained_fp8._load_finegrained_fp8_kernel is " + f"missing on transformers {transformers.__version__}; the FP8 BF16 fallback " + "was NOT installed. Block-scaled FP8 models may fail to run calibration.", + UserWarning, + stacklevel=2, + ) + return + + try: + _ff8._load_finegrained_fp8_kernel() + except ImportError: + warnings.warn( + "finegrained-fp8 kernel unavailable; using a lossy BF16 dequant fallback " + "for FP8 matmul. Suitable for calibration amax collection only.", + UserWarning, + stacklevel=2, + ) + _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback def run_nemotron_vl_preview( @@ -741,8 +768,10 @@ def get_model( offload_folder=None, max_cpu_memory_gb=None, max_gpu_memory_gb=None, + allow_compat_shims=False, ): - _install_transformers_compat_shims() + if allow_compat_shims: + _install_transformers_compat_shims() print(f"Initializing model from {ckpt_path}") _disk_offload = offload_folder is not None diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index a57c6e1d84a..b19d7eef735 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -581,6 +581,7 @@ def load_model(args: argparse.Namespace): offload_folder=args.offload_folder, max_cpu_memory_gb=args.max_cpu_memory_gb, max_gpu_memory_gb=args.max_gpu_memory_gb, + allow_compat_shims=args.allow_compat_shims, ) else: assert args.qformat in QUANT_CFG_CHOICES, ( @@ -1656,6 +1657,17 @@ def parse_args() -> argparse.Namespace: "Defaults to 80%% of available GPU memory when not specified." ), ) + parser.add_argument( + "--allow_compat_shims", + action="store_true", + help=( + "Patch transformers so block-scaled FP8 checkpoints (e.g. DeepSeek-R1) load on " + "partial installs: disables flash-attention availability checks when flash_attn " + "is broken, and substitutes a lossy BF16 dequant for the finegrained-FP8 matmul " + "when the kernels package is missing. The FP8 fallback degrades numerics and is " + "only suitable for calibration amax collection, so this is opt-in." + ), + ) args = parser.parse_args() if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index cbac80eeff3..f26004438df 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -873,6 +873,8 @@ def __init__(self, export_dir: Path | str, max_shard_size: int) -> None: self._total_bytes: int = 0 # Maps tensor key → part-file index (recorded at flush time) self._key_to_part: dict[str, int] = {} + # Storage identity of every buffered tensor, so aliases never reach save_file. + self._buffer_storage: dict[tuple[int, tuple[int, ...], torch.dtype], str] = {} def _flush(self) -> None: if not self._buffer: @@ -885,10 +887,28 @@ def _flush(self) -> None: self._part_files.append(part_path) self._total_bytes += self._buffer_bytes self._buffer = {} + self._buffer_storage = {} self._buffer_bytes = 0 def add(self, key: str, tensor: torch.Tensor) -> None: - """Buffer a tensor, flushing the current shard to disk when it is full.""" + """Buffer a tensor, flushing the current shard to disk when it is full. + + ``safetensors.save_file`` rejects a dict containing two tensors that share + storage, so aliases are resolved here. ``data_ptr()`` is unreliable across the + whole export (offloaded weights are materialized and freed per layer), but it is + reliable *within* one buffer because every buffered tensor is kept alive until + :meth:`_flush`. An exact match on (pointer, shape, dtype) is a genuine tied + weight and is dropped, matching the batch path's dedup; a partial match is a + distinct view onto shared storage and is copied so no tensor is silently lost. + """ + storage_id = (tensor.data_ptr(), tuple(tensor.shape), tensor.dtype) + if storage_id in self._buffer_storage: + return + if any(ptr == tensor.data_ptr() for ptr, _, _ in self._buffer_storage): + tensor = tensor.clone() + storage_id = (tensor.data_ptr(), tuple(tensor.shape), tensor.dtype) + + self._buffer_storage[storage_id] = key self._buffer[key] = tensor self._buffer_bytes += tensor.nbytes if self._buffer_bytes >= self._max_shard_size: @@ -1211,18 +1231,21 @@ def _export_transformers_checkpoint( # TODO: Handle mixed precision requantize_resmooth_fused_llm_layers(model) - # Detect accelerate offload before removing hooks; offloaded models need weights - # materialized layer-by-layer during export (hooks must stay alive for that pass). - _offloaded = _has_accelerate_offload(model) + # Offloaded models need their weights materialized layer-by-layer, which this + # whole-state-dict path cannot do; export_hf_checkpoint() streams them instead. + if _has_accelerate_offload(model): + raise NotImplementedError( + "_export_transformers_checkpoint does not support disk/CPU-offloaded models. " + "Use export_hf_checkpoint() which dispatches to _export_transformers_checkpoint_streaming." + ) - # Remove all hooks from the model (deferred for offloaded models) - if not _offloaded: - try: - from accelerate.hooks import remove_hook_from_module + # Remove all hooks from the model + try: + from accelerate.hooks import remove_hook_from_module - remove_hook_from_module(model, recurse=True) - except ImportError: - pass # no accelerate installed → no offload hooks exist to remove + remove_hook_from_module(model, recurse=True) + except ImportError: + pass # no accelerate installed → no offload hooks exist to remove quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) @@ -1262,23 +1285,18 @@ def _export_transformers_checkpoint( # Process all quantized modules and export weights from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear - if _offloaded: - raise NotImplementedError( - "_export_transformers_checkpoint does not support disk/CPU-offloaded models. " - "Use export_hf_checkpoint() which dispatches to _export_transformers_checkpoint_streaming." + _process_quantized_modules(model, dtype, is_modelopt_qlora) + _reconstruct_fused_moe_linear(model) + + if is_fsdp2_model(model): + # FSDP2: gather the full (unsharded) state_dict to CPU on rank 0. + quantized_state_dict = get_model_state_dict( + model, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), ) else: - _reconstruct_fused_moe_linear(model) - - if is_fsdp2_model(model): - # FSDP2: gather the full (unsharded) state_dict to CPU on rank 0. - quantized_state_dict = get_model_state_dict( - model, - options=StateDictOptions(full_state_dict=True, cpu_offload=True), - ) - else: - # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict). - quantized_state_dict = model.state_dict() + # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict). + quantized_state_dict = model.state_dict() # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization. kv_cache_max_bound = 448 diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index d6f5747ca69..6059617445d 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -203,6 +203,42 @@ def test_streaming_shard_writer_tensors_readable(): assert torch.allclose(recovered, t), "recovered tensor does not match original" +def test_streaming_shard_writer_drops_tied_alias(): + """Two keys sharing storage must not both reach save_file, which rejects aliases. + + The name-based _tied_weights_keys filter misses ties that transformers does not + declare (e.g. tie_word_embeddings=False but shared storage), so the writer needs + its own guard. + """ + with tempfile.TemporaryDirectory() as tmpdir: + shared = torch.ones(4, 4) + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("embed_tokens.weight", shared) + writer.add("lm_head.weight", shared) + weight_map = writer.finalize() + + assert set(weight_map) == {"embed_tokens.weight"}, ( + "tied alias should be dropped, keeping only the first key" + ) + + +def test_streaming_shard_writer_copies_aliased_view(): + """A distinct view onto shared storage must be copied, not dropped.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = torch.arange(16, dtype=torch.float32).reshape(4, 4) + view = base.view(16) # same data_ptr, different shape + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("base", base) + writer.add("view", view) + weight_map = writer.finalize() + + assert set(weight_map) == {"base", "view"}, "aliased view must be kept, not dropped" + shard_file = Path(tmpdir) / weight_map["view"] + with safe_open(str(shard_file), framework="pt") as f: + assert torch.equal(f.get_tensor("view"), view) + assert torch.equal(f.get_tensor("base"), base) + + # --------------------------------------------------------------------------- # _postprocess_single_tensor # --------------------------------------------------------------------------- From 96b2d900d215a396fc9d81a00e07c4e515265f55 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:12:14 +0000 Subject: [PATCH 11/21] fix(ptq): scope the DeepSeek load-path change and respect --trust_remote_code The removal of the "Deepseek" clause from the architecture dispatch was an undocumented drive-by in dd76c1b2fc. It had two side effects beyond its intent: --trust_remote_code was silently ignored for the model class, and DeepseekV2 and DeepseekVL were swept onto the built-in path alongside DeepseekV3 despite neither being validated here. Restore the clause, gated on the flag. --trust_remote_code now selects the bundled modeling code as before; without it the built-in class is used, which is what the disk-offload and streaming-export paths are validated against and which previously raised AssertionError. Non-DeepSeek architectures are untouched. Also reword the compat-shim docstring: two of its three patches target bundled modeling files, not the built-in path. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 20 +++++--- modelopt/torch/export/unified_export_hf.py | 11 ++--- tests/examples/hf_ptq/test_example_utils.py | 52 +++++++++++++++++++++ 3 files changed, 70 insertions(+), 13 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 5a6f5273d65..cd709c5cddb 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -139,7 +139,10 @@ def matmul( def _install_transformers_compat_shims() -> None: - """Patch transformers so large FP8 checkpoints (e.g. DeepSeek-R1) load on partial installs. + """Patch transformers so DeepSeek-R1-style checkpoints load on newer/partial installs. + + Mostly aimed at bundled (``--trust_remote_code``) modeling files, which import symbols + newer transformers dropped; the FP8 shim also covers the built-in loading path. Opt-in only (``--allow_compat_shims``): these are process-wide monkeypatches, and the FP8 one degrades numerics, so they must never be applied on a user's behalf. @@ -886,11 +889,16 @@ def has_pack_quantized_config(config): raise ValueError(f"Model config at {ckpt_path} has no architectures defined") architecture = hf_config.architectures[0] - if not hasattr(transformers, architecture): - warnings.warn( - f"Architecture {architecture} not found in transformers: {transformers.__version__}. " - "Falling back to AutoModelForCausalLM (or AutoModel for non-causal architectures)." - ) + # DeepSeek ships bundled modeling code, but the built-in class is what the + # disk-offload and streaming-export paths are validated against. + use_bundled_code = trust_remote_code and "Deepseek" in architecture + + if not hasattr(transformers, architecture) or use_bundled_code: + if not hasattr(transformers, architecture): + warnings.warn( + f"Architecture {architecture} not found in transformers: {transformers.__version__}. " + "Falling back to AutoModelForCausalLM (or AutoModel for non-causal architectures)." + ) assert trust_remote_code, ( "Please set trust_remote_code to True if you want to use this architecture" ) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f26004438df..398eab83be7 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -893,13 +893,10 @@ def _flush(self) -> None: def add(self, key: str, tensor: torch.Tensor) -> None: """Buffer a tensor, flushing the current shard to disk when it is full. - ``safetensors.save_file`` rejects a dict containing two tensors that share - storage, so aliases are resolved here. ``data_ptr()`` is unreliable across the - whole export (offloaded weights are materialized and freed per layer), but it is - reliable *within* one buffer because every buffered tensor is kept alive until - :meth:`_flush`. An exact match on (pointer, shape, dtype) is a genuine tied - weight and is dropped, matching the batch path's dedup; a partial match is a - distinct view onto shared storage and is copied so no tensor is silently lost. + ``save_file`` rejects tensors sharing storage. ``data_ptr()`` is only meaningful + within one buffer (entries stay alive until :meth:`_flush`), so dedup here: an + exact (pointer, shape, dtype) match is a real tie and is dropped; a partial match + is a distinct view and is copied rather than lost. """ storage_id = (tensor.data_ptr(), tuple(tensor.shape), tensor.dtype) if storage_id in self._buffer_storage: diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 00621ec6125..39ef53b7ca9 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -316,3 +316,55 @@ def from_pretrained(*args, **kwargs): else: assert "trust_remote_code" not in calls["from_config"] assert calls["from_pretrained"]["trust_remote_code"] is True + + +@pytest.mark.parametrize( + ("trust_remote_code", "expect_bundled_code"), + [(True, True), (False, False)], +) +def test_get_model_deepseek_honors_trust_remote_code( + monkeypatch, trust_remote_code, expect_bundled_code +): + """DeepSeek ships bundled modeling code; --trust_remote_code selects it, else built-in.""" + used = {} + hf_config = SimpleNamespace( + architectures=["DeepseekV3ForCausalLM"], + dtype=torch.bfloat16, + model_type="deepseek_v3", + torch_dtype=torch.bfloat16, + ) + + class FakeModel: + def eval(self): + return None + + def _record(tag): + class Fake: + @staticmethod + def from_config(config, **kwargs): + used["path"] = tag + return FakeModel() + + _from_config = from_config + + @staticmethod + def from_pretrained(*args, **kwargs): + used["path"] = tag + return FakeModel() + + return Fake + + monkeypatch.setattr(example_utils.AutoConfig, "from_pretrained", lambda *a, **k: hf_config) + monkeypatch.setattr(example_utils, "AutoModelForCausalLM", _record("bundled")) + monkeypatch.setattr( + example_utils.transformers, "DeepseekV3ForCausalLM", _record("builtin"), raising=False + ) + monkeypatch.setattr(example_utils, "is_nemotron_vl", lambda config: False) + monkeypatch.setattr(example_utils, "is_speculative", lambda config: False) + monkeypatch.setattr(example_utils, "init_empty_weights", lambda include_buffers: nullcontext()) + monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 1024}) + monkeypatch.setattr(example_utils, "infer_auto_device_map", lambda model, max_memory: {"": 0}) + + example_utils.get_model("checkpoint", device="cpu", trust_remote_code=trust_remote_code) + + assert used["path"] == ("bundled" if expect_bundled_code else "builtin") From fe09a30d442b3fd9a46e643ec77963d3fe6ffd1c Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:27:31 +0000 Subject: [PATCH 12/21] revert(ptq): drop transformers compat shims in favor of a correct environment These monkeypatched transformers process-wide to paper over environment problems, which does not belong in an example -- especially a quantization one. - _FP8BF16Fallback substituted a lossy BF16 dequant for the block-scaled FP8 matmul, degrading the very forward pass used for calibration amax collection. It also patched _load_finegrained_fp8_kernel, a private symbol that no longer exists in transformers 5.7 (now lazy_load_kernel), so within the supported range (>=4.56,<5.13) it silently no-ops. The fix is `pip install kernels`. - The flash_attn shim forced availability checks to False when the package was installed but unimportable. Uninstalling or repairing the broken install reaches the same state without a global patch. - The is_torch_fx_available shim covered old bundled modeling files on transformers 5; not worth a process-wide patch on its own. Also removes --allow_compat_shims, added in the previous commit to gate them. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 96 -------------------------------- examples/hf_ptq/hf_ptq.py | 12 ---- 2 files changed, 108 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index cd709c5cddb..d9a9f62a7a4 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -110,99 +110,6 @@ def validate_fsdp2_supported(args, config): + "\n - ".join(issues) + "\nRemove --use_fsdp2 or use a standard causal-LM checkpoint." ) - -class _FP8BF16Fallback: - """BF16 dequant fallback for block-scaled FP8 matmul when the kernels package is absent. - - Calibration amax collection only — not accurate for production inference. - """ - - @staticmethod - def matmul( - input, weight, weight_scale_inv, block_size, output_dtype=None, activation_scale=None - ): - # Expand by block_size, not by out_f // nb_out: the scale grid is ceil-divided, so - # deriving the factor from the ratio mis-groups (or truncates) the last block - # whenever a dimension is not an exact multiple of block_size. - block_out, block_in = ( - (block_size, block_size) if isinstance(block_size, int) else block_size - ) - out_f, in_f = weight.shape[-2], weight.shape[-1] - scale = ( - weight_scale_inv.to(torch.bfloat16) - .repeat_interleave(block_out, -2) - .repeat_interleave(block_in, -1)[..., :out_f, :in_f] - ) - w_bf16 = weight.to(torch.bfloat16) * scale - out = torch.nn.functional.linear(input.to(torch.bfloat16), w_bf16) - return out if output_dtype is None else out.to(output_dtype) - - -def _install_transformers_compat_shims() -> None: - """Patch transformers so DeepSeek-R1-style checkpoints load on newer/partial installs. - - Mostly aimed at bundled (``--trust_remote_code``) modeling files, which import symbols - newer transformers dropped; the FP8 shim also covers the built-in loading path. - - Opt-in only (``--allow_compat_shims``): these are process-wide monkeypatches, and the - FP8 one degrades numerics, so they must never be applied on a user's behalf. - """ - import transformers.utils as _tu - import transformers.utils.import_utils as _tui - - # transformers >=5 removed is_torch_fx_available; older bundled model files still import it. - if not hasattr(_tui, "is_torch_fx_available"): - _tui.is_torch_fx_available = lambda: False - - # Broken flash_attn installs (.so undefined-symbol) crash at import time, not find_spec time. - # Force transformers' availability checks to False so bundled models skip the flash-attn path. - try: - import flash_attn # noqa: F401 - except Exception as exc: - warnings.warn( - f"flash_attn is unavailable ({exc}); forcing transformers' flash-attention " - "availability checks to False for this process.", - UserWarning, - stacklevel=2, - ) - for _mod in (_tu, _tui): - for _fn in ( - "is_flash_attn_2_available", - "is_flash_attn_available", - "is_flash_attn_greater_or_equal_2_10", - ): - setattr(_mod, _fn, lambda: False) - - # No `kernels` package → block-scaled FP8 matmul fails; swap in lossy BF16 fallback. - # Only ImportError is suppressed: a missing _load_finegrained_fp8_kernel means the - # transformers internals moved, which must surface rather than silently no-op. - try: - import transformers.integrations.finegrained_fp8 as _ff8 - except ImportError: - return - - if not hasattr(_ff8, "_load_finegrained_fp8_kernel"): - warnings.warn( - "transformers.integrations.finegrained_fp8._load_finegrained_fp8_kernel is " - f"missing on transformers {transformers.__version__}; the FP8 BF16 fallback " - "was NOT installed. Block-scaled FP8 models may fail to run calibration.", - UserWarning, - stacklevel=2, - ) - return - - try: - _ff8._load_finegrained_fp8_kernel() - except ImportError: - warnings.warn( - "finegrained-fp8 kernel unavailable; using a lossy BF16 dequant fallback " - "for FP8 matmul. Suitable for calibration amax collection only.", - UserWarning, - stacklevel=2, - ) - _ff8._load_finegrained_fp8_kernel = lambda: _FP8BF16Fallback - - def run_nemotron_vl_preview( full_model, tokenizer, @@ -771,10 +678,7 @@ def get_model( offload_folder=None, max_cpu_memory_gb=None, max_gpu_memory_gb=None, - allow_compat_shims=False, ): - if allow_compat_shims: - _install_transformers_compat_shims() print(f"Initializing model from {ckpt_path}") _disk_offload = offload_folder is not None diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index b19d7eef735..a57c6e1d84a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -581,7 +581,6 @@ def load_model(args: argparse.Namespace): offload_folder=args.offload_folder, max_cpu_memory_gb=args.max_cpu_memory_gb, max_gpu_memory_gb=args.max_gpu_memory_gb, - allow_compat_shims=args.allow_compat_shims, ) else: assert args.qformat in QUANT_CFG_CHOICES, ( @@ -1657,17 +1656,6 @@ def parse_args() -> argparse.Namespace: "Defaults to 80%% of available GPU memory when not specified." ), ) - parser.add_argument( - "--allow_compat_shims", - action="store_true", - help=( - "Patch transformers so block-scaled FP8 checkpoints (e.g. DeepSeek-R1) load on " - "partial installs: disables flash-attention availability checks when flash_attn " - "is broken, and substitutes a lossy BF16 dequant for the finegrained-FP8 matmul " - "when the kernels package is missing. The FP8 fallback degrades numerics and is " - "only suitable for calibration amax collection, so this is opt-in." - ), - ) args = parser.parse_args() if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): From 7e1de0475061d2efc38e7c16c558d34d245c5b6e Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:00:50 +0000 Subject: [PATCH 13/21] fix(export): scope data_ptr tie identity to resident tensors The offloaded export path used data_ptr() as a stable tensor identity, but it only identifies a tensor while that tensor is resident. Two failure modes, both silent, both verified on Qwen3.6-35B-A3B with a 30/20 GB budget: Recycled addresses. ExportContext.tied_cache / moe_tied_cache persisted for the whole export while the streaming path materialises and frees one module at a time, so the allocator handed a later expert the address of a freed earlier one. The alias step then re-pointed its weight and scales at the wrong module. 1536 false hits, leaving 60% of expert weights (18440 tensors) as byte-identical copies of unrelated experts, and short-circuiting half the export (_export_quantized_weight ran 14592 times instead of 30720). Null addresses. sync_tied_input_amax groups by weight data_ptr, and every meta tensor reports 0, so all offloaded modules collapsed into two buckets whose amaxes were max-merged model-wide. 16896 expert input_scale values inflated, median 90% relative error, worst 12x. Fixes: ExportContext.reset_tied_caches() drops dedup state at the end of each materialization window, and sync_tied_input_amax skips modules whose weights are not resident, warning only when the model declares ties it could not honour. Verified against the batch export path, which is unaffected (it keeps every weight resident): all 123846 shared tensors now match bitwise, versus 70580 differing before. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/quant_utils.py | 20 ++++++- modelopt/torch/export/registry.py | 10 ++++ modelopt/torch/export/unified_export_hf.py | 2 + .../unit/torch/export/test_offload_export.py | 52 +++++++++++++++++++ 4 files changed, 83 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index d4f7199b7d3..18206ca9bbf 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -1604,11 +1604,19 @@ def sync_tied_input_amax(model: nn.Module) -> int: from collections import defaultdict by_dp: dict = defaultdict(list) + skipped_candidates = 0 for _, m in model.named_modules(): # Fused MoE: 3-D source tensors with shared input quantizers first_proj_attr = getattr(m, "_first_proj_attr", "gate_up_proj") - first_proj = getattr(m, first_proj_attr, None) first_proj_input_quantizer_attr = f"{first_proj_attr}_input_quantizer" + # data_ptr() only identifies a tensor that is resident: every meta tensor reports + # 0, which would collapse unrelated modules into a single tied group and merge + # their amaxes model-wide. + if any(p.is_meta for p in m.parameters(recurse=False)): + if hasattr(m, "input_quantizer") or hasattr(m, first_proj_input_quantizer_attr): + skipped_candidates += 1 + continue + first_proj = getattr(m, first_proj_attr, None) if ( hasattr(m, first_proj_input_quantizer_attr) and first_proj is not None @@ -1625,6 +1633,16 @@ def sync_tied_input_amax(model: nn.Module) -> int: ): by_dp[("dense", m.weight.data_ptr())].append(m) + # Only meaningful when the model declares ties: without them there is nothing to lose + # by skipping non-resident modules, and an unconditional warning on every offloaded + # export would be noise. + if skipped_candidates and getattr(model, "_tied_weights_keys", None): + warn( + f"sync_tied_input_amax: {skipped_candidates} quantized module(s) have offloaded " + "weights and were skipped, so ties among them were not merged; those modules keep " + "their per-side input_scale." + ) + def _merge(quantizers: list) -> bool: """Max-merge amaxes across the quantizer list. Returns True on merge.""" valid = [ diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 260cb32eea3..c99efb7807b 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -68,6 +68,16 @@ def __post_init__(self) -> None: self.tied_cache = None self.moe_tied_cache = None + def reset_tied_caches(self) -> None: + """Drop dedup state between materialization windows. + + The streaming export frees each module after exporting it, so a recycled + address would alias the next module to the wrong weights. Genuine sharing is + always within one window, so nothing is lost. + """ + self.tied_cache.clear() + self.moe_tied_cache.clear() + ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 398eab83be7..17ec9ad3008 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1118,6 +1118,7 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: continue seen_keys.add(full_key) _stream_tensor(full_key, tensor) + ctx.reset_tied_caches() # Non-decoder modules with offload hooks (embed_tokens, norm, lm_head, etc.) for name, module in model.named_modules(): @@ -1143,6 +1144,7 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: continue seen_keys.add(full_key) _stream_tensor(full_key, tensor) + ctx.reset_tied_caches() # GPU-resident parameters and persistent buffers (not covered by the above loops). # named_buffers() includes non-persistent buffers that state_dict() excludes; filter them. diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 6059617445d..0326fb89040 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -309,3 +309,55 @@ def test_postprocess_vision_model_summary_idxs_dropped(): "vision_model.radio_model.summary_idxs", torch.tensor([0, 1]), 448.0, None ) assert key is None + + +# --------------------------------------------------------------------------- +# data_ptr identity under offload +# +# Both guards below exist because ``data_ptr()`` only identifies a tensor while +# that tensor is resident. Getting this wrong silently exported wrong weights: +# meta tensors all report 0, and freed addresses are recycled by the allocator. +# --------------------------------------------------------------------------- + + +class _FakeAmaxQuantizer(nn.Module): + def __init__(self, value: float): + super().__init__() + self.register_buffer("_amax", torch.tensor(value)) + self.is_enabled = True + + @property + def amax(self): + return self._amax + + @amax.setter + def amax(self, v): + self._amax = v + + +class _MetaLinearWithInputQuantizer(nn.Module): + def __init__(self, amax: float): + super().__init__() + self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) + self.input_quantizer = _FakeAmaxQuantizer(amax) + + +def test_sync_tied_input_amax_skips_offloaded_modules(): + """Untied modules whose weights are offloaded must not be merged together. + + Every meta tensor reports ``data_ptr() == 0``, so without the residency guard + these two unrelated Linears land in one group and both get amax 9.0. + """ + from modelopt.torch.export.quant_utils import sync_tied_input_amax + + model = nn.Module() + model._tied_weights_keys = {"b.weight": "a.weight"} + model.a = _MetaLinearWithInputQuantizer(1.0) + model.b = _MetaLinearWithInputQuantizer(9.0) + + with pytest.warns(UserWarning, match="offloaded weights"): + merged = sync_tied_input_amax(model) + + assert merged == 0 + assert model.a.input_quantizer._amax.item() == 1.0 + assert model.b.input_quantizer._amax.item() == 9.0 From d000b18f34d7b7b5a3a44fef4216f865f5a4ac84 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:02:37 +0000 Subject: [PATCH 14/21] fix(export): stream extra_state_dict tensors in the offload path The streaming path warned that extra_state_dict "is not supported" and dropped it. That silently lost the MTP weights: HF builds only num_hidden_layers decoders, so multi-token-prediction tensors are orphaned and reach export solely through extra_state_dict (hf_ptq.py passes load_mtp_weights' output there). Exporting Qwen3.6-35B-A3B produced 19 fewer tensors than the batch path, all mtp.* -- mtp.fc.weight, mtp.norm.weight, and the whole mtp.layers.0 block including its fused experts. A checkpoint missing them cannot serve speculative decoding. Feed them to the shard writer after the resident-tensor pass. They are already materialized and skip per-tensor postprocessing, matching the batch path which merges them after postprocess_state_dict; only the hub-name reversal applies. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/unified_export_hf.py | 20 ++++++++++++++----- .../unit/torch/export/test_offload_export.py | 18 +++++++++++++++++ 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 17ec9ad3008..52b245bc87d 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -969,6 +969,7 @@ def _export_transformers_checkpoint_streaming( is_modelopt_qlora: bool = False, export_dir: Path | str = ".", max_shard_size: int | str = "10GB", + extra_state_dict: dict[str, torch.Tensor] | None = None, **kwargs, ) -> tuple[None, dict[str, Any]]: """Export a disk/CPU-offloaded model by streaming tensors layer-by-layer to shard files. @@ -1157,6 +1158,19 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: seen_keys.add(name) _stream_tensor(name, tensor) + # Tensors the model never held — e.g. MTP weights, which HF leaves orphaned because it + # builds only num_hidden_layers decoders. They are already materialized and skip the + # per-tensor postprocessing, matching how the batch path merges them after + # postprocess_state_dict; only the hub-name reversal applies. + for name, tensor in (extra_state_dict or {}).items(): + if name in seen_keys: + continue + seen_keys.add(name) + writer.add( + name_mapper(name) if name_mapper is not None else name, + tensor.detach().contiguous().cpu(), + ) + writer.finalize() # Write non-weight artifacts: config.json, generation_config.json, and the custom @@ -1847,16 +1861,12 @@ def export_hf_checkpoint( "save_modelopt_state=True is not supported in the streaming offload export " "path and will be ignored." ) - if extra_state_dict: - warnings.warn( - "extra_state_dict is not supported in the streaming offload export path " - "and will be ignored." - ) _, hf_quant_config = _export_transformers_checkpoint_streaming( model, dtype, export_dir=export_dir, max_shard_size=max_shard_size, + extra_state_dict=extra_state_dict, **kwargs, ) if getattr(model, "hf_quantizer", None) is not None: diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 0326fb89040..2d7962f9aa8 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -361,3 +361,21 @@ def test_sync_tied_input_amax_skips_offloaded_modules(): assert merged == 0 assert model.a.input_quantizer._amax.item() == 1.0 assert model.b.input_quantizer._amax.item() == 9.0 + + +def test_streaming_shard_writer_accepts_extra_tensors(): + """extra_state_dict tensors must land in the shards. + + MTP weights are orphaned — HF builds only num_hidden_layers decoders, so they are + never in model.state_dict() and reach export only via extra_state_dict. The streaming + path used to drop them, silently losing 19 tensors relative to the batch export. + """ + with tempfile.TemporaryDirectory() as tmpdir: + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("model.layers.0.weight", torch.ones(4, 4)) + writer.add("mtp.fc.weight", torch.full((2, 2), 7.0)) + weight_map = writer.finalize() + + assert "mtp.fc.weight" in weight_map + with safe_open(str(Path(tmpdir) / weight_map["mtp.fc.weight"]), framework="pt") as f: + assert torch.equal(f.get_tensor("mtp.fc.weight"), torch.full((2, 2), 7.0)) From b2289de8c1599a77cbc46063b2ea28d770f496ab Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:43:15 -0700 Subject: [PATCH 15/21] fix(export): remove _tied_cache from fused-expert export to prevent cross-layer aliasing Per-expert weight wrappers inside _export_fused_experts are freshly allocated .contiguous() copies. Their data_ptr() is ephemeral: it is freed when packing replaces the tensor, so the GPU allocator can recycle that address for an unrelated expert in a later layer. Passing those ephemeral addresses into the caller-owned _tied_cache caused false-positive cache hits, silently aliasing the packed weight/scale tensors of one MoE layer onto another (manifesting as 35 k weight_scale mismatches between layers 15-60 in the DSR1 NVFP4 export). Changes: - Drop _tied_cache param from _export_fused_experts signature and all call sites (_export_fused_experts_module in hf_export_handlers.py). - Extend docstring to explain WHY _tied_cache is intentionally excluded. - Update test_fused_experts.py: remove _tied_cache from both dedup test call sites (they only need _moe_tied_cache for module-level dedup, which is unaffected). Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/hf_export_handlers.py | 1 - modelopt/torch/export/moe_utils.py | 21 +++++++++----- modelopt/torch/export/unified_export_hf.py | 29 +++++++++++++++++++ .../plugins/test_fused_experts.py | 16 ++++------ 4 files changed, 47 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 800b51daca9..f8bf883c9eb 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -134,7 +134,6 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex module, ctx.dtype, _moe_tied_cache=ctx.moe_tied_cache, - _tied_cache=ctx.tied_cache, ) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 787e173959e..eada3a635c9 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -78,7 +78,6 @@ def _export_fused_experts( module: nn.Module, dtype: torch.dtype, _moe_tied_cache: dict[tuple[int, int], nn.Module] | None = None, - _tied_cache: dict[int, nn.Module] | None = None, ) -> None: """Split fused MoE expert weights and export per-expert quantization scales. @@ -107,12 +106,18 @@ def _export_fused_experts( ``(.data_ptr(), down_proj.data_ptr())``), the alias step at the end re-points the per-expert ``weight`` / ``weight_scale`` / ``weight_scale_2`` / ``input_scale`` buffers at a previously-processed - module sharing the same source memory. ``_tied_cache`` (int-keyed) is - threaded through to the per-projection ``_export_quantized_weight`` - calls so wrapper-level dedup uses the same scope as standalone Linears. - Both caches are owned by the caller (typically - ``_export_transformers_checkpoint``) and scoped to one export - invocation; when ``None`` the corresponding alias step is skipped. + module sharing the same source memory. ``_moe_tied_cache`` is owned by + the caller and scoped to one export invocation; when ``None`` the alias + step is skipped. + + ``_tied_cache`` is intentionally NOT threaded through to the inner + ``_export_quantized_weight`` calls. Each per-expert weight wrapper is a + freshly allocated ``.contiguous()`` copy whose ``data_ptr()`` is + ephemeral — it is freed as soon as packing replaces the tensor, so the + allocator can recycle that address for a completely unrelated expert in a + later layer. Putting ephemeral addresses into a caller-owned + ``_tied_cache`` causes cross-layer false-positive hits and silently + aliases the packed weight / scale tensors of one layer onto another. """ from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim @@ -271,7 +276,7 @@ def _export_fused_experts( wrapper.weight_quantizer = w_quantizer wrapper.input_quantizer = i_quantizer - _export_quantized_weight(wrapper, dtype, _tied_cache=_tied_cache) + _export_quantized_weight(wrapper, dtype) proj = nn.Module() proj.weight = wrapper.weight diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 52b245bc87d..ac0f8a955d1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1119,7 +1119,36 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: continue seen_keys.add(full_key) _stream_tensor(full_key, tensor) + # Release GPU tensors added by export handlers before hook.post_forward + # runs, to prevent cross-layer accumulation on disk-offloaded models. + # + # Two categories accumulate without explicit cleanup: + # + # 1. CUDA *buffers* on any sub-module (weight_scale, weight_scale_2, + # input_scale): AlignDevicesHook.post_forward uses offload_buffers=False + # by default, so it never offloads buffers. Pre-existing buffers in + # disk-offloaded layers live on CPU, so any CUDA buffer encountered here + # was registered by the export handlers and is safe to drop. + # + # 2. CUDA *parameters* on modules WITHOUT _hf_hook: _export_fused_experts + # creates fresh nn.Module objects (one per expert × projection) and adds + # them to the layer via add_module() *after* weight_access_and_writeback + # captured its materialized list. hook.post_forward never visits these + # new modules, so their packed NVFP4 weight parameters (~5 GB per MoE + # layer) stay live on GPU. Modules WITH _hf_hook are original model + # modules whose parameters hook.post_forward will meta-ify; leave those + # alone. + for sub_mod in layer_module.modules(): + for buf_name in list(sub_mod._buffers): + buf = sub_mod._buffers[buf_name] + if buf is not None and buf.device.type == "cuda": + sub_mod._buffers[buf_name] = None + if not hasattr(sub_mod, "_hf_hook"): + for param_name, param in list(sub_mod._parameters.items()): + if param is not None and param.device.type == "cuda": + sub_mod._parameters[param_name] = None ctx.reset_tied_caches() + torch.cuda.empty_cache() # Non-decoder modules with offload hooks (embed_tokens, norm, lm_head, etc.) for name, module in model.named_modules(): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 9fa836bb620..d5d708ab189 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -693,21 +693,18 @@ def test_per_expert_buffers_share_data_ptr_for_tied_fused_experts(self): try: _calibrate_two_moe_blocks(parent) - # Per-call dedup caches threaded through both export calls; int keys - # for per-expert wrapper dedup, tuple keys for module-level dedup. - tied_cache: dict = {} + # Module-level dedup cache; per-expert wrapper addresses are ephemeral + # so _tied_cache is intentionally not passed (see _export_fused_experts). moe_tied_cache: dict = {} _export_fused_experts( parent.encoder.experts, torch.float16, _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, ) _export_fused_experts( parent.decoder.experts, torch.float16, _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, ) for idx in range(NUM_EXPERTS): @@ -734,22 +731,19 @@ def test_per_expert_buffers_have_independent_data_ptrs_for_untied_fused_experts( try: _calibrate_two_moe_blocks(parent) - # Same fresh caches as the positive case — confirms that even with - # dedup enabled, untied modules with distinct source data_ptrs do - # not get falsely aliased. - tied_cache: dict = {} + # Same fresh cache as the positive case — confirms that even with + # module-level dedup enabled, untied modules with distinct source + # data_ptrs do not get falsely aliased. moe_tied_cache: dict = {} _export_fused_experts( parent.encoder.experts, torch.float16, _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, ) _export_fused_experts( parent.decoder.experts, torch.float16, _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, ) for idx in range(NUM_EXPERTS): From dcc0d7d1b46dc5eadf8a7bf7ee28b661c5dcb7ce Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:11:22 +0000 Subject: [PATCH 16/21] fix(export): reconcile tied-cache reset with FSDP2 dedup opt-out main disables pointer-keyed dedup under FSDP2 by setting ExportContext.tied_cache and moe_tied_cache to None, for the same reason this branch scopes them per materialization window: FSDP2 recycles data_ptr() values as modules are resharded. reset_tied_caches() called .clear() unconditionally and would raise AttributeError on an FSDP2 export, so make it a no-op when dedup is already disabled. Also fixes two artifacts of the rebase resolution: a missing blank line after the FSDP2 helpers in example_utils.py, and an ambiguous multiplication sign in a comment that RUF003 rejects. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/example_utils.py | 2 ++ modelopt/torch/export/registry.py | 9 ++++++--- modelopt/torch/export/unified_export_hf.py | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d9a9f62a7a4..5bf6dbf9958 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -110,6 +110,8 @@ def validate_fsdp2_supported(args, config): + "\n - ".join(issues) + "\nRemove --use_fsdp2 or use a standard causal-LM checkpoint." ) + + def run_nemotron_vl_preview( full_model, tokenizer, diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index c99efb7807b..2be393c6603 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -73,10 +73,13 @@ def reset_tied_caches(self) -> None: The streaming export frees each module after exporting it, so a recycled address would alias the next module to the wrong weights. Genuine sharing is - always within one window, so nothing is lost. + always within one window, so nothing is lost. No-op when dedup is already + disabled (FSDP2), which recycles addresses for the same reason. """ - self.tied_cache.clear() - self.moe_tied_cache.clear() + if self.tied_cache is not None: + self.tied_cache.clear() + if self.moe_tied_cache is not None: + self.moe_tied_cache.clear() ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index ac0f8a955d1..3fd34f73e61 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1131,7 +1131,7 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: # was registered by the export handlers and is safe to drop. # # 2. CUDA *parameters* on modules WITHOUT _hf_hook: _export_fused_experts - # creates fresh nn.Module objects (one per expert × projection) and adds + # creates fresh nn.Module objects (one per expert x projection) and adds # them to the layer via add_module() *after* weight_access_and_writeback # captured its materialized list. hook.post_forward never visits these # new modules, so their packed NVFP4 weight parameters (~5 GB per MoE From 3175d37a5a449192a04590b15fbf6858e8c77f42 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:57:24 +0000 Subject: [PATCH 17/21] refactor(export): decide materialization via the shared dispatcher The streaming exporter decided which modules to materialize by testing accelerate internals directly (_hf_hook, _get_offload_hook, meta params), so it only ever worked for accelerate offload even though the mechanism it drives -- enable_weight_access_and_writeback -- already dispatches over FSDP2, HF TP and accelerate alike. Add requires_weight_materialization() beside that dispatcher, mirroring its branches, and have the exporter ask it instead. Two conditions must hold: the module owns tensors that are not readable right now (meta, or a sharded DTensor), and a context exists that can materialize them. Also pass the real root model and a cached name_to_module to both windows -- FSDP2 detection walks up from root_model, so passing the layer as its own root could never detect sharding, and the cache avoids the O(N^2) rescan the parameter was added for. Fixes a bug this surfaced: the non-decoder pass skipped decoder layers by id but not their children. An offloaded layer's children return to meta when its window closes, so they passed the filter and were exported a second time, hitting a None amax on weights the first pass had already packed. Only reproduced when the device map placed a layer on CPU, which is why the GPU tests caught it and the unit tests did not. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/unified_export_hf.py | 35 +++++++++++-------- .../torch/quantization/utils/core_utils.py | 30 +++++++++++++++- 2 files changed, 49 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 3fd34f73e61..d0d29cd0112 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -982,12 +982,17 @@ def _export_transformers_checkpoint_streaming( responsible for writing ``hf_quant_config.json`` and updating ``config.json`` with ``quantization_config``. """ - from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear - from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback + from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + requires_weight_materialization, + ) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector export_dir = Path(export_dir) + # Materialization dispatch walks the module tree from the root; without this cache each + # call re-derives it, which is O(N^2) over a MoE model's expert modules. + name_to_module = dict(model.named_modules()) # --- Same model-level setup as _export_transformers_checkpoint --- if dtype is None: @@ -1076,6 +1081,10 @@ def _export_transformers_checkpoint_streaming( "The model architecture is not supported by LayerActivationCollector." ) decoder_layer_ids = {id(m) for m in decoder_layers} + # Descendants too, not just the layers: an offloaded layer's children return to meta + # when its window closes, so a child-level check would re-enter and re-export weights + # this pass already packed. + decoder_owned_ids = {id(m) for layer in decoder_layers for m in layer.modules()} # --- Persistent-buffer predicate (mirrors state_dict() which excludes non-persistent) --- def _is_persistent_buffer(name: str) -> bool: @@ -1103,11 +1112,13 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: return writer.add(new_key, new_value.detach().contiguous().cpu()) - # Decoder layers (offloaded: materialize one at a time) + # Decoder layers: materialize one at a time for layer_name, layer_module in model.named_modules(): if id(layer_module) not in decoder_layer_ids: continue - with enable_weight_access_and_writeback(layer_module, layer_module, writeback=False): + with enable_weight_access_and_writeback( + layer_module, model, name_to_module, writeback=False + ): for sub_name, sub_mod in layer_module.named_modules(): full_name = f"{layer_name}.{sub_name}" if sub_name else layer_name _dispatch_export_handler(full_name, sub_mod, ctx) @@ -1150,20 +1161,14 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: ctx.reset_tied_caches() torch.cuda.empty_cache() - # Non-decoder modules with offload hooks (embed_tokens, norm, lm_head, etc.) + # Non-decoder modules whose weights are not directly readable (embed_tokens, norm, + # lm_head, ...). Containers are skipped: their children get their own window. for name, module in model.named_modules(): - if id(module) in decoder_layer_ids: - continue - if not hasattr(module, "_hf_hook"): + if id(module) in decoder_owned_ids: continue - if _get_offload_hook(module._hf_hook) is None: - continue - if not ( - any(p is not None and p.is_meta for p in module._parameters.values()) - or any(b is not None and b.is_meta for b in module._buffers.values()) - ): + if not requires_weight_materialization(module, model, name_to_module): continue - with enable_weight_access_and_writeback(module, module, writeback=False): + with enable_weight_access_and_writeback(module, model, name_to_module, writeback=False): for sub_name, sub_mod in module.named_modules(): full_name = f"{name}.{sub_name}" if sub_name else name _dispatch_export_handler(full_name, sub_mod, ctx) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 1bdf23da64a..b58991278d4 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -16,6 +16,7 @@ """Quantization utilities.""" import copy +import itertools from collections import namedtuple from contextlib import ExitStack, contextmanager, nullcontext from typing import TYPE_CHECKING, Any @@ -25,7 +26,7 @@ import torch.nn.functional as F from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam -from torch.distributed.tensor import Replicate +from torch.distributed.tensor import DTensor, Replicate from modelopt.torch.quantization.config import QuantizerCfgEntry from modelopt.torch.utils import get_unwrapped_name, print_rank_0 @@ -619,6 +620,33 @@ def enable_weight_access_and_writeback( yield +def requires_weight_materialization(module, root_model, name_to_module: dict | None = None) -> bool: + """Whether ``module``'s own weights are currently unreadable and need a window. + + Mirrors the dispatch in :func:`enable_weight_access_and_writeback`, so callers + deciding *whether* to open a window agree with what opening one would do. Two things + must hold: the module owns tensors that are not directly readable right now + (offloaded to meta, or a sharded ``DTensor``), and a context exists that can + materialize them. Modules already materialized are excluded -- re-entering a window + would re-run export handlers over already-packed weights. + """ + if not any( + t is not None and (t.is_meta or isinstance(t, DTensor)) + for t in itertools.chain(module._parameters.values(), module._buffers.values()) + ): + return False + if _get_enclosing_fsdp_module(module, root_model, name_to_module) is not None: + return True + if is_quantized_parallel_linear(module) and hasattr(module, "_hf_tp_plan"): + return True + hook = getattr(module, "_hf_hook", None) + if hook is None: + return False + from ..plugins.accelerate import _get_offload_hook + + return _get_offload_hook(hook) is not None + + @contextmanager def persistent_materialization(layer, writeback: bool = True): """Keep all layer weights materialized on GPU for the duration. From 57ab0abc838dc4cf9db9e8b8247166a08dc694c9 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:54:05 +0000 Subject: [PATCH 18/21] fix(export): skip sharded DTensors in sync_tied_input_amax, not just meta The residency guard only skipped meta tensors, so under FSDP2 -- where params are DTensors rather than meta -- modules fell through and were grouped by data_ptr() anyway. That is the same assumption main already rejected for the dedup caches: a DTensor's address belongs to a local shard the allocator recycles on reshard, which is why ExportContext.__post_init__ sets tied_cache and moe_tied_cache to None for FSDP2 models. Match requires_weight_materialization(), which already tests `is_meta or isinstance(t, DTensor)`, so the two agree on what "not readable right now" means. Note the DTensor branch is not covered by tests: exercising it needs a live process group, and the existing FSDP2 GPU tests cover weight writeback rather than this path. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/quant_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 18206ca9bbf..c0323cabc63 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -23,6 +23,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor from modelopt import __version__ from modelopt.torch.quantization.model_calib import ( @@ -1611,8 +1612,10 @@ def sync_tied_input_amax(model: nn.Module) -> int: first_proj_input_quantizer_attr = f"{first_proj_attr}_input_quantizer" # data_ptr() only identifies a tensor that is resident: every meta tensor reports # 0, which would collapse unrelated modules into a single tied group and merge - # their amaxes model-wide. - if any(p.is_meta for p in m.parameters(recurse=False)): + # their amaxes model-wide. A sharded DTensor is no better -- its address belongs + # to a local shard the allocator recycles on reshard, which is why FSDP2 disables + # pointer-keyed dedup outright (see ExportContext.__post_init__). + if any(p.is_meta or isinstance(p, DTensor) for p in m.parameters(recurse=False)): if hasattr(m, "input_quantizer") or hasattr(m, first_proj_input_quantizer_attr): skipped_candidates += 1 continue From d5615d70dd7795f8d69e267eeddc21ed8744910e Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:37:46 +0000 Subject: [PATCH 19/21] fix(export): disable tie dedup whenever weights are not resident Pointer-keyed tie dedup assumes data_ptr() identifies a tensor for the whole export. main already disabled it for FSDP2 on those grounds; accelerate offload breaks the same assumption a different way -- a module's weights are freed when its materialization window closes, so the allocator can hand that address to an unrelated module in a later layer. Replace the FSDP2-only check with has_non_resident_weights(), a structural test (FSDP2 wrapping or accelerate offload hooks) that lives in core_utils beside requires_weight_materialization, so both residency questions are answered in one place. Tied-weight export (DiffusionGemma) is now explicitly resident-path only, and the per-window cache reset this branch previously carried is deleted -- with the caches disabled there is nothing left to scope. Two consequences worth stating: - moe_tied_cache is disabled under offload too, not just tied_cache. It is a fast path as well as a dedup, so a genuinely tied fused-expert block now repeats unpack/pack work there. Correct, since its (first_proj, down_proj) key degenerates to (0, 0) once those params go meta. - The check is model-level, so a partially offloaded model loses dedup for its resident modules too. Conservative, and consistent with how FSDP2 is treated. Measured on Qwen3.5-27B dense + FP8 under offload: 400 calls reached the cached path and none matched, on a model with no tied weights -- the cache could only ever have produced false positives there. Offload and batch exports of that model are bitwise identical across all 2031 tensors. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/registry.py | 30 ++++++++----------- modelopt/torch/export/unified_export_hf.py | 2 -- .../torch/quantization/utils/core_utils.py | 21 +++++++++++++ .../unit/torch/export/test_offload_export.py | 19 ++++++++++++ 4 files changed, 52 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 2be393c6603..8815bd48c05 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -32,7 +32,7 @@ import torch import torch.nn as nn -from modelopt.torch.utils.distributed import is_fsdp2_model +from modelopt.torch.quantization.utils.core_utils import has_non_resident_weights __all__ = [ "ExportContext", @@ -51,6 +51,9 @@ class ExportContext: recycled by PyTorch's allocator across exports, causing silent false-positive aliasing. ``tied_cache`` (int keys) holds dense Linear / per-expert wrapper dedup; ``moe_tied_cache`` (tuple keys) holds MoE fused-experts module dedup. + + Both are ``None`` when the model's weights are not resident for the whole export + (FSDP2 or accelerate offload) — see :meth:`__post_init__`. """ model: nn.Module @@ -60,27 +63,18 @@ class ExportContext: moe_tied_cache: dict[tuple[int, int], nn.Module] | None = field(default_factory=dict) def __post_init__(self) -> None: - # FSDP2 may recycle data_ptr() values as modules are resharded, so pointer-keyed dedup can - # falsely alias distinct weights. Disable it for FSDP2; consequently, legitimately tied - # packed weights and scale buffers are not re-aliased and may be stored as duplicates. + # Pointer-keyed dedup needs data_ptr() to identify a tensor for the whole export. + # That only holds while weights stay resident: FSDP2 recycles addresses as modules + # are resharded, and accelerate frees a module's weights when its materialization + # window closes, leaving the allocator free to hand the address to an unrelated + # module. Disable dedup whenever weights move; tied weights are then written as + # duplicates rather than re-aliased, so tied-weight export (DiffusionGemma) is + # supported on the resident path only. # TODO: replace this with stable, name-based tied-group deduplication. - if is_fsdp2_model(self.model): + if has_non_resident_weights(self.model): self.tied_cache = None self.moe_tied_cache = None - def reset_tied_caches(self) -> None: - """Drop dedup state between materialization windows. - - The streaming export frees each module after exporting it, so a recycled - address would alias the next module to the wrong weights. Genuine sharing is - always within one window, so nothing is lost. No-op when dedup is already - disabled (FSDP2), which recycles addresses for the same reason. - """ - if self.tied_cache is not None: - self.tied_cache.clear() - if self.moe_tied_cache is not None: - self.moe_tied_cache.clear() - ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index d0d29cd0112..e3bc2a600ad 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1158,7 +1158,6 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: for param_name, param in list(sub_mod._parameters.items()): if param is not None and param.device.type == "cuda": sub_mod._parameters[param_name] = None - ctx.reset_tied_caches() torch.cuda.empty_cache() # Non-decoder modules whose weights are not directly readable (embed_tokens, norm, @@ -1179,7 +1178,6 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: continue seen_keys.add(full_key) _stream_tensor(full_key, tensor) - ctx.reset_tied_caches() # GPU-resident parameters and persistent buffers (not covered by the above loops). # named_buffers() includes non-persistent buffers that state_dict() excludes; filter them. diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index b58991278d4..9dd06f804d4 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -30,6 +30,7 @@ from modelopt.torch.quantization.config import QuantizerCfgEntry from modelopt.torch.utils import get_unwrapped_name, print_rank_0 +from modelopt.torch.utils.distributed import is_fsdp2_model from modelopt.torch.utils.network import temporarily_remove_accelerate_hook if TYPE_CHECKING: @@ -647,6 +648,26 @@ def requires_weight_materialization(module, root_model, name_to_module: dict | N return _get_offload_hook(hook) is not None +def has_non_resident_weights(module: nn.Module) -> bool: + """Whether any weight under ``module`` lives outside it for part of the export. + + Structural (FSDP2 wrapping, accelerate offload hooks) rather than a snapshot of + current placement: offloaded weights come and go as materialization windows open and + close, so a point-in-time check answers differently depending on when it runs. + + Callers use this for decisions that must hold for a whole export — notably + pointer-keyed dedup, which needs ``data_ptr()`` to identify a tensor throughout. + """ + if is_fsdp2_model(module): + return True + from ..plugins.accelerate import _get_offload_hook + + return any( + (hook := getattr(m, "_hf_hook", None)) is not None and _get_offload_hook(hook) is not None + for m in module.modules() + ) + + @contextmanager def persistent_materialization(layer, writeback: bool = True): """Keep all layer weights materialized on GPU for the duration. diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 2d7962f9aa8..0d8dfb5ac27 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -379,3 +379,22 @@ def test_streaming_shard_writer_accepts_extra_tensors(): assert "mtp.fc.weight" in weight_map with safe_open(str(Path(tmpdir) / weight_map["mtp.fc.weight"]), framework="pt") as f: assert torch.equal(f.get_tensor("mtp.fc.weight"), torch.full((2, 2), 7.0)) + + +def test_export_context_disables_tie_dedup_when_weights_move(): + """Offloaded models get no pointer-keyed dedup — data_ptr is not a stable identity. + + An offloaded module's weights are freed when its materialization window closes, so a + recycled address would alias an unrelated module. Tied-weight export is therefore + supported on the resident path only, matching what FSDP2 already does. + """ + from modelopt.torch.export.registry import ExportContext + + resident = nn.Linear(8, 8) + offloaded, _ = _make_offloaded_linear() + + assert ExportContext(model=resident, dtype=torch.float16).tied_cache == {} + ctx = ExportContext(model=offloaded, dtype=torch.float16) + + assert ctx.tied_cache is None + assert ctx.moe_tied_cache is None From 8d51c8c03d8783795cadad3ca939bbda9f15a63a Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:07:58 +0000 Subject: [PATCH 20/21] refactor(export): consolidate offload detection into one helper has_non_resident_weights() duplicated _has_accelerate_offload()'s module scan and dropped its try/except ImportError, so constructing an ExportContext raised on installs without accelerate -- an optional 'hf' extra -- including on the fully resident path. Promote the scan to core_utils as has_accelerate_offload(), keeping the import guard, and define has_non_resident_weights() as its composition with is_fsdp2_model(). unified_export_hf and the four detection tests now share the one implementation, so the copies can no longer drift. Also trim the __post_init__ comment now that the data_ptr rationale lives on has_non_resident_weights(), and point the TODO at the _tied_weights_keys resolution in _collect_canonical_tied_patterns, which already does the name-based tie identity the fix needs. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/registry.py | 16 ++++----- modelopt/torch/export/unified_export_hf.py | 18 ++-------- .../torch/quantization/utils/core_utils.py | 21 ++++++----- .../unit/torch/export/test_offload_export.py | 36 +++++++++---------- 4 files changed, 39 insertions(+), 52 deletions(-) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 8815bd48c05..6f4d4be0a88 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -53,7 +53,8 @@ class ExportContext: dedup; ``moe_tied_cache`` (tuple keys) holds MoE fused-experts module dedup. Both are ``None`` when the model's weights are not resident for the whole export - (FSDP2 or accelerate offload) — see :meth:`__post_init__`. + (FSDP2 or accelerate offload), since ``data_ptr`` keys are meaningless once weights + move. """ model: nn.Module @@ -63,14 +64,11 @@ class ExportContext: moe_tied_cache: dict[tuple[int, int], nn.Module] | None = field(default_factory=dict) def __post_init__(self) -> None: - # Pointer-keyed dedup needs data_ptr() to identify a tensor for the whole export. - # That only holds while weights stay resident: FSDP2 recycles addresses as modules - # are resharded, and accelerate frees a module's weights when its materialization - # window closes, leaving the allocator free to hand the address to an unrelated - # module. Disable dedup whenever weights move; tied weights are then written as - # duplicates rather than re-aliased, so tied-weight export (DiffusionGemma) is - # supported on the resident path only. - # TODO: replace this with stable, name-based tied-group deduplication. + # data_ptr() only identifies a tensor while it stays resident, so dedup is unsafe + # once weights move. Tied weights are then written as duplicates rather than + # re-aliased, making tied-weight export (DiffusionGemma) resident-path only. + # TODO: dedup by tied-group name instead, reusing the _tied_weights_keys + # resolution in _collect_canonical_tied_patterns, which survives weight moves. if has_non_resident_weights(self.model): self.tied_cache = None self.moe_tied_cache = None diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index e3bc2a600ad..4df8df7b0d8 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -61,6 +61,7 @@ from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, 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 @@ -841,19 +842,6 @@ def _process_quantized_modules( _dispatch_export_handler(name, sub_module, ctx) -def _has_accelerate_offload(model: nn.Module) -> bool: - """Return True if any module in model has a CPU- or disk-offload accelerate hook.""" - try: - from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook - except ImportError: - return False - for mod in model.modules(): - hook = getattr(mod, "_hf_hook", None) - if hook is not None and _get_offload_hook(hook) is not None: - return True - return False - - class _StreamingShardWriter: """Write tensors to safetensors shard files without accumulating the full state dict. @@ -1278,7 +1266,7 @@ def _export_transformers_checkpoint( # Offloaded models need their weights materialized layer-by-layer, which this # whole-state-dict path cannot do; export_hf_checkpoint() streams them instead. - if _has_accelerate_offload(model): + if has_accelerate_offload(model): raise NotImplementedError( "_export_transformers_checkpoint does not support disk/CPU-offloaded models. " "Use export_hf_checkpoint() which dispatches to _export_transformers_checkpoint_streaming." @@ -1884,7 +1872,7 @@ def export_hf_checkpoint( ) # Streaming path writes shard files layer-by-layer without accumulating the full # state dict in RAM (peak = 1 layer + 1 shard buffer vs. ~764 GiB for Ultra 550B). - _offloaded = _has_accelerate_offload(model) + _offloaded = has_accelerate_offload(model) try: if _offloaded: diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 9dd06f804d4..4d77607936c 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -648,6 +648,18 @@ def requires_weight_materialization(module, root_model, name_to_module: dict | N return _get_offload_hook(hook) is not None +def has_accelerate_offload(module: nn.Module) -> bool: + """Return True if any module in ``module`` has a CPU- or disk-offload accelerate hook.""" + try: + from ..plugins.accelerate import _get_offload_hook + except ImportError: + return False + + return any( + _get_offload_hook(getattr(m, "_hf_hook", None)) is not None for m in module.modules() + ) + + def has_non_resident_weights(module: nn.Module) -> bool: """Whether any weight under ``module`` lives outside it for part of the export. @@ -658,14 +670,7 @@ def has_non_resident_weights(module: nn.Module) -> bool: Callers use this for decisions that must hold for a whole export — notably pointer-keyed dedup, which needs ``data_ptr()`` to identify a tensor throughout. """ - if is_fsdp2_model(module): - return True - from ..plugins.accelerate import _get_offload_hook - - return any( - (hook := getattr(m, "_hf_hook", None)) is not None and _get_offload_hook(hook) is not None - for m in module.modules() - ) + return has_accelerate_offload(module) or is_fsdp2_model(module) @contextmanager diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 0d8dfb5ac27..911b81f4ada 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -32,11 +32,9 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.quant_utils import _postprocess_single_tensor -from modelopt.torch.export.unified_export_hf import ( - _export_quantized_weight, - _has_accelerate_offload, - _StreamingShardWriter, -) +from modelopt.torch.export.registry import ExportContext +from modelopt.torch.export.unified_export_hf import _export_quantized_weight, _StreamingShardWriter +from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload # --------------------------------------------------------------------------- # Helpers @@ -54,18 +52,18 @@ def _make_offloaded_linear(dim: int = 16): # --------------------------------------------------------------------------- -# _has_accelerate_offload +# has_accelerate_offload # --------------------------------------------------------------------------- def test_has_accelerate_offload_true(): linear, _ = _make_offloaded_linear() - assert _has_accelerate_offload(linear) is True + assert has_accelerate_offload(linear) is True def test_has_accelerate_offload_false_no_hooks(): linear = nn.Linear(16, 16) - assert _has_accelerate_offload(linear) is False + assert has_accelerate_offload(linear) is False def test_has_accelerate_offload_false_non_offload_hook(): @@ -73,7 +71,7 @@ def test_has_accelerate_offload_false_non_offload_hook(): linear = nn.Linear(16, 16) hook = AlignDevicesHook(execution_device="cpu", offload=False) add_hook_to_module(linear, hook) - assert _has_accelerate_offload(linear) is False + assert has_accelerate_offload(linear) is False def test_has_accelerate_offload_detects_nested_module(): @@ -93,7 +91,7 @@ def forward(self, x): add_hook_to_module(parent.child, hook) set_module_tensor_to_device(parent.child, "weight", "meta") - assert _has_accelerate_offload(parent) is True + assert has_accelerate_offload(parent) is True # --------------------------------------------------------------------------- @@ -381,20 +379,18 @@ def test_streaming_shard_writer_accepts_extra_tensors(): assert torch.equal(f.get_tensor("mtp.fc.weight"), torch.full((2, 2), 7.0)) -def test_export_context_disables_tie_dedup_when_weights_move(): - """Offloaded models get no pointer-keyed dedup — data_ptr is not a stable identity. +def test_export_context_dedup_follows_weight_residency(): + """Pointer-keyed dedup is enabled only while weights stay resident. An offloaded module's weights are freed when its materialization window closes, so a recycled address would alias an unrelated module. Tied-weight export is therefore supported on the resident path only, matching what FSDP2 already does. """ - from modelopt.torch.export.registry import ExportContext + resident_ctx = ExportContext(model=nn.Linear(8, 8), dtype=torch.float16) + assert resident_ctx.tied_cache == {} + assert resident_ctx.moe_tied_cache == {} - resident = nn.Linear(8, 8) offloaded, _ = _make_offloaded_linear() - - assert ExportContext(model=resident, dtype=torch.float16).tied_cache == {} - ctx = ExportContext(model=offloaded, dtype=torch.float16) - - assert ctx.tied_cache is None - assert ctx.moe_tied_cache is None + offloaded_ctx = ExportContext(model=offloaded, dtype=torch.float16) + assert offloaded_ctx.tied_cache is None + assert offloaded_ctx.moe_tied_cache is None From f5badf396db97a885b0bfa032553057fb503bf45 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:25:41 +0000 Subject: [PATCH 21/21] test(export): pin the dedup opt-outs, and correct the DTensor rationale The comment on the sync_tied_input_amax guard claimed a sharded DTensor's address "belongs to a local shard the allocator recycles on reshard". That is wrong: DTensor.data_ptr() is unconditionally 0 for both Shard and Replicate placements, with the real address reachable only via to_local(). The failure mode is therefore deterministic rather than occasional -- every DTensor collides on key 0, collapsing the model into one tied group -- which is what makes the guard load-bearing. Correct the rationale so the TODO'd name-based dedup is designed against the real behavior. Add two tests for behavior nothing covered: - the FSDP2 leg of has_non_resident_weights, which could be deleted outright without failing a test while silently restoring pointer-keyed aliasing; - that tied modules each pack their own weight when dedup is off. Verified non-vacuous: with a live cache the two weights alias (equal data_ptr). Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/quant_utils.py | 10 +++---- .../unit/torch/export/test_offload_export.py | 28 +++++++++++++++++++ 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index c0323cabc63..2886ad52511 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -1610,11 +1610,11 @@ def sync_tied_input_amax(model: nn.Module) -> int: # Fused MoE: 3-D source tensors with shared input quantizers first_proj_attr = getattr(m, "_first_proj_attr", "gate_up_proj") first_proj_input_quantizer_attr = f"{first_proj_attr}_input_quantizer" - # data_ptr() only identifies a tensor that is resident: every meta tensor reports - # 0, which would collapse unrelated modules into a single tied group and merge - # their amaxes model-wide. A sharded DTensor is no better -- its address belongs - # to a local shard the allocator recycles on reshard, which is why FSDP2 disables - # pointer-keyed dedup outright (see ExportContext.__post_init__). + # data_ptr() only identifies a tensor that is resident. Meta tensors and DTensors + # both report 0 -- a DTensor's real address is reachable only through to_local() -- + # so without this guard every such module collapses into one tied group and has its + # amax merged model-wide. This is why FSDP2 disables pointer-keyed dedup outright + # (see ExportContext.__post_init__). if any(p.is_meta or isinstance(p, DTensor) for p in m.parameters(recurse=False)): if hasattr(m, "input_quantizer") or hasattr(m, first_proj_input_quantizer_attr): skipped_candidates += 1 diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 911b81f4ada..d8e68af324c 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -34,6 +34,7 @@ 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, _StreamingShardWriter +from modelopt.torch.quantization.utils import core_utils from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload # --------------------------------------------------------------------------- @@ -394,3 +395,30 @@ def test_export_context_dedup_follows_weight_residency(): offloaded_ctx = ExportContext(model=offloaded, dtype=torch.float16) assert offloaded_ctx.tied_cache is None assert offloaded_ctx.moe_tied_cache is None + + +def test_export_context_dedup_disabled_for_fsdp2(monkeypatch): + """FSDP2 shards recycle addresses, so its dedup opt-out must survive the offload rework.""" + monkeypatch.setattr(core_utils, "is_fsdp2_model", lambda _: True) + + ctx = ExportContext(model=nn.Linear(8, 8), dtype=torch.float16) + assert ctx.tied_cache is None + assert ctx.moe_tied_cache is None + + +def test_tied_weights_exported_independently_without_cache(): + """With dedup off, tied modules each pack their own weight instead of aliasing. + + Guards the offload path: an alias would make two shard entries share storage, which + the writer must then drop or copy. Independent tensors keep both keys intact. + """ + shared = nn.Parameter(torch.randn(16, 16)) + first, second = nn.Linear(16, 16, bias=False), nn.Linear(16, 16, bias=False) + first.weight = second.weight = shared + + for linear in (first, second): + mtq.quantize(linear, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(1, 16))) + _export_quantized_weight(linear, torch.float16, _tied_cache=None) + + assert first.weight.data_ptr() != second.weight.data_ptr() + assert torch.equal(first.weight, second.weight)