| """
|
| SmartCore V1 — Mamba-3 (SISO) + GQA 5:1 hibrit, 48K EN+TR hedefi.
|
|
|
| Referans iskelet: referans_kod/mamba3_minimal.py (saf-PyTorch Mamba-3, Triton'suz).
|
| Bu dosya o modelin üstüne GQA attention mixer'ı ekler ve 5:1 hibrit kurar:
|
| her (attn_every)'inci katmanın SSM mixer'ı yerine GQAttention konur; ilk ve son
|
| katman Mamba kalır (induction head'ler derinlik ortasında oluşsun diye).
|
|
|
| Forward arayüzü minimal modelle birebir: mixer(x, h) -> (y, h). GQA, h'i yok sayar
|
| (eğitim/forward yolu); KV-cache'li decode bu dosyada gerekli değil (param + smoke).
|
|
|
| CPU'da çalışır (Ryzen 5500). Amaç: hibrit öğreniyor mu + kesin ~180M param-sayımı.
|
| """
|
| import sys, time, gc
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
| sys.path.insert(0, "referans_kod")
|
| import mamba3_minimal as m3
|
| from mamba3_minimal import Mamba3Config, Mamba3LMHeadModel, RMSNorm
|
|
|
| torch.manual_seed(0)
|
| torch.set_num_threads(6)
|
| DEV = torch.device("cpu")
|
|
|
|
|
| def rotate_half(x):
|
| x1, x2 = x.chunk(2, dim=-1)
|
| return torch.cat((-x2, x1), dim=-1)
|
|
|
|
|
| class GQAttention(nn.Module):
|
| """Grouped-Query Attention mixer (QK-norm + RoPE, causal). minimal-Mamba3 ile
|
| aynı (y, h) arayüzü. n_heads q-head, n_kv_heads kv-head (GQA sıkıştırma)."""
|
|
|
| def __init__(self, d_model, n_heads=12, n_kv_heads=3, rope_base=10000.0, device=None):
|
| super().__init__()
|
| assert d_model % n_heads == 0, "d_model n_heads'e bölünmeli"
|
| assert n_heads % n_kv_heads == 0, "n_heads n_kv_heads'e bölünmeli (GQA)"
|
| self.n_heads, self.n_kv, self.hd = n_heads, n_kv_heads, d_model // n_heads
|
| self.rep = n_heads // n_kv_heads
|
| self.q_proj = nn.Linear(d_model, n_heads * self.hd, bias=False, device=device)
|
| self.k_proj = nn.Linear(d_model, n_kv_heads * self.hd, bias=False, device=device)
|
| self.v_proj = nn.Linear(d_model, n_kv_heads * self.hd, bias=False, device=device)
|
| self.o_proj = nn.Linear(n_heads * self.hd, d_model, bias=False, device=device)
|
| self.q_norm = RMSNorm(self.hd, device=device)
|
| self.k_norm = RMSNorm(self.hd, device=device)
|
| inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.hd, 2, device=device).float() / self.hd))
|
| self.register_buffer("inv_freq", inv_freq, persistent=False)
|
|
|
| def _rope(self, x, T):
|
| t = torch.arange(T, device=x.device, dtype=torch.float32)
|
| freqs = torch.outer(t, self.inv_freq)
|
| emb = torch.cat((freqs, freqs), dim=-1)
|
| cos, sin = emb.cos()[None, None], emb.sin()[None, None]
|
| return (x * cos + rotate_half(x) * sin).to(x.dtype)
|
|
|
| def forward(self, x, h=None):
|
| B, T, _ = x.shape
|
| q = self.q_proj(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
|
| k = self.k_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
|
| v = self.v_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
|
| q, k = self.q_norm(q), self.k_norm(k)
|
| q, k = self._rope(q, T), self._rope(k, T)
|
| k = k.repeat_interleave(self.rep, dim=1)
|
| v = v.repeat_interleave(self.rep, dim=1)
|
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| y = y.transpose(1, 2).contiguous().view(B, T, -1)
|
| return self.o_proj(y), h
|
|
|
|
|
| def init_weights(model):
|
| """SSM-güvenli init (minimal demo.py'den): A_log, D, dt_bias kritik —
|
| constructor bunları set etmiyor; atlanırsa SSM NaN üretir. GQA projeksiyonları
|
| (2D) std=0.02, RMSNorm ağırlıkları (1D) dokunulmaz (default 1.0)."""
|
| for name, p in model.named_parameters():
|
| if "A_log" in name:
|
| torch.nn.init.uniform_(p, -4, -1)
|
| elif "D" in name and p.dim() == 1:
|
| torch.nn.init.ones_(p)
|
| elif "dt_bias" in name:
|
| torch.nn.init.uniform_(p, 0.001, 0.1)
|
| elif p.dim() >= 2:
|
| torch.nn.init.normal_(p, std=0.02)
|
|
|
|
|
| def build_hybrid(args, attn_every=6, n_heads=12, n_kv_heads=3, device=None):
|
| """minimal Mamba3LMHeadModel kur, sonra her attn_every'inci katmanın mixer'ını
|
| GQAttention ile değiştir (ilk/son katman Mamba kalır)."""
|
| model = Mamba3LMHeadModel(args, device=device)
|
| attn_idx = []
|
| for i in range(args.n_layer):
|
| if (i + 1) % attn_every == 0 and i != 0 and i != args.n_layer - 1:
|
| model.backbone.layers[i].mixer = GQAttention(
|
| args.d_model, n_heads, n_kv_heads, device=device)
|
| attn_idx.append(i)
|
| init_weights(model)
|
| return model, attn_idx
|
|
|
|
|
| def n_params(model):
|
| seen, tot = set(), 0
|
| for p in model.parameters():
|
| if id(p) in seen:
|
| continue
|
| seen.add(id(p)); tot += p.numel()
|
| return tot
|
|
|
|
|
| def make_config(d_model, n_layer, vocab, d_mlp_inner=1500, d_state=128, headdim=64, chunk_size=64):
|
| c = Mamba3Config(d_model=d_model, n_layer=n_layer, d_state=d_state,
|
| headdim=headdim, chunk_size=chunk_size, vocab_size=vocab)
|
| c.d_mlp_inner = d_mlp_inner
|
| return c
|
|
|
|
|
| def param_sweep():
|
| """180M hedef için 48K vocab'ta katman sweep'i (import edilmez, sadece doğrudan çalıştırınca)."""
|
| print(f" {'n_layer':>7} {'Mamba':>6} {'GQA':>4} {'embed(M)':>9} {'TOPLAM(M)':>10}")
|
| for nlayer in [18, 19, 20, 21, 22]:
|
| c = make_config(d_model=768, n_layer=nlayer, vocab=48000, d_mlp_inner=1500)
|
| mdl, ai = build_hybrid(c, attn_every=6, n_heads=12, n_kv_heads=3, device=DEV)
|
| tot = n_params(mdl) / 1e6
|
| embed_m = c.vocab_size * c.d_model / 1e6
|
| mark = " <== 180M bandi" if 175 <= tot <= 185 else ""
|
| print(f" {nlayer:>7} {nlayer-len(ai):>6} {len(ai):>4} {embed_m:>9.1f} {tot:>10.1f}{mark}")
|
| del mdl; gc.collect()
|
|
|
|
|
| if __name__ == "__main__":
|
|
|
| param_sweep()
|
|
|
