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..f5f92b905a 100644 --- a/auto_round/compressors/orchestrator.py +++ b/auto_round/compressors/orchestrator.py @@ -13,15 +13,17 @@ # limitations under the License. import copy import gc +import os import time from functools import partial -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union import accelerate import torch 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, @@ -58,6 +60,9 @@ from auto_round.utils.device_manager import device_manager from auto_round.wrapper import WrapperMultiblock +if TYPE_CHECKING: + from auto_round.utils.resume import ResumeState + # TODO wenhuach align all the API args class CompressionOrchestrator(BaseOrchestrator): @@ -178,6 +183,8 @@ def _quantize_blocks( pbar: tqdm | None = None, input_others_extra_blocks: dict | None = None, token_ids: list[torch.Tensor] | None = None, + resume_state: Optional["ResumeState"] = None, + resume_input_ids=None, ): """Quantize and dequantize the weights of the specified blocks in the model. @@ -187,6 +194,20 @@ def _quantize_blocks( block_names: The names of the blocks to be quantized and dequantized. nblocks: The number of blocks to quantize and dequantize. device: The device for quantization and dequantization. + resume_state: when set and already partway through this block group + (`resume_state.resume_index > 0`), the caller has already + substituted `inputs`/`q_input` for the first not-yet-done block; + this method just needs to start its loop there instead of at + index 0, and record each block as done afterward. See + auto_round/utils/resume.py. + resume_input_ids: the exact `input_ids` value the interrupted run had + live for the first not-yet-done block (cached by + `ResumeState.mark_block_done`). `inputs` still supplies + `input_others` (legitimately re-sourced from the same pre-cache + every iteration regardless of resuming), but the chained + hidden-state tensor itself must come from here, not be re-derived + from `inputs` -- see auto_round/utils/resume.py's module + docstring for why those two aren't interchangeable. Returns: None @@ -196,11 +217,14 @@ def _quantize_blocks( m.requires_grad_(False) input_ids, input_others = self._preprocess_block_inputs(inputs) + if resume_input_ids is not None: + input_ids = resume_input_ids if pbar is None: pbar = tqdm(range(0, len(block_names), nblocks)) - for i in range(0, len(block_names), nblocks): + start_index = resume_state.resume_index if resume_state is not None and nblocks == 1 else 0 + for i in range(start_index, len(block_names), nblocks): if input_others_extra_blocks and block_names[i] in input_others_extra_blocks: input_others = input_others_extra_blocks[block_names[i]] _, input_others = self._preprocess_block_inputs(input_others) @@ -217,7 +241,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: @@ -304,6 +340,17 @@ def _quantize_blocks( if self.compress_context.is_immediate_saving: self.shard_writer.write(m, is_finalize=False) + # ShardWriter only actually flushes to disk once its + # shard-size budget is reached (`_flush_shard`, private but + # there's no public equivalent) -- `write()` above may just + # buffer this block's tensors in memory. Force a flush here + # whenever resumability is active, since marking a block + # "done" in the resume manifest is a lie if a crash before + # the next natural flush would lose its tensors entirely. + # Only pay this extra small-shard-fragmentation cost when + # AR_RESUME_DIR is actually set. + if resume_state is not None: + self.shard_writer._flush_shard() if self.compress_context.low_cpu_mem_usage and not self.compress_context.is_immediate_saving: if nblocks == 1: @@ -311,6 +358,19 @@ def _quantize_blocks( else: for name in names: self._offloader(model, name, overwrite=True) + + # Record this block as durably done (its quantized weights are + # either flushed to a shard on disk via ShardWriter, or saved to + # the offloader's temp dir) only now, after that write has + # happened -- so a crash before this point correctly re-does the + # block on resume instead of skipping it with incomplete/missing + # output. See auto_round/utils/resume.py. + if resume_state is not None and nblocks == 1: + # `input_ids` was already reassigned to `next_input_ids` + # above -- it now holds the value the *next* block should use + # as its chained hidden-state input, which is exactly what + # needs to be persisted here. + resume_state.mark_block_done(n, q_input, input_ids) if pbar is not None: pbar.update(1) @@ -550,6 +610,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: @@ -561,40 +632,155 @@ def _quantize_data_driven(self) -> tuple[torch.nn.Module, dict[str, Any]]: self.alg_composer.prepare_run() - for block_names in all_blocks: - inputs = all_inputs[block_names[0]] - all_inputs.pop(block_names[0]) - q_inputs = None - if all_q_inputs is not None: - q_inputs = all_q_inputs[block_names[0]] - all_q_inputs.pop(block_names[0]) - - inputs, q_inputs = _update_inputs(inputs, q_inputs) - - clear_memory(self.inputs) - - self._quantize_blocks( - self.model_context.model, - inputs, - block_names, - q_input=q_inputs if q_inputs is not None else None, - nblocks=self.nblocks, - pbar=pbar, - input_others_extra_blocks=all_inputs, - token_ids=input_ids_cache, + # Build one ResumeState per block group (almost always just one group + # for text-only dense models) when AR_RESUME_DIR is set, so a + # crash/kill mid-tuning can resume from the first not-yet-quantized + # block instead of restarting from block 0. See auto_round/utils/resume.py. + resume_states = None + if envs.AR_RESUME_DIR: + if not self.compress_context.is_immediate_saving and not self.compress_context.low_cpu_mem_usage: + logger.warning( + "AR_RESUME_DIR is set but neither immediate saving nor " + "low_cpu_mem_usage is active. Without low_cpu_mem_usage, " + "already-quantized blocks are never offloaded anywhere a " + "resumed process could find them (see OffloadManager's " + "deterministic resume directory in offload.py), so a " + "resumed run's in-memory model will have meta/empty " + "weights for blocks completed in a PRIOR process. Pass " + "low_cpu_mem_usage=True (or a format= to quantize_and_save) " + "for resumability to be meaningful here." + ) + from auto_round.utils.resume import ResumeState, compute_run_signature, layer_config_fingerprint + + model_dir = getattr(self.model_context, "disk_stream_model_dir", None) or getattr( + getattr(self.model_context.model, "config", None), "_name_or_path", None ) - if self.compress_context.is_immediate_packing and len(self.formats) != 1: - raise ValueError( - f"Expected exactly one packing format when 'immediate_packing' is True, " - f"but got {len(self.formats)} formats." + dataset_desc = str(getattr(self, "dataset", None)) + # str(self.scheme) alone is bits-blind for AutoScheme runs: two runs + # with different avg_bits share it, so include the resolved + # per-layer allocation (see layer_config_fingerprint docstring). + scheme_desc = ( + str(self.scheme) + + "|" + + layer_config_fingerprint(getattr(getattr(self, "quantizer", None), "layer_config", None)) + ) + resume_states = [] + for group_idx, block_names in enumerate(all_blocks): + sig = compute_run_signature( + model_dir, + scheme_desc, + dataset_desc, + self.calibration_context.nsamples, + self.calibration_context.seqlen, + block_names, + ) + resume_states.append( + ResumeState(os.path.join(envs.AR_RESUME_DIR, f"group_{group_idx}"), sig, block_names) ) - # ── Pipeline lifecycle: finalize_quantization (model-level teardown) - self.alg_composer.finalize_run() + try: + for group_idx, block_names in enumerate(all_blocks): + inputs = all_inputs[block_names[0]] + all_inputs.pop(block_names[0]) + q_inputs = None + if all_q_inputs is not None: + q_inputs = all_q_inputs[block_names[0]] + all_q_inputs.pop(block_names[0]) + + inputs, q_inputs = _update_inputs(inputs, q_inputs) + + clear_memory(self.inputs) + + resume_state = resume_states[group_idx] if resume_states is not None else None + resume_input_ids = None + if resume_state is not None and resume_state.resume_index > 0: + if self.nblocks != 1: + logger.warning( + "AR_RESUME_DIR is set but nblocks != 1; resuming mid-group is only " + "supported for nblocks=1 -- restarting this group from block 0." + ) + resume_state = None + else: + resume_name = block_names[resume_state.resume_index] + # Only used here for `input_others` (position/mask info, + # which is legitimately re-sourced from this same cache + # every iteration regardless of resuming); the actual + # chained `input_ids` comes from `resume_input_ids` + # below, not this cache -- see + # auto_round/utils/resume.py's module docstring for why + # the two aren't interchangeable. + if resume_name in all_inputs: + inputs = all_inputs.pop(resume_name) + q_inputs = resume_state.load_q_input() + resume_input_ids = resume_state.load_input_ids() + if resume_input_ids is None: + logger.warning( + "AR_RESUME_DIR manifest is missing its cached input_ids tensor; " + "restarting this group from block 0 instead of resuming with a " + "possibly-inconsistent chain value." + ) + resume_state = None + else: + pbar.update(resume_state.resume_index) + + self._quantize_blocks( + self.model_context.model, + inputs, + block_names, + q_input=q_inputs if q_inputs is not None else None, + nblocks=self.nblocks, + pbar=pbar, + input_others_extra_blocks=all_inputs, + token_ids=input_ids_cache, + resume_state=resume_state, + resume_input_ids=resume_input_ids, + ) + if self.compress_context.is_immediate_packing and len(self.formats) != 1: + raise ValueError( + f"Expected exactly one packing format when 'immediate_packing' is True, " + f"but got {len(self.formats)} formats." + ) + if resume_states is not None: + if self.compress_context.is_immediate_saving: + # Don't clear resume state yet when exporting to shards -- + # a crash in the save/export step that follows this method + # returning (config writing, tokenizer copy, format-specific + # global packing pass) would otherwise force a full + # re-tune from block 0 on the next attempt, even though + # every block's weights are already correctly flushed to + # disk. quantize_and_save() clears these once + # save_quantized() actually succeeds. + self._resume_states = resume_states + else: + for rs in resume_states: + rs.clear() + finally: + # ── Pipeline lifecycle: finalize_quantization (model-level teardown) + self.alg_composer.finalize_run() pbar.set_description("Quantizing done") pbar.close() if self.compress_context.low_cpu_mem_usage: - self._offloader.reload(self.model_context.model) + if envs.AR_RESUME_DIR and not self.compress_context.is_immediate_saving: + # `reload(names=None)` only reloads names in + # `self._offloader._saved` -- populated by THIS process's own + # offload() calls. A resumed process never touches blocks it + # skipped via ResumeState (they're left exactly as the meta + # skeleton started), so they'd never be in `_saved` and would + # stay meta in the returned model. Request every block + # explicitly so _reload()'s discovery check (see offload.py) + # gets a chance to pull each skipped block's real quantized + # weights back from a prior crashed process's offload dir. + # + # Skipped entirely under is_immediate_saving: ShardWriter has + # already flushed every block's packed weights to disk (both + # this run's and, via its own discovery, a prior crashed run's), + # and the `shard_writer.write(is_finalize=True)` call right + # below would treat any block reloaded back to real memory here + # as newly-dirty and re-emit its raw, unpacked weight tensor + # alongside the already-packed one. + self._offloader.reload(self.model_context.model, flatten_list(all_blocks)) + elif not self.compress_context.is_immediate_saving: + self._offloader.reload(self.model_context.model) self._quantize_layers_outside_blocks(layer_names, all_inputs, token_ids=input_ids_cache) convert_module_to_hp_if_necessary( diff --git a/auto_round/compressors/shard_writer.py b/auto_round/compressors/shard_writer.py index e95468fbe3..1ca0395c7a 100644 --- a/auto_round/compressors/shard_writer.py +++ b/auto_round/compressors/shard_writer.py @@ -14,11 +14,13 @@ import json import os +import re from collections import OrderedDict from typing import Optional, Union import torch +from auto_round import envs from auto_round.compressors.utils import _get_save_folder_name from auto_round.context.compress import CompressContext from auto_round.context.model import ModelContext @@ -95,8 +97,77 @@ def __init__( self.total_param_size_bytes = 0 self.skipped_meta_tensors = [] + # When resumability is active + # (AR_RESUME_DIR set), a fresh process's ShardWriter otherwise has no + # idea a previous, crashed process already flushed some shards to + # output_dir -- it would restart shard_counter at 0 and overwrite + # `model-shard-00001...`, and finalize()'s index would only cover the + # tensors written by *this* process, producing a corrupt/incomplete + # checkpoint. Gated on AR_RESUME_DIR so normal (non-resuming) runs + # never change behavior even if output_dir happens to be reused. + # + # Deliberately NOT run here in __init__: at construction time (from + # post_init(), inside quantize_and_save()) `self.output_dir` still + # reflects the pre-`_get_export_dir()` path -- the final subfolder + # (e.g. `-w4g128/`) hasn't been appended yet, so discovery + # would silently look in the wrong directory and find nothing + # (confirmed empirically: this exact ordering bug let a resumed run's + # blocks 0-2 through tuning correctly, then still lose their output, + # since `_flush_shard`/`finalize` never learned about the crashed + # run's shards). Deferred to the first real `_flush_shard()` call + # instead, by which point `output_dir` is always the final path. + self._existing_shards_discovered = False + ShardWriter._initialized = True + def _discover_existing_shards(self) -> None: + """Recover shard-writer state from shard files a previous (crashed) + process already flushed to ``output_dir``, so this process continues + shard numbering instead of colliding with them, and ``finalize()``'s + index covers tensors from both processes. + + Only files still in the pre-``finalize()`` temp naming + (``model-shard-NNNNN.``) are considered: once ``finalize()`` runs + it renames everything to the final HF layout, so a directory with no + such temp files means either nothing has been flushed yet, or a prior + run already finished -- neither should be treated as in-progress + shards to adopt. + """ + output_dir = self.output_dir + if not os.path.isdir(output_dir): + return + pattern = re.compile(rf"^model-shard-(\d+)\.{re.escape(self.shard_suffix)}$") + found = [] + for fname in os.listdir(output_dir): + m = pattern.match(fname) + if m: + found.append((int(m.group(1)), fname)) + if not found: + return + found.sort() + for _, fname in found: + path = os.path.join(output_dir, fname) + params = self._read_shard_tensor_names(path) + self.shard_meta.append({"tmp_file": fname, "params": params, "dir": output_dir}) + self._all_saved.update(params) + self.shard_counter = found[-1][0] + logger.info( + f"ShardWriter: discovered {len(found)} already-flushed shard(s) in {output_dir} " + f"from a previous run; resuming shard numbering from {self.shard_counter}." + ) + + def _read_shard_tensor_names(self, path: str) -> list[str]: + """Read only the tensor-name header of an already-flushed shard file, + without materializing any tensor data.""" + if self.use_safetensors: + from safetensors import safe_open + + with safe_open(path, framework="pt") as f: + return list(f.keys()) + else: + sd = torch.load(path, map_location="meta") + return list(sd.keys()) + @property def output_dir(self) -> str: """Derive the output directory from the current CompressContext at access time. @@ -258,6 +329,10 @@ def _flush_shard(self): if not self.current_shard_tensors: return + if envs.AR_RESUME_DIR and not self._existing_shards_discovered: + self._discover_existing_shards() + self._existing_shards_discovered = True + self.shard_counter += 1 output_dir = self.output_dir os.makedirs(output_dir, exist_ok=True) 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/export/export_to_autoround/export.py b/auto_round/export/export_to_autoround/export.py index 32f688ab1f..4b9c9e30e7 100644 --- a/auto_round/export/export_to_autoround/export.py +++ b/auto_round/export/export_to_autoround/export.py @@ -164,6 +164,18 @@ def pack_layer(layer_name, model, backend, device=None): if type(layer) not in SUPPORTED_LAYER_TYPES: ##already packed return + # A resumed disk-streamed run only + # materializes/quantizes the blocks it didn't already finish in a prior + # (crashed) process. Blocks it skipped are never touched in *this* + # process and stay on the meta device, while their packed weights + # already live in shard files the previous process flushed to disk (see + # ShardWriter._discover_existing_shards). There is nothing to pack here + # -- attempting to would fail (no real weight data to read `.scale` + # from) and would be redundant even if it didn't, since the on-disk + # export for this layer is already complete. + if layer.weight.device.type == "meta": + return + if int(layer.act_bits) <= 8: return pack_qact_layer(layer_name, model) 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/auto_round/utils/resume.py b/auto_round/utils/resume.py new file mode 100644 index 0000000000..9299cbb6c0 --- /dev/null +++ b/auto_round/utils/resume.py @@ -0,0 +1,222 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +# Checkpoints the tuning loop's per-block +# progress to disk so a crash or kill mid-run doesn't require restarting from +# block 0 -- important for large models where each block's tuning can take +# minutes and a full run spans hours. +# +# What this caches and why: AutoRound's block-sequential tuning chains two +# things forward from one block to the next -- ``input_ids`` (the current +# block's reference/FP output, used as the next block's FP reference input) +# and, when ``enable_quanted_input=True`` (SignRound's default), ``q_input`` +# (the *quantized* block's output, used as the next block's quantized-input +# companion). Both are cached here, not just ``q_input``: an earlier version of +# this patch assumed ``input_ids`` could be cheaply regenerated from the +# per-block inputs already pre-cached by `cache_inter_data`/ +# `try_cache_inter_data_gpucpu` before tuning starts (reasoning that it's a +# pure function of unmodified weights, so it shouldn't matter which code path +# computed it) -- but that pre-cache pass and the in-loop reference forward +# turned out not to be numerically identical (confirmed by a resume test +# producing a 20x-larger tuning loss on the first resumed block vs. an +# uninterrupted control run), so the actual live chain value has to be +# persisted and reloaded verbatim, the same way ``q_input`` already is. +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from pathlib import Path +from typing import Optional + +import torch + +from auto_round.logger import logger + +__all__ = ["ResumeState"] + +_MANIFEST_NAME = "resume_manifest.json" +_Q_INPUT_NAME = "resume_q_input.pt" +_INPUT_IDS_NAME = "resume_input_ids.pt" + + +def _to_cpu_recursive(obj): + """`q_input` (SignRound's `enable_quanted_input` chain value) may be a + plain Tensor or a nested list/tuple/dict of tensors depending on the + algorithm/batching path; detach + move to cpu recursively so it's both + safely picklable and reloadable regardless of which GPU produced it.""" + if isinstance(obj, torch.Tensor): + return obj.detach().to("cpu") + if isinstance(obj, dict): + return {k: _to_cpu_recursive(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_to_cpu_recursive(v) for v in obj] + if isinstance(obj, tuple): + return tuple(_to_cpu_recursive(v) for v in obj) + return obj + + +def _atomic_write_json(path: Path, data: dict) -> None: + fd, tmp_path = tempfile.mkstemp(dir=str(path.parent), prefix=".tmp_resume_") + try: + with os.fdopen(fd, "w") as f: + json.dump(data, f) + os.replace(tmp_path, path) + except Exception: + try: + os.remove(tmp_path) + except OSError: + pass + raise + + +class ResumeState: + """Tracks which blocks of a tuning run have already been quantized and + written to disk, plus the one small tensor (``q_input``) needed to resume + the sequential calibration chain correctly. Keyed by a signature over the + run's identifying configuration, so a resume directory reused for a + different model/scheme/dataset is detected and ignored rather than + silently misapplied. + """ + + def __init__(self, resume_dir: str, signature: str, block_names: list[str]): + self.dir = Path(resume_dir) + self.dir.mkdir(parents=True, exist_ok=True) + self.signature = signature + self.block_names = list(block_names) + self.manifest_path = self.dir / _MANIFEST_NAME + self.q_input_path = self.dir / _Q_INPUT_NAME + self.input_ids_path = self.dir / _INPUT_IDS_NAME + self.completed_blocks: list[str] = [] + self._load() + + def _load(self) -> None: + if not self.manifest_path.exists(): + return + try: + with open(self.manifest_path) as f: + data = json.load(f) + except Exception as e: + logger.warning(f"ResumeState: failed to read {self.manifest_path}: {e}; starting fresh.") + return + if data.get("signature") != self.signature: + logger.info( + "ResumeState: existing resume manifest is for a different run " + "(model/scheme/dataset/block list changed); ignoring it and starting fresh." + ) + return + completed = data.get("completed_blocks", []) + # Only trust a prefix of block_names in order -- anything else means the + # manifest was corrupted or hand-edited; safer to restart from scratch + # than resume from a possibly-inconsistent point. + if completed != self.block_names[: len(completed)]: + logger.warning( + "ResumeState: completed_blocks in manifest is not a prefix of the " + "current block order; ignoring it and starting fresh." + ) + return + self.completed_blocks = completed + logger.info(f"ResumeState: resuming after {len(completed)}/{len(self.block_names)} already-quantized blocks") + + @property + def resume_index(self) -> int: + """Index of the first block not yet completed.""" + return len(self.completed_blocks) + + def load_q_input(self): + return self._load_tensor(self.q_input_path) + + def load_input_ids(self): + return self._load_tensor(self.input_ids_path) + + def _load_tensor(self, path: Path): + if not self.completed_blocks or not path.exists(): + return None + try: + return torch.load(path, map_location="cpu") + except Exception as e: + logger.warning(f"ResumeState: failed to load {path.name} ({e}); resuming without it.") + return None + + def mark_block_done(self, block_name: str, q_input, input_ids) -> None: + expected = self.block_names[len(self.completed_blocks)] + assert ( + block_name == expected + ), f"ResumeState.mark_block_done called out of order: expected {expected!r}, got {block_name!r}" + if q_input is not None: + torch.save(_to_cpu_recursive(q_input), self.q_input_path) + elif self.q_input_path.exists(): + self.q_input_path.unlink() + # input_ids is required (unlike q_input, which is legitimately None + # when enable_quanted_input=False) -- the FP reference chain always + # exists. + torch.save(_to_cpu_recursive(input_ids), self.input_ids_path) + self.completed_blocks.append(block_name) + _atomic_write_json( + self.manifest_path, + {"signature": self.signature, "completed_blocks": self.completed_blocks}, + ) + + def clear(self) -> None: + """Remove the resume manifest/cache -- call after a full run completes + successfully, so a later unrelated run doesn't mistake stale state for + an in-progress one (only possible if reusing the exact same + model/scheme/dataset/block list, but still worth cleaning up).""" + for p in (self.manifest_path, self.q_input_path, self.input_ids_path): + if p.exists(): + try: + p.unlink() + except OSError: + pass + + +def compute_run_signature( + model_dir: Optional[str], + scheme_desc: str, + dataset_desc: str, + nsamples: int, + seqlen: int, + block_names: list[str], +) -> str: + """Hash the run's identifying configuration. Any change here (different + model, scheme, dataset, calibration size, or block set) must produce a + different signature so `ResumeState` refuses to reuse a stale manifest.""" + h = hashlib.sha256() + for part in (model_dir or "", scheme_desc, dataset_desc, str(nsamples), str(seqlen), "|".join(block_names)): + h.update(part.encode("utf-8")) + h.update(b"\x00") + return h.hexdigest() + + +def layer_config_fingerprint(layer_config) -> str: + """Deterministic string over the resolved per-layer quantization config. + + ``str(self.scheme)`` (and the literal ``"rtn_with_imatrix"`` used by the + imatrix-RTN path) does not capture the per-layer bit allocation that + AutoScheme resolves from ``avg_bits`` -- two runs with different avg_bits + targets would otherwise produce identical run signatures and silently + resume each other's state (observed: a quality sweep's avg_bits=4.8 + candidate resuming the 4.2 candidate's 40/40-complete manifest, saving an + output with no layer tensors at all). Fold this into ``scheme_desc`` when + calling :func:`compute_run_signature`. + + Only scalar config values (bits, group_size, sym, data_type, ...) are + included; anything non-scalar is ignored so the fingerprint stays cheap + and deterministic. + """ + if not layer_config: + return "" + parts = [] + for name in sorted(layer_config): + cfg = layer_config[name] + if isinstance(cfg, dict): + desc = ",".join( + f"{k}={v}" + for k, v in sorted(cfg.items(), key=lambda kv: str(kv[0])) + if isinstance(v, (int, float, bool, str, type(None))) + ) + else: + desc = str(cfg) + parts.append(f"{name}:{desc}") + return ";".join(parts) diff --git a/test/test_cpu/core/test_resume_integration.py b/test/test_cpu/core/test_resume_integration.py new file mode 100644 index 0000000000..6396a65aa5 --- /dev/null +++ b/test/test_cpu/core/test_resume_integration.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Integration test for AR_RESUME_DIR: simulates a crash partway through the +tuning loop and verifies a fresh AutoRound run against the same resume +directory picks up from the first not-yet-completed block instead of +restarting from block 0. +""" + +import os +from unittest import mock + +import pytest + +from auto_round import AutoRound +from auto_round.utils.resume import ResumeState + + +@pytest.fixture(autouse=True) +def _clean_resume_env(): + previous_resume_dir = os.environ.get("AR_RESUME_DIR") + previous_disk_stream = os.environ.get("AR_DISK_STREAM_MODEL") + yield + for key, previous in (("AR_RESUME_DIR", previous_resume_dir), ("AR_DISK_STREAM_MODEL", previous_disk_stream)): + if previous is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous + + +class TestResumeIntegration: + def test_resume_skips_already_completed_blocks(self, tiny_opt_model_path, tmp_path): + resume_dir = str(tmp_path / "resume") + # AR_DISK_STREAM_MODEL keeps low_cpu_mem_usage active for this plain + # dense (non-MoE) model -- without it, DataDrivenCompressor.quantize() + # silently disables low_cpu_mem_usage, and a resumed process's + # in-memory model would have meta/empty weights for blocks completed + # by the prior (crashed) process. + os.environ["AR_RESUME_DIR"] = resume_dir + os.environ["AR_DISK_STREAM_MODEL"] = "1" + + original_mark_block_done = ResumeState.mark_block_done + crashed_after = [] + + def crash_after_first_block(self, block_name, q_input, input_ids): + original_mark_block_done(self, block_name, q_input, input_ids) + crashed_after.append(block_name) + if len(crashed_after) == 1: + raise RuntimeError("simulated crash") + + with mock.patch.object(ResumeState, "mark_block_done", crash_after_first_block): + with pytest.raises(RuntimeError, match="simulated crash"): + ar = AutoRound(model=tiny_opt_model_path, scheme="W4A16", iters=1, nsamples=1) + ar.quantize() + + assert crashed_after == ["model.decoder.layers.0"] + + processed_on_resume = [] + + def track_block(self, block_name, q_input, input_ids): + original_mark_block_done(self, block_name, q_input, input_ids) + processed_on_resume.append(block_name) + + with mock.patch.object(ResumeState, "mark_block_done", track_block): + ar = AutoRound(model=tiny_opt_model_path, scheme="W4A16", iters=1, nsamples=1) + _, layer_config = ar.quantize() + + # Only the block that never completed before the crash should be + # (re-)processed -- block 0's already-durable result must be reused, + # not redone. + assert processed_on_resume == ["model.decoder.layers.1"] + + # The final layer_config still covers every quantized layer in both + # blocks, proving the resumed run produced a complete result. + quantized_layers = {name for name, cfg in layer_config.items() if "bits" in cfg} + assert any(name.startswith("model.decoder.layers.0.") for name in quantized_layers) + assert any(name.startswith("model.decoder.layers.1.") for name in quantized_layers) diff --git a/test/test_cpu/utils/test_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" diff --git a/test/test_cpu/utils/test_resume.py b/test/test_cpu/utils/test_resume.py new file mode 100644 index 0000000000..67e73054cc --- /dev/null +++ b/test/test_cpu/utils/test_resume.py @@ -0,0 +1,166 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for auto_round.utils.resume (crash/resume checkpointing, AR_RESUME_DIR).""" + +import pytest +import torch + +from auto_round.utils.resume import ( + ResumeState, + compute_run_signature, + layer_config_fingerprint, +) + + +class TestResumeState: + def test_fresh_state_has_no_completed_blocks(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + assert state.resume_index == 0 + assert state.load_q_input() is None + assert state.load_input_ids() is None + + def test_mark_block_done_advances_resume_index(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + state.mark_block_done("b0", q_input=torch.ones(2), input_ids=torch.zeros(2)) + assert state.resume_index == 1 + assert state.completed_blocks == ["b0"] + + def test_mark_block_done_out_of_order_raises(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + with pytest.raises(AssertionError): + state.mark_block_done("b1", q_input=None, input_ids=torch.zeros(2)) + + def test_q_input_and_input_ids_round_trip(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + q_input = {"a": torch.arange(4).reshape(2, 2)} + input_ids = torch.arange(6).reshape(2, 3) + state.mark_block_done("b0", q_input=q_input, input_ids=input_ids) + + # A fresh ResumeState pointed at the same dir with the same signature + # picks the saved tensors back up, exactly like a resumed process would. + resumed = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + assert resumed.resume_index == 1 + loaded_q_input = resumed.load_q_input() + assert torch.equal(loaded_q_input["a"], q_input["a"]) + assert torch.equal(resumed.load_input_ids(), input_ids) + + def test_q_input_none_clears_any_previous_q_input_file(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + state.mark_block_done("b0", q_input=torch.ones(2), input_ids=torch.zeros(2)) + assert state.q_input_path.exists() + + # Simulate a second run reusing the dir where enable_quanted_input is + # off (q_input legitimately None) -- the stale file must not linger + # and be mistaken for this block's q_input on a later resume. + state2 = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + state2.completed_blocks = [] + state2.mark_block_done("b0", q_input=None, input_ids=torch.zeros(2)) + assert not state2.q_input_path.exists() + + def test_mismatched_signature_starts_fresh(self, tmp_path): + state = ResumeState(str(tmp_path), "sig-a", ["b0", "b1"]) + state.mark_block_done("b0", q_input=None, input_ids=torch.zeros(2)) + + other = ResumeState(str(tmp_path), "sig-b", ["b0", "b1"]) + assert other.resume_index == 0 + assert other.completed_blocks == [] + + def test_completed_blocks_not_a_prefix_starts_fresh(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1", "b2"]) + state.mark_block_done("b0", q_input=None, input_ids=torch.zeros(2)) + + # Different block order for the same signature (e.g. block list changed + # some other way the signature didn't capture) -- "b0" is no longer a + # valid prefix of this new order, so the manifest must be distrusted. + other = ResumeState(str(tmp_path), "sig", ["b1", "b0", "b2"]) + assert other.resume_index == 0 + assert other.completed_blocks == [] + + def test_clear_removes_all_state(self, tmp_path): + state = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + state.mark_block_done("b0", q_input=torch.ones(2), input_ids=torch.zeros(2)) + assert state.manifest_path.exists() + + state.clear() + assert not state.manifest_path.exists() + assert not state.q_input_path.exists() + assert not state.input_ids_path.exists() + + fresh = ResumeState(str(tmp_path), "sig", ["b0", "b1"]) + assert fresh.resume_index == 0 + + def test_full_completion_round_trip(self, tmp_path): + block_names = ["b0", "b1", "b2"] + state = ResumeState(str(tmp_path), "sig", block_names) + for name in block_names: + state.mark_block_done(name, q_input=None, input_ids=torch.zeros(2)) + assert state.resume_index == len(block_names) + assert state.completed_blocks == block_names + + +class TestComputeRunSignature: + def test_identical_inputs_produce_identical_signature(self): + sig1 = compute_run_signature("m", "scheme", "dataset", 8, 2048, ["b0", "b1"]) + sig2 = compute_run_signature("m", "scheme", "dataset", 8, 2048, ["b0", "b1"]) + assert sig1 == sig2 + + @pytest.mark.parametrize( + "kwargs", + [ + {"model_dir": "other"}, + {"scheme_desc": "other"}, + {"dataset_desc": "other"}, + {"nsamples": 16}, + {"seqlen": 4096}, + {"block_names": ["b0", "b1", "b2"]}, + ], + ) + def test_any_changed_field_changes_signature(self, kwargs): + base = dict( + model_dir="m", + scheme_desc="scheme", + dataset_desc="dataset", + nsamples=8, + seqlen=2048, + block_names=["b0", "b1"], + ) + sig1 = compute_run_signature(**base) + sig2 = compute_run_signature(**{**base, **kwargs}) + assert sig1 != sig2 + + def test_none_model_dir_does_not_crash(self): + sig = compute_run_signature(None, "scheme", "dataset", 8, 2048, ["b0"]) + assert isinstance(sig, str) and len(sig) == 64 # sha256 hexdigest + + +class TestLayerConfigFingerprint: + def test_empty_layer_config(self): + assert layer_config_fingerprint(None) == "" + assert layer_config_fingerprint({}) == "" + + def test_deterministic_regardless_of_dict_order(self): + cfg_a = {"layer.1": {"bits": 4, "sym": True}, "layer.0": {"bits": 8, "sym": False}} + cfg_b = {"layer.0": {"bits": 8, "sym": False}, "layer.1": {"bits": 4, "sym": True}} + assert layer_config_fingerprint(cfg_a) == layer_config_fingerprint(cfg_b) + + def test_different_bits_produce_different_fingerprint(self): + cfg_4bit = {"layer.0": {"bits": 4}} + cfg_8bit = {"layer.0": {"bits": 8}} + assert layer_config_fingerprint(cfg_4bit) != layer_config_fingerprint(cfg_8bit) + + def test_non_scalar_values_are_ignored(self): + cfg_with_tensor = {"layer.0": {"bits": 4, "weight": torch.ones(2, 2)}} + cfg_without_tensor = {"layer.0": {"bits": 4}} + assert layer_config_fingerprint(cfg_with_tensor) == layer_config_fingerprint(cfg_without_tensor)