llm.quantization — Model Quantization

Post-training quantization (PTQ), GPTQ, and AWQ support for reducing model size and inference latency. The CLI entry point is llm-quantize.

Overview

Method	Description
GPTQ	Greedy row-wise pruning with second-order information
AWQ	Activation-aware per-channel scales with grid search
SmoothQuant	INT8 weights + activations with activation smoothing
Mixed-Precision	Per-layer quantization dispatch via LayerQuantPolicy

SmoothQuant

smooth

SmoothQuant (Xiao et al., ICML 2023) post-training quantization.

SmoothQuant makes INT8 weight+activation quantization tractable for LLMs by migrating the quantization difficulty from activations to weights: per-input- channel smoothing factors s_j = act_max[j]**alpha / w_max[j]**(1-alpha) are folded into the weights (W·s), and the input is divided by s before its per-tensor INT8 fake quantization.

This module mirrors the GPTQ / AWQ paths: a frozen config, a stateful per-layer quantizer, and model-level entry points that capture per-layer calibration activations through forward hooks.

SmoothQuantConfig dataclass

Configuration for SmoothQuant quantization.

属性：

名称	类型	描述
alpha	float	Smoothing strength in [0, 1]. alpha=0 pushes all quantization difficulty onto activations (weights normalized), alpha=1 pushes it onto weights (activations normalized); the paper's default is 0.5 (balanced).
search_alpha	bool	If True, grid-search alpha per layer over {0.25, 0.5, 0.75, 1.0} using the calibration activations, picking the value with the lowest output reconstruction error. Requires retaining the calibration batches (more memory).
bits	int	Weight bit width. SmoothQuant is an INT8 method in v1 — only 8 is accepted; sub-8-bit weight variants are a follow-up.
group_size	int	Always -1 (per-channel) in v1 — SmoothQuant weights are quantized per output row by design. Present only so LayerQuantPolicy overrides can be validated uniformly.
sym	bool	If True, symmetric quantization (no zero-point). Asymmetric SmoothQuant is not yet implemented.
act_order	bool	Accepted for LayerQuantPolicy compatibility but ignored — SmoothQuant has no column-reordering step.
layer_policies	tuple[LayerQuantPolicy, ...]	Atomic per-layer override policies. For v1 the only meaningful overrides are bits=8 and group_size=-1 (SmoothQuant weights are per-channel by design); anything else fails loudly at quantize time.

源代码位于： src/llm/quantization/smooth.py

@dataclass(frozen=True)
class SmoothQuantConfig:
    """Configuration for SmoothQuant quantization.

    Attributes:
        alpha: Smoothing strength in [0, 1]. alpha=0 pushes all quantization
            difficulty onto activations (weights normalized), alpha=1 pushes
            it onto weights (activations normalized); the paper's default is
            0.5 (balanced).
        search_alpha: If True, grid-search alpha per layer over
            {0.25, 0.5, 0.75, 1.0} using the calibration activations,
            picking the value with the lowest output reconstruction error.
            Requires retaining the calibration batches (more memory).
        bits: Weight bit width. SmoothQuant is an INT8 method in v1 —
            only 8 is accepted; sub-8-bit weight variants are a follow-up.
        group_size: Always -1 (per-channel) in v1 — SmoothQuant weights are
            quantized per output row by design. Present only so
            ``LayerQuantPolicy`` overrides can be validated uniformly.
        sym: If True, symmetric quantization (no zero-point). Asymmetric
            SmoothQuant is not yet implemented.
        act_order: Accepted for ``LayerQuantPolicy`` compatibility but
            ignored — SmoothQuant has no column-reordering step.
        layer_policies: Atomic per-layer override policies. For v1 the only
            meaningful overrides are ``bits=8`` and ``group_size=-1``
            (SmoothQuant weights are per-channel by design); anything else
            fails loudly at quantize time.
    """

    alpha: float = 0.5
    search_alpha: bool = False
    bits: int = 8
    group_size: int = -1
    sym: bool = True
    act_order: bool = False

    # Per-layer atomic override policies (additive; empty tuple = no override).
    layer_policies: tuple[LayerQuantPolicy, ...] = ()

    def __post_init__(self):
        if not (0.0 <= self.alpha <= 1.0):
            raise ValueError(f"SmoothQuantConfig.alpha must be in [0, 1], got {self.alpha}.")
        if self.bits != 8:
            raise ValueError(
                f"SmoothQuantConfig.bits must be 8 in v1 (SmoothQuant is an INT8 "
                f"weight+activation method), got {self.bits}. Use AWQ or GPTQ for "
                "sub-8-bit weight-only quantization."
            )
        if self.group_size != -1:
            raise ValueError(f"SmoothQuantConfig.group_size must be -1 (per-channel) in v1, got {self.group_size}.")
        for i, p in enumerate(self.layer_policies):
            if not isinstance(p, LayerQuantPolicy):
                raise TypeError(
                    f"SmoothQuantConfig.layer_policies[{i}] must be LayerQuantPolicy; got {type(p).__name__}."
                )

SmoothQuantQuantizer

Stateful per-layer SmoothQuant processor.

Lifecycle

q = SmoothQuantQuantizer(layer, config) for batch in calib_iter_for_this_layer: q.add_batch(batch) components = q.quantize() # uses config.alpha or searches alpha

源代码位于： src/llm/quantization/smooth.py

