02_linearization_sweep.py

"""02: Lineerleştirme oran taraması — eğitimsiz katman-bazlı damıtmanın PPL degradasyon eğrisi (@512/@2048)."""
import torch, torch.nn as nn, numpy as np, math, time
from transformers import AutoModelForCausalLM, AutoTokenizer
from fla.layers import GatedDeltaNet
torch.manual_seed(0); np.random.seed(0)
T0=time.time()
def el(): return f"[{(time.time()-T0)/60:.1f}dk]"

MID="ecloudtech/Erk-14B"
KEEP=[0,9,28,29,30,31,32,33,34,35]
LIN_ORDER=[i for i in range(40) if i not in KEEP]   # 30 katman, sirali=derinlik=kolay->zor
assert len(LIN_ORDER)==30 and sorted(LIN_ORDER)==sorted(set(range(40))-set(KEEP))
CUTOFFS=[6,12,18,24,30]; VAL_N=16

tok=AutoTokenizer.from_pretrained(MID)
model=AutoModelForCausalLM.from_pretrained(MID, torch_dtype=torch.bfloat16).to("cuda").eval()
for p in model.parameters(): p.requires_grad=False
H=model.config.hidden_size
print(f"{el()} Erk yuklendi hidden={H}, LIN_ORDER={LIN_ORDER}", flush=True)

# GDN shape sanity
_g=GatedDeltaNet(hidden_size=H, head_dim=128, num_heads=40, use_gate=True, use_short_conv=True, mode="chunk").to("cuda").to(torch.bfloat16)
with torch.no_grad():
    _o=_g(torch.randn(1,64,H,device="cuda",dtype=torch.bfloat16)); _y=_o[0] if isinstance(_o,tuple) else _o
assert tuple(_y.shape)==(1,64,H), f"GDN shape {_y.shape}"
del _g,_o,_y; torch.cuda.empty_cache()
print(f"{el()} GDN shape sanity OK", flush=True)

# veri havuzlari (ayrik)
txt=open("./data/eval_tr.txt",encoding="utf-8",errors="ignore").read(4000000)
paras=[p.strip() for p in txt.split("\n") if 200<len(p.strip())<1500]
print(f"{el()} {len(paras)} paragraf havuzu", flush=True)
cal=paras[0:48]; val=paras[48:64]; test=paras[64:164]
drift=paras[164:176]; abl_pool=paras[176:272]; abl_val=paras[272:288]; longsrc=paras[288:]

def ids_of(p, ml=512): return tok(p, return_tensors="pt", truncation=True, max_length=ml).input_ids
def build_long(plist, block=2048, n=8):
    if not plist: return []
    s=tok("\n\n".join(plist), return_tensors="pt").input_ids[0]; out=[]
    for i in range(0, max(0,len(s)-block), block):
        out.append(s[i:i+block].unsqueeze(0))
        if len(out)>=n: break
    return out
test_short=[ids_of(p) for p in test]; test_long=build_long(longsrc)
print(f"{el()} test_short={len(test_short)} para, test_long={len(test_long)} blok@2048", flush=True)

def ppl(id_list):
    if not id_list: return (float("nan"),float("nan"),float("nan"))
    per=[]
    with torch.no_grad():
        for ids in id_list:
            ids=ids.to("cuda")
            if ids.shape[1]<8: continue
            out=model(ids, labels=ids, use_cache=False); n=ids.shape[1]-1
            per.append((out.loss.item()*n, n))
    arr=np.array(per, dtype=np.float64)
    p=math.exp(arr[:,0].sum()/max(arr[:,1].sum(),1))
    rng=np.random.default_rng(0); B=len(arr); boot=[]
    for _ in range(400):
        s=arr[rng.integers(0,B,B)]; boot.append(math.exp(s[:,0].sum()/max(s[:,1].sum(),1)))
    lo,hi=np.percentile(boot,[2.5,97.5]); return p,lo,hi

