Add crash/resume support for tuning runs (AR_RESUME_DIR)

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/shard_writer.py

@@ -14,11 +14,13 @@
 
 import json
 import os
+import re
 from collections import OrderedDict
 from typing import Optional, Union
 
 import torch
 
+from auto_round import envs
 from auto_round.compressors.utils import _get_save_folder_name
 from auto_round.context.compress import CompressContext
 from auto_round.context.model import ModelContext
@@ -95,8 +97,77 @@ def __init__(
         self.total_param_size_bytes = 0
         self.skipped_meta_tensors = []
 
+        # When resumability is active
+        # (AR_RESUME_DIR set), a fresh process's ShardWriter otherwise has no
+        # idea a previous, crashed process already flushed some shards to
+        # output_dir -- it would restart shard_counter at 0 and overwrite
+        # `model-shard-00001...`, and finalize()'s index would only cover the
+        # tensors written by *this* process, producing a corrupt/incomplete
+        # checkpoint. Gated on AR_RESUME_DIR so normal (non-resuming) runs
+        # never change behavior even if output_dir happens to be reused.
+        #
+        # Deliberately NOT run here in __init__: at construction time (from
+        # post_init(), inside quantize_and_save()) `self.output_dir` still
+        # reflects the pre-`_get_export_dir()` path -- the final subfolder
+        # (e.g. `<model>-w4g128/`) hasn't been appended yet, so discovery
+        # would silently look in the wrong directory and find nothing
+        # (confirmed empirically: this exact ordering bug let a resumed run's
+        # blocks 0-2 through tuning correctly, then still lose their output,
+        # since `_flush_shard`/`finalize` never learned about the crashed
+        # run's shards). Deferred to the first real `_flush_shard()` call
+        # instead, by which point `output_dir` is always the final path.
+        self._existing_shards_discovered = False
+
         ShardWriter._initialized = True
 
+    def _discover_existing_shards(self) -> None:
+        """Recover shard-writer state from shard files a previous (crashed)
+        process already flushed to ``output_dir``, so this process continues
+        shard numbering instead of colliding with them, and ``finalize()``'s
+        index covers tensors from both processes.
+
+        Only files still in the pre-``finalize()`` temp naming
+        (``model-shard-NNNNN.<ext>``) are considered: once ``finalize()`` runs
+        it renames everything to the final HF layout, so a directory with no
+        such temp files means either nothing has been flushed yet, or a prior
+        run already finished -- neither should be treated as in-progress
+        shards to adopt.
+        """
+        output_dir = self.output_dir
+        if not os.path.isdir(output_dir):
+            return
+        pattern = re.compile(rf"^model-shard-(\d+)\.{re.escape(self.shard_suffix)}$")
+        found = []
+        for fname in os.listdir(output_dir):
+            m = pattern.match(fname)
+            if m:
+                found.append((int(m.group(1)), fname))
+        if not found:
+            return
+        found.sort()
+        for _, fname in found:
+            path = os.path.join(output_dir, fname)
+            params = self._read_shard_tensor_names(path)
+            self.shard_meta.append({"tmp_file": fname, "params": params, "dir": output_dir})
+            self._all_saved.update(params)
+        self.shard_counter = found[-1][0]
+        logger.info(
+            f"ShardWriter: discovered {len(found)} already-flushed shard(s) in {output_dir} "
+            f"from a previous run; resuming shard numbering from {self.shard_counter}."
+        )
+
+    def _read_shard_tensor_names(self, path: str) -> list[str]:
+        """Read only the tensor-name header of an already-flushed shard file,
+        without materializing any tensor data."""
+        if self.use_safetensors:
+            from safetensors import safe_open
+
+            with safe_open(path, framework="pt") as f:
+                return list(f.keys())
+        else:
+            sd = torch.load(path, map_location="meta")
+            return list(sd.keys())
+
     @property
     def output_dir(self) -> str:
         """Derive the output directory from the current CompressContext at access time.
@@ -258,6 +329,10 @@ def _flush_shard(self):
         if not self.current_shard_tensors:
             return
 
+        if envs.AR_RESUME_DIR and not self._existing_shards_discovered:
+            self._discover_existing_shards()
+            self._existing_shards_discovered = True
+
         self.shard_counter += 1
         output_dir = self.output_dir
         os.makedirs(output_dir, exist_ok=True)

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),
 }
 
 

auto_round/export/export_to_autoround/export.py

@@ -164,6 +164,18 @@ def pack_layer(layer_name, model, backend, device=None):
     if type(layer) not in SUPPORTED_LAYER_TYPES:  ##already packed
         return
 
+    # A resumed disk-streamed run only
+    # materializes/quantizes the blocks it didn't already finish in a prior
+    # (crashed) process. Blocks it skipped are never touched in *this*
+    # process and stay on the meta device, while their packed weights
+    # already live in shard files the previous process flushed to disk (see
+    # ShardWriter._discover_existing_shards). There is nothing to pack here
+    # -- attempting to would fail (no real weight data to read `.scale`
+    # from) and would be redundant even if it didn't, since the on-disk
+    # export for this layer is already complete.
+    if layer.weight.device.type == "meta":
+        return
+
     if int(layer.act_bits) <= 8:
         return pack_qact_layer(layer_name, model)