From 70580f10f0b05f3e1e8f06d3c8a0b12e110f7a46 Mon Sep 17 00:00:00 2001 From: Yuxuan Han Date: Mon, 18 May 2026 17:06:05 +0100 Subject: [PATCH 1/6] Phase split quantisation support (prefill side) + Testing --- .../validate_llama_quant_gptq_phase_split.py | 404 +++++++++++ src/chop/nn/quantized/modules/linear.py | 156 ++++- .../nn/quantized/modules/llama/attention.py | 650 ++++++++---------- src/chop/nn/quantized/modules/llama/mlp.py | 53 +- .../nn/quantized/modules/llama/rms_norm.py | 98 ++- src/chop/nn/quantized/modules/phase_config.py | 57 ++ .../nn/quantized/modules/phase_context.py | 134 ++++ .../passes/module/module_modify_helper.py | 50 ++ src/chop/passes/module/transforms/gptq/run.py | 16 +- .../module/transforms/quantize/quantize.py | 183 ++++- .../modules/test_llama_phase_split_step1.py | 213 ++++++ .../quantize/test_quantize_phase_config.py | 107 +++ 12 files changed, 1665 insertions(+), 456 deletions(-) create mode 100644 scripts/validate_llama_quant_gptq_phase_split.py create mode 100644 src/chop/nn/quantized/modules/phase_config.py create mode 100644 src/chop/nn/quantized/modules/phase_context.py create mode 100644 test/nn/quantized/modules/test_llama_phase_split_step1.py create mode 100644 test/passes/module/transforms/quantize/test_quantize_phase_config.py 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..d375aeb62 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_active_phase +from chop.nn.quantized.modules.phase_config import ( + get_phase_subconfig, + normalize_phase_q_config, +) from ..utils import get_stats, quantiser_passthrough @@ -819,6 +824,17 @@ def forward(self, x): class LinearMXFP(_LinearBase): + """MXFP linear with prefill/decode phase-aware execution. + + Step-1 policy for maintainable migration: + - Prefill can run quantized path. + - Decode is forced to full precision (`decode_policy='fp_only'`). + + Why dual weight representation exists: + - `self.weight` stores prefill weight (PTQ/GPTQ output). + - `_decode_weight_fp` stores original FP weight snapshot for decode path. + """ + # NOTE: backward is not supported — inference only (PTQ) def __init__( self, @@ -831,10 +847,32 @@ 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 != "fp_only": + raise ValueError( + "Step-1 integration only supports decode_policy='fp_only' " + f"for {self.__class__.__name__}, got {self.decode_policy!r}." + ) + + 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) + + # Non-persistent buffers keep decode FP weights without polluting + # serialized checkpoints. + self.register_buffer("_decode_weight_fp", torch.empty(0), persistent=False) + self.register_buffer("_decode_bias_fp", torch.empty(0), persistent=False) + + def _capture_decode_fp_snapshot(self, state_dict) -> None: + """Save FP weights for decode before any prefill quantization mutates data.""" + + 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": @@ -895,15 +933,16 @@ def from_linear(cls, linear: torch.nn.Linear, config: dict) -> "LinearMXFP": def load_state_dict(self, state_dict, strict=True, assign=False): """Load pretrained weights, then quantize them in place.""" + self._capture_decode_fp_snapshot(state_dict) result = super().load_state_dict(state_dict, strict=strict, assign=assign) if self.bypass or self.gptq: 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 +954,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 @@ -935,15 +974,37 @@ def load_state_dict(self, state_dict, strict=True, assign=False): return result + def _is_decode_fp_mode(self) -> bool: + """Return True when current call must use decode full-precision path.""" + + phase = get_active_phase() + return phase == "decode" and self.decode_policy == "fp_only" + @torch.no_grad() def forward(self, x): - if self.bypass: + if self._is_decode_fp_mode(): + # Why this branch exists: + # decode must stay full precision in step-1 even when prefill is + # quantized/GPTQ. This guarantees deterministic behavior for staged + # rollout and easier upstream rebases. + 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) + + prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") + if prefill_cfg.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_cfg.get("data_in_block_size") + x_exp_bits = prefill_cfg.get("data_in_exponent_width") + x_frac_bits = prefill_cfg.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 +1018,12 @@ def forward(self, x): class LinearMXInt(_LinearBase): + """MXInt linear with prefill/decode phase-aware execution. + + Behavior mirrors `LinearMXFP`: prefill may be quantized, decode is forced to + FP using cached original weights/bias. + """ + # NOTE: backward is not supported — inference only (PTQ) def __init__( self, @@ -969,10 +1036,30 @@ 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 != "fp_only": + raise ValueError( + "Step-1 integration only supports decode_policy='fp_only' " + f"for {self.__class__.__name__}, got {self.decode_policy!r}." + ) + + 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) + + def _capture_decode_fp_snapshot(self, state_dict) -> None: + """Save FP weights for decode before prefill quantization mutates data.""" + + 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) -> "LinearMXInt": @@ -1043,14 +1130,15 @@ def from_linear(cls, linear: torch.nn.Linear, config: dict) -> "LinearMXInt": def load_state_dict(self, state_dict, strict=True, assign=False): """Load pretrained weights, then quantize them in place.""" + self._capture_decode_fp_snapshot(state_dict) result = super().load_state_dict(state_dict, strict=strict, assign=assign) if self.bypass or self.gptq: 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 +1150,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 @@ -1081,14 +1169,32 @@ def load_state_dict(self, state_dict, strict=True, assign=False): return result + def _is_decode_fp_mode(self) -> bool: + """Return True when current call must use decode full-precision path.""" + + phase = get_active_phase() + return phase == "decode" and self.decode_policy == "fp_only" + @torch.no_grad() def forward(self, x): - if self.bypass: + if self._is_decode_fp_mode(): + 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) + + prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") + if prefill_cfg.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_cfg.get("data_in_block_size") + x_element_bits = prefill_cfg.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( diff --git a/src/chop/nn/quantized/modules/llama/attention.py b/src/chop/nn/quantized/modules/llama/attention.py index 54525848d..f7e01b0fe 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,11 @@ 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 +231,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 +247,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 +306,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 +352,11 @@ 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 +365,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 +381,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 +435,127 @@ 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) - """ - - 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) - - 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) - - 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, - ) - - 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, - ) + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + 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, + ) -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: + 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) - 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_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights - 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) - 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) +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) - 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) + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - 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, - ) + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask - 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, - ) + 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, + ) - 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, - ) + 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) - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights diff --git a/src/chop/nn/quantized/modules/llama/mlp.py b/src/chop/nn/quantized/modules/llama/mlp.py index 533504da5..5b67837fc 100644 --- a/src/chop/nn/quantized/modules/llama/mlp.py +++ b/src/chop/nn/quantized/modules/llama/mlp.py @@ -1,8 +1,23 @@ +"""Llama MLP quantization with phase-aware dispatch. + +Step-1 policy is explicit: +- prefill: quantization path can run +- decode: full precision only + +Phase source: +- Runtime phase is written by decoder-layer pre-hooks before layer execution. +""" + 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_active_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 +66,24 @@ 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"] + if self.decode_policy != "fp_only": + raise ValueError( + "Step-1 integration only supports decode_policy='fp_only' " + f"for {self.__class__.__name__}, got {self.decode_policy!r}." + ) def forward(self, x: Tensor) -> Tensor: - if self.bypass: + phase = get_active_phase() + sub_cfg, decode_policy = get_phase_subconfig(self.phase_q_config, phase) + bypass = sub_cfg.get("bypass", False) + if 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), sub_cfg) * self.up_proj(x) return self.down_proj(x) @@ -67,11 +93,22 @@ 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"] + if self.decode_policy != "fp_only": + raise ValueError( + "Step-1 integration only supports decode_policy='fp_only' " + f"for {self.__class__.__name__}, got {self.decode_policy!r}." + ) def forward(self, x: Tensor) -> Tensor: - if self.bypass: + phase = get_active_phase() + sub_cfg, decode_policy = get_phase_subconfig(self.phase_q_config, phase) + bypass = sub_cfg.get("bypass", False) + if 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), sub_cfg) * 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..633f9422d 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. + +Decode is intentionally forced to FP in step-1 for deterministic integration. +Runtime phase is supplied by decoder-layer pre-hooks via phase context. +""" + 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_active_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,69 @@ 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"), - ), + self.phase_q_config = normalize_phase_q_config(q_config) + self.decode_policy = self.phase_q_config["decode_policy"] + if self.decode_policy != "fp_only": + raise ValueError( + "Step-1 integration only supports decode_policy='fp_only' " + f"for {self.__class__.__name__}, got {self.decode_policy!r}." ) - else: - self.x_quantizer = None + + @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): + phase = get_active_phase() + sub_cfg, decode_policy = get_phase_subconfig(self.phase_q_config, phase) + bypass = sub_cfg.get("bypass", False) + if phase == "decode" and decode_policy == "fp_only": + # Why force bypass here: + # step-1 intentionally keeps decode fully FP for stability and + # backward compatibility while still accepting phase-shaped configs. + bypass = True + + w_quantizer = self._build_weight_quantizer(sub_cfg, bypass) + x_quantizer = self._build_input_quantizer(sub_cfg, 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 + 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..69a5dd74d --- /dev/null +++ b/src/chop/nn/quantized/modules/phase_context.py @@ -0,0 +1,134 @@ +"""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"] + + +_ACTIVE_PHASE: ContextVar[Phase] = ContextVar("active_quant_phase", default="prefill") +_DECODE_POLICY: ContextVar[DecodePolicy] = ContextVar( + "active_decode_policy", default="fp_only" +) + + +def set_active_phase(phase: Phase) -> None: + """Set the current runtime phase for quantized module dispatch.""" + + if phase not in ("prefill", "decode"): + raise ValueError(f"Unsupported phase: {phase}") + _ACTIVE_PHASE.set(phase) + + +def get_active_phase() -> Phase: + """Return the current runtime phase. + + Defaults to `prefill` when no explicit phase has been written yet. + """ + + return _ACTIVE_PHASE.get() + + +def set_decode_policy(policy: DecodePolicy) -> None: + """Set decode policy for current runtime context. + + Current step intentionally supports only `fp_only` to keep the first + migration minimal and deterministic. + """ + + if policy != "fp_only": + raise ValueError(f"Unsupported decode policy: {policy}") + _DECODE_POLICY.set(policy) + + +def get_decode_policy() -> DecodePolicy: + """Return decode policy for current runtime context. + + Defaults to `fp_only`, matching step-1 product decision. + """ + + return _DECODE_POLICY.get() + + +def infer_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_past_cache_from_decoder_layer_inputs( + 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_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_past_cache_from_decoder_layer_inputs(args, kwargs) + return infer_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..8bf751f43 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,40 @@ 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. + - Step-1 requires decode to stay FP (including weight path). + - We therefore preserve original FP snapshots on source modules and + transfer them here when the target quantized module exposes decode buffers. + """ + + 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 + + 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, + ) + + def get_module_by_name(network, name): return network.get_submodule(name) # names = name.split(sep='.') @@ -272,6 +307,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..5d0ffeda9 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,16 @@ 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) diff --git a/src/chop/passes/module/transforms/quantize/quantize.py b/src/chop/passes/module/transforms/quantize/quantize.py index c7272d27f..03aa8a3bf 100644 --- a/src/chop/passes/module/transforms/quantize/quantize.py +++ b/src/chop/passes/module/transforms/quantize/quantize.py @@ -1,10 +1,64 @@ +"""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: return config[name]["config"] @@ -12,6 +66,102 @@ 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 +174,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 +195,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 +225,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 +281,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 +303,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..b33d384b4 --- /dev/null +++ b/test/nn/quantized/modules/test_llama_phase_split_step1.py @@ -0,0 +1,213 @@ +"""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 + +import torch +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 +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_active_phase, + infer_phase_from_hidden_and_cache, + get_active_phase, + infer_phase_from_decoder_layer_inputs, +) +from chop.passes.module.transforms.quantize.quantize import ( + _llama_decoder_layer_phase_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_phase_from_hidden_and_cache(hidden_prefill, None) == "prefill" + assert ( + infer_phase_from_hidden_and_cache(hidden_prefill, _DummyCache(seq_len=0)) + == "prefill" + ) + assert ( + infer_phase_from_hidden_and_cache(hidden_decode, _DummyCache(seq_len=0)) + == "decode" + ) + assert ( + infer_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_active_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(_llama_decoder_layer_phase_pre_hook, with_kwargs=True) + + hidden = torch.randn(1, 1, cfg.hidden_size) + set_active_phase("prefill") + layer( + hidden_states=hidden, + past_key_values=_DummyCache(seq_len=8), + ) + + assert observed["phase_at_input_ln"] == "decode" + + +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_phase_from_decoder_layer_inputs( + args=(), + kwargs={"hidden_states": hidden, "past_key_values": cache}, + ) + phase_old = infer_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_active_phase(" not in src + + +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_active_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_active_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) + + # Keep global test context deterministic for subsequent tests. + set_active_phase("prefill") 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..962ba3a65 --- /dev/null +++ b/test/passes/module/transforms/quantize/test_quantize_phase_config.py @@ -0,0 +1,107 @@ +"""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 +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_pre_hooks, +) + + +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_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 forward(self, x): + return x + + class TinyNetwork(torch.nn.Module): + def __init__(self, with_quantized_marker: bool): + 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 = LlamaAttentionMXFP() + + net_without_marker = TinyNetwork(with_quantized_marker=False) + before_no_marker = len(net_without_marker.layer._forward_pre_hooks) + _install_llama_phase_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_pre_hooks(net_with_marker) + after_first = len(net_with_marker.layer._forward_pre_hooks) + _install_llama_phase_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 From 2256d6791e1872bb4db10e7a6015629c305e8bcf Mon Sep 17 00:00:00 2001 From: Yuxuan Han Date: Thu, 21 May 2026 16:38:23 +0100 Subject: [PATCH 2/6] Added support for decode phase quantisation --- src/chop/nn/quantized/modules/linear.py | 480 +++++++----------- src/chop/nn/quantized/modules/llama/mlp.py | 44 +- .../nn/quantized/modules/llama/rms_norm.py | 28 +- .../nn/quantized/modules/phase_context.py | 43 +- .../passes/module/module_modify_helper.py | 49 +- .../modules/test_llama_phase_split_step1.py | 199 +++++++- .../quantize/test_quantize_phase_config.py | 64 ++- third_party/gorilla/gorilla | 1 + 8 files changed, 519 insertions(+), 389 deletions(-) create mode 160000 third_party/gorilla/gorilla diff --git a/src/chop/nn/quantized/modules/linear.py b/src/chop/nn/quantized/modules/linear.py index d375aeb62..8505a0b84 100644 --- a/src/chop/nn/quantized/modules/linear.py +++ b/src/chop/nn/quantized/modules/linear.py @@ -16,7 +16,7 @@ import torch from torch import Tensor from torch.nn import functional as F -from chop.nn.quantized.modules.phase_context import get_active_phase +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, @@ -42,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 @@ -826,13 +822,14 @@ def forward(self, x): class LinearMXFP(_LinearBase): """MXFP linear with prefill/decode phase-aware execution. - Step-1 policy for maintainable migration: - - Prefill can run quantized path. - - Decode is forced to full precision (`decode_policy='fp_only'`). + Runtime phase policy: + - `fp_only`: decode uses preserved FP snapshots. + - `quantized`: decode uses decode-specific quantized weight bank. - Why dual weight representation exists: - - `self.weight` stores prefill weight (PTQ/GPTQ output). - - `_decode_weight_fp` stores original FP weight snapshot for decode path. + 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) @@ -849,10 +846,10 @@ def __init__( assert config is not None, "config is None!" self.phase_config = normalize_phase_q_config(config) self.decode_policy = self.phase_config["decode_policy"] - if self.decode_policy != "fp_only": + if self.decode_policy not in ("fp_only", "quantized"): raise ValueError( - "Step-1 integration only supports decode_policy='fp_only' " - f"for {self.__class__.__name__}, got {self.decode_policy!r}." + "Unsupported decode_policy " + f"{self.decode_policy!r} for {self.__class__.__name__}." ) prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") @@ -861,82 +858,88 @@ def __init__( self.gptq = prefill_cfg.get("gptq", False) self.clip_search = prefill_cfg.get("clip_search", False) - # Non-persistent buffers keep decode FP weights without polluting - # serialized checkpoints. + # 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_snapshot(self, state_dict) -> None: - """Save FP weights for decode before any prefill quantization mutates data.""" + 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 + + 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) - return new + 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_snapshot(state_dict) + 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 @@ -972,21 +975,15 @@ def load_state_dict(self, state_dict, strict=True, assign=False): ) ) + self._build_decode_quantized_bank() return result - def _is_decode_fp_mode(self) -> bool: - """Return True when current call must use decode full-precision path.""" - - phase = get_active_phase() - return phase == "decode" and self.decode_policy == "fp_only" - @torch.no_grad() def forward(self, x): - if self._is_decode_fp_mode(): + runtime_phase = get_runtime_phase() + if runtime_phase == "decode" and self.decode_policy == "fp_only": # Why this branch exists: - # decode must stay full precision in step-1 even when prefill is - # quantized/GPTQ. This guarantees deterministic behavior for staged - # rollout and easier upstream rebases. + # Default path must remain fully backward-compatible. decode_weight = ( self._decode_weight_fp if self._decode_weight_fp.numel() > 0 else self.weight ) @@ -997,14 +994,41 @@ def forward(self, x): ) return F.linear(x, decode_weight, decode_bias) - prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") - if prefill_cfg.get("bypass", False): + 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 = prefill_cfg.get("data_in_block_size") - x_exp_bits = prefill_cfg.get("data_in_exponent_width") - x_frac_bits = prefill_cfg.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, @@ -1020,8 +1044,7 @@ def forward(self, x): class LinearMXInt(_LinearBase): """MXInt linear with prefill/decode phase-aware execution. - Behavior mirrors `LinearMXFP`: prefill may be quantized, decode is forced to - FP using cached original weights/bias. + Behavior mirrors `LinearMXFP` with the same low-memory decode-bank policy. """ # NOTE: backward is not supported — inference only (PTQ) @@ -1038,10 +1061,10 @@ def __init__( assert config is not None, "config is None!" self.phase_config = normalize_phase_q_config(config) self.decode_policy = self.phase_config["decode_policy"] - if self.decode_policy != "fp_only": + if self.decode_policy not in ("fp_only", "quantized"): raise ValueError( - "Step-1 integration only supports decode_policy='fp_only' " - f"for {self.__class__.__name__}, got {self.decode_policy!r}." + "Unsupported decode_policy " + f"{self.decode_policy!r} for {self.__class__.__name__}." ) prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") @@ -1052,88 +1075,72 @@ def __init__( 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_snapshot(self, state_dict) -> None: - """Save FP weights for decode before prefill quantization mutates data.""" + 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() - @classmethod - def from_linear(cls, linear: torch.nn.Linear, config: dict) -> "LinearMXInt": - """Create a LinearMXInt that REUSES the original Linear's Parameters. + def _build_decode_quantized_bank(self) -> None: + """Build decode quantized bank once from FP source snapshot.""" - 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. + 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 - 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, - ) - ) + 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 - 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, - ) - ) + 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) - return new + 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_snapshot(state_dict) + 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 @@ -1167,17 +1174,13 @@ def load_state_dict(self, state_dict, strict=True, assign=False): ) ) + self._build_decode_quantized_bank() return result - def _is_decode_fp_mode(self) -> bool: - """Return True when current call must use decode full-precision path.""" - - phase = get_active_phase() - return phase == "decode" and self.decode_policy == "fp_only" - @torch.no_grad() def forward(self, x): - if self._is_decode_fp_mode(): + 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 ) @@ -1188,139 +1191,42 @@ def forward(self, x): ) return F.linear(x, decode_weight, decode_bias) - prefill_cfg, _ = get_phase_subconfig(self.phase_config, "prefill") - if prefill_cfg.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 = prefill_cfg.get("data_in_block_size") - x_element_bits = prefill_cfg.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, - element_bits=x_element_bits, - block_dim=-1, - ) - - return F.linear(x, self.weight, self.bias) - - -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``. - - 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_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), + 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) - return F.linear(x, self.weight, self.bias) - - -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: + 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) - x_block_size = self.config.get("data_in_block_size") - x_element_bits = self.config.get("data_in_width") + # Only quantize activations; weights/bias already quantized in load_state_dict + 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: - from chop.nn.quantizers.rotation import mxint_rotate_quantizer - - x = mxint_rotate_quantizer( + x = mxint_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), - ) - - 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. - """ - - @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) diff --git a/src/chop/nn/quantized/modules/llama/mlp.py b/src/chop/nn/quantized/modules/llama/mlp.py index 5b67837fc..b34b55105 100644 --- a/src/chop/nn/quantized/modules/llama/mlp.py +++ b/src/chop/nn/quantized/modules/llama/mlp.py @@ -1,11 +1,7 @@ """Llama MLP quantization with phase-aware dispatch. -Step-1 policy is explicit: -- prefill: quantization path can run -- decode: full precision only - -Phase source: -- Runtime phase is written by decoder-layer pre-hooks before layer execution. +Default policy keeps decode full precision (`fp_only`) for compatibility. +`quantized` decode is opt-in and consumes decode-phase sub-config directly. """ import torch @@ -13,7 +9,7 @@ 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_active_phase +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, @@ -68,22 +64,19 @@ def __init__(self, config, layer_idx=None, q_config: dict = None): self.layer_idx = layer_idx self.phase_q_config = normalize_phase_q_config(q_config) self.decode_policy = self.phase_q_config["decode_policy"] - if self.decode_policy != "fp_only": - raise ValueError( - "Step-1 integration only supports decode_policy='fp_only' " - f"for {self.__class__.__name__}, got {self.decode_policy!r}." - ) def forward(self, x: Tensor) -> Tensor: - phase = get_active_phase() - sub_cfg, decode_policy = get_phase_subconfig(self.phase_q_config, phase) - bypass = sub_cfg.get("bypass", False) - if phase == "decode" and decode_policy == "fp_only": + 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), sub_cfg) * self.up_proj(x) + x = silu_minifloat(self.gate_proj(x), phase_subconfig) * self.up_proj(x) return self.down_proj(x) @@ -95,20 +88,17 @@ def __init__(self, config, layer_idx=None, q_config: dict = None): self.layer_idx = layer_idx self.phase_q_config = normalize_phase_q_config(q_config) self.decode_policy = self.phase_q_config["decode_policy"] - if self.decode_policy != "fp_only": - raise ValueError( - "Step-1 integration only supports decode_policy='fp_only' " - f"for {self.__class__.__name__}, got {self.decode_policy!r}." - ) def forward(self, x: Tensor) -> Tensor: - phase = get_active_phase() - sub_cfg, decode_policy = get_phase_subconfig(self.phase_q_config, phase) - bypass = sub_cfg.get("bypass", False) - if phase == "decode" and decode_policy == "fp_only": + 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), sub_cfg) * 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 633f9422d..cca762b8a 100644 --- a/src/chop/nn/quantized/modules/llama/rms_norm.py +++ b/src/chop/nn/quantized/modules/llama/rms_norm.py @@ -1,7 +1,7 @@ """Llama RMSNorm quantization with phase-aware dispatch. -Decode is intentionally forced to FP in step-1 for deterministic integration. -Runtime phase is supplied by decoder-layer pre-hooks via phase context. +Default decode behavior remains FP-only for backward compatibility. +`quantized` decode is opt-in via phase config and policy. """ from functools import partial @@ -11,7 +11,7 @@ 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_active_phase +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, @@ -51,11 +51,6 @@ def __init__(self, config=None, layer_idx=None, q_config: dict = None): self.variance_epsilon = config.rms_norm_eps self.phase_q_config = normalize_phase_q_config(q_config) self.decode_policy = self.phase_q_config["decode_policy"] - if self.decode_policy != "fp_only": - raise ValueError( - "Step-1 integration only supports decode_policy='fp_only' " - f"for {self.__class__.__name__}, got {self.decode_policy!r}." - ) @staticmethod def _build_weight_quantizer(sub_cfg: dict, bypass: bool): @@ -90,17 +85,16 @@ def _build_input_quantizer(sub_cfg: dict, bypass: bool): ) def forward(self, hidden_states): - phase = get_active_phase() - sub_cfg, decode_policy = get_phase_subconfig(self.phase_q_config, phase) - bypass = sub_cfg.get("bypass", False) - if phase == "decode" and decode_policy == "fp_only": - # Why force bypass here: - # step-1 intentionally keeps decode fully FP for stability and - # backward compatibility while still accepting phase-shaped configs. + 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(sub_cfg, bypass) - x_quantizer = self._build_input_quantizer(sub_cfg, bypass) + w_quantizer = self._build_weight_quantizer(phase_subconfig, bypass) + x_quantizer = self._build_input_quantizer(phase_subconfig, bypass) input_dtype = hidden_states.dtype if x_quantizer is not None: diff --git a/src/chop/nn/quantized/modules/phase_context.py b/src/chop/nn/quantized/modules/phase_context.py index 69a5dd74d..884d75eb8 100644 --- a/src/chop/nn/quantized/modules/phase_context.py +++ b/src/chop/nn/quantized/modules/phase_context.py @@ -19,7 +19,7 @@ from typing import Any, Literal Phase = Literal["prefill", "decode"] -DecodePolicy = Literal["fp_only"] +DecodePolicy = Literal["fp_only", "quantized"] _ACTIVE_PHASE: ContextVar[Phase] = ContextVar("active_quant_phase", default="prefill") @@ -28,16 +28,16 @@ ) -def set_active_phase(phase: Phase) -> None: - """Set the current runtime phase for quantized module dispatch.""" +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_active_phase() -> Phase: - """Return the current runtime phase. +def get_runtime_phase() -> Phase: + """Return current runtime phase. Defaults to `prefill` when no explicit phase has been written yet. """ @@ -45,28 +45,35 @@ def get_active_phase() -> Phase: return _ACTIVE_PHASE.get() -def set_decode_policy(policy: DecodePolicy) -> None: - """Set decode policy for current runtime context. +def set_runtime_decode_policy(policy: DecodePolicy) -> None: + """Set decode policy in runtime phase context. - Current step intentionally supports only `fp_only` to keep the first - migration minimal and deterministic. + 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 != "fp_only": + if policy not in ("fp_only", "quantized"): raise ValueError(f"Unsupported decode policy: {policy}") _DECODE_POLICY.set(policy) -def get_decode_policy() -> DecodePolicy: - """Return decode policy for current runtime context. +def get_runtime_decode_policy() -> DecodePolicy: + """Return decode policy from runtime phase context. - Defaults to `fp_only`, matching step-1 product decision. + Defaults to `fp_only` to preserve backward compatibility. """ return _DECODE_POLICY.get() -def infer_phase_from_hidden_and_cache(hidden_states: Any, past_cache: Any) -> Phase: +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: @@ -102,7 +109,7 @@ def infer_phase_from_hidden_and_cache(hidden_states: Any, past_cache: Any) -> Ph return "prefill" -def extract_past_cache_from_decoder_layer_inputs( +def extract_decoder_layer_past_cache( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: @@ -123,12 +130,12 @@ def extract_past_cache_from_decoder_layer_inputs( return None -def infer_phase_from_decoder_layer_inputs( +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_past_cache_from_decoder_layer_inputs(args, kwargs) - return infer_phase_from_hidden_and_cache(hidden_states, past_cache) + 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 8bf751f43..b5c337210 100644 --- a/src/chop/passes/module/module_modify_helper.py +++ b/src/chop/passes/module/module_modify_helper.py @@ -120,9 +120,8 @@ def _restore_decode_fp_snapshot_if_available(source_module, target_module): Why this hook exists: - GPTQ pre-pass rewrites `nn.Linear.weight` in-place before replacement. - - Step-1 requires decode to stay FP (including weight path). - - We therefore preserve original FP snapshots on source modules and - transfer them here when the target quantized module exposes decode buffers. + - 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) @@ -133,20 +132,38 @@ def _restore_decode_fp_snapshot_if_available(source_module, target_module): if not target_has_decode_weight: return - 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, + 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): diff --git a/test/nn/quantized/modules/test_llama_phase_split_step1.py b/test/nn/quantized/modules/test_llama_phase_split_step1.py index b33d384b4..c8fea0fa8 100644 --- a/test/nn/quantized/modules/test_llama_phase_split_step1.py +++ b/test/nn/quantized/modules/test_llama_phase_split_step1.py @@ -13,8 +13,10 @@ 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 @@ -43,16 +45,18 @@ 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_active_phase, - infer_phase_from_hidden_and_cache, - get_active_phase, - infer_phase_from_decoder_layer_inputs, + 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_decoder_layer_phase_pre_hook, + _llama_phase_context_pre_hook, ) @@ -82,17 +86,23 @@ def test_infer_runtime_phase_from_cache_state(): hidden_prefill = torch.randn(1, 8, 16) hidden_decode = torch.randn(1, 1, 16) - assert infer_phase_from_hidden_and_cache(hidden_prefill, None) == "prefill" + assert infer_runtime_phase_from_hidden_and_cache(hidden_prefill, None) == "prefill" assert ( - infer_phase_from_hidden_and_cache(hidden_prefill, _DummyCache(seq_len=0)) + infer_runtime_phase_from_hidden_and_cache( + hidden_prefill, _DummyCache(seq_len=0) + ) == "prefill" ) assert ( - infer_phase_from_hidden_and_cache(hidden_decode, _DummyCache(seq_len=0)) + infer_runtime_phase_from_hidden_and_cache( + hidden_decode, _DummyCache(seq_len=0) + ) == "decode" ) assert ( - infer_phase_from_hidden_and_cache(hidden_prefill, _DummyCache(seq_len=32)) + infer_runtime_phase_from_hidden_and_cache( + hidden_prefill, _DummyCache(seq_len=32) + ) == "decode" ) @@ -113,7 +123,7 @@ def test_decoder_layer_pre_hook_sets_phase_before_input_layernorm(): orig_ln_forward = layer.input_layernorm.forward def _spy_input_ln(hidden_states): - observed["phase_at_input_ln"] = get_active_phase() + observed["phase_at_input_ln"] = get_runtime_phase() return orig_ln_forward(hidden_states) layer.input_layernorm.forward = _spy_input_ln @@ -129,16 +139,20 @@ def forward(self, hidden_states): layer.self_attn = _DummySelfAttn() layer.mlp = _DummyMLP() - layer.register_forward_pre_hook(_llama_decoder_layer_phase_pre_hook, with_kwargs=True) + 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_active_phase("prefill") + 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(): @@ -147,11 +161,11 @@ def test_decoder_layer_input_extraction_supports_old_and_new_cache_names(): hidden = torch.randn(1, 1, 16) cache = _DummyCache(seq_len=3) - phase_new = infer_phase_from_decoder_layer_inputs( + phase_new = infer_runtime_phase_from_decoder_layer_inputs( args=(), kwargs={"hidden_states": hidden, "past_key_values": cache}, ) - phase_old = infer_phase_from_decoder_layer_inputs( + phase_old = infer_runtime_phase_from_decoder_layer_inputs( args=(), kwargs={"hidden_states": hidden, "past_key_value": cache}, ) @@ -164,7 +178,40 @@ def test_attention_no_longer_writes_runtime_phase(): """Guardrail: attention should consume phase context, not mutate it.""" src = inspect.getsource(LlamaAttentionMXFP.forward) - assert "set_active_phase(" not in src + 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(): @@ -197,17 +244,133 @@ def test_linear_mxfp_decode_uses_fp_snapshot_after_weight_replacement(): x = torch.randn(2, 4) - set_active_phase("prefill") + 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_active_phase("decode") + 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_active_phase("prefill") + 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 index 962ba3a65..9ec378da9 100644 --- a/test/passes/module/transforms/quantize/test_quantize_phase_config.py +++ b/test/passes/module/transforms/quantize/test_quantize_phase_config.py @@ -8,6 +8,7 @@ from pathlib import Path import torch +import pytest from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaDecoderLayer @@ -29,7 +30,8 @@ from chop.nn.quantized.modules.linear import LinearMXFP from chop.passes.module.transforms.quantize.quantize import ( quantize_module_transform_pass, - _install_llama_phase_pre_hooks, + _install_llama_phase_context_pre_hooks, + _infer_llama_decode_policy_from_quantized_modules, ) @@ -70,15 +72,47 @@ def test_quantize_pass_accepts_phase_config_for_mxfp_linear(): 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 TinyNetwork(torch.nn.Module): - def __init__(self, with_quantized_marker: bool): + def __init__(self, with_quantized_marker: bool, decode_policy: str = "fp_only"): super().__init__() cfg = LlamaConfig( hidden_size=16, @@ -89,19 +123,37 @@ def __init__(self, with_quantized_marker: bool): ) self.layer = LlamaDecoderLayer(cfg, layer_idx=0) if with_quantized_marker: - self.quant_marker = LlamaAttentionMXFP() + self.quant_marker = LlamaAttentionMXFP(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_pre_hooks(net_without_marker) + _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_pre_hooks(net_with_marker) + _install_llama_phase_context_pre_hooks(net_with_marker) after_first = len(net_with_marker.layer._forward_pre_hooks) - _install_llama_phase_pre_hooks(net_with_marker) + _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 + + +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()) diff --git a/third_party/gorilla/gorilla b/third_party/gorilla/gorilla new file mode 160000 index 000000000..6ea57973c --- /dev/null +++ b/third_party/gorilla/gorilla @@ -0,0 +1 @@ +Subproject commit 6ea57973c7a6097fd7c5915698c54c17c5b1b6c8 From e07c287a1c5442d732b8f557838964d66fcf734a Mon Sep 17 00:00:00 2001 From: Yuxuan Han Date: Thu, 21 May 2026 20:46:34 +0100 Subject: [PATCH 3/6] bug fix --- .../quantize/test_quantize_phase_config.py | 28 +++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/test/passes/module/transforms/quantize/test_quantize_phase_config.py b/test/passes/module/transforms/quantize/test_quantize_phase_config.py index 9ec378da9..5baafd0cd 100644 --- a/test/passes/module/transforms/quantize/test_quantize_phase_config.py +++ b/test/passes/module/transforms/quantize/test_quantize_phase_config.py @@ -111,8 +111,21 @@ def __init__(self, decode_policy: str = "fp_only"): 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, decode_policy: str = "fp_only"): + 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, @@ -123,7 +136,7 @@ def __init__(self, with_quantized_marker: bool, decode_policy: str = "fp_only"): ) self.layer = LlamaDecoderLayer(cfg, layer_idx=0) if with_quantized_marker: - self.quant_marker = LlamaAttentionMXFP(decode_policy=decode_policy) + 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) @@ -140,6 +153,17 @@ def __init__(self, with_quantized_marker: bool, decode_policy: str = "fp_only"): 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.""" From 71ed8ebe28bc716a9f612e7afce497ccd988ce9c Mon Sep 17 00:00:00 2001 From: Yuxuan Han Date: Thu, 21 May 2026 21:01:35 +0100 Subject: [PATCH 4/6] Drop local gorilla gitlink from PR payload --- third_party/gorilla/gorilla | 1 - 1 file changed, 1 deletion(-) delete mode 160000 third_party/gorilla/gorilla diff --git a/third_party/gorilla/gorilla b/third_party/gorilla/gorilla deleted file mode 160000 index 6ea57973c..000000000 --- a/third_party/gorilla/gorilla +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 6ea57973c7a6097fd7c5915698c54c17c5b1b6c8 From ae6ec775c2b7538131585186a52a8ad31f8e4d44 Mon Sep 17 00:00:00 2001 From: Yuxuan Han Date: Thu, 21 May 2026 21:05:12 +0100 Subject: [PATCH 5/6] Restore rotate symbol compatibility after mainline migration --- src/chop/nn/quantized/modules/linear.py | 21 +++++++++++++++++++ .../nn/quantized/modules/llama/attention.py | 21 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/src/chop/nn/quantized/modules/linear.py b/src/chop/nn/quantized/modules/linear.py index 8505a0b84..935b26143 100644 --- a/src/chop/nn/quantized/modules/linear.py +++ b/src/chop/nn/quantized/modules/linear.py @@ -1230,3 +1230,24 @@ def forward(self, x): ) return F.linear(x, self.weight, self.bias) + + +class RotateMXFPLinear(LinearMXFP): + """Compatibility alias for rotate-aware MXFP linear. + + 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. + """ + + pass + + +class RotateMXIntLinear(LinearMXInt): + """Compatibility alias for rotate-aware MXInt linear. + + This keeps the public import contract intact while split-phase work is + focused on decode-policy wiring and bank selection. + """ + + pass diff --git a/src/chop/nn/quantized/modules/llama/attention.py b/src/chop/nn/quantized/modules/llama/attention.py index f7e01b0fe..a387bc138 100644 --- a/src/chop/nn/quantized/modules/llama/attention.py +++ b/src/chop/nn/quantized/modules/llama/attention.py @@ -559,3 +559,24 @@ def _eager_attention_forward_mxint( attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights + + +class LlamaAttentionMXFPRotate(LlamaAttentionMXFP): + """Rotation-variant compatibility alias for MXFP attention. + + 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. + """ + + pass + + +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 From b77c9579db3a6ae188671ba7a0ed638e78bec794 Mon Sep 17 00:00:00 2001 From: Yuxuan Han Date: Fri, 22 May 2026 19:38:54 +0100 Subject: [PATCH 6/6] Fix python-format for phase-split quant files --- src/chop/nn/quantized/modules/linear.py | 28 +++++++++++++++---- .../nn/quantized/modules/llama/attention.py | 8 ++++-- .../nn/quantized/modules/llama/rms_norm.py | 6 +--- src/chop/passes/module/transforms/gptq/run.py | 7 +++-- .../module/transforms/quantize/quantize.py | 6 +--- .../modules/test_llama_phase_split_step1.py | 12 ++++---- .../quantize/test_quantize_phase_config.py | 4 ++- 7 files changed, 45 insertions(+), 26 deletions(-) diff --git a/src/chop/nn/quantized/modules/linear.py b/src/chop/nn/quantized/modules/linear.py index 935b26143..12914ca87 100644 --- a/src/chop/nn/quantized/modules/linear.py +++ b/src/chop/nn/quantized/modules/linear.py @@ -876,7 +876,11 @@ def _capture_decode_fp_bank_snapshot(self, state_dict) -> None: 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: + 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: @@ -985,7 +989,9 @@ def forward(self, x): # 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 + self._decode_weight_fp + if self._decode_weight_fp.numel() > 0 + else self.weight ) decode_bias = ( self._decode_bias_fp @@ -1013,7 +1019,9 @@ def forward(self, x): block_dim=-1, ) decode_weight = ( - self._decode_weight_q if self._decode_weight_q.numel() > 0 else self.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 @@ -1088,7 +1096,11 @@ def _capture_decode_fp_bank_snapshot(self, state_dict) -> None: 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: + 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: @@ -1182,7 +1194,9 @@ def forward(self, x): 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 + self._decode_weight_fp + if self._decode_weight_fp.numel() > 0 + else self.weight ) decode_bias = ( self._decode_bias_fp @@ -1206,7 +1220,9 @@ def forward(self, x): block_dim=-1, ) decode_weight = ( - self._decode_weight_q if self._decode_weight_q.numel() > 0 else self.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 diff --git a/src/chop/nn/quantized/modules/llama/attention.py b/src/chop/nn/quantized/modules/llama/attention.py index a387bc138..848209263 100644 --- a/src/chop/nn/quantized/modules/llama/attention.py +++ b/src/chop/nn/quantized/modules/llama/attention.py @@ -221,7 +221,9 @@ def forward( 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) + 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) @@ -355,7 +357,9 @@ def forward( 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) + 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) diff --git a/src/chop/nn/quantized/modules/llama/rms_norm.py b/src/chop/nn/quantized/modules/llama/rms_norm.py index cca762b8a..3cb5424ab 100644 --- a/src/chop/nn/quantized/modules/llama/rms_norm.py +++ b/src/chop/nn/quantized/modules/llama/rms_norm.py @@ -102,9 +102,5 @@ def forward(self, 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 = ( - w_quantizer(self.weight) - if 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/passes/module/transforms/gptq/run.py b/src/chop/passes/module/transforms/gptq/run.py index 5d0ffeda9..eb26f8a2a 100644 --- a/src/chop/passes/module/transforms/gptq/run.py +++ b/src/chop/passes/module/transforms/gptq/run.py @@ -98,7 +98,9 @@ def run_gptq(network, gptq_config): 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"): + 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 @@ -146,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 03aa8a3bf..3a4303a2f 100644 --- a/src/chop/passes/module/transforms/quantize/quantize.py +++ b/src/chop/passes/module/transforms/quantize/quantize.py @@ -23,7 +23,6 @@ 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", @@ -76,10 +75,7 @@ def _has_llama_quantized_runtime_modules(network) -> bool: """ for module in network.modules(): - if ( - module.__class__.__name__ - in _LLAMA_PHASE_CONTEXT_POLICY_MODULE_CLASS_NAMES - ): + if module.__class__.__name__ in _LLAMA_PHASE_CONTEXT_POLICY_MODULE_CLASS_NAMES: return True return False diff --git a/test/nn/quantized/modules/test_llama_phase_split_step1.py b/test/nn/quantized/modules/test_llama_phase_split_step1.py index c8fea0fa8..e8e9cedc1 100644 --- a/test/nn/quantized/modules/test_llama_phase_split_step1.py +++ b/test/nn/quantized/modules/test_llama_phase_split_step1.py @@ -29,7 +29,9 @@ 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()) +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")] @@ -94,9 +96,7 @@ def test_infer_runtime_phase_from_cache_state(): == "prefill" ) assert ( - infer_runtime_phase_from_hidden_and_cache( - hidden_decode, _DummyCache(seq_len=0) - ) + infer_runtime_phase_from_hidden_and_cache(hidden_decode, _DummyCache(seq_len=0)) == "decode" ) assert ( @@ -291,7 +291,9 @@ def test_linear_mxfp_decode_quantized_uses_decode_quant_bank(): 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) + expected_decode = torch.nn.functional.linear( + x, target._decode_weight_q, target.bias + ) assert out_decode.shape == expected_decode.shape set_runtime_phase("prefill") diff --git a/test/passes/module/transforms/quantize/test_quantize_phase_config.py b/test/passes/module/transforms/quantize/test_quantize_phase_config.py index 5baafd0cd..3e005825d 100644 --- a/test/passes/module/transforms/quantize/test_quantize_phase_config.py +++ b/test/passes/module/transforms/quantize/test_quantize_phase_config.py @@ -22,7 +22,9 @@ # 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()) + 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