Implement Warmup-Stable-Decay (WSD) Learning Rate Schedule

src/MaxText/configs/base.yml

@@ -607,7 +607,7 @@ grain_file_type: 'arrayrecord' # arrayrecord or parquet
 grain_packing_type: 'first_fit' # 'first_fit' or 'concat_then_split'. See details of the corresponding module in https://google-grain.readthedocs.io/en/latest/grain.experimental.html
 grain_worker_count: 1 # Set to -1 to enable auto-tuning: automatically determines optimal worker count. See https://google-grain.readthedocs.io/en/latest/_autosummary/grain.experimental.pick_performance_config.html
 grain_per_worker_buffer_size: 1
-# num_threads and prefetch_buffer_size are per-worker per-dataset. 
+# num_threads and prefetch_buffer_size are per-worker per-dataset.
 # When using array_records, they are used in ReadOptions (https://google-grain.readthedocs.io/en/latest/tutorials/data_loader_tutorial.html#per-worker-readoptions)
 # The default value matches that in the Grain package. If mixing multiple data sources, consider lowering these values to reduce memory usage.
 # When using parquet, grain_num_threads is the number of files to read and interleave in parallel
@@ -635,15 +635,29 @@ skip_jax_distributed_system: False # If True we will not initialize the jax dist
 # However when run on google internal TPUs the coordination service is started automatically
 # and we should set this to True so we won't try to initialize a second time manually.
 
-# We take inspiration from Llama2's learning rate (LR) schedule, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
-# Learning rate schedule has either two or three parts:
+# Learning rate schedule structure depends on lr_schedule_type:
+#
+# Cosine schedule (lr_schedule_type='cosine'):
+# Inspired by Llama2's learning rate schedule, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
+# 1) Linear warmup from 0 to [learning_rate] over steps 0 to [learning_rate_schedule_steps * warmup_steps_fraction]
+# 2) Cosine decay from [learning_rate] to [learning_rate * cosine_learning_rate_final_fraction] until learning_rate_schedule_steps
+# 3) Constant learning rate of 0 from learning_rate_schedule_steps to steps (if steps > learning_rate_schedule_steps)
+#
+# WSD schedule (lr_schedule_type='wsd', Warmup-Stable-Decay):
 # 1) Linear warmup from 0 to [learning_rate] over steps 0 to [learning_rate_schedule_steps * warmup_steps_fraction]
-# 2) Cosine decay from [learning_rate] to [learning_rate * cosine_learning_rate_final_fraction] from warmup to learning_rate_schedule_steps
-# 3) Constant learning rate of 0 from learning_rate_schedule_steps to steps.
+# 2) Stable phase at [learning_rate] for the majority of training
+# 3) Decay from [learning_rate] to [learning_rate * wsd_learning_rate_final_fraction] over [learning_rate_schedule_steps * wsd_decay_steps_fraction] steps
+#    The decay can be either linear or cosine based on wsd_decay_style
+# 4) Constant learning rate of 0 from learning_rate_schedule_steps to steps (if steps > learning_rate_schedule_steps)
+#
 # The zero learning rate section can be used to more accurately measure the fully trained model's performance.
 learning_rate: 3.e-5
-cosine_learning_rate_final_fraction: 0.1
-warmup_steps_fraction: 0.1
+lr_schedule_type: 'cosine'  # Options: 'cosine' or 'wsd'
+cosine_learning_rate_final_fraction: 0.1  # Final LR as fraction of peak LR for cosine schedule
+wsd_learning_rate_final_fraction: 0.1  # Final LR as fraction of peak LR for WSD schedule
+wsd_decay_steps_fraction: 0.1  # Fraction of learning_rate_schedule_steps used for decay phase in WSD (e.g., 0.1 = 10%)
+wsd_decay_style: 'linear'  # Decay style for WSD schedule: 'linear' or 'cosine'
+warmup_steps_fraction: 0.1  # Fraction of learning_rate_schedule_steps used for warmup phase (applies to both schedules)
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.
 # However you may choose a longer schedule (learning_rate_schedule_steps > steps), in which case the training will end before
 # dropping fully down. Or you may choose a shorter schedule, where the unspecified steps will have a learning rate of 0.

