Stream targeted block re-quantization's calibration pass

auto_round/auto_scheme/delta_loss.py

@@ -651,6 +651,9 @@ def __init__(self, message):
         super().__init__(message)
 
 
+last_grad_input = None
+
+
 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

auto_round/calibration/llm.py

@@ -81,16 +81,108 @@ def calibration(self, block_names, nsamples, layer_names=None, last_cache_name=N
         ):
             # low_gpu_mem_usage or calibrate only the embedding layer (also fast on CPU)
             calibrate_on_cpu = True
+            # When AR_DISK_STREAM_MODEL built a
+            # meta-device skeleton and to_quant_block_names restricts
+            # quant_block_list to fewer than all decoder blocks (a targeted
+            # re-quantization of specific blocks in an already-quantized
+            # checkpoint, not the normal full-model run), this forward pass
+            # still needs REAL weights in every block leading up to (and, when
+            # there's only one target block so last_cache_name ends up None
+            # and no early-stop applies, all the way through) the target
+            # block(s) -- but nothing else materializes blocks outside
+            # quant_block_list for this specific forward pass. They stay meta
+            # forever, and the forward silently propagates meta-ness through
+            # them until it collides with a genuinely-materialized module
+            # (the final norm, or any block actually in quant_block_list that
+            # already went through the compressor's own offload/reload cycle)
+            # -- confirmed via a reproduction against the tiny hybrid-MoE test
+            # fixture (checkpoint_full_arch_test, to_quant_block_names="model.
+            # layers.7"): "Tensor on device meta is not on the expected device
+            # cpu!" inside Qwen3_5MoeRMSNorm.forward. Full (unrestricted) runs
+            # never hit this: quant_block_list already covers every block in
+            # that case, so there's nothing left outside it to leave meta.
+            # Fix: stream-materialize (then free) any block that's still meta
+            # for the duration of this one forward pass, reusing the same
+            # stream_block_forward primitive _streaming_eval_model() already
+            # uses in bin/el_quantize_autoround_mixed.py for the analogous
+            # held-out-loss-eval case.
+            stream_ctx = None
+            _moved_tensors = []
+            if envs.AR_DISK_STREAM_MODEL:
+                disk_index = getattr(self.model, "_disk_stream_index", None)
+                if disk_index is not None:
+                    from auto_round.utils import get_block_names, get_module
+
+                    meta_block_names = [
+                        name
+                        for name in flatten_list(get_block_names(self.model))
+                        if any(p.device.type == "meta" for p in get_module(self.model, name).parameters())
+                    ]
+                    if meta_block_names:
+                        from auto_round.utils.disk_stream_util import stream_block_forward
+
+                        # This whole branch is calibrate_on_cpu -- every other
+                        # tensor in this forward pass (hidden states,
+                        # already-materialized non-block params) lives on cpu,
+                        # not device_manager.device (the GPU tuning device).
+                        # Materializing ONLY the streamed blocks on GPU caused
+                        # a real cuda:0/cpu mismatch crash the first time it
+                        # was tried against the full 397B-scale checkpoint --
+                        # hence the cpu default. But a cpu forward through
+                        # every pre-target block of a 100B+ model is unusably
+                        # slow for the targeted re-quantization use case, so
+                        # AR_CALIB_STREAM_DEVICE (set by el_requantize_blocks
+                        # .py) opts the WHOLE pass onto one device coherently:
+                        # every already-real (non-meta) param/buffer is moved
+                        # there for the duration (the same recipe bin/
+                        # el_quantize_autoround_mixed.py's _streaming_eval_
+                        # model() already proved at 207GB scale, including its
+                        # stray-buffer sweep for e.g. RoPE inv_freq), blocks
+                        # stream-materialize there, and calib() batches follow
+                        # model.device automatically. Everything is moved back
+                        # afterwards so the tuning phase sees the exact layout
+                        # it would have without this.
+                        calib_stream_device = envs.AR_CALIB_STREAM_DEVICE or "cpu"
+                        if calib_stream_device != "cpu":
+                            for module in self.model.modules():
+                                for _pname, _t in list(module.named_parameters(recurse=False)) + list(
+                                    module.named_buffers(recurse=False)
+                                ):
+                                    if _t.device.type != "meta" and str(_t.device) != calib_stream_device:
+                                        _moved_tensors.append((_t, str(_t.device)))
+                                        _t.data = _t.data.to(calib_stream_device)
+                            logger.info(
+                                "AR_CALIB_STREAM_DEVICE=%s: moved %d non-block tensors for the "
+                                "calibration forward; decoder blocks stream through the same device.",
+                                calib_stream_device,
+                                len(_moved_tensors),
+                            )
+
+                        stream_ctx = stream_block_forward(
+                            self.model,
+                            disk_index,
+                            device=calib_stream_device,
+                            block_names=meta_block_names,
+                        )
             try:
