pare-quant 0.1.0__py3-none-any.whl

pare/__init__.py

@@ -0,0 +1,99 @@
+"""Pare — production-ready quantization for large language and multimodal models."""
+
+from __future__ import annotations
+
+import torch.nn as nn
+
+from pare.config import QuantConfig
+from pare.core.dtype import QuantDtype
+from pare.model.io import load_quantized, save_quantized
+
+__version__ = "0.1.0"
+__all__ = [
+    "QuantConfig", "QuantDtype", "quantize",
+    "save_quantized", "load_quantized",
+    "__version__",
+]
+
+
+def quantize(
+    model: nn.Module,
+    config: QuantConfig,
+    calibration_data: "list | None" = None,
+    device: "str" = "cpu",
+) -> nn.Module:
+    """Quantize a model in-place using the scheme specified in ``config``.
+
+    Args:
+        model:            Any ``nn.Module`` (HuggingFace, custom, etc.).
+        config:           A ``QuantConfig`` describing the target dtype, scheme,
+                          granularity, and scheme-specific hyperparameters.
+        calibration_data: Required for GPTQ, AWQ, and SmoothQuant.
+                          A list of input_ids tensors, each shaped [batch, seq_len].
+                          Ignored for RTN.
+        device:           Device to run calibration forward passes on (GPTQ only).
+
+    Returns:
+        The same model object with ``nn.Linear`` layers replaced by
+        ``QuantizedLinear`` instances.
+
+    Examples::
+
+        # RTN — no calibration needed
+        from pare import quantize, QuantConfig
+        config = QuantConfig(bits=4, scheme="rtn", group_size=128)
+        model = quantize(model, config)
+
+        # GPTQ — calibration data required
+        config = QuantConfig(bits=4, scheme="gptq", group_size=128)
+        model = quantize(model, config, calibration_data=calib_ids, device="cuda")
+    """
+    # Mixed-precision sensitivity: score layers BEFORE quantization so we
+    # still have the original FP16 weights.  RTN is used as a fast proxy.
+    layer_bits_override: dict[str, int] = {}
+    if config.sensitive_bits is not None and calibration_data is not None:
+        from pare.sensitivity import score_layers
+        scores = score_layers(
+            model,
+            calibration_data,
+            bits=config.bits,
+            granularity=config.granularity,
+            group_size=config.group_size,
+            device=device,
+        )
+        layer_bits_override = {
+            name: config.sensitive_bits
+            for name, err in scores.items()
+            if err > config.sensitivity_threshold
+        }
+        n_total = len(scores)
+        n_sensitive = len(layer_bits_override)
+        print(
+            f"[pare] Sensitivity: {n_total} layers scored, "
+            f"{n_sensitive} above {config.sensitivity_threshold:.0%} threshold "
+            f"→ {config.bits}-bit→{config.sensitive_bits}-bit"
+        )
+
+    scheme = config.scheme
+    if scheme == "rtn":
+        from pare.schemes.rtn import RTNQuantizer
+        quantizer = RTNQuantizer(config, layer_bits_override=layer_bits_override)
+        return quantizer.quantize_model(model)
+
+    elif scheme == "gptq":
+        from pare.schemes.gptq import GPTQQuantizer
+        quantizer = GPTQQuantizer(config, layer_bits_override=layer_bits_override)
+        return quantizer.quantize_model(model, calibration_data=calibration_data, device=device)
+
+    elif scheme == "awq":
+        from pare.schemes.awq import AWQQuantizer
+        quantizer = AWQQuantizer(config, layer_bits_override=layer_bits_override)
+        return quantizer.quantize_model(model, calibration_data=calibration_data, device=device)
+
+    elif scheme == "smoothquant":
+        from pare.schemes.smoothquant import SmoothQuantQuantizer
+        quantizer = SmoothQuantQuantizer(config, layer_bits_override=layer_bits_override)
+        return quantizer.quantize_model(model, calibration_data=calibration_data, device=device)
+
+    else:
+        raise ValueError(f"Unknown scheme: {scheme!r}")

pare/__main__.py

