File size: 7,039 Bytes
b4f5f26 | 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 | """Текстовый путь UDM: замороженный Ettin + обучаемый выбор слоёв (по идее moss/adapter).
Почему не один слой. Скрытые состояния растут по норме с глубиной (в Qwen3.5-9B — в 40 раз,
13 у слоя 2 против 557 у слоя 27), поэтому «взять последний слой» теряет лексику и цифры,
а «взять ранний» — семантику. В moss/adapter это решили иначе: несколько срезов по K
нормируются ПОШТУЧНО, идут последовательностью длины K через 2 блока self-attention, затем
`Linear(K→1)` схлопывает ось срезов, а результат zero-init'ом добавляется к основному пути.
Тот же приём (плюс 0.8B вместо 9B) дал cos 0.924 к учителю.
Здесь то же самое, но цель — не регрессия к чужому пространству, а обусловливание ДиТ:
выход миксера идёт в `xkv` (кросс-внимание) и в `pmlp` (adaLN).
Срезы Ettin-150M: 22 слоя, global attention каждые 3 → берём выходы после глобальных слоёв.
"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
# выходы после глобальных слоёв ModernBERT (22 слоя, global_attn_every_n_layers=3)
ETTIN_150M_SLICES = (2, 5, 8, 11, 14, 17, 20)
class AttnBlock(nn.Module):
"""Pre-norm self-attention + MLP (для последовательности длины K — оси срезов)."""
def __init__(self, d: int, heads: int = 8, ffn_mult: float = 2.0):
super().__init__()
self.h = heads
self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
self.qkv = nn.Linear(d, 3 * d)
self.proj = nn.Linear(d, d)
self.fc1 = nn.Linear(d, int(d * ffn_mult))
self.fc2 = nn.Linear(int(d * ffn_mult), d)
def forward(self, x):
b, n, d = x.shape
q, k, v = self.qkv(self.n1(x)).reshape(b, n, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
a = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(b, n, d)
x = x + self.proj(a)
return x + self.fc2(F.gelu(self.fc1(self.n2(x))))
class SliceMixer(nn.Module):
"""Выбор слоя: пер-срезный LayerNorm → attention по оси срезов → Linear(K→1) → zero-init out.
blocks=0 оставляет только обученные веса слоёв (Linear(K→1) поверх нормированных срезов) —
дешёвый вариант «какой слой важнее», без внимания между срезами.
"""
def __init__(self, n_slices: int, d: int, blocks: int = 2, heads: int = 8):
super().__init__()
self.n, self.d = n_slices, d
self.per = nn.LayerNorm(d, elementwise_affine=False)
self.blocks = nn.ModuleList([AttnBlock(d, heads) for _ in range(blocks)])
self.proj = nn.Linear(n_slices, 1, bias=False)
self.out = nn.Linear(d, d)
nn.init.zeros_(self.out.weight)
nn.init.zeros_(self.out.bias)
def forward(self, slices): # (B, L, K, D)
b, l, k, d = slices.shape
normed = self.per(slices) # пер-срезная нормировка ОБЯЗАТЕЛЬНА:
h = normed.reshape(b * l, k, d) # сырые нормы срезов растут с глубиной (52…548)
for blk in self.blocks:
h = blk(h)
mixed = self.out(self.proj(h.permute(0, 2, 1)).squeeze(-1)).reshape(b, l, d)
return normed.mean(2) + mixed # среднее по нормированным срезам + zero-init добавка
def layer_weights(self) -> torch.Tensor:
"""Что микшер выучил: ненормированные веса слоёв (для логов)."""
return self.proj.weight.detach().squeeze()
class EttinText(nn.Module):
"""Замороженный энкодер Ettin-150M + обучаемый SliceMixer. Возвращает (B,L,768) и маску."""
def __init__(self, path: str, device, max_len: int = 192, slices=ETTIN_150M_SLICES,
blocks: int = 2, dtype=torch.bfloat16):
super().__init__()
from transformers import AutoModel, AutoTokenizer
self.max_len, self.layer_idx = max_len, tuple(slices)
self.tok = AutoTokenizer.from_pretrained(path)
self.model = AutoModel.from_pretrained(path, dtype=dtype).to(device).eval().requires_grad_(False)
self.dim = self.model.config.hidden_size
self.mixer = SliceMixer(len(self.layer_idx), self.dim, blocks=blocks).to(device)
self.device = torch.device(device)
def tokenize(self, caps: list[str], pad_to: int = 8, pad_to_len: int | None = None):
# padding="longest" (а не до max_len): короткие капшены стоят дешевле
t = self.tok(caps, padding="longest", truncation=True, max_length=self.max_len,
return_tensors="pt")
ids, mask = t["input_ids"], t["attention_mask"].bool()
into = pad_to_len or ids.shape[1]
extra = into - ids.shape[1]
if extra > 0:
ids = F.pad(ids, (0, extra), value=self.tok.pad_token_id)
mask = F.pad(mask, (0, extra))
elif extra < 0:
ids, mask = ids[:, :into], mask[:, :into]
else:
extra = (-ids.shape[1]) % pad_to
if extra:
ids = F.pad(ids, (0, extra), value=self.tok.pad_token_id)
mask = F.pad(mask, (0, extra))
return ids.to(self.device), mask.to(self.device)
@torch.no_grad()
def slices(self, caps: list[str], pad_to_len: int | None = None):
"""Срезы слоёв — считаются ОДИН раз на капшен (энкодер заморожен)."""
ids, mask = self.tokenize(caps, pad_to_len=pad_to_len)
out = self.model(input_ids=ids, attention_mask=mask, output_hidden_states=True)
return torch.stack([out.hidden_states[i] for i in self.layer_idx], dim=2), mask # (B,L,K,D)
def mix(self, hs):
"""Обучаемый путь: среднее по нормированным срезам + zero-init добавка микшера."""
return self.mixer(hs.float())
def forward(self, caps: list[str]):
hs, mask = self.slices(caps)
return self.mix(hs), mask, hs
def trainable(self):
return list(self.mixer.parameters())
def empty_text(enc: EttinText):
"""Пустой промпт = [CLS][SEP] — безусловный контекст для CFG."""
return enc([""])
|