File size: 5,309 Bytes
e9c8366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """ssm_patterns — exclude patterns for State-Space Model / linear-attention
layers (Mamba, KDA, RWKV, gated linear attention, etc.).
SSM layers contain parameters that must NOT be quantized through the standard
weight-quantization walker:
- `A_log`, `dt_bias`, `D` — nn.Parameter (not nn.Module): the walker does not
touch them automatically. Documented here for clarity.
- `ShortConvolution` / `conv1d` — custom 1D-conv wrappers (not nn.Conv1d):
not in SUPPORTED_MODULE_TYPES, walker skips them.
- Linear projections feeding the SSM core (`q_proj`, `k_proj`, `b_proj`,
`f_a_proj`, `f_b_proj`) — distorting these breaks the delta rule / recurrent
dynamics. Should be excluded for low-bit formats.
- Output projections (`v_proj`, `o_proj`, `g_proj`, `g_a_proj`, `g_b_proj`)
— less critical: quantizing them is generally safe.
Usage:
from agiws_neural_quant.ssm_patterns import get_ssm_exclude_patterns
excl = get_ssm_exclude_patterns("conservative")
quantize_model(model, format="int4", exclude_modules=excl)
Levels:
conservative — quantize only o_proj (output) inside SSM modules.
moderate — quantize o_proj + v_proj + g_* (output path), keep q/k/b/f.
aggressive — quantize all Linear inside SSM except q_proj/k_proj
(preserve SSM input, quantize output + gating).
full — quantize every Linear (no SSM exclusions; for fp8 / high-bit
where distortion is minimal).
subtree — skip whole SSM modules entirely (for Mamba/RWKV whose module
name contains 'mamba'/'rwkv'/'kda'/'delta_attn'). NOT suitable
for Kimi-K3 where the SSM layer is named `self_attn` (shared
with MLA layers) — use level-based exclude there.
All patterns use fnmatch glob syntax and are matched against the full dotted
module path (e.g. "layers.5.self_attn.q_proj"). Bare names without a leading
`*.` are intentionally omitted — they would never match the dotted path and
would be dead patterns.
"""
from __future__ import annotations
# Full exclusion: skip the entire SSM subtree (any Linear inside).
# Matched against the SSM module name itself. Use for Mamba/RWKV where the
# module has a distinctive name. For Kimi-K3 (SSM layer named `self_attn`,
# shared with MLA), use level-based exclude instead.
SSM_SUBTREE_PATTERNS: tuple[str, ...] = (
"mamba", "*_mamba", "*mamba*",
"rwkv", "*_rwkv", "*rwkv*",
"*delta_attention*", "*delta_attn*",
"*linear_attention*", "*linear_attn*",
"*kda*", "*_kda",
)
# Conservative: quantize only the final output projection inside SSM.
# All other SSM Linear (q/k/v/b/f/g inputs and gating) are excluded.
SSM_CONSERVATIVE_EXCLUDE: tuple[str, ...] = (
"*.self_attn.q_proj",
"*.self_attn.k_proj",
"*.self_attn.v_proj",
"*.self_attn.b_proj",
"*.self_attn.f_a_proj",
"*.self_attn.f_b_proj",
"*.self_attn.g_a_proj",
"*.self_attn.g_b_proj",
"*.self_attn.in_proj",
"*.self_attn.dt_proj",
)
# Moderate: quantize output path (o/v/g), keep SSM input + gating (q/k/b/f).
SSM_MODERATE_EXCLUDE: tuple[str, ...] = (
"*.self_attn.q_proj",
"*.self_attn.k_proj",
"*.self_attn.b_proj",
"*.self_attn.f_a_proj",
"*.self_attn.f_b_proj",
"*.self_attn.in_proj",
"*.self_attn.dt_proj",
)
# Aggressive: quantize all SSM Linear except q/k (preserve SSM input only).
SSM_AGGRESSIVE_EXCLUDE: tuple[str, ...] = (
"*.self_attn.q_proj",
"*.self_attn.k_proj",
"*.self_attn.in_proj",
)
_LEVELS: dict[str, tuple[str, ...] | None] = {
"full": None, # no exclusions
"aggressive": SSM_AGGRESSIVE_EXCLUDE,
"moderate": SSM_MODERATE_EXCLUDE,
"conservative": SSM_CONSERVATIVE_EXCLUDE,
"subtree": SSM_SUBTREE_PATTERNS, # skip whole SSM modules
}
def get_ssm_exclude_patterns(level: str = "conservative") -> list[str]:
"""Return glob patterns to pass to quantize_model(exclude_modules=...).
Args:
level: one of 'conservative', 'moderate', 'aggressive', 'full', 'subtree'.
- conservative (default): keep SSM core + gating, quantize o_proj only.
- moderate: quantize output path (o/v/g), keep SSM input + gating.
- aggressive: quantize all SSM Linear except q/k (SSM input).
- full: no SSM exclusions (empty list).
- subtree: skip whole SSM modules entirely (most conservative,
for Mamba/RWKV; NOT for Kimi-K3 `self_attn`).
Returns:
list of fnmatch glob patterns; empty for 'full'.
"""
if level not in _LEVELS:
raise ValueError(
f"Unknown SSM level {level!r}. Available: {sorted(_LEVELS)}"
)
pats = _LEVELS[level]
return list(pats) if pats is not None else []
def get_ssm_subtree_patterns() -> list[str]:
"""Patterns to skip whole SSM modules (do not recurse into them at all).
Use for Mamba/RWKV. For Kimi-K3 (SSM layer named `self_attn`, shared with
MLA), prefer the level-based exclude instead.
"""
return list(SSM_SUBTREE_PATTERNS)
__all__ = [
"SSM_SUBTREE_PATTERNS",
"SSM_CONSERVATIVE_EXCLUDE",
"SSM_MODERATE_EXCLUDE",
"SSM_AGGRESSIVE_EXCLUDE",
"get_ssm_exclude_patterns",
"get_ssm_subtree_patterns",
] |