fix: OOM at batch=256 — cap batch by logits memory, enable grad ckpt

Root cause: batch=256 × seq=16 × vocab=200073 × 4B = 3.28 GB for logits
alone. Plus backward gradients of same size + activations across 28
layers → exceeds 32 GB RAM → system swaps → appears frozen.

Fix in loops.py:
- Add _safe_batch() that caps eff_batch so logits tensor stays under
max_logits_gb (default 2 GB). With vocab=200073:
batch=256,seq=16: logits=3.28GB → capped to batch=156,seq=16: 2.0GB
batch=128,seq=32: logits=3.28GB → capped to batch=78,seq=32: 2.0GB
batch=64,seq=64: logits=3.28GB → capped to batch=39,seq=64: 2.0GB
batch=32,seq=128: logits=3.28GB → capped to batch=19,seq=128: 1.95GB
- Enable gradient_checkpointing on model (recompute activations during
backward, saves ~60% activation memory at cost of ~30% more compute)

Fix in train_hyper.py:
- batch_size default 32→4 (base batch; GrowLength scales up with cap)
- GrowLength stages use fixed safe batch sizes directly"

chimera/training/loops.py

@@ -13,6 +13,24 @@ from .common import save_final_checkpoint, save_training_checkpoint
 from .hyper import ProgressiveLoopScheduler
 
 
+def _safe_batch(desired_batch: int, seq_len: int, vocab_size: int,
+                max_logits_gb: float = 2.0) -> int:
+    """Cap batch size so the logits tensor fits in memory.
+
+    Logits shape: [batch, seq, vocab] at fp32 = batch * seq * vocab * 4 bytes.
+    With vocab=200073, batch=256, seq=16: 3.28 GB just for logits.
+    Backward doubles this. Must stay well under 32 GB total.
+    """
+    bytes_per_sample = seq_len * vocab_size * 4  # fp32 logits
+    max_bytes = int(max_logits_gb * 1024**3)
+    max_batch = max(1, max_bytes // bytes_per_sample)
+    capped = min(desired_batch, max_batch)
+    if capped < desired_batch:
+        print(f"  [MEM] Batch {desired_batch} → {capped} (logits would be "
+              f"{desired_batch * seq_len * vocab_size * 4 / 1e9:.1f} GB, cap={max_logits_gb} GB)")
+    return capped
+
+
 def train_fast_loop(args, model, config, loader, compute_loss) -> str:
     optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.98))
     os.makedirs(args.output_dir, exist_ok=True)
@@ -52,41 +70,40 @@ def train_standard_loop(args, model, config, loader, compute_loss, optimizer, us
 
 def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer):
     use_compile = getattr(args, "compile", False)
+    vocab_size = int(config.get("vocab_size", 200073))
 
     # ── Muon LR for ternary BitLinear ──
-    # v12.1: Raised from 0.008 to 0.012. The clamp-aware STE in BitLinear
-    # gates gradients to zero for weights outside [-1, 1], so the effective
-    # learning signal is self-limiting. 0.012 is the highest rate before
-    # NS-orthogonalized momentum causes oscillation at the STE boundary.
-    # At 300 steps, every step counts — 0.008 converges too slowly.
     muon_lr = 0.012
-    muon_warmup = 30  # 10% of 300 steps; was min(args.warmup, 100)
+    muon_warmup = 30
 
     model, optimizer, scheduler, extras = chimera_turbo.apply(
         model,
         max_steps=args.max_steps,
         lr=muon_lr,
-        weight_decay=0.02,         # was 0.01; BitNet SLM: wd=0.05 optimal
+        weight_decay=0.02,
         warmup_steps=muon_warmup,
         use_compile=use_compile,
-        mtp_heads=0,               # vocab/hidden=781:1 → MTP noisy + slow
-        llrd_decay=0.90,           # was 0.92; 0.90^27=0.058 → more bottom grad
-        grokfast_alpha=0.95,       # was 0.98; shorter EMA window for 300 steps
-        grokfast_lambda=1.5,       # was 1.0; amplify slow grads more aggressively
+        mtp_heads=0,
+        llrd_decay=0.90,
+        grokfast_alpha=0.95,
+        grokfast_lambda=1.5,
     )
     model.train()
 
+    # ── Gradient checkpointing: saves ~60% activation memory ──
+    # Critical with vocab=200K: without it, activations across 28 layers
+    # at batch=32 can consume several GB.
+    raw_model = getattr(model, "_orig_mod", model)
+    if hasattr(raw_model, "enable_gradient_checkpointing"):
+        raw_model.enable_gradient_checkpointing()
+        print(f"[OPT] Gradient checkpointing: ON")
+
     # ── Looping: force loops=1 for all 300 steps ──
-    # Progressive 1→2→3 doubles/triples forward cost. At 300 steps,
-    # throughput (more tokens seen) beats iterative refinement (same
-    # tokens processed multiple times). Each extra loop adds ~18 layers
-    # of compute through the loop trunk for diminishing convergence gain.
     cur_loops = 1
-    raw_model = getattr(model, "_orig_mod", model)
     if hasattr(raw_model, "loop_controller"):
         raw_model.loop_controller.loop_default = 1
         raw_model.loop_controller.loop_min = 1
-        raw_model.loop_controller.loop_max = 1  # Lock to 1
+        raw_model.loop_controller.loop_max = 1
 
     use_bf16 = bool(args.bf16)
 
@@ -95,7 +112,11 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
     step, total_loss, valid_count, best_loss, toks = 0, 0.0, 0, float("inf"), 0
     t0 = time.time()
     cur_seq = initial_seq
-    eff_batch = args.batch_size * max(1, args.seq_len // max(1, cur_seq))
+
+    # ── Compute memory-safe batch size ──
+    desired_batch = args.batch_size * max(1, args.seq_len // max(1, cur_seq))
+    eff_batch = _safe_batch(desired_batch, cur_seq, vocab_size)
+
     loader = torch.utils.data.DataLoader(
         dataset, batch_size=eff_batch, shuffle=True, num_workers=0, drop_last=True)
     data_iter = iter(loader)
@@ -110,13 +131,13 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
             if ns != cur_seq:
                 cur_seq = ns
                 dataset.set_seq_len(cur_seq)
-                eff_batch = args.batch_size * max(1, args.seq_len // max(1, cur_seq))
+                desired_batch = args.batch_size * max(1, args.seq_len // max(1, cur_seq))
+                eff_batch = _safe_batch(desired_batch, cur_seq, vocab_size)
                 loader = torch.utils.data.DataLoader(
                     dataset, batch_size=eff_batch, shuffle=True, num_workers=0, drop_last=True)
                 data_iter = iter(loader)
                 print(f"  [P1] seq -> {cur_seq}  batch -> {eff_batch}")
 
-        # Loops locked to 1 — no progressive schedule
         if unfreezer:
             unfreezer.update(step)