24.5-dars: Transformer bloki

POST-LN (asl maqola, 2017):
  x = LN(x + Dropout(Attention(x)))
  x = LN(x + Dropout(FFN(x)))

PRE-LN (GPT-2 va ko'p zamonaviy modellar):
  x = x + Dropout(Attention(LN(x)))
  x = x + Dropout(FFN(LN(x)))
  ... oxirgi blokdan keyin yana bitta LN

FFN (har pozitsiyaga ALOHIDA, bir xil og'irliklar bilan):
  FFN(x) = W_2 * aktivatsiya(W_1 x + b_1) + b_2
  W_1: (d_ff, d),  W_2: (d, d_ff),  odatda d_ff = 4d
  aktivatsiya: ReLU (asl), GELU (BERT, GPT)

SHAKLLAR:
  kirish (B, L, d)  ->  chiqish (B, L, d)   - o'lcham o'zgarmaydi
  shuning uchun bloklarni istalgancha ketma-ket qo'yish mumkin

PYTORCH:
  nn.TransformerEncoderLayer(d, h, d_ff, dropout, activation,
                             batch_first=True, norm_first=False)
  norm_first=False -> post-LN,  norm_first=True -> pre-LN
  1-misol: og'irliklarni ko'chirsak, o'z blokimiz bilan farq 0.0e+00

POST-LN:  x_{l+1} = LN(x_l + f(x_l))
  residual yo'l har blokda LN dan O'TADI
  -> "toza" to'g'ri yo'l yo'q: gradient har LN da qayta masshtablanadi
  -> chiqishga yaqin bloklarda gradient kattaroq (1-misol: 12/1 nisbati 1.16)
  -> katta lr da warmup kerak; chuqur stekda beqaror

PRE-LN:  x_{l+1} = x_l + f(LN(x_l))
  residual yo'l LN siz: x_L = x_0 + yig'indi f_l(...)
  -> gradient chiqishdan kirishgacha to'g'ridan-to'g'ri o'tadi
  -> bloklar gradienti bir tekis (12/1 nisbati 0.86)
  -> katta lr va warmup siz ham o'rganadi
  -> oxirida qo'shimcha LN kerak (x_L normallanmagan)

3-MISOL (lr = 0.006, warmup yo'q, 100 qadam, 3 seed):
            2 blok   12 blok
  pre-LN    1.000    1.000
  post-LN   1.000    0.060     (tasodif 1/16 = 0.0625)
  12 blokda post-LN loss 2.774 da qotib qoldi (ln 16 = 2.773)

NARX:
  post-LN yaxshi sozlanganda ba'zan sal yaxshiroq natija beradi,
  lekin sozlash qiyin; amaliyotda pre-LN - xavfsiz tanlov

RESIDUAL SIZ:  x_{l+1} = f(x_l)
  zanjir qoidasi: dL/dx_0 = J_1^T J_2^T ... J_L^T dL/dx_L
  12 ta Yakobian ko'paytmasi - har biri kichraytirsa, gradient so'nadi

RESIDUAL BILAN:  x_{l+1} = x_l + f(x_l)
  dx_{l+1}/dx_l = I + J_l     <- "I" - to'g'ri yo'l
  gradient hech bo'lmaganda birlik matritsa orqali o'tadi

2-MISOL (12 blok, boshlang'ich holat, blok kirishidagi gradient):
                     1-blok     12-blok    nisbat
  residual bilan    1.24e-02   1.12e-02    1.11
  residual siz      2.71e-03   2.54e-02    0.11

YANA BIR MUAMMO - TOKENLAR "QO'SHILIB KETADI":
  attention - vaznli o'rtacha; o'rtachaning o'rtachasi ... -> hamma bir xil
  2-misol: tokenlar orasidagi o'rtacha cos
    residual siz: 1-blokdan keyin 0.924, 2-blokdan keyin 1.000
    residual bilan: 12-blokdan keyin ham 0.284
  residual har blokda token o'z vektorini "eslab qolishini" ta'minlaydi

O'RGATISH (2-misol): residual bilan 1.000, residual siz 0.061 (tasodif 0.0625)

ATTENTION NIMA QILADI:
  boshqa pozitsiyalardagi vektorlarning vaznli o'rtachasini olib keladi
  "ma'lumotni tashish" - lekin qiymatlar ustida chiziqli (w V)

FFN NIMA QILADI:
  har pozitsiyada alohida, keng (d_ff = 4d) nochiziqli almashtirish
  "yig'ilgan ma'lumot ustida hisoblash"
  talqinlardan biri: kalit-qiymat xotirasi (W_1 qatorlari - naqshlar,
  W_2 ustunlari - ularga mos chiqishlar)

4-MISOL: y_t = (x_t + x_{t-1}) mod 10
  1 attention:          0.736   (5386 parametr)
  3 attention:          0.714   (13962 parametr; seedlar 0.245 ... 0.987)
  1 attention + FFN:    1.000   (13802 parametr; 3 seedda ham)
  -> parametr soni emas, FFN ning nochiziqli hisobi hal qildi

PARAMETR VA HISOB ULUSHI (d_ff = 4d):
  parametrlarning ~2/3 qismi FFN da (4-misol: 66.7%)
  FLOP: L < 4d bo'lsa FFN ko'proq; L = 4d da skor qismi FFN ga tenglashadi

nn.TransformerEncoderLayer ichida (1-misol):
  self_attn.dropout - attention OG'IRLIKLARIGA (softmax dan keyin)
  dropout1          - attention chiqishiga, residual qo'shishdan OLDIN
  dropout           - FFN ichida, aktivatsiyadan keyin
  dropout2          - FFN chiqishiga, residual qo'shishdan OLDIN

QOIDA:
  dropout f(x) ga qo'llanadi, residual yo'lga EMAS
  x + Dropout(f(x))  -  to'g'ri
  Dropout(x + f(x))  -  to'g'ri yo'lni ham "teshadi"

TRAIN VA EVAL:
  train: ikki chaqiruv farqi 0.796 (1-misol) - har safar boshqa niqob
  eval:  0.0e+00 - dropout o'chadi
  baholash va attention tahlilidan oldin model.eval() shart

BITTA BLOK (biaslar bilan):
  attention:  4 d^2 + 4 d              (W_q, W_k, W_v, W_o)
  FFN:        2 d d_ff + d_ff + d
  LayerNorm:  2 * 2d = 4d              (gamma, beta)
  jami:       4 d^2 + 2 d d_ff + 9d + d_ff
  bosh soni h parametrga TA'SIR QILMAYDI (d = h * d_h)

1-MISOL (d = 64, d_ff = 256):
  16640 + 33088 + 256 = 49984 = sum(p.numel()) = nn.TransformerEncoderLayer
  d_ff = 4d da:  ~12 d^2
  BERT-base (d = 768):  bir blok 7 087 872, 12 blok 85.1 mln
  BERT-large (d = 1024): bir blok 12 596 224, 24 blok 302.3 mln
  (qolgani - token va pozitsiya embeddinglari)

FLOP (bir token, bitta blok, ko'paytirish + qo'shish = 2 FLOP):
  proyeksiyalar  8 d^2
  Q K^T va w V   4 L d            <- L ga chiziqli (butun ketma-ketlik uchun L^2)
  FFN            4 d d_ff = 16 d^2
  4-misol, d = 768: L = 512 da FFN ulushi 60.0%, L = 4096 da 35.3%

import torch
import torch.nn as nn
import torch.nn.functional as F

class Blok(nn.Module):
    def __init__(self, d, h, d_ff, p=0.1):
        super().__init__()
        self.att = nn.MultiheadAttention(d, h, dropout=p, batch_first=True)
        self.ffn = nn.Sequential(nn.Linear(d, d_ff), nn.GELU(),
                                 nn.Dropout(p), nn.Linear(d_ff, d))
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.drop1, self.drop2 = nn.Dropout(p), nn.Dropout(p)

    def forward(self, x):                                   # pre-LN
        y = self.ln1(x)
        x = x + self.drop1(self.att(y, y, y, need_weights=False)[0])
        return x + self.drop2(self.ffn(self.ln2(x)))

# tayyor modul
qatlam = nn.TransformerEncoderLayer(d, h, d_ff, dropout=0.1, activation="gelu",
                                    batch_first=True, norm_first=True)
stek = nn.TransformerEncoder(qatlam, num_layers=6, enable_nested_tensor=False)

# parametrlar: formula va tekshiruv
jami = 4 * d * d + 2 * d * d_ff + 9 * d + d_ff
assert jami == sum(p.numel() for p in qatlam.parameters())

# bloklar bo'yicha gradient normasi
for i, b in enumerate(model.bloklar):
    print(i, torch.cat([p.grad.flatten() for p in b.parameters()]).norm())

blok = attention + FFN, har biri residual + LayerNorm bilan
pre-LN: x + f(LN(x)) - chuqur stekda barqaror, oxirida LN
residual - gradient va token ma'lumoti uchun to'g'ri yo'l
FFN - har pozitsiyada nochiziqli hisob, parametrlarning ~2/3
parametr: 4d^2 + 2 d d_ff + 9d + d_ff  (d_ff = 4d da ~12 d^2)

"""Transformer bloki noldan: parametr formulasi va nn.TransformerEncoderLayer bilan moslik."""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class Blok(nn.Module):
    """attention + FFN + residual + LayerNorm; pre_ln=True -> pre-LN, aks holda post-LN."""

    def __init__(self, d, h, d_ff, dropout=0.1, pre_ln=False):
        super().__init__()
        self.h, self.pre_ln, self.p = h, pre_ln, dropout
        self.W_qkv = nn.Linear(d, 3 * d)           # W_q, W_k, W_v bitta matritsada
        self.W_o = nn.Linear(d, d)
        self.ff1, self.ff2 = nn.Linear(d, d_ff), nn.Linear(d_ff, d)
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)

    def attention(self, x):
        B, L, d = x.shape
        q, k, v = self.W_qkv(x).reshape(B, L, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        w = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(d // self.h), dim=-1)
        w = F.dropout(w, self.p, self.training)                     # 1: og'irliklarda
        a = (w @ v).transpose(1, 2).reshape(B, L, d)
        return F.dropout(self.W_o(a), self.p, self.training)       # 2: residual dan oldin

    def ffn(self, x):
        z = F.dropout(F.relu(self.ff1(x)), self.p, self.training)  # 3: FFN ichida
        return F.dropout(self.ff2(z), self.p, self.training)       # 4: residual dan oldin

    def forward(self, x):
        if self.pre_ln:
            x = x + self.attention(self.ln1(x))
            return x + self.ffn(self.ln2(x))
        x = self.ln1(x + self.attention(x))
        return self.ln2(x + self.ffn(x))


def formula(d, d_ff):
    attention = 4 * d * d + 4 * d          # W_q, W_k, W_v, W_o va biaslar
    ffn = 2 * d * d_ff + d_ff + d          # ikki Linear
    ln = 2 * 2 * d                         # ikki LayerNorm: gamma va beta
    return attention, ffn, ln


def main() -> None:
    torch.manual_seed(0)
    d, h, d_ff = 64, 4, 256

    print("=== 1. Parametrlar: formula va sum(p.numel()) ===")
    blok = Blok(d, h, d_ff)
    a, f, n = formula(d, d_ff)
    print(f"  attention 4d^2 + 4d        = {a:>7}")
    print(f"  FFN       2*d*d_ff + d_ff + d = {f:>7}")
    print(f"  LayerNorm 4d               = {n:>7}")
    print(f"  jami (formula)             = {a + f + n:>7}")
    print(f"  jami (o'z blokimiz)        = {sum(p.numel() for p in blok.parameters()):>7}")
    tel = nn.TransformerEncoderLayer(d, h, d_ff, dropout=0.1, batch_first=True)
    print(f"  jami (nn.TransformerEncoderLayer) = {sum(p.numel() for p in tel.parameters())}")
    print(f"  ulushlar: attention {a / (a + f + n):.1%}, FFN {f / (a + f + n):.1%}, "
          f"LayerNorm {n / (a + f + n):.1%}")

    print("\n=== 2. Katta o'lchamlar (d_ff = 4d) ===")
    for nom, dd, hh, qatlam in [("BERT-base", 768, 12, 12), ("BERT-large", 1024, 16, 24)]:
        a, f, n = formula(dd, 4 * dd)
        haqiqiy = sum(p.numel() for p in
                      nn.TransformerEncoderLayer(dd, hh, 4 * dd, batch_first=True).parameters())
        print(f"  {nom:<11} d={dd:<5} bir blok {a + f + n:>9} (numel {haqiqiy:>9}),"
              f" {qatlam} blok {qatlam * (a + f + n) / 1e6:.1f} mln")

    print("\n=== 3. Og'irliklarni ko'chirib natijani solishtirish (eval) ===")
    X = torch.randn(4, 10, d)
    for pre_ln in (False, True):
        torch.manual_seed(1)
        mening = Blok(d, h, d_ff, pre_ln=pre_ln).eval()
        tayyor = nn.TransformerEncoderLayer(d, h, d_ff, dropout=0.1, batch_first=True,
                                            norm_first=pre_ln).eval()
        with torch.no_grad():
            oldin = (mening(X) - tayyor(X)).abs().max().item()
            tayyor.self_attn.in_proj_weight.copy_(mening.W_qkv.weight)
            tayyor.self_attn.in_proj_bias.copy_(mening.W_qkv.bias)
            tayyor.self_attn.out_proj.weight.copy_(mening.W_o.weight)
            tayyor.self_attn.out_proj.bias.copy_(mening.W_o.bias)
            for a_, b_ in [(tayyor.linear1, mening.ff1), (tayyor.linear2, mening.ff2),
                           (tayyor.norm1, mening.ln1), (tayyor.norm2, mening.ln2)]:
                a_.weight.copy_(b_.weight)
                a_.bias.copy_(b_.bias)
            farq = (mening(X) - tayyor(X)).abs().max().item()
        nom = "pre-LN (norm_first=True) " if pre_ln else "post-LN (norm_first=False)"
        print(f"  {nom}: ko'chirishdan oldin {oldin:.3f}, keyin {farq:.1e}  "
              f"{'MOS' if farq < 1e-4 else 'MOS EMAS'}")

    print("\n=== 4. Dropout joylari: train va eval ===")
    print("  nn.TransformerEncoderLayer ichidagi dropoutlar:")
    for nom, mod in tel.named_modules():
        if isinstance(mod, nn.Dropout):
            print(f"    {nom:<9} p = {mod.p}")
    print(f"    self_attn.dropout (og'irliklarda) p = {tel.self_attn.dropout}")
    blok = Blok(d, h, d_ff, dropout=0.1)
    torch.manual_seed(2)
    blok.train()
    y1, y2 = blok(X), blok(X)
    blok.eval()
    y3, y4 = blok(X), blok(X)
    print(f"  train: ikki chaqiruv farqi {(y1 - y2).abs().max().item():.3f}")
    print(f"  eval:  ikki chaqiruv farqi {(y3 - y4).abs().max().item():.1e}")
    print("  ⭐ Blok = attention + FFN, har biri residual va LayerNorm bilan o'ralgan")


if __name__ == "__main__":
    main()

=== 1. Parametrlar: formula va sum(p.numel()) ===
  attention 4d^2 + 4d        =   16640
  FFN       2*d*d_ff + d_ff + d =   33088
  LayerNorm 4d               =     256
  jami (formula)             =   49984
  jami (o'z blokimiz)        =   49984
  jami (nn.TransformerEncoderLayer) = 49984
  ulushlar: attention 33.3%, FFN 66.2%, LayerNorm 0.5%

=== 2. Katta o'lchamlar (d_ff = 4d) ===
  BERT-base   d=768   bir blok   7087872 (numel   7087872), 12 blok 85.1 mln
  BERT-large  d=1024  bir blok  12596224 (numel  12596224), 24 blok 302.3 mln

=== 3. Og'irliklarni ko'chirib natijani solishtirish (eval) ===
  post-LN (norm_first=False): ko'chirishdan oldin 1.618, keyin 0.0e+00  MOS
  pre-LN (norm_first=True) : ko'chirishdan oldin 1.637, keyin 0.0e+00  MOS

=== 4. Dropout joylari: train va eval ===
  nn.TransformerEncoderLayer ichidagi dropoutlar:
    dropout   p = 0.1
    dropout1  p = 0.1
    dropout2  p = 0.1
    self_attn.dropout (og'irliklarda) p = 0.1
  train: ikki chaqiruv farqi 0.796
  eval:  ikki chaqiruv farqi 0.0e+00
  ⭐ Blok = attention + FFN, har biri residual va LayerNorm bilan o'ralgan

"""Residual ulanishsiz chuqur stek: qatlamlar bo'yicha gradient normasi va tokenlar 'qo'shilib ketishi'."""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

V, SILJISH, L, N = 16, 3, 16, 12


class Blok(nn.Module):
    """Pre-LN blok; residual=False bo'lsa x + f(x) o'rniga faqat f(x)."""

    def __init__(self, d, h, residual):
        super().__init__()
        self.h, self.residual = h, residual
        self.qkv, self.o = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.ffn = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def attention(self, x):
        B, L_, d = x.shape
        q, k, v = self.qkv(x).reshape(B, L_, 3, self.h, -1).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v)
        return self.o(a.transpose(1, 2).reshape(B, L_, d))

    def forward(self, x):
        if self.residual:
            x = x + self.attention(self.ln1(x))
            return x + self.ffn(self.ln2(x))
        x = self.attention(self.ln1(x))
        return self.ffn(self.ln2(x))


class Model(nn.Module):
    def __init__(self, residual, n=N, d=32, h=2):
        super().__init__()
        self.emb, self.pos = nn.Embedding(V, d), nn.Embedding(L, d)
        self.bloklar = nn.ModuleList([Blok(d, h, residual) for _ in range(n)])
        self.ln, self.chiq = nn.LayerNorm(d), nn.Linear(d, V)

    def forward(self, X, holatlar=None):
        x = self.emb(X) + self.pos(torch.arange(X.shape[1]))
        for b in self.bloklar:
            if holatlar is not None:
                x.retain_grad()
                holatlar.append(x)
            x = b(x)
        if holatlar is not None:
            holatlar.append(x)
        return self.chiq(self.ln(x))


def batch(n, g):
    X = torch.randint(0, V, (n, L), generator=g)
    Y = torch.full((n, L), -100)
    Y[:, SILJISH:] = X[:, :-SILJISH]
    return X, Y


def tashxis(residual):
    """Boshlang'ich holatda: har blok KIRISHIGA kelgan gradient va tokenlar o'xshashligi."""
    torch.manual_seed(0)
    m = Model(residual)
    X, Y = batch(64, torch.Generator().manual_seed(5))
    holatlar = []
    F.cross_entropy(m(X, holatlar).reshape(-1, V), Y.reshape(-1)).backward()
    grad = [h.grad.norm().item() for h in holatlar[:-1]]
    kos = []
    for h in holatlar:
        z = F.normalize(h.detach(), dim=-1)
        kos.append((z @ z.transpose(1, 2)).mean().item())   # tokenlar orasidagi o'rtacha cos
    return grad, kos


def orgat(residual, seed, qadam=120):
    torch.manual_seed(seed)
    m = Model(residual)
    opt = torch.optim.Adam(m.parameters(), lr=0.003)
    g = torch.Generator().manual_seed(seed)
    for _ in range(qadam):
        X, Y = batch(32, g)
        loss = F.cross_entropy(m(X).reshape(-1, V), Y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    m.eval()
    X, Y = batch(500, torch.Generator().manual_seed(99))
    with torch.no_grad():
        p = m(X).argmax(-1)
    return (p[Y >= 0] == Y[Y >= 0]).float().mean().item()


def main() -> None:
    torch.set_num_threads(1)
    print(f"=== 1. Gradient normasi blok kirishida ({N} blok, boshlang'ich holat) ===")
    g_r, k_r = tashxis(True)
    g_y, k_y = tashxis(False)
    print(f"  {'blok':>5} {'residual bilan':>15} {'residual siz':>13}")
    for i in (0, 2, 4, 6, 8, 10, 11):
        print(f"  {i + 1:>5} {g_r[i]:>15.2e} {g_y[i]:>13.2e}")
    print(f"  1-blok / 12-blok nisbati: residual bilan {g_r[0] / g_r[-1]:.2f}, "
          f"residual siz {g_y[0] / g_y[-1]:.2f}")

    print("\n=== 2. Tokenlar bir-biriga o'xshab ketadimi (o'rtacha cos) ===")
    print(f"  {'qatlam':>7} {'residual bilan':>15} {'residual siz':>13}")
    for i in (0, 1, 2, 4, 8, 12):
        nom = "kirish" if i == 0 else f"{i}-blok"
        print(f"  {nom:>7} {k_r[i]:>15.3f} {k_y[i]:>13.3f}")
    if k_y[-1] > 0.99:
        print("  residual siz: hamma pozitsiya deyarli BIR XIL vektorga aylandi")
        print("  -> qaysi token qayerda turgani haqidagi ma'lumot yo'qoldi")

    print("\n=== 3. O'rgatish: 3 qadam oldingi tokenni topish (120 qadam, 3 seed) ===")
    nat = {True: [], False: []}
    for s in range(3):
        for r in (True, False):
            nat[r].append(orgat(r, s))
    for r in (True, False):
        nom = "residual bilan" if r else "residual siz"
        print(f"  {nom:<15} " + " ".join(f"{a:.3f}" for a in nat[r])
              + f"   o'rtacha {np.mean(nat[r]):.3f}")
    f = np.array(nat[True]) - np.array(nat[False])
    se = f.std(ddof=1) / np.sqrt(len(f))
    print(f"  farq {f.mean():+.4f}, SE {se:.4f}, sezilarli: {abs(f.mean()) > 2 * se}")
    print(f"  tasodif darajasi 1/{V} = {1 / V:.3f}")
    print("  ⭐ Residual - gradient va ma'lumot uchun 'to'g'ri yo'l'; usiz chuqur stek o'rganmaydi")


if __name__ == "__main__":
    main()

=== 1. Gradient normasi blok kirishida (12 blok, boshlang'ich holat) ===
   blok  residual bilan  residual siz
      1        1.24e-02      2.71e-03
      3        1.20e-02      7.65e-03
      5        1.16e-02      4.04e-03
      7        1.16e-02      6.15e-03
      9        1.13e-02      9.40e-03
     11        1.12e-02      1.33e-02
     12        1.12e-02      2.54e-02
  1-blok / 12-blok nisbati: residual bilan 1.11, residual siz 0.11

=== 2. Tokenlar bir-biriga o'xshab ketadimi (o'rtacha cos) ===
   qatlam  residual bilan  residual siz
   kirish           0.079         0.079
   1-blok           0.095         0.924
   2-blok           0.105         1.000
   4-blok           0.126         1.000
   8-blok           0.219         1.000
  12-blok           0.284         1.000
  residual siz: hamma pozitsiya deyarli BIR XIL vektorga aylandi
  -> qaysi token qayerda turgani haqidagi ma'lumot yo'qoldi

=== 3. O'rgatish: 3 qadam oldingi tokenni topish (120 qadam, 3 seed) ===
  residual bilan  1.000 1.000 1.000   o'rtacha 1.000
  residual siz    0.059 0.065 0.058   o'rtacha 0.061
  farq +0.9391, SE 0.0023, sezilarli: True
  tasodif darajasi 1/16 = 0.062
  ⭐ Residual - gradient va ma'lumot uchun 'to'g'ri yo'l'; usiz chuqur stek o'rganmaydi

"""Pre-LN va post-LN: sayoz va chuqur stekda o'rgatish barqarorligi."""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

V, SILJISH, L, LR = 16, 3, 16, 0.006


class Blok(nn.Module):
    def __init__(self, d, h, tur):
        super().__init__()
        self.h, self.tur = h, tur
        self.qkv, self.o = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.ffn = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def attention(self, x):
        B, L_, d = x.shape
        q, k, v = self.qkv(x).reshape(B, L_, 3, self.h, -1).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v)
        return self.o(a.transpose(1, 2).reshape(B, L_, d))

    def forward(self, x):
        if self.tur == "pre":                            # x + f(LN(x))
            x = x + self.attention(self.ln1(x))
            return x + self.ffn(self.ln2(x))
        x = self.ln1(x + self.attention(x))              # LN(x + f(x))
        return self.ln2(x + self.ffn(x))


class Model(nn.Module):
    def __init__(self, tur, n, d=32, h=2):
        super().__init__()
        self.emb, self.pos = nn.Embedding(V, d), nn.Embedding(L, d)
        self.bloklar = nn.ModuleList([Blok(d, h, tur) for _ in range(n)])
        self.ln = nn.LayerNorm(d) if tur == "pre" else nn.Identity()   # pre-LN: yakuniy LN
        self.chiq = nn.Linear(d, V)

    def forward(self, X):
        x = self.emb(X) + self.pos(torch.arange(X.shape[1]))
        for b in self.bloklar:
            x = b(x)
        return self.chiq(self.ln(x))


def batch(n, g):
    X = torch.randint(0, V, (n, L), generator=g)
    Y = torch.full((n, L), -100)
    Y[:, SILJISH:] = X[:, :-SILJISH]
    return X, Y


def orgat(tur, n, seed, qadam=100):
    torch.manual_seed(seed)
    m = Model(tur, n)
    opt = torch.optim.Adam(m.parameters(), lr=LR)
    g = torch.Generator().manual_seed(seed)
    tarix, normalar = [], []
    for _ in range(qadam):
        X, Y = batch(32, g)
        loss = F.cross_entropy(m(X).reshape(-1, V), Y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        normalar.append(torch.cat([p.grad.flatten() for p in m.parameters()]).norm().item())
        opt.step()
        tarix.append(loss.item())
    m.eval()
    X, Y = batch(500, torch.Generator().manual_seed(99))
    with torch.no_grad():
        p = m(X).argmax(-1)
    return (p[Y >= 0] == Y[Y >= 0]).float().mean().item(), tarix, normalar


def blok_gradlari(tur, n=12):
    torch.manual_seed(0)
    m = Model(tur, n)
    X, Y = batch(64, torch.Generator().manual_seed(5))
    F.cross_entropy(m(X).reshape(-1, V), Y.reshape(-1)).backward()
    return [torch.cat([p.grad.flatten() for p in b.parameters()]).norm().item()
            for b in m.bloklar]


def main() -> None:
    torch.set_num_threads(1)
    print("=== 1. Boshlang'ich holatda blok parametrlari gradienti (12 blok) ===")
    gp, gq = blok_gradlari("pre"), blok_gradlari("post")
    print(f"  {'blok':>5} {'pre-LN':>9} {'post-LN':>9}")
    for i in (0, 3, 7, 11):
        print(f"  {i + 1:>5} {gp[i]:>9.3f} {gq[i]:>9.3f}")
    print(f"  12-blok / 1-blok: pre-LN {gp[-1] / gp[0]:.2f}, post-LN {gq[-1] / gq[0]:.2f}")

    print(f"\n=== 2. O'rgatish: lr = {LR}, warmup yo'q, 100 qadam, 3 seed ===")
    nat = {}
    for n in (2, 12):
        for tur in ("pre", "post"):
            nat[(n, tur)] = [orgat(tur, n, s) for s in range(3)]
    print(f"  {'bloklar':>7} {'tur':<5} {'seedlar':<20}  o'rtacha")
    for (n, tur), r in nat.items():
        a = [x[0] for x in r]
        print(f"  {n:>7} {tur:<5} {' '.join(f'{v:.3f}' for v in a):<20} {np.mean(a):>9.3f}")
    for n in (2, 12):
        f = np.array([x[0] for x in nat[(n, "pre")]]) - np.array([x[0] for x in nat[(n, "post")]])
        se = f.std(ddof=1) / np.sqrt(len(f))
        print(f"  {n:>2} blok: pre - post = {f.mean():+.4f}, SE {se:.4f}, "
              f"sezilarli: {abs(f.mean()) > 2 * se}")

    print("\n=== 3. 12 blok: loss va gradient normasi jarayonda (seed 0) ===")
    print(f"  {'qadam':>6} {'pre loss':>9} {'post loss':>10} {'pre |g|':>8} {'post |g|':>9}")
    _, tp, np_ = nat[(12, "pre")][0]
    _, tq, nq = nat[(12, "post")][0]
    for q in (1, 10, 20, 40, 70, 100):
        print(f"  {q:>6} {tp[q - 1]:>9.3f} {tq[q - 1]:>10.3f} {np_[q - 1]:>8.2f} {nq[q - 1]:>9.2f}")
    print(f"  ln({V}) = {np.log(V):.3f} - 'hech narsa o'rganmagan' loss")
    if tq[-1] > np.log(V) - 0.1 > tp[-1]:
        print("  post-LN loss tasodif darajasida qotib qoldi, pre-LN esa nolga yaqin")
    print("  ⭐ Chuqur stekda pre-LN katta lr va warmup siz ham barqaror o'rganadi")


if __name__ == "__main__":
    main()

=== 1. Boshlang'ich holatda blok parametrlari gradienti (12 blok) ===
   blok    pre-LN   post-LN
      1     0.127     0.240
      4     0.110     0.219
      8     0.115     0.266
     12     0.108     0.278
  12-blok / 1-blok: pre-LN 0.86, post-LN 1.16

=== 2. O'rgatish: lr = 0.006, warmup yo'q, 100 qadam, 3 seed ===
  bloklar tur   seedlar               o'rtacha
        2 pre   1.000 1.000 1.000        1.000
        2 post  1.000 1.000 1.000        1.000
       12 pre   1.000 1.000 1.000        1.000
       12 post  0.058 0.058 0.063        0.060
   2 blok: pre - post = +0.0000, SE 0.0000, sezilarli: False
  12 blok: pre - post = +0.9403, SE 0.0018, sezilarli: True

=== 3. 12 blok: loss va gradient normasi jarayonda (seed 0) ===
   qadam  pre loss  post loss  pre |g|  post |g|
       1     2.932      2.899     0.56      0.88
      10     2.793      2.801     0.46      0.49
      20     2.753      2.776     0.28      0.25
      40     2.297      2.779     0.48      0.35
      70     0.062      2.776     0.11      0.31
     100     0.006      2.774     0.01      0.24
  ln(16) = 2.773 - 'hech narsa o'rganmagan' loss
  post-LN loss tasodif darajasida qotib qoldi, pre-LN esa nolga yaqin
  ⭐ Chuqur stekda pre-LN katta lr va warmup siz ham barqaror o'rganadi

"""FFN roli: attention yig'adi, FFN qayta ishlaydi; parametr va FLOP ulushi."""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

V, L = 10, 12


class Attention(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.qkv, self.o = nn.Linear(d, 3 * d), nn.Linear(d, d)

    def forward(self, x):
        B, L_, d = x.shape
        q, k, v = self.qkv(x).reshape(B, L_, 3, self.h, -1).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v)
        return self.o(a.transpose(1, 2).reshape(B, L_, d))


class Model(nn.Module):
    """n_att ta pre-LN attention qatlami, ixtiyoriy bitta FFN qatlami."""

    def __init__(self, n_att, ffn, d=32, h=2):
        super().__init__()
        self.emb, self.pos = nn.Embedding(V, d), nn.Embedding(L, d)
        self.lnlar = nn.ModuleList([nn.LayerNorm(d) for _ in range(n_att)])
        self.attlar = nn.ModuleList([Attention(d, h) for _ in range(n_att)])
        self.ffn = None
        if ffn:
            self.ln_f = nn.LayerNorm(d)
            self.ffn = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.ln, self.chiq = nn.LayerNorm(d), nn.Linear(d, V)

    def forward(self, X):
        x = self.emb(X) + self.pos(torch.arange(X.shape[1]))
        for ln, att in zip(self.lnlar, self.attlar):
            x = x + att(ln(x))
        if self.ffn is not None:
            x = x + self.ffn(self.ln_f(x))
        return self.chiq(self.ln(x))


def batch(n, g):
    """Nishon: (joriy token + oldingi token) mod 10 - ikki tokenni 'birlashtirish'."""
    X = torch.randint(0, V, (n, L), generator=g)
    Y = torch.full((n, L), -100)
    Y[:, 1:] = (X[:, 1:] + X[:, :-1]) % V
    return X, Y


def orgat(n_att, ffn, seed, qadam=400):
    torch.manual_seed(seed)
    m = Model(n_att, ffn)
    opt = torch.optim.Adam(m.parameters(), lr=0.003)
    g = torch.Generator().manual_seed(seed)
    for _ in range(qadam):
        X, Y = batch(64, g)
        loss = F.cross_entropy(m(X).reshape(-1, V), Y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    m.eval()
    X, Y = batch(1000, torch.Generator().manual_seed(99))
    with torch.no_grad():
        p = m(X).argmax(-1)
    return (p[Y >= 0] == Y[Y >= 0]).float().mean().item(), sum(p.numel() for p in m.parameters())


def floplar(d, L_, d_ff):
    """Bir token uchun ko'paytirish-qo'shish x2 (FLOP), bitta blok."""
    proyeksiya = 2 * 4 * d * d               # W_q, W_k, W_v, W_o
    skor = 2 * 2 * L_ * d                    # Q K^T va w V
    ffn = 2 * 2 * d * d_ff
    return proyeksiya, skor, ffn


def main() -> None:
    torch.set_num_threads(1)
    print("=== 1. Vazifa: y_t = (x_t + x_{t-1}) mod 10 ===")
    X, Y = batch(1, torch.Generator().manual_seed(0))
    print(f"  x: {X[0].tolist()}")
    print(f"  y: {['-' if t < 0 else t for t in Y[0].tolist()]}")
    print("  attention x_{t-1} ni olib keladi; 'mod 10 yig'indi' esa har token ichidagi hisob")

    variantlar = [("1 attention", 1, False), ("3 attention", 3, False),
                  ("1 attention + FFN", 1, True)]
    nat = {}
    for nom, n_att, ffn in variantlar:
        nat[nom] = [orgat(n_att, ffn, s) for s in range(3)]

    print("\n=== 2. Aniqlik (400 qadam, 3 seed) ===")
    print(f"  {'model':<19} {'parametr':>8}  {'seedlar':<20} o'rtacha")
    for nom, r in nat.items():
        a = [x[0] for x in r]
        print(f"  {nom:<19} {r[0][1]:>8}  {' '.join(f'{v:.3f}' for v in a):<20} "
              f"{np.mean(a):>8.3f}")
    print(f"  tasodif 1/{V} = {1 / V:.3f}")
    ffn_a = np.array([x[0] for x in nat["1 attention + FFN"]])
    for nom in ["1 attention", "3 attention"]:
        f = ffn_a - np.array([x[0] for x in nat[nom]])
        se = f.std(ddof=1) / np.sqrt(len(f))
        print(f"  FFN - ({nom}): {f.mean():+.4f}, SE {se:.4f}, "
              f"sezilarli: {abs(f.mean()) > 2 * se}")

    print("\n=== 3. Parametr ulushi (d_ff = 4d) ===")
    for d in (32, 768):
        att, ffn = 4 * d * d + 4 * d, 2 * d * 4 * d + 4 * d + d
        print(f"  d = {d:>4}: attention {att:>9}, FFN {ffn:>9}, FFN ulushi {ffn / (att + ffn):.1%}")

    print("\n=== 4. Bir token uchun FLOP, bitta blok (d = 768, d_ff = 3072) ===")
    print(f"  {'L':>6} {'proyeksiya':>11} {'QK^T, wV':>10} {'FFN':>10} {'FFN ulushi':>11}")
    for L_ in (128, 512, 1536, 4096):
        p, s, f = floplar(768, L_, 3072)
        print(f"  {L_:>6} {p / 1e6:>10.1f}M {s / 1e6:>9.1f}M {f / 1e6:>9.1f}M "
              f"{f / (p + s + f):>11.1%}")
    print("  skor qismi L ga proporsional; L = 4d = 3072 da FFN bilan tenglashadi")
    print("  ⭐ Attention ma'lumotni pozitsiyalar orasida tashiydi, FFN har pozitsiyada hisoblaydi")


if __name__ == "__main__":
    main()

=== 1. Vazifa: y_t = (x_t + x_{t-1}) mod 10 ===
  x: [4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6]
  y: ['-', 3, 2, 3, 3, 2, 6, 0, 0, 0, 4, 7]
  attention x_{t-1} ni olib keladi; 'mod 10 yig'indi' esa har token ichidagi hisob

=== 2. Aniqlik (400 qadam, 3 seed) ===
  model               parametr  seedlar              o'rtacha
  1 attention             5386  0.837 0.764 0.607       0.736
  3 attention            13962  0.245 0.987 0.910       0.714
  1 attention + FFN      13802  1.000 1.000 1.000       1.000
  tasodif 1/10 = 0.100
  FFN - (1 attention): +0.2639, SE 0.0678, sezilarli: True
  FFN - (3 attention): +0.2861, SE 0.2355, sezilarli: False

=== 3. Parametr ulushi (d_ff = 4d) ===
  d =   32: attention      4224, FFN      8352, FFN ulushi 66.4%
  d =  768: attention   2362368, FFN   4722432, FFN ulushi 66.7%

=== 4. Bir token uchun FLOP, bitta blok (d = 768, d_ff = 3072) ===
       L  proyeksiya   QK^T, wV        FFN  FFN ulushi
     128        4.7M       0.4M       9.4M       64.9%
     512        4.7M       1.6M       9.4M       60.0%
    1536        4.7M       4.7M       9.4M       50.0%
    4096        4.7M      12.6M       9.4M       35.3%
  skor qismi L ga proporsional; L = 4d = 3072 da FFN bilan tenglashadi
  ⭐ Attention ma'lumotni pozitsiyalar orasida tashiydi, FFN har pozitsiyada hisoblaydi

return self.chiq(x)                                              # ⚠️
return self.chiq(self.ln_oxirgi(x))                              # ✅

x = self.ln1(x); x = x + self.att(x)          # asl x yo'qoldi   # ⚠️
x = x + self.att(self.ln1(x))                 # pre-LN           # ✅

x = self.drop(x + self.att(self.ln1(x)))                         # ⚠️
x = x + self.drop(self.att(self.ln1(x)))                         # ✅

nn.TransformerEncoderLayer(d, h, d_ff)        # kutadi (L, B, d) # ⚠️
nn.TransformerEncoderLayer(d, h, d_ff, batch_first=True)         # ✅

nn.TransformerEncoderLayer(d, h, d_ff, norm_first=False)  # 12 qatlam, lr 6e-3 # ⚠️
nn.TransformerEncoderLayer(d, h, d_ff, norm_first=True)   # yoki warmup      # ✅

with torch.no_grad(): p = model(X)            # train rejimida   # ⚠️
model.eval()
with torch.no_grad(): p = model(X)                               # ✅

print("model tayyor")                                            # ⚠️
assert sum(p.numel() for p in blok.parameters()) == formula(d, d_ff)  # ✅

1.  # post-LN qism-blok formulasi?
2.  # pre-LN qism-blok formulasi?
3.  # pre-LN da oxirida nima qo'shiladi?
4.  # FFN formulasi va odatdagi d_ff?
5.  # d = 64, d_ff = 256 da blok parametrlari?
6.  # h = 4 dan h = 8 ga o'tsa parametrlar?
7.  # d_ff = 4d da bitta blok taxminan necha d^2?
8.  # residual siz 12 blokdan keyin tokenlar o'xshashligi?
9.  # dropout nechta joyda?
10. # dx_{l+1}/dx_l residual bilan?
11. # bir token uchun FFN FLOP?
12. # qaysi L da skor qismi FFN ga tenglashadi?

1.  def forward(self, x):                 # pre-LN bo'lishi kerak edi
        x = self.ln1(x)
        x = x + self.att(x)
        x = self.ln2(x)
        return x + self.ffn(x)

2.  qatlam = nn.TransformerEncoderLayer(64, 4, 256)
    y = qatlam(torch.randn(8, 20, 64))    # (B, L, d)

3.  x = self.drop(x + self.ffn(self.ln2(x)))

4.  stek = nn.TransformerEncoder(nn.TransformerEncoderLayer(
        64, 4, 256, batch_first=True, norm_first=True), 6)
    y = self.chiq(stek(x))

5.  p = model(X_test).argmax(-1)          # o'rgatishdan keyin darhol