[update] minimind-3

.gitignore

@@ -1,4 +1,6 @@
-model/__pycache__
+__pycache__/
+*.pyc
+.DS_Store
 out
 website/
 docs-minimind/

dataset/lm_dataset.py

@@ -1,11 +1,15 @@
 from torch.utils.data import Dataset
 import torch
+import json
 import os
 import random
-from datasets import load_dataset
+from datasets import load_dataset, Features, Sequence, Value
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 def pre_processing_chat(conversations, add_system_ratio=0.2):
+    # tool use 数据完整保留不做处理
+    if any(conv.get('tools') for conv in conversations): return conversations
+
     SYSTEM_PROMPTS = [
         "你是一个知识丰富的AI，尽力为用户提供准确的信息。",
         "你是minimind，一个小巧但有用的语言模型。",
@@ -18,12 +22,14 @@ def pre_processing_chat(conversations, add_system_ratio=0.2):
         "You are a knowledgeable AI. Try your best to provide accurate information.",
         "You are minimind, a small but useful language model."
     ]
-    if conversations and conversations[0].get('role') != 'system':
+    # 概率性添加system
+    if conversations[0].get('role') != 'system':
         if random.random() < add_system_ratio:
             return [{'role': 'system', 'content': random.choice(SYSTEM_PROMPTS)}] + conversations
     return conversations
 
-def post_processing_chat(prompt_content, empty_think_ratio=0.05):
+def post_processing_chat(prompt_content, empty_think_ratio=0.2):
+    # 以80%概率移除空思考标签
     if '<think>\n\n</think>\n\n' in prompt_content and random.random() > empty_think_ratio:
         prompt_content = prompt_content.replace('<think>\n\n</think>\n\n', '')
     return prompt_content
@@ -54,7 +60,8 @@ class SFTDataset(Dataset):
         super().__init__()
         self.tokenizer = tokenizer
         self.max_length = max_length
-        self.samples = load_dataset('json', data_files=jsonl_path, split='train')
+        features = Features({'conversations': [{'role': Value('string'), 'content': Value('string'), 'reasoning_content': Value('string'), 'tools': Value('string'), 'tool_calls': Value('string')}]})
+        self.samples = load_dataset('json', data_files=jsonl_path, split='train', features=features)
         self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
         self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
 
@@ -62,8 +69,15 @@ class SFTDataset(Dataset):
         return len(self.samples)
 
     def create_chat_prompt(self, conversations):
-        messages = conversations.copy()
-        tools = conversations[0]["functions"] if (conversations and conversations[0]["role"] == "system" and conversations[0].get("functions")) else None
+        messages = []
+        tools = None
+        for message in conversations:
+            message = dict(message)
+            if message.get("role") == "system" and message.get("tools"):
+                tools = json.loads(message["tools"]) if isinstance(message["tools"], str) else message["tools"]
+            if message.get("tool_calls") and isinstance(message["tool_calls"], str):
+                message["tool_calls"] = json.loads(message["tool_calls"])
+            messages.append(message)
         return self.tokenizer.apply_chat_template(
             messages,
             tokenize=False,
@@ -179,10 +193,11 @@ class DPODataset(Dataset):
 
 
 class RLAIFDataset(Dataset):
-    def __init__(self, jsonl_path, tokenizer, max_length=1024):
+    def __init__(self, jsonl_path, tokenizer, max_length=1024, thinking_ratio=0.5):
         super().__init__()
         self.tokenizer = tokenizer
         self.max_length = max_length
+        self.thinking_ratio = thinking_ratio  # 按概率开启 thinking
         self.samples = load_dataset('json', data_files=jsonl_path, split='train')
         self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
         self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids
@@ -191,28 +206,51 @@ class RLAIFDataset(Dataset):
         return len(self.samples)
 
     def create_chat_prompt(self, conversations):
-        messages = []
-        answer = ''
-        for i, turn in enumerate(conversations):
-            role = 'user' if i % 2 == 0 else 'assistant'
-            messages.append({"role": role, "content": turn['content']})
-            answer = turn['content']
-        prompt = self.tokenizer.apply_chat_template(
-            messages[:-1],
+        conversations = pre_processing_chat(conversations)
+        use_thinking = random.random() < self.thinking_ratio
+        return self.tokenizer.apply_chat_template(
+            conversations[:-1],
             tokenize=False,
-            add_generation_prompt=True  # 这里需要True
+            open_thinking=use_thinking,
+            add_generation_prompt=True
         )
-        prompt = post_processing_chat(prompt)
-        return prompt, answer
-
     def __getitem__(self, index):
         sample = self.samples[index]
-        prompt, answer = self.create_chat_prompt(sample['conversations'])
+        prompt = self.create_chat_prompt(sample['conversations'])
 
         return {
             'prompt': prompt,
-            'answer': answer
+            'answer': ""
         }
 
+class AgentRLDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=1024):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.samples = []
+        with open(jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                self.samples.append(json.loads(line.strip()))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def parse_conversations(self, conversations):
+        messages = []
+        tools = None
+        for message in conversations:
+            message = dict(message)
+            if message.get("role") == "system" and message.get("tools"):
+                tools = json.loads(message["tools"]) if isinstance(message["tools"], str) else message["tools"]
+            messages.append(message)
+        return messages[:-1], tools
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        messages, tools = self.parse_conversations(sample['conversations'])
+        return {'messages': messages, 'tools': tools, 'gt': sample['gt']}
+
+
 if __name__ == "__main__":
     pass

eval_llm.py

@@ -23,7 +23,7 @@ def init_model(args):
         model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
         if args.lora_weight != 'None':
             apply_lora(model)
-            load_lora(model, f'./{args.save_dir}/lora/{args.lora_weight}_{args.hidden_size}.pth')
+            load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}.pth')
     else:
         model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
     get_model_params(model, model.config)
@@ -35,13 +35,14 @@ def main():
     parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
     parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, rlhf, reason, ppo_actor, grpo, spo）")
     parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称（None表示不使用，可选：lora_identity, lora_medical）")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度（512=Small-26M, 640=MoE-145M, 768=Base-104M）")
-    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量（Small/MoE=8, Base=16）")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, 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('--max_new_tokens', default=8192, type=int, help="最大生成长度（注意：并非模型实际长文本能力）")
     parser.add_argument('--temperature', default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
-    parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument('--top_p', default=0.95, type=float, help="nucleus采样阈值（0-1）")
+    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="运行设备")
@@ -65,23 +66,25 @@ def main():
     
     prompt_iter = prompts if input_mode == 0 else iter(lambda: input('💬: '), '')
     for prompt in prompt_iter:
-        setup_seed(2026) # or setup_seed(random.randint(0, 2048))
+        setup_seed(random.randint(0, 31415926))
+        setup_seed(42)
         if input_mode == 0: print(f'💬: {prompt}')
         conversation = conversation[-args.historys:] if args.historys else []
         conversation.append({"role": "user", "content": prompt})
-
-        templates = {"conversation": conversation, "tokenize": False, "add_generation_prompt": True}
-        if args.weight == 'reason': templates["enable_thinking"] = True # 仅Reason模型使用
-        inputs = tokenizer.apply_chat_template(**templates) if args.weight != 'pretrain' else (tokenizer.bos_token + prompt)
+        if 'pretrain' in args.weight:
+            inputs = tokenizer.bos_token + prompt
+        else:
+            inputs = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
+        
         inputs = tokenizer(inputs, return_tensors="pt", truncation=True).to(args.device)
 
-        print('🤖: ', end='')
+        print('🧠: ', end='')
         st = time.time()
         generated_ids = model.generate(
             inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
             max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
             pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
-            top_p=args.top_p, temperature=args.temperature, repetition_penalty=1.0
+            top_p=args.top_p, temperature=args.temperature, repetition_penalty=1
         )
         response = tokenizer.decode(generated_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
         conversation.append({"role": "assistant", "content": response})

model/model_lora.py

@@ -18,7 +18,7 @@ class LoRA(nn.Module):
         return self.B(self.A(x))
 
 
-def apply_lora(model, rank=8):
+def apply_lora(model, rank=16):
     for name, module in model.named_modules():
         if isinstance(module, nn.Linear) and module.weight.shape[0] == module.weight.shape[1]:
             lora = LoRA(module.weight.shape[0], module.weight.shape[1], rank=rank).to(model.device)
@@ -48,6 +48,18 @@ def save_lora(model, path):
     for name, module in raw_model.named_modules():
         if hasattr(module, 'lora'):
             clean_name = name[7:] if name.startswith("module.") else name
-            lora_state = {f'{clean_name}.lora.{k}': v for k, v in module.lora.state_dict().items()}
+            lora_state = {f'{clean_name}.lora.{k}': v.cpu().half() for k, v in module.lora.state_dict().items()}
             state_dict.update(lora_state)
     torch.save(state_dict, path)
+
+
+def merge_lora(model, lora_path, save_path):
+    load_lora(model, lora_path)
+    raw_model = getattr(model, '_orig_mod', model)
+    state_dict = {k: v.cpu().half() for k, v in raw_model.state_dict().items() if '.lora.' not in k}
+    for name, module in raw_model.named_modules():
+        if isinstance(module, nn.Linear) and '.lora.' not in name:
+            state_dict[f'{name}.weight'] = module.weight.data.clone().cpu().half()
+            if hasattr(module, 'lora'):
+                state_dict[f'{name}.weight'] += (module.lora.B.weight.data @ module.lora.A.weight.data).cpu().half()
+    torch.save(state_dict, save_path)

model/model_minimind.py

@@ -1,60 +1,33 @@
-# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
-#                                             MiniMind Config
-# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
-
-from transformers import PretrainedConfig
-
+import math, torch, torch.nn.functional as F
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
 
+# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
+#                                     MiniMind Config
+# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
 class MiniMindConfig(PretrainedConfig):
     model_type = "minimind"
-
-    def __init__(
-            self,
-            dropout: float = 0.0,
-            bos_token_id: int = 1,
-            eos_token_id: int = 2,
-            hidden_act: str = 'silu',
-            hidden_size: int = 512,
-            intermediate_size: int = None,
-            max_position_embeddings: int = 32768,
-            num_attention_heads: int = 8,
-            num_hidden_layers: int = 8,
-            num_key_value_heads: int = 2,
-            vocab_size: int = 6400,
-            rms_norm_eps: float = 1e-05,
-            rope_theta: int = 1000000.0,
-            inference_rope_scaling: bool = False,
-            flash_attn: bool = True,
-            ####################################################
-            # Here are the specific configurations of MOE
-            # When use_moe is false, the following is invalid
-            ####################################################
-            use_moe: bool = False,
-            num_experts_per_tok: int = 2,
-            n_routed_experts: int = 4,
-            n_shared_experts: int = 1,
-            scoring_func: str = 'softmax',
-            aux_loss_alpha: float = 0.01,
-            seq_aux: bool = True,
-            norm_topk_prob: bool = True,
-            **kwargs
-    ):
+    def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
         super().__init__(**kwargs)
-        self.dropout = dropout
-        self.bos_token_id = bos_token_id
-        self.eos_token_id = eos_token_id
-        self.hidden_act = hidden_act
         self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.max_position_embeddings = max_position_embeddings
-        self.num_attention_heads = num_attention_heads
         self.num_hidden_layers = num_hidden_layers
-        self.num_key_value_heads = num_key_value_heads
-        self.vocab_size = vocab_size
-        self.rms_norm_eps = rms_norm_eps
-        self.rope_theta = rope_theta
-        self.inference_rope_scaling = inference_rope_scaling
-        # 外推长度 = factor * original_max_position_embeddings = 32768
+        self.use_moe = use_moe
+        self.dropout = kwargs.get("dropout", 0.0)
+        self.vocab_size = kwargs.get("vocab_size", 6400)
+        self.bos_token_id = kwargs.get("bos_token_id", 1)
+        self.eos_token_id = kwargs.get("eos_token_id", 2)
+        self.flash_attn = kwargs.get("flash_attn", True)
+        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
+        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
+        self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
+        self.hidden_act = kwargs.get("hidden_act", 'silu')
+        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
+        self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
+        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
+        self.rope_theta = kwargs.get("rope_theta", 1e6)
+        self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
         self.rope_scaling = {
             "beta_fast": 32,
             "beta_slow": 1,
@@ -63,301 +36,147 @@ class MiniMindConfig(PretrainedConfig):
             "attention_factor": 1.0,
             "type": "yarn"
         } if self.inference_rope_scaling else None
-        self.flash_attn = flash_attn
-        ####################################################
-        # Here are the specific configurations of MOE
-        # When use_moe is false, the following is invalid
-        ####################################################
-        self.use_moe = use_moe
-        self.num_experts_per_tok = num_experts_per_tok  # 每个token选择的专家数量
-        self.n_routed_experts = n_routed_experts  # 总的专家数量
-        self.n_shared_experts = n_shared_experts  # 共享专家
-        self.scoring_func = scoring_func  # 评分函数，默认为'softmax'
-        self.aux_loss_alpha = aux_loss_alpha  # 辅助损失的alpha参数
-        self.seq_aux = seq_aux  # 是否在序列级别上计算辅助损失
-        self.norm_topk_prob = norm_topk_prob  # 是否标准化top-k概率
-
-
-# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
-#                                             MiniMind Model
-# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
-
-import math
-import torch
-import torch.nn.init as init
-import torch.nn.functional as F
-from torch import nn
-from transformers.activations import ACT2FN
-from typing import Optional, Tuple, List, Union
-from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
+        ### MoE specific configs (ignored if use_moe = False)
+        self.num_experts = kwargs.get("num_experts", 4)
+        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
+        self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
+        self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
+        self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)
 
+# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
+#                                     MiniMind Model
+# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
 class RMSNorm(torch.nn.Module):
     def __init__(self, dim: int, eps: float = 1e-5):
         super().__init__()
         self.eps = eps
         self.weight = nn.Parameter(torch.ones(dim))
 
-    def _norm(self, x):
+    def norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
 
     def forward(self, x):
-        return self.weight * self._norm(x.float()).type_as(x)
+        return (self.weight * self.norm(x.float())).type_as(x)
 
-
-def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6,
-                         rope_scaling: Optional[dict] = None):
+def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
     freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
-    if rope_scaling is not None:
+    if rope_scaling is not None: # YaRN: f'(i) = f(i)((1-γ) + γ/s), where γ∈[0,1] is linear ramp
         orig_max, factor, beta_fast, beta_slow, attn_factor = (
             rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
             rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
         )
         if end / orig_max > 1.0:
-            # YaRN: f'(i) = f(i)((1-γ) + γ/s), where γ∈[0,1] is linear ramp
             inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
             low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
             ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
             freqs = freqs * (1 - ramp + ramp / factor)
-
     t = torch.arange(end, device=freqs.device)
     freqs = torch.outer(t, freqs).float()
     freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
     freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
     return freqs_cos, freqs_sin
 
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    def rotate_half(x):
-        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
-
-    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
-    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
+    q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
+    k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
     return q_embed, k_embed
 
-
 def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
     bs, slen, num_key_value_heads, head_dim = x.shape
-    if n_rep == 1:
-        return x
-    return (
-        x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
-    )
-
+    if n_rep == 1: return x
+    return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim))
 
 class Attention(nn.Module):
-    def __init__(self, args: MiniMindConfig):
+    def __init__(self, config: MiniMindConfig):
         super().__init__()
-        self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads
-        assert args.num_attention_heads % self.num_key_value_heads == 0
-        self.n_local_heads = args.num_attention_heads
+        self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
+        self.n_local_heads = config.num_attention_heads
         self.n_local_kv_heads = self.num_key_value_heads
         self.n_rep = self.n_local_heads // self.n_local_kv_heads
-        self.head_dim = args.hidden_size // args.num_attention_heads
-        self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
-        self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
-        self.attn_dropout = nn.Dropout(args.dropout)
-        self.resid_dropout = nn.Dropout(args.dropout)
-        self.dropout = args.dropout
-        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
-        # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+        self.head_dim = config.head_dim
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.dropout = config.dropout
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn
 
-    def forward(self,
-                x: torch.Tensor,
-                position_embeddings: Tuple[torch.Tensor, torch.Tensor],  # 修改为接收cos和sin
-                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-                use_cache=False,
-                attention_mask: Optional[torch.Tensor] = None):
+    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
         bsz, seq_len, _ = x.shape
         xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
         xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
         xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
         xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
-
+        xq, xk = self.q_norm(xq), self.k_norm(xk)
         cos, sin = position_embeddings
         xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
-
-        # kv_cache实现
         if past_key_value is not None:
             xk = torch.cat([past_key_value[0], xk], dim=1)
             xv = torch.cat([past_key_value[1], xv], dim=1)
         past_kv = (xk, xv) if use_cache else None
-
-        xq, xk, xv = (
-            xq.transpose(1, 2),
-            repeat_kv(xk, self.n_rep).transpose(1, 2),
-            repeat_kv(xv, self.n_rep).transpose(1, 2)
-        )
-
+        xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
         if self.flash and (seq_len > 1) and (past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
             output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
         else:
             scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
-            scores[:, :, :, -seq_len:] += torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=scores.device), diagonal=1)
-
-            if attention_mask is not None:
-                extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-                extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
-                scores = scores + extended_attention_mask
-
-            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
-            scores = self.attn_dropout(scores)
-            output = scores @ xv
-
+            scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
+            if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
+            output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
         output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
         output = self.resid_dropout(self.o_proj(output))
         return output, past_kv
 
-
 class FeedForward(nn.Module):
-    def __init__(self, config: MiniMindConfig):
+    def __init__(self, config: MiniMindConfig, intermediate_size: int = None):
         super().__init__()
-        if config.intermediate_size is None:
-            intermediate_size = int(config.hidden_size * 8 / 3)
-            config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)
-        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
-        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.dropout = nn.Dropout(config.dropout)
+        intermediate_size = intermediate_size or config.intermediate_size
+        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
         self.act_fn = ACT2FN[config.hidden_act]
 
     def forward(self, x):
-        return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
-
-
-class MoEGate(nn.Module):
-    def __init__(self, config: MiniMindConfig):
-        super().__init__()
-        self.config = config
-        self.top_k = config.num_experts_per_tok
-        self.n_routed_experts = config.n_routed_experts
-
-        self.scoring_func = config.scoring_func
-        self.alpha = config.aux_loss_alpha
-        self.seq_aux = config.seq_aux
-
-        self.norm_topk_prob = config.norm_topk_prob
-        self.gating_dim = config.hidden_size
-        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-
-    def forward(self, hidden_states):
-        bsz, seq_len, h = hidden_states.shape
-        hidden_states = hidden_states.view(-1, h)
-        logits = F.linear(hidden_states, self.weight, None)
-        if self.scoring_func == 'softmax':
-            scores = logits.softmax(dim=-1)
-        else:
-            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
-
-        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
-
-        if self.top_k > 1 and self.norm_topk_prob:
-            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
-            topk_weight = topk_weight / denominator
-
-        if self.training and self.alpha > 0.0:
-            scores_for_aux = scores
-            aux_topk = self.top_k
-            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
-            if self.seq_aux:
-                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
-                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
-                ce.scatter_add_(1, topk_idx_for_aux_loss,
-                                torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
-                    seq_len * aux_topk / self.n_routed_experts)
-                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
-            else:
-                mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
-                ce = mask_ce.float().mean(0)
-                Pi = scores_for_aux.mean(0)
-                fi = ce * self.n_routed_experts
-                aux_loss = (Pi * fi).sum() * self.alpha
-        else:
-            aux_loss = scores.new_zeros(1).squeeze()
-        return topk_idx, topk_weight, aux_loss
-
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
 
 class MOEFeedForward(nn.Module):
     def __init__(self, config: MiniMindConfig):
         super().__init__()
         self.config = config
-        self.experts = nn.ModuleList([
-            FeedForward(config)
-            for _ in range(config.n_routed_experts)
-        ])
-        self.gate = MoEGate(config)
-        if config.n_shared_experts > 0:
-            self.shared_experts = nn.ModuleList([
-                FeedForward(config)
-                for _ in range(config.n_shared_experts)
-            ])
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+        self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)])
+        self.act_fn = ACT2FN[config.hidden_act]
 
     def forward(self, x):
-        identity = x
-        orig_shape = x.shape
-        bsz, seq_len, _ = x.shape
-        # 使用门控机制选择专家
-        topk_idx, topk_weight, aux_loss = self.gate(x)
-        x = x.view(-1, x.shape[-1])
-        flat_topk_idx = topk_idx.view(-1)
-        if self.training:
-            x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
-            y = torch.empty_like(x, dtype=x.dtype)
-            for i, expert in enumerate(self.experts):
-                expert_out = expert(x[flat_topk_idx == i])
-                if expert_out.shape[0] > 0: y[flat_topk_idx == i] = expert_out.to(y.dtype)
-                else: y[flat_topk_idx == i] = expert_out.to(y.dtype) + 0 * sum(p.sum() for p in expert.parameters())
-            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
-            y = y.view(*orig_shape)
+        batch_size, seq_len, hidden_dim = x.shape
+        x_flat = x.view(-1, hidden_dim)
+        scores = F.softmax(self.gate(x_flat), dim=-1)
+        topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
+        if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
+        y = torch.zeros_like(x_flat)
+        for i, expert in enumerate(self.experts):
+            mask = (topk_idx == i)
+            if mask.any():
+                token_idx = mask.any(dim=-1).nonzero().flatten()
+                weight = topk_weight[mask].view(-1, 1)
+                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
+            elif self.training:
+                y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
+        if self.training and self.config.router_aux_loss_coef > 0:
+            load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
+            self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
         else:
-            y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
-        if self.config.n_shared_experts > 0:
-            for expert in self.shared_experts:
-                y = y + expert(identity)
-        self.aux_loss = aux_loss
-        return y
-
-    @torch.no_grad()
-    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
-        expert_cache = torch.zeros_like(x)
-        idxs = flat_expert_indices.argsort()
-        tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
-        token_idxs = idxs // self.config.num_experts_per_tok
-        # 当tokens_per_expert = [6, 15, 20, 26]，tokens_per_expert.shape[0]即为专家数量（此时为4）
-        # 且token_idxs = [3, 7, 19, 21, 24, 25,  4,  5,  6, 10, 11, 12...] 时
-        # 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token（每个token有可能被多个专家处理，这取决于num_experts_per_tok）
-        # 接下来9个位置token_idxs[6:15] -> [4,  5,  6, 10, 11, 12...]属于专家1处理的token...依此类推
-        for i, end_idx in enumerate(tokens_per_expert):
-            start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
-            if start_idx == end_idx:
-                continue
-            expert = self.experts[i]
-            exp_token_idx = token_idxs[start_idx:end_idx]
-            expert_tokens = x[exp_token_idx]
-            expert_out = expert(expert_tokens).to(expert_cache.dtype)
-            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
-            expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
-
-        return expert_cache
-
+            self.aux_loss = scores.new_zeros(1).squeeze()
+        return y.view(batch_size, seq_len, hidden_dim)
 
 class MiniMindBlock(nn.Module):
     def __init__(self, layer_id: int, config: MiniMindConfig):
         super().__init__()
-        self.num_attention_heads = config.num_attention_heads
-        self.hidden_size = config.hidden_size
-        self.head_dim = config.hidden_size // config.num_attention_heads
         self.self_attn = Attention(config)
-
-        self.layer_id = layer_id
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
@@ -372,7 +191,6 @@ class MiniMindBlock(nn.Module):
         hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
         return hidden_states, present_key_value
 
-
 class MiniMindModel(nn.Module):
     def __init__(self, config: MiniMindConfig):
         super().__init__()
@@ -382,33 +200,19 @@ class MiniMindModel(nn.Module):
         self.dropout = nn.Dropout(config.dropout)
         self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads,
-                                                    end=config.max_position_embeddings, rope_base=config.rope_theta,
-                                                    rope_scaling=config.rope_scaling)
+        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
         self.register_buffer("freqs_cos", freqs_cos, persistent=False)
         self.register_buffer("freqs_sin", freqs_sin, persistent=False)
 
-    def forward(self,
-                input_ids: Optional[torch.Tensor] = None,
-                attention_mask: Optional[torch.Tensor] = None,
-                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
-                use_cache: bool = False,
-                **kwargs):
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
         batch_size, seq_length = input_ids.shape
         if hasattr(past_key_values, 'layers'): past_key_values = None
         past_key_values = past_key_values or [None] * len(self.layers)
         start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
-
         hidden_states = self.dropout(self.embed_tokens(input_ids))
-
-        position_embeddings = (
-            self.freqs_cos[start_pos:start_pos + seq_length],
-            self.freqs_sin[start_pos:start_pos + seq_length]
-        )
-
+        position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
         presents = []
-        for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
+        for layer, past_key_value in zip(self.layers, past_key_values):
             hidden_states, present = layer(
                 hidden_states,
                 position_embeddings,
@@ -417,47 +221,59 @@ class MiniMindModel(nn.Module):
                 attention_mask=attention_mask
             )
             presents.append(present)
-
         hidden_states = self.norm(hidden_states)
-
         aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
         return hidden_states, presents, aux_loss
 
-
 class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = MiniMindConfig
-
     def __init__(self, config: MiniMindConfig = None):
         self.config = config or MiniMindConfig()
         super().__init__(self.config)
         self.model = MiniMindModel(self.config)
         self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
         self.model.embed_tokens.weight = self.lm_head.weight
