diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 83a54849110..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 @@ -53,6 +54,61 @@ 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, @@ -587,6 +643,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 +661,22 @@ 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, ): + _install_transformers_compat_shims() 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 +780,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 +816,41 @@ 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: + 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 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/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..912c0003484 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") +_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, + 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 = _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": + 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): + 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(): + 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 + return new_key, _maybe_squeeze_scale(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, @@ -976,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 = {} @@ -1036,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 cee64c22c05..359e42811e8 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,6 +15,7 @@ """Code that export quantized Hugging Face models for deployment.""" +import itertools import json import re import tempfile @@ -95,6 +96,7 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import ( + _postprocess_single_tensor, fuse_prequant_layernorm, fuse_prequant_to_linear, get_activation_scaling_factor, @@ -568,6 +570,14 @@ def _export_quantized_weight( quantizer_attrs = quantizer_attr_names(weight_name) weight: nn.Parameter = getattr(sub_module, weight_name) + if weight.is_meta: + raise RuntimeError( + f"Weight '{weight_name}' of {type(sub_module).__name__} is a meta tensor during " + "export. If the model was loaded with disk/CPU offload, export must run inside an " + "enable_weight_access_and_writeback context. Use the offload-aware export path " + "(_process_quantized_modules_offloaded) rather than _process_quantized_modules." + ) + # Capture source identity BEFORE any tensor-creating operation below. # For HF-tied weights this matches across all modules sharing the # underlying Parameter; the cache lookup at the end of this function @@ -778,6 +788,20 @@ def _export_quantized_weight( torch.cuda.empty_cache() +def _dispatch_export_handler(name: str, sub_module: nn.Module, ctx: ExportContext) -> None: + """QLoRA skip, unpack-weight preprocessing, and handler dispatch for one module.""" + if ctx.is_modelopt_qlora and hasattr(sub_module, "base_layer"): + return + # Restore unpacked weight so the export path can read the live quantizer state. + if hasattr(sub_module, "weight_packed") or ( + "QuantFP8Linear" in type(sub_module).__name__ and sub_module.weight.element_size() <= 1 + ): + sub_module.unpack_weight() + handler = ExportModuleRegistry.match(sub_module) + if handler is not None: + handler(name, sub_module, ctx) + + def _process_quantized_modules( model: nn.Module, dtype: torch.dtype, @@ -811,20 +835,420 @@ 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"): + _dispatch_export_handler(name, sub_module, ctx) + + +def _has_accelerate_offload(model: nn.Module) -> bool: + """Return True if any module in model has a CPU- or disk-offload accelerate hook.""" + try: + from modelopt.torch.quantization.plugins.accelerate import _get_offload_hook + except ImportError: + return False + for mod in model.modules(): + hook = getattr(mod, "_hf_hook", None) + if hook is not None and _get_offload_hook(hook) is not None: + return True + return False + + +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() - # 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 + # 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()) ): - sub_module.unpack_weight() + 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. + + 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. + """ - handler = ExportModuleRegistry.match(sub_module) - if handler is not None: - handler(name, sub_module, ctx) + 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._total_bytes: int = 0 + # 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._total_bytes += 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 = 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) + 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: + 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, 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. + 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} + + # --- 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) + 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().contiguous().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 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) + _stream_tensor(name, tensor) + + 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) + + _sanitize_generation_config_for_save(model) + _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 def _export_transformers_checkpoint( @@ -874,13 +1298,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 +1347,20 @@ 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: + # MoE reconstruction happens per-layer inside _process_quantized_modules_offloaded. + quantized_state_dict = _process_quantized_modules_offloaded(model, dtype, is_modelopt_qlora) 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 @@ -1430,14 +1861,39 @@ 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) + export_state_dict = None + 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 @@ -1452,7 +1908,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: @@ -1481,23 +1940,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/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..05f9974df6e --- /dev/null +++ b/tests/unit/torch/export/test_offload_export.py @@ -0,0 +1,318 @@ +# 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 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 + 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.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 +# --------------------------------------------------------------------------- + + +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) + + +# --------------------------------------------------------------------------- +# _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 +# --------------------------------------------------------------------------- + + +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