fix: print every step + first-step timing to diagnose slow forward

With 227M params on CPU, each step can take 30-60s. With log_every=10,
the first output wouldn't appear for 5-10 minutes — looks like a hang.

Changes:
- Print step 1 timing immediately after first forward+backward
- Log every step for the first 5 steps, then every log_every
- Flush stdout after every print to ensure immediate display
- Add sys.stdout.flush() calls"

chimera/training/loops.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import json
 import math
 import os
+import sys
 import time
 
 import torch
@@ -28,6 +29,7 @@ def _safe_batch(desired_batch: int, seq_len: int, vocab_size: int,
     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)")
+        sys.stdout.flush()
     return capped
 
 
@@ -91,14 +93,12 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
     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 ──
+    # ── Looping: force loops=1 ──
     cur_loops = 1
     if hasattr(raw_model, "loop_controller"):
         raw_model.loop_controller.loop_default = 1
@@ -111,9 +111,10 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
     log_f = open(os.path.join(args.output_dir, "log_hyper.jsonl"), "w")
     step, total_loss, valid_count, best_loss, toks = 0, 0.0, 0, float("inf"), 0
     t0 = time.time()
+    t_start = t0
     cur_seq = initial_seq
 
-    # ── Compute memory-safe batch size ──
+    # ── 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)
 
@@ -123,7 +124,9 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
 
     print(f"\n{'=' * 65}")
     print(f"Training  batch={eff_batch}  seq={cur_seq}  loops={cur_loops}")
-    print(f"{'=' * 65}\n")
+    print(f"Starting first step (may take 30-60s on CPU with 227M params)...")
+    print(f"{'=' * 65}")
+    sys.stdout.flush()
 
     while step < args.max_steps:
         if grow:
@@ -137,6 +140,7 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
                     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}")
+                sys.stdout.flush()
 
         if unfreezer:
             unfreezer.update(step)
@@ -147,20 +151,34 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
             data_iter = iter(loader)
             batch = next(data_iter)
 
+        step_t0 = time.time()
+
         loss_val = chimera_turbo.training_step(
             model, batch, optimizer, scheduler,
             extras=extras, grad_accum_steps=1, step=step,
             autocast_dtype=torch.bfloat16 if use_bf16 else None,
         )
 
+        step_dt = time.time() - step_t0
+
         cur_lr = optimizer.param_groups[0]["lr"] * optimizer.param_groups[0].get("lr_scale", 1.0)
         if math.isfinite(loss_val):
             total_loss += loss_val
             valid_count += 1
-        toks += batch["input_ids"].numel()
+        step_toks = batch["input_ids"].numel()
+        toks += step_toks
         step += 1
 
-        if step % args.log_every == 0:
+        # Print every step for the first 5 steps, then every log_every
+        should_log = (step <= 5) or (step % args.log_every == 0)
+
+        if step == 1:
+            step_tps = step_toks / step_dt if step_dt > 0 else 0
+            print(f"  ✓ Step 1 completed in {step_dt:.1f}s "
+                  f"({step_tps:.0f} tok/s, loss={loss_val:.4f})")
+            sys.stdout.flush()
+
+        if should_log:
             dt = time.time() - t0
             if valid_count > 0:
                 avg = total_loss / valid_count
@@ -169,28 +187,39 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
                 avg = float("nan")
                 ppl = float("nan")
             tps = toks / dt if dt > 0 else 0
-            eta = (args.max_steps - step) / max(1, step) * (time.time() - t0) / 3600 if step > 0 else 0
+            elapsed = time.time() - t_start
+            eta_s = (args.max_steps - step) * (elapsed / max(1, step))
             log_f.write(json.dumps({
                 "step": step, "loss": round(avg, 4) if math.isfinite(avg) else None,
                 "ppl": round(ppl, 2) if math.isfinite(ppl) else None,
                 "lr": round(cur_lr, 6), "tok/s": round(tps),
                 "seq": cur_seq, "loops": cur_loops,
+                "step_time": round(step_dt, 2),
             }) + "\n")
             log_f.flush()
             print(
                 f"  step {step:>6}/{args.max_steps} | loss {avg:.4f} | ppl {ppl:>8.2f} "
-                f"| lr {cur_lr:.2e} | {tps:,.0f} tok/s | seq {cur_seq} | L{cur_loops} | ETA {eta:.1f}h"
+                f"| {tps:,.0f} tok/s | {step_dt:.1f}s/step | seq {cur_seq} "
+                f"| ETA {eta_s / 60:.0f}m"
             )
-            best_loss = min(best_loss, avg) if math.isfinite(avg) else best_loss
-            total_loss, valid_count, toks, t0 = 0.0, 0, 0, time.time()
+            sys.stdout.flush()
+
+            if step > 5:
+                # Reset counters for clean averages
+                best_loss = min(best_loss, avg) if math.isfinite(avg) else best_loss
+                total_loss, valid_count, toks, t0 = 0.0, 0, 0, time.time()
 
         if step % args.save_every == 0:
             d = save_training_checkpoint(model, config, step,
                                           os.path.join(args.output_dir, f"ckpt-{step}"))
             print(f"  [SAVE] {d}")
+            sys.stdout.flush()
 
     d = save_final_checkpoint(model, config, step, best_loss,
                                os.path.join(args.output_dir, "final"))
     log_f.close()
-    print(f"\nDONE -- best loss {best_loss:.4f}  ppl {math.exp(min(best_loss, 20)):.2f}")
+    total_time = time.time() - t_start
+    print(f"\nDONE -- best loss {best_loss:.4f}  ppl {math.exp(min(best_loss, 20)):.2f}"
+          f"  total time {total_time / 60:.1f}m")
+    sys.stdout.flush()
     return d