From f81a948b52a62681dad36e742faab6975e3d95f2 Mon Sep 17 00:00:00 2001
From: Linxiushen <1787979356@qq.com>
Date: Sun, 20 Sep 2026 21:41:23 +0800
Subject: [PATCH] =?UTF-8?q?[feat]=20get=5Flr=20=E5=A2=9E=E5=8A=A0=E5=8F=AF?=
 =?UTF-8?q?=E9=80=89=20warmup=5Ffrac=20=E5=8F=82=E6=95=B0=EF=BC=88?=
 =?UTF-8?q?=E9=BB=98=E8=AE=A4=E5=85=B3=E9=97=AD=EF=BC=8C=E8=A1=8C=E4=B8=BA?=
 =?UTF-8?q?=E9=80=90=E5=80=BC=E4=B8=8D=E5=8F=98=EF=BC=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

这是 #862 的收窄重提。#862 同时改了 5 个训练脚本的 argparse，面太宽；
这一版只改 trainer_utils.py 里的 get_lr 本身，不动任何训练脚本、不加任何命令行参数。

get_lr(current_step, total_steps, lr) 在 current_step=0 时返回满 lr，没有 warmup。
需要 warmup 的使用者（例如从随机初始化跑较大 batch 的预训练）现在可以传第四个
参数；不传则走原来那一行表达式。

在 total_steps ∈ {100, 1000, 7339, 39634, 150000} 上取 588 个采样点，
warmup_frac=0 与原实现的最大差异是 0.000e+00。
---
 trainer/trainer_utils.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py
index b171289..78e3eb2 100644
--- a/trainer/trainer_utils.py
+++ b/trainer/trainer_utils.py
@@ -39,7 +39,15 @@ def Logger(content):
         print(content)
 
 
-def get_lr(current_step, total_steps, lr):
+def get_lr(current_step, total_steps, lr, warmup_frac=0.0):
+    # warmup_frac=0（默认）时与原实现逐值相同。需要 warmup 的场景可以传入一个比例，
+    # 前 warmup_frac*total_steps 步把 lr 线性升到满值，其后在剩余步数上走原来的余弦。
+    if warmup_frac > 0:
+        warmup_steps = max(1, int(warmup_frac * total_steps))
+        if current_step < warmup_steps:
+            return lr * (current_step + 1) / warmup_steps
+        current_step -= warmup_steps
+        total_steps = max(1, total_steps - warmup_steps)
     return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))