feat(batch11): WSD LR scheduler, KV CPU offloading, GQA KV head expansion

- WSD scheduler (Warmup-Stable-Decay / trapezoidal): linear warmup → constant plateau → cosine/linear/sqrt decay; extend_stable() adds steps mid-run without restart; phase_at()/decay_progress() introspection; 26 tests + 2 doctests - KV CPU offloading: KvCpuOffloadManager LRU-based GPU→CPU page spill with on-demand prefetch; insert() auto-offloads when at gpu_page_limit; stats() with hit rate and utilization; 14 tests - GQA KV head expansion: GqaConfig validates num_q_heads/num_kv_heads divisibility; expand_kv_heads() tiles KV [batch,kv_heads,seq,dim]→[batch,q_heads,seq,dim]; gqa_attention_cpu() with numerically stable softmax + causal mask; 15 tests Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-27 05:28:27 +00:00
parent 2360d7bffc
commit 4ae0c34537
6 changed files with 2034 additions and 2 deletions
@@ -12,6 +12,7 @@ pub mod polynomial_decay;
 pub mod reduce_lr_on_plateau;
 pub mod step_lr;
 pub mod warmup;
+pub mod wsd;

 #[cfg(all(test, feature = "disabled_tests"))]
 pub mod scheduler_integration_tests;
@@ -25,6 +26,7 @@ pub use reduce_lr_on_plateau::{
 };
 pub use step_lr::StepLRScheduler;
 pub use warmup::WarmupScheduler;