-
-    def forward(self,
-                input_ids: Optional[torch.Tensor] = None,
-                attention_mask: Optional[torch.Tensor] = None,
-                labels: Optional[torch.Tensor] = None,
-                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
-                use_cache: bool = False,
-                logits_to_keep: Union[int, torch.Tensor] = 0,
-                **args):
-        hidden_states, past_key_values, aux_loss = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            past_key_values=past_key_values,
-            use_cache=use_cache,
-            **args
-        )
+    
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
+        hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
         logits = self.lm_head(hidden_states[:, slice_indices, :])
-
         loss = None
         if labels is not None:
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
-
-        output = CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
-        output.aux_loss = aux_loss
-        return output
+            x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
+            loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
+        return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
+    
+    # https://github.com/jingyaogong/minimind/discussions/611
+    @torch.inference_mode()
+    def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
+        input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
+        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
+        past_key_values = kwargs.pop("past_key_values", None)
+        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
+        if streamer: streamer.put(input_ids.cpu())
+        for _ in range(max_new_tokens):
+            past_len = past_key_values[0][0].shape[1] if past_key_values else 0
+            outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
+            attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
+            logits = outputs.logits[:, -1, :] / temperature
+            if repetition_penalty != 1.0:
+                for i in range(input_ids.shape[0]): logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+            if top_k > 0: 
+                logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
+                mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
+                logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
+            next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
+            if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            past_key_values = outputs.past_key_values if use_cache else None
+            if streamer: streamer.put(next_token.cpu())
+            if eos_token_id is not None:
+                finished |= next_token.squeeze(-1).eq(eos_token_id)
+                if finished.all(): break
+        if streamer: streamer.end()
+        if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
+        return input_ids

model/tokenizer_config.json

@@ -1,43 +1,335 @@
 {
-  "add_bos_token": false,
-  "add_eos_token": false,
-  "add_prefix_space": false,
-  "added_tokens_decoder": {
-    "0": {
-      "content": "<|endoftext|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
+    "add_bos_token": false,
+    "add_eos_token": false,
+    "add_prefix_space": false,
+    "added_tokens_decoder": {
+        "0": {
+            "content": "<|endoftext|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "1": {
+            "content": "<|im_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "2": {
+            "content": "<|im_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "3": {
+            "content": "<|object_ref_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "4": {
+            "content": "<|object_ref_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "5": {
+            "content": "<|box_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "6": {
+            "content": "<|box_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "7": {
+            "content": "<|quad_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "8": {
+            "content": "<|quad_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "9": {
+            "content": "<|vision_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "10": {
+            "content": "<|vision_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "11": {
+            "content": "<|vision_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "12": {
+            "content": "<|image_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "13": {
+            "content": "<|video_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "14": {
+            "content": "<|audio_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "15": {
+            "content": "<|audio_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "16": {
+            "content": "<|audio_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "17": {
+            "content": "<tts_pad>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "18": {
+            "content": "<tts_text_bos>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "19": {
+            "content": "<tts_text_eod>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "20": {
+            "content": "<tts_text_bos_single>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "21": {
+            "content": "<tool_call>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "22": {
+            "content": "</tool_call>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "23": {
+            "content": "<tool_response>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "24": {
+            "content": "</tool_response>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "25": {
+            "content": "<think>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "26": {
+            "content": "</think>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "27": {
+            "content": "<|buffer1|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "28": {
+            "content": "<|buffer2|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "29": {
+            "content": "<|buffer3|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "30": {
+            "content": "<|buffer4|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "31": {
+            "content": "<|buffer5|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "32": {
+            "content": "<|buffer6|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "33": {
+            "content": "<|buffer7|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "34": {
+            "content": "<|buffer8|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "35": {
+            "content": "<|buffer9|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        }
     },
-    "1": {
-      "content": "<|im_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "2": {
-      "content": "<|im_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
-  "additional_special_tokens": [],
-  "bos_token": "<|im_start|>",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|im_end|>",
-  "legacy": true,
-  "model_max_length": 32768,
-  "pad_token": "<|endoftext|>",
-  "sp_model_kwargs": {},
-  "spaces_between_special_tokens": false,
-  "tokenizer_class": "PreTrainedTokenizerFast",
-  "unk_token": "<|endoftext|>",
-  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0]['role'] == 'system' -%}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else -%}\n        {{- '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n   {{- '<|im_start|>' + message.role + '\\n' + content }}\n  {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}"
+    "additional_special_tokens": [
+        "<|im_start|>",
+        "<|im_end|>",
+        "<|object_ref_start|>",
+        "<|object_ref_end|>",
+        "<|box_start|>",
+        "<|box_end|>",
+        "<|quad_start|>",
+        "<|quad_end|>",
+        "<|vision_start|>",
+        "<|vision_end|>",
+        "<|vision_pad|>",
+        "<|image_pad|>",
+        "<|video_pad|>",
+        "<|audio_start|>",
+        "<|audio_end|>",
+        "<|audio_pad|>",
+        "<tts_pad>",
+        "<tts_text_bos>",
+        "<tts_text_eod>",
+        "<tts_text_bos_single>"
+    ],
+    "bos_token": "<|im_start|>",
+    "clean_up_tokenization_spaces": false,
+    "eos_token": "<|im_end|>",
+    "legacy": true,
+    "model_max_length": 131072,
+    "pad_token": "<|endoftext|>",
+    "sp_model_kwargs": {},
+    "spaces_between_special_tokens": false,
+    "unk_token": "<|endoftext|>",
+    "image_token": "<|image_pad|>",
+    "audio_token": "<|audio_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>",
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if true %}\n            {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if open_thinking is defined and open_thinking is true %}\n        {{- '<think>\\n' }}\n    {%- else %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+    "tokenizer_class": "PreTrainedTokenizerFast"
 }

requirements.txt

@@ -5,12 +5,12 @@ Flask_Cors==4.0.0
 jieba==0.42.1
 jsonlines==4.0.0
 marshmallow==3.22.0
-matplotlib==3.10.0
+# matplotlib==3.10.0
 ngrok==1.4.0
 nltk==3.8
 numpy==1.26.4
 openai==1.59.6
-peft==0.7.1
+# peft==0.7.1
 psutil==5.9.8
 pydantic==2.11.5
 rich==13.7.1
@@ -18,7 +18,7 @@ scikit_learn==1.5.1
 sentence_transformers==2.3.1
 simhash==2.1.2
 tiktoken==0.10.0
-transformers==4.57.1
+transformers==4.57.6
 jinja2==3.1.2
 jsonlines==4.0.0
 trl==0.13.0
@@ -26,6 +26,7 @@ ujson==5.1.0
 wandb==0.18.3
 streamlit==1.50.0
 einops==0.8.1
-swanlab==0.6.8
-torch==2.6.0
-torchvision==0.21.0
+swanlab==0.7.11
+modelscope==1.30.0
+# torch==2.6.0
+# torchvision==0.21.0

scripts/chat_api.py

@@ -1,8 +1,8 @@
 from openai import OpenAI
 
 client = OpenAI(
-    api_key="ollama",
-    base_url="http://127.0.0.1:8998/v1"
+    api_key="sk-123",
+    base_url="http://localhost:11434/v1"
 )
 stream = True
 conversation_history_origin = []
@@ -12,22 +12,29 @@ while True:
     query = input('[Q]: ')
     conversation_history.append({"role": "user", "content": query})
     response = client.chat.completions.create(
-        model="minimind",
+        model="minimind-local:latest",
         messages=conversation_history[-(history_messages_num or 1):],
         stream=stream,
-        temperature=0.7,
+        temperature=0.8,
         max_tokens=2048,
-        top_p=0.9
+        top_p=0.8,
+        extra_body={"chat_template_kwargs": {"open_thinking": True}, "reasoning_effort": "medium"} # 思考开关
     )
     if not stream:
         assistant_res = response.choices[0].message.content
         print('[A]: ', assistant_res)
     else:
-        print('[A]: ', end='')
+        print('[A]: ', end='', flush=True)
         assistant_res = ''
         for chunk in response:
-            print(chunk.choices[0].delta.content or "", end="")
-            assistant_res += chunk.choices[0].delta.content or ""
+            delta = chunk.choices[0].delta
+            r = getattr(delta, 'reasoning_content', None) or ""
+            c = delta.content or ""
+            if r:
+                print(f'\033[90m{r}\033[0m', end="", flush=True)
+            if c:
+                print(c, end="", flush=True)
+            assistant_res += c
 
     conversation_history.append({"role": "assistant", "content": assistant_res})
-    print('\n\n')
+    print('\n\n')

scripts/convert_model.py

@@ -5,14 +5,14 @@ import json
 __package__ = "scripts"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 import torch
+import transformers
 import warnings
-from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, Qwen3Config, Qwen3ForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_lora import apply_lora, merge_lora
 
 warnings.filterwarnings('ignore', category=UserWarning)
 
-
-# MoE模型需使用此函数转换
 def convert_torch2transformers_minimind(torch_path, transformers_path, dtype=torch.float16):
     MiniMindConfig.register_for_auto_class()
     MiniMindForCausalLM.register_for_auto_class("AutoModelForCausalLM")
@@ -26,52 +26,118 @@ def convert_torch2transformers_minimind(torch_path, transformers_path, dtype=tor
     lm_model.save_pretrained(transformers_path, safe_serialization=False)
     tokenizer = AutoTokenizer.from_pretrained('../model/')
     tokenizer.save_pretrained(transformers_path)
-    # 兼容transformers-5.0的写法
-    config_path = os.path.join(transformers_path, "tokenizer_config.json")
-    json.dump({**json.load(open(config_path, 'r', encoding='utf-8')), "tokenizer_class": "PreTrainedTokenizerFast", "extra_special_tokens": {}}, open(config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+    # ======= transformers-5.0的兼容低版本写法 =======
+    if int(transformers.__version__.split('.')[0]) >= 5:
+        tokenizer_config_path, config_path = os.path.join(transformers_path, "tokenizer_config.json"), os.path.join(transformers_path, "config.json")
+        json.dump({**json.load(open(tokenizer_config_path, 'r', encoding='utf-8')), "tokenizer_class": "PreTrainedTokenizerFast", "extra_special_tokens": {}}, open(tokenizer_config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+        config = json.load(open(config_path, 'r', encoding='utf-8'))
+        config['rope_theta'] = lm_config.rope_theta; config['rope_scaling'] = None; del config['rope_parameters']
+        json.dump(config, open(config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
     print(f"模型已保存为 Transformers-MiniMind 格式: {transformers_path}")
 
 
-# LlamaForCausalLM结构兼容第三方生态
-def convert_torch2transformers_llama(torch_path, transformers_path, dtype=torch.float16):
+# QwenForCausalLM/LlamaForCausalLM结构兼容生态
+def convert_torch2transformers(torch_path, transformers_path, dtype=torch.float16):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     state_dict = torch.load(torch_path, map_location=device)
-    llama_config = LlamaConfig(
-        vocab_size=lm_config.vocab_size,
-        hidden_size=lm_config.hidden_size,
-        intermediate_size=64 * ((int(lm_config.hidden_size * 8 / 3) + 64 - 1) // 64),
-        num_hidden_layers=lm_config.num_hidden_layers,
-        num_attention_heads=lm_config.num_attention_heads,
-        num_key_value_heads=lm_config.num_key_value_heads,
-        max_position_embeddings=lm_config.max_position_embeddings,
-        rms_norm_eps=lm_config.rms_norm_eps,
-        rope_theta=lm_config.rope_theta,
-        tie_word_embeddings=True
-    )
-    llama_model = LlamaForCausalLM(llama_config)
-    llama_model.load_state_dict(state_dict, strict=False)
-    llama_model = llama_model.to(dtype)  # 转换模型权重精度
-    llama_model.save_pretrained(transformers_path)
-    model_params = sum(p.numel() for p in llama_model.parameters() if p.requires_grad)
+    common_config = {
+        "vocab_size": lm_config.vocab_size,
+        "hidden_size": lm_config.hidden_size,
+        "intermediate_size": lm_config.intermediate_size,
+        "num_hidden_layers": lm_config.num_hidden_layers,
+        "num_attention_heads": lm_config.num_attention_heads,
+        "num_key_value_heads": lm_config.num_key_value_heads,
+        "head_dim": lm_config.hidden_size // lm_config.num_attention_heads,
+        "max_position_embeddings": lm_config.max_position_embeddings,
+        "rms_norm_eps": lm_config.rms_norm_eps,
+        "rope_theta": lm_config.rope_theta,
+        "tie_word_embeddings": True
+    }
+    if not lm_config.use_moe:
+        qwen_config = Qwen3Config(
+            **common_config, 
+            use_sliding_window=False, 
+            sliding_window=None
+        )
+        qwen_model = Qwen3ForCausalLM(qwen_config)
+    else:
+        qwen_config = Qwen3MoeConfig(
+            **common_config,
+            num_experts=lm_config.num_experts,
+            num_experts_per_tok=lm_config.num_experts_per_tok,
+            moe_intermediate_size=lm_config.moe_intermediate_size,
+            norm_topk_prob=lm_config.norm_topk_prob
+        )
+        qwen_model = Qwen3MoeForCausalLM(qwen_config)
+        # ======= transformers-5.0的兼容低版本写法 =======
+        if int(transformers.__version__.split('.')[0]) >= 5:
+            new_sd = {k: v for k, v in state_dict.items() if 'experts.' not in k or 'gate.weight' in k}
+            for l in range(lm_config.num_hidden_layers):
+                p = f'model.layers.{l}.mlp.experts'
+                new_sd[f'{p}.gate_up_proj'] = torch.cat([torch.stack([state_dict[f'{p}.{e}.gate_proj.weight'] for e in range(lm_config.num_experts)]), torch.stack([state_dict[f'{p}.{e}.up_proj.weight'] for e in range(lm_config.num_experts)])], dim=1)
+                new_sd[f'{p}.down_proj'] = torch.stack([state_dict[f'{p}.{e}.down_proj.weight'] for e in range(lm_config.num_experts)])
+            state_dict = new_sd
+
+    qwen_model.load_state_dict(state_dict, strict=True)
+    qwen_model = qwen_model.to(dtype)  # 转换模型权重精度
+    qwen_model.save_pretrained(transformers_path)
+    model_params = sum(p.numel() for p in qwen_model.parameters() if p.requires_grad)
     print(f'模型参数: {model_params / 1e6} 百万 = {model_params / 1e9} B (Billion)')
     tokenizer = AutoTokenizer.from_pretrained('../model/')
     tokenizer.save_pretrained(transformers_path)
-    # 兼容transformers-5.0的写法
-    config_path = os.path.join(transformers_path, "tokenizer_config.json")
-    json.dump({**json.load(open(config_path, 'r', encoding='utf-8')), "tokenizer_class": "PreTrainedTokenizerFast", "extra_special_tokens": {}}, open(config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
-    print(f"模型已保存为 Transformers-Llama 格式: {transformers_path}")
+
+    # ======= transformers-5.0的兼容低版本写法 =======
+    if int(transformers.__version__.split('.')[0]) >= 5:
+        tokenizer_config_path, config_path = os.path.join(transformers_path, "tokenizer_config.json"), os.path.join(transformers_path, "config.json")
+        json.dump({**json.load(open(tokenizer_config_path, 'r', encoding='utf-8')), "tokenizer_class": "PreTrainedTokenizerFast", "extra_special_tokens": {}}, open(tokenizer_config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+        config = json.load(open(config_path, 'r', encoding='utf-8'))
+        config['rope_theta'] = lm_config.rope_theta; config['rope_scaling'] = None; del config['rope_parameters']
+        json.dump(config, open(config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+    print(f"模型已保存为 Transformers 格式: {transformers_path}")
 
 
 def convert_transformers2torch(transformers_path, torch_path):
     model = AutoModelForCausalLM.from_pretrained(transformers_path, trust_remote_code=True)
     torch.save({k: v.cpu().half() for k, v in model.state_dict().items()}, torch_path)
-    print(f"模型已保存为 PyTorch 格式 (half精度): {torch_path}")
+    print(f"模型已保存为 PyTorch 格式: {torch_path}")
+
+
+def convert_merge_base_lora(base_torch_path, lora_path, merged_torch_path):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    lm_model = MiniMindForCausalLM(lm_config).to(device)
+    state_dict = torch.load(base_torch_path, map_location=device)
+    lm_model.load_state_dict(state_dict, strict=False)
+    apply_lora(lm_model)
+    merge_lora(lm_model, lora_path, merged_torch_path)
+    print(f"LoRA 已合并并保存为基模结构 PyTorch 格式: {merged_torch_path}")
+
+
+def convert_jinja_to_json(jinja_path):
+    with open(jinja_path, 'r') as f: template = f.read()
+    escaped = json.dumps(template)
+    print(f'"chat_template": {escaped}')
+
+
+def convert_json_to_jinja(json_file_path, output_path):
+    with open(json_file_path, 'r') as f: config = json.load(f)
+    template = config['chat_template']
+    with open(output_path, 'w') as f: f.write(template)
+    print(f"模板已保存为 jinja 文件: {output_path}")
 
 
 if __name__ == '__main__':
-    lm_config = MiniMindConfig(hidden_size=512, num_hidden_layers=8, max_seq_len=8192, use_moe=False)
+    lm_config = MiniMindConfig(hidden_size=768, num_hidden_layers=8, max_seq_len=8192, use_moe=True)
+    # convert torch to transformers
     torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
-    transformers_path = '../MiniMind2-Small'
-    convert_torch2transformers_llama(torch_path, transformers_path)
-    # # convert transformers to torch model
+    transformers_path = '../minimind-3-moe'
+    convert_torch2transformers(torch_path, transformers_path)
+
+    # # merge lora
+    # base_torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # lora_path = f"../out/lora_identity_{lm_config.hidden_size}.pth"
+    # merged_torch_path = f"../out/merge_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # convert_merge_base_lora(base_torch_path, lora_path, merged_torch_path)
+
     # convert_transformers2torch(transformers_path, torch_path)
+    # convert_json_to_jinja('../model/tokenizer_config.json', '../model/chat_template.jinja')
+    # convert_jinja_to_json('../model/chat_template.jinja')

scripts/eval_toolcall.py

@@ -0,0 +1,240 @@
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import re
+import json
+import time
+import random
+import argparse
+import warnings
+import torch
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from openai import OpenAI
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from trainer.trainer_utils import setup_seed, get_model_params
+warnings.filterwarnings('ignore')
+
+TOOLS = [
+    {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式的结果，支持加减乘除、幂运算、开方等", "parameters": {"type": "object", "properties": {"expression": {"type": "string", "description": "数学表达式，如123+456、2**10、sqrt(144)"}}, "required": ["expression"]}}},
+    {"type": "function", "function": {"name": "get_current_time", "description": "获取当前日期和时间，支持指定时区", "parameters": {"type": "object", "properties": {"timezone": {"type": "string", "description": "时区名称，如Asia/Shanghai、America/New_York", "default": "Asia/Shanghai"}}, "required": []}}},
+    {"type": "function", "function": {"name": "random_number", "description": "生成指定范围内的随机数", "parameters": {"type": "object", "properties": {"min": {"type": "integer", "description": "最小值", "default": 0}, "max": {"type": "integer", "description": "最大值", "default": 100}}, "required": []}}},
+    {"type": "function", "function": {"name": "text_length", "description": "计算文本的字符数和单词数", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "要统计的文本"}}, "required": ["text"]}}},
+    {"type": "function", "function": {"name": "unit_converter", "description": "进行单位换算，支持长度、重量、温度等", "parameters": {"type": "object", "properties": {"value": {"type": "number", "description": "要转换的数值"}, "from_unit": {"type": "string", "description": "源单位，如km、miles、kg、pounds、celsius、fahrenheit"}, "to_unit": {"type": "string", "description": "目标单位"}}, "required": ["value", "from_unit", "to_unit"]}}},
+    {"type": "function", "function": {"name": "get_current_weather", "description": "获取指定城市的当前天气信息，包括温度、湿度和天气状况", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "城市名称，如北京、上海、New York"}, "unit": {"type": "string", "description": "温度单位，celsius或fahrenheit", "enum": ["celsius", "fahrenheit"], "default": "celsius"}}, "required": ["location"]}}},
+    {"type": "function", "function": {"name": "get_exchange_rate", "description": "查询两种货币之间的实时汇率", "parameters": {"type": "object", "properties": {"from_currency": {"type": "string", "description": "源货币代码，如USD、CNY、EUR"}, "to_currency": {"type": "string", "description": "目标货币代码，如USD、CNY、EUR"}}, "required": ["from_currency", "to_currency"]}}},
+    {"type": "function", "function": {"name": "translate_text", "description": "将文本翻译成目标语言", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "要翻译的文本"}, "target_language": {"type": "string", "description": "目标语言，如english、chinese、japanese、french"}}, "required": ["text", "target_language"]}}},
+]
+
+MOCK_RESULTS = {
+    "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("²", "**2").replace("³", "**3").replace("（", "(").replace("）", ")")))},
+    "get_current_time": lambda args: {"datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "timezone": args.get("timezone", "Asia/Shanghai")},
+    "random_number": lambda args: {"result": random.randint(int(args.get("min", 0)), int(args.get("max", 100)))},
+    "text_length": lambda args: {"characters": len(args.get("text", "")), "words": len(args.get("text", "").split())},
+    "unit_converter": lambda args: {"result": round(float(args.get("value", 0)) * 0.621371, 2), "from": f"{args.get('value', 0)} {args.get('from_unit', '')}", "to": args.get("to_unit", "")},
+    "get_current_weather": lambda args: {"city": args.get("location"), "temperature": "22°C", "humidity": "65%", "condition": "晴"},
+    "get_exchange_rate": lambda args: {"from": args.get("from_currency", ""), "to": args.get("to_currency", ""), "rate": 7.15},
+    "translate_text": lambda args: {"translated": "hello world"},
+}
+
+TOOL_MAP = {t["function"]["name"]: t for t in TOOLS}
+
+def get_tools(names):
+    return [TOOL_MAP[n] for n in names]
+
+TEST_CASES = [
+    {"prompt": "帮我算一下 256 乘以 37 等于多少", "tools": ["calculate_math", "get_current_time"]},
+    {"prompt": "现在几点了？", "tools": ["get_current_time", "random_number"]},
+    {"prompt": "帮我把100公里换算成英里", "tools": ["unit_converter", "calculate_math"]},
+    {"prompt": "帮我生成一个1到1000的随机数，然后计算它的平方", "tools": ["random_number", "calculate_math", "text_length"]},
+    {"prompt": "北京今天天气怎么样？", "tools": ["get_current_weather", "get_current_time"]},
+    {"prompt": "查一下美元兑人民币汇率", "tools": ["get_exchange_rate", "get_current_time"]},
+    {"prompt": "把'你好世界'翻译成英文", "tools": ["translate_text", "text_length"]},
+    {"prompt": "What is the weather in Tokyo? Also convert 30 celsius to fahrenheit.", "tools": ["get_current_weather", "unit_converter", "get_current_time"]},
+]
+
+
+def init_model(args):
+    tokenizer = AutoTokenizer.from_pretrained(args.load_from)
+    if 'model' in args.load_from:
+        model = MiniMindForCausalLM(MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe)))
+        moe_suffix = '_moe' if args.use_moe else ''
+        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
+    get_model_params(model, model.config)
+    return model.eval().to(args.device), tokenizer
+
+
+def parse_tool_calls(text):
+    matches = re.findall(r'<tool_call>(.*?)</tool_call>', text, re.DOTALL)
+    calls = []
+    for m in matches:
+        try:
+            calls.append(json.loads(m.strip()))
+        except Exception:
+            pass
+    return calls
+
+
+def parse_tool_call_from_text(content):
+    pattern = r'<tool_call>\s*(\{.*?\})\s*</tool_call>'
+    matches = re.findall(pattern, content, re.DOTALL)
+    if not matches:
+        return None
+    tool_calls = []
+    for i, match in enumerate(matches):
+        try:
+            data = json.loads(match)
+            tool_calls.append({
+                "id": f"call_{i}",
+                "function": {"name": data.get("name", ""), "arguments": json.dumps(data.get("arguments", {}), ensure_ascii=False)}
+            })
+        except Exception:
+            pass
+    return tool_calls if tool_calls else None
+
+
+def execute_tool(call, arguments=None):
+    name = call.get("name", "") if isinstance(call, dict) else call
+    try:
+        raw_args = call.get("arguments", {}) if isinstance(call, dict) else arguments
+        args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
+    except Exception:
+        args = {}
+    fn = MOCK_RESULTS.get(name)
+    if not fn:
+        return {"error": f"未知工具: {name}"}
+    try:
+        return fn(args)
+    except Exception as e:
+        return {"error": f"工具执行失败: {str(e)[:80]}"}
+
+
+def generate(model, tokenizer, messages, tools, args):
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=False)
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True).to(args.device)
+    st = time.time()
+    print('🧠: ', end='')
+    generated_ids = model.generate(
+        inputs["input_ids"], attention_mask=inputs["attention_mask"],
+        max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+        pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+        top_p=args.top_p, temperature=args.temperature
+    )
+    response = tokenizer.decode(generated_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+    gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
+    print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s') if args.show_speed else print()
+    return response
+
+
+def chat_api(client, messages, tools, args, stream=True):
+    response = client.chat.completions.create(
+        model=args.api_model, messages=messages, tools=tools,
+        stream=stream, temperature=args.temperature,
+        max_tokens=8192, top_p=args.top_p
+    )
+    if not stream:
+        choice = response.choices[0]
+        content = choice.message.content or ""
+        tool_calls = choice.message.tool_calls
+        if not tool_calls:
+            tool_calls = parse_tool_call_from_text(content)
+        print(f'🧠: {content}')
+        return content, tool_calls
+    print('🧠: ', end='', flush=True)
+    content, tool_calls = "", None
+    for chunk in response:
+        delta = chunk.choices[0].delta
+        if delta.content:
+            print(delta.content, end="", flush=True)
+            content += delta.content
+        if delta.tool_calls:
+            if tool_calls is None:
+                tool_calls = []
+            for tc_chunk in delta.tool_calls:
+                idx = tc_chunk.index if tc_chunk.index is not None else len(tool_calls)
+                while len(tool_calls) <= idx:
+                    tool_calls.append({
+                        "id": "",
+                        "function": {"name": "", "arguments": ""}
+                    })
+                if tc_chunk.id:
+                    tool_calls[idx]["id"] += tc_chunk.id
+                if tc_chunk.function:
+                    if tc_chunk.function.name:
+                        tool_calls[idx]["function"]["name"] += tc_chunk.function.name
+                    if tc_chunk.function.arguments:
+                        tool_calls[idx]["function"]["arguments"] += tc_chunk.function.arguments
+    print()
+    if not tool_calls:
+        tool_calls = parse_tool_call_from_text(content)
+    return content, tool_calls
+
+
+def run_case(prompt, tools, args, model=None, tokenizer=None, client=None):
+    messages = [{"role": "user", "content": prompt}]
+    while True:
+        if args.backend == 'local':
+            content = generate(model, tokenizer, messages, tools, args)
+            tool_calls = parse_tool_calls(content)
+        else:
+            content, tool_calls = chat_api(client, messages, tools, args, stream=bool(args.stream))
+        if not tool_calls:
+            break
+        tool_calls = [{
+            "id": tc.id if hasattr(tc, 'id') else tc.get("id", ""),
+            "name": tc.function.name if hasattr(tc, 'function') else tc["function"]["name"],
+            "arguments": tc.function.arguments if hasattr(tc, 'function') else tc["function"]["arguments"]
+        } for tc in tool_calls] if args.backend == 'api' else tool_calls
+        messages.append({"role": "assistant", "content": content} if args.backend == 'local' else {"role": "assistant", "content": content, "tool_calls": [{"id": tc["id"], "type": "function", "function": {"name": tc["name"], "arguments": tc["arguments"]}} for tc in tool_calls]})
+        for tc in tool_calls:
+            name = tc["name"]
+            arguments = tc["arguments"]
+            print(f'📞 [Tool Calling]: {name} | args={arguments}')
+            result = execute_tool(tc if args.backend == 'local' else name, arguments)
+            print(f'✅ [Tool Called]: {json.dumps(result, ensure_ascii=False)}')
+            messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)} if args.backend == 'local' else {"role": "tool", "content": json.dumps(result, ensure_ascii=False), "tool_call_id": tc["id"]})
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MiniMind ToolCall评测")
+    parser.add_argument('--backend', default='local', choices=['local', 'api'], type=str, help="推理后端（local=本地模型，api=OpenAI兼容接口）")
+    parser.add_argument('--load_from', default='../model', type=str, help="模型加载路径（model=原生torch权重，其他路径=transformers格式）")
+    parser.add_argument('--save_dir', default='../out', type=str, help="模型权重目录")
+    parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, rlhf, reason, ppo_actor, grpo, spo）")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度")
+    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=是）")
+    args = parser.parse_args()
+
+    model = tokenizer = client = None
+    if args.backend == 'local': model, tokenizer = init_model(args)
+    else: client = OpenAI(api_key=args.api_key, base_url=args.api_base_url)
+
+    input_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
+
+    cases = [{"prompt": case["prompt"], "tools": get_tools(case["tools"]), "tool_names": case["tools"]} for case in TEST_CASES] if input_mode == 0 else iter(lambda: {"prompt": input('💬: '), "tools": TOOLS, "tool_names": [t["function"]["name"] for t in TOOLS]}, {"prompt": "", "tools": TOOLS, "tool_names": []})
+    for case in cases:
+        if not case["prompt"]: break
+        setup_seed(random.randint(0, 31415926))
+        if input_mode == 0:
+            print(f'📦 可用工具: {case["tool_names"]}\n')
+            print(f'💬: {case["prompt"]}')
+        run_case(case["prompt"], case["tools"], args, model=model, tokenizer=tokenizer, client=client)
+        print('\n' + '-' * 50 + '\n')
+
+
+if __name__ == "__main__":
+    main()

scripts/serve_openai_api.py

@@ -1,5 +1,6 @@
 import argparse
 import json
+import re
 import os
 import sys
 
@@ -52,8 +53,19 @@ class ChatRequest(BaseModel):
     temperature: float = 0.7
     top_p: float = 0.92
     max_tokens: int = 8192
-    stream: bool = False
+    stream: bool = True
     tools: list = []
+    open_thinking: bool = False
+    chat_template_kwargs: dict = None
+    
+    def get_open_thinking(self) -> bool:
+        """兼容多种方式开启 thinking"""
+        if self.open_thinking:
+            return True
+        if self.chat_template_kwargs:
+            return self.chat_template_kwargs.get('open_thinking', False) or \
+                   self.chat_template_kwargs.get('enable_thinking', False)
+        return False
 
 
 class CustomStreamer(TextStreamer):
@@ -68,9 +80,31 @@ class CustomStreamer(TextStreamer):
             self.queue.put(None)
 
 
-def generate_stream_response(messages, temperature, top_p, max_tokens):
+def parse_response(text):
+    reasoning_content = None
+    think_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
+    if think_match:
+        reasoning_content = think_match.group(1).strip()
+        text = re.sub(r'<think>.*?</think>\s*', '', text, flags=re.DOTALL)
+    elif '</think>' in text:
+        parts = text.split('</think>', 1)
+        reasoning_content = parts[0].strip()
+        text = parts[1].strip() if len(parts) > 1 else ''
+    tool_calls = []
+    for i, m in enumerate(re.findall(r'<tool_call>(.*?)</tool_call>', text, re.DOTALL)):
+        try:
+            call = json.loads(m.strip())
+            tool_calls.append({"id": f"call_{int(time.time())}_{i}", "type": "function", "function": {"name": call.get("name", ""), "arguments": json.dumps(call.get("arguments", {}), ensure_ascii=False)}})
+        except Exception:
+            pass
+    if tool_calls:
+        text = re.sub(r'<tool_call>.*?</tool_call>', '', text, flags=re.DOTALL)
+    return text.strip(), reasoning_content, tool_calls or None
+
+
+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)[-max_tokens:]
+        new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools or None, open_thinking=open_thinking)[-max_tokens:]
         inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
 
         queue = Queue()
@@ -91,20 +125,44 @@ def generate_stream_response(messages, temperature, top_p, max_tokens):
 
         Thread(target=_generate).start()
 
+        full_text = ""
+        emitted = 0
+        thinking_ended = not bool(open_thinking)
+
         while True:
             text = queue.get()
             if text is None:
-                yield json.dumps({
-                    "choices": [{
-                        "delta": {},
-                        "finish_reason": "stop"
-                    }]
-                }, ensure_ascii=False)
                 break
+            full_text += text
 
-            yield json.dumps({
-                "choices": [{"delta": {"content": text}}]
-            }, ensure_ascii=False)
+            if not thinking_ended:
+                pos = full_text.find('</think>')
+                if pos >= 0:
+                    thinking_ended = True
+                    new_r = full_text[emitted:pos]
+                    if new_r:
+                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                    emitted = pos + len('</think>')
+                    after = full_text[emitted:].lstrip('\n')
+                    emitted = len(full_text) - len(after)
+                    if after:
+                        yield json.dumps({"choices": [{"delta": {"content": after}}]}, ensure_ascii=False)
+                        emitted = len(full_text)
+                else:
+                    new_r = full_text[emitted:]
+                    if new_r:
+                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                        emitted = len(full_text)
+            else:
+                new_c = full_text[emitted:]
+                if new_c:
+                    yield json.dumps({"choices": [{"delta": {"content": new_c}}]}, ensure_ascii=False)
+                    emitted = len(full_text)
+
+        _, _, tool_calls = parse_response(full_text)
+        if tool_calls:
+            yield json.dumps({"choices": [{"delta": {"tool_calls": tool_calls}}]}, ensure_ascii=False)
+        yield json.dumps({"choices": [{"delta": {}, "finish_reason": "tool_calls" if tool_calls else "stop"}]}, ensure_ascii=False)
 
     except Exception as e:
         yield json.dumps({"error": str(e)})
@@ -119,7 +177,9 @@ async def chat_completions(request: ChatRequest):
                     messages=request.messages,
                     temperature=request.temperature,
                     top_p=request.top_p,
-                    max_tokens=request.max_tokens
+                    max_tokens=request.max_tokens,
+                    tools=request.tools,
+                    open_thinking=request.get_open_thinking()
                 )),
                 media_type="text/event-stream"
             )
