Fix final value of exponential decay lr

EleutherAI · Sep 23, 2023 · 1f832c1 · 1f832c1
1 parent 390d37c
commit 1f832c1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/megatron/learning_rates.py b/megatron/learning_rates.py
@@ -82,7 +82,8 @@ def get_lr(self):
             )
         elif self.decay_style == "exponential":
             # exp(-0.693) = 1/2
-            lr = self.start_lr * math.exp(-0.693 * num_iters_ / self.end_iter)
+            end_iter = self.end_iter - self.warmup_iter
+            lr = self.start_lr * math.exp(-0.693 * num_iters_ / end_iter)
         else:
             lr = self.start_lr
         return max(lr, self.min_lr)