"""dispatch — unified quantize_model router + recursive walker. quantize_model(model, format=..., skip_types=..., exclude_modules=..., dual_path=False, teacher_format=None, chunk_size=1024, adaptive=False, **kwargs) Routes format string → Quantizer preset + walker. Supports dual-path (teacher QuantizedModule for cross-quantization distillation). """ from __future__ import annotations import fnmatch from typing import Any import torch import torch.nn as nn from agiws_neural_quant.base import QuantizedModule, SUPPORTED_MODULE_TYPES from agiws_neural_quant.quantizer import Quantizer from agiws_neural_quant.presets import get_preset def make_quantizer(format: str, **kwargs) -> Quantizer: """Build a Quantizer from a format string + kwargs.""" preset_kwargs = get_preset(format, **kwargs) return Quantizer(**preset_kwargs) def _matches_any(name: str, patterns: set[str]) -> bool: for p in patterns: if fnmatch.fnmatch(name, p): return True return False def _walk( module: nn.Module, quantizer: Quantizer, teacher_quantizer: Quantizer | None, skip_types: tuple[type, ...], exclude_modules: set[str], compute_dtype: str, chunk_size: int | None, adaptive: bool, prefix: str, replaced: list[str], per_head_set: set[str] | None = None, num_heads: int = 0, head_dim: int = 0, ) -> None: for name, child in list(module.named_children()): full = f"{prefix}.{name}" if prefix else name if _matches_any(full, exclude_modules) or _matches_any(name, exclude_modules): continue if isinstance(child, skip_types): continue if isinstance(child, QuantizedModule): continue if isinstance(child, SUPPORTED_MODULE_TYPES) and hasattr(child, "weight") and child.weight is not None: # Per-head override for matching module names. use_quantizer = quantizer use_teacher = teacher_quantizer if per_head_set and (name in per_head_set or _matches_any(name, per_head_set)): qcfg = quantizer.to_config() qcfg["scale_mode"] = "per-head" qcfg["num_heads"] = num_heads qcfg["head_dim"] = head_dim use_quantizer = Quantizer(**qcfg) if teacher_quantizer is not None: tcfg = teacher_quantizer.to_config() tcfg["scale_mode"] = "per-head" tcfg["num_heads"] = num_heads tcfg["head_dim"] = head_dim use_teacher = Quantizer(**tcfg) teacher = None if use_teacher is not None: teacher = QuantizedModule.from_module( child, use_teacher, compute_dtype=compute_dtype, chunk_size=chunk_size, adaptive=adaptive, ) qm = QuantizedModule.from_module( child, use_quantizer, compute_dtype=compute_dtype, chunk_size=chunk_size, adaptive=adaptive, teacher=teacher, ) setattr(module, name, qm) replaced.append(full) continue _walk(child, quantizer, teacher_quantizer, skip_types, exclude_modules, compute_dtype, chunk_size, adaptive, full, replaced, per_head_set, num_heads, head_dim) def quantize_model( model: nn.Module, format: str = "int8", skip_types: tuple[type, ...] | None = None, exclude_modules: list[str] | None = None, compute_dtype: str = "fp32", chunk_size: int | None = 1024, adaptive: bool = False, dual_path: bool = False, teacher_format: str | None = None, per_head_modules: list[str] | None = None, num_heads: int = 0, head_dim: int = 0, **kwargs, ) -> nn.Module: """Quantize a PyTorch model in-place via the unified QuantizedModule. Args: model: any nn.Module (walked recursively). format: quantization format string (see presets.FORMAT_PRESETS). skip_types: module types to leave untouched. exclude_modules: glob patterns of module names to skip. compute_dtype: 'fp32' | 'fp16' | 'bf16' for dequantized matmul. chunk_size: output-dim chunk for dequant (None = no chunking). adaptive: if True, AdaptiveChunkSize monitor adjusts chunk_size. dual_path: if True, create teacher QuantizedModule alongside student. teacher_format: format string for teacher (if dual_path). If None and dual_path=True, uses "fp16" (passthrough). per_head_modules: module name patterns to use per-head scale_mode (e.g. ["q_proj", "k_proj", "v_proj", "o_proj"]). num_heads: number of attention heads (for per-head scale_mode). head_dim: dimension per head (for per-head scale_mode). **kwargs: forwarded to Quantizer (override preset values). """ quantizer = make_quantizer(format, **kwargs) teacher_quantizer = None if dual_path: t_fmt = teacher_format or "fp16" teacher_quantizer = make_quantizer(t_fmt) skip_types_t = tuple(skip_types) if skip_types else () exclude = set(exclude_modules) if exclude_modules else set() per_head_set = set(per_head_modules) if per_head_modules else set() replaced: list[str] = [] _walk(model, quantizer, teacher_quantizer, skip_types_t, exclude, compute_dtype, chunk_size, adaptive, "", replaced, per_head_set, num_heads, head_dim) model._quantized_replaced = replaced # type: ignore[attr-defined] model._quantized_format = format # type: ignore[attr-defined] # Log summary. qinfo = quantizer.info() print(f"[NeuralQuant] quantize_model: format={format} repr={qinfo['repr']} " f"bits={qinfo['bits']} scale={qinfo['scale']} group={qinfo['group']} " f"w=True a={qinfo['a']} learnable={qinfo['learnable']} " f"replaced={len(replaced)} modules", flush=True) if per_head_set: ph_count = sum(1 for r in replaced if any(p in r for p in per_head_set)) print(f"[NeuralQuant] per-head: {ph_count} modules (head_dim={head_dim}, " f"num_heads={num_heads})", flush=True) return model def count_quantizable_layers(model: nn.Module) -> dict[str, int]: counts: dict[str, int] = {} for module in model.modules(): if isinstance(module, SUPPORTED_MODULE_TYPES) and hasattr(module, "weight") and module.weight is not None: key = type(module).__name__ counts[key] = counts.get(key, 0) + 1 return counts # --------------------------------------------------------------------------- # Save / Load — persist a FULL quantized model (v3_hybrid_state format). # # save_model(model, path): saves QuantizedModule.to_dict() for quantized # layers + native_state (all non-quantized params/buffers) for the rest. # # load_model(model, path): fills an ALREADY created + quantized model with # weights from the file. Replaces QuantizedModule instances via from_dict # (packed buffers are format-specific, load_state_dict cannot handle them). # Fills native params/buffers via load_state_dict(strict=False). # # Format: "agiws_neural_quant_v3_hybrid_state" (breaking change from v2). # v2 is rejected with a clear error (no backward compat). # --------------------------------------------------------------------------- # Module types registered as quantizable (used to distinguish weight buffers). _WEIGHT_MODULE_TYPES = (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d, nn.Embedding, nn.LayerNorm, nn.Bilinear) def _collect_quantized_paths(model: nn.Module) -> list[str]: """Return dotted paths of all QuantizedModule instances in the model.""" paths: list[str] = [] for name, mod in model.named_modules(): if isinstance(mod, QuantizedModule): paths.append(name) return paths def _walk_quantized( module: nn.Module, prefix: str, out: dict[str, dict], ) -> None: """Collect to_dict() for every QuantizedModule, keyed by dotted path.""" for name, child in list(module.named_children()): full = f"{prefix}.{name}" if prefix else name if isinstance(child, QuantizedModule): out[full] = child.to_dict() else: _walk_quantized(child, full, out) def _collect_native_state(model: nn.Module, quantized_paths: set[str]) -> dict[str, torch.Tensor]: """Collect all non-quantized parameters and buffers. Excludes any param/buffer that belongs to a QuantizedModule (those are serialized separately in quantized_modules). Uses standard PyTorch dotted-path keys (same as model.state_dict()). """ native: dict[str, torch.Tensor] = {} # Parameters. for name, param in model.named_parameters(): # Skip if this param belongs to a QuantizedModule. if _belongs_to_quantized(name, quantized_paths): continue native[name] = param.detach().cpu().clone() # Buffers. for name, buf in model.named_buffers(): if _belongs_to_quantized(name, quantized_paths): continue if buf is None: continue native[name] = buf.detach().cpu().clone() return native def _belongs_to_quantized(dotted_name: str, quantized_paths: set[str]) -> bool: """Check if a dotted-path param/buffer name belongs to a QuantizedModule. A param at "blocks.0.attn.q_proj.weight" belongs to the QuantizedModule at "blocks.0.attn.q_proj" if that path is in quantized_paths. """ for qpath in quantized_paths: if dotted_name == qpath or dotted_name.startswith(qpath + "."): return True return False def save_model(model: nn.Module, path: str) -> None: """Save a FULL quantized model to a .pt file (torch.save). Saves: - quantized_modules: {dotted_path: QuantizedModule.to_dict()} for every QuantizedModule in the model (packed weight buffers + meta + config). - native_state: {dotted_path: tensor} for all non-quantized parameters and buffers (custom layers, embeddings, positional encodings, etc.). - quantized_paths: list of dotted paths of all QuantizedModule instances. - quant_format: the format string used for quantization. - compute_dtype: target compute dtype (from first QuantizedModule). Format: "agiws_neural_quant_v3_hybrid_state". Args: model: a quantized model (after quantize_model). path: output .pt file path. """ quantized: dict[str, dict] = {} _walk_quantized(model, "", quantized) quantized_paths = _collect_quantized_paths(model) native_state = _collect_native_state(model, set(quantized_paths)) # Compute dtype from first QuantizedModule (all should match). compute_dtype = "fp32" if quantized: first_path = next(iter(quantized)) compute_dtype = quantized[first_path].get("compute_dtype", "fp32") quant_format = getattr(model, "_quantized_format", "unknown") payload = { "format": "agiws_neural_quant_v3_hybrid_state", "quant_format": quant_format, "compute_dtype": compute_dtype, "quantized_paths": quantized_paths, "quantized_modules": quantized, "native_state": native_state, "model_config": getattr(model, "_nq_model_config", None), } torch.save(payload, path) def load_model(model: nn.Module, path: str) -> None: """Load weights from a save file into an ALREADY created + quantized model. The model must be created by the user (via their factory) and quantized via quantize_model() BEFORE calling load_model. This function: 1. Replaces each QuantizedModule in the model with QuantizedModule.from_dict() from the file (correct packed weight buffers). 2. Fills non-quantized parameters/buffers via load_state_dict(native_state, strict=False) and reports missing/unexpected keys. Why from_dict + setattr instead of load_state_dict into existing QuantizedModule: QuantizedModule stores packed weight buffers (int4 codes, FP8 scale codes, codebook indices, etc.) — these are NOT raw float weights. load_state_dict cannot reconstruct the packed format from a state_dict. from_dict rebuilds the QuantizedModule with the exact buffers + meta + quantizer config from the file, producing identical dequantized weights. Args: model: an already-created, already-quantized model. Will be modified in-place (QuantizedModule instances replaced, native params filled). path: .pt file saved by save_model or convert_model. Raises: ValueError: if the file is not a v3_hybrid_state file, or if it's an old v2 file (with guidance to re-quantize). """ payload = torch.load(path, map_location="cpu", weights_only=False) if not isinstance(payload, dict): raise ValueError(f"load_model: not a valid NeuralQuant save file: {path}") fmt = payload.get("format") if fmt == "agiws_neural_quant_v2": raise ValueError( f"load_model: file {path} uses old format 'agiws_neural_quant_v2' " f"(pre-0.3.0). This format is no longer supported. " f"Re-quantize your model with NeuralQuant >= 0.3.0 and save again." ) if fmt != "agiws_neural_quant_v3_hybrid_state": raise ValueError( f"load_model: unsupported format {fmt!r}. " f"Expected 'agiws_neural_quant_v3_hybrid_state'." ) quantized_modules = payload["quantized_modules"] native_state = payload["native_state"] # 1. Replace QuantizedModule instances with from_dict versions. for dotted_path, qm_dict in quantized_modules.items(): qm = QuantizedModule.from_dict(qm_dict) _set_module_by_path(model, dotted_path, qm) # 2. Fill native params/buffers. if native_state: missing, unexpected = model.load_state_dict(native_state, strict=False) if missing: print(f"[NeuralQuant] load_model: {len(missing)} missing keys " f"(not in saved native_state): {missing[:5]}{'...' if len(missing) > 5 else ''}", flush=True) if unexpected: print(f"[NeuralQuant] load_model: {len(unexpected)} unexpected keys " f"(in saved file but not in model): {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}", flush=True) def _set_module_by_path(root: nn.Module, dotted_path: str, new_module: nn.Module) -> None: """Set a module at a dotted path within root (e.g. 'blocks.0.attn.q_proj'). Uses standard PyTorch dotted-path convention: split by '.', traverse parent modules, setattr on the parent. """ parts = dotted_path.split(".") parent = root for part in parts[:-1]: # ModuleList indices are accessed via int indexing. if part.isdigit(): parent = parent[int(part)] else: parent = getattr(parent, part) last = parts[-1] if last.isdigit(): parent[int(last)] = new_module # type: ignore[index] else: setattr(parent, last, new_module)