From ccec7f5e328a7fa47fca54250ca24a36b3bf0f0b Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Wed, 20 Nov 2024 09:33:09 -0800 Subject: [PATCH] Fix lr schedule (#813) --- src/levanter/optim/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/levanter/optim/config.py b/src/levanter/optim/config.py index d1f4841fc..d814a6b64 100644 --- a/src/levanter/optim/config.py +++ b/src/levanter/optim/config.py @@ -180,7 +180,7 @@ def lr_scheduler(self, num_train_steps): if stable_steps != 0: stable = optax.constant_schedule(self.learning_rate) schedules.append(stable) - boundaries.append(start + stable_steps) + boundaries.append(start + warmup_steps + stable_steps) match self.lr_schedule: case "constant":