@@ -127,7 +187,9 @@ async def chat_completions(request: ChatRequest):
             new_prompt = tokenizer.apply_chat_template(
                 request.messages,
                 tokenize=False,
-                add_generation_prompt=True
+                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():
@@ -142,6 +204,12 @@ async def chat_completions(request: ChatRequest):
                     temperature=request.temperature
                 )
                 answer = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+            content, reasoning_content, tool_calls = parse_response(answer)
+            message = {"role": "assistant", "content": content}
+            if reasoning_content:
+                message["reasoning_content"] = reasoning_content
+            if tool_calls:
+                message["tool_calls"] = tool_calls
             return {
                 "id": f"chatcmpl-{int(time.time())}",
                 "object": "chat.completion",
@@ -150,8 +218,8 @@ async def chat_completions(request: ChatRequest):
                 "choices": [
                     {
                         "index": 0,
-                        "message": {"role": "assistant", "content": answer},
-                        "finish_reason": "stop"
+                        "message": message,
+                        "finish_reason": "tool_calls" if tool_calls else "stop"
                     }
                 ]
             }
@@ -165,8 +233,8 @@ if __name__ == "__main__":
     parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
     parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, dpo, reason, ppo_actor, grpo, spo）")
     parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称（None表示不使用，可选：lora_identity, lora_medical）")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度（512=Small-26M, 640=MoE-145M, 768=Base-104M）")
-    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量（Small/MoE=8, Base=16）")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     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倍，仅解决位置编码问题）")

scripts/web_demo.py

@@ -1,10 +1,13 @@
 import random
 import re
+import json
+import os
 from threading import Thread
 
 import torch
 import numpy as np
 import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 st.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed")
 
@@ -64,33 +67,130 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-system_prompt = []
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# 多语言文本
+LANG_TEXTS = {
+    'zh': {
+        'settings': '模型设定调整',
+        'history_rounds': '历史对话轮次',
+        'max_length': '最大生成长度',
+        'temperature': '温度',
+        'thinking': '思考',
+        'tools': '工具',
+        'language': '语言',
+        'send': '给 MiniMind 发送消息',
+        'disclaimer': 'AI 生成内容可能存在错误，请仔细核实',
+        'think_tip': '自适应思考，目前多轮对话或Tool Call共存时思考不稳定',
+        'tool_select': '工具选择（最多4个）',
+    },
+    'en': {
+        'settings': 'Model Settings',
+        'history_rounds': 'History Rounds',
+        'max_length': 'Max Length',
+        'temperature': 'Temperature',
+        'thinking': 'Thinking',
+        'tools': 'Tools',
+        'language': 'Language',
+        'send': 'Send a message to MiniMind',
+        'disclaimer': 'AI-generated content may be inaccurate, please verify',
+        'think_tip': 'Adaptive thinking; may be unstable with multi-turn or Tool Call',
+        'tool_select': 'Tool Selection (max 4)',
+    }
+}
 
-def process_assistant_content(content):
-    if model_source == "API" and 'R1' not in api_model_name:
-        return content
-    if model_source != "API" and 'R1' not in MODEL_PATHS[selected_model][1]:
-        return content
+def get_text(key):
+    lang = st.session_state.get('lang', 'en')
+    return LANG_TEXTS.get(lang, {}).get(key, LANG_TEXTS['zh'].get(key, key))
+
+# 工具定义
+TOOLS = [
+    {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式", "parameters": {"type": "object", "properties": {"expression": {"type": "string", "description": "数学表达式"}}, "required": ["expression"]}}},
+    {"type": "function", "function": {"name": "get_current_time", "description": "获取当前时间", "parameters": {"type": "object", "properties": {"timezone": {"type": "string", "default": "Asia/Shanghai"}}, "required": []}}},
+    {"type": "function", "function": {"name": "random_number", "description": "生成随机数", "parameters": {"type": "object", "properties": {"min": {"type": "integer"}, "max": {"type": "integer"}}, "required": ["min", "max"]}}},
+    {"type": "function", "function": {"name": "text_length", "description": "计算文本长度", "parameters": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]}}},
+    {"type": "function", "function": {"name": "unit_converter", "description": "单位转换", "parameters": {"type": "object", "properties": {"value": {"type": "number"}, "from_unit": {"type": "string"}, "to_unit": {"type": "string"}}, "required": ["value", "from_unit", "to_unit"]}}},
+    {"type": "function", "function": {"name": "get_current_weather", "description": "获取天气", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}},
+    {"type": "function", "function": {"name": "get_exchange_rate", "description": "获取汇率", "parameters": {"type": "object", "properties": {"from_currency": {"type": "string"}, "to_currency": {"type": "string"}}, "required": ["from_currency", "to_currency"]}}},
+    {"type": "function", "function": {"name": "translate_text", "description": "翻译文本", "parameters": {"type": "object", "properties": {"text": {"type": "string"}, "target_lang": {"type": "string"}}, "required": ["text", "target_lang"]}}},
+]
+
+TOOL_SHORT_NAMES = {
+    'calculate_math': '数学', 'get_current_time': '时间', 'random_number': '随机',
+    'text_length': '字数', 'unit_converter': '单位', 'get_current_weather': '天气',
+    'get_exchange_rate': '汇率', 'translate_text': '翻译'
+}
+
+def execute_tool(tool_name, args):
+    import datetime
+    try:
+        if tool_name == 'calculate_math':
+            return {"result": eval(args.get('expression', '0'))}
+        elif tool_name == 'get_current_time':
+            tz = args.get('timezone', 'Asia/Shanghai')
+            return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
+        elif tool_name == 'random_number':
+            return {"result": random.randint(args.get('min', 0), args.get('max', 100))}
+        elif tool_name == 'text_length':
+            return {"result": len(args.get('text', ''))}
+        elif tool_name == 'unit_converter':
+            return {"result": f"{args.get('value', 0)} {args.get('from_unit', '')} = ? {args.get('to_unit', '')}"}
+        elif tool_name == 'get_current_weather':
+            return {"result": f"{args.get('city', 'Unknown')}: 晴, 7~10°C"}
+        elif tool_name == 'get_exchange_rate':
+            return {"result": f"1 {args.get('from_currency', 'USD')} = 7.2 {args.get('to_currency', 'CNY')}"}
+        elif tool_name == 'translate_text':
+            return {"result": f"[翻译结果]: hello world"}
+        return {"result": "Unknown tool"}
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def process_assistant_content(content, is_streaming=False):
+    # 处理tool_call标签，格式化显示
+    if '<tool_call>' in content:
+        def format_tool_call(match):
+            try:
+                tc = json.loads(match.group(1))
+                name = tc.get('name', 'unknown')
+                args = tc.get('arguments', {})
+                return f'<div style="background: rgba(80, 110, 150, 0.20); border: 1px solid rgba(140, 170, 210, 0.30); padding: 10px 12px; border-radius: 12px; margin: 6px 0;"><div style="font-size:12px;opacity:.75;display:block;margin:0 0 6px 0;line-height:1;">ToolCalling</div><div><b>{name}</b>: {json.dumps(args, ensure_ascii=False)}</div></div>'
+            except:
+                return match.group(0)
+        content = re.sub(r'<tool_call>(.*?)</tool_call>', format_tool_call, content, flags=re.DOTALL)
+    
+    # 流式生成且开启思考时，一开始就放到折叠里
+    if is_streaming and st.session_state.get('enable_thinking', False) and '</think>' not in content and '<think>' not in content:
+        m = re.search(r'(\n\n(?:我是|您好|你好)[^\n]*)', content)
+        if m and m.start(1) > 5:
+            i = m.start(1)
+            think_part = content[:i]
+            answer_part = content[i:]
+            return f'<details open style="border-left: 2px solid #666; padding-left: 12px; margin: 8px 0;"><summary style="cursor: pointer; color: #888;">已思考</summary><div style="color: #aaa; font-size: 0.95em; margin-top: 8px; max-height: 100px; overflow-y: auto;">{think_part.strip()}</div></details>{answer_part}'
+        elif len(content) > 5:
+            return f'<details open style="border-left: 2px solid #666; padding-left: 12px; margin: 8px 0;"><summary style="cursor: pointer; color: #888;">思考中...</summary><div style="color: #aaa; font-size: 0.95em; margin-top: 8px; max-height: 100px; overflow-y: auto; display: flex; flex-direction: column-reverse;"><div style="margin-bottom: auto;">{content.strip().replace(chr(10), "<br>")}</div></div></details>'
 
     if '<think>' in content and '</think>' in content:
-        content = re.sub(r'(<think>)(.*?)(</think>)',
-                         r'<details style="font-style: italic; background: rgba(222, 222, 222, 0.5); padding: 10px; border-radius: 10px;"><summary style="font-weight:bold;">推理内容（展开）</summary>\2</details>',
-                         content,
-                         flags=re.DOTALL)
+        def format_think(match):
+            think_content = match.group(2)
+            if think_content.replace('\n', '').strip():  # 不是全换行
+                return f'<details open style="border-left: 2px solid #666; padding-left: 12px; margin: 8px 0;"><summary style="cursor: pointer; color: #888;">已思考</summary><div style="color: #aaa; font-size: 0.95em; margin-top: 8px; max-height: 100px; overflow-y: auto;">{think_content.strip()}</div></details>'
+            return ''
+        content = re.sub(r'(<think>)(.*?)(</think>)', format_think, content, flags=re.DOTALL)
 
     if '<think>' in content and '</think>' not in content:
-        content = re.sub(r'<think>(.*?)$',
-                         r'<details open style="font-style: italic; background: rgba(222, 222, 222, 0.5); padding: 10px; border-radius: 10px;"><summary style="font-weight:bold;">推理中...</summary>\1</details>',
-                         content,
-                         flags=re.DOTALL)
+        def format_think_in_progress(match):
+            tc = match.group(1)
+            return f'<details open style="border-left: 2px solid #666; padding-left: 12px; margin: 8px 0;"><summary style="cursor: pointer; color: #888;">思考中...</summary><div style="color: #aaa; font-size: 0.95em; margin-top: 8px; max-height: 100px; overflow-y: auto; display: flex; flex-direction: column-reverse;"><div style="margin-bottom: auto;">{tc.strip().replace(chr(10), "<br>")}</div></div></details>'
+        content = re.sub(r'<think>(.*?)$', format_think_in_progress, content, flags=re.DOTALL)
 
     if '<think>' not in content and '</think>' in content:
-        content = re.sub(r'(.*?)</think>',
-                         r'<details style="font-style: italic; background: rgba(222, 222, 222, 0.5); padding: 10px; border-radius: 10px;"><summary style="font-weight:bold;">推理内容（展开）</summary>\1</details>',
-                         content,
-                         flags=re.DOTALL)
+        def format_think_no_start(match):
+            think_content = match.group(1)
+            if think_content.replace('\n', '').strip():
+                return f'<details open style="border-left: 2px solid #666; padding-left: 12px; margin: 8px 0;"><summary style="cursor: pointer; color: #888;">已思考</summary><div style="color: #aaa; font-size: 0.95em; margin-top: 8px; max-height: 100px; overflow-y: auto;">{think_content.strip()}</div></details>'
+            return ''
+        content = re.sub(r'(.*?)</think>', format_think_no_start, content, flags=re.DOTALL)
 
     return content
 
@@ -118,17 +218,10 @@ def init_chat_messages():
     if "messages" in st.session_state:
         for i, message in enumerate(st.session_state.messages):
             if message["role"] == "assistant":
-                with st.chat_message("assistant", avatar=image_url):
-                    st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
-                    if st.button("🗑", key=f"delete_{i}"):
-                        st.session_state.messages.pop(i)
-                        st.session_state.messages.pop(i - 1)
-                        st.session_state.chat_messages.pop(i)
-                        st.session_state.chat_messages.pop(i - 1)
-                        st.rerun()
+                st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
             else:
                 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: #ddd; border-radius: 10px; color: black;">{message["content"]}</div></div>',
+                    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;">{message["content"]}</div></div>',
                     unsafe_allow_html=True)
 
     else:
@@ -143,52 +236,64 @@ def regenerate_answer(index):
     st.rerun()
 
 
-def delete_conversation(index):
-    st.session_state.messages.pop(index)
-    st.session_state.messages.pop(index - 1)
-    st.session_state.chat_messages.pop(index)
-    st.session_state.chat_messages.pop(index - 1)
+# 动态扫描模型目录
+script_dir = os.path.dirname(os.path.abspath(__file__))
+MODEL_PATHS = {}
+for d in sorted(os.listdir(script_dir), reverse=True):
+    full_path = os.path.join(script_dir, d)
+    if os.path.isdir(full_path) and not d.startswith('.') and not d.startswith('_'):
+        if any(f.endswith(('.bin', '.safetensors', '.pt')) or os.path.exists(os.path.join(full_path, 'model.safetensors.index.json')) for f in os.listdir(full_path) if os.path.isfile(os.path.join(full_path, f))):
+            MODEL_PATHS[d] = [d, d]
+if not MODEL_PATHS:
+    MODEL_PATHS = {"No models found": ["", "No models"]}
+
+# 模型选择
+selected_model = st.sidebar.selectbox('Model', list(MODEL_PATHS.keys()), index=0)
+model_path = MODEL_PATHS[selected_model][0]
+slogan = f"我是 {MODEL_PATHS[selected_model][1]}，有什么可以帮你的？" if st.session_state.get('lang', 'en') == 'zh' else f"I am {MODEL_PATHS[selected_model][1]}, how can I help you?"
+
+st.sidebar.markdown('<hr style="margin: 12px 0 16px 0;">', unsafe_allow_html=True)
+
+# 语言选择
+lang_options = {'中文': 'zh', 'English': 'en'}
+current_lang = st.session_state.get('lang', 'en')
+lang_index = 0 if current_lang == 'zh' else 1
+lang_label = st.sidebar.radio('Language / 语言', list(lang_options.keys()), index=lang_index, horizontal=True)
+if lang_options[lang_label] != current_lang:
+    st.session_state.lang = lang_options[lang_label]
     st.rerun()
 
+st.sidebar.markdown('<hr style="margin: 12px 0 16px 0;">', unsafe_allow_html=True)
 
-st.sidebar.title("模型设定调整")
+# 参数设置
+st.session_state.history_chat_num = st.sidebar.slider(get_text('history_rounds'), 0, 8, 0, step=2)
+st.session_state.max_new_tokens = st.sidebar.slider(get_text('max_length'), 256, 8192, 8192, step=1)
+st.session_state.temperature = st.sidebar.slider(get_text('temperature'), 0.6, 1.2, 0.90, step=0.01)
 
-# st.sidebar.text("训练数据偏差，增加上下文记忆时\n多轮对话（较单轮）容易出现能力衰减")
-st.session_state.history_chat_num = st.sidebar.slider("Number of Historical Dialogues", 0, 6, 0, step=2)
-# st.session_state.history_chat_num = 0
-st.session_state.max_new_tokens = st.sidebar.slider("Max Sequence Length", 256, 8192, 8192, step=1)
-st.session_state.temperature = st.sidebar.slider("Temperature", 0.6, 1.2, 0.85, step=0.01)
+st.sidebar.markdown('<hr style="margin: 12px 0 16px 0;">', unsafe_allow_html=True)
 
