Add VRAM-first cache release hook

README.md

@@ -73,7 +73,6 @@ Structured-mask update 2026-05-29: the sublinear backend now accepts symbolic ca
 SAT block-causal, and unrestricted/NAT mask rules. This removes dense O(T^2) mask
 allocation for long context, and also gathers ALiBi bias directly for selected
 local/anchor keys instead of materializing dense `[heads x T x T]` bias tensors.
-A trainer heartbeat and post-checkpoint CUDA cache clear were added for easier
-long-running Vast monitoring.
+A trainer heartbeat, post-checkpoint CUDA cache clear, and optional `--empty_cache_every_steps` hook were added for easier long-running Vast monitoring and VRAM-first allocator behavior.
 
 License: Apache-2.0 (matching the upstream method).

nB300_agillm4_vram_dblock.py

@@ -2376,6 +2376,12 @@ def _train_phase(
         seen_tok += toks_processed
         pbar.set_postfix(loss=f"{loss_value:.3f}", B=BATCH, L=BLOCK)
         pbar.update(toks_processed)
+        empty_cache_every = int(getattr(args, "empty_cache_every_steps", 0) or 0)
+        if DEV.type == "cuda" and empty_cache_every > 0 and (step % empty_cache_every) == 0:
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
         heartbeat_every = int(getattr(args, "heartbeat_every_sec", 300) or 0)
         now_mono = time.monotonic()
         if heartbeat_every > 0 and now_mono - last_heartbeat_mono >= heartbeat_every:
@@ -2893,6 +2899,8 @@ def main():
     tr.add_argument("--save_every_sec", type=int, default=DEFAULT_SAVE_SEC)
     tr.add_argument("--heartbeat_every_sec", type=int, default=300,
                     help="Print lightweight trainer heartbeat/status lines every N seconds; 0 disables.")
+    tr.add_argument("--empty_cache_every_steps", type=int, default=0,
+                    help="Call torch.cuda.empty_cache() every N train steps; useful for VRAM-first runs where lower reserved VRAM matters more than speed.")
     tr.add_argument("--delta_every_steps", type=int, default=DEFAULT_DELTA_STEPS, help="Weight-only delta save every N steps (0=off)")
     tr.add_argument("--delta_max_keep", type=int, default=DEFAULT_MAX_DELTAS, help="Max delta checkpoints to keep")
     tr.add_argument("--resume_delta", type=str, help="Resume from a delta (weight-only, no optimizer state)")

relaunch_agillm4_dblock.sh

@@ -20,4 +20,5 @@ exec python -u nB300_agillm4.py train --preset agillm4_floor --resume "$CKPT" \
   --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 --heartbeat_every_sec "${AGILLM4_HEARTBEAT_EVERY_SEC:-300}" \
+  --empty_cache_every_steps "${AGILLM4_EMPTY_CACHE_EVERY_STEPS:-1}" \
   --delta_every_steps 25000 --delta_max_keep 1 --max_ckpts 1

relaunch_agillm4_dblock_tied.sh

@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+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=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
+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_TIED $(date -u +%Y-%m-%dT%H:%M:%SZ) resume=$CKPT --dblock --tie_weights --attn_backend sublinear (fused_ce fixed)"
+exec python -u nB300_agillm4.py train --preset agillm4_floor --resume "$CKPT" \
+  --dblock --dblock_blocks "${AGILLM4_DBLOCKS:-4}" --dblock_schedule "${AGILLM4_DBLOCK_SCHEDULE:-loss_balanced}" \
+  --dblock_warmup_steps "${AGILLM4_DBLOCK_WARMUP:-16}" --dblock_sigma_curriculum_steps "${AGILLM4_DBLOCK_SIGMA_CURRICULUM:-2000}" \
+  --dblock_log_every "${AGILLM4_DBLOCK_LOG_EVERY:-25}" --tie_weights \
+  --batch_size 1 --block 1280 --amp --attn_backend sublinear --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 --heartbeat_every_sec "${AGILLM4_HEARTBEAT_EVERY_SEC:-300}" \
+  --empty_cache_every_steps "${AGILLM4_EMPTY_CACHE_EVERY_STEPS:-1}" \
+  --delta_every_steps 25000 --delta_max_keep 1 --max_ckpts 1