[Train] Add warmup-stable-decay learning rate scheduler

[qianlihuang]

Purpose

Add an independent warmup-stable-decay (WSD) learning-rate scheduler to the trainer.

The scheduler provides:

• linear warmup from a configurable initial LR ratio
• a stable learning-rate phase
• configurable final-decay duration and minimum LR ratio
• linear, cosine, exponential, and minus-square-root decay curves
• validation that rejects overlapping warmup and decay phases

Tests

make quality
All checks passed
Success: no issues found in 215 source files

CUDA_VISIBLE_DEVICES= python -m pytest tests/unit/train/test_trainer_scheduler.py -q
13 passed

Checklist

I have filled in:

• The purpose of the PR and the RFC it relates to.
• The test plan and results.
• Optional documentation update.
• I (a human) have written or reviewed the code in this PR to the best of my ability.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Team

Run ID: 70b73699-12e5-4bef-a34f-8e96cb2473b1

📥 Commits

Reviewing files that changed from the base of the PR and between aadce61 and ce59c9e.

📒 Files selected for processing (4)

• src/speculators/train/cli.py
• src/speculators/train/config/schema.py
• src/speculators/train/trainer.py
• tests/unit/train/test_trainer_scheduler.py

Included review availability: Your plan provides up to 4 included reviews per hour; 3 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

WSD scheduler

Layer / File(s)	Summary
WSD configuration and CLI wiring src/speculators/train/config/schema.py, src/speculators/train/trainer.py, src/speculators/train/cli.py	The configuration accepts scheduler_type="wsd" and four WSD parameters. The CLI passes these parameters to TrainerConfig.
WSD schedule implementation src/speculators/train/trainer.py	The trainer validates schedule parameters, supports four decay styles, builds a LambdaLR schedule, and selects it through make_scheduler.
WSD scheduler validation and tests tests/unit/train/test_trainer_scheduler.py	Tests cover the warmup, stable, and decay phases, all decay styles, and overlapping-phase validation.

Merge Risk: ⚪ Minimal · up to ce59c

The PR adds a validated WSD learning-rate schedule without introducing an actionable merge-blocking risk; it is merge-ready after normal checks and review.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 15 functions across 4 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the main change: adding a warmup-stable-decay learning-rate scheduler.
Description check	✅ Passed	The description directly explains the WSD scheduler, its configuration options, validation behavior, and test results.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @qianlihuang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Copilot]

🟢 Approval recommended

The new scheduler is gated behind an explicit config choice, includes input validation, and is covered by focused unit tests for core behaviors and edge cases.

Pull request overview

Adds a new warmup-stable-decay (WSD) learning-rate scheduler option to the training stack, integrating it into the existing trainer scheduler selection and exposing configuration/CLI schema knobs while backing the new behavior with unit tests.

Changes:

• Introduces a WSD LR schedule implementation (warmup → stable → final decay) with multiple decay curves and overlap validation.
• Extends trainer/config schema and CLI wiring to support scheduler_type="wsd" and its parameters.
• Adds unit tests covering the schedule shape, decay-style coefficients, and invalid overlapping phase configurations.

File summaries

File	Description
tests/unit/train/test_trainer_scheduler.py	Adds unit coverage for WSD schedule progression, decay curve coefficients, and overlap rejection.
src/speculators/train/trainer.py	Implements WSD scheduler logic and wires it into TrainerConfig + scheduler creation.
src/speculators/train/config/schema.py	Exposes WSD scheduler options in the config schema (type + ratios + decay style).
src/speculators/train/cli.py	Passes new scheduler args through to TrainerConfig so CLI/config values take effect.

Review details

• Files reviewed: 4/4 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.