-model_source = st.sidebar.radio("选择模型来源", ["本地模型", "API"], index=0)
-
-if model_source == "API":
-    api_url = st.sidebar.text_input("API URL", value="http://127.0.0.1:8000/v1")
-    api_model_id = st.sidebar.text_input("Model ID", value="minimind")
-    api_model_name = st.sidebar.text_input("Model Name", value="MiniMind2")
-    api_key = st.sidebar.text_input("API Key", value="none", type="password")
-    slogan = f"Hi, I'm {api_model_name}"
-else:
-    MODEL_PATHS = {
-        "MiniMind2-R1 (0.1B)": ["../MiniMind2-R1", "MiniMind2-R1"],
-        "MiniMind2-Small-R1 (0.02B)": ["../MiniMind2-Small-R1", "MiniMind2-Small-R1"],
-        "MiniMind2 (0.1B)": ["../MiniMind2", "MiniMind2"],
-        "MiniMind2-MoE (0.15B)": ["../MiniMind2-MoE", "MiniMind2-MoE"],
-        "MiniMind2-Small (0.02B)": ["../MiniMind2-Small", "MiniMind2-Small"]
-    }
-
-    selected_model = st.sidebar.selectbox('Models', list(MODEL_PATHS.keys()), index=2)  # 默认选择 MiniMind2
-    model_path = MODEL_PATHS[selected_model][0]
-    slogan = f"Hi, I'm {MODEL_PATHS[selected_model][1]}"
+# 功能开关
+st.session_state.enable_thinking = st.sidebar.checkbox(get_text('thinking'), value=False, help=get_text('think_tip'))
+st.session_state.selected_tools = []
+with st.sidebar.expander(get_text('tools')):
+    st.caption(get_text('tool_select'))
+    selected_count = sum(1 for tool in TOOLS if st.session_state.get(f"tool_{tool['function']['name']}", False))
+    for tool in TOOLS:
+        name = tool['function']['name']
+        short_name = TOOL_SHORT_NAMES.get(name, name)
+        checked = st.checkbox(short_name, key=f"tool_{name}", disabled=(selected_count >= 4 and not st.session_state.get(f"tool_{name}", False)))
+        if checked and len(st.session_state.selected_tools) < 4:
+            st.session_state.selected_tools.append(name)
 
 image_url = "https://www.modelscope.cn/api/v1/studio/gongjy/MiniMind/repo?Revision=master&FilePath=images%2Flogo2.png&View=true"
 
 st.markdown(
     f'<div style="display: flex; flex-direction: column; align-items: center; text-align: center; margin: 0; padding: 0;">'
     '<div style="font-style: italic; font-weight: 900; margin: 0; padding-top: 4px; display: flex; align-items: center; justify-content: center; flex-wrap: wrap; width: 100%;">'
-    f'<img src="{image_url}" style="width: 45px; height: 45px; "> '
+    f'<img src="{image_url}" style="width: 40px; height: 40px; "> '
     f'<span style="font-size: 26px; margin-left: 10px;">{slogan}</span>'
     '</div>'
-    '<span style="color: #bbb; font-style: italic; margin-top: 6px; margin-bottom: 10px;">内容完全由AI生成，请务必仔细甄别<br>Content AI-generated, please discern with care</span>'
+    f'<span style="color: #bbb; font-style: italic; margin-top: 6px; margin-bottom: 10px;">{get_text("disclaimer")}</span>'
     '</div>',
     unsafe_allow_html=True
 )
@@ -205,10 +310,7 @@ def setup_seed(seed):
 
 
 def main():
-    if model_source == "本地模型":
-        model, tokenizer = load_model_tokenizer(model_path)
-    else:
-        model, tokenizer = None, None
+    model, tokenizer = load_model_tokenizer(model_path)
 
     if "messages" not in st.session_state:
         st.session_state.messages = []
@@ -218,18 +320,13 @@ def main():
 
     for i, message in enumerate(messages):
         if message["role"] == "assistant":
-            with st.chat_message("assistant", avatar=image_url):
-                st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
-                if st.button("×", key=f"delete_{i}"):
-                    st.session_state.messages = st.session_state.messages[:i - 1]
-                    st.session_state.chat_messages = st.session_state.chat_messages[:i - 1]
-                    st.rerun()
+            st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
         else:
             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: gray; border-radius: 10px; color:white; ">{message["content"]}</div></div>',
+                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;">{message["content"]}</div></div>',
                 unsafe_allow_html=True)
 
-    prompt = st.chat_input(key="input", placeholder="给 MiniMind 发送消息")
+    prompt = st.chat_input(key="input", placeholder=get_text('send'))
 
     if hasattr(st.session_state, 'regenerate') and st.session_state.regenerate:
         prompt = st.session_state.last_user_message
@@ -240,89 +337,84 @@ def main():
 
     if prompt:
         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: gray; border-radius: 10px; color:white; ">{prompt}</div></div>',
+            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:]})
 
-        with st.chat_message("assistant", avatar=image_url):
-            placeholder = st.empty()
+        placeholder = st.empty()
 
-            if model_source == "API":
-                try:
-                    from openai import OpenAI
+        random_seed = random.randint(0, 2 ** 32 - 1)
+        setup_seed(random_seed)
 
-                    client = OpenAI(
-                        api_key=api_key,
-                        base_url=api_url
-                    )
-                    history_num = st.session_state.history_chat_num + 1  # +1 是为了包含当前的用户消息
-                    conversation_history = system_prompt + st.session_state.chat_messages[-history_num:]
-                    answer = ""
-                    response = client.chat.completions.create(
-                        model=api_model_id,
-                        messages=conversation_history,
-                        stream=True,
-                        temperature=st.session_state.temperature
-                    )
+        tools = [t for t in TOOLS if t['function']['name'] in st.session_state.get('selected_tools', [])] or None
+        sys_prompt = [] if tools else [{"role": "system", "content": "你是MiniMind，一个乐于助人、知识渊博的AI助手。请用完整且友好的方式回答用户问题。"}]
+        st.session_state.chat_messages = sys_prompt + st.session_state.chat_messages[-(st.session_state.history_chat_num + 1):]
+        template_kwargs = {"tokenize": False, "add_generation_prompt": True}
+        if st.session_state.get('enable_thinking', False):
+            template_kwargs["open_thinking"] = True
+        if tools:
+            template_kwargs["tools"] = tools
+        new_prompt = tokenizer.apply_chat_template(st.session_state.chat_messages, **template_kwargs)
 
-                    for chunk in response:
-                        content = chunk.choices[0].delta.content or ""
-                        answer += content
-                        placeholder.markdown(process_assistant_content(answer), unsafe_allow_html=True)
+        inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
 
-                except Exception as e:
-                    answer = f"API调用出错: {str(e)}"
-                    placeholder.markdown(answer, unsafe_allow_html=True)
-            else:
-                random_seed = random.randint(0, 2 ** 32 - 1)
-                setup_seed(random_seed)
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {
+            "input_ids": inputs.input_ids,
+            "max_length": inputs.input_ids.shape[1] + st.session_state.max_new_tokens,
+            "num_return_sequences": 1,
+            "do_sample": True,
+            "attention_mask": inputs.attention_mask,
+            "pad_token_id": tokenizer.pad_token_id,
+            "eos_token_id": tokenizer.eos_token_id,
+            "temperature": st.session_state.temperature,
+            "top_p": 0.85,
+            "streamer": streamer,
+        }
 
-                st.session_state.chat_messages = system_prompt + st.session_state.chat_messages[
-                                                                 -(st.session_state.history_chat_num + 1):]
-                new_prompt = tokenizer.apply_chat_template(
-                    st.session_state.chat_messages,
-                    tokenize=False,
-                    add_generation_prompt=True
-                )
+        Thread(target=model.generate, kwargs=generation_kwargs).start()
 
-                inputs = tokenizer(
-                    new_prompt,
-                    return_tensors="pt",
-                    truncation=True
-                ).to(device)
+        answer = ""
+        for new_text in streamer:
+            answer += new_text
+            placeholder.markdown(process_assistant_content(answer, is_streaming=True), unsafe_allow_html=True)
 
-                streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-                generation_kwargs = {
-                    "input_ids": inputs.input_ids,
-                    "max_length": inputs.input_ids.shape[1] + st.session_state.max_new_tokens,
-                    "num_return_sequences": 1,
-                    "do_sample": True,
-                    "attention_mask": inputs.attention_mask,
-                    "pad_token_id": tokenizer.pad_token_id,
-                    "eos_token_id": tokenizer.eos_token_id,
-                    "temperature": st.session_state.temperature,
-                    "top_p": 0.85,
-                    "streamer": streamer,
-                }
-
-                Thread(target=model.generate, kwargs=generation_kwargs).start()
-
-                answer = ""
-                for new_text in streamer:
-                    answer += new_text
-                    placeholder.markdown(process_assistant_content(answer), unsafe_allow_html=True)
-
-            messages.append({"role": "assistant", "content": answer})
+        full_answer = answer
+        for _ in range(16):
+            tool_calls = re.findall(r'<tool_call>(.*?)</tool_call>', answer, re.DOTALL)
+            if not tool_calls:
+                break
             st.session_state.chat_messages.append({"role": "assistant", "content": answer})
-            with st.empty():
-                if st.button("×", key=f"delete_{len(messages) - 1}"):
-                    st.session_state.messages = st.session_state.messages[:-2]
-                    st.session_state.chat_messages = st.session_state.chat_messages[:-2]
-                    st.rerun()
+            tool_results = []
+            for tc_str in tool_calls:
+                try:
+                    tc = json.loads(tc_str.strip())
+                    result = execute_tool(tc.get('name', ''), tc.get('arguments', {}))
+                    st.session_state.chat_messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)})
+                    tool_results.append(f'<div style="background: rgba(90, 130, 110, 0.20); border: 1px solid rgba(150, 200, 170, 0.30); padding: 10px 12px; border-radius: 12px; margin: 6px 0;"><div style="font-size:12px;opacity:.75;display:block;margin:0 0 6px 0;line-height:1;">ToolCalled</div><div><b>{tc.get("name", "")}</b>: {json.dumps(result, ensure_ascii=False)}</div></div>')
+                except:
+                    pass
+            full_answer += "\n" + "\n".join(tool_results) + "\n"
+            placeholder.markdown(process_assistant_content(full_answer, is_streaming=True), unsafe_allow_html=True)
+            new_prompt = tokenizer.apply_chat_template(st.session_state.chat_messages, **template_kwargs)
+            inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
+            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+            generation_kwargs["input_ids"] = inputs.input_ids
+            generation_kwargs["attention_mask"] = inputs.attention_mask
+            generation_kwargs["max_length"] = inputs.input_ids.shape[1] + st.session_state.max_new_tokens
+            generation_kwargs["streamer"] = streamer
+            Thread(target=model.generate, kwargs=generation_kwargs).start()
+            answer = ""
+            for new_text in streamer:
+                answer += new_text
+                placeholder.markdown(process_assistant_content(full_answer + answer, is_streaming=True), unsafe_allow_html=True)
+            full_answer += answer
+        answer = full_answer
+
+        messages.append({"role": "assistant", "content": answer})
+        st.session_state.chat_messages.append({"role": "assistant", "content": answer})
 
 
 if __name__ == "__main__":
-    from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-
     main()

trainer/rollout_engine.py

@@ -0,0 +1,212 @@
+"""Rollout Engine - 可插拔的推理引擎
+python -m sglang.launch_server --model-path ./minimind-3 --attention-backend triton --host 0.0.0.0 --port 8998
+"""
+import os
+import sys
+
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import requests
+import torch
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+from torch import Tensor
+from torch.nn.parallel import DistributedDataParallel
+from transformers import AutoTokenizer
+
+
+# ===== 计算每个 token 的 logprob =====
+def compute_per_token_logps(model, input_ids: Tensor, n_keep: int, attention_mask: Optional[Tensor] = None) -> Tensor:
+    if n_keep <= 0:
+        return input_ids.new_empty((input_ids.size(0), 0), dtype=torch.float32)
+    unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
+    input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
+    logits = unwrapped(input_ids, attention_mask=attention_mask, logits_to_keep=n_keep + 1).logits[:, :-1, :]
+    per_token_logps = []
+    for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
+        ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
+        per_token_logps.append(
+            torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1)
+        )
+    return torch.stack(per_token_logps)
+
+
+# ===== Rollout 结果 =====
+@dataclass
+class RolloutResult:
+    output_ids: Tensor
+    completion_ids: Tensor
+    per_token_logps: Tensor
+    completions: List[str]
+
+
+# ===== Rollout 引擎抽象基类 =====
+class RolloutEngine(ABC):
+    tokenizer = None
+    
+    @abstractmethod
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+        pass
+    
+    @abstractmethod
+    def update_policy(self, model: torch.nn.Module):
+        pass
+
+
+# ===== PyTorch 原生推理引擎 =====
+class TorchRolloutEngine(RolloutEngine):
+    def __init__(self, policy_model: torch.nn.Module, tokenizer, device: str = "cuda", autocast_ctx=None):
+        self.policy_model = policy_model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.autocast_ctx = autocast_ctx
+    
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+        model = self.policy_model.module if isinstance(self.policy_model, DistributedDataParallel) else self.policy_model
+        
+        with torch.no_grad():
+            output_ids = model.generate(
+                input_ids=prompt_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                num_return_sequences=num_generations,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )  # [B*num_gen, P+R]
+        
+        prompt_len = prompt_ids.size(1)
+        completion_ids = output_ids[:, prompt_len:]  # [B*num_gen, R]
+        
+        from contextlib import nullcontext
+        ctx = self.autocast_ctx if self.autocast_ctx else nullcontext()
+        with ctx:
+            per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1))
+        
+        completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
+        return RolloutResult(output_ids, completion_ids, per_token_logps, completions)
+    
+    def update_policy(self, model: torch.nn.Module):
+        self.policy_model = model
+
+
+# ===== SGLang HTTP API 推理引擎 =====
+class SGLangRolloutEngine(RolloutEngine):
+    def __init__(self, base_url: str, model_path: str, shared_ckpt_path: str = "./sglang_ckpt", timeout: int = 120):
+        self.base_url = base_url.rstrip('/')
+        self.shared_ckpt_path = shared_ckpt_path
+        self.timeout = timeout
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.http = requests
+    
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+        # 去除左侧 padding tokens，只保留有效 token
+        input_ids_list = []
+        for ids, mask in zip(prompt_ids, attention_mask):
+            valid_ids = ids[mask.bool()].tolist()
+            input_ids_list.append(valid_ids)
+        all_input_ids = [ids for ids in input_ids_list for _ in range(num_generations)]
+        
+        payload = {
+            "input_ids": all_input_ids,
+            "sampling_params": {
+                "temperature": temperature,
+                "max_new_tokens": max_new_tokens,
+                "stop_token_ids": [self.tokenizer.eos_token_id] if self.tokenizer.eos_token_id else [],
+            },
+            "return_logprob": True,
+        }
+        
+        resp = self.http.post(f"{self.base_url}/generate", json=payload, timeout=self.timeout)
+        resp.raise_for_status()
+        
+        results = resp.json()
+        if not isinstance(results, list):
+            results = [results]
+        
+        all_output_ids, all_completion_ids, all_logprobs = [], [], []
+        completions = []
+        prompt_len = prompt_ids.size(1)
+        
+        for i, result in enumerate(results):
+            meta = result.get("meta_info", {})
+            completion_ids = meta.get("output_ids", result.get("output_ids", []))
+            raw_logprobs = meta.get("output_token_logprobs", [])
+            
+            logprobs = []
+            for item in raw_logprobs:
+                if isinstance(item, (list, tuple)) and len(item) >= 1:
+                    logprobs.append(item[0])
+                elif isinstance(item, (int, float)):
+                    logprobs.append(item)
+            
+            prompt = all_input_ids[i]
+            full_output = prompt + completion_ids
+            all_output_ids.append(full_output)
+            all_completion_ids.append(completion_ids)
+            all_logprobs.append(logprobs)
+            completions.append(self.tokenizer.decode(completion_ids, skip_special_tokens=True))
+        
+        device = prompt_ids.device
+        max_out_len = max(len(ids) for ids in all_output_ids)
+        max_comp_len = max(len(ids) for ids in all_completion_ids)
+        max_logp_len = max(len(lp) for lp in all_logprobs)
+        
+        def pad_to_tensor(seqs, max_len, pad_val=0):
+            return torch.tensor([s + [pad_val] * (max_len - len(s)) for s in seqs], device=device)
+        
+        return RolloutResult(
+            output_ids=pad_to_tensor(all_output_ids, max_out_len),
+            completion_ids=pad_to_tensor(all_completion_ids, max_comp_len),
+            per_token_logps=pad_to_tensor(all_logprobs, max_logp_len, pad_val=0.0),
+            completions=completions,
+        )
+    
+    def update_policy(self, model: torch.nn.Module):
+        unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
+        abs_path = os.path.abspath(self.shared_ckpt_path)
+        unwrapped.lm_head.weight = torch.nn.Parameter(unwrapped.lm_head.weight.clone())
+        state_dict = {k: v.detach().half().cpu() for k, v in unwrapped.state_dict().items()}
+        unwrapped.save_pretrained(abs_path, state_dict=state_dict, safe_serialization=False)
+        unwrapped.model.embed_tokens.weight = unwrapped.lm_head.weight
+        self.tokenizer.save_pretrained(abs_path)
+        resp = self.http.post(
+            f"{self.base_url}/update_weights_from_disk",
+            json={"model_path": abs_path},
+            timeout=self.timeout
+        )
+        if resp.status_code != 200: print(f"[SGLANG WARNING] update_weights 失败: {resp.status_code}, {resp.text}")
+        return resp.status_code == 200
+    
+    def flush_cache(self) -> bool:
+        resp = self.http.post(f"{self.base_url}/flush_cache", timeout=30)
+        return resp.status_code == 200
+    
+    def health(self) -> bool:
+        try:
+            resp = self.http.get(f"{self.base_url}/health", timeout=5)
+            return resp.status_code == 200
+        except:
+            return False
+
+
+# ===== 工厂函数 =====
+def create_rollout_engine(
+    engine_type: str = "torch",
+    policy_model: torch.nn.Module = None,
+    tokenizer = None,
+    device: str = "cuda",
+    autocast_ctx = None,
+    sglang_base_url: str = None,
+    sglang_model_path: str = None,
+    sglang_shared_path: str = None,
+) -> RolloutEngine:
+    if engine_type == "torch":
+        return TorchRolloutEngine(policy_model, tokenizer, device, autocast_ctx)
+    elif engine_type == "sglang":
+        return SGLangRolloutEngine(sglang_base_url, sglang_model_path, sglang_shared_path)
+    else:
+        raise ValueError(f"不支持的引擎类型: {engine_type}")

trainer/train_agent.py

