fix: loops.py — use chimera_turbo v8 defaults (wd=0.01, warmup=750, β2=0.98) instead of hardcoded values"

chimera/training/loops.py

@@ -13,7 +13,7 @@ from .common import cosine_lr, save_final_checkpoint, save_training_checkpoint
 
 
 def train_fast_loop(args, model, config, loader, compute_loss) -> str:
-    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95))
+    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")
 
@@ -147,19 +147,21 @@ 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=True now works: STE uses detach() trick = zero graph breaks
-    use_compile = getattr(args, "compile", True)
+    use_compile = getattr(args, "compile", False)
+
+    # Let chimera_turbo.apply() use its v8 BitNet-aligned defaults:
+    # lr=1.5e-3, wd=0.01, β=(0.9,0.98), warmup=750
+    # Only override lr if the user explicitly set it via CLI
     model, optimizer, scheduler = chimera_turbo.apply(
         model,
         max_steps=args.max_steps,
-        lr=args.lr,
-        weight_decay=0.05,
-        warmup_steps=min(500, args.max_steps // 10),
+        lr=args.lr,               # CLI default is 1e-3; chimera_turbo default is 1.5e-3
+        # weight_decay and warmup_steps use chimera_turbo defaults (0.01, 750)
         use_compile=use_compile,
         use_ipex=True,
     )
     model.train()
-    print(f"[P5] Train mode: BitLinear STE path (detach trick, compile-friendly)")
+    print(f"[P5] Train mode: BitLinear STE (clamp-aware, NaN-safe)")
     use_bf16 = bool(args.bf16)
 
     os.makedirs(args.output_dir, exist_ok=True)
@@ -200,37 +202,29 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
             data_iter = iter(loader)
             batch = next(data_iter)
 
-        # grad_accum_steps=1: DataLoader already provides eff_batch items.
         loss_val = chimera_turbo.training_step(
-            model,
-            batch,
-            optimizer,
-            scheduler,
-            grad_accum_steps=1,
-            step=step,
+            model, batch, optimizer, scheduler,
+            grad_accum_steps=1, step=step,
             autocast_dtype=torch.bfloat16 if use_bf16 else None,
         )
 
         cur_lr = optimizer.param_groups[0]["lr"]
-        total_loss += loss_val
+        total_loss += loss_val if math.isfinite(loss_val) else 0.0
         toks += batch["input_ids"].numel()
         step += 1
 
         if step % args.log_every == 0:
             dt = time.time() - t0
             avg = total_loss / args.log_every
-            ppl = math.exp(min(avg, 20))
+            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
             log_f.write(
                 json.dumps({
-                    "step": step,
-                    "loss": round(avg, 4),
-                    "ppl": round(ppl, 2),
-                    "lr": round(cur_lr, 6),
-                    "tok/s": round(tps),
-                    "seq_len": cur_seq,
-                    "eff_batch": eff_batch,
+                    "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, "eff_batch": eff_batch,
                 }) + "\n"
             )
             log_f.flush()
@@ -238,7 +232,7 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
                 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} | ETA {eta:.1f}h"
             )
-            best_loss = min(best_loss, avg)
+            best_loss = min(best_loss, avg) if math.isfinite(avg) else best_loss
             total_loss = 0.0
             toks = 0
             t0 = time.time()