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

[bzantium]

Description

This PR implements the Warmup-Stable-Decay (WSD) learning rate schedule as a configurable option alongside the existing Cosine schedule. This allows users to choose between the standard cosine decay and a schedule that maintains a stable peak learning rate for the majority of training before a rapid decay.

Additionally, this implementation introduces a wsd_decay_style parameter, giving users the flexibility to choose the decay profile (linear or cosine) for the final annealing phase.

Details and Context:

• Why: The WSD schedule is a widely adopted training strategy where the learning rate warms up, stays constant (stable) to maximize training throughput, and then decays rapidly to converge. Separating the stable and decay phases allows for "infinite" training horizons and flexible checkpointing.
• Implementation:
  • Configuration (src/MaxText/configs/base.yml):
    • Added lr_schedule_type (options: 'cosine', 'wsd').
    • Added WSD-specific parameters: wsd_learning_rate_final_fraction, wsd_decay_steps_fraction.
    • Added wsd_decay_style: Supports 'linear' (default, standard for WSD) or 'cosine' decay for the final phase.
  • Types (src/MaxText/configs/types.py):
    • Added LearningRateScheduleType and WsdDecayStyle Enums.
    • Updated the Optimizer class to include validation for these new fields.
  • Logic (src/MaxText/maxtext_utils.py):
    • Refactored create_learning_rate_schedule to switch between Cosine and WSD logic.
    • Implemented WSD construction: Linear Warmup -> Constant Stable -> Decay.
    • The decay phase dynamically selects between optax.linear_schedule and a custom cosine schedule based on wsd_decay_style.
    • Added validation to ensure warmup_steps_fraction + wsd_decay_steps_fraction <= 1.0.

Tests

I have added a comprehensive test suite, TestLearningRateSchedules, in tests/maxtext_utils_test.py.

• Unit Tests:
  • Cosine Schedule: Verified standard behavior (Warmup -> Cosine Decay).
  • WSD Schedule: Verified the 3-phase structure (Warmup -> Stable -> Decay) for both linear and cosine decay styles.
  • Checked that the learning rate hits the correct peak, stable values, and final fraction values.
• Edge Cases: Verified that invalid configurations (e.g., sum of fractions > 1.0) raise a ValueError.

To reproduce/test:

python3 -m unittest tests/maxtext_utils_test.py

Fixes: #2882

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 23 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/maxtext_utils.py	0.00%	23 Missing ⚠️

📢 Thoughts on this report? Let us know!