@@ -0,0 +1,487 @@
+import os
+import sys
+
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import re
+import gc
+import json
+import math
+import random
+import signal
+import argparse
+import warnings
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from transformers import AutoTokenizer
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from dataset.lm_dataset import AgentRLDataset
+from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
+
+warnings.filterwarnings('ignore')
+
+# ================================ 工具与 Reward = Start ================================
+
+def rep_penalty(text, n=3, cap=0.5):
+    toks = re.findall(r"\w+|[^\w\s]", text.lower())
+    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
+
+# ======== 工具定义 ========
+TOOLS = [
+    {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式", "parameters": {"type": "object", "properties": {"expression": {"type": "string"}}, "required": ["expression"]}}},
+    {"type": "function", "function": {"name": "unit_converter", "description": "单位换算", "parameters": {"type": "object", "properties": {"value": {"type": "number"}, "from_unit": {"type": "string"}, "to_unit": {"type": "string"}}, "required": ["value", "from_unit", "to_unit"]}}},
+    {"type": "function", "function": {"name": "get_current_weather", "description": "获取天气", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}},
+    {"type": "function", "function": {"name": "get_current_time", "description": "获取时间", "parameters": {"type": "object", "properties": {"timezone": {"type": "string", "default": "Asia/Shanghai"}}, "required": []}}},
+    {"type": "function", "function": {"name": "get_exchange_rate", "description": "查询汇率", "parameters": {"type": "object", "properties": {"from_currency": {"type": "string"}, "to_currency": {"type": "string"}}, "required": ["from_currency", "to_currency"]}}},
+    {"type": "function", "function": {"name": "translate_text", "description": "翻译文本", "parameters": {"type": "object", "properties": {"text": {"type": "string"}, "target_language": {"type": "string"}}, "required": ["text", "target_language"]}}},
+]
+
+# ======== 模拟数据 ========
+WEATHER_DATA = {"北京": ("28°C", "晴"), "上海": ("15°C", "多云"), "广州": ("32°C", "闷热"), "深圳": ("30°C", "晴"), "杭州": ("22°C", "阴"), "成都": ("18°C", "小雨"), "武汉": ("25°C", "多云"), "南京": ("20°C", "晴"), "西安": ("16°C", "大风"), "重庆": ("26°C", "阴"), "Tokyo": ("12°C", "晴"), "New York": ("8°C", "多云"), "London": ("5°C", "小雨"), "Paris": ("10°C", "阴"), "Sydney": ("25°C", "晴朗")}
+TIME_DATA = {"Asia/Shanghai": "2025-03-07 14:30:00", "America/New_York": "2025-03-07 01:30:00", "Europe/London": "2025-03-07 06:30:00", "Asia/Tokyo": "2025-03-07 15:30:00", "Europe/Paris": "2025-03-07 07:30:00", "Australia/Sydney": "2025-03-07 17:30:00"}
+EXCHANGE_DATA = {("USD", "CNY"): 7.21, ("EUR", "CNY"): 7.85, ("GBP", "CNY"): 9.12, ("JPY", "CNY"): 0.048, ("USD", "EUR"): 0.92, ("USD", "GBP"): 0.79, ("CNY", "JPY"): 20.83, ("AUD", "CNY"): 4.72}
+TRANSLATE_DATA = {("你好世界", "english"): "Hello World", ("Good morning", "chinese"): "早上好", ("今天天气真好", "english"): "The weather is nice today", ("I love programming", "chinese"): "我喜欢编程", ("机器学习很有趣", "english"): "Machine learning is interesting", ("Happy birthday", "chinese"): "生日快乐"}
+UNIT_DATA = {"km_miles": 0.621371, "miles_km": 1.60934, "kg_pounds": 2.20462, "pounds_kg": 0.453592, "meters_feet": 3.28084, "feet_meters": 0.3048, "celsius_fahrenheit": 1.8, "fahrenheit_celsius": 0.5556}
+
+# ======== 模拟执行 ========
+MOCK_RESULTS = {
+    "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("（", "(").replace("）", ")"), {"__builtins__": {}, "math": math}))},
+    "unit_converter": lambda args: {"result": round(float(args.get("value", 0)) * UNIT_DATA.get(f"{args.get('from_unit', '').lower()}_{args.get('to_unit', '').lower()}", 1), 4)},
+    "get_current_weather": lambda args: (lambda w: {"city": args.get("location"), "temperature": w[0], "humidity": "65%", "condition": w[1]})(WEATHER_DATA.get(args.get("location"), ("22°C", "晴"))),
+    "get_current_time": lambda args: {"datetime": TIME_DATA.get(args.get("timezone", "Asia/Shanghai"), "2025-03-07 14:30:00"), "timezone": args.get("timezone", "Asia/Shanghai")},
+    "get_exchange_rate": lambda args: {"from": args.get("from_currency"), "to": args.get("to_currency"), "rate": EXCHANGE_DATA.get((args.get("from_currency"), args.get("to_currency")), 1.0)},
+    "translate_text": lambda args: {"translated_text": TRANSLATE_DATA.get((args.get("text"), args.get("target_language")), args.get("text", ""))},
+}
+
+# ======== 参数校验 ========
+CHECK_ARGS = {
+    "calculate_math": lambda a: bool(a.get("expression")),
+    "unit_converter": lambda a: a.get("value") is not None and a.get("from_unit") and a.get("to_unit"),
+    "get_current_weather": lambda a: bool(a.get("location")),
+    "get_current_time": lambda a: True,
+    "get_exchange_rate": lambda a: bool(a.get("from_currency")) and bool(a.get("to_currency")),
+    "translate_text": lambda a: bool(a.get("text")) and bool(a.get("target_language")),
+}
+
+# ======== 工具调用解析与执行 ========
+def parse_tool_calls(text):
+    calls = []
+    for m in re.findall(r'<tool_call>(.*?)</tool_call>', text, re.DOTALL):
+        try: calls.append(json.loads(m.strip()))
+        except: pass
+    return calls
+
+def execute_tool(name, args):
+    fn = MOCK_RESULTS.get(name)
+    if not fn: return None
+    try:
+        signal.signal(signal.SIGALRM, lambda *_: (_ for _ in ()).throw(TimeoutError()))
+        signal.alarm(1)
+        return fn(args)
+    except:
+        return None
+    finally:
+        try: signal.alarm(0)
+        except: pass
+
+# ======== 多轮 Rollout ========
+def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"):
+    all_outputs = []
+    prompt_ids = None
+    response_ids = []
+    response_mask = []
+    response_old_logps = []
+    final_context = ""
+    unfinished = False
+    open_thinking = random.random() < thinking_ratio
+    for turn in range(max_turns):
+        context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=open_thinking)
+        inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False).to(device)
+        context_ids = inputs["input_ids"][0].tolist()
+        if prompt_ids is None:
+            prompt_ids = context_ids
+        rollout_result = rollout_engine.rollout(
+            prompt_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            num_generations=1,
+            max_new_tokens=max_new_tokens,
+            temperature=0.8,
+        )
+        new_ids = rollout_result.completion_ids[0].tolist()
+        new_logps = rollout_result.per_token_logps[0].tolist()
+        if len(new_ids) != len(new_logps): Logger(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
+        pairs = [(t, lp) for t, lp in zip(new_ids, new_logps) if t != tokenizer.pad_token_id and t != tokenizer.eos_token_id]
+        new_ids = [t for t, _ in pairs]
+        new_logps = [lp for _, lp in pairs]
+        new_text = rollout_result.completions[0]
+        all_outputs.append(new_text)
+        response_ids.extend(new_ids)
+        response_mask.extend([1] * len(new_ids))
+        response_old_logps.extend(new_logps)
+        final_context = context + new_text
+        calls = parse_tool_calls(new_text)
+        if not calls:
+            break
+        unfinished = turn == max_turns - 1
+        messages.append({"role": "assistant", "content": new_text})
+        for call in calls:
+            name, raw = call.get("name", ""), call.get("arguments", {})
+            if isinstance(raw, str):
+                try: raw = json.loads(raw)
+                except: raw = {}
+            result = execute_tool(name, raw)
+            result_str = (json.dumps(result, ensure_ascii=False) if result else '{"error": "tool not found"}')[:2048]  # 防止天文数字撑爆tokenizer
+            messages.append({"role": "tool", "content": result_str})
+
+        observe_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
+        observe_ids = tokenizer(observe_context, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
+        current_len = len(prompt_ids) + len(response_ids)
+        obs_delta = observe_ids[current_len:]
+        response_ids.extend(obs_delta)
+        response_mask.extend([0] * len(obs_delta))
+        response_old_logps.extend([0.0] * len(obs_delta))
+        final_context = observe_context
+
+    final_output = all_outputs[-1] if all_outputs else ""
+    prompt_ids = prompt_ids or []
+    return final_output, final_context, prompt_ids, response_ids, response_mask, response_old_logps, list(all_outputs), unfinished
+
+def rollout_batch(rollout_engine, tokenizer, messages_batch, tools_batch, num_gen, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"):
+    all_completions = []
+    all_contexts = []
+    all_prompt_ids = []
+    all_response_ids = []
+    all_response_masks = []
+    all_response_old_logps = []
+    all_turn_outputs = []
+    all_unfinished = []
+    for messages, tools in zip(messages_batch, tools_batch):
+        for _ in range(num_gen):
+            msgs_copy = [dict(m) for m in messages]
+            completion, context, prompt_ids, response_ids, response_mask, response_old_logps, turn_outputs, unfinished = rollout_single(rollout_engine, tokenizer, msgs_copy, tools, max_turns, max_new_tokens, thinking_ratio, device)
+            all_completions.append(completion)
+            all_contexts.append(context)
+            all_prompt_ids.append(prompt_ids)
+            all_response_ids.append(response_ids)
+            all_response_masks.append(response_mask)
+            all_response_old_logps.append(response_old_logps)
+            all_turn_outputs.append(turn_outputs)
+            all_unfinished.append(unfinished)
+    return all_completions, all_contexts, all_prompt_ids, all_response_ids, all_response_masks, all_response_old_logps, all_turn_outputs, all_unfinished
+
+# ======== Reward 计算 ========
+def validate_gt_in_text(text, gt_list):
+    text, text_num = str(text), str(text).replace(',', '')
+    nums = [float(x) for x in re.findall(r'(?<![\w.])[-+]?\d+(?:\.\d+)?(?![\w.])', text_num)]
+    return {g for g in gt_list if ((s := str(g).strip()) and s.lower() in text.lower()) or (re.fullmatch(r'[-+]?\d+(?:\.\d+)?', str(g).strip().replace(',', '')) and any(abs(float(str(g).strip().replace(',', '')) - n) < 1e-6 for n in nums))}
+
+def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, reward_model=None, device="cuda", turn_outputs_batch=None, unfinished_batch=None):
+    rewards = torch.zeros(len(completions), device=device)
+    for idx, response in enumerate(completions):
+        reward, answer = 0.0, response
+        sample_idx = idx // num_gen
+        tools = tools_batch[sample_idx]
+        turn_outputs = turn_outputs_batch[idx] if turn_outputs_batch is not None else [response]
+        unfinished = unfinished_batch[idx] if unfinished_batch is not None else False
+        turn_answers = [turn.split('</think>', 1)[-1].strip() if '</think>' in turn else turn.strip() for turn in turn_outputs]
+        answer = turn_answers[-1] if turn_answers else response.strip()
+        valid_names = {t['function']['name'] for t in tools} if tools else set()
+        tool_calls = []
+        for turn_answer in turn_answers: tool_calls.extend(parse_tool_calls(turn_answer))  # 解析tool调用
+        reward -= 0.5 * sum(abs(turn.count('<tool_call>') - turn.count('</tool_call>')) for turn in turn_answers)  # 标签扣分
+        # -------- 无工具调用：格式+reward奖励 --------
+        if not tool_calls:
+            reward += 0.5 if 5 <= len(response.strip()) <= 800 else -0.5  # 长度分
+            if '</think>' in response:
+                think, answer = response.split('</think>', 1)
+                reward += 1.0 if 20 <= len(think.strip()) <= 300 else -0.5  # 思考长度分
+                reward += 0.25 if response.count('</think>') == 1 else -0.25  # 思考闭合分
+                answer = answer.strip()
+            if reward_model is not None:
+                prompt = prompts[sample_idx]
+                pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
+                matches = re.findall(pattern, prompt, re.DOTALL)
+                messages = [{"role": role, "content": content.strip()} for role, content in matches]
+                score = reward_model.get_score(messages, answer)
+                reward += score  # RM分
+            reward -= rep_penalty(answer)
+            rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
+        # -------- 有工具调用：执行结果奖励 --------
+        else:
+            gt = gt_batch[sample_idx]
+            valid_call_count = 0
+            for tool_call in tool_calls:
+                name, raw = tool_call.get("name", ""), tool_call.get("arguments", {})
+                if isinstance(raw, str):
+                    try: raw = json.loads(raw)
+                    except: raw = {}
+                check = CHECK_ARGS.get(name)
+                valid_call_count += int(bool(name in valid_names and check and check(raw)))
+            tool_gap = abs(valid_call_count - len(gt)) + max(0, len(tool_calls) - valid_call_count)  # tool数差值
+            reward += 0.5 if tool_gap == 0 else -0.5 * tool_gap  # tool对齐分
+            
+            final_text = "" if unfinished else (answer.split('</tool_call>')[-1] if '</tool_call>' in answer else answer)
+            verified = validate_gt_in_text(final_text, gt) if gt else set()
+            if gt: reward += 2.5 * len(verified) / len(gt)  # GT分
+            if unfinished: reward -= 0.5  # 未完成扣分
+            reward -= rep_penalty(final_text if final_text else answer)
+            rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
+    return rewards
+
+# ================================ 工具与 Reward = End ================================
+def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model=None, start_step=0, wandb=None, use_sglang=False):
+    last_step = start_step
+    for step, batch in enumerate(loader, start=start_step + 1):
+        messages_batch = batch['messages']
+        tools_batch = batch['tools']
+        gt_batch = batch['gt']
+        last_step = step
+
+        with torch.no_grad():
+            completions, contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch = rollout_batch(rollout_engine, tokenizer, messages_batch, tools_batch, args.num_generations, max_turns=3, max_new_tokens=args.max_gen_len, thinking_ratio=args.thinking_ratio, device=args.device)
+
+        prompts = [tokenizer.apply_chat_template(m, tokenize=False, add_generation_prompt=True, tools=t) for m, t in zip(messages_batch, tools_batch)]
+        packed_samples = []
+        for p, r, m, old_lp in zip(prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch):
+            ids = p + r
+            mask = [0] * len(p) + m
+            old_logps = [0.0] * max(len(p) - 1, 0) + old_lp
+            if len(ids) > args.max_total_len:
+                ids = ids[-args.max_total_len:]
+                mask = mask[-args.max_total_len:]
+                old_logps = old_logps[-(len(ids) - 1):]
+            prompt_len = next((i for i, v in enumerate(mask) if v == 1), len(mask))
+            packed_samples.append((ids, mask, prompt_len, old_logps))
+        seq_lens = torch.tensor([len(ids) for ids, _, _, _ in packed_samples], device=args.device)
+        max_len = seq_lens.max().item()
+        input_ids = torch.tensor([ids + [tokenizer.pad_token_id] * (max_len - len(ids)) for ids, _, _, _ in packed_samples], device=args.device)
+        prompt_lens = torch.tensor([prompt_len for _, _, prompt_len, _ in packed_samples], device=args.device)
+        full_response_masks = torch.tensor([mask + [0] * (max_len - len(mask)) for _, mask, _, _ in packed_samples], device=args.device, dtype=torch.float32)
+        old_per_token_logps = torch.tensor([old_logps + [0.0] * ((max_len - 1) - len(old_logps)) for _, _, _, old_logps in packed_samples], device=args.device, dtype=torch.float32)
+
+        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
+        with autocast_ctx:
+            res = model_unwrapped(input_ids)
+            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
+            logits = res.logits[:, :-1, :]
+            per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
+
+        with torch.no_grad():
+            ref_per_token_logps = compute_per_token_logps(ref_model, input_ids, input_ids.size(1) - 1)
+
+        completion_mask = full_response_masks[:, 1:]
+        is_eos = (input_ids[:, 1:] == tokenizer.eos_token_id) & completion_mask.bool()
+        eos_idx = torch.full((completion_mask.size(0),), completion_mask.size(1) - 1, device=args.device, dtype=torch.long)
+        has_eos = is_eos.any(dim=1)
+        eos_idx[has_eos] = is_eos.int().argmax(dim=1)[has_eos]
+        pos = torch.arange(completion_mask.size(1), device=args.device).unsqueeze(0)
+        completion_mask = completion_mask * (pos <= eos_idx.unsqueeze(1)).float()
+        token_counts = completion_mask.sum(dim=1)
+        valid_rows = token_counts > 0
+        rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch)
+
+        if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+            for i in range(len(messages_batch)):
+                Logger(f"[DEBUG] step={step}, gt[{i}]: {repr(gt_batch[i])}")
+                Logger('-'*100)
+                for j in range(args.num_generations):
+                    idx = i * args.num_generations + j
+                    plen, slen = prompt_lens[idx].item(), seq_lens[idx].item()
+                    Logger(f"{'=' * 30} [DEBUG] gen[{i}][{j}] CONTEXT_BEGIN {'=' * 30}")
+                    Logger(contexts[idx])
+                    Logger(f"{'=' * 31} [DEBUG] gen[{i}][{j}] CONTEXT_END {'=' * 31}")
+                    Logger(f"[DEBUG] gen[{i}][{j}] prompt_len={plen}, seq_len={slen}")
+                    tokens = input_ids[idx, plen:slen].tolist()
+                    text = tokenizer.decode(tokens, skip_special_tokens=False)
+                    Logger(f"{'=' * 28} [DEBUG] gen[{i}][{j}] COMPLETION_BEGIN [{plen}:{slen}] {'=' * 28}")
+                    Logger(text)
+                    Logger(f"{'=' * 29} [DEBUG] gen[{i}][{j}] COMPLETION_END {'=' * 29}")
+                    Logger(f"[DEBUG] gen[{i}][{j}] reward={rewards[idx].item():.4f}")
+                    Logger('='*100)
+
+        grouped_rewards = rewards.view(-1, args.num_generations)
+        mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)
+        std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)
+        advantages = (rewards - mean_r) / (std_r + 1e-4)
+
+        kl_div = ref_per_token_logps - per_token_logps
+        per_token_kl = torch.exp(kl_div) - kl_div - 1
+        ratio = torch.exp(per_token_logps - old_per_token_logps)
+        if args.loss_type == "cispo":
+            clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
+            per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
+        else:
+            clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
+            per_token_loss1 = ratio * advantages.unsqueeze(1)
+            per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
+            per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
+        policy_loss = (((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean()
+                       if valid_rows.any() else per_token_loss.sum() * 0.0)
+        loss = (policy_loss + aux_loss) / args.accumulation_steps
+        loss.backward()
+
+        if step % args.accumulation_steps == 0:
+            if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step(); scheduler.step(); optimizer.zero_grad()
+            if is_main_process() and step % args.save_interval == 0: rollout_engine.update_policy(model)
+
+        if step % args.log_interval == 0 or step == iters:
+            pl = loss.item() * args.accumulation_steps
+            ar = rewards.mean().item()
+            al = token_counts.float().mean().item()
+            kl = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(token_counts.sum().item(), 1)
+            gs = grouped_rewards.std(dim=1, unbiased=False).mean().item()
+            am, ast = advantages.mean().item(), advantages.std().item()
+            lr = optimizer.param_groups[0]['lr']
+            Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}), Reward:{ar:.4f}, KL:{kl:.4f}, GrpStd:{gs:.4f}, AdvStd:{ast:.4f}, Loss:{pl:.4f}, AvgLen:{al:.2f}, AdvMean:{am:.4f}, LR:{lr:.8f}')
+            if wandb and is_main_process():
+                wandb.log({"reward":ar,"kl_ref":kl,"group_reward_std":gs,"advantages_std":ast,"policy_loss":pl,"avg_response_len":al,"advantages_mean":am,"learning_rate":lr})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
+                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
+            model.train()
+            del state_dict
+
+        del per_token_logps, ref_per_token_logps
+        del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        optimizer.step(); scheduler.step(); optimizer.zero_grad()
+        if is_main_process() and last_step % args.save_interval == 0: rollout_engine.update_policy(model)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind Agent RL")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='agent', type=str, help="保存权重名称")
+    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="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="模型隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="模型层数")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE")
+    parser.add_argument('--max_seq_len', default=1024, type=int, help="最大序列长度")
+    parser.add_argument("--max_gen_len", type=int, default=768, help="单次最大生成长度")
+    parser.add_argument("--max_total_len", type=int, default=2500, help="训练侧最终总长度上界")
+    parser.add_argument("--data_path", type=str, default="../dataset/agent_rl.jsonl", help="训练数据路径")
+    parser.add_argument("--num_generations", type=int, default=4, help="每个prompt生成数量")
+    parser.add_argument("--beta", type=float, default=0.1, help="KL散度惩罚系数")
+    parser.add_argument("--loss_type", type=str, default="cispo", choices=["grpo", "cispo"], help="loss类型")
+    parser.add_argument("--epsilon", type=float, default=0.2, help="GRPO的PPO clip epsilon")
+    parser.add_argument("--epsilon_high", type=float, default=5.0, help="epsilon上界")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="加载预训练权重名称")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否从checkpoint恢复")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb记录")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-Agent-RL", help="wandb项目名称")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile")
+    parser.add_argument("--debug_mode", action="store_true", help="调试模式")
+    parser.add_argument("--debug_interval", type=int, default=20, help="调试日志间隔")
+    parser.add_argument("--thinking_ratio", type=float, default=0.1, help="按概率开启thinking（0.0~1.0）")
+    parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
+    parser.add_argument("--rollout_engine", type=str, default="sglang", choices=["torch", "sglang"], help="rollout引擎类型")
+    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共享存储路径")
+    args = parser.parse_args()
+
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                               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"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb.init(project=args.wandb_project, name=f"Agent-RL-E{args.epochs}-B{args.batch_size}-LR{args.learning_rate}", id=wandb_id, resume=resume)
+
+    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+
+    ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
+    ref_model = ref_model.eval().requires_grad_(False)
+
+    reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
+    Logger(f'Loaded reward model from {args.reward_model_path}')
+    # Rollout引擎
+    rollout_engine = create_rollout_engine(
+        engine_type=args.rollout_engine,
+        policy_model=model,
+        tokenizer=tokenizer,
+        device=args.device,
+        autocast_ctx=autocast_ctx,
+        sglang_base_url=args.sglang_base_url,
+        sglang_model_path=args.sglang_model_path,
+        sglang_shared_path=args.sglang_shared_path,
+    )
+    train_ds = AgentRLDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    def collate_fn(batch): return {'messages': [b['messages'] for b in batch], 'tools': [b['tools'] for b in batch], 'gt': [b['gt'] for b in batch]}
+    loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler, collate_fn=collate_fn)
+    iters = len(loader_for_count)
+    total_optimizer_steps = math.ceil(iters / args.accumulation_steps) * args.epochs
+    scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
+
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'])
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scheduler.load_state_dict(ckp_data['scheduler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+    if dist.is_initialized():
+        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    if is_main_process(): rollout_engine.update_policy(model)
+
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn)
+        if skip > 0:
+            Logger(f'Epoch [{epoch+1}/{args.epochs}]: skip {start_step} steps')
+            rl_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
+        else:
+            rl_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
+
+    if dist.is_initialized(): dist.destroy_process_group()

trainer/train_distillation.py

@@ -37,12 +37,14 @@ def distillation_loss(student_logits, teacher_logits, temperature=1.0, reduction
 
 def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_step=0, wandb=None, alpha=0.0, temperature=1.0):
     start_time = time.time()
+    last_step = start_step
     
     if teacher_model is not None:
         teacher_model.eval()
         teacher_model.requires_grad_(False)
 
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+        last_step = step
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
         loss_mask = (labels[..., 1:] != -100).float()
@@ -104,7 +106,7 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
             current_ce_loss = ce_loss_raw.item()
             current_aux_loss = res.aux_loss.item() if lm_config_student.use_moe else 0.0
             current_lr = optimizer.param_groups[-1]['lr']
-            eta_min = spend_time / step * iters // 60 - spend_time // 60
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
             
             Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, ce: {current_ce_loss:.4f}, aux_loss: {current_aux_loss:.4f}, distill: {distill_loss.item():.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min')
             
@@ -132,8 +134,16 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
 
         del input_ids, labels, loss_mask, res, student_logits, ce_loss, distill_loss, loss
 
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
 
 if __name__ == "__main__":
+    # 模拟用moe模型蒸馏dense模型，也可以用更大teacher_hidden_size模型蒸馏更小student_hidden_size的
     parser = argparse.ArgumentParser(description="MiniMind Knowledge Distillation")
     parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
     parser.add_argument('--save_weight', default='full_dist', type=str, help="保存权重的前缀名")
@@ -148,12 +158,13 @@ if __name__ == "__main__":
     parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
     parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
     parser.add_argument("--max_seq_len", type=int, default=340, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
-    parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl", help="训练数据路径")
-    parser.add_argument('--student_hidden_size', default=512, type=int, help="学生模型隐藏层维度")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
+    parser.add_argument('--student_hidden_size', default=768, type=int, help="学生模型隐藏层维度")
     parser.add_argument('--student_num_layers', default=8, type=int, help="学生模型隐藏层数量")
     parser.add_argument('--teacher_hidden_size', default=768, type=int, help="教师模型隐藏层维度")
-    parser.add_argument('--teacher_num_layers', default=16, type=int, help="教师模型隐藏层数量")
-    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--teacher_num_layers', default=8, type=int, help="教师模型隐藏层数量")
+    parser.add_argument('--student_use_moe', default=0, type=int, choices=[0, 1], help="学生模型是否使用MoE（0=否，1=是）")
+    parser.add_argument('--teacher_use_moe', default=1, type=int, choices=[0, 1], help="教师模型是否使用MoE（0=否，1=是）")
     parser.add_argument('--from_student_weight', default='full_sft', type=str, help="学生模型基于哪个权重")
     parser.add_argument('--from_teacher_weight', default='full_sft', type=str, help="教师模型基于哪个权重")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
@@ -171,8 +182,8 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
-    lm_config_student = MiniMindConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.use_moe))
-    lm_config_teacher = MiniMindConfig(hidden_size=args.teacher_hidden_size, num_hidden_layers=args.teacher_num_layers, use_moe=bool(args.use_moe))
+    lm_config_student = MiniMindConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.student_use_moe))
+    lm_config_teacher = MiniMindConfig(hidden_size=args.teacher_hidden_size, num_hidden_layers=args.teacher_num_layers, use_moe=bool(args.teacher_use_moe))
     ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========
@@ -191,9 +202,6 @@ if __name__ == "__main__":
     
     # ========== 5. 定义学生和教师模型 ==========
     model, tokenizer = init_model(lm_config_student, args.from_student_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
     Logger(f'学生模型总参数量：{sum(p.numel() for p in model.parameters()) / 1e6:.3f} M')
     teacher_model, _ = init_model(lm_config_teacher, args.from_teacher_weight, device=args.device)
     teacher_model.eval()
@@ -213,7 +221,10 @@ if __name__ == "__main__":
         start_epoch = ckp_data['epoch']
         start_step = ckp_data.get('step', 0)
     
-    # ========== 7. DDP包模型 ==========
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
     if dist.is_initialized():
         model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         model = DistributedDataParallel(model, device_ids=[local_rank])

trainer/train_dpo.py

@@ -32,10 +32,8 @@ def logits_to_log_probs(logits, labels):
 
 def dpo_loss(ref_log_probs, policy_log_probs, mask, beta):
     # ref_log_probs 和 policy_log_probs 都是 shape: (batch_size, seq_len)
-    # https://github.com/jingyaogong/minimind/issues/298
-    seq_lengths = mask.sum(dim=1, keepdim=True).clamp_min(1e-8)  # 防止零长度mask导致除零NaN
-    ref_log_probs = (ref_log_probs * mask).sum(dim=1) / seq_lengths.squeeze()
-    policy_log_probs = (policy_log_probs * mask).sum(dim=1) / seq_lengths.squeeze()
+    ref_log_probs = (ref_log_probs * mask).sum(dim=1)
+    policy_log_probs = (policy_log_probs * mask).sum(dim=1)
 
     # 将 chosen 和 rejected 数据分开
     batch_size = ref_log_probs.shape[0]
@@ -53,8 +51,10 @@ def dpo_loss(ref_log_probs, policy_log_probs, mask, beta):
 
 def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=None, beta=0.1):
     start_time = time.time()
-    
+    last_step = start_step
+
     for step, batch in enumerate(loader, start=start_step + 1):
+        last_step = step
         x_chosen = batch['x_chosen'].to(args.device)
         x_rejected = batch['x_rejected'].to(args.device)
         y_chosen = batch['y_chosen'].to(args.device)
@@ -98,7 +98,7 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
             current_dpo_loss = dpo_loss_val.item()
             current_aux_loss = outputs.aux_loss.item()
             current_lr = optimizer.param_groups[-1]['lr']
-            eta_min = spend_time / step * iters // 60 - spend_time // 60
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
             
             Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, dpo_loss: {current_dpo_loss:.4f}, aux_loss: {current_aux_loss:.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min')
             
@@ -119,6 +119,13 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
         del x_chosen, x_rejected, y_chosen, y_rejected, mask_chosen, mask_rejected, x, y, mask
         del ref_outputs, ref_logits, ref_log_probs, outputs, logits, policy_log_probs, loss
 
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind DPO (Direct Preference Optimization)")
@@ -134,14 +141,14 @@ if __name__ == "__main__":
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
     parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=1024, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
     parser.add_argument("--data_path", type=str, default="../dataset/dpo.jsonl", help="DPO训练数据路径")
     parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
-    parser.add_argument('--beta', default=0.1, type=float, help="DPO中的beta参数")
+    parser.add_argument('--beta', default=0.15, type=float, help="DPO中的beta参数")
     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=是）")
@@ -173,9 +180,6 @@ if __name__ == "__main__":
     
     # ========== 5. 定义模型和参考模型 ==========
     model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
     Logger(f'策略模型总参数量：{sum(p.numel() for p in model.parameters()) / 1e6:.3f} M')
     # 初始化参考模型（ref_model冻结）
     ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
@@ -197,7 +201,10 @@ if __name__ == "__main__":
         start_epoch = ckp_data['epoch']
         start_step = ckp_data.get('step', 0)
     
-    # ========== 7. DDP包模型 ==========
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
     if dist.is_initialized():
         model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         model = DistributedDataParallel(model, device_ids=[local_rank])

trainer/train_full_sft.py

@@ -22,9 +22,11 @@ warnings.filterwarnings('ignore')
 
 def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
     start_time = time.time()
+    last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
+        last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -51,7 +53,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
             current_logits_loss = current_loss - current_aux_loss
             current_lr = optimizer.param_groups[-1]['lr']
-            eta_min = spend_time / step * iters // 60 - spend_time // 60
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
             Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
             if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
 
@@ -70,6 +72,13 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         del input_ids, labels, res, loss
 
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Full SFT")
@@ -77,7 +86,7 @@ if __name__ == "__main__":
     parser.add_argument('--save_weight', default='full_sft', type=str, help="保存权重的前缀名")
     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-6, help="初始学习率")
