Fix NaN loss reporting: show nan instead of 0.0 when all steps in window are NaN

chimera/training/loops.py

@@ -55,9 +55,9 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
 
     # FIX: Use args.lr instead of hardcoded 0.02.
     # FIX: Use args.warmup instead of hardcoded 200.
-    # FIX: Reduce MTP heads from 3→2 to cut 51M params of overhead.
-    # FIX: Soften LLRD decay (0.85→0.92) so early layers still learn.
-    # FIX: Lower Grokfast lambda (2.0→1.0) to reduce gradient amplification noise.
+    # FIX: Reduce MTP heads from 3->2 to cut 51M params of overhead.
+    # FIX: Soften LLRD decay (0.85->0.92) so early layers still learn.
+    # FIX: Lower Grokfast lambda (2.0->1.0) to reduce gradient amplification noise.
     model, optimizer, scheduler, extras = chimera_turbo.apply(
         model,
         max_steps=args.max_steps,
@@ -102,7 +102,7 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
                 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}")
+                print(f"  [P1] seq -> {cur_seq}  batch -> {eff_batch}")
 
         new_loops = loop_sched.get_loops(step)
         if new_loops != cur_loops:
@@ -110,7 +110,7 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
             raw = getattr(model, "_orig_mod", model)
             if hasattr(raw, "loop_controller"):
                 raw.loop_controller.loop_default = cur_loops
-            print(f"  [LOOP] → {cur_loops}")
+            print(f"  [LOOP] -> {cur_loops}")
 
         if unfreezer:
             unfreezer.update(step)
@@ -136,8 +136,12 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
 
         if step % args.log_every == 0:
             dt = time.time() - t0
-            avg = total_loss / max(1, valid_count)
-            ppl = math.exp(min(avg, 20)) if math.isfinite(avg) else float("nan")
+            if valid_count > 0:
+                avg = total_loss / valid_count
+                ppl = math.exp(min(avg, 20)) if math.isfinite(avg) else float("nan")
+            else:
+                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
             log_f.write(json.dumps({
@@ -162,5 +166,5 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
     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}")
+    print(f"\nDONE -- best loss {best_loss:.4f}  ppl {math.exp(min(best_loss, 20)):.2f}")
     return d