diff --git a/auto_round/auto_scheme/delta_loss.py b/auto_round/auto_scheme/delta_loss.py index 8d23ff1524..33d052374c 100644 --- a/auto_round/auto_scheme/delta_loss.py +++ b/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 diff --git a/auto_round/calibration/llm.py b/auto_round/calibration/llm.py index 9bd45ae2d7..08f1965f9f 100644 --- a/auto_round/calibration/llm.py +++ b/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: diff --git a/docs/environments_CN.md b/docs/environments_CN.md index 27ac26672d..09c0d4750b 100644 --- a/docs/environments_CN.md +++ b/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=, ...)` 会将模型构建为 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 +``` + ## 使用示例 ### 设置环境变量 diff --git a/test/test_cpu/core/test_targeted_block_calib_stream.py b/test/test_cpu/core/test_targeted_block_calib_stream.py new file mode 100644 index 0000000000..16dfe8bfb8 --- /dev/null +++ b/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) diff --git a/test/test_cpu/schemes/test_auto_scheme_disk_stream.py b/test/test_cpu/schemes/test_auto_scheme_disk_stream.py new file mode 100644 index 0000000000..4205bde954 --- /dev/null +++ b/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