src/MaxText/configs/types.py

@@ -124,6 +124,20 @@ class OptimizerType(str, Enum):
   MUON = "muon"
 
 
+class LearningRateScheduleType(str, Enum):
+  """Supported learning rate schedule types."""
+
+  COSINE = "cosine"
+  WSD = "wsd"
+
+
+class WsdDecayStyle(str, Enum):
+  """Supported decay styles for WSD schedule."""
+
+  LINEAR = "linear"
+  COSINE = "cosine"
+
+
 class RopeType(str, Enum):
   """Supported Rotary Positional Embedding (RoPE) implementations."""
 
@@ -1005,9 +1019,21 @@ class Optimizer(BaseModel):
       1.0, description="The threshold for gradient clipping. 0 disables clipping."
   )
   learning_rate: NonNegativeFloat = Field(3.0e-5, description="The peak learning rate.")
+  lr_schedule_type: LearningRateScheduleType = Field(
+      LearningRateScheduleType.COSINE, description="The type of learning rate schedule to use."
+  )
   cosine_learning_rate_final_fraction: float = Field(
       0.1, description="Final LR as a fraction of peak LR in cosine decay."
   )
+  wsd_learning_rate_final_fraction: float = Field(
+      0.1, description="Final LR as a fraction of peak LR in WSD decay phase."
+  )
+  wsd_decay_steps_fraction: float = Field(
+      0.1, ge=0.0, le=1.0, description="Fraction of total steps for decay phase in WSD schedule."
+  )
+  wsd_decay_style: WsdDecayStyle = Field(
+      WsdDecayStyle.LINEAR, description="The decay style for WSD schedule ('linear' or 'cosine')."
+  )
   warmup_steps_fraction: float = Field(0.1, ge=0.0, le=1.0, description="Fraction of total steps for LR warmup.")
   learning_rate_schedule_steps: int = Field(
       -1,

src/MaxText/maxtext_utils.py

@@ -1103,16 +1103,25 @@ def create_device_mesh(config, devices=None):
 
 
 def create_learning_rate_schedule(config):
-  """Creates a warmup and cosine decay learning rate schedule:
-  We take inspiration from Llama2's learning rate (LR) schedule, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
-  Learning rate schedule has either two or three parts:
+  """Creates a learning rate schedule with warmup and decay.
+
+  Supports two schedule types:
+  - Cosine: Inspired by Llama2's learning rate schedule, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
+  - WSD (Warmup-Stable-Decay): Maintains constant learning rate for most of training before final decay
+
+  Schedule structure:
   1) Linear warmup from 0 to [learning_rate] over steps 0 to [learning_rate_schedule_steps * warmup_steps_fraction]