+pub use wsd::{WsdDecayType, WsdPhase, WsdScheduler};

 /// Trait for learning rate schedulers
 pub trait LearningRateScheduler: Send + Sync {
@@ -0,0 +1,745 @@
+//! Warmup-Stable-Decay (WSD) learning rate scheduler
+//!
+//! Implements the trapezoidal LR schedule used by Mistral, MiniMax-Text-01, MegaMath,
+//! and other modern open-source LLMs. The schedule has three phases:
+//!
+//! 1. **Warmup** — linear ramp from `0` to `peak_lr` over `warmup_steps`
+//! 2. **Stable** — constant at `peak_lr` for `stable_steps` (extendable mid-training)
+//! 3. **Decay** — cosine, linear, or sqrt decay from `peak_lr` to `min_lr` over `decay_steps`
+//!
+//! # Key property
+//!
+//! `stable_steps` can be extended at any point during training (via [`WsdScheduler::extend_stable`])
+//! without resetting the scheduler or restarting training. This enables flexible compute budgets
+//! where you decide when to stop spending tokens on the plateau.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::schedulers::wsd::{WsdScheduler, WsdDecayType};
+//! use rtx_transformers::schedulers::LearningRateScheduler;
+//!
+//! let mut sched = WsdScheduler::cosine(3e-4, 1e-5, 100, 900, 200).unwrap();
+//!
+//! // Warmup: LR ramps up
+//! let lr_step0 = sched.get_lr(0, 0);
+//! assert!(lr_step0 < 3e-4);
+//!
+//! // Stable: LR stays at peak
+//! let lr_stable = sched.get_lr(0, 500);
+//! assert!((lr_stable - 3e-4).abs() < 1e-12);
+//!
+//! // Extend stable phase by 500 more steps mid-run
+//! sched.extend_stable(500);
+//! assert_eq!(sched.total_steps(), 100 + 1400 + 200);
+//! ```
+
+use std::f64::consts::PI;
+
+use crate::schedulers::LearningRateScheduler;
+use crate::{Result, TransformerError};
+use serde::{Deserialize, Serialize};
+use tracing::{debug, trace};
+
+// ---------------------------------------------------------------------------
+// Decay variant
+// ---------------------------------------------------------------------------
+
+/// Decay function applied during the third (decay) phase.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum WsdDecayType {
+    /// Cosine decay from `peak_lr` to `min_lr` (smoother, generally recommended).
+    ///
+    /// `lr(p) = min_lr + (peak_lr - min_lr) * 0.5 * (1 + cos(π * p))`
+    Cosine,
+    /// Linear decay from `peak_lr` to `min_lr`.
+    ///
+    /// `lr(p) = peak_lr + (min_lr - peak_lr) * p`
+    Linear,
+    /// Square-root decay from `peak_lr` to `min_lr`.
+    ///
+    /// `lr(p) = min_lr + (peak_lr - min_lr) * (1 - sqrt(p))`
+    Sqrt,
+}
+
+// ---------------------------------------------------------------------------
+// Phase enum
+// ---------------------------------------------------------------------------
+
+/// Phase the scheduler is currently in, for external introspection.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum WsdPhase {
+    /// Linear warmup — LR is still ramping up.
+    Warmup,
+    /// Plateau — LR is constant at `peak_lr`.
+    Stable,
+    /// Decay — LR is decreasing toward `min_lr`.
+    Decay,
+    /// Past the end of the schedule — LR is clamped to `min_lr`.
+    Complete,
+}
+
+// ---------------------------------------------------------------------------
+// Scheduler struct
+// ---------------------------------------------------------------------------
+
+/// Warmup-Stable-Decay learning rate scheduler.
+///
+/// Three-phase schedule: linear warmup → constant plateau → cosine/linear/sqrt decay.
+///
+/// See the [module-level documentation](self) for a full description and usage example.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WsdScheduler {
+    /// Target learning rate at the peak (end of warmup / entire stable phase).
+    pub peak_lr: f64,
+    /// Minimum learning rate at the end of the decay phase (and beyond).
+    pub min_lr: f64,
+    /// Number of steps in the linear warmup phase.
+    pub warmup_steps: usize,
+    /// Number of steps in the stable (constant) phase. Can be extended at any time.
+    pub stable_steps: usize,
+    /// Number of steps in the decay phase.
+    pub decay_steps: usize,
+    /// Which decay function to use during the third phase.
+    pub decay_type: WsdDecayType,
+    /// Internal step counter; advanced by [`Self::step`].
+    current_step: usize,
+}
+
+impl WsdScheduler {
+    // ------------------------------------------------------------------
+    // Constructors
+    // ------------------------------------------------------------------
+
+    /// Create a new WSD scheduler with explicit parameters.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`TransformerError`] when:
+    /// - `peak_lr` is not positive, or `peak_lr <= min_lr`
+    /// - `min_lr` is negative
+    /// - `warmup_steps` is zero
+    /// - `decay_steps` is zero
+    pub fn new(
+        peak_lr: f64,
+        min_lr: f64,
+        warmup_steps: usize,
+        stable_steps: usize,
+        decay_steps: usize,
+        decay_type: WsdDecayType,
+    ) -> Result<Self> {
+        if peak_lr <= 0.0 {
+            return Err(TransformerError::generic(format!(
+                "peak_lr {peak_lr} must be positive"
+            )));
+        }
+        if min_lr < 0.0 {
+            return Err(TransformerError::generic(format!(
+                "min_lr {min_lr} must be non-negative"
+            )));
+        }
+        if peak_lr <= min_lr {
+            return Err(TransformerError::generic(format!(
+                "peak_lr {peak_lr} must be greater than min_lr {min_lr}"
+            )));
+        }
+        if warmup_steps == 0 {
+            return Err(TransformerError::generic(
+                "warmup_steps must be greater than 0".to_string(),
+            ));
+        }
+        if decay_steps == 0 {
+            return Err(TransformerError::generic(
+                "decay_steps must be greater than 0".to_string(),
+            ));
+        }
+
+        debug!(
+            "Creating WsdScheduler: peak_lr={peak_lr}, min_lr={min_lr}, \
+             warmup={warmup_steps}, stable={stable_steps}, decay={decay_steps}, \
+             decay_type={decay_type:?}"
+        );
+
+        Ok(Self {
+            peak_lr,
+            min_lr,
+            warmup_steps,
+            stable_steps,
+            decay_steps,
+            decay_type,
+            current_step: 0,
+        })
+    }
+
+    /// Convenience constructor that uses cosine decay (the most common choice).
+    ///
+    /// # Errors
+    ///
+    /// Propagates validation errors from [`Self::new`].
+    pub fn cosine(
+        peak_lr: f64,
+        min_lr: f64,
+        warmup_steps: usize,
+        stable_steps: usize,
+        decay_steps: usize,
+    ) -> Result<Self> {
+        Self::new(
+            peak_lr,
+            min_lr,
+            warmup_steps,
+            stable_steps,
+            decay_steps,
+            WsdDecayType::Cosine,
+        )
+    }
+
+    // ------------------------------------------------------------------
+    // Mutation helpers
+    // ------------------------------------------------------------------
+
+    /// Extend the stable phase by `extra_steps` without resetting the scheduler.
+    ///
+    /// This is the key flexibility property of WSD: you can decide mid-training
+    /// to spend more compute at the peak learning rate before decaying.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use rtx_transformers::schedulers::wsd::WsdScheduler;
+    ///
+    /// let mut sched = WsdScheduler::cosine(1e-3, 1e-5, 100, 500, 100).unwrap();
+    /// assert_eq!(sched.total_steps(), 700);
+    /// sched.extend_stable(200);
+    /// assert_eq!(sched.total_steps(), 900);
+    /// ```
+    pub fn extend_stable(&mut self, extra_steps: usize) {
+        self.stable_steps = self.stable_steps.saturating_add(extra_steps);
+        debug!("Extended stable phase by {extra_steps}; new stable_steps={}", self.stable_steps);
+    }
+
+    // ------------------------------------------------------------------
+    // Queries
+    // ------------------------------------------------------------------
+
+    /// Total steps across all three phases.
+    #[must_use]
+    pub fn total_steps(&self) -> usize {
+        self.warmup_steps
+            .saturating_add(self.stable_steps)
+            .saturating_add(self.decay_steps)
+    }
+
+    /// Which phase the scheduler is in at `step`.
+    ///
+    /// Uses the *configured* `stable_steps` (including any extensions).
+    #[must_use]
+    pub fn phase_at(&self, step: usize) -> WsdPhase {
+        if step < self.warmup_steps {
+            WsdPhase::Warmup
+        } else if step < self.warmup_steps.saturating_add(self.stable_steps) {
+            WsdPhase::Stable
+        } else if step < self.total_steps() {
+            WsdPhase::Decay
+        } else {
+            WsdPhase::Complete
+        }
+    }
+
+    /// Which phase the scheduler is currently in (based on `current_step`).
+    #[must_use]
+    pub fn current_phase(&self) -> WsdPhase {
+        self.phase_at(self.current_step)
+    }
+
+    /// LR at the current internal step.
+    #[must_use]
+    pub fn current_lr(&self) -> f64 {
+        self.get_lr(0, self.current_step)
+    }
+
+    /// Fraction through the decay phase at `step`, in `0.0..=1.0`.
+    ///
+    /// Returns `0.0` for steps outside the decay phase.
+    #[must_use]
+    pub fn decay_progress_at(&self, step: usize) -> f64 {
+        let decay_start = self.warmup_steps.saturating_add(self.stable_steps);
+        if step < decay_start || step >= self.total_steps() {
+            0.0
+        } else {
+            let t = step - decay_start;
+            // Clamp to [0, 1] — last decay step maps to exactly 1.0
+            (t as f64 / self.decay_steps as f64).min(1.0)
+        }
+    }
+
+    /// Fraction through the decay phase at `current_step`.
+    #[must_use]
+    pub fn decay_progress(&self) -> f64 {
+        self.decay_progress_at(self.current_step)
+    }
+
+    // ------------------------------------------------------------------
+    // Core LR computation (pure, no mutation)
+    // ------------------------------------------------------------------
+
+    fn compute_lr(&self, step: usize) -> f64 {
+        let decay_start = self.warmup_steps.saturating_add(self.stable_steps);
+
+        if step < self.warmup_steps {
+            // Phase 1: linear warmup — lr = peak_lr * (step + 1) / warmup_steps
+            // Using (step + 1) so step=0 gives a non-zero but tiny LR, and
+            // step = warmup_steps - 1 gives peak_lr * (warmup_steps / warmup_steps) = peak_lr.
+            let progress = (step + 1) as f64 / self.warmup_steps as f64;
+            let lr = self.peak_lr * progress;
+            trace!("WSD warmup  step={step}: progress={progress:.4}, lr={lr:.6e}");
+            lr
+        } else if step < decay_start {
+            // Phase 2: stable plateau
+            trace!("WSD stable  step={step}: lr={:.6e}", self.peak_lr);
+            self.peak_lr
+        } else if step < self.total_steps() {
+            // Phase 3: decay
+            let t = step - decay_start;
+            let p = (t as f64 / self.decay_steps as f64).clamp(0.0, 1.0);
+            let lr = match self.decay_type {
+                WsdDecayType::Cosine => {
+                    self.min_lr + (self.peak_lr - self.min_lr) * 0.5 * (1.0 + (PI * p).cos())
+                }
+                WsdDecayType::Linear => self.peak_lr + (self.min_lr - self.peak_lr) * p,
+                WsdDecayType::Sqrt => {
+                    self.min_lr + (self.peak_lr - self.min_lr) * (1.0 - p.sqrt())
+                }
+            };
+            trace!("WSD decay   step={step}: p={p:.4}, lr={lr:.6e}");
+            lr
+        } else {
+            // Phase 4: complete — clamp to min_lr
+            trace!("WSD complete step={step}: lr={:.6e}", self.min_lr);
+            self.min_lr
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Trait implementation
+// ---------------------------------------------------------------------------
+
+impl LearningRateScheduler for WsdScheduler {
+    /// Compute LR at the given `step` (the `epoch` parameter is ignored — WSD
+    /// is step-based, consistent with how it is used in large LLM training).
+    fn get_lr(&self, _epoch: usize, step: usize) -> f64 {
+        self.compute_lr(step)
+    }
+
+    /// Advance the internal step counter by one.
+    fn step(&mut self) {
+        self.current_step += 1;
+        trace!("WSD scheduler stepped to {}", self.current_step);
+    }
+
+    fn current_step(&self) -> usize {
+        self.current_step
+    }
+
+    fn reset(&mut self) {
+        self.current_step = 0;
+        debug!("Reset WsdScheduler");
+    }
+
+    fn scheduler_type(&self) -> &'static str {
+        "WarmupStableDecay"
+    }
+
+    fn base_lr(&self) -> f64 {
+        self.peak_lr
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const PEAK: f64 = 3e-4;
+    const MIN: f64 = 1e-5;
+    const WARMUP: usize = 100;
+    const STABLE: usize = 900;
+    const DECAY: usize = 200;
+
+    fn default_sched() -> WsdScheduler {
+        WsdScheduler::cosine(PEAK, MIN, WARMUP, STABLE, DECAY).unwrap()
+    }
+
+    // ------------------------------------------------------------------
+    // Warmup phase
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_warmup_phase_starts_at_zero() {
+        let s = default_sched();
+        // step=0 gives peak_lr * 1/100, which is very small but non-zero
+        let lr0 = s.get_lr(0, 0);
+        assert!(lr0 > 0.0, "lr at step 0 should be > 0");
+        assert!(lr0 < PEAK, "lr at step 0 should be less than peak_lr");
+
+        // step warmup_steps-1 should be strictly less than peak_lr
+        let lr_last_warmup = s.get_lr(0, WARMUP - 1);
+        // (100/100) * PEAK = PEAK — actually at WARMUP-1 = 99: progress = 100/100 = 1.0 → PEAK
+        // Wait: step=99, (99+1)/100 = 100/100 = 1.0 → lr = PEAK
+        // So last warmup step IS peak. Let's check step 98 instead.
+        let lr_98 = s.get_lr(0, 98);
+        assert!(lr_98 < PEAK);
+        // At step 99 (last warmup step), progress = 100/100 = 1.0, so lr = PEAK
+        assert!((lr_last_warmup - PEAK).abs() < 1e-12);
+    }
+
+    #[test]
+    fn test_warmup_phase_reaches_peak() {
+        let s = default_sched();
+        // First step of stable phase = warmup_steps (step == WARMUP)
+        let lr = s.get_lr(0, WARMUP);
+        assert!(
+            (lr - PEAK).abs() < 1e-12,
+            "LR at step=warmup_steps must equal peak_lr, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_warmup_linear_increase() {
+        let s = default_sched();
+        // LR must be monotonically increasing during warmup
+        let mut prev = s.get_lr(0, 0);
+        for step in 1..WARMUP {
+            let cur = s.get_lr(0, step);
+            assert!(
+                cur >= prev,
+                "warmup should be monotonically non-decreasing: step {step} lr={cur} < prev={prev}"
+            );
+            prev = cur;
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // Stable phase
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_stable_phase_constant() {
+        let s = default_sched();
+        for step in WARMUP..=(WARMUP + STABLE - 1) {
+            let lr = s.get_lr(0, step);
+            assert!(
+                (lr - PEAK).abs() < 1e-12,
+                "stable phase step {step}: expected {PEAK} got {lr}"
+            );
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // Decay phase — cosine
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_decay_cosine_at_start() {
+        let s = default_sched();
+        // First decay step: p = 1/200 = 0.005; cos(π * 0.005) is very close to 1.
+        // lr ≈ peak_lr — the cosine has barely moved away from the peak.
+        let decay_start = WARMUP + STABLE;
+        let lr = s.get_lr(0, decay_start);
+        assert!(
+            (lr - PEAK).abs() < 1e-3,
+            "first decay step should be ≈ peak_lr (within 1e-3), got {lr}"
+        );
+        // Must be less than or equal to peak_lr — the decay cannot overshoot.
+        assert!(
+            lr <= PEAK,
+            "first decay step must not exceed peak_lr, got {lr}"
+        );
+        // And it must be greater than min_lr — we are only at the very start.
+        assert!(
+            lr > MIN,
+            "first decay step must still be far above min_lr, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_decay_cosine_at_end() {
+        let s = default_sched();
+        // Last decay step: step = total_steps - 1, p = (DECAY-1)/DECAY
+        let last = s.total_steps() - 1;
+        let lr = s.get_lr(0, last);
+        // cos(π * (199/200)) ≈ -1 → lr ≈ min_lr
+        assert!(
+            (lr - MIN).abs() < 1e-3,
+            "last decay step should be ≈ min_lr, got {lr}"
+        );
+        assert!(lr > MIN, "last decay step should be strictly greater than min_lr");
+    }
+
+    #[test]
+    fn test_decay_cosine_monotone_decreasing() {
+        let s = default_sched();
+        let decay_start = WARMUP + STABLE;
+        let mut prev = s.get_lr(0, decay_start);
+        for step in (decay_start + 1)..s.total_steps() {
+            let cur = s.get_lr(0, step);
+            assert!(
+                cur <= prev + 1e-15,
+                "cosine decay should be monotone non-increasing: step {step} lr={cur} > prev={prev}"
+            );
+            prev = cur;
+        }
+    }
+
+    #[test]
+    fn test_decay_linear_midpoint() {
+        let mut s = WsdScheduler::new(PEAK, MIN, WARMUP, STABLE, DECAY, WsdDecayType::Linear)
+            .unwrap();
+        // extend stable to 0 so midpoint is clearly mid-decay
+        s.stable_steps = 0;
+        let decay_start = WARMUP;
+        let mid = decay_start + DECAY / 2;
+        let lr = s.get_lr(0, mid);
+        let expected = PEAK + (MIN - PEAK) * 0.5;
+        assert!(
+            (lr - expected).abs() < 1e-10,
+            "linear decay midpoint: expected {expected}, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_decay_sqrt_midpoint() {
+        let s =
+            WsdScheduler::new(PEAK, MIN, WARMUP, STABLE, DECAY, WsdDecayType::Sqrt).unwrap();
+        let decay_start = WARMUP + STABLE;
+        let mid = decay_start + DECAY / 2;
+        let p = 0.5_f64;
+        let lr = s.get_lr(0, mid);
+        let expected = MIN + (PEAK - MIN) * (1.0 - p.sqrt());
+        assert!(
+            (lr - expected).abs() < 1e-8,
+            "sqrt decay midpoint: expected {expected}, got {lr}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Complete phase
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_complete_phase_returns_min_lr() {
+        let s = default_sched();
+        let total = s.total_steps();
+        for step in [total, total + 1, total + 1000] {
+            let lr = s.get_lr(0, step);
+            assert!(
+                (lr - MIN).abs() < 1e-12,
+                "step {step} (past end) should return min_lr={MIN}, got {lr}"
+            );
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // total_steps
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_total_steps() {
+        let s = default_sched();
+        assert_eq!(
+            s.total_steps(),
+            WARMUP + STABLE + DECAY,
+            "total_steps() must equal warmup + stable + decay"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Phase introspection
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_current_phase_warmup() {
+        let mut s = default_sched();
+        // current_step starts at 0 → Warmup
+        assert_eq!(s.current_phase(), WsdPhase::Warmup);
+        assert_eq!(s.phase_at(0), WsdPhase::Warmup);
+        assert_eq!(s.phase_at(WARMUP - 1), WsdPhase::Warmup);
+        // Step past warmup
+        for _ in 0..WARMUP {
+            s.step();
+        }
+        assert_ne!(s.current_phase(), WsdPhase::Warmup);
+    }
+
+    #[test]
+    fn test_current_phase_stable() {
+        let s = default_sched();
+        assert_eq!(s.phase_at(WARMUP), WsdPhase::Stable);
+        assert_eq!(s.phase_at(WARMUP + STABLE - 1), WsdPhase::Stable);
+    }
+
+    #[test]
+    fn test_current_phase_decay() {
+        let s = default_sched();
+        let decay_start = WARMUP + STABLE;
+        assert_eq!(s.phase_at(decay_start), WsdPhase::Decay);
+        assert_eq!(s.phase_at(decay_start + DECAY - 1), WsdPhase::Decay);
+    }
+
+    #[test]
+    fn test_current_phase_complete() {
+        let s = default_sched();
+        assert_eq!(s.phase_at(s.total_steps()), WsdPhase::Complete);
+        assert_eq!(s.phase_at(s.total_steps() + 999), WsdPhase::Complete);
+    }
+
+    // ------------------------------------------------------------------
+    // extend_stable
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_extend_stable_delays_decay() {
+        let mut s = default_sched();
+        let original_total = s.total_steps();
+        s.extend_stable(100);
+        assert_eq!(s.stable_steps, STABLE + 100);
+        assert_eq!(s.total_steps(), original_total + 100);
+
+        // The step that used to be the first decay step is now still stable
+        let old_decay_start = WARMUP + STABLE;
+        assert_eq!(
+            s.phase_at(old_decay_start),
+            WsdPhase::Stable,
+            "old decay start should now be in Stable after extension"
+        );
+        // New decay start
+        let new_decay_start = WARMUP + STABLE + 100;
+        assert_eq!(s.phase_at(new_decay_start), WsdPhase::Decay);
+    }
+
+    #[test]
+    fn test_extend_stable_lr_still_peak_at_extended_steps() {
+        let mut s = default_sched();
+        s.extend_stable(300);
+        // Any step in [WARMUP, WARMUP + STABLE + 300) should return PEAK
+        for step in [WARMUP, WARMUP + STABLE, WARMUP + STABLE + 299] {
+            let lr = s.get_lr(0, step);
+            assert!(
+                (lr - PEAK).abs() < 1e-12,
+                "extended stable step {step} should return peak_lr, got {lr}"
+            );
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // decay_progress
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_decay_progress_zero_in_stable() {
+        let s = default_sched();
+        // During warmup and stable, decay_progress should be 0.0
+        for step in [0, WARMUP / 2, WARMUP, WARMUP + STABLE - 1] {
+            assert_eq!(
+                s.decay_progress_at(step),
+                0.0,
+                "step {step} is not in decay, progress should be 0.0"
+            );
+        }
+    }
+
+    #[test]
+    fn test_decay_progress_one_at_end() {
+        let s = default_sched();
+        // Last decay step: step = total_steps - 1
+        // t = total_steps - 1 - (warmup + stable) = DECAY - 1
+        // p = (DECAY - 1) / DECAY
+        let last_decay = s.total_steps() - 1;
+        let p = s.decay_progress_at(last_decay);
+        let expected = (DECAY - 1) as f64 / DECAY as f64;
+        assert!(
+            (p - expected).abs() < 1e-12,
+            "decay progress at last decay step: expected {expected}, got {p}"
+        );
+        // And complete phase (past total_steps) gives 0.0 (outside decay)
+        assert_eq!(s.decay_progress_at(s.total_steps()), 0.0);
+    }
+
+    // ------------------------------------------------------------------
+    // Trait method wiring
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_step_advances_current_step() {
+        let mut s = default_sched();
+        assert_eq!(s.current_step(), 0);
+        s.step();
+        assert_eq!(s.current_step(), 1);
+        s.step();
+        assert_eq!(s.current_step(), 2);
+    }
+
+    #[test]
+    fn test_reset_returns_to_zero() {
+        let mut s = default_sched();
+        for _ in 0..500 {
+            s.step();
+        }
+        assert_eq!(s.current_step(), 500);
+        s.reset();
+        assert_eq!(s.current_step(), 0);
+    }
+
+    #[test]
+    fn test_scheduler_type_str() {
+        let s = default_sched();
+        assert_eq!(s.scheduler_type(), "WarmupStableDecay");
+    }
+
+    #[test]
+    fn test_base_lr_returns_peak() {
+        let s = default_sched();
+        assert_eq!(s.base_lr(), PEAK);
+    }
+
+    #[test]
+    fn test_current_lr_matches_get_lr() {
+        let mut s = default_sched();
+        for _ in 0..750 {
+            s.step();
+        }
+        assert!((s.current_lr() - s.get_lr(0, s.current_step())).abs() < 1e-15);
+    }
+
+    // ------------------------------------------------------------------
+    // Validation
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_invalid_params_rejected() {
+        // negative peak_lr
+        assert!(WsdScheduler::cosine(-1e-3, 1e-5, 100, 100, 100).is_err());
+        // zero peak_lr
+        assert!(WsdScheduler::cosine(0.0, 1e-5, 100, 100, 100).is_err());
+        // peak_lr <= min_lr
+        assert!(WsdScheduler::cosine(1e-5, 1e-5, 100, 100, 100).is_err());
+        assert!(WsdScheduler::cosine(1e-6, 1e-5, 100, 100, 100).is_err());
+        // negative min_lr
+        assert!(WsdScheduler::cosine(1e-3, -1.0, 100, 100, 100).is_err());
+        // zero warmup_steps
+        assert!(WsdScheduler::cosine(1e-3, 1e-5, 0, 100, 100).is_err());
+        // zero decay_steps
+        assert!(WsdScheduler::cosine(1e-3, 1e-5, 100, 100, 0).is_err());
+    }
+
+    #[test]
+    fn test_zero_stable_steps_allowed() {
+        // stable_steps = 0 is valid (warmup directly into decay)
+        let s = WsdScheduler::cosine(PEAK, MIN, WARMUP, 0, DECAY).unwrap();
+        assert_eq!(s.total_steps(), WARMUP + DECAY);
+        assert_eq!(s.phase_at(WARMUP), WsdPhase::Decay);
+    }
+}