+    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="数据加载线程数")
@@ -85,11 +94,11 @@ if __name__ == "__main__":
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
     parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl", help="训练数据路径")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
     parser.add_argument('--from_weight', default='pretrain', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
@@ -123,9 +132,6 @@ if __name__ == "__main__":
     
     # ========== 5. 定义模型、数据、优化器 ==========
     model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
@@ -140,7 +146,10 @@ if __name__ == "__main__":
         start_epoch = ckp_data['epoch']
         start_step = ckp_data.get('step', 0)
     
-    # ========== 7. DDP包模型 ==========
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
     if dist.is_initialized():
         model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         model = DistributedDataParallel(model, device_ids=[local_rank])

trainer/train_grpo.py

@@ -5,10 +5,12 @@ __package__ = "trainer"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 import argparse
+import math
 import re
 import gc
 import warnings
 import torch
+import torch.nn.functional as F
 import torch.distributed as dist
 from transformers import AutoTokenizer
 from contextlib import nullcontext
@@ -19,47 +21,24 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
 from transformers import AutoModel
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import RLAIFDataset
-from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model
+from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
 
 warnings.filterwarnings('ignore')
 
 
-def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
-    """整合所有奖励函数计算总奖励"""
-    def reasoning_model_reward(rewards):
-        pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
-        pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
-        matches_pattern = [re.match(pattern, response, re.S) for response in responses]
-        matches_pattern2 = [re.match(pattern2, response, re.S) for response in responses]
+def rep_penalty(text, n=3, cap=0.5):
+    toks = re.findall(r"\w+|[^\w\s]", text.lower())
+    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
 
-        format_rewards = []
-        for match_pattern, match_pattern2 in zip(matches_pattern, matches_pattern2):
-            if match_pattern or match_pattern2:
-                format_rewards.append(0.5)
-            else:
-                format_rewards.append(0.0)
-        rewards += torch.tensor(format_rewards, device=args.device)
-
-        def mark_num(text):
-            reward = 0
-            if text.count("<think>") == 1: reward += 0.25
-            if text.count("</think>") == 1: reward += 0.25
-            if text.count("<answer>") == 1: reward += 0.25
-            if text.count("</answer>") == 1: reward += 0.25
-            return reward
-
-        mark_rewards = [mark_num(response) for response in responses]
-        rewards += torch.tensor(mark_rewards, device=args.device)
-        return rewards
 
+def calculate_rewards(prompts, responses, reward_model):
     rewards = torch.zeros(len(responses), device=args.device)
-    if args.reasoning == 1:
-        rewards = reasoning_model_reward(rewards)
 
     with torch.no_grad():
         reward_model_scores = []
         batch_size = len(prompts)
-        scale = 3.0
 
         for i in range(batch_size):
             for j in range(args.num_generations):
@@ -70,20 +49,16 @@ def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
                 pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
                 matches = re.findall(pattern, prompt, re.DOTALL)
                 messages = [{"role": role, "content": content.strip()} for role, content in matches]
+                answer = response
+                rewards[response_idx] += 0.5 if 20 <= len(response.strip()) <= 800 else -0.5
+                if '</think>' in response:
+                    thinking_content, answer_content = response.split('</think>', 1)
+                    rewards[response_idx] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5
+                    rewards[response_idx] += 0.25 if response.count('</think>') == 1 else -0.25
+                    answer = answer_content.strip()
+                rewards[response_idx] -= rep_penalty(answer)
 
-                tmp_chat = messages + [{"role": "assistant", "content": response}]
-                score = reward_model.get_score(reward_tokenizer, tmp_chat)
-                score = max(min(score, scale), -scale)
-
-                if args.reasoning == 1:
-                    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
-                    if answer_match:
-                        answer_content = answer_match.group(1).strip()
-                        tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
-                        answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
-                        answer_score = max(min(answer_score, scale), -scale)
-                        score = score * 0.4 + answer_score * 0.6
-
+                score = reward_model.get_score(messages, answer)
                 reward_model_scores.append(score)
 
         reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
@@ -92,49 +67,61 @@ def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
     return rewards
 
 
-def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokenizer, start_step=0, wandb=None):
+def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model, start_step=0, wandb=None, use_sglang=False):
     for step, batch in enumerate(loader, start=start_step + 1):
         prompts = batch['prompt']  # list[str], length B
         prompt_inputs = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False,
-                                  padding_side="left", add_special_tokens=False).to(args.device)  # input_ids: [B, P], attention_mask: [B, P]
+                                  padding_side="left", add_special_tokens=False).to(args.device)
         if args.max_seq_len:
             prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -args.max_seq_len:]
             prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -args.max_seq_len:]
 
-        with torch.no_grad():
-            # DDP 模型需要使用 .module 访问 generate 方法
-            model_for_gen = model.module if isinstance(model, DistributedDataParallel) else model
-            outputs = model_for_gen.generate(
-                **prompt_inputs, max_new_tokens=args.max_gen_len, do_sample=True, temperature=0.8,
-                num_return_sequences=args.num_generations, pad_token_id=tokenizer.pad_token_id)  # [B*num_gen, P+R]
-
-        completion_ids = outputs[:, prompt_inputs["input_ids"].size(1):]  # [B*num_gen, R]
-        
-        def get_per_token_logps(mdl, input_ids, n_keep):
-            input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
-            logits = mdl(input_ids, logits_to_keep=n_keep + 1).logits[:, :-1, :]
-            per_token_logps = []
-            for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
-                ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
-                per_token_logps.append(torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1))
-            return torch.stack(per_token_logps)
+        rollout_result = rollout_engine.rollout(
+            prompt_ids=prompt_inputs["input_ids"],
+            attention_mask=prompt_inputs["attention_mask"],
+            num_generations=args.num_generations,
+            max_new_tokens=args.max_gen_len,
+            temperature=0.8,
+        )
+        outputs = rollout_result.output_ids
+        completion_ids = rollout_result.completion_ids
+        completions = rollout_result.completions
+        old_per_token_logps = rollout_result.per_token_logps.to(args.device)
 
+        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
         with autocast_ctx:
-            per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1))  # [B*num_gen, R]
-            res = model(outputs) if lm_config.use_moe else None
-            aux_loss = res.aux_loss if res is not None else torch.tensor(0.0, device=args.device)
+            if use_sglang or lm_config.use_moe:
+                res = model_unwrapped(outputs)
+                aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
+                logits = res.logits[:, :-1, :]
+                per_token_logps = F.log_softmax(logits, dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1)[:, -completion_ids.size(1):]
+            else:
+                aux_loss = torch.tensor(0.0, device=args.device)
+                per_token_logps = rollout_result.per_token_logps
         
         with torch.no_grad():
-            ref_per_token_logps = get_per_token_logps(ref_model, outputs, completion_ids.size(1))  # [B*num_gen, R]
+            ref_per_token_logps = compute_per_token_logps(ref_model, outputs, completion_ids.size(1))
+        rewards = calculate_rewards(prompts, completions, reward_model).to(args.device)  # [B*num_gen]
 
-        completions = tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
-        rewards = calculate_rewards(prompts, completions, reward_model, reward_tokenizer).to(args.device)  # [B*num_gen]
+        if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+            for i in range(len(prompts)):
+                Logger(f"[DEBUG] step={step}, sample[{i}]")
+                Logger('-'*100)
+                Logger(f"{'=' * 30} [DEBUG] sample[{i}] CONTEXT_BEGIN {'=' * 30}")
+                Logger(prompts[i])
+                Logger(f"{'=' * 31} [DEBUG] sample[{i}] CONTEXT_END {'=' * 31}")
+                for j in range(args.num_generations):
+                    idx = i * args.num_generations + j
+                    Logger(f"{'=' * 28} [DEBUG] gen[{j}] RESPONSE_BEGIN {'=' * 28}")
+                    Logger(completions[idx])
+                    Logger(f"{'=' * 29} [DEBUG] gen[{j}] RESPONSE_END {'=' * 29}")
+                    Logger(f"[DEBUG] gen[{j}] reward={rewards[idx].item():.4f}")
+                Logger('='*100)
 
         grouped_rewards = rewards.view(-1, args.num_generations)  # [B, num_gen]
         mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)  # [B*num_gen]
         std_r = grouped_rewards.std(dim=1).repeat_interleave(args.num_generations)  # [B*num_gen]
-        advantages = torch.clamp((rewards - mean_r) / (std_r + 1e-4), -10, 10)
-        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)  # [B*num_gen]
+        advantages = (rewards - mean_r) / (std_r + 1e-4)  # [B*num_gen]
 
         is_eos = completion_ids == tokenizer.eos_token_id  # [B*num_gen, R]
         eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=args.device)
@@ -143,7 +130,15 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token
 
         kl_div = ref_per_token_logps - per_token_logps
         per_token_kl = torch.exp(kl_div) - kl_div - 1  # [B*num_gen, R]
-        per_token_loss = -(torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) - args.beta * per_token_kl)  # [B*num_gen, R]
+        ratio = torch.exp(per_token_logps - old_per_token_logps)  # [B*num_gen, R]
+        if args.loss_type == "cispo":
+            clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
+            per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
+        else:
+            clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
+            per_token_loss1 = ratio * advantages.unsqueeze(1)
+            per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
+            per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
         policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
         loss = (policy_loss + aux_loss) / args.accumulation_steps  # scalar
         loss.backward()
@@ -154,25 +149,31 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token
             optimizer.step()
             scheduler.step()
             optimizer.zero_grad()
+            if is_main_process() and step % args.save_interval == 0: rollout_engine.update_policy(model)
 
         if step % args.log_interval == 0 or step == iters:
             policy_loss_val = loss.item() * args.accumulation_steps
             current_aux_loss = aux_loss.item()
             avg_reward_val = rewards.mean().item()
             avg_len_val = completion_mask.sum(dim=1).float().mean().item()
+            kl_ref_val = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / completion_mask.sum().item()
+            advantages_mean_val = advantages.mean().item()
+            advantages_std_val = advantages.std().item()
             current_lr = optimizer.param_groups[0]['lr']
 
             Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
-                   f'Actor Loss: {policy_loss_val:.4f}, Aux Loss: {current_aux_loss:.4f}, Reward: {avg_reward_val:.4f}, '
-                   f'Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}')
+                   f'Reward: {avg_reward_val:.4f}, KL_ref: {kl_ref_val:.4f}, '
+                   f'Adv Std: {advantages_std_val:.4f}, Adv Mean: {advantages_mean_val:.4f}, '
+                   f'Actor Loss: {policy_loss_val:.4f}, Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}')
 
             if wandb and is_main_process():
                 wandb.log({
-                    "policy_loss": policy_loss_val,
-                    "aux_loss": current_aux_loss,
                     "reward": avg_reward_val,
+                    "kl_ref": kl_ref_val,
+                    "advantages_std": advantages_std_val,
+                    "advantages_mean": advantages_mean_val,
+                    "policy_loss": policy_loss_val,
                     "avg_response_len": avg_len_val,
-                    "advantages_mean": advantages.mean().item(),
                     "learning_rate": current_lr
                 })
 
@@ -189,6 +190,14 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token
             model.train()
             del state_dict
 
+    if step > start_step and step % args.accumulation_steps != 0:
+        if args.grad_clip > 0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+        if is_main_process() and step % args.save_interval == 0: rollout_engine.update_policy(model)
+
         del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
         del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
 
@@ -199,7 +208,7 @@ if __name__ == "__main__":
     parser.add_argument('--save_weight', default='grpo', type=str, help="保存权重的前缀名")
     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=8e-8, 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="混合精度类型")
     parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
@@ -207,20 +216,30 @@ if __name__ == "__main__":
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
     parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument('--max_seq_len', default=66, type=int, help="Prompt最大长度")
-    parser.add_argument("--max_gen_len", type=int, default=1536, help="生成的最大长度")
-    parser.add_argument("--data_path", type=str, default="../dataset/rlaif-mini.jsonl", help="RLAIF数据路径")
-    parser.add_argument("--num_generations", type=int, default=8, help="每个prompt生成的样本数")
-    parser.add_argument("--beta", type=float, default=0.02, help="KL惩罚系数")
-    parser.add_argument("--reasoning", type=int, default=1, choices=[0, 1], help='推理模型类型（0=普通模型，1=推理模型）')
+    parser.add_argument('--max_seq_len', default=768, type=int, help="Prompt最大长度")
+    parser.add_argument("--max_gen_len", type=int, default=1024, help="生成的最大长度")
+    parser.add_argument("--data_path", type=str, default="../dataset/rlaif.jsonl", help="RLAIF数据路径")
+    parser.add_argument("--num_generations", type=int, default=6, help="每个prompt生成的样本数")
+    parser.add_argument("--beta", type=float, default=0.1, help="KL惩罚系数")
+    parser.add_argument("--loss_type", type=str, default="cispo", choices=["grpo", "cispo"], help="loss类型")
+    parser.add_argument("--epsilon", type=float, default=0.2, help="GRPO的PPO clip epsilon")
+    parser.add_argument("--epsilon_high", type=float, default=5.0, help="epsilon上界")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
     parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-GRPO", help="wandb项目名")
     parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--debug_mode", action="store_true", help="是否打印训练调试采样")
+    parser.add_argument("--debug_interval", type=int, default=20, help="debug模式下每隔多少step打印一次采样")
+    parser.add_argument("--thinking_ratio", type=float, default=0.9, help="按概率开启thinking（0.0~1.0）")
+    parser.add_argument("--rollout_engine", type=str, default="sglang", choices=["torch", "sglang"], help="rollout引擎类型")
+    parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8996", 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共享存储路径")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
@@ -249,28 +268,32 @@ if __name__ == "__main__":
         wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
     
     # ========== 5. 初始化模型和数据 ==========
-    base_weight = "reason" if args.reasoning == 1 else "full_sft"
+    base_weight = args.from_weight
     # Policy模型
     model, tokenizer = init_model(lm_config, base_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
     # Reference模型
     ref_model, _ = init_model(lm_config, base_weight, device=args.device)
     ref_model = ref_model.eval().requires_grad_(False)
     # Reward模型
-    reward_model = AutoModel.from_pretrained(
-        args.reward_model_path, torch_dtype=torch.float16, trust_remote_code=True
+    reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
+    # Rollout引擎（可插拔替换，只负责 policy 推理）
+    rollout_engine = create_rollout_engine(
+        engine_type=args.rollout_engine,
+        policy_model=model,
+        tokenizer=tokenizer,
+        device=args.device,
+        autocast_ctx=autocast_ctx,
+        sglang_base_url=args.sglang_base_url,
+        sglang_model_path=args.sglang_model_path,
+        sglang_shared_path=args.sglang_shared_path,
     )
-    reward_model = reward_model.to(args.device).eval().requires_grad_(False)
-    reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True)
     # 数据和优化器
-    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
+    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len, thinking_ratio=args.thinking_ratio)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
     loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
     iters = len(loader_for_count)
-    total_optimizer_steps = (iters // args.accumulation_steps) * args.epochs
+    total_optimizer_steps = math.ceil(iters / args.accumulation_steps) * args.epochs
     scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
     
     # ========== 6. 从ckp恢复状态 ==========
@@ -282,10 +305,15 @@ if __name__ == "__main__":
         start_epoch = ckp_data['epoch']
         start_step = ckp_data.get('step', 0)
     
-    # ========== 7. DDP包模型 ==========
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+        rollout_engine.update_policy(model)
     if dist.is_initialized():
         model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         model = DistributedDataParallel(model, device_ids=[local_rank])
+    if is_main_process(): rollout_engine.update_policy(model)
     
     # ========== 8. 开始训练 ==========
     for epoch in range(start_epoch, args.epochs):
@@ -296,9 +324,9 @@ if __name__ == "__main__":
         loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
         if skip > 0: 
             Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
-            grpo_train_epoch(epoch, loader, len(loader) + skip, ref_model, reward_model, reward_tokenizer, start_step, wandb)
+            grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
         else:
-            grpo_train_epoch(epoch, loader, len(loader), ref_model, reward_model, reward_tokenizer, 0, wandb)
+            grpo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
     
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized(): dist.destroy_process_group()

trainer/train_lora.py

@@ -23,9 +23,11 @@ warnings.filterwarnings('ignore')
 
 def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
     start_time = time.time()
+    last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
+        last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -50,7 +52,7 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
             current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
             current_logits_loss = current_loss - current_aux_loss
             current_lr = optimizer.param_groups[-1]['lr']
-            eta_min = spend_time / step * iters // 60 - spend_time // 60
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
             Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
             if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
 
@@ -64,12 +66,18 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
 
         del input_ids, labels, res, loss
 
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning")
-    parser.add_argument("--save_dir", type=str, default="../out/lora", help="模型保存目录")
-    parser.add_argument("--lora_name", type=str, default="lora_identity", help="LoRA权重名称(如lora_identity/lora_medical等)")
-    parser.add_argument("--epochs", type=int, default=50, help="训练轮数")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument("--lora_name", type=str, default="lora_medical", help="LoRA权重名称(如lora_identity/lora_medical等)")
+    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="训练设备")
@@ -79,11 +87,11 @@ if __name__ == "__main__":
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=10, help="日志打印间隔")
     parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument("--data_path", type=str, default="../dataset/lora_identity.jsonl", help="LoRA训练数据路径")
+    parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl", help="LoRA训练数据路径")
     parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练，默认full_sft")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
@@ -117,9 +125,6 @@ if __name__ == "__main__":
     
     # ========== 5. 定义模型、应用LoRA、冻结非LoRA参数 ==========
     model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
     apply_lora(model)
     
     # 统计参数
@@ -153,7 +158,10 @@ if __name__ == "__main__":
         start_epoch = ckp_data['epoch']
         start_step = ckp_data.get('step', 0)
     
-    # ========== 8. DDP包模型 ==========
+    # ========== 8. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
     if dist.is_initialized():
         model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         model = DistributedDataParallel(model, device_ids=[local_rank])

trainer/train_ppo.py

@@ -5,6 +5,7 @@ __package__ = "trainer"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 import argparse
+import math
 import re
 import warnings
 import torch
@@ -17,14 +18,20 @@ from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from torch.nn.utils import clip_grad_norm_
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from transformers import AutoModel
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import RLAIFDataset
-from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model
+from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from trainer.rollout_engine import create_rollout_engine
 
 warnings.filterwarnings('ignore')
 
 
+def rep_penalty(text, n=3, cap=0.5):
+    toks = re.findall(r"\w+|[^\w\s]", text.lower())
+    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
+
+
 # 自定义的Critic模型，继承自MiniMindLM
 class CriticModel(MiniMindForCausalLM):
     def __init__(self, params):
@@ -41,73 +48,25 @@ class CriticModel(MiniMindForCausalLM):
         return values
 
 
-def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
-    """整合所有奖励函数计算总奖励"""
-    def reasoning_model_reward(rewards):
-        # 1. 格式奖励（仅针对训练推理模型时使用）
-        pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
-        pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
-
-        matches_pattern = [re.match(pattern, response, re.S) for response in responses]
-        matches_pattern2 = [re.match(pattern2, response, re.S) for response in responses]
-
-        format_rewards = []
-        for match_pattern, match_pattern2 in zip(matches_pattern, matches_pattern2):
-            if match_pattern:
-                format_rewards.append(0.5)
-            elif match_pattern2:
-                format_rewards.append(0.5)
-            else:
-                format_rewards.append(0.0)
-        rewards += torch.tensor(format_rewards, device=args.device)
-
-        # 2. 标记奖励（防止严格奖励稀疏，仅针对训练推理模型时使用）
-        def mark_num(text):
-            reward = 0
-            if text.count("<think>") == 1:
-                reward += 0.25
-            if text.count("</think>") == 1:
-                reward += 0.25
-            if text.count("<answer>") == 1:
-                reward += 0.25
-            if text.count("</answer>") == 1:
-                reward += 0.25
-            return reward
-
-        mark_rewards = [mark_num(response) for response in responses]
-        rewards += torch.tensor(mark_rewards, device=args.device)
-        return rewards
-
+def calculate_rewards(prompts, responses, reward_model):
     rewards = torch.zeros(len(responses), device=args.device)
 
-    # 格式奖励
-    if args.reasoning == 1:
-        rewards = reasoning_model_reward(rewards)
-
-    # 使用reward model计算整个response的奖励
     with torch.no_grad():
         reward_model_scores = []
-        for prompt, response in zip(prompts, responses):
+        for i, (prompt, response) in enumerate(zip(prompts, responses)):
             pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
             matches = re.findall(pattern, prompt, re.DOTALL)
             messages = [{"role": role, "content": content.strip()} for role, content in matches]
+            answer = response
+            rewards[i] += 0.5 if 20 <= len(response.strip()) <= 800 else -0.5
+            if '</think>' in response:
+                thinking_content, answer_content = response.split('</think>', 1)
+                rewards[i] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5
+                rewards[i] += 0.25 if response.count('</think>') == 1 else -0.25
+                answer = answer_content.strip()
+            rewards[i] -= rep_penalty(answer)
 
-            tmp_chat = messages + [{"role": "assistant", "content": response}]
-            score = reward_model.get_score(reward_tokenizer, tmp_chat)
-
-            scale = 3.0
-            score = max(min(score, scale), -scale)
-
-            # 当args.reasoning=1时，额外计算<answer>内容的奖励
-            if args.reasoning == 1:
-                answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
-                if answer_match:
-                    answer_content = answer_match.group(1).strip()
-                    # 对answer内容单独计算reward
-                    tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
-                    answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
-                    answer_score = max(min(answer_score, scale), -scale)
-                    score = score * 0.4 + answer_score * 0.6
+            score = reward_model.get_score(messages, answer)
             reward_model_scores.append(score)
 
         reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
@@ -116,65 +75,172 @@ def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
     return rewards
 
 
-def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_scheduler, critic_scheduler, reward_model, reward_tokenizer, start_step=0, wandb=None):
+def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step=0, wandb=None, use_sglang=False):
     actor_model.train()
     critic_model.train()
+    grad_accum_step = 0
 
     for step, batch in enumerate(loader, start=start_step + 1):
         prompts = batch["prompt"]  # list[str], length B
-        enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, 
-                       max_length=args.max_seq_len, padding_side="left").to(args.device)  # input_ids: [B, P], attention_mask: [B, P]
+        enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, max_length=args.max_seq_len,
+                        padding_side="left").to(args.device)  # input_ids: [B, P], attention_mask: [B, P]
         prompt_length = enc.input_ids.shape[1]
 
-        with torch.no_grad():
-            # DDP 模型需要使用 .module 访问 generate 方法
-            model_for_gen = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
-            gen_out = model_for_gen.generate(
-                input_ids=enc.input_ids, attention_mask=enc.attention_mask,
-                max_new_tokens=args.max_gen_len, do_sample=True, temperature=0.8,
-                pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)  # [B, P+R]
+        rollout_result = rollout_engine.rollout(
+            prompt_ids=enc.input_ids,
+            attention_mask=enc.attention_mask,
+            num_generations=1,
+            max_new_tokens=args.max_gen_len,
+            temperature=0.8,
+        )
+        gen_out = rollout_result.output_ids
+        responses_text = rollout_result.completions
+        rewards = calculate_rewards(prompts, responses_text, reward_model)  # [B]
 
-        responses_text = [tokenizer.decode(gen_out[i, prompt_length:], skip_special_tokens=True) for i in range(len(prompts))]
-        rewards = calculate_rewards(prompts, responses_text, reward_model, reward_tokenizer)  # [B]
+        if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+            for i in range(len(prompts)):
+                Logger(f"[DEBUG] step={step}, sample[{i}]")
+                Logger('-'*100)
+                Logger(f"{'=' * 30} [DEBUG] sample[{i}] CONTEXT_BEGIN {'=' * 30}")
+                Logger(prompts[i])
+                Logger(f"{'=' * 31} [DEBUG] sample[{i}] CONTEXT_END {'=' * 31}")
+                Logger(f"[DEBUG] prompt_len={prompt_length}, response_len={len(responses_text[i])}")
+                Logger(f"{'=' * 28} [DEBUG] sample[{i}] RESPONSE_BEGIN {'=' * 28}")
+                Logger(responses_text[i])
+                Logger(f"{'=' * 29} [DEBUG] sample[{i}] RESPONSE_END {'=' * 29}")
+                Logger(f"[DEBUG] reward={rewards[i].item():.4f}")
+                Logger('='*100)
 
         full_mask = (gen_out != tokenizer.pad_token_id).long()  # [B, P+R]
-        values_seq = critic_model(input_ids=gen_out, attention_mask=full_mask)  # [B, P+R]
-        last_indices = (full_mask * torch.arange(full_mask.size(1), device=gen_out.device)).argmax(dim=1)
-        values = values_seq[torch.arange(values_seq.size(0), device=values_seq.device), last_indices]  # [B]
-        advantages = rewards - values.detach()  # [B]
-
-        with autocast_ctx:
-            res = actor_model(input_ids=gen_out, attention_mask=full_mask)
-            logits = res.logits  # [B, P+R, V]
-            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
-        
         labels = gen_out[:, 1:].clone()  # [B, P+R-1]
-        logp_tokens = F.log_softmax(logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)  # [B, P+R-1]
-        seq_len = gen_out.size(1) - 1
-        resp_mask = torch.arange(seq_len, device=gen_out.device).unsqueeze(0) >= prompt_length - 1
-        final_mask = resp_mask & (~labels.eq(tokenizer.pad_token_id))  # [B, P+R-1]
-        actor_logp = (logp_tokens * final_mask).sum(dim=1)  # [B]
+        seq_len, resp_start = gen_out.size(1) - 1, prompt_length - 1
+        resp_mask = torch.arange(seq_len, device=gen_out.device).unsqueeze(0) >= resp_start
+        final_mask = (resp_mask & (~labels.eq(tokenizer.pad_token_id))).float()  # [B, P+R-1]
+        B = len(prompts)
+        resp_labels = labels[:, resp_start:]  # [B, R]
+        resp_idx = torch.arange(resp_labels.size(1), device=gen_out.device).unsqueeze(0)
+        resp_pad_mask = ~resp_labels.eq(tokenizer.pad_token_id)
+        resp_lengths = resp_pad_mask.sum(dim=1); eos_mask = resp_labels.eq(tokenizer.eos_token_id) & resp_pad_mask
+        has_eos = eos_mask.any(dim=1); eos_pos = torch.argmax(eos_mask.int(), dim=1)
+        resp_lengths = torch.where(has_eos, eos_pos + 1, resp_lengths).long().clamp(min=1)
+        resp_policy_mask = ((resp_idx < resp_lengths.unsqueeze(1)) & resp_pad_mask).float()
+        resp_value_mask = resp_policy_mask.clone()
 