-  2) Cosine from [learning_rate] to [learning_rate * cosine_learning_rate_final_fraction] until learning_rate_schedule_steps
+  2) Decay from [learning_rate] to a final value until learning_rate_schedule_steps
+     - Cosine: decays to [learning_rate * cosine_learning_rate_final_fraction]
+     - WSD: maintains [learning_rate] for a stable phase, then decays to [learning_rate * wsd_learning_rate_final_fraction]
+       using either linear or cosine decay based on wsd_decay_style
   3) Constant learning rate of 0 from learning_rate_schedule_steps to steps.
   The zero learning rate section can be used to more accurately measure the fully trained model's performance.
   """
 
   def make_cos_schedule(init_lr, final_lr, len_steps):
+    """Creates a cosine decay schedule from init_lr to final_lr over len_steps."""
+
     def schedule(step):
       pct = (step) / len_steps
       a = 0.5 * (jnp.cos(jnp.pi * pct) + 1)
@@ -1122,25 +1131,50 @@ def schedule(step):
     return schedule
 
   lr = config.learning_rate
-  cos_final_lr = lr * config.cosine_learning_rate_final_fraction
-
   warmup_steps = int(config.learning_rate_schedule_steps * config.warmup_steps_fraction)
-  cos_steps = config.learning_rate_schedule_steps - warmup_steps
   constant_zero_steps = config.steps - config.learning_rate_schedule_steps
-
   warmup_schedule = optax.linear_schedule(init_value=0.0, end_value=lr, transition_steps=warmup_steps)
-  cos_schedule = make_cos_schedule(lr, cos_final_lr, cos_steps)
-  constant_schedule = optax.constant_schedule(0.0)
 
-  pieces = [warmup_schedule, cos_schedule]
-  boundaries = [
-      warmup_steps,
-      warmup_steps + cos_steps,
-  ]
+  if config.lr_schedule_type == "cosine":
+    cos_final_lr = lr * config.cosine_learning_rate_final_fraction
+    cos_steps = config.learning_rate_schedule_steps - warmup_steps
+    cos_schedule = make_cos_schedule(lr, cos_final_lr, cos_steps)
+
+    pieces = [warmup_schedule, cos_schedule]
+    boundaries = [warmup_steps, warmup_steps + cos_steps]
+
+  elif config.lr_schedule_type == "wsd":
+    wsd_final_lr = lr * config.wsd_learning_rate_final_fraction
+    decay_steps = int(config.learning_rate_schedule_steps * config.wsd_decay_steps_fraction)
+    stable_steps = config.learning_rate_schedule_steps - warmup_steps - decay_steps
+
+    if stable_steps < 0:
+      raise ValueError(
+          f"Invalid WSD schedule: warmup_steps_fraction ({config.warmup_steps_fraction}) + "
+          f"wsd_decay_steps_fraction ({config.wsd_decay_steps_fraction}) must not exceed 1.0. "
+          f"Current sum: {config.warmup_steps_fraction + config.wsd_decay_steps_fraction}"
+      )
+
+    stable_schedule = optax.constant_schedule(lr)
+
+    # Create decay schedule based on wsd_decay_style
+    if config.wsd_decay_style == "linear":
+      decay_schedule = optax.linear_schedule(init_value=lr, end_value=wsd_final_lr, transition_steps=decay_steps)
+    elif config.wsd_decay_style == "cosine":
+      decay_schedule = make_cos_schedule(lr, wsd_final_lr, decay_steps)
+    else:
+      raise ValueError(f"Invalid wsd_decay_style: {config.wsd_decay_style}. " "Must be either 'linear' or 'cosine'.")
+
+    pieces = [warmup_schedule, stable_schedule, decay_schedule]
+    boundaries = [warmup_steps, warmup_steps + stable_steps]
+
+  else:
+    raise ValueError(f"Invalid lr_schedule_type: {config.lr_schedule_type}. " "Must be either 'cosine' or 'wsd'.")
 
   if constant_zero_steps > 0:
+    constant_schedule = optax.constant_schedule(0.0)
     pieces.append(constant_schedule)
-    boundaries.append(warmup_steps + cos_steps + constant_zero_steps)
+    boundaries.append(boundaries[-1] + constant_zero_steps)
 
   return optax.join_schedules(pieces, boundaries)
 

tests/maxtext_utils_test.py

@@ -682,5 +682,101 @@ def test_bytes_from_pytree_empty_dict(self):
     self.assertEqual(max_utils.calculate_bytes_from_pytree({}), 0)
 
 
+class TestLearningRateSchedules(unittest.TestCase):
+  """Test suite for learning rate schedule functions."""
+
+  def test_cosine_schedule(self):
+    """Tests cosine learning rate schedule."""
+    config = pyconfig.initialize(
+        [None, os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")],
+        enable_checkpointing=False,
+        learning_rate=1e-3,
+        learning_rate_schedule_steps=1000,
+        steps=1200,
+        warmup_steps_fraction=0.1,
+        lr_schedule_type="cosine",
+        cosine_learning_rate_final_fraction=0.1,
+    )
+
+    schedule_fn = maxtext_utils.create_learning_rate_schedule(config)
+    warmup_steps = int(config.learning_rate_schedule_steps * config.warmup_steps_fraction)
+
+    # Warmup phase: 0 -> peak
+    self.assertAlmostEqual(float(schedule_fn(0)), 0.0, places=6)
+    self.assertAlmostEqual(float(schedule_fn(warmup_steps)), config.learning_rate, places=6)
+
+    # Cosine decay phase
+    lr_end = schedule_fn(config.learning_rate_schedule_steps - 1)
+    expected_final = config.learning_rate * config.cosine_learning_rate_final_fraction
+    self.assertLess(float(lr_end), config.learning_rate)
+    self.assertGreater(float(lr_end), expected_final * 0.9)
+
+    # Zero phase
+    self.assertAlmostEqual(float(schedule_fn(config.steps - 1)), 0.0, places=6)
+
+  def test_wsd_schedule(self):
+    """Tests WSD learning rate schedule with both linear and cosine decay styles."""
+    learning_rate = 1e-3
+    learning_rate_schedule_steps = 1000
+    steps = 1200
+    warmup_steps_fraction = 0.1
+    wsd_learning_rate_final_fraction = 0.1
+    wsd_decay_steps_fraction = 0.1
+
+    warmup_steps = int(learning_rate_schedule_steps * warmup_steps_fraction)
+    decay_steps = int(learning_rate_schedule_steps * wsd_decay_steps_fraction)
+    stable_steps = learning_rate_schedule_steps - warmup_steps - decay_steps
+    decay_start = warmup_steps + stable_steps
+
+    # Test both decay styles: linear and cosine
+    for decay_style in ["linear", "cosine"]:
+      config = pyconfig.initialize(
+          [None, os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")],
+          enable_checkpointing=False,
+          learning_rate=learning_rate,
+          learning_rate_schedule_steps=learning_rate_schedule_steps,
+          steps=steps,
+          warmup_steps_fraction=warmup_steps_fraction,
+          lr_schedule_type="wsd",
+          wsd_learning_rate_final_fraction=wsd_learning_rate_final_fraction,
+          wsd_decay_steps_fraction=wsd_decay_steps_fraction,
+          wsd_decay_style=decay_style,
+      )
+      schedule_fn = maxtext_utils.create_learning_rate_schedule(config)
+
+      # 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)
+
+      # Stable phase: constant at peak
+      self.assertAlmostEqual(float(schedule_fn(warmup_steps + 10)), learning_rate, places=6)
+      self.assertAlmostEqual(float(schedule_fn(warmup_steps + stable_steps // 2)), learning_rate, places=6)
+      self.assertAlmostEqual(float(schedule_fn(decay_start - 1)), learning_rate, places=6)
+
+      # Decay phase: peak -> final
+      lr_mid_decay = schedule_fn(decay_start + decay_steps // 2)
+      expected_final = learning_rate * wsd_learning_rate_final_fraction
+      self.assertLess(float(lr_mid_decay), learning_rate)
+      self.assertGreater(float(lr_mid_decay), expected_final)
+
+      # Zero phase
+      self.assertAlmostEqual(float(schedule_fn(steps - 1)), 0.0, places=6)
+
+    # Test invalid fractions
+    config_invalid_fractions = pyconfig.initialize(
+        [None, os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")],
+        enable_checkpointing=False,
+        learning_rate=learning_rate,
+        learning_rate_schedule_steps=learning_rate_schedule_steps,
+        steps=steps,
+        warmup_steps_fraction=0.6,
+        lr_schedule_type="wsd",
+        wsd_learning_rate_final_fraction=wsd_learning_rate_final_fraction,
+        wsd_decay_steps_fraction=0.5,  # Sum > 1.0
+    )
+    with self.assertRaises(ValueError):
+      maxtext_utils.create_learning_rate_schedule(config_invalid_fractions)
+
+
 if __name__ == "__main__":
   unittest.main()