Single gpu disk offload PTQ for DSR1/Ultra

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,16 +643,40 @@ 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",
     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,
 ):
+    _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,

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
 
 

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 = []