-        with torch.no_grad():
-            old_logits = old_actor_model(input_ids=gen_out, attention_mask=full_mask).logits  # [B, P+R, V]
-            old_logp_tokens = F.log_softmax(old_logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)  # [B, P+R-1]
-            old_logp = (old_logp_tokens * final_mask).sum(dim=1)  # [B]
+        with torch.no_grad():  # Rollout阶段只需推理获取old_logp和old_values，切断梯度省显存
+            critic_for_rollout = critic_model.module if isinstance(critic_model, DistributedDataParallel) else critic_model
+            values_seq = critic_for_rollout(input_ids=gen_out, attention_mask=full_mask)
+            old_resp_values = values_seq[:, resp_start:-1] * resp_value_mask
             
-            ref_logits = ref_model(input_ids=gen_out, attention_mask=full_mask).logits  # [B, P+R, V]
-            ref_logp_tokens = F.log_softmax(ref_logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)  # [B, P+R-1]
-            ref_logp = (ref_logp_tokens * final_mask).sum(dim=1)  # [B]
+            actor_for_rollout = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
+            with autocast_ctx:
+                logits = actor_for_rollout(input_ids=gen_out, attention_mask=full_mask).logits
+            
+            old_resp_logp = F.log_softmax(logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)[:, resp_start:]
+            
+            ref_logp_all = F.log_softmax(ref_model(input_ids=gen_out, attention_mask=full_mask).logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)
+            ref_resp_logp = ref_logp_all[:, resp_start:]
+            token_rewards = torch.zeros_like(old_resp_logp)
+            last_idx = resp_lengths - 1  # [B]
+            token_rewards[torch.arange(B, device=args.device), last_idx] += rewards  # 末尾加外部奖励
 
-        kl = (actor_logp - old_logp).mean()  # scalar
-        kl_ref = (actor_logp - ref_logp).mean()  # scalar
-        ratio = torch.exp(actor_logp - old_logp)  # [B]
-        surr1 = ratio * advantages  # [B]
-        surr2 = torch.clamp(ratio, 1.0 - args.clip_epsilon, 1.0 + args.clip_epsilon) * advantages  # [B]
-        policy_loss = -torch.min(surr1, surr2).mean()  # scalar
-        value_loss = F.mse_loss(values, rewards)  # scalar
-        loss = (policy_loss + args.vf_coef * value_loss + args.kl_coef * kl_ref + aux_loss) / args.accumulation_steps  # scalar
-        loss.backward()
+            gen_len = old_resp_values.size(1); lastgaelam = torch.zeros(B, device=args.device); advs_rev = []
+            for t in reversed(range(gen_len)):
+                nv = old_resp_values[:, t + 1] if t < gen_len - 1 else 0.0
+                delta = token_rewards[:, t] + args.gamma * nv - old_resp_values[:, t]
+                lastgaelam = delta + args.gamma * args.lam * lastgaelam
+                advs_rev.append(lastgaelam)
+            advantages = torch.stack(advs_rev[::-1], dim=1)  # [B, R]
+            returns = advantages + old_resp_values  # [B, R]
 
-        if step % args.accumulation_steps == 0:
+            adv_mean = (advantages * resp_policy_mask).sum() / resp_policy_mask.sum().clamp(min=1)
+            adv_var = ((advantages - adv_mean) ** 2 * resp_policy_mask).sum() / resp_policy_mask.sum().clamp(min=1)
+            advantages = (advantages - adv_mean) * torch.rsqrt(adv_var + 1e-8) * resp_policy_mask
+
+        mb_size = max(1, min(args.mini_batch_size, B))
+        stop_ppo = False
+        policy_loss_sum = 0.0
+        value_loss_sum = 0.0
+        kl_sum = 0.0
+        kl_ref_sum = 0.0
+        clipfrac_sum = 0.0
+        aux_loss_sum = 0.0
+        log_count = 0
+        actor_unwrapped = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
+        critic_unwrapped = critic_model.module if isinstance(critic_model, DistributedDataParallel) else critic_model
+        for ppo_epoch in range(args.ppo_update_iters):
+            if stop_ppo:
+                break
+            b_inds = torch.randperm(B, device=args.device)
+            for i in range(0, B, mb_size):
+                inds = b_inds[i:i + mb_size]
+                
+                mb_values_seq = critic_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds])
+                mb_resp_values = mb_values_seq[:, resp_start:-1]
+
+                with autocast_ctx:
+                    res = actor_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds])
+                    aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
+
+                mb_logp_all = F.log_softmax(res.logits[:, :-1], dim=-1).gather(2, labels[inds].unsqueeze(-1)).squeeze(-1)
+                mb_resp_logp = mb_logp_all[:, resp_start:]
+                
+                log_ratio = mb_resp_logp - old_resp_logp[inds]
+                approx_kl = (0.5 * (log_ratio ** 2) * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
+                
+                # 同步各卡的 approx_kl，防止某卡 break 而其它卡继续导致 DDP 死锁
+                approx_kl_val = approx_kl.detach().clone()
+                if dist.is_initialized():
+                    dist.all_reduce(approx_kl_val, op=dist.ReduceOp.AVG)
+                    
+                if approx_kl_val > args.early_stop_kl:
+                    stop_ppo = True
+                
+                ratio = torch.exp(log_ratio)
+                clipfrac = ((((ratio - 1.0).abs() > args.clip_epsilon).float() * resp_policy_mask[inds]).sum()
+                            / resp_policy_mask[inds].sum().clamp(min=1))
+                kl_ref_penalty = ((torch.exp(ref_resp_logp[inds] - mb_resp_logp) - (ref_resp_logp[inds] - mb_resp_logp) - 1.0)
+                                  * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
+                policy_loss = ((torch.max(-advantages[inds] * ratio,
+                                          -advantages[inds] * torch.clamp(ratio, 1.0 - args.clip_epsilon, 1.0 + args.clip_epsilon))
+                               * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
+                               + args.kl_coef * kl_ref_penalty)
+                value_loss = 0.5 * (torch.max((mb_resp_values - returns[inds]) ** 2,
+                                              (torch.clamp(mb_resp_values, old_resp_values[inds] - args.cliprange_value,
+                                                           old_resp_values[inds] + args.cliprange_value) - returns[inds]) ** 2)
+                                    * resp_value_mask[inds]).sum() / resp_value_mask[inds].sum().clamp(min=1)
+
+                kl = approx_kl_val
+                kl_ref = kl_ref_penalty.detach()
+
+                # 早停时必须保证 forward-backward 闭环，故只截断 loss 不中断 DDP 通信
+                if stop_ppo:
+                    loss = (policy_loss + args.vf_coef * value_loss + aux_loss) * 0.0
+                else:
+                    loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps
+                
+                loss.backward()
+
+                policy_loss_sum += policy_loss.item()
+                value_loss_sum += value_loss.item()
+                kl_sum += kl.item()
+                kl_ref_sum += kl_ref.item()
+                clipfrac_sum += clipfrac.item()
+                aux_loss_sum += aux_loss.item()
+                log_count += 1
+
+                grad_accum_step += 1
+
+                if grad_accum_step % args.accumulation_steps == 0:
+                    clip_grad_norm_(actor_model.parameters(), args.grad_clip)
+                    clip_grad_norm_(critic_model.parameters(), args.grad_clip)
+                    actor_optimizer.step()
+                    critic_optimizer.step()
+                    actor_scheduler.step()
+                    critic_scheduler.step()
+                    actor_optimizer.zero_grad()
+                    critic_optimizer.zero_grad()
+
+        if grad_accum_step % args.accumulation_steps != 0:
             clip_grad_norm_(actor_model.parameters(), args.grad_clip)
             clip_grad_norm_(critic_model.parameters(), args.grad_clip)
             actor_optimizer.step()
@@ -183,49 +249,35 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche
             critic_scheduler.step()
             actor_optimizer.zero_grad()
             critic_optimizer.zero_grad()
+        
+        if is_main_process() and step % args.save_interval == 0: rollout_engine.update_policy(actor_model)
 
         if is_main_process():
-            response_ids = gen_out[:, enc.input_ids.shape[1]:]
-            is_eos = (response_ids == tokenizer.eos_token_id)
-            eos_indices = torch.argmax(is_eos.int(), dim=1)
-            has_eos = is_eos.any(dim=1)
-            lengths = torch.where(has_eos, eos_indices + 1, torch.tensor(response_ids.shape[1], device=is_eos.device))
-            avg_len = lengths.float().mean()
-
-            actor_loss_val = policy_loss.item()
-            critic_loss_val = value_loss.item()
-            current_aux_loss = aux_loss.item()
+            critic_loss_val = value_loss_sum / max(log_count, 1)
             reward_val = rewards.mean().item()
-            kl_val = kl.item()
-            kl_ref_val = kl_ref.item()
-            avg_len_val = avg_len.item()
-            actor_lr = actor_optimizer.param_groups[0]['lr']
-            critic_lr = critic_optimizer.param_groups[0]['lr']
+            approx_kl_val = kl_sum / max(log_count, 1)
+            kl_ref_val = kl_ref_sum / max(log_count, 1)
+            clipfrac_val = clipfrac_sum / max(log_count, 1)
+            avg_len_val = resp_lengths.float().mean().item()
+            actor_lr, critic_lr = actor_optimizer.param_groups[0]['lr'], critic_optimizer.param_groups[0]['lr']
 
             if wandb is not None:
                 wandb.log({
-                    "actor_loss": actor_loss_val,
-                    "critic_loss": critic_loss_val,
-                    "aux_loss": current_aux_loss,
                     "reward": reward_val,
-                    "kl": kl_val,
                     "kl_ref": kl_ref_val,
+                    "approx_kl": approx_kl_val,
+                    "clipfrac": clipfrac_val,
+                    "critic_loss": critic_loss_val,
                     "avg_response_len": avg_len_val,
                     "actor_lr": actor_lr,
+                    "critic_lr": critic_lr,
                 })
 
             Logger(f"Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), "
-                   f"Actor Loss: {actor_loss_val:.4f}, Critic Loss: {critic_loss_val:.4f}, Aux Loss: {current_aux_loss:.4f}, "
-                   f"Reward: {reward_val:.4f}, KL: {kl_val:.4f}, KL_ref: {kl_ref_val:.4f}, "
+                   f"Reward: {reward_val:.4f}, KL_ref: {kl_ref_val:.4f}, Approx KL: {approx_kl_val:.4f}, "
+                   f"ClipFrac: {clipfrac_val:.4f}, Critic Loss: {critic_loss_val:.4f}, "
                    f"Avg Response Len: {avg_len_val:.2f}, Actor LR: {actor_lr:.8f}, Critic LR: {critic_lr:.8f}")
 
-        if (step + 1) % args.update_old_actor_freq == 0:
-            raw_actor = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
-            raw_actor = getattr(raw_actor, '_orig_mod', raw_actor)
-            state_dict = raw_actor.state_dict()
-            old_actor_model.load_state_dict({k: v.detach().cpu() for k, v in state_dict.items()})
-            old_actor_model.to(args.device)
-
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             actor_model.eval()
             moe_suffix = '_moe' if lm_config.use_moe else ''
@@ -243,9 +295,9 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche
             actor_model.train()
             del actor_state
 
-        del enc, gen_out, responses_text, rewards, full_mask, values_seq, values, advantages
-        del logits, labels, logp_tokens, final_mask, actor_logp, old_logits, old_logp, ref_logits, ref_logp
-        del kl, kl_ref, ratio, surr1, surr2, policy_loss, value_loss, loss
+        del enc, gen_out, responses_text, rewards, full_mask, values_seq, advantages
+        del logits, labels, final_mask, resp_labels, resp_idx, resp_pad_mask, eos_mask, has_eos, eos_pos, resp_lengths, resp_policy_mask, resp_value_mask, old_resp_logp, ref_logp_all, ref_resp_logp
+        del kl, kl_ref, policy_loss, value_loss, loss, token_rewards, returns, old_resp_values
 
 
 if __name__ == "__main__":
@@ -254,8 +306,8 @@ if __name__ == "__main__":
     parser.add_argument('--save_weight', default='ppo_actor', type=str, help="保存权重的前缀名")
     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=8e-8, help="Actor学习率")
-    parser.add_argument("--critic_learning_rate", type=float, default=8e-8, help="Critic学习率")
+    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="数据加载线程数")
@@ -263,22 +315,34 @@ if __name__ == "__main__":
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
     parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument('--max_seq_len', default=66, type=int, help="Prompt最大长度")
-    parser.add_argument("--max_gen_len", type=int, default=1536, help="生成的最大长度")
-    parser.add_argument("--data_path", type=str, default="../dataset/rlaif-mini.jsonl", help="RLAIF数据路径")
-    parser.add_argument("--clip_epsilon", type=float, default=0.1, help="PPO裁剪参数")
+    parser.add_argument('--max_seq_len', default=768, type=int, help="Prompt最大长度")
+    parser.add_argument("--max_gen_len", type=int, default=1024, help="生成的最大长度")
+    parser.add_argument("--data_path", type=str, default="../dataset/rlaif.jsonl", help="RLAIF数据路径")
+    parser.add_argument("--clip_epsilon", type=float, default=0.2, help="PPO裁剪参数")
     parser.add_argument("--vf_coef", type=float, default=0.5, help="Value function系数")
     parser.add_argument("--kl_coef", type=float, default=0.02, help="KL散度惩罚系数")
-    parser.add_argument("--reasoning", type=int, default=1, choices=[0, 1], help='推理模型类型（0=普通模型，1=推理模型）')
-    parser.add_argument("--update_old_actor_freq", type=int, default=4, help="更新old_actor_model的频率")
+    parser.add_argument("--gamma", type=float, default=1.0, help="GAE折扣因子")
+    parser.add_argument("--lam", type=float, default=0.95, help="GAE lambda参数")
+    parser.add_argument("--cliprange_value", type=float, default=0.2, help="Value function裁剪范围")
+    parser.add_argument("--ppo_update_iters", type=int, default=2, help="同一批rollout重复更新次数")
+    parser.add_argument("--early_stop_kl", type=float, default=0.25, help="PPO early stop 的 KL 阈值")
+    parser.add_argument("--mini_batch_size", type=int, default=2, help="PPO每次更新的minibatch大小")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
     parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-PPO", help="wandb项目名")
     parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--debug_mode", action="store_true", help="是否打印训练调试采样")
+    parser.add_argument("--debug_interval", type=int, default=20, help="debug模式下每隔多少step打印一次采样")
+    parser.add_argument("--thinking_ratio", type=float, default=0.9, help="按概率开启thinking（0.0~1.0）")
+    parser.add_argument("--rollout_engine", type=str, default="sglang", choices=["torch", "sglang"], help="rollout引擎类型")
+    parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8997", 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共享存储路径")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
@@ -306,43 +370,40 @@ if __name__ == "__main__":
         wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
     
     # ========== 5. 初始化模型和数据 ==========
-    base_weight = "reason" if args.reasoning == 1 else "full_sft"
+    base_weight = args.from_weight
     # Actor模型
     actor_model, tokenizer = init_model(lm_config, base_weight, device=args.device)
-    if args.use_compile == 1:
-        actor_model = torch.compile(actor_model)
-        Logger('torch.compile enabled')
-    # Old Actor模型
-    old_actor_model, _ = init_model(lm_config, base_weight, device=args.device)
-    old_actor_model = old_actor_model.eval().requires_grad_(False)
-    # Reference模型
     ref_model, _ = init_model(lm_config, base_weight, device=args.device)
     ref_model = ref_model.eval().requires_grad_(False)
-    # Critic模型
     moe_suffix = '_moe' if lm_config.use_moe else ''
     ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
     state_dict = torch.load(ckp, map_location=args.device)
     critic_model = CriticModel(lm_config)
     critic_model.load_state_dict(state_dict, strict=False)
     critic_model = critic_model.to(args.device)
-    # Reward模型
-    reward_model = AutoModel.from_pretrained(
-        args.reward_model_path, torch_dtype=torch.float16, trust_remote_code=True
+    reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
+    # Rollout引擎
+    rollout_engine = create_rollout_engine(
+        engine_type=args.rollout_engine,
+        policy_model=actor_model,
+        tokenizer=tokenizer,
+        device=args.device,
+        autocast_ctx=autocast_ctx,
+        sglang_base_url=args.sglang_base_url,
+        sglang_model_path=args.sglang_model_path,
+        sglang_shared_path=args.sglang_shared_path,
     )
-    reward_model = reward_model.to(args.device).eval().requires_grad_(False)
-    reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True)
-    # 数据和优化器
-    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=(args.max_seq_len + args.max_gen_len))
+    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=(args.max_seq_len + args.max_gen_len), thinking_ratio=args.thinking_ratio)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     actor_optimizer = optim.AdamW(actor_model.parameters(), lr=args.learning_rate)
     critic_optimizer = optim.AdamW(critic_model.parameters(), lr=args.critic_learning_rate)
     loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
     iters = len(loader_for_count)
