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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] [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/12] 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/12] 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/12] 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/12] Re-trigger CI (retest after flaky VRAM/RAM perf checks and llmc_format flake) Signed-off-by: Fabrizio del Tin