File size: 6,156 Bytes
7ccb33d | 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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """In-place model conversion utilities for Sherry and RTN QAT."""
from __future__ import annotations
import fnmatch
import re
from collections.abc import Sequence
import torch.nn as nn
from .sherry_quant import SherryLinear
from .rtn_quant import RTNEmbedding, RTNLinear
DEFAULT_TARGET_PATTERNS = (
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
)
DEFAULT_EXCLUDE_PATTERNS = (
"audio",
"encoder",
"embed",
"lm_head",
"projector",
)
def _matches(path: str, leaf: str, patterns: Sequence[str]) -> bool:
"""Match exact leaves, path substrings, shell globs, or explicit regexes."""
for pattern in patterns:
if pattern.startswith("re:"):
if re.search(pattern[3:], path):
return True
elif any(character in pattern for character in "*?["):
if fnmatch.fnmatchcase(path, pattern) or fnmatch.fnmatchcase(leaf, pattern):
return True
elif pattern == leaf or pattern in path:
return True
return False
def wrap_sherry(
model: nn.Module,
target_patterns: Sequence[str] | None = None,
exclude_patterns: Sequence[str] | None = None,
*,
group_size: int = 128,
) -> list[str]:
"""Replace matching decoder ``nn.Linear`` children in place.
The default targets are Qwen3's seven attention/MLP projections and require
a ``layers.<integer>`` decoder path, while audio towers, encoders, embeddings,
projectors, and ``lm_head`` are excluded. Plain strings match path substrings
or exact leaf names; glob patterns and ``re:...`` regexes are also accepted.
Differences from AngelSlim Sherry: this avoids a version-pinned fork of the
Transformers modeling file, preserves weight/bias key paths, copies weights
into fp32 masters, and returns an auditable replacement list.
"""
use_default_targets = target_patterns is None
targets = tuple(DEFAULT_TARGET_PATTERNS if target_patterns is None else target_patterns)
excludes = tuple(DEFAULT_EXCLUDE_PATTERNS if exclude_patterns is None else exclude_patterns)
replaced: list[str] = []
# Materialize before mutation: named_modules otherwise observes replacements.
for parent_path, parent in list(model.named_modules()):
for leaf, child in list(parent.named_children()):
if not isinstance(child, nn.Linear):
continue
path = f"{parent_path}.{leaf}" if parent_path else leaf
if use_default_targets and re.search(r"(?:^|\.)layers\.\d+\.", path) is None:
continue
if excludes and _matches(path, leaf, excludes):
continue
if not _matches(path, leaf, targets):
continue
if isinstance(child, RTNLinear):
raise ValueError(
f"quantization wrapper conflict at {path!r}: already wrapped by RTN"
)
if isinstance(child, SherryLinear):
continue
converted = SherryLinear.from_linear(child, group_size=group_size)
setattr(parent, leaf, converted)
replaced.append(path)
return replaced
def wrap_rtn(
model: nn.Module,
target_patterns: Sequence[str],
exclude_patterns: Sequence[str] | None = None,
*,
bits: int = 8,
granularity: str = "per_channel",
group_size: int = 128,
) -> list[str]:
"""Replace matching linear/embedding children with fp32-master RTN layers.
Patterns have the same exact-leaf, substring, glob, and ``re:`` semantics as
:func:`wrap_sherry`. Calling this over a path already wrapped by Sherry is
also an explicit conflict.
"""
targets = tuple(target_patterns)
excludes = tuple(() if exclude_patterns is None else exclude_patterns)
replaced: list[str] = []
for parent_path, parent in list(model.named_modules()):
for leaf, child in list(parent.named_children()):
if not isinstance(child, (nn.Linear, nn.Embedding)):
continue
path = f"{parent_path}.{leaf}" if parent_path else leaf
if excludes and _matches(path, leaf, excludes):
continue
if not _matches(path, leaf, targets):
continue
if isinstance(child, SherryLinear):
raise ValueError(
f"quantization wrapper conflict at {path!r}: already wrapped by Sherry"
)
if isinstance(child, (RTNLinear, RTNEmbedding)):
continue
if isinstance(child, nn.Embedding):
converted = RTNEmbedding.from_embedding(
child,
bits=bits,
granularity=granularity,
group_size=group_size,
)
else:
converted = RTNLinear.from_linear(
child,
bits=bits,
granularity=granularity,
group_size=group_size,
)
setattr(parent, leaf, converted)
replaced.append(path)
return replaced
def freeze_non_quantized(model: nn.Module) -> list[str]:
"""Freeze parameters outside Sherry layers and return trainable key names.
AngelSlim constructs a fully quantized modeling fork and leaves global
freezing to its trainer. The wrapper approach needs an explicit helper;
parameters owned by each Sherry module (fp32 master and optional bias) remain
trainable while embeddings, norms, heads, and skipped towers are frozen.
"""
for parameter in model.parameters():
parameter.requires_grad_(False)
trainable: list[str] = []
for module_name, module in model.named_modules():
if not isinstance(module, SherryLinear):
continue
for parameter_name, parameter in module.named_parameters(recurse=False):
parameter.requires_grad_(True)
path = f"{module_name}.{parameter_name}" if module_name else parameter_name
trainable.append(path)
return trainable
|