Tune DBlock backward-math speed line

README.md

@@ -94,3 +94,6 @@ Quality target update 2026-05-29: Scott clarified the previous AGILLM run was ro
 
 
 License: Apache-2.0 (matching the upstream method).
+
+Backward-math speed update 2026-05-29: live profiling confirmed the expensive path is AR transformer backward/recompute, not fused CE, data loading, or optimizer step. A clean context/batch sweep found B6/L1024 gives the best raw throughput among quality-preserving candidates (6144 tokens/step, ~3.3k tok/s profiler math) while keeping a 1024-token context. A second objective-mix sweep tested a harder experimental path: reducing full causal AR steps from 85% to 70% and moving the probability mass to SAT/NAT (`--dblock_ar_prob 0.70 --dblock_sat_prob 0.15 --dblock_nat_prob 0.15`). Clean profiler math reached ~3440 tok/s, and the live relaunch now uses B6/L1024 + ar70. The relaunch script intentionally unsets `PYTORCH_CUDA_ALLOC_CONF` because the previous expandable-segments allocator reserved ~23 GB and caused memory-pressure slowdown on this shape; without it, the live line returns to ~15.1 GB tensor peak / ~16.5-17.4 GB observed VRAM. This is experimental by design: it trades some AR-step density for faster multi-objective DBlock training while preserving checkpoint warm start and the 55 tokens/param target.
+

dblocks_train.py

@@ -182,21 +182,50 @@ def _update_stats(state, bi, loss_value):
     state["step"] = int(state.get("step", 0)) + 1
 
 
-def _run_block(block, x, mask, use_checkpoint):
+def _activation_offload_enabled(args):
+    return bool(getattr(args, "dblock_activation_offload", False)) and torch.cuda.is_available()
+
+
+def _activation_offload_hooks(args):
+    min_bytes = int(float(getattr(args, "dblock_activation_offload_min_mb", 1.0) or 1.0) * 1024 * 1024)
+
+    def pack(t):
+        if not torch.is_tensor(t) or not t.is_cuda or not t.is_floating_point() or t.numel() * t.element_size() < min_bytes:
+            return t
+        return ("cpu_offload", t.device, t.detach().to("cpu", non_blocking=True))
+
+    def unpack(x):
+        if isinstance(x, tuple) and len(x) == 3 and x[0] == "cpu_offload":
+            _, dev, cpu_t = x
+            return cpu_t.to(dev, non_blocking=True)
+        return x
+
+    return torch.autograd.graph.saved_tensors_hooks(pack, unpack)
+
+
+def _run_block(block, x, mask, use_checkpoint, args=None):
     if use_checkpoint:
         return _ck.checkpoint(lambda y, block=block: block(y, mask), x, use_reentrant=False)
+    if args is not None and _activation_offload_enabled(args):
+        with _activation_offload_hooks(args):
+            return block(x, mask)
     return block(x, mask)
 
 
-def _dblock_checkpoint_this_layer(args, base_enabled, layer_pos):
+def _dblock_checkpoint_this_layer(args, base_enabled, layer_pos, layer_count=None):
     if not base_enabled:
         return False
+    pos = int(layer_pos)
+    count = int(layer_count or 0)
+    skip_tail = max(0, int(getattr(args, "dblock_checkpoint_skip_tail", 0) or 0))
+    if skip_tail > 0 and count > 0 and pos >= max(0, count - skip_tail):
+        return False
     stride = int(getattr(args, "dblock_checkpoint_stride", 1) or 1)
     if stride <= 0:
         return False
     if stride == 1:
         return True
-    return (int(layer_pos) % stride) == 0
+    return (pos % stride) == 0
 
 
 def _sample_token_loss_inputs(hidden, targets, max_tokens):
@@ -293,7 +322,7 @@ def _dblock_step(core, ar_h, sat_h, nat_h, opt, scaler, args, ids, state):
             zt = emb + sig[:, None, None] * torch.randn_like(emb)
             h = ci * zt
             for lpos, li in enumerate(layers):
-                h = _run_block(core.blocks[li], h, causal, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos))
+                h = _run_block(core.blocks[li], h, causal, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args)
             Dn = core.ln(cs * zt + co * h)
         _profile_toc(state, "ar_forward", _t)
         _t = _profile_tic(prof)
@@ -316,7 +345,7 @@ def _dblock_step(core, ar_h, sat_h, nat_h, opt, scaler, args, ids, state):
             zt2 = emb2 + sig[:, None, None] * torch.randn_like(emb2)
             h2 = ci * zt2
             for lpos, li in enumerate(layers):
