Single gpu disk offload PTQ for DSR1/Ultra #2008

coderabbitai · 2026-07-30T19:24:15Z

🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash set -euo pipefail ast-grep outline modelopt/torch/export/quant_utils.py --view expanded | sed -n '1,240p' printf '\n--- slice around reported lines ---\n' sed -n '1000,1060p' modelopt/torch/export/quant_utils.py printf '\n--- nearby references to kv_cache_format and value ---\n' rg -n "kv_cache_format|value\.item\(|value\b" modelopt/torch/export/quant_utils.py

Repository: NVIDIA/Model-Optimizer

Length of output: 7807

🏁 Script executed:

#!/bin/bash set -euo pipefail python3 - <<'PY' from pathlib import Path p = Path("modelopt/torch/export/quant_utils.py") lines = p.read_text().splitlines() for start, end in [(1015, 1055), (920, 980)]: print(f"\n--- {p}:{start}-{end} ---") for i in range(start-1, min(end, len(lines))): print(f"{i+1:5d}: {lines[i]}") PY

Repository: NVIDIA/Model-Optimizer

Length of output: 4885

🏁 Script executed:

#!/bin/bash set -euo pipefail printf '%s\n' '--- uses of _postprocess_single_tensor ---' rg -n "_postprocess_single_tensor\(" -S modelopt/torch/export/quant_utils.py modelopt || true printf '\n%s\n' '--- streaming/state-dict path around postprocess_state_dict ---' sed -n '1050,1145p' modelopt/torch/export/quant_utils.py printf '\n%s\n' '--- any explicit cpu()/to(\"cpu\") around this path ---' rg -n "to\\(['\"]cpu['\"]\\)|\\.cpu\\(" modelopt/torch/export/quant_utils.py

Repository: NVIDIA/Model-Optimizer

Length of output: 4390

🏁 Script executed:

#!/bin/bash set -euo pipefail sed -n '1060,1115p' modelopt/torch/export/unified_export_hf.py printf '\n%s\n' '--- surrounding helpers/entrypoints ---' rg -n "stream|state_dict|_postprocess_single_tensor|postprocess_state_dict" modelopt/torch/export/unified_export_hf.py

Repository: NVIDIA/Model-Optimizer

Length of output: 7192

🏁 Script executed:

#!/bin/bash set -euo pipefail printf '%s\n' '--- enable_weight_access_and_writeback definition ---' rg -n "def enable_weight_access_and_writeback|class enable_weight_access_and_writeback|enable_weight_access_and_writeback\(" modelopt -S printf '\n%s\n' '--- streaming export around the offload context ---' sed -n '1088,1188p' modelopt/torch/export/unified_export_hf.py printf '\n%s\n' '--- surrounding offload helper docs ---' sed -n '1,260p' modelopt/torch/export/unified_export_hf.py

Repository: NVIDIA/Model-Optimizer

Length of output: 19523

