diff --git a/auto_round/compressors/base.py b/auto_round/compressors/base.py index f0060b54c8..0efa0dd205 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 diff --git a/auto_round/compressors/orchestrator.py b/auto_round/compressors/orchestrator.py index 30bad9d2ab..679b8e4785 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: diff --git a/auto_round/context/model.py b/auto_round/context/model.py index 285f6f5df5..a8002d9697 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 + # Remember the original string model name + # so OffloadManager can later materialize blocks directly from this + # checkpoint. + 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 - ) + # 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,87 @@ 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): + """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``. + """ + 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 + # 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. + 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``.""" + 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") + + # 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. + 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) diff --git a/auto_round/envs.py b/auto_round/envs.py index 0769754cfe..17444e71c5 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"), + # 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), } diff --git a/auto_round/utils/disk_stream_util.py b/auto_round/utils/disk_stream_util.py new file mode 100644 index 0000000000..5115b300f6 --- /dev/null +++ b/auto_round/utils/disk_stream_util.py @@ -0,0 +1,303 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +# 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 +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). + + 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 + + # 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, 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": + 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, tensor.dtype)) + + 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, 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: + """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. + """ + 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) diff --git a/auto_round/utils/offload.py b/auto_round/utils/offload.py index 975b5f5075..bc3d72f351 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: + """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) + # `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). + 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,36 @@ 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: + # 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. + 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. + 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) @@ -703,7 +783,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: + # 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. + 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 @@ -721,6 +817,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: + # 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. + return safe_save_file(state_dict, save_path) self._saved[name] = {"save_path": save_path} del state_dict 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 0000000000..bfc9949f8e --- /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"