gb=[ppl(test_short[i::3])[0] for i in range(3)]
print(f"{el()} GURULTU TABANI PPL@512: {gb[0]:.3f}/{gb[1]:.3f}/{gb[2]:.3f} std={np.std(gb):.4f}", flush=True)
p512=ppl(test_short); p2048=ppl(test_long)
ppl_o512=p512[0]; ppl_o2048=p2048[0]
print(f"{el()} ORACLE @512={p512[0]:.3f}[{p512[1]:.3f},{p512[2]:.3f}] | @2048={p2048[0]:.3f}[{p2048[1]:.3f},{p2048[2]:.3f}]", flush=True)

# TUM 30 katman I/O (orijinal modelden TEK gecis; bf16 cpu)
cache={li:{"in":[],"out":[]} for li in LIN_ORDER}
def mk_hook(li):
    def hook(m,args,kwargs,output):
        inp=args[0] if len(args)>0 else kwargs.get("hidden_states")
        out=output[0] if isinstance(output,tuple) else output
        cache[li]["in"].append(inp.detach().to(torch.bfloat16).cpu())
        cache[li]["out"].append(out.detach().to(torch.bfloat16).cpu())
    return hook
hks=[model.model.layers[li].self_attn.register_forward_hook(mk_hook(li), with_kwargs=True) for li in LIN_ORDER]
with torch.no_grad():
    for p in cal+val:
        ids=ids_of(p).to("cuda")
        if ids.shape[1]<32: continue
        model(ids, use_cache=False)
for h in hks: h.remove()
M=len(cache[LIN_ORDER[0]]["in"]); n_cal=M-VAL_N
cal_idx=list(range(n_cal)); val_idx=list(range(n_cal,M))
print(f"{el()} I/O: {M} dizi/katman (cal={n_cal}, val={VAL_N})", flush=True)

def nmse(y,t): return (((y-t)**2).sum()/((t)**2).sum())
def val_nmse(gdn,li,vi):
    tot=0.0
    with torch.no_grad():
        for i in vi:
            x=cache[li]["in"][i].to("cuda").float(); t=cache[li]["out"][i].to("cuda").float()
            o=gdn(x); y=(o[0] if isinstance(o,tuple) else o).float(); tot+=nmse(y,t).item()
    return tot/len(vi)

def distill(li, max_steps=1500, patience=5, evk=100, lr=1e-3):
    torch.manual_seed(100+li)
    gdn=GatedDeltaNet(hidden_size=H, head_dim=128, num_heads=40, use_gate=True, use_short_conv=True, mode="chunk").to("cuda").float()
    opt=torch.optim.AdamW(gdn.parameters(), lr=lr); sch=torch.optim.lr_scheduler.CosineAnnealingLR(opt,T_max=max_steps)
    best=9e9; best_sd=None; bad=0; N=len(cal_idx)
    for s in range(max_steps):
        i=cal_idx[s%N]
        x=cache[li]["in"][i].to("cuda").float(); t=cache[li]["out"][i].to("cuda").float()
        o=gdn(x); y=(o[0] if isinstance(o,tuple) else o).float(); loss=nmse(y,t)
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(gdn.parameters(),1.0); opt.step(); sch.step()
        if (s+1)%evk==0:
            v=val_nmse(gdn,li,val_idx)
            if v<best-1e-4: best=v; best_sd={k:vv.detach().cpu().clone() for k,vv in gdn.state_dict().items()}; bad=0
            else:
                bad+=1
                if bad>=patience: break
    if best_sd: gdn.load_state_dict({k:vv.to("cuda") for k,vv in best_sd.items()})
    return gdn.eval().to(torch.bfloat16), best

class GDNWrap(nn.Module):
    def __init__(self, gdn): super().__init__(); self.gdn=gdn
    def forward(self, hidden_states, *a, **k):
        o=self.gdn(hidden_states); y=o[0] if isinstance(o,tuple) else o; return (y,None)

orig_attn={li:model.model.layers[li].self_attn for li in LIN_ORDER}
wraps={}; nmse_tf={}
print(f"{el()} === 30 KATMAN DAMITMA (held-out val-nMSE) ===", flush=True)
for li in LIN_ORDER:
    gdn,v=distill(li); wraps[li]=GDNWrap(gdn).to("cuda").to(torch.bfloat16); nmse_tf[li]=v
    print(f"{el()}   L{li:2d} val-nMSE {v:.3f}", flush=True)

