chore: update .gitignore 和 trainer 相关脚本，完善文档

.gitignore

@@ -1,3 +1,12 @@
 /model/__pycache__
 /dataset
 /out
+# Python 缓存文件
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+# macOS 系统文件
+.DS_Store
+# Jupyter Notebook 检查点
+.ipynb_checkpoints/

README.md

@@ -308,6 +308,10 @@ print(torch.cuda.is_available())
 
 ```bash
 python train_pretrain.py
+
+# or
+nohup torchrun --nproc_per_node=8 train_pretrain.py --use_wandb --wandb_api_key=你的key --ddp --batch_size=32 > train_pretrain.log 2>&1
+
 ```
 
 > 执行预训练，得到 `pretrain_*.pth` 作为预训练的输出权重（其中*为模型的dimension，默认为512）
@@ -342,6 +346,13 @@ python train_full_sft.py
 python eval_model.py --model_mode 1 # 默认为0：测试pretrain模型效果，设置为1：测试full_sft模型效果
 ```
 
+例如，如果你想加载`./out/pretrain_512_moe.pth`，应该这样写：
+```bash
+python eval_model.py --model_mode 0 --out_dir out --hidden_size 512 --use_moe True
+```
+这样脚本会自动拼出`./out/pretrain_512_moe.pth`这个路径。
+
+
 <details style="color:rgb(128,128,128)">
 <summary>注：测试须知</summary>
 
@@ -375,6 +386,10 @@ deepspeed --master_port 29500 --num_gpus=N train_xxx.py
 ```bash
 # 需要登录: wandb login
 torchrun --nproc_per_node N train_xxx.py --use_wandb