-    total_optimizer_steps = (iters // args.accumulation_steps) * args.epochs
+    mb_factor = max(1, math.ceil(args.batch_size / args.mini_batch_size))
+    total_optimizer_steps = math.ceil(iters * args.epochs * args.ppo_update_iters * mb_factor / args.accumulation_steps)
     actor_scheduler = CosineAnnealingLR(actor_optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
     critic_scheduler = CosineAnnealingLR(critic_optimizer, T_max=total_optimizer_steps, eta_min=args.critic_learning_rate / 10)
-    
-    # ========== 6. 从ckp恢复状态 ==========
+
     start_epoch, start_step = 0, 0
     if ckp_data:
         actor_model.load_state_dict(ckp_data['model'])
@@ -354,13 +415,17 @@ if __name__ == "__main__":
         start_epoch = ckp_data['epoch']
         start_step = ckp_data.get('step', 0)
     
-    # ========== 7. DDP包模型 ==========
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        actor_model = torch.compile(actor_model)
+        Logger('torch.compile enabled')
+        rollout_engine.update_policy(actor_model)
     if dist.is_initialized():
         actor_model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         critic_model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
         critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
-        old_actor_model.to(args.device)
+    if is_main_process(): rollout_engine.update_policy(actor_model)
     
     # ========== 8. 开始训练 ==========
     for epoch in range(start_epoch, args.epochs):
@@ -371,11 +436,9 @@ if __name__ == "__main__":
         loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
         if skip > 0: 
             Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
-            ppo_train_epoch(epoch, loader, len(loader) + skip, old_actor_model, ref_model, 
-                           actor_scheduler, critic_scheduler, reward_model, reward_tokenizer, start_step, wandb)
+            ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
         else:
-            ppo_train_epoch(epoch, loader, len(loader), old_actor_model, ref_model, 
-                           actor_scheduler, critic_scheduler, reward_model, reward_tokenizer, 0, wandb)
+            ppo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
     
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized(): dist.destroy_process_group()

trainer/train_pretrain.py

@@ -22,9 +22,11 @@ warnings.filterwarnings('ignore')
 
 def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
     start_time = time.time()
+    last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
+        last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -51,7 +53,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
             current_logits_loss = current_loss - current_aux_loss
             current_lr = optimizer.param_groups[-1]['lr']
-            eta_min = spend_time / step * iters // 60 - spend_time // 60
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
             Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
             if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
 
@@ -69,12 +71,19 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         del input_ids, labels, res, loss
 
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Pretraining")
     parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
     parser.add_argument('--save_weight', default='pretrain', type=str, help="保存权重的前缀名")
-    parser.add_argument("--epochs", type=int, default=1, help="训练轮数（建议1轮zero或2-6轮充分训练）")
+    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="训练设备")
@@ -84,11 +93,11 @@ if __name__ == "__main__":
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
     parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl", help="预训练数据路径")
+    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl", help="预训练数据路径")
     parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练，为none则从头开始")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
@@ -122,9 +131,6 @@ if __name__ == "__main__":
     
     # ========== 5. 定义模型、数据、优化器 ==========
     model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
@@ -139,7 +145,10 @@ if __name__ == "__main__":
         start_epoch = ckp_data['epoch']
         start_step = ckp_data.get('step', 0)
     
-    # ========== 7. DDP包模型 ==========
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
     if dist.is_initialized():
         model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
         model = DistributedDataParallel(model, device_ids=[local_rank])

trainer/train_reason.py

@@ -1,179 +0,0 @@
-import os
-import sys
-
-__package__ = "trainer"
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
-import argparse
-import time
-import warnings
-import torch
-import torch.distributed as dist
-from contextlib import nullcontext
-from torch import optim, nn
-from torch.nn.parallel import DistributedDataParallel
-from torch.utils.data import DataLoader, DistributedSampler
-from model.model_minimind import MiniMindConfig
-from dataset.lm_dataset import SFTDataset
-from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
-
-warnings.filterwarnings('ignore')
-
-
-def train_epoch(epoch, loader, iters, tokenizer, lm_config, start_step=0, wandb=None):
-    start_of_think_ids = tokenizer('<think>').input_ids
-    end_of_think_ids = tokenizer('</think>').input_ids
-    start_of_answer_ids = tokenizer('<answer>').input_ids
-    end_of_answer_ids = tokenizer('</answer>').input_ids
-    loss_fct = nn.CrossEntropyLoss(reduction='none')
-    start_time = time.time()
-    
-    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
-        input_ids = input_ids.to(args.device)
-        labels = labels.to(args.device)
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
-        for param_group in optimizer.param_groups:
-            param_group['lr'] = lr
-
-        with autocast_ctx:
-            res = model(input_ids)
-            shift_logits = res.logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)).view(shift_labels.size())
-
-            loss_mask = (shift_labels != -100).float()
-            sp_ids = torch.isin(shift_labels.view(-1),
-                                torch.tensor(start_of_think_ids + end_of_think_ids
-                                             + start_of_answer_ids + end_of_answer_ids
-                                             ).to(args.device))
-            loss_mask_flat = loss_mask.view(-1)
-            loss_mask_sum = loss_mask_flat.sum()
-            loss_mask_flat[sp_ids] = 10
-            loss_mask = loss_mask_flat.view(shift_labels.size())
-            logits_loss = (loss * loss_mask).sum() / loss_mask_sum
-            loss = logits_loss + res.aux_loss
-            loss = loss / args.accumulation_steps
-
-        scaler.scale(loss).backward()
-
-        if step % args.accumulation_steps == 0:
-            scaler.unscale_(optimizer)
-            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-            scaler.step(optimizer)
-            scaler.update()
-            optimizer.zero_grad(set_to_none=True)
-
-        if step % args.log_interval == 0 or step == iters:
-            spend_time = time.time() - start_time
-            current_loss = loss.item() * args.accumulation_steps
-            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
-            current_logits_loss = logits_loss.item()
-            current_lr = optimizer.param_groups[-1]['lr']
-            eta_min = spend_time / step * iters // 60 - spend_time // 60
-            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
-            if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
-
-        if (step % args.save_interval == 0 or step == iters) and is_main_process():
-            model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
-            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
-            raw_model = getattr(raw_model, '_orig_mod', raw_model)
-            state_dict = raw_model.state_dict()
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
-            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
-            model.train()
-            del state_dict
-
-        del input_ids, labels, res, loss
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="MiniMind Reasoning Distillation")
-    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
-    parser.add_argument('--save_weight', default='reason', type=str, help="保存权重的前缀名")
-    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
-    parser.add_argument("--batch_size", type=int, default=8, help="batch size")
-    parser.add_argument("--learning_rate", type=float, default=1e-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="梯度累积步数")
-    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
-    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
-    parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
-    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
-    parser.add_argument('--max_seq_len', default=720, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
-    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument("--data_path", type=str, default="../dataset/r1_mix_1024.jsonl", help="推理蒸馏数据路径")
-    parser.add_argument('--from_weight', default='dpo', type=str, help="基于哪个权重训练，默认dpo")
-    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
-    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-Reasoning", help="wandb项目名")
-    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
-    args = parser.parse_args()
-
-    # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
-    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
-    
-    # ========== 2. 配置目录、模型参数、检查ckp ==========
-    os.makedirs(args.save_dir, exist_ok=True)
-    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, 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
-    
-    # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" 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)
-    
-    # ========== 4. 配wandb ==========
-    wandb = None
-    if args.use_wandb and is_main_process():
-        import swanlab as wandb
-        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
-        resume = 'must' if wandb_id else None
-        wandb_run_name = f"MiniMind-Reasoning-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LR-{args.learning_rate}"
-        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
-    
-    # ========== 5. 定义模型、数据、优化器 ==========
-    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
-    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
-    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
-    
-    # ========== 6. 从ckp恢复状态 ==========
-    start_epoch, start_step = 0, 0
-    if ckp_data:
-        model.load_state_dict(ckp_data['model'])
-        optimizer.load_state_dict(ckp_data['optimizer'])
-        scaler.load_state_dict(ckp_data['scaler'])
-        start_epoch = ckp_data['epoch']
-        start_step = ckp_data.get('step', 0)
-    
-    # ========== 7. DDP包模型 ==========
-    if dist.is_initialized():
-        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
-        model = DistributedDataParallel(model, device_ids=[local_rank])
-    
-    # ========== 8. 开始训练 ==========
-    for epoch in range(start_epoch, args.epochs):
-        train_sampler and train_sampler.set_epoch(epoch)
-        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
-        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
-        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
-        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
-        if skip > 0: 
-            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
-            train_epoch(epoch, loader, len(loader) + skip, tokenizer, lm_config, start_step, wandb)
-        else:
-            train_epoch(epoch, loader, len(loader), tokenizer, lm_config, 0, wandb)
-    
-    # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()

trainer/train_spo.py

@@ -1,354 +0,0 @@
-import os
-import sys
-
-__package__ = "trainer"
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
-import argparse
-import re
-import gc
-import warnings
-import torch
-import torch.distributed as dist
-from transformers import AutoTokenizer
-from contextlib import nullcontext
-from torch import optim
-from torch.nn.parallel import DistributedDataParallel
-from torch.utils.data import DataLoader, DistributedSampler
-from torch.optim.lr_scheduler import CosineAnnealingLR
-from transformers import AutoModel
-from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
-from dataset.lm_dataset import RLAIFDataset
-from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model
-
-warnings.filterwarnings('ignore')
-
-
-class AutoAdaptiveValueTracker:
-    """SPO自适应价值追踪器"""
-    def __init__(self, rho_mode='kl', rho_const=0.9, D_half=0.06, clip_lower=0.5, clip_upper=0.96):
-        self.rho_mode = rho_mode
-        self.rho_const = rho_const
-        self.D_half = D_half
-        self.clip_lower = clip_lower
-        self.clip_upper = clip_upper
-        N_init = 1.0 / (1.0 - self.clip_lower)
-        self.alpha = 0.5 * N_init
-        self.beta = 0.5 * N_init
-        self.old_mean_logprob = None
-
-    def get_baselines(self, batch_size):
-        baseline = self.alpha / (self.alpha + self.beta)
-        return torch.full((batch_size,), baseline, dtype=torch.float32)
-
-    def compute_rho(self, cur_mean_logprob):
-        if self.rho_mode == 'constant':
-            return self.rho_const
-        if self.old_mean_logprob is None:
-            return self.rho_const
-        kl = abs(self.old_mean_logprob - cur_mean_logprob)
-        rho = 2 ** (-kl / self.D_half)
-        return max(min(rho, self.clip_upper), self.clip_lower)
-
-    def update(self, rewards, cur_logprobs=None, response_masks=None):
-        if cur_logprobs is not None and response_masks is not None:
-            mean_logprob = ((cur_logprobs * response_masks).sum() / response_masks.sum()).item()
-            rho = self.compute_rho(mean_logprob)
-            self.old_mean_logprob = mean_logprob
-        else:
-            rho = self.rho_const
-
-        scale = 3.0
-        normalized_rewards = (rewards + scale) / (2 * scale)
-        avg_normalized_reward = normalized_rewards.mean().item()
-        self.alpha = rho * self.alpha + avg_normalized_reward
-        self.beta = rho * self.beta + (1 - avg_normalized_reward)
-        return rho
-
-
-def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
-    """整合所有奖励函数计算总奖励"""
-    def reasoning_model_reward(rewards):
-        pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
-        pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
-        matches_pattern = [re.match(pattern, response, re.S) for response in responses]
-        matches_pattern2 = [re.match(pattern2, response, re.S) for response in responses]
-
-        format_rewards = []
-        for match_pattern, match_pattern2 in zip(matches_pattern, matches_pattern2):
-            if match_pattern or match_pattern2:
-                format_rewards.append(0.5)
-            else:
-                format_rewards.append(0.0)
-        rewards += torch.tensor(format_rewards, device=args.device)
-
-        def mark_num(text):
-            reward = 0
-            if text.count("<think>") == 1: reward += 0.25
-            if text.count("</think>") == 1: reward += 0.25
-            if text.count("<answer>") == 1: reward += 0.25
-            if text.count("</answer>") == 1: reward += 0.25
-            return reward
-
-        mark_rewards = [mark_num(response) for response in responses]
-        rewards += torch.tensor(mark_rewards, device=args.device)
-        return rewards
-
-    rewards = torch.zeros(len(responses), device=args.device)
-    if args.reasoning == 1:
-        rewards = reasoning_model_reward(rewards)
-
-    with torch.no_grad():
-        reward_model_scores = []
-        scale = 3.0
-
-        for i, (prompt, response) in enumerate(zip(prompts, responses)):
-            pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
-            matches = re.findall(pattern, prompt, re.DOTALL)
-            messages = [{"role": role, "content": content.strip()} for role, content in matches]
-
-            tmp_chat = messages + [{"role": "assistant", "content": response}]
-            score = reward_model.get_score(reward_tokenizer, tmp_chat)
-            score = max(min(score, scale), -scale)
-
-            if args.reasoning == 1:
-                answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
-                if answer_match:
-                    answer_content = answer_match.group(1).strip()
-                    tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
-                    answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
-                    answer_score = max(min(answer_score, scale), -scale)
-                    score = score * 0.4 + answer_score * 0.6
-
-            reward_model_scores.append(score)
-
-        reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
-        rewards += reward_model_scores
-
-    return rewards
-
-
-def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokenizer, value_tracker, start_step=0, wandb=None):
-    for step, batch in enumerate(loader, start=start_step + 1):
-        prompts = batch['prompt']  # list[str], length B
-        prompt_inputs = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False,
-                                  padding_side="left", add_special_tokens=False).to(args.device)  # input_ids: [B, P], attention_mask: [B, P]
-        if args.max_seq_len:
-            prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -args.max_seq_len:]
-            prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -args.max_seq_len:]
-
-        with torch.no_grad():
-            # DDP 模型需要使用 .module 访问 generate 方法
-            model_for_gen = model.module if isinstance(model, DistributedDataParallel) else model
-            outputs = model_for_gen.generate(
-                **prompt_inputs, max_new_tokens=args.max_gen_len, do_sample=True, temperature=0.8,
-                num_return_sequences=1, pad_token_id=tokenizer.pad_token_id)  # [B, P+R]
-
-        completion_ids = outputs[:, prompt_inputs["input_ids"].size(1):]  # [B, R]
-
-        def get_per_token_logps(mdl, input_ids, n_keep):
-            input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
-            logits = mdl(input_ids, logits_to_keep=n_keep + 1).logits[:, :-1, :]
-            per_token_logps = []
-            for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
-                ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
-                per_token_logps.append(torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1))
-            return torch.stack(per_token_logps)
-
-        with autocast_ctx:
-            per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1))  # [B, R]
-            res = model(outputs) if lm_config.use_moe else None
-            aux_loss = res.aux_loss if res is not None else torch.tensor(0.0, device=args.device)
-        
-        with torch.no_grad():
-            ref_per_token_logps = get_per_token_logps(ref_model, outputs, completion_ids.size(1))  # [B, R]
-
-        completions = tokenizer.batch_decode(completion_ids, skip_special_tokens=True)  # list[str], length B
-        rewards = calculate_rewards(prompts, completions, reward_model, reward_tokenizer).to(args.device)  # [B]
-
-        baselines = value_tracker.get_baselines(len(prompts)).to(args.device)  # [B]
-
-        scale = 3.0
-        # Un-normalize baselines to be in the same scale as raw rewards [-3, 3]
-        unnormalized_baselines = baselines * (2 * scale) - scale  # [B]
-        advantages = rewards - unnormalized_baselines  # [B]
-
-        # 直接使用 baseline 提供的优势估计，只做裁剪防止梯度爆炸。不再做 batch 内归一化，因为 baseline 已经提供了跨 batch 的稳定基线
-        advantages = advantages.clamp(-5.0, 5.0)
-
-        is_eos = completion_ids == tokenizer.eos_token_id  # [B, R]
-        eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=args.device)  # [B]
-        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
-        completion_mask = (torch.arange(is_eos.size(1), device=args.device).expand(is_eos.size(0), -1) <= eos_idx.unsqueeze(1)).int()  # [B, R]
-
-        kl_div = ref_per_token_logps - per_token_logps  # [B, R]
-        per_token_kl = torch.exp(kl_div) - kl_div - 1  # [B, R]
-        per_token_loss = -per_token_logps * advantages.unsqueeze(1) + args.beta * per_token_kl  # [B, R]
-        policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
-        loss = (policy_loss + aux_loss) / args.accumulation_steps  # scalar
-        loss.backward()
-
-        response_masks = completion_mask.float()  # [B, R]
-        rho = value_tracker.update(rewards, per_token_logps.detach(), response_masks)
-
-        if step % args.accumulation_steps == 0:
-            if args.grad_clip > 0:
-                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-            optimizer.step()
-            scheduler.step()
-            optimizer.zero_grad()
-
-        if step % args.log_interval == 0 or step == iters:
-            policy_loss_val = loss.item() * args.accumulation_steps
-            current_aux_loss = aux_loss.item()
-            avg_reward_val = rewards.mean().item()
-            avg_len_val = completion_mask.sum(dim=1).float().mean().item()
-            kl_val = ((per_token_kl * completion_mask).sum() / (completion_mask.sum() + 1e-8)).item()
-            avg_baseline_val = baselines.mean().item()
-            current_lr = optimizer.param_groups[0]['lr']
-
-            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
-                   f'Actor Loss: {policy_loss_val:.4f}, Aux Loss: {current_aux_loss:.4f}, Reward: {avg_reward_val:.4f}, '
-                   f'Baseline: {avg_baseline_val:.4f}, KL: {kl_val:.4f}, Rho: {rho:.4f}, '
-                   f'Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}')
-
-            if wandb and is_main_process():
-                wandb.log({
-                    "policy_loss": policy_loss_val,
-                    "aux_loss": current_aux_loss,
-                    "reward": avg_reward_val,
-                    "kl": kl_val,
-                    "rho": float(rho),
-                    "baseline": avg_baseline_val,
-                    "advantages_mean": advantages.mean().item(),
-                    "learning_rate": current_lr
-                })
-
-        if (step % args.save_interval == 0 or step == iters) and is_main_process():
-            model.eval()
-            moe_suffix = '_moe' if lm_config.use_moe else ''
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
-            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
-            raw_model = getattr(raw_model, '_orig_mod', raw_model)
-            state_dict = raw_model.state_dict()
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
-            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, 
-                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
-            model.train()
-            del state_dict
-
-        del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
-        del completions, rewards, advantages, completion_mask, baselines, response_masks
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="MiniMind SPO (Self-Play Optimization)")
-    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
-    parser.add_argument('--save_weight', default='spo', type=str, help="保存权重的前缀名")
-    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=1e-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=4, help="梯度累积步数")
-    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
-    parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
-    parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
-    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
-    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
-    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument('--max_seq_len', default=66, type=int, help="Prompt最大长度")
-    parser.add_argument("--max_gen_len", type=int, default=1536, help="生成的最大长度")
-    parser.add_argument("--data_path", type=str, default="../dataset/rlaif-mini.jsonl", help="RLAIF数据路径")
-    parser.add_argument("--beta", type=float, default=0.02, help="KL惩罚系数")
-    parser.add_argument("--reasoning", type=int, default=1, choices=[0, 1], help='推理模型类型（0=普通模型，1=推理模型）')
-    parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
-    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
-    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-SPO", help="wandb项目名")
-    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
-    args = parser.parse_args()
-
-    # ========== 1. 初始化环境和随机种子 ==========
-    local_rank = init_distributed_mode()
-    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
-    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
-    
-    # ========== 2. 配置目录、模型参数、检查ckp ==========
-    os.makedirs(args.save_dir, exist_ok=True)
-    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
-                               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
-    
-    # ========== 3. 设置混合精度 ==========
-    device_type = "cuda" if "cuda" 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)
-    
-    # ========== 4. 配wandb ==========
-    wandb = None
-    if args.use_wandb and is_main_process():
-        import swanlab as wandb
-        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
-        resume = 'must' if wandb_id else None
-        wandb_run_name = f"MiniMind-SPO-Epoch-{args.epochs}-BS-{args.batch_size}-LR-{args.learning_rate}"
-        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
-    
-    # ========== 5. 初始化模型（Policy, Ref, Reward）和Value Tracker、数据 ==========
-    base_weight = "reason" if args.reasoning == 1 else "full_sft"
-    # Policy模型
-    model, tokenizer = init_model(lm_config, base_weight, device=args.device)
-    if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
-    # Reference模型
-    ref_model, _ = init_model(lm_config, base_weight, device=args.device)
-    ref_model = ref_model.eval().requires_grad_(False)
-    # Reward模型
-    reward_model = AutoModel.from_pretrained(
-        args.reward_model_path, torch_dtype=torch.float16, trust_remote_code=True
-    )
-    reward_model = reward_model.to(args.device).eval().requires_grad_(False)
-    reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True)
-    # Value Tracker
-    value_tracker = AutoAdaptiveValueTracker(rho_mode='kl', rho_const=0.9, D_half=0.06, clip_lower=0.5, clip_upper=0.96)
-    
-    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
-    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
-    
-    loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
-    iters = len(loader_for_count)
-    total_optimizer_steps = (iters // args.accumulation_steps) * args.epochs
-    scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
-    
-    # ========== 6. 从ckp恢复状态 ==========
-    start_epoch, start_step = 0, 0
-    if ckp_data:
-        model.load_state_dict(ckp_data['model'])
-        optimizer.load_state_dict(ckp_data['optimizer'])
-        scheduler.load_state_dict(ckp_data['scheduler'])
-        start_epoch = ckp_data['epoch']
-        start_step = ckp_data.get('step', 0)
-    
-    # ========== 7. DDP包模型 ==========
-    if dist.is_initialized():
-        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
-        model = DistributedDataParallel(model, device_ids=[local_rank])
-    
-    # ========== 8. 开始训练 ==========
-    for epoch in range(start_epoch, args.epochs):
-        train_sampler and train_sampler.set_epoch(epoch)
-        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
-        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
-        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
-        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
-        if skip > 0: 
-            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
-            spo_train_epoch(epoch, loader, len(loader) + skip, ref_model, reward_model, reward_tokenizer, value_tracker, start_step, wandb)
-        else:
-            spo_train_epoch(epoch, loader, len(loader), ref_model, reward_model, reward_tokenizer, value_tracker, 0, wandb)
-    
-    # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()

trainer/train_tokenizer.py

@@ -4,92 +4,115 @@ import os
 import json
 from tokenizers import decoders, models, pre_tokenizers, trainers, Tokenizer
 
-DATA_PATH = '../dataset/pretrain_hq.jsonl'
+DATA_PATH = '../dataset/sft_t2t_mini.jsonl'
 TOKENIZER_DIR = '../model_learn_tokenizer/'
 VOCAB_SIZE = 6400
+SPECIAL_TOKENS_NUM = 36
 
 def get_texts(data_path):
-    with open(data_path, 'r', encoding='utf-8') as f:
+    with open(data_path, 'r', encoding='utf-8', errors='ignore') as f:
         for i, line in enumerate(f):
-            if i >= 10000: break # 实验性，可只用前10000行测试
-            data = json.loads(line)
-            yield data['text']
+            if i >= 10000: break # 选10000行测试
+            try:
+                data = json.loads(line)
+                contents = [item.get('content') for item in data.get('conversations', []) if item.get('content')]
+                if contents:
+                    yield "\n".join(contents)
+            except json.JSONDecodeError:
+                continue
 
-def train_tokenizer(data_path, tokenizer_dir, vocab_size):
+def train_tokenizer(data_path, tokenizer_dir, vocab_size, special_tokens_num=SPECIAL_TOKENS_NUM):
     tokenizer = Tokenizer(models.BPE())
     tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    
+    special_tokens_list = [
+        "<|endoftext|>", "<|im_start|>", "<|im_end|>", 
+        "<|object_ref_start|>", "<|object_ref_end|>", "<|box_start|>", "<|box_end|>", "<|quad_start|>", "<|quad_end|>", 
+        "<|vision_start|>", "<|vision_end|>", "<|vision_pad|>", "<|image_pad|>", "<|video_pad|>", 
+        "<|audio_start|>", "<|audio_end|>", "<|audio_pad|>", "<tts_pad>", "<tts_text_bos>", "<tts_text_eod>", "<tts_text_bos_single>"
+    ]
+    
+    additional_tokens_list = [
+        "<tool_call>", "</tool_call>",
+        "<tool_response>", "</tool_response>",
+        "<think>", "</think>"
+    ]
+    num_buffer = special_tokens_num - len(special_tokens_list + additional_tokens_list)
+    buffer_tokens = [f"<|buffer{i}|>" for i in range(1, num_buffer + 1)] # 预留一定数量的token位置
+    all_special_tokens = special_tokens_list + additional_tokens_list + buffer_tokens
     trainer = trainers.BpeTrainer(
         vocab_size=vocab_size,
-        special_tokens=["<|endoftext|>", "<|im_start|>", "<|im_end|>"],
         show_progress=True,
-        initial_alphabet=pre_tokenizers.ByteLevel.alphabet()
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+        special_tokens=all_special_tokens
     )
     texts = get_texts(data_path)
     tokenizer.train_from_iterator(texts, trainer=trainer)
     tokenizer.decoder = decoders.ByteLevel()
-
-    assert tokenizer.token_to_id("<|endoftext|>") == 0
-    assert tokenizer.token_to_id("<|im_start|>") == 1
-    assert tokenizer.token_to_id("<|im_end|>") == 2
+    tokenizer.add_special_tokens(special_tokens_list)
 
     os.makedirs(tokenizer_dir, exist_ok=True)
     tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
     tokenizer.model.save(tokenizer_dir)
+    tokenizer_json_path = os.path.join(tokenizer_dir, "tokenizer.json")
+    with open(tokenizer_json_path, 'r', encoding='utf-8') as f:
+        tokenizer_data = json.load(f)
+    for token_info in tokenizer_data.get('added_tokens', []):
+        if token_info['content'] not in special_tokens_list:
+            token_info['special'] = False
+    with open(tokenizer_json_path, 'w', encoding='utf-8') as f:
+        json.dump(tokenizer_data, f, ensure_ascii=False, indent=2)
+    
+    added_tokens_decoder = {}
+    for i, token in enumerate(all_special_tokens):
+        idx = tokenizer.token_to_id(token)
+        added_tokens_decoder[str(idx)] = {
+            "content": token,
+            "lstrip": False,
+            "normalized": False,
+            "rstrip": False,
+            "single_word": False,
+            "special": True if token in special_tokens_list else False
+        }
+
     config = {
         "add_bos_token": False,
         "add_eos_token": False,
         "add_prefix_space": False,
-        "added_tokens_decoder": {
-            "0": {
-                "content": "<|endoftext|>",
-                "lstrip": False,
-                "normalized": False,
-                "rstrip": False,
-                "single_word": False,
-                "special": True
-            },
-            "1": {
-                "content": "<|im_start|>",
-                "lstrip": False,
-                "normalized": False,
-                "rstrip": False,
-                "single_word": False,
-                "special": True
-            },
-            "2": {
-                "content": "<|im_end|>",
-                "lstrip": False,
-                "normalized": False,
-                "rstrip": False,
-                "single_word": False,
-                "special": True
-            }
-        },
-        "additional_special_tokens": [],
+        "added_tokens_decoder": added_tokens_decoder,
+        "additional_special_tokens": [t for t in special_tokens_list if t not in ["<|endoftext|>"]],
         "bos_token": "<|im_start|>",
         "clean_up_tokenization_spaces": False,
         "eos_token": "<|im_end|>",
         "legacy": True,
-        "model_max_length": 32768,
+        "model_max_length": 131072,
         "pad_token": "<|endoftext|>",
         "sp_model_kwargs": {},
         "spaces_between_special_tokens": False,
-        "tokenizer_class": "PreTrainedTokenizerFast",
         "unk_token": "<|endoftext|>",
-        "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0]['role'] == 'system' -%}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else -%}\n        {{- '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n   {{- '<|im_start|>' + message.role + '\\n' + content }}\n  {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}"
+        "image_token": "<|image_pad|>",
+        "audio_token": "<|audio_pad|>",
+        "video_token": "<|video_pad|>",
+        "vision_bos_token": "<|vision_start|>",
+        "vision_eos_token": "<|vision_end|>",
+        "audio_bos_token": "<|audio_start|>",
+        "audio_eos_token": "<|audio_end|>",
+        "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if true %}\n            {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if open_thinking is defined and open_thinking is true %}\n        {{- '<think>\\n' }}\n    {%- else %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+        "tokenizer_class": "PreTrainedTokenizerFast"
     }
 
     with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w", encoding="utf-8") as f:
         json.dump(config, f, ensure_ascii=False, indent=4)
     print("Tokenizer training completed.")
 
-
 def eval_tokenizer(tokenizer_dir):
     from transformers import AutoTokenizer
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
     messages = [
         {"role": "system", "content": "你是一个优秀的聊天机器人，总是给我正确的回应！"},
         {"role": "user", "content": '你来自哪里？'},
+        {"role": "assistant", "content": '我来自月球'},
+        {"role": "user", "content": '你到底来自哪里？'},
         {"role": "assistant", "content": '我来自地球'}
     ]
     new_prompt = tokenizer.apply_chat_template(
@@ -98,16 +121,35 @@ def eval_tokenizer(tokenizer_dir):
     )
     print('-'*100)
     print(new_prompt)
-
-
     print('-'*100)
     print('tokenizer词表长度：', len(tokenizer))
     model_inputs = tokenizer(new_prompt)
     print('encoder长度：', len(model_inputs['input_ids']))
     response = tokenizer.decode(model_inputs['input_ids'], skip_special_tokens=False)
     print('decoder一致性：', response == new_prompt, "\n")
-
-
+    print('-'*100)
+    print('压缩率测试（Chars/Tokens）：')
+    test_texts = [
+        # 中文样本 (约200字)
+        "人工智能是计算机科学的一个分支，它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器，该领域的研究包括机器人、语言识别、图像识别、自然语言处理和专家系统等。人工智能从诞生以来，理论和技术日益成熟，应用领域也不断扩大，可以设想，未来人工智能带来的科技产品，将会是人类智慧的“容器”。人工智能可以对人的意识、思维的信息过程的模拟。人工智能不是人的智能，但能像人那样思考、也可能超过人的智能。",
+        "星际航行是指在星系内甚至星系间的空间中进行的航行。由于宇宙空间极其广阔，传统的化学火箭动力在恒星间航行时显得力不从心。科学家们提出了多种方案，包括离子推进器、核热火箭、甚至是利用反物质作为能源的设想。此外，曲率驱动和虫洞旅行等科幻概念也在理论物理研究中被反复探讨。尽管目前人类的足迹仅限于月球，但随着核聚变技术和材料科学的突破，前往火星乃至更遥远的太阳系边缘将成为可能。",
+        # 英文样本 (约200词/字符)
+        "Large language models (LLMs) are a type of artificial intelligence (AI) trained on vast amounts of text data to understand and generate human-like language. These models use deep learning techniques, specifically transformers, to process and predict the next word in a sequence. LLMs like GPT-4, Llama, and Claude have demonstrated remarkable capabilities in coding, translation, and creative writing. However, they also face challenges such as hallucinations, where the model generates factually incorrect information, and the need for significant computational resources.",
+        "The development of sustainable energy is crucial for the future of our planet. As climate change continues to impact global weather patterns, transitioning from fossil fuels to renewable sources like solar, wind, and hydroelectric power has become an urgent priority. Innovations in battery storage technology and smart grid management are essential to ensure a reliable energy supply. International cooperation and policy frameworks are also necessary to drive the global shift towards a greener economy and reduce carbon emissions.",
+        # 混合样本
+        "Python 是一种高级编程语言，以其简洁的语法和强大的生态系统而闻名。It is widely used in data science, machine learning, and web development. 开发者可以利用 NumPy, Pandas, and PyTorch 等库快速构建复杂的应用。学习 Python 的过程非常愉快，因为它的代码读起来就像英语一样。Whether you are a beginner or an expert, Python offers something for everyone.",
+    ]
+    
+    total_compression = 0
+    for i, text in enumerate(test_texts):
+        encoded = tokenizer.encode(text)
+        token_count = len(encoded)
+        char_count = len(text)
+        compression_ratio = char_count / token_count
+        total_compression += compression_ratio
+        print(f"样本 {i+1} | 字符数: {char_count:4} | Tokens: {token_count:3} | 压缩率: {compression_ratio:.2f}")
+    
+    print(f"平均压缩率: {total_compression / len(test_texts):.2f}")
     print('-'*100)
     print('流式解码（字节缓冲）测试：')
     input_ids = model_inputs['input_ids']

trainer/trainer_utils.py

@@ -12,7 +12,7 @@ import torch
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import Sampler
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
 from model.model_minimind import MiniMindForCausalLM
 
 def get_model_params(model, config):
@@ -154,4 +154,24 @@ class SkipBatchSampler(Sampler):
 
     def __len__(self):
         total_batches = (len(self.sampler) + self.batch_size - 1) // self.batch_size
-        return max(0, total_batches - self.skip_batches)
+        return max(0, total_batches - self.skip_batches)
+
+
+class LMForRewardModel:
+    def __init__(self, model_path, device="cuda", dtype=torch.float16):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model = AutoModel.from_pretrained(model_path, torch_dtype=dtype, trust_remote_code=True)
+        self.model = self.model.to(device).eval()
+        self.device = device
+
+    @torch.no_grad()
+    def get_score(self, messages, response):
+        history_text = "\n".join([f"{m['role']}: {m['content']}" for m in messages[:-1]])
+        last_query = messages[-1]['content'] if messages else ""
+        message_context = f"{history_text}\n以上是对话历史。我的新问题是：\n{last_query}" if history_text else last_query
+        eval_messages = [
+            {"role": "user", "content": message_context},
+            {"role": "assistant", "content": response}
+        ]
+        score = self.model.get_score(self.tokenizer, eval_messages)
+        return max(min(score, 3.0), -3.0)