fix: train_hyper_loop grad_accum=1 (DataLoader already batches), better tok/s logging

chimera/training/loops.py

@@ -150,10 +150,10 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
     model, optimizer, scheduler = chimera_turbo.apply(
         model,
         max_steps=args.max_steps,
-        lr=1e-3,
+        lr=args.lr,
         weight_decay=0.05,
         warmup_steps=min(500, args.max_steps // 10),
-        use_compile=True,
+        use_compile=False,  # ← disabled: 84 graph breaks from STE
         use_ipex=True,
     )
     model.train()
@@ -169,7 +169,9 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
     t0 = time.time()
     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)
+    loader = torch.utils.data.DataLoader(
+        dataset, batch_size=eff_batch, shuffle=True, num_workers=0, drop_last=True
+    )
     data_iter = iter(loader)
 
     print(f"\n{'=' * 65}")
@@ -183,9 +185,11 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
                 cur_seq = ns
                 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)
+                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}")
         if unfreezer:
             unfreezer.update(step)
         try:
@@ -193,32 +197,60 @@ def train_hyper_loop(args, model, config, dataset, initial_seq, grow, unfreezer)
         except StopIteration:
             data_iter = iter(loader)
             batch = next(data_iter)
-        grad_accum_steps = max(1, eff_batch // max(1, args.batch_size))
+
+        # grad_accum_steps=1: DataLoader already provides eff_batch items.
+        # The effective batch IS eff_batch. No need to accumulate further.
         loss_val = chimera_turbo.training_step(
-            model, batch, optimizer, scheduler, grad_accum_steps=grad_accum_steps, step=step, autocast_dtype=torch.bfloat16 if use_bf16 else None
+            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
         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))
             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": cur_lr, "tok/s": round(tps), "seq_len": cur_seq, "eff_batch": eff_batch}) + "\n")
+            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,
+                }) + "\n"
+            )
             log_f.flush()
-            print(f"  step {step:>6}/{args.max_steps} | loss {avg:.4f} | ppl {ppl:>8.2f} | {tps:,.0f} tok/s | seq {cur_seq} | ETA {eta:.1f}h")
+            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} | ETA {eta:.1f}h"
+            )
             best_loss = min(best_loss, avg)
             total_loss = 0.0
             toks = 0
             t0 = 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