Fold VRAM-first dblock mode into AGILLM4

README.md

@@ -37,6 +37,14 @@ whose released code is ViT/classification only.
 | 4096 | 11.08 GB | 8.68 GB |
 | 8192 | 22.11 GB | 19.37 GB |
 
+## Official-line integration snapshot (2026-05-29)
+- `nB300_agillm4_vram_dblock.py` is the patched official trainer snapshot from the Vast box.
+- `dblocks_train.py` is the folded-in low-VRAM training step used by `--dblock`.
+- `relaunch_agillm4_dblock.sh` restarts the official line with `--dblock --tie_weights --attn_backend sublinear`.
+- `--tie_weights` now means AR, SAT, and NAT share the embedding projection tensor. This drops the live parameter count from 1,213,418,242 to 716,595,202.
+- Old untied checkpoint head matrices are intentionally skipped under tied mode; core weights still warm-start and the optimizer can rebuild.
+- SAT now uses fused vocab-streaming CE in the dblock path, and the dblock step releases AR/SAT activations before moving to the next objective.
+
 ## 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.
@@ -44,8 +52,11 @@ whose released code is ViT/classification only.
 - 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.
+Status update 2026-05-29: Scott chose the VRAM-first route. The official AGILLM-4
+training line is now the DiffusionBlocks mode folded into `nB300_agillm4.py`, with
+`--dblock --tie_weights --attn_backend sublinear`. Checkpoint compatibility is best-effort
+only: old untied AR/SAT/NAT head tensors are skipped when tied heads are active, and the
+optimizer state is allowed to reset. The priority is lower VRAM over preserving every
+old training assumption.
 
 License: Apache-2.0 (matching the upstream method).

dblocks_train.py

