AGILLM4-DiffusionBlocks: block-wise AR+SAT+NAT denoising, fused CE, tied heads

README.md

@@ -0,0 +1,42 @@
+# AGILLM4-DiffusionBlocks
+
+Block-wise (DiffusionBlocks-style) training adapted to the **AGILLM-4** decoder-only
+LLM, with all three AGILLM-4 heads (AR + SAT fixed/variable + NAT) and a set of
+long-context memory upgrades. Adapts & improves on
+[SakanaAI/DiffusionBlocks](https://github.com/SakanaAI/DiffusionBlocks) (ICLR 2026),
+whose released code is ViT/classification only.
+
+## What's here
+- `dblocks_agillm4.py` — self-contained DiffusionBlocks prototype (proves the
+  per-step 1/B gradient locality on a small transformer).
+- `dblocks_agillm4_lm.py` — real-architecture trainer: reuses AGILLM-4's `Encoder`
+  blocks (ALiBi, causal `Block`), partitions 28 layers into B EDM-noise blocks,
+  trains one block per step as an EDM denoiser, supervised by:
+    * **AR** — causal next-token CE
+    * **SAT** — fixed (proj CE over SAT_BLOCK, block-causal mask) **and** variable (gate CE)
+    * **NAT** — bidirectional mask-predict CE
+  Memory upgrades: EDM-weight clamp, AMP bf16, **tied shared vocab projection**
+  (1.21B -> 0.72B params), grad-checkpointed blocks, and `fused_ce`.
+- `fused_ce.py` — fused cross-entropy that streams over the 129k vocab via
+  online-softmax with a custom backward, never materializing the `[T x 129280]`
+  logit matrix (the DiffusionBlocks "process in chunks" idea applied to the head).
+
+## Measured (real AGILLM-4 floor, 28L, 0.72B tied)
+| ctx | full (28L) | one block (7L) |
+|----:|-----------:|---------------:|
+| 1280 | 7.48 GB | 5.53 GB |
+| 4096 | 11.08 GB | 8.68 GB |
+| 8192 | 22.11 GB | 19.37 GB |
+
+## Honest findings
+- DiffusionBlocks and gradient-checkpointing are **substitutes** for activation
+  memory; with checkpointing on, the 28->7 layer saving is only ~1.1-1.3x.
+- The big fixed-cost win was **tied heads + fused CE** (ctx-1280 ~24 GB -> ~5.5 GB).
+- The long-context ceiling (16k+) is **attention-bound**, so the next lever is the
+  sublinear attention backend, not more block/CE work.
+
+Status: validated research prototype. The official AGILLM-4 training line remains the
+proven AR/SAT/NAT end-to-end trainer; these wins (tied heads, fused-CE, sublinear
+attention) are intended to be folded into it for long context.
+
+License: Apache-2.0 (matching the upstream method).

dblocks_agillm4.py

@@ -0,0 +1,82 @@
+"""DiffusionBlocks prototype for AGILLM-4 (block-wise denoising training).
+
+Faithful to SakanaAI/DiffusionBlocks: EDM (Karras) diffusion, equi-probability
+block sigma partitioning, and the key property -- each step trains ONE block as a
+denoiser for its noise range, so backprop touches only that block => training
+memory ~ 1/B of end-to-end.  Here we reuse a transformer-block stack (stand-in for
+AGILLM-4's Encoder.blocks: same pre-norm residual Block API x = blk(x, mask)).
+"""
+import math, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+def _cdf(x): return 0.5*(1+math.erf(x/math.sqrt(2)))
+def _ppf(p): return float(torch.erfinv(torch.tensor(2*p-1.0))*math.sqrt(2))
+
+# ---- EDM block partitioning (verbatim mechanism from dblock_modules.py) ----
+def get_block_sigmas(num_blocks, sigma_min=0.002, sigma_max=80.0, p_mean=-1.2, p_std=1.2):
+    cdf_min = _cdf((np.log(sigma_min)-p_mean)/p_std)
+    cdf_max = _cdf((np.log(sigma_max)-p_mean)/p_std)
+    out=[]
+    for i in range(num_blocks+1):
+        p = cdf_min + (cdf_max-cdf_min)*(i/num_blocks)
+        out.append(float(np.exp(p_mean+p_std*_ppf(p))))
+    return out   # descending->ascending sigma boundaries, equal CDF mass per block
+
+class DBlockLM(nn.Module):
+    """Shared token emb + output proj; n_layers split into B independent blocks."""
+    def __init__(self, vocab=2048, d=256, n_layers=8, heads=4, num_blocks=4):
+        super().__init__()
+        self.emb = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList([nn.TransformerEncoderLayer(d, heads, 4*d,
+                          batch_first=True, norm_first=True) for _ in range(n_layers)])
+        self.ln = nn.LayerNorm(d); self.out = nn.Linear(d, vocab)
+        self.num_blocks = num_blocks
+        split = n_layers // num_blocks
+        self.assign = [list(range(i*split,(i+1)*split)) for i in range(num_blocks)]
+        self.block_sigmas = get_block_sigmas(num_blocks)
+        self.sigma_data = 0.5
+
+    def run_block(self, b, zt, sigma):
+        # EDM preconditioning (Karras 2022), exactly as DiffusionBlocks.denoise()
+        s = sigma[:,None,None]
+        c_skip = self.sigma_data**2/(s**2+self.sigma_data**2)
+        c_out  = s*self.sigma_data/(s**2+self.sigma_data**2)**0.5
+        c_in   = 1/(s**2+self.sigma_data**2)**0.5
+        h = zt*c_in
+        for li in self.assign[b]:          # <-- ONLY this block's layers run
+            h = self.blocks[li](h)
+        return c_skip*zt + c_out*h         # denoiser D_theta(zt,sigma) -> predicts z0
+
+def edm_weight(sigma, sigma_data=0.5):
+    return (sigma**2+sigma_data**2)/(sigma*sigma_data)**2
+
+def train_step(model, opt, ids):
+    z0 = F.normalize(model.emb(ids), p=2, dim=-1).detach()      # clean target embeds
+    b = np.random.randint(model.num_blocks)                     # pick one block
+    lo, hi = model.block_sigmas[b], model.block_sigmas[b+1]
+    lo, hi = min(lo,hi), max(lo,hi)
+    sigma = torch.from_numpy(np.exp(np.random.uniform(np.log(lo),np.log(hi),
+                              size=ids.shape[0])).astype('float32'))
+    zt = z0 + sigma[:,None,None]*torch.randn_like(z0)
+    D = model.run_block(b, zt, sigma)                           # backprop only block b
+    loss = (edm_weight(sigma)[:,None,None]*(D-z0)**2).mean()
+    opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
+    return b, loss.item()
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    NB=4; m=DBlockLM(num_blocks=NB); opt=torch.optim.AdamW(m.parameters(),1e-3)
+    ids=torch.randint(0,2048,(8,64))
+    print("layers:",len(m.blocks),"| blocks:",NB,"| layer assignment:",m.assign)
+    print("block sigma boundaries:",[round(s,3) for s in m.block_sigmas])
+    # one step, then verify ONLY the chosen block's layer params got gradients
+    b,loss=train_step(m,opt,ids)
+    def has_grad(mod): return any(p.grad is not None and p.grad.abs().sum()>0 for p in mod.parameters())
+    grad_blocks=[i for i in range(NB) if any(has_grad(m.blocks[li]) for li in m.assign[i])]
+    tot=sum(p.numel() for p in m.parameters())
+    blk_params=sum(p.numel() for li in m.assign[b] for p in m.blocks[li].parameters())
+    print(f"\\nstep trained block {b}  loss={loss:.4f}")
+    print("blocks whose layers received gradients:",grad_blocks,"(expect just",[b],")")
+    print(f"layer-params with grad this step: {blk_params:,} / {sum(p.numel() for blk in m.blocks for p in blk.parameters()):,}"
+          f"  ({100*blk_params/sum(p.numel() for blk in m.blocks for p in blk.parameters()):.0f}% = ~1/{NB})")
+    print("\\n--- a few more steps (each trains one block independently) ---")
+    for _ in range(6):
+        b,loss=train_step(m,opt,ids); print(f"  block {b}: loss={loss:.4f}")

dblocks_agillm4_lm.py

@@ -0,0 +1,110 @@
+"""AR + SAT(fixed+variable) + NAT DiffusionBlocks for AGILLM-4 (v3, long-context).
+
+Upgrades over v2:
+  * SAT fixed (proj CE over SAT_BLOCK, block-causal sat_mask) AND variable (gate CE)
+  * tied shared vocab projection (emb.weight) across AR/SAT/NAT -> drops ~0.5B params
+  * chunked cut-CE: logits computed per token-chunk on a detached hidden, backward
+    incrementally, then ONE block backward -> [T x 129280] never fully materialized
+  * sequential per-objective backward (one block-forward graph live at a time)
+  * AMP bf16 + clamped EDM weight (stable)
+"""
+import sys, math, argparse, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+import torch.utils.checkpoint as _ck
+from fused_ce import fused_ce
+sys.path.insert(0,"/workspace/agillm-4"); import nB300_agillm4 as M
+V=M.VOCAB; SATB=M.SAT_BLOCK; EMIT=getattr(M,"EMIT_LAMBDA",0.1); SD=0.5
+def _cdf(x): return 0.5*(1+math.erf(x/math.sqrt(2)))
+def _ppf(p): return float(torch.erfinv(torch.tensor(2*p-1.0))*math.sqrt(2))
+def bsig(B,smin=0.002,smax=80.0,pm=-1.2,ps=1.2):
+    a,b=_cdf((math.log(smin)-pm)/ps),_cdf((math.log(smax)-pm)/ps)
+    return [float(np.exp(pm+ps*_ppf(a+(b-a)*(i/B)))) for i in range(B+1)]
+def edm_pre(s): s=s[:,None,None]; return SD**2/(s**2+SD**2), s*SD/(s**2+SD**2)**0.5, 1/(s**2+SD**2)**0.5
+def edm_w(s,wmax=5.0): return float(((s**2+SD**2)/(s*SD)**2).clamp(max=wmax).mean())
+
+def make(B):
+    cfg=M.PRESETS["agillm4_floor"].copy()
+    core=M.Encoder(cfg, attn_backend="sdpa", grad_checkpoint=False).to(M.DEV)
+    sat_gate=nn.Linear(cfg["d"],2).to(M.DEV)          # SAT *variable* gate (tiny)
+    L=cfg["layers"]; sp=L//B
+    asg=[list(range(i*sp,(i+1)*sp)) for i in range(B)]; asg[-1]=list(range((B-1)*sp,L))
+    return core,sat_gate,asg,bsig(B)
+
+def runblk(core,h,layers,mask):
+    for li in layers: h=_ck.checkpoint(core.blocks[li], h, mask, use_reentrant=False)  # recompute in backward
+    return h
+
+def cut_ce_backward(core, hidden, targets, scale, chunk=1024, mask_pos=None):
+    """chunked CE on detached hidden -> grad; caller backprops block once."""
+    hd=hidden.detach().requires_grad_(True); T=hd.size(1)
+    if mask_pos is None: N=targets.numel()
+    else: N=int(mask_pos.sum().item()) or 1
+    tot=0.0
+    for s in range(0,T,chunk):
+        lg=F.linear(hd[:,s:s+chunk], core.emb.weight).float()      # tied proj, transient
+        tg=targets[:,s:s+chunk]
+        if mask_pos is None:
+            l=F.cross_entropy(lg.reshape(-1,V),tg.reshape(-1),reduction="sum")*scale/N
+        else:
+            mm=mask_pos[:,s:s+chunk]
+            if not bool(mm.any()): continue
+            l=F.cross_entropy(lg[mm],tg[mm],reduction="sum")*scale/N
+        l.backward(); tot+=float(l.detach())
+    return tot, hd.grad
+
+def step(core,sat_gate,ids,layers,lo,hi,nat_ratio=0.5,bw=True):
+    T=ids.size(1); ar=sat=nat=0.0
+    sig=torch.from_numpy(np.exp(np.random.uniform(math.log(lo),math.log(hi),ids.size(0))).astype("float32")).to(M.DEV)
+    cs,co,ci=edm_pre(sig); w=edm_w(sig)
+    # ---- AR: causal diffusion denoise ----
+    with torch.autocast("cuda",dtype=torch.bfloat16):
+        emb=core.emb(ids); zt=emb+sig[:,None,None]*torch.randn_like(emb)
+        Dn=core.ln(cs*zt+co*runblk(core,ci*zt,layers,M.causal_mask(T)))
+    ar_loss=w*fused_ce(Dn[:,:-1].contiguous(), core.emb.weight, ids[:,1:].contiguous())
+    if bw: ar_loss.backward()
+    ar=float(ar_loss.detach())
+    # ---- SAT: block-causal diffusion; fixed proj + variable gate ----
+    with torch.autocast("cuda",dtype=torch.bfloat16):
+        zt2=core.emb(ids)+sig[:,None,None]*torch.randn_like(emb)
+        Ds=core.ln(cs*zt2+co*runblk(core,ci*zt2,layers,M.sat_mask(T)))
+        last=Ds[:,-SATB:]
+        satf=F.cross_entropy(F.linear(last,core.emb.weight).float().reshape(-1,V), ids[:,1:SATB+1].reshape(-1))
+        satv=EMIT*F.cross_entropy(sat_gate(Ds[:,0].float()), torch.ones(ids.size(0),dtype=torch.long,device=M.DEV))
+        sat_loss=w*(satf+satv)
+    if bw: sat_loss.backward()
+    sat=float(sat_loss)
+    # ---- NAT: bidirectional mask-predict ----
+    with torch.autocast("cuda",dtype=torch.bfloat16):
+        nat_ids=ids.clone(); m=torch.rand(ids.shape,device=M.DEV)<nat_ratio
+        if not bool(m.any()): m[...,-1]=True
+        nat_ids[m]=M.BLANK
+        Dnat=core.ln(runblk(core,core.emb(nat_ids),layers,None))
+    nat_loss=fused_ce(Dnat[m], core.emb.weight, ids[m])
+    if bw: nat_loss.backward()
+    nat=float(nat_loss.detach())
+    return ar,sat,nat
+
+def peak(fn):
+    torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats(); fn(); torch.cuda.synchronize()
+    return torch.cuda.max_memory_allocated()/1e9
+
+if __name__=="__main__":
+    ap=argparse.ArgumentParser(); ap.add_argument("--blocks",type=int,default=4); a=ap.parse_args()
+    core,sg,asg,bs=make(a.blocks); nL=len(core.blocks); bi=a.blocks//2
+    P=sum(p.numel() for p in core.parameters())+sum(p.numel() for p in sg.parameters())
+    print(f"AGILLM-4 floor {nL}L, {P/1e9:.2f}B params (tied heads) | blocks={a.blocks} assign={asg}")
+    def full(ctx): step(core,sg,torch.randint(0,V,(1,ctx),device=M.DEV),list(range(nL)),bs[0],bs[-1]); core.zero_grad(set_to_none=True); sg.zero_grad(set_to_none=True)
+    def blk(ctx):  step(core,sg,torch.randint(0,V,(1,ctx),device=M.DEV),asg[bi],bs[bi],bs[bi+1]); core.zero_grad(set_to_none=True); sg.zero_grad(set_to_none=True)
+    print("\\n=== v3: tied-heads + cut-CE + AMP + SAT(fixed+var)+NAT ===")
+    for ctx in (1280,4096,8192,16384):
+        try: mf=peak(lambda:full(ctx)); sf=f"{mf:.2f} GB"
+        except RuntimeError as e: torch.cuda.empty_cache(); sf=("OOM" if "out of memory" in str(e).lower() else "ERR")
+        try: mb=peak(lambda:blk(ctx)); sb=f"{mb:.2f} GB"
+        except RuntimeError as e: torch.cuda.empty_cache(); sb=("OOM" if "out of memory" in str(e).lower() else "ERR")
+        print(f"  ctx={ctx:>6}: full(28L)={sf:>9}   one block({len(asg[bi])}L)={sb:>9}")
+    print("\\n--- smoke train (block-wise AR+SAT(fixed+var)+NAT) ---")
+    for st in range(4):
+        b=np.random.randint(a.blocks)
+        ar,sat,nat=step(core,sg,torch.randint(0,V,(1,1280),device=M.DEV),asg[b],bs[b],bs[b+1])
+        opt=torch.optim.SGD([p for li in asg[b] for p in core.blocks[li].parameters()]+list(sg.parameters()),1e-3)
+        opt.step(); opt.zero_grad(); core.zero_grad(set_to_none=True)
+        print(f"  step {st} block {b}: AR={ar:.2f} SAT={sat:.2f} NAT={nat:.2f}")

fused_ce.py

@@ -0,0 +1,31 @@
+"""Fused cross-entropy: streams over the VOCAB dimension (online-softmax) so the
+[N x V] logit matrix is NEVER materialized -- only [N x vchunk]. Custom backward
+recomputes softmax per vocab-chunk (grad = softmax - onehot). This is the
+DiffusionBlocks 'process in chunks, don't hold the whole thing' idea applied to
+the output head instead of network depth."""
+import torch
+class FusedCE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, h, W, tgt, vchunk=16384):
+        N,d=h.shape; V=W.shape[0]; hf=h.float()
+        m=torch.full((N,),-1e30,device=h.device); s=torch.zeros(N,device=h.device); zt=torch.zeros(N,device=h.device)
+        for c in range(0,V,vchunk):
+            lg=hf@W[c:c+vchunk].float().T                    # [N,vchunk] transient only
+            cm=lg.max(1).values; nm=torch.maximum(m,cm)
+            s=s*torch.exp(m-nm)+torch.exp(lg-nm[:,None]).sum(1); m=nm
+            ic=(tgt>=c)&(tgt<c+vchunk)
+            if ic.any(): zt[ic]=lg[ic,tgt[ic]-c]
+        lse=m+torch.log(s); ctx.save_for_backward(h,W,tgt,lse); ctx.vchunk=vchunk
+        return (lse-zt).mean()
+    @staticmethod
+    def backward(ctx, go):
+        h,W,tgt,lse=ctx.saved_tensors; vc=ctx.vchunk; N,d=h.shape; V=W.shape[0]; hf=h.float()
+        gh=torch.zeros_like(hf); gW=torch.zeros(W.shape,device=W.device,dtype=torch.float32); sc=float(go)/N
+        for c in range(0,V,vc):
+            Wc=W[c:c+vc].float(); p=torch.exp(hf@Wc.T-lse[:,None])     # softmax chunk [N,vchunk]
+            ic=(tgt>=c)&(tgt<c+vc)
+            if ic.any(): p[ic,tgt[ic]-c]-=1.0
+            p*=sc; gh+=p@Wc; gW[c:c+vc]+=p.T@hf
+        return gh.to(h.dtype), gW.to(W.dtype), None, None
+def fused_ce(h, W, tgt, vchunk=16384):
+    return FusedCE.apply(h.reshape(-1,h.size(-1)), W, tgt.reshape(-1), vchunk)