feat: loops.py v11 — aligned with GENESIS engine, no distiller overhead"

chimera/training/loops.py

@@ -9,20 +9,17 @@ import torch
 
 import chimera_turbo
 
-from .common import cosine_lr, save_final_checkpoint, save_training_checkpoint
+from .common import save_final_checkpoint, save_training_checkpoint
 from .hyper import ProgressiveLoopScheduler
 
 
 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)
-    log_f = open(os.path.join(args.output_dir, "log.jsonl"), "w", encoding="utf-8")
     model.train()
     step, total_loss, best_loss, toks = 0, 0.0, float("inf"), 0
     t0 = time.time()
     data_iter = iter(loader)
-    warmup = min(args.warmup, max(1, args.max_steps // 10))
-
     while step < args.max_steps:
         try:
             batch = next(data_iter)
@@ -33,9 +30,6 @@ def train_fast_loop(args, model, config, loader, compute_loss) -> str:
         loss.backward()
         total_loss += float(loss.item())
         torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
-        cur_lr = cosine_lr(step, warmup, args.max_steps, args.lr, args.lr * 0.1)
-        for pg in optimizer.param_groups:
-            pg["lr"] = cur_lr
         optimizer.step()
         optimizer.zero_grad(set_to_none=True)
         toks += batch["input_ids"].numel()
@@ -53,33 +47,29 @@ def train_fast_loop(args, model, config, loader, compute_loss) -> str:
 
 
 def train_standard_loop(args, model, config, loader, compute_loss, optimizer, use_mezo):
-    # Legacy — unchanged
     pass
 
 
 def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer):
     use_compile = getattr(args, "compile", False)
 
-    # Apply all paradigms: Muon + MTP + EMA Distillation
     model, optimizer, scheduler, extras = chimera_turbo.apply(
         model,
         max_steps=args.max_steps,
-        lr=0.02,          # Muon default LR (10× higher than AdamW, paper-standard)
+        lr=0.02,
         weight_decay=0.01,
-        warmup_steps=200,  # Short warmup for fast ramp
+        warmup_steps=200,
         use_compile=use_compile,
-        use_muon=True,
-        use_mtp=True,
-        use_distill=True,
-        mtp_heads=3,       # Predict next 3 tokens
+        mtp_heads=3,
+        llrd_decay=0.85,
+        grokfast_alpha=0.98,
+        grokfast_lambda=2.0,
     )
     model.train()
 
     # Progressive looping
     loop_sched = ProgressiveLoopScheduler(args.max_steps, max_loops=3)
     cur_loops = 1
-    print(f"[LOOP] Progressive looping: 1→2→3 over {args.max_steps} steps")
-    print(f"[P5] Train mode: BitLinear STE (clamp-aware, NaN-safe)")
 
     use_bf16 = bool(args.bf16)
 
@@ -90,16 +80,14 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
     cur_seq = initial_seq
     eff_batch = args.batch_size * max(1, args.seq_len // max(1, cur_seq))
     loader = torch.utils.data.DataLoader(
-        dataset, batch_size=eff_batch, shuffle=True, num_workers=0, drop_last=True
-    )
+        dataset, batch_size=eff_batch, shuffle=True, num_workers=0, drop_last=True)
     data_iter = iter(loader)
 
     print(f"\n{'=' * 65}")
-    print(f"Training  eff_batch={eff_batch}  seq={cur_seq}  loops={cur_loops}")
+    print(f"Training  batch={eff_batch}  seq={cur_seq}  loops={cur_loops}")
     print(f"{'=' * 65}\n")
 
     while step < args.max_steps:
-        # ── Seq length scheduling ──
         if grow:
             ns = grow.get_seq_len(step)
             if ns != cur_seq:
@@ -107,19 +95,17 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
                 dataset.set_seq_len(cur_seq)
                 eff_batch = args.batch_size * max(1, args.seq_len // max(1, cur_seq))
                 loader = torch.utils.data.DataLoader(
-                    dataset, batch_size=eff_batch, shuffle=True, num_workers=0, drop_last=True
-                )
+                    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}")
 
-        # ── Loop scheduling ──
         new_loops = loop_sched.get_loops(step)
         if new_loops != cur_loops:
             cur_loops = new_loops
             raw = getattr(model, "_orig_mod", model)
             if hasattr(raw, "loop_controller"):
                 raw.loop_controller.loop_default = cur_loops
-            print(f"  [LOOP] loops → {cur_loops}")
+            print(f"  [LOOP] → {cur_loops}")
 
         if unfreezer:
             unfreezer.update(step)
@@ -132,12 +118,11 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
 
         loss_val = chimera_turbo.training_step(
             model, batch, optimizer, scheduler,
-            extras=extras,
-            grad_accum_steps=1, step=step,
+            extras=extras, grad_accum_steps=1, step=step,
             autocast_dtype=torch.bfloat16 if use_bf16 else None,
         )
 
-        cur_lr = optimizer.param_groups[0]["lr"]
+        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
@@ -149,12 +134,12 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
             avg = total_loss / max(1, valid_count)
             ppl = math.exp(min(avg, 20)) if math.isfinite(avg) else float("nan")
             tps = toks / dt if dt > 0 else 0
-            eta = (args.max_steps - step) / (step / dt) / 3600 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({
                 "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_len": cur_seq, "loops": cur_loops,
+                "seq": cur_seq, "loops": cur_loops,
             }) + "\n")
             log_f.flush()
             print(
@@ -165,10 +150,12 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
             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}"))
+            d = save_training_checkpoint(model, config, step,
+                                          os.path.join(args.output_dir, f"ckpt-{step}"))
             print(f"  [SAVE] {d}")
 
-    d = save_final_checkpoint(model, config, step, best_loss, os.path.join(args.output_dir, "final"))
+    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}")
     return d