"""PG-v2: SP8192+ParallelRes+DepthRec+TTT+Int6GPTQ+EMA"""
from __future__ import annotations
import copy,glob,io,math,os,random,subprocess,sys,time,uuid,zlib
from pathlib import Path
import numpy as np
import sentencepiece as spm
import torch,torch.distributed as dist,torch.nn.functional as F
from torch import Tensor,nn
from torch.nn.parallel import DistributedDataParallel as DDP

class H:
    dp=os.environ.get("DATA_PATH","./data/datasets/fineweb10B_sp8192")
    tf=os.path.join(dp,"fineweb_train_*.bin")
    vf=os.path.join(dp,"fineweb_val_*.bin")
    tp=os.environ.get("TOKENIZER_PATH","./data/tokenizers/fineweb_8192_bpe.model")
    rid=os.environ.get("RUN_ID",str(uuid.uuid4()))
    seed=int(os.environ.get("SEED","1337"))
    vbs=int(os.environ.get("VBS","524288"));vle=int(os.environ.get("VLE","1000"))
    tle=int(os.environ.get("TLE","200"))
    iters=int(os.environ.get("ITERS","20000"))
    wdi=int(os.environ.get("WDI","3500"));wui=int(os.environ.get("WUI","20"))
    tbt=int(os.environ.get("TBT","524288"));tsl=int(os.environ.get("TSL","1024"))
    mws=float(os.environ.get("MWS","600.0"))
    V=int(os.environ.get("V","8192"));D=int(os.environ.get("D","768"))
    nh=int(os.environ.get("NH","12"));nkv=int(os.environ.get("NKV","4"))
    mm=int(os.environ.get("MM","4"))
    nul=int(os.environ.get("NUL","3"));nr=int(os.environ.get("NR","8"))
    ner=int(os.environ.get("NER","0"))
    rb=float(os.environ.get("RB","10000.0"))
    lsc=float(os.environ.get("LSC","30.0"))
    qkg=float(os.environ.get("QKG","5.25"))
    sws=int(os.environ.get("SWS","64"));swl=int(os.environ.get("SWL","1024"))
    tte=int(os.environ.get("TTE","1"))
    ttlr=float(os.environ.get("TTLR","0.01"))
    ttcs=int(os.environ.get("TTCS","64"))
    ttly=os.environ.get("TTLY","all")
    elr=float(os.environ.get("ELR","0.05"))
    mlr=float(os.environ.get("MLR","0.04"))
    slr=float(os.environ.get("SLR","0.04"))
    mmo=float(os.environ.get("MMO","0.95"))
    mbs=int(os.environ.get("MBS","5"))
    mwd=float(os.environ.get("MWD","0.09"))
    b1=float(os.environ.get("B1","0.9"))
    b2=float(os.environ.get("B2","0.95"))
    ae=float(os.environ.get("AE","1e-8"))
    gb=int(os.environ.get("GB","6"))
    sdn=float(os.environ.get("SDN","2.5"))
    esf=float(os.environ.get("ESF","0.4"))

CP=tuple(p for p in "attn_scale,mlp_scale,resid_mix,q_gain".split(",") if p)

def zp5(G,s=10,e=1e-7):
    a,b,c=3.4445,-4.7750,2.0315
    X=G.bfloat16();X/=X.norm()+e
    tr=G.size(0)>G.size(1)
    if tr:X=X.T
    for _ in range(s):
        A=X@X.T;B=b*A+c*A@A;X=a*X+B@X
    return X.T if tr else X

