24.11-dars: Samaradorlik va masshtab

BITTA BLOK FORWARD I, n token, d kenglik (B = 1), matmul FLOP = 2 * m * k * n:
  Q, K, V proyeksiyasi   2 * n * d * 3d  =  6 n d^2
  chiqish proyeksiyasi   2 * n * d * d   =  2 n d^2
  FFN (d -> 4d -> d)     2 * 2 * n * d * 4d = 16 n d^2
  -------------------------------------------- chiziqli qism: 24 n d^2
  skorlar Q K^T          2 * n * n * d   =  2 n^2 d
  A @ V                  2 * n * n * d   =  2 n^2 d
  -------------------------------------------- attention:     4 n^2 d

TENG BO'LADIGAN NUQTA: 4 n^2 d = 24 n d^2  ->  n = 6d
  d = 768 (BERT-base) -> n = 4608 token
  n < 6d: chiziqli qatlamlar qimmatroq;  n > 6d: attention ustunlik qiladi

XOTIRA (skor matritsasi, har qatlam): B * h * n^2 * 4 bayt (float32)
  1-misol, d=768, h=12:  n=512 -> 12.0 MB,  n=8192 -> 3.0 GB,  n=32768 -> 48.0 GB

ODDIY (qo'lda) ATTENTION:
  S = Q K^T / sqrt(d_h)        (B, h, n, n)  <- xotirada TO'LIQ
  P = softmax(S)               (B, h, n, n)  <- backward uchun SAQLANADI
  O = P V
  backward uchun xotira: O(n^2)

FLASH / MEMORY-EFFICIENT G'OYASI:
  Q, K, V ni bloklarga bo'lib, softmax ni "onlayn" hisoblash:
    har K-blok uchun: joriy maksimum va yig'indini yangilash
    (softmax(x) = exp(x - max) / sum - bo'laklab hisoblasa bo'ladi)
  n x n matritsa HECH QACHON to'liq yaratilmaydi
  backward uchun faqat O va logsumexp (har qator uchun 1 son) saqlanadi
  backward da kerakli bloklar QAYTA hisoblanadi
  FLOP bir xil (yoki biroz ko'p), xotira O(n), GPU da xotiraga murojaat kam -> tez

PYTORCH DA: F.scaled_dot_product_attention(q, k, v, attn_mask=..., is_causal=...)
  mavjud yadrolardan (flash, memory-efficient, oddiy) mosini o'zi tanlaydi
  CPU da ham flash uslubidagi yadro bor - 1-misol

1-MISOL (backward uchun saqlangan baytlar, h=4, d_h=16):
  n       qo'lda      SDPA     nisbat
  128    608.0 KB   130.0 KB     4.7
  2048   129.5 MB     2.0 MB    63.8
  n 2x oshganda: qo'lda 3.95x, SDPA 2.00x

KESHSIZ GENERATSIYA (24.9-darsdagi sodda sikl):
  qadam t: model(X[:, :k+t]) - BUTUN prefiks qaytadan
  har qadamda eski tokenlarning K, V si QAYTA hisoblanadi (natija o'sha-o'sha!)
  modeldan o'tgan pozitsiyalar: k + (k+1) + ... + (k+n-1)  ~  O(n^2)

KUZATUV: kauzal niqob tufayli eski tokenning K, V si yangi token qo'shilganda
         O'ZGARMAYDI -> ularni saqlab qo'ysa bo'ladi

KESHLI GENERATSIYA:
  1. prefiks bir marta:  K, V -> kesh (har blok uchun)
  2. har qadam: faqat YANGI token -> q, k, v
     k, v keshga qo'shiladi:  K = cat(K_kesh, k),  V = cat(V_kesh, v)
     attention: yangi q  x  butun K  -> (1 x t) skor, (t x t) emas
     pozitsiya kodlash: boshi = joriy uzunlik - 1  (yangi token pozitsiyasi!)
  modeldan o'tgan pozitsiyalar: k + (n - 1)  ~  O(n)

2-MISOL (prefiks k=10, 40 belgi):
  pozitsiyalar: keshsiz 1180, keshli 49;  matmul FLOP 24.7 barobar kam
  tokenlar aynan bir xil, logitlar farqi 7.6e-06

NARXI - XOTIRA: 2 (K va V) * L * n * d * bayt  (har ketma-ketlik uchun)
  GPT-2 small (L=12, d=768, n=1024, float16):   36 MB
  7B model    (L=32, d=4096, n=4096, float16): 2048 MB

BITTA BLOK (d kenglik, FFN 4d):
  qkv:   3d^2 + 3d        proj:  d^2 + d
  ffn1:  4d^2 + 4d        ffn2:  4d^2 + d        2 x LayerNorm: 4d
  jami:  12 d^2 + 13 d   ~  12 d^2

BUTUN MODEL:  L * (12 d^2 + 13 d)  +  V * d  (token emb)  +  T * d  (pozitsiya)
              + 2d (oxirgi LN)  +  d * V + V  (chiqish; bog'langan bo'lsa yo'q)

3-MISOL (sanalgan == formula, uchala konfiguratsiyada):
  V=27,   d=48,  L=2:  bloklar 56544,  12*L*d^2 = 55296  (xato 2.2%)
  V=1000, d=128, L=6:  bloklar 1189632, 12*L*d^2 = 1179648 (xato 0.8%)

MA'LUM MODELLAR (bloklar + token embedding):
  BERT-base   L=12 d=768:    84.9M + 23.4M = 108.4M   (e'lon qilingan ~110M)
  GPT-2 small L=12 d=768:    84.9M + 38.6M = 123.5M   (e'lon qilingan ~124M)
  GPT-2 XL    L=48 d=1600: 1474.6M + 80.4M = 1555.0M  (e'lon qilingan ~1.5B)

DA'VO (katta tadqiqotlar, N 10^6 dan 10^10 gacha):
  loss(N) ~ (N_c / N)^alfa     yetarli ma'lumot va hisob bo'lsa
  log-log grafikda TO'G'RI CHIZIQ, alfa ~ 0.07-0.1 (til modellari)
  shuningdek ma'lumot (D) va hisob (C ~ 6ND) bo'yicha ham
  amaliy xulosa: kichik modellarda o'lchab, kattasining loss ini bashorat qilish

BIZNING TAJRIBA (3-misol):
  belgi darajasidagi GPT, L=2, d = 8, 16, 24, 32, 48
  bir xil matn, bir xil 250 qadam, bir xil lr
     d    N (bloklar)   val loss
     8        1744       1.9800
    16        6560       1.6611
    24       14448       1.1540
    32       25408       0.8357
    48       56544       0.4737
  umumiy qiyalik -0.407 (SE 0.077)
  mahalliy qiyaliklar: -0.133, -0.461, -0.572, -0.709  <- TO'G'RI CHIZIQ EMAS

                bit   maks qiymat    eps (1 dan keyingi qadam)
  float32        32    3.4e38         1.2e-07
  float16        16    65504          9.8e-04
  bfloat16       16    3.4e38         7.8e-03

4-MISOL:  70000 -> float16: inf,  bfloat16: 70144.0
          1 + 0.001 -> float16: 1.0009766,  bfloat16: 1.0000000

ARALASH ANIQLIK (torch.autocast):
  parametrlar va optimizator holati - float32 ("master" nusxa)
  matmul va konvolyutsiyalar - float16 / bfloat16 (tez, aktivatsiyalar 2x kichik)
  softmax, loss, normlashtirish - float32 da qoladi (autocast o'zi hal qiladi)
  float16 + GradScaler: loss ni katta songa ko'paytirib gradient "yo'qolishini" oldini olish
  bfloat16: diapazon float32 niki -> scaler kerak emas; aniqlik pastligi o'rgatishda
            odatda muammo emas

4-MISOL (4 blokli model, bfloat16 autocast):
  loss 3.5084 (float32: 3.5082), saqlangan aktivatsiyalar 55% (baytlarda)

ODDIY BACKWARD: forward dagi har oraliq tenzor (aktivatsiya) saqlanadi
  xotira ~ L * (blokdagi aktivatsiyalar)

CHECKPOINT: blok ichidagi aktivatsiyalar SAQLANMAYDI, faqat blok KIRISHI
  backward da blok forward i QAYTA bajariladi, keyin gradient
  torch.utils.checkpoint.checkpoint(blok, x, use_reentrant=False)

4-MISOL (4 blok, B=8, T=64):
                 saqlangan tenzor      bayt      FLOP (forward + backward)
  oddiy               113          11310212      419364864
  checkpoint           17            723076      520028160
  + blok kirishlari                  393216
  xotira 9.9%, hisob 1.24x, gradientlar farqi 0.0

G'OYA: fine-tuning dagi vazn o'zgarishi dW "past rangli" deb faraz qilinadi
  W' = W + dW,   dW = (alfa / r) * B A,   B: (d_chiq x r),  A: (r x d_kir),  r << d
  W muzlatilgan; faqat A va B o'rgatiladi
  B = 0 dan boshlanadi -> boshida W' = W (model aynan pretrained holatda)

PARAMETRLAR (d x d qatlam uchun): d^2 o'rniga 2 r d
  bizning model (d=48, r=4, qkv + proj + ffn da): 6144 / 62428 = 9.8%
  katta modelda (d=4096, r=8): bir qatlamda 16.8M o'rniga 65536 - 0.4%

AFZALLIKLAR:
  optimizator holati faqat A, B uchun -> xotira keskin kam
  asl model o'zgarmaydi -> katastrofik unutish yo'q 24.10-bob, bitta asos + ko'p adapter
  inference: W' = W + k B A ni BIRLASHTIRISH -> qo'shimcha hisob yo'q

4-MISOL (sharhlar modeli -> klinika sharhlari, 100 qadam, 3 urug'):
  usul         o'rgatiladi   ulush   val loss   - to'liq    2*SE
  faqat bosh       1372       2.2%    1.3857    +0.9880    0.0100
  LoRA r=4         6144       9.8%    0.5700    +0.1723    0.0106
  to'liq FT       62428     100.0%    0.3977
  moslashdan oldin: 3.2853;  birlashtirishdan keyin chiqish farqi 8.1e-06

import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from torch.utils.flop_counter import FlopCounterMode

# FLOP ni sanash (vaqt emas)
fc = FlopCounterMode(display=False)
with fc:
    model(X)
fc.get_total_flops()                               # blok: 24 n d^2 + 4 n^2 d

# samarali attention (flash / memory-efficient)
o = F.scaled_dot_product_attention(q, k, v, attn_mask=maska, is_causal=True)

# KV-cache: yangi tokenning k, v si keshga qo'shiladi
k = torch.cat([kesh["k"], k], 2); v = torch.cat([kesh["v"], v], 2)
lg = model(t, keshlar, boshi=X.shape[1] - 1)       # faqat yangi token

# parametrlar
N = sum(p.numel() for p in model.parameters())     # ~ 12 L d^2 + V d + ...

# aralash aniqlik (CPU: bfloat16; GPU: float16 + torch.amp.GradScaler)
with torch.autocast("cpu", dtype=torch.bfloat16):
    loss = loss_fn()

# gradient checkpointing
x = checkpoint(blok, x, use_reentrant=False)

# LoRA
class LoRA(nn.Module):                             # W x + (alfa/r) B A x
    ...
yangi.weight += k * B @ A                          # birlashtirish

blok narxi: 24nd^2 (chiziqli) + 4n^2d (attention); teng nuqta n = 6d
SDPA: xotira O(n^2) -> O(n), natija bir xil
KV-cache: hisob O(n^2) -> O(n), xotira 2*L*n*d; natija aynan bir xil
N ~ 12*L*d^2 + embeddinglar; o'rgatish ~ 16 bayt/parametr, 6N FLOP/token
masshtablash qonuni: keng diapazon va mos byudjet kerak; kichik tajribada egri
bfloat16 autocast, checkpoint (bir xil grad), LoRA (kam parametr, unutishsiz)

"""Attention narxi: O(n^2) FLOP va xotira, formulani sanagich bilan tekshirish, SDPA."""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from torch.utils.flop_counter import FlopCounterMode

MAVZU = ["telefon", "kitob", "kafe", "mehmonxona", "avtobus", "dastur", "kurs", "film",
         "restoran", "kamera", "taksi", "televizor"]
IJOBIY = ["a'lo", "ajoyib", "zo'r", "yaxshi", "qulay", "chiroyli", "tez", "sifatli",
          "mazali", "toza", "arzon", "ishonchli"]
SALBIY = ["yomon", "sekin", "qimmat", "iflos", "noqulay", "xunuk", "sifatsiz", "bemaza",
          "eski", "zerikarli", "ishonchsiz", "buzuq"]
DUM = {1: ["hammaga tavsiya qilaman", "yana kelaman", "juda mamnunman"],
       0: ["hech kimga tavsiya qilmayman", "boshqa kelmayman", "vaqtim behuda ketdi"]}


def sharhlar(seed, n, mavzu=MAVZU, dum=DUM):
    """Sintetik sharhlar matni (belgi darajasida til modeli uchun)."""
    rng = np.random.default_rng(seed)
    qatorlar = []
    for _ in range(n):
        q = int(rng.integers(2))
        s = [mavzu[int(rng.integers(len(mavzu)))], ["juda", "ancha", ""][int(rng.integers(3))],
             (IJOBIY if q else SALBIY)[int(rng.integers(12))], ["ekan", "edi", ""][int(rng.integers(3))]]
        if rng.random() < 0.6:
            s.append(dum[q][int(rng.integers(3))])
        qatorlar.append(" ".join(" ".join(s).split()))
    return " . ".join(qatorlar) + " . "


def qolda_attention(q, k, v, kauzal=False):
    """Oddiy attention: (B, h, n, n) skor matritsasi xotirada TO'LIQ yaratiladi."""
    skor = q @ k.transpose(-1, -2) / math.sqrt(q.shape[-1])
    if kauzal:
        n, m = skor.shape[-2:]
        maska = torch.ones(n, m, dtype=torch.bool).tril(m - n)
        skor = skor.masked_fill(~maska, float("-inf"))
    return torch.softmax(skor, -1) @ v


class Blok(nn.Module):
    """Pre-LN decoder bloki; ixtiyoriy KV-cache (kesh lug'ati)."""

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

    def forward(self, x, kesh=None):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).view(B, T, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        if kesh is not None:
            if "k" in kesh:                          # oldingi tokenlarning K, V si
                k = torch.cat([kesh["k"], k], 2)
                v = torch.cat([kesh["v"], v], 2)
            kesh["k"], kesh["v"] = k, v
        a = qolda_attention(q, k, v, kauzal=True)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, d))
        return x + self.ffn(self.ln2(x))


class GPT(nn.Module):
    def __init__(self, V, d=48, L=2, h=4, T=64):
        super().__init__()
        self.emb, self.poz = nn.Embedding(V, d), nn.Embedding(T, d)
        self.bloklar = nn.ModuleList([Blok(d, h) for _ in range(L)])
        self.ln = nn.LayerNorm(d)
        self.chiqish = nn.Linear(d, V)

    def forward(self, X, keshlar=None, boshi=0, ckpt=False):
        x = self.emb(X) + self.poz(torch.arange(boshi, boshi + X.shape[1]))
        for i, b in enumerate(self.bloklar):
            if ckpt:
                x = checkpoint(b, x, use_reentrant=False)
            else:
                x = b(x, None if keshlar is None else keshlar[i])
        return self.chiqish(self.ln(x))


def batch(data, g, B=32, T=32):
    i = torch.randint(0, len(data) - T - 1, (B,), generator=g).tolist()
    return (torch.stack([data[j:j + T] for j in i]),
            torch.stack([data[j + 1:j + T + 1] for j in i]))


def orgat(model, data, qadam, lr=3e-3, seed=0, params=None):
    params = [p for p in model.parameters() if p.requires_grad] if params is None else params
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.01)
    g = torch.Generator().manual_seed(seed)
    V = model.chiqish.out_features
    for _ in range(qadam):
        xb, yb = batch(data, g)
        loss = F.cross_entropy(model(xb).reshape(-1, V), yb.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    return model


@torch.no_grad()
def val_loss(model, X, Y):
    return F.cross_entropy(model(X).reshape(-1, model.chiqish.out_features), Y.reshape(-1)).item()


def saqlangan_baytlar(fn):
    """Backward uchun saqlangan tenzorlar: (soni, baytlar)."""
    hajm = []

    def pack(t):
        hajm.append(t.numel() * t.element_size())
        return t

    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        natija = fn()
    return natija, len(hajm), sum(hajm)


def blok_flop(n, d):
    """Bitta blok forward i (B=1): chiziqli qismlar va attention matmullari."""
    chiziqli = 2 * n * d * 3 * d + 2 * n * d * d + 2 * (2 * n * d * 4 * d)   # qkv, proj, ffn
    attention = 2 * n * n * d + 2 * n * n * d                                # QK^T va A@V
    return chiziqli, attention


def hajm(b):
    for birlik, k in [("GB", 2 ** 30), ("MB", 2 ** 20), ("KB", 2 ** 10)]:
        if b >= k:
            return f"{b / k:.1f} {birlik}"
    return f"{b} B"


def main() -> None:
    torch.set_num_threads(1)
    print("=== 1. Formula: BERT-base o'lchamidagi bitta qatlam (d=768, h=12, B=1) ===")
    d, h = 768, 12
    print(f"  {'n':>6} {'chiziqli GFLOP':>15} {'attention GFLOP':>16} {'attention ulushi':>17} "
          f"{'skor matritsasi':>16}")
    for n in [128, 512, 2048, 8192, 32768]:
        c, a = blok_flop(n, d)
        print(f"  {n:>6} {c / 1e9:>15.1f} {a / 1e9:>16.1f} {a / (a + c):>17.1%} "
              f"{hajm(h * n * n * 4):>16}")
    print(f"  attention = chiziqli bo'ladigan nuqta: 4n^2 d = 24 n d^2 -> n = 6d = {6 * d}")
    print("  n 4 barobar oshsa: chiziqli qism 4x, attention 16x, skor matritsasi 16x")

    print("\n=== 2. Formulani FLOP sanagich bilan tekshirish (d=64, h=4) ===")
    torch.manual_seed(0)
    blok = Blok(64, 4)
    print(f"  {'n':>5} {'sanagich':>12} {'formula':>12} {'mos':>5}")
    for n in [32, 128, 512]:
        x = torch.randn(1, n, 64)
        fc = FlopCounterMode(display=False)
        with fc, torch.no_grad():
            blok(x)
        c, a = blok_flop(n, 64)
        print(f"  {n:>5} {fc.get_total_flops():>12} {c + a:>12} "
              f"{str(fc.get_total_flops() == c + a):>5}")

    print("\n=== 3. Backward uchun saqlanadigan xotira: qo'lda va SDPA (h=4, d_h=16) ===")
    print(f"  {'n':>5} {'qo`lda':>10} {'SDPA':>10} {'nisbat':>7}".replace("`", "'"))
    oldingi = None
    for n in [128, 256, 512, 1024, 2048]:
        q, k, v = [torch.randn(1, 4, n, 16, requires_grad=True) for _ in range(3)]
        _, _, b1 = saqlangan_baytlar(lambda: qolda_attention(q, k, v))
        _, _, b2 = saqlangan_baytlar(lambda: F.scaled_dot_product_attention(q, k, v))
        print(f"  {n:>5} {hajm(b1):>10} {hajm(b2):>10} {b1 / b2:>7.1f}")
        if oldingi is not None and n == 2048:
            print(f"  n 2x oshganda: qo'lda {b1 / oldingi[0]:.2f}x, SDPA {b2 / oldingi[1]:.2f}x")
        oldingi = (b1, b2)
    print("  SDPA (flash/memory-efficient yadro) n x n matritsani saqlamaydi:")
    print("  bloklab hisoblaydi va backward da qayta hisoblaydi")

    print("\n=== 4. Natija bir xilmi? ===")
    torch.manual_seed(1)
    q, k, v = [torch.randn(2, 4, 256, 16, requires_grad=True) for _ in range(3)]
    for kauzal in [False, True]:
        a1 = qolda_attention(q, k, v, kauzal)
        a2 = F.scaled_dot_product_attention(q, k, v, is_causal=kauzal)
        g1 = torch.autograd.grad(a1.square().sum(), (q, k, v))
        g2 = torch.autograd.grad(a2.square().sum(), (q, k, v))
        fg = max((x - y).abs().max().item() for x, y in zip(g1, g2))
        print(f"  kauzal={str(kauzal):<5}  chiqish farqi {(a1 - a2).abs().max().item():.1e}  "
              f"gradient farqi {fg:.1e}  allclose: {torch.allclose(a1, a2, atol=1e-5)}")
    print("  farq faqat float32 yaxlitlash darajasida - matematik jihatdan bir xil amal")


if __name__ == "__main__":
    main()

=== 1. Formula: BERT-base o'lchamidagi bitta qatlam (d=768, h=12, B=1) ===
       n  chiziqli GFLOP  attention GFLOP  attention ulushi  skor matritsasi
     128             1.8              0.1              2.7%         768.0 KB
     512             7.2              0.8             10.0%          12.0 MB
    2048            29.0             12.9             30.8%         192.0 MB
    8192           116.0            206.2             64.0%           3.0 GB
   32768           463.9           3298.5             87.7%          48.0 GB
  attention = chiziqli bo'ladigan nuqta: 4n^2 d = 24 n d^2 -> n = 6d = 4608
  n 4 barobar oshsa: chiziqli qism 4x, attention 16x, skor matritsasi 16x

=== 2. Formulani FLOP sanagich bilan tekshirish (d=64, h=4) ===
      n     sanagich      formula   mos
     32      3407872      3407872  True
    128     16777216     16777216  True
    512    117440512    117440512  True

=== 3. Backward uchun saqlanadigan xotira: qo'lda va SDPA (h=4, d_h=16) ===
      n     qo'lda       SDPA  nisbat
    128   608.0 KB   130.0 KB     4.7
    256     2.2 MB   260.0 KB     8.6
    512     8.4 MB   520.0 KB    16.5
   1024    32.8 MB     1.0 MB    32.2
   2048   129.5 MB     2.0 MB    63.8
  n 2x oshganda: qo'lda 3.95x, SDPA 2.00x
  SDPA (flash/memory-efficient yadro) n x n matritsani saqlamaydi:
  bloklab hisoblaydi va backward da qayta hisoblaydi

=== 4. Natija bir xilmi? ===
  kauzal=False  chiqish farqi 5.4e-07  gradient farqi 6.0e-07  allclose: True
  kauzal=True   chiqish farqi 4.8e-07  gradient farqi 4.8e-06  allclose: True
  farq faqat float32 yaxlitlash darajasida - matematik jihatdan bir xil amal

"""KV-cache noldan: keshsiz va keshli generatsiya - natija bir xil, hisob keskin kam."""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from torch.utils.flop_counter import FlopCounterMode

MAVZU = ["telefon", "kitob", "kafe", "mehmonxona", "avtobus", "dastur", "kurs", "film",
         "restoran", "kamera", "taksi", "televizor"]
IJOBIY = ["a'lo", "ajoyib", "zo'r", "yaxshi", "qulay", "chiroyli", "tez", "sifatli",
          "mazali", "toza", "arzon", "ishonchli"]
SALBIY = ["yomon", "sekin", "qimmat", "iflos", "noqulay", "xunuk", "sifatsiz", "bemaza",
          "eski", "zerikarli", "ishonchsiz", "buzuq"]
DUM = {1: ["hammaga tavsiya qilaman", "yana kelaman", "juda mamnunman"],
       0: ["hech kimga tavsiya qilmayman", "boshqa kelmayman", "vaqtim behuda ketdi"]}


def sharhlar(seed, n, mavzu=MAVZU, dum=DUM):
    """Sintetik sharhlar matni (belgi darajasida til modeli uchun)."""
    rng = np.random.default_rng(seed)
    qatorlar = []
    for _ in range(n):
        q = int(rng.integers(2))
        s = [mavzu[int(rng.integers(len(mavzu)))], ["juda", "ancha", ""][int(rng.integers(3))],
             (IJOBIY if q else SALBIY)[int(rng.integers(12))], ["ekan", "edi", ""][int(rng.integers(3))]]
        if rng.random() < 0.6:
            s.append(dum[q][int(rng.integers(3))])
        qatorlar.append(" ".join(" ".join(s).split()))
    return " . ".join(qatorlar) + " . "


def qolda_attention(q, k, v, kauzal=False):
    """Oddiy attention: (B, h, n, n) skor matritsasi xotirada TO'LIQ yaratiladi."""
    skor = q @ k.transpose(-1, -2) / math.sqrt(q.shape[-1])
    if kauzal:
        n, m = skor.shape[-2:]
        maska = torch.ones(n, m, dtype=torch.bool).tril(m - n)
        skor = skor.masked_fill(~maska, float("-inf"))
    return torch.softmax(skor, -1) @ v


class Blok(nn.Module):
    """Pre-LN decoder bloki; ixtiyoriy KV-cache (kesh lug'ati)."""

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

    def forward(self, x, kesh=None):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).view(B, T, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        if kesh is not None:
            if "k" in kesh:                          # oldingi tokenlarning K, V si
                k = torch.cat([kesh["k"], k], 2)
                v = torch.cat([kesh["v"], v], 2)
            kesh["k"], kesh["v"] = k, v
        a = qolda_attention(q, k, v, kauzal=True)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, d))
        return x + self.ffn(self.ln2(x))


class GPT(nn.Module):
    def __init__(self, V, d=48, L=2, h=4, T=64):
        super().__init__()
        self.emb, self.poz = nn.Embedding(V, d), nn.Embedding(T, d)
        self.bloklar = nn.ModuleList([Blok(d, h) for _ in range(L)])
        self.ln = nn.LayerNorm(d)
        self.chiqish = nn.Linear(d, V)

    def forward(self, X, keshlar=None, boshi=0, ckpt=False):
        x = self.emb(X) + self.poz(torch.arange(boshi, boshi + X.shape[1]))
        for i, b in enumerate(self.bloklar):
            if ckpt:
                x = checkpoint(b, x, use_reentrant=False)
            else:
                x = b(x, None if keshlar is None else keshlar[i])
        return self.chiqish(self.ln(x))


def batch(data, g, B=32, T=32):
    i = torch.randint(0, len(data) - T - 1, (B,), generator=g).tolist()
    return (torch.stack([data[j:j + T] for j in i]),
            torch.stack([data[j + 1:j + T + 1] for j in i]))


def orgat(model, data, qadam, lr=3e-3, seed=0, params=None):
    params = [p for p in model.parameters() if p.requires_grad] if params is None else params
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.01)
    g = torch.Generator().manual_seed(seed)
    V = model.chiqish.out_features
    for _ in range(qadam):
        xb, yb = batch(data, g)
        loss = F.cross_entropy(model(xb).reshape(-1, V), yb.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    return model


@torch.no_grad()
def val_loss(model, X, Y):
    return F.cross_entropy(model(X).reshape(-1, model.chiqish.out_features), Y.reshape(-1)).item()


def saqlangan_baytlar(fn):
    """Backward uchun saqlangan tenzorlar: (soni, baytlar)."""
    hajm = []

    def pack(t):
        hajm.append(t.numel() * t.element_size())
        return t

    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        natija = fn()
    return natija, len(hajm), sum(hajm)


@torch.no_grad()
def keshsiz(model, prefiks, n):
    """Har qadamda BUTUN ketma-ketlik qaytadan o'tkaziladi."""
    X, logitlar, pozitsiyalar = prefiks.clone(), [], 0
    for _ in range(n):
        pozitsiyalar += X.shape[1]
        lg = model(X)[:, -1]
        logitlar.append(lg)
        X = torch.cat([X, lg.argmax(-1, keepdim=True)], 1)
    return X, torch.stack(logitlar, 1), pozitsiyalar


@torch.no_grad()
def keshli(model, prefiks, n):
    """Prefiks bir marta; keyin har qadamda faqat YANGI token, K va V keshdan."""
    keshlar = [{} for _ in model.bloklar]
    lg = model(prefiks, keshlar)[:, -1]
    X, logitlar, pozitsiyalar = prefiks.clone(), [], prefiks.shape[1]
    for _ in range(n):
        logitlar.append(lg)
        t = lg.argmax(-1, keepdim=True)
        X = torch.cat([X, t], 1)
        if len(logitlar) < n:
            pozitsiyalar += 1
            lg = model(t, keshlar, boshi=X.shape[1] - 1)[:, -1]
    return X, torch.stack(logitlar, 1), pozitsiyalar, keshlar


def main() -> None:
    torch.set_num_threads(1)
    matn = sharhlar(0, 4000)
    belgilar = sorted(set(matn))
    s2i = {c: i for i, c in enumerate(belgilar)}
    data = torch.tensor([s2i[c] for c in matn])

    print("=== 1. Kichik GPT (belgi darajasida) ===")
    torch.manual_seed(0)
    model = orgat(GPT(len(belgilar), d=48, L=2, h=4, T=64), data, 400)
    print(f"  matn {len(data)} belgi, lug'at {len(belgilar)}, "
          f"parametrlar {sum(p.numel() for p in model.parameters())}")

    print("\n=== 2. Generatsiya: keshsiz va keshli (greedy, 40 belgi) ===")
    prefiks = torch.tensor([[s2i[c] for c in "kafe juda "]])
    n = 40
    X1, L1, p1 = keshsiz(model, prefiks, n)
    X2, L2, p2, keshlar = keshli(model, prefiks, n)
    print(f"  keshsiz: {''.join(belgilar[i] for i in X1[0].tolist())!r}")
    print(f"  keshli:  {''.join(belgilar[i] for i in X2[0].tolist())!r}")
    print(f"  tokenlar aynan bir xil: {torch.equal(X1, X2)}")
    print(f"  logitlar max farqi: {(L1 - L2).abs().max().item():.1e}")

    print("\n=== 3. Qancha hisob bajarildi ===")
    natija = {}
    for nom, fn in [("keshsiz", keshsiz), ("keshli", keshli)]:
        fc = FlopCounterMode(display=False)
        with fc:
            fn(model, prefiks, n)
        natija[nom] = fc.get_total_flops()
    print(f"  {'':<9} {'modeldan o`tgan pozitsiyalar':>29} {'matmul FLOP':>13}".replace("`", "'"))
    print(f"  {'keshsiz':<9} {p1:>29} {natija['keshsiz']:>13}")
    print(f"  {'keshli':<9} {p2:>29} {natija['keshli']:>13}")
    print(f"  tejash: {natija['keshsiz'] / natija['keshli']:.1f} barobar kam FLOP")
    k = prefiks.shape[1]
    print(f"  formula: keshsiz sum(k..k+n-1) = {sum(range(k, k + n))}, "
          f"keshli k + n - 1 = {k + n - 1}")

    print("\n=== 4. Keshning o'zi qancha xotira oladi ===")
    kv = sum(t.numel() * t.element_size() for c in keshlar for t in c.values())
    print(f"  bizning model: {len(keshlar)} blok x (K, V) x {tuple(keshlar[0]['k'].shape)} "
          f"= {kv} bayt")
    print(f"  formula 2 * L * n * d * 4 bayt = {2 * 2 * keshlar[0]['k'].shape[2] * 48 * 4}")
    for nom, L, dm, n_ in [("GPT-2 small", 12, 768, 1024), ("7B model", 32, 4096, 4096)]:
        b = 2 * L * n_ * dm * 2
        print(f"  {nom:<12} L={L}, d={dm}, n={n_}, float16: {b / 2 ** 20:.0f} MB "
              "bitta ketma-ketlik uchun")
    print("  kesh FLOP ni kamaytiradi, lekin xotirani n ga chiziqli oshiradi")


if __name__ == "__main__":
    main()

=== 1. Kichik GPT (belgi darajasida) ===
  matn 138379 belgi, lug'at 27, parametrlar 62331

=== 2. Generatsiya: keshsiz va keshli (greedy, 40 belgi) ===
  keshsiz: "kafe juda a'lo edi . kamera ancha jules a a a ajud"
  keshli:  "kafe juda a'lo edi . kamera ancha jules a a a ajud"
  tokenlar aynan bir xil: True
  logitlar max farqi: 7.6e-06

=== 3. Qancha hisob bajarildi ===
             modeldan o'tgan pozitsiyalar   matmul FLOP
  keshsiz                            1180     148970880
  keshli                               49       6033696
  tejash: 24.7 barobar kam FLOP
  formula: keshsiz sum(k..k+n-1) = 1180, keshli k + n - 1 = 49

=== 4. Keshning o'zi qancha xotira oladi ===
  bizning model: 2 blok x (K, V) x (1, 4, 49, 12) = 37632 bayt
  formula 2 * L * n * d * 4 bayt = 37632
  GPT-2 small  L=12, d=768, n=1024, float16: 36 MB bitta ketma-ketlik uchun
  7B model     L=32, d=4096, n=4096, float16: 2048 MB bitta ketma-ketlik uchun
  kesh FLOP ni kamaytiradi, lekin xotirani n ga chiziqli oshiradi

"""Parametr soni formulasi (12*L*d^2) va kichik modellar oilasida masshtablash."""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from torch.utils.flop_counter import FlopCounterMode

MAVZU = ["telefon", "kitob", "kafe", "mehmonxona", "avtobus", "dastur", "kurs", "film",
         "restoran", "kamera", "taksi", "televizor"]
IJOBIY = ["a'lo", "ajoyib", "zo'r", "yaxshi", "qulay", "chiroyli", "tez", "sifatli",
          "mazali", "toza", "arzon", "ishonchli"]
SALBIY = ["yomon", "sekin", "qimmat", "iflos", "noqulay", "xunuk", "sifatsiz", "bemaza",
          "eski", "zerikarli", "ishonchsiz", "buzuq"]
DUM = {1: ["hammaga tavsiya qilaman", "yana kelaman", "juda mamnunman"],
       0: ["hech kimga tavsiya qilmayman", "boshqa kelmayman", "vaqtim behuda ketdi"]}


def sharhlar(seed, n, mavzu=MAVZU, dum=DUM):
    """Sintetik sharhlar matni (belgi darajasida til modeli uchun)."""
    rng = np.random.default_rng(seed)
    qatorlar = []
    for _ in range(n):
        q = int(rng.integers(2))
        s = [mavzu[int(rng.integers(len(mavzu)))], ["juda", "ancha", ""][int(rng.integers(3))],
             (IJOBIY if q else SALBIY)[int(rng.integers(12))], ["ekan", "edi", ""][int(rng.integers(3))]]
        if rng.random() < 0.6:
            s.append(dum[q][int(rng.integers(3))])
        qatorlar.append(" ".join(" ".join(s).split()))
    return " . ".join(qatorlar) + " . "


def qolda_attention(q, k, v, kauzal=False):
    """Oddiy attention: (B, h, n, n) skor matritsasi xotirada TO'LIQ yaratiladi."""
    skor = q @ k.transpose(-1, -2) / math.sqrt(q.shape[-1])
    if kauzal:
        n, m = skor.shape[-2:]
        maska = torch.ones(n, m, dtype=torch.bool).tril(m - n)
        skor = skor.masked_fill(~maska, float("-inf"))
    return torch.softmax(skor, -1) @ v


class Blok(nn.Module):
    """Pre-LN decoder bloki; ixtiyoriy KV-cache (kesh lug'ati)."""

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

    def forward(self, x, kesh=None):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).view(B, T, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        if kesh is not None:
            if "k" in kesh:                          # oldingi tokenlarning K, V si
                k = torch.cat([kesh["k"], k], 2)
                v = torch.cat([kesh["v"], v], 2)
            kesh["k"], kesh["v"] = k, v
        a = qolda_attention(q, k, v, kauzal=True)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, d))
        return x + self.ffn(self.ln2(x))


class GPT(nn.Module):
    def __init__(self, V, d=48, L=2, h=4, T=64):
        super().__init__()
        self.emb, self.poz = nn.Embedding(V, d), nn.Embedding(T, d)
        self.bloklar = nn.ModuleList([Blok(d, h) for _ in range(L)])
        self.ln = nn.LayerNorm(d)
        self.chiqish = nn.Linear(d, V)

    def forward(self, X, keshlar=None, boshi=0, ckpt=False):
        x = self.emb(X) + self.poz(torch.arange(boshi, boshi + X.shape[1]))
        for i, b in enumerate(self.bloklar):
            if ckpt:
                x = checkpoint(b, x, use_reentrant=False)
            else:
                x = b(x, None if keshlar is None else keshlar[i])
        return self.chiqish(self.ln(x))


def batch(data, g, B=32, T=32):
    i = torch.randint(0, len(data) - T - 1, (B,), generator=g).tolist()
    return (torch.stack([data[j:j + T] for j in i]),
            torch.stack([data[j + 1:j + T + 1] for j in i]))


def orgat(model, data, qadam, lr=3e-3, seed=0, params=None):
    params = [p for p in model.parameters() if p.requires_grad] if params is None else params
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.01)
    g = torch.Generator().manual_seed(seed)
    V = model.chiqish.out_features
    for _ in range(qadam):
        xb, yb = batch(data, g)
        loss = F.cross_entropy(model(xb).reshape(-1, V), yb.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    return model


@torch.no_grad()
def val_loss(model, X, Y):
    return F.cross_entropy(model(X).reshape(-1, model.chiqish.out_features), Y.reshape(-1)).item()


def saqlangan_baytlar(fn):
    """Backward uchun saqlangan tenzorlar: (soni, baytlar)."""
    hajm = []

    def pack(t):
        hajm.append(t.numel() * t.element_size())
        return t

    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        natija = fn()
    return natija, len(hajm), sum(hajm)


def formula(V, d, L, T):
    """Aniq hisob: bloklar + embeddinglar + oxirgi LayerNorm + chiqish qatlami."""
    blok = (3 * d * d + 3 * d) + (d * d + d) + (4 * d * d + 4 * d) + (4 * d * d + d) + 4 * d
    return L * blok + V * d + T * d + 2 * d + (d * V + V), L * blok


def main() -> None:
    torch.set_num_threads(1)
    print("=== 1. Parametr soni: sanash va formula ===")
    print(f"  {'V':>4} {'d':>4} {'L':>3} {'sanalgan':>10} {'formula':>10} {'bloklar':>9} "
          f"{'12*L*d^2':>9} {'xato':>6}")
    for V, d, L in [(27, 48, 2), (27, 96, 4), (1000, 128, 6)]:
        m = GPT(V, d=d, L=L, h=4, T=64)
        jami, blok = formula(V, d, L, 64)
        sanalgan = sum(p.numel() for p in m.parameters())
        print(f"  {V:>4} {d:>4} {L:>3} {sanalgan:>10} {jami:>10} {blok:>9} {12 * L * d * d:>9} "
              f"{(blok - 12 * L * d * d) / blok:>6.1%}")
    print("  bitta blok: attention 4d^2 (Q, K, V, O) + FFN 8d^2 (d -> 4d -> d) = 12d^2")
    print("\n  Ma'lum modellar (bloklar 12*L*d^2 + token embedding V*d):")
    for nom, L, d, V in [("BERT-base", 12, 768, 30522), ("GPT-2 small", 12, 768, 50257),
                         ("GPT-2 XL", 48, 1600, 50257)]:
        b = 12 * L * d * d
        print(f"    {nom:<12} L={L:>2} d={d:>4}: bloklar {b / 1e6:>6.1f}M + embedding "
              f"{V * d / 1e6:>5.1f}M = {(b + V * d) / 1e6:>6.1f}M")

    print("\n=== 2. Masshtablash: bir xil ma'lumot va qadamlar, d o'zgaradi (L=2) ===")
    matn = sharhlar(0, 6000)
    belgilar = sorted(set(matn))
    s2i = {c: i for i, c in enumerate(belgilar)}
    data = torch.tensor([s2i[c] for c in matn])
    n = int(len(data) * 0.9)
    oquv, val = data[:n], data[n:]
    Xv, Yv = batch(val, torch.Generator().manual_seed(123), B=256)
    N, loss = [], []
    print(f"  {'d':>4} {'N (bloklar)':>12} {'val loss':>9}")
    for d in [8, 16, 24, 32, 48]:
        torch.manual_seed(0)
        m = orgat(GPT(len(belgilar), d=d, L=2, h=4, T=64), oquv, 250)
        N.append(formula(len(belgilar), d, 2, 64)[1])
        loss.append(val_loss(m, Xv, Yv))
        print(f"  {d:>4} {N[-1]:>12} {loss[-1]:>9.4f}")

    print("\n=== 3. Log-log qiyalik: loss ~ N^(-alfa) ===")
    lx, ly = np.log(N), np.log(loss)
    (k, b), cov = np.polyfit(lx, ly, 1, cov=True)
    print(f"  umumiy qiyalik: {k:.3f} (SE {np.sqrt(cov[0, 0]):.3f})  -> alfa = {-k:.3f}")
    qiya = np.diff(ly) / np.diff(lx)
    for i in range(len(qiya)):
        print(f"    N {N[i]:>6} -> {N[i + 1]:>6}: mahalliy qiyalik {qiya[i]:+.3f}")
    qoldiq = ly - (k * lx + b)
    print(f"  to'g'ri chiziqdan eng katta chetlanish: {np.abs(qoldiq).max():.3f} (log birlikda)")
    if np.ptp(qiya) > 0.2:
        print("  mahalliy qiyaliklar bir-biridan ancha farq qiladi - bu diapazonda")
        print("  toza darajali qonun yo'q: kichik modellar va qat'iy qadamlar soni")
    print("  eslatma: katta tadqiqotlarda N 6-7 tartibga o'zgaradi, bu yerda 1.5 tartib")


if __name__ == "__main__":
    main()

=== 1. Parametr soni: sanash va formula ===
     V    d   L   sanalgan    formula   bloklar  12*L*d^2   xato
    27   48   2      62331      62331     56544     55296   2.2%
    27   96   4     458907     458907    447360    442368   1.1%
  1000  128   6    1455080    1455080   1189632   1179648   0.8%
  bitta blok: attention 4d^2 (Q, K, V, O) + FFN 8d^2 (d -> 4d -> d) = 12d^2

  Ma'lum modellar (bloklar 12*L*d^2 + token embedding V*d):
    BERT-base    L=12 d= 768: bloklar   84.9M + embedding  23.4M =  108.4M
    GPT-2 small  L=12 d= 768: bloklar   84.9M + embedding  38.6M =  123.5M
    GPT-2 XL     L=48 d=1600: bloklar 1474.6M + embedding  80.4M = 1555.0M

=== 2. Masshtablash: bir xil ma'lumot va qadamlar, d o'zgaradi (L=2) ===
     d  N (bloklar)  val loss
     8         1744    1.9800
    16         6560    1.6611
    24        14448    1.1540
    32        25408    0.8357
    48        56544    0.4737

=== 3. Log-log qiyalik: loss ~ N^(-alfa) ===
  umumiy qiyalik: -0.407 (SE 0.077)  -> alfa = 0.407
    N   1744 ->   6560: mahalliy qiyalik -0.133
    N   6560 ->  14448: mahalliy qiyalik -0.461
    N  14448 ->  25408: mahalliy qiyalik -0.572
    N  25408 ->  56544: mahalliy qiyalik -0.709
  to'g'ri chiziqdan eng katta chetlanish: 0.193 (log birlikda)
  mahalliy qiyaliklar bir-biridan ancha farq qiladi - bu diapazonda
  toza darajali qonun yo'q: kichik modellar va qat'iy qadamlar soni
  eslatma: katta tadqiqotlarda N 6-7 tartibga o'zgaradi, bu yerda 1.5 tartib

"""Aralash aniqlik, gradient checkpointing va LoRA noldan."""

import copy
import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from torch.utils.flop_counter import FlopCounterMode

MAVZU = ["telefon", "kitob", "kafe", "mehmonxona", "avtobus", "dastur", "kurs", "film",
         "restoran", "kamera", "taksi", "televizor"]
IJOBIY = ["a'lo", "ajoyib", "zo'r", "yaxshi", "qulay", "chiroyli", "tez", "sifatli",
          "mazali", "toza", "arzon", "ishonchli"]
SALBIY = ["yomon", "sekin", "qimmat", "iflos", "noqulay", "xunuk", "sifatsiz", "bemaza",
          "eski", "zerikarli", "ishonchsiz", "buzuq"]
DUM = {1: ["hammaga tavsiya qilaman", "yana kelaman", "juda mamnunman"],
       0: ["hech kimga tavsiya qilmayman", "boshqa kelmayman", "vaqtim behuda ketdi"]}


def sharhlar(seed, n, mavzu=MAVZU, dum=DUM):
    """Sintetik sharhlar matni (belgi darajasida til modeli uchun)."""
    rng = np.random.default_rng(seed)
    qatorlar = []
    for _ in range(n):
        q = int(rng.integers(2))
        s = [mavzu[int(rng.integers(len(mavzu)))], ["juda", "ancha", ""][int(rng.integers(3))],
             (IJOBIY if q else SALBIY)[int(rng.integers(12))], ["ekan", "edi", ""][int(rng.integers(3))]]
        if rng.random() < 0.6:
            s.append(dum[q][int(rng.integers(3))])
        qatorlar.append(" ".join(" ".join(s).split()))
    return " . ".join(qatorlar) + " . "


def qolda_attention(q, k, v, kauzal=False):
    """Oddiy attention: (B, h, n, n) skor matritsasi xotirada TO'LIQ yaratiladi."""
    skor = q @ k.transpose(-1, -2) / math.sqrt(q.shape[-1])
    if kauzal:
        n, m = skor.shape[-2:]
        maska = torch.ones(n, m, dtype=torch.bool).tril(m - n)
        skor = skor.masked_fill(~maska, float("-inf"))
    return torch.softmax(skor, -1) @ v


class Blok(nn.Module):
    """Pre-LN decoder bloki; ixtiyoriy KV-cache (kesh lug'ati)."""

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

    def forward(self, x, kesh=None):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).view(B, T, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        if kesh is not None:
            if "k" in kesh:                          # oldingi tokenlarning K, V si
                k = torch.cat([kesh["k"], k], 2)
                v = torch.cat([kesh["v"], v], 2)
            kesh["k"], kesh["v"] = k, v
        a = qolda_attention(q, k, v, kauzal=True)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, d))
        return x + self.ffn(self.ln2(x))