@@ -0,0 +1,191 @@
+"""CLI entry point: python -m pare <command> [options]
+
+Commands
+--------
+quantize   Quantize a HuggingFace model and save to disk.
+eval       Evaluate a quantized (or FP16) model's perplexity.
+
+Examples
+--------
+# GPTQ INT4, 128 calibration sequences, save to ./llama2-7b-int4/
+python -m pare quantize \\
+    --model  meta-llama/Llama-2-7b-hf \\
+    --bits   4 \\
+    --scheme gptq \\
+    --output ./llama2-7b-int4
+
+# AWQ INT4 with custom alpha
+python -m pare quantize \\
+    --model  meta-llama/Llama-2-7b-hf \\
+    --bits   4 \\
+    --scheme awq \\
+    --output ./llama2-7b-awq
+
+# SmoothQuant INT8
+python -m pare quantize \\
+    --model  meta-llama/Llama-2-7b-hf \\
+    --bits   8 \\
+    --scheme smoothquant \\
+    --output ./llama2-7b-sq-int8
+
+# Evaluate a saved quantized model
+python -m pare eval \\
+    --model  meta-llama/Llama-2-7b-hf \\
+    --quantized ./llama2-7b-int4 \\
+    --dataset wikitext2
+
+# Evaluate FP16 baseline (no --quantized flag)
+python -m pare eval \\
+    --model  meta-llama/Llama-2-7b-hf \\
+    --dataset wikitext2
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m pare",
+        description="Pare — post-training quantization for LLMs",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    # ── quantize ──────────────────────────────────────────────────────
+    p_q = sub.add_parser("quantize", help="Quantize a model and save to disk")
+    p_q.add_argument("--model",    required=True,  help="HuggingFace model ID or local path")
+    p_q.add_argument("--output",   required=True,  help="Directory to save the quantized model")
+    p_q.add_argument("--bits",     type=int, default=4, choices=[2, 3, 4, 8],
+                     help="Target bit-width (default: 4)")
+    p_q.add_argument("--scheme",   default="awq",
+                     choices=["rtn", "gptq", "awq", "smoothquant"],
+                     help="Quantization scheme (default: awq)")
+    p_q.add_argument("--granularity", default="per_group",
+                     choices=["per_tensor", "per_channel", "per_group"],
+                     help="Scale granularity (default: per_group)")
+    p_q.add_argument("--group-size", type=int, default=128,
+                     help="Group size for per_group granularity (default: 128)")
+    p_q.add_argument("--sym", action="store_true",
+                     help="Symmetric quantization (no zero-point)")
+    p_q.add_argument("--smooth-alpha", type=float, default=0.5,
+                     help="SmoothQuant migration strength α (default: 0.5)")
+    p_q.add_argument("--n-calib", type=int, default=128,
+                     help="Number of calibration sequences (default: 128)")
+    p_q.add_argument("--calib-seqlen", type=int, default=2048,
+                     help="Length of each calibration sequence (default: 2048)")
+    p_q.add_argument("--device", default="cuda",
+                     help="Device for calibration forward passes (default: cuda)")
+    p_q.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"],
+                     help="Model load dtype (default: float16)")
+    p_q.add_argument("--token", default=None,
+                     help="HuggingFace access token for gated models")
+
+    # ── eval ──────────────────────────────────────────────────────────
+    p_e = sub.add_parser("eval", help="Evaluate perplexity of a model")
+    p_e.add_argument("--model",    required=True,  help="HuggingFace model ID or local path")
+    p_e.add_argument("--quantized",   default=None,
+                     help="Path to a Pare-saved quantized model (omit for FP16 baseline)")
+    p_e.add_argument("--dataset",  default="wikitext2",
+                     choices=["wikitext2", "c4"],
+                     help="Evaluation dataset (default: wikitext2)")
+    p_e.add_argument("--seq-len",  type=int, default=2048,
+                     help="Sequence length (default: 2048)")
+    p_e.add_argument("--n-samples", type=int, default=None,
+                     help="Number of sequences to evaluate (default: full dataset)")
+    p_e.add_argument("--device",   default="cuda")
+    p_e.add_argument("--dtype",    default="float16", choices=["float16", "bfloat16", "float32"])
+    p_e.add_argument("--token",    default=None)
+
+    return parser
+
+
+def cmd_quantize(args: argparse.Namespace) -> None:
+    import torch
+    from datasets import load_dataset
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    from pare import QuantConfig, quantize, save_quantized
+
+    dtype_map = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
+    torch_dtype = dtype_map[args.dtype]
+
+    print(f"[pare] Loading {args.model} ...", flush=True)
+    tokenizer = AutoTokenizer.from_pretrained(args.model, token=args.token)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, torch_dtype=torch_dtype, token=args.token, device_map="auto",
+    )
+    model.eval()
+
+    config = QuantConfig(
+        bits=args.bits,
+        scheme=args.scheme,
+        granularity=args.granularity,
+        group_size=args.group_size,
+        sym=args.sym,
+        smooth_alpha=args.smooth_alpha,
+    )
+
+    calib_data = None
+    if args.scheme != "rtn":
+        print(f"[pare] Preparing {args.n_calib} calibration sequences ...", flush=True)
+        ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train")
+        text = "\n\n".join(ds["text"])
+        tokens = tokenizer(text, return_tensors="pt").input_ids
+        seq_len = args.calib_seqlen
+        calib_data = [tokens[:, i * seq_len : (i + 1) * seq_len] for i in range(args.n_calib)]
+
+    print(f"[pare] Quantizing with {args.scheme} INT{args.bits} ...", flush=True)
+    quantize(model, config, calibration_data=calib_data, device=args.device)
+
+    save_quantized(model, args.output)
+    print(f"[pare] Done. Model saved to {args.output}", flush=True)
+
+
+def cmd_eval(args: argparse.Namespace) -> None:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    from pare import load_quantized
+    from pare.eval.perplexity import evaluate_perplexity
+
+    dtype_map = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
+    torch_dtype = dtype_map[args.dtype]
+
+    print(f"[pare] Loading {args.model} ...", flush=True)
+    tokenizer = AutoTokenizer.from_pretrained(args.model, token=args.token)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, torch_dtype=torch_dtype, token=args.token, device_map="auto",
+    )
+    model.eval()
+
+    if args.quantized:
+        print(f"[pare] Loading quantized weights from {args.quantized} ...", flush=True)
+        load_quantized(model, args.quantized)
+
+    print(f"[pare] Evaluating PPL on {args.dataset} ...", flush=True)
+    ppl = evaluate_perplexity(
+        model, tokenizer, dataset=args.dataset,
+        seq_len=args.seq_len, n_samples=args.n_samples, device=args.device,
+    )
+
+    tag = f"{args.quantized or args.dtype}"
+    print(f"\n  {args.model}  [{tag}]  {args.dataset} PPL: {ppl:.2f}")
+
+
+def main() -> None:
+    parser = _build_parser()
+    args = parser.parse_args()
+
+    if args.command == "quantize":
+        cmd_quantize(args)
+    elif args.command == "eval":
+        cmd_eval(args)
+    else:
+        parser.print_help()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

pare/calibration/hessian.py

@@ -0,0 +1,73 @@
+"""Per-layer Hessian accumulation for GPTQ.
+
+The GPTQ objective for one weight row w is:
+
+    (w - q)^T H (w - q)
+
+where H = 2/n * X X^T, X shaped [in_features, n_tokens].
+
+We accumulate H incrementally over batches so we never store all
+activations in memory — only the running [in_features, in_features] sum.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+
+class HessianAccumulator:
+    """Accumulates the per-layer Hessian H = 2/n * X X^T online.
+
+    Usage::
+
+        acc = HessianAccumulator()
+        for batch_activations in ...:
+            acc.accumulate(batch_activations)
+        H = acc.finalize()   # [in_features, in_features]
+    """
+
+    def __init__(self) -> None:
+        self.H: Tensor | None = None
+        self.n_tokens: int = 0
+
+    def accumulate(self, x: Tensor) -> None:
+        """Add one batch of activations to the running sum.
+
+        Args:
+            x: Input activations, shape [batch, seq_len, in_features]
+               or [batch, in_features].  Will be cast to float32.
+        """
+        x = x.detach().float()
+
+        # Flatten batch and sequence dims → [n_tokens, in_features]
+        if x.dim() == 3:
+            x = x.reshape(-1, x.shape[-1])
+        elif x.dim() != 2:
+            raise ValueError(f"Expected 2-D or 3-D activation, got shape {tuple(x.shape)}")
+
+        n = x.shape[0]  # number of tokens in this batch
+
+        if self.H is None:
+            in_features = x.shape[1]
+            self.H = torch.zeros(
+                in_features, in_features,
+                device=x.device, dtype=torch.float32,
+            )
+
+        # H_unnorm += X_row^T @ X_row  (= X_col @ X_col^T  where X_col = x.T)
+        self.H.addmm_(x.T, x)   # in-place fused multiply-add, no extra alloc
+        self.n_tokens += n
+
+    def finalize(self) -> Tensor:
+        """Return the normalised Hessian H = 2/n * ΣX^TX.
+
+        Raises RuntimeError if no samples have been accumulated.
+        """
+        if self.H is None or self.n_tokens == 0:
+            raise RuntimeError("HessianAccumulator has no samples — did you forget to call accumulate()?")
+        return self.H * (2.0 / self.n_tokens)
+
+    def reset(self) -> None:
+        self.H = None
+        self.n_tokens = 0

