phase3: task_mode (is_prime), train_frac, lr_decay flags + tests (41 green); Addendum 5 locked; 11-job batch

1 files changed, 5 insertions, 0 deletions

diff --git a/src/train.py b/src/train.py
index aa0fb4b..406defd 100644
--- a/src/train.py
+++ b/src/train.py
@@ -166,6 +166,11 @@ def main() -> None:
             halt = out["halt_steps"]
             penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0
             loss = loss_tokens + penalty
+            if cfg.lr_decay:
+                # cosine 1e-3 -> 1e-4 over the full run (E3, locked in Addendum 4)
+                frac = min(1.0, step / cfg.max_train_steps)
+                lr_t = cfg.lr * 0.1 + 0.5 * (cfg.lr - cfg.lr * 0.1) * (1 + math.cos(math.pi * frac))
+                opt.param_groups[0]["lr"] = lr_t
             opt.zero_grad()
             loss.backward()
             opt.step()