From 7227c6a21c1570229ac4c1496749a6c0f3b21aea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B9=BF=E5=B3=B0?= Date: Sat, 16 May 2026 01:26:42 +0800 Subject: [PATCH] =?UTF-8?q?chore(args):=20=E6=94=AF=E6=8C=81Mac=20MPS?= =?UTF-8?q?=E8=AE=BE=E5=A4=87=E4=BD=9C=E4=B8=BA=E8=BF=90=E8=A1=8C=E5=92=8C?= =?UTF-8?q?=E8=AE=AD=E7=BB=83=E8=AE=BE=E5=A4=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 修改eval_llm.py脚本支持MPS设备（macOS GPU加速） - 修改eval_toolcall.py脚本支持MPS设备 - 修改serve_openai_api.py脚本支持MPS设备 - 修改多个训练脚本（train_agent.py、train_distillation.py、train_dpo.py、train_full_sft.py、train_grpo.py、 train_lora.py、train_ppo.py、train_pretrain.py）支持MPS设备 - 优先判断CUDA，其次判断MPS，最后回退到CPU设备 - 统一设备选择逻辑以适配不同硬件加速环境 --- eval_llm.py | 2 +- scripts/eval_toolcall.py | 2 +- scripts/serve_openai_api.py | 2 +- trainer/train_agent.py | 2 +- trainer/train_distillation.py | 2 +- trainer/train_dpo.py | 2 +- trainer/train_full_sft.py | 2 +- trainer/train_grpo.py | 2 +- trainer/train_lora.py | 2 +- trainer/train_ppo.py | 2 +- trainer/train_pretrain.py | 2 +- 11 files changed, 11 insertions(+), 11 deletions(-) diff --git a/eval_llm.py b/eval_llm.py index 9359ca8..b303a10 100755 --- a/eval_llm.py +++ b/eval_llm.py @@ -45,7 +45,7 @@ def main(): parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）") parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）") parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）") - parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") + parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() and torch.backends.mps.is_built() else 'cpu', type=str, help="运行设备") args = parser.parse_args() prompts = [ diff --git a/scripts/eval_toolcall.py b/scripts/eval_toolcall.py index c4a07b9..bc0c2fc 100644 --- a/scripts/eval_toolcall.py +++ b/scripts/eval_toolcall.py @@ -224,7 +224,7 @@ def main(): parser.add_argument('--temperature', default=0.9, type=float, help="生成温度，控制随机性（0-1，越大越随机）") parser.add_argument('--top_p', default=0.9, type=float, help="nucleus采样阈值（0-1）") parser.add_argument('--show_speed', default=0, type=int, help="显示decode速度（tokens/s）") - parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") + parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() and torch.backends.mps.is_built() else 'cpu', type=str, help="运行设备") parser.add_argument('--api_base_url', default="http://localhost:11434/v1", type=str, help="OpenAI兼容接口的base_url") parser.add_argument('--api_key', default='sk-123', type=str, help="OpenAI兼容接口的api_key") parser.add_argument('--api_model', default='jingyaogong/minimind-3:latest', type=str, help="API请求时使用的模型名称") diff --git a/scripts/serve_openai_api.py b/scripts/serve_openai_api.py index 0550d4a..280f4fc 100644 --- a/scripts/serve_openai_api.py +++ b/scripts/serve_openai_api.py @@ -238,7 +238,7 @@ if __name__ == "__main__": parser.add_argument('--max_seq_len', default=8192, type=int, help="最大序列长度") parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）") - parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") + parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() and torch.backends.mps.is_built() else 'cpu', type=str, help="运行设备") args = parser.parse_args() device = args.device model, tokenizer = init_model(args) diff --git a/trainer/train_agent.py b/trainer/train_agent.py index e652c3e..0f69221 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -377,7 +377,7 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--batch_size", type=int, default=2, help="批次大小") parser.add_argument("--learning_rate", type=float, default=3e-7, help="学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 63b6025..e01dd51 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -150,7 +150,7 @@ 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("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 558d108..a02176e 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -134,7 +134,7 @@ 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("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 51c755b..d47de0b 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -87,7 +87,7 @@ 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("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 52f7c7d..e2da32f 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -209,7 +209,7 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--batch_size", type=int, default=2, help="batch size") parser.add_argument("--learning_rate", type=float, default=3e-7, help="初始学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 974bf37..5decedb 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -81,7 +81,7 @@ 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("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 228fb68..883d4b2 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -300,7 +300,7 @@ if __name__ == "__main__": parser.add_argument("--batch_size", type=int, default=2, help="batch size") parser.add_argument("--learning_rate", type=float, default=3e-7, help="Actor学习率") parser.add_argument("--critic_learning_rate", type=float, default=5e-7, help="Critic学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 5d33de4..9d23f52 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -86,7 +86,7 @@ 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("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() and torch.backends.mps.is_built() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")