diff --git a/.gitignore b/.gitignore index de6bd7a..0518e52 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,12 @@ /model/__pycache__ /dataset -/out \ No newline at end of file +/out +# Python 缓存文件 +__pycache__/ +*.pyc +*.pyo +*.pyd +# macOS 系统文件 +.DS_Store +# Jupyter Notebook 检查点 +.ipynb_checkpoints/ \ No newline at end of file diff --git a/README.md b/README.md index 8ef9609..e5c3c83 100644 --- a/README.md +++ b/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`这个路径。 + +
注：测试须知 @@ -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 ``` diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index 4a62cd3..1a977aa 100644 --- a/trainer/train_distill_reason.py +++ b/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 diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index c0607de..cf72962 100644 --- a/trainer/train_distillation.py +++ b/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 diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 7987d93..9cf9cfb 100644 --- a/trainer/train_dpo.py +++ b/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 diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 981d926..65d6555 100644 --- a/trainer/train_full_sft.py +++ b/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 diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 04fefcf..1209ff5 100644 --- a/trainer/train_lora.py +++ b/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 diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index b276eb7..6672e55 100644 --- a/trainer/train_pretrain.py +++ b/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