Merge 5b26153b75 into 21ec325dba

[fix] 用 AST 求值器替换工具里的 eval()，并修复 Windows 上工具调用全部静默失败

scripts/eval_toolcall.py

@@ -12,7 +12,7 @@ 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
+from trainer.trainer_utils import setup_seed, get_model_params, safe_math_eval
 warnings.filterwarnings('ignore')
 
 TOOLS = [
@@ -27,7 +27,7 @@ TOOLS = [
 ]
 
 MOCK_RESULTS = {
-    "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("²", "**2").replace("³", "**3").replace("（", "(").replace("）", ")")))},
+    "calculate_math": lambda args: {"result": str(safe_math_eval(args.get("expression", "0")))},
     "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())},

scripts/web_demo.py

@@ -2,6 +2,7 @@ import random
 import re
 import json
 import os
+import sys
 from threading import Thread
 
 import torch
@@ -9,6 +10,9 @@ import numpy as np
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from trainer.trainer_utils import safe_math_eval
+
 st.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed")
 
 st.markdown("""
@@ -125,7 +129,7 @@ def execute_tool(tool_name, args):
     import datetime
     try:
         if tool_name == 'calculate_math':
-            return {"result": eval(args.get('expression', '0'))}
+            return {"result": safe_math_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')}

tests/test_reward_utils.py

@@ -0,0 +1,43 @@
+import unittest
+
+from trainer.reward_utils import rep_penalty
+
+
+class StubTokenizer:
+    def __init__(self, token_ids):
+        self.token_ids = token_ids
+        self.calls = []
+
+    def encode(self, text, add_special_tokens=True):
+        self.calls.append((text, add_special_tokens))
+        return self.token_ids
+
+
+class RepetitionPenaltyTests(unittest.TestCase):
+    def test_repeated_chinese_tokens_receive_penalty(self):
+        tokenizer = StubTokenizer([1, 2, 1, 2, 1, 2, 1, 2, 1, 2])
+
+        self.assertEqual(rep_penalty("你好你好你好你好你好", tokenizer), 0.5)
+        self.assertEqual(tokenizer.calls, [("你好你好你好你好你好", False)])
+
+    def test_repeated_english_tokens_receive_penalty_case_insensitively(self):
+        tokenizer = StubTokenizer([3, 3, 3, 3])
+
+        self.assertEqual(rep_penalty("Hello HELLO hello hello", tokenizer), 0.5)
+        self.assertEqual(tokenizer.calls, [("hello hello hello hello", False)])
+
+    def test_unique_and_short_sequences_have_no_penalty(self):
+        self.assertEqual(rep_penalty("unique", StubTokenizer([1, 2, 3, 4, 5])), 0.0)
+        self.assertEqual(rep_penalty("short", StubTokenizer([1, 2])), 0.0)
+
+    def test_invalid_configuration_is_rejected(self):
+        tokenizer = StubTokenizer([])
+
+        with self.assertRaisesRegex(ValueError, "n must be greater than zero"):
+            rep_penalty("text", tokenizer, n=0)
+        with self.assertRaisesRegex(ValueError, "cap must be non-negative"):
+            rep_penalty("text", tokenizer, cap=-0.1)
+
+
+if __name__ == "__main__":
+    unittest.main()

trainer/reward_utils.py

@@ -0,0 +1,19 @@
+def rep_penalty(text, tokenizer, n=3, cap=0.5):
+    r"""Return a capped repetition penalty over model-token n-grams.
+
+    Token IDs keep the calculation consistent across languages. In particular,
+    contiguous CJK text must not be treated as one word, which is what happens
+    with a ``\w+``-based regular expression.
+    """
+    if n <= 0:
+        raise ValueError("n must be greater than zero")
+    if cap < 0:
+        raise ValueError("cap must be non-negative")
+
+    token_ids = tokenizer.encode(text.lower(), add_special_tokens=False)
+    grams = [tuple(token_ids[i:i + n]) for i in range(len(token_ids) - n + 1)]
+    if not grams:
+        return 0.0
+
+    repeated = len(grams) - len(set(grams))
+    return min(cap, repeated * cap * 2 / len(grams))

trainer/train_agent.py

@@ -10,7 +10,6 @@ import gc
 import json
 import math
 import random
-import signal
 import argparse
 import warnings
 import torch
@@ -24,18 +23,14 @@ 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.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel, safe_math_eval
 from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
+from trainer.reward_utils import rep_penalty
 
 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"]}}},
@@ -55,7 +50,7 @@ UNIT_DATA = {"km_miles": 0.621371, "miles_km": 1.60934, "kg_pounds": 2.20462, "p
 
 # ======== 模拟执行 ========
 MOCK_RESULTS = {
-    "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("（", "(").replace("）", ")"), {"__builtins__": {}, "math": math}))},
+    "calculate_math": lambda args: {"result": str(safe_math_eval(args.get("expression", "0")))},
     "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")},
@@ -85,14 +80,9 @@ 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:
+    except Exception:
         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"):
@@ -220,7 +210,7 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa
                 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)
+            reward -= rep_penalty(answer, tokenizer)
             rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
         # -------- 有工具调用：执行结果奖励 --------
         else:
@@ -240,7 +230,7 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa
             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)
+            reward -= rep_penalty(final_text if final_text else answer, tokenizer)
             rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
     return rewards
 
@@ -278,9 +268,10 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
 
         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)
 
-        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
         with autocast_ctx:
-            res = model_unwrapped(input_ids, attention_mask=full_mask)
+            # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
+            # DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
+            res = model(input_ids, attention_mask=full_mask)
             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)

trainer/train_grpo.py

@@ -24,16 +24,11 @@ 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, LMForRewardModel
 from trainer.rollout_engine import create_rollout_engine
+from trainer.reward_utils import rep_penalty
 
 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
-
-
 def calculate_rewards(prompts, responses, reward_model):
     rewards = torch.zeros(len(responses), device=args.device)
 
@@ -57,7 +52,7 @@ def calculate_rewards(prompts, responses, reward_model):
                     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)
+                rewards[response_idx] -= rep_penalty(answer, tokenizer)
 
                 score = reward_model.get_score(messages, answer)
                 reward_model_scores.append(score)
@@ -95,9 +90,10 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
 
         rewards = calculate_rewards(prompts, completions, reward_model).to(args.device)  # [B*num_gen]
 
-        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
         with autocast_ctx:
-            res = model_unwrapped(outputs, attention_mask=full_mask)
+            # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
+            # DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
+            res = model(outputs, attention_mask=full_mask)
             aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
             per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
 

trainer/train_ppo.py

@@ -22,16 +22,11 @@ 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, LMForRewardModel
 from trainer.rollout_engine import create_rollout_engine
+from trainer.reward_utils import rep_penalty
 
 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):
@@ -64,7 +59,7 @@ def calculate_rewards(prompts, responses, reward_model):
                 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)
+            rewards[i] -= rep_penalty(answer, tokenizer)
 
             score = reward_model.get_score(messages, answer)
             reward_model_scores.append(score)
@@ -159,8 +154,6 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
         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
@@ -168,11 +161,13 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
             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])
+                # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
+                # DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
+                mb_values_seq = critic_model(input_ids=gen_out[inds], attention_mask=full_mask[inds])
                 mb_resp_values = mb_values_seq.gather(1, logp_pos[inds])
 
                 with autocast_ctx:
-                    res = actor_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds])
+                    res = actor_model(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)
                     # 在 autocast 内计算 log_softmax，避免直接对 fp16/bf16 logits
                     # 计算造成额外数值偏差。
@@ -191,7 +186,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
                                f"ratio_max={torch.exp(_lrv).max().item():.6f} "
                                f"ratio_min={torch.exp(_lrv).min().item():.6f} "
                                f"dropout={getattr(lm_config, 'dropout', None)} "
-                               f"training={actor_unwrapped.training}")
+                               f"training={actor_model.training}")
                 approx_kl = (0.5 * (log_ratio ** 2) * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
                 
                 # 同步各卡的 approx_kl，防止某卡 break 而其它卡继续导致 DDP 死锁

trainer/trainer_utils.py

@@ -5,6 +5,8 @@ import os
 import sys
 __package__ = "trainer"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import ast
+import operator
 import random
 import math
 import numpy as np
@@ -175,3 +177,33 @@ class LMForRewardModel:
         ]
         score = self.model.get_score(self.tokenizer, eval_messages)
         return max(min(score, 3.0), -3.0)
+
+
+# ===== 数学表达式安全求值：替代 eval，只放行算术运算与 math 白名单（长度上限 512） =====
+def safe_math_eval(expression):
+    """对模型生成的数学表达式求值：支持 + - * / // % **、math 白名单函数与 pi/e/tau 常量。"""
+    def pow_guard(base, exp):  # 幂运算统一走这里，拦截 9**9**9 / 10**99999 这类把进程算死的输入
+        if abs(exp) > 1e4 or (abs(base) > 1 and abs(exp) * math.log10(abs(base)) > 100): raise ValueError('幂运算结果过大')
+        return base ** exp
+    def resolve(node):  # 同时兼容 sqrt(4) 与 math.sqrt(4) 两种写法
+        if isinstance(node, ast.Name): return node.id
+        if isinstance(node, ast.Attribute) and getattr(node.value, 'id', '') == 'math': return node.attr
+    def walk(node):
+        if isinstance(node, ast.Constant) and type(node.value) in (int, float): return node.value
+        name = resolve(node)
+        if name in consts: return consts[name]
+        if isinstance(node, ast.UnaryOp): return unary_ops[type(node.op)](walk(node.operand))
+        if isinstance(node, ast.BinOp): return bin_ops[type(node.op)](walk(node.left), walk(node.right))
+        if isinstance(node, ast.Call) and not node.keywords: return funcs[resolve(node.func)](*map(walk, node.args))
+        raise ValueError(f'不支持的表达式语法: {type(node).__name__}')
+    consts = {'pi': math.pi, 'e': math.e, 'tau': math.tau}
+    funcs = {'pow': pow_guard, **{n: getattr(math, n) for n in 'sqrt exp log log2 log10 sin cos tan asin acos atan atan2 sinh cosh tanh floor ceil trunc fabs fmod hypot gcd degrees radians'.split()}}
+    bin_ops = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod, ast.Pow: pow_guard}
+    unary_ops = {ast.UAdd: operator.pos, ast.USub: operator.neg}
+    chars = str.maketrans({'^': '**', '×': '*', '÷': '/', '−': '-', '²': '**2', '³': '**3', '（': '(', '）': ')'})
+    expr = str(expression).translate(chars).strip()
+    if not expr or len(expr) > 512: raise ValueError('表达式为空或过长')
+    try:
+        return walk(ast.parse(expr, mode='eval').body)
+    except (KeyError, SyntaxError, TypeError):
+        raise ValueError('不支持的表达式语法') from None