diff --git a/scripts/validate_llama_quant_gptq_phase_split.py b/scripts/validate_llama_quant_gptq_phase_split.py new file mode 100644 index 000000000..38c0a05de --- /dev/null +++ b/scripts/validate_llama_quant_gptq_phase_split.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 +"""Validate Llama GPTQ + phase-split integration in the module quantize flow. + +This script is a *standalone integration validator* for the current step-1 +quantization design: +1) GPTQ quantizes Llama linear weights with calibration data (wikitext2). +2) Quantized Llama wrapper modules are installed through + ``quantize_module_transform_pass``. +3) Runtime phase is inferred by decoder-layer pre-hooks. +4) Decode path stays full precision by policy (``decode_policy=fp_only``). + +The script intentionally prioritizes robustness and maintainability over +aggressive compression. It uses conservative quantization settings and prints +clear runtime summaries for human inspection. +""" + +from __future__ import annotations + +import argparse +import random +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from chop.passes.module.transforms.quantize.quantize import ( + quantize_module_transform_pass, +) + + +@dataclass(frozen=True) +class ValidationConfig: + """Runtime configuration for this integration validator.""" + + model_name: str + device: str + dtype: str + local_files_only: bool + trust_remote_code: bool + prompt: str + max_new_tokens: int + seed: int + print_chars: int + # GPTQ settings + gptq_dataset: str + gptq_nsamples: int + gptq_seqlen: int + gptq_cali_batch_size: int + gptq_max_layers: int | None + gptq_checkpoint_dir: str | None + hf_token: str | None + + +def _parse_args() -> ValidationConfig: + """Parse CLI arguments and return an immutable runtime config.""" + + parser = argparse.ArgumentParser( + description=( + "Validate Llama GPTQ + phase-split module quantization " + "using local HF cache and conservative settings." + ) + ) + parser.add_argument( + "--model_name", + type=str, + default="meta-llama/Llama-3.1-8B-Instruct", + help="HuggingFace model id. Must exist in local cache when offline.", + ) + parser.add_argument( + "--device", + type=str, + default="cuda:0", + help=( + "Torch device string. Use with CUDA_VISIBLE_DEVICES=1 to pin this " + "process to physical GPU1." + ), + ) + parser.add_argument( + "--dtype", + type=str, + choices=("float16", "bfloat16", "float32"), + default="bfloat16", + help="Model load dtype.", + ) + parser.add_argument( + "--local_files_only", + type=lambda x: str(x).lower() in {"1", "true", "yes", "y"}, + default=True, + help="Force loading model/tokenizer/dataset from local cache only.", + ) + parser.add_argument( + "--trust_remote_code", + type=lambda x: str(x).lower() in {"1", "true", "yes", "y"}, + default=False, + help="Forwarded to HF model/tokenizer loaders.", + ) + parser.add_argument( + "--prompt", + type=str, + default=( + "You are a helpful assistant. Briefly explain three practical " + "steps to debug slow Python code, then end with one-line advice." + ), + help="Validation prompt used for generation output inspection.", + ) + parser.add_argument( + "--max_new_tokens", + type=int, + default=96, + help="Maximum generated tokens for the validation output.", + ) + parser.add_argument("--seed", type=int, default=42, help="Global RNG seed.") + parser.add_argument( + "--print_chars", + type=int, + default=1200, + help="Print at most this many output characters (human inspection).", + ) + + # GPTQ + parser.add_argument( + "--gptq_dataset", + type=str, + default="wikitext2", + choices=("wikitext2", "c4", "ptb"), + help="Calibration dataset for GPTQ.", + ) + parser.add_argument("--gptq_nsamples", type=int, default=128) + parser.add_argument("--gptq_seqlen", type=int, default=256) + parser.add_argument("--gptq_cali_batch_size", type=int, default=32) + parser.add_argument( + "--gptq_max_layers", + type=int, + default=None, + help="Optional limit for GPTQ layers; None means full-layer quantization.", + ) + parser.add_argument( + "--gptq_checkpoint_dir", + type=str, + default=None, + help="Optional checkpoint directory for GPTQ resume.", + ) + parser.add_argument( + "--hf_token", + type=str, + default=None, + help="Optional HF token if your cached model metadata still requires it.", + ) + + args = parser.parse_args() + return ValidationConfig(**vars(args)) + + +def _resolve_torch_dtype(name: str) -> torch.dtype: + """Map CLI dtype name to torch dtype.""" + + mapping = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + } + return mapping[name] + + +def _set_deterministic_seed(seed: int) -> None: + """Set deterministic seeds for reproducibility of this validator.""" + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _build_conservative_pass_args(cfg: ValidationConfig) -> dict[str, Any]: + """Build conservative quantization config for module quantization pass. + + Design intent: + - Keep dynamic quantization paths in attention/MLP/RMS mostly bypassed to + reduce quality regression risk in this validator. + - Let GPTQ provide the primary compression effect on linear weights. + - Preserve phase-structured config fields to validate prefill/decode wiring. + """ + + # Conservative phase sub-config for attention kernels. + # We keep bypass=True to prioritize text readability while still exercising + # phase-structured config parsing and hook-based runtime phase selection. + attn_phase_cfg = { + "qk_matmul": {"bypass": True}, + "av_matmul": {"bypass": True}, + "rope": {"bypass": True}, + "softmax": {"bypass": True}, + "kv_cache": {"bypass": True}, + } + + # For linear_mxfp: + # - prefill uses GPTQ weight result (gptq=True prevents re-PTQ in load_state_dict) + # - decode remains fp-only via global decode policy + linear_prefill_cfg = { + "bypass": True, + "gptq": True, + "clip_search": False, + "weight_block_size": 16, + "weight_exponent_width": 4, + "weight_frac_width": 3, + "data_in_block_size": 16, + "data_in_exponent_width": 4, + "data_in_frac_width": 3, + "bias_block_size": 16, + "bias_exponent_width": 4, + "bias_frac_width": 3, + } + linear_decode_cfg = { + "bypass": True, + "weight_block_size": 16, + "weight_exponent_width": 4, + "weight_frac_width": 3, + "data_in_block_size": 16, + "data_in_exponent_width": 4, + "data_in_frac_width": 3, + "bias_block_size": 16, + "bias_exponent_width": 4, + "bias_frac_width": 3, + } + + # MLP/RMS minifloat phase cfg: conservative bypass while preserving schema. + mlp_rms_phase_cfg = { + "prefill": {"bypass": True}, + "decode": {"bypass": True}, + "decode_policy": "fp_only", + } + + return { + "by": "regex_name", + "gptq": { + "model_name": cfg.model_name, + "device": cfg.device, + "dataset": cfg.gptq_dataset, + "nsamples": cfg.gptq_nsamples, + "seqlen": cfg.gptq_seqlen, + "format": "mxfp", + "weight_config": { + "weight_block_size": 16, + "weight_exponent_width": 4, + "weight_frac_width": 3, + }, + "quantile_search": True, + "clip_search_y": False, + "cali_batch_size": cfg.gptq_cali_batch_size, + "checkpoint_dir": cfg.gptq_checkpoint_dir, + "hf_token": cfg.hf_token, + "max_layers": cfg.gptq_max_layers, + }, + r"model\.layers\.\d+\.self_attn$": { + "config": { + "name": "mxfp", + "decode_policy": "fp_only", + "prefill": attn_phase_cfg, + "decode": attn_phase_cfg, + } + }, + r"model\.layers\.\d+\.(input_layernorm|post_attention_layernorm)$": { + "config": { + "name": "minifloat", + "decode_policy": "fp_only", + **mlp_rms_phase_cfg, + } + }, + r"model\.layers\.\d+\.mlp$": { + "config": { + "name": "mxfp", + "decode_policy": "fp_only", + **mlp_rms_phase_cfg, + } + }, + r"model\.layers\.\d+\.(self_attn\.(q_proj|k_proj|v_proj|o_proj)|mlp\.(gate_proj|up_proj|down_proj))$": { + "config": { + "name": "mxfp", + "decode_policy": "fp_only", + "prefill": linear_prefill_cfg, + "decode": linear_decode_cfg, + } + }, + } + + +def _load_model_and_tokenizer(cfg: ValidationConfig) -> tuple[Any, Any]: + """Load model/tokenizer from local cache (offline by default).""" + + torch_dtype = _resolve_torch_dtype(cfg.dtype) + load_kwargs = { + "local_files_only": cfg.local_files_only, + "trust_remote_code": cfg.trust_remote_code, + } + if cfg.hf_token: + load_kwargs["token"] = cfg.hf_token + + tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, **load_kwargs) + model = AutoModelForCausalLM.from_pretrained( + cfg.model_name, + torch_dtype=torch_dtype, + **load_kwargs, + ) + model = model.to(cfg.device) + model.eval() + return model, tokenizer + + +def _print_runtime_header(cfg: ValidationConfig) -> None: + """Print a compact runtime summary for reproducibility.""" + + print("=" * 88) + print("Llama GPTQ + phase-split validator") + print(f"model_name : {cfg.model_name}") + print(f"device : {cfg.device}") + print(f"dtype : {cfg.dtype}") + print(f"local_files_only : {cfg.local_files_only}") + print(f"gptq_dataset : {cfg.gptq_dataset}") + print(f"gptq_nsamples : {cfg.gptq_nsamples}") + print(f"gptq_seqlen : {cfg.gptq_seqlen}") + print(f"gptq_max_layers : {cfg.gptq_max_layers}") + print(f"max_new_tokens : {cfg.max_new_tokens}") + print("=" * 88) + + +def _print_cuda_mapping() -> None: + """Print effective CUDA mapping for traceability. + + Expected setup for this script: + - launch with CUDA_VISIBLE_DEVICES=1 + - pass --device cuda:0 + """ + + if not torch.cuda.is_available(): + print("CUDA not available; validation cannot run on GPU path.") + return + idx = torch.cuda.current_device() + print(f"torch.cuda.current_device(): {idx}") + print(f"torch.cuda.get_device_name : {torch.cuda.get_device_name(idx)}") + + +@torch.no_grad() +def _generate_preview(model: Any, tokenizer: Any, cfg: ValidationConfig) -> str: + """Generate one sample output for human readability inspection.""" + + inputs = tokenizer(cfg.prompt, return_tensors="pt").to(cfg.device) + generated = model.generate( + **inputs, + max_new_tokens=cfg.max_new_tokens, + do_sample=False, + temperature=0.0, + ) + text = tokenizer.decode(generated[0], skip_special_tokens=True) + return text + + +def main() -> None: + cfg = _parse_args() + _set_deterministic_seed(cfg.seed) + _print_runtime_header(cfg) + _print_cuda_mapping() + + print("[1/4] Loading model and tokenizer from local cache...") + t0 = time.time() + model, tokenizer = _load_model_and_tokenizer(cfg) + print(f" done in {time.time() - t0:.2f}s") + + print("[2/4] Building conservative quantization config...") + pass_args = _build_conservative_pass_args(cfg) + print(" done") + + print("[3/4] Running quantize_module_transform_pass (includes GPTQ)...") + t1 = time.time() + model, _ = quantize_module_transform_pass(model, pass_args) + # Replacement may instantiate some wrapper modules on CPU by default. + # We force a post-pass device sync so generation uses one consistent device. + model = model.to(cfg.device) + model.eval() + print(f" done in {time.time() - t1:.2f}s") + + print("[4/4] Running generation preview...") + out = _generate_preview(model, tokenizer, cfg) + + print("\n" + "=" * 88) + print("Prompt:") + print(cfg.prompt) + print("-" * 88) + print("Generated text preview (manual readability inspection):") + if cfg.print_chars > 0: + print(out[: cfg.print_chars]) + if len(out) > cfg.print_chars: + print(f"\n... [truncated to {cfg.print_chars} chars]") + else: + print(out) + print("=" * 88) + + +if __name__ == "__main__": + # Example: + # CUDA_VISIBLE_DEVICES=1 PYTHONPATH=src python scripts/validate_llama_quant_gptq_phase_split.py --device cuda:0 + main() diff --git a/src/chop/nn/quantized/modules/linear.py b/src/chop/nn/quantized/modules/linear.py index 59a6a0162..12914ca87 100644 --- a/src/chop/nn/quantized/modules/linear.py +++ b/src/chop/nn/quantized/modules/linear.py @@ -16,6 +16,11 @@ import torch from torch import Tensor from torch.nn import functional as F +from chop.nn.quantized.modules.phase_context import get_runtime_phase +from chop.nn.quantized.modules.phase_config import ( + get_phase_subconfig, + normalize_phase_q_config, +) from ..utils import get_stats, quantiser_passthrough @@ -37,10 +42,6 @@ mxfp_quantizer, ) -# `chop.nn.quantizers.rotation` triggers an `import fast_hadamard_transform` -# at load time. We defer that import to RotateMXIntLinear.forward so users -# that never touch the rotate path don't need the CUDA extension installed. - # LUTNet import numpy as np from typing import Type @@ -819,6 +820,18 @@ def forward(self, x): class LinearMXFP(_LinearBase): + """MXFP linear with prefill/decode phase-aware execution. + + Runtime phase policy: + - `fp_only`: decode uses preserved FP snapshots. + - `quantized`: decode uses decode-specific quantized weight bank. + + Memory policy (this change): + - We do not keep all decode banks alive at once. + - Active policy keeps only one decode bank type (`_fp` or `_q`), so common + runtime memory is bounded to `self.weight + one decode bank`. + """ + # NOTE: backward is not supported — inference only (PTQ) def __init__( self, @@ -831,79 +844,112 @@ def __init__( ) -> None: super().__init__(in_features, out_features, bias, device, dtype) assert config is not None, "config is None!" - self.config = config - self.bypass = config.get("bypass", False) - self.gptq = config.get("gptq", False) - self.clip_search = config.get("clip_search", False) + self.phase_config = normalize_phase_q_config(config) + self.decode_policy = self.phase_config["decode_policy"] + if self.decode_policy not in ("fp_only", "quantized"): + raise ValueError( + "Unsupported decode_policy " + f"{self.decode_policy!r} for {self.__class__.__name__}." + ) + + prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") + self.prefill_config = prefill_cfg + self.bypass = prefill_cfg.get("bypass", False) + self.gptq = prefill_cfg.get("gptq", False) + self.clip_search = prefill_cfg.get("clip_search", False) + + # Decode banks are non-persistent runtime buffers; they are excluded + # from checkpoints so saved artifacts remain upstream-compatible. + self.register_buffer("_decode_weight_fp", torch.empty(0), persistent=False) + self.register_buffer("_decode_bias_fp", torch.empty(0), persistent=False) + self.register_buffer("_decode_weight_q", torch.empty(0), persistent=False) + self.register_buffer("_decode_bias_q", torch.empty(0), persistent=False) + + def _capture_decode_fp_bank_snapshot(self, state_dict) -> None: + """Save FP snapshot only when fp-only decode policy actually needs it.""" + + if self.decode_policy != "fp_only": + # Quantized decode should not keep an extra FP bank in steady-state. + self._decode_weight_fp = torch.empty(0, device=self.weight.device) + self._decode_bias_fp = torch.empty(0, device=self.weight.device) + return + + if "weight" in state_dict: + self._decode_weight_fp = state_dict["weight"].detach().clone() + if ( + self.bias is not None + and "bias" in state_dict + and state_dict["bias"] is not None + ): + self._decode_bias_fp = state_dict["bias"].detach().clone() - @classmethod - def from_linear(cls, linear: torch.nn.Linear, config: dict) -> "LinearMXFP": - """Create a LinearMXFP that REUSES the original Linear's Parameters. + def _build_decode_quantized_bank(self) -> None: + """Build decode quantized bank once from current model weights. - Mirrors ``LinearMXInt.from_linear`` for MXFP: shares the existing - ``weight`` / ``bias`` Parameter (preserving DTensor sharding), then - quantizes them in place. Used by the rotation-search swap path so - ``LinearMXFP`` ↔ ``RotateMXFPLinear`` swaps stay zero-copy. + Why source is `self.weight`: + - This mode targets low-memory operation with no extra FP decode bank. + - `self.weight` already represents current prefill bank (FP/PTQ/GPTQ), + so decode quant bank stays deterministic without a third copy. """ - assert config is not None, "config is None!" - new = cls.__new__(cls) - torch.nn.Module.__init__(new) - - new.in_features = linear.in_features - new.out_features = linear.out_features - new.weight = linear.weight - new.bias = linear.bias - - new.pruning_masks = None - - new.config = config - new.bypass = config.get("bypass", False) - new.gptq = config.get("gptq", False) - new.clip_search = config.get("clip_search", False) - - if not new.bypass and not new.gptq: - with torch.no_grad(): - new.weight.data.copy_( - mxfp_quantizer( - new.weight.data, - block_size=config["weight_block_size"], - element_exp_bits=config["weight_exponent_width"], - element_frac_bits=config["weight_frac_width"], - block_dim=1, - ) - ) - b_block_size = config.get("bias_block_size") - b_exp_bits = config.get("bias_exponent_width") - b_frac_bits = config.get("bias_frac_width") - if ( - new.bias is not None - and b_block_size is not None - and b_exp_bits is not None - ): - new.bias.data.copy_( - mxfp_quantizer( - new.bias.data, - block_size=b_block_size, - element_exp_bits=b_exp_bits, - element_frac_bits=b_frac_bits, - block_dim=0, - ) - ) + if self.decode_policy != "quantized": + self._decode_weight_q = torch.empty(0, device=self.weight.device) + self._decode_bias_q = torch.empty(0, device=self.weight.device) + return - return new + decode_cfg, _ = get_phase_subconfig(self.phase_config, "decode") + if decode_cfg.get("bypass", False): + self._decode_weight_q = torch.empty(0, device=self.weight.device) + self._decode_bias_q = torch.empty(0, device=self.weight.device) + return + + source_weight = self.weight + source_bias = self.bias + source_weight = source_weight.to(self.weight.device, self.weight.dtype) + self._decode_weight_q = mxfp_quantizer( + source_weight, + block_size=decode_cfg["weight_block_size"], + element_exp_bits=decode_cfg["weight_exponent_width"], + element_frac_bits=decode_cfg["weight_frac_width"], + block_dim=1, + ) + + if ( + source_bias is not None + and decode_cfg.get("bias_block_size") is not None + and decode_cfg.get("bias_exponent_width") is not None + ): + self._decode_bias_q = mxfp_quantizer( + source_bias.to(self.weight.device, self.weight.dtype), + block_size=decode_cfg["bias_block_size"], + element_exp_bits=decode_cfg["bias_exponent_width"], + element_frac_bits=decode_cfg.get("bias_frac_width"), + block_dim=0, + ) + else: + self._decode_bias_q = torch.empty(0, device=self.weight.device) + # Quantized mode keeps only quantized decode bank, not FP decode bank. + self._decode_weight_fp = torch.empty(0, device=self.weight.device) + self._decode_bias_fp = torch.empty(0, device=self.weight.device) + + def refresh_decode_runtime_bank(self) -> None: + """Refresh decode bank according to active decode policy.""" + + self._build_decode_quantized_bank() def load_state_dict(self, state_dict, strict=True, assign=False): """Load pretrained weights, then quantize them in place.""" + self._capture_decode_fp_bank_snapshot(state_dict) result = super().load_state_dict(state_dict, strict=strict, assign=assign) if self.bypass or self.gptq: + self._build_decode_quantized_bank() return result # Quantize weight - w_block_size = self.config["weight_block_size"] - w_exp_bits = self.config["weight_exponent_width"] - w_frac_bits = self.config["weight_frac_width"] + w_block_size = self.prefill_config["weight_block_size"] + w_exp_bits = self.prefill_config["weight_exponent_width"] + w_frac_bits = self.prefill_config["weight_frac_width"] self.weight.data.copy_( mxfp_quantizer( self.weight.data, @@ -915,9 +961,9 @@ def load_state_dict(self, state_dict, strict=True, assign=False): ) # Quantize bias - b_block_size = self.config.get("bias_block_size") - b_exp_bits = self.config.get("bias_exponent_width") - b_frac_bits = self.config.get("bias_frac_width") + b_block_size = self.prefill_config.get("bias_block_size") + b_exp_bits = self.prefill_config.get("bias_exponent_width") + b_frac_bits = self.prefill_config.get("bias_frac_width") if ( self.bias is not None and b_block_size is not None @@ -933,17 +979,64 @@ def load_state_dict(self, state_dict, strict=True, assign=False): ) ) + self._build_decode_quantized_bank() return result @torch.no_grad() def forward(self, x): - if self.bypass: + runtime_phase = get_runtime_phase() + if runtime_phase == "decode" and self.decode_policy == "fp_only": + # Why this branch exists: + # Default path must remain fully backward-compatible. + decode_weight = ( + self._decode_weight_fp + if self._decode_weight_fp.numel() > 0 + else self.weight + ) + decode_bias = ( + self._decode_bias_fp + if self.bias is not None and self._decode_bias_fp.numel() > 0 + else self.bias + ) + return F.linear(x, decode_weight, decode_bias) + + if runtime_phase == "decode" and self.decode_policy == "quantized": + decode_phase_config, _ = get_phase_subconfig(self.phase_config, "decode") + if decode_phase_config.get("bypass", False): + # No fallback bank: bypass means decode consumes current bank + # directly (typically FP if prefill is FP, or PTQ/GPTQ bank). + return F.linear(x, self.weight, self.bias) + + x_block_size = decode_phase_config.get("data_in_block_size") + x_exp_bits = decode_phase_config.get("data_in_exponent_width") + x_frac_bits = decode_phase_config.get("data_in_frac_width") + if x_block_size is not None and x_exp_bits is not None: + x = mxfp_quantizer( + x, + block_size=x_block_size, + element_exp_bits=x_exp_bits, + element_frac_bits=x_frac_bits, + block_dim=-1, + ) + decode_weight = ( + self._decode_weight_q + if self._decode_weight_q.numel() > 0 + else self.weight + ) + if self.bias is not None and self._decode_bias_q.numel() > 0: + decode_bias = self._decode_bias_q + else: + decode_bias = self.bias + return F.linear(x, decode_weight, decode_bias) + + prefill_phase_config, _ = get_phase_subconfig(self.phase_config, "prefill") + if prefill_phase_config.get("bypass", False): return F.linear(x, self.weight, self.bias) # Only quantize activations; weights/bias already quantized in load_state_dict - x_block_size = self.config.get("data_in_block_size") - x_exp_bits = self.config.get("data_in_exponent_width") - x_frac_bits = self.config.get("data_in_frac_width") + x_block_size = prefill_phase_config.get("data_in_block_size") + x_exp_bits = prefill_phase_config.get("data_in_exponent_width") + x_frac_bits = prefill_phase_config.get("data_in_frac_width") if x_block_size is not None and x_exp_bits is not None: x = mxfp_quantizer( x, @@ -957,6 +1050,11 @@ def forward(self, x): class LinearMXInt(_LinearBase): + """MXInt linear with prefill/decode phase-aware execution. + + Behavior mirrors `LinearMXFP` with the same low-memory decode-bank policy. + """ + # NOTE: backward is not supported — inference only (PTQ) def __init__( self, @@ -969,88 +1067,97 @@ def __init__( ) -> None: super().__init__(in_features, out_features, bias, device, dtype) assert config is not None, "config is None!" - self.config = config - self.bypass = config.get("bypass", False) - self.gptq = config.get("gptq", False) - self.clip_search = config.get("clip_search", False) - - @classmethod - def from_linear(cls, linear: torch.nn.Linear, config: dict) -> "LinearMXInt": - """Create a LinearMXInt that REUSES the original Linear's Parameters. - - Unlike ``__init__`` which allocates a fresh weight tensor via - ``torch.nn.Linear.__init__``, this classmethod shares the original - module's ``weight`` and ``bias`` Parameters directly. This is - critical for tensor-parallel models: if ``linear.weight`` is a - DTensor shard (from HF's ``tp_plan="auto"``), re-allocating would - lose the sharding and fail the state-dict copy. By sharing the - Parameter, the DTensor is preserved and TP dispatch keeps working. - - The weight is then quantized in place. ``mxint_quantizer`` handles - DTensor inputs transparently (each rank quantizes its local shard). - """ - assert config is not None, "config is None!" - new = cls.__new__(cls) - torch.nn.Module.__init__(new) - - # nn.Linear attributes - new.in_features = linear.in_features - new.out_features = linear.out_features - new.weight = linear.weight # share Parameter (may be DTensor) - new.bias = linear.bias # share Parameter (may be DTensor) - - # _LinearBase attributes - new.pruning_masks = None - - # LinearMXInt attributes - new.config = config - new.bypass = config.get("bypass", False) - new.gptq = config.get("gptq", False) - new.clip_search = config.get("clip_search", False) - - # Quantize weight (and optionally bias) in place. This mirrors the - # logic in load_state_dict() but skips the state-dict copy. - if not new.bypass and not new.gptq: - with torch.no_grad(): - new.weight.data.copy_( - mxint_quantizer( - new.weight.data, - block_size=config["weight_block_size"], - element_bits=config["weight_width"], - block_dim=1, - quantile_search=new.clip_search, - ) - ) + self.phase_config = normalize_phase_q_config(config) + self.decode_policy = self.phase_config["decode_policy"] + if self.decode_policy not in ("fp_only", "quantized"): + raise ValueError( + "Unsupported decode_policy " + f"{self.decode_policy!r} for {self.__class__.__name__}." + ) - b_block_size = config.get("bias_block_size") - b_element_bits = config.get("bias_width") - if ( - new.bias is not None - and b_block_size is not None - and b_element_bits is not None - ): - new.bias.data.copy_( - mxint_quantizer( - new.bias.data, - block_size=b_block_size, - element_bits=b_element_bits, - block_dim=0, - quantile_search=new.clip_search, - ) - ) + prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") + self.prefill_config = prefill_cfg + self.bypass = prefill_cfg.get("bypass", False) + self.gptq = prefill_cfg.get("gptq", False) + self.clip_search = prefill_cfg.get("clip_search", False) + + self.register_buffer("_decode_weight_fp", torch.empty(0), persistent=False) + self.register_buffer("_decode_bias_fp", torch.empty(0), persistent=False) + self.register_buffer("_decode_weight_q", torch.empty(0), persistent=False) + self.register_buffer("_decode_bias_q", torch.empty(0), persistent=False) + + def _capture_decode_fp_bank_snapshot(self, state_dict) -> None: + """Save FP snapshot only when fp-only decode policy actually needs it.""" + + if self.decode_policy != "fp_only": + self._decode_weight_fp = torch.empty(0, device=self.weight.device) + self._decode_bias_fp = torch.empty(0, device=self.weight.device) + return + + if "weight" in state_dict: + self._decode_weight_fp = state_dict["weight"].detach().clone() + if ( + self.bias is not None + and "bias" in state_dict + and state_dict["bias"] is not None + ): + self._decode_bias_fp = state_dict["bias"].detach().clone() + + def _build_decode_quantized_bank(self) -> None: + """Build decode quantized bank once from FP source snapshot.""" + + if self.decode_policy != "quantized": + self._decode_weight_q = torch.empty(0, device=self.weight.device) + self._decode_bias_q = torch.empty(0, device=self.weight.device) + return + + decode_cfg, _ = get_phase_subconfig(self.phase_config, "decode") + if decode_cfg.get("bypass", False): + self._decode_weight_q = torch.empty(0, device=self.weight.device) + self._decode_bias_q = torch.empty(0, device=self.weight.device) + return + + source_weight = self.weight + source_bias = self.bias + source_weight = source_weight.to(self.weight.device, self.weight.dtype) + self._decode_weight_q = mxint_quantizer( + source_weight, + block_size=decode_cfg["weight_block_size"], + element_bits=decode_cfg["weight_width"], + block_dim=1, + quantile_search=decode_cfg.get("clip_search", False), + ) + if source_bias is not None and decode_cfg.get("bias_block_size") is not None: + self._decode_bias_q = mxint_quantizer( + source_bias.to(self.weight.device, self.weight.dtype), + block_size=decode_cfg["bias_block_size"], + element_bits=decode_cfg.get("bias_width"), + block_dim=0, + quantile_search=decode_cfg.get("clip_search", False), + ) + else: + self._decode_bias_q = torch.empty(0, device=self.weight.device) + # Quantized mode keeps only quantized decode bank, not FP decode bank. + self._decode_weight_fp = torch.empty(0, device=self.weight.device) + self._decode_bias_fp = torch.empty(0, device=self.weight.device) + + def refresh_decode_runtime_bank(self) -> None: + """Refresh decode bank according to active decode policy.""" - return new + self._build_decode_quantized_bank() def load_state_dict(self, state_dict, strict=True, assign=False): """Load pretrained weights, then quantize them in place.""" + self._capture_decode_fp_bank_snapshot(state_dict) result = super().load_state_dict(state_dict, strict=strict, assign=assign) if self.bypass or self.gptq: + self._build_decode_quantized_bank() return result # Quantize weight - w_block_size = self.config["weight_block_size"] - w_element_bits = self.config["weight_width"] + w_block_size = self.prefill_config["weight_block_size"] + w_element_bits = self.prefill_config["weight_width"] self.weight.data.copy_( mxint_quantizer( self.weight.data, @@ -1062,8 +1169,8 @@ def load_state_dict(self, state_dict, strict=True, assign=False): ) # Quantize bias - b_block_size = self.config.get("bias_block_size") - b_element_bits = self.config.get("bias_width") + b_block_size = self.prefill_config.get("bias_block_size") + b_element_bits = self.prefill_config.get("bias_width") if ( self.bias is not None and b_block_size is not None @@ -1079,18 +1186,58 @@ def load_state_dict(self, state_dict, strict=True, assign=False): ) ) + self._build_decode_quantized_bank() return result @torch.no_grad() def forward(self, x): - if self.bypass: + runtime_phase = get_runtime_phase() + if runtime_phase == "decode" and self.decode_policy == "fp_only": + decode_weight = ( + self._decode_weight_fp + if self._decode_weight_fp.numel() > 0 + else self.weight + ) + decode_bias = ( + self._decode_bias_fp + if self.bias is not None and self._decode_bias_fp.numel() > 0 + else self.bias + ) + return F.linear(x, decode_weight, decode_bias) + + if runtime_phase == "decode" and self.decode_policy == "quantized": + decode_phase_config, _ = get_phase_subconfig(self.phase_config, "decode") + if decode_phase_config.get("bypass", False): + return F.linear(x, self.weight, self.bias) + + x_block_size = decode_phase_config.get("data_in_block_size") + x_element_bits = decode_phase_config.get("data_in_width") + if x_block_size is not None and x_element_bits is not None: + x = mxint_quantizer( + x, + block_size=x_block_size, + element_bits=x_element_bits, + block_dim=-1, + ) + decode_weight = ( + self._decode_weight_q + if self._decode_weight_q.numel() > 0 + else self.weight + ) + if self.bias is not None and self._decode_bias_q.numel() > 0: + decode_bias = self._decode_bias_q + else: + decode_bias = self.bias + return F.linear(x, decode_weight, decode_bias) + + prefill_phase_config, _ = get_phase_subconfig(self.phase_config, "prefill") + if prefill_phase_config.get("bypass", False): return F.linear(x, self.weight, self.bias) # Only quantize activations; weights/bias already quantized in load_state_dict - x_block_size = self.config.get("data_in_block_size") - x_element_bits = self.config.get("data_in_width") + x_block_size = prefill_phase_config.get("data_in_block_size") + x_element_bits = prefill_phase_config.get("data_in_width") if x_block_size is not None and x_element_bits is not None: - # mxint_quantizer natively handles DTensor inputs (TP-compatible). x = mxint_quantizer( x, block_size=x_block_size, @@ -1102,119 +1249,21 @@ def forward(self, x): class RotateMXFPLinear(LinearMXFP): - """LinearMXFP + exact Hadamard rotation around the activation quantizer. - - MXFP equivalent of ``RotateMXIntLinear``: identical weight/bias quantize - pipeline (inherits ``__init__`` / ``from_linear`` / ``load_state_dict``), - only the forward swaps the plain MXFP activation quantize for - ``mxfp_rotate_quantizer``. + """Compatibility alias for rotate-aware MXFP linear. - Extra config keys (optional): - force_fp32_had: run the Hadamard multiplications in fp32. + Phase-split migration does not modify rotate-specific linear behavior in + this PR. We keep the exported symbol so module-map and package imports from + mainline remain stable after cherry-pick migration. """ - @torch.no_grad() - def forward(self, x): - if self.bypass: - return F.linear(x, self.weight, self.bias) - - x_block_size = self.config.get("data_in_block_size") - x_exp_bits = self.config.get("data_in_exponent_width") - x_frac_bits = self.config.get("data_in_frac_width") - if x_block_size is not None and x_exp_bits is not None: - from chop.nn.quantizers.rotation import mxfp_rotate_quantizer - - x = mxfp_rotate_quantizer( - x, - hadamard_dim=self.in_features, - block_size=x_block_size, - element_exp_bits=x_exp_bits, - element_frac_bits=x_frac_bits, - block_dim=-1, - quantile_search=self.clip_search, - force_fp32=self.config.get("force_fp32_had", False), - ) - - return F.linear(x, self.weight, self.bias) + pass class RotateMXIntLinear(LinearMXInt): - """LinearMXInt + exact Hadamard rotation around the activation quantizer. - - Identical to ``LinearMXInt`` for weight/bias quantization (and reuses its - ``__init__`` / ``from_linear`` / ``load_state_dict``). The forward path - differs by replacing the plain MXINT activation quantize with - ``mxint_rotate_quantizer``: the input is rotated by an exact Hadamard, - quantized, and rotated back. - - The rotation pair is mathematically a no-op in fp; for the round trip to - cancel correctly under the surrounding linears, the upstream weights must - be offline-rotated to match (run the rotate pass with - ``online_rotate=True``). Typical use: layers whose inputs feed - ``o_proj`` / ``down_proj``. - - Extra config keys (optional): - force_fp32_had: run the Hadamard multiplications in fp32. - """ - - @torch.no_grad() - def forward(self, x): - if self.bypass: - return F.linear(x, self.weight, self.bias) - - x_block_size = self.config.get("data_in_block_size") - x_element_bits = self.config.get("data_in_width") - if x_block_size is not None and x_element_bits is not None: - from chop.nn.quantizers.rotation import mxint_rotate_quantizer - - x = mxint_rotate_quantizer( - x, - hadamard_dim=self.in_features, - block_size=x_block_size, - element_bits=x_element_bits, - block_dim=-1, - quantile_search=self.clip_search, - force_fp32=self.config.get("force_fp32_had", False), - ) + """Compatibility alias for rotate-aware MXInt linear. - return F.linear(x, self.weight, self.bias) - - -class RotateMXFPLinear(LinearMXFP): - """LinearMXFP + exact Hadamard rotation around the activation quantizer. - - Mirrors ``RotateMXIntLinear`` for MXFP. Reuses ``LinearMXFP`` for weight - quantization (and its ``__init__`` / ``from_linear`` / ``load_state_dict``); - only the forward activation quantize swaps to ``mxfp_rotate_quantizer``. - - Extra config keys (optional): - force_fp32_had: run the Hadamard multiplications in fp32. + This keeps the public import contract intact while split-phase work is + focused on decode-policy wiring and bank selection. """ - @torch.no_grad() - def forward(self, x): - if self.bypass: - return F.linear(x, self.weight, self.bias) - - x_block_size = self.config.get("data_in_block_size") - x_exp_bits = self.config.get("data_in_exponent_width") - x_frac_bits = self.config.get("data_in_frac_width") - if ( - x_block_size is not None - and x_exp_bits is not None - and x_frac_bits is not None - ): - from chop.nn.quantizers.rotation import mxfp_rotate_quantizer - - x = mxfp_rotate_quantizer( - x, - hadamard_dim=self.in_features, - block_size=x_block_size, - element_exp_bits=x_exp_bits, - element_frac_bits=x_frac_bits, - block_dim=-1, - quantile_search=self.clip_search, - force_fp32=self.config.get("force_fp32_had", False), - ) - - return F.linear(x, self.weight, self.bias) + pass diff --git a/src/chop/nn/quantized/modules/llama/attention.py b/src/chop/nn/quantized/modules/llama/attention.py index 54525848d..848209263 100644 --- a/src/chop/nn/quantized/modules/llama/attention.py +++ b/src/chop/nn/quantized/modules/llama/attention.py @@ -1,3 +1,17 @@ +"""Llama attention quantization modules with phase-aware runtime dispatch. + +Compatibility boundaries: +1. Keep module replacement API stable for existing passes. +2. Accept phase-structured configs (`prefill`/`decode`) without forcing new + caller-side changes. +3. Preserve default decode behavior (`decode_policy='fp_only'`) while allowing + explicit opt-in decode quantization (`decode_policy='quantized'`). +4. Keep a single local eager attention path (no backend fallback dispatch). + +Runtime phase is set by decoder-layer pre-hooks in quantize pass. Attention +consumes context phase to select per-phase quantization settings. +""" + from typing import Optional, Tuple import torch @@ -8,9 +22,7 @@ Cache, repeat_kv, LlamaAttention, - eager_attention_forward as _hf_eager_attention_forward, ) -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from functools import partial @@ -18,17 +30,13 @@ from chop.nn.quantizers import mxfp_quantizer, mxint_quantizer from chop.nn.quantized.functional.rope import rope_minifloat from chop.nn.quantized.functional.softmax import softmax_minifloat -from chop.nn.quantized.functional.kvcache import ( - kv_cache_mxfp, - kv_cache_mxfp_rotate, - kv_cache_mxint, - kv_cache_mxint_rotate, +from chop.nn.quantized.functional.kvcache import kv_cache_mxfp, kv_cache_mxint +from chop.nn.quantized.modules.phase_context import ( + get_runtime_phase, ) -from chop.nn.quantized.functional.attention import ( - eager_attention_forward_mxfp as _eager_attention_forward_mxfp, - eager_attention_forward_mxfp_rotate as _eager_attention_forward_mxfp_rotate, - eager_attention_forward_mxint as _eager_attention_forward_mxint, - eager_attention_forward_mxint_rotate as _eager_attention_forward_mxint_rotate, +from chop.nn.quantized.modules.phase_config import ( + get_phase_subconfig, + normalize_phase_q_config, ) import logging @@ -36,31 +44,21 @@ logger = logging.getLogger(__name__) -def _hf_attention_dispatch( - module, - query_states, - key_states, - value_states, - attention_mask, - **kwargs, -): - """Call HF's configured attention backend (sdpa / flash / eager / ...). +def _apply_decode_fp_only_bypass( + phase_subconfig: dict, decode_policy: str, runtime_phase: str +) -> dict: + """Apply `fp_only` decode policy to a phase sub-config. - Used when all in-attention quant stages (qk / av / softmax) are bypassed — - in that case the wrapper has nothing to inject inside the attention compute, - so we shouldn't force eager. Mirrors HF Llama's own dispatch line. + Why this helper exists: + - We keep phase policy handling explicit and centralized instead of + duplicating `if decode/fp_only` branches for each attention sub-stage. """ - attention_interface = ALL_ATTENTION_FUNCTIONS.get( - module.config._attn_implementation, _hf_eager_attention_forward - ) - return attention_interface( - module, - query_states, - key_states, - value_states, - attention_mask, - **kwargs, - ) + + if runtime_phase == "decode" and decode_policy == "fp_only": + cfg = dict(phase_subconfig) + cfg["bypass"] = True + return cfg + return phase_subconfig class LlamaAttentionLSQInteger(nn.Module): @@ -111,10 +109,12 @@ def forward( hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], - past_key_value: Optional[Cache] = None, + past_key_values: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + # HF compatibility: tolerate legacy alias when callers still pass it. + past_key_value = kwargs.pop("past_key_value", past_key_values) input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) @@ -166,39 +166,48 @@ def forward( return attn_output, attn_weights - @classmethod - def from_self(cls, attention: LlamaAttention, q_config: dict = None): - new_attn = cls( - config=attention.config, - layer_idx=attention.layer_idx, - q_config=q_config, - ) - device, dtype = ( - next(attention.parameters()).device, - next(attention.parameters()).dtype, - ) - new_attn = new_attn.to(dtype=dtype, device=device) - # strict=False: LSQInteger quantizer submodules add params not in base LlamaAttention - new_attn.load_state_dict(attention.state_dict(), strict=False) - return new_attn - class LlamaAttentionMXFP(LlamaAttention): """MXFP-quantized LlamaAttention.""" def __init__(self, config, layer_idx, q_config: dict = None): super().__init__(config, layer_idx) - q_config = q_config or {} - self.qk_config = q_config.get("qk_matmul", {}) - self.av_config = q_config.get("av_matmul", {}) - self.rope_config = q_config.get("rope", {}) - self.softmax_config = q_config.get("softmax", {}) - self.kv_cache_config = q_config.get("kv_cache", {}) - self.qk_bypass = self.qk_config.get("bypass", False) - self.av_bypass = self.av_config.get("bypass", False) - self.rope_bypass = self.rope_config.get("bypass", False) - self.softmax_bypass = self.softmax_config.get("bypass", False) - self.kv_cache_bypass = self.kv_cache_config.get("bypass", False) + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + + def _resolve_phase_attention_quant_config(self, runtime_phase: str) -> dict: + """Resolve attention quant config for the current runtime phase.""" + + phase_subconfig, decode_policy = get_phase_subconfig( + self.phase_q_config, runtime_phase + ) + qk_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("qk_matmul", {}), decode_policy, runtime_phase + ) + av_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("av_matmul", {}), decode_policy, runtime_phase + ) + rope_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("rope", {}), decode_policy, runtime_phase + ) + softmax_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("softmax", {}), decode_policy, runtime_phase + ) + kv_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("kv_cache", {}), decode_policy, runtime_phase + ) + return { + "qk_config": qk_cfg, + "av_config": av_cfg, + "rope_config": rope_cfg, + "softmax_config": softmax_cfg, + "kv_cache_config": kv_cfg, + "qk_bypass": qk_cfg.get("bypass", False), + "av_bypass": av_cfg.get("bypass", False), + "rope_bypass": rope_cfg.get("bypass", False), + "softmax_bypass": softmax_cfg.get("bypass", False), + "kv_cache_bypass": kv_cfg.get("bypass", False), + } def forward( self, @@ -209,6 +218,13 @@ def forward( cache_position: Optional[LongTensor] = None, **kwargs, ) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor]]]: + past_key_value = kwargs.pop("past_key_value", past_key_values) + # Phase is written by decoder-layer pre-hooks before input_layernorm. + runtime_phase = get_runtime_phase() + phase_attention_config = self._resolve_phase_attention_quant_config( + runtime_phase + ) + input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) @@ -217,13 +233,13 @@ def forward( value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings - if not self.rope_bypass: + if not phase_attention_config["rope_bypass"]: query_states, key_states = rope_minifloat( query_states, key_states, cos, sin, - self.rope_config, + phase_attention_config["rope_config"], ) else: query_states, key_states = apply_rotary_pos_emb( @@ -233,63 +249,46 @@ def forward( sin, ) - if past_key_values is not None: + if past_key_value is not None: cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if not self.kv_cache_bypass: + if not phase_attention_config["kv_cache_bypass"]: key_states, value_states = kv_cache_mxfp( key_states, value_states, - self.kv_cache_config, + phase_attention_config["kv_cache_config"], ) - key_states, value_states = past_key_values.update( + key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, cache_kwargs, ) - # Skip-replace: if no in-attention quant stage is active, fall back to - # whatever attention backend HF was configured for (sdpa / fa2 / eager). - # KV-cache / RoPE quant happens above and is unaffected. - if self.qk_bypass and self.av_bypass and self.softmax_bypass: - attn_output, attn_weights = _hf_attention_dispatch( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) - else: - assert self.config._attn_implementation == "eager", ( - "MXFP-quantized eager attention requires _attn_implementation='eager' " - "when any of qk/av/softmax stages are active." - ) - attn_output, attn_weights = _eager_attention_forward_mxfp( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - qk_bypass=self.qk_bypass, - qk_config=self.qk_config, - av_bypass=self.av_bypass, - av_config=self.av_config, - softmax_bypass=self.softmax_bypass, - softmax_config=self.softmax_config, - **kwargs, - ) + # Keep a single eager path for maintainability: all-bypass still uses + # this function but bypass flags short-circuit quantization stages. + attn_output, attn_weights = _eager_attention_forward_mxfp( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + qk_bypass=phase_attention_config["qk_bypass"], + qk_config=phase_attention_config["qk_config"], + av_bypass=phase_attention_config["av_bypass"], + av_config=phase_attention_config["av_config"], + softmax_bypass=phase_attention_config["softmax_bypass"], + softmax_config=phase_attention_config["softmax_config"], + **kwargs, + ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights @classmethod - def from_self(cls, attention: LlamaAttention, q_config: dict = None): + def from_attention(cls, attention: LlamaAttention, q_config: dict = None): new_attn = cls( config=attention.config, layer_idx=attention.layer_idx, @@ -309,17 +308,42 @@ class LlamaAttentionMXInt(LlamaAttention): def __init__(self, config, layer_idx, q_config: dict = None): super().__init__(config, layer_idx) - q_config = q_config or {} - self.qk_config = q_config.get("qk_matmul", {}) - self.av_config = q_config.get("av_matmul", {}) - self.rope_config = q_config.get("rope", {}) - self.softmax_config = q_config.get("softmax", {}) - self.kv_cache_config = q_config.get("kv_cache", {}) - self.qk_bypass = self.qk_config.get("bypass", False) - self.av_bypass = self.av_config.get("bypass", False) - self.rope_bypass = self.rope_config.get("bypass", False) - self.softmax_bypass = self.softmax_config.get("bypass", False) - self.kv_cache_bypass = self.kv_cache_config.get("bypass", False) + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + + def _resolve_phase_attention_quant_config(self, runtime_phase: str) -> dict: + """Resolve attention quant config for the current runtime phase.""" + + phase_subconfig, decode_policy = get_phase_subconfig( + self.phase_q_config, runtime_phase + ) + qk_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("qk_matmul", {}), decode_policy, runtime_phase + ) + av_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("av_matmul", {}), decode_policy, runtime_phase + ) + rope_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("rope", {}), decode_policy, runtime_phase + ) + softmax_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("softmax", {}), decode_policy, runtime_phase + ) + kv_cfg = _apply_decode_fp_only_bypass( + phase_subconfig.get("kv_cache", {}), decode_policy, runtime_phase + ) + return { + "qk_config": qk_cfg, + "av_config": av_cfg, + "rope_config": rope_cfg, + "softmax_config": softmax_cfg, + "kv_cache_config": kv_cfg, + "qk_bypass": qk_cfg.get("bypass", False), + "av_bypass": av_cfg.get("bypass", False), + "rope_bypass": rope_cfg.get("bypass", False), + "softmax_bypass": softmax_cfg.get("bypass", False), + "kv_cache_bypass": kv_cfg.get("bypass", False), + } def forward( self, @@ -330,6 +354,13 @@ def forward( cache_position: Optional[LongTensor] = None, **kwargs, ) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor]]]: + past_key_value = kwargs.pop("past_key_value", past_key_values) + # Phase is written by decoder-layer pre-hooks before input_layernorm. + runtime_phase = get_runtime_phase() + phase_attention_config = self._resolve_phase_attention_quant_config( + runtime_phase + ) + input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) @@ -338,13 +369,13 @@ def forward( value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings - if not self.rope_bypass: + if not phase_attention_config["rope_bypass"]: query_states, key_states = rope_minifloat( query_states, key_states, cos, sin, - self.rope_config, + phase_attention_config["rope_config"], ) else: query_states, key_states = apply_rotary_pos_emb( @@ -354,60 +385,46 @@ def forward( sin, ) - if past_key_values is not None: + if past_key_value is not None: cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if not self.kv_cache_bypass: + if not phase_attention_config["kv_cache_bypass"]: key_states, value_states = kv_cache_mxint( key_states, value_states, - self.kv_cache_config, + phase_attention_config["kv_cache_config"], ) - key_states, value_states = past_key_values.update( + key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, cache_kwargs, ) - if self.qk_bypass and self.av_bypass and self.softmax_bypass: - attn_output, attn_weights = _hf_attention_dispatch( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) - else: - assert self.config._attn_implementation == "eager", ( - "MXInt-quantized eager attention requires _attn_implementation='eager' " - "when any of qk/av/softmax stages are active." - ) - attn_output, attn_weights = _eager_attention_forward_mxint( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - qk_bypass=self.qk_bypass, - qk_config=self.qk_config, - av_bypass=self.av_bypass, - av_config=self.av_config, - softmax_bypass=self.softmax_bypass, - softmax_config=self.softmax_config, - **kwargs, - ) + # Keep a single eager path for maintainability: all-bypass still uses + # this function but bypass flags short-circuit quantization stages. + attn_output, attn_weights = _eager_attention_forward_mxint( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + qk_bypass=phase_attention_config["qk_bypass"], + qk_config=phase_attention_config["qk_config"], + av_bypass=phase_attention_config["av_bypass"], + av_config=phase_attention_config["av_config"], + softmax_bypass=phase_attention_config["softmax_bypass"], + softmax_config=phase_attention_config["softmax_config"], + **kwargs, + ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights @classmethod - def from_self(cls, attention: LlamaAttention, q_config: dict = None): + def from_attention(cls, attention: LlamaAttention, q_config: dict = None): new_attn = cls( config=attention.config, layer_idx=attention.layer_idx, @@ -422,230 +439,148 @@ def from_self(cls, attention: LlamaAttention, q_config: dict = None): return new_attn +def _eager_attention_forward_mxfp( + module, + query, + key, + value, + attention_mask, + scaling, + dropout=0.0, + qk_bypass=False, + qk_config=None, + av_bypass=False, + av_config=None, + softmax_bypass=False, + softmax_config=None, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + if not qk_bypass: + q_quantizer = partial( + mxfp_quantizer, + block_size=qk_config["data_in_block_size"], + element_exp_bits=qk_config["data_in_exponent_width"], + element_frac_bits=qk_config["data_in_frac_width"], + block_dim=-1, + ) + query = q_quantizer(query) -class LlamaAttentionMXIntRotate(LlamaAttentionMXInt): - """LlamaAttentionMXInt with online Hadamard rotation around the - activation quantizers (KV-cache, Q@K, A@V). Mirrors - ``Qwen3AttentionMXIntRotate``. - - Per-stage rotate toggles let the rotation search swap stages - independently. Each defaults True (preserves "all rotated" baseline if - the user doesn't pre-set them); the search drives them via the matching - block in q_config: - - q_config["qk_matmul"]["rotate"] = True | False (default True) - q_config["av_matmul"]["rotate"] = True | False (default True) - q_config["kv_cache"]["rotate"] = True | False (default True) - """ + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - def __init__(self, config, layer_idx, q_config: dict = None): - super().__init__(config, layer_idx, q_config=q_config) - self.qk_use_rotate = self.qk_config.get("rotate", True) - self.av_use_rotate = self.av_config.get("rotate", True) - self.kv_cache_use_rotate = self.kv_cache_config.get("rotate", True) + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask - def forward( - self, - hidden_states: Tensor, - position_embeddings: Tuple[Tensor, Tensor], - attention_mask: Optional[Tensor], - past_key_values: Optional[Cache] = None, - cache_position: Optional[LongTensor] = None, - **kwargs, - ) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor]]]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) + if not softmax_bypass: + attn_weights = softmax_minifloat(attn_weights, softmax_config, dim=-1) + else: + attn_weights = nn.functional.softmax( + attn_weights.to(torch.float32), + dim=-1, + ).to(attn_weights.dtype) + + attn_weights = nn.functional.dropout( + attn_weights, + p=dropout, + training=module.training, + ) - query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) - key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + if not av_bypass: + a_quantizer = partial( + mxfp_quantizer, + block_size=av_config["data_in_block_size"], + element_exp_bits=av_config["data_in_exponent_width"], + element_frac_bits=av_config["data_in_frac_width"], + block_dim=-1, + ) + attn_weights = a_quantizer(attn_weights) - cos, sin = position_embeddings - if not self.rope_bypass: - query_states, key_states = rope_minifloat( - query_states, - key_states, - cos, - sin, - self.rope_config, - ) - else: - query_states, key_states = apply_rotary_pos_emb( - query_states, - key_states, - cos, - sin, - ) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if not self.kv_cache_bypass: - if self.kv_cache_use_rotate: - key_states, value_states = kv_cache_mxint_rotate( - key_states, - value_states, - self.kv_cache_config, - ) - else: - key_states, value_states = kv_cache_mxint( - key_states, - value_states, - self.kv_cache_config, - ) - key_states, value_states = past_key_values.update( - key_states, - value_states, - self.layer_idx, - cache_kwargs, - ) - if self.qk_bypass and self.av_bypass and self.softmax_bypass: - attn_output, attn_weights = _hf_attention_dispatch( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) - else: - assert self.config._attn_implementation == "eager", ( - "MXInt-rotate-quantized attention requires _attn_implementation='eager' " - "when any of qk/av/softmax stages are active." - ) - attn_output, attn_weights = _eager_attention_forward_mxint_rotate( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - qk_bypass=self.qk_bypass, - qk_config=self.qk_config, - av_bypass=self.av_bypass, - av_config=self.av_config, - softmax_bypass=self.softmax_bypass, - softmax_config=self.softmax_config, - qk_use_rotate=self.qk_use_rotate, - av_use_rotate=self.av_use_rotate, - **kwargs, - ) +def _eager_attention_forward_mxint( + module, + query, + key, + value, + attention_mask, + scaling, + dropout=0.0, + qk_bypass=False, + qk_config=None, + av_bypass=False, + av_config=None, + softmax_bypass=False, + softmax_config=None, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + if not qk_bypass: + q_quantizer = partial( + mxint_quantizer, + block_size=qk_config["data_in_block_size"], + element_bits=qk_config["data_in_width"], + block_dim=-1, + ) + query = q_quantizer(query) - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask -class LlamaAttentionMXFPRotate(LlamaAttentionMXFP): - """LlamaAttentionMXFP with online Hadamard rotation around the activation - quantizers (KV-cache, Q@K, A@V). MXFP analogue of - ``LlamaAttentionMXIntRotate``; per-stage rotate toggles keep the - rotation-search machinery uniform across formats: - - q_config["qk_matmul"]["rotate"] = True | False (default True) - q_config["av_matmul"]["rotate"] = True | False (default True) - q_config["kv_cache"]["rotate"] = True | False (default True) - """ + if not softmax_bypass: + attn_weights = softmax_minifloat(attn_weights, softmax_config, dim=-1) + else: + attn_weights = nn.functional.softmax( + attn_weights.to(torch.float32), + dim=-1, + ).to(attn_weights.dtype) + + attn_weights = nn.functional.dropout( + attn_weights, + p=dropout, + training=module.training, + ) - def __init__(self, config, layer_idx, q_config: dict = None): - super().__init__(config, layer_idx, q_config=q_config) - self.qk_use_rotate = self.qk_config.get("rotate", True) - self.av_use_rotate = self.av_config.get("rotate", True) - self.kv_cache_use_rotate = self.kv_cache_config.get("rotate", True) + if not av_bypass: + a_quantizer = partial( + mxint_quantizer, + block_size=av_config["data_in_block_size"], + element_bits=av_config["data_in_width"], + block_dim=-1, + ) + attn_weights = a_quantizer(attn_weights) - def forward( - self, - hidden_states: Tensor, - position_embeddings: Tuple[Tensor, Tensor], - attention_mask: Optional[Tensor], - past_key_values: Optional[Cache] = None, - cache_position: Optional[LongTensor] = None, - **kwargs, - ) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor]]]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights - query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) - key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) - cos, sin = position_embeddings - if not self.rope_bypass: - query_states, key_states = rope_minifloat( - query_states, - key_states, - cos, - sin, - self.rope_config, - ) - else: - query_states, key_states = apply_rotary_pos_emb( - query_states, - key_states, - cos, - sin, - ) +class LlamaAttentionMXFPRotate(LlamaAttentionMXFP): + """Rotation-variant compatibility alias for MXFP attention. - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if not self.kv_cache_bypass: - if self.kv_cache_use_rotate: - key_states, value_states = kv_cache_mxfp_rotate( - key_states, - value_states, - self.kv_cache_config, - ) - else: - key_states, value_states = kv_cache_mxfp( - key_states, - value_states, - self.kv_cache_config, - ) - key_states, value_states = past_key_values.update( - key_states, - value_states, - self.layer_idx, - cache_kwargs, - ) + The phase-split migration intentionally keeps one phase-aware eager path in + this file. We retain the rotate class symbol to preserve mainline import + contracts (`llama/__init__.py`) and downstream module map wiring. + """ - if self.qk_bypass and self.av_bypass and self.softmax_bypass: - attn_output, attn_weights = _hf_attention_dispatch( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) - else: - assert self.config._attn_implementation == "eager", ( - "MXFP-rotate-quantized attention requires _attn_implementation='eager' " - "when any of qk/av/softmax stages are active." - ) - attn_output, attn_weights = _eager_attention_forward_mxfp_rotate( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - qk_bypass=self.qk_bypass, - qk_config=self.qk_config, - av_bypass=self.av_bypass, - av_config=self.av_config, - softmax_bypass=self.softmax_bypass, - softmax_config=self.softmax_config, - qk_use_rotate=self.qk_use_rotate, - av_use_rotate=self.av_use_rotate, - **kwargs, - ) + pass - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights + +class LlamaAttentionMXIntRotate(LlamaAttentionMXInt): + """Rotation-variant compatibility alias for MXInt attention. + + Keeping this symbol avoids breaking existing registry/import paths while we + focus this PR on phase-split decode policy behavior. + """ + + pass diff --git a/src/chop/nn/quantized/modules/llama/mlp.py b/src/chop/nn/quantized/modules/llama/mlp.py index 533504da5..b34b55105 100644 --- a/src/chop/nn/quantized/modules/llama/mlp.py +++ b/src/chop/nn/quantized/modules/llama/mlp.py @@ -1,8 +1,19 @@ +"""Llama MLP quantization with phase-aware dispatch. + +Default policy keeps decode full precision (`fp_only`) for compatibility. +`quantized` decode is opt-in and consumes decode-phase sub-config directly. +""" + import torch from torch import nn, Tensor from chop.nn.quantizers.SNN.LSQ import LSQInteger from chop.nn.quantized.functional.silu import silu_minifloat +from chop.nn.quantized.modules.phase_context import get_runtime_phase +from chop.nn.quantized.modules.phase_config import ( + get_phase_subconfig, + normalize_phase_q_config, +) from transformers.models.llama.modeling_llama import LlamaMLP, ACT2FN @@ -51,13 +62,21 @@ class LlamaMLPMXFP(LlamaMLP): def __init__(self, config, layer_idx=None, q_config: dict = None): super().__init__(config) self.layer_idx = layer_idx - self.q_config = q_config or {} - self.bypass = self.q_config.get("bypass", False) + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] def forward(self, x: Tensor) -> Tensor: - if self.bypass: + runtime_phase = get_runtime_phase() + phase_subconfig, decode_policy = get_phase_subconfig( + self.phase_q_config, runtime_phase + ) + bypass = phase_subconfig.get("bypass", False) + if runtime_phase == "decode" and decode_policy == "fp_only": + bypass = True + + if bypass: return super().forward(x) - x = silu_minifloat(self.gate_proj(x), self.q_config) * self.up_proj(x) + x = silu_minifloat(self.gate_proj(x), phase_subconfig) * self.up_proj(x) return self.down_proj(x) @@ -67,11 +86,19 @@ class LlamaMLPMXInt(LlamaMLP): def __init__(self, config, layer_idx=None, q_config: dict = None): super().__init__(config) self.layer_idx = layer_idx - self.q_config = q_config or {} - self.bypass = self.q_config.get("bypass", False) + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] def forward(self, x: Tensor) -> Tensor: - if self.bypass: + runtime_phase = get_runtime_phase() + phase_subconfig, decode_policy = get_phase_subconfig( + self.phase_q_config, runtime_phase + ) + bypass = phase_subconfig.get("bypass", False) + if runtime_phase == "decode" and decode_policy == "fp_only": + bypass = True + + if bypass: return super().forward(x) - x = silu_minifloat(self.gate_proj(x), self.q_config) * self.up_proj(x) + x = silu_minifloat(self.gate_proj(x), phase_subconfig) * self.up_proj(x) return self.down_proj(x) diff --git a/src/chop/nn/quantized/modules/llama/rms_norm.py b/src/chop/nn/quantized/modules/llama/rms_norm.py index f777ef8a7..3cb5424ab 100644 --- a/src/chop/nn/quantized/modules/llama/rms_norm.py +++ b/src/chop/nn/quantized/modules/llama/rms_norm.py @@ -1,3 +1,9 @@ +"""Llama RMSNorm quantization with phase-aware dispatch. + +Default decode behavior remains FP-only for backward compatibility. +`quantized` decode is opt-in via phase config and policy. +""" + from functools import partial import torch @@ -5,6 +11,11 @@ from chop.nn.quantizers.SNN.LSQ import LSQInteger from chop.nn.quantizers._minifloat_mx import MinifloatMeta, minifloat_quantizer_sim +from chop.nn.quantized.modules.phase_context import get_runtime_phase +from chop.nn.quantized.modules.phase_config import ( + get_phase_subconfig, + normalize_phase_q_config, +) from transformers.models.llama.modeling_llama import LlamaRMSNorm @@ -37,48 +48,59 @@ class LlamaRMSNormMinifloat(LlamaRMSNorm): def __init__(self, config=None, layer_idx=None, q_config: dict = None): super().__init__(hidden_size=config.hidden_size, eps=config.rms_norm_eps) self.layer_idx = layer_idx - self.q_config = q_config or {} self.variance_epsilon = config.rms_norm_eps - self.bypass = self.q_config.get("bypass", False) - self.weight_bypass = self.q_config.get("weight_bypass", False) - self.data_in_bypass = self.q_config.get("data_in_bypass", False) - - if not self.bypass and not self.weight_bypass: - self.w_quantizer = partial( - minifloat_quantizer_sim, - minifloat_meta=MinifloatMeta( - exp_bits=self.q_config["weight_exponent_width"], - frac_bits=self.q_config["weight_frac_width"], - is_finite=self.q_config.get("weight_is_finite", True), - round_mode=self.q_config.get("weight_round_mode", "rn"), - ), - ) - else: - self.w_quantizer = None - - if not self.bypass and not self.data_in_bypass: - self.x_quantizer = partial( - minifloat_quantizer_sim, - minifloat_meta=MinifloatMeta( - exp_bits=self.q_config["data_in_exponent_width"], - frac_bits=self.q_config["data_in_frac_width"], - is_finite=self.q_config.get("data_in_is_finite", True), - round_mode=self.q_config.get("data_in_round_mode", "rn"), - ), - ) - else: - self.x_quantizer = None + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + + @staticmethod + def _build_weight_quantizer(sub_cfg: dict, bypass: bool): + """Build weight quantizer for current phase config.""" + + if bypass or sub_cfg.get("weight_bypass", False): + return None + return partial( + minifloat_quantizer_sim, + minifloat_meta=MinifloatMeta( + exp_bits=sub_cfg["weight_exponent_width"], + frac_bits=sub_cfg["weight_frac_width"], + is_finite=sub_cfg.get("weight_is_finite", True), + round_mode=sub_cfg.get("weight_round_mode", "rn"), + ), + ) + + @staticmethod + def _build_input_quantizer(sub_cfg: dict, bypass: bool): + """Build input quantizer for current phase config.""" + + if bypass or sub_cfg.get("data_in_bypass", False): + return None + return partial( + minifloat_quantizer_sim, + minifloat_meta=MinifloatMeta( + exp_bits=sub_cfg["data_in_exponent_width"], + frac_bits=sub_cfg["data_in_frac_width"], + is_finite=sub_cfg.get("data_in_is_finite", True), + round_mode=sub_cfg.get("data_in_round_mode", "rn"), + ), + ) def forward(self, hidden_states): + runtime_phase = get_runtime_phase() + phase_subconfig, decode_policy = get_phase_subconfig( + self.phase_q_config, runtime_phase + ) + bypass = phase_subconfig.get("bypass", False) + if runtime_phase == "decode" and decode_policy == "fp_only": + bypass = True + + w_quantizer = self._build_weight_quantizer(phase_subconfig, bypass) + x_quantizer = self._build_input_quantizer(phase_subconfig, bypass) + input_dtype = hidden_states.dtype - if self.x_quantizer is not None: - hidden_states = self.x_quantizer(hidden_states) + if x_quantizer is not None: + hidden_states = x_quantizer(hidden_states) hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - weight = ( - self.w_quantizer(self.weight) - if self.w_quantizer is not None - else self.weight - ) + weight = w_quantizer(self.weight) if w_quantizer is not None else self.weight return weight * hidden_states.to(input_dtype) diff --git a/src/chop/nn/quantized/modules/phase_config.py b/src/chop/nn/quantized/modules/phase_config.py new file mode 100644 index 000000000..ea98ab2cf --- /dev/null +++ b/src/chop/nn/quantized/modules/phase_config.py @@ -0,0 +1,57 @@ +"""Shared phase-config normalization helpers for quantized modules. + +Step-1 integration contracts: +1. Accept both legacy flat configs and phase-structured configs. +2. Return a stable normalized shape consumed by runtime modules. +3. Do not decide runtime policy here (that remains module-side logic). +""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + + +def normalize_phase_q_config(q_config: dict[str, Any] | None) -> dict[str, Any]: + """Normalize config into `{decode_policy, prefill, decode}` shape. + + Compatibility behavior: + - Legacy flat config -> copied into both prefill and decode buckets. + - Explicit phase config -> missing decode bucket falls back to prefill. + + Why we keep this helper minimal: + - It preserves old config semantics without enforcing model-specific + policy in the parsing layer. + """ + + cfg = deepcopy(q_config or {}) + decode_policy = cfg.get("decode_policy", "fp_only") + + if "prefill" in cfg or "decode" in cfg: + prefill = deepcopy(cfg.get("prefill", {})) + decode = deepcopy(cfg.get("decode", prefill)) + else: + legacy = deepcopy(cfg) + legacy.pop("decode_policy", None) + prefill = legacy + decode = deepcopy(legacy) + + return { + "decode_policy": decode_policy, + "prefill": prefill, + "decode": decode, + } + + +def get_phase_subconfig( + normalized: dict[str, Any], phase: str +) -> tuple[dict[str, Any], str]: + """Select phase sub-config from normalized phase config. + + Returns: + (sub_config, decode_policy) + """ + + if phase == "decode": + return deepcopy(normalized.get("decode", {})), normalized["decode_policy"] + return deepcopy(normalized.get("prefill", {})), normalized["decode_policy"] diff --git a/src/chop/nn/quantized/modules/phase_context.py b/src/chop/nn/quantized/modules/phase_context.py new file mode 100644 index 000000000..884d75eb8 --- /dev/null +++ b/src/chop/nn/quantized/modules/phase_context.py @@ -0,0 +1,141 @@ +"""Shared runtime phase context for quantized decoder-only inference. + +This module intentionally centralizes phase state (`prefill` vs `decode`) in a +single place so quantized modules can stay loosely coupled: + +1. Llama decoder-layer pre-hooks detect runtime phase from cache semantics. +2. Downstream quantized modules (attention/linear/mlp/rms) read the same phase + without changing + their public `forward(...)` signatures. + +Why `ContextVar` is used: +- It avoids global mutable state bleeding across threads/tasks. +- It keeps the integration minimally invasive for the existing MASE module API. +""" + +from __future__ import annotations + +from contextvars import ContextVar +from typing import Any, Literal + +Phase = Literal["prefill", "decode"] +DecodePolicy = Literal["fp_only", "quantized"] + + +_ACTIVE_PHASE: ContextVar[Phase] = ContextVar("active_quant_phase", default="prefill") +_DECODE_POLICY: ContextVar[DecodePolicy] = ContextVar( + "active_decode_policy", default="fp_only" +) + + +def set_runtime_phase(phase: Phase) -> None: + """Set runtime phase used by phase-aware quantized modules.""" + + if phase not in ("prefill", "decode"): + raise ValueError(f"Unsupported phase: {phase}") + _ACTIVE_PHASE.set(phase) + + +def get_runtime_phase() -> Phase: + """Return current runtime phase. + + Defaults to `prefill` when no explicit phase has been written yet. + """ + + return _ACTIVE_PHASE.get() + + +def set_runtime_decode_policy(policy: DecodePolicy) -> None: + """Set decode policy in runtime phase context. + + Supported policies: + - `fp_only`: decode path is forced to full precision. + - `quantized`: decode path may consume decode-phase quantized configs. + + Why we validate here: + - This context is shared by all quantized modules during forward. + - Central validation prevents silent fallback to unintended behavior. + """ + + if policy not in ("fp_only", "quantized"): + raise ValueError(f"Unsupported decode policy: {policy}") + _DECODE_POLICY.set(policy) + + +def get_runtime_decode_policy() -> DecodePolicy: + """Return decode policy from runtime phase context. + + Defaults to `fp_only` to preserve backward compatibility. + """ + + return _DECODE_POLICY.get() + + +def infer_runtime_phase_from_hidden_and_cache( + hidden_states: Any, past_cache: Any +) -> Phase: + """Infer runtime phase from hidden-state shape and cache state. + + Rules are intentionally unchanged from step-1 attention-local logic: + 1) `past_cache is None` -> `prefill` + 2) `past_cache.get_seq_length() > 0` -> `decode` + 3) Fallback: if current query length is 1 -> `decode`, else `prefill` + + Why this helper exists: + - The same heuristic must be shared by all Llama decoder layers. + - We keep it in the shared phase module so hooks and tests stay consistent. + """ + + if past_cache is None: + return "prefill" + + past_len = 0 + get_seq_length = getattr(past_cache, "get_seq_length", None) + if callable(get_seq_length): + try: + past_len = int(get_seq_length()) + except Exception: # pragma: no cover - defensive for custom Cache impls + past_len = 0 + + if past_len > 0: + return "decode" + + q_len = 0 + shape = getattr(hidden_states, "shape", None) + if shape is not None and len(shape) >= 2: + q_len = int(shape[-2]) + if q_len == 1: + return "decode" + return "prefill" + + +def extract_decoder_layer_past_cache( + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> Any: + """Extract past-cache object from decoder-layer call arguments. + + Compatibility: + - New HF naming: `past_key_values` + - Legacy alias: `past_key_value` + - Positional call fallback: 4th argument for `LlamaDecoderLayer.forward` + """ + + if "past_key_values" in kwargs: + return kwargs["past_key_values"] + if "past_key_value" in kwargs: + return kwargs["past_key_value"] + if len(args) >= 4: + return args[3] + return None + + +def infer_runtime_phase_from_decoder_layer_inputs( + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> Phase: + """Infer runtime phase from decoder-layer forward inputs.""" + + hidden_states = kwargs.get("hidden_states", args[0] if args else None) + past_cache = extract_decoder_layer_past_cache(args, kwargs) + return infer_runtime_phase_from_hidden_and_cache(hidden_states, past_cache) diff --git a/src/chop/passes/module/module_modify_helper.py b/src/chop/passes/module/module_modify_helper.py index 441457d9b..b5c337210 100644 --- a/src/chop/passes/module/module_modify_helper.py +++ b/src/chop/passes/module/module_modify_helper.py @@ -103,6 +103,7 @@ def check_module_instance(module, prefix_map): def weight_replacement(x, y): target_state_dict = deepcopy(x.state_dict()) missing_keys, unexpected_keys = y.load_state_dict(target_state_dict, strict=False) + _restore_decode_fp_snapshot_if_available(x, y) if missing_keys: logging.warning( f"Missing keys when loading state_dict: {missing_keys} from {x} to {y}" @@ -114,6 +115,57 @@ def weight_replacement(x, y): return y +def _restore_decode_fp_snapshot_if_available(source_module, target_module): + """Restore decode FP snapshot onto phase-aware linear targets. + + Why this hook exists: + - GPTQ pre-pass rewrites `nn.Linear.weight` in-place before replacement. + - In `fp_only` mode, decode needs preserved FP snapshots. + - In `quantized` mode, we avoid keeping FP decode banks to save memory. + """ + + decode_weight = getattr(source_module, "_mase_decode_weight_fp", None) + decode_bias = getattr(source_module, "_mase_decode_bias_fp", None) + + target_has_decode_weight = hasattr(target_module, "_decode_weight_fp") + target_has_decode_bias = hasattr(target_module, "_decode_bias_fp") + if not target_has_decode_weight: + return + + decode_policy = getattr(target_module, "decode_policy", None) + if decode_policy == "quantized": + # Memory policy: quantized decode should not keep an additional FP + # decode bank alive after replacement. + target_module._decode_weight_fp = torch.empty( + 0, device=target_module.weight.device + ) + if target_has_decode_bias: + target_module._decode_bias_fp = torch.empty( + 0, device=target_module.weight.device + ) + elif decode_policy == "fp_only": + if decode_weight is not None: + target_module._decode_weight_fp = decode_weight.to( + device=target_module.weight.device, + dtype=target_module.weight.dtype, + ) + if ( + target_has_decode_bias + and decode_bias is not None + and getattr(target_module, "bias", None) is not None + ): + target_module._decode_bias_fp = decode_bias.to( + device=target_module.weight.device, + dtype=target_module.weight.dtype, + ) + + refresh_decode_runtime_bank = getattr( + target_module, "refresh_decode_runtime_bank", None + ) + if callable(refresh_decode_runtime_bank): + refresh_decode_runtime_bank() + + def get_module_by_name(network, name): return network.get_submodule(name) # names = name.split(sep='.') @@ -272,6 +324,21 @@ def instantiate_llama_module( layer_idx=module.layer_idx if hasattr(module, "layer_idx") else None, q_config=module_args, ) + # Keep replacement dtype/device aligned with the original HF module. + # Why this is required: + # - Llama quantized wrappers (RMS/MLP/Attention) can be constructed with + # default FP32 parameters. + # - In mixed replacement flows, downstream quantized linear modules may be + # FP16/BF16, so FP32 activations from wrappers can trigger runtime dtype + # mismatches at F.linear boundaries. + # - Aligning here makes replacement behavior consistent with instantiate_linear + # and avoids scattering ad-hoc casts in forward paths. + ref_param = next(module.parameters(), None) + if ref_param is not None: + llama_module = llama_module.to( + device=ref_param.device, + dtype=ref_param.dtype, + ) return llama_module diff --git a/src/chop/passes/module/transforms/gptq/run.py b/src/chop/passes/module/transforms/gptq/run.py index f4873f5f3..eb26f8a2a 100644 --- a/src/chop/passes/module/transforms/gptq/run.py +++ b/src/chop/passes/module/transforms/gptq/run.py @@ -32,7 +32,7 @@ def run_gptq(network, gptq_config): device: str - e.g. "cuda:0". dataset: str - "wikitext2" | "c4" | "ptb". nsamples: int - calibration samples (default 128). - seqlen: int - sequence length (default 2048). + seqlen: int - sequence length (default 256). format: str - "mxfp" | "mxint". weight_config: dict - Mase-style weight config, e.g. {"weight_block_size": 32, "weight_exponent_width": 2, "weight_frac_width": 1} @@ -51,7 +51,9 @@ def run_gptq(network, gptq_config): dev = gptq_config.get("device", "cuda:0") dataset = gptq_config.get("dataset", "wikitext2") nsamples = gptq_config.get("nsamples", 128) - seqlen = gptq_config.get("seqlen", 2048) + # Keep a conservative default to avoid unnecessary memory pressure during + # integration validation and small-run experiments. + seqlen = gptq_config.get("seqlen", 256) fmt = gptq_config["format"] weight_config = gptq_config["weight_config"] quantile_search = gptq_config.get("quantile_search", True) @@ -89,6 +91,18 @@ def run_gptq(network, gptq_config): layers = network.model.layers + # Preserve decode-time FP weights before GPTQ mutates any linear weights. + # This supports step-1 phase policy: prefill may use GPTQ, decode must use FP. + for layer in layers: + for module in layer.modules(): + if isinstance(module, nn.Linear): + if not hasattr(module, "_mase_decode_weight_fp"): + module._mase_decode_weight_fp = module.weight.detach().clone().cpu() + if module.bias is not None and not hasattr( + module, "_mase_decode_bias_fp" + ): + module._mase_decode_bias_fp = module.bias.detach().clone().cpu() + # Move embedding + norm + rope to device network.model.embed_tokens = network.model.embed_tokens.to(dev) network.model.norm = network.model.norm.to(dev) @@ -134,7 +148,8 @@ def forward(self, inp, **kwargs): logging.warning( "GPTQ requested nsamples=%d but only collected=%d; " "continuing with collected samples.", - nsamples, collected, + nsamples, + collected, ) nsamples = collected inps = inps[:nsamples] diff --git a/src/chop/passes/module/transforms/quantize/quantize.py b/src/chop/passes/module/transforms/quantize/quantize.py index c7272d27f..3a4303a2f 100644 --- a/src/chop/passes/module/transforms/quantize/quantize.py +++ b/src/chop/passes/module/transforms/quantize/quantize.py @@ -1,9 +1,62 @@ +"""Quantization transform pass. + +Llama phase-split integration note: +- This pass remains a thin wiring layer. +- Runtime phase execution stays in quantized modules. +- This pass only normalizes configs and writes a consistent runtime decode + policy so downstream modules do not diverge silently. +""" + +from copy import deepcopy +from functools import partial + import torch +from transformers.models.llama.modeling_llama import LlamaDecoderLayer from chop.nn.quantized.modules import quantized_module_map +from chop.nn.quantized.modules.phase_config import normalize_phase_q_config +from chop.nn.quantized.modules.phase_context import ( + infer_runtime_phase_from_decoder_layer_inputs, + set_runtime_phase, + set_runtime_decode_policy, +) from ...module_modify_helper import replace_by_name, instantiate_module from ...state_dict_map import match_a_pattern, check_is_huggingface_model +_LLAMA_PHASE_CONTEXT_POLICY_MODULE_CLASS_NAMES = { + "LlamaAttentionMXFP", + "LlamaAttentionMXInt", + "LlamaMLPMXFP", + "LlamaMLPMXInt", + "LlamaRMSNormMinifloat", + "LinearMXFP", + "LinearMXInt", +} +"""Class names that participate in Llama phase-policy runtime wiring. + +Why include LinearMX*: +- Linear modules consume runtime phase/decode policy during forward. +- Hook installation must trigger even for linear-only quantization configs. +- Reusing one list for detection + policy inference prevents drift. +""" + + +def _normalize_quantize_module_config(config: dict, postfix: str) -> dict: + """Prepare module config before module instantiation. + + Why this helper exists: + - `quantize` pass historically expected a flat config dict. + - Step-1 phase split introduces optional `{prefill, decode}` structure. + - We normalize only for phase-aware quantizers so other quantizers remain + behavior-identical. + """ + + cfg = deepcopy(config) + phase_aware_postfixes = {"mxfp", "mxint", "minifloat"} + if postfix in phase_aware_postfixes: + return normalize_phase_q_config(cfg) + return cfg + def get_config(config: dict, name: str): if name in config: @@ -12,6 +65,99 @@ def get_config(config: dict, name: str): return config["default"]["config"] +def _has_llama_quantized_runtime_modules(network) -> bool: + """Return True when network contains quantized Llama modules. + + Why name-based detection: + - It avoids additional direct imports from quantized module files. + - It keeps this pass decoupled from specific class symbols while still + matching the module-replacement products used in step-1. + """ + + for module in network.modules(): + if module.__class__.__name__ in _LLAMA_PHASE_CONTEXT_POLICY_MODULE_CLASS_NAMES: + return True + return False + + +def _infer_llama_decode_policy_from_quantized_modules(network) -> str: + """Infer a single decode policy from quantized Llama modules. + + Why this is fail-fast: + - Phase context is global per forward call. + - If different Llama quantized modules disagree on decode policy, runtime + behavior becomes order-dependent and difficult to debug. + """ + + observed_policies = set() + for module in network.modules(): + if ( + module.__class__.__name__ + not in _LLAMA_PHASE_CONTEXT_POLICY_MODULE_CLASS_NAMES + ): + continue + policy = getattr(module, "decode_policy", None) + if policy is None: + continue + if policy not in ("fp_only", "quantized"): + raise ValueError( + f"Unsupported decode_policy={policy!r} on " + f"{module.__class__.__name__}." + ) + observed_policies.add(policy) + + if not observed_policies: + return "fp_only" + if len(observed_policies) > 1: + raise ValueError( + "Mixed decode policies detected across quantized Llama modules: " + f"{sorted(observed_policies)}. Use a single decode_policy." + ) + return next(iter(observed_policies)) + + +def _llama_phase_context_pre_hook(module, args, kwargs, decode_policy): + """Set runtime phase before decoder-layer body executes. + + Hook timing is critical: it runs before `input_layernorm`, ensuring modules + that execute before attention still observe correct phase in step-1. + """ + + phase = infer_runtime_phase_from_decoder_layer_inputs(args, kwargs) + set_runtime_phase(phase) + # Keep decode policy explicit in context so all downstream modules apply + # the same policy in this forward call. + set_runtime_decode_policy(decode_policy) + return None + + +def _install_llama_phase_context_pre_hooks(network) -> None: + """Install idempotent phase pre-hooks on all Llama decoder layers.""" + + if not _has_llama_quantized_runtime_modules(network): + return + decode_policy = _infer_llama_decode_policy_from_quantized_modules(network) + + for module in network.modules(): + if not isinstance(module, LlamaDecoderLayer): + continue + if getattr(module, "_mase_phase_hook_installed", False): + prev_policy = getattr(module, "_mase_phase_hook_decode_policy", None) + if prev_policy is not None and prev_policy != decode_policy: + raise ValueError( + "Llama phase hook already installed with decode_policy=" + f"{prev_policy!r}, but current network resolves to " + f"{decode_policy!r}. Rebuild model to avoid mixed policies." + ) + continue + module.register_forward_pre_hook( + partial(_llama_phase_context_pre_hook, decode_policy=decode_policy), + with_kwargs=True, + ) + module._mase_phase_hook_installed = True + module._mase_phase_hook_decode_policy = decode_policy + + def quantize_by_type(network, pass_args): for type_name, config in pass_args.items(): n_m = {} @@ -24,8 +170,9 @@ def quantize_by_type(network, pass_args): module = torch.nn.Conv2d else: raise ValueError(f"{type_name} is not supported!") - config = config["config"] + config = deepcopy(config["config"]) postfix = config.pop("name") + config = _normalize_quantize_module_config(config, postfix) for n, m in n_m.items(): if isinstance(m, module): new_m = instantiate_module( @@ -44,10 +191,9 @@ def quantize_by_name(network, pass_args): n_m[n] = m for n, m in n_m.items(): if n in quantize_names: - quan_config = pass_args[n] - - quan_config = quan_config["config"] + quan_config = deepcopy(pass_args[n]["config"]) postfix = quan_config.pop("name") + quan_config = _normalize_quantize_module_config(quan_config, postfix) additional_module_args = ( {"config": quan_config, "network_config": network.config} @@ -75,8 +221,9 @@ def quantize_by_regex_name(network, pass_args): if not matched_pattern: continue - quan_config = pass_args[matched_pattern]["config"] + quan_config = deepcopy(pass_args[matched_pattern]["config"]) postfix = quan_config["name"] + quan_config = _normalize_quantize_module_config(quan_config, postfix) additional_module_args = ( {"config": quan_config, "network_config": network.config} @@ -130,19 +277,11 @@ def quantize_module_transform_pass(network, pass_args): :raises ValueError: If the quantize "by" argument is unsupported. """ - # If TOML has a [rotation_search] block, route the WHOLE quantize step - # through the rotation search pass — it handles GPTQ, baseline module - # replacement, and per-matmul rotate flag tuning end-to-end. Decisions - # are cached to disk (default /rotation_decisions.json) - # so a re-run skips the calib forwards entirely (mirrors GPTQ's - # checkpoint resume). - if "rotation_search" in pass_args: - from .rotation_search import dispatch_rotation_search_block - - rot_cfg = pass_args.pop("rotation_search") - return dispatch_rotation_search_block(network, pass_args, rot_cfg) - - # GPTQ pre-pass: quantize linear weights before module replacement + # Defensive copy avoids mutating caller-owned pass_args, which is + # important for reproducible experiment runners that reuse config dicts. + pass_args = deepcopy(pass_args) + + # GPTQ pre-pass: quantize linear weights before module replacement. gptq_config = pass_args.pop("gptq", None) if gptq_config is not None: from ..gptq import run_gptq @@ -160,4 +299,8 @@ def quantize_module_transform_pass(network, pass_args): case _: raise ValueError(f'Unsupported quantize "by": {by}') + # Install phase hooks only after module replacement, so detection sees + # quantized Llama modules rather than original HF modules. + _install_llama_phase_context_pre_hooks(network) + return network, {} diff --git a/test/nn/quantized/modules/test_llama_phase_split_step1.py b/test/nn/quantized/modules/test_llama_phase_split_step1.py new file mode 100644 index 000000000..e8e9cedc1 --- /dev/null +++ b/test/nn/quantized/modules/test_llama_phase_split_step1.py @@ -0,0 +1,378 @@ +"""Step-1 phase-split regression tests for quantized Llama integration. + +These tests intentionally stay small and deterministic: +- no model downloads +- no full-model forward +- only boundary behavior needed to protect step-1 invariants +""" + +from __future__ import annotations + +import sys +import types +import importlib.util +import inspect +from pathlib import Path +from functools import partial + +import torch +import pytest +from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.models.llama.modeling_llama import LlamaDecoderLayer + +# Keep unit tests lightweight in environments without tensorboard dependency. +if "torch.utils.tensorboard" not in sys.modules: + tensorboard_stub = types.ModuleType("torch.utils.tensorboard") + tensorboard_stub.SummaryWriter = object + sys.modules["torch.utils.tensorboard"] = tensorboard_stub +if "cvxpy" not in sys.modules: + sys.modules["cvxpy"] = types.ModuleType("cvxpy") + +# Bypass heavyweight `chop/__init__.py` side effects during focused unit tests. +repo_root = next( + p for p in Path(__file__).resolve().parents if (p / "src/chop").exists() +) +if "chop" not in sys.modules: + chop_stub = types.ModuleType("chop") + chop_stub.__path__ = [str(repo_root / "src/chop")] + sys.modules["chop"] = chop_stub + +_helper_spec = importlib.util.spec_from_file_location( + "mase_module_modify_helper", + repo_root / "src/chop/passes/module/module_modify_helper.py", +) +_helper_module = importlib.util.module_from_spec(_helper_spec) +assert _helper_spec is not None and _helper_spec.loader is not None +_helper_spec.loader.exec_module(_helper_module) +weight_replacement = _helper_module.weight_replacement + +from chop.nn.quantized.modules.linear import LinearMXFP +import chop.nn.quantized.modules.llama.attention as llama_attention_module +from chop.nn.quantized.modules.llama.attention import LlamaAttentionMXFP +from chop.nn.quantized.modules.phase_config import normalize_phase_q_config +from chop.nn.quantized.modules.phase_context import ( + set_runtime_phase, + infer_runtime_phase_from_hidden_and_cache, + get_runtime_phase, + infer_runtime_phase_from_decoder_layer_inputs, + get_runtime_decode_policy, +) +from chop.passes.module.transforms.quantize.quantize import ( + _llama_phase_context_pre_hook, +) + + +class _DummyCache: + """Small cache stub for phase inference tests.""" + + def __init__(self, seq_len: int): + self._seq_len = seq_len + + def get_seq_length(self) -> int: + return self._seq_len + + +def test_normalize_phase_q_config_legacy_compatibility(): + """Legacy flat config should map to both prefill/decode buckets.""" + + legacy = {"data_in_block_size": 16, "bypass": False} + normalized = normalize_phase_q_config(legacy) + assert normalized["decode_policy"] == "fp_only" + assert normalized["prefill"]["data_in_block_size"] == 16 + assert normalized["decode"]["data_in_block_size"] == 16 + + +def test_infer_runtime_phase_from_cache_state(): + """Shared phase inference should follow cache semantics.""" + + hidden_prefill = torch.randn(1, 8, 16) + hidden_decode = torch.randn(1, 1, 16) + + assert infer_runtime_phase_from_hidden_and_cache(hidden_prefill, None) == "prefill" + assert ( + infer_runtime_phase_from_hidden_and_cache( + hidden_prefill, _DummyCache(seq_len=0) + ) + == "prefill" + ) + assert ( + infer_runtime_phase_from_hidden_and_cache(hidden_decode, _DummyCache(seq_len=0)) + == "decode" + ) + assert ( + infer_runtime_phase_from_hidden_and_cache( + hidden_prefill, _DummyCache(seq_len=32) + ) + == "decode" + ) + + +def test_decoder_layer_pre_hook_sets_phase_before_input_layernorm(): + """Pre-hook should set phase before `input_layernorm` executes.""" + + cfg = LlamaConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + ) + layer = LlamaDecoderLayer(cfg, layer_idx=0) + + observed = {"phase_at_input_ln": None} + orig_ln_forward = layer.input_layernorm.forward + + def _spy_input_ln(hidden_states): + observed["phase_at_input_ln"] = get_runtime_phase() + return orig_ln_forward(hidden_states) + + layer.input_layernorm.forward = _spy_input_ln + + # Keep attention/MLP simple to avoid requiring full model plumbing. + class _DummySelfAttn(torch.nn.Module): + def forward(self, hidden_states, **kwargs): + return hidden_states, None + + class _DummyMLP(torch.nn.Module): + def forward(self, hidden_states): + return hidden_states + + layer.self_attn = _DummySelfAttn() + layer.mlp = _DummyMLP() + layer.register_forward_pre_hook( + partial(_llama_phase_context_pre_hook, decode_policy="fp_only"), + with_kwargs=True, + ) + + hidden = torch.randn(1, 1, cfg.hidden_size) + set_runtime_phase("prefill") + layer( + hidden_states=hidden, + past_key_values=_DummyCache(seq_len=8), + ) + + assert observed["phase_at_input_ln"] == "decode" + assert get_runtime_decode_policy() == "fp_only" + + +def test_decoder_layer_input_extraction_supports_old_and_new_cache_names(): + """Phase extraction must support both HF naming variants.""" + + hidden = torch.randn(1, 1, 16) + cache = _DummyCache(seq_len=3) + + phase_new = infer_runtime_phase_from_decoder_layer_inputs( + args=(), + kwargs={"hidden_states": hidden, "past_key_values": cache}, + ) + phase_old = infer_runtime_phase_from_decoder_layer_inputs( + args=(), + kwargs={"hidden_states": hidden, "past_key_value": cache}, + ) + + assert phase_new == "decode" + assert phase_old == "decode" + + +def test_attention_no_longer_writes_runtime_phase(): + """Guardrail: attention should consume phase context, not mutate it.""" + + src = inspect.getsource(LlamaAttentionMXFP.forward) + assert "set_runtime_phase(" not in src + + +def test_decoder_layer_pre_hook_can_propagate_quantized_decode_policy(): + """Hook should write explicit decode policy for downstream modules.""" + + cfg = LlamaConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + ) + layer = LlamaDecoderLayer(cfg, layer_idx=0) + layer.register_forward_pre_hook( + partial(_llama_phase_context_pre_hook, decode_policy="quantized"), + with_kwargs=True, + ) + + class _DummySelfAttn(torch.nn.Module): + def forward(self, hidden_states, **kwargs): + return hidden_states, None + + class _DummyMLP(torch.nn.Module): + def forward(self, hidden_states): + return hidden_states + + layer.self_attn = _DummySelfAttn() + layer.mlp = _DummyMLP() + layer( + hidden_states=torch.randn(1, 1, cfg.hidden_size), + past_key_values=_DummyCache(seq_len=8), + ) + assert get_runtime_decode_policy() == "quantized" + + +def test_linear_mxfp_decode_uses_fp_snapshot_after_weight_replacement(): + """Decode path must use preserved FP weight even when prefill weights differ. + + Why this matters: + GPTQ mutates source linear weights before module replacement. This test + verifies the replacement seam restores the original FP snapshot for decode. + """ + + source = torch.nn.Linear(4, 3, bias=True) + source_weight_quantized = torch.full_like(source.weight, 5.0) + source_bias_quantized = torch.full_like(source.bias, -2.0) + source.weight.data.copy_(source_weight_quantized) + source.bias.data.copy_(source_bias_quantized) + + # Simulate snapshots captured before GPTQ in-place mutation. + fp_weight = torch.full_like(source.weight, 2.0) + fp_bias = torch.full_like(source.bias, 1.0) + source._mase_decode_weight_fp = fp_weight.detach().clone() + source._mase_decode_bias_fp = fp_bias.detach().clone() + + target = LinearMXFP( + in_features=4, + out_features=3, + bias=True, + config={"bypass": True}, + ) + target = weight_replacement(source, target) + + x = torch.randn(2, 4) + + set_runtime_phase("prefill") + out_prefill = target(x) + expected_prefill = torch.nn.functional.linear( + x, source_weight_quantized, source_bias_quantized + ) + assert torch.allclose(out_prefill, expected_prefill, atol=1e-6, rtol=0) + + set_runtime_phase("decode") + out_decode = target(x) + expected_decode = torch.nn.functional.linear(x, fp_weight, fp_bias) + assert torch.allclose(out_decode, expected_decode, atol=1e-6, rtol=0) + # fp_only mode should not build a quantized decode bank. + assert target._decode_weight_q.numel() == 0 + + # Keep global test context deterministic for subsequent tests. + set_runtime_phase("prefill") + + +def test_linear_mxfp_decode_quantized_uses_decode_quant_bank(): + """`decode_policy=quantized` should route decode through decode quant bank.""" + + target = LinearMXFP( + in_features=8, + out_features=4, + bias=True, + config={ + "decode_policy": "quantized", + "prefill": {"bypass": True}, + "decode": { + "bypass": False, + "data_in_block_size": 8, + "data_in_exponent_width": 4, + "data_in_frac_width": 3, + "weight_block_size": 8, + "weight_exponent_width": 4, + "weight_frac_width": 3, + }, + }, + ) + state = target.state_dict() + target.load_state_dict(state, strict=True) + assert target._decode_weight_q.numel() > 0 + assert target._decode_weight_fp.numel() == 0 + + x = torch.randn(2, 8) + set_runtime_phase("decode") + out_decode = target(x) + expected_decode = torch.nn.functional.linear( + x, target._decode_weight_q, target.bias + ) + assert out_decode.shape == expected_decode.shape + + set_runtime_phase("prefill") + + +def test_linear_mxfp_quantized_decode_bypass_uses_current_weight_bank(): + """Quantized policy + decode bypass should not fallback to FP snapshot.""" + + target = LinearMXFP( + in_features=8, + out_features=4, + bias=True, + config={ + "decode_policy": "quantized", + "prefill": {"bypass": True}, + "decode": {"bypass": True}, + }, + ) + state = target.state_dict() + target.load_state_dict(state, strict=True) + assert target._decode_weight_fp.numel() == 0 + assert target._decode_weight_q.numel() == 0 + + x = torch.randn(2, 8) + set_runtime_phase("decode") + out_decode = target(x) + expected_decode = torch.nn.functional.linear(x, target.weight, target.bias) + assert torch.allclose(out_decode, expected_decode, atol=1e-6, rtol=0) + set_runtime_phase("prefill") + + +def test_attention_all_bypass_still_uses_local_eager_path(monkeypatch): + """All-bypass case should stay on local eager path (no backend fallback).""" + + cfg = LlamaConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + attention_dropout=0.0, + ) + attn = LlamaAttentionMXFP( + config=cfg, + layer_idx=0, + q_config={ + "decode_policy": "fp_only", + "prefill": { + "qk_matmul": {"bypass": True}, + "av_matmul": {"bypass": True}, + "rope": {"bypass": True}, + "softmax": {"bypass": True}, + "kv_cache": {"bypass": True}, + }, + "decode": {"bypass": True}, + }, + ) + + called = {"eager": 0} + + def _fake_eager(*args, **kwargs): + called["eager"] += 1 + q = args[1] + out = torch.zeros_like(q.transpose(1, 2)) + return out, None + + monkeypatch.setattr( + llama_attention_module, "_eager_attention_forward_mxfp", _fake_eager + ) + + set_runtime_phase("prefill") + x = torch.randn(1, 2, cfg.hidden_size) + cos = torch.ones(1, 2, attn.head_dim) + sin = torch.zeros(1, 2, attn.head_dim) + attn( + hidden_states=x, + position_embeddings=(cos, sin), + attention_mask=None, + past_key_values=None, + ) + + assert called["eager"] == 1 diff --git a/test/passes/module/transforms/quantize/test_quantize_phase_config.py b/test/passes/module/transforms/quantize/test_quantize_phase_config.py new file mode 100644 index 000000000..3e005825d --- /dev/null +++ b/test/passes/module/transforms/quantize/test_quantize_phase_config.py @@ -0,0 +1,185 @@ +"""Regression tests for phase-structured quantize pass wiring.""" + +from __future__ import annotations + +from copy import deepcopy +import sys +import types +from pathlib import Path + +import torch +import pytest +from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.models.llama.modeling_llama import LlamaDecoderLayer + +# Keep unit tests lightweight in environments without tensorboard dependency. +if "torch.utils.tensorboard" not in sys.modules: + tensorboard_stub = types.ModuleType("torch.utils.tensorboard") + tensorboard_stub.SummaryWriter = object + sys.modules["torch.utils.tensorboard"] = tensorboard_stub +if "cvxpy" not in sys.modules: + sys.modules["cvxpy"] = types.ModuleType("cvxpy") + +# Bypass heavyweight `chop/__init__.py` side effects during focused unit tests. +if "chop" not in sys.modules: + repo_root = next( + p for p in Path(__file__).resolve().parents if (p / "src/chop").exists() + ) + chop_stub = types.ModuleType("chop") + chop_stub.__path__ = [str(repo_root / "src/chop")] + sys.modules["chop"] = chop_stub + +from chop.nn.quantized.modules.linear import LinearMXFP +from chop.passes.module.transforms.quantize.quantize import ( + quantize_module_transform_pass, + _install_llama_phase_context_pre_hooks, + _infer_llama_decode_policy_from_quantized_modules, +) + + +class _TinyMLP(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc1 = torch.nn.Linear(8, 8) + self.fc2 = torch.nn.Linear(8, 8) + + def forward(self, x): + return self.fc2(self.fc1(x)) + + +def test_quantize_pass_accepts_phase_config_for_mxfp_linear(): + """Phase config should be forwarded to phase-aware linear module.""" + + model = _TinyMLP() + pass_args = { + "by": "name", + "fc1": { + "config": { + "name": "mxfp", + "decode_policy": "fp_only", + "prefill": {"bypass": True}, + "decode": {"bypass": True}, + } + }, + } + pass_args_before = deepcopy(pass_args) + + quantized_model, _ = quantize_module_transform_pass(model, pass_args) + assert isinstance(quantized_model.fc1, LinearMXFP) + assert quantized_model.fc1.phase_config["decode_policy"] == "fp_only" + assert quantized_model.fc1.phase_config["prefill"]["bypass"] is True + assert quantized_model.fc1.phase_config["decode"]["bypass"] is True + + # Caller-owned config should remain unchanged after the pass. + assert pass_args == pass_args_before + + +def test_quantize_pass_accepts_quantized_decode_policy(): + """`decode_policy=quantized` should pass parsing and instantiation.""" + + model = _TinyMLP() + pass_args = { + "by": "name", + "fc1": { + "config": { + "name": "mxfp", + "decode_policy": "quantized", + "prefill": {"bypass": True}, + "decode": { + "bypass": False, + "data_in_block_size": 16, + "data_in_exponent_width": 4, + "data_in_frac_width": 3, + "weight_block_size": 16, + "weight_exponent_width": 4, + "weight_frac_width": 3, + }, + } + }, + } + quantized_model, _ = quantize_module_transform_pass(model, pass_args) + assert isinstance(quantized_model.fc1, LinearMXFP) + assert quantized_model.fc1.decode_policy == "quantized" + + +def test_llama_phase_pre_hook_installation_is_gated_and_idempotent(): + """Hooks should install only for quantized Llama runs and only once.""" + + class LlamaAttentionMXFP(torch.nn.Module): + def __init__(self, decode_policy: str = "fp_only"): + super().__init__() + self.decode_policy = decode_policy + + def forward(self, x): + return x + + class LinearMXFP(torch.nn.Module): + def __init__(self, decode_policy: str = "fp_only"): + super().__init__() + self.decode_policy = decode_policy + + def forward(self, x): + return x + + class TinyNetwork(torch.nn.Module): + def __init__( + self, + with_quantized_marker: bool, + marker_class: type[torch.nn.Module] = LlamaAttentionMXFP, + decode_policy: str = "fp_only", + ): + super().__init__() + cfg = LlamaConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + ) + self.layer = LlamaDecoderLayer(cfg, layer_idx=0) + if with_quantized_marker: + self.quant_marker = marker_class(decode_policy=decode_policy) + + net_without_marker = TinyNetwork(with_quantized_marker=False) + before_no_marker = len(net_without_marker.layer._forward_pre_hooks) + _install_llama_phase_context_pre_hooks(net_without_marker) + assert len(net_without_marker.layer._forward_pre_hooks) == before_no_marker + + net_with_marker = TinyNetwork(with_quantized_marker=True) + before_with_marker = len(net_with_marker.layer._forward_pre_hooks) + _install_llama_phase_context_pre_hooks(net_with_marker) + after_first = len(net_with_marker.layer._forward_pre_hooks) + _install_llama_phase_context_pre_hooks(net_with_marker) + after_second = len(net_with_marker.layer._forward_pre_hooks) + + assert after_first == before_with_marker + 1 + assert after_second == after_first + + # Linear-only quantization must also install the hook, because runtime + # phase/decode policy is consumed inside quantized Linear forward paths. + net_with_linear_marker = TinyNetwork( + with_quantized_marker=True, marker_class=LinearMXFP + ) + before_linear_marker = len(net_with_linear_marker.layer._forward_pre_hooks) + _install_llama_phase_context_pre_hooks(net_with_linear_marker) + assert len(net_with_linear_marker.layer._forward_pre_hooks) == ( + before_linear_marker + 1 + ) + + +def test_decode_policy_inference_fails_fast_on_mixed_llama_policies(): + """Mixed decode policies must raise instead of silently choosing one.""" + + class LlamaAttentionMXFP(torch.nn.Module): + def __init__(self, decode_policy: str): + super().__init__() + self.decode_policy = decode_policy + + class TinyNetwork(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = LlamaAttentionMXFP("fp_only") + self.b = LlamaAttentionMXFP("quantized") + + with pytest.raises(ValueError, match="Mixed decode policies"): + _infer_llama_decode_policy_from_quantized_modules(TinyNetwork())