class SmoothQuantQuantizer:
    """Stateful per-layer SmoothQuant processor.

    Lifecycle:
        q = SmoothQuantQuantizer(layer, config)
        for batch in calib_iter_for_this_layer:
            q.add_batch(batch)
        components = q.quantize()   # uses config.alpha or searches alpha
    """

    def __init__(self, layer: nn.Linear, config: SmoothQuantConfig):
        self.config = config
        self.layer = layer
        self.device = layer.weight.device

        self.in_features = layer.weight.shape[1]
        self.act_max = torch.zeros(self.in_features, dtype=torch.float32, device=self.device)
        self.n_samples = 0
        self._batches: list[torch.Tensor] = []

    def add_batch(self, x: torch.Tensor) -> None:
        """Accumulate per-channel max abs activation (and optionally batches)."""
        with _single_thread_reductions():
            x = x.to(device=self.device, dtype=torch.float32)
            if x.dim() == 1:
                x = x.unsqueeze(0)
            x = x.reshape(-1, x.shape[-1])
            if x.shape[0] == 0:
                return
            self.act_max = torch.maximum(self.act_max, x.abs().max(dim=0)[0])
            self.n_samples += x.shape[0]
            if self.config.search_alpha:
                self._batches.append(x.detach().clone())

    def quantize(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Quantize the layer; returns the packed component tuple."""
        if self.n_samples == 0:
            raise RuntimeError(
                "No calibration data accumulated (n_samples=0). Feed at least one calibration batch via add_batch()."
            )
        if not self.config.sym:
            raise NotImplementedError("Asymmetric SmoothQuant not yet implemented. Use sym=True.")

        with _single_thread_reductions():
            w = self.layer.weight.detach().to(device=self.device, dtype=torch.float32)
            bias = self.layer.bias.detach() if self.layer.bias is not None else None

            if self.config.search_alpha:
                if not self._batches:
                    raise RuntimeError("search_alpha=True requires calibration batches to evaluate alpha candidates.")
                best_alpha = min(
                    ALPHA_SEARCH_GRID,
                    key=lambda a: _eval_layer_error(w, self.act_max, self._batches, a, bias),
                )
                logger.info(f"Layer alpha search picked {best_alpha}")
                alpha = best_alpha
            else:
                alpha = self.config.alpha

            return _quantize_layer_components(w, self.act_max, alpha)

add_batch

add_batch(x)

Accumulate per-channel max abs activation (and optionally batches).

源代码位于： src/llm/quantization/smooth.py

def add_batch(self, x: torch.Tensor) -> None:
    """Accumulate per-channel max abs activation (and optionally batches)."""
    with _single_thread_reductions():
        x = x.to(device=self.device, dtype=torch.float32)
        if x.dim() == 1:
            x = x.unsqueeze(0)
        x = x.reshape(-1, x.shape[-1])
        if x.shape[0] == 0:
            return
        self.act_max = torch.maximum(self.act_max, x.abs().max(dim=0)[0])
        self.n_samples += x.shape[0]
        if self.config.search_alpha:
            self._batches.append(x.detach().clone())

quantize

quantize()

Quantize the layer; returns the packed component tuple.

源代码位于： src/llm/quantization/smooth.py

def quantize(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Quantize the layer; returns the packed component tuple."""
    if self.n_samples == 0:
        raise RuntimeError(
            "No calibration data accumulated (n_samples=0). Feed at least one calibration batch via add_batch()."
        )
    if not self.config.sym:
        raise NotImplementedError("Asymmetric SmoothQuant not yet implemented. Use sym=True.")

    with _single_thread_reductions():
        w = self.layer.weight.detach().to(device=self.device, dtype=torch.float32)
        bias = self.layer.bias.detach() if self.layer.bias is not None else None

        if self.config.search_alpha:
            if not self._batches:
                raise RuntimeError("search_alpha=True requires calibration batches to evaluate alpha candidates.")
            best_alpha = min(
                ALPHA_SEARCH_GRID,
                key=lambda a: _eval_layer_error(w, self.act_max, self._batches, a, bias),
            )
            logger.info(f"Layer alpha search picked {best_alpha}")
            alpha = best_alpha
        else:
            alpha = self.config.alpha

        return _quantize_layer_components(w, self.act_max, alpha)

quantize_model_smoothquant

quantize_model_smoothquant(model, calib_iter, config=None, target_modules=None, device=None)

Quantize a model with SmoothQuant (INT8 weights + activations).

参数：

名称	类型	描述	默认
model	Module	nn.Module containing nn.Linear layers to quantize.	必需
calib_iter	Iterator[Tensor]	Iterator yielding input tensors for the model forward pass.	必需
config	SmoothQuantConfig | None	SmoothQuantConfig (default: alpha=0.5, INT8 symmetric).	None
target_modules	Iterable[str] | None	Iterable of fully-qualified layer names to quantize. If None, all nn.Linear layers are quantized.	None
device	device | str | None	Device to run calibration on (default: model's device).	None

返回：

类型	描述
Module	The model with nn.Linear layers replaced by SmoothQuantLinear.

引发：

类型	描述
ValueError	If model has no nn.Linear, target_modules unmatched, layer already quantized, calibration empty, or a layer policy targets an unsupported override.

源代码位于： src/llm/quantization/smooth.py

def quantize_model_smoothquant(
    model: nn.Module,
    calib_iter: Iterator[torch.Tensor],
    config: SmoothQuantConfig | None = None,
    target_modules: Iterable[str] | None = None,
    device: torch.device | str | None = None,
) -> nn.Module:
    """Quantize a model with SmoothQuant (INT8 weights + activations).

    Args:
        model: nn.Module containing nn.Linear layers to quantize.
        calib_iter: Iterator yielding input tensors for the model forward pass.
        config: SmoothQuantConfig (default: alpha=0.5, INT8 symmetric).
        target_modules: Iterable of fully-qualified layer names to quantize.
            If None, all nn.Linear layers are quantized.
        device: Device to run calibration on (default: model's device).

    Returns:
        The model with nn.Linear layers replaced by SmoothQuantLinear.

    Raises:
        ValueError: If model has no nn.Linear, target_modules unmatched, layer
            already quantized, calibration empty, or a layer policy targets
            an unsupported override.
    """
    config = config or SmoothQuantConfig()
    if device is not None:
        model = model.to(device)

    for n, m in model.named_modules():
        if isinstance(m, SmoothQuantLinear):
            raise ValueError(f"Layer {n} is already SmoothQuant-quantized. Pass a fresh model or unquantize first.")

    linear_layers = [(n, m) for n, m in model.named_modules() if isinstance(m, nn.Linear)]
    if not linear_layers:
        raise ValueError("model has no nn.Linear modules; nothing to quantize.")

    if target_modules is not None:
        target_set = set(target_modules)
        all_names = {n for n, _ in linear_layers}
        matched = target_set & all_names
        if not matched:
            available = sorted(all_names)[:10]
            raise ValueError(
                f"target_modules {list(target_set)} matched no nn.Linear. "
                f"Available: {available}{'...' if len(all_names) > 10 else ''}"
            )
        targets = [(n, m) for n, m in linear_layers if n in target_set]
    else:
        targets = linear_layers

    calib_batches = list(calib_iter)
    if not calib_batches:
        raise ValueError("calib_iter is empty; need at least 1 batch for activation statistics.")

    captured: dict[str, list[torch.Tensor]] = {n: [] for n, _ in targets}
    hooks = []

    def make_hook(name: str):
        def hook(_module, inputs, _output):
            captured[name].append(inputs[0].detach().clone())

        return hook

    for n, m in targets:
        hooks.append(m.register_forward_hook(make_hook(n)))

    model.eval()
    with torch.no_grad():
        param_device = next(model.parameters()).device
        try:
            # Feed EVERY calibration batch so each layer's activation stats
            # cover the full calibration set.  Previously only
            # calib_batches[0] was forwarded, silently dropping later batches
            # while the fallback path below used all of them.
            for batch in calib_batches:
                _ = model(batch.to(param_device))
        except (RuntimeError, ValueError, TypeError) as e:
            logger.debug(f"Model forward failed during calibration: {e}; falling back to direct layer calls.")

    # A mid-loop forward failure must not leave the per-layer capture counts
    # diverged (a layer that never got its first batch holds ZERO captures and
    # later crashes with "No calibration data accumulated (n_samples=0)" after
    # earlier layers were already replaced). Any inconsistency -> rebuild every
    # target's captures from direct calls over the full set (RIL ISS-136).
    expected_captures = len(calib_batches)
    capt_sizes = {n: len(v) for n, v in captured.items()}
    any_captured = any(s > 0 for s in capt_sizes.values())
    consistent = bool(capt_sizes) and all(s == expected_captures for s in capt_sizes.values())
    if not any_captured or not consistent:
        if any_captured and not consistent:
            logger.warning(
                "Per-layer calibration captures diverged after a partial forward "
                "failure (%s); falling back to direct layer calls for ALL targets "
                "so every layer quantizes over the same calibration set.",
                capt_sizes,
            )
        elif not any_captured:
            logger.warning(
                "Model forward failed on EVERY calibration batch (see the "
                "DEBUG log above for the first error); no per-layer inputs "
                "were captured. Falling back to feeding the raw calibration "
                "batches directly as each target layer's inputs — this is only "
                "valid when those tensors ARE the layers' activations (e.g. a "
                "bare sequence of target layers). If the model embeds or "
                "reshapes inputs first (a real decoder), the quantized weights "
                "will be garbage; fix the model forward signature or use "
                "quantize_model_with_collector."
            )
        for h in hooks:
            h.remove()
        for n, _m in targets:
            captured[n] = [batch.detach().clone() for batch in calib_batches]

    for h in hooks:
        h.remove()

    available_layer_names = {n for n, _ in targets}
    effective_configs = resolve_layer_policies(
        config.layer_policies,
        available_layer_names,
        config,
    )

    for name, layer in targets:
        effective_config = effective_configs.get(name, config)
        # v1 policy constraints (bits=8, group_size=-1) are enforced by
        # SmoothQuantConfig.__post_init__, which dataclasses.replace runs
        # when resolve_layer_policies builds the effective config.
        new_layer = _quantize_linear_with_smoothquant(layer, captured[name], effective_config)
        if layer.bias is not None:
            with torch.no_grad():
                new_layer.bias.copy_(layer.bias.data)
        # Adopt the replaced layer's dtype (RIL ISS-191): selective
        # quantization over an already fp16/bf16 base must yield a half
        # quant layer, otherwise the forward emits fp32 into the remaining
        # half-precision linears and crashes.
        new_layer = new_layer.to(layer.weight.dtype)
        _replace_module(model, name, new_layer)
        logger.info(
            f"Quantized layer {name}: {layer.weight.shape} → INT8 weight+activation "
            f"(alpha={effective_config.alpha}{'+search' if effective_config.search_alpha else ''})"
        )

    return model

quantize_model_smoothquant_with_collector

quantize_model_smoothquant_with_collector(model, collector, n_samples, config=None, target_modules=None, device=None)

Quantize a model using an existing calibration batch source.

Mirrors quantize_model_with_collector / quantize_model_awq_with_collector: materializes up to n_samples batches, then funnels them into quantize_model_smoothquant.

源代码位于： src/llm/quantization/smooth.py

def quantize_model_smoothquant_with_collector(
    model: nn.Module,
    collector: CalibrationDataCollector | Iterable[torch.Tensor],
    n_samples: int,
    config: SmoothQuantConfig | None = None,
    target_modules: Iterable[str] | None = None,
    device: torch.device | str | None = None,
) -> nn.Module:
    """Quantize a model using an existing calibration batch source.

    Mirrors ``quantize_model_with_collector`` / ``quantize_model_awq_with_collector``:
    materializes up to ``n_samples`` batches, then funnels them into
    ``quantize_model_smoothquant``.
    """
    if n_samples <= 0:
        raise ValueError(f"n_samples must be positive, got {n_samples}.")

    if not isinstance(collector, Iterable):
        raise TypeError(
            "collector must be an iterable of calibration batches; "
            "CalibrationDataCollector stores activation statistics, not batches"
        )
    batches: list[torch.Tensor] = []
    for i, batch in enumerate(collector):
        if not isinstance(batch, torch.Tensor):
            raise TypeError(f"calibration batches must be tensors, got {type(batch).__name__}")
        batches.append(batch)
        if i + 1 >= n_samples:
            break

    return quantize_model_smoothquant(
        model,
        calib_iter=iter(batches),
        config=config,
        target_modules=target_modules,
        device=device,
    )

AWQ

awq

AWQ (Lin et al., MLSys 2024) post-training quantization.

Activation-aware Weight Quantization protects the ~1% of weight channels that dominate quantization error: per-input-channel scales s are searched over a power-of-two grid to minimize the activation-weighted reconstruction error of the layer output, then the layer is group-quantized from W·s with the scale compensation x/s applied at forward time.

This module mirrors the GPTQ path in gptq.py: a frozen config, a stateful per-layer quantizer, and model-level entry points that capture per-layer calibration activations through forward hooks.

AWQConfig dataclass

Configuration for AWQ quantization.

属性：

名称	类型	描述
bits	int	Quantization bit width (4 or 8).
group_size	int	Quantization group size along input dim. -1 means per-channel (one scale per output row).
sym	bool	If True, symmetric quantization (no zero-point). Asymmetric AWQ (zero points) is not yet implemented.
n_grid	int	Number of power-of-two scale candidates in the grid search, centered on 1 (ratios 2**(-n_grid//2) ... 2**(n_grid//2)). Larger grids find better scales at more search cost.
clip_ratio	float | None	Optional weight clipping ratio in (0, 0.5]. When set, each layer's weights are clipped to [min + rho*(max-min), max - rho*(max-min)] before the scale search, suppressing outlier magnitudes. None = no clipping.
layer_policies	tuple[LayerQuantPolicy, ...]	Atomic per-layer override policies (algorithm-agnostic LayerQuantPolicy tuples), same semantics as GPTQ (see ADR-008).

源代码位于： src/llm/quantization/awq.py

@dataclass(frozen=True)
class AWQConfig:
    """Configuration for AWQ quantization.

    Attributes:
        bits: Quantization bit width (4 or 8).
        group_size: Quantization group size along input dim.
            -1 means per-channel (one scale per output row).
        sym: If True, symmetric quantization (no zero-point). Asymmetric
            AWQ (zero points) is not yet implemented.
        n_grid: Number of power-of-two scale candidates in the grid search,
            centered on 1 (ratios ``2**(-n_grid//2) ... 2**(n_grid//2)``).
            Larger grids find better scales at more search cost.
        clip_ratio: Optional weight clipping ratio in (0, 0.5]. When set,
            each layer's weights are clipped to
            ``[min + rho*(max-min), max - rho*(max-min)]`` before the scale
            search, suppressing outlier magnitudes. None = no clipping.
        layer_policies: Atomic per-layer override policies (algorithm-agnostic
            ``LayerQuantPolicy`` tuples), same semantics as GPTQ (see ADR-008).
    """

    bits: int = 4
    group_size: int = 128
    sym: bool = True
    n_grid: int = 20
    clip_ratio: float | None = None

    # Per-layer atomic override policies (additive; empty tuple = no override).
    layer_policies: tuple[LayerQuantPolicy, ...] = ()

    def __post_init__(self):
        if self.bits not in (4, 8):
            raise ValueError(
                f"AWQConfig.bits must be 4 or 8, got {self.bits}. "
                "For mixed precision, use target_modules to skip sensitive layers."
            )
        if self.group_size != -1 and self.group_size <= 0:
            raise ValueError(f"group_size must be -1 (per-channel) or positive, got {self.group_size}.")
        if self.n_grid < 1:
            raise ValueError(f"n_grid must be >= 1, got {self.n_grid}.")
        if self.clip_ratio is not None and not (0.0 < self.clip_ratio <= 0.5):
            raise ValueError(f"clip_ratio must be in (0, 0.5] or None, got {self.clip_ratio}.")
        for i, p in enumerate(self.layer_policies):
            if not isinstance(p, LayerQuantPolicy):
                raise TypeError(f"AWQConfig.layer_policies[{i}] must be LayerQuantPolicy; got {type(p).__name__}.")

AWQQuantizer

Stateful per-layer AWQ processor.

Lifecycle

q = AWQQuantizer(layer, config) for batch in calib_iter_for_this_layer: q.add_batch(batch) scale = q.search_scale() # per-input-channel AWQ scale packed, scales, effective_group_size = q.quantize(scale)

源代码位于： src/llm/quantization/awq.py

class AWQQuantizer:
    """Stateful per-layer AWQ processor.

    Lifecycle:
        q = AWQQuantizer(layer, config)
        for batch in calib_iter_for_this_layer:
            q.add_batch(batch)
        scale = q.search_scale()   # per-input-channel AWQ scale
        packed, scales, effective_group_size = q.quantize(scale)
    """

    def __init__(self, layer: nn.Linear, config: AWQConfig):
        self.config = config
        self.layer = layer
        self.device = layer.weight.device
        self.compute_dtype = torch.float32

        self.out_features, self.in_features = layer.weight.shape

        # Group quantization requires the *effective* group size (clamped to
        # in_features, so a group larger than the row is a single group) to
        # divide in_features (the packing path builds
        # ``in_features // group_size`` groups and slices columns
        # [g*gs:(g+1)*gs)); reject non-divisible effective group sizes up
        # front with a clear error instead of a late packing bug.
        gs = min(self.config.group_size, self.in_features)
        if self.config.group_size != -1 and self.in_features % gs != 0:
            raise ValueError(
                f"group_size ({self.config.group_size}) must divide in_features "
                f"({self.in_features}); got remainder {self.in_features % gs}. "
                "Use group_size=-1 (per-channel) or a divisor of in_features."
            )

        self.act_abs_sum = torch.zeros(
            self.in_features,
            dtype=self.compute_dtype,
            device=self.device,
        )
        self.n_samples = 0

    def add_batch(self, x: torch.Tensor) -> None:
        """Accumulate per-input-channel absolute activation sums.

        Args:
            x: Input activations to `self.layer`, shape [..., in_features].
                Leading dims are flattened; only the mean absolute activation
                per channel is retained (that is all the scale search needs).
        """
        x = x.to(device=self.device, dtype=self.compute_dtype)
        if x.dim() == 1:
            x = x.unsqueeze(0)
        x = x.reshape(-1, x.shape[-1])  # flatten leading dims

        n = x.shape[0]
        if n == 0:
            return
        self.act_abs_sum += x.abs().sum(dim=0)
        self.n_samples += n

    def act_mean(self) -> torch.Tensor:
        """Mean absolute activation per input channel [in_features]."""
        if self.n_samples == 0:
            raise RuntimeError(
                "No calibration data accumulated (n_samples=0). Feed at least one calibration batch via add_batch()."
            )
        return self.act_abs_sum / self.n_samples

    def search_scale(self) -> torch.Tensor:
        """Run the activation-aware grid search; returns per-channel scale [in_f]."""
        if not self.config.sym:
            raise NotImplementedError("Asymmetric AWQ not yet implemented. Use sym=True.")

        w = self.layer.weight.detach().clone().to(device=self.device, dtype=self.compute_dtype)
        act = self.act_mean()

        if self.config.clip_ratio is not None:
            rho = self.config.clip_ratio
            w_min = w.min()
            w_max = w.max()
            lower = w_min + rho * (w_max - w_min)
            upper = w_max - rho * (w_max - w_min)
            w = w.clamp(lower, upper)

        return _search_scale(
            w,
            act,
            bits=self.config.bits,
            group_size=self.config.group_size,
            n_grid=self.config.n_grid,
        )

    def quantize(self, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
        """Quantize ``W·scale`` into packed storage.

        Returns:
            weight_packed: int8 packed weights.
            scales: per-group (or per-channel) fp32 scales.
            effective_group_size: group size actually used for packing
                (never larger than ``in_features``; -1 for per-channel).
        """
        w = self.layer.weight.detach().clone().to(device=self.device, dtype=self.compute_dtype)
        if self.config.clip_ratio is not None:
            rho = self.config.clip_ratio
            w_min = w.min()
            w_max = w.max()
            lower = w_min + rho * (w_max - w_min)
            upper = w_max - rho * (w_max - w_min)
            w = w.clamp(lower, upper)
        w = w * scale  # [out_f, in_f] * [in_f] broadcasts per input channel
        packed, scales, effective_group_size = _pack_weights(w, self.config.bits, self.config.group_size)
        return packed, scales, effective_group_size

add_batch

add_batch(x)

Accumulate per-input-channel absolute activation sums.

参数：

名称	类型	描述	默认
x	Tensor	Input activations to self.layer, shape [..., in_features]. Leading dims are flattened; only the mean absolute activation per channel is retained (that is all the scale search needs).	必需

源代码位于： src/llm/quantization/awq.py

def add_batch(self, x: torch.Tensor) -> None:
    """Accumulate per-input-channel absolute activation sums.

    Args:
        x: Input activations to `self.layer`, shape [..., in_features].
            Leading dims are flattened; only the mean absolute activation
            per channel is retained (that is all the scale search needs).
    """
    x = x.to(device=self.device, dtype=self.compute_dtype)
    if x.dim() == 1:
        x = x.unsqueeze(0)
    x = x.reshape(-1, x.shape[-1])  # flatten leading dims

    n = x.shape[0]
    if n == 0:
        return
    self.act_abs_sum += x.abs().sum(dim=0)
    self.n_samples += n

act_mean

act_mean()

Mean absolute activation per input channel [in_features].

源代码位于： src/llm/quantization/awq.py

def act_mean(self) -> torch.Tensor:
    """Mean absolute activation per input channel [in_features]."""
    if self.n_samples == 0:
        raise RuntimeError(
            "No calibration data accumulated (n_samples=0). Feed at least one calibration batch via add_batch()."
        )
    return self.act_abs_sum / self.n_samples

search_scale

search_scale()

Run the activation-aware grid search; returns per-channel scale [in_f].

源代码位于： src/llm/quantization/awq.py

def search_scale(self) -> torch.Tensor:
    """Run the activation-aware grid search; returns per-channel scale [in_f]."""
    if not self.config.sym:
        raise NotImplementedError("Asymmetric AWQ not yet implemented. Use sym=True.")

    w = self.layer.weight.detach().clone().to(device=self.device, dtype=self.compute_dtype)
    act = self.act_mean()

    if self.config.clip_ratio is not None:
        rho = self.config.clip_ratio
        w_min = w.min()
        w_max = w.max()
        lower = w_min + rho * (w_max - w_min)
        upper = w_max - rho * (w_max - w_min)
        w = w.clamp(lower, upper)

    return _search_scale(
        w,
        act,
        bits=self.config.bits,
        group_size=self.config.group_size,
        n_grid=self.config.n_grid,
    )

quantize

quantize(scale)

Quantize W·scale into packed storage.

返回：

名称	类型	描述
weight_packed	Tensor	int8 packed weights.
scales	Tensor	per-group (or per-channel) fp32 scales.
effective_group_size	int	group size actually used for packing (never larger than in_features; -1 for per-channel).

源代码位于： src/llm/quantization/awq.py

def quantize(self, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
    """Quantize ``W·scale`` into packed storage.

    Returns:
        weight_packed: int8 packed weights.
        scales: per-group (or per-channel) fp32 scales.
        effective_group_size: group size actually used for packing
            (never larger than ``in_features``; -1 for per-channel).
    """
    w = self.layer.weight.detach().clone().to(device=self.device, dtype=self.compute_dtype)
    if self.config.clip_ratio is not None:
        rho = self.config.clip_ratio
        w_min = w.min()
        w_max = w.max()
        lower = w_min + rho * (w_max - w_min)
        upper = w_max - rho * (w_max - w_min)
        w = w.clamp(lower, upper)
    w = w * scale  # [out_f, in_f] * [in_f] broadcasts per input channel
    packed, scales, effective_group_size = _pack_weights(w, self.config.bits, self.config.group_size)
    return packed, scales, effective_group_size

quantize_model_awq

quantize_model_awq(model, calib_iter, config=None, target_modules=None, device=None)

Quantize a model with AWQ.

参数：

名称	类型	描述	默认
model	Module	nn.Module containing nn.Linear layers to quantize.	必需
calib_iter	Iterator[Tensor]	Iterator yielding input tensors for the model forward pass.	必需
config	AWQConfig | None	AWQConfig (default: 4-bit, group_size=128, symmetric, n_grid=20).	None
target_modules	Iterable[str] | None	Iterable of fully-qualified layer names to quantize. If None, all nn.Linear layers are quantized.	None
device	device | str | None	Device to run calibration on (default: model's device).	None

返回：

类型	描述
Module	The model with nn.Linear layers replaced by AWQQuantizedLinear.

引发：

类型	描述
ValueError	If model has no nn.Linear, target_modules unmatched, layer already quantized, or calibration is empty.

源代码位于： src/llm/quantization/awq.py

def quantize_model_awq(
    model: nn.Module,
    calib_iter: Iterator[torch.Tensor],
    config: AWQConfig | None = None,
    target_modules: Iterable[str] | None = None,
    device: torch.device | str | None = None,
) -> nn.Module:
    """Quantize a model with AWQ.

    Args:
        model: nn.Module containing nn.Linear layers to quantize.
        calib_iter: Iterator yielding input tensors for the model forward pass.
        config: AWQConfig (default: 4-bit, group_size=128, symmetric, n_grid=20).
        target_modules: Iterable of fully-qualified layer names to quantize.
            If None, all nn.Linear layers are quantized.
        device: Device to run calibration on (default: model's device).

    Returns:
        The model with nn.Linear layers replaced by AWQQuantizedLinear.

    Raises:
        ValueError: If model has no nn.Linear, target_modules unmatched, layer
            already quantized, or calibration is empty.
    """
    config = config or AWQConfig()
    if device is not None:
        model = model.to(device)

    for n, m in model.named_modules():
        if isinstance(m, AWQQuantizedLinear):
            raise ValueError(f"Layer {n} is already AWQ-quantized. Pass a fresh model or unquantize first.")

    linear_layers = [(n, m) for n, m in model.named_modules() if isinstance(m, nn.Linear)]
    if not linear_layers:
        raise ValueError("model has no nn.Linear modules; nothing to quantize.")

    if target_modules is not None:
        target_set = set(target_modules)
        all_names = {n for n, _ in linear_layers}
        matched = target_set & all_names
        if not matched:
            available = sorted(all_names)[:10]
            raise ValueError(
                f"target_modules {list(target_set)} matched no nn.Linear. "
                f"Available: {available}{'...' if len(all_names) > 10 else ''}"
            )
        targets = [(n, m) for n, m in linear_layers if n in target_set]
    else:
        targets = linear_layers

    calib_batches = list(calib_iter)
    if not calib_batches:
        raise ValueError("calib_iter is empty; need at least 1 batch for activation statistics.")

    # Per-layer input capture via forward hooks (same mechanism as GPTQ).
    captured: dict[str, list[torch.Tensor]] = {n: [] for n, _ in targets}
    hooks = []

    def make_hook(name: str):
        def hook(_module, inputs, _output):
            captured[name].append(inputs[0].detach().clone())

        return hook

    for n, m in targets:
        hooks.append(m.register_forward_hook(make_hook(n)))

    model.eval()
    with torch.no_grad():
        param_device = next(model.parameters()).device
        try:
            # Feed EVERY calibration batch so each layer's activation stats
            # cover the full calibration set.  Previously only
            # calib_batches[0] was forwarded, silently dropping later batches
            # while the fallback path below used all of them.
            for batch in calib_batches:
                _ = model(batch.to(param_device))
        except (RuntimeError, ValueError, TypeError) as e:
            logger.debug(f"Model forward failed during calibration: {e}; falling back to direct layer calls.")

    # A mid-loop forward failure must not leave the per-layer capture counts
    # diverged (a layer that never got its first batch holds ZERO captures and
    # later crashes with "No calibration data accumulated (n_samples=0)" after
    # earlier layers were already replaced). Any inconsistency -> rebuild every
    # target's captures from direct calls over the full set (RIL ISS-136).
    expected_captures = len(calib_batches)
    capt_sizes = {n: len(v) for n, v in captured.items()}
    any_captured = any(s > 0 for s in capt_sizes.values())
    consistent = bool(capt_sizes) and all(s == expected_captures for s in capt_sizes.values())
    if not any_captured or not consistent:
        if any_captured and not consistent:
            logger.warning(
                "Per-layer calibration captures diverged after a partial forward "
                "failure (%s); falling back to direct layer calls for ALL targets "
                "so every layer quantizes over the same calibration set.",
                capt_sizes,
            )
        elif not any_captured:
            logger.warning(
                "Model forward failed on EVERY calibration batch (see the "
                "DEBUG log above for the first error); no per-layer inputs "
                "were captured. Falling back to feeding the raw calibration "
                "batches directly as each target layer's inputs — this is only "
                "valid when those tensors ARE the layers' activations (e.g. a "
                "bare sequence of target layers). If the model embeds or "
                "reshapes inputs first (a real decoder), the quantized weights "
                "will be garbage; fix the model forward signature or use "
                "quantize_model_with_collector."
            )
        for h in hooks:
            h.remove()
        for n, _m in targets:
            captured[n] = [batch.detach().clone() for batch in calib_batches]

    for h in hooks:
        h.remove()

    available_layer_names = {n for n, _ in targets}
    effective_configs = resolve_layer_policies(
        config.layer_policies,
        available_layer_names,
        config,
    )

    for name, layer in targets:
        effective_config = effective_configs.get(name, config)
        new_layer = _quantize_linear_with_awq(layer, captured[name], effective_config)
        if layer.bias is not None:
            with torch.no_grad():
                new_layer.bias.copy_(layer.bias.data)
        # Adopt the replaced layer's dtype (RIL ISS-191): selective
        # quantization over an already fp16/bf16 base must yield a half
        # quant layer, otherwise the forward emits fp32 into the remaining
        # half-precision linears and crashes.
        new_layer = new_layer.to(layer.weight.dtype)
        _replace_module(model, name, new_layer)
        logger.info(
            f"Quantized layer {name}: {layer.weight.shape} → "
            f"{effective_config.bits}-bit, group_size={effective_config.group_size}"
        )

    return model

quantize_model_awq_with_collector

quantize_model_awq_with_collector(model, collector, n_samples, config=None, target_modules=None, device=None)

Quantize a model using an existing calibration batch source.

Trainer-loop entry point mirroring quantize_model_with_collector: materializes up to n_samples batches, then funnels them into quantize_model_awq.

参数：

名称	类型	描述	默认
model	Module	nn.Module to quantize.	必需
collector	CalibrationDataCollector | Iterable[Tensor]	Iterable yielding Tensor batches. Up to n_samples batches are consumed.	必需
n_samples	int	Maximum number of batches to use for calibration.	必需
config	AWQConfig | None	AWQConfig (default: 4-bit, group_size=128, symmetric).	None
target_modules	Iterable[str] | None	Optional layer-name filter forwarded to quantize_model_awq.	None
device	device | str | None	Target device forwarded to quantize_model_awq.	None

返回：

类型	描述
Module	The quantized model (same instance as model).

源代码位于： src/llm/quantization/awq.py

def quantize_model_awq_with_collector(
    model: nn.Module,
    collector: CalibrationDataCollector | Iterable[torch.Tensor],
    n_samples: int,
    config: AWQConfig | None = None,
    target_modules: Iterable[str] | None = None,
    device: torch.device | str | None = None,
) -> nn.Module:
    """Quantize a model using an existing calibration batch source.

    Trainer-loop entry point mirroring ``quantize_model_with_collector``:
    materializes up to ``n_samples`` batches, then funnels them into
    ``quantize_model_awq``.

    Args:
        model: nn.Module to quantize.
        collector: Iterable yielding Tensor batches. Up to ``n_samples``
            batches are consumed.
        n_samples: Maximum number of batches to use for calibration.
        config: AWQConfig (default: 4-bit, group_size=128, symmetric).
        target_modules: Optional layer-name filter forwarded to
            ``quantize_model_awq``.
        device: Target device forwarded to ``quantize_model_awq``.

    Returns:
        The quantized model (same instance as ``model``).
    """
    if n_samples <= 0:
        raise ValueError(f"n_samples must be positive, got {n_samples}.")

    if not isinstance(collector, Iterable):
        raise TypeError(
            "collector must be an iterable of calibration batches; "
            "CalibrationDataCollector stores activation statistics, not batches"
        )
    batches: list[torch.Tensor] = []
    for i, batch in enumerate(collector):
        if not isinstance(batch, torch.Tensor):
            raise TypeError(f"calibration batches must be tensors, got {type(batch).__name__}")
        batches.append(batch)
        if i + 1 >= n_samples:
            break

    return quantize_model_awq(
        model,
        calib_iter=iter(batches),
        config=config,
        target_modules=target_modules,
        device=device,
    )

GPTQ

gptq

GPTQ (Frantar et al. 2022) post-training quantization.

Provides 4-bit / 8-bit Hessian-aware quantization orthogonal to the simple-PTQ path in ptq.py.

GPTQConfig dataclass

Configuration for GPTQ quantization.

属性：

名称	类型	描述
bits	int	Quantization bit width (4 or 8).
group_size	int	Quantization group size along input dim. -1 means per-channel (one scale per output row). Positive integer g means one scale per g consecutive input cols.
sym	bool	If True, symmetric quantization (no zero-point).
percdamp	float	Hessian damping as a fraction of mean(diag(H)). Prevents numerical issues when H is near-singular.
blocksize	int	Number of weight columns processed per Cholesky block. Larger = faster but more memory. Must be divisible by group_size when group_size > 0.
act_order	bool	If True, sort weight columns by diag(H) descending before quantization. Improves accuracy at slight cost.
static_groups	bool	If True, compute group partitions once and reuse across all layers. Faster, slight accuracy loss.
layer_policies	tuple[LayerQuantPolicy, ...]	Atomic per-layer override policies (algorithm-agnostic LayerQuantPolicy tuples). Empty tuple (default) → all layers use the base config; otherwise each policy dispatches its overrides to its target_modules subset. See ADR-008 and docs/superpowers/specs/2026-07-22-mixed-precision-quantization-design.md.

源代码位于： src/llm/quantization/gptq.py

@dataclass(frozen=True)
class GPTQConfig:
    """Configuration for GPTQ quantization.

    Attributes:
        bits: Quantization bit width (4 or 8).
        group_size: Quantization group size along input dim.
            -1 means per-channel (one scale per output row).
            Positive integer g means one scale per g consecutive input cols.
        sym: If True, symmetric quantization (no zero-point).
        percdamp: Hessian damping as a fraction of mean(diag(H)).
            Prevents numerical issues when H is near-singular.
        blocksize: Number of weight columns processed per Cholesky block.
            Larger = faster but more memory. Must be divisible by group_size
            when group_size > 0.
        act_order: If True, sort weight columns by diag(H) descending
            before quantization. Improves accuracy at slight cost.
        static_groups: If True, compute group partitions once and reuse
            across all layers. Faster, slight accuracy loss.
        layer_policies: Atomic per-layer override policies (algorithm-agnostic
            LayerQuantPolicy tuples). Empty tuple (default) → all layers use
            the base config; otherwise each policy dispatches its overrides
            to its target_modules subset. See ADR-008 and
            docs/superpowers/specs/2026-07-22-mixed-precision-quantization-design.md.
    """

    bits: int = 4
    group_size: int = 128
    sym: bool = True
    percdamp: float = 0.01
    blocksize: int = 128
    act_order: bool = False
    static_groups: bool = False

    # Per-layer atomic override policies (additive; empty tuple = no override).
    # Each LayerQuantPolicy is validated at its own __post_init__; this field
    # only enforces that all elements are LayerQuantPolicy instances.
    layer_policies: tuple[LayerQuantPolicy, ...] = ()

    def __post_init__(self):
        if self.bits not in (4, 8):
            raise ValueError(
                f"GPTQConfig.bits must be 4 or 8, got {self.bits}. "
                f"For mixed precision, use target_modules to skip sensitive layers."
            )
        if self.group_size != -1 and self.group_size <= 0:
            raise ValueError(f"group_size must be -1 (per-channel) or positive, got {self.group_size}.")
        if not (0.0 < self.percdamp < 1.0):
            raise ValueError(f"percdamp must be in (0, 1), got {self.percdamp}.")
        if self.blocksize <= 0:
            raise ValueError(f"blocksize must be positive, got {self.blocksize}.")
        if self.group_size > 0 and self.blocksize % self.group_size != 0:
            raise ValueError(
                f"blocksize ({self.blocksize}) must be divisible by "
                f"group_size ({self.group_size}) for correct packing alignment."
            )
        # Validate layer_policies contains only LayerQuantPolicy instances.
        # Per-policy field validation (bits, group_size, etc.) is done in
        # LayerQuantPolicy.__post_init__ at construction time.
        for i, p in enumerate(self.layer_policies):
            if not isinstance(p, LayerQuantPolicy):
                raise TypeError(f"GPTQConfig.layer_policies[{i}] must be LayerQuantPolicy; got {type(p).__name__}.")

GPTQQuantizer

Stateful per-layer GPTQ processor.

Lifecycle

q = GPTQQuantizer(layer, config) for batch in calib_iter_for_this_layer: q.add_batch(batch) W_packed, scales, zeros = q.quantize()

源代码位于： src/llm/quantization/gptq.py

class GPTQQuantizer:
    """Stateful per-layer GPTQ processor.

    Lifecycle:
        q = GPTQQuantizer(layer, config)
        for batch in calib_iter_for_this_layer:
            q.add_batch(batch)
        W_packed, scales, zeros = q.quantize()
    """

    def __init__(self, layer: nn.Linear, config: GPTQConfig):
        self.config = config
        self.layer = layer
        self.device = layer.weight.device
        # Compute in float32 for numerical stability of Cholesky
        self.compute_dtype = torch.float32

        # Weight dimensions
        self.out_features, self.in_features = layer.weight.shape

        # Group quantization requires the *effective* group size (clamped to
        # in_features, so a group larger than the row is a single group) to
        # divide in_features: the packing / dequant paths assume an integral
        # number of full groups (``in_features // group_size`` then
        # ``repeat_interleave(group_size)``). Reject a non-divisible effective
        # group up front with a clear error instead of a late
        # ``repeat_interleave`` broadcast crash deep in packing.
        gs = min(self.config.group_size, self.in_features)
        if self.config.group_size != -1 and self.in_features % gs != 0:
            raise ValueError(
                f"group_size ({self.config.group_size}) must divide in_features "
                f"({self.in_features}); got remainder {self.in_features % gs}. "
                "Use group_size=-1 (per-channel) or a divisor of in_features."
            )
        # 4-bit packing stores two weights per int8 byte over the whole
        # tensor, so an odd total weight count would crash mid-pipeline in
        # ``_pack_4bit`` (after Hessian accumulation) for BOTH per-channel and
        # per-group 4-bit. Fail fast with the same "clear error instead of a
        # late crash" spirit as the group-size check above (round-81 quant
        # deep-dive F3 — the guard previously covered only group_size == -1).
        if self.config.bits == 4 and (self.out_features * self.in_features) % 2 != 0:
            raise ValueError(
                f"4-bit quantization requires an even total weight count; "
                f"{self.layer.weight.shape} has an odd product "
                f"({self.out_features * self.in_features}). Use an "
                "architecture with even in/out features."
            )

        # Hessian accumulator
        self.H = torch.zeros(
            (self.in_features, self.in_features),
            dtype=self.compute_dtype,
            device=self.device,
        )
        self.n_samples = 0

    def add_batch(self, x: torch.Tensor) -> None:
        """Accumulate Hessian contribution from a calibration batch.

        Maintains the invariant H == (2 / N_total) · Σ X_b^T X_b so that
        multiple mini-batches produce the same H as a single concatenated
        add_batch (Frantar 2022, eq. 3). Uses the canonical EMA-style
        rescale: H_new = (N_old / N_new) · H_old + (2 / N_new) · X^T X.

        Args:
            x: Input activations to `self.layer`, shape [..., in_features].
                Will be flattened to [N, in_features] internally.
        """
        x = x.to(device=self.device, dtype=self.compute_dtype)
        if x.dim() == 1:
            x = x.unsqueeze(0)
        x = x.reshape(-1, x.shape[-1])  # flatten leading dims

        n = x.shape[0]
        if n == 0:
            return

        new_total = self.n_samples + n
        # Rescale previous contribution to its raw Σ X^T X form, then
        # re-apply the (2 / N_new) factor on the new total. This makes
        # H the exact (2 / N_total) · Σ X^T X across any batch partition.
        self.H *= self.n_samples / new_total
        self.n_samples = new_total
        self.H += (2.0 / new_total) * (x.t() @ x)

    def quantize(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
        """Run GPTQ on accumulated Hessian.

        Returns:
            W_q: Quantized weights (integer-valued, stored as fp32),
                 shape [out_features, in_features]. Multiply by `scales`
                 to dequantize: W_recon = W_q * scales.
            scales: Per-row scale if group_size=-1 [out_features, 1],
                    else per-group scale [out_features, in_features // group_size].
            zeros: Per-group zero-points, or None (symmetric only in v1).

        Raises:
            RuntimeError: If calibration is empty or Hessian is ill-conditioned
                          (rank-deficient with insufficient damping).
        """
        # Guard: zero calibration data
        if self.n_samples == 0:
            raise RuntimeError(
                "No calibration data accumulated (n_samples=0). "
                "Hessian is empty. Try increasing percdamp or check calibration data quality."
            )

        w = self.layer.weight.detach().clone().to(device=self.device, dtype=self.compute_dtype)
        h = self.H.clone()

        # Detect rank-deficient Hessian BEFORE dead handling
        n_dead = int(torch.sum(torch.diag(h) == 0).item())

        # Handle all-zero columns (degenerate / unused features)
        dead = torch.diag(h) == 0
        h[dead, dead] = 1.0
        w[:, dead] = 0.0

        # Damping for numerical stability (Frantar 2022, eq. 4)
        damp = self.config.percdamp * torch.mean(torch.diag(h))
        diag_idx = torch.arange(self.in_features, device=self.device)
        h[diag_idx, diag_idx] += damp

        # Guard: rank-deficient Hessian with insufficient damping.
        # If most columns are dead (no variance in calibration), low damping
        # leaves the live columns effectively unscaled relative to the dead
        # ones (which dead-handling set to 1.0). Reject and tell the user.
        if n_dead > self.in_features // 2 and damp < 0.25:
            raise RuntimeError(
                f"Hessian is rank-deficient ({n_dead}/{self.in_features} "
                f"columns have zero variance). Damping (percdamp="
                f"{self.config.percdamp}) is insufficient. Try increasing "
                f"percdamp (e.g. 0.5) or check calibration data quality."
            )

        # Cholesky inverse: only H^-1 is needed for canonical Frantar error correction.
        try:
            h_inv = torch.linalg.inv(h)
        except RuntimeError as e:
            raise RuntimeError(
                f"Hessian is not positive-definite even after damping "
                f"(percdamp={self.config.percdamp}). "
                f"Try increasing percdamp (e.g. 0.1) or check calibration data quality."
            ) from e

        # Optional act-order: sort columns by diag(H_inv) descending
        if self.config.act_order:
            perm = torch.argsort(torch.diag(h_inv), descending=True)
            w = w[:, perm]
            h_inv = h_inv[perm][:, perm]

        # Asymmetric not yet implemented — fail fast before entering column loop
        if not self.config.sym:
            raise NotImplementedError("Asymmetric GPTQ not yet implemented. Use sym=True.")

        # Compute per-row scale ONCE for per-channel (group_size=-1).
        # Canonical GPTQ (Frantar 2022, eq. 5): scale = w.abs().max() / qmax.
        qmax = 2 ** (self.config.bits - 1) - 1
        if self.config.group_size == -1:
            row_scales = w.abs().max(dim=1)[0] / qmax  # [out_f]
            row_scales = row_scales.clamp(min=1e-8)

        # Quantize column-by-column with error correction (Frantar 2022, eq. 5)
        q_out = torch.zeros_like(w)

        for i in range(0, self.in_features, self.config.blocksize):
            i_end = min(i + self.config.blocksize, self.in_features)
            count = i_end - i

            w1 = w[:, i:i_end].clone()
            q1 = torch.zeros_like(w1)
            err1 = torch.zeros_like(w1)
            hinv1 = h_inv[i:i_end, i:i_end]

            # Per-group scale cache: skip recomputation within the same group
            cached_group_idx: int | None = None
            cached_scale: torch.Tensor | None = None

            for j in range(count):
                col = w1[:, j]
                d = hinv1[j, j]

                if self.config.group_size == -1:
                    scale = row_scales  # [out_f], same for all columns
                else:
                    # Per-group scale: cache by group_idx to avoid redundant
                    # .abs().max(). The scale MUST be PER OUTPUT ROW
                    # (``max(dim=1)``), matching exactly the per-row scales
                    # stored for dequantization below — a whole-group scalar
                    # (``.max()`` over all rows) systematically shrank the
                    # reconstruction of any row whose magnitude was below the
                    # group max (RIL ISS-057).
                    gs = self.config.group_size
                    group_idx = (i + j) // gs
                    if group_idx != cached_group_idx:
                        group_start = group_idx * gs
                        group_end = group_start + gs
                        w_group = w[:, group_start:group_end]
                        cached_scale = (w_group.abs().max(dim=1)[0] / qmax).clamp(min=1e-8)  # [out_f]
                        cached_group_idx = group_idx
                    scale = cached_scale
                    if scale is None:
                        raise RuntimeError("group scale was not computed")

                # Quantize to INTEGER (clamped to symmetric range)
                q_int = torch.round(col / scale).clamp(-qmax - 1, qmax)
                q1[:, j] = q_int  # store integer-valued fp32

                # Error correction: propagate quantization error to remaining columns.
                # Canonical GPTQ (Frantar 2022, eq. 5): the error vector is
                # scaled by H^-1 and propagated via the row of H^-1 starting
                # at the current column. Note: H^-1 (not its Cholesky factor U)
                # must be used in BOTH the denominator and the propagation to
                # match the canonical formula exactly.
                err = (col - q_int * scale) / d
                if j + 1 < count:
                    w1[:, j + 1 :] -= err.unsqueeze(1) * hinv1[j, j + 1 :].unsqueeze(0)

                err1[:, j] = err

            q_out[:, i:i_end] = q1

            # Propagate block errors to all remaining columns (canonical GPTQ).
            # Uses H^-1 (not U) to match the canonical Frantar 2022 formula.
            if i_end < self.in_features:
                w[:, i_end:] -= err1 @ h_inv[i:i_end, i_end:]

        # Undo act-order permutation if applied
        if self.config.act_order:
            invperm = torch.argsort(perm)
            q_out = q_out[:, invperm]

        # Return per-row or per-group scales matching the scales used in the loop
        if self.config.group_size != -1:
            # Clamp to in_features: a group larger than the row is a single
            # group (matches the packing step's ``effective_group_size``).
            gs = min(self.config.group_size, self.in_features)
            n_groups = self.in_features // gs
            scales = torch.zeros(
                self.out_features,
                n_groups,
                dtype=torch.float32,
                device=self.device,
            )
            for g in range(n_groups):
                s = g * gs
                e = s + gs
                w_g = w[:, s:e]
                scales[:, g] = w_g.abs().max(dim=1)[0] / qmax
                scales[:, g] = scales[:, g].clamp(min=1e-8)

            if self.config.act_order:
                # ``q_out`` above was restored to the *original* column order
                # (``q_out = q_out[:, invperm]``), but the per-group scales
                # were derived from the *permuted* columns ``w`` (which is
                # never un-permuted).  Each original column ``j`` sat at
                # permuted position ``invperm[j]``, so it must be scaled by
                # the group that position belonged to — i.e. the scales have
                # to be gathered per-column through the inverse permutation.
                # Leaving them in the contiguous ``[out, in_f // gs]`` layout
                # and letting the caller ``repeat_interleave`` would apply the
                # wrong group's scale to every column whose permutation
                # crossed a group boundary.
                scales = scales[:, invperm // gs]  # [out_f, in_f] per-column
        else:
            scales = row_scales.unsqueeze(1)  # [out_f, 1]

        zeros: torch.Tensor | None = None
        return q_out, scales, zeros

add_batch

add_batch(x)

Accumulate Hessian contribution from a calibration batch.

Maintains the invariant H == (2 / N_total) · Σ X_b^T X_b so that multiple mini-batches produce the same H as a single concatenated add_batch (Frantar 2022, eq. 3). Uses the canonical EMA-style rescale: H_new = (N_old / N_new) · H_old + (2 / N_new) · X^T X.

参数：

名称	类型	描述	默认
x	Tensor	Input activations to self.layer, shape [..., in_features]. Will be flattened to [N, in_features] internally.	必需

源代码位于： src/llm/quantization/gptq.py

def add_batch(self, x: torch.Tensor) -> None:
    """Accumulate Hessian contribution from a calibration batch.

    Maintains the invariant H == (2 / N_total) · Σ X_b^T X_b so that
    multiple mini-batches produce the same H as a single concatenated
    add_batch (Frantar 2022, eq. 3). Uses the canonical EMA-style
    rescale: H_new = (N_old / N_new) · H_old + (2 / N_new) · X^T X.

    Args:
        x: Input activations to `self.layer`, shape [..., in_features].
            Will be flattened to [N, in_features] internally.
    """
    x = x.to(device=self.device, dtype=self.compute_dtype)
    if x.dim() == 1:
        x = x.unsqueeze(0)
    x = x.reshape(-1, x.shape[-1])  # flatten leading dims

    n = x.shape[0]
    if n == 0:
        return

    new_total = self.n_samples + n
    # Rescale previous contribution to its raw Σ X^T X form, then
    # re-apply the (2 / N_new) factor on the new total. This makes
    # H the exact (2 / N_total) · Σ X^T X across any batch partition.
    self.H *= self.n_samples / new_total
    self.n_samples = new_total
    self.H += (2.0 / new_total) * (x.t() @ x)

quantize

quantize()

Run GPTQ on accumulated Hessian.

返回：

名称	类型	描述
W_q	Tensor	Quantized weights (integer-valued, stored as fp32), shape [out_features, in_features]. Multiply by scales to dequantize: W_recon = W_q * scales.
scales	Tensor	Per-row scale if group_size=-1 [out_features, 1], else per-group scale [out_features, in_features // group_size].
zeros	Tensor | None	Per-group zero-points, or None (symmetric only in v1).

引发：

类型	描述
RuntimeError	If calibration is empty or Hessian is ill-conditioned (rank-deficient with insufficient damping).

源代码位于： src/llm/quantization/gptq.py

def quantize(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """Run GPTQ on accumulated Hessian.

    Returns:
        W_q: Quantized weights (integer-valued, stored as fp32),
             shape [out_features, in_features]. Multiply by `scales`
             to dequantize: W_recon = W_q * scales.
        scales: Per-row scale if group_size=-1 [out_features, 1],
                else per-group scale [out_features, in_features // group_size].
        zeros: Per-group zero-points, or None (symmetric only in v1).

    Raises:
        RuntimeError: If calibration is empty or Hessian is ill-conditioned
                      (rank-deficient with insufficient damping).
    """
    # Guard: zero calibration data
    if self.n_samples == 0:
        raise RuntimeError(
            "No calibration data accumulated (n_samples=0). "
            "Hessian is empty. Try increasing percdamp or check calibration data quality."
        )

    w = self.layer.weight.detach().clone().to(device=self.device, dtype=self.compute_dtype)
    h = self.H.clone()

    # Detect rank-deficient Hessian BEFORE dead handling
    n_dead = int(torch.sum(torch.diag(h) == 0).item())

    # Handle all-zero columns (degenerate / unused features)
    dead = torch.diag(h) == 0
    h[dead, dead] = 1.0
    w[:, dead] = 0.0

    # Damping for numerical stability (Frantar 2022, eq. 4)
    damp = self.config.percdamp * torch.mean(torch.diag(h))
    diag_idx = torch.arange(self.in_features, device=self.device)
    h[diag_idx, diag_idx] += damp

    # Guard: rank-deficient Hessian with insufficient damping.
    # If most columns are dead (no variance in calibration), low damping
    # leaves the live columns effectively unscaled relative to the dead
    # ones (which dead-handling set to 1.0). Reject and tell the user.
    if n_dead > self.in_features // 2 and damp < 0.25:
        raise RuntimeError(
            f"Hessian is rank-deficient ({n_dead}/{self.in_features} "
            f"columns have zero variance). Damping (percdamp="
            f"{self.config.percdamp}) is insufficient. Try increasing "
            f"percdamp (e.g. 0.5) or check calibration data quality."
        )

    # Cholesky inverse: only H^-1 is needed for canonical Frantar error correction.
    try:
        h_inv = torch.linalg.inv(h)
    except RuntimeError as e:
        raise RuntimeError(
            f"Hessian is not positive-definite even after damping "
            f"(percdamp={self.config.percdamp}). "
            f"Try increasing percdamp (e.g. 0.1) or check calibration data quality."
        ) from e

    # Optional act-order: sort columns by diag(H_inv) descending
    if self.config.act_order:
        perm = torch.argsort(torch.diag(h_inv), descending=True)
        w = w[:, perm]
        h_inv = h_inv[perm][:, perm]

    # Asymmetric not yet implemented — fail fast before entering column loop
    if not self.config.sym:
        raise NotImplementedError("Asymmetric GPTQ not yet implemented. Use sym=True.")

    # Compute per-row scale ONCE for per-channel (group_size=-1).
    # Canonical GPTQ (Frantar 2022, eq. 5): scale = w.abs().max() / qmax.
    qmax = 2 ** (self.config.bits - 1) - 1
    if self.config.group_size == -1:
        row_scales = w.abs().max(dim=1)[0] / qmax  # [out_f]
        row_scales = row_scales.clamp(min=1e-8)

    # Quantize column-by-column with error correction (Frantar 2022, eq. 5)
    q_out = torch.zeros_like(w)

    for i in range(0, self.in_features, self.config.blocksize):
        i_end = min(i + self.config.blocksize, self.in_features)
        count = i_end - i

        w1 = w[:, i:i_end].clone()
        q1 = torch.zeros_like(w1)
        err1 = torch.zeros_like(w1)
        hinv1 = h_inv[i:i_end, i:i_end]

        # Per-group scale cache: skip recomputation within the same group
        cached_group_idx: int | None = None
        cached_scale: torch.Tensor | None = None

        for j in range(count):
            col = w1[:, j]
            d = hinv1[j, j]

            if self.config.group_size == -1:
                scale = row_scales  # [out_f], same for all columns
            else:
                # Per-group scale: cache by group_idx to avoid redundant
                # .abs().max(). The scale MUST be PER OUTPUT ROW
                # (``max(dim=1)``), matching exactly the per-row scales
                # stored for dequantization below — a whole-group scalar
                # (``.max()`` over all rows) systematically shrank the
                # reconstruction of any row whose magnitude was below the
                # group max (RIL ISS-057).
                gs = self.config.group_size
                group_idx = (i + j) // gs
                if group_idx != cached_group_idx:
                    group_start = group_idx * gs
                    group_end = group_start + gs
                    w_group = w[:, group_start:group_end]
                    cached_scale = (w_group.abs().max(dim=1)[0] / qmax).clamp(min=1e-8)  # [out_f]
                    cached_group_idx = group_idx
                scale = cached_scale
                if scale is None:
                    raise RuntimeError("group scale was not computed")

            # Quantize to INTEGER (clamped to symmetric range)
            q_int = torch.round(col / scale).clamp(-qmax - 1, qmax)
            q1[:, j] = q_int  # store integer-valued fp32

            # Error correction: propagate quantization error to remaining columns.
            # Canonical GPTQ (Frantar 2022, eq. 5): the error vector is
            # scaled by H^-1 and propagated via the row of H^-1 starting
            # at the current column. Note: H^-1 (not its Cholesky factor U)
            # must be used in BOTH the denominator and the propagation to
            # match the canonical formula exactly.
            err = (col - q_int * scale) / d
            if j + 1 < count:
                w1[:, j + 1 :] -= err.unsqueeze(1) * hinv1[j, j + 1 :].unsqueeze(0)

            err1[:, j] = err

        q_out[:, i:i_end] = q1

        # Propagate block errors to all remaining columns (canonical GPTQ).
        # Uses H^-1 (not U) to match the canonical Frantar 2022 formula.
        if i_end < self.in_features:
            w[:, i_end:] -= err1 @ h_inv[i:i_end, i_end:]

    # Undo act-order permutation if applied
    if self.config.act_order:
        invperm = torch.argsort(perm)
        q_out = q_out[:, invperm]

    # Return per-row or per-group scales matching the scales used in the loop
    if self.config.group_size != -1:
        # Clamp to in_features: a group larger than the row is a single
        # group (matches the packing step's ``effective_group_size``).
        gs = min(self.config.group_size, self.in_features)
        n_groups = self.in_features // gs
        scales = torch.zeros(
            self.out_features,
            n_groups,
            dtype=torch.float32,
            device=self.device,
        )
        for g in range(n_groups):
            s = g * gs
            e = s + gs
            w_g = w[:, s:e]
            scales[:, g] = w_g.abs().max(dim=1)[0] / qmax
            scales[:, g] = scales[:, g].clamp(min=1e-8)

        if self.config.act_order:
            # ``q_out`` above was restored to the *original* column order
            # (``q_out = q_out[:, invperm]``), but the per-group scales
            # were derived from the *permuted* columns ``w`` (which is
            # never un-permuted).  Each original column ``j`` sat at
            # permuted position ``invperm[j]``, so it must be scaled by
            # the group that position belonged to — i.e. the scales have
            # to be gathered per-column through the inverse permutation.
            # Leaving them in the contiguous ``[out, in_f // gs]`` layout
            # and letting the caller ``repeat_interleave`` would apply the
            # wrong group's scale to every column whose permutation
            # crossed a group boundary.
            scales = scales[:, invperm // gs]  # [out_f, in_f] per-column
    else:
        scales = row_scales.unsqueeze(1)  # [out_f, 1]

    zeros: torch.Tensor | None = None
    return q_out, scales, zeros

quantize_model_gptq

quantize_model_gptq(model, calib_iter, config=None, target_modules=None, device=None)

Quantize a model with GPTQ.

参数：

名称	类型	描述	默认
model	Module	nn.Module containing nn.Linear layers to quantize.	必需
calib_iter	Iterator[Tensor]	Iterator yielding input tensors for the model forward pass.	必需
config	GPTQConfig | None	GPTQConfig (default: 4-bit, group_size=128, symmetric).	None
target_modules	Iterable[str] | None	Iterable of fully-qualified layer names to quantize. If None, all nn.Linear layers are quantized.	None
device	device | str | None	Device to run calibration on (default: model's device).	None

返回：

类型	描述
Module	The model with nn.Linear layers replaced by GPTQQuantizedLinear.

引发：

类型	描述
ValueError	If model has no nn.Linear, target_modules unmatched, or layer already quantized.

源代码位于： src/llm/quantization/gptq.py

def quantize_model_gptq(
    model: nn.Module,
    calib_iter: Iterator[torch.Tensor],
    config: GPTQConfig | None = None,
    target_modules: Iterable[str] | None = None,
    device: torch.device | str | None = None,
) -> nn.Module:
    """Quantize a model with GPTQ.

    Args:
        model: nn.Module containing nn.Linear layers to quantize.
        calib_iter: Iterator yielding input tensors for the model forward pass.
        config: GPTQConfig (default: 4-bit, group_size=128, symmetric).
        target_modules: Iterable of fully-qualified layer names to quantize.
            If None, all nn.Linear layers are quantized.
        device: Device to run calibration on (default: model's device).

    Returns:
        The model with nn.Linear layers replaced by GPTQQuantizedLinear.

    Raises:
        ValueError: If model has no nn.Linear, target_modules unmatched, or layer already quantized.
    """
    config = config or GPTQConfig()
    if device is not None:
        model = model.to(device)

    # Check for already-quantized layers FIRST so that a model with
    # only GPTQQuantizedLinear surfaces the actionable error instead of
    # the generic "no nn.Linear" message.
    for n, m in model.named_modules():
        if isinstance(m, GPTQQuantizedLinear):
            raise ValueError(f"Layer {n} is already GPTQ-quantized. Pass a fresh model or unquantize first.")

    linear_layers = [(n, m) for n, m in model.named_modules() if isinstance(m, nn.Linear)]
    if not linear_layers:
        raise ValueError("model has no nn.Linear modules; nothing to quantize.")

    if target_modules is not None:
        target_set = set(target_modules)
        all_names = {n for n, _ in linear_layers}
        matched = target_set & all_names
        if not matched:
            available = sorted(all_names)[:10]
            raise ValueError(
                f"target_modules {list(target_set)} matched no nn.Linear. "
                f"Available: {available}{'...' if len(all_names) > 10 else ''}"
            )
        targets = [(n, m) for n, m in linear_layers if n in target_set]
    else:
        targets = linear_layers

    calib_batches = list(calib_iter)
    if not calib_batches:
        raise ValueError("calib_iter is empty; need at least 1 batch for Hessian accumulation.")

    # Per-layer input capture: register hooks on target modules to capture their inputs.
    # Each target layer's input is what GPTQ needs for Hessian accumulation.
    captured: dict[str, list[torch.Tensor]] = {n: [] for n, _ in targets}
    hooks = []

    def make_hook(name: str):
        def hook(_module, inputs, _output):
            captured[name].append(inputs[0].detach().clone())

        return hook

    for n, m in targets:
        hooks.append(m.register_forward_hook(make_hook(n)))

    # Try to capture per-layer inputs via model forward pass.
    # If forward fails (shape mismatch etc), fall back to direct layer calls.
    model.eval()
    with torch.no_grad():
        param_device = next(model.parameters()).device
        try:
            # Feed EVERY calibration batch through the model so each hook
            # captures one input per batch and each layer's Hessian/activation
            # stats accumulate over the full calibration set (matching the
            # ``quantize_model_with_collector`` contract of "up to n_samples
            # batches").  Previously only calib_batches[0] was forwarded,
            # silently dropping all later batches from calibration while the
            # direct-layer-call fallback below used every batch.
            for batch in calib_batches:
                _ = model(batch.to(param_device))
        except (RuntimeError, ValueError, TypeError) as e:
            logger.debug(f"Model forward failed during calibration: {e}; falling back to direct layer calls.")

    # If hooks captured nothing, fall back to calling each target layer directly
    # (circumventing the model graph). A capture FAILURE mid-loop must be
    # handled uniformly: hooks fire in graph order, so a forward that raises on
    # batch k leaves layers before the failure with k+1 captures and layers at/
    # after it with only k — the per-layer batch counts would DIVERGE, and a
    # layer whose first batch never made it holds ZERO captures (its Hessian
    # later crashes with "No calibration data accumulated (n_samples=0)">
    # after earlier layers were already replaced (partial in-place mutation).
    # Any inconsistency -> rebuild every layer's captures from the direct
    # calls over the full calibration set (RIL ISS-136).
    expected_captures = len(calib_batches)
    capt_sizes = {n: len(v) for n, v in captured.items()}
    any_captured = any(s > 0 for s in capt_sizes.values())
    consistent = bool(capt_sizes) and all(s == expected_captures for s in capt_sizes.values())
    if not any_captured or not consistent:
        if any_captured and not consistent:
            logger.warning(
                "Per-layer calibration captures diverged after a partial forward "
                "failure (%s); falling back to direct layer calls for ALL targets "
                "so every layer quantizes over the same calibration set.",
                capt_sizes,
            )
        elif not any_captured:
            logger.warning(
                "Model forward failed on EVERY calibration batch (see the "
                "DEBUG log above for the first error); no per-layer inputs "
                "were captured. Falling back to feeding the raw calibration "
                "batches directly as each target layer's inputs — this is only "
                "valid when those tensors ARE the layers' activations (e.g. a "
                "bare sequence of target layers). If the model embeds or "
                "reshapes inputs first (a real decoder), the quantized weights "
                "will be garbage; fix the model forward signature or use "
                "quantize_model_with_collector."
            )
        for h in hooks:
            h.remove()
        for n, _m in targets:
            captured[n] = [batch.detach().clone() for batch in calib_batches]

    for h in hooks:
        h.remove()

    # Resolve per-layer effective configs (orthogonal to target_modules filter).
    # available_layer_names is the post-filter set so policy targets are
    # strictly validated against layers that will actually be quantized
    # (strict mode: fails fast if a policy references a filtered-out layer).
    available_layer_names = {n for n, _ in targets}
    effective_configs = resolve_layer_policies(
        config.layer_policies,
        available_layer_names,
        config,
    )

    for name, layer in targets:
        effective_config = effective_configs.get(name, config)
        new_layer = _quantize_linear_with_gptq(layer, captured[name], effective_config)
        if layer.bias is not None:
            with torch.no_grad():
                new_layer.bias.copy_(layer.bias.data)
        # Adopt the replaced layer's dtype (RIL ISS-191): selective
        # quantization over an already fp16/bf16 base must yield an fp16
        # quant layer, otherwise the forward emits fp32 into the remaining
        # half-precision linears and crashes.
        new_layer = new_layer.to(layer.weight.dtype)
        _replace_module(model, name, new_layer)
        logger.info(
            f"Quantized layer {name}: {layer.weight.shape} → "
            f"{effective_config.bits}-bit, group_size={effective_config.group_size}"
        )

    return model

quantize_model_with_collector

quantize_model_with_collector(model, collector, n_samples, config=None, target_modules=None, device=None)

Quantize a model using an existing CalibrationDataCollector.

Trainer-loop entry point: reuse the same calibration batches already collected during training (e.g. for activation stats). Materializes up to n_samples batches, then funnels into quantize_model_gptq.

参数：

名称	类型	描述	默认
model	Module	nn.Module to quantize.	必需
collector	CalibrationDataCollector | Iterable[Tensor]	CalibrationDataCollector (or any iterable yielding Tensor batches). Up to n_samples batches are consumed.	必需
n_samples	int	Maximum number of batches to use for calibration.	必需
config	GPTQConfig | None	GPTQConfig (default: 4-bit, group_size=128, symmetric).	None
target_modules	Iterable[str] | None	Optional layer-name filter forwarded to quantize_model_gptq.	None
device	device | str | None	Target device forwarded to quantize_model_gptq.	None

返回：

类型	描述
Module	The quantized model (same instance as model, with nn.Linear replaced).

引发：

类型	描述
ValueError	Forwarded from quantize_model_gptq (no nn.Linear, unmatched target_modules, etc.).

源代码位于： src/llm/quantization/gptq.py

def quantize_model_with_collector(
    model: nn.Module,
    collector: CalibrationDataCollector | Iterable[torch.Tensor],
    n_samples: int,
    config: GPTQConfig | None = None,
    target_modules: Iterable[str] | None = None,
    device: torch.device | str | None = None,
) -> nn.Module:
    """Quantize a model using an existing CalibrationDataCollector.

    Trainer-loop entry point: reuse the same calibration batches already
    collected during training (e.g. for activation stats). Materializes
    up to `n_samples` batches, then funnels into `quantize_model_gptq`.

    Args:
        model: nn.Module to quantize.
        collector: CalibrationDataCollector (or any iterable yielding Tensor
            batches). Up to `n_samples` batches are consumed.
        n_samples: Maximum number of batches to use for calibration.
        config: GPTQConfig (default: 4-bit, group_size=128, symmetric).
        target_modules: Optional layer-name filter forwarded to
            `quantize_model_gptq`.
        device: Target device forwarded to `quantize_model_gptq`.

    Returns:
        The quantized model (same instance as `model`, with nn.Linear replaced).

    Raises:
        ValueError: Forwarded from `quantize_model_gptq` (no nn.Linear,
            unmatched target_modules, etc.).
    """
    if n_samples <= 0:
        raise ValueError(f"n_samples must be positive, got {n_samples}.")

    # Materialize up to n_samples batches from the collector. We stop early
    # so collectors backed by expensive iterators (e.g. dataset streams)
    # don't pull more data than needed.
    if not isinstance(collector, Iterable):
        raise TypeError(
            "collector must be an iterable of calibration batches; "
            "CalibrationDataCollector stores activation statistics, not batches"
        )
    batches: list[torch.Tensor] = []
    for i, batch in enumerate(collector):
        if not isinstance(batch, torch.Tensor):
            raise TypeError(f"calibration batches must be tensors, got {type(batch).__name__}")
        batches.append(batch)
        if i + 1 >= n_samples:
            break

    return quantize_model_gptq(
        model,
        calib_iter=iter(batches),
        config=config,
        target_modules=target_modules,
        device=device,
    )

Calibration

calibration

Calibration for Quantization.

Collects activation statistics for quantization scale computation.

ActivationStats dataclass

Statistics for a single layer's activations.

源代码位于： src/llm/quantization/calibration.py

@dataclass
class ActivationStats:
    """Statistics for a single layer's activations."""

    name: str
    min_val: float = float("inf")
    max_val: float = float("-inf")
    abs_max: float = 0.0
    mean: float = 0.0
    std: float = 0.0
    num_samples: int = 0

    def update(self, tensor: torch.Tensor) -> None:
        """Update statistics with new tensor."""
        tensor = tensor.detach().float()

        batch_min = tensor.min().item()
        batch_max = tensor.max().item()
        batch_abs_max = tensor.abs().max().item()
        batch_mean = tensor.mean().item()
        batch_size = tensor.numel()

        # Running statistics
        total_samples = self.num_samples + batch_size

        # Update min/max
        self.min_val = min(self.min_val, batch_min)
        self.max_val = max(self.max_val, batch_max)
        self.abs_max = max(self.abs_max, batch_abs_max)

        # Welford's online algorithm for mean and variance
        old_mean = self.mean
        self.mean = old_mean + (batch_mean - old_mean) * batch_size / total_samples

        self.num_samples = total_samples

    def compute_scale(self, bits: int = 8, symmetric: bool = True) -> float:
        """Compute quantization scale."""
        qmax = 2 ** (bits - 1) - 1 if symmetric else 2**bits - 1

        if symmetric:
            scale = self.abs_max / qmax if self.abs_max > 0 else 1.0
        else:
            scale = (self.max_val - self.min_val) / qmax if self.max_val > self.min_val else 1.0

        return max(scale, 1e-8)

update

update(tensor)

Update statistics with new tensor.

源代码位于： src/llm/quantization/calibration.py

def update(self, tensor: torch.Tensor) -> None:
    """Update statistics with new tensor."""
    tensor = tensor.detach().float()

    batch_min = tensor.min().item()
    batch_max = tensor.max().item()
    batch_abs_max = tensor.abs().max().item()
    batch_mean = tensor.mean().item()
    batch_size = tensor.numel()

    # Running statistics
    total_samples = self.num_samples + batch_size

    # Update min/max
    self.min_val = min(self.min_val, batch_min)
    self.max_val = max(self.max_val, batch_max)
    self.abs_max = max(self.abs_max, batch_abs_max)

    # Welford's online algorithm for mean and variance
    old_mean = self.mean
    self.mean = old_mean + (batch_mean - old_mean) * batch_size / total_samples

    self.num_samples = total_samples

compute_scale

compute_scale(bits=8, symmetric=True)

Compute quantization scale.

源代码位于： src/llm/quantization/calibration.py

def compute_scale(self, bits: int = 8, symmetric: bool = True) -> float:
    """Compute quantization scale."""
    qmax = 2 ** (bits - 1) - 1 if symmetric else 2**bits - 1

    if symmetric:
        scale = self.abs_max / qmax if self.abs_max > 0 else 1.0
    else:
        scale = (self.max_val - self.min_val) / qmax if self.max_val > self.min_val else 1.0

    return max(scale, 1e-8)

CalibrationDataCollector

Collects activation statistics for quantization calibration.

Hooks into model forward passes to record min/max/mean/std of activations at each layer.

源代码位于： src/llm/quantization/calibration.py

class CalibrationDataCollector:
    """
    Collects activation statistics for quantization calibration.

    Hooks into model forward passes to record min/max/mean/std
    of activations at each layer.
    """

    def __init__(self, model: nn.Module):
        """
        Initialize collector.

        Args:
            model: Model to collect statistics from.
        """
        self.model = model
        self.stats: dict[str, ActivationStats] = {}
        self.hooks: list[Any] = []

    def register_hooks(self, layer_types: tuple = (nn.Linear,)) -> None:
        """
        Register forward hooks on specified layer types.

        Args:
            layer_types: Tuple of layer types to hook.
        """
        for name, module in self.model.named_modules():
            if isinstance(module, layer_types):
                self.stats[name] = ActivationStats(name=name)
                hook = module.register_forward_hook(self._make_hook(name))
                self.hooks.append(hook)

        logger.info(f"Registered {len(self.hooks)} calibration hooks")

    def _make_hook(self, name: str):
        """Create a forward hook for a named layer."""

        def hook(_module, _input, output):
            if isinstance(output, tuple):
                output = output[0]
            self.stats[name].update(output)

        return hook

    def collect(
        self,
        dataloader: DataLoader,
        num_batches: int | None = None,
        device: str | torch.device = "cuda",
    ) -> dict[str, ActivationStats]:
        """
        Collect activation statistics from calibration data.

        Args:
            dataloader: Calibration data loader.
            num_batches: Maximum number of batches to process.
            device: Device to run on.

        Returns:
            Dictionary of layer name to activation stats.
        """
        self.model.eval()
        self.model.to(device)

        with torch.no_grad():
            for i, batch in enumerate(dataloader):
                if num_batches and i >= num_batches:
                    break

                # Handle different batch formats
                if isinstance(batch, dict):
                    input_ids = batch.get("input_ids", batch.get("inputs"))
                elif isinstance(batch, (tuple, list)):
                    input_ids = batch[0]
                else:
                    input_ids = batch

                input_ids = input_ids.to(device)
                self.model(input_ids)

                if (i + 1) % 10 == 0:
                    logger.info(f"Calibrated {i + 1} batches")

        logger.info(f"Calibration complete: {sum(s.num_samples for s in self.stats.values())} samples")
        return self.stats

    def remove_hooks(self) -> None:
        """Remove all registered hooks."""
        for hook in self.hooks:
            hook.remove()
        self.hooks.clear()

    def get_scales(self, bits: int = 8, symmetric: bool = True) -> dict[str, float]:
        """
        Compute quantization scales for all layers.

        Args:
            bits: Quantization bit width.
            symmetric: Whether to use symmetric quantization.

        Returns:
            Dictionary of layer name to scale.
        """
        return {name: stats.compute_scale(bits, symmetric) for name, stats in self.stats.items()}

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.remove_hooks()

register_hooks

register_hooks(layer_types=(nn.Linear,))

Register forward hooks on specified layer types.

参数：

名称	类型	描述	默认
layer_types	tuple	Tuple of layer types to hook.	(Linear,)

源代码位于： src/llm/quantization/calibration.py

def register_hooks(self, layer_types: tuple = (nn.Linear,)) -> None:
    """
    Register forward hooks on specified layer types.

    Args:
        layer_types: Tuple of layer types to hook.
    """
    for name, module in self.model.named_modules():
        if isinstance(module, layer_types):
            self.stats[name] = ActivationStats(name=name)
            hook = module.register_forward_hook(self._make_hook(name))
            self.hooks.append(hook)

    logger.info(f"Registered {len(self.hooks)} calibration hooks")

collect

collect(dataloader, num_batches=None, device='cuda')

Collect activation statistics from calibration data.

参数：

名称	类型	描述	默认
dataloader	DataLoader	Calibration data loader.	必需
num_batches	int | None	Maximum number of batches to process.	None
device	str | device	Device to run on.	'cuda'

返回：

类型	描述
dict[str, ActivationStats]	Dictionary of layer name to activation stats.

源代码位于： src/llm/quantization/calibration.py

def collect(
    self,
    dataloader: DataLoader,
    num_batches: int | None = None,
    device: str | torch.device = "cuda",
) -> dict[str, ActivationStats]:
    """
    Collect activation statistics from calibration data.

    Args:
        dataloader: Calibration data loader.
        num_batches: Maximum number of batches to process.
        device: Device to run on.

    Returns:
        Dictionary of layer name to activation stats.
    """
    self.model.eval()
    self.model.to(device)

    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            if num_batches and i >= num_batches:
                break

            # Handle different batch formats
            if isinstance(batch, dict):
                input_ids = batch.get("input_ids", batch.get("inputs"))
            elif isinstance(batch, (tuple, list)):
                input_ids = batch[0]
            else:
                input_ids = batch

            input_ids = input_ids.to(device)
            self.model(input_ids)

            if (i + 1) % 10 == 0:
                logger.info(f"Calibrated {i + 1} batches")

    logger.info(f"Calibration complete: {sum(s.num_samples for s in self.stats.values())} samples")
    return self.stats

remove_hooks

remove_hooks()

Remove all registered hooks.

源代码位于： src/llm/quantization/calibration.py

def remove_hooks(self) -> None:
    """Remove all registered hooks."""
    for hook in self.hooks:
        hook.remove()
    self.hooks.clear()

get_scales

get_scales(bits=8, symmetric=True)

Compute quantization scales for all layers.

参数：

名称	类型	描述	默认
bits	int	Quantization bit width.	8
symmetric	bool	Whether to use symmetric quantization.	True

返回：

类型	描述
dict[str, float]	Dictionary of layer name to scale.

源代码位于： src/llm/quantization/calibration.py

def get_scales(self, bits: int = 8, symmetric: bool = True) -> dict[str, float]:
    """
    Compute quantization scales for all layers.

    Args:
        bits: Quantization bit width.
        symmetric: Whether to use symmetric quantization.

    Returns:
        Dictionary of layer name to scale.
    """
    return {name: stats.compute_scale(bits, symmetric) for name, stats in self.stats.items()}

Policies

_policy

Per-layer quantization policy — algorithm-agnostic.

LayerQuantPolicy binds a set of target layer names to a bundle of override fields (bits / group_size / sym / act_order). All override fields are optional; None means "inherit from the algorithm's base config".

This module is intentionally algorithm-agnostic: the four override fields are the public subset shared by all PTQ-style quantization algorithms (GPTQ today, AWQ / SmoothQuant / QAT in future slices). The resolve_layer_policies helper is generic over the base config dataclass, so future algorithms reuse it without modification.

See ADR-008 and docs/superpowers/specs/2026-07-22-mixed-precision-quantization-design.md.

LayerQuantPolicy dataclass

Atomic per-layer quantization override policy. Algorithm-agnostic.

A LayerQuantPolicy binds a set of target layer names to a bundle of override fields. Fields set to None mean "inherit from the algorithm's base config". Multiple LayerQuantPolicy in a config are additive; each target module must appear in at most one policy (overlap raises ValueError at resolve time).

Field semantics are universal across PTQ-style algorithms

bits: 4 or 8 — quantization bit-width group_size: -1 (per-channel) or positive int — quantization group size sym: True (symmetric) or False (asymmetric) act_order: True (sort columns by diag(H) descending) or False

The last three fields are the FP8 knobs (RIL TASK-203). FP8's knobs are NOT expressible in the shared fields above (an FP8 layer has no bits / group_size / symmetry in the int-quant sense), so the policy model is extended instead of abusing bits: weight_dtype: "e4m3" (E4M3FN, default) or "e5m2" — the FP8 format per_channel: True (per-output-row weight scale) or False (per-tensor) activation: "static" (calibration-captured per-layer scale) or "dynamic" (per-forward absmax, no calibration)

resolve_layer_policies applies ONLY the override fields the target algorithm's base config actually has; an FP8 field on an int algorithm (or an int field on FP8) is rejected loudly rather than silently dropped.

属性：

名称	类型	描述
target_modules	tuple[str, ...]	Tuple of fully-qualified layer names (dotted notation, matching the target_modules arg style of quantize_model_gptq).
bits	int | None	Override bit-width (None = inherit).
group_size	int | None	Override group size (None = inherit).
sym	bool | None	Override symmetry (None = inherit).
act_order	bool | None	Override act-order (None = inherit).
weight_dtype	str | None	Override FP8 format (None = inherit).
per_channel	bool | None	Override FP8 weight-scale granularity (None = inherit).
activation	str | None	Override FP8 activation mode (None = inherit).

源代码位于： src/llm/quantization/_policy.py

@dataclass(frozen=True)
class LayerQuantPolicy:
    """Atomic per-layer quantization override policy. Algorithm-agnostic.

    A LayerQuantPolicy binds a set of target layer names to a bundle of
    override fields. Fields set to None mean "inherit from the algorithm's
    base config". Multiple LayerQuantPolicy in a config are additive; each
    target module must appear in at most one policy (overlap raises
    ValueError at resolve time).

    Field semantics are universal across PTQ-style algorithms:
        bits:       4 or 8 — quantization bit-width
        group_size: -1 (per-channel) or positive int — quantization group size
        sym:        True (symmetric) or False (asymmetric)
        act_order:  True (sort columns by diag(H) descending) or False

    The last three fields are the FP8 knobs (RIL TASK-203). FP8's knobs are
    NOT expressible in the shared fields above (an FP8 layer has no bits /
    group_size / symmetry in the int-quant sense), so the policy model is
    extended instead of abusing ``bits``:
        weight_dtype: "e4m3" (E4M3FN, default) or "e5m2" — the FP8 format
        per_channel: True (per-output-row weight scale) or False (per-tensor)
        activation:  "static" (calibration-captured per-layer scale) or
                     "dynamic" (per-forward absmax, no calibration)

    ``resolve_layer_policies`` applies ONLY the override fields the target
    algorithm's base config actually has; an FP8 field on an int algorithm
    (or an int field on FP8) is rejected loudly rather than silently dropped.

    Attributes:
        target_modules: Tuple of fully-qualified layer names (dotted notation,
            matching the `target_modules` arg style of quantize_model_gptq).
        bits: Override bit-width (None = inherit).
        group_size: Override group size (None = inherit).
        sym: Override symmetry (None = inherit).
        act_order: Override act-order (None = inherit).
        weight_dtype: Override FP8 format (None = inherit).
        per_channel: Override FP8 weight-scale granularity (None = inherit).
        activation: Override FP8 activation mode (None = inherit).
    """

    target_modules: tuple[str, ...]
    bits: int | None = None
    group_size: int | None = None
    sym: bool | None = None
    act_order: bool | None = None
    weight_dtype: str | None = None
    per_channel: bool | None = None
    activation: str | None = None

    def __post_init__(self):
        # target_modules: non-empty, no duplicates within one policy
        if not self.target_modules:
            raise ValueError("LayerQuantPolicy.target_modules cannot be empty; specify at least one layer name.")
        if len(set(self.target_modules)) != len(self.target_modules):
            duplicates = sorted({n for n in self.target_modules if list(self.target_modules).count(n) > 1})
            raise ValueError(f"LayerQuantPolicy.target_modules has duplicates within a single policy: {duplicates}.")
        # bits: None or {4, 8}
        if self.bits is not None and self.bits not in (4, 8):
            raise ValueError(f"LayerQuantPolicy.bits must be 4, 8, or None (inherit); got {self.bits}.")
        # group_size: None, -1, or positive int
        if self.group_size is not None:
            if not isinstance(self.group_size, int) or isinstance(self.group_size, bool):
                raise ValueError(
                    f"LayerQuantPolicy.group_size must be int or None; got {type(self.group_size).__name__}."
                )
            if self.group_size != -1 and self.group_size <= 0:
                raise ValueError(
                    f"LayerQuantPolicy.group_size must be -1 (per-channel) or positive; got {self.group_size}."
                )
        # sym / act_order: None or bool (dataclass rejects other types at
        # construction, but we keep the explicit check for symmetric error
        # messages with group_size/bits paths)
        if self.sym is not None and not isinstance(self.sym, bool):
            raise ValueError(f"LayerQuantPolicy.sym must be bool or None; got {type(self.sym).__name__}.")
        if self.act_order is not None and not isinstance(self.act_order, bool):
            raise ValueError(f"LayerQuantPolicy.act_order must be bool or None; got {type(self.act_order).__name__}.")
        # FP8-specific knobs.
        if self.weight_dtype is not None and self.weight_dtype not in ("e4m3", "e5m2"):
            raise ValueError(
                f"LayerQuantPolicy.weight_dtype must be 'e4m3', 'e5m2', or None (inherit); got {self.weight_dtype!r}."
            )
        if self.per_channel is not None and not isinstance(self.per_channel, bool):
            raise ValueError(
                f"LayerQuantPolicy.per_channel must be bool or None; got {type(self.per_channel).__name__}."
            )
        if self.activation is not None and self.activation not in ("static", "dynamic"):
            raise ValueError(
                f"LayerQuantPolicy.activation must be 'static', 'dynamic', or None (inherit); got {self.activation!r}."
            )

resolve_layer_policies

resolve_layer_policies(policies, available_names, base_config)

Build layer-name -> effective config map from policies.

Generic over the base config type (T). Works for GPTQConfig and AWQ / SmoothQuantConfig today, and for Fp8Config since TASK-203 extended LayerQuantPolicy with the FP8 knobs. Each override field is applied only when the base config dataclass actually has that field — an FP8 override on an int algorithm (or an int override on FP8) is rejected with a clear error instead of crashing deep in dataclasses.replace or being silently dropped.

参数：

名称	类型	描述	默认
policies	tuple[LayerQuantPolicy, ...]	Tuple of LayerQuantPolicy to resolve. Empty tuple is a no-op.	必需
available_names	set[str]	Set of layer names that are actually eligible to be quantized (typically the post-target_modules filter set).	必需
base_config	T	The base algorithm config to inherit from.	必需

返回：

类型	描述
dict[str, T]	Dict mapping each policy-targeted layer name to its effective config
dict[str, T]	(= base_config with non-None policy fields applied via
dict[str, T]	dataclasses.replace). Empty dict if no policies.

引发：

类型	描述
ValueError	If any policy targets a name not in available_names, if the same layer name appears in multiple policies, or if a non-None policy field is not a field of base_config.

源代码位于： src/llm/quantization/_policy.py

def resolve_layer_policies[T](
    policies: tuple[LayerQuantPolicy, ...],
    available_names: set[str],
    base_config: T,
) -> dict[str, T]:
    """Build layer-name -> effective config map from policies.

    Generic over the base config type (T). Works for GPTQConfig and AWQ /
    SmoothQuantConfig today, and for Fp8Config since TASK-203 extended
    LayerQuantPolicy with the FP8 knobs. Each override field is applied only
    when the base config dataclass actually has that field — an FP8 override
    on an int algorithm (or an int override on FP8) is rejected with a clear
    error instead of crashing deep in ``dataclasses.replace`` or being
    silently dropped.

    Args:
        policies: Tuple of LayerQuantPolicy to resolve. Empty tuple is a no-op.
        available_names: Set of layer names that are actually eligible to be
            quantized (typically the post-`target_modules` filter set).
        base_config: The base algorithm config to inherit from.

    Returns:
        Dict mapping each policy-targeted layer name to its effective config
        (= base_config with non-None policy fields applied via
        `dataclasses.replace`). Empty dict if no policies.

    Raises:
        ValueError: If any policy targets a name not in `available_names`,
            if the same layer name appears in multiple policies, or if a
            non-None policy field is not a field of `base_config`.
    """
    if not policies:
        return {}

    # Phase 1: validate every policy's targets exist in available_names,
    # AND collect into a single map (later writes win on intra-iteration
    # duplicates; cross-policy overlap is detected in Phase 2).
    name_to_policy: dict[str, LayerQuantPolicy] = {}
    for i, policy in enumerate(policies):
        unmatched = set(policy.target_modules) - available_names
        if unmatched:
            sample_available = sorted(available_names)[:10]
            more = "..." if len(available_names) > 10 else ""
            raise ValueError(
                f"LayerQuantPolicy[{i}].target_modules {sorted(unmatched)} "
                f"not found in available layers. Available: "
                f"{sample_available}{more}"
            )
        for name in policy.target_modules:
            name_to_policy[name] = policy

    # Phase 2: detect cross-policy overlaps (fail-fast).
    target_counts: dict[str, int] = {}
    for policy in policies:
        for name in policy.target_modules:
            target_counts[name] = target_counts.get(name, 0) + 1
    duplicates = sorted(n for n, c in target_counts.items() if c > 1)
    if duplicates:
        raise ValueError(
            f"LayerQuantPolicy.target_modules overlap detected across "
            f"policies: {duplicates}. Each layer name must appear in at "
            f"most one policy."
        )

    # Phase 3: build effective configs (base + non-None overrides).
    base_fields: set[str] = set(getattr(base_config, "__dataclass_fields__", {}))
    effective_map: dict[str, T] = {}
    for name, policy in name_to_policy.items():
        overrides: dict[str, object] = {}
        # Every optional override, in policy-declaration order. Only fields
        # the base config actually HAS are applied; a non-None field the base
        # lacks is a config/model mismatch and must fail loudly, never be a
        # silently-ignored option (TASK-203).
        for key, value in (
            ("bits", policy.bits),
            ("group_size", policy.group_size),
            ("sym", policy.sym),
            ("act_order", policy.act_order),
            ("weight_dtype", policy.weight_dtype),
            ("per_channel", policy.per_channel),
            ("activation", policy.activation),
        ):
            if value is None:
                continue
            if key not in base_fields:
                raise ValueError(
                    f"LayerQuantPolicy for {sorted(policy.target_modules)} sets "
                    f"{key}={value!r}, but {type(base_config).__name__} has no {key} field. "
                    f"Available override fields: {sorted(base_fields)}"
                )
            overrides[key] = value
        # Strip recursion-vector field if base config has one (e.g.
        # GPTQConfig.layer_policies). Effective configs must NOT carry the
        # policies that produced them, or infinite recursion ensues.
        # We explicitly set it to () (not just remove from overrides) so
        # `dataclasses.replace` resets it instead of preserving the base's
        # value.
        if hasattr(base_config, "layer_policies"):
            overrides["layer_policies"] = ()
        # `base_config` is T (any dataclass by convention); the `replace()` signature
        # bounds its first arg to `DataclassInstance` (typeshed-only Protocol not
        # available at runtime in Python 3.14). Suppress: the helper is documented
        # as requiring a dataclass with the four override fields.
        effective_map[name] = replace(base_config, **overrides)  # ty: ignore[invalid-argument-type]

    return effective_map

PTQ

ptq

Post-Training Quantization (PTQ).

Provides utilities for quantizing models after training. Supports symmetric (scale-only) and asymmetric (scale + zero-point) 8-bit weight quantization, per-tensor or per-channel. The asymmetric path stores q - 128 in the int8 buffer and folds the offset into weight_zero_point so dequantization stays (q - zp) * scale; it is exact on the grid and beats symmetric on skewed (all-positive / all-negative) weight distributions.

QuantConfig dataclass

Configuration for quantization.

源代码位于： src/llm/quantization/ptq.py

@dataclass
class QuantConfig:
    """Configuration for quantization."""

    bits: int = 8
    symmetric: bool = True  # False = asymmetric (scale + zero-point) 8-bit
    per_channel: bool = False
    dynamic: bool = False  # Dynamic vs static quantization

    def __post_init__(self):
        if self.bits not in (4, 8):
            raise ValueError(f"Unsupported bit width: {self.bits}. Use 4 or 8.")

QuantizedLinear

Bases: Module

Quantized Linear layer with INT8/INT4 weights.

Stores quantized weights and scales, dequantizes during forward.

源代码位于： src/llm/quantization/ptq.py

class QuantizedLinear(nn.Module):
    """
    Quantized Linear layer with INT8/INT4 weights.

    Stores quantized weights and scales, dequantizes during forward.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        config: QuantConfig | None = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.config = config or QuantConfig()

        # Type annotations for registered buffers (prevents ty from inferring
        # the broader `Tensor | Module` union that register_buffer creates).
        self.weight_quantized: torch.Tensor
        self.weight_scale: torch.Tensor

        # Quantized weights (stored as int8)
        self.register_buffer(
            "weight_quantized",
            torch.zeros(out_features, in_features, dtype=torch.int8),
        )

        # Scales for dequantization
        if self.config.per_channel:
            self.register_buffer("weight_scale", torch.ones(out_features))
        else:
            self.register_buffer("weight_scale", torch.ones(1))

        # Zero point for asymmetric quantization
        self.weight_zero_point: torch.Tensor | None
        if not self.config.symmetric:
            self.register_buffer("weight_zero_point", torch.zeros_like(self.weight_scale))
        else:
            self.weight_zero_point = None

        # Bias remains in fp32
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass with dequantized weights.

        Dequantizes to fp32, computes in fp32 for accuracy, and returns in
        the layer's effective dtype (native ``nn.Linear`` semantics).  This
        keeps the layer a faithful drop-in even after the model is cast to
        fp16/bf16 — the serving engine's ``model.to(device, dtype=fp16)`` or
        selective quantization over a half base converts ``bias`` to half,
        and passing it straight into ``F.linear`` against fp32 weights
        crashed with a dtype mismatch (RIL TASK-196 / ISS-236, quant
        deep-dive F1; same fix as the GPTQ/AWQ/SmoothQuant layers in
        ISS-191).
        """
        weight = self._dequantize_weight()
        dtype = self.bias.dtype if self.bias is not None else x.dtype
        out = nn.functional.linear(
            x.to(torch.float32),
            weight,
            self.bias.to(torch.float32) if self.bias is not None else self.bias,
        )
        return out.to(dtype)

    def _dequantize_weight(self) -> torch.Tensor:
        """Dequantize stored weights."""
        weight = self.weight_quantized.float()

        if self.weight_zero_point is not None:
            weight = weight - self.weight_zero_point.view(-1, 1)

        weight = weight * self.weight_scale.view(-1, 1)

        return weight

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear,
        config: QuantConfig | None = None,
        scale: float | torch.Tensor | None = None,
    ) -> QuantizedLinear:
        """
        Create QuantizedLinear from a regular Linear layer.

        Args:
            linear: Source Linear layer.
            config: Quantization configuration.
            scale: Pre-computed scale (optional).

        Returns:
            Quantized layer.
        """
        config = config or QuantConfig()
        if config.bits != 8:
            # The simple-PTQ layer stores weights in an int8 buffer. A
            # ``bits=4`` config quantized to a 4-bit *range* but persisted
            # them as full int8 values — no packing, no memory saving, while
            # still advertising a 4-bit model (RIL ISS-197). Real 4-bit
            # packing lives in the GPTQ/AWQ layers; fail fast instead of
            # silently storing 8-bit weights behind a 4-bit claim.
            raise NotImplementedError(
                f"Simple-PTQ only supports bits=8 (got bits={config.bits}). "
                "4-bit simple-PTQ would store the 4-bit grid as int8 with no "
                "memory saving; use the GPTQ or AWQ path "
                "(llm.quantization.gptq / llm.quantization.awq) for packed 4-bit."
            )
        quant_linear = cls(
            in_features=linear.in_features,
            out_features=linear.out_features,
            bias=linear.bias is not None,
            config=config,
        )

        # Quantize weights
        weight = linear.weight.data

        if not config.symmetric:
            # Asymmetric 8-bit weight quantization (scale + zero-point).
            # The 8-bit unsigned grid is [0, 255]; we store ``q - 128`` in the
            # int8 buffer and fold the offset into ``weight_zero_point`` so
            # dequantization still reads ``(q - zp) * scale`` (QuantizedLinear
            # already subtracts the zero point before scaling).
            if quant_linear.weight_zero_point is None:
                raise ValueError("asymmetric layer must have a zero-point buffer")
            qmax = (1 << config.bits) - 1  # 255
            if config.per_channel:
                wmin = weight.min(dim=1)[0]
                wmax = weight.max(dim=1)[0]
                scale_t = ((wmax - wmin) / qmax).clamp(min=1e-8)
                # ``zp = round(-min/scale)`` may be negative (all-positive rows)
                # or exceed qmax (all-negative rows); only ``q`` is clamped to
                # the 8-bit grid. The offset is folded into the int8 storage.
                zp = torch.round(-wmin / scale_t)
                quant_linear.weight_scale.copy_(scale_t)
                quant_linear.weight_zero_point.copy_(zp - 128.0)
                q = ((weight / scale_t.view(-1, 1)).round() + zp.view(-1, 1)).clamp(0, qmax)
            else:
                wmin = weight.min()
                wmax = weight.max()
                scale_v = max(((wmax - wmin) / qmax).item(), 1e-8)
                zp = 0 if scale_v == 1e-8 else round(-wmin.item() / scale_v)
                quant_linear.weight_scale.fill_(scale_v)
                quant_linear.weight_zero_point.fill_(float(zp - 128))
                q = ((weight / scale_v).round() + zp).clamp(0, qmax)
            weight_quantized = (q - 128.0).to(torch.int8)
        elif config.per_channel:
            # Per-channel quantization: a per-row scale vector.
            if scale is None:
                abs_max = weight.abs().max(dim=1)[0]
                qmax = 2 ** (config.bits - 1) - 1
                scale = abs_max / qmax
                scale = scale.clamp(min=1e-8)
            elif isinstance(scale, (int, float)):
                # A caller-supplied scalar scale has no per-channel meaning —
                # as a 0-dim tensor it broadcasts to every row, silently making
                # the dequantization PER-TENSOR while the layer still declares
                # per-channel (every row shares one scale, inflating the
                # quantization-error profile for no structural benefit). Fail
                # fast instead of silently mis-scaling (RIL ISS-135).
                raise ValueError(
                    "per_channel=True requires a per-row scale tensor (or None to compute "
                    f"it), got a scalar {scale!r}. Pass a scale with one value per output "
                    "channel, or set config.per_channel=False."
                )

            # scale is a per-row Tensor (computed or caller-supplied).
            scale_tensor = torch.as_tensor(scale, device=weight.device, dtype=weight.dtype)
            if scale_tensor.numel() != weight.shape[0]:
                raise ValueError(
                    f"per_channel scale must have one value per output channel "
                    f"({weight.shape[0]}), got {scale_tensor.numel()}"
                )
            weight_quantized = (weight / scale_tensor.view(-1, 1)).round().clamp(-128, 127).to(torch.int8)
            quant_linear.weight_scale.copy_(scale_tensor)
        else:
            # Per-tensor quantization
            if scale is None:
                abs_max = weight.abs().max()
                qmax = 2 ** (config.bits - 1) - 1
                scale = abs_max / qmax
                scale = max(scale.item(), 1e-8)

            weight_quantized = (weight / scale).round().clamp(-128, 127).to(torch.int8)
            quant_linear.weight_scale.fill_(scale)

        quant_linear.weight_quantized.copy_(weight_quantized)

        if linear.bias is not None:
            quant_linear.bias.data.copy_(linear.bias.data)

        return quant_linear

forward

forward(x)

Forward pass with dequantized weights.

Dequantizes to fp32, computes in fp32 for accuracy, and returns in the layer's effective dtype (native nn.Linear semantics). This keeps the layer a faithful drop-in even after the model is cast to fp16/bf16 — the serving engine's model.to(device, dtype=fp16) or selective quantization over a half base converts bias to half, and passing it straight into F.linear against fp32 weights crashed with a dtype mismatch (RIL TASK-196 / ISS-236, quant deep-dive F1; same fix as the GPTQ/AWQ/SmoothQuant layers in ISS-191).

源代码位于： src/llm/quantization/ptq.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass with dequantized weights.

    Dequantizes to fp32, computes in fp32 for accuracy, and returns in
    the layer's effective dtype (native ``nn.Linear`` semantics).  This
    keeps the layer a faithful drop-in even after the model is cast to
    fp16/bf16 — the serving engine's ``model.to(device, dtype=fp16)`` or
    selective quantization over a half base converts ``bias`` to half,
    and passing it straight into ``F.linear`` against fp32 weights
    crashed with a dtype mismatch (RIL TASK-196 / ISS-236, quant
    deep-dive F1; same fix as the GPTQ/AWQ/SmoothQuant layers in
    ISS-191).
    """
    weight = self._dequantize_weight()
    dtype = self.bias.dtype if self.bias is not None else x.dtype
    out = nn.functional.linear(
        x.to(torch.float32),
        weight,
        self.bias.to(torch.float32) if self.bias is not None else self.bias,
    )
    return out.to(dtype)

from_linear classmethod

from_linear(linear, config=None, scale=None)

Create QuantizedLinear from a regular Linear layer.

参数：

名称	类型	描述	默认
linear	Linear	Source Linear layer.	必需
config	QuantConfig | None	Quantization configuration.	None
scale	float | Tensor | None	Pre-computed scale (optional).	None

返回：

类型	描述
QuantizedLinear	Quantized layer.

源代码位于： src/llm/quantization/ptq.py

@classmethod
def from_linear(
    cls,
    linear: nn.Linear,
    config: QuantConfig | None = None,
    scale: float | torch.Tensor | None = None,
) -> QuantizedLinear:
    """
    Create QuantizedLinear from a regular Linear layer.

    Args:
        linear: Source Linear layer.
        config: Quantization configuration.
        scale: Pre-computed scale (optional).

    Returns:
        Quantized layer.
    """
    config = config or QuantConfig()
    if config.bits != 8:
        # The simple-PTQ layer stores weights in an int8 buffer. A
        # ``bits=4`` config quantized to a 4-bit *range* but persisted
        # them as full int8 values — no packing, no memory saving, while
        # still advertising a 4-bit model (RIL ISS-197). Real 4-bit
        # packing lives in the GPTQ/AWQ layers; fail fast instead of
        # silently storing 8-bit weights behind a 4-bit claim.
        raise NotImplementedError(
            f"Simple-PTQ only supports bits=8 (got bits={config.bits}). "
            "4-bit simple-PTQ would store the 4-bit grid as int8 with no "
            "memory saving; use the GPTQ or AWQ path "
            "(llm.quantization.gptq / llm.quantization.awq) for packed 4-bit."
        )
    quant_linear = cls(
        in_features=linear.in_features,
        out_features=linear.out_features,
        bias=linear.bias is not None,
        config=config,
    )

    # Quantize weights
    weight = linear.weight.data

    if not config.symmetric:
        # Asymmetric 8-bit weight quantization (scale + zero-point).
        # The 8-bit unsigned grid is [0, 255]; we store ``q - 128`` in the
        # int8 buffer and fold the offset into ``weight_zero_point`` so
        # dequantization still reads ``(q - zp) * scale`` (QuantizedLinear
        # already subtracts the zero point before scaling).
        if quant_linear.weight_zero_point is None:
            raise ValueError("asymmetric layer must have a zero-point buffer")
        qmax = (1 << config.bits) - 1  # 255
        if config.per_channel:
            wmin = weight.min(dim=1)[0]
            wmax = weight.max(dim=1)[0]
            scale_t = ((wmax - wmin) / qmax).clamp(min=1e-8)
            # ``zp = round(-min/scale)`` may be negative (all-positive rows)
            # or exceed qmax (all-negative rows); only ``q`` is clamped to
            # the 8-bit grid. The offset is folded into the int8 storage.
            zp = torch.round(-wmin / scale_t)
            quant_linear.weight_scale.copy_(scale_t)
            quant_linear.weight_zero_point.copy_(zp - 128.0)
            q = ((weight / scale_t.view(-1, 1)).round() + zp.view(-1, 1)).clamp(0, qmax)
        else:
            wmin = weight.min()
            wmax = weight.max()
            scale_v = max(((wmax - wmin) / qmax).item(), 1e-8)
            zp = 0 if scale_v == 1e-8 else round(-wmin.item() / scale_v)
            quant_linear.weight_scale.fill_(scale_v)
            quant_linear.weight_zero_point.fill_(float(zp - 128))
            q = ((weight / scale_v).round() + zp).clamp(0, qmax)
        weight_quantized = (q - 128.0).to(torch.int8)
    elif config.per_channel:
        # Per-channel quantization: a per-row scale vector.
        if scale is None:
            abs_max = weight.abs().max(dim=1)[0]
            qmax = 2 ** (config.bits - 1) - 1
            scale = abs_max / qmax
            scale = scale.clamp(min=1e-8)
        elif isinstance(scale, (int, float)):
            # A caller-supplied scalar scale has no per-channel meaning —
            # as a 0-dim tensor it broadcasts to every row, silently making
            # the dequantization PER-TENSOR while the layer still declares
            # per-channel (every row shares one scale, inflating the
            # quantization-error profile for no structural benefit). Fail
            # fast instead of silently mis-scaling (RIL ISS-135).
            raise ValueError(
                "per_channel=True requires a per-row scale tensor (or None to compute "
                f"it), got a scalar {scale!r}. Pass a scale with one value per output "
                "channel, or set config.per_channel=False."
            )

        # scale is a per-row Tensor (computed or caller-supplied).
        scale_tensor = torch.as_tensor(scale, device=weight.device, dtype=weight.dtype)
        if scale_tensor.numel() != weight.shape[0]:
            raise ValueError(
                f"per_channel scale must have one value per output channel "
                f"({weight.shape[0]}), got {scale_tensor.numel()}"
            )
        weight_quantized = (weight / scale_tensor.view(-1, 1)).round().clamp(-128, 127).to(torch.int8)
        quant_linear.weight_scale.copy_(scale_tensor)
    else:
        # Per-tensor quantization
        if scale is None:
            abs_max = weight.abs().max()
            qmax = 2 ** (config.bits - 1) - 1
            scale = abs_max / qmax
            scale = max(scale.item(), 1e-8)

        weight_quantized = (weight / scale).round().clamp(-128, 127).to(torch.int8)
        quant_linear.weight_scale.fill_(scale)

    quant_linear.weight_quantized.copy_(weight_quantized)

    if linear.bias is not None:
        quant_linear.bias.data.copy_(linear.bias.data)

    return quant_linear

quantize_linear_layer

quantize_linear_layer(layer, config=None, scale=None)

Quantize a single Linear layer.

参数：

名称	类型	描述	默认
layer	Linear	Linear layer to quantize.	必需
config	QuantConfig | None	Quantization configuration.	None
scale	float | Tensor | None	Pre-computed scale.	None

返回：

类型	描述
QuantizedLinear	Quantized layer.

源代码位于： src/llm/quantization/ptq.py

def quantize_linear_layer(
    layer: nn.Linear,
    config: QuantConfig | None = None,
    scale: float | torch.Tensor | None = None,
) -> QuantizedLinear:
    """
    Quantize a single Linear layer.

    Args:
        layer: Linear layer to quantize.
        config: Quantization configuration.
        scale: Pre-computed scale.

    Returns:
        Quantized layer.
    """
    return QuantizedLinear.from_linear(layer, config, scale)

quantize_model

quantize_model(model, config=None, scales=None, inplace=False)

Quantize all Linear layers in a model.

参数：

名称	类型	描述	默认
model	Module	Model to quantize.	必需
config	QuantConfig | None	Quantization configuration.	None
scales	dict[str, float] | None	Pre-computed scales per layer name.	None
inplace	bool	Whether to modify model in-place.	False

返回：

类型	描述
Module	Quantized model.

源代码位于： src/llm/quantization/ptq.py

def quantize_model(
    model: nn.Module,
    config: QuantConfig | None = None,
    scales: dict[str, float] | None = None,
    inplace: bool = False,
) -> nn.Module:
    """
    Quantize all Linear layers in a model.

    Args:
        model: Model to quantize.
        config: Quantization configuration.
        scales: Pre-computed scales per layer name.
        inplace: Whether to modify model in-place.

    Returns:
        Quantized model.
    """
    config = config or QuantConfig()
    scales = scales or {}

    if not inplace:
        import copy

        model = copy.deepcopy(model)

    # Track replacements
    replacements = []

    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            scale = scales.get(name)
            quant_layer = QuantizedLinear.from_linear(module, config, scale)
            replacements.append((name, quant_layer))

    # Apply replacements
    for name, quant_layer in replacements:
        _replace_module(model, name, quant_layer)

    logger.info(f"Quantized {len(replacements)} linear layers")

    return model

compute_model_size

compute_model_size(model)

Compute model size statistics.

Recognizes all quantized layer flavors in the library: :class:QuantizedLinear (simple PTQ), and the GPTQ / AWQ / SmoothQuant / FP8 layers (:class:~llm.quantization._gptq_layer.GPTQQuantizedLinear, :class:~llm.quantization._awq_layer.AWQQuantizedLinear, :class:~llm.quantization._smooth_layer.SmoothQuantLinear, :class:~llm.quantization._fp8_layer.Fp8QuantizedLinear). Those replace nn.Linear entirely, so without explicit handling a GPTQ/AWQ/Smooth/FP8-quantized model reported zero parameters and zero bytes.

total_params counts true weights: for 4-bit GPTQ/AWQ layers each packed int8 byte stores two int4 weights, so total_params is the unpacked weight count while total_bytes is the actual (packed) on-disk size. Use total_params for parameter counts and total_bytes / size_mb for footprint.

返回：

类型	描述
dict[str, Any]	Dictionary with size information.

源代码位于： src/llm/quantization/ptq.py

def compute_model_size(model: nn.Module) -> dict[str, Any]:
    """
    Compute model size statistics.

    Recognizes all quantized layer flavors in the library:
    :class:`QuantizedLinear` (simple PTQ), and the GPTQ / AWQ / SmoothQuant /
    FP8 layers (:class:`~llm.quantization._gptq_layer.GPTQQuantizedLinear`,
    :class:`~llm.quantization._awq_layer.AWQQuantizedLinear`,
    :class:`~llm.quantization._smooth_layer.SmoothQuantLinear`,
    :class:`~llm.quantization._fp8_layer.Fp8QuantizedLinear`). Those replace
    ``nn.Linear`` entirely, so without explicit handling a
    GPTQ/AWQ/Smooth/FP8-quantized model reported zero parameters and zero bytes.

    ``total_params`` counts **true weights**: for 4-bit GPTQ/AWQ layers each
    packed int8 byte stores two int4 weights, so ``total_params`` is the
    unpacked weight count while ``total_bytes`` is the actual (packed) on-disk
    size. Use ``total_params`` for parameter counts and ``total_bytes`` /
    ``size_mb`` for footprint.

    Returns:
        Dictionary with size information.
    """
    # Lazy import to keep this module import-light and avoid a circular
    # dependency (the layer modules import from llm.quantization too).
    from llm.quantization._awq_layer import AWQQuantizedLinear
    from llm.quantization._fp8_layer import Fp8QuantizedLinear
    from llm.quantization._gptq_layer import GPTQQuantizedLinear
    from llm.quantization._smooth_layer import SmoothQuantLinear

    total_params = 0
    total_bytes = 0
    quantized_layers = 0

    for module in model.modules():
        if isinstance(module, (GPTQQuantizedLinear, AWQQuantizedLinear)):
            quantized_layers += 1
            # Packed int8 storage: ``weight_packed`` holds two int4 values
            # per byte for bits=4, or one int8 value per byte for bits=8.
            # ``numel()`` counts *bytes*; the true parameter count is
            # bytes * weights-per-byte (otherwise bits=4 reports half the
            # real weight count — ISS-94). ``total_bytes`` stays the actual
            # packed storage either way.
            packed_attr = cast(torch.Tensor, module.weight_packed)
            weights_per_byte = 2 if getattr(module, "bits", 4) == 4 else 1
            scales_attr = cast(torch.Tensor, module.scales)
            total_params += packed_attr.numel() * weights_per_byte
            total_bytes += packed_attr.numel() * packed_attr.element_size()
            total_bytes += scales_attr.numel() * scales_attr.element_size()
            input_scales = cast(torch.Tensor | None, getattr(module, "input_scales", None))
            if input_scales is not None:
                total_bytes += input_scales.numel() * input_scales.element_size()
            zeros_attr = cast(torch.Tensor | None, getattr(module, "zeros", None))
            if zeros_attr is not None:
                total_bytes += zeros_attr.numel() * zeros_attr.element_size()
            if module.bias is not None:
                total_bytes += module.bias.numel() * module.bias.element_size()
        elif isinstance(module, Fp8QuantizedLinear):
            quantized_layers += 1
            # FP8 weights are real float8 storage (1 byte/weight). The
            # ``weight_scale`` is fp32 per-tensor (1 value) or per-channel.
            weight_attr = module.weight_fp8
            total_params += weight_attr.numel()
            total_bytes += weight_attr.numel() * weight_attr.element_size()  # 1 byte
            scales_attr = module.weight_scale
            total_bytes += scales_attr.numel() * scales_attr.element_size()
            act_attr = cast(torch.Tensor | None, getattr(module, "activation_scale", None))
            if act_attr is not None:
                total_bytes += act_attr.numel() * act_attr.element_size()
            if module.bias is not None:
                total_bytes += module.bias.numel() * module.bias.element_size()
        elif isinstance(module, SmoothQuantLinear):
            quantized_layers += 1
            packed_attr = cast(torch.Tensor, module.weight_packed)
            weight_scales_attr = cast(torch.Tensor, module.weight_scales)
            act_scale_attr = cast(torch.Tensor, module.act_scale)
            total_params += packed_attr.numel()
            total_bytes += packed_attr.numel() * packed_attr.element_size()
            total_bytes += weight_scales_attr.numel() * weight_scales_attr.element_size()
            total_bytes += act_scale_attr.numel() * act_scale_attr.element_size()
            if module.input_scales is not None:
                total_bytes += module.input_scales.numel() * module.input_scales.element_size()
            if module.bias is not None:
                total_bytes += module.bias.numel() * module.bias.element_size()
        elif isinstance(module, QuantizedLinear):
            quantized_layers += 1
            # INT8 weights
            total_params += module.weight_quantized.numel()
            total_bytes += module.weight_quantized.numel()  # 1 byte per int8
            # FP32 scales
            total_bytes += module.weight_scale.numel() * 4
            if module.bias is not None:
                total_bytes += module.bias.numel() * 4
        elif isinstance(module, nn.Linear):
            total_params += module.weight.numel()
            total_bytes += module.weight.numel() * module.weight.element_size()
            if module.bias is not None:
                total_bytes += module.bias.numel() * module.bias.element_size()

    return {
        "total_params": total_params,
        "total_bytes": total_bytes,
        "size_mb": total_bytes / (1024 * 1024),
        "quantized_layers": quantized_layers,
    }

GPTQ Layer

_gptq_layer

GPTQQuantizedLinear: GPTQ-quantized Linear with packed 4-bit (or 8-bit) storage.

Storage convention for bits=4: - weight_packed: int8 tensor, two int4 values per byte. Pair (w[2i], w[2i+1]) packed as (w[2i] << 4) | (w[2i+1] & 0x0F). - scales: float16 tensor, shape [out_features, in_features // group_size]. - zeros: int8 tensor (or None if sym=True), shape [out_features, in_features // group_size]. - group_size=-1: scales shape [out_features, 1] (per-channel).

GPTQQuantizedLinear

Bases: Module

GPTQ-quantized Linear with packed 4-bit (or 8-bit) weight storage.

源代码位于： src/llm/quantization/_gptq_layer.py

class GPTQQuantizedLinear(nn.Module):
    """GPTQ-quantized Linear with packed 4-bit (or 8-bit) weight storage."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool,
        weight_packed: torch.Tensor,
        scales: torch.Tensor,
        zeros: torch.Tensor | None,
        bits: int = 4,
        group_size: int = 128,
        sym: bool = True,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.bits = bits
        self.group_size = group_size
        self.sym = sym

        # Register packed weights and scales as buffers (not Parameters — no grad)
        self.register_buffer("weight_packed", weight_packed)
        self.register_buffer("scales", scales)
        if zeros is not None:
            self.register_buffer("zeros", zeros)
        else:
            self.zeros = None

        # Bias remains fp32 / Parameter (only if original layer had bias)
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)

    def _unpack_weights(self) -> torch.Tensor:
        """Unpack int8 storage to int4 (or int8) tensor of shape [out_features, in_features]."""
        weight_packed = self.weight_packed
        if not isinstance(weight_packed, torch.Tensor):
            raise RuntimeError("GPTQ packed weights were not initialized")
        if self.bits == 4:
            unpacked = _unpack_4bit(weight_packed, numel=self.out_features * self.in_features)
            return unpacked.reshape(self.out_features, self.in_features)
        else:
            # 8-bit: weight_packed stores int8 values directly
            return weight_packed.reshape(self.out_features, self.in_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass with dequantized weights.

        Dequantizes to fp32 and computes the matmul in fp32 for accuracy,
        then returns in the input's dtype — native ``nn.Linear`` semantics,
        keeping the layer a faithful drop-in replacement even after the
        model is cast to fp16/bf16 (serving default) or selectively
        quantized over a half base (RIL ISS-191). Asymmetric quantization
        (`sym=False`) is not yet implemented — raises NotImplementedError.

        Args:
            x: Input tensor of shape [..., in_features].

        Returns:
            Output tensor of shape [..., out_features].

        Raises:
            NotImplementedError: If `sym=False` was passed at construction.
        """
        if not self.sym:
            raise NotImplementedError("Asymmetric GPTQ forward is not yet implemented. Construct with sym=True.")

        scales = self.scales
        if not isinstance(scales, torch.Tensor):
            raise RuntimeError("GPTQ scales were not initialized")

        # Always dequantize from int4/int8 storage: trades compute for memory.
        # (Caching fp32 weights would double storage; deferred to future optimization.)
        w_int = self._unpack_weights()  # [out_features, in_features]
        # 4-bit storage is unsigned [0, 15] → shift to signed [-8, 7];
        # 8-bit storage is already signed int8 [-128, 127] (no shift).
        w_int_signed = w_int.to(torch.float32) - 8.0 if self.bits == 4 else w_int.to(torch.float32)

        if self.group_size == -1:
            # Per-channel: scales shape [out_features, 1] broadcasts across input dim.
            w_fp = w_int_signed * scales.to(torch.float32)
        else:
            # Per-group: scales shape [out_features, in_features // group_size].
            # Expand to [out_features, in_features] by repeating within each group.
            gs = self.group_size
            scales_expanded = scales.to(torch.float32).repeat_interleave(gs, dim=1)
            w_fp = w_int_signed * scales_expanded

        # Dequantize in fp32 for accuracy, then compute in fp32 and return
        # in the layer's effective dtype (RIL ISS-191). A post-quant
        # fp16/bf16 cast — the serving engine's
        # ``model.to(device, dtype=torch.float16)``, or selective
        # quantization over an already-half base — converts ``bias``/scales
        # to half; computing in fp32 avoids mixing dtypes inside
        # ``F.linear`` (crash) and returning the layer's dtype avoids
        # emitting fp32 into half-precision residual linears (crash). An
        # fp32 model is unchanged (fp32 in -> fp32 out).
        dtype = self.bias.dtype if self.bias is not None else x.dtype
        out = torch.nn.functional.linear(
            x.to(torch.float32),
            w_fp,
            self.bias.to(torch.float32) if self.bias is not None else self.bias,
        )
        return out.to(dtype)

forward

forward(x)

Forward pass with dequantized weights.

Dequantizes to fp32 and computes the matmul in fp32 for accuracy, then returns in the input's dtype — native nn.Linear semantics, keeping the layer a faithful drop-in replacement even after the model is cast to fp16/bf16 (serving default) or selectively quantized over a half base (RIL ISS-191). Asymmetric quantization (sym=False) is not yet implemented — raises NotImplementedError.

参数：

名称	类型	描述	默认
x	Tensor	Input tensor of shape [..., in_features].	必需

返回：

类型	描述
Tensor	Output tensor of shape [..., out_features].

引发：

类型	描述
NotImplementedError	If sym=False was passed at construction.

源代码位于： src/llm/quantization/_gptq_layer.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass with dequantized weights.

    Dequantizes to fp32 and computes the matmul in fp32 for accuracy,
    then returns in the input's dtype — native ``nn.Linear`` semantics,
    keeping the layer a faithful drop-in replacement even after the
    model is cast to fp16/bf16 (serving default) or selectively
    quantized over a half base (RIL ISS-191). Asymmetric quantization
    (`sym=False`) is not yet implemented — raises NotImplementedError.

    Args:
        x: Input tensor of shape [..., in_features].

    Returns:
        Output tensor of shape [..., out_features].

    Raises:
        NotImplementedError: If `sym=False` was passed at construction.
    """
    if not self.sym:
        raise NotImplementedError("Asymmetric GPTQ forward is not yet implemented. Construct with sym=True.")

    scales = self.scales
    if not isinstance(scales, torch.Tensor):
        raise RuntimeError("GPTQ scales were not initialized")

    # Always dequantize from int4/int8 storage: trades compute for memory.
    # (Caching fp32 weights would double storage; deferred to future optimization.)
    w_int = self._unpack_weights()  # [out_features, in_features]
    # 4-bit storage is unsigned [0, 15] → shift to signed [-8, 7];
    # 8-bit storage is already signed int8 [-128, 127] (no shift).
    w_int_signed = w_int.to(torch.float32) - 8.0 if self.bits == 4 else w_int.to(torch.float32)

    if self.group_size == -1:
        # Per-channel: scales shape [out_features, 1] broadcasts across input dim.
        w_fp = w_int_signed * scales.to(torch.float32)
    else:
        # Per-group: scales shape [out_features, in_features // group_size].
        # Expand to [out_features, in_features] by repeating within each group.
        gs = self.group_size
        scales_expanded = scales.to(torch.float32).repeat_interleave(gs, dim=1)
        w_fp = w_int_signed * scales_expanded

    # Dequantize in fp32 for accuracy, then compute in fp32 and return
    # in the layer's effective dtype (RIL ISS-191). A post-quant
    # fp16/bf16 cast — the serving engine's
    # ``model.to(device, dtype=torch.float16)``, or selective
    # quantization over an already-half base — converts ``bias``/scales
    # to half; computing in fp32 avoids mixing dtypes inside
    # ``F.linear`` (crash) and returning the layer's dtype avoids
    # emitting fp32 into half-precision residual linears (crash). An
    # fp32 model is unchanged (fp32 in -> fp32 out).
    dtype = self.bias.dtype if self.bias is not None else x.dtype
    out = torch.nn.functional.linear(
        x.to(torch.float32),
        w_fp,
        self.bias.to(torch.float32) if self.bias is not None else self.bias,
    )
    return out.to(dtype)

AWQ Layer

_awq_layer

AWQQuantizedLinear: activation-aware quantized Linear with packed storage.

Storage convention mirrors :class:GPTQQuantizedLinear (symmetric INT4/INT8 group quantization with per-group scales), plus one AWQ-specific buffer:

• input_scales: per-input-channel scaling factors s of shape [in_features] (fp16). The layer was quantized from W * s and compensates at forward time by dividing the input: y = Q(W·s)·(x/s).

Keeping the compensation in the layer (rather than folding it into the preceding layer) is exact and needs no graph analysis; cross-layer folding is a follow-up optimization (see ADR-009).

AWQQuantizedLinear

Bases: Module

Activation-aware weight-quantized Linear (symmetric group quantization).

属性：

名称	类型	描述
in_features	/ out_features	linear geometry.
bits		4 or 8.
group_size		-1 (per-channel) or positive int (per-group).
sym		must be True (asymmetric AWQ is a follow-up).
input_scales		per-input-channel AWQ scale s, shape [in_features], or None when the caller already folded the scales upstream.

源代码位于： src/llm/quantization/_awq_layer.py

class AWQQuantizedLinear(nn.Module):
    """Activation-aware weight-quantized Linear (symmetric group quantization).

    Attributes:
        in_features / out_features: linear geometry.
        bits: 4 or 8.
        group_size: -1 (per-channel) or positive int (per-group).
        sym: must be True (asymmetric AWQ is a follow-up).
        input_scales: per-input-channel AWQ scale ``s``, shape [in_features],
            or None when the caller already folded the scales upstream.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool,
        weight_packed: torch.Tensor,
        scales: torch.Tensor,
        bits: int = 4,
        group_size: int = 128,
        sym: bool = True,
        input_scales: torch.Tensor | None = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.bits = bits
        self.group_size = group_size
        self.sym = sym

        self.register_buffer("weight_packed", weight_packed)
        self.register_buffer("scales", scales)
        if input_scales is not None:
            self.register_buffer("input_scales", input_scales)
        else:
            self.input_scales = None

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)

    def _unpack_weights(self) -> torch.Tensor:
        """Unpack int8 storage to int4 (or int8) tensor [out_features, in_features]."""
        weight_packed = self.weight_packed
        if not isinstance(weight_packed, torch.Tensor):
            raise RuntimeError("AWQ packed weights were not initialized")
        if self.bits == 4:
            unpacked = _unpack_4bit(weight_packed, numel=self.out_features * self.in_features)
            return unpacked.reshape(self.out_features, self.in_features)
        return weight_packed.reshape(self.out_features, self.in_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass: ``Q(W·s)·(x/s)`` with exact AWQ scale compensation.

        Computes the matmul in fp32 for accuracy, then returns in the input's
        dtype — native ``nn.Linear`` semantics — so post-quant fp16/bf16
        casts (serving engine default, or selective quantization over a half
        base) don't mix dtypes inside ``F.linear`` or emit fp32 into
        half-precision residual linears (RIL ISS-191). An fp32 model is
        unchanged.

        Args:
            x: Input tensor of shape [..., in_features].

        Returns:
            Output tensor of shape [..., out_features].

        Raises:
            NotImplementedError: If ``sym=False`` was passed at construction.
        """
        if not self.sym:
            raise NotImplementedError("Asymmetric AWQ forward is not yet implemented. Construct with sym=True.")

        if self.input_scales is not None:
            x = x / self.input_scales.to(x.dtype)

        w_int = self._unpack_weights()
        # 4-bit storage is unsigned [0, 15] → shift to signed [-8, 7];
        # 8-bit storage is already signed int8 [-128, 127] (no shift).
        w_int_signed = w_int.to(torch.float32) - 8.0 if self.bits == 4 else w_int.to(torch.float32)

        scales = self.scales
        if not isinstance(scales, torch.Tensor):
            raise RuntimeError("AWQ scales were not initialized")

        if self.group_size == -1:
            w_fp = w_int_signed * scales.to(torch.float32)
        else:
            gs = self.group_size
            scales_expanded = scales.to(torch.float32).repeat_interleave(gs, dim=1)
            w_fp = w_int_signed * scales_expanded

        # Compute in fp32 for accuracy, return in the layer's effective
        # dtype (RIL ISS-191) — see GPTQQuantizedLinear for the same
        # reasoning; the output must follow the surrounding model's dtype
        # (fp32 model stays fp32; a post-quant fp16/bf16 cast produces
        # half output), not always be fp32.
        dtype = self.bias.dtype if self.bias is not None else x.dtype
        out = torch.nn.functional.linear(
            x.to(torch.float32),
            w_fp,
            self.bias.to(torch.float32) if self.bias is not None else self.bias,
        )
        return out.to(dtype)

forward

forward(x)

Forward pass: Q(W·s)·(x/s) with exact AWQ scale compensation.

Computes the matmul in fp32 for accuracy, then returns in the input's dtype — native nn.Linear semantics — so post-quant fp16/bf16 casts (serving engine default, or selective quantization over a half base) don't mix dtypes inside F.linear or emit fp32 into half-precision residual linears (RIL ISS-191). An fp32 model is unchanged.

参数：

名称	类型	描述	默认
x	Tensor	Input tensor of shape [..., in_features].	必需

返回：

类型	描述
Tensor	Output tensor of shape [..., out_features].

引发：

类型	描述
NotImplementedError	If sym=False was passed at construction.

源代码位于： src/llm/quantization/_awq_layer.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass: ``Q(W·s)·(x/s)`` with exact AWQ scale compensation.

    Computes the matmul in fp32 for accuracy, then returns in the input's
    dtype — native ``nn.Linear`` semantics — so post-quant fp16/bf16
    casts (serving engine default, or selective quantization over a half
    base) don't mix dtypes inside ``F.linear`` or emit fp32 into
    half-precision residual linears (RIL ISS-191). An fp32 model is
    unchanged.

    Args:
        x: Input tensor of shape [..., in_features].

    Returns:
        Output tensor of shape [..., out_features].

    Raises:
        NotImplementedError: If ``sym=False`` was passed at construction.
    """
    if not self.sym:
        raise NotImplementedError("Asymmetric AWQ forward is not yet implemented. Construct with sym=True.")

    if self.input_scales is not None:
        x = x / self.input_scales.to(x.dtype)

    w_int = self._unpack_weights()
    # 4-bit storage is unsigned [0, 15] → shift to signed [-8, 7];
    # 8-bit storage is already signed int8 [-128, 127] (no shift).
    w_int_signed = w_int.to(torch.float32) - 8.0 if self.bits == 4 else w_int.to(torch.float32)

    scales = self.scales
    if not isinstance(scales, torch.Tensor):
        raise RuntimeError("AWQ scales were not initialized")

    if self.group_size == -1:
        w_fp = w_int_signed * scales.to(torch.float32)
    else:
        gs = self.group_size
        scales_expanded = scales.to(torch.float32).repeat_interleave(gs, dim=1)
        w_fp = w_int_signed * scales_expanded

    # Compute in fp32 for accuracy, return in the layer's effective
    # dtype (RIL ISS-191) — see GPTQQuantizedLinear for the same
    # reasoning; the output must follow the surrounding model's dtype
    # (fp32 model stays fp32; a post-quant fp16/bf16 cast produces
    # half output), not always be fp32.
    dtype = self.bias.dtype if self.bias is not None else x.dtype
    out = torch.nn.functional.linear(
        x.to(torch.float32),
        w_fp,
        self.bias.to(torch.float32) if self.bias is not None else self.bias,
    )
    return out.to(dtype)

SmoothQuant Layer

_smooth_layer

SmoothQuantLinear: weight+activation INT8 Linear with activation smoothing.

Storage convention: - weight_packed: int8 weights, shape [out_features * in_features] (SmoothQuant is an INT8 method; no nibble packing). - weight_scales: per-output-row fp16 scales [out_features] — w_int8 * weight_scales dequantizes the smoothed weights. - act_scale: per-tensor fp16 activation scale (max abs / 127). - input_scales: per-input-channel smoothing factors s [in_features] (fp16). The layer was quantized from W·s and compensates at forward time by dividing the input: y = Q8(W·s)·Q8(x/s).

Keeping the smoothing compensation in the layer (rather than folding it into the preceding layer) is exact and needs no graph analysis; cross-layer folding is a follow-up optimization (see ADR-010).

SmoothQuantLinear

Bases: Module

INT8 weight+activation quantized Linear with per-channel smoothing.

属性：

名称	类型	描述
in_features	/ out_features	linear geometry.
sym		must be True (asymmetric SmoothQuant is a follow-up).
input_scales		per-input-channel smoothing scale s, shape [in_features], or None when the caller already folded the scales upstream.

源代码位于： src/llm/quantization/_smooth_layer.py

class SmoothQuantLinear(nn.Module):
    """INT8 weight+activation quantized Linear with per-channel smoothing.

    Attributes:
        in_features / out_features: linear geometry.
        sym: must be True (asymmetric SmoothQuant is a follow-up).
        input_scales: per-input-channel smoothing scale ``s``, shape
            [in_features], or None when the caller already folded the
            scales upstream.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool,
        weight_packed: torch.Tensor,
        weight_scales: torch.Tensor,
        act_scale: torch.Tensor,
        sym: bool = True,
        input_scales: torch.Tensor | None = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.sym = sym

        self.register_buffer("weight_packed", weight_packed)
        self.register_buffer("weight_scales", weight_scales)
        self.register_buffer("act_scale", act_scale)
        if input_scales is not None:
            self.register_buffer("input_scales", input_scales)
        else:
            self.input_scales = None

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)

    def _dequantize_weights(self) -> torch.Tensor:
        """Dequantize the smoothed int8 weights: [out_features, in_features] fp32."""
        weight_packed = self.weight_packed
        if not isinstance(weight_packed, torch.Tensor):
            raise RuntimeError("SmoothQuant packed weights were not initialized")
        weight_scales = self.weight_scales
        if not isinstance(weight_scales, torch.Tensor):
            raise RuntimeError("SmoothQuant weight scales were not initialized")
        w_int = weight_packed.reshape(self.out_features, self.in_features).to(torch.float32)
        return w_int * weight_scales.to(torch.float32)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass: ``Q8(W·s)·Q8(x/s)`` with INT8 fake quantization.

        Activations are quantized per-tensor (the SmoothQuant contract), then
        multiplied by the dequantized smoothed weights. The matmul runs in
        fp32 for accuracy and the result returns in the input's dtype —
        native ``nn.Linear`` semantics — so post-quant fp16/bf16 casts don't
        mix dtypes inside ``F.linear`` or emit fp32 into half-precision
        residual linears (RIL ISS-191). An fp32 model is unchanged.

        Args:
            x: Input tensor of shape [..., in_features].

        Returns:
            Output tensor of shape [..., out_features].

        Raises:
            NotImplementedError: If ``sym=False`` was passed at construction.
        """
        if not self.sym:
            raise NotImplementedError("Asymmetric SmoothQuant forward is not yet implemented. Construct with sym=True.")

        if self.input_scales is not None:
            x = x / self.input_scales.to(x.dtype)

        # Per-tensor INT8 activation fake-quantization.
        act_scale = self.act_scale
        if not isinstance(act_scale, torch.Tensor):
            raise RuntimeError("SmoothQuant activation scale was not initialized")
        act_scale = act_scale.to(x.dtype)
        x_q = torch.clamp(torch.round(x / act_scale), -128, 127) * act_scale

        w_fp = self._dequantize_weights()
        # Compute in fp32 for accuracy, return in the layer's effective
        # dtype (RIL ISS-191) — see GPTQQuantizedLinear for the same
        # reasoning; the output must follow the surrounding model's dtype
        # (fp32 model stays fp32; a post-quant fp16/bf16 cast produces
        # half output), not always be fp32.
        dtype = self.bias.dtype if self.bias is not None else x_q.dtype
        out = torch.nn.functional.linear(
            x_q.to(torch.float32),
            w_fp,
            self.bias.to(torch.float32) if self.bias is not None else self.bias,
        )
        return out.to(dtype)

forward

forward(x)

Forward pass: Q8(W·s)·Q8(x/s) with INT8 fake quantization.

Activations are quantized per-tensor (the SmoothQuant contract), then multiplied by the dequantized smoothed weights. The matmul runs in fp32 for accuracy and the result returns in the input's dtype — native nn.Linear semantics — so post-quant fp16/bf16 casts don't mix dtypes inside F.linear or emit fp32 into half-precision residual linears (RIL ISS-191). An fp32 model is unchanged.

参数：

名称	类型	描述	默认
x	Tensor	Input tensor of shape [..., in_features].	必需

返回：

类型	描述
Tensor	Output tensor of shape [..., out_features].

引发：

类型	描述
NotImplementedError	If sym=False was passed at construction.

源代码位于： src/llm/quantization/_smooth_layer.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass: ``Q8(W·s)·Q8(x/s)`` with INT8 fake quantization.

    Activations are quantized per-tensor (the SmoothQuant contract), then
    multiplied by the dequantized smoothed weights. The matmul runs in
    fp32 for accuracy and the result returns in the input's dtype —
    native ``nn.Linear`` semantics — so post-quant fp16/bf16 casts don't
    mix dtypes inside ``F.linear`` or emit fp32 into half-precision
    residual linears (RIL ISS-191). An fp32 model is unchanged.

    Args:
        x: Input tensor of shape [..., in_features].

    Returns:
        Output tensor of shape [..., out_features].

    Raises:
        NotImplementedError: If ``sym=False`` was passed at construction.
    """
    if not self.sym:
        raise NotImplementedError("Asymmetric SmoothQuant forward is not yet implemented. Construct with sym=True.")

    if self.input_scales is not None:
        x = x / self.input_scales.to(x.dtype)

    # Per-tensor INT8 activation fake-quantization.
    act_scale = self.act_scale
    if not isinstance(act_scale, torch.Tensor):
        raise RuntimeError("SmoothQuant activation scale was not initialized")
    act_scale = act_scale.to(x.dtype)
    x_q = torch.clamp(torch.round(x / act_scale), -128, 127) * act_scale

    w_fp = self._dequantize_weights()
    # Compute in fp32 for accuracy, return in the layer's effective
    # dtype (RIL ISS-191) — see GPTQQuantizedLinear for the same
    # reasoning; the output must follow the surrounding model's dtype
    # (fp32 model stays fp32; a post-quant fp16/bf16 cast produces
    # half output), not always be fp32.
    dtype = self.bias.dtype if self.bias is not None else x_q.dtype
    out = torch.nn.functional.linear(
        x_q.to(torch.float32),
        w_fp,
        self.bias.to(torch.float32) if self.bias is not None else self.bias,
    )
    return out.to(dtype)