From 01d2da28935d1553605556025fc51540dd921fa0 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 22 Jun 2026 10:35:14 +0800
Subject: [PATCH] =?UTF-8?q?feat=20:=20=E8=AE=AD=E7=BB=83=E6=94=AF=E6=8C=81?=
 =?UTF-8?q?=20--schedule=5Ftype=20=E5=8F=8A=E5=AF=B9=E5=BA=94=E8=B0=83?=
 =?UTF-8?q?=E5=BA=A6=E5=99=A8=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- --schedule_type 可选 cosine/sgdr/wsd，默认 cosine
- --min_rate 统一控制最小 LR 比率
- --cycle_length / --t_mult 用于 sgdr
- --stable_steps / --decay_steps 用于 wsd，自动计算默认值
---
 scripts/tools/train.py | 70 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 65 insertions(+), 5 deletions(-)

diff --git a/scripts/tools/train.py b/scripts/tools/train.py
index 999ff0a..2f3752c 100644
--- a/scripts/tools/train.py
+++ b/scripts/tools/train.py
@@ -221,6 +221,44 @@ def parse_args() -> argparse.Namespace:
         help="NEFTune noise alpha (0=disabled, typical: 5.0).",
     )
 
+    parser.add_argument(
+        "--schedule_type",
+        type=str,
+        default="cosine",
+        choices=["cosine", "sgdr", "wsd"],
+        help="Learning rate scheduler type.",
+    )
+    parser.add_argument(
+        "--min_rate",
+        type=float,
+        default=None,
+        help="Minimum LR as fraction of base LR. Uses scheduler default if not set (cosine/sgdr: 0.05, wsd: 0.0).",
+    )
+    parser.add_argument(
+        "--cycle_length",
+        type=int,
+        default=None,
+        help="SGDR first cycle length in steps. Defaults to total_steps - warmup_steps.",
+    )
+    parser.add_argument(
+        "--t_mult",
+        type=int,
+        default=2,
+        help="SGDR cycle length multiplier per restart.",
+    )
+    parser.add_argument(
+        "--stable_steps",
+        type=int,
+        default=None,
+        help="WSD stable plateau steps. Required when --schedule_type wsd.",
+    )
+    parser.add_argument(
+        "--decay_steps",
+        type=int,
+        default=None,
+        help="WSD decay steps. Defaults to total_steps - warmup_steps - stable_steps.",
+    )
+
     args = parser.parse_args()
 
     return args
@@ -300,6 +338,12 @@ def train(
     master_port: str,
     start_method: str,
     neftune_alpha: float,
+    schedule_type: str,
+    min_rate: float,
+    cycle_length: int,
+    t_mult: int,
+    stable_steps: int,
+    decay_steps: int,
 ):
     assert train_type in ["seq", "sft", "dpo", "grpo"]
     assert os.path.exists(param_path)
@@ -349,14 +393,30 @@ def train(
         len(dataset), n_epoch, batch_per_device, nprocs, grad_accum_steps
     )
     warmup_steps = int(warmup_ratio * total_steps)
+    warmup_steps = min(warmup_steps, total_steps)
+
+    scheduler_kwargs = {"warmup_steps": warmup_steps}
+
+    if schedule_type == "cosine":
+        scheduler_kwargs["lr_decay_steps"] = total_steps - warmup_steps
+    elif schedule_type == "sgdr":
+        scheduler_kwargs["cycle_length"] = cycle_length or (total_steps - warmup_steps)
+        scheduler_kwargs["t_mult"] = t_mult
+    elif schedule_type == "wsd":
+        remaining = total_steps - warmup_steps
+        stable_steps_ = stable_steps or max(1, int(remaining * 0.8))
+        scheduler_kwargs["stable_steps"] = stable_steps_
+        scheduler_kwargs["decay_steps"] = max(
+            1, decay_steps or (remaining - stable_steps_)
+        )
+
+    if min_rate is not None:
+        scheduler_kwargs["min_rate"] = min_rate
 
     scheduler_fn = partial(
         create_scheduler,
-        **{
-            "schedule_type": "cosine",
-            "warmup_steps": min(warmup_steps, total_steps),
-            "lr_decay_steps": total_steps - min(warmup_steps, total_steps),
-        },
+        schedule_type=schedule_type,
+        **scheduler_kwargs,
     )
 
     grad_ckpt_modules = [DecoderBlock] if gradient_checkpointing else []