Merge efd8073450 into 512eed0b65

README.md

@@ -765,16 +765,16 @@ LoRA 是一种常见的参数高效微调（Parameter-Efficient Fine-Tuning, PEF
 ```bash
 # train_lora.py 在 CPU 上通常也能比较轻快地完成
 # 方式1
-torchrun --nproc_per_node 1 train_lora.py
+cd trainer && torchrun --nproc_per_node 1 train_lora.py
 # 方式2
-python train_lora.py
+cd trainer && python train_lora.py
 ```
 
 > 训练后的模型权重文件默认每隔`save_interval步`保存为: `lora_xxx_*.pth`（*为模型具体dimension，每次保存时新文件会覆盖旧文件）
 
 
 LoRA 很适合处理“如何在尽量保留通用能力的前提下，让模型快速适应私有领域或垂直场景”这类问题。例如基础模型医学知识不足时，就可以在原有模型之上叠加一层面向医疗场景的 LoRA 权重，以较小代价获得更好的领域表现。
-通常只需要准备同样的多轮对话格式数据，放置到 `lora_xxx.jsonl`，再执行 `python train_lora.py`，即可得到新的 `LoRA` 模型权重。
+通常只需要准备同样的多轮对话格式数据，放置到 `lora_xxx.jsonl`，再从仓库根目录执行 `cd trainer && python train_lora.py`，即可得到新的 `LoRA` 模型权重。
 
 例1：垂域数据
 

README_en.md

@@ -765,16 +765,16 @@ Its core idea is to introduce low-rank incremental branches alongside the origin
 ```bash
 # train_lora.py can usually be completed fairly quickly even on CPU
 # Method 1
-torchrun --nproc_per_node 1 train_lora.py
+cd trainer && torchrun --nproc_per_node 1 train_lora.py
 # Method 2
-python train_lora.py
+cd trainer && python train_lora.py
 ```
 
 > The trained model weight files are saved by default every `save_interval steps` as: `lora_xxx_*.pth` (* is the specific model dimension, each save overwrites the previous file)
 
 
 LoRA is well-suited for handling problems like "how to let the model quickly adapt to private domains or vertical scenarios while preserving general capabilities as much as possible." For example, when the base model lacks medical knowledge, a medical-oriented LoRA weight layer can be stacked on top of the original model to achieve better domain performance at relatively small cost.
-Usually you only need to prepare multi-turn dialogue format data in the same way, place it in `lora_xxx.jsonl`, and then run `python train_lora.py` to obtain new `LoRA` model weights.
+Usually you only need to prepare multi-turn dialogue format data in the same way, place it in `lora_xxx.jsonl`, and then run `cd trainer && python train_lora.py` from the repository root to obtain new `LoRA` model weights.
 
 Example 1: Vertical domain data
 

eval_llm.py

@@ -3,6 +3,10 @@ import argparse
 import random
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from model.model_lora import *
@@ -45,7 +49,10 @@ 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 ('npu' if torch_npu and torch_npu.npu.is_available() else 'cpu'),
+                        type=str,
+                        help="运行设备")
     args = parser.parse_args()
     
     prompts = [

scripts/chat_api.py

@@ -21,13 +21,19 @@ while True:
         extra_body={"chat_template_kwargs": {"open_thinking": True}, "reasoning_effort": "medium"} # 思考开关
     )
     if not stream:
+        if not response.choices or response.choices[0].message is None:
+            raise ValueError("LLM returned empty or filtered response")
         assistant_res = response.choices[0].message.content
         print('[A]: ', assistant_res)
     else:
         print('[A]: ', end='', flush=True)
         assistant_res = ''
         for chunk in response:
+            if not chunk.choices:
+                continue
             delta = chunk.choices[0].delta
+            if delta is None:
+                continue
             r = getattr(delta, 'reasoning_content', None) or ""
             c = delta.content or ""
             if r:

scripts/eval_toolcall.py