class Muon(torch.optim.Optimizer):
    def __init__(s,p,lr,mom,bs,wd=0.,nest=True):
        super().__init__(p,dict(lr=lr,mom=mom,bs=bs,wd=wd,nest=nest))
    @torch.no_grad()
    def step(s,cl=None):
        lo=None
        if cl:
            with torch.enable_grad():lo=cl()
        dd=dist.is_available() and dist.is_initialized()
        ws=dist.get_world_size() if dd else 1
        rk=dist.get_rank() if dd else 0
        for g in s.param_groups:
            ps=g["params"];lr=g["lr"];mo=g["mom"];bs=g["bs"];wd=g["wd"];ne=g["nest"]
            tot=sum(int(p.numel()) for p in ps)
            fl=torch.zeros(tot,device=ps[0].device,dtype=torch.bfloat16)
            cur=0
            for i,p in enumerate(ps):
                if i%ws==rk and p.grad is not None:
                    gr=p.grad
                    if wd:gr=gr+wd*p.data.to(gr.dtype)
                    st=s.state[p]
                    if "mb" not in st:st["mb"]=torch.zeros_like(gr)
                    buf=st["mb"];buf.mul_(mo).add_(gr)
                    if ne:gr=gr.add(buf,alpha=mo)
                    gr=zp5(gr,steps=bs)
                    gr*=max(1,gr.size(0)/gr.size(1))**0.5
                    fl[cur:cur+p.numel()]=gr.reshape(-1)
                cur+=p.numel()
            if dd:dist.all_reduce(fl,op=dist.ReduceOp.SUM)
            cur=0
            for p in ps:
                gr=fl[cur:cur+p.numel()].view_as(p).to(dtype=p.dtype)
                p.add_(gr,alpha=-lr);cur+=p.numel()
        return lo

def build_sp_luts(sp,vs,dev):
    sv=int(sp.vocab_size());sz=max(sv,vs)
    bb=np.zeros(sz,dtype=np.int16);hs=np.zeros(sz,dtype=bool);ib=np.ones(sz,dtype=bool)
    for t in range(sv):
        if sp.is_control(t) or sp.is_unknown(t) or sp.is_unused(t):continue
        ib[t]=False
        if sp.is_byte(t):bb[t]=1;continue
        pc=sp.id_to_piece(t)
        if pc.startswith("\u2581"):hs[t]=True;pc=pc[1:]
        bb[t]=len(pc.encode("utf-8"))
    return(torch.tensor(bb,dtype=torch.int16,device=dev),
           torch.tensor(hs,dtype=torch.bool,device=dev),
           torch.tensor(ib,dtype=torch.bool,device=dev))

def eval_sw(a,mdl,rk,ws,dev,vt,bbl,hsl,ibl,ttt=False):
    sl=a.swl;st=a.sws;T=vt.numel()
    starts=list(range(0,T-sl-1,st))
    my=starts[rk::ws]
    ls=torch.zeros((),device=dev,dtype=torch.float64)
    tc=torch.zeros((),device=dev,dtype=torch.float64)
    bc=torch.zeros((),device=dev,dtype=torch.float64)
    rm=mdl
    while hasattr(rm,'module'):rm=rm.module
    if hasattr(rm,'_orig_mod'):rm=rm._orig_mod
    rm.eval()
    ctx=torch.no_grad if ttt else torch.inference_mode
    with ctx():
        for s in my:
            e=s+sl
            x=vt[s:e].unsqueeze(0).to(dev,dtype=torch.int64)
            y=vt[s+1:e+1].unsqueeze(0).to(dev,dtype=torch.int64)
            with torch.autocast("cuda",dtype=torch.bfloat16):
                if ttt and a.tte:ptl=rm.ptl_ttt(x,y,a)
                else:ptl=rm.ptl(x,y)
            lo=sl-st;ps=ptl[0,lo:];ys=y[0,lo:];xs=x[0,lo:]
            ls+=ps.to(torch.float64).sum();tc+=ps.numel()
            tb=bbl[ys].to(torch.float64)
            tb+=(hsl[ys]&~ibl[xs]).to(torch.float64)
            bc+=tb.sum()
    if dist.is_available() and dist.is_initialized():
        for t in(ls,tc,bc):dist.all_reduce(t,op=dist.ReduceOp.SUM)
    vl=float((ls/tc).item());bpb=float((ls/math.log(2)/bc).item())
    rm.train();return vl,bpb

def sdclip(t,n=2.5):
    m=t.float().mean();s=t.float().std()
    return t.clamp((m-n*s).item(),(m+n*s).item())

