| |
| """Схема адаптера и его загрузка — в одном месте. |
| |
| Раньше класс адаптера был скопирован в train_adapter.py и test_klein_adapter.py |
| с комментарием «схемы обязаны совпадать»: любая правка в одном файле ломала |
| второй молча, потому что веса грузятся по именам и совпадают по shape. |
| Теперь обучение, A/B-тест, экспорт в safetensors и пример с klein ходят сюда. |
| |
| Схема: MLP поточечно (каждый токен независимо) |
| in -> [norm] -> hidden -> GELU(tanh) -> ... -> out |
| norm="ln-per-layer" — каждый срез слоя ТЕ нормируется отдельно, затем склейка |
| (нормы скрытых состояний растут с глубиной в 40 раз: 13 у слоя 2, 557 у слоя 27, |
| поэтому без пер-слойной нормировки ранние слои не влияют на выход). |
| """ |
| import torch |
| import torch.nn as nn |
|
|
|
|
| class PerLayerNorm(nn.Module): |
| """Нормировка каждого среза слоя отдельно + общая нормировка склейки. |
| |
| Срезы учителя имеют разный масштаб (у Qwen3.5-9B норма растёт с глубиной в десятки раз), |
| поэтому без пер-слойной нормировки ранние слои почти не влияют на выход. |
| """ |
|
|
| def __init__(self, in_dim, n_slices): |
| super().__init__() |
| assert in_dim % n_slices == 0, f"{in_dim} не делится на {n_slices}" |
| self.n, self.d = n_slices, in_dim // n_slices |
| self.per = nn.LayerNorm(self.d, elementwise_affine=False) |
| self.all = nn.LayerNorm(in_dim) |
|
|
| def forward(self, x): |
| b, l, _ = x.shape |
| return self.all(self.per(x.view(b, l, self.n, self.d)).reshape(b, l, -1)) |
|
|
|
|
| def build_adapter(in_dim, out_dim, hidden=4096, layers=2, norm="none", n_slices=1, |
| attention=0, attn_dim=1024, attn_heads=8, max_len=256, |
| mixer=0, mixer_heads=8, mixer_ffn=2): |
| """MLP-проектор: `layers` линейных слоёв, `layers-1` GELU(tanh) между ними. |
| |
| attention>0 — добавить residual-ветку self-attention на `attention` блоков |
| (выход ветки инициализирован нулями, поэтому на старте функция та же, что у чистого MLP). |
| """ |
| if layers < 2: |
| raise SystemExit("нужно хотя бы 2 линейных слоя") |
| mods = [] |
| if norm == "ln": |
| mods.append(nn.LayerNorm(in_dim)) |
| elif norm == "ln-per-layer": |
| mods.append(PerLayerNorm(in_dim, n_slices)) |
| elif norm == "rms": |
| mods.append(nn.RMSNorm(in_dim)) |
| elif norm != "none": |
| raise SystemExit(f"неизвестный --norm: {norm}") |
| mods += [nn.Linear(in_dim, hidden), nn.GELU(approximate="tanh")] |
| for _ in range(layers - 2): |
| mods += [nn.Linear(hidden, hidden), nn.GELU(approximate="tanh")] |
| mods.append(nn.Linear(hidden, out_dim)) |
| mix = (SliceMixer(n_slices, in_dim // n_slices, out_dim, mixer_heads, mixer, mixer_ffn) |
| if mixer else None) |
| attn = (AttnMixer(in_dim, out_dim, attn_dim, attn_heads, attention, max_len) |
| if attention else None) |
| if mix is not None or attn is not None: |
| return AttnAdapter(mods, attn, mix) |
| return nn.Sequential(*mods) |
|
|
|
|
| class AttnBlock(nn.Module): |
| """Pre-norm self-attention + FFN. Без причинности: текст не автогенеративный, |
| а DiT всё равно видит всю последовательность целиком.""" |
|
|
| def __init__(self, d_model, n_heads, ffn_mult=4): |
| super().__init__() |
| self.n1 = nn.LayerNorm(d_model) |
| self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True) |
| self.n2 = nn.LayerNorm(d_model) |
| self.ffn = nn.Sequential(nn.Linear(d_model, d_model * ffn_mult), |
| nn.GELU(approximate="tanh"), |
| nn.Linear(d_model * ffn_mult, d_model)) |
|
|
| def forward(self, x, key_padding_mask=None): |
| h = self.n1(x) |
| x = x + self.attn(h, h, h, need_weights=False, key_padding_mask=key_padding_mask)[0] |
| return x + self.ffn(self.n2(x)) |
|
|
|
|
| class AttnMixer(nn.Module): |
| """Ветка внимания: своя нормировка -> проекция в d_model -> N блоков -> проекция в out_dim. |
| |
| Последняя проекция инициализирована нулями: на старте ветка не добавляет ничего (см. AttnAdapter). |
| """ |
|
|
| def __init__(self, in_dim, out_dim, d_model=1024, n_heads=8, blocks=2, max_len=256): |
| super().__init__() |
| self.norm = nn.LayerNorm(in_dim) |
| self.inp = nn.Linear(in_dim, d_model) |
| self.pos = nn.Parameter(torch.zeros(1, max_len, d_model)) |
| self.blocks = nn.ModuleList([AttnBlock(d_model, n_heads) for _ in range(blocks)]) |
| self.out = nn.Linear(d_model, out_dim) |
| nn.init.zeros_(self.out.weight) |
| nn.init.zeros_(self.out.bias) |
|
|
| def forward(self, x, key_padding_mask=None): |
| h = self.inp(self.norm(x)) |
| h = h + self.pos[:, :h.shape[1]] |
| for b in self.blocks: |
| h = b(h, key_padding_mask) |
| return self.out(h) |
|
|
|
|
| class SliceMixer(nn.Module): |
| """Attention ПО ОСИ СРЕЗОВ (как txtfusion Krea 2) + добавка к выходу адаптера. |
| |
| Порядок как у Krea: срезы нормируются поштучно (нормы растут с глубиной), идут |
| последовательностью длины K через 2 блока self-attention, затем `Linear(K->1)` |
| схлопывает ось — как `projector` в Krea. |
| |
| Отличие от Krea — точка инъекции: там fusion стоит на входе ДиТ'a, у нас обученная МЛП |
| слишком чувствительна к входу (добавка 0.4% к входу меняет выход в разы — проверено), |
| поэтому результат добавляется к ВЫХОДУ (zero-init), как ветка внимания по токенам. |
| """ |
|
|
| def __init__(self, n_slices, d, out_dim, n_heads=8, blocks=2, ffn_mult=2): |
| super().__init__() |
| self.n, self.d = n_slices, d |
| self.per = nn.LayerNorm(d, elementwise_affine=False) |
| self.blocks = nn.ModuleList([AttnBlock(d, n_heads, ffn_mult) for _ in range(blocks)]) |
| self.proj = nn.Linear(n_slices, 1, bias=False) |
| self.out = nn.Linear(d, out_dim) |
| nn.init.zeros_(self.out.weight) |
| nn.init.zeros_(self.out.bias) |
|
|
| def forward(self, x): |
| b, l, _ = x.shape |
| h = self.per(x.view(b, l, self.n, self.d)).reshape(b * l, self.n, self.d) |
| for blk in self.blocks: |
| h = blk(h) |
| h = self.proj(h.permute(0, 2, 1)).squeeze(-1) |
| return self.out(h).reshape(b, l, -1) |
|
|
|
|
| class AttnAdapter(nn.Module): |
| """Поточечная MLP (self.mlp) + residual-ветки: по токенам (self.attn) и по срезам (self.mixer). |
| |
| Обе ветки инициализированы нулём и добавляются к **выходу** MLP: на вход добавлять |
| нельзя — обученная MLP разгоняет крошечную добавку в разы (см. SliceMixer). |
| Ключи MLP получают префикс "mlp.", ветки — "attn." и "mixer."; чекпоинт без |
| префиксов грузится с переносом (см. train_adapter). |
| """ |
|
|
| def __init__(self, mods, attn=None, mixer=None): |
| super().__init__() |
| self.mlp = nn.Sequential(*mods) |
| self.attn = attn |
| self.mixer = mixer |
|
|
| def forward(self, x, mask=None): |
| """mask: (B, L) bool, True = реальный токен. В attention уходит key_padding_mask = ~mask, |
| иначе (как было в klein до v14) на длине 256 внимание смотрит в паддинг.""" |
| out = self.mlp(x) |
| if self.attn is not None: |
| kpm = None if mask is None else ~mask.bool() |
| out = out + self.attn(x, kpm) |
| if self.mixer is not None: |
| out = out + self.mixer(x) |
| return out |
|
|
|
|
| def infer_arch(state_dict): |
| """Размеры и тип нормировки прямо из весов — чтобы примеру не нужны были флаги. |
| |
| Нормировка опознаётся по ключам: `0.per.weight` — ln-per-layer, `0.weight`+`0.bias` |
| (1-D) — LayerNorm, `0.weight` (1-D) без bias — RMSNorm, `0.weight` (2-D) — нормировки нет. |
| """ |
| mlp = {k[len("mlp."):]: v for k, v in state_dict.items() if k.startswith("mlp.")} |
| if not mlp: |
| mlp = {k: v for k, v in state_dict.items() if not k.startswith("attn.")} |
| linears = sorted((int(k.split(".")[0]), k) for k, v in mlp.items() |
| if k.endswith(".weight") and v.dim() == 2 and k.split(".")[0].isdigit()) |
| if not linears: |
| raise SystemExit("в весах нет линейных слоёв") |
| first, last = mlp[linears[0][1]], mlp[linears[-1][1]] |
| if "0.per.weight" in mlp or "0.all.weight" in mlp: |
| norm = "ln-per-layer" |
| in_dim = mlp["0.all.weight" if "0.all.weight" in mlp else "0.per.weight"].shape[0] |
| elif any(k.startswith("0.") and v.dim() == 1 for k, v in mlp.items()): |
| norm = "ln" if "0.bias" in mlp else "rms" |
| in_dim = first.shape[1] |
| else: |
| norm = "none" |
| in_dim = first.shape[1] |
| attention = len({k.split(".")[2] for k in state_dict if k.startswith("attn.blocks.")}) |
| mixer = len({k.split(".")[2] for k in state_dict if k.startswith("mixer.blocks.")}) |
| arch = {"in_dim": in_dim, "out_dim": last.shape[0], "hidden": first.shape[0], |
| "proj_layers": len(linears), "norm": norm, "attention": attention, "mixer": mixer} |
| if attention: |
| arch["attn_dim"] = state_dict["attn.out.weight"].shape[1] |
| |
| |
| arch["attn_max_len"] = state_dict["attn.pos"].shape[1] |
| return arch |
|
|
|
|
| def load_adapter(path, device="cuda", dtype=None, n_slices=None): |
| """Веса -> (adapter в eval, meta). Понимает .pt (голые веса или полный чекпоинт) и .safetensors. |
| |
| n_slices (число слоёв ТЕ в склейке) нельзя вытащить из весов: у PerLayerNorm |
| нормировка без обучаемых параметров. Поэтому он берётся из meta, а если её нет — |
| из аргумента; при norm="ln-per-layer" и отсутствии обоих это ошибка, а не догадка. |
| """ |
| meta = {} |
| if path.endswith(".safetensors"): |
| from safetensors.torch import load_file |
| from safetensors import safe_open |
| sd = load_file(path) |
| with safe_open(path, framework="pt") as f: |
| meta = dict(f.metadata() or {}) |
| else: |
| sd = torch.load(path, map_location="cpu", weights_only=False) |
| if isinstance(sd, dict) and "state_dict" in sd: |
| meta = {k: v for k, v in sd.items() if k != "state_dict"} |
| sd = sd["state_dict"] |
|
|
| arch = infer_arch(sd) |
| if meta.get("student_layers"): |
| layers = [int(x) for x in str(meta["student_layers"]).split(",")] |
| n_slices = n_slices or len(layers) |
| if n_slices is None: |
| if arch["norm"] == "ln-per-layer": |
| raise SystemExit("нужен n_slices: в весах его нет, а norm=ln-per-layer без него " |
| "загрузится с неверной формой и будет считать мусор") |
| n_slices = 1 |
| adapter = build_adapter(arch["in_dim"], arch["out_dim"], arch["hidden"], |
| arch["proj_layers"], arch["norm"], n_slices=n_slices, |
| attention=arch["attention"], |
| attn_dim=arch.get("attn_dim", 1024), |
| attn_heads=int(meta.get("attn_heads", 8)), |
| max_len=arch.get("attn_max_len", 256), |
| mixer=arch.get("mixer", 0), |
| mixer_heads=int(meta.get("mixer_heads", 8))) |
| adapter.load_state_dict(sd) |
| if dtype is not None: |
| adapter = adapter.to(dtype) |
| adapter = adapter.to(device).eval() |
| meta.update(arch | {"n_slices": n_slices}) |
| return adapter, meta |
|
|
|
|
| def save_safetensors(adapter, path, meta=None): |
| """Сохранить адаптер в safetensors вместе с метаданными (строки, как требует формат).""" |
| from safetensors.torch import save_file |
| md = {k: ("".join(map(str, v)) if isinstance(v, (list, tuple)) else str(v)) |
| for k, v in (meta or {}).items()} |
| save_file({k: v.contiguous() for k, v in adapter.state_dict().items()}, path, metadata=md) |
|
|
