From 67debd36325a40a852f12256e4b37016c42668aa Mon Sep 17 00:00:00 2001
From: yangnianboy <yangnianboy@gmail.com>
Date: Fri, 27 Jun 2025 01:00:46 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=A8=A1=E5=9E=8B=E5=8A=A0?=
 =?UTF-8?q?=E8=BD=BD=E5=92=8C=E8=AE=AD=E7=BB=83=E8=BF=87=E7=A8=8B=EF=BC=8C?=
 =?UTF-8?q?=E6=B7=BB=E5=8A=A0torch.compile=E4=BB=A5=E6=8F=90=E5=8D=87?=
 =?UTF-8?q?=E6=80=A7=E8=83=BD=EF=BC=8C=E8=B0=83=E6=95=B4=E8=AE=AD=E7=BB=83?=
 =?UTF-8?q?=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 eval_model.py             |  3 ++-
 trainer/train_full_sft.py | 14 ++++++++++++--
 trainer/train_pretrain.py | 11 +++++++----
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/eval_model.py b/eval_model.py
index f0dad41..07c2c5e 100644
--- a/eval_model.py
+++ b/eval_model.py
@@ -21,7 +21,8 @@ def init_model(args):
             num_hidden_layers=args.num_hidden_layers,
             use_moe=args.use_moe
         ))
-
+        model = torch.compile(model)
+        
         model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
 
         if args.lora_name != 'None':
diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py
index e2a9a2f..5e8f95e 100644
--- a/trainer/train_full_sft.py
+++ b/trainer/train_full_sft.py
@@ -103,6 +103,15 @@ def init_model(lm_config):
 
     Logger(f'LLM可训练总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
     model = model.to(args.device)
+
+    # 使用 torch.compile 优化模型以提升速度 (需要 PyTorch 2.0+)
+    try:
+        # 使用 "reduce-overhead" 模式来降低编译时的内存消耗
+        model = torch.compile(model, mode="reduce-overhead")
+        Logger("Model compiled with torch.compile() for optimization.")
+    except Exception as e:
+        Logger(f"torch.compile() failed: {e}")
+
     return model, tokenizer
 
 
@@ -130,7 +139,7 @@ if __name__ == "__main__":
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
-    parser.add_argument("--accumulation_steps", type=int, default=1)
+    parser.add_argument("--accumulation_steps", type=int, default=2)
     parser.add_argument("--grad_clip", type=float, default=1.0)
     parser.add_argument("--warmup_iters", type=int, default=0)
     parser.add_argument("--log_interval", type=int, default=100)
@@ -191,7 +200,8 @@ if __name__ == "__main__":
     )
 
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    # 在 CUDA 上使用 fused AdamW 以获得潜在的速度提升
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=(device_type == 'cuda'))
 
     if ddp:
         model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py
index eb45963..d3f78c1 100644
--- a/trainer/train_pretrain.py
+++ b/trainer/train_pretrain.py
@@ -18,6 +18,7 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import PretrainDataset
 
 warnings.filterwarnings('ignore')
+torch.set_float32_matmul_precision('high') # <--- 添加此行
 
 
 def Logger(content):
@@ -118,8 +119,8 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Pretraining")
     parser.add_argument("--out_dir", type=str, default="../out")
     # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
-    parser.add_argument("--epochs", type=int, default=1)
-    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--epochs", type=int, default=2)
+    parser.add_argument("--batch_size", type=int, default=32)
     parser.add_argument("--learning_rate", type=float, default=5e-4)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--dtype", type=str, default="bfloat16")
@@ -174,6 +175,7 @@ if __name__ == "__main__":
         wandb = None
 
     model, tokenizer = init_model(lm_config)
+    model = torch.compile(model) # <--- 添加此行
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(
@@ -187,11 +189,12 @@ if __name__ == "__main__":
     )
 
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=True if device_type == 'cuda' else False) # <--- 修改此行
 
     if ddp:
         model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
-        model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
+        # 注意：torch.compile 应该在 DDP 封装之前调用
+        model = DistributedDataParallel(model, device_ids=[ddp_local_rank], static_graph=True) # <--- 修改此行
 
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):