@@ -8,6 +8,10 @@ import random
 import argparse
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 from datetime import datetime
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from openai import OpenAI
@@ -212,11 +216,14 @@ 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('--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请求时使用的模型名称")
     parser.add_argument('--stream', default=1, type=int, help="API模式下是否流式输出（0=否，1=是）")
+    parser.add_argument('--device',
+                        default='cuda' if torch.cuda.is_available() else ('npu' if torch_npu and torch_npu.npu.is_available() else 'cpu'),
+                        type=str,
+                        help="运行设备")
     args = parser.parse_args()
 
     model = tokenizer = client = None

scripts/serve_openai_api.py

@@ -8,6 +8,10 @@ __package__ = "scripts"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 import time
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import warnings
 import uvicorn
 
@@ -15,7 +19,7 @@ from threading import Thread
 from queue import Queue
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from model.model_lora import apply_lora, load_lora
@@ -54,7 +58,7 @@ class ChatRequest(BaseModel):
     top_p: float = 0.92
     max_tokens: int = 8192
     stream: bool = True
-    tools: list = []
+    tools: list = Field(default_factory=list)
     open_thinking: bool = False
     chat_template_kwargs: dict = None
     
@@ -104,13 +108,14 @@ def parse_response(text):
 
 def generate_stream_response(messages, temperature, top_p, max_tokens, tools=None, open_thinking=False):
     try:
-        new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools or None, open_thinking=open_thinking)[-max_tokens:]
+        new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools or None, open_thinking=open_thinking)
         inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
 
         queue = Queue()
         streamer = CustomStreamer(tokenizer, queue)
 
         def _generate():
+            try:
                 model.generate(
                     inputs.input_ids,
                     max_new_tokens=max_tokens,
@@ -122,6 +127,9 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
                     eos_token_id=tokenizer.eos_token_id,
                     streamer=streamer
                 )
+            except Exception as e:
+                queue.put({"error": str(e)})
+                queue.put(None)
 
         Thread(target=_generate).start()
 
@@ -133,6 +141,9 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
             text = queue.get()
             if text is None:
                 break
+            if isinstance(text, dict):
+                yield json.dumps(text, ensure_ascii=False)
+                continue
             full_text += text
 
             if not thinking_ended:
@@ -190,7 +201,7 @@ async def chat_completions(request: ChatRequest):
                 add_generation_prompt=True,
                 tools=request.tools or None,
                 open_thinking=request.get_open_thinking()
-            )[-request.max_tokens:]
+            )
             inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
             with torch.no_grad():
                 generated_ids = model.generate(
@@ -238,7 +249,10 @@ 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 ('npu' if torch_npu and torch_npu.npu.is_available() else 'cpu'),
+                        type=str,
+                        help="运行设备")
     args = parser.parse_args()
     device = args.device
     model, tokenizer = init_model(args)

scripts/web_demo.py

