From dac780c344926315ad851ea873b8d068713e4eee Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:36:55 +0200 Subject: [PATCH 01/25] Add AR_DISK_STREAM_MODEL and AR_RESUME_DIR env vars Two opt-in switches used by the disk-streaming and resumability work that follows in later commits. Default off (unset) preserves upstream behavior exactly. Signed-off-by: Fabrizio del Tin --- auto_round/envs.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/auto_round/envs.py b/auto_round/envs.py index 0769754cf..4f4798b95 100644 --- a/auto_round/envs.py +++ b/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"), + # Local addition: 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), } From 083bb4405ebfab431256be2befd823ca1845e7a1 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:36:55 +0200 Subject: [PATCH 02/25] Add disk-streaming primitive for per-block checkpoint materialization New auto_round/utils/disk_stream_util.py, no upstream equivalent. Lazy, mmap-backed reads of individual tensors by name straight from a checkpoint's safetensors shards, plus meta<->real materialize/free for a whole module. Also provides build_meta_model() (a meta skeleton + tokenizer + SafetensorsIndex, narrower than llm_load_model -- no bagel/glm/mxfp4/HPU special-casing) and materialize_non_block_params() (real-loads everything outside the decoder blocks: embeddings/ lm_head/final norm). Both materialize functions pass dtype=values[full_name].dtype explicitly to accelerate's set_module_tensor_to_device(): without it, accelerate casts real checkpoint data to whatever dtype the meta skeleton's parameter happened to declare, not the checkpoint's real dtype -- silently wrong for any module built without a matching dtype context (e.g. an unfused-MoE replacement module's per-expert nn.Linears, built under torch.device("meta") alone with no dtype, which default to float32 regardless of the checkpoint's actual dtype). This is the streaming primitive; it isn't wired into AutoRound's own model loading or tuning loop yet -- that follows in later commits. Signed-off-by: Fabrizio del Tin --- auto_round/utils/disk_stream_util.py | 302 +++++++++++++++++++++++++++ 1 file changed, 302 insertions(+) create mode 100644 auto_round/utils/disk_stream_util.py diff --git a/auto_round/utils/disk_stream_util.py b/auto_round/utils/disk_stream_util.py new file mode 100644 index 000000000..3fecc26d2 --- /dev/null +++ b/auto_round/utils/disk_stream_util.py @@ -0,0 +1,302 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +# Local addition (not upstream AutoRound). Ported verbatim from this project's +# vendor/reap/src/reap/disk_stream_util.py -- architecture-agnostic (only needs +# module dotted names + accelerate's set_module_tensor_to_device), so it carries +# over unchanged. Used by auto_round/auto_scheme/delta_loss.py's streaming scoring +# path (get_score_for_scheme_streaming) to materialize one decoder block's real +# tensors from the checkpoint's safetensors shards right before scoring it and +# release them back to meta right after -- instead of AutoScheme's original +# approach of loading the entire model onto CPU via llm_load_model(device_map="cpu") +# (~207GB+ for a model like command-a-translate, not available on hardware with far +# less RAM+VRAM combined than the checkpoint size). See LOCAL_PATCHES.md. +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Dict + +import torch +import torch.nn as nn +from accelerate.utils import set_module_tensor_to_device +from safetensors import safe_open + +logger = logging.getLogger(__name__) + + +class SafetensorsIndex: + """Lazy, mmap-backed access to a checkpoint's tensors by name. + + Deliberately does NOT cache open safe_open() handles across calls: mmap'd pages + stay resident (counted in RSS) for as long as the mapping is open, even after + the torch tensor copied out of them is freed. For a checkpoint far larger than + available RAM, caching handles indefinitely would silently re-create the exact + problem this module exists to avoid (RSS creeping up by however much of the + file has been touched so far, instead of staying bounded to one block at a + time). Every read (or batch of reads via read_tensors) opens a shard, reads, + and closes/unmaps before returning. + """ + + def __init__(self, checkpoint_dir: str): + self.checkpoint_dir = Path(checkpoint_dir) + index_path = self.checkpoint_dir / "model.safetensors.index.json" + if index_path.exists(): + with open(index_path) as f: + self.weight_map: Dict[str, str] = json.load(f)["weight_map"] + else: + # Small, unsharded checkpoint: one model.safetensors file. + single_file = self.checkpoint_dir / "model.safetensors" + with safe_open(str(single_file), framework="pt") as f: + self.weight_map = {name: single_file.name for name in f.keys()} + + def has_tensor(self, name: str) -> bool: + return name in self.weight_map + + def read_tensor(self, name: str, device: str = "cpu") -> torch.Tensor: + return self.read_tensors([name], device=device)[name] + + def read_tensors(self, names: list[str], device: str = "cpu") -> Dict[str, torch.Tensor]: + """Read several tensors, grouped by shard file so each shard is opened and + closed (unmapped) once regardless of how many tensors are pulled from it.""" + by_shard: Dict[str, list[str]] = {} + for name in names: + by_shard.setdefault(self.weight_map[name], []).append(name) + + result: Dict[str, torch.Tensor] = {} + for shard_name, shard_tensor_names in by_shard.items(): + with safe_open(str(self.checkpoint_dir / shard_name), framework="pt") as f: + for name in shard_tensor_names: + tensor = f.get_tensor(name) + if device != "cpu": + tensor = tensor.to(device) + result[name] = tensor + return result + + def tensor_names_with_prefix(self, prefix: str) -> list[str]: + dotted = prefix if prefix.endswith(".") else prefix + "." + return [n for n in self.weight_map if n == prefix or n.startswith(dotted)] + + +def materialize_module( + module: nn.Module, module_name: str, index: SafetensorsIndex, device: str +) -> None: + """Populate `module`'s (currently meta) parameters/buffers with real data read + directly from the checkpoint, onto `device`. `module_name` is `module`'s dotted + path in the full model (used as the tensor-name prefix in the checkpoint). + + Local addition (auto_round only, not in REAP's original): AutoScheme's scoring + wraps quantized layers in ``AutoSchemeWrapperLinear``, which replaces a plain + ``nn.Linear`` with a wrapper holding the real layer at ``.orig_layer`` -- + inserting an extra ``.orig_layer`` path segment that doesn't exist in the + checkpoint's own tensor names. Strip it back out before looking up the name. + """ + import re as _re + + # Local addition (not upstream): fused-MoE replacement modules + # (SequentialQwen3_5MoeExperts and friends) expose UNFUSED per-expert + # parameter names (experts.{i}.gate_proj.weight ...) that don't exist in a + # checkpoint whose on-disk layout is the fused 3D one + # (experts.gate_up_proj [N, 2*inter, hidden] / experts.down_proj). The + # compressor's own tuning loop handles this via OffloadManager.reload + + # materialize_model_, but every bare materialize_module() consumer + # (AutoScheme's delta_loss scoring streams, stream_block_forward for the + # calibration/eval forwards) previously left those params on meta -- the + # meta-ness then propagated silently until a crash far downstream. Map + # each unfused name onto its fused on-disk tensor and slice. + _FUSED_RE = _re.compile(r"^(.*\.experts)\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$") + _fused_cache: dict = {} + + def _fused_lookup(full_name: str): + m = _FUSED_RE.match(full_name) + if not m: + return None + prefix, expert_idx, proj = m.group(1), int(m.group(2)), m.group(3) + fused_name = f"{prefix}.gate_up_proj" if proj in ("gate_proj", "up_proj") else f"{prefix}.down_proj" + if not index.has_tensor(fused_name): + return None + if fused_name not in _fused_cache: + _fused_cache[fused_name] = index.read_tensors([fused_name], device=device)[fused_name] + fused = _fused_cache[fused_name][expert_idx] + if proj == "down_proj": + return fused.contiguous() + inter = fused.shape[0] // 2 + return (fused[:inter] if proj == "gate_proj" else fused[inter:]).contiguous() + + targets = [] # (param_name, full_checkpoint_name) + fused_targets = [] # (param_name, sliced_value) + for name, tensor in list(module.named_parameters()) + list(module.named_buffers()): + if str(tensor.device) != "meta": + continue # already materialized (e.g. shared/tied weights) + full_name = f"{module_name}.{name}".replace(".orig_layer.", ".") + if not index.has_tensor(full_name): + sliced = _fused_lookup(full_name) + if sliced is not None: + fused_targets.append((name, sliced)) + continue + logger.warning("No checkpoint tensor found for %s, leaving on meta", full_name) + continue + targets.append((name, full_name)) + + for name, value in fused_targets: + set_module_tensor_to_device(module, name, device, value=value, dtype=value.dtype) + _fused_cache.clear() + + if not targets: + return + values = index.read_tensors([full_name for _, full_name in targets], device=device) + for name, full_name in targets: + # Local addition (not upstream): explicit dtype= is required here. + # accelerate's set_module_tensor_to_device(), when dtype isn't passed, + # casts `value` to the *existing* (meta) parameter's declared dtype -- + # not the checkpoint's real dtype. A meta skeleton built without an + # enclosing dtype context (e.g. Qwen3_5MoeExperts' per-expert Linears, + # built under `torch.device("meta")` alone) defaults that declared + # dtype to float32 regardless of the checkpoint being bf16, so without + # this, real bf16 weights silently get upcast to float32 on + # materialization -- found via a real 8-layer MoE fixture crashing + # with "expected m1 and m2 to have the same dtype" inside AutoScheme + # scoring. The checkpoint's own dtype must always win. + set_module_tensor_to_device(module, name, device, value=values[full_name], dtype=values[full_name].dtype) + + +def free_module(module: nn.Module) -> None: + """Release a module's real tensors back to the meta device, freeing memory.""" + for name, tensor in list(module.named_parameters()) + list(module.named_buffers()): + if str(tensor.device) == "meta": + continue + set_module_tensor_to_device(module, name, "meta") + + +def total_resident_bytes(model: nn.Module) -> int: + """Debug helper: sum the byte size of every non-meta parameter/buffer in + `model`. Used to diagnose whether blocks are genuinely returning to meta + after free_module(), or something else is holding real memory.""" + total = 0 + for _, tensor in list(model.named_parameters()) + list(model.named_buffers()): + if str(tensor.device) != "meta": + total += tensor.numel() * tensor.element_size() + return total + + +def build_meta_model(model_name: str, trust_remote_code: bool = True): + """Build a meta-device model skeleton (~0 RAM) plus its tokenizer and a + SafetensorsIndex for on-demand materialization, instead of AutoRound's own + ``llm_load_model(model_name, device_map="cpu")`` which fully materializes the + checkpoint on CPU RAM in one shot. Deliberately narrower than ``llm_load_model``: + only covers the common local-directory ``AutoModelForCausalLM`` case (no + bagel/glm/mxfp4/HPU special-casing) -- callers should fall back to + ``llm_load_model`` for anything this doesn't handle. See LOCAL_PATCHES.md. + """ + from accelerate import init_empty_weights + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code) + # Prefer the exact class named in config.architectures: AutoModelForCausalLM + # cannot resolve multimodal architectures (e.g. + # Qwen3_5MoeForConditionalGeneration), which previously forced VLM + # checkpoints down the full-CPU-load path. Same resolution strategy as + # reap/layerwise_prune.py's disk-streamed model builder. + import transformers as _transformers + + archs = getattr(config, "architectures", None) or [] + model_cls = next((getattr(_transformers, a) for a in archs if hasattr(_transformers, a)), None) + with init_empty_weights(): + if model_cls is not None: + model = model_cls(config) + else: + model = AutoModelForCausalLM.from_config(config, trust_remote_code=trust_remote_code) + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code) + index = SafetensorsIndex(model_name) + return model, tokenizer, index + + +def materialize_non_block_params( + model: nn.Module, block_prefixes: list[str], index: SafetensorsIndex, device: str +) -> None: + """Materialize every real (non-meta) parameter/buffer NOT under one of + ``block_prefixes`` -- i.e. embeddings, final norm, lm_head, and similar small + top-level modules -- leaving the (typically 100+GB combined) decoder blocks on + meta for later per-block materialize/free. These non-block modules are needed + continuously throughout scoring and are comparatively small even for large + vocabularies, so it's simplest to load them once, for real, up front. + """ + + def _in_block(name: str) -> bool: + return any(name == p or name.startswith(p + ".") for p in block_prefixes) + + targets = [] # (param_name, full_checkpoint_name) + for name, tensor in list(model.named_parameters()) + list(model.named_buffers()): + if str(tensor.device) != "meta" or _in_block(name): + continue + full_name = name.replace(".orig_layer.", ".") + if not index.has_tensor(full_name): + logger.warning("No checkpoint tensor found for %s, leaving on meta", full_name) + continue + targets.append((name, full_name)) + + if not targets: + return + values = index.read_tensors([full_name for _, full_name in targets], device=device) + for name, full_name in targets: + # See the matching comment in materialize_module() -- explicit dtype= + # is required so the checkpoint's real dtype wins over whatever the + # meta skeleton happened to declare. + set_module_tensor_to_device(model, name, device, value=values[full_name], dtype=values[full_name].dtype) + + +class stream_block_forward: + """Context manager: wrap every top-level decoder block's ``forward`` so it + materializes its own real weights from ``index`` right before running and + frees them back to meta right after -- letting a plain ``model(...)`` call + (e.g. for computing held-out loss) drive the model exactly as normal while + only ever one block's weights are resident at a time. + + Deliberately much lighter than the auto_scheme delta_loss.py streaming + forward it's modeled on (``prepare_model_low_gpu``/``model_forward_low_gpu``): + no input-caching for later backward replay, no grad-mode bookkeeping -- this + is for a plain inference-only forward pass (e.g. eval loss), not tuning. + """ + + def __init__(self, model: nn.Module, index: SafetensorsIndex, device: str, block_names: list[str] = None): + self.model = model + self.index = index + self.device = device + self.block_names = block_names if block_names is not None else _default_block_names(model) + self._originals: Dict[str, "callable"] = {} + + def __enter__(self): + for block_name in self.block_names: + module = _get_module(self.model, block_name) + self._originals[block_name] = module.forward + + def make_wrapped(module=module, block_name=block_name, original_forward=module.forward): + def wrapped(*args, **kwargs): + materialize_module(module, block_name, self.index, device=self.device) + try: + return original_forward(*args, **kwargs) + finally: + free_module(module) + + return wrapped + + module.forward = make_wrapped() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + for block_name, original_forward in self._originals.items(): + _get_module(self.model, block_name).forward = original_forward + return False + + +def _default_block_names(model: nn.Module) -> list[str]: + from auto_round.utils import get_block_names + + return get_block_names(model)[0] + + +def _get_module(model: nn.Module, name: str) -> nn.Module: + from auto_round.utils import get_module + + return get_module(model, name) From eaaf80c37561ed55035d0fa74802a914db25a271 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:37:36 +0200 Subject: [PATCH 03/25] context/model: build meta-device skeleton under AR_DISK_STREAM_MODEL ModelContext._load_model() unconditionally called llm_load_model(..., device="cpu") (or mllm_load_model for multimodal checkpoints) whenever model was a string, fully materializing the checkpoint on CPU RAM before any block-wise AutoScheme/tuning logic ever ran. This is the fix for the initial-load problem: nothing downstream can be memory-safe if the model is already 100%+ resident before either ever runs. When AR_DISK_STREAM_MODEL=1 and model is a string (not diffusion), build an all-meta skeleton via the new disk_stream_util.build_meta_model() instead, for both the plain-text and multimodal (mllm_load_model) paths. Sets model.path = model_name (satisfies the existing but previously-dead unsupported_meta_device() escape hatch, which only allows an all-meta model) and stashes the SafetensorsIndex both on self._disk_stream_index and on model._disk_stream_index, so code that only has the model object (e.g. AutoScheme's gen_layer_config, which runs after ModelContext has already turned a string into an object) can still find it. Materializes non-block params (embeddings/lm_head/ final norm) for real right after the meta-device guard passes, leaving the (typically 100+GB combined) decoder blocks meta for later per-block materialize/free. Falls back to the original full CPU load on any exception. Also re-ties output embeddings via model.tie_weights() right after materializing: a tied lm_head.weight has no entry of its own in the checkpoint's safetensors index (relies on the model re-establishing the tie at load time, which a normal from_pretrained() does automatically but per-tensor materialization does not), so without this the tied module is silently left on meta. Signed-off-by: Fabrizio del Tin --- auto_round/context/model.py | 133 +++++++++++++++++++++++++++++++++--- 1 file changed, 123 insertions(+), 10 deletions(-) diff --git a/auto_round/context/model.py b/auto_round/context/model.py index 285f6f5df..2dee5d01d 100644 --- a/auto_round/context/model.py +++ b/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 + # Local addition (not upstream): remember the original string model name + # so OffloadManager can later materialize blocks directly from this + # checkpoint (see disk_stream_model_dir below and LOCAL_PATCHES.md). + 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 - ) + # Local addition (not upstream): 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,89 @@ 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): + """Local addition (not upstream): 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``. + See LOCAL_PATCHES.md. + """ + 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 + # Local addition (not upstream): 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. See + # LOCAL_PATCHES.md. + 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``. See LOCAL_PATCHES.md.""" + 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") + + # Local addition (not upstream): 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. See LOCAL_PATCHES.md. + 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) From c90acd4dce5dff1c3b93d61a76927d74a4891a77 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:38:38 +0200 Subject: [PATCH 04/25] offload: materialize/reload blocks from meta correctly OffloadManager's existing per-block offload/reload cycle assumed every block started CPU-resident; starting from a meta skeleton (AR_DISK_STREAM_MODEL=1) broke it in three places: - _load_state_dict_into_module() copied a freshly-read real tensor onto the target parameter's existing device -- but for first-time materialization from meta, that existing device IS meta, so the copy silently discarded the real data instead of landing it on cpu. Now targets "cpu" specifically when the existing parameter is meta. - _save_to_disk() unconditionally recorded a block as saved even when its state_dict was empty (an all-meta block that hasn't been materialized yet has nothing real to persist). A later reload() then trusted that record and loaded an empty file, leaving the block meta. Now skips recording in that case. - _reload(), in "offload" mode, silently did nothing for a block not in self._saved (true for a still-meta block, or one _save_to_disk just started correctly skipping). Now falls back to load_block_from_model_files(self.model_dir, name, module) -- an existing upstream function, previously only used by "clean" mode -- reading the block directly from the original checkpoint. Requires compressors/base.py to propagate model_dir onto the offloader (next commit). Also fixes a real-scale bug found against qwen3.5-397b-base: when a checkpoint's on-disk MoE layout uses fused 3D expert tensors (experts.gate_up_proj/down_proj) but the in-memory module tree has already been replaced by unfused per-expert nn.Linears, assigning the fused key resolves to nothing and the experts stay meta. Added _maybe_split_fused_expert_keys(), which detects that mismatch and splits the fused tensor into per-expert keys via the existing missing_tensors.split_fused_expert_tensors() helper before assignment. Signed-off-by: Fabrizio del Tin --- auto_round/utils/offload.py | 93 ++++++++++++++++++++++++++++++++++++- 1 file changed, 92 insertions(+), 1 deletion(-) diff --git a/auto_round/utils/offload.py b/auto_round/utils/offload.py index 975b5f507..32f075603 100644 --- a/auto_round/utils/offload.py +++ b/auto_round/utils/offload.py @@ -65,12 +65,55 @@ # ===================================================================== +def _maybe_split_fused_expert_keys(state_dict: dict, module: torch.nn.Module) -> dict: + """Local addition (not upstream): a checkpoint whose on-disk MoE layout is + the fused 3D one (``...experts.gate_up_proj [N, 2*inter, hidden]`` / + ``...experts.down_proj``, e.g. the real qwen3.5-397b-base) cannot be + assigned onto a module tree whose experts were already replaced by the + unfused ``SequentialQwen3_5MoeExperts`` (per-expert ``nn.Linear``s) -- + the fused key silently resolves to nothing and the experts stay meta, + crashing far downstream at ``block.to(device)``. Detect that exact case + (fused key present, target tree lacks the fused attribute) and split the + fused tensor into per-expert keys via the same + ``split_fused_expert_tensors`` helper the missing-tensors export pass + uses. Trees that still hold the original fused module are left alone. + Never triggered before because every prior fixture's on-disk layout was + already unfused (transformers >=5.10 unfuses on save); the real + checkpoint is fused.""" + fused_keys = [ + k + for k, v in state_dict.items() + if k.endswith((".experts.gate_up_proj", ".experts.down_proj")) and torch.is_tensor(v) and v.dim() == 3 + ] + if not fused_keys: + return state_dict + to_split = {} + for key in fused_keys: + parts = key.split(".") + target = module + resolved = True + for part in parts[:-1]: + if not hasattr(target, part): + resolved = False + break + target = getattr(target, part) + if resolved and hasattr(target, parts[-1]): + continue # original fused module still in the tree; assign as-is + to_split[key] = state_dict.pop(key) + if to_split: + from auto_round.utils.missing_tensors import split_fused_expert_tensors + + state_dict.update(split_fused_expert_tensors(to_split)) + return state_dict + + def _load_state_dict_into_module(state_dict: dict, module: torch.nn.Module) -> None: """Assign every key in *state_dict* to the corresponding sub-module. Handles cleared parameters (empty tensors) and wrapper objects that store the original layer in an ``orig_layer`` attribute. """ + state_dict = _maybe_split_fused_expert_keys(state_dict, module) for name, param in state_dict.items(): parts = name.split(".") target = module @@ -85,7 +128,14 @@ def _load_state_dict_into_module(state_dict: dict, module: torch.nn.Module) -> N if hasattr(target, param_name): old_param = getattr(target, param_name) if isinstance(old_param, torch.nn.Parameter): - param = param.to(dtype=old_param.dtype, device=old_param.device) + # Local addition (not upstream): `old_param` is on meta when the + # module started as a meta skeleton (AR_DISK_STREAM_MODEL=1) and + # is being materialized for the first time, rather than reloaded + # onto a previously-cleared real (cpu/cuda) tensor. Target "cpu" + # in that case instead of literally copying to meta (which would + # silently discard the just-read real data). See LOCAL_PATCHES.md. + target_device = "cpu" if old_param.device.type == "meta" else old_param.device + param = param.to(dtype=old_param.dtype, device=target_device) setattr(target, param_name, torch.nn.Parameter(param, requires_grad=old_param.requires_grad)) else: setattr(target, param_name, param) @@ -502,6 +552,38 @@ def _reload(self, model: torch.nn.Module, name: str) -> None: if module is None: return if self.mode == "offload": + if name not in self._saved: + # Local addition (not upstream): before falling back to the + # original checkpoint, check whether a *prior process* already + # offloaded this block to the deterministic resume directory + # (see _ensure_dir()) -- this is what makes bare in-memory + # .quantize() resumability actually work: a resumed process's + # self._saved starts empty (it's an in-memory dict, not + # persisted), so without this check it would always look like + # nothing was ever offloaded, even when a crashed prior + # process's real work is sitting right there on disk. See + # LOCAL_PATCHES.md. + from auto_round import envs + + if envs.AR_RESUME_DIR: + safe_name = name.replace(".", "_") + candidate_path = os.path.join(self._ensure_dir(), f"{safe_name}.safetensors") + if os.path.exists(candidate_path): + self._saved[name] = {"save_path": candidate_path} + self._load_from_disk(name, module) + if not self.retain_saved_entries: + self._remove_saved_entry(name) + return + # This block was never actually offloaded with real data + # (either it's still on meta, or `_save_to_disk` found nothing + # real to save), which happens when the model started as a + # meta skeleton (AR_DISK_STREAM_MODEL=1) instead of a full CPU + # load. There is nothing on the temp dir to load from -- read + # straight from the original checkpoint instead. See + # LOCAL_PATCHES.md. + if self.model_dir is not None: + load_block_from_model_files(self.model_dir, name, module) + return self._load_from_disk(name, module) if not self.retain_saved_entries: self._remove_saved_entry(name) @@ -721,6 +803,15 @@ def _save_to_disk(self, name: str, module: torch.nn.Module) -> None: for k, v in module.state_dict().items() if isinstance(v, torch.Tensor) and v.device.type != "meta" } + if not state_dict: + # Local addition (not upstream): nothing real to save -- the + # module was still on meta (e.g. the model started as a meta + # skeleton under AR_DISK_STREAM_MODEL=1 and this block hasn't + # been touched yet). Do NOT record it in self._saved: a later + # reload() must fall through to materializing straight from + # the original checkpoint (see _reload), not silently load an + # empty file and leave the block on meta. See LOCAL_PATCHES.md. + return safe_save_file(state_dict, save_path) self._saved[name] = {"save_path": save_path} del state_dict From a72847beb7112f51421c7f7d8f567869981f57be Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:38:38 +0200 Subject: [PATCH 05/25] offload: use a deterministic offload dir when AR_RESUME_DIR is set OffloadManager._ensure_dir() always used tempfile.mkdtemp() -- a fresh, uniquely-named directory every process, impossible for a resumed process to ever find again. Whenever AR_RESUME_DIR is set, use a stable path (/offload/_resume/) instead, so a resumed process's OffloadManager can find and reuse whatever a prior crashed process already offloaded there (see the companion discovery check in _reload(), previous commit). Signed-off-by: Fabrizio del Tin --- auto_round/utils/offload.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/auto_round/utils/offload.py b/auto_round/utils/offload.py index 32f075603..8307d74c7 100644 --- a/auto_round/utils/offload.py +++ b/auto_round/utils/offload.py @@ -785,7 +785,23 @@ def _ensure_dir(self) -> str: base_dir = os.path.join(envs.AR_WORK_SPACE, "offload") os.makedirs(base_dir, exist_ok=True) - self._tempdir = tempfile.mkdtemp(prefix=f"{self._prefix}_", dir=base_dir) + if envs.AR_RESUME_DIR: + # Local addition (not upstream): a fresh tempfile.mkdtemp() + # directory is unique to this process and can never be found + # again by a resumed run in a new process -- that's the whole + # reason bare in-memory .quantize() (no format=) resumability + # didn't actually work: ResumeState correctly skipped + # re-tuning already-done blocks, but their quantized weights, + # offloaded here, were unreachable from the resumed process, + # leaving those blocks on meta in the returned model. Use a + # stable, deterministic path instead whenever AR_RESUME_DIR is + # set, so a resumed process's OffloadManager can find (see + # _reload()'s discovery check below) and reuse what a prior + # crashed process already saved here. See LOCAL_PATCHES.md. + self._tempdir = os.path.join(base_dir, f"{self._prefix}_resume") + os.makedirs(self._tempdir, exist_ok=True) + else: + self._tempdir = tempfile.mkdtemp(prefix=f"{self._prefix}_", dir=base_dir) logger.info(f"OffloadManager ({self._prefix}): tempdir = {self._tempdir}") return self._tempdir From 011874a102b8b76c43fa4d03549d14b7a306cb8a Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:39:12 +0200 Subject: [PATCH 06/25] compressors/base: propagate model_dir to offloader; defer resume clear Two small additions supporting the disk-streaming/resumability work in adjacent commits: - After constructing self.model_context, if it was built as a disk- streamed meta skeleton, propagate its checkpoint path onto self._offloader.model_dir, so OffloadManager can materialize never-yet-offloaded blocks directly from disk (see the reload fix in utils/offload.py). - New self._resume_states, cleared by quantize_and_save() only after save_quantized() actually returns successfully -- not right after the tuning loop finishes, since a crash during the export/packing step that follows would otherwise wipe resumability for no reason. Populated by DataDrivenCompressor.quantize() in the next commit. Signed-off-by: Fabrizio del Tin --- auto_round/compressors/base.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/auto_round/compressors/base.py b/auto_round/compressors/base.py index f0060b54c..0efa0dd20 100644 --- a/auto_round/compressors/base.py +++ b/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 From 97696fe70f550bd8fd8af94f1f5eb12fba61bd41 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:42:09 +0200 Subject: [PATCH 07/25] compressors/data_driven: honor disk streaming outside resumability Three fixes needed for the tuning/RTN loops to work correctly when a model was built as a meta skeleton (AR_DISK_STREAM_MODEL=1), unrelated to resumability: - The standard tuning loop's per-block reload only fired when low_cpu_mem_usage was true. GGUF export forces low_cpu_mem_usage False for reasons of its own (unrelated to disk streaming), so a streamed block was never materialized before GGUF's tuning loop touched it. Now also reloads when AR_DISK_STREAM_MODEL is set, regardless of low_cpu_mem_usage. - configure_layer_config() disables low_cpu_mem_usage for any non-MoE-patched (dense) model on the assumption that the whole model is already CPU-resident, so per-block offload/reload buys nothing. False once the initial load is itself no longer full-residency: keep it enabled when self.model_context._disk_stream_index is not None. - CalibratedRTNCompressor (--iters 0 path)'s safe_to_cpu_() call tries to consolidate the whole model onto CPU, including decoder blocks intentionally still on meta -- crashing with "Cannot copy out of meta tensor". Skipped in both the normal and OOM-fallback branches when AR_DISK_STREAM_MODEL is set. Signed-off-by: Fabrizio del Tin --- auto_round/compressors/orchestrator.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/auto_round/compressors/orchestrator.py b/auto_round/compressors/orchestrator.py index 30bad9d2a..679b8e478 100644 --- a/auto_round/compressors/orchestrator.py +++ b/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: From a863e84726fc91e0c55c010aa9da8d3b428cdbb4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 14:32:18 +0000 Subject: [PATCH 08/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- auto_round/utils/disk_stream_util.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/auto_round/utils/disk_stream_util.py b/auto_round/utils/disk_stream_util.py index 3fecc26d2..064ed310f 100644 --- a/auto_round/utils/disk_stream_util.py +++ b/auto_round/utils/disk_stream_util.py @@ -79,9 +79,7 @@ def tensor_names_with_prefix(self, prefix: str) -> list[str]: return [n for n in self.weight_map if n == prefix or n.startswith(dotted)] -def materialize_module( - module: nn.Module, module_name: str, index: SafetensorsIndex, device: str -) -> None: +def materialize_module(module: nn.Module, module_name: str, index: SafetensorsIndex, device: str) -> None: """Populate `module`'s (currently meta) parameters/buffers with real data read directly from the checkpoint, onto `device`. `module_name` is `module`'s dotted path in the full model (used as the tensor-name prefix in the checkpoint). From 933cb6b3cbf5256cb51dec11a4d7ac8a74c616cd Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 12:36:30 +0200 Subject: [PATCH 09/25] Strip leaked internal-repo comments; fix meta-materialize dtype bug - Remove "Local addition"/LOCAL_PATCHES.md/vendor/reap references from every touched file -- local-only tooling metadata that doesn't apply upstream. - materialize_module() in disk_stream_util.py always forced the checkpoint's raw on-disk dtype onto rematerialized decoder blocks, fighting the compute dtype ModelContext._set_amp_dtype() had already promoted the meta skeleton to. This crashes with a dtype mismatch (e.g. BFloat16 vs Half) the moment a checkpoint's native dtype differs from the chosen amp dtype. Now prefers the meta parameter's already-declared (promoted) dtype, only falling back to the checkpoint's dtype for the one case that motivated the original behavior: an untyped meta context defaulting to float32. Signed-off-by: Fabrizio del Tin --- auto_round/context/model.py | 20 ++++----- auto_round/envs.py | 2 +- auto_round/utils/disk_stream_util.py | 61 +++++++++++++++------------- auto_round/utils/offload.py | 22 +++++----- 4 files changed, 52 insertions(+), 53 deletions(-) diff --git a/auto_round/context/model.py b/auto_round/context/model.py index 2dee5d01d..a8002d969 100644 --- a/auto_round/context/model.py +++ b/auto_round/context/model.py @@ -106,9 +106,9 @@ def __init__( self.need_calib = need_calib self.quant_nontext_module = quant_nontext_module - # Local addition (not upstream): remember the original string model name + # Remember the original string model name # so OffloadManager can later materialize blocks directly from this - # checkpoint (see disk_stream_model_dir below and LOCAL_PATCHES.md). + # checkpoint. self.disk_stream_model_dir = model if isinstance(model, str) else None self._disk_stream_index = None @@ -157,7 +157,7 @@ def _load_model(self): if is_mllm_model(self.model, platform=self.platform): self.is_mllm = True if isinstance(self.model, str): - # Local addition (not upstream): multimodal checkpoints used to + # 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 @@ -273,25 +273,23 @@ def _load_model(self): self._model_loaded = True def _build_disk_stream_model(self, model_name: str): - """Local addition (not upstream): build an all-meta skeleton instead of + """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``. - See LOCAL_PATCHES.md. """ 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 - # Local addition (not upstream): stash the index on the model object + # 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. See - # LOCAL_PATCHES.md. + # -- can still find it instead of re-scanning the checkpoint. model._disk_stream_index = index return model, tokenizer, index @@ -300,7 +298,7 @@ def _materialize_disk_stream_non_block_params(self) -> None: 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``. See LOCAL_PATCHES.md.""" + ``_build_disk_stream_model``.""" if self._disk_stream_index is None: return from auto_round.utils import flatten_list, get_block_names @@ -309,7 +307,7 @@ def _materialize_disk_stream_non_block_params(self) -> None: 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") - # Local addition (not upstream): tied output embeddings (e.g. lm_head.weight + # 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 @@ -318,7 +316,7 @@ def _materialize_disk_stream_non_block_params(self) -> None: # 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. See LOCAL_PATCHES.md. + # same (now real) Parameter object. if hasattr(self.model, "tie_weights"): self.model.tie_weights() diff --git a/auto_round/envs.py b/auto_round/envs.py index 4f4798b95..17444e71c 100644 --- a/auto_round/envs.py +++ b/auto_round/envs.py @@ -58,7 +58,7 @@ 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"), - # Local addition: device for the disk-streamed calibration forward pass in + # 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. diff --git a/auto_round/utils/disk_stream_util.py b/auto_round/utils/disk_stream_util.py index 064ed310f..5115b300f 100644 --- a/auto_round/utils/disk_stream_util.py +++ b/auto_round/utils/disk_stream_util.py @@ -1,16 +1,13 @@ # Copyright (C) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -# Local addition (not upstream AutoRound). Ported verbatim from this project's -# vendor/reap/src/reap/disk_stream_util.py -- architecture-agnostic (only needs -# module dotted names + accelerate's set_module_tensor_to_device), so it carries -# over unchanged. Used by auto_round/auto_scheme/delta_loss.py's streaming scoring -# path (get_score_for_scheme_streaming) to materialize one decoder block's real -# tensors from the checkpoint's safetensors shards right before scoring it and -# release them back to meta right after -- instead of AutoScheme's original -# approach of loading the entire model onto CPU via llm_load_model(device_map="cpu") -# (~207GB+ for a model like command-a-translate, not available on hardware with far -# less RAM+VRAM combined than the checkpoint size). See LOCAL_PATCHES.md. +# Lazy, mmap-backed reads of individual tensors by name straight from a +# checkpoint's safetensors shards, plus meta<->real materialize/free for a +# whole module. Used by auto_round/auto_scheme/delta_loss.py's streaming +# scoring path (get_score_for_scheme_streaming) to materialize one decoder +# block's real tensors right before scoring it and release them back to meta +# right after -- instead of loading the entire model onto CPU up front, which +# doesn't fit when the checkpoint is larger than available RAM+VRAM combined. from __future__ import annotations import json @@ -84,7 +81,7 @@ def materialize_module(module: nn.Module, module_name: str, index: SafetensorsIn directly from the checkpoint, onto `device`. `module_name` is `module`'s dotted path in the full model (used as the tensor-name prefix in the checkpoint). - Local addition (auto_round only, not in REAP's original): AutoScheme's scoring + AutoScheme's scoring wraps quantized layers in ``AutoSchemeWrapperLinear``, which replaces a plain ``nn.Linear`` with a wrapper holding the real layer at ``.orig_layer`` -- inserting an extra ``.orig_layer`` path segment that doesn't exist in the @@ -92,7 +89,7 @@ def materialize_module(module: nn.Module, module_name: str, index: SafetensorsIn """ import re as _re - # Local addition (not upstream): fused-MoE replacement modules + # Fused-MoE replacement modules # (SequentialQwen3_5MoeExperts and friends) expose UNFUSED per-expert # parameter names (experts.{i}.gate_proj.weight ...) that don't exist in a # checkpoint whose on-disk layout is the fused 3D one @@ -122,7 +119,7 @@ def _fused_lookup(full_name: str): inter = fused.shape[0] // 2 return (fused[:inter] if proj == "gate_proj" else fused[inter:]).contiguous() - targets = [] # (param_name, full_checkpoint_name) + targets = [] # (param_name, full_checkpoint_name, declared_meta_dtype) fused_targets = [] # (param_name, sliced_value) for name, tensor in list(module.named_parameters()) + list(module.named_buffers()): if str(tensor.device) != "meta": @@ -135,7 +132,7 @@ def _fused_lookup(full_name: str): continue logger.warning("No checkpoint tensor found for %s, leaving on meta", full_name) continue - targets.append((name, full_name)) + targets.append((name, full_name, tensor.dtype)) for name, value in fused_targets: set_module_tensor_to_device(module, name, device, value=value, dtype=value.dtype) @@ -143,20 +140,26 @@ def _fused_lookup(full_name: str): if not targets: return - values = index.read_tensors([full_name for _, full_name in targets], device=device) - for name, full_name in targets: - # Local addition (not upstream): explicit dtype= is required here. - # accelerate's set_module_tensor_to_device(), when dtype isn't passed, - # casts `value` to the *existing* (meta) parameter's declared dtype -- - # not the checkpoint's real dtype. A meta skeleton built without an - # enclosing dtype context (e.g. Qwen3_5MoeExperts' per-expert Linears, - # built under `torch.device("meta")` alone) defaults that declared - # dtype to float32 regardless of the checkpoint being bf16, so without - # this, real bf16 weights silently get upcast to float32 on - # materialization -- found via a real 8-layer MoE fixture crashing - # with "expected m1 and m2 to have the same dtype" inside AutoScheme - # scoring. The checkpoint's own dtype must always win. - set_module_tensor_to_device(module, name, device, value=values[full_name], dtype=values[full_name].dtype) + values = index.read_tensors([full_name for _, full_name, _ in targets], device=device) + for name, full_name, declared_dtype in targets: + # Prefer the meta parameter's already-declared dtype: it reflects + # whatever compute dtype the caller already promoted the (still-meta) + # model to (e.g. ModelContext._set_amp_dtype()'s `model.to(amp_dtype)`), + # and materializing to a different dtype than sibling non-block params + # that were promoted while still real breaks ops mixing the two (e.g. + # LayerNorm on bf16 activations with fp16 weight/bias). The one + # exception: a meta skeleton built without an enclosing dtype context + # (e.g. Qwen3_5MoeExperts' per-expert Linears, built under + # `torch.device("meta")` alone) defaults the declared dtype to + # float32 regardless of the checkpoint's real dtype -- found via a + # real 8-layer MoE fixture crashing with "expected m1 and m2 to have + # the same dtype" inside AutoScheme scoring. Detect that case (declared + # float32 but checkpoint isn't) and fall back to the checkpoint's own + # dtype instead. + target_dtype = declared_dtype + if declared_dtype == torch.float32 and values[full_name].dtype != torch.float32: + target_dtype = values[full_name].dtype + set_module_tensor_to_device(module, name, device, value=values[full_name], dtype=target_dtype) def free_module(module: nn.Module) -> None: @@ -185,7 +188,7 @@ def build_meta_model(model_name: str, trust_remote_code: bool = True): checkpoint on CPU RAM in one shot. Deliberately narrower than ``llm_load_model``: only covers the common local-directory ``AutoModelForCausalLM`` case (no bagel/glm/mxfp4/HPU special-casing) -- callers should fall back to - ``llm_load_model`` for anything this doesn't handle. See LOCAL_PATCHES.md. + ``llm_load_model`` for anything this doesn't handle. """ from accelerate import init_empty_weights from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer diff --git a/auto_round/utils/offload.py b/auto_round/utils/offload.py index 8307d74c7..bc3d72f35 100644 --- a/auto_round/utils/offload.py +++ b/auto_round/utils/offload.py @@ -66,7 +66,7 @@ def _maybe_split_fused_expert_keys(state_dict: dict, module: torch.nn.Module) -> dict: - """Local addition (not upstream): a checkpoint whose on-disk MoE layout is + """A checkpoint whose on-disk MoE layout is the fused 3D one (``...experts.gate_up_proj [N, 2*inter, hidden]`` / ``...experts.down_proj``, e.g. the real qwen3.5-397b-base) cannot be assigned onto a module tree whose experts were already replaced by the @@ -128,12 +128,12 @@ def _load_state_dict_into_module(state_dict: dict, module: torch.nn.Module) -> N if hasattr(target, param_name): old_param = getattr(target, param_name) if isinstance(old_param, torch.nn.Parameter): - # Local addition (not upstream): `old_param` is on meta when the + # `old_param` is on meta when the # module started as a meta skeleton (AR_DISK_STREAM_MODEL=1) and # is being materialized for the first time, rather than reloaded # onto a previously-cleared real (cpu/cuda) tensor. Target "cpu" # in that case instead of literally copying to meta (which would - # silently discard the just-read real data). See LOCAL_PATCHES.md. + # silently discard the just-read real data). target_device = "cpu" if old_param.device.type == "meta" else old_param.device param = param.to(dtype=old_param.dtype, device=target_device) setattr(target, param_name, torch.nn.Parameter(param, requires_grad=old_param.requires_grad)) @@ -553,7 +553,7 @@ def _reload(self, model: torch.nn.Module, name: str) -> None: return if self.mode == "offload": if name not in self._saved: - # Local addition (not upstream): before falling back to the + # Before falling back to the # original checkpoint, check whether a *prior process* already # offloaded this block to the deterministic resume directory # (see _ensure_dir()) -- this is what makes bare in-memory @@ -561,8 +561,7 @@ def _reload(self, model: torch.nn.Module, name: str) -> None: # self._saved starts empty (it's an in-memory dict, not # persisted), so without this check it would always look like # nothing was ever offloaded, even when a crashed prior - # process's real work is sitting right there on disk. See - # LOCAL_PATCHES.md. + # process's real work is sitting right there on disk. from auto_round import envs if envs.AR_RESUME_DIR: @@ -579,8 +578,7 @@ def _reload(self, model: torch.nn.Module, name: str) -> None: # real to save), which happens when the model started as a # meta skeleton (AR_DISK_STREAM_MODEL=1) instead of a full CPU # load. There is nothing on the temp dir to load from -- read - # straight from the original checkpoint instead. See - # LOCAL_PATCHES.md. + # straight from the original checkpoint instead. if self.model_dir is not None: load_block_from_model_files(self.model_dir, name, module) return @@ -786,7 +784,7 @@ def _ensure_dir(self) -> str: base_dir = os.path.join(envs.AR_WORK_SPACE, "offload") os.makedirs(base_dir, exist_ok=True) if envs.AR_RESUME_DIR: - # Local addition (not upstream): a fresh tempfile.mkdtemp() + # A fresh tempfile.mkdtemp() # directory is unique to this process and can never be found # again by a resumed run in a new process -- that's the whole # reason bare in-memory .quantize() (no format=) resumability @@ -797,7 +795,7 @@ def _ensure_dir(self) -> str: # stable, deterministic path instead whenever AR_RESUME_DIR is # set, so a resumed process's OffloadManager can find (see # _reload()'s discovery check below) and reuse what a prior - # crashed process already saved here. See LOCAL_PATCHES.md. + # crashed process already saved here. self._tempdir = os.path.join(base_dir, f"{self._prefix}_resume") os.makedirs(self._tempdir, exist_ok=True) else: @@ -820,13 +818,13 @@ def _save_to_disk(self, name: str, module: torch.nn.Module) -> None: if isinstance(v, torch.Tensor) and v.device.type != "meta" } if not state_dict: - # Local addition (not upstream): nothing real to save -- the + # Nothing real to save -- the # module was still on meta (e.g. the model started as a meta # skeleton under AR_DISK_STREAM_MODEL=1 and this block hasn't # been touched yet). Do NOT record it in self._saved: a later # reload() must fall through to materializing straight from # the original checkpoint (see _reload), not silently load an - # empty file and leave the block on meta. See LOCAL_PATCHES.md. + # empty file and leave the block on meta. return safe_save_file(state_dict, save_path) self._saved[name] = {"save_path": save_path} From 09450c8ec91577715fc0afddfbadb99f3d9997cc Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 19:38:50 +0200 Subject: [PATCH 10/25] Add unit tests for disk_stream_util.py primitives - Materialize/free round-trip for a real decoder block, and a re-materialize-is-a-no-op check for already-real (e.g. tied) params. - Regression coverage for the meta-materialize dtype bug: materialize_module must prefer the meta parameter's already-declared dtype (reflecting whatever compute dtype the caller promoted the model to) over the checkpoint's raw on-disk dtype, except when the declared dtype is an untyped-context float32 default. Verified this test fails against the pre-fix code with the exact reported "BFloat16 vs Half"-style mismatch. - build_meta_model + materialize_non_block_params: non-block params (embeddings) become real while decoder blocks stay meta. Signed-off-by: Fabrizio del Tin --- test/test_cpu/utils/test_disk_stream_util.py | 122 +++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 test/test_cpu/utils/test_disk_stream_util.py diff --git a/test/test_cpu/utils/test_disk_stream_util.py b/test/test_cpu/utils/test_disk_stream_util.py new file mode 100644 index 000000000..bfc9949f8 --- /dev/null +++ b/test/test_cpu/utils/test_disk_stream_util.py @@ -0,0 +1,122 @@ +# 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 auto_round.utils.disk_stream_util (AR_DISK_STREAM_MODEL primitives).""" + +import torch +import torch.nn as nn +from accelerate import init_empty_weights + +from auto_round.utils.disk_stream_util import ( + build_meta_model, + free_module, + materialize_module, + materialize_non_block_params, + total_resident_bytes, +) + + +class TestMaterializeModuleRoundTrip: + 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 + + def test_already_materialized_params_are_left_alone(self, tiny_opt_model_path): + """Shared/tied weights may already be real by the time materialize_module + runs on a second block referencing them; it must not touch (or crash on) + a parameter that's already off meta.""" + model, _tokenizer, index = build_meta_model(tiny_opt_model_path) + block = model.model.decoder.layers[0] + materialize_module(block, "model.decoder.layers.0", index, device="cpu") + weight_before = block.self_attn.k_proj.weight.data.clone() + + # Re-materializing an already-real block must be a no-op, not an error. + materialize_module(block, "model.decoder.layers.0", index, device="cpu") + assert torch.equal(block.self_attn.k_proj.weight.data, weight_before) + + +class TestMaterializeModuleDtype: + """Regression coverage for the dtype-promotion bug: materialize_module used + to always force the checkpoint's raw on-disk dtype, fighting whatever compute + dtype the caller had already promoted the (still-meta) model to.""" + + def test_prefers_declared_meta_dtype_over_checkpoint_dtype(self, tiny_opt_model_path): + model, _tokenizer, index = build_meta_model(tiny_opt_model_path) + block = model.model.decoder.layers[0] + + # Simulate ModelContext._set_amp_dtype()'s `model.to(amp_dtype)` promoting + # the still-meta block to bf16, even though the checkpoint itself is fp16 + # (facebook/opt-125m's native dtype). + block.to(torch.bfloat16) + for _, tensor in block.named_parameters(): + assert tensor.dtype == torch.bfloat16 + + materialize_module(block, "model.decoder.layers.0", index, device="cpu") + for name, tensor in block.named_parameters(): + assert tensor.dtype == torch.bfloat16, f"{name} was materialized as {tensor.dtype}, expected bfloat16" + + def test_falls_back_to_checkpoint_dtype_when_meta_declared_float32(self, tiny_opt_model_path): + """The one case the checkpoint's dtype must still win: a meta skeleton + built without an enclosing dtype context (e.g. a module built under + `torch.device("meta")` alone) defaults to float32 regardless of the + checkpoint's real dtype.""" + with init_empty_weights(): + linear = nn.Linear(64, 64, bias=False) + assert linear.weight.dtype == torch.float32 + + from auto_round.utils.disk_stream_util import SafetensorsIndex + + with torch.no_grad(): + real_weight = torch.randn(64, 64, dtype=torch.float16) + + class _FakeIndex(SafetensorsIndex): + def __init__(self): + self.weight_map = {"fc.weight": "dummy"} + + def has_tensor(self, name): + return name in self.weight_map + + def read_tensors(self, names, device="cpu"): + return {n: real_weight.clone() for n in names} + + materialize_module(linear, "fc", _FakeIndex(), device="cpu") + assert linear.weight.dtype == torch.float16 + + +class TestBuildMetaModel: + def test_non_block_params_are_real_blocks_stay_meta(self, tiny_opt_model_path): + model, tokenizer, index = build_meta_model(tiny_opt_model_path) + assert tokenizer is not None + + block_names = ["model.decoder.layers.0", "model.decoder.layers.1"] + materialize_non_block_params(model, block_names, index, device="cpu") + + # Embeddings (a non-block module) must be real now. + assert str(model.model.decoder.embed_tokens.weight.device) != "meta" + # Decoder blocks must still be untouched (meta). + for name, tensor in model.model.decoder.layers[0].named_parameters(): + assert str(tensor.device) == "meta", f"{name} was unexpectedly materialized" From de3894505be464d6dd6f6de9ed4b0411709a2f7b Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 20:36:30 +0200 Subject: [PATCH 11/25] Re-trigger CI (unrelated HF Hub 503 flake on test_gptoss) Signed-off-by: Fabrizio del Tin From b60e8646db2aba43af768dcdc4f3bcd222be7c6c Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Mon, 27 Jul 2026 11:23:00 +0200 Subject: [PATCH 12/25] Re-trigger CI (retest after flaky VRAM/RAM perf checks and llmc_format flake) Signed-off-by: Fabrizio del Tin From 22e0c2dcb7830be5197dcbb42e5c421d74622e4f Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 20:36:30 +0200 Subject: [PATCH 13/25] Re-trigger CI (unrelated HF Hub 503 flake on test_gptoss) Signed-off-by: Fabrizio del Tin From 06f1a0b921df55f811f580304ba1aad5264002e6 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Mon, 27 Jul 2026 11:23:00 +0200 Subject: [PATCH 14/25] Re-trigger CI (retest after flaky VRAM/RAM perf checks and llmc_format flake) Signed-off-by: Fabrizio del Tin From 949ff2ca0def7aa1930bb61f5bab2ab9424b3c62 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:36:55 +0200 Subject: [PATCH 15/25] Add AR_DISK_STREAM_MODEL and AR_RESUME_DIR env vars Two opt-in switches used by the disk-streaming and resumability work that follows in later commits. Default off (unset) preserves upstream behavior exactly. Signed-off-by: Fabrizio del Tin --- auto_round/envs.py | 32 -------------------------------- 1 file changed, 32 deletions(-) diff --git a/auto_round/envs.py b/auto_round/envs.py index 17444e71c..0e3366cbc 100644 --- a/auto_round/envs.py +++ b/auto_round/envs.py @@ -27,38 +27,6 @@ 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]: - """Read an optional env var that must be a positive integer when set.""" - raw = os.getenv(name) - if raw is None: - return None - try: - value = int(raw) - except ValueError as exc: - raise ValueError(f"{name} must be a positive integer, got {raw!r}") from exc - if value < 1: - raise ValueError(f"{name} must be a positive integer, got {value}") - return value - - -environment_variables: dict[str, Callable[[], Any]] = { - # this is used for configuring the default logging level - "AR_LOG_LEVEL": lambda: os.getenv("AR_LOG_LEVEL", "INFO").upper(), - "AR_ENABLE_COMPILE_PACKING": lambda: os.getenv("AR_ENABLE_COMPILE_PACKING", "0").lower() in ("1", "true", "yes"), - "AR_USE_MODELSCOPE": lambda: os.getenv("AR_USE_MODELSCOPE", "False").lower() in ["1", "true"], - "AR_WORK_SPACE": lambda: os.getenv("AR_WORK_SPACE", "ar_work_space").lower(), - "AR_ENABLE_UNIFY_MOE_INPUT_SCALE": lambda: os.getenv("AR_ENABLE_UNIFY_MOE_INPUT_SCALE", "False").lower() - in ["1", "true"], - "AR_OMP_NUM_THREADS": lambda: os.getenv("AR_OMP_NUM_THREADS", None), - "AR_DISABLE_OFFLOAD": lambda: os.getenv("AR_DISABLE_OFFLOAD", "0").lower() in ("1", "true", "yes"), - "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. From 0050b4dde0a4b6e41f0b3697f0a88a4a3e9a2a1a Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 20:36:30 +0200 Subject: [PATCH 16/25] Re-trigger CI (unrelated HF Hub 503 flake on test_gptoss) Signed-off-by: Fabrizio del Tin From 11d1d45d7544cbfdd6cccdcb3bcc774815b0b8f7 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:42:30 +0200 Subject: [PATCH 17/25] compressors/data_driven: resumable standard tuning loop (AR_RESUME_DIR) When AR_RESUME_DIR is set, DataDrivenCompressor.quantize() now builds one ResumeState per block group (auto_round/utils/resume.py, added in the next commit), keyed by a signature over model path + scheme + dataset + nsamples/seqlen + block list. On a partial resume, the group's first not-yet-done block substitutes its cached input_others from the pre-existing all_inputs cache, but the chained input_ids/ q_input come from the ResumeState's cached tensors, not that cache -- the pre-cache pass and the in-loop reference forward aren't numerically identical, so reusing the wrong one produced a 20x larger tuning loss on the first resumed block in testing. _quantize_blocks() starts its loop at resume_state.resume_index instead of 0 (nblocks=1 only), forces shard_writer._flush_shard() after each block when resuming is active (write() alone only buffers until the shard-size budget is hit -- a lie about durability that a real crash-and-resume test exposed as zero files on disk), and calls resume_state.mark_block_done(...) only after that write, so a crash before it correctly re-does the block rather than skipping it with incomplete output. Clearing the resume manifest is deferred to quantize_and_save() (see compressors/base.py, previous commit) rather than done right after the tuning loop, for the shard-export path specifically: quantize() returning successfully isn't the end of the pipeline there, and a crash during the packing/config-write step that follows would otherwise wipe resumability for no reason. The final "reload everything before returning" call now passes the full flattened block list explicitly when AR_RESUME_DIR is set (skipped entirely under shard-export/is_immediate_saving): reload(names=None) only reloads names already in the offloader's own _saved dict, which never includes a block a resumed process skipped entirely via ResumeState. Under shard export, reloading those blocks back to real memory is actively harmful, not just unnecessary -- the shard_writer's subsequent is_finalize=True write would re-emit their raw, unpacked weights alongside the already-correct packed ones already flushed by a prior process, producing duplicate/inconsistent tensors for the same layer (confirmed by diffing tensor names against an uninterrupted control run). Signed-off-by: Fabrizio del Tin --- auto_round/compressors/orchestrator.py | 222 +++++++++++++++++++++---- 1 file changed, 192 insertions(+), 30 deletions(-) diff --git a/auto_round/compressors/orchestrator.py b/auto_round/compressors/orchestrator.py index 679b8e478..f5f92b905 100644 --- a/auto_round/compressors/orchestrator.py +++ b/auto_round/compressors/orchestrator.py @@ -13,9 +13,10 @@ # limitations under the License. import copy import gc +import os import time from functools import partial -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union import accelerate import torch @@ -59,6 +60,9 @@ from auto_round.utils.device_manager import device_manager from auto_round.wrapper import WrapperMultiblock +if TYPE_CHECKING: + from auto_round.utils.resume import ResumeState + # TODO wenhuach align all the API args class CompressionOrchestrator(BaseOrchestrator): @@ -179,6 +183,8 @@ def _quantize_blocks( pbar: tqdm | None = None, input_others_extra_blocks: dict | None = None, token_ids: list[torch.Tensor] | None = None, + resume_state: Optional["ResumeState"] = None, + resume_input_ids=None, ): """Quantize and dequantize the weights of the specified blocks in the model. @@ -188,6 +194,20 @@ def _quantize_blocks( block_names: The names of the blocks to be quantized and dequantized. nblocks: The number of blocks to quantize and dequantize. device: The device for quantization and dequantization. + resume_state: when set and already partway through this block group + (`resume_state.resume_index > 0`), the caller has already + substituted `inputs`/`q_input` for the first not-yet-done block; + this method just needs to start its loop there instead of at + index 0, and record each block as done afterward. See + auto_round/utils/resume.py. + resume_input_ids: the exact `input_ids` value the interrupted run had + live for the first not-yet-done block (cached by + `ResumeState.mark_block_done`). `inputs` still supplies + `input_others` (legitimately re-sourced from the same pre-cache + every iteration regardless of resuming), but the chained + hidden-state tensor itself must come from here, not be re-derived + from `inputs` -- see auto_round/utils/resume.py's module + docstring for why those two aren't interchangeable. Returns: None @@ -197,11 +217,14 @@ def _quantize_blocks( m.requires_grad_(False) input_ids, input_others = self._preprocess_block_inputs(inputs) + if resume_input_ids is not None: + input_ids = resume_input_ids if pbar is None: pbar = tqdm(range(0, len(block_names), nblocks)) - for i in range(0, len(block_names), nblocks): + start_index = resume_state.resume_index if resume_state is not None and nblocks == 1 else 0 + for i in range(start_index, len(block_names), nblocks): if input_others_extra_blocks and block_names[i] in input_others_extra_blocks: input_others = input_others_extra_blocks[block_names[i]] _, input_others = self._preprocess_block_inputs(input_others) @@ -317,6 +340,17 @@ def _quantize_blocks( if self.compress_context.is_immediate_saving: self.shard_writer.write(m, is_finalize=False) + # ShardWriter only actually flushes to disk once its + # shard-size budget is reached (`_flush_shard`, private but + # there's no public equivalent) -- `write()` above may just + # buffer this block's tensors in memory. Force a flush here + # whenever resumability is active, since marking a block + # "done" in the resume manifest is a lie if a crash before + # the next natural flush would lose its tensors entirely. + # Only pay this extra small-shard-fragmentation cost when + # AR_RESUME_DIR is actually set. + if resume_state is not None: + self.shard_writer._flush_shard() if self.compress_context.low_cpu_mem_usage and not self.compress_context.is_immediate_saving: if nblocks == 1: @@ -324,6 +358,19 @@ def _quantize_blocks( else: for name in names: self._offloader(model, name, overwrite=True) + + # Record this block as durably done (its quantized weights are + # either flushed to a shard on disk via ShardWriter, or saved to + # the offloader's temp dir) only now, after that write has + # happened -- so a crash before this point correctly re-does the + # block on resume instead of skipping it with incomplete/missing + # output. See auto_round/utils/resume.py. + if resume_state is not None and nblocks == 1: + # `input_ids` was already reassigned to `next_input_ids` + # above -- it now holds the value the *next* block should use + # as its chained hidden-state input, which is exactly what + # needs to be persisted here. + resume_state.mark_block_done(n, q_input, input_ids) if pbar is not None: pbar.update(1) @@ -585,40 +632,155 @@ def _quantize_data_driven(self) -> tuple[torch.nn.Module, dict[str, Any]]: self.alg_composer.prepare_run() - for block_names in all_blocks: - inputs = all_inputs[block_names[0]] - all_inputs.pop(block_names[0]) - q_inputs = None - if all_q_inputs is not None: - q_inputs = all_q_inputs[block_names[0]] - all_q_inputs.pop(block_names[0]) - - inputs, q_inputs = _update_inputs(inputs, q_inputs) - - clear_memory(self.inputs) - - self._quantize_blocks( - self.model_context.model, - inputs, - block_names, - q_input=q_inputs if q_inputs is not None else None, - nblocks=self.nblocks, - pbar=pbar, - input_others_extra_blocks=all_inputs, - token_ids=input_ids_cache, + # Build one ResumeState per block group (almost always just one group + # for text-only dense models) when AR_RESUME_DIR is set, so a + # crash/kill mid-tuning can resume from the first not-yet-quantized + # block instead of restarting from block 0. See auto_round/utils/resume.py. + resume_states = None + if envs.AR_RESUME_DIR: + if not self.compress_context.is_immediate_saving and not self.compress_context.low_cpu_mem_usage: + logger.warning( + "AR_RESUME_DIR is set but neither immediate saving nor " + "low_cpu_mem_usage is active. Without low_cpu_mem_usage, " + "already-quantized blocks are never offloaded anywhere a " + "resumed process could find them (see OffloadManager's " + "deterministic resume directory in offload.py), so a " + "resumed run's in-memory model will have meta/empty " + "weights for blocks completed in a PRIOR process. Pass " + "low_cpu_mem_usage=True (or a format= to quantize_and_save) " + "for resumability to be meaningful here." + ) + from auto_round.utils.resume import ResumeState, compute_run_signature, layer_config_fingerprint + + model_dir = getattr(self.model_context, "disk_stream_model_dir", None) or getattr( + getattr(self.model_context.model, "config", None), "_name_or_path", None ) - if self.compress_context.is_immediate_packing and len(self.formats) != 1: - raise ValueError( - f"Expected exactly one packing format when 'immediate_packing' is True, " - f"but got {len(self.formats)} formats." + dataset_desc = str(getattr(self, "dataset", None)) + # str(self.scheme) alone is bits-blind for AutoScheme runs: two runs + # with different avg_bits share it, so include the resolved + # per-layer allocation (see layer_config_fingerprint docstring). + scheme_desc = ( + str(self.scheme) + + "|" + + layer_config_fingerprint(getattr(getattr(self, "quantizer", None), "layer_config", None)) + ) + resume_states = [] + for group_idx, block_names in enumerate(all_blocks): + sig = compute_run_signature( + model_dir, + scheme_desc, + dataset_desc, + self.calibration_context.nsamples, + self.calibration_context.seqlen, + block_names, + ) + resume_states.append( + ResumeState(os.path.join(envs.AR_RESUME_DIR, f"group_{group_idx}"), sig, block_names) ) - # ── Pipeline lifecycle: finalize_quantization (model-level teardown) - self.alg_composer.finalize_run() + try: + for group_idx, block_names in enumerate(all_blocks): + inputs = all_inputs[block_names[0]] + all_inputs.pop(block_names[0]) + q_inputs = None + if all_q_inputs is not None: + q_inputs = all_q_inputs[block_names[0]] + all_q_inputs.pop(block_names[0]) + + inputs, q_inputs = _update_inputs(inputs, q_inputs) + + clear_memory(self.inputs) + + resume_state = resume_states[group_idx] if resume_states is not None else None + resume_input_ids = None + if resume_state is not None and resume_state.resume_index > 0: + if self.nblocks != 1: + logger.warning( + "AR_RESUME_DIR is set but nblocks != 1; resuming mid-group is only " + "supported for nblocks=1 -- restarting this group from block 0." + ) + resume_state = None + else: + resume_name = block_names[resume_state.resume_index] + # Only used here for `input_others` (position/mask info, + # which is legitimately re-sourced from this same cache + # every iteration regardless of resuming); the actual + # chained `input_ids` comes from `resume_input_ids` + # below, not this cache -- see + # auto_round/utils/resume.py's module docstring for why + # the two aren't interchangeable. + if resume_name in all_inputs: + inputs = all_inputs.pop(resume_name) + q_inputs = resume_state.load_q_input() + resume_input_ids = resume_state.load_input_ids() + if resume_input_ids is None: + logger.warning( + "AR_RESUME_DIR manifest is missing its cached input_ids tensor; " + "restarting this group from block 0 instead of resuming with a " + "possibly-inconsistent chain value." + ) + resume_state = None + else: + pbar.update(resume_state.resume_index) + + self._quantize_blocks( + self.model_context.model, + inputs, + block_names, + q_input=q_inputs if q_inputs is not None else None, + nblocks=self.nblocks, + pbar=pbar, + input_others_extra_blocks=all_inputs, + token_ids=input_ids_cache, + resume_state=resume_state, + resume_input_ids=resume_input_ids, + ) + if self.compress_context.is_immediate_packing and len(self.formats) != 1: + raise ValueError( + f"Expected exactly one packing format when 'immediate_packing' is True, " + f"but got {len(self.formats)} formats." + ) + if resume_states is not None: + if self.compress_context.is_immediate_saving: + # Don't clear resume state yet when exporting to shards -- + # a crash in the save/export step that follows this method + # returning (config writing, tokenizer copy, format-specific + # global packing pass) would otherwise force a full + # re-tune from block 0 on the next attempt, even though + # every block's weights are already correctly flushed to + # disk. quantize_and_save() clears these once + # save_quantized() actually succeeds. + self._resume_states = resume_states + else: + for rs in resume_states: + rs.clear() + finally: + # ── Pipeline lifecycle: finalize_quantization (model-level teardown) + self.alg_composer.finalize_run() pbar.set_description("Quantizing done") pbar.close() if self.compress_context.low_cpu_mem_usage: - self._offloader.reload(self.model_context.model) + if envs.AR_RESUME_DIR and not self.compress_context.is_immediate_saving: + # `reload(names=None)` only reloads names in + # `self._offloader._saved` -- populated by THIS process's own + # offload() calls. A resumed process never touches blocks it + # skipped via ResumeState (they're left exactly as the meta + # skeleton started), so they'd never be in `_saved` and would + # stay meta in the returned model. Request every block + # explicitly so _reload()'s discovery check (see offload.py) + # gets a chance to pull each skipped block's real quantized + # weights back from a prior crashed process's offload dir. + # + # Skipped entirely under is_immediate_saving: ShardWriter has + # already flushed every block's packed weights to disk (both + # this run's and, via its own discovery, a prior crashed run's), + # and the `shard_writer.write(is_finalize=True)` call right + # below would treat any block reloaded back to real memory here + # as newly-dirty and re-emit its raw, unpacked weight tensor + # alongside the already-packed one. + self._offloader.reload(self.model_context.model, flatten_list(all_blocks)) + elif not self.compress_context.is_immediate_saving: + self._offloader.reload(self.model_context.model) self._quantize_layers_outside_blocks(layer_names, all_inputs, token_ids=input_ids_cache) convert_module_to_hp_if_necessary( From d9cc0fce9a877179678362045b1c4516b3cec0ce Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:45:31 +0200 Subject: [PATCH 18/25] Add ResumeState: crash/resume checkpointing for the tuning loop New auto_round/utils/resume.py, no upstream equivalent. Tracks completed blocks plus cached chain tensors (resume_q_input.pt/ resume_input_ids.pt) for a tuning run, keyed by a signature hash over model path + scheme + dataset + nsamples/seqlen + block list, so a resume directory reused for a different run is detected and ignored rather than silently misapplied. Both chain tensors (q_input and input_ids) are cached, not just q_input: the FP reference chain (input_ids) is not numerically identical between AutoRound's pre-tuning cache pass and the in-loop reference forward, so reconstructing it from the pre-cache instead of persisting the live value produced a 20x larger tuning loss on the first resumed block in testing. Also adds layer_config_fingerprint(), folded into the run signature by this file's callers (auto_round/compressors/data_driven.py, previous commits): str(self.scheme) (or the literal "rtn_with_imatrix") alone is bits-blind for AutoScheme runs -- two runs against the same model/dataset/nsamples/seqlen but different avg_bits targets produced identical signatures, so the second run silently resumed the first's already-complete manifest and saved an output containing no layer tensors at all. Folding the resolved per-layer bit allocation into the signature fixes this. Signed-off-by: Fabrizio del Tin --- auto_round/utils/resume.py | 222 +++++++++++++++++++++++++++++++++++++ 1 file changed, 222 insertions(+) create mode 100644 auto_round/utils/resume.py diff --git a/auto_round/utils/resume.py b/auto_round/utils/resume.py new file mode 100644 index 000000000..ec4661e6f --- /dev/null +++ b/auto_round/utils/resume.py @@ -0,0 +1,222 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +# Local addition (not upstream). Checkpoints the tuning loop's per-block +# progress to disk so a crash or kill mid-run doesn't require restarting from +# block 0 -- important for large models where each block's tuning can take +# minutes and a full run spans hours. See LOCAL_PATCHES.md. +# +# What this caches and why: AutoRound's block-sequential tuning chains two +# things forward from one block to the next -- ``input_ids`` (the current +# block's reference/FP output, used as the next block's FP reference input) +# and, when ``enable_quanted_input=True`` (SignRound's default), ``q_input`` +# (the *quantized* block's output, used as the next block's quantized-input +# companion). Both are cached here, not just ``q_input``: an earlier version of +# this patch assumed ``input_ids`` could be cheaply regenerated from the +# per-block inputs already pre-cached by `cache_inter_data`/ +# `try_cache_inter_data_gpucpu` before tuning starts (reasoning that it's a +# pure function of unmodified weights, so it shouldn't matter which code path +# computed it) -- but that pre-cache pass and the in-loop reference forward +# turned out not to be numerically identical (confirmed by a resume test +# producing a 20x-larger tuning loss on the first resumed block vs. an +# uninterrupted control run), so the actual live chain value has to be +# persisted and reloaded verbatim, the same way ``q_input`` already is. +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from pathlib import Path +from typing import Optional + +import torch + +from auto_round.logger import logger + +__all__ = ["ResumeState"] + +_MANIFEST_NAME = "resume_manifest.json" +_Q_INPUT_NAME = "resume_q_input.pt" +_INPUT_IDS_NAME = "resume_input_ids.pt" + + +def _to_cpu_recursive(obj): + """`q_input` (SignRound's `enable_quanted_input` chain value) may be a + plain Tensor or a nested list/tuple/dict of tensors depending on the + algorithm/batching path; detach + move to cpu recursively so it's both + safely picklable and reloadable regardless of which GPU produced it.""" + if isinstance(obj, torch.Tensor): + return obj.detach().to("cpu") + if isinstance(obj, dict): + return {k: _to_cpu_recursive(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_to_cpu_recursive(v) for v in obj] + if isinstance(obj, tuple): + return tuple(_to_cpu_recursive(v) for v in obj) + return obj + + +def _atomic_write_json(path: Path, data: dict) -> None: + fd, tmp_path = tempfile.mkstemp(dir=str(path.parent), prefix=".tmp_resume_") + try: + with os.fdopen(fd, "w") as f: + json.dump(data, f) + os.replace(tmp_path, path) + except Exception: + try: + os.remove(tmp_path) + except OSError: + pass + raise + + +class ResumeState: + """Tracks which blocks of a tuning run have already been quantized and + written to disk, plus the one small tensor (``q_input``) needed to resume + the sequential calibration chain correctly. Keyed by a signature over the + run's identifying configuration, so a resume directory reused for a + different model/scheme/dataset is detected and ignored rather than + silently misapplied. + """ + + def __init__(self, resume_dir: str, signature: str, block_names: list[str]): + self.dir = Path(resume_dir) + self.dir.mkdir(parents=True, exist_ok=True) + self.signature = signature + self.block_names = list(block_names) + self.manifest_path = self.dir / _MANIFEST_NAME + self.q_input_path = self.dir / _Q_INPUT_NAME + self.input_ids_path = self.dir / _INPUT_IDS_NAME + self.completed_blocks: list[str] = [] + self._load() + + def _load(self) -> None: + if not self.manifest_path.exists(): + return + try: + with open(self.manifest_path) as f: + data = json.load(f) + except Exception as e: + logger.warning(f"ResumeState: failed to read {self.manifest_path}: {e}; starting fresh.") + return + if data.get("signature") != self.signature: + logger.info( + "ResumeState: existing resume manifest is for a different run " + "(model/scheme/dataset/block list changed); ignoring it and starting fresh." + ) + return + completed = data.get("completed_blocks", []) + # Only trust a prefix of block_names in order -- anything else means the + # manifest was corrupted or hand-edited; safer to restart from scratch + # than resume from a possibly-inconsistent point. + if completed != self.block_names[: len(completed)]: + logger.warning( + "ResumeState: completed_blocks in manifest is not a prefix of the " + "current block order; ignoring it and starting fresh." + ) + return + self.completed_blocks = completed + logger.info(f"ResumeState: resuming after {len(completed)}/{len(self.block_names)} already-quantized blocks") + + @property + def resume_index(self) -> int: + """Index of the first block not yet completed.""" + return len(self.completed_blocks) + + def load_q_input(self): + return self._load_tensor(self.q_input_path) + + def load_input_ids(self): + return self._load_tensor(self.input_ids_path) + + def _load_tensor(self, path: Path): + if not self.completed_blocks or not path.exists(): + return None + try: + return torch.load(path, map_location="cpu") + except Exception as e: + logger.warning(f"ResumeState: failed to load {path.name} ({e}); resuming without it.") + return None + + def mark_block_done(self, block_name: str, q_input, input_ids) -> None: + expected = self.block_names[len(self.completed_blocks)] + assert block_name == expected, ( + f"ResumeState.mark_block_done called out of order: expected {expected!r}, got {block_name!r}" + ) + if q_input is not None: + torch.save(_to_cpu_recursive(q_input), self.q_input_path) + elif self.q_input_path.exists(): + self.q_input_path.unlink() + # input_ids is required (unlike q_input, which is legitimately None + # when enable_quanted_input=False) -- the FP reference chain always + # exists. + torch.save(_to_cpu_recursive(input_ids), self.input_ids_path) + self.completed_blocks.append(block_name) + _atomic_write_json( + self.manifest_path, + {"signature": self.signature, "completed_blocks": self.completed_blocks}, + ) + + def clear(self) -> None: + """Remove the resume manifest/cache -- call after a full run completes + successfully, so a later unrelated run doesn't mistake stale state for + an in-progress one (only possible if reusing the exact same + model/scheme/dataset/block list, but still worth cleaning up).""" + for p in (self.manifest_path, self.q_input_path, self.input_ids_path): + if p.exists(): + try: + p.unlink() + except OSError: + pass + + +def compute_run_signature( + model_dir: Optional[str], + scheme_desc: str, + dataset_desc: str, + nsamples: int, + seqlen: int, + block_names: list[str], +) -> str: + """Hash the run's identifying configuration. Any change here (different + model, scheme, dataset, calibration size, or block set) must produce a + different signature so `ResumeState` refuses to reuse a stale manifest.""" + h = hashlib.sha256() + for part in (model_dir or "", scheme_desc, dataset_desc, str(nsamples), str(seqlen), "|".join(block_names)): + h.update(part.encode("utf-8")) + h.update(b"\x00") + return h.hexdigest() + + +def layer_config_fingerprint(layer_config) -> str: + """Deterministic string over the resolved per-layer quantization config. + + ``str(self.scheme)`` (and the literal ``"rtn_with_imatrix"`` used by the + imatrix-RTN path) does not capture the per-layer bit allocation that + AutoScheme resolves from ``avg_bits`` -- two runs with different avg_bits + targets would otherwise produce identical run signatures and silently + resume each other's state (observed: a quality sweep's avg_bits=4.8 + candidate resuming the 4.2 candidate's 40/40-complete manifest, saving an + output with no layer tensors at all). Fold this into ``scheme_desc`` when + calling :func:`compute_run_signature`. + + Only scalar config values (bits, group_size, sym, data_type, ...) are + included; anything non-scalar is ignored so the fingerprint stays cheap + and deterministic. + """ + if not layer_config: + return "" + parts = [] + for name in sorted(layer_config): + cfg = layer_config[name] + if isinstance(cfg, dict): + desc = ",".join( + f"{k}={v}" + for k, v in sorted(cfg.items(), key=lambda kv: str(kv[0])) + if isinstance(v, (int, float, bool, str, type(None))) + ) + else: + desc = str(cfg) + parts.append(f"{name}:{desc}") + return ";".join(parts) From f595ce1adb6d0969b0b6c58f75c5f80963dbfcb8 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:45:51 +0200 Subject: [PATCH 19/25] shard_writer: discover pre-existing shards to resume across processes A fresh process's ShardWriter has no memory of shards a previous, crashed process already flushed to output_dir -- it would restart shard_counter at 0, collide with existing shard filenames, and finalize()'s index would only cover this process's tensors, producing a corrupt/incomplete checkpoint. Adds _discover_existing_shards(): when AR_RESUME_DIR is set, on the first real _flush_shard() call (not __init__ -- see below), scans output_dir for leftover pre-rename model-shard-NNNNN. files, reads each one's tensor names straight from its safetensors/torch header (no data materialization needed), and seeds shard_counter/shard_meta/ _all_saved from them so numbering doesn't collide and finalize()'s index covers both processes' shards. Discovery has to be deferred past __init__: ShardWriter.__init__ runs during post_init(), before quantize_and_save()'s _get_export_dir() appends the final subfolder (e.g. -w4g128/) to output_dir -- discovering at construction time silently looked in the wrong directory and found nothing, confirmed by a real crash-and-resume test where blocks 0-2 resumed correctly through tuning but still lost their output. Fixed by running discovery lazily inside _flush_shard() itself, guarded by a self._existing_shards_discovered flag, by which point output_dir is always the final path. Gated on AR_RESUME_DIR throughout, so normal non-resuming runs never change behavior even if output_dir happens to be reused. Signed-off-by: Fabrizio del Tin --- auto_round/compressors/shard_writer.py | 75 ++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/auto_round/compressors/shard_writer.py b/auto_round/compressors/shard_writer.py index e95468fbe..41447c077 100644 --- a/auto_round/compressors/shard_writer.py +++ b/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 = [] + # Local addition (not upstream): 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. `-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.``) 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) From 3140a5f4627daf7d6c40b79fe7fcb42b0e4087b8 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:46:06 +0200 Subject: [PATCH 20/25] export: skip pack_layer() for a resumed, still-meta layer A resumed disk-streamed run only materializes/quantizes the blocks it didn't already finish in a prior (crashed) process; blocks it skipped are untouched 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, earlier commit). The global post-tuning packing pass otherwise crashed trying to read .scale off such a layer. Early-return when the layer's weight is still on meta: there is nothing to pack, and the on-disk export for it is already complete. Signed-off-by: Fabrizio del Tin --- auto_round/export/export_to_autoround/export.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/auto_round/export/export_to_autoround/export.py b/auto_round/export/export_to_autoround/export.py index 32f688ab1..a4e9bc9a7 100644 --- a/auto_round/export/export_to_autoround/export.py +++ b/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 + # Local addition (not upstream): 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) From e0ca142aaf810736fda4274e4b2463a7a799ed38 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 14:48:19 +0000 Subject: [PATCH 21/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- auto_round/utils/resume.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/auto_round/utils/resume.py b/auto_round/utils/resume.py index ec4661e6f..93c3bb853 100644 --- a/auto_round/utils/resume.py +++ b/auto_round/utils/resume.py @@ -141,9 +141,9 @@ def _load_tensor(self, path: Path): def mark_block_done(self, block_name: str, q_input, input_ids) -> None: expected = self.block_names[len(self.completed_blocks)] - assert block_name == expected, ( - f"ResumeState.mark_block_done called out of order: expected {expected!r}, got {block_name!r}" - ) + assert ( + block_name == expected + ), f"ResumeState.mark_block_done called out of order: expected {expected!r}, got {block_name!r}" if q_input is not None: torch.save(_to_cpu_recursive(q_input), self.q_input_path) elif self.q_input_path.exists(): From d7f0e36ce28f7a98e012713564993583c8920be3 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 14:46:53 +0200 Subject: [PATCH 22/25] Strip leaked internal-repo comment markers Remove "Local addition"/LOCAL_PATCHES.md references from every touched file -- local-only tooling metadata that doesn't apply upstream. No logic change. Signed-off-by: Fabrizio del Tin --- auto_round/compressors/shard_writer.py | 2 +- auto_round/export/export_to_autoround/export.py | 2 +- auto_round/utils/resume.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/auto_round/compressors/shard_writer.py b/auto_round/compressors/shard_writer.py index 41447c077..1ca0395c7 100644 --- a/auto_round/compressors/shard_writer.py +++ b/auto_round/compressors/shard_writer.py @@ -97,7 +97,7 @@ def __init__( self.total_param_size_bytes = 0 self.skipped_meta_tensors = [] - # Local addition (not upstream): when resumability is active + # 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 diff --git a/auto_round/export/export_to_autoround/export.py b/auto_round/export/export_to_autoround/export.py index a4e9bc9a7..4b9c9e30e 100644 --- a/auto_round/export/export_to_autoround/export.py +++ b/auto_round/export/export_to_autoround/export.py @@ -164,7 +164,7 @@ def pack_layer(layer_name, model, backend, device=None): if type(layer) not in SUPPORTED_LAYER_TYPES: ##already packed return - # Local addition (not upstream): a resumed disk-streamed run only + # 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 diff --git a/auto_round/utils/resume.py b/auto_round/utils/resume.py index 93c3bb853..9299cbb6c 100644 --- a/auto_round/utils/resume.py +++ b/auto_round/utils/resume.py @@ -1,10 +1,10 @@ # Copyright (C) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -# Local addition (not upstream). Checkpoints the tuning loop's per-block +# Checkpoints the tuning loop's per-block # progress to disk so a crash or kill mid-run doesn't require restarting from # block 0 -- important for large models where each block's tuning can take -# minutes and a full run spans hours. See LOCAL_PATCHES.md. +# minutes and a full run spans hours. # # What this caches and why: AutoRound's block-sequential tuning chains two # things forward from one block to the next -- ``input_ids`` (the current From 995afeaf0c339429808221d91ec1f2f80acd27b8 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 15:34:15 +0200 Subject: [PATCH 23/25] Add unit and integration tests for AR_RESUME_DIR - test/test_cpu/utils/test_resume.py: unit tests for ResumeState (mark_block_done ordering, q_input/input_ids round-trip, signature mismatch and non-prefix manifest handling both correctly discard stale state, clear()), compute_run_signature, and layer_config_fingerprint. - test/test_cpu/core/test_resume_integration.py: end-to-end test that simulates a crash after the first block (injected via a ResumeState.mark_block_done wrapper that raises right after persisting state) and verifies a fresh AutoRound run against the same AR_RESUME_DIR resumes from the second block only, producing a complete layer_config for both blocks. Signed-off-by: Fabrizio del Tin --- test/test_cpu/core/test_resume_integration.py | 89 ++++++++++ test/test_cpu/utils/test_resume.py | 166 ++++++++++++++++++ 2 files changed, 255 insertions(+) create mode 100644 test/test_cpu/core/test_resume_integration.py create mode 100644 test/test_cpu/utils/test_resume.py diff --git a/test/test_cpu/core/test_resume_integration.py b/test/test_cpu/core/test_resume_integration.py new file mode 100644 index 000000000..6396a65aa --- /dev/null +++ b/test/test_cpu/core/test_resume_integration.py @@ -0,0 +1,89 @@ +# 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 AR_RESUME_DIR: simulates a crash partway through the +tuning loop and verifies a fresh AutoRound run against the same resume +directory picks up from the first not-yet-completed block instead of +restarting from block 0. +""" + +import os +from unittest import mock + +import pytest + +from auto_round import AutoRound +from auto_round.utils.resume import ResumeState + + +@pytest.fixture(autouse=True) +def _clean_resume_env(): + previous_resume_dir = os.environ.get("AR_RESUME_DIR") + previous_disk_stream = os.environ.get("AR_DISK_STREAM_MODEL") + yield + for key, previous in (("AR_RESUME_DIR", previous_resume_dir), ("AR_DISK_STREAM_MODEL", previous_disk_stream)): + if previous is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous + + +class TestResumeIntegration: + def test_resume_skips_already_completed_blocks(self, tiny_opt_model_path, tmp_path): + resume_dir = str(tmp_path / "resume") + # AR_DISK_STREAM_MODEL keeps low_cpu_mem_usage active for this plain + # dense (non-MoE) model -- without it, DataDrivenCompressor.quantize() + # silently disables low_cpu_mem_usage, and a resumed process's + # in-memory model would have meta/empty weights for blocks completed + # by the prior (crashed) process. + os.environ["AR_RESUME_DIR"] = resume_dir + os.environ["AR_DISK_STREAM_MODEL"] = "1" + + original_mark_block_done = ResumeState.mark_block_done + crashed_after = [] + + def crash_after_first_block(self, block_name, q_input, input_ids): + original_mark_block_done(self, block_name, q_input, input_ids) + crashed_after.append(block_name) + if len(crashed_after) == 1: + raise RuntimeError("simulated crash") + + with mock.patch.object(ResumeState, "mark_block_done", crash_after_first_block): + with pytest.raises(RuntimeError, match="simulated crash"): + ar = AutoRound(model=tiny_opt_model_path, scheme="W4A16", iters=1, nsamples=1) + ar.quantize() + + assert crashed_after == ["model.decoder.layers.0"] + + processed_on_resume = [] + + def track_block(self, block_name, q_input, input_ids): + original_mark_block_done(self, block_name, q_input, input_ids) + processed_on_resume.append(block_name) + + with mock.patch.object(ResumeState, "mark_block_done", track_block): + ar = AutoRound(model=tiny_opt_model_path, scheme="W4A16", iters=1, nsamples=1) + _, layer_config = ar.quantize() + + # Only the block that never completed before the crash should be + # (re-)processed -- block 0's already-durable result must be reused, + # not redone. + assert processed_on_resume == ["model.decoder.layers.1"] + + # The final layer_config still covers every quantized layer in both + # blocks, proving the resumed run produced a complete result. + quantized_layers = {name for name, cfg in layer_config.items() if "bits" in cfg} + assert any(name.startswith("model.decoder.layers.0.") for name in quantized_layers) + assert any(name.startswith("model.decoder.layers.1.") for name in quantized_layers) diff --git a/test/test_cpu/utils/test_resume.py b/test/test_cpu/utils/test_resume.py new file mode 100644 index 000000000..67e73054c --- /dev/null +++ b/test/test_cpu/utils/test_resume.py @@ -0,0 +1,166 @@ +# 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 auto_round.utils.resume (crash/resume checkpointing, AR_RESUME_DIR).""" + +import pytest +import torch + +from auto_round.utils.resume import ( + ResumeState, + compute_run_signature, + layer_config_fingerprint, +) + + +class TestResumeState: + def test_fresh_state_has_no_completed_blocks(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + assert state.resume_index == 0 + assert state.load_q_input() is None + assert state.load_input_ids() is None + + def test_mark_block_done_advances_resume_index(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + state.mark_block_done("b0", q_input=torch.ones(2), input_ids=torch.zeros(2)) + assert state.resume_index == 1 + assert state.completed_blocks == ["b0"] + + def test_mark_block_done_out_of_order_raises(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + with pytest.raises(AssertionError): + state.mark_block_done("b1", q_input=None, input_ids=torch.zeros(2)) + + def test_q_input_and_input_ids_round_trip(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + q_input = {"a": torch.arange(4).reshape(2, 2)} + input_ids = torch.arange(6).reshape(2, 3) + state.mark_block_done("b0", q_input=q_input, input_ids=input_ids) + + # A fresh ResumeState pointed at the same dir with the same signature + # picks the saved tensors back up, exactly like a resumed process would. + resumed = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + assert resumed.resume_index == 1 + loaded_q_input = resumed.load_q_input() + assert torch.equal(loaded_q_input["a"], q_input["a"]) + assert torch.equal(resumed.load_input_ids(), input_ids) + + def test_q_input_none_clears_any_previous_q_input_file(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + state.mark_block_done("b0", q_input=torch.ones(2), input_ids=torch.zeros(2)) + assert state.q_input_path.exists() + + # Simulate a second run reusing the dir where enable_quanted_input is + # off (q_input legitimately None) -- the stale file must not linger + # and be mistaken for this block's q_input on a later resume. + state2 = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + state2.completed_blocks = [] + state2.mark_block_done("b0", q_input=None, input_ids=torch.zeros(2)) + assert not state2.q_input_path.exists() + + def test_mismatched_signature_starts_fresh(self, tmp_path): + state = ResumeState(str(tmp_path), "sig-a", ["b0", "b1"]) + state.mark_block_done("b0", q_input=None, input_ids=torch.zeros(2)) + + other = ResumeState(str(tmp_path), "sig-b", ["b0", "b1"]) + assert other.resume_index == 0 + assert other.completed_blocks == [] + + def test_completed_blocks_not_a_prefix_starts_fresh(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + state.mark_block_done("b0", q_input=None, input_ids=torch.zeros(2)) + + # Different block order for the same signature (e.g. block list changed + # some other way the signature didn't capture) -- "b0" is no longer a + # valid prefix of this new order, so the manifest must be distrusted. + other = ResumeState(str(tmp_path), "sig", ["b1", "b0", "b2"]) + assert other.resume_index == 0 + assert other.completed_blocks == [] + + def test_clear_removes_all_state(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + state.mark_block_done("b0", q_input=torch.ones(2), input_ids=torch.zeros(2)) + assert state.manifest_path.exists() + + state.clear() + assert not state.manifest_path.exists() + assert not state.q_input_path.exists() + assert not state.input_ids_path.exists() + + fresh = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + assert fresh.resume_index == 0 + + def test_full_completion_round_trip(self, tmp_path): + block_names = ["b0", "b1", "b2"] + state = ResumeState(str(tmp_path), "sig", block_names) + for name in block_names: + state.mark_block_done(name, q_input=None, input_ids=torch.zeros(2)) + assert state.resume_index == len(block_names) + assert state.completed_blocks == block_names + + +class TestComputeRunSignature: + def test_identical_inputs_produce_identical_signature(self): + sig1 = compute_run_signature("m", "scheme", "dataset", 8, 2048, ["b0", "b1"]) + sig2 = compute_run_signature("m", "scheme", "dataset", 8, 2048, ["b0", "b1"]) + assert sig1 == sig2 + + @pytest.mark.parametrize( + "kwargs", + [ + {"model_dir": "other"}, + {"scheme_desc": "other"}, + {"dataset_desc": "other"}, + {"nsamples": 16}, + {"seqlen": 4096}, + {"block_names": ["b0", "b1", "b2"]}, + ], + ) + def test_any_changed_field_changes_signature(self, kwargs): + base = dict( + model_dir="m", + scheme_desc="scheme", + dataset_desc="dataset", + nsamples=8, + seqlen=2048, + block_names=["b0", "b1"], + ) + sig1 = compute_run_signature(**base) + sig2 = compute_run_signature(**{**base, **kwargs}) + assert sig1 != sig2 + + def test_none_model_dir_does_not_crash(self): + sig = compute_run_signature(None, "scheme", "dataset", 8, 2048, ["b0"]) + assert isinstance(sig, str) and len(sig) == 64 # sha256 hexdigest + + +class TestLayerConfigFingerprint: + def test_empty_layer_config(self): + assert layer_config_fingerprint(None) == "" + assert layer_config_fingerprint({}) == "" + + def test_deterministic_regardless_of_dict_order(self): + cfg_a = {"layer.1": {"bits": 4, "sym": True}, "layer.0": {"bits": 8, "sym": False}} + cfg_b = {"layer.0": {"bits": 8, "sym": False}, "layer.1": {"bits": 4, "sym": True}} + assert layer_config_fingerprint(cfg_a) == layer_config_fingerprint(cfg_b) + + def test_different_bits_produce_different_fingerprint(self): + cfg_4bit = {"layer.0": {"bits": 4}} + cfg_8bit = {"layer.0": {"bits": 8}} + assert layer_config_fingerprint(cfg_4bit) != layer_config_fingerprint(cfg_8bit) + + def test_non_scalar_values_are_ignored(self): + cfg_with_tensor = {"layer.0": {"bits": 4, "weight": torch.ones(2, 2)}} + cfg_without_tensor = {"layer.0": {"bits": 4}} + assert layer_config_fingerprint(cfg_with_tensor) == layer_config_fingerprint(cfg_without_tensor) From ead1edac20fbe30837a3500495cb5ff22b6ed551 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Mon, 27 Jul 2026 11:23:16 +0200 Subject: [PATCH 24/25] Re-trigger CI (retest after flaky VRAM/RAM perf checks and llmc_format flake) Signed-off-by: Fabrizio del Tin From 2092f07ce88361089c283719dc5a27f936ebb895 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 29 Jul 2026 08:29:48 +0200 Subject: [PATCH 25/25] Fix envs.py syntax corrupted during rebase conflict resolution Signed-off-by: Fabrizio del Tin --- auto_round/envs.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/auto_round/envs.py b/auto_round/envs.py index 0e3366cbc..17444e71c 100644 --- a/auto_round/envs.py +++ b/auto_round/envs.py @@ -27,6 +27,38 @@ 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]: + """Read an optional env var that must be a positive integer when set.""" + raw = os.getenv(name) + if raw is None: + return None + try: + value = int(raw) + except ValueError as exc: + raise ValueError(f"{name} must be a positive integer, got {raw!r}") from exc + if value < 1: + raise ValueError(f"{name} must be a positive integer, got {value}") + return value + + +environment_variables: dict[str, Callable[[], Any]] = { + # this is used for configuring the default logging level + "AR_LOG_LEVEL": lambda: os.getenv("AR_LOG_LEVEL", "INFO").upper(), + "AR_ENABLE_COMPILE_PACKING": lambda: os.getenv("AR_ENABLE_COMPILE_PACKING", "0").lower() in ("1", "true", "yes"), + "AR_USE_MODELSCOPE": lambda: os.getenv("AR_USE_MODELSCOPE", "False").lower() in ["1", "true"], + "AR_WORK_SPACE": lambda: os.getenv("AR_WORK_SPACE", "ar_work_space").lower(), + "AR_ENABLE_UNIFY_MOE_INPUT_SCALE": lambda: os.getenv("AR_ENABLE_UNIFY_MOE_INPUT_SCALE", "False").lower() + in ["1", "true"], + "AR_OMP_NUM_THREADS": lambda: os.getenv("AR_OMP_NUM_THREADS", None), + "AR_DISABLE_OFFLOAD": lambda: os.getenv("AR_DISABLE_OFFLOAD", "0").lower() in ("1", "true", "yes"), + "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.