#!/usr/bin/env python3 """Схема адаптера и его загрузка — в одном месте. Раньше класс адаптера был скопирован в 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] # таблица позиций ветки: длина входа на инференсе может быть больше, чем при обучении, # но собрать адаптер надо ровно под сохранённые веса, иначе load_state_dict падает 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)