Add step_offset parameter to ema_schedule (#1516)

* Add step_offset parameter to ema_schedule

* Fix formatting

GitOrigin-RevId: 00eaebd

Author: bleykher

axlearn/common/schedule.py

@@ -417,14 +417,18 @@ def linear_schedule_with_warmup(
     )
 
 
-def ema_schedule(decay: float = 0.9999, *, warmup_steps: int = 1) -> ScheduleFn:
+def ema_schedule(
+    decay: float = 0.9999, *, warmup_steps: int = 1, step_offset: int = 0
+) -> ScheduleFn:
     """Ema decay schedule with warm-up.
 
     The ema decay is 0, 1/2, 2/3, 3/4, 4/5, ... during warm-up, and then is constant at decay.
 
     Args:
         decay: ema decay.
         warmup_steps: The number of steps of the warm-up schedule.
+        step_offset: The initial step number.
+            If `step` is less than or equal to `step_offset`, we clamp to `step - step_offset` to 0.
 
     Returns:
         A ema decay schedule.
@@ -436,6 +440,7 @@ def ema_schedule(decay: float = 0.9999, *, warmup_steps: int = 1) -> ScheduleFn:
         raise ValueError("warmup_steps must be > 0.")
 
     def fn(step):
+        step = jnp.maximum(step - step_offset, 0)
         return step / (1.0 + step) * (step < warmup_steps) + decay * (step >= warmup_steps)
 
     return fn

axlearn/common/schedule_test.py

@@ -241,19 +241,19 @@ def test_adafactor_decay_rate(self):
         self.assertAlmostEqual(fn(200), 1 - (100) ** (-0.8))
 
     def test_ema_schedule(self):
-        warmup_steps = 5
-        s = jax.jit(
-            schedule.ema_schedule(
-                warmup_steps=warmup_steps,
-            )
-        )
+        warmup_steps, step_offset = 5, 3
+        s = jax.jit(schedule.ema_schedule(warmup_steps=warmup_steps, step_offset=step_offset))
+        expected_offset = 0.0
         expected_warmup = [0.0, 1.0 / 2, 2.0 / 3, 3.0 / 4, 4.0 / 5]
         expected_decay = 0.9999
         for step in range(10):
             value = s(step)
-            if step < warmup_steps:
+            if step < step_offset:
+                # Test offset.
+                self.assertAlmostEqual(expected_offset, value)
+            elif step < warmup_steps + step_offset:
                 # Test warmup.
-                self.assertAlmostEqual(expected_warmup[step], value)
+                self.assertAlmostEqual(expected_warmup[step - step_offset], value)
             else:
                 # Test inverse sqrt schedule.
                 self.assertAlmostEqual(expected_decay, value)