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 linear decay.

Details and Context:

• Why: The WSD schedule is an alternative training strategy where the learning rate warms up, stays constant (stable) for a defined period, and then linearly decays. This separates the estimation of the optimization landscape from the convergence phase.
• Implementation:
  • Modified src/MaxText/configs/base.yml to introduce lr_schedule_type (defaulting to 'cosine') and WSD-specific parameters: wsd_learning_rate_final_fraction and wsd_decay_steps_fraction.
  • Updated src/MaxText/configs/types.py to include a LearningRateScheduleType Enum and validation logic in the Optimizer class.
  • Refactored create_learning_rate_schedule in src/MaxText/maxtext_utils.py to generate the schedule using optax based on the selected type.
  • Added validation to ensure the sum of warmup and decay fractions does not exceed 1.0.

Tests

I have added a new test suite, TestLearningRateSchedules, in tests/maxtext_utils_test.py to verify the logic.

• Unit Tests: Verified that both cosine and wsd schedules generate correct learning rate values at specific steps (warmup, stable, decay, and end phases).
• Edge Cases: Verified that invalid configurations (e.g., total fractions > 1.0) raise 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 27 lines in your changes missing coverage. Please review.

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

📢 Thoughts on this report? Let us know!