-                h2 = _run_block(core.blocks[li], h2, smask, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos))
+                h2 = _run_block(core.blocks[li], h2, smask, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args)
             Ds = core.ln(cs * zt2 + co * h2)
             last = Ds[:, -SATB:]
         _profile_toc(state, "sat_forward", _t)
@@ -355,7 +384,7 @@ def _dblock_step(core, ar_h, sat_h, nat_h, opt, scaler, args, ids, state):
             nat_in[m] = M.BLANK
             hn = core.emb(nat_in)
             for lpos, li in enumerate(layers):
-                hn = _run_block(core.blocks[li], hn, None, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos))
+                hn = _run_block(core.blocks[li], hn, None, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args)
             Dnat = core.ln(hn)
         _profile_toc(state, "nat_forward", _t)
         _t = _profile_tic(prof)

nB300_agillm4_vram_dblock.py

@@ -3039,6 +3039,12 @@ def main():
                     help="Print DBlock block/loss/VRAM diagnostics every N DBlock steps; 0 disables.")
     tr.add_argument("--dblock_checkpoint_stride", type=int, default=1,
                     help="With --grad_checkpoint in --dblock mode, checkpoint one layer every N selected block layers; 1=all layers, 2=alternate, 0=off.")
+    tr.add_argument("--dblock_checkpoint_skip_tail", type=int, default=0,
+                    help="Experimental DBlock speed knob: do not checkpoint this many final layers in the selected block, reducing backward recompute at higher VRAM cost.")
+    tr.add_argument("--dblock_activation_offload", action="store_true",
+                    help="Experimental DBlock speed knob: for non-checkpointed block layers, offload saved backward tensors to CPU RAM instead of recomputing.")
+    tr.add_argument("--dblock_activation_offload_min_mb", type=float, default=1.0,
+                    help="Minimum CUDA tensor size in MB to offload under --dblock_activation_offload.")
     tr.add_argument("--dblock_sigma_curriculum_steps", type=int, default=2000,
                     help="Warm sigma ranges from easy to full span over this many DBlock steps; 0 disables.")
     tr.add_argument("--dblock_edm_wmax", type=float, default=5.0,

relaunch_agillm4_dblock_sg2.sh

@@ -4,20 +4,20 @@ 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
+unset PYTORCH_CUDA_ALLOC_CONF  # B6/L1024 ar70: avoid near-full allocator reservation slowdown
 export AGILLM_ATTN_BACKEND=sublinear
 [ -f /root/.cache/huggingface/token ] && { export HF_TOKEN="$(tr -d '\r\n' </root/.cache/huggingface/token)"; export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"; }
 SAVE_DIR=/workspace/agillm4_4090_ckpts
 TOKEN_PARAM_RATIO="${TOKEN_PARAM_RATIO:-55}"
 CKPT="$(ls -1t "$SAVE_DIR"/pretrain_step*.pt 2>/dev/null | head -1)"
 exec >> /workspace/agillm4_floor_train.log 2>&1
-echo "RELAUNCH_AGILLM4_DBLOCK_SG2 $(date -u +%Y-%m-%dT%H:%M:%SZ) resume=$CKPT (quality ratio55 + batch4 + sublinear v2)"
+echo "RELAUNCH_AGILLM4_DBLOCK_SG2 $(date -u +%Y-%m-%dT%H:%M:%SZ) resume=$CKPT (quality ratio55 + B6/L1024 + ar70 backward-math + sublinear v2)"
 exec python -u nB300_agillm4.py train --preset agillm4_floor --resume "$CKPT" \
   --dblock --dblock_blocks 4 --dblock_schedule loss_balanced --dblock_warmup_steps 16 \
   --dblock_sigma_curriculum_steps 2000 --dblock_log_every 25 --dblock_objective_mode stochastic \
-  --dblock_ar_prob 0.85 --dblock_sat_prob 0.075 --dblock_nat_prob 0.075 \
+  --dblock_ar_prob 0.70 --dblock_sat_prob 0.15 --dblock_nat_prob 0.15 \
   --dblock_ar_loss_tokens 512 --dblock_sat_loss_tokens 0 --dblock_nat_loss_tokens 512 \
-  --tie_weights --batch_size 4 --block 1280 --amp --attn_backend sublinear \
+  --tie_weights --batch_size 6 --block 1024 --amp --attn_backend sublinear \
   --sublinear_window 128 --sublinear_stride 128 --sublinear_max_anchors 128 --sublinear_chunk 128 \
   --sublinear_sinks 4 --sublinear_recent_anchors 64 --no-sublinear_pooled_landmarks \
   --grad_checkpoint --dblock_checkpoint_stride 1 --optimizer paged_adamw8bit --sat_every 4 --nat_every 4 --nat_max_tokens 768 --nat_mask_ratio 0.5 \