Stream AutoScheme's per-layer sensitivity scoring block-by-block

auto_round/auto_scheme/delta_loss.py

@@ -563,14 +563,20 @@ def __init__(self, message):
 last_grad_input = None
 
 
-def prepare_model_low_gpu(model, block_inputs: dict = None, pbar=None, major_device="cpu"):
+def prepare_model_low_gpu(model, block_inputs: dict = None, pbar=None, major_device="cpu", disk_index=None):
     """Wrap every block's forward so that, for one calibration batch, it (1) moves itself to
     ``major_device`` on demand, (2) records its own inputs into ``block_inputs`` (on CPU) so
     they can be replayed later, and (3) moves itself back to CPU once done.
 
     Called once per calibration batch before ``model_forward_low_gpu`` runs the actual
     forward+backward -- the recorded ``block_inputs`` are what let the backward pass be
     replayed manually, one block at a time, without keeping every block resident on GPU.
+
+    When ``disk_index`` is set (streaming mode -- the model is a meta-device skeleton,
+    see ``gen_layer_config``/``disk_stream_util.py``), each block's real weights are
+    materialized from the checkpoint right before its own forward and released back to
+    meta right after, instead of assuming the block already has real CPU-resident weights
+    to shuffle to GPU and back.
     """
     block_inputs.clear()
     for n, m in model.named_modules():
@@ -590,6 +596,10 @@ def new_forward(*args, **kwargs):
             """Move the block to device, run its original forward, cache its (CPU) inputs
             for later replay, then move the block back to CPU.
             """
+            if disk_index is not None:
+                from auto_round.utils.disk_stream_util import materialize_module
+
+                materialize_module(module, module_name, disk_index, device=major_device)
             move_module_to_tuning_device(module, major_device=major_device)
             # for n,m in module.named_modules():
             #     if hasattr(m, "post_init_qdqw"):
@@ -608,7 +618,12 @@ def new_forward(*args, **kwargs):
             }
             block_inputs[module_name] = input_info
 
-            module.to("cpu")
+            if disk_index is not None:
+                from auto_round.utils.disk_stream_util import free_module
+
+                free_module(module)
+            else:
+                module.to("cpu")
             memory_monitor.update(device_list=major_device)
             # clear_memory(device_list=major_device) #slow
             # memory_monitor.log_summary()
@@ -685,14 +700,19 @@ def model_forward(model, data, **forward_kwargs):
     return model(**prepared, **forward_kwargs), prepared
 
 
-def model_forward_low_gpu(model, dataloader, major_device="cuda", pbar=None, scheme_tag=None):
+def model_forward_low_gpu(model, dataloader, major_device="cuda", pbar=None, scheme_tag=None, disk_index=None):
     """Run one full scoring pass (all calibration batches) in low-GPU-memory mode.
 
     For each batch: capture per-block inputs via ``prepare_model_low_gpu``, run a forward
     pass whose backward is deliberately interrupted at the last block (``backward_pre_hook``
     raising ``MyCustomError``), then manually replay the backward pass block-by-block
     (moving each block to ``major_device`` only for its own recompute + backward, then back
     to CPU) so only one block's weights need to be resident on GPU at a time.
+
+    When ``disk_index`` is set (streaming mode -- the model is a meta-device skeleton),
+    each block's real weights are materialized from the checkpoint right before use and
+    released back to meta right after, both here (the manual reverse-order backward
+    replay) and in ``prepare_model_low_gpu`` (the initial forward capture pass).
     """
     block_inputs = {}
     total_batches = len(dataloader) if hasattr(dataloader, "__len__") else None
@@ -710,7 +730,7 @@ def backward_pre_hook(module, grad_input):
         raise MyCustomError("Interrupt backward pass")
 
     for batch_idx, data in enumerate(dataloader, start=1):
-        prepare_model_low_gpu(model, block_inputs, major_device=major_device, pbar=pbar)
+        prepare_model_low_gpu(model, block_inputs, major_device=major_device, pbar=pbar, disk_index=disk_index)
 
         # lm_head sits outside every decoder block, so it never gets `grad_mode=True`
         # in the manual block-by-block backward below. Scope the fix narrowly to
@@ -777,6 +797,10 @@ def backward_pre_hook(module, grad_input):
             for n, m in block_module.named_modules():
                 if hasattr(m, "grad_mode"):
                     m.grad_mode = True
+            if disk_index is not None:
+                from auto_round.utils.disk_stream_util import materialize_module
+
+                materialize_module(block_module, block_name, disk_index, device=major_device)
             move_module_to_tuning_device(block_module, major_device=major_device)
 
             # Set the block to eval mode while enabling gradient computation
@@ -816,7 +840,12 @@ def backward_pre_hook(module, grad_input):
                 break
 
             del block_output, main_output, block_input_args, block_input_kwargs
