From dd76c1b2fcccb83778bbd3e91362b200acb827bb 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 1/8] 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 | 136 ++++++++++++++---- .../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, 556 insertions(+), 50 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 83a54849110..4a27780cd28 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -28,6 +28,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 @@ -587,6 +641,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", @@ -594,9 +659,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" @@ -700,12 +777,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" ) @@ -737,24 +813,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 57a3dd6e264..21e6dc9ac06 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -538,6 +538,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, ( @@ -1563,6 +1566,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): @@ -1584,6 +1619,16 @@ def parse_args() -> argparse.Namespace: "the low-memory loader initializes quantizers from --qformat/--kv_cache_qformat." ) + 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 cee64c22c05..7bab9c95349 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -568,6 +568,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 @@ -778,6 +786,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, @@ -811,20 +833,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( @@ -874,13 +948,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) @@ -918,18 +997,21 @@ 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 accelerator is not None: - # Gather state_dict from all ranks - quantized_state_dict = accelerator.get_state_dict(model) + 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: - quantized_state_dict = model.state_dict() + _process_quantized_modules(model, dtype, is_modelopt_qlora) + _reconstruct_fused_moe_linear(model) + + if accelerator is not None: + # Gather state_dict from all ranks + quantized_state_dict = accelerator.get_state_dict(model) + else: + 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 178d7c354f66431bdc1e59e86c62359f173eb470 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 2/8] 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 7bab9c95349..4be6dfb5b6a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -894,8 +894,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 87ccd1a0be3eb83be05ebe1cf84e8840b4a74796 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 3/8] =?UTF-8?q?fix(export):=20correct=20offloaded=20export?= =?UTF-8?q?=20path=20=E2=80=94=20MoE=20ordering,=20lm=5Fhead=20handler,=20?= =?UTF-8?q?memory?= 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 4a27780cd28..06c815f7286 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -815,10 +815,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 4be6dfb5b6a..d03e625d2e5 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -854,16 +854,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 @@ -881,27 +882,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 @@ -910,20 +913,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 @@ -1032,9 +1035,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: _process_quantized_modules(model, dtype, is_modelopt_qlora) _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 808eff9364b29b7678ed833fc14051e6d0867eb6 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 4/8] 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 | 350 +++++++++++++++++- .../unit/torch/export/test_offload_export.py | 131 ++++++- 3 files changed, 545 insertions(+), 19 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 d03e625d2e5..8a411d73c07 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -95,6 +95,7 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import ( + _postprocess_single_tensor, fuse_prequant_layernorm, fuse_prequant_to_linear, get_activation_scaling_factor, @@ -936,6 +937,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, @@ -1546,14 +1832,38 @@ def export_hf_checkpoint( ) return + # 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 @@ -1568,7 +1878,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: @@ -1597,23 +1910,24 @@ 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) - 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) + 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 80a3f652cc4d33df8fdee9183ff99429ba2d46ea 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 5/8] 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 8a411d73c07..b2d6674a2e1 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 @@ -953,7 +955,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] = {} @@ -966,7 +968,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 @@ -1000,7 +1002,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) @@ -1013,7 +1015,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 @@ -1195,23 +1197,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) @@ -1835,6 +1829,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 8b6201b8b133b57d25fe9ca8454c992ce0091ab6 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 6/8] 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 | 111 +++++++++++---------- modelopt/torch/export/unified_export_hf.py | 59 ++++++++--- 2 files changed, 104 insertions(+), 66 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 06c815f7286..0bb97fd8300 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 @@ -28,60 +29,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 @@ -107,6 +54,61 @@ def matmul(input, weight, weight_scale_inv, block_size, output_dtype=None, activ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] +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, tokenizer, @@ -663,6 +665,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 b2d6674a2e1..359e42811e8 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 @@ -1109,8 +1108,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. @@ -1136,6 +1142,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) @@ -1152,7 +1166,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(): @@ -1196,8 +1210,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) @@ -1205,13 +1223,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 e4d69c7bc8556fe81b0b9c085633154c81b67037 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 7/8] 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 359e42811e8..e9d589a8395 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 @@ -1223,30 +1225,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 f76231394f93a6a9b946cc1ac3aefbacc37a9393 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 8/8] 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 | 229 ++++++------------ .../unit/torch/export/test_offload_export.py | 107 +++----- 2 files changed, 104 insertions(+), 232 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index e9d589a8395..4ba8e16394b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -575,9 +575,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. @@ -853,93 +852,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. @@ -1310,7 +1222,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) @@ -1351,8 +1263,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: _process_quantized_modules(model, dtype, is_modelopt_qlora) _reconstruct_fused_moe_linear(model) @@ -1814,6 +1728,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, @@ -1865,7 +1811,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: @@ -1886,15 +1831,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 @@ -1909,10 +1866,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: @@ -1921,64 +1875,25 @@ def export_hf_checkpoint( "names may not match the original HF hub checkpoint." ) - # 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) - - 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 + 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