class GPT(nn.Module):
    def __init__(self, V, d=48, L=2, h=4, T=64):
        super().__init__()
        self.emb, self.poz = nn.Embedding(V, d), nn.Embedding(T, d)
        self.bloklar = nn.ModuleList([Blok(d, h) for _ in range(L)])
        self.ln = nn.LayerNorm(d)
        self.chiqish = nn.Linear(d, V)

    def forward(self, X, keshlar=None, boshi=0, ckpt=False):
        x = self.emb(X) + self.poz(torch.arange(boshi, boshi + X.shape[1]))
        for i, b in enumerate(self.bloklar):
            if ckpt:
                x = checkpoint(b, x, use_reentrant=False)
            else:
                x = b(x, None if keshlar is None else keshlar[i])
        return self.chiqish(self.ln(x))


def batch(data, g, B=32, T=32):
    i = torch.randint(0, len(data) - T - 1, (B,), generator=g).tolist()
    return (torch.stack([data[j:j + T] for j in i]),
            torch.stack([data[j + 1:j + T + 1] for j in i]))


def orgat(model, data, qadam, lr=3e-3, seed=0, params=None):
    params = [p for p in model.parameters() if p.requires_grad] if params is None else params
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.01)
    g = torch.Generator().manual_seed(seed)
    V = model.chiqish.out_features
    for _ in range(qadam):
        xb, yb = batch(data, g)
        loss = F.cross_entropy(model(xb).reshape(-1, V), yb.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    return model


@torch.no_grad()
def val_loss(model, X, Y):
    return F.cross_entropy(model(X).reshape(-1, model.chiqish.out_features), Y.reshape(-1)).item()


def saqlangan_baytlar(fn):
    """Backward uchun saqlangan tenzorlar: (soni, baytlar)."""
    hajm = []

    def pack(t):
        hajm.append(t.numel() * t.element_size())
        return t

    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        natija = fn()
    return natija, len(hajm), sum(hajm)


class LoRA(nn.Module):
    """W x + (alfa/r) * B A x;  W muzlatilgan, A va B o'rgatiladi, B = 0 dan boshlanadi."""

    def __init__(self, asl, r=4, alfa=8):
        super().__init__()
        self.asl = asl
        for p in self.asl.parameters():
            p.requires_grad_(False)
        self.A = nn.Parameter(torch.randn(r, asl.in_features) / math.sqrt(asl.in_features))
        self.B = nn.Parameter(torch.zeros(asl.out_features, r))
        self.k = alfa / r

    def forward(self, x):
        return self.asl(x) + (x @ self.A.T @ self.B.T) * self.k

    def birlashtir(self):
        """Inference uchun: W' = W + k * B A - qo'shimcha hisobsiz oddiy Linear."""
        yangi = copy.deepcopy(self.asl)
        with torch.no_grad():
            yangi.weight += self.k * self.B @ self.A
        return yangi


def lora_qosh(model, r=4):
    for p in model.parameters():
        p.requires_grad_(False)
    for b in model.bloklar:
        b.qkv, b.proj = LoRA(b.qkv, r), LoRA(b.proj, r)
        b.ffn[0], b.ffn[2] = LoRA(b.ffn[0], r), LoRA(b.ffn[2], r)
    return model


def main() -> None:
    torch.set_num_threads(1)
    print("=== 1. Aralash aniqlik: float16 va bfloat16 ===")
    for nom, dt in [("float32", torch.float32), ("float16", torch.float16),
                    ("bfloat16", torch.bfloat16)]:
        katta = torch.tensor(70000.0).to(dt).item()
        kichik = (torch.tensor(1.0, dtype=dt) + torch.tensor(1e-3, dtype=dt)).item()
        print(f"  {nom:<9} {torch.finfo(dt).bits:>2} bit  70000 -> {katta:<8}  "
              f"1 + 0.001 = {kichik:<10.7f}  eps {torch.finfo(dt).eps:.1e}")
    print("  float16: diapazon tor (maks 65504) -> gradient scaling kerak")
    print("  bfloat16: diapazon float32 niki, aniqlik past -> scaling kerak emas")

    matn = sharhlar(0, 5000)
    shifo = ["shifoxona", "poliklinika", "dorixona", "shifokor", "hamshira", "laboratoriya"]
    shifo_dum = {1: ["davolanishdan mamnunman", "tez tuzaldim", "rahmat shifokorlarga"],
                 0: ["davolanish foyda bermadi", "soatlab kutdim", "boshqa bormayman"]}
    yangi_matn = sharhlar(1, 3000, shifo, shifo_dum)
    belgilar = sorted(set(matn + yangi_matn))
    s2i = {c: i for i, c in enumerate(belgilar)}
    A = torch.tensor([s2i[c] for c in matn])
    B_ = torch.tensor([s2i[c] for c in yangi_matn])
    V = len(belgilar)
    torch.manual_seed(0)
    model = GPT(V, d=48, L=4, h=4, T=64)
    xb, yb = batch(A, torch.Generator().manual_seed(0), B=8, T=64)

    def loss_fn(ckpt=False):
        return F.cross_entropy(model(xb, ckpt=ckpt).reshape(-1, V), yb.reshape(-1))

    l32, _, b32 = saqlangan_baytlar(loss_fn)
    with torch.autocast("cpu", dtype=torch.bfloat16):
        l16, _, b16 = saqlangan_baytlar(loss_fn)
    print(f"  autocast(bfloat16): loss {l16.item():.4f} (float32: {l32.item():.4f}), "
          f"saqlangan aktivatsiyalar {b16 / b32:.0%} (baytlarda)")
    print(f"  parametrlar turi autocast ichida ham: {next(model.parameters()).dtype}")

    print("\n=== 2. Gradient checkpointing (4 blok, B=8, T=64) ===")
    natija = {}
    for ckpt in [False, True]:
        model.zero_grad()
        fc = FlopCounterMode(display=False)
        with fc:
            loss, soni, bayt = saqlangan_baytlar(lambda: loss_fn(ckpt))
            loss.backward()
        natija[ckpt] = (loss.item(), soni, bayt, fc.get_total_flops(),
                        [p.grad.clone() for p in model.parameters()])
    print(f"  {'':<14} {'loss':>8} {'saqlangan tenzor':>17} {'bayt':>10} {'FLOP (f+b)':>12}")
    for ckpt, nom in [(False, "oddiy"), (True, "checkpoint")]:
        l, s, b, f, _ = natija[ckpt]
        print(f"  {nom:<14} {l:>8.4f} {s:>17} {b:>10} {f:>12}")
    farq = max((a - b).abs().max().item() for a, b in zip(natija[False][4], natija[True][4]))
    print(f"  gradientlar max farqi: {farq}")
    kirishlar = len(model.bloklar) * xb.numel() * 48 * 4       # har blok kirishi (B, T, d)
    jami = natija[True][2] + kirishlar
    print(f"  checkpoint har blok KIRISHINI ham ushlab turadi: {kirishlar} bayt")
    print(f"  xotira {jami} / {natija[False][2]} = {jami / natija[False][2]:.1%}, "
          f"hisob {natija[True][3] / natija[False][3]:.2f}x (forward qayta hisoblanadi)")

    print("\n=== 3. LoRA: 'sharhlar' da o'rgatilgan modelni 'klinika sharhlari' ga moslash ===")
    torch.manual_seed(0)
    asos = orgat(GPT(V, d=48, L=2, h=4, T=64), A, 300)
    n = int(len(B_) * 0.8)
    Xv, Yv = batch(B_[n:], torch.Generator().manual_seed(7), B=256)
    print(f"  moslashdan oldin val loss (yangi domen): {val_loss(asos, Xv, Yv):.4f}")
    usullar = [("faqat bosh", 3e-3), ("LoRA r=4", 1e-2), ("to'liq FT", 3e-3)]
    natija = {}
    for s in range(3):
        for nom, lr in usullar:
            torch.manual_seed(s)
            m = copy.deepcopy(asos)
            if nom == "LoRA r=4":
                lora_qosh(m)
            elif nom == "faqat bosh":
                for p in m.parameters():
                    p.requires_grad_(False)
                for p in m.chiqish.parameters():
                    p.requires_grad_(True)
            orgat(m, B_[:n], 100, lr=lr, seed=s)
            natija[(nom, s)] = val_loss(m, Xv, Yv)
            if s == 0:
                natija[nom] = (sum(p.numel() for p in m.parameters() if p.requires_grad),
                               sum(p.numel() for p in asos.parameters()), m)
    toliq = np.array([natija[("to'liq FT", s)] for s in range(3)])
    print(f"  {'usul':<11} {'o`rgatiladi':>11} {'ulush':>7} {'val loss':>9} {'- to`liq':>9} "
          f"{'2*SE':>7}".replace("`", "'"))
    for nom, _ in usullar:
        a = np.array([natija[(nom, s)] for s in range(3)])
        d = a - toliq
        se = d.std(ddof=1) / np.sqrt(3)
        o, j, _ = natija[nom]
        qator = f"  {nom:<11} {o:>11} {o / j:>7.1%} {a.mean():>9.4f}"
        if nom != "to'liq FT":
            qator += f" {d.mean():>+9.4f} {2 * se:>7.4f}"
        print(qator)
    lora = np.array([natija[("LoRA r=4", s)] for s in range(3)])
    bosh = np.array([natija[("faqat bosh", s)] for s in range(3)])
    if (bosh - lora).mean() > 2 * (bosh - lora).std(ddof=1) / np.sqrt(3):
        print("  LoRA faqat boshdan sezilarli yaxshi")
    if (lora - toliq).mean() > 2 * (lora - toliq).std(ddof=1) / np.sqrt(3):
        print("  lekin to'liq fine-tuning ga yetmadi (bu kichik modelda, 100 qadamda)")

    m = natija["LoRA r=4"][2]
    with torch.no_grad():
        oldin = m(Xv[:4])
        for b in m.bloklar:
            b.qkv, b.proj = b.qkv.birlashtir(), b.proj.birlashtir()
            b.ffn[0], b.ffn[2] = b.ffn[0].birlashtir(), b.ffn[2].birlashtir()
        keyin = m(Xv[:4])
    print(f"  birlashtirilgandan keyin parametrlar {sum(p.numel() for p in m.parameters())}, "
          f"chiqish farqi {(oldin - keyin).abs().max().item():.1e}")
    print("  asl model vaznlari o'zgarmadi: har domen uchun faqat A, B ni saqlash yetadi")


if __name__ == "__main__":
    main()

=== 1. Aralash aniqlik: float16 va bfloat16 ===
  float32   32 bit  70000 -> 70000.0   1 + 0.001 = 1.0010000   eps 1.2e-07
  float16   16 bit  70000 -> inf       1 + 0.001 = 1.0009766   eps 9.8e-04
  bfloat16  16 bit  70000 -> 70144.0   1 + 0.001 = 1.0000000   eps 7.8e-03
  float16: diapazon tor (maks 65504) -> gradient scaling kerak
  bfloat16: diapazon float32 niki, aniqlik past -> scaling kerak emas
  autocast(bfloat16): loss 3.5084 (float32: 3.5082), saqlangan aktivatsiyalar 55% (baytlarda)
  parametrlar turi autocast ichida ham: torch.float32

=== 2. Gradient checkpointing (4 blok, B=8, T=64) ===
                     loss  saqlangan tenzor       bayt   FLOP (f+b)
  oddiy            3.5082               113   11310212    419364864
  checkpoint       3.5082                17     723076    520028160
  gradientlar max farqi: 0.0
  checkpoint har blok KIRISHINI ham ushlab turadi: 393216 bayt
  xotira 1116292 / 11310212 = 9.9%, hisob 1.24x (forward qayta hisoblanadi)

=== 3. LoRA: 'sharhlar' da o'rgatilgan modelni 'klinika sharhlari' ga moslash ===
  moslashdan oldin val loss (yangi domen): 3.2853
  usul        o'rgatiladi   ulush  val loss  - to'liq    2*SE
  faqat bosh         1372    2.2%    1.3857   +0.9880  0.0100
  LoRA r=4           6144    9.8%    0.5700   +0.1723  0.0106
  to'liq FT         62428  100.0%    0.3977
  LoRA faqat boshdan sezilarli yaxshi
  lekin to'liq fine-tuning ga yetmadi (bu kichik modelda, 100 qadamda)
  birlashtirilgandan keyin parametrlar 62428, chiqish farqi 8.1e-06
  asl model vaznlari o'zgarmadi: har domen uchun faqat A, B ni saqlash yetadi

t0 = time.time(); model(X); print(time.time() - t0)                # ⚠️ mashinaga bog'liq
fc = FlopCounterMode(display=False)                                # ✅ FLOP, bayt, parametr
with fc: model(X)

lg = model(t, keshlar)                     # boshi=0                  # ⚠️ har token 0-pozitsiyada
lg = model(t, keshlar, boshi=X.shape[1] - 1)                       # ✅

X2 = keshli(model, prefiks, n)             # "tezroq ishlayapti"      # ⚠️
assert torch.equal(keshsiz(model, prefiks, n)[0], X2[0])           # ✅

w = torch.softmax(q @ k.transpose(-1, -2) / d_h ** 0.5, -1); o = w @ v   # ⚠️ O(n^2) xotira
o = F.scaled_dot_product_attention(q, k, v, is_causal=True)        # ✅

with torch.autocast("cuda", dtype=torch.float16): loss = f()       # ⚠️ kichik gradientlar 0 ga
loss.backward()
scaler = torch.amp.GradScaler(); scaler.scale(loss).backward()     # ✅ (yoki bfloat16)
scaler.step(opt); scaler.update()

x = checkpoint(blok, x)                                            # ⚠️ use_reentrant ogohlantirishi
x = checkpoint(blok, x, use_reentrant=False)                       # ✅

self.B = nn.Parameter(torch.randn(d_chiq, r))                      # ⚠️ boshidan W' != W
self.B = nn.Parameter(torch.zeros(d_chiq, r))                      # ✅ boshida aynan pretrained

1.  # bitta blokning chiziqli qismi FLOP formulasi?
2.  # attention qismi FLOP formulasi?
3.  # d = 1024 da attention qaysi n dan ustun keladi?
4.  # n 2 barobar oshsa, skor matritsasi xotirasi necha barobar?
5.  # SDPA backward uchun n x n o'rniga nima saqlaydi?
6.  # keshsiz generatsiyada k=10, n=40 bo'lsa modeldan nechta pozitsiya o'tadi?
7.  # keshli generatsiyada-chi?
8.  # L=24, d=1024, n=2048, float16 da KV-cache hajmi?
9.  # d=1024, L=24 modelning bloklardagi parametrlari taxminan?
10. # 1B parametrli modelni AdamW bilan float32 da o'rgatish uchun taxminiy xotira?
11. # gradient checkpointing gradientni o'zgartiradimi?
12. # d_kir = d_chiq = 1024, r = 8 LoRA qatlamida nechta o'rgatiladigan parametr?

1.  for _ in range(n):
        lg = model(t, keshlar)[:, -1]
        t = lg.argmax(-1, keepdim=True)

2.  self.B = nn.Parameter(torch.randn(asl.out_features, r) * 0.01)

3.  N = 12 * L * d * d
    print(f"model {N / 1e6:.1f}M parametr")          # V=50000, d=256, L=4

4.  x = checkpoint(blok, x)

5.  with torch.autocast("cpu", dtype=torch.bfloat16):
        model = model.to(torch.bfloat16)
        loss = loss_fn()