-            block_module.to("cpu")
+            if disk_index is not None:
+                from auto_round.utils.disk_stream_util import free_module
+
+                free_module(block_module)
+            else:
+                block_module.to("cpu")
 
             # clear_memory(device_list=major_device) # this one is very slow and seems does not affect max ram usage
             memory_monitor.update()
@@ -856,6 +885,7 @@ def get_score_for_scheme(
     force_mllm: bool = False,
     model_name: Optional[str] = None,
     scheme_tag: Optional[str] = None,
+    disk_index=None,
 ):
     """Wrap every quantizable layer in ``quant_layer_names`` with a scoring wrapper, run
     forward(+backward, unless RTN-only) calibration over ``nsamples`` examples from
@@ -1078,11 +1108,20 @@ def _build_mllm_calib_dataloader():
                     "AutoScheme(force_mllm): cannot build mllm dataloader. "
                     "Provide a `processor` and a multimodal `dataset`."
                 )
-            model_forward_low_gpu(model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag)
+            model_forward_low_gpu(
+                model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag, disk_index=disk_index
+            )
         else:
             try:
                 dataloader = _build_calib_dataloader()
-                model_forward_low_gpu(model, dataloader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag)
+                model_forward_low_gpu(
+                    model,
+                    dataloader,
+                    major_device=major_device,
+                    pbar=pbar,
+                    scheme_tag=scheme_tag,
+                    disk_index=disk_index,
+                )
             except Exception as exc:  # noqa: BLE001
                 if not is_vlm:
                     raise
@@ -1094,7 +1133,14 @@ def _build_mllm_calib_dataloader():
                 batch_size = 1
                 if mllm_loader is None:
                     raise
-                model_forward_low_gpu(model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag)
+                model_forward_low_gpu(
+                    model,
+                    mllm_loader,
+                    major_device=major_device,
+                    pbar=pbar,
+                    scheme_tag=scheme_tag,
+                    disk_index=disk_index,
+                )
     else:
         for n, m in model.named_modules():
             if hasattr(m, "grad_mode"):
@@ -1861,6 +1907,7 @@ def _gen_layer_config(
     device_list=None,
     processor=None,
     is_vlm: bool = False,
+    disk_index=None,
 ):
     """Score every candidate scheme in ``auto_scheme.options`` against ``quant_layer_names``
     and return per-layer per-scheme losses used by the caller to pick a final bit-width
@@ -1878,8 +1925,15 @@ def _gen_layer_config(
     # Create offload context for CPU RAM optimization
     # Note: low_cpu_mem_usage only works when low_gpu_mem_usage is also enabled,
     # because it requires layer-by-layer processing
+    #
+    # When disk_index is set, gen_layer_config already built the model as a
+    # meta-device skeleton and materialize_module/free_module (called directly
+    # around each block's use, see get_score_for_scheme/model_forward_low_gpu/
+    # prepare_model_low_gpu above) are the actual streaming mechanism --
+    # OffloadManager's hook-based approach doesn't apply to a model that never
+    # had real CPU-resident weights to begin with.
     offload_context = None
-    if auto_scheme.low_cpu_mem_usage and auto_scheme.low_gpu_mem_usage:
+    if disk_index is None and auto_scheme.low_cpu_mem_usage and auto_scheme.low_gpu_mem_usage:
         _model_dir = model_name
         if _model_dir is None and hasattr(model, "config"):
             _model_dir = getattr(model.config, "_name_or_path", None)
@@ -2225,6 +2279,13 @@ def _save_per_op_scores(index, scheme, cache_key, cache_path, per_op_scores):
             and num_gpus >= 1
             and len(uncached_indices) >= 2
             and not need_imatrix
+            # Each parallel worker fully loads its own copy of the model
+            # (_load_scheme_worker_model) in a separate process -- incompatible
+            # with a meta-device streaming skeleton, whose entire point is to
+            # avoid ever materializing a full copy. Force serial scoring
+            # (which honors disk_index via materialize_module/free_module)
+            # instead when streaming is active.
+            and disk_index is None
         )
         if not parallel_enabled and len(uncached_indices) >= 2:
             logger.info(
@@ -2360,7 +2421,13 @@ def _serialize_scheme(scheme):
                 pbar.reset(total=pbar_cnt)
 
         if not parallel_done:
-            if uncached_indices:
+            if uncached_indices and disk_index is None:
+                # Skipped in streaming mode: materialize_model_ only acts on
+                # ReplacementModuleBase (fused-MoE) instances -- a no-op for
+                # dense models like ours -- but it also warns once per
+                # still-meta parameter/buffer, which would flood the log with
+                # one warning per decoder-block tensor (intentionally still
+                # meta, to be streamed on demand later).
                 from auto_round.modeling.fused_moe.replace_modules import materialize_model_
 
                 materialize_model_(model)
@@ -2422,6 +2489,7 @@ def _serialize_scheme(scheme):
                             force_mllm=force_mllm,
                             model_name=model_name,
                             scheme_tag=scheme_tag,
+                            disk_index=disk_index,
                         )
                     memory_monitor.update()
                     memory_monitor.log_summary()
