Upload chimera/training/loops.py

chimera/training/loops.py

@@ -53,17 +53,22 @@ 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 = getattr(args, "compile", False)
 
+    # 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.
     model, optimizer, scheduler, extras = chimera_turbo.apply(
         model,
         max_steps=args.max_steps,
-        lr=0.02,
+        lr=args.lr,
         weight_decay=0.01,
-        warmup_steps=200,
+        warmup_steps=args.warmup,
         use_compile=use_compile,
-        mtp_heads=3,
-        llrd_decay=0.85,
+        mtp_heads=2,
+        llrd_decay=0.92,
         grokfast_alpha=0.98,
-        grokfast_lambda=2.0,
+        grokfast_lambda=1.0,
     )
     model.train()