| """Текстовый путь 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 |
|
|
| |
| 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 = slices.shape |
| normed = self.per(slices) |
| h = normed.reshape(b * l, k, d) |
| 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 |
|
|
| 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): |
| |
| 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 |
|
|
| 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([""]) |
|
|