-                all_inputs = self.cache_inter_data(
-                    block_names, nsamples, layer_names=[], last_cache_name=last_cache_name
-                )
+                if stream_ctx is not None:
+                    with stream_ctx:
+                        all_inputs = self.cache_inter_data(
+                            block_names, nsamples, layer_names=[], last_cache_name=last_cache_name
+                        )
+                else:
+                    all_inputs = self.cache_inter_data(
+                        block_names, nsamples, layer_names=[], last_cache_name=last_cache_name
+                    )
             except NotImplementedError as error:
                 error_msg = str(error)
                 if "flash_attn::" in error_msg and "CPU" in error_msg:
                     cannot_calibrate_on_cpu = True
                 else:
                     raise error
+            finally:
+                for _t, _orig_device in _moved_tensors:
+                    _t.data = _t.data.to(_orig_device)
 
         if not calibrate_on_cpu or cannot_calibrate_on_cpu:
             try:

docs/environments_CN.md

@@ -251,6 +251,26 @@ AR_ALLOW_W8_ASYM=1 python -m auto_round --model ... --scheme W8A16 --asym --form
 export AR_RESUME_DIR=/path/to/resume/state
 ```
 
+### AR_DISK_STREAM_MODEL
+- **描述**：启用后，`AutoRound(model=<path>, ...)` 会将模型构建为 meta 设备骨架，而不是先把整个 checkpoint 完全加载到 CPU 内存；随后按需从 checkpoint 的 safetensors 分片中流式加载每个解码器块的真实权重——在该块被使用前（校准、调优或 `AutoScheme` 敏感度评分）才实体化，用完后立即释放回 meta。这样峰值 CPU 内存基本保持平稳，而不会随 checkpoint 大小成比例增长。非块参数(embedding、`lm_head`、最终归一化层)体积通常较小，仍会一次性加载。
+- **默认值**：`False`
+- **有效值**：`"1"`、`"true"`、`"yes"`(不区分大小写)表示启用；其他任何值表示禁用
+- **用途**：用于量化体积超过可用 CPU 内存 + GPU 显存总和的 checkpoint。仅在 `model` 为字符串(本地目录)路径时生效，对已加载的模型对象无效。
+
+```bash
+export AR_DISK_STREAM_MODEL=1
+```
+
+### AR_RESUME_DIR
+- **描述**：设置为目录路径后，逐块调优循环会在每完成一个块后将进度写入该目录，并在针对同一目录的新一次运行中从第一个未完成的块继续——而不是在崩溃或被杀死后从第 0 块重新开始整个调优过程。
+- **默认值**：未设置(不支持断点续跑)
+- **有效值**：任意可写目录路径
+- **用途**：用于大 checkpoint 的长时间量化任务，避免运行中途崩溃导致从头重跑的高昂代价。
+
+```bash
+export AR_RESUME_DIR=/path/to/resume/state
+```
+
 ## 使用示例
 
 ### 设置环境变量

test/test_cpu/core/test_targeted_block_calib_stream.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2026 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Integration test for targeted block re-quantization (``to_quant_block_names``)
+combined with disk streaming (``AR_DISK_STREAM_MODEL``).
+
+Regression coverage for: restricting tuning to a block that isn't the first
+one left every block *before* the target on the meta device, and the
+"cache block inputs" calibration forward silently propagated that meta-ness
+until it collided with a genuinely-materialized module -- "Tensor on device
+meta is not on the expected device cpu!". A full (unrestricted) run never
+hits this, since every block is already covered by quant_block_list in that
+case.
+"""
+
+import os
+
+import pytest
+
+from auto_round import AutoRound
+
+
+@pytest.fixture(autouse=True)
+def _clean_disk_stream_env():
+    previous = os.environ.get("AR_DISK_STREAM_MODEL")
+    yield
+    if previous is None:
+        os.environ.pop("AR_DISK_STREAM_MODEL", None)
+    else:
+        os.environ["AR_DISK_STREAM_MODEL"] = previous
+
+
+@pytest.fixture(scope="module")
+def tiny_opt_3layer_model_path():
+    from test.helpers import save_tiny_model
+
+    path = save_tiny_model("facebook/opt-125m", "./tmp/tiny_opt_3layer_model_path", num_layers=3)
+    yield path
+    import shutil
+
+    shutil.rmtree(path, ignore_errors=True)
+
+
+class TestTargetedBlockCalibStream:
+    def test_targeted_block_with_disk_streaming_does_not_crash(self, tiny_opt_3layer_model_path):
+        """Restricting to_quant_block_names to the LAST of 3 blocks leaves
+        blocks 0 and 1 meta-only (never in quant_block_list) while the
+        calibration forward pass still needs to run through them to reach
+        block 2 -- exactly the scenario that used to crash."""
+        os.environ["AR_DISK_STREAM_MODEL"] = "1"
+
+        ar = AutoRound(
+            model=tiny_opt_3layer_model_path,
+            scheme="W4A16",
+            iters=1,
+            nsamples=1,
+            to_quant_block_names="model.decoder.layers.2",
+        )
+        _, layer_config = ar.quantize()
+
+        quantized_layers = {name for name, cfg in layer_config.items() if "bits" in cfg}
+        assert quantized_layers, "expected the target block's layers to be quantized"
+        assert all(name.startswith("model.decoder.layers.2.") for name in quantized_layers)
+        assert not any(
+            name.startswith(("model.decoder.layers.0.", "model.decoder.layers.1.")) for name in quantized_layers
+        )
+
+    def test_targeted_block_without_disk_streaming_still_works(self, tiny_opt_3layer_model_path):
+        """Baseline: the same targeted re-quantization without disk streaming
+        never hit this bug (nothing is meta), so it must keep working too."""
+        os.environ.pop("AR_DISK_STREAM_MODEL", None)
+
+        ar = AutoRound(
+            model=tiny_opt_3layer_model_path,
+            scheme="W4A16",
+            iters=1,
+            nsamples=1,
+            to_quant_block_names="model.decoder.layers.2",
+        )
+        _, layer_config = ar.quantize()
+
+        quantized_layers = {name for name, cfg in layer_config.items() if "bits" in cfg}
+        assert quantized_layers
+        assert all(name.startswith("model.decoder.layers.2.") for name in quantized_layers)

test/test_cpu/schemes/test_auto_scheme_disk_stream.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2026 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Unit tests for AutoScheme's disk-streaming mode (AR_DISK_STREAM_MODEL).
+
+Verifies that streaming per-block sensitivity scoring from disk (instead of
+fully materializing the checkpoint on CPU RAM up front) produces the same
+mixed-bit layer_config as the non-streaming baseline, and that the underlying
+materialize/free primitives round-trip correctly.
+"""
+
+import os
+import shutil
+
+import pytest
+import torch
+
+from auto_round import AutoRound, AutoScheme
+from auto_round.utils.disk_stream_util import build_meta_model, free_module, materialize_module, total_resident_bytes
+
+
+@pytest.fixture(autouse=True)
+def _clean_disk_stream_env():
+    # AR_DISK_STREAM_MODEL is read lazily by auto_round.envs; make sure a test
+    # that sets it can't leak into whichever test runs next.
+    previous = os.environ.get("AR_DISK_STREAM_MODEL")
+    yield
+    if previous is None:
+        os.environ.pop("AR_DISK_STREAM_MODEL", None)
+    else:
+        os.environ["AR_DISK_STREAM_MODEL"] = previous
+
+
+class TestAutoSchemeDiskStream:
+    @pytest.fixture(autouse=True)
+    def setup_save_dir(self, tmp_path):
+        self.save_dir = str(tmp_path / "saved")
+        yield
+        shutil.rmtree(self.save_dir, ignore_errors=True)
+
+    def _gen_layer_config(self, model_name, target_bits=3.5):
+        # iters=1 (the standard tuning loop) rather than iters=0 (RTN): RTN's
+        # separate block-materialization path doesn't support disk streaming yet
+        # and is unrelated to this PR, which only streams AutoScheme's own
+        # sensitivity-scoring pass.
+        scheme = AutoScheme(avg_bits=target_bits, options=("W2A16", "W4A16", "BF16"), nsamples=1)
+        ar = AutoRound(model=model_name, scheme=scheme, iters=1, nsamples=1)
+        _, layer_config = ar.quantize()
+        return {name: cfg["bits"] for name, cfg in layer_config.items() if "bits" in cfg}
+
+    def test_disk_stream_matches_baseline_layer_config(self, tiny_opt_model_path):
+        """AR_DISK_STREAM_MODEL=1 must select the exact same per-layer bits as the
+        non-streaming baseline -- streaming changes *how* weights are loaded during
+        scoring, not the scores themselves."""
+        os.environ.pop("AR_DISK_STREAM_MODEL", None)
+        baseline_bits = self._gen_layer_config(tiny_opt_model_path)
+
+        os.environ["AR_DISK_STREAM_MODEL"] = "1"
+        streamed_bits = self._gen_layer_config(tiny_opt_model_path)
+
+        assert streamed_bits == baseline_bits
+
+    def test_disk_stream_default_off(self):
+        """With AR_DISK_STREAM_MODEL unset, behavior must be the unstreamed default."""
+        os.environ.pop("AR_DISK_STREAM_MODEL", None)
+        from auto_round import envs
+
+        assert envs.AR_DISK_STREAM_MODEL is False
+
+
+class TestDiskStreamUtilRoundTrip:
+    """Tests for the materialize/free primitives directly, independent of AutoScheme."""
+
+    def test_materialize_then_free_round_trip(self, tiny_opt_model_path):
+        model, _tokenizer, index = build_meta_model(tiny_opt_model_path)
+        block = model.model.decoder.layers[0]
+
+        for _, tensor in list(block.named_parameters()):
+            assert str(tensor.device) == "meta"
+
+        materialize_module(block, "model.decoder.layers.0", index, device="cpu")
+        for _, tensor in list(block.named_parameters()):
+            assert str(tensor.device) != "meta"
+        assert total_resident_bytes(block) > 0
+
+        free_module(block)
+        for _, tensor in list(block.named_parameters()):
+            assert str(tensor.device) == "meta"
+        assert total_resident_bytes(block) == 0