rnd_cache={}
def set_hybrid(k, use_random=False):
    for i,li in enumerate(LIN_ORDER):
        if i<k:
            if use_random:
                if li not in rnd_cache:
                    torch.manual_seed(7000+li)
                    rg=GatedDeltaNet(hidden_size=H, head_dim=128, num_heads=40, use_gate=True, use_short_conv=True, mode="chunk").to("cuda").to(torch.bfloat16).eval()
                    rnd_cache[li]=GDNWrap(rg).to("cuda").to(torch.bfloat16)
                model.model.layers[li].self_attn=rnd_cache[li]
            else:
                model.model.layers[li].self_attn=wraps[li]
        else:
            model.model.layers[li].self_attn=orig_attn[li]

print(f"{el()} === TARAMA (kumulatif free-running PPL) ===", flush=True)
sweep=[]
for k in CUTOFFS:
    set_hybrid(k); a=ppl(test_short); b=ppl(test_long)
    avn=float(np.mean([nmse_tf[li] for li in LIN_ORDER[:k]]))
    sweep.append((k,a,b,avn))
    print(f"{el()}   k={k:2d} -> PPL@512 {a[0]:.3f}(x{a[0]/ppl_o512:.3f})[{a[1]:.2f},{a[2]:.2f}] | @2048 {b[0]:.3f}(x{b[0]/ppl_o2048:.3f}) | ort.tf-nMSE {avn:.3f}", flush=True)

print(f"{el()} === ALT-SINIR (rastgele-init GDN) ===", flush=True)
for k in [6,30]:
    set_hybrid(k, use_random=True); a=ppl(test_short)
    print(f"{el()}   k={k} rastgele -> PPL@512 {a[0]:.3f} (x{a[0]/ppl_o512:.2f})", flush=True)

# DRIFT: tam-swap modelde her katmanin gercek girdisi vs orijinal
print(f"{el()} === DRIFT (free-running birikme) ===", flush=True)
set_hybrid(0); orig_in={li:[] for li in LIN_ORDER}
def mk_in_hook(li, store):
    def h(m,args,kwargs,output):
        store[li].append((args[0] if len(args)>0 else kwargs.get("hidden_states")).detach().to(torch.bfloat16).cpu())
    return h
hks=[model.model.layers[li].self_attn.register_forward_hook(mk_in_hook(li,orig_in), with_kwargs=True) for li in LIN_ORDER]
with torch.no_grad():
    for p in drift:
        ids=ids_of(p).to("cuda")
        if ids.shape[1]<32: continue
        model(ids, use_cache=False)
for h in hks: h.remove()
set_hybrid(30); sw_in={li:[] for li in LIN_ORDER}
def mk_pre(li):
    def h(m,args): sw_in[li].append(args[0].detach().to(torch.bfloat16).cpu())
    return h
hks=[model.model.layers[li].self_attn.gdn.register_forward_pre_hook(mk_pre(li)) for li in LIN_ORDER]
with torch.no_grad():
    for p in drift:
        ids=ids_of(p).to("cuda")
        if ids.shape[1]<32: continue
        model(ids, use_cache=False)
for h in hks: h.remove()
drift_nmse={}
for li in LIN_ORDER:
    ds=[(((b.float()-a.float())**2).sum()/((a.float())**2).sum()).item() for a,b in zip(orig_in[li],sw_in[li])]
    drift_nmse[li]=float(np.mean(ds)) if ds else float("nan")
    print(f"{el()}   L{li:2d} tf-nMSE {nmse_tf[li]:.3f} | drift-nMSE {drift_nmse[li]:.3f}", flush=True)