pare/calibration/layerwise.py

@@ -0,0 +1,241 @@
+"""Layerwise GPTQ for transformer block models (Llama, Mistral, Qwen, etc.).
+
+Problem: For 7B+ models, accumulating Hessians for all layers simultaneously
+requires ~39 GB on top of the 14 GB model — OOM on 48 GB GPUs.
+
+Solution: process one transformer block at a time:
+  1. Intercept the first block with a Catcher to capture embedding outputs.
+  2. Precompute position_embeddings once from the shared rotary_emb.
+  3. For each block i:
+     a. Load block to GPU.
+     b. Register hooks on its linear sublayers; run all calibration inputs.
+     c. GPTQ-quantize each sublayer immediately (peak VRAM: ~2 GB).
+     d. Collect outputs (= inputs to block i+1).
+     e. Offload block to CPU.
+
+Supported architectures: any HF model with model.model.layers (Llama, Mistral,
+Qwen2, Phi-3, Gemma, ...).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from pare.calibration.hessian import HessianAccumulator
+
+
+def is_supported(model: nn.Module) -> bool:
+    """Return True if this model has the block structure we support."""
+    return (
+        hasattr(model, "model")
+        and hasattr(model.model, "layers")
+        and len(model.model.layers) > 0
+        and isinstance(model.model.layers[0], nn.Module)
+    )
+
+
+class LayerwiseGPTQ:
+    """Block-by-block GPTQ. Keeps peak GPU usage to O(one_block + calib_acts)."""
+
+    def run(
+        self,
+        model: nn.Module,
+        calibration_data: list[Tensor],
+        quantizer,
+        device: str | torch.device,
+    ) -> nn.Module:
+        """Quantize model in-place using layerwise Hessian collection.
+
+        Args:
+            model:            Model to quantize (Llama/Mistral-family HF model).
+            calibration_data: List of input_ids tensors [1, seq_len].
+            quantizer:        GPTQQuantizer instance; quantize_layer is called
+                              for each sublayer.
+            device:           GPU device string, e.g. "cuda".
+
+        Returns:
+            The same model, quantized in-place.
+        """
+        device = torch.device(device)
+        layers = model.model.layers
+        seq_len = calibration_data[0].shape[1]
+
+        # ── 1. Capture embedding outputs (inputs to layer 0) ───────────────
+        inps = _capture_embeddings(model, layers, calibration_data, device)
+
+        # ── 2. Precompute shared position_embeddings ───────────────────────
+        pe = _precompute_pe(model, layers, inps, seq_len, device)
+        model.cpu()
+        torch.cuda.empty_cache()
+
+        # ── 3. Resolve the dotted path to model.model.layers ──────────────
+        layers_path = _find_layers_path(model)
+
+        # ── 4. Process one block at a time ─────────────────────────────────
+        n_layers = len(layers)
+        for li, layer in enumerate(layers):
+            layer_prefix = f"{layers_path}.{li}"
+            layer.to(device)
+
+            # Register hooks on every quantizable linear sublayer in this block.
+            subs: dict[str, nn.Linear] = {}
+            accs: dict[str, HessianAccumulator] = {}
+            hooks = []
+
+            for rel_name, mod in layer.named_modules():
+                if not isinstance(mod, nn.Linear):
+                    continue
+                full_name = f"{layer_prefix}.{rel_name}"
+                if not quantizer._should_quantize(full_name, mod):
+                    continue
+                acc = HessianAccumulator()
+                subs[full_name] = mod
+                accs[full_name] = acc
+                hooks.append(mod.register_forward_hook(_make_hook(acc)))
+
+            with torch.no_grad():
+                for x in inps:
+                    _call_layer(layer, x.to(device), pe, device)
+
+            for h in hooks:
+                h.remove()
+
+            # Quantize each sublayer immediately — no simultaneous Hessian storage.
+            for full_name, linear in subs.items():
+                H = accs[full_name].finalize()
+                # Expose the Hessian so quantize_layer can find it.
+                quantizer._hessians[full_name] = H.cpu()
+                q_layer = quantizer.quantize_layer(linear, full_name)
+                _set_submodule(layer, full_name[len(layer_prefix) + 1:], q_layer)
+                # Clean up immediately to keep memory flat.
+                del quantizer._hessians[full_name], H
+                torch.cuda.empty_cache()
+
+            # Collect this block's outputs → inputs for the next block.
+            outs: list[Tensor] = []
+            with torch.no_grad():
+                for x in inps:
+                    out = _call_layer(layer, x.to(device), pe, device)
+                    outs.append(out.cpu())
+            inps = outs
+
+            layer.cpu()
+            torch.cuda.empty_cache()
+            print(f"[pare] layer {li + 1}/{n_layers} quantized", flush=True)
+
+        return model
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+class _Catcher(nn.Module):
+    """Replaces layer 0 temporarily to intercept embedding outputs."""
+    def __init__(self, layer: nn.Module, store: list[Tensor]) -> None:
+        super().__init__()
+        self._layer = layer
+        self._store = store
+
+    def forward(self, x: Tensor, **_kw) -> Tensor:
+        self._store.append(x.cpu())
+        raise StopIteration
+
+
+def _capture_embeddings(
+    model: nn.Module,
+    layers: nn.ModuleList,
+    calibration_data: list[Tensor],
+    device: torch.device,
+) -> list[Tensor]:
+    """Run the embedding + pre-block path, capture inputs to layer 0."""
+    inps: list[Tensor] = []
+    original = layers[0]
+    layers[0] = _Catcher(original, inps)
+    model.to(device)
+    try:
+        with torch.no_grad():
+            for input_ids in calibration_data:
+                try:
+                    model(input_ids.to(device))
+                except StopIteration:
+                    pass
+    finally:
+        layers[0] = original
+    return inps
+
+
+def _precompute_pe(
+    model: nn.Module,
+    layers: nn.ModuleList,
+    inps: list[Tensor],
+    seq_len: int,
+    device: torch.device,
+) -> tuple[Tensor, Tensor] | None:
+    """Precompute (cos, sin) position embeddings for seq_len tokens.
+
+    In transformers >= 4.43, rotary_emb is a shared module at model.model
+    level. Older versions attach it per attention layer. Returns None if
+    the model doesn't use RoPE.
+    """
+    rotary = getattr(model.model, "rotary_emb", None)
+    if rotary is None:
+        # Try old-style per-layer rotary (transformers < 4.43).
+        rotary = getattr(getattr(layers[0], "self_attn", None), "rotary_emb", None)
+    if rotary is None:
+        return None
+
+    with torch.no_grad():
+        position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
+        cos, sin = rotary(inps[0].to(device), position_ids)
+    return cos.cpu(), sin.cpu()
+
+
+def _call_layer(
+    layer: nn.Module,
+    x: Tensor,
+    pe: tuple[Tensor, Tensor] | None,
+    device: torch.device,
+) -> Tensor:
+    """Call a decoder layer, handling both new and old transformers APIs.
+
+    - New (>= 4.46): requires position_embeddings=(cos, sin) kwarg; returns Tensor.
+    - Old (< 4.46): uses position_ids internally; returns (hidden_states, ...) tuple.
+    """
+    kwargs: dict = {"use_cache": False}
+    if pe is not None:
+        kwargs["position_embeddings"] = (pe[0].to(device), pe[1].to(device))
+
+    out = layer(x, **kwargs)
+
+    # Normalise output: newer transformers returns a plain Tensor,
+    # older versions return a tuple (hidden_states, ...).
+    if isinstance(out, tuple):
+        out = out[0]
+    return out
+
+
+def _find_layers_path(model: nn.Module) -> str:
+    """Return the dotted name of model.model.layers in named_modules()."""
+    target = model.model.layers
+    for name, mod in model.named_modules():
+        if mod is target:
+            return name
+    raise RuntimeError("Could not find model.model.layers in named_modules()")
+
+
+def _set_submodule(parent: nn.Module, rel_path: str, new_mod: nn.Module) -> None:
+    """Set parent.a.b.c = new_mod given rel_path='a.b.c'."""
+    parts = rel_path.split(".")
+    m = parent
+    for part in parts[:-1]:
+        m = getattr(m, part)
+    setattr(m, parts[-1], new_mod)
+
+
+def _make_hook(acc: HessianAccumulator):
+    def hook(module: nn.Module, inputs: tuple, output: Tensor) -> None:
+        acc.accumulate(inputs[0].detach())
+    return hook

pare/calibration/observer.py

@@ -0,0 +1,72 @@
+"""Per-channel activation magnitude observer for AWQ.
+
+AWQ needs to know which input channels carry large activations so it can
+protect those weight columns during quantization.  This module accumulates
+the per-channel mean |x| online — just like HessianAccumulator but O(in)
+instead of O(in²).
+
+Usage::
+
+    obs = ActivationObserver()
+    hook = layer.register_forward_hook(
+        lambda m, inp, out: obs.accumulate(inp[0])
+    )
+    # ... run calibration forward passes ...
+    hook.remove()
+    x_max = obs.finalize()   # [in_features]  mean |x| per input channel
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+
+class ActivationObserver:
+    """Online accumulator for per-channel activation statistics.
+
+    Tracks both mean |x| (used by AWQ) and max |x| (used by SmoothQuant)
+    in a single pass over calibration data.
+
+    Accumulates over any number of batches of shape
+    ``[batch, seq_len, in_features]`` or ``[batch, in_features]``.
+    """
+
+    def __init__(self) -> None:
+        self._sum: Tensor | None = None
+        self._max: Tensor | None = None
+        self._n_tokens: int = 0
+
+    def accumulate(self, x: Tensor) -> None:
+        x = x.detach().float()
+        if x.dim() == 3:
+            x = x.reshape(-1, x.shape[-1])   # [n_tokens, in_features]
+        elif x.dim() != 2:
+            raise ValueError(f"Expected 2-D or 3-D activation, got {tuple(x.shape)}")
+
+        x_abs = x.abs()
+
+        if self._sum is None:
+            self._sum = torch.zeros(x.shape[1], device=x.device, dtype=torch.float32)
+            self._max = torch.zeros(x.shape[1], device=x.device, dtype=torch.float32)
+
+        self._sum.add_(x_abs.sum(dim=0))
+        torch.maximum(self._max, x_abs.amax(dim=0), out=self._max)
+        self._n_tokens += x.shape[0]
+
+    def finalize(self) -> Tensor:
+        """Return mean |x| per input channel, shape [in_features]. Used by AWQ."""
+        if self._sum is None or self._n_tokens == 0:
+            raise RuntimeError("ActivationObserver has no samples")
+        return self._sum / self._n_tokens
+
+    def max_abs(self) -> Tensor:
+        """Return max |x| per input channel, shape [in_features]. Used by SmoothQuant."""
+        if self._max is None or self._n_tokens == 0:
+            raise RuntimeError("ActivationObserver has no samples")
+        return self._max.clone()
+
+    def reset(self) -> None:
+        self._sum = None
+        self._max = None
+        self._n_tokens = 0