Avoid the .item() sync in the streaming KV-cache warning path. This branch runs before the tensor is moved to CPU, so value.item() adds a device-host sync just to decide whether to log. Move the check after .cpu() or batch it with other export-side warnings.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/quant_utils.py` at line 1039, Update the KV-cache warning condition in the export path around kv_cache_format and value so it does not call value.item() while the tensor is still device-resident. Move the scalar comparison until after value has been transferred to CPU, or batch it with the existing export-side warning checks, while preserving the warning threshold and behavior.

Source: Coding guidelines

-Original file line number
+Diff line change
@@ Expand Up @@
         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",
         gpu_mem_percentage=0.8,
         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"
@@ Expand Down Expand Up / @@ -728,6 +751,20 @@ def has_pack_quantized_config(config): @@
                     return True
             return False
+        # Only the general load path below threads max_memory/offload_folder into
+        # from_pretrained; the specialized loaders build their own calls.
+        if _disk_offload and (
+            is_speculative(hf_config)
+            or has_pack_quantized_config(hf_config)
+            or get_original_hf_quant_method(hf_config) == "mxfp4"
+        ):
+            warnings.warn(
+                "offload_folder is ignored for speculative, pack-quantized, and MXFP4 "
+                "checkpoints: these use dedicated load paths that cannot offload. The model "
+                "will be loaded fully resident.",
+                UserWarning,
+            )
         if is_speculative(hf_config):
             model = AutoModelForCausalLM.from_pretrained(
                 ckpt_path,
@@ Expand Down Expand Up / @@ -772,7 +809,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:
+            # DeepSeek ships bundled modeling code, but the built-in class is what the
+            # disk-offload and streaming-export paths are validated against.
+            use_bundled_code = trust_remote_code and "Deepseek" in architecture
+            if not hasattr(transformers, architecture) or use_bundled_code:
                 if not hasattr(transformers, architecture):
                     warnings.warn(
                         f"Architecture {architecture} not found in transformers: {transformers.__version__}. "
@@ Expand Down Expand Up / @@ -809,24 +850,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,
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -578,6 +578,9 @@ def load_model(args: argparse.Namespace): @@
                 trust_remote_code=args.trust_remote_code,
                 use_seq_device_map=args.use_seq_device_map,
                 attn_implementation=args.attn_implementation,
+                offload_folder=args.offload_folder,
+                max_cpu_memory_gb=args.max_cpu_memory_gb,
+                max_gpu_memory_gb=args.max_gpu_memory_gb,
             )
         else:
             assert args.qformat in QUANT_CFG_CHOICES, (
@@ Expand Down Expand Up / @@ -1621,6 +1624,38 @@ def parse_args() -> argparse.Namespace: @@
                 "openai/gpt-oss-20b) and the target qformat is NVFP4-family."
             ),
         )
+        parser.add_argument(
+            "--offload_folder",
+            type=str,
+            default=None,
+            help=(
+                "Path to a local folder for disk-offloaded model weights. "
+                "When set, activates disk-offload mode: model weights that exceed the GPU+CPU "
+                "budgets are streamed from disk during calibration and export. "
+                "Pair with --max_cpu_memory_gb to cap CPU RAM usage. "
+                "Incompatible with --low_memory_mode and --use_seq_device_map."
+            ),
+        )
+        parser.add_argument(
+            "--max_cpu_memory_gb",
+            type=float,
+            default=None,
+            help=(
+                "Maximum CPU RAM budget in GiB for disk-offload model loading. "
+                "Only effective when --offload_folder is set. "
+                "Weights beyond this limit are streamed from disk."
+            ),
+        )
+        parser.add_argument(
+            "--max_gpu_memory_gb",
+            type=float,
+            default=None,
+            help=(
+                "Maximum GPU memory budget per device in GiB for disk-offload model loading. "
+                "Only effective when --offload_folder is set. "
+                "Defaults to 80%% of available GPU memory when not specified."
+            ),
+        )
         args = parser.parse_args()
         if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0):
@@ Expand Down Expand Up / @@ -1655,6 +1690,31 @@ def parse_args() -> argparse.Namespace: @@
         if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4:
             parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.")
+        if args.offload_folder is not None and args.low_memory_mode:
+            parser.error("--offload_folder (disk-offload) is not compatible with --low_memory_mode.")
+        if args.offload_folder is not None and args.use_seq_device_map:
+            parser.error(
+                "--offload_folder (disk-offload) is not compatible with --use_seq_device_map; "
+                "device_map=auto is used for disk-offload to let accelerate place layers across "
+                "GPU, CPU, and disk."
+            )
+        if args.offload_folder is not None and args.device == "cpu":
+            parser.error(
+                "--offload_folder (disk-offload) is not compatible with --device cpu; "
+                "device_map=cpu makes accelerate ignore the memory budgets and offload folder, "
+                "loading the whole model into RAM."
+            )
+        if args.offload_folder is None and (
+            args.max_cpu_memory_gb is not None or args.max_gpu_memory_gb is not None
+        ):
+            parser.error(
+                "--max_cpu_memory_gb/--max_gpu_memory_gb only apply to disk-offload loading; "
+                "pass --offload_folder to enable it."
+            )
         return args
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -959,6 +959,94 @@ 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,
@@ Expand All / @@ -976,31 +1064,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 = {}
@@ Expand Down Expand Up / @@ -1036,15 +1101,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 = []
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -32,7 +32,7 @@ @@
     import torch
     import torch.nn as nn
-    from modelopt.torch.utils.distributed import is_fsdp2_model
+    from modelopt.torch.quantization.utils.core_utils import has_non_resident_weights
     __all__ = [
         "ExportContext",
@@ Expand All / @@ -51,6 +51,10 @@ class ExportContext: @@
         recycled by PyTorch's allocator across exports, causing silent false-positive
         aliasing. ``tied_cache`` (int keys) holds dense Linear / per-expert wrapper
         dedup; ``moe_tied_cache`` (tuple keys) holds MoE fused-experts module dedup.
+        Both are ``None`` when the model's weights are not resident for the whole export
+        (FSDP2 or accelerate offload), since ``data_ptr`` keys are meaningless once weights
+        move.
         """
         model: nn.Module
@@ Expand All / @@ -60,11 +64,12 @@ class ExportContext: @@
         moe_tied_cache: dict[tuple[int, int], nn.Module] | None = field(default_factory=dict)
         def __post_init__(self) -> None:
-            # FSDP2 may recycle data_ptr() values as modules are resharded, so pointer-keyed dedup can
-            # falsely alias distinct weights. Disable it for FSDP2; consequently, legitimately tied
-            # packed weights and scale buffers are not re-aliased and may be stored as duplicates.
-            # TODO: replace this with stable, name-based tied-group deduplication.
-            if is_fsdp2_model(self.model):
+            # data_ptr() only identifies a tensor while it stays resident, so dedup is unsafe
+            # once weights move. Tied weights are then written as duplicates rather than
+            # re-aliased, making tied-weight export (DiffusionGemma) resident-path only.
+            # TODO: dedup by tied-group name instead, reusing the _tied_weights_keys
+            # resolution in _collect_canonical_tied_patterns, which survives weight moves.
+            if has_non_resident_weights(self.model):
                 self.tied_cache = None
                 self.moe_tied_cache = None
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Single gpu disk offload PTQ for DSR1/Ultra #2008

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot Jul 30, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Single gpu disk offload PTQ for DSR1/Ultra #2008

Are you sure you want to change the base?

Uh oh!

Single gpu disk offload PTQ for DSR1/Ultra #2008

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot Jul 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!