# VERI-OLCEKLEME ABLASYONU (L2 kolay, L39 zor)
print(f"{el()} === VERI-OLCEKLEME ABLASYONU ===", flush=True)
try:
    set_hybrid(0); ABL=[2,39]; abl={li:{"in":[],"out":[]} for li in ABL}
    def mk_abl(li):
        def h(m,args,kwargs,output):
            inp=args[0] if len(args)>0 else kwargs.get("hidden_states")
            out=output[0] if isinstance(output,tuple) else output
            abl[li]["in"].append(inp.detach().to(torch.bfloat16).cpu()); abl[li]["out"].append(out.detach().to(torch.bfloat16).cpu())
        return h
    hks=[model.model.layers[li].self_attn.register_forward_hook(mk_abl(li), with_kwargs=True) for li in ABL]
    with torch.no_grad():
        for p in abl_pool+abl_val:
            ids=ids_of(p).to("cuda")
            if ids.shape[1]<32: continue
            model(ids, use_cache=False)
    for h in hks: h.remove()
    MA=len(abl[ABL[0]]["in"]); VN=16; vi=list(range(MA-VN,MA))
    def distill_abl(li,ntrain):
        ci=list(range(ntrain)); torch.manual_seed(100+li)
        gdn=GatedDeltaNet(hidden_size=H, head_dim=128, num_heads=40, use_gate=True, use_short_conv=True, mode="chunk").to("cuda").float()
        opt=torch.optim.AdamW(gdn.parameters(), lr=1e-3); sch=torch.optim.lr_scheduler.CosineAnnealingLR(opt,T_max=1500)
        best=9e9; bad=0; N=len(ci)
        for s in range(1500):
            i=ci[s%N]; x=abl[li]["in"][i].to("cuda").float(); t=abl[li]["out"][i].to("cuda").float()
            o=gdn(x); y=(o[0] if isinstance(o,tuple) else o).float(); loss=(((y-t)**2).sum()/((t)**2).sum())
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(gdn.parameters(),1.0); opt.step(); sch.step()
            if (s+1)%100==0:
                tot=0.0
                with torch.no_grad():
                    for j in vi:
                        x=abl[li]["in"][j].to("cuda").float(); t=abl[li]["out"][j].to("cuda").float()
                        o=gdn(x); y=(o[0] if isinstance(o,tuple) else o).float(); tot+=(((y-t)**2).sum()/((t)**2).sum()).item()
                v=tot/len(vi)
                if v<best-1e-4: best=v; bad=0
                else:
                    bad+=1
                    if bad>=5: break
        return best
    big=min(96,MA-VN)
    for li in ABL:
        v48=distill_abl(li,48); vb=distill_abl(li,big)
        print(f"{el()}   L{li:2d}: 48para={v48:.3f} -> {big}para={vb:.3f} (dusus={v48-vb:+.3f} => {'VERI-SINIRLI' if (v48-vb)>0.03 else 'doygun'})", flush=True)
except Exception as e:
    import traceback; print(f"{el()}   ablasyon atlandi: {repr(e)}", flush=True)

# OZET
print("\n"+"="*66+"\nOZET - lineerlestirme degradasyon egrisi\n"+"="*66, flush=True)
print(f"  Oracle (0 lineer): PPL@512 {ppl_o512:.3f} | @2048 {ppl_o2048:.3f}  (gurultu std {np.std(gb):.4f})")
for k,a,b,vn in sweep:
    print(f"  {k:2d}/30 lineer (~%{100*(len(KEEP)+ (0))/40:.0f} full kalir): PPL@512 x{a[0]/ppl_o512:.3f} | @2048 x{b[0]/ppl_o2048:.3f} | ort.tf-nMSE {vn:.3f}")
wd=max(v for v in drift_nmse.values() if not math.isnan(v))
print(f"  En yuksek drift-nMSE (birikme): {wd:.3f}  (tf-nMSE dusuk kalip drift buyurse = hata birikmesi belirtisi)")
f512=sweep[-1][1][0]/ppl_o512; f2048=sweep[-1][2][0]/ppl_o2048
print(f"\n  YORUM: 30/30 lineer -> PPL@512 x{f512:.2f}, @2048 x{f2048:.2f} (EGITIMSIZ, LoRA yok = degradasyon TABANI).")
print(f"  Final kalite LoRA ko-egitimi sonrasi olculur. Uzun-baglam (@2048) egrisi karar icin belirleyici.")
print(f"  Kosullar: {len(test_short)}-para @512 + {len(test_long)}-blok @2048, tek seed, bootstrap %95 GA.")
print(f"{el()} BITTI", flush=True)