def qi6(t,ns=2.5):
    t32=t.float();mx=31
    if t32.ndim==2:
        m=t32.mean(1,keepdim=True);s=t32.std(1,keepdim=True).clamp_min(1e-9)
        lo=m-ns*s;hi=m+ns*s
        tc=t32.clamp(lo.expand_as(t32),hi.expand_as(t32))
        cv=tc.abs().amax(1).clamp_min(1e-9);sc=cv/mx
        q=torch.clamp(torch.round(tc/sc[:,None]),-mx,mx).to(torch.int8)
        return q.contiguous(),sc.to(torch.float16).contiguous()
    tc=sdclip(t32,ns);cv=float(tc.abs().max().item())
    sc=torch.tensor(max(cv/mx,1./mx),dtype=torch.float32)
    q=torch.clamp(torch.round(tc/sc),-mx,mx).to(torch.int8)
    return q.contiguous(),sc

def qi8(t,ns=2.5):
    t32=t.float()
    if t32.ndim==2:
        m=t32.mean(1,keepdim=True);s=t32.std(1,keepdim=True).clamp_min(1e-9)
        lo=m-ns*s;hi=m+ns*s
        tc=t32.clamp(lo.expand_as(t32),hi.expand_as(t32))
        cv=tc.abs().amax(1).clamp_min(1e-9);sc=cv/127.
        q=torch.clamp(torch.round(tc/sc[:,None]),-127,127).to(torch.int8)
        return q.contiguous(),sc.to(torch.float16).contiguous()
    cv=float(sdclip(t32,ns).abs().max().item())
    sc=torch.tensor(max(cv/127.,1./127.),dtype=torch.float32)
    q=torch.clamp(torch.round(t32.clamp(-cv,cv)/sc),-127,127).to(torch.int8)
    return q.contiguous(),sc

def qsd(sd,gb=6,ns=2.5):
    qf=qi6 if gb==6 else qi8
    qu,sc,dt,pt,po,qm={},{},{},{},{},{}
    st={k:0 for k in("pc","nt","bb","qb")}
    for n,t in sd.items():
        t=t.detach().cpu().contiguous()
        st["pc"]+=t.numel();st["nt"]+=1;st["bb"]+=t.numel()*t.element_size()
        if not t.is_floating_point():pt[n]=t;st["qb"]+=t.numel()*t.element_size();continue
        ic=any(p in n for p in CP);ism=t.numel()<=65536
        if "tok_emb" in n:
            po[n]=str(t.dtype).removeprefix("torch.")
            q,s=qi8(t,ns);qu[n]=q;sc[n]=s;dt[n]=po[n]
            if s.ndim>0:qm[n]={"scheme":"per_row","axis":0,"bits":8}
            st["qb"]+=q.numel()+s.numel()*s.element_size();continue
        if ic or ism:
            if t.dtype in(torch.float32,torch.bfloat16):po[n]=str(t.dtype).removeprefix("torch.")
            pt[n]=t.float() if ic else t.to(torch.float16)
            pt[n]=pt[n].contiguous();st["qb"]+=pt[n].numel()*pt[n].element_size();continue
        q,s=qf(t,ns)
        if s.ndim>0:qm[n]={"scheme":"per_row","axis":0,"bits":gb}
        qu[n]=q;sc[n]=s;dt[n]=str(t.dtype).removeprefix("torch.")
        st["qb"]+=q.numel()+s.numel()*s.element_size()
    obj={"__qf__":f"i{gb}sd","q":qu,"s":sc,"d":dt,"p":pt}
    if qm:obj["m"]=qm
    if po:obj["o"]=po
    return obj,st

def dqsd(obj):
    out={};qm=obj.get("m",{});po=obj.get("o",{})
    for n,q in obj["q"].items():
        dt=getattr(torch,obj["d"][n]);s=obj["s"][n]
        if qm.get(n,{}).get("scheme")=="per_row" or s.ndim>0:
            s=s.to(torch.float32)
            out[n]=(q.float()*s.view(q.shape[0],*([1]*(q.ndim-1)))).to(dt).contiguous()
        else:out[n]=(q.float()*float(s.item())).to(dt).contiguous()
    for n,t in obj["p"].items():
        ot=t.detach().cpu().contiguous();od=po.get(n)
        if isinstance(od,str):ot=ot.to(dtype=getattr(torch,od)).contiguous()
        out[n]=ot
    return out