@@ -2689,12 +2757,48 @@ def gen_layer_config(
     """
     model_name = None
     is_vlm = False
+    disk_index = None
     if isinstance(model, str):
         model_name = model
-        model, tokenizer, processor, _, _, is_vlm, _ = load_model(model_name, device="cpu", use_auto_mapping=False)
+        is_vlm = is_mllm_model(model_name)
+        if not is_vlm and auto_scheme.low_cpu_mem_usage and low_gpu_mem_usage:
+            # Disk-streamed load (meta-device skeleton + on-demand per-block
+            # materialize/free, see disk_stream_util.py) instead of
+            # load_model()'s full-checkpoint CPU RAM load -- infeasible for a
+            # checkpoint bigger than available RAM. Falls back to load_model()
+            # for anything build_meta_model doesn't cover.
+            try:
+                from auto_round.utils.disk_stream_util import build_meta_model, materialize_non_block_params
+
+                model, tokenizer, disk_index = build_meta_model(model_name)
+                block_prefixes = flatten_list(get_block_names(model, quant_vision=is_vlm))
+                materialize_non_block_params(model, block_prefixes, disk_index, device="cpu")
+            except Exception as exc:  # noqa: BLE001
+                logger.warning(
+                    f"AutoScheme streaming load failed ({exc}); falling back to "
+                    f"load_model() (needs the whole checkpoint resident in RAM)."
+                )
+                disk_index = None
+                model, tokenizer, processor, _, _, is_vlm, _ = load_model(
+                    model_name, device="cpu", use_auto_mapping=False
+                )
+        else:
+            model, tokenizer, processor, _, _, is_vlm, _ = load_model(model_name, device="cpu", use_auto_mapping=False)
     else:
         # Object passed in: still try to detect VLM so we can pick the right dataloader later.
-        _, _, _, _, _, is_vlm, _ = load_model(model)
+        try:
+            _, _, _, _, _, is_vlm, _ = load_model(model)
+        except Exception:  # noqa: BLE001
+            is_vlm = False
+        # By the time AutoRound's compressor calls into AutoScheme, ModelContext
+        # has already turned a string model into a real object -- meaning the
+        # `isinstance(model, str)` branch above never actually runs in the
+        # standard `AutoRound(model=path, ...)` API flow. When ModelContext
+        # itself built the object as a meta skeleton (AR_DISK_STREAM_MODEL=1),
+        # it stashes the SafetensorsIndex on the model so we can pick it up
+        # here instead of re-detecting streaming mode from scratch (or,
+        # worse, silently treating a meta model as if it were fully real).
+        disk_index = getattr(model, "_disk_stream_index", None)
 
     # ---- Vision-tower scoring requires a full backward ---- #
     # ``model_forward_low_gpu`` only walks the language tower (it uses
@@ -2789,6 +2893,7 @@ def _enable_gc(mod):
             min_avg_bit_scheme=min_avg_bit_scheme,
             processor=processor,
             is_vlm=is_vlm,
+            disk_index=disk_index,
         )
     except torch.OutOfMemoryError:
         logger.warning(
@@ -2820,6 +2925,7 @@ def _enable_gc(mod):
             min_avg_bit_scheme=min_avg_bit_scheme,
             processor=processor,
             is_vlm=is_vlm,
+            disk_index=disk_index,
         )
 
     return res

auto_round/auto_scheme/gen_auto_scheme.py

@@ -100,9 +100,14 @@ def __init__(
         processor=None,
     ):
         self.auto_scheme = auto_scheme
-        if self.auto_scheme.low_cpu_mem_usage:
-            logger.info("force not using `low_cpu_mem_usage` in AutoScheme")
-            self.auto_scheme.low_cpu_mem_usage = False
+        # Upstream unconditionally forced low_cpu_mem_usage=False here
+        # (commit 0c9c5b1d, "reduce memory consumption... for gguf") because of an
+        # acknowledged bug for mixed INT4/INT8 schemes under the old OffloadManager-
+        # based streaming path (see an earlier commit's "low_cpu_mem_usage: bool =
+        # False # TODO bug for INT4 INT8 mixed bug"). We don't use that old path at
+        # all -- see delta_loss.py's materialize_module/free_module-based streaming,
+        # gated on this same flag -- so the bug that motivated disabling it doesn't
+        # apply here.
         self.model = model
         self.tokenizer = tokenizer
         self.processor = processor

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