diff --git a/onecomp/quantized_model_loader.py b/onecomp/quantized_model_loader.py index 8b05e5f..422e67c 100644 --- a/onecomp/quantized_model_loader.py +++ b/onecomp/quantized_model_loader.py @@ -97,11 +97,18 @@ def load_quantized_model( # model.language_model.layers.* directly. state_dict = cls._remap_state_dict_keys(state_dict, model) - # Replace quantized layers with empty modules - cls._replace_quantized_layers(model, state_dict, quant_config) - - # Load all weights (quantized + non-quantized) in one go - model.load_state_dict(state_dict, strict=False, assign=True) + # Replace quantized layers with empty modules and align quantized + # tensor keys with the actual module names in the model built from + # config. This is required when the saved checkpoint and the + # from_config model use different wrapper prefixes, e.g. + # model.language_model.layers.* vs model.layers.*. + state_dict = cls._replace_quantized_layers(model, state_dict, quant_config) + + # Load all weights (quantized + non-quantized) in one go. strict=False + # is intentional because some wrapper-only components may be absent, but + # critical language-model and quantized-buffer mismatches must fail fast. + incompat = model.load_state_dict(state_dict, strict=False, assign=True) + cls._check_load_state_dict_result(incompat) # ``assign=True`` swaps Parameter objects in place, which breaks the # weight sharing established by ``from_config`` for models with @@ -148,6 +155,8 @@ def load_quantized_model( converted, ) + cls._assert_quantized_modules_loaded(model) + cls._load_generation_config(model, save_directory) # Register Hadamard hooks for rotation-preprocessed models @@ -546,29 +555,61 @@ def _remap_state_dict_keys(cls, state_dict: dict, model: torch.nn.Module) -> dic return remapped @staticmethod - def _apply_known_state_dict_key_rewrites(ckpt_key: str) -> Optional[str]: - """Return a rewritten key for known save/load prefix drift, else None.""" + def _known_state_dict_key_rewrite_candidates(ckpt_key: str) -> List[str]: + """Return candidate keys for known save/load prefix drift.""" + candidates: List[str] = [] + + # Gemma-like: + # model.language_model.model.layers.* -> model.language_model.layers.* if ".language_model.model." in ckpt_key: - return ckpt_key.replace(".language_model.model.", ".language_model.", 1) + candidates.append( + ckpt_key.replace(".language_model.model.", ".language_model.", 1) + ) + + # Composite wrapper -> text-only: + # model.language_model.layers.* -> model.layers.* + if ckpt_key.startswith("model.language_model."): + candidates.append("model." + ckpt_key[len("model.language_model.") :]) + if ckpt_key.startswith("language_model.model."): - return "model." + ckpt_key.replace("language_model.model.", "language_model.", 1) - return None + candidates.append("model." + ckpt_key[len("language_model.model.") :]) + + if ckpt_key.startswith("language_model."): + candidates.append("model." + ckpt_key[len("language_model.") :]) + + return list(dict.fromkeys(candidates)) + + @staticmethod + def _apply_known_state_dict_key_rewrites(ckpt_key: str) -> Optional[str]: + """Return a rewritten key for known save/load prefix drift, else None.""" + candidates = QuantizedModelLoader._known_state_dict_key_rewrite_candidates( + ckpt_key + ) + return candidates[0] if candidates else None @staticmethod def _resolve_state_dict_key(ckpt_key: str, model_keys: set) -> Optional[str]: - """Return the remapped key for ckpt_key, or None if unknown.""" - rewritten = QuantizedModelLoader._apply_known_state_dict_key_rewrites(ckpt_key) - if rewritten is not None: - return rewritten - - # Suffix fallback for other prefix drift: only when the suffix maps - # uniquely onto the current model (non-quantized params/buffers). - match = re.search(r"(layers\.\d+(?:\..+)*)$", ckpt_key) - if match: - suffix = match.group(1) + """Return the remapped key for ckpt_key, or None if unknown. + + Important: + Only return a candidate if it exists in model_keys. Quantized + buffers do not exist before _replace_quantized_layers(), so they + are handled in _replace_quantized_layers() instead. + """ + for candidate in QuantizedModelLoader._known_state_dict_key_rewrite_candidates( + ckpt_key + ): + if candidate in model_keys: + return candidate + + # Generic unique suffix fallback for non-quantized params/buffers. + parts = ckpt_key.split(".") + for start in range(max(0, len(parts) - 8), len(parts)): + suffix = ".".join(parts[start:]) hits = [name for name in model_keys if name.endswith(suffix)] if len(hits) == 1: return hits[0] + return None @staticmethod @@ -589,6 +630,183 @@ def _load_state_dict_from_dir(directory: str) -> dict: ) return state_dict + @staticmethod + def _flatten_module_names(module_list) -> List[str]: + """Flatten modules_in_block_to_quantize. + + Supports both: + ["model.layers.0.mlp.up_proj", ...] + and nested forms: + [["...q_proj", "...k_proj"], ["...up_proj"]] + """ + names: List[str] = [] + + def rec(x): + if isinstance(x, str): + names.append(x) + elif isinstance(x, (list, tuple)): + for y in x: + rec(y) + + rec(module_list) + return names + + @staticmethod + def _resolve_module_name( + name: str, + name_to_module: Dict[str, torch.nn.Module], + ) -> Optional[str]: + """Resolve a quantized module name to the actual model module path.""" + if name in name_to_module: + return name + + match = re.search(r"(layers\.\d+\..+)$", name) + if not match: + return None + + suffix = match.group(1) + hits = [n for n in name_to_module if n.endswith(suffix)] + + if len(hits) == 1: + return hits[0] + + if len(hits) > 1: + lang_hits = [h for h in hits if "language_model" in h] + if len(lang_hits) == 1: + return lang_hits[0] + raise RuntimeError( + f"Ambiguous module suffix match for {name}: {hits[:20]}" + ) + + return None + + @staticmethod + def _build_state_dict_prefix_map(state_dict: dict) -> Dict[str, List[str]]: + """Build prefix -> full keys map. + + Example: + model.layers.0.mlp.up_proj.qweight + -> prefix: model.layers.0.mlp.up_proj + """ + prefix_map: Dict[str, List[str]] = {} + + for key in state_dict: + prefix, sep, _field = key.rpartition(".") + if not sep: + continue + prefix_map.setdefault(prefix, []).append(key) + + return prefix_map + + @staticmethod + def _find_layer_state( + target_name: str, + state_dict: dict, + sd_prefix_map: Dict[str, List[str]], + ) -> Tuple[dict, Optional[str]]: + """Find tensors belonging to a quantized layer. + + Returns: + (layer_sd, source_prefix) + + layer_sd is field-name based: + {"qweight": tensor, "scales": tensor, ...} + {"scaling0": tensor, "bp": tensor, ...} + + This is intentionally quantizer-agnostic. + """ + exact_keys = sd_prefix_map.get(target_name) + if exact_keys: + return ( + { + k[len(target_name) + 1 :] : state_dict[k] + for k in exact_keys + }, + target_name, + ) + + match = re.search(r"(layers\.\d+\..+)$", target_name) + if not match: + return {}, None + + suffix = match.group(1) + hits = [prefix for prefix in sd_prefix_map if prefix.endswith(suffix)] + + if len(hits) == 1: + source_prefix = hits[0] + return ( + { + k[len(source_prefix) + 1 :] : state_dict[k] + for k in sd_prefix_map[source_prefix] + }, + source_prefix, + ) + + if len(hits) > 1: + lang_hits = [h for h in hits if "language_model" in h] + if len(lang_hits) == 1: + source_prefix = lang_hits[0] + return ( + { + k[len(source_prefix) + 1 :] : state_dict[k] + for k in sd_prefix_map[source_prefix] + }, + source_prefix, + ) + + raise RuntimeError( + f"Ambiguous state_dict prefix for {target_name}, " + f"suffix={suffix}: {hits[:20]}" + ) + + return {}, None + + @staticmethod + def _materialize_layer_state_dict( + state_dict: dict, + *, + source_prefix: str, + target_prefix: str, + layer_sd: dict, + ) -> dict: + """Move one layer's tensors from source_prefix to target_prefix. + + This is the key generic fix. + + GPTQ: + source.qweight -> target.qweight + source.scales -> target.scales + + DBF: + source.scaling0 -> target.scaling0 + source.bp -> target.bp + + Future quantizers: + source. -> target. + """ + if not layer_sd: + raise RuntimeError(f"No layer state found for {target_prefix}") + + if source_prefix == target_prefix: + return state_dict + + for field, tensor in layer_sd.items(): + source_key = f"{source_prefix}.{field}" + target_key = f"{target_prefix}.{field}" + + if target_key in state_dict and target_key != source_key: + raise RuntimeError( + "State dict key collision while remapping quantized layer: " + f"{source_key} -> {target_key}" + ) + + state_dict[target_key] = tensor + + if source_key in state_dict and source_key != target_key: + del state_dict[source_key] + + return state_dict + @staticmethod def _resolve_name_by_layer_suffix( name: str, @@ -646,13 +864,12 @@ def _resolve_name_by_layer_suffix( return hits[0] if hits else None @staticmethod - def _replace_quantized_layers(model, state_dict: dict, quant_config: dict): - """Replace ``nn.Linear`` with empty quantized modules for layers in config. + def _replace_quantized_layers(model, state_dict: dict, quant_config: dict) -> dict: + """Replace ``nn.Linear`` with empty quantized modules. - *quant_config* must contain ``modules_in_block_to_quantize`` (list of layer - names). Modules are created with zero buffers of the right shape; *state_dict* - is left unchanged so the caller can ``load_state_dict(state_dict)`` once to - fill all weights. + In addition, materialize quantized tensor keys from checkpoint source + prefixes to actual model module prefixes. This avoids GPTQ/DBF/OneBit + buffers staying all-zero when config and checkpoint prefixes differ. """ quant_method = quant_config["quant_method"] # mixed_* use the same tensor format as the base method (e.g. mixed_gptq -> gptq) @@ -684,12 +901,14 @@ def _replace_quantized_layers(model, state_dict: dict, quant_config: dict): ) module_list = quant_config["modules_in_block_to_quantize"] if not module_list: - return # nothing to replace + return state_dict + flat_module_list = QuantizedModelLoader._flatten_module_names(module_list) quantization_bits_list = quant_config.get("quantization_bits") or [] if quant_method and quant_method.startswith("mixed_") and quantization_bits_list: - # Build from quantization_bits; use module_list[0] to infer layer name prefix - first_name = module_list[0] + # Build from quantization_bits; use the first module name to infer + # the layer prefix, while still supporting nested module_list forms. + first_name = flat_module_list[0] if flat_module_list else "model.layers.0" prefix_match = re.match(r"^(.+\.layers)\.\d+\.", first_name) prefix = prefix_match.group(1) if prefix_match else "model.layers" quantized_names = sorted( @@ -699,61 +918,76 @@ def _replace_quantized_layers(model, state_dict: dict, quant_config: dict): for suffix in layer_cfg ) else: - quantized_names = sorted(module_list) + quantized_names = sorted(flat_module_list) name_to_module = dict(model.named_modules()) + sd_prefix_map = QuantizedModelLoader._build_state_dict_prefix_map(state_dict) - # For VLMs with tied/shared submodules (e.g. Gemma3), the - # named_modules() path may differ from the state_dict key prefix. - # Build a suffix -> state_dict prefix map to handle this. - sd_prefix_map: dict[str, str] = {} - for key in state_dict: - parts = key.rsplit(".", 1) - if len(parts) == 2: - sd_prefix_map.setdefault(parts[0], parts[0]) - - def _get_layer_sd(name: str) -> dict: - prefix = name + "." - result = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} - if result: - return result - # Fallback: match by layer suffix (e.g. "layers.0.self_attn.q_proj") - alt_name = QuantizedModelLoader._resolve_name_by_layer_suffix(name, sd_prefix_map) - if alt_name: - alt_prefix = alt_name + "." - return { - k[len(alt_prefix) :]: v - for k, v in state_dict.items() - if k.startswith(alt_prefix) - } - return {} - - for name in quantized_names: - if name not in name_to_module: + replaced = 0 + missing_modules = [] + missing_states = [] + + for saved_name in quantized_names: + target_name = QuantizedModelLoader._resolve_module_name( + saved_name, + name_to_module, + ) + + if target_name is None: + missing_modules.append(saved_name) continue - layer_sd = _get_layer_sd(name) + layer_sd, source_prefix = QuantizedModelLoader._find_layer_state( + target_name, + state_dict, + sd_prefix_map, + ) + + if not layer_sd: + layer_sd, source_prefix = QuantizedModelLoader._find_layer_state( + saved_name, + state_dict, + sd_prefix_map, + ) - linear = name_to_module[name] + if not layer_sd or source_prefix is None: + missing_states.append((saved_name, target_name)) + continue + + state_dict = QuantizedModelLoader._materialize_layer_state_dict( + state_dict, + source_prefix=source_prefix, + target_prefix=target_name, + layer_sd=layer_sd, + ) + + linear = name_to_module[target_name] in_features, out_features = linear.in_features, linear.out_features if effective_method == "gptq": - layer_wbits = resolve_gptq_layer_wbits(name, quant_config) - layer_groupsize = resolve_gptq_layer_group_size(name, quant_config) + layer_wbits = resolve_gptq_layer_wbits(saved_name, quant_config) + layer_groupsize = resolve_gptq_layer_group_size(saved_name, quant_config) quantized_module = GPTQLinear.from_saved_state( layer_sd, in_features=in_features, out_features=out_features, wbits=layer_wbits, groupsize=layer_groupsize, - actorder=get_quant_param(quant_config, "desc_act", "actorder", default=False), + actorder=get_quant_param( + quant_config, + "desc_act", + "actorder", + default=False, + ), empty=True, checkpoint_format=get_quant_param( - quant_config, "checkpoint_format", default="gptq" + quant_config, + "checkpoint_format", + default="gptq", ), ) elif effective_method == "dbf": - layer_target_bits = resolve_dbf_layer_bits(name, quant_config) + layer_target_bits = resolve_dbf_layer_bits(saved_name, quant_config) quantized_module = DoubleBinaryLinear.from_saved_state( layer_sd, in_features=in_features, @@ -773,7 +1007,124 @@ def _get_layer_sd(name: str) -> dict: f"Unknown quant_method: {quant_method} (effective: {effective_method})" ) - QuantizedModelLoader._set_module_by_name(model, name, quantized_module) + QuantizedModelLoader._set_module_by_name(model, target_name, quantized_module) + replaced += 1 + + if missing_modules or missing_states: + raise RuntimeError( + "Failed to replace/load all quantized layers.\n" + f"expected={len(quantized_names)}, replaced={replaced}\n" + f"missing_modules={missing_modules[:50]}\n" + f"missing_states={missing_states[:50]}" + ) + + logger.info( + "Replaced %d %s quantized layer(s)", + replaced, + effective_method, + ) + + return state_dict + + @staticmethod + def _check_load_state_dict_result(incompat) -> None: + """Raise if critical keys were not loaded.""" + missing = list(getattr(incompat, "missing_keys", [])) + unexpected = list(getattr(incompat, "unexpected_keys", [])) + + critical_patterns = ( + "embed_tokens", + "lm_head", + ".qweight", + ".qzeros", + ".scales", + ".g_idx", + ".scaling0", + ".bp", + ) + + critical_missing = [ + k + for k in missing + if ( + k.endswith("norm.weight") + or any(p in k for p in critical_patterns) + ) + ] + + critical_unexpected = [ + k + for k in unexpected + if ( + k.endswith("norm.weight") + or any(p in k for p in critical_patterns) + ) + ] + + if critical_missing or critical_unexpected: + raise RuntimeError( + "Critical state_dict mismatch after quantized model loading.\n" + f"critical_missing={len(critical_missing)}\n" + + "\n".join(f" MISSING: {k}" for k in critical_missing[:80]) + + "\n" + f"critical_unexpected={len(critical_unexpected)}\n" + + "\n".join(f" UNEXPECTED: {k}" for k in critical_unexpected[:80]) + ) + + if missing: + logger.warning("Non-critical missing keys: %d", len(missing)) + for k in missing[:20]: + logger.warning(" missing: %s", k) + + if unexpected: + logger.warning("Non-critical unexpected keys: %d", len(unexpected)) + for k in unexpected[:20]: + logger.warning(" unexpected: %s", k) + + @staticmethod + def _assert_quantized_modules_loaded(model: torch.nn.Module) -> None: + """Detect all-zero or invalid quantized buffers after loading.""" + bad = [] + + for name, module in model.named_modules(): + cls_name = module.__class__.__name__ + + if cls_name == "GPTQLinear": + required_attrs = ["qweight", "qzeros", "scales", "g_idx"] + nonzero_attrs = {"qweight", "scales"} + elif cls_name == "DoubleBinaryLinear": + required_attrs = ["scaling0", "bp"] + nonzero_attrs = {"scaling0", "bp"} + else: + continue + + for attr in required_attrs: + if not hasattr(module, attr): + bad.append((name, attr, "missing")) + continue + + tensor = getattr(module, attr) + + if not isinstance(tensor, torch.Tensor): + bad.append((name, attr, "not_tensor")) + continue + + if tensor.numel() == 0: + bad.append((name, attr, "empty")) + continue + + if not torch.isfinite(tensor.detach().float()).all().item(): + bad.append((name, attr, "non_finite")) + continue + + if attr in nonzero_attrs and torch.count_nonzero(tensor.detach()).item() == 0: + bad.append((name, attr, "all_zero")) + + if bad: + raise RuntimeError( + f"Invalid quantized module buffers detected: {len(bad)}\n" + + "\n".join(f" {x}" for x in bad[:80]) + ) @staticmethod def _apply_lora_adapters_from_sidecar(model, save_directory: str) -> int: diff --git a/onecomp/runner.py b/onecomp/runner.py index 05d02ab..8fa4c36 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -8,24 +8,22 @@ # pylint: disable=too-many-arguments, too-many-positional-arguments import copy +import math import gc import json -import math import os +from typing import Optional, Literal import time from logging import getLogger from pathlib import Path -from typing import Optional import torch -import torch.nn as nn from .__version__ import __version__ from .calibration import CalibrationConfig, prepare_calibration_dataset -from .log import setup_logger -from .lpcd import LPCDConfig from .model_config import ModelConfig from .qep import QEPConfig +from .lpcd import LPCDConfig from .quantizer import GPTQ, Quantizer from .quantizer.autobit import AssignmentStrategy, AutoBitQuantizer from .quantizer.autobit.dbf_fallback import MPS_DBF_FALLBACK_ERROR @@ -35,6 +33,7 @@ from .utils.device import is_mps_device from .utils.lora import LORA_ADAPTER_SUBDIR from .utils.quantization_progress import QuantizationProgressTracker +from .log import setup_logger class Runner: @@ -134,17 +133,9 @@ def __init__( post_processes (list[PostQuantizationProcess] or None): Optional list of post-quantization processes to execute after the main quantization step. Each process receives - a packed quantized model on CPU (built via - ``create_quantized_model(pack_weights=True, use_gemlite=False)``) - and may modify it in-place. Processes preserve the - incoming pack state, so the final ``self.quantized_model`` - remains packed in the production path. Processes are - executed in order. Default is None. - report_progress (bool): - When ``True`` (default), emit ``[progress]`` log lines with - completed steps, elapsed time, and a linear ETA estimate - during long quantization (calibration, chunked, multi-GPU, - QEP). Set to ``False`` for quiet runs (e.g. CI). + a quantized model on CPU (built via + ``create_quantized_model``) and may modify it in-place. + Processes are executed in order. Default is None. Note: For zero-config quantization (VRAM auto-estimation + @@ -311,31 +302,20 @@ def check(self): raise ValueError( f"Quantizer '{type(self.quantizer).__name__}' " f"(or one of its candidate quantizers) does not support " - f"QEP (Quantization Error Propagation). " + f"QEP (Quantization Error Propagation).\n" f"Set qep=False, or use a QEP-compatible quantizer " f"(e.g., GPTQ, DBF, AutoBitQuantizer with " f"QEP-compatible candidates)." ) - # Cross-validate calibration_dataset when AutoBitQuantizer is used - quantizer = self.quantizer or (self.quantizers[0] if self.quantizers else None) - if isinstance(quantizer, AutoBitQuantizer) and quantizer.calibration_config is not None: - runner_ds = self.calibration_config.calibration_dataset - quantizer_ds = quantizer.calibration_config.calibration_dataset - if runner_ds != quantizer_ds: - raise ValueError( - f"Calibration dataset mismatch: Runner uses " - f"{runner_ds!r} but quantizer uses {quantizer_ds!r}. " - f"Set the same calibration_dataset in both " - f"CalibrationConfig objects." - ) - - # MPS device validation: only GPTQ (or AutoBitQuantizer whose - # candidates are all GPTQ, without DBF fallback) is supported on MPS + # MPS device validation: only GPTQ is supported on MPS. + # AutoBitQuantizer is allowed only when all candidates are GPTQ and + # DBF fallback cannot be selected implicitly. device = self.model_config.device if is_mps_device(device): if self.multi_gpu: raise ValueError("multi_gpu is not supported on MPS device.") + all_quantizers = self.quantizers if self.quantizers is not None else [self.quantizer] for i, q in enumerate(all_quantizers): label = f"quantizers[{i}]" if self.quantizers else "quantizer" @@ -344,7 +324,7 @@ def check(self): if not isinstance(cand, GPTQ): raise ValueError( f"{label}.quantizers[{j}] ({type(cand).__name__}) " - f"is not supported on MPS device. " + f"is not supported on MPS device.\n" "AutoBitQuantizer on MPS requires all candidate " "quantizers to be GPTQ." ) @@ -360,6 +340,19 @@ def check(self): "Only GPTQ quantization supports MPS." ) + # Cross-validate calibration_dataset when AutoBitQuantizer is used + quantizer = self.quantizer or (self.quantizers[0] if self.quantizers else None) + if isinstance(quantizer, AutoBitQuantizer) and quantizer.calibration_config is not None: + runner_ds = self.calibration_config.calibration_dataset + quantizer_ds = quantizer.calibration_config.calibration_dataset + if runner_ds != quantizer_ds: + raise ValueError( + f"Calibration dataset mismatch: Runner uses " + f"{runner_ds!r} but quantizer uses {quantizer_ds!r}. " + f"Set the same calibration_dataset in both " + f"CalibrationConfig objects." + ) + def _exclude_moe_router_if_needed(self): """Exclude MoE router layers from quantization. @@ -369,19 +362,27 @@ def _exclude_moe_router_if_needed(self): config = self.model_config.load_config() num_experts = ( getattr(config, "num_experts", 0) - or getattr(getattr(config, "text_config", None), "num_experts", 0) - or 0 + or getattr( + getattr(config, "text_config", None), "num_experts", 0 + ) or + 0 ) if num_experts == 0: return keyword = "router" - target_quantizers = self.quantizers if self.quantizers is not None else [self.quantizer] + target_quantizers = ( + self.quantizers + if self.quantizers is not None + else [self.quantizer] + ) for q in target_quantizers: if q.exclude_layer_keywords is None: q.exclude_layer_keywords = [keyword] elif keyword not in q.exclude_layer_keywords: - q.exclude_layer_keywords = list(q.exclude_layer_keywords) + [keyword] + q.exclude_layer_keywords = list(q.exclude_layer_keywords) + [ + keyword + ] self.logger.info( "MoE model (num_experts=%d): excluding '%s' layers from " @@ -558,9 +559,12 @@ def auto_run( uniform_bit = max(valid_wbits) if save_dir == "auto": model_name = model_id.rstrip("/").split("/")[-1] - save_dir = f"{model_name}-gptq-{uniform_bit}bit" + save_dir = ( + f"{model_name}-gptq-{uniform_bit}bit" + ) logger.warning( - "Gemma 4 detected → falling back to uniform GPTQ %d-bit " "(target wbits=%.2f)", + "Gemma 4 detected → falling back to uniform GPTQ %d-bit " + "(target wbits=%.2f)", uniform_bit, wbits, ) @@ -571,7 +575,6 @@ def auto_run( save_dir = f"{model_name}-autobit-{wbits}bit" from .quantizer.autobit import AutoBitQuantizer - candidate_quantizers = [ GPTQ(wbits=b, groupsize=groupsize, **kwargs) for b in candidate_bits ] @@ -766,6 +769,7 @@ def quantize_without_calibration(self): len(self.quantizer.module_to_name), "Quantization without calibration (layers)", ) + for module in self.quantizer.module_to_name.keys(): self.quantizer.quantize(module, None, None) if progress: @@ -900,12 +904,8 @@ def quantize_with_jointq_error_propagation( def run_post_processes(self): """Execute post-quantization processes. - Builds a packed quantized model on CPU from ``quantizer.results`` - and passes it to each :class:`PostQuantizationProcess` in order. Each - process preserves the incoming pack state (packed-in -> packed-out), so - ``self.quantized_model`` stays packed for memory-efficient eval reuse. - Processes that train (e.g. :class:`PostProcessLoraSFT`) unpack the base - weights internally for the duration of training and re-pack on exit. + Builds a quantized model on CPU from ``quantizer.results`` and + passes it to each :class:`PostQuantizationProcess` in order. Raises: ValueError: If ``self.quantizer`` is ``None`` @@ -920,8 +920,6 @@ def run_post_processes(self): ) logger.info("Building quantized model for post-quantization processes...") - # Pass a packed model to post-processes; each one keeps the incoming - # pack state (packed-in -> packed-out) so the result stays packed. # use_gemlite=False: GemLite uses fp16-only Triton kernels that break when # LoRA SFT runs with bfloat16 autocast. Plain buffers (qweight/scales) are # needed so training can call base_layer.forward() without dtype mismatch. @@ -944,7 +942,7 @@ def prepare_calibration_dataset(self, device, model=None): Args: device (torch.device): Device to place tensors on (CPU or GPU) - model: Model instance (optional). Add model-specific fields + model: Model instance (optional). Add model-specific fields (e.g. mm_token_type_ids for Gemma 4). Returns: @@ -1068,7 +1066,10 @@ def save_quantization_statistics(self, path: str, quantizer=None): logger.info("Saving the quantization statistics to %s", path) - statistics = {key: result.get_statistics() for key, result in quantizer.results.items()} + statistics = { + key: result.get_statistics() + for key, result in quantizer.results.items() + } with open(path, "w", encoding="utf-8") as f: json.dump(statistics, f, indent=4) @@ -1677,16 +1678,9 @@ def create_quantized_model(self, pack_weights: bool = True, quantizer=None, use_ With post-process: - >>> model, tokenizer = runner.create_quantized_model( - ... pack_weights=True, - ... use_gemlite=False, - ... ) + >>> model, tokenizer = runner.create_quantized_model(pack_weights=False) >>> post_process = PostProcessLoraSFT(data_files="train.jsonl") >>> post_process.run(model, runner.model_config) - - Post-processes preserve the incoming pack state. Use - ``pack_weights=False`` only when intentionally debugging an - unpacked-buffer path. """ if quantizer is None: quantizer = self.quantizer @@ -1732,8 +1726,23 @@ def create_quantized_model(self, pack_weights: bool = True, quantizer=None, use_ fp32_had, ) - # Build modules_in_block_to_quantize from actually-quantized layer names. - quantized_names = sorted(quantizer.results.keys()) + # Build modules_in_block_to_quantize from the actual model that will + # be saved, not from quantizer.results.keys(). + # + # quantizer.results.keys() reflects the names seen during quantization. + # For wrapper/composite models, the actual save-time model may expose + # different module names, e.g. + # + # quantization-time: model.layers.0.... + # save-time: model.language_model.layers.0.... + # + # The saved quantization_config must match save-time module names. + quantized_names = self._collect_quantized_module_names(model) + + if not quantized_names: + # Fallback for future quantizers that do not replace modules. + quantized_names = sorted(quantizer.results.keys()) + modules_in_block = list(quantized_names) quant_config["modules_in_block_to_quantize"] = modules_in_block quant_config["quantized_layer_names"] = modules_in_block @@ -1755,10 +1764,15 @@ def create_quantized_model(self, pack_weights: bool = True, quantizer=None, use_ # cf) https://docs.vllm.ai/en/stable/features/quantization/#implementing-a-quantized-moe-method num_experts = ( getattr(model.config, "num_experts", None) - or getattr(getattr(model.config, "text_config", None), "num_experts", None) + or getattr( + getattr(model.config, "text_config", None), "num_experts", None + ) or 0 ) - if quant_config.get("quant_method") == "gptq" and num_experts > 0: + if ( + quant_config.get("quant_method") == "gptq" + and num_experts > 0 + ): quant_config["quant_method"] = "mixed_gptq" self.logger.info( "MoE model detected (num_experts=%d): " @@ -1782,13 +1796,15 @@ def _patch_k_eq_v_for_vllm(self, model, quant_config: dict) -> None: Gemma4 full-attention layers with attention_k_eq_v=True have no v_proj weight — the model reuses key states as value states. vLLM fuses q/k/v into a single qkv_proj and requires all shards - to share the same quantization status. + to share the same quantization status. """ text_cfg = getattr(model.config, "text_config", None) if text_cfg is None or not getattr(text_cfg, "attention_k_eq_v", False): return layer_types = getattr(text_cfg, "layer_types", []) - k_eq_v_indices = {i for i, lt in enumerate(layer_types) if lt == "full_attention"} + k_eq_v_indices = { + i for i, lt in enumerate(layer_types) if lt == "full_attention" + } if not k_eq_v_indices: return @@ -1826,7 +1842,9 @@ def _patch_k_eq_v_for_vllm(self, model, quant_config: dict) -> None: and "self_attn.k_proj" in layer_cfg and "self_attn.v_proj" not in layer_cfg ): - layer_cfg["self_attn.v_proj"] = copy.deepcopy(layer_cfg["self_attn.k_proj"]) + layer_cfg["self_attn.v_proj"] = copy.deepcopy( + layer_cfg["self_attn.k_proj"] + ) for key in ("modules_in_block_to_quantize", "quantized_layer_names"): names = quant_config.get(key, []) @@ -1840,9 +1858,10 @@ def _patch_k_eq_v_for_vllm(self, model, quant_config: dict) -> None: quant_config[key] = sorted(names + added) # ======================================== - # Unified Save/Load Methods + # Unified Save/Load Methods (Using quantizer.results) # ======================================== + @staticmethod def _packable_gptq_wbits(wbits: int) -> bool: """Return whether OneComp can export GPTQ tensors in packed format.""" @@ -1853,7 +1872,6 @@ def _packable_gptq_wbits(wbits: int) -> bool: @staticmethod def _collect_lora_gptq_modules(model) -> list[tuple[str, torch.nn.Module]]: """Return ``LoRAGPTQLinear`` modules contained in *model*.""" - # Avoid importing post_process_lora_sft here; it pulls training deps. return [ (name, mod) for name, mod in model.named_modules() @@ -1870,6 +1888,7 @@ def _iter_gptq_export_layers( layers: list[tuple[str, torch.nn.Module]] = [] lora_base_names = set() + for name, mod in lora_modules: base_layer = getattr(mod, "base_layer", None) if isinstance(base_layer, GPTQLinear): @@ -1881,6 +1900,7 @@ def _iter_gptq_export_layers( continue if isinstance(mod, GPTQLinear): layers.append((name, mod)) + return layers def _build_base_quantized_state_dict( @@ -1889,28 +1909,25 @@ def _build_base_quantized_state_dict( lora_modules: list[tuple[str, torch.nn.Module]], pack_weights: bool = True, ) -> dict[str, torch.Tensor]: - """Build a base-model state_dict for HF/vLLM-compatible export. - - ``LoRAGPTQLinear`` wrappers are flattened back to the base GPTQLinear - key layout, and LoRA tensors are omitted from the base weights. If a - GPTQLinear is unpacked and the bit-width is packable, only the exported - tensors are packed; the in-memory model is left unchanged. - """ + """Build a base-model state_dict for HF/vLLM-compatible export.""" export_state_dict: dict[str, torch.Tensor] = {} lora_modules = sorted(lora_modules, key=lambda item: len(item[0]), reverse=True) for key, tensor in model.state_dict().items(): skip = False export_key = key + for lora_name, _mod in lora_modules: prefix = f"{lora_name}." if lora_name else "" if key.startswith(f"{prefix}lora_A.") or key.startswith(f"{prefix}lora_B."): skip = True break + base_prefix = f"{prefix}base_layer." if key.startswith(base_prefix): export_key = f"{prefix}{key[len(base_prefix):]}" break + if not skip: export_state_dict[export_key] = tensor @@ -1935,6 +1952,7 @@ def _build_base_quantized_state_dict( prefix = f"{layer_name}." if layer_name else "" qweight_key = f"{prefix}qweight" qzeros_key = f"{prefix}qzeros" + if qweight_key not in export_state_dict or qzeros_key not in export_state_dict: self.logger.warning( "Skipping GPTQ export packing for %s because qweight/qzeros " @@ -1965,62 +1983,24 @@ def _build_base_quantized_state_dict( len(skipped_layers), skipped_layers, ) + return export_state_dict def _save_lora_adapter_sidecar(self, save_directory: str, model=None) -> bool: - """Write a PEFT-compatible LoRA adapter sidecar if *model* - contains ``LoRAGPTQLinear`` modules (typically produced by - ``PostProcessLoraSFT``). - - The sidecar is placed in a ``lora_adapter/`` subdirectory rather than - directly in ``save_directory``. Reason: vLLM's base-model safetensors - loader globs ``*.safetensors`` at the top level of the model directory - and would otherwise try to load ``adapter_model.safetensors`` as - base-model weights, crashing with ``"no module or parameter named - 'base_model' in LlamaForCausalLM"``. Keeping the adapter under a - subdirectory avoids that collision while still keeping the whole model - self-contained under one directory tree. - - The subdirectory contains: - - ``adapter_model.safetensors`` - - ``adapter_config.json`` - - The format matches what vLLM's native PEFT LoRA loader expects, so:: - - LLM(model=save_dir, enable_lora=True) - LoRARequest(..., lora_path=os.path.join(save_dir, "lora_adapter")) - - will load and apply the adapter without any OneComp-specific changes - to the vLLM plugin. - - Returns: - bool: True iff an adapter was written. False if there is no - in-memory LoRA state to save (e.g. no post-process ran). - """ + """Write a PEFT-compatible LoRA adapter sidecar if present.""" if model is None: model = self.quantized_model if model is None: return False - # Inline imports keep runner.py import-time cheap and avoid any - # circular-import risk with the post_process package. from safetensors.torch import save_file as _st_save_file lora_modules = self._collect_lora_gptq_modules(model) if not lora_modules: return False - # Save in the base model's runtime dtype so the round-trip is a single - # fp32(train) -> base-dtype rounding. Hardcoding float16 would add a - # needless fp16 intermediate for bf16 models (fp32 -> fp16 -> bf16 in - # vLLM). This save path expects model_config.dtype to be a concrete - # "float16"/"bfloat16" that maps directly to a torch dtype; an - # unexpected value (e.g. "auto") is out of scope and intentionally - # raises via getattr rather than silently falling back. save_dtype = getattr(torch, self.model_config.dtype) - # PEFT convention: keys are prefixed with "base_model.model." and the - # module path matches what we will see on the loaded HF model. state_dict = {} for name, mod in lora_modules: state_dict[f"base_model.model.{name}.lora_A.weight"] = ( @@ -2032,7 +2012,6 @@ def _save_lora_adapter_sidecar(self, save_directory: str, model=None) -> bool: first = lora_modules[0][1] lora_r = int(first.lora_r) - # scaling = alpha / r is stored as float; round-trip back to int alpha. lora_alpha = int(round(float(first.scaling) * float(first.lora_r))) lora_dropout = ( float(first.dropout.p) if isinstance(first.dropout, torch.nn.Dropout) else 0.0 @@ -2076,18 +2055,7 @@ def _save_lora_adapter_sidecar(self, save_directory: str, model=None) -> bool: return True def _resolve_source_model_dir(self) -> Optional[str]: - """Resolve the original model directory for auxiliary file copy. - - Returns the local directory of the source model used by this runner. - If ``ModelConfig`` points at a Hugging Face Hub ID rather than a local - directory, attempts to locate the snapshot via - ``huggingface_hub.snapshot_download(local_files_only=True)``. - - Returns: - The absolute path to the source model directory, or ``None`` if it - could not be resolved (in which case auxiliary copying is skipped - and a warning is logged by the caller). - """ + """Resolve the original model directory for auxiliary file copy.""" src = self.model_config.get_model_id_or_path() if not src: return None @@ -2101,10 +2069,6 @@ def _resolve_source_model_dir(self) -> Optional[str]: self.logger.warning("Could not resolve source model dir for %s: %s", src, exc) return None - # File patterns excluded from the auxiliary-config copy in - # ``save_quantized_model``. Weight tensors are written by HF - # ``save_pretrained`` directly, so copying the originals would either - # collide with the quantized weights or balloon the save directory. _AUX_COPY_EXCLUDE_FILES = frozenset( { "config.json", @@ -2117,20 +2081,7 @@ def _resolve_source_model_dir(self) -> Optional[str]: _AUX_COPY_INCLUDE_SUFFIXES = (".json", ".jinja") def _copy_auxiliary_files(self, src_dir: str, save_directory: str) -> int: - """Copy auxiliary ``*.json`` / ``*.jinja`` files from ``src_dir``. - - Files already present in ``save_directory`` are left untouched so that - the artifacts written by ``model.save_pretrained`` / - ``tokenizer.save_pretrained`` (and the Gemma BOS post-processing) are - never overwritten. Weight tensors and weight index files are skipped. - - Args: - src_dir: Resolved original model directory. - save_directory: Destination directory. - - Returns: - Number of files actually copied. - """ + """Copy auxiliary ``*.json`` / ``*.jinja`` files from source model.""" import shutil copied = 0 @@ -2153,80 +2104,50 @@ def _copy_auxiliary_files(self, src_dir: str, save_directory: str) -> int: continue if not lower.endswith(self._AUX_COPY_INCLUDE_SUFFIXES): continue + dst = os.path.join(save_directory, name) if os.path.exists(dst): - # Don't clobber files already in the save directory. - # Typically these were just written by - # ``model.save_pretrained`` / ``tokenizer.save_pretrained`` - # (and possibly post-edited, e.g. the Gemma 4 - # ``add_bos_token`` patch on ``tokenizer_config.json``); - # we deliberately keep that fresh copy instead of - # overwriting it with the original-model file. Logging - # the skip simply makes the auxiliary-copy step easier - # to follow alongside the ``Copied %s`` entries below. self.logger.info("Using existing %s in save directory", name) continue + shutil.copy2(src, dst) copied += 1 self.logger.info("Copied %s to save directory", name) - return copied - def save_quantized_model(self, save_directory: str, pack_weights: bool = True): - """Save the quantized model to the specified directory + return copied - If ``self.quantized_model`` is available, that in-place updated model - is saved. Otherwise, the base quantized model is built from - ``quantizer.results`` via :meth:`create_quantized_model`. The result - is saved in HuggingFace-compatible safetensors format. - If the selected model contains ``LoRAGPTQLinear`` wrappers, this method - saves base weights with LoRA tensors excluded and additionally writes a - PEFT-compatible LoRA adapter sidecar - (``adapter_model.safetensors`` + ``adapter_config.json``) into the same - directory. The resulting directory can then be loaded back with - :func:`onecomp.load_quantized_model` (which auto-detects the sidecar - and re-wraps the layers) or served by vLLM via ``enable_lora=True``. + def save_quantized_model( + self, + save_directory: str, + pack_weights: bool = True, + save_format: Literal["auto", "native", "full_wrapper"] = "auto", + ): + """Save the quantized model to the specified directory Args: save_directory (str): The path to save the quantized model. pack_weights (bool): - Whether to pack quantized weights when ``self.quantized_model`` - is not set and a model is built from ``quantizer.results``. - When saving an existing ``self.quantized_model``, packable - unpacked GPTQ buffers are packed only in the export - ``state_dict`` without mutating the in-memory model. + Whether to pack quantized weights for more memory/storage-efficient + representation. Examples: Single quantizer mode: >>> runner.save_quantized_model("./quantized_model") - - GPTQ + LoRA SFT: - - >>> runner = Runner( - ... model_config=model_config, - ... quantizer=GPTQ(wbits=4, groupsize=128), - ... post_processes=[PostProcessLoraSFT(data_files="train.jsonl")], - ... ) - >>> runner.run() - >>> runner.save_quantized_model("./quantized_model_lora") """ logger = self.logger logger.info("Saving quantized model to %s", save_directory) if self.quantized_model is not None: - logger.info( - "Saving in-memory quantized model from self.quantized_model; " - "quantization results will not be used to rebuild it" - ) + logger.info("Using existing quantized model (post-process results preserved)") model = self.quantized_model tokenizer = self.model_config.load_tokenizer() else: - # Disable GemLite when saving to avoid extra params in safetensors. + # Disable GemLite when saving to avoid extra params in safetensors model, tokenizer = self.create_quantized_model( - pack_weights=pack_weights, - use_gemlite=False, + pack_weights=pack_weights, use_gemlite=False ) # Save model and tokenizer @@ -2238,14 +2159,42 @@ def save_quantized_model(self, save_directory: str, pack_weights: bool = True): needs_export_state_dict = bool(lora_modules) or any( not getattr(layer, "_weight_is_packed", False) for _name, layer in gptq_layers ) + export_state_dict = None if needs_export_state_dict: - base_state_dict = self._build_base_quantized_state_dict( + export_state_dict = self._build_base_quantized_state_dict( model, lora_modules, pack_weights=pack_weights ) - model.save_pretrained(save_directory, state_dict=base_state_dict) - else: - model.save_pretrained(save_directory) + + orig_model_config_for_restore = model.config + + try: + save_state_dict = self._prepare_model_for_quantized_save( + model, + save_format=save_format, + state_dict=export_state_dict, + ) + + if save_state_dict is None: + model.save_pretrained(save_directory) + else: + model.save_pretrained( + save_directory, + state_dict=save_state_dict, + ) + + if save_format == "full_wrapper": + config_path = Path(save_directory) / "config.json" + config_path.write_text( + json.dumps(model.config.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + finally: + if save_format == "full_wrapper": + model.config = orig_model_config_for_restore + tokenizer.save_pretrained(save_directory) + if save_format == "full_wrapper": + self._save_processor_files_if_available(save_directory) # Gemma 4 PT models require BOS token for coherent generation but the # upstream tokenizer_config.json omits add_bos_token. Ensure it is @@ -2260,38 +2209,29 @@ def save_quantized_model(self, save_directory: str, pack_weights: bool = True): logger.info("Set add_bos_token=true in tokenizer_config.json") # Copy auxiliary config / template files from the original model so the - # save directory is self-contained for VLM (image/audio) inference and - # for runtimes that expect e.g. ``preprocessor_config.json``, - # ``processor_config.json``, ``special_tokens_map.json``, - # ``chat_template.jinja`` to live next to the weights. + # save directory is self-contained for VLM inference and runtimes that + # expect e.g. processor_config.json, preprocessor_config.json, + # special_tokens_map.json, or chat_template.jinja next to weights. src_dir = self._resolve_source_model_dir() if src_dir and os.path.isdir(src_dir): self._copy_auxiliary_files(src_dir, save_directory) else: logger.warning("Source model dir not resolvable; skipping auxiliary file copy.") - # LoRA sidecar (only if the selected model contains LoRAGPTQLinear). + # LoRA sidecar: only written if selected model contains LoRAGPTQLinear. wrote_adapter = self._save_lora_adapter_sidecar(save_directory, model=model) if not wrote_adapter: - # Remove any stale sidecar from a previous run so the directory is - # self-consistent and load_quantized_model does not pick up an - # adapter that no longer matches the saved base model. stale_adapter_dir = save_path / LORA_ADAPTER_SUBDIR if stale_adapter_dir.is_dir(): - for stale in ( - "adapter_model.safetensors", - "adapter_config.json", - ): + for stale in ("adapter_model.safetensors", "adapter_config.json"): stale_path = stale_adapter_dir / stale if stale_path.exists(): stale_path.unlink() - # Remove the (now-empty) subdirectory if nothing else lives there. try: stale_adapter_dir.rmdir() except OSError: pass - # Also remove any top-level adapter files left by older versions of - # this helper (previous layout put the sidecar directly in save_dir). + for legacy in ("adapter_model.safetensors", "adapter_config.json"): legacy_path = save_path / legacy if legacy_path.exists(): @@ -2397,8 +2337,9 @@ def analyze_cumulative_error( """ # Lazy import: load submodule only when needed # pylint: disable-next=import-outside-toplevel - # pylint: disable-next=import-outside-toplevel from .analyzer.cumulative_error import analyze_cumulative_error as _analyze + + # pylint: disable-next=import-outside-toplevel from .analyzer.cumulative_error import plot_cumulative_error as _plot logger = self.logger @@ -2499,3 +2440,337 @@ def analyze_cumulative_error( json.dump(json_results, f, indent=2, ensure_ascii=False) return all_results + + def _save_processor_files_if_available(self, save_directory: str) -> None: + """Save processor / image processor files required by full-wrapper VLM checkpoints. + + vLLM loads multimodal/full-wrapper checkpoints through Transformers + processor utilities. For VLM configs, tokenizer files alone are not + enough: preprocessor_config.json is required for the image processor. + """ + model_id_or_path = self.model_config.get_model_id_or_path() + + try: + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained( + model_id_or_path, + trust_remote_code=True, + ) + processor.save_pretrained(save_directory) + self.logger.info("Saved processor files from %s", model_id_or_path) + except Exception as exc: + self.logger.warning( + "Could not save AutoProcessor files from %s: %s", + model_id_or_path, + exc, + ) + + try: + from transformers import AutoImageProcessor + + image_processor = AutoImageProcessor.from_pretrained( + model_id_or_path, + trust_remote_code=True, + ) + image_processor.save_pretrained(save_directory) + self.logger.info("Saved image processor files from %s", model_id_or_path) + except Exception as exc: + self.logger.warning( + "Could not save AutoImageProcessor files from %s: %s", + model_id_or_path, + exc, + ) + + # ------------------------------------------------------------------ + # Save namespace helpers + # ------------------------------------------------------------------ + + def _is_quantized_module(self, module) -> bool: + """Return True for OneCompression quantized inference modules. + + Keep this name-based to avoid import cycles and to remain extensible + for future quantizers. + """ + return module.__class__.__name__ in { + "GPTQLinear", + "DoubleBinaryLinear", + } + + def _collect_quantized_module_names(self, model) -> list[str]: + """Collect actual quantized module names from the model to be saved. + + This is more reliable than quantizer.results.keys() because wrapper + models may expose different names during quantization and save. + """ + return sorted( + name + for name, module in model.named_modules() + if self._is_quantized_module(module) + ) + + def _detect_weight_namespace(self, model) -> str: + """Detect whether state_dict is text-only or full-wrapper style.""" + keys = list(model.state_dict().keys()) + + if any(k.startswith("model.language_model.model.layers.") for k in keys): + return "full_language_model" + + if any(k.startswith("model.language_model.layers.") for k in keys): + return "full_language_model" + + if any(k.startswith("language_model.layers.") for k in keys): + return "full_language_model" + + if any(".language_model.layers." in k for k in keys): + return "full_language_model" + + if any(k.startswith("model.layers.") for k in keys): + return "text_only" + + return "unknown" + + def _detect_config_namespace(self, model) -> str: + """Detect whether config is text-only or full/composite style.""" + cfg = model.config + model_type = getattr(cfg, "model_type", None) + + if model_type in { + "qwen3_5_text", + "qwen3_5_moe_text", + }: + return "text_only" + + if model_type in { + "qwen3_5", + "qwen3_5_moe", + }: + return "full_language_model" + + if getattr(cfg, "text_config", None) is not None: + return "full_language_model" + + return "unknown" + + def _restore_original_composite_config_if_needed(self, model) -> None: + """Restore outer/composite config if state_dict uses wrapper prefix. + + Example bad checkpoint: + config: + model_type = qwen3_5_text + architectures = Qwen3_5ForCausalLM + + state_dict: + model.language_model.layers.0.... + + For such a model, save the original outer config instead: + model_type = qwen3_5 + architectures = Qwen3_5ForConditionalGeneration + text_config = {...} + vision_config = {...} + """ + cfg_ns = self._detect_config_namespace(model) + sd_ns = self._detect_weight_namespace(model) + + if cfg_ns != "text_only" or sd_ns != "full_language_model": + return + + orig_config = self.model_config.load_config() + + # Only restore when the original config is actually composite. + if getattr(orig_config, "text_config", None) is None: + return + + quant_config = getattr(model.config, "quantization_config", None) + + self.logger.warning( + "Detected text-only config with full language_model state_dict. " + "Restoring original composite config before save_pretrained()." + ) + + model.config = copy.deepcopy(orig_config) + + if quant_config is not None: + model.config.quantization_config = quant_config + + def _assert_config_state_dict_namespace_consistent(self, model) -> None: + """Fail before saving a checkpoint whose config and state_dict disagree.""" + cfg_ns = self._detect_config_namespace(model) + sd_ns = self._detect_weight_namespace(model) + + if cfg_ns == "unknown" or sd_ns == "unknown": + return + + if cfg_ns != sd_ns: + raise RuntimeError( + "config/state_dict namespace mismatch before save_pretrained().\n" + f" config namespace: {cfg_ns}\n" + f" state_dict namespace: {sd_ns}\n" + f" config model_type: {getattr(model.config, 'model_type', None)}\n" + "This would create a checkpoint that may load with missing or " + "zero-filled quantized buffers." + ) + + def _prepare_model_for_quantized_save( + self, + model, + *, + save_format: str, + state_dict: dict | None = None, + ) -> dict | None: + """Prepare model.config and optional state_dict for save_pretrained. + + Returns: + None: + Use model.state_dict() as-is. + + dict: + Pass this remapped state_dict to save_pretrained(). + """ + if save_format not in {"auto", "native", "full_wrapper"}: + raise ValueError( + f"Unknown save_format={save_format!r}. " + "Expected one of: auto, native, full_wrapper." + ) + + if save_format in {"auto", "native"}: + self._restore_original_composite_config_if_needed(model) + self._assert_config_state_dict_namespace_consistent(model) + self._assert_quant_config_matches_model_namespace(model) + return state_dict + + # save_format == "full_wrapper" + return self._prepare_full_wrapper_quantized_save(model, state_dict=state_dict) + + def _assert_quant_config_matches_model_namespace(self, model) -> None: + quant_config = getattr(model.config, "quantization_config", None) + if not quant_config: + return + + names = quant_config.get("modules_in_block_to_quantize") or [] + names = [n for n in names if isinstance(n, str)] + if not names: + return + + model_module_names = set(dict(model.named_modules()).keys()) + + missing = [n for n in names if n not in model_module_names] + + if missing: + raise RuntimeError( + "quantization_config contains module names that do not exist " + "in the model being saved.\n" + f"missing={missing[:50]}" + ) + + def _prepare_full_wrapper_quantized_save(self, model, state_dict: dict | None = None) -> dict: + """Prepare full-wrapper checkpoint for vLLM-like runtimes. + + This does not mutate tensor objects; it only remaps state_dict keys and + replaces model.config with the original composite config. + """ + cfg_ns = self._detect_config_namespace(model) + sd_ns = self._detect_weight_namespace(model) + + if cfg_ns == "full_language_model" and sd_ns == "full_language_model": + self._assert_config_state_dict_namespace_consistent(model) + self._assert_quant_config_matches_model_namespace(model) + return state_dict + + orig_config = self.model_config.load_config() + + if getattr(orig_config, "text_config", None) is None: + raise RuntimeError( + "save_format='full_wrapper' was requested, but the original " + "model config is not composite and has no text_config." + ) + + if cfg_ns != "text_only" or sd_ns != "text_only": + raise RuntimeError( + "save_format='full_wrapper' currently supports converting " + "consistent text-only checkpoints only.\n" + f"config namespace: {cfg_ns}\n" + f"state_dict namespace: {sd_ns}" + ) + + old_quant_config = getattr(model.config, "quantization_config", None) + if old_quant_config is None: + raise RuntimeError("model.config has no quantization_config") + + full_quant_config = self._remap_text_only_quant_config_to_full_wrapper( + old_quant_config + ) + + # Replace config with original composite config. + model.config = copy.deepcopy(orig_config) + model.config.quantization_config = full_quant_config + + # Remap tensors to match the full-wrapper config. + source_state_dict = state_dict if state_dict is not None else model.state_dict() + full_state_dict = self._remap_text_only_state_dict_to_full_wrapper( + source_state_dict + ) + + self.logger.info( + "Prepared full-wrapper quantized save: model_type=%s, first_quantized=%s", + getattr(model.config, "model_type", None), + full_quant_config.get("modules_in_block_to_quantize", [""])[0], + ) + + return full_state_dict + + @staticmethod + def _remap_text_only_quant_config_to_full_wrapper(quant_config: dict) -> dict: + quant_config = copy.deepcopy(quant_config) + + def remap_name(name: str) -> str: + if name.startswith("model.layers."): + return "model.language_model.layers." + name[len("model.layers.") :] + if name.startswith("model.embed_tokens."): + return "model.language_model.embed_tokens." + name[len("model.embed_tokens.") :] + if name.startswith("model.norm."): + return "model.language_model.norm." + name[len("model.norm.") :] + return name + + for key in ("modules_in_block_to_quantize", "quantized_layer_names"): + names = quant_config.get(key) + if isinstance(names, list): + quant_config[key] = [ + remap_name(n) if isinstance(n, str) else n + for n in names + ] + + return quant_config + + @staticmethod + def _remap_text_only_state_dict_to_full_wrapper(state_dict: dict) -> dict: + """Remap text-only CausalLM state_dict to composite language_model prefix. + + Example: + model.layers.0... -> model.language_model.layers.0... + model.embed_tokens... -> model.language_model.embed_tokens... + model.norm... -> model.language_model.norm... + + Keep top-level lm_head.* unchanged. + """ + remapped = {} + + for key, tensor in state_dict.items(): + new_key = key + + if key.startswith("model.") and not key.startswith("model.language_model."): + new_key = "model.language_model." + key[len("model.") :] + + # lm_head usually stays top-level. + if key.startswith("lm_head."): + new_key = key + + if new_key in remapped: + raise RuntimeError( + f"State dict key collision during full-wrapper remap: " + f"{key} -> {new_key}" + ) + + remapped[new_key] = tensor + + return remapped \ No newline at end of file diff --git a/onecomp/utils/blockwise.py b/onecomp/utils/blockwise.py index 3efcfbb..aa17b2d 100644 --- a/onecomp/utils/blockwise.py +++ b/onecomp/utils/blockwise.py @@ -177,12 +177,25 @@ def get_blocks_and_inputs( blocks = _get_blocks(model) + + # Detect models with heterogeneous layer types (e.g. Gemma4 with # full_attention / sliding_attention) blocks_parent = _find_blocks_parent(model, blocks) layer_types = getattr(getattr(blocks_parent, "config", None), "layer_types", None) unique_layer_types = set(layer_types) if layer_types else set() - has_mixed_types = len(unique_layer_types) > 1 + + is_gemma_like_mixed_attention = unique_layer_types <= { + "full_attention", + "sliding_attention", + } and len(unique_layer_types) > 1 + + is_qwen35_like_hybrid = unique_layer_types <= { + "linear_attention", + "full_attention", + } and "linear_attention" in unique_layer_types + + has_mixed_types = is_gemma_like_mixed_attention or is_qwen35_like_hybrid rotary_hook_handle = None pos_emb_map: dict[str, tuple[torch.Tensor, ...]] = {} @@ -258,13 +271,41 @@ def _capture_rotary(_mod, args, output): return (blocks, inps, kwargs) +def _create_linear_attention_mask( + config, + inputs_embeds, + attention_mask, + cache_position=None, + *, + past_key_values=None, + position_ids=None, + **kwargs, +): + if attention_mask is None: + return None + + if past_key_values is not None and past_key_values.has_previous_state(): + return None + + attention_mask = attention_mask.to(device=inputs_embeds.device, dtype=torch.bool) + + if torch.all(attention_mask): + return None + + return attention_mask + + def _compute_per_type_attention_masks(blocks_parent, kwargs, unique_layer_types): """Compute attention masks for each layer type. Uses create_causal_mask / create_sliding_window_causal_mask from transformers to produce the correct mask per layer type. """ - from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask + + from transformers.masking_utils import ( + create_causal_mask, + create_sliding_window_causal_mask, + ) position_ids = kwargs.get("position_ids") if position_ids is None: @@ -277,12 +318,18 @@ def _compute_per_type_attention_masks(blocks_parent, kwargs, unique_layer_types) dummy_embeds = torch.zeros(1, seq_len, config.hidden_size, device=device, dtype=dtype) attn_mask_1d = torch.ones(1, seq_len, device=device, dtype=torch.long) - # tmp: Gemma4 only has full_attention and sliding_attention layer types. - _mask_creators = { + # tmp: Qwen3.5 has linear_attention layer type. + if unique_layer_types <= {"linear_attention", "full_attention"}: + _mask_creators = { + "linear_attention": _create_linear_attention_mask, + "full_attention": create_causal_mask, + } + else: + _mask_creators = { "full_attention": create_causal_mask, "sliding_attention": create_sliding_window_causal_mask, - } - + } + mask_map = {} for lt in unique_layer_types: creator = _mask_creators.get(lt) @@ -370,24 +417,28 @@ def prepare_block_kwargs(batch_kwargs, block, pli, offset, batch_size, device): # 2) Per-type position embeddings pos_map = batch_kwargs.pop(_POS_EMB_MAP_KEY, None) if pos_map is not None: - layer_type = getattr(block, "layer_type", None) or getattr( - getattr(block, "self_attn", None), "layer_type", None - ) + layer_type = _get_block_layer_type(block) if layer_type and layer_type in pos_map: batch_kwargs["position_embeddings"] = pos_map[layer_type] # 3) Per-type attention mask mask_map = batch_kwargs.pop(_ATTN_MASK_MAP_KEY, None) if mask_map is not None: - layer_type = getattr(block, "layer_type", None) or getattr( - getattr(block, "self_attn", None), "layer_type", None - ) + layer_type = _get_block_layer_type(block) if layer_type and layer_type in mask_map: batch_kwargs["attention_mask"] = mask_map[layer_type] return batch_kwargs +def _get_block_layer_type(block: nn.Module) -> str | None: + return ( + getattr(block, "layer_type", None) + or getattr(block, "block_type", None) + or getattr(getattr(block, "self_attn", None), "layer_type", None) + or getattr(getattr(block, "linear_attn", None), "layer_type", None) + ) + @torch.no_grad() def forward_input( inps: torch.Tensor,