diff options
Diffstat (limited to 'src/train.py')
| -rw-r--r-- | src/train.py | 5 |
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() |