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="梯度累积步数")