def lds(f):
    h=np.fromfile(f,dtype="<i4",count=256)
    if h.size!=256 or int(h[0])!=20240520 or int(h[1])!=1:raise ValueError(f"Bad:{f}")
    n=int(h[2]);t=np.fromfile(f,dtype="<u2",count=n,offset=256*4)
    return torch.from_numpy(t.astype(np.uint16,copy=False))

def lvt(pat,sl):
    fs=[Path(p) for p in sorted(glob.glob(pat))]
    if not fs:raise FileNotFoundError(f"No val:{pat}")
    t=torch.cat([lds(f) for f in fs]).contiguous()
    u=((t.numel()-1)//sl)*sl;return t[:u+1]

class TS:
    def __init__(s,pat):
        fs=[Path(p) for p in sorted(glob.glob(pat))]
        if not fs:raise FileNotFoundError(f"No:{pat}")
        s.fs=fs;s.i=0;s.t=lds(fs[0]);s.p=0
    def take(s,n):
        ch=[];r=n
        while r>0:
            av=s.t.numel()-s.p
            if av<=0:s.i=(s.i+1)%len(s.fs);s.t=lds(s.fs[s.i]);s.p=0;av=s.t.numel()
            k=min(r,av);ch.append(s.t[s.p:s.p+k]);s.p+=k;r-=k
        return ch[0] if len(ch)==1 else torch.cat(ch)

class DTL:
    def __init__(s,pat,rk,ws,dev):s.rk=rk;s.ws=ws;s.dev=dev;s.st=TS(pat)
    def nb(s,gt,sl,ga):
        lt=gt//(s.ws*ga);ps=lt+1
        ch=s.st.take(ps*s.ws);st=s.rk*ps
        lo=ch[st:st+ps].to(torch.int64)
        x=lo[:-1].reshape(-1,sl);y=lo[1:].reshape(-1,sl)
        return x.to(s.dev,non_blocking=True),y.to(s.dev,non_blocking=True)

class RN(nn.Module):
    def __init__(s,eps=None):super().__init__();s.eps=eps
    def forward(s,x):return F.rms_norm(x,(x.size(-1),),eps=s.eps)

class Rot(nn.Module):
    def __init__(s,d,b=10000.):
        super().__init__()
        s.register_buffer("if_",1./(b**(torch.arange(0,d,2,dtype=torch.float32)/d)),persistent=False)
        s._cl=0;s._c=None;s._s=None
    def forward(s,sl,dev,dt):
        if s._c is None or s._cl!=sl or s._c.device!=dev:
            t=torch.arange(sl,device=dev,dtype=s.if_.dtype)
            fr=torch.outer(t,s.if_.to(dev))
            s._c=fr.cos()[None,None,:,:];s._s=fr.sin()[None,None,:,:];s._cl=sl
        return s._c.to(dtype=dt),s._s.to(dtype=dt)

def arot(x,c,si):
    h=x.size(-1)//2;x1,x2=x[...,:h],x[...,h:]
    return torch.cat((x1*c+x2*si,x1*(-si)+x2*c),dim=-1)

class CSA(nn.Module):
    def __init__(s,d,nh,nk,rb,qkg):
        super().__init__()
        assert d%nh==0 and nh%nk==0
        s.nh=nh;s.nk=nk;s.hd=d//nh;kd=nk*s.hd
        s.cq=nn.Linear(d,d,bias=False);s.ck=nn.Linear(d,kd,bias=False)
        s.cv=nn.Linear(d,kd,bias=False);s.pr=nn.Linear(d,d,bias=False)
        s.qg=nn.Parameter(torch.full((nh,),qkg,dtype=torch.float32))
        s.rot=Rot(s.hd,base=rb)
    def forward(s,x):
        B,T,_=x.shape
        q=s.cq(x).reshape(B,T,s.nh,s.hd).transpose(1,2)
        k=s.ck(x).reshape(B,T,s.nk,s.hd).transpose(1,2)
        v=s.cv(x).reshape(B,T,s.nk,s.hd).transpose(1,2)
        q=F.rms_norm(q,(q.size(-1),));k=F.rms_norm(k,(k.size(-1),))
        c,si=s.rot(T,x.device,q.dtype)
        q=arot(q,c,si);k=arot(k,c,si)
        q=q*s.qg.to(dtype=q.dtype)[None,:,None,None]
        y=F.scaled_dot_product_attention(q,k,v,attn_mask=None,is_causal=True,
            enable_gqa=(s.nk!=s.nh))
        return s.pr(y.transpose(1,2).contiguous().reshape(B,T,-1))

class MLP(nn.Module):
    def __init__(s,d,m):
        super().__init__()
        h=d*m;s.fc=nn.Linear(d,h,bias=False);s.pr=nn.Linear(h,d,bias=False)
    def forward(s,x):return s.pr(torch.relu(s.fc(x)).square())

class PB(nn.Module):
    """Parallel residual block."""
    def __init__(s,d,nh,nk,mm,rb,qkg):
        super().__init__()
        s.n=RN();s.a=CSA(d,nh,nk,rb,qkg);s.m=MLP(d,mm)
        s.as_=nn.Parameter(torch.ones(d,dtype=torch.float32))
        s.ms=nn.Parameter(torch.ones(d,dtype=torch.float32))
        s.rm=nn.Parameter(torch.stack([torch.ones(d),torch.zeros(d)]).float())
    def forward(s,x,x0):
        mx=s.rm.to(x.dtype)
        x=mx[0][None,None,:]*x+mx[1][None,None,:]*x0
        h=s.n(x)
        x=x+s.as_.to(x.dtype)[None,None,:]*s.a(h)+s.ms.to(x.dtype)[None,None,:]*s.m(h)
        return x

class RGPT(nn.Module):
    def __init__(s,a):
        super().__init__()
        s.lsc=a.lsc;s._tr=a.nr;s._er=a.ner or a.nr*2;s._V=a.V
        s.te=nn.Embedding(a.V,a.D)
        s.bl=nn.ModuleList([PB(a.D,a.nh,a.nkv,a.mm,a.rb,a.qkg) for _ in range(a.nul)])
        s.fn=RN();nn.init.normal_(s.te.weight,std=0.005)
    def _fh(s,ids):
        x=F.rms_norm(s.te(ids),(s.te.embedding_dim,));x0=x
        n=s._tr if s.training else s._er
        for _ in range(n):
            for b in s.bl:x=b(x,x0)
        return s.fn(x)
    def forward(s,ids,tgt):
        h=s._fh(ids);lo=F.linear(h.reshape(-1,h.size(-1)),s.te.weight)
        lo=s.lsc*torch.tanh(lo/s.lsc)
        return F.cross_entropy(lo.float(),tgt.reshape(-1),reduction="mean")
    def ptl(s,ids,tgt):
        h=s._fh(ids);B,T,D=h.shape
        lo=F.linear(h.reshape(B*T,D),s.te.weight)
        lo=s.lsc*torch.tanh(lo/s.lsc)
        return F.cross_entropy(lo.float(),tgt.reshape(B*T),reduction="none").reshape(B,T)
    @torch.no_grad()
    def ptl_ttt(s,ids,tgt,a):
        """Score-first TTT: score chunk, then update MLP W_down for next chunk."""
        cs=a.ttcs;lr=a.ttlr;B,T=ids.shape
        if a.ttly=="all":li=list(range(len(s.bl)))
        else:li=[int(x) for x in a.ttly.split(",")]
        ow={i:s.bl[i].m.pr.weight.data.clone() for i in li}
        ap=[];nc=(T+cs-1)//cs
        for ci in range(nc):
            lo=ci*cs;hi=min((ci+1)*cs,T)
            h=s._fh(ids);hc=h[:,lo:hi,:];yc=tgt[:,lo:hi]
            lg=F.linear(hc.reshape(-1,hc.size(-1)),s.te.weight)
            lg=s.lsc*torch.tanh(lg/s.lsc)
            pt=F.cross_entropy(lg.float(),yc.reshape(-1),reduction="none").reshape(B,hi-lo)
            ap.append(pt)
            if ci<nc-1:
                for i in li:
                    blk=s.bl[i];hn=F.rms_norm(hc.reshape(-1,hc.size(-1)).float(),(hc.size(-1),))
                    z=torch.relu(blk.m.fc(hn.to(hc.dtype))).square()
                    pred=z@blk.m.pr.weight.T
                    res=pred-hn.to(pred.dtype)
                    gw=res.T@z/z.size(0)
                    blk.m.pr.weight.data-=lr*gw.to(blk.m.pr.weight.dtype)
        for i in li:s.bl[i].m.pr.weight.data=ow[i]
        return torch.cat(ap,dim=1)

class EMA:
    def __init__(s,m,d=0.999):s.m=m;s.d=d;s.sh={n:p.data.clone() for n,p in m.named_parameters()};s.bk={}
    def up(s):
        for n,p in s.m.named_parameters():s.sh[n].mul_(s.d).add_(p.data,alpha=1.-s.d)
    def ap(s):
        s.bk={};
        for n,p in s.m.named_parameters():s.bk[n]=p.data.clone();p.data.copy_(s.sh[n])
    def re(s):
        for n,p in s.m.named_parameters():p.data.copy_(s.bk[n])
        s.bk={}

def main():
    global zp5
    code=Path(__file__).read_text(encoding="utf-8")
    a=H();zp5=torch.compile(zp5)
    dd="RANK" in os.environ and "WORLD_SIZE" in os.environ
    rk=int(os.environ.get("RANK","0"));ws=int(os.environ.get("WORLD_SIZE","1"))
    lr_=int(os.environ.get("LOCAL_RANK","0"))
    ga=max(1,8//ws);gs=1./ga
    if not torch.cuda.is_available():raise RuntimeError("CUDA required")
    dev=torch.device("cuda",lr_);torch.cuda.set_device(dev)
    if dd:dist.init_process_group("nccl",device_id=dev);dist.barrier()
    ma=rk==0
    torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True
    from torch.backends.cuda import enable_flash_sdp,enable_math_sdp,enable_mem_efficient_sdp,enable_cudnn_sdp
    enable_flash_sdp(True);enable_math_sdp(False);enable_mem_efficient_sdp(False);enable_cudnn_sdp(False)
    lf=None
    if ma:os.makedirs("logs",exist_ok=True);lf=f"logs/{a.rid}.txt";print(lf)
    def l0(m,c=True):
        if not ma:return
        if c:print(m)
        if lf:
            with open(lf,"a") as f:print(m,file=f)
    l0(code,console=False);l0(f"Python {sys.version}",console=False);l0(f"PyTorch {torch.__version__}",console=False)
    try:l0(subprocess.run(["nvidia-smi"],capture_output=True,text=True,check=False).stdout,console=False)
    except FileNotFoundError:pass
    random.seed(a.seed);np.random.seed(a.seed);torch.manual_seed(a.seed);torch.cuda.manual_seed_all(a.seed)
    sp=spm.SentencePieceProcessor(model_file=a.tp)
    bbl,hsl,ibl=build_sp_luts(sp,a.V,dev)
    vt=lvt(a.vf,a.swl);l0(f"val_tokens:{vt.numel()}")
    bm=RGPT(a).to(dev).bfloat16()
    cm=torch.compile(bm,dynamic=False,fullgraph=True)
    mdl=DDP(cm,device_ids=[lr_],broadcast_buffers=False) if dd else cm
    nu=sum(p.numel() for p in bm.parameters())
    ed=a.nul*a.nr
    l0(f"params:{nu} eff_depth:{ed} train_loops:{a.nr} eval_loops:{bm._er}")
    l0(f"ws:{ws} ga:{ga}")
    bp=list(bm.bl.named_parameters())
    mp=[p for n,p in bp if p.ndim==2 and not any(c in n for c in CP)]
    sp_=[p for n,p in bp if p.ndim<2 or any(c in n for c in CP)]
    ot=torch.optim.Adam([{"params":[bm.te.weight],"lr":a.elr,"base_lr":a.elr}],betas=(a.b1,a.b2),eps=a.ae,fused=True)
    om=Muon(mp,lr=a.mlr,mom=a.mmo,bs=a.mbs,wd=a.mwd)
    for g in om.param_groups:g["base_lr"]=a.mlr
    os_=torch.optim.Adam([{"params":sp_,"lr":a.slr,"base_lr":a.slr}],betas=(a.b1,a.b2),eps=a.ae,fused=True)
    opts=[ot,om,os_]
    ema=EMA(bm,d=0.999);ess=int(a.iters*a.esf)
    mms=1000.*a.mws if a.mws>0 else None
    def lrm(st,el):
        if a.wdi<=0:return 1.
        if mms is None:
            w=max(a.iters-a.wdi,0)
            return max((a.iters-st)/max(a.wdi,1),0.) if w<=st<a.iters else 1.
        sm=el/max(st,1);rm=max(mms-el,0.);wm=a.wdi*sm
        return rm/max(wm,1e-9) if rm<=wm else 1.
    def za():[o.zero_grad(set_to_none=True) for o in opts]
    if a.wui>0:
        im={n:t.detach().cpu().clone() for n,t in bm.state_dict().items()}
        io_=[copy.deepcopy(o.state_dict()) for o in opts]
        mdl.train();tw=DTL(a.tf,rk,ws,dev)
        for _ in range(a.wui):
            za()
            for mi in range(ga):
                if dd:mdl.require_backward_grad_sync=(mi==ga-1)
                x,y=tw.nb(a.tbt,a.tsl,ga)
                with torch.autocast("cuda",torch.bfloat16):(mdl(x,y)*gs).backward()
            for o in opts:o.step()
            za()
        bm.load_state_dict(im,strict=True)
        for o,s in zip(opts,io_):o.load_state_dict(s)
        za()
        if dd:mdl.require_backward_grad_sync=True
    tl=DTL(a.tf,rk,ws,dev);tms=0.;ss=None
    torch.cuda.synchronize();t0=time.perf_counter();step=0
    while True:
        ls=step==a.iters or(ss is not None and step>=ss)
        dv=ls or(a.vle>0 and step%a.vle==0)
        if dv:
            torch.cuda.synchronize();tms+=1000.*(time.perf_counter()-t0)
            vl,vb=eval_sw(a,mdl,rk,ws,dev,vt,bbl,hsl,ibl,ttt=False)
            l0(f"step:{step}/{a.iters} val_loss:{vl:.4f} val_bpb:{vb:.4f} train_ms:{tms:.0f} step_avg:{tms/max(step,1):.2f}ms")
            torch.cuda.synchronize();t0=time.perf_counter()
        if ls:
            if ma:
                ema.ap();l0("EMA+TTT eval...")
                vle,vbe=eval_sw(a,bm,rk,ws,dev,vt,bbl,hsl,ibl,ttt=True)
                l0(f"ema_ttt val_loss:{vle:.4f} val_bpb:{vbe:.4f}")
                sd=bm.state_dict();obj,st=qsd(sd,a.gb,a.sdn)
                buf=io.BytesIO();torch.save(obj,buf)
                cmp=zlib.compress(buf.getvalue(),level=9)
                cb=len(code.encode());mb=len(cmp);tb=cb+mb
                l0(f"artifact code:{cb} model:{mb} total:{tb} ({tb/1e6:.3f}MB) params:{st['pc']}")
                sd2=dqsd(obj);bm.load_state_dict(sd2,strict=True)
                vl2,vb2=eval_sw(a,bm,rk,ws,dev,vt,bbl,hsl,ibl,ttt=True)
                l0(f"quant+ttt val_loss:{vl2:.4f} val_bpb:{vb2:.4f}")
                ema.re()
            break
        if ss is None and mms is not None:
            torch.cuda.synchronize()
            el=1000.*(time.perf_counter()-t0)+tms
            if el>=mms:ss=step+1
        za()
        for mi in range(ga):
            if dd:mdl.require_backward_grad_sync=(mi==ga-1)
            x,y=tl.nb(a.tbt,a.tsl,ga)
            with torch.autocast("cuda",torch.bfloat16):(mdl(x,y)*gs).backward()
        torch.cuda.synchronize();el=1000.*(time.perf_counter()-t0)+tms
        m=lrm(step,el)
        for o in opts:
            for g in o.param_groups:g["lr"]=g["base_lr"]*m
        for o in opts:o.step()
        if step>=ess:ema.up()
        if step%a.tle==0 and ma:l0(f"step:{step} lr_mul:{m:.4f}")
        step+=1

if __name__=="__main__":main()