diff --git a/checkpoints/wu_128.pth b/checkpoints/wu_128.pth
new file mode 100644
index 0000000..ca95b06
Binary files /dev/null and b/checkpoints/wu_128.pth differ
diff --git a/checkpoints/wu_128_resume.pth b/checkpoints/wu_128_resume.pth
new file mode 100644
index 0000000..e8d67cb
Binary files /dev/null and b/checkpoints/wu_128_resume.pth differ
diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py
index 085b2a7..88f215d 100644
--- a/trainer/train_distillation.py
+++ b/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="数据加载线程数")
diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py
index 479eb85..98b24dc 100644
--- a/trainer/train_dpo.py
+++ b/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="数据加载线程数")
diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py
index c22eb80..216fd36 100644
--- a/trainer/train_full_sft.py
+++ b/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="数据加载线程数")
diff --git a/trainer/train_lora.py b/trainer/train_lora.py
index 19a1499..cf730a2 100644
--- a/trainer/train_lora.py
+++ b/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="数据加载线程数")
diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py
index 83a0382..67e2261 100644
--- a/trainer/train_pretrain.py
+++ b/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="数据加载线程数")
diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py
index b171289..ad99ee4 100644
--- a/trainer/trainer_utils.py
+++ b/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)))