Stream large checkpoints from disk during quantization

auto_round/compressors/base.py

@@ -393,6 +393,12 @@ def __init__(
         # second AutoRound(...) call reuses the previous instance and silently keeps
         # stale values (e.g. low_cpu_mem_usage=True from a prior run).
         CompressContext.reset_context()
+        # When the model was built as a meta skeleton (AR_DISK_STREAM_MODEL=1),
+        # give the offloader the original checkpoint path so it can materialize
+        # each block on first touch directly from disk instead of assuming
+        # blocks already hold real weights (see OffloadManager._reload).
+        if self.model_context.disk_stream_model_dir is not None:
+            self._offloader.model_dir = self.model_context.disk_stream_model_dir
         # Alternatively, you can use CompressContext.create_context
         self.compress_context = CompressContext(
             low_cpu_mem_usage,
@@ -403,6 +409,11 @@ def __init__(
             static_attention_dtype=self.static_attention_dtype,
         )
         self.shard_writer = None
+        # Resumability state deferred from Orchestrator._quantize_data_driven() until
+        # quantize_and_save()'s save_quantized() call actually succeeds; see the
+        # comment in quantize() near "is_immediate_saving" for why clearing is
+        # deferred.
+        self._resume_states = None
 
         # Flag for post_init idempotency.  Set to False here so post_init() can be called
         # either via quantize_and_save() (preferred, outside inference_mode) or directly
@@ -1625,6 +1636,14 @@ def quantize_and_save(
         model, folders = self.save_quantized(output_dir, inplace=inplace, return_folders=True, **kwargs)
         memory_monitor.log_summary()
 
+        # Only now -- after the full export (packing pass, config/tokenizer
+        # writes) has actually succeeded -- is it safe to drop the resume
+        # manifest. See the deferral comment in Orchestrator._quantize_data_driven().
+        if self._resume_states:
+            for rs in self._resume_states:
+                rs.clear()
+            self._resume_states = None
+
         return model, folders
 
 

auto_round/compressors/orchestrator.py

@@ -22,6 +22,7 @@
 from accelerate.big_modeling import dispatch_model
 from tqdm import tqdm
 
+from auto_round import envs
 from auto_round.calibration import CalibrationContext
 from auto_round.calibration.utils import (
     _update_inputs,
@@ -217,7 +218,19 @@ def _quantize_blocks(
                 modules = [get_module(model, n) for n in names]
                 m = WrapperMultiblock(modules)
 
-            if self.compress_context.low_cpu_mem_usage:
+            # Also reload when `AR_DISK_STREAM_MODEL` is set even if
+            # `low_cpu_mem_usage` has been forced False (e.g. GGUF export --
+            # see base.py's `_finalize_compress_context`, which disables
+            # `low_cpu_mem_usage` for gguf formats for reasons unrelated to disk
+            # streaming). Under streaming, a block starts on the meta device
+            # regardless of `low_cpu_mem_usage`, which only ever controlled whether
+            # to *free* it again after use -- without this, the block below is never
+            # materialized at all and `m.to(device)` crashes with "Cannot copy out
+            # of meta tensor". The block intentionally stays real afterward (no
+            # matching post-tune offload runs when `low_cpu_mem_usage` is False --
+            # see the `is_immediate_saving`-adjacent offload call further down),
+            # matching upstream's own choice not to cycle blocks for these formats.
+            if self.compress_context.low_cpu_mem_usage or envs.AR_DISK_STREAM_MODEL:
                 if nblocks == 1:
                     self._offloader.reload(model, n)
                 else:
@@ -550,6 +563,17 @@ def _quantize_data_driven(self) -> tuple[torch.nn.Module, dict[str, Any]]:
                 )
                 if not self._offloader.enabled:
                     self.compress_context.low_cpu_mem_usage = False
+            elif self.model_context._disk_stream_index is not None:
+                # Dense (non-MoE-patched) models normally get low_cpu_mem_usage
+                # disabled here because the per-block offload/reload dance is
+                # pointless when the whole model is already CPU-resident from
+                # the initial full load -- there's no memory to save. That
+                # assumption doesn't hold when the model started as a meta
+                # skeleton (AR_DISK_STREAM_MODEL=1): blocks are still on meta
+                # and must go through the same reload()-before/offload()-after
+                # cycle to get materialized from disk one at a time and freed
+                # again, so keep it enabled here.
+                pass
             else:
                 self.compress_context.low_cpu_mem_usage = False
         if len(all_blocks) > 1:

auto_round/context/model.py

@@ -14,6 +14,7 @@
 
 import gc
 import importlib
+import os
 from typing import Any, Callable, Optional, Union
 
 import torch
@@ -105,6 +106,12 @@ def __init__(
         self.need_calib = need_calib
         self.quant_nontext_module = quant_nontext_module
 
+        # Remember the original string model name
+        # so OffloadManager can later materialize blocks directly from this
+        # checkpoint.
+        self.disk_stream_model_dir = model if isinstance(model, str) else None
+        self._disk_stream_index = None
+
         # Load model and run basic initialization eagerly so the model is ready
         # by the time BaseCompressor.post_init() runs.
         self._load_model()
@@ -114,6 +121,8 @@ def __init__(
                 "AutoRound does not support parameters on meta device. "
                 "Please use more GPUs by setting `--device 0,1,2,3` or just place the model on CPU."
             )
+        if self._disk_stream_index is not None:
+            self._materialize_disk_stream_non_block_params()
         check_and_mark_quantized_module(self.model)
         self.model = self.model.eval()
         self.shared_cache_keys = get_shared_keys(self.model)
@@ -148,9 +157,42 @@ def _load_model(self):
         if is_mllm_model(self.model, platform=self.platform):
             self.is_mllm = True
             if isinstance(self.model, str):
-                self.model, self.processor, self.tokenizer, self.image_processor = mllm_load_model(
-                    self.model, platform=self.platform, device="cpu", model_dtype=self.model_dtype
-                )
+                # Multimodal checkpoints used to
+                # bypass disk streaming entirely -- mllm_load_model fully
+                # materializes the checkpoint on CPU, infeasible for a 100B+
+                # VLM (e.g. Ornith). Build the same meta skeleton the text
+                # path uses (build_meta_model resolves the multimodal class
+                # from config.architectures) and load the processor stack
+                # cheaply; non-block params INCLUDING the whole vision tower
+                # (small, and needed real for calibration forwards and RTN)
+                # are materialized right after the meta-device guard in
+                # __init__, exactly like the text path. Falls back to the
+                # full mllm_load_model on any failure.
+                loaded_via_meta = False
+                if envs.AR_DISK_STREAM_MODEL and os.path.isdir(self.model):
+                    try:
+                        self.model, self.tokenizer, self._disk_stream_index = self._build_disk_stream_model(
+                            self.disk_stream_model_dir
+                        )
+                        from transformers import AutoProcessor
+
+                        self.processor = AutoProcessor.from_pretrained(
+                            self.disk_stream_model_dir, trust_remote_code=self.trust_remote_code
+                        )
+                        self.image_processor = getattr(self.processor, "image_processor", None)
+                        loaded_via_meta = True
+                    except Exception:
+                        logger.warning(
+                            "AR_DISK_STREAM_MODEL requested but building a multimodal meta "
+                            "skeleton for %s failed; falling back to a full CPU load.",
+                            self.disk_stream_model_dir,
+                            exc_info=True,
+                        )
+                        self._disk_stream_index = None
+                if not loaded_via_meta:
+                    self.model, self.processor, self.tokenizer, self.image_processor = mllm_load_model(
+                        self.model, platform=self.platform, device="cpu", model_dtype=self.model_dtype
+                    )
         elif is_diffusion_model(self.model):
             self.is_diffusion = True
             self.pipe, self.model = diffusion_load_model(
@@ -197,18 +239,87 @@ def _load_model(self):
             gc.collect()
             _force_trim_malloc()
 
-            self.model, self.tokenizer = llm_load_model(
-                self.model,
-                platform=self.platform,
-                device="cpu",  # always load cpu first
-                model_dtype=self.model_dtype,
-                trust_remote_code=self.trust_remote_code,
-            )
+            if envs.AR_DISK_STREAM_MODEL:
+                try:
+                    self.model, self.tokenizer, self._disk_stream_index = self._build_disk_stream_model(
+                        self.disk_stream_model_dir
+                    )
+                except Exception:
+                    logger.warning(
+                        "AR_DISK_STREAM_MODEL requested but building a meta skeleton for %s failed; "
+                        "falling back to a normal full CPU load.",
+                        self.disk_stream_model_dir,
+                        exc_info=True,
+                    )
+                    self._disk_stream_index = None
+                    self.model, self.tokenizer = llm_load_model(
+                        self.model,
+                        platform=self.platform,
+                        device="cpu",
+                        model_dtype=self.model_dtype,
+                        trust_remote_code=self.trust_remote_code,
+                    )
+            else:
+                self.model, self.tokenizer = llm_load_model(
+                    self.model,
+                    platform=self.platform,
+                    device="cpu",  # always load cpu first
+                    model_dtype=self.model_dtype,
+                    trust_remote_code=self.trust_remote_code,
+                )
         elif self.tokenizer is None and not self.is_diffusion and self.need_calib:
             raise ValueError("A tokenizer must be set for non-str model input")
 
         self._model_loaded = True
 
+    def _build_disk_stream_model(self, model_name: str):
+        """Build an all-meta skeleton instead of
+        fully materializing the checkpoint on CPU RAM. Left fully meta here
+        (not even embeddings/lm_head materialized yet) so it passes the
+        existing ``unsupported_meta_device`` guard, which only allows models
+        that are either fully real or fully meta (with ``model.path`` set).
+        Non-block params are materialized for real right after that guard
+        runs, in ``__init__``, via ``_materialize_disk_stream_non_block_params``.
+        """
+        from auto_round.utils.disk_stream_util import build_meta_model
+
+        model, tokenizer, index = build_meta_model(model_name, trust_remote_code=self.trust_remote_code)
+        model.path = model_name
+        # Stash the index on the model object
+        # itself so downstream code that only has a reference to the model
+        # (not this ModelContext) -- e.g. AutoScheme's gen_layer_config, which
+        # runs after ModelContext has already turned the string into an object
+        # -- can still find it instead of re-scanning the checkpoint.
+        model._disk_stream_index = index
+        return model, tokenizer, index
+
+    def _materialize_disk_stream_non_block_params(self) -> None:
+        """Materialize embeddings/lm_head/norm (everything outside the
+        quantizable decoder blocks) for real, leaving the (typically 100+GB
+        combined) decoder blocks on meta for later per-block materialize/free
+        by OffloadManager. No-op unless the model was built via
+        ``_build_disk_stream_model``."""
+        if self._disk_stream_index is None:
+            return
+        from auto_round.utils import flatten_list, get_block_names
+        from auto_round.utils.disk_stream_util import materialize_non_block_params
+
+        block_prefixes = flatten_list(get_block_names(self.model, quant_vision=self.quant_nontext_module))
+        materialize_non_block_params(self.model, block_prefixes, self._disk_stream_index, device="cpu")
+
+        # Tied output embeddings (e.g. lm_head.weight
+        # tied to embed_tokens.weight) have no entry of their own in the checkpoint's
+        # safetensors index -- the on-disk format only stores the input embedding and
+        # relies on the model re-establishing the tie at load time. A normal
+        # from_pretrained() handles this automatically; our meta-skeleton +
+        # per-tensor materialize path does not, so without this the tied output
+        # module is silently left on meta (materialize_non_block_params only logs a
+        # warning and moves on). Re-tying now, after the real input embedding has
+        # been materialized above, makes the tied module real too by sharing the
+        # same (now real) Parameter object.
+        if hasattr(self.model, "tie_weights"):
+            self.model.tie_weights()
+
     def _import_custom_moe_replacements(self, model_or_config) -> None:
         model_type = getattr(model_or_config, "model_type", None)
         module_name = _CUSTOM_MOE_REPLACEMENT_MODULES.get(model_type)

auto_round/envs.py

@@ -27,6 +27,8 @@
     AR_AUTO_SCHEME_NSAMPLES: Optional[int] = None
     AR_AUTO_SCHEME_BATCH_SIZE: Optional[int] = None
     AR_ENABLE_AUTO_SCHEME_PARALLEL: bool = False
+    AR_DISK_STREAM_MODEL: bool = False
+    AR_RESUME_DIR: Optional[str] = None
 
 
 def _get_optional_positive_int_env(name: str) -> Optional[int]:
@@ -56,6 +58,11 @@ def _get_optional_positive_int_env(name: str) -> Optional[int]:
     "AR_DISABLE_DATASET_SUBPROCESS": lambda: os.getenv("AR_DISABLE_DATASET_SUBPROCESS", "0").lower() in ("1", "true"),
     "AR_DISABLE_COPY_MTP_WEIGHTS": lambda: os.getenv("AR_DISABLE_COPY_MTP_WEIGHTS", "0").lower()
     in ("1", "true", "yes"),
+    # Device for the disk-streamed calibration forward pass in
+    # LLMCalibrator.collect()'s calibrate_on_cpu branch (targeted block
+    # re-quantization). Unset = upstream behavior (cpu). Set to e.g. "cuda:0"
+    # to run the whole pass on GPU -- see el_requantize_blocks.py.
+    "AR_CALIB_STREAM_DEVICE": lambda: os.getenv("AR_CALIB_STREAM_DEVICE", None),
     "AR_ACT_SCALE": lambda: float(os.getenv("AR_ACT_SCALE", "1.0")),
     "AR_ENABLE_ACT_MINMAX_TUNING": lambda: os.getenv("AR_ENABLE_ACT_MINMAX_TUNING", "0").lower()
     in ("1", "true", "yes"),
@@ -90,6 +97,16 @@ def _get_optional_positive_int_env(name: str) -> Optional[int]:
     # avoid multiple model-loading workers exhausting host RAM or device memory.
     "AR_ENABLE_AUTO_SCHEME_PARALLEL": lambda: os.getenv("AR_ENABLE_AUTO_SCHEME_PARALLEL", "0").lower()
     in ("1", "true", "yes"),
+    # When set, the model is built as a meta-device skeleton and streamed
+    # block-by-block from disk during quantization instead of being fully
+    # materialized on CPU RAM up front.
+    "AR_DISK_STREAM_MODEL": lambda: os.getenv("AR_DISK_STREAM_MODEL", "0").lower() in ("1", "true", "yes"),
+    # When set to a directory path, the per-block tuning loop checkpoints its
+    # progress there after each completed block, and resumes from the first
+    # not-yet-completed block on a fresh run against the same directory --
+    # instead of restarting the whole tuning pass from block 0 after a
+    # crash/kill. See auto_round/utils/resume.py.
+    "AR_RESUME_DIR": lambda: os.getenv("AR_RESUME_DIR", None),
 }