Merge pull request #3095 from AI-Hypercomputer:jimmytsai/fix-learning-rate-schedule

PiperOrigin-RevId: 878971906

Author: Google-ML-Automation

src/maxtext/utils/maxtext_utils.py

@@ -1432,7 +1432,7 @@ def schedule(step):
   boundaries = []
 
   if warmup_steps > 0:
-    warmup_schedule = optax.linear_schedule(init_value=0.0, end_value=lr, transition_steps=warmup_steps - 1)
+    warmup_schedule = optax.linear_schedule(init_value=0.0, end_value=lr, transition_steps=warmup_steps)
     pieces.append(warmup_schedule)
     boundaries.append(warmup_steps)
 

tests/unit/maxtext_utils_test.py

@@ -816,6 +816,11 @@ def test_cosine_schedule(self):
     # Warmup phase: 0 -> peak
     self.assertAlmostEqual(float(schedule_fn(0)), 0.0, places=6)
     self.assertAlmostEqual(float(schedule_fn(warmup_steps)), learning_rate, places=6)
+    # Ensure delta is constant
+    expected_slope = learning_rate / warmup_steps
+    for i in range(1, warmup_steps + 1):
+      current_lr = float(schedule_fn(i))
+      self.assertAlmostEqual(current_lr - float(schedule_fn(i - 1)), expected_slope, places=6)
 
     # Cosine decay phase
     lr_end = schedule_fn(learning_rate_schedule_steps - 1)
@@ -859,6 +864,11 @@ def test_wsd_schedule(self):
       # Warmup phase: 0 -> peak
       self.assertAlmostEqual(float(schedule_fn(0)), 0.0, places=6)
       self.assertAlmostEqual(float(schedule_fn(warmup_steps)), learning_rate, places=6)
+      # Ensure delta is constant
+      expected_slope = learning_rate / warmup_steps
+      for i in range(1, warmup_steps + 1):
+        current_lr = float(schedule_fn(i))
+        self.assertAlmostEqual(current_lr - float(schedule_fn(i - 1)), expected_slope, places=6)
 
       # Stable phase: constant at peak
       self.assertAlmostEqual(float(schedule_fn(warmup_steps + 10)), learning_rate, places=6)