Fix: probe with return_logits=False (was wasting 3GB on logits tensor), show peak VRAM

train.py

@@ -240,7 +240,7 @@ def probe_max_batch(model, device, dtype, seq_len, vocab_size, max_search=4096):
             x = torch.randint(0, vocab_size, (mid, seq_len), device=device)
             t = torch.randint(0, vocab_size, (mid, seq_len), device=device)
             with autocast(device_type="cuda", dtype=dtype):
-                _, loss = model(x, t)
+                _, loss = model(x, t, return_logits=False)
             loss.backward()
             tmp_opt.step()
             tmp_opt.zero_grad(set_to_none=True)
@@ -563,7 +563,7 @@ def train(cfg: dict):
             tps = tokens_per_step / dt
             steps_left = total_steps - step
             eta_h = steps_left * dt / 3600
-            vram = torch.cuda.memory_allocated() / 1024**3 if device.type == "cuda" else 0
+            vram = torch.cuda.max_memory_allocated() / 1024**3 if device.type == "cuda" else 0
             print(f"  step {step:6d}/{total_steps} | loss {total_loss:.4f} | "
                   f"lr {lr_now:.2e} | {tps:,.0f} tok/s | VRAM {vram:.1f}GB | ETA {eta_h:.1f}h")