@@ -5,6 +5,10 @@ import os
 from threading import Thread
 
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import numpy as np
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -67,7 +71,7 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else ("npu" if torch_npu and torch_npu.npu.is_available() else "cpu")
 
 # 多语言文本
 LANG_TEXTS = {
@@ -339,8 +343,8 @@ def main():
         st.markdown(
             f'<div style="display: flex; justify-content: flex-end;"><div style="display: inline-block; margin: 10px 0; padding: 8px 12px 8px 12px; background-color: #3d4450; border-radius: 22px; color: white;">{prompt}</div></div>',
             unsafe_allow_html=True)
-        messages.append({"role": "user", "content": prompt[-st.session_state.max_new_tokens:]})
-        st.session_state.chat_messages.append({"role": "user", "content": prompt[-st.session_state.max_new_tokens:]})
+        messages.append({"role": "user", "content": prompt})
+        st.session_state.chat_messages.append({"role": "user", "content": prompt})
 
         placeholder = st.empty()
 

trainer/train_agent.py

@@ -14,6 +14,10 @@ import signal
 import argparse
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.nn.functional as F
 import torch.distributed as dist
 from contextlib import nullcontext
@@ -378,7 +382,6 @@ 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("--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="梯度累积步数")
@@ -410,10 +413,14 @@ if __name__ == "__main__":
     parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998", help="SGLang服务器URL")
     parser.add_argument("--sglang_model_path", type=str, default="../model", help="SGLang tokenizer路径")
     parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agent", help="SGLang共享存储路径")
+    parser.add_argument("--device",
+                        type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
 
     os.makedirs(args.save_dir, exist_ok=True)
@@ -421,7 +428,7 @@ if __name__ == "__main__":
                                max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
 
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
 

trainer/train_distillation.py

@@ -9,6 +9,10 @@ import argparse
 import time
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.nn.functional as F
 import torch.distributed as dist
 from contextlib import nullcontext
@@ -151,7 +155,6 @@ 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("--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="梯度累积步数")
@@ -174,11 +177,15 @@ if __name__ == "__main__":
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Distillation", help="wandb项目名")
     parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--device",
+                        type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
@@ -188,7 +195,7 @@ if __name__ == "__main__":
     ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
     

trainer/train_dpo.py

@@ -9,6 +9,10 @@ import argparse
 import time
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.nn.functional as F
 import torch.distributed as dist
 from contextlib import nullcontext
@@ -135,7 +139,6 @@ 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("--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="梯度累积步数")
@@ -153,11 +156,15 @@ if __name__ == "__main__":
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-DPO", help="wandb项目名")
     parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--device",
+                        type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
@@ -166,7 +173,7 @@ if __name__ == "__main__":
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
     

trainer/train_full_sft.py

@@ -9,6 +9,10 @@ import argparse
 import time
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.distributed as dist
 from contextlib import nullcontext
 from torch import optim, nn
@@ -88,7 +92,6 @@ 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("--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="梯度累积步数")
@@ -105,11 +108,15 @@ if __name__ == "__main__":
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT", help="wandb项目名")
     parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--device",
+                        type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
@@ -118,7 +125,7 @@ if __name__ == "__main__":
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
     

trainer/train_grpo.py

@@ -11,6 +11,10 @@ import re
 import gc
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.nn.functional as F
 import torch.distributed as dist
 from transformers import AutoTokenizer
@@ -210,7 +214,6 @@ 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("--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="梯度累积步数")
@@ -241,11 +244,15 @@ if __name__ == "__main__":
     parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998", help="SGLang服务器URL")
     parser.add_argument("--sglang_model_path", type=str, default="../model", help="SGLang tokenizer路径")
     parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_grpo", help="SGLang共享存储路径")
+    parser.add_argument("--device",
+                        type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
@@ -255,7 +262,7 @@ if __name__ == "__main__":
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
     

trainer/train_lora.py

@@ -9,6 +9,10 @@ import argparse
 import time
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.distributed as dist
 from contextlib import nullcontext
 from torch import optim, nn
@@ -82,7 +86,6 @@ 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("--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="梯度累积步数")
@@ -99,11 +102,15 @@ if __name__ == "__main__":
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-LoRA", help="wandb项目名")
     parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--device",
+                        type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
@@ -112,7 +119,7 @@ if __name__ == "__main__":
     ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
     

trainer/train_ppo.py

@@ -10,6 +10,10 @@ import math
 import re
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.distributed as dist
 import torch.nn.functional as F
 from transformers import AutoTokenizer
@@ -301,7 +305,6 @@ 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("--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="梯度累积步数")
@@ -336,11 +339,15 @@ if __name__ == "__main__":
     parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998", help="SGLang服务器URL")
     parser.add_argument("--sglang_model_path", type=str, default="../model", help="SGLang tokenizer路径")
     parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_ppo", help="SGLang共享存储路径")
+    parser.add_argument("--device",
+                        type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
@@ -349,7 +356,7 @@ if __name__ == "__main__":
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
     

trainer/train_pretrain.py

@@ -9,6 +9,10 @@ import argparse
 import time
 import warnings
 import torch
+try:
+    import torch_npu
+except:
+    torch_npu = None
 import torch.distributed as dist
 from contextlib import nullcontext
 from torch import optim, nn
@@ -87,7 +91,6 @@ 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("--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="梯度累积步数")
@@ -104,11 +107,14 @@ if __name__ == "__main__":
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名")
     parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--device", type=str,
+                        default="cuda:0" if torch.cuda.is_available() else ("npu:0" if torch_npu and torch_npu.npu.is_available() else "cpu"),
+                        help="训练设备")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    local_rank = init_distributed_mode("nccl" if torch.cuda.is_available() else "hccl")
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else f"npu:{local_rank}"
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
@@ -117,7 +123,7 @@ if __name__ == "__main__":
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    device_type = "cuda" if "cuda" in args.device else ("npu" if "npu" in args.device else "cpu")
     dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
     autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
     

trainer/trainer_utils.py

@@ -41,11 +41,11 @@ def get_lr(current_step, total_steps, lr):
     return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
 
 
-def init_distributed_mode():
+def init_distributed_mode(backend="nccl"):
     if int(os.environ.get("RANK", -1)) == -1:
         return 0  # 非DDP模式
 
-    dist.init_process_group(backend="nccl")
+    dist.init_process_group(backend=backend)
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     return local_rank