@@ -0,0 +1,67 @@
+"""DiffusionBlocks training mode folded into AGILLM-4 (gated by --dblock).
+
+Block-wise EDM denoising on the real Encoder blocks, supervising AR + SAT(fixed+var)
++ NAT each step on ONE block, with grad-checkpointed layers and fused vocab-streaming
+CE. Reuses the live data stream / optimizer / checkpointing of nB300_agillm4.
+Lazy-imports nB300 inside functions to avoid a circular import.
+"""
+import math, random, 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
+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 _block_sigmas(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 _dblock_init(core, args):
+    B=int(getattr(args,"dblock_blocks",4)); L=len(core.blocks); sp=max(1,L//B)
+    asg=[list(range(i*sp,(i+1)*sp)) for i in range(B)]; asg[-1]=list(range((B-1)*sp,L))
+    print(f"[dblock] DiffusionBlocks mode: {L} layers -> {B} blocks {asg}")
+    print(f"[dblock] equi-prob sigma boundaries: {[round(x,3) for x in _block_sigmas(B)]}")
+    return {"B":B,"assign":asg,"bsig":_block_sigmas(B)}
+
+def _dblock_step(core, ar_h, sat_h, nat_h, opt, scaler, args, ids, state):
+    import nB300_agillm4 as M
+    B=state["B"]; asg=state["assign"]; bs=state["bsig"]; T=ids.size(1)
+    bi=random.randrange(B); lo,hi=sorted([bs[bi],bs[bi+1]]); layers=asg[bi]
+    sig=torch.from_numpy(np.exp(np.random.uniform(math.log(max(lo,1e-4)),math.log(hi),ids.size(0))).astype("float32")).to(ids.device)
+    cs,co,ci=_edm_pre(sig); w=_edm_w(sig); SATB=M.SAT_BLOCK
+    # ---- AR: causal diffusion denoise ----
+    with M.amp(args.amp):
+        emb=core.emb(ids); zt=emb+sig[:,None,None]*torch.randn_like(emb); h=ci*zt
+        for li in layers: h=_ck.checkpoint(core.blocks[li], h, M.causal_mask(T), use_reentrant=False)
+        Dn=core.ln(cs*zt+co*h)
+    ar=w*fused_ce(Dn[:,:-1].contiguous(), ar_h.proj.weight, ids[:,1:].contiguous())
+    scaler.scale(ar).backward()
+    ar_val=float(ar.detach())
+    del emb, zt, h, Dn, ar
+    # ---- SAT: block-causal diffusion; fixed proj + variable gate ----
+    with M.amp(args.amp):
+        emb2=core.emb(ids); zt2=emb2+sig[:,None,None]*torch.randn_like(emb2); h2=ci*zt2
+        for li in layers: h2=_ck.checkpoint(core.blocks[li], h2, M.sat_mask(T), use_reentrant=False)
+        Ds=core.ln(cs*zt2+co*h2); last=Ds[:,-SATB:]
+        satf=fused_ce(last.contiguous(), sat_h.proj.weight, ids[:,1:SATB+1].contiguous())
+        satv=(M.EMIT_LAMBDA*F.cross_entropy(sat_h.gate(Ds[:,0].float()), torch.ones(ids.size(0),dtype=torch.long,device=ids.device))) if sat_h.gate is not None else 0.0
+        sat=w*(satf+satv)
+    scaler.scale(sat).backward()
+    sat_val=float(sat.detach())
+    del emb2, zt2, h2, Ds, last, satf, satv, sat
+    # ---- NAT: bidirectional mask-predict ----
+    nat_val=0.0
+    if nat_h is not None:
+        ratio=min(max(float(getattr(args,"nat_mask_ratio",0.5)),0.05),0.95)
+        with M.amp(args.amp):
+            nat_ids=ids.clone(); m=torch.rand(ids.shape,device=ids.device)<ratio
+            if not bool(m.any()): m[...,-1]=True
+            nat_ids[m]=M.BLANK; hn=core.emb(nat_ids)
+            for li in layers: hn=_ck.checkpoint(core.blocks[li], hn, None, use_reentrant=False)
+            Dnat=core.ln(hn)
+        nat=fused_ce(Dnat[m], nat_h.proj.weight, ids[m]); scaler.scale(nat).backward(); nat_val=float(nat.detach()); del nat_ids, m, hn, Dnat, nat
+    scaler.unscale_(opt)
+    nn.utils.clip_grad_norm_([p for g in opt.param_groups for p in g["params"]],1.0)
+    scaler.step(opt); scaler.update(); opt.zero_grad(set_to_none=True)
+    return ar_val+sat_val+nat_val

relaunch_agillm4_dblock.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# OFFICIAL LINE: DiffusionBlocks block-wise denoising (low VRAM). Resumes newest ckpt.
+set -Eeuo pipefail
+cd /workspace/agillm-4
+export TOKENIZERS_PARALLELISM=false
+export TOKENIZER_ID="${TOKENIZER_ID:-deepseek-ai/DeepSeek-V4-Pro}"
+export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512,expandable_segments:True
+export AGILLM_ATTN_BACKEND="${AGILLM_ATTN_BACKEND:-sublinear}"
+if [ -f /root/.cache/huggingface/token ]; then export HF_TOKEN="$(tr -d '\r\n' </root/.cache/huggingface/token)"; export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"; fi
+SAVE_DIR=/workspace/agillm4_4090_ckpts
+CKPT="$(ls -1t "$SAVE_DIR"/pretrain_step*.pt 2>/dev/null | head -1)"
+[ -n "$CKPT" ] || { echo "no ckpt" >&2; exit 1; }
+exec >> /workspace/agillm4_floor_train.log 2>&1
+echo "RELAUNCH_AGILLM4_DBLOCK $(date -u +%Y-%m-%dT%H:%M:%SZ) resume=$CKPT --dblock blocks=${AGILLM4_DBLOCKS:-4} tie_weights=1 attn=${AGILLM_ATTN_BACKEND}"
+exec python -u nB300_agillm4.py train --preset agillm4_floor --resume "$CKPT" \
+  --dblock --dblock_blocks "${AGILLM4_DBLOCKS:-4}" --tie_weights \
+  --batch_size 1 --block "${AGILLM4_BLOCK:-1280}" --amp --attn_backend "${AGILLM_ATTN_BACKEND}" --grad_checkpoint \
+  --optimizer paged_adamw8bit --sat_every 1 --nat_every 1 --nat_max_tokens 768 --nat_mask_ratio 0.5 \
+  --token_param_ratio 100 --save_dir "$SAVE_DIR" \
+  --save_every_sec 86400 --delta_every_steps 25000 --delta_max_keep 1 --max_ckpts 1