+
+# 通过py 脚本登录wandb
+nohup torchrun --nproc_per_node=8 train_pretrain.py --use_wandb --wandb_api_key=你的key --ddp --batch_size=32 > train_pretrain.log 2>&1
+
 # and
 python train_xxx.py --use_wandb
 ```

trainer/train_distill_reason.py

@@ -218,6 +218,8 @@ if __name__ == "__main__":
     parser.add_argument('--use_moe', default=False, type=bool)
     # 训练数据路径
     parser.add_argument("--data_path", type=str, default="../dataset/r1_mix_1024.jsonl")
+    # wandb API Key
+    parser.add_argument("--wandb_api_key", type=str, default=None, help="WandB API Key，用于无交互环境自动登录")
 
     args = parser.parse_args()
 
@@ -248,7 +250,12 @@ if __name__ == "__main__":
 
     if args.use_wandb and (not ddp or ddp_local_rank == 0):
         import wandb
-
+        # 优先使用命令行参数，其次环境变量
+        api_key = args.wandb_api_key or os.environ.get("WANDB_API_KEY", None)
+        if api_key is not None:
+            # 自动登录，适用于无交互环境
+            wandb.login(key=api_key)
+        # 初始化 wandb 任务
         wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
         wandb = None

trainer/train_distillation.py

@@ -230,6 +230,8 @@ if __name__ == "__main__":
     parser.add_argument("--max_seq_len", type=int, default=512)
     parser.add_argument('--local_rank', type=int, default=-1)
     parser.add_argument("--data_path", type=str, default="../dataset/sft_xxx.jsonl")
+    parser.add_argument("--wandb_api_key", type=str, default=None, help="WandB API Key，用于无交互环境自动登录")
+
 
     args = parser.parse_args()
     # 定义学生模型和教师模型
@@ -260,7 +262,12 @@ if __name__ == "__main__":
 
     if args.use_wandb and (not ddp or ddp_local_rank == 0):
         import wandb
-
+        # 优先使用命令行参数，其次环境变量
+        api_key = args.wandb_api_key or os.environ.get("WANDB_API_KEY", None)
+        if api_key is not None:
+            # 自动登录，适用于无交互环境
+            wandb.login(key=api_key)
+        # 初始化 wandb 任务
         wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
         wandb = None

trainer/train_dpo.py

@@ -241,6 +241,8 @@ if __name__ == "__main__":
     # 数据参数
     parser.add_argument("--data_path", type=str, default="../dataset/dpo.jsonl", help="训练数据路径，包含人类偏好对")
 
+    parser.add_argument("--wandb_api_key", type=str, default=None, help="WandB API Key，用于无交互环境自动登录")
+
     args = parser.parse_args()
 
     lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=args.use_moe)
@@ -269,7 +271,12 @@ if __name__ == "__main__":
 
     if args.use_wandb and (not ddp or ddp_local_rank == 0):
         import wandb
-
+        # 优先使用命令行参数，其次环境变量
+        api_key = args.wandb_api_key or os.environ.get("WANDB_API_KEY", None)
+        if api_key is not None:
+            # 自动登录，适用于无交互环境
+            wandb.login(key=api_key)
+        # 初始化 wandb 任务
         wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
         wandb = None

trainer/train_full_sft.py

@@ -191,6 +191,7 @@ if __name__ == "__main__":
     parser.add_argument('--max_seq_len', default=512, type=int, help="最大序列长度")
     parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE")
     parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl", help="训练数据路径")
+    parser.add_argument("--wandb_api_key", type=str, default=None, help="WandB API Key，用于无交互环境自动登录")
 
     args = parser.parse_args()
 
@@ -231,6 +232,12 @@ if __name__ == "__main__":
     # 初始化wandb（如果启用）
     if args.use_wandb and (not ddp or ddp_local_rank == 0):
         import wandb
+        # 优先使用命令行参数，其次环境变量
+        api_key = args.wandb_api_key or os.environ.get("WANDB_API_KEY", None)
+        if api_key is not None:
+            # 自动登录，适用于无交互环境
+            wandb.login(key=api_key)
+        # 初始化 wandb 任务
         wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
         wandb = None

trainer/train_lora.py

@@ -159,6 +159,7 @@ if __name__ == "__main__":
     # 数据和LoRA参数
     parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl", help="训练数据路径")
     parser.add_argument("--lora_name", type=str, default="lora_medical", help="根据任务保存成lora_(英文/医学/心理...)")
+    parser.add_argument("--wandb_api_key", type=str, default=None, help="WandB API Key，用于无交互环境自动登录")
     args = parser.parse_args()
 
     lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
@@ -187,7 +188,12 @@ if __name__ == "__main__":
     args.wandb_run_name = f"MiniMind-Lora-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
     if args.use_wandb and (not ddp or ddp_local_rank == 0):
         import wandb
-
+        # 优先使用命令行参数，其次环境变量
+        api_key = args.wandb_api_key or os.environ.get("WANDB_API_KEY", None)
+        if api_key is not None:
+            # 自动登录，适用于无交互环境
+            wandb.login(key=api_key)
+        # 初始化 wandb 任务
         wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
         wandb = None

trainer/train_pretrain.py

@@ -154,6 +154,7 @@ if __name__ == "__main__":
     parser.add_argument('--max_seq_len', default=512, type=int)
     parser.add_argument('--use_moe', default=False, type=bool)
     parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl")
+    parser.add_argument("--wandb_api_key", type=str, default=None, help="WandB API Key，用于无交互环境自动登录")
     args = parser.parse_args()
 
     lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=args.use_moe)
@@ -184,7 +185,12 @@ if __name__ == "__main__":
 
     if args.use_wandb and (not ddp or ddp_local_rank == 0):
         import wandb
-
+        # 优先使用命令行参数，其次环境变量
+        api_key = args.wandb_api_key or os.environ.get("WANDB_API_KEY", None)
+        if api_key is not None:
+            # 自动登录，适用于无交互环境
+            wandb.login(key=api_key)
+        # 初始化 wandb 任务
         wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
         wandb = None