[feat] get_lr 支持可选的线性 warmup（默认关闭，行为不变）

get_lr 在 current_step=0 时返回 lr*(0.1+0.45*2) = 满 lr，没有 warmup。
从随机初始化开始的预训练在前几十步梯度范数通常是稳态的若干倍，
满 lr 容易把刚初始化的权重推坏。

改动保持默认行为逐值不变：
- get_lr 新增 warmup_frac=0.0 参数；为 0 时走原来的表达式，不做任何改动
- 5 个共用该函数的脚本（train_pretrain / train_full_sft / train_lora /
  train_dpo / train_distillation）各加一个 --warmup_frac，默认 0.0

验证：在 total_steps ∈ {100, 1000, 7339, 39634} 上取 398 个点，
warmup_frac=0 与原实现的最大差异是 0.000e+00。

trainer/train_distillation.py

@@ -49,7 +49,7 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
         loss_mask = (labels[..., 1:] != -100).float()
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_frac)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -151,6 +151,8 @@ if __name__ == "__main__":
     parser.add_argument("--epochs", type=int, default=6, help="训练轮数")
     parser.add_argument("--batch_size", type=int, default=32, help="batch size")
     parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率")
+    parser.add_argument("--warmup_frac", type=float, default=0.0,
+                        help="线性 warmup 占总步数的比例，0 表示不用（保持原行为）")
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
     parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
     parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")

trainer/train_dpo.py

@@ -66,7 +66,7 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
         y = torch.cat([y_chosen, y_rejected], dim=0)
         mask = torch.cat([mask_chosen, mask_rejected], dim=0)
 
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_frac)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -135,6 +135,8 @@ if __name__ == "__main__":
     parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
     parser.add_argument("--batch_size", type=int, default=4, help="batch size")
     parser.add_argument("--learning_rate", type=float, default=4e-8, help="初始学习率（建议<=5e-8避免遗忘）")
+    parser.add_argument("--warmup_frac", type=float, default=0.0,
+                        help="线性 warmup 占总步数的比例，0 表示不用（保持原行为）")
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
     parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
     parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")

trainer/train_full_sft.py

@@ -28,7 +28,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
         last_step = step
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_frac)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -88,6 +88,8 @@ if __name__ == "__main__":
     parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
     parser.add_argument("--batch_size", type=int, default=16, help="batch size")
     parser.add_argument("--learning_rate", type=float, default=1e-5, help="初始学习率")
+    parser.add_argument("--warmup_frac", type=float, default=0.0,
+                        help="线性 warmup 占总步数的比例，0 表示不用（保持原行为）")
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
     parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
     parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")

trainer/train_lora.py

@@ -29,7 +29,7 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
         last_step = step
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_frac)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -82,6 +82,8 @@ if __name__ == "__main__":
     parser.add_argument("--epochs", type=int, default=10, help="训练轮数")
     parser.add_argument("--batch_size", type=int, default=32, help="batch size")
     parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
+    parser.add_argument("--warmup_frac", type=float, default=0.0,
+                        help="线性 warmup 占总步数的比例，0 表示不用（保持原行为）")
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
     parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
     parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")

trainer/train_pretrain.py

@@ -28,7 +28,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
         last_step = step
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_frac)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -87,6 +87,8 @@ if __name__ == "__main__":
     parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
     parser.add_argument("--batch_size", type=int, default=32, help="batch size")
     parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
+    parser.add_argument("--warmup_frac", type=float, default=0.0,
+                        help="线性 warmup 占总步数的比例，0 表示不用（保持原行为）")
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
     parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
     parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")

trainer/trainer_utils.py

@@ -39,7 +39,22 @@ 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。
+
+    warmup_frac=0（默认）时与原实现逐值相同，不改变任何已有脚本的行为。
+
+    warmup_frac>0 时，前 warmup_frac*total_steps 步把 lr 从 lr/warmup_steps
+    线性升到满值，之后在剩余步数上走原来的余弦衰减（1.0x -> 0.1x）。
+    从随机初始化开始时 step 0 就给满 lr，前几十步的梯度范数通常是稳态的
+    若干倍，容易把刚初始化的权重推坏；warmup 是标准做法。
+    """
+    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)))