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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] [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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] [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/41] 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/41] 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/41] 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/41] 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. From 7434f229b092c5df40387058298dd6f0e43ecf5e Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 29 Jul 2026 09:23:44 +0200 Subject: [PATCH 26/41] auto_scheme: disk-streaming support prep (gen_auto_scheme low_cpu_mem_usage, docs, tests) Signed-off-by: Fabrizio del Tin --- auto_round/auto_scheme/gen_auto_scheme.py | 11 +- docs/environments.md | 20 ++++ docs/environments_CN.md | 20 ++++ .../schemes/test_auto_scheme_disk_stream.py | 101 ++++++++++++++++++ 4 files changed, 149 insertions(+), 3 deletions(-) create mode 100644 test/test_cpu/schemes/test_auto_scheme_disk_stream.py diff --git a/auto_round/auto_scheme/gen_auto_scheme.py b/auto_round/auto_scheme/gen_auto_scheme.py index a65fc05ba..f3235c9da 100644 --- a/auto_round/auto_scheme/gen_auto_scheme.py +++ b/auto_round/auto_scheme/gen_auto_scheme.py @@ -100,9 +100,14 @@ def __init__( processor=None, ): self.auto_scheme = auto_scheme - if self.auto_scheme.low_cpu_mem_usage: - logger.info("force not using `low_cpu_mem_usage` in AutoScheme") - self.auto_scheme.low_cpu_mem_usage = False + # Upstream unconditionally forced low_cpu_mem_usage=False here + # (commit 0c9c5b1d, "reduce memory consumption... for gguf") because of an + # acknowledged bug for mixed INT4/INT8 schemes under the old OffloadManager- + # based streaming path (see an earlier commit's "low_cpu_mem_usage: bool = + # False # TODO bug for INT4 INT8 mixed bug"). We don't use that old path at + # all -- see delta_loss.py's materialize_module/free_module-based streaming, + # gated on this same flag -- so the bug that motivated disabling it doesn't + # apply here. self.model = model self.tokenizer = tokenizer self.processor = processor diff --git a/docs/environments.md b/docs/environments.md index 3ccd400c0..c27591934 100644 --- a/docs/environments.md +++ b/docs/environments.md @@ -151,6 +151,26 @@ export AR_AUTO_SCHEME_BATCH_SIZE=1 export AR_ENABLE_AUTO_SCHEME_PARALLEL=1 ``` +### AR_DISK_STREAM_MODEL +- **Description**: When enabled, `AutoRound(model=, ...)` builds the model as a meta-device skeleton instead of fully materializing the checkpoint on CPU RAM up front, and streams each decoder block's real weights from the checkpoint's safetensors shards on demand -- materializing right before a block is used (calibration, tuning, or `AutoScheme` sensitivity scoring) and freeing it back to meta right after. This keeps peak CPU RAM roughly flat regardless of checkpoint size, instead of proportional to it. Non-block parameters (embeddings, `lm_head`, final norm) are still loaded up front, since they are typically small. +- **Default**: `False` +- **Valid Values**: `"1"`, `"true"`, `"yes"` (case-insensitive) for enabling; any other value for disabling +- **Usage**: Enable this to quantize checkpoints larger than available CPU RAM + GPU VRAM combined. Only applies when `model` is a string (local directory) path; has no effect on already-loaded model objects. + +```bash +export AR_DISK_STREAM_MODEL=1 +``` + +### AR_RESUME_DIR +- **Description**: 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 or kill. +- **Default**: unset (no resumability) +- **Valid Values**: any writable directory path +- **Usage**: Set this for long-running quantization jobs on large checkpoints where a mid-run crash would otherwise be expensive to restart from scratch. + +```bash +export AR_RESUME_DIR=/path/to/resume/state +``` + ## Usage Examples ### Setting Environment Variables diff --git a/docs/environments_CN.md b/docs/environments_CN.md index e2100d60a..935accd39 100644 --- a/docs/environments_CN.md +++ b/docs/environments_CN.md @@ -151,6 +151,26 @@ export AR_AUTO_SCHEME_BATCH_SIZE=1 export AR_ENABLE_AUTO_SCHEME_PARALLEL=1 ``` +### AR_DISK_STREAM_MODEL +- **描述**：启用后，`AutoRound(model=, ...)` 会将模型构建为 meta 设备骨架，而不是先把整个 checkpoint 完全加载到 CPU 内存；随后按需从 checkpoint 的 safetensors 分片中流式加载每个解码器块的真实权重——在该块被使用前（校准、调优或 `AutoScheme` 敏感度评分）才实体化，用完后立即释放回 meta。这样峰值 CPU 内存基本保持平稳，而不会随 checkpoint 大小成比例增长。非块参数(embedding、`lm_head`、最终归一化层)体积通常较小，仍会一次性加载。 +- **默认值**：`False` +- **有效值**：`"1"`、`"true"`、`"yes"`(不区分大小写)表示启用；其他任何值表示禁用 +- **用途**：用于量化体积超过可用 CPU 内存 + GPU 显存总和的 checkpoint。仅在 `model` 为字符串(本地目录)路径时生效，对已加载的模型对象无效。 + +```bash +export AR_DISK_STREAM_MODEL=1 +``` + +### AR_RESUME_DIR +- **描述**：设置为目录路径后，逐块调优循环会在每完成一个块后将进度写入该目录，并在针对同一目录的新一次运行中从第一个未完成的块继续——而不是在崩溃或被杀死后从第 0 块重新开始整个调优过程。 +- **默认值**：未设置(不支持断点续跑) +- **有效值**：任意可写目录路径 +- **用途**：用于大 checkpoint 的长时间量化任务，避免运行中途崩溃导致从头重跑的高昂代价。 + +```bash +export AR_RESUME_DIR=/path/to/resume/state +``` + ## 使用示例 ### 设置环境变量 diff --git a/test/test_cpu/schemes/test_auto_scheme_disk_stream.py b/test/test_cpu/schemes/test_auto_scheme_disk_stream.py new file mode 100644 index 000000000..4205bde95 --- /dev/null +++ b/test/test_cpu/schemes/test_auto_scheme_disk_stream.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unit tests for AutoScheme's disk-streaming mode (AR_DISK_STREAM_MODEL). + +Verifies that streaming per-block sensitivity scoring from disk (instead of +fully materializing the checkpoint on CPU RAM up front) produces the same +mixed-bit layer_config as the non-streaming baseline, and that the underlying +materialize/free primitives round-trip correctly. +""" + +import os +import shutil + +import pytest +import torch + +from auto_round import AutoRound, AutoScheme +from auto_round.utils.disk_stream_util import build_meta_model, free_module, materialize_module, total_resident_bytes + + +@pytest.fixture(autouse=True) +def _clean_disk_stream_env(): + # AR_DISK_STREAM_MODEL is read lazily by auto_round.envs; make sure a test + # that sets it can't leak into whichever test runs next. + previous = os.environ.get("AR_DISK_STREAM_MODEL") + yield + if previous is None: + os.environ.pop("AR_DISK_STREAM_MODEL", None) + else: + os.environ["AR_DISK_STREAM_MODEL"] = previous + + +class TestAutoSchemeDiskStream: + @pytest.fixture(autouse=True) + def setup_save_dir(self, tmp_path): + self.save_dir = str(tmp_path / "saved") + yield + shutil.rmtree(self.save_dir, ignore_errors=True) + + def _gen_layer_config(self, model_name, target_bits=3.5): + # iters=1 (the standard tuning loop) rather than iters=0 (RTN): RTN's + # separate block-materialization path doesn't support disk streaming yet + # and is unrelated to this PR, which only streams AutoScheme's own + # sensitivity-scoring pass. + scheme = AutoScheme(avg_bits=target_bits, options=("W2A16", "W4A16", "BF16"), nsamples=1) + ar = AutoRound(model=model_name, scheme=scheme, iters=1, nsamples=1) + _, layer_config = ar.quantize() + return {name: cfg["bits"] for name, cfg in layer_config.items() if "bits" in cfg} + + def test_disk_stream_matches_baseline_layer_config(self, tiny_opt_model_path): + """AR_DISK_STREAM_MODEL=1 must select the exact same per-layer bits as the + non-streaming baseline -- streaming changes *how* weights are loaded during + scoring, not the scores themselves.""" + os.environ.pop("AR_DISK_STREAM_MODEL", None) + baseline_bits = self._gen_layer_config(tiny_opt_model_path) + + os.environ["AR_DISK_STREAM_MODEL"] = "1" + streamed_bits = self._gen_layer_config(tiny_opt_model_path) + + assert streamed_bits == baseline_bits + + def test_disk_stream_default_off(self): + """With AR_DISK_STREAM_MODEL unset, behavior must be the unstreamed default.""" + os.environ.pop("AR_DISK_STREAM_MODEL", None) + from auto_round import envs + + assert envs.AR_DISK_STREAM_MODEL is False + + +class TestDiskStreamUtilRoundTrip: + """Tests for the materialize/free primitives directly, independent of AutoScheme.""" + + def test_materialize_then_free_round_trip(self, tiny_opt_model_path): + model, _tokenizer, index = build_meta_model(tiny_opt_model_path) + block = model.model.decoder.layers[0] + + for _, tensor in list(block.named_parameters()): + assert str(tensor.device) == "meta" + + materialize_module(block, "model.decoder.layers.0", index, device="cpu") + for _, tensor in list(block.named_parameters()): + assert str(tensor.device) != "meta" + assert total_resident_bytes(block) > 0 + + free_module(block) + for _, tensor in list(block.named_parameters()): + assert str(tensor.device) == "meta" + assert total_resident_bytes(block) == 0 From b9eb98ff6b48a6c882e377442eeda429300ae882 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 29 Jul 2026 09:31:00 +0200 Subject: [PATCH 27/41] auto_scheme: stream per-layer sensitivity scoring block-by-block Reconstructed against the post-#2083 caching/parallel-scoring rewrite of delta_loss.py: threads disk_index through the serial scoring path (prepare_model_low_gpu, model_forward_low_gpu, get_score_for_scheme, gen_layer_config/_gen_layer_config) the same way as before, but now explicitly excludes streaming from the parallel multi-process scoring path added by #2083 -- each parallel worker fully loads its own copy of the model in a separate process, which defeats disk streaming's entire purpose. Per-scheme score caching is unaffected either way. The two model.to("cpu") calls that used to need an explicit disk_index-aware skip are now handled for free by safe_to_cpu_() (added upstream independently), which already checks for meta tensors before moving -- no manual guard needed at either call site anymore. Signed-off-by: Fabrizio del Tin --- auto_round/auto_scheme/delta_loss.py | 125 ++++++++++++++++++++++++--- 1 file changed, 113 insertions(+), 12 deletions(-) diff --git a/auto_round/auto_scheme/delta_loss.py b/auto_round/auto_scheme/delta_loss.py index 7dbf1a6c3..4592d1f98 100644 --- a/auto_round/auto_scheme/delta_loss.py +++ b/auto_round/auto_scheme/delta_loss.py @@ -563,7 +563,7 @@ def __init__(self, message): last_grad_input = None -def prepare_model_low_gpu(model, block_inputs: dict = None, pbar=None, major_device="cpu"): +def prepare_model_low_gpu(model, block_inputs: dict = None, pbar=None, major_device="cpu", disk_index=None): """Wrap every block's forward so that, for one calibration batch, it (1) moves itself to ``major_device`` on demand, (2) records its own inputs into ``block_inputs`` (on CPU) so they can be replayed later, and (3) moves itself back to CPU once done. @@ -571,6 +571,12 @@ def prepare_model_low_gpu(model, block_inputs: dict = None, pbar=None, major_dev Called once per calibration batch before ``model_forward_low_gpu`` runs the actual forward+backward -- the recorded ``block_inputs`` are what let the backward pass be replayed manually, one block at a time, without keeping every block resident on GPU. + + When ``disk_index`` is set (streaming mode -- the model is a meta-device skeleton, + see ``gen_layer_config``/``disk_stream_util.py``), each block's real weights are + materialized from the checkpoint right before its own forward and released back to + meta right after, instead of assuming the block already has real CPU-resident weights + to shuffle to GPU and back. """ block_inputs.clear() for n, m in model.named_modules(): @@ -590,6 +596,10 @@ def new_forward(*args, **kwargs): """Move the block to device, run its original forward, cache its (CPU) inputs for later replay, then move the block back to CPU. """ + if disk_index is not None: + from auto_round.utils.disk_stream_util import materialize_module + + materialize_module(module, module_name, disk_index, device=major_device) move_module_to_tuning_device(module, major_device=major_device) # for n,m in module.named_modules(): # if hasattr(m, "post_init_qdqw"): @@ -608,7 +618,12 @@ def new_forward(*args, **kwargs): } block_inputs[module_name] = input_info - module.to("cpu") + if disk_index is not None: + from auto_round.utils.disk_stream_util import free_module + + free_module(module) + else: + module.to("cpu") memory_monitor.update(device_list=major_device) # clear_memory(device_list=major_device) #slow # memory_monitor.log_summary() @@ -685,7 +700,7 @@ def model_forward(model, data, **forward_kwargs): return model(**prepared, **forward_kwargs), prepared -def model_forward_low_gpu(model, dataloader, major_device="cuda", pbar=None, scheme_tag=None): +def model_forward_low_gpu(model, dataloader, major_device="cuda", pbar=None, scheme_tag=None, disk_index=None): """Run one full scoring pass (all calibration batches) in low-GPU-memory mode. For each batch: capture per-block inputs via ``prepare_model_low_gpu``, run a forward @@ -693,6 +708,11 @@ def model_forward_low_gpu(model, dataloader, major_device="cuda", pbar=None, sch raising ``MyCustomError``), then manually replay the backward pass block-by-block (moving each block to ``major_device`` only for its own recompute + backward, then back to CPU) so only one block's weights need to be resident on GPU at a time. + + When ``disk_index`` is set (streaming mode -- the model is a meta-device skeleton), + each block's real weights are materialized from the checkpoint right before use and + released back to meta right after, both here (the manual reverse-order backward + replay) and in ``prepare_model_low_gpu`` (the initial forward capture pass). """ block_inputs = {} total_batches = len(dataloader) if hasattr(dataloader, "__len__") else None @@ -710,7 +730,7 @@ def backward_pre_hook(module, grad_input): raise MyCustomError("Interrupt backward pass") for batch_idx, data in enumerate(dataloader, start=1): - prepare_model_low_gpu(model, block_inputs, major_device=major_device, pbar=pbar) + prepare_model_low_gpu(model, block_inputs, major_device=major_device, pbar=pbar, disk_index=disk_index) # lm_head sits outside every decoder block, so it never gets `grad_mode=True` # in the manual block-by-block backward below. Scope the fix narrowly to @@ -777,6 +797,10 @@ def backward_pre_hook(module, grad_input): for n, m in block_module.named_modules(): if hasattr(m, "grad_mode"): m.grad_mode = True + if disk_index is not None: + from auto_round.utils.disk_stream_util import materialize_module + + materialize_module(block_module, block_name, disk_index, device=major_device) move_module_to_tuning_device(block_module, major_device=major_device) # Set the block to eval mode while enabling gradient computation @@ -816,7 +840,12 @@ def backward_pre_hook(module, grad_input): break del block_output, main_output, block_input_args, block_input_kwargs - block_module.to("cpu") + if disk_index is not None: + from auto_round.utils.disk_stream_util import free_module + + free_module(block_module) + else: + block_module.to("cpu") # clear_memory(device_list=major_device) # this one is very slow and seems does not affect max ram usage memory_monitor.update() @@ -856,6 +885,7 @@ def get_score_for_scheme( force_mllm: bool = False, model_name: Optional[str] = None, scheme_tag: Optional[str] = None, + disk_index=None, ): """Wrap every quantizable layer in ``quant_layer_names`` with a scoring wrapper, run forward(+backward, unless RTN-only) calibration over ``nsamples`` examples from @@ -1078,11 +1108,20 @@ def _build_mllm_calib_dataloader(): "AutoScheme(force_mllm): cannot build mllm dataloader. " "Provide a `processor` and a multimodal `dataset`." ) - model_forward_low_gpu(model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag) + model_forward_low_gpu( + model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag, disk_index=disk_index + ) else: try: dataloader = _build_calib_dataloader() - model_forward_low_gpu(model, dataloader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag) + model_forward_low_gpu( + model, + dataloader, + major_device=major_device, + pbar=pbar, + scheme_tag=scheme_tag, + disk_index=disk_index, + ) except Exception as exc: # noqa: BLE001 if not is_vlm: raise @@ -1094,7 +1133,9 @@ def _build_mllm_calib_dataloader(): batch_size = 1 if mllm_loader is None: raise - model_forward_low_gpu(model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag) + model_forward_low_gpu( + model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag, disk_index=disk_index + ) else: for n, m in model.named_modules(): if hasattr(m, "grad_mode"): @@ -1861,6 +1902,7 @@ def _gen_layer_config( device_list=None, processor=None, is_vlm: bool = False, + disk_index=None, ): """Score every candidate scheme in ``auto_scheme.options`` against ``quant_layer_names`` and return per-layer per-scheme losses used by the caller to pick a final bit-width @@ -1878,8 +1920,15 @@ def _gen_layer_config( # Create offload context for CPU RAM optimization # Note: low_cpu_mem_usage only works when low_gpu_mem_usage is also enabled, # because it requires layer-by-layer processing + # + # When disk_index is set, gen_layer_config already built the model as a + # meta-device skeleton and materialize_module/free_module (called directly + # around each block's use, see get_score_for_scheme/model_forward_low_gpu/ + # prepare_model_low_gpu above) are the actual streaming mechanism -- + # OffloadManager's hook-based approach doesn't apply to a model that never + # had real CPU-resident weights to begin with. offload_context = None - if auto_scheme.low_cpu_mem_usage and auto_scheme.low_gpu_mem_usage: + if disk_index is None and auto_scheme.low_cpu_mem_usage and auto_scheme.low_gpu_mem_usage: _model_dir = model_name if _model_dir is None and hasattr(model, "config"): _model_dir = getattr(model.config, "_name_or_path", None) @@ -2225,6 +2274,13 @@ def _save_per_op_scores(index, scheme, cache_key, cache_path, per_op_scores): and num_gpus >= 1 and len(uncached_indices) >= 2 and not need_imatrix + # Each parallel worker fully loads its own copy of the model + # (_load_scheme_worker_model) in a separate process -- incompatible + # with a meta-device streaming skeleton, whose entire point is to + # avoid ever materializing a full copy. Force serial scoring + # (which honors disk_index via materialize_module/free_module) + # instead when streaming is active. + and disk_index is None ) if not parallel_enabled and len(uncached_indices) >= 2: logger.info( @@ -2360,7 +2416,13 @@ def _serialize_scheme(scheme): pbar.reset(total=pbar_cnt) if not parallel_done: - if uncached_indices: + if uncached_indices and disk_index is None: + # Skipped in streaming mode: materialize_model_ only acts on + # ReplacementModuleBase (fused-MoE) instances -- a no-op for + # dense models like ours -- but it also warns once per + # still-meta parameter/buffer, which would flood the log with + # one warning per decoder-block tensor (intentionally still + # meta, to be streamed on demand later). from auto_round.modeling.fused_moe.replace_modules import materialize_model_ materialize_model_(model) @@ -2422,6 +2484,7 @@ def _serialize_scheme(scheme): force_mllm=force_mllm, model_name=model_name, scheme_tag=scheme_tag, + disk_index=disk_index, ) memory_monitor.update() memory_monitor.log_summary() @@ -2689,12 +2752,48 @@ def gen_layer_config( """ model_name = None is_vlm = False + disk_index = None if isinstance(model, str): model_name = model - model, tokenizer, processor, _, _, is_vlm, _ = load_model(model_name, device="cpu", use_auto_mapping=False) + is_vlm = is_mllm_model(model_name) + if not is_vlm and auto_scheme.low_cpu_mem_usage and low_gpu_mem_usage: + # Disk-streamed load (meta-device skeleton + on-demand per-block + # materialize/free, see disk_stream_util.py) instead of + # load_model()'s full-checkpoint CPU RAM load -- infeasible for a + # checkpoint bigger than available RAM. Falls back to load_model() + # for anything build_meta_model doesn't cover. + try: + from auto_round.utils.disk_stream_util import build_meta_model, materialize_non_block_params + + model, tokenizer, disk_index = build_meta_model(model_name) + block_prefixes = flatten_list(get_block_names(model, quant_vision=is_vlm)) + materialize_non_block_params(model, block_prefixes, disk_index, device="cpu") + except Exception as exc: # noqa: BLE001 + logger.warning( + f"AutoScheme streaming load failed ({exc}); falling back to " + f"load_model() (needs the whole checkpoint resident in RAM)." + ) + disk_index = None + model, tokenizer, processor, _, _, is_vlm, _ = load_model( + model_name, device="cpu", use_auto_mapping=False + ) + else: + model, tokenizer, processor, _, _, is_vlm, _ = load_model(model_name, device="cpu", use_auto_mapping=False) else: # Object passed in: still try to detect VLM so we can pick the right dataloader later. - _, _, _, _, _, is_vlm, _ = load_model(model) + try: + _, _, _, _, _, is_vlm, _ = load_model(model) + except Exception: # noqa: BLE001 + is_vlm = False + # By the time AutoRound's compressor calls into AutoScheme, ModelContext + # has already turned a string model into a real object -- meaning the + # `isinstance(model, str)` branch above never actually runs in the + # standard `AutoRound(model=path, ...)` API flow. When ModelContext + # itself built the object as a meta skeleton (AR_DISK_STREAM_MODEL=1), + # it stashes the SafetensorsIndex on the model so we can pick it up + # here instead of re-detecting streaming mode from scratch (or, + # worse, silently treating a meta model as if it were fully real). + disk_index = getattr(model, "_disk_stream_index", None) # ---- Vision-tower scoring requires a full backward ---- # # ``model_forward_low_gpu`` only walks the language tower (it uses @@ -2789,6 +2888,7 @@ def _enable_gc(mod): min_avg_bit_scheme=min_avg_bit_scheme, processor=processor, is_vlm=is_vlm, + disk_index=disk_index, ) except torch.OutOfMemoryError: logger.warning( @@ -2820,6 +2920,7 @@ def _enable_gc(mod): min_avg_bit_scheme=min_avg_bit_scheme, processor=processor, is_vlm=is_vlm, + disk_index=disk_index, ) return res From df209eeb199934712277db3f5b40364b8f3ffe71 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 20:36:30 +0200 Subject: [PATCH 28/41] Re-trigger CI (unrelated HF Hub 503 flake on test_gptoss) Signed-off-by: Fabrizio del Tin From 5ccb5feacf6243f8b82716390200d8d3cb3bfb0b Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:47:00 +0200 Subject: [PATCH 29/41] calibration: stream targeted block re-quantization's calibration pass to_quant_block_names lets a caller restrict tuning to a subset of decoder blocks -- useful for cheaply re-quantizing just a couple of blocks in an already-produced checkpoint at higher precision instead of redoing a full multi-hour run. Combined with AR_DISK_STREAM_MODEL, this crashed: the "cache block inputs" forward pass (calibrate_on_cpu branch) needs real weights in every block leading up to (and, when there's only one target block, all the way through) the target block(s), but nothing materializes blocks outside quant_block_list for this specific forward pass -- they stay meta forever, and the forward silently propagates meta-ness through them until it collides with a genuinely-materialized module. Reproduced against a tiny hybrid-MoE fixture with to_quant_block_names restricted to one block: "Tensor on device meta is not on the expected device cpu!" inside the final norm. Full (unrestricted) runs never hit this, since quant_block_list already covers every block in that case. Fix: when disk streaming is active and any decoder block still has meta parameters at this point, wrap the calibration forward with the existing stream_block_forward primitive (already used by _streaming_eval_model() in bin/el_quantize_autoround_mixed.py for the analogous held-out-loss-eval case), scoped to just the still-meta blocks -- materializing each on demand and freeing it right after. Only activates when there's something left meta to fix, so it's a no-op for the normal full-quantization path. Also adds an AR_CALIB_STREAM_DEVICE env-gated fast path: the default keeps this forward entirely on cpu (mixing a GPU-streamed block with cpu-resident hidden states crashed with a device mismatch the first time this was tried at full 397B scale), but a full cpu forward through every pre-target block of a 100B+ model is unusably slow for this specific targeted-requant use case. When set, every already-real (non-meta) param/buffer is moved to that device for the duration of the pass (including stray non-parameter buffers like RoPE inv_freq), blocks stream-materialize there too, and everything is moved back afterward so the tuning phase sees the exact layout it would have without this. Signed-off-by: Fabrizio del Tin --- auto_round/calibration/llm.py | 98 +++++++++++++++++++++++++++++++++-- 1 file changed, 95 insertions(+), 3 deletions(-) diff --git a/auto_round/calibration/llm.py b/auto_round/calibration/llm.py index 9bd45ae2d..c593c4f20 100644 --- a/auto_round/calibration/llm.py +++ b/auto_round/calibration/llm.py @@ -81,16 +81,108 @@ def calibration(self, block_names, nsamples, layer_names=None, last_cache_name=N ): # low_gpu_mem_usage or calibrate only the embedding layer (also fast on CPU) calibrate_on_cpu = True + # Local addition (not upstream): when AR_DISK_STREAM_MODEL built a + # meta-device skeleton and to_quant_block_names restricts + # quant_block_list to fewer than all decoder blocks (a targeted + # re-quantization of specific blocks in an already-quantized + # checkpoint, not the normal full-model run), this forward pass + # still needs REAL weights in every block leading up to (and, when + # there's only one target block so last_cache_name ends up None + # and no early-stop applies, all the way through) the target + # block(s) -- but nothing else materializes blocks outside + # quant_block_list for this specific forward pass. They stay meta + # forever, and the forward silently propagates meta-ness through + # them until it collides with a genuinely-materialized module + # (the final norm, or any block actually in quant_block_list that + # already went through the compressor's own offload/reload cycle) + # -- confirmed via a reproduction against the tiny hybrid-MoE test + # fixture (checkpoint_full_arch_test, to_quant_block_names="model. + # layers.7"): "Tensor on device meta is not on the expected device + # cpu!" inside Qwen3_5MoeRMSNorm.forward. Full (unrestricted) runs + # never hit this: quant_block_list already covers every block in + # that case, so there's nothing left outside it to leave meta. + # Fix: stream-materialize (then free) any block that's still meta + # for the duration of this one forward pass, reusing the same + # stream_block_forward primitive _streaming_eval_model() already + # uses in bin/el_quantize_autoround_mixed.py for the analogous + # held-out-loss-eval case. + stream_ctx = None + _moved_tensors = [] + if envs.AR_DISK_STREAM_MODEL: + disk_index = getattr(self.model, "_disk_stream_index", None) + if disk_index is not None: + from auto_round.utils import get_block_names, get_module + + meta_block_names = [ + name + for name in flatten_list(get_block_names(self.model)) + if any(p.device.type == "meta" for p in get_module(self.model, name).parameters()) + ] + if meta_block_names: + from auto_round.utils.disk_stream_util import stream_block_forward + + # This whole branch is calibrate_on_cpu -- every other + # tensor in this forward pass (hidden states, + # already-materialized non-block params) lives on cpu, + # not device_manager.device (the GPU tuning device). + # Materializing ONLY the streamed blocks on GPU caused + # a real cuda:0/cpu mismatch crash the first time it + # was tried against the full 397B-scale checkpoint -- + # hence the cpu default. But a cpu forward through + # every pre-target block of a 100B+ model is unusably + # slow for the targeted re-quantization use case, so + # AR_CALIB_STREAM_DEVICE (set by el_requantize_blocks + # .py) opts the WHOLE pass onto one device coherently: + # every already-real (non-meta) param/buffer is moved + # there for the duration (the same recipe bin/ + # el_quantize_autoround_mixed.py's _streaming_eval_ + # model() already proved at 207GB scale, including its + # stray-buffer sweep for e.g. RoPE inv_freq), blocks + # stream-materialize there, and calib() batches follow + # model.device automatically. Everything is moved back + # afterwards so the tuning phase sees the exact layout + # it would have without this. + calib_stream_device = envs.AR_CALIB_STREAM_DEVICE or "cpu" + if calib_stream_device != "cpu": + for module in self.model.modules(): + for _pname, _t in list(module.named_parameters(recurse=False)) + list( + module.named_buffers(recurse=False) + ): + if _t.device.type != "meta" and str(_t.device) != calib_stream_device: + _moved_tensors.append((_t, str(_t.device))) + _t.data = _t.data.to(calib_stream_device) + logger.info( + "AR_CALIB_STREAM_DEVICE=%s: moved %d non-block tensors for the " + "calibration forward; decoder blocks stream through the same device.", + calib_stream_device, + len(_moved_tensors), + ) + + stream_ctx = stream_block_forward( + self.model, + disk_index, + device=calib_stream_device, + block_names=meta_block_names, + ) try: - all_inputs = self.cache_inter_data( - block_names, nsamples, layer_names=[], last_cache_name=last_cache_name - ) + if stream_ctx is not None: + with stream_ctx: + all_inputs = self.cache_inter_data( + block_names, nsamples, layer_names=[], last_cache_name=last_cache_name + ) + else: + all_inputs = self.cache_inter_data( + block_names, nsamples, layer_names=[], last_cache_name=last_cache_name + ) except NotImplementedError as error: error_msg = str(error) if "flash_attn::" in error_msg and "CPU" in error_msg: cannot_calibrate_on_cpu = True else: raise error + finally: + for _t, _orig_device in _moved_tensors: + _t.data = _t.data.to(_orig_device) if not calibrate_on_cpu or cannot_calibrate_on_cpu: try: From 8c7593800392a6212412117f18b96716bd8a7859 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 14:46:14 +0200 Subject: [PATCH 30/41] Strip leaked internal-repo comment marker Remove a "Local addition (not upstream)" comment prefix -- local-only tooling metadata that doesn't apply upstream. No logic change. Signed-off-by: Fabrizio del Tin --- auto_round/calibration/llm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auto_round/calibration/llm.py b/auto_round/calibration/llm.py index c593c4f20..08f1965f9 100644 --- a/auto_round/calibration/llm.py +++ b/auto_round/calibration/llm.py @@ -81,7 +81,7 @@ def calibration(self, block_names, nsamples, layer_names=None, last_cache_name=N ): # low_gpu_mem_usage or calibrate only the embedding layer (also fast on CPU) calibrate_on_cpu = True - # Local addition (not upstream): when AR_DISK_STREAM_MODEL built a + # When AR_DISK_STREAM_MODEL built a # meta-device skeleton and to_quant_block_names restricts # quant_block_list to fewer than all decoder blocks (a targeted # re-quantization of specific blocks in an already-quantized From a2bfa508e16b7de3568dcb992d1199c0f35d6a67 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 19:21:27 +0200 Subject: [PATCH 31/41] Add regression test for targeted block re-quantization + disk streaming test_targeted_block_with_disk_streaming_does_not_crash reproduces the exact crash this PR fixes: restricting to_quant_block_names to the last of 3 blocks leaves the earlier blocks meta-only (never in quant_block_list), and the calibration forward pass used to propagate that meta-ness until it hit "Tensor on device meta is not on the expected device cpu!". Verified this test fails with that error against the pre-fix commit (15aa8854) and passes with the fix. test_targeted_block_without_disk_streaming_still_works is the baseline: the same targeted re-quantization without streaming never hit this bug, so it must keep working unchanged. Signed-off-by: Fabrizio del Tin --- .../core/test_targeted_block_calib_stream.py | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 test/test_cpu/core/test_targeted_block_calib_stream.py diff --git a/test/test_cpu/core/test_targeted_block_calib_stream.py b/test/test_cpu/core/test_targeted_block_calib_stream.py new file mode 100644 index 000000000..16dfe8bfb --- /dev/null +++ b/test/test_cpu/core/test_targeted_block_calib_stream.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Integration test for targeted block re-quantization (``to_quant_block_names``) +combined with disk streaming (``AR_DISK_STREAM_MODEL``). + +Regression coverage for: restricting tuning to a block that isn't the first +one left every block *before* the target on the meta device, and the +"cache block inputs" calibration forward silently propagated that meta-ness +until it collided with a genuinely-materialized module -- "Tensor on device +meta is not on the expected device cpu!". A full (unrestricted) run never +hits this, since every block is already covered by quant_block_list in that +case. +""" + +import os + +import pytest + +from auto_round import AutoRound + + +@pytest.fixture(autouse=True) +def _clean_disk_stream_env(): + previous = os.environ.get("AR_DISK_STREAM_MODEL") + yield + if previous is None: + os.environ.pop("AR_DISK_STREAM_MODEL", None) + else: + os.environ["AR_DISK_STREAM_MODEL"] = previous + + +@pytest.fixture(scope="module") +def tiny_opt_3layer_model_path(): + from test.helpers import save_tiny_model + + path = save_tiny_model("facebook/opt-125m", "./tmp/tiny_opt_3layer_model_path", num_layers=3) + yield path + import shutil + + shutil.rmtree(path, ignore_errors=True) + + +class TestTargetedBlockCalibStream: + def test_targeted_block_with_disk_streaming_does_not_crash(self, tiny_opt_3layer_model_path): + """Restricting to_quant_block_names to the LAST of 3 blocks leaves + blocks 0 and 1 meta-only (never in quant_block_list) while the + calibration forward pass still needs to run through them to reach + block 2 -- exactly the scenario that used to crash.""" + os.environ["AR_DISK_STREAM_MODEL"] = "1" + + ar = AutoRound( + model=tiny_opt_3layer_model_path, + scheme="W4A16", + iters=1, + nsamples=1, + to_quant_block_names="model.decoder.layers.2", + ) + _, layer_config = ar.quantize() + + quantized_layers = {name for name, cfg in layer_config.items() if "bits" in cfg} + assert quantized_layers, "expected the target block's layers to be quantized" + assert all(name.startswith("model.decoder.layers.2.") for name in quantized_layers) + assert not any( + name.startswith(("model.decoder.layers.0.", "model.decoder.layers.1.")) for name in quantized_layers + ) + + def test_targeted_block_without_disk_streaming_still_works(self, tiny_opt_3layer_model_path): + """Baseline: the same targeted re-quantization without disk streaming + never hit this bug (nothing is meta), so it must keep working too.""" + os.environ.pop("AR_DISK_STREAM_MODEL", None) + + ar = AutoRound( + model=tiny_opt_3layer_model_path, + scheme="W4A16", + iters=1, + nsamples=1, + to_quant_block_names="model.decoder.layers.2", + ) + _, layer_config = ar.quantize() + + quantized_layers = {name for name, cfg in layer_config.items() if "bits" in cfg} + assert quantized_layers + assert all(name.startswith("model.decoder.layers.2.") for name in quantized_layers) From 7c709d0989c37e59dad21926bfd0242d056ac640 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 20:37:10 +0200 Subject: [PATCH 32/41] Re-trigger CI (unrelated HF Hub 503 flake on test_gptoss) Signed-off-by: Fabrizio del Tin From 3841bb37660c10d249de6b3055afe8f046402c4a Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Mon, 27 Jul 2026 11:23:19 +0200 Subject: [PATCH 33/41] Re-trigger CI (retest after flaky VRAM/RAM perf checks and llmc_format flake) Signed-off-by: Fabrizio del Tin From 007507628e29a51bdc37b900bc15eb13fd4536ae Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 15:47:09 +0200 Subject: [PATCH 34/41] calibration: allow non-MLLM calibration datasets when not quantizing vision MLLMCalibrator unconditionally routed to get_mllm_dataloader whenever the model was multimodal, which KeyErrors on any dataset not in its MLLM_DATASET registry (its own "os.path.isfile(dataset) or dataset in MLLM_DATASET" guard still indexes MLLM_DATASET with the raw file path). But a local text dataset (file/dir/HF text set) is perfectly valid calibration for a VLM whose non-text modules are NOT being quantized (quant_nontext_module=False): the full-model forward runs text-only, and the loop's generic-dict branch already just feeds model(**batch) regardless of dataset shape. When quant_nontext_module is False and the dataset isn't a known MLLM_DATASET entry, use the standard text get_dataloader() instead of falling through to get_mllm_dataloader. Signed-off-by: Fabrizio del Tin --- auto_round/calibration/mllm.py | 65 +++++++++++++++++++++++----------- 1 file changed, 44 insertions(+), 21 deletions(-) diff --git a/auto_round/calibration/mllm.py b/auto_round/calibration/mllm.py index 6ea2d786e..b3990a7c5 100644 --- a/auto_round/calibration/mllm.py +++ b/auto_round/calibration/mllm.py @@ -108,27 +108,50 @@ def calib(self, nsamples: int, bs: int) -> None: " will use liuhaotian/llava_conv_58k with default config as an alternative." ) dataset = "liuhaotian/llava_conv_58k" - orig_bs = self.batch_size - ( - self.dataloader, - self.batch_size, - self.seqlen, - ) = get_mllm_dataloader( - template=template_obj, - model=self.model, - tokenizer=tokenizer, - processor=self.processor, - image_processor=image_processor, - dataset=dataset, - extra_data_dir=self.extra_data_dir, - seqlen=self.seqlen, - bs=bs, - seed=self.seed, - nsamples=nsamples, - quant_nontext_module=self.quant_nontext_module, - ) - if orig_bs != 1 and self.batch_size == 1: - self.is_only_supported_bs1 = True + from auto_round.compressors.mllm.dataset import MLLM_DATASET + + if not self.quant_nontext_module and dataset not in MLLM_DATASET: + # Local patch (see LOCAL_PATCHES.md, "Text-only calibration of + # multimodal checkpoints"): a local text dataset (file/dir/HF + # text set) is valid calibration for a VLM whose non-text + # modules are NOT being quantized -- the full-model forward + # runs text-only and the generic-dict branch of the loop below + # feeds model(**batch). Upstream instead falls through to + # get_mllm_dataloader, which KeyErrors on any dataset not in + # its MLLM_DATASET registry (its `os.path.isfile(dataset) or + # dataset in MLLM_DATASET` guard still indexes MLLM_DATASET + # with the file path). + from auto_round.calib_dataset import get_dataloader + + logger.info( + f"Multimodal model with non-MLLM calibration dataset {dataset!r} and " + "quant_nontext_module=False: using the standard text dataloader " + "(vision/audio towers are not being quantized, so text-only " + "calibration through the full-model forward is sufficient)." + ) + self.dataloader = get_dataloader(tokenizer, self.seqlen, dataset, self.seed, bs, nsamples) + else: + orig_bs = self.batch_size + ( + self.dataloader, + self.batch_size, + self.seqlen, + ) = get_mllm_dataloader( + template=template_obj, + model=self.model, + tokenizer=tokenizer, + processor=self.processor, + image_processor=image_processor, + dataset=dataset, + extra_data_dir=self.extra_data_dir, + seqlen=self.seqlen, + bs=bs, + seed=self.seed, + nsamples=nsamples, + quant_nontext_module=self.quant_nontext_module, + ) + if orig_bs != 1 and self.batch_size == 1: + self.is_only_supported_bs1 = True else: self.dataloader = self.dataset From a95a3498757ac0d6d556b33b991bdc0efb3ccb31 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Sun, 19 Jul 2026 18:23:40 +0200 Subject: [PATCH 35/41] Re-trigger CI The prior /azp run Performance-Test-AutoRound comment failed with 'Commenter does not have sufficient privileges' (external contributors can't trigger Azure Pipelines re-runs by comment on this repo). Pushing an empty commit instead, since a new commit re-runs CI regardless of commenter privileges. See PR discussion: the only failing check (Qwen3_W4A16 peak VRAM, +4.11%% vs baseline) is on a non-multimodal model that never exercises this PR's is_mllm=True-gated change in calibration/mllm.py. Signed-off-by: Fabrizio del Tin From b34e11a26732b1b1a69f6c5683031a8f1f729cf9 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 14:31:27 +0200 Subject: [PATCH 36/41] Strip leaked internal-repo comment reference Remove a "Local patch (see LOCAL_PATCHES.md, ...)" comment marker -- local-only tooling metadata that doesn't apply upstream. No logic change. Signed-off-by: Fabrizio del Tin --- auto_round/calibration/mllm.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/auto_round/calibration/mllm.py b/auto_round/calibration/mllm.py index b3990a7c5..4eb3ea80b 100644 --- a/auto_round/calibration/mllm.py +++ b/auto_round/calibration/mllm.py @@ -111,8 +111,7 @@ def calib(self, nsamples: int, bs: int) -> None: from auto_round.compressors.mllm.dataset import MLLM_DATASET if not self.quant_nontext_module and dataset not in MLLM_DATASET: - # Local patch (see LOCAL_PATCHES.md, "Text-only calibration of - # multimodal checkpoints"): a local text dataset (file/dir/HF + # A local text dataset (file/dir/HF # text set) is valid calibration for a VLM whose non-text # modules are NOT being quantized -- the full-model forward # runs text-only and the generic-dict branch of the loop below From 11449fa6e44790f9d79b494e4ead2746070e6203 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 19:54:05 +0200 Subject: [PATCH 37/41] Add regression test for VLM calibration with a local text dataset test_local_text_dataset_with_quant_nontext_module_false reproduces the exact crash this PR fixes: a local text calibration file used to KeyError inside get_mllm_dataloader (os.path.isfile(dataset) makes it enter that branch, but the file path is never actually a key in MLLM_DATASET). Verified this test fails with that exact KeyError against the pre-fix commit (c774c83b) and passes with the fix, which routes quant_nontext_module=False + non-MLLM dataset through the standard text dataloader instead. Signed-off-by: Fabrizio del Tin --- .../models/test_mllm_text_only_calib.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 test/test_cpu/models/test_mllm_text_only_calib.py diff --git a/test/test_cpu/models/test_mllm_text_only_calib.py b/test/test_cpu/models/test_mllm_text_only_calib.py new file mode 100644 index 000000000..40e3cf0de --- /dev/null +++ b/test/test_cpu/models/test_mllm_text_only_calib.py @@ -0,0 +1,71 @@ +# 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. + +""" +Regression test for calibrating a multimodal (VLM) checkpoint with a +non-MLLM (plain text) local calibration dataset when quant_nontext_module +is False. + +Before this fix, any string dataset unconditionally fell through to +``get_mllm_dataloader``, which indexes ``MLLM_DATASET`` with the dataset's +value even when ``os.path.isfile(dataset)`` is what made it enter that +branch (a local file path is never actually a key in that registry) -- +``KeyError: ''``. Since the vision/audio towers aren't +being quantized in this scenario, plain text-only calibration through the +standard (non-MLLM) dataloader is sufficient and now used instead. +""" + +import json + +import pytest + +from auto_round import AutoRound + + +class TestMllmTextOnlyCalibration: + @pytest.fixture(autouse=True) + def setup_save_dir(self, tmp_path): + self.save_dir = str(tmp_path / "saved") + yield + import shutil + + shutil.rmtree(self.save_dir, ignore_errors=True) + + @pytest.fixture + def local_text_calib_file(self, tmp_path): + data = [{"text": "The quick brown fox jumps over the lazy dog. " * 40}] * 20 + path = tmp_path / "calib.json" + with open(path, "w") as f: + json.dump(data, f) + return str(path) + + def test_local_text_dataset_with_quant_nontext_module_false(self, tiny_qwen_vl_model_path, local_text_calib_file): + """A local (non-MLLM-registered) text file used to KeyError inside + get_mllm_dataloader; it must now route through the standard text + dataloader instead, since the vision tower isn't being quantized.""" + ar = AutoRound( + model=tiny_qwen_vl_model_path, + scheme="W4A16", + iters=1, + nsamples=1, + seqlen=32, + dataset=local_text_calib_file, + quant_nontext_module=False, + ) + _, layer_config = ar.quantize() + + quantized_language_layers = [ + name for name, cfg in layer_config.items() if "bits" in cfg and "language_model" in name + ] + assert quantized_language_layers, "expected language-model layers to be quantized" From 6602c7f142ca5de6155898ed1ed9ea94e2ad53cb Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Wed, 22 Jul 2026 20:37:26 +0200 Subject: [PATCH 38/41] Re-trigger CI (unrelated HF Hub 503 flake on test_gptoss) Signed-off-by: Fabrizio del Tin From 7cc83e3b2cc4823061909aa16964425ecbf66734 Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Mon, 27 Jul 2026 11:23:21 +0200 Subject: [PATCH 39/41] Re-trigger CI (retest after flaky VRAM/RAM perf checks and llmc_format flake) Signed-off-by: Fabrizio del Tin From 0d56e632b7c7ea21e77122fd91cbb39fa021d77e Mon Sep 17 00:00:00 2001 From: Fabrizio del Tin Date: Mon, 27 Jul 2026 13:55:28 +0200 Subject: [PATCH 40/41] Re-trigger CI (retest after flaky perf VRAM check on Qwen3_FP8_STATIC) Signed-off-by: Fabrizio del Tin From 6e64458fa30051181ce496c8ade7db6a26dd509b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 07:37:12 +0000 Subject: [PATCH 41/41] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- auto_round/auto_scheme/delta_loss.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/auto_round/auto_scheme/delta_loss.py b/auto_round/auto_scheme/delta_loss.py index 4592d1f98..20a2427f9 100644 --- a/auto_round/auto_scheme/delta_loss.py +++ b/auto_round/auto_scheme/delta_loss.py @@ -1134,7 +1134,12 @@ def _build_mllm_calib_dataloader(): if mllm_loader is None: raise model_forward_low_gpu( - model, mllm_loader, major_device=major_device, pbar=pbar, scheme_tag=scheme_tag, disk_index=disk_index + model, + mllm_loader, + major_device=major_device, + pbar=pbar, + scheme_tag=scheme_tag, + disk_index=disk_index, ) else: for n, m in model.named_modules():