Blame for adapter/src/adapter_lib.py

#!/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)