From a4e5a8c81c389bb549eabec495cc82a075ddccf5 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Thu, 18 Jun 2026 15:55:15 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=20WSD=20=E5=AD=A6?=
 =?UTF-8?q?=E4=B9=A0=E7=8E=87=E8=B0=83=E5=BA=A6=E5=99=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 支持 Warmup-Stable-Decay 三段式调度
- stable 阶段保持最高 lr，decay 阶段 sqrt 衰减
- 适用于持续预训练、SFT、RLHF 场景
---
 astrai/trainer/schedule.py | 63 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/astrai/trainer/schedule.py b/astrai/trainer/schedule.py
index 5c0225e..b49cca5 100644
--- a/astrai/trainer/schedule.py
+++ b/astrai/trainer/schedule.py
@@ -164,3 +164,66 @@ class SGDRScheduler(BaseScheduler):
         self.min_rate = state_dict.pop("min_rate")
         self.t_mult = state_dict.pop("t_mult")
         super().load_state_dict(state_dict)
+
+
+@SchedulerFactory.register("wsd")
+class WSDScheduler(BaseScheduler):
+    """WSD (Warmup-Stable-Decay) scheduler with sqrt cooldown.
+
+    warmup_steps:  linear warmup from min_rate to 1.0
+    stable_steps:   constant at base_lr
+    decay_steps:    sqrt decay from base_lr to min_rate
+    min_rate:       minimum lr as fraction of base_lr (default 0.0)
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_steps: int,
+        stable_steps: int,
+        decay_steps: int,
+        min_rate: float = 0.0,
+        last_epoch: int = -1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.stable_steps = stable_steps
+        self.decay_steps = decay_steps
+        self.min_rate = min_rate
+        self.total_steps = warmup_steps + stable_steps + decay_steps
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> List[float]:
+        if self.last_epoch < self.warmup_steps:
+            factor = self.last_epoch / max(self.warmup_steps, 1)
+            return [base_lr * factor for base_lr in self.base_lrs]
+
+        offset = self.last_epoch - self.warmup_steps
+
+        if offset < self.stable_steps:
+            return list(self.base_lrs)
+
+        decay_ratio = (offset - self.stable_steps) / max(self.decay_steps, 1)
+        decay_ratio = min(decay_ratio, 1.0)
+        factor = (1.0 - self.min_rate) * (1.0 - decay_ratio) ** 2 + self.min_rate
+        return [base_lr * factor for base_lr in self.base_lrs]
+
+    def state_dict(self):
+        state = super().state_dict()
+        state.update(
+            {
+                "warmup_steps": self.warmup_steps,
+                "stable_steps": self.stable_steps,
+                "decay_steps": self.decay_steps,
+                "min_rate": self.min_rate,
+                "total_steps": self.total_steps,
+            }
+        )
+        return state
+
+    def load_state_dict(self, state_dict):
+        self.warmup_steps = state_dict.pop("warmup_steps")
+        self.stable_steps = state_dict.pop("stable_steps")
+        self.decay_steps = state_dict.pop("decay_steps")
+        self.min_rate = state_dict.pop("min_rate")
+        self.total_steps = state_dict.pop("total_steps")
+        super().load_state_dict(state_dict)