Merge 6b2f4c939a into cdf0a93e02

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

model/model_minimind.py

@@ -159,15 +159,30 @@ class MOEFeedForward(nn.Module):
         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)
+        # 先把 token 按专家排成连续段，每个专家取一段连续切片。
+        # 直觉上更自然的写法是每个专家做一次 mask：`if mask.any(): x_flat[mask.any(-1).nonzero()]`，
+        # 但 mask.any() 的取值、nonzero() 的输出长度、布尔索引的结果长度都必须先回到主机才知道，
+        # 于是每个专家都要 3 次 device→host 同步（每层 3E 次，默认配置下每次前向 96 次）。
+        # 同步会清空 CUDA 流：GPU 停下来等 CPU 读值再重新下发，而这些开销全花在路由的簿记上。
+        # 排序之后每段的边界一次就能算完，整层只剩 ends.tolist() 一次同步，且与专家数无关。
+        # stable=True 保证段内 token 仍按原顺序，因此结果与逐专家 mask 的写法逐位相同
+        #（掩码版本作为参考实现保留在 scripts/test_moe_sorted.py，每次跑测试都会断言两者一致）。
+        e, k = self.config.num_experts, self.config.num_experts_per_tok
+        flat = topk_idx.reshape(-1)
+        idx_sorted, order = torch.sort(flat, stable=True)
+        ends = torch.searchsorted(idx_sorted, torch.arange(e, device=x_flat.device, dtype=flat.dtype), right=True)
+        token_idx = order.div(k, rounding_mode='floor')  # 每个槽位对应的原 token
+        x_sorted = x_flat.index_select(0, token_idx)
+        weight = topk_weight.reshape(-1).index_select(0, order).unsqueeze(1)
+        bounds = [0] + ends.tolist()  # 整层唯一一次 device→host 同步
+        y, outs = 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))
+            lo, hi = bounds[i], bounds[i + 1]
+            if hi > lo:
+                outs.append(expert(x_sorted[lo:hi]))  # 主机侧的 python int 比较，不会触发同步
             elif self.training:
                 y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
+        if outs: y.index_add_(0, token_idx, (torch.cat(outs) * weight).to(y.dtype))
         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

scripts/bench_moe.py

@@ -0,0 +1,115 @@
+"""MoE 分发的性能基准 / MoE dispatch benchmark. 自包含，不引入任何新依赖。
+
+  --mode sync-probe   测本机一次 device→host 同步的往返开销
+  --mode layer        单个 MoE 层：当前实现 vs 改动前的掩码写法
+
+为什么要先测同步开销：掩码写法每层要做 3E 次主机往返，一次往返的代价由主机 CPU/PCIe
+延迟决定，而不是 GPU。我实测过从裸机主机的几微秒到 WSL2 笔记本的约 120 微秒，
+所以同一份基准在同型号 GPU、不同主机上可以相差 2-3 倍。先看自己机器的往返开销，
+再看下面的加速比，才知道该期待哪一端。
+
+计时协议：≥10 次预热，计时区间前后 torch.cuda.synchronize()，取中位数并报告 min/max。
+"""
+import argparse
+import os
+import statistics
+import sys
+import time
+
+import torch
+
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+from model.model_minimind import MiniMindConfig, MOEFeedForward
+from scripts.test_moe_sorted import reference_masked_forward
+
+
+def timed(fn, warmup=10, reps=20):
+    for _ in range(warmup): fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(reps):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        fn()
+        torch.cuda.synchronize()
+        samples.append(time.perf_counter() - t0)
+    return {'median': statistics.median(samples), 'min': min(samples), 'max': max(samples)}
+
+
+def mode_sync_probe(args):
+    flag = torch.zeros(1, dtype=torch.bool, device='cuda')
+    x = torch.randn(1024, 1024, device='cuda')
+
+    def probe(fn, n):
+        for _ in range(20): fn()
+        torch.cuda.synchronize()
+        s = []
+        for _ in range(args.reps):
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            for _ in range(n): fn()
+            torch.cuda.synchronize()
+            s.append((time.perf_counter() - t0) / n)
+        return statistics.median(s)
+
+    launch = probe(lambda: x.add_(0.0), 200)
+    rt = probe(lambda: flag.any().item(), 100)
+    print(f'  kernel launch (无同步) : {launch * 1e6:8.2f} us')
+    print(f'  device->host 往返      : {rt * 1e6:8.2f} us   <- 掩码写法每层要做 3E 次')
+    e, l = args.experts[0], args.num_hidden_layers
+    print(f'\n  按 3 x {e} experts x {l} layers 估算，仅路由簿记的主机开销下限约 '
+          f'{3 * e * l * rt * 1e3:.2f} ms/前向')
+
+
+def mode_layer(args):
+    print(f'{"E":>4}  {"masked (改动前)":>18}  {"sorted (当前)":>18}  {"加速":>7}  {"峰值显存":>16}')
+    for e in args.experts:
+        torch.manual_seed(0)
+        mod = MOEFeedForward(MiniMindConfig(
+            hidden_size=args.hidden_size, use_moe=True, num_experts=e,
+            num_experts_per_tok=args.top_k, moe_intermediate_size=args.hidden_size * 2)).cuda().train()
+        x = torch.randn(args.batch_size, args.seq_len, args.hidden_size, device='cuda')
+
+        def run(fn):
+            def step():
+                with torch.autocast('cuda', dtype=torch.bfloat16):
+                    fn(x).sum().backward()
+                mod.zero_grad(set_to_none=True)
+            torch.cuda.reset_peak_memory_stats()
+            r = timed(step, args.warmup, args.reps)
+            r['peak'] = torch.cuda.max_memory_allocated() / 1e6
+            return r
+
+        try:
+            a = run(lambda t: reference_masked_forward(mod, t))
+            b = run(mod)
+        except torch.cuda.OutOfMemoryError:
+            print(f'{e:>4}  OOM'); continue
+        print(f'{e:>4}  {a["median"] * 1e3:8.2f} ms [{a["min"] * 1e3:6.2f},{a["max"] * 1e3:6.2f}]  '
+              f'{b["median"] * 1e3:8.2f} ms [{b["min"] * 1e3:6.2f},{b["max"] * 1e3:6.2f}]  '
+              f'{a["median"] / b["median"]:6.2f}x  {a["peak"]:7.0f} -> {b["peak"]:5.0f} MB')
+        del mod; torch.cuda.empty_cache()
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--mode', default='sync-probe', choices=['sync-probe', 'layer'])
+    p.add_argument('--hidden_size', default=768, type=int)
+    p.add_argument('--num_hidden_layers', default=8, type=int)
+    p.add_argument('--batch_size', default=8, type=int)
+    p.add_argument('--seq_len', default=512, type=int)
+    p.add_argument('--top_k', default=1, type=int)
+    p.add_argument('--experts', default='2,4,8,16', type=lambda s: [int(x) for x in s.split(',')])
+    p.add_argument('--warmup', default=10, type=int)
+    p.add_argument('--reps', default=20, type=int)
+    args = p.parse_args()
+
+    if not torch.cuda.is_available():
+        print('本脚本需要 CUDA。/ This benchmark requires CUDA.'); sys.exit(1)
+    print(f'torch {torch.__version__}  {torch.cuda.get_device_name(0)}  '
+          f'sm_{"".join(map(str, torch.cuda.get_device_capability(0)))}\n')
+    {'sync-probe': mode_sync_probe, 'layer': mode_layer}[args.mode](args)
+
+
+if __name__ == '__main__':
+    main()

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/test_moe_sorted.py

@@ -0,0 +1,186 @@
+"""MoE 排序分发的正确性测试 / correctness tests for the sorted MoE dispatch.
+
+本仓库没有 pytest / CI，所以这里不引入任何测试框架，直接 `python scripts/test_moe_sorted.py`
+运行，失败返回非零。
+
+reference_masked_forward 是改动之前逐专家做 mask 的写法，保留在这里作为参考实现：
+它每次运行都会被执行，并断言与当前实现逐位相同，因此既能当作对照读，也不会随代码腐烂。
+"""
+import os
+import sys
+import traceback
+import warnings
+
+import torch
+import torch.nn.functional as F
+
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+from model.model_minimind import MiniMindConfig, MOEFeedForward
+
+
+class Skip(Exception): pass
+
+
+# ============================== 参考实现 / reference ==============================
+
+def reference_masked_forward(mod, x):
+    """改动前的逐专家掩码写法，逐位对照用。
+
+    每个专家都要 3 次 device→host 同步：mask.any() 的取值、nonzero() 的输出长度、
+    布尔索引的结果长度，全都得先回到主机才知道。
+    """
+    batch_size, seq_len, hidden_dim = x.shape
+    x_flat = x.view(-1, hidden_dim)
+    scores = F.softmax(mod.gate(x_flat), dim=-1)
+    topk_weight, topk_idx = torch.topk(scores, k=mod.config.num_experts_per_tok, dim=-1, sorted=False)
+    if mod.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(mod.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 mod.training:
+            y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
+    return y.view(batch_size, seq_len, hidden_dim)
+
+
+# ============================== 工具 / helpers ==============================
+
+def build(seed=0, **kw):
+    base = dict(hidden_size=64, use_moe=True, num_experts=4, num_experts_per_tok=1, moe_intermediate_size=128)
+    base.update(kw)
+    torch.manual_seed(seed)
+    return MOEFeedForward(MiniMindConfig(**base))
+
+
+def force_routing(mod, assign, hidden_size=64):
+    """构造 x 使得 token t 被路由到 assign[t]：gate 取单位阵，x 在目标维上给大值。"""
+    with torch.no_grad():
+        mod.gate.weight.zero_()
+        for i in range(mod.config.num_experts): mod.gate.weight[i, i] = 1.0
+    x = torch.zeros(1, len(assign), hidden_size)
+    for t, a in enumerate(assign): x[0, t, a] = 10.0
+    return x
+
+
+# ============================== 用例 / tests ==============================
+
+def test_bit_identical_to_masked_reference():
+    """与掩码参考实现对照：前向和参数梯度在所有 k 下逐位相同；输入梯度在 k<=2 下逐位相同。
+
+    k>=3 时输入梯度只能做到 1 ulp 以内，这不是 bug：每个 token 要把 k 份专家贡献加起来，
+    浮点加法不满足结合律，(a+b)+c 与 a+(b+c) 可以差 1 ulp；两种写法的求和顺序不同。
+    k<=2 时任何顺序都精确相等，所以那里仍然要求 torch.equal。实测最大相对偏差 ~1.2e-7
+    （fp32 eps ≈ 1.2e-7），这里的阈值 1e-6 只是留一点余量。
+    """
+    for e, k in ((1, 1), (2, 1), (4, 1), (4, 2), (4, 3), (4, 4), (8, 2), (8, 3), (8, 8)):
+        mod = build(seed=3, num_experts=e, num_experts_per_tok=k)
+        torch.manual_seed(21)
+        x = torch.randn(3, 11, 64)
+        for train in (True, False):
+            mod.train(train)
+            xa, xb = x.clone().requires_grad_(), x.clone().requires_grad_()
+            ya = reference_masked_forward(mod, xa)
+            ya.sum().backward()
+            ga = {n: (p.grad.clone() if p.grad is not None else None) for n, p in mod.named_parameters()}
+            mod.zero_grad(set_to_none=True)
+            yb = mod(xb)
+            yb.sum().backward()
+            assert torch.equal(ya, yb), f'E={e} k={k} train={train}: 前向不是逐位相同'
+            if k <= 2:
+                assert torch.equal(xa.grad, xb.grad), f'E={e} k={k} train={train}: 输入梯度不是逐位相同'
+            else:  # 求和顺序不同，只能要求 1 ulp 量级
+                rel = (xa.grad - xb.grad).abs().max() / xa.grad.abs().max()
+                assert rel < 1e-6, f'E={e} k={k} train={train}: 输入梯度相对偏差 {rel:.2e} 超出浮点求和顺序可解释的范围'
+            for n, p in mod.named_parameters():
+                assert (ga[n] is None) == (p.grad is None), f'{n}: 梯度有无不一致'
+                if p.grad is not None:
+                    assert torch.equal(ga[n], p.grad), f'E={e} k={k} train={train} {n}: 参数梯度不是逐位相同'
+            mod.zero_grad(set_to_none=True)
+
+
+def test_edge_cases():
+    """零token专家 / 全部token到同一专家 / batch 小于专家数。"""
+    for name, assign in {'零token专家': [0, 1, 0, 1, 0, 1], '全部到专家0': [0] * 6, 'batch<专家数': [1, 2]}.items():
+        mod = build(seed=5); mod.train()
+        x = force_routing(mod, assign)
+        assert torch.equal(reference_masked_forward(mod, x), mod(x)), f'{name}: 与掩码参考实现不一致'
+
+
+def test_zero_token_expert_still_gets_grad():
+    """DDP 以 find_unused_parameters=False 构造，空专家也必须出现在自动求导图里。"""
+    mod = build(seed=5); mod.train()
+    mod(force_routing(mod, [0] * 6)).sum().backward()
+    missing = [n for n, p in mod.named_parameters() if p.grad is None]
+    assert not missing, f'{missing} 没有梯度 -> DDP(find_unused_parameters=False) 会报错'
+
+
+def test_aux_loss_and_state_dict_unchanged():
+    """aux_loss 与 state_dict 结构不受本次改动影响。"""
+    mod = build(seed=7); mod.train()
+    torch.manual_seed(11)
+    x = torch.randn(2, 16, 64)
+    mod(x); got = mod.aux_loss.clone()
+    scores = F.softmax(mod.gate(x.view(-1, 64)), dim=-1)
+    _, topk_idx = torch.topk(scores, k=1, dim=-1, sorted=False)
+    load = F.one_hot(topk_idx, 4).float().mean(0)
+    want = (load * scores.mean(0)).sum() * 4 * mod.config.router_aux_loss_coef
+    assert torch.equal(got, want), 'aux_loss 变了'
+    expect = {'gate.weight'} | {f'experts.{i}.{p}_proj.weight' for i in range(4) for p in ('gate', 'up', 'down')}
+    assert set(mod.state_dict()) == expect, 'state_dict 键名变了，旧 checkpoint 将无法 strict 加载'
+
+
+def test_host_sync_count():
+    """本次改动的核心：整层同步次数恒为 1，不随专家数增长。"""
+    if not torch.cuda.is_available(): raise Skip('需要 CUDA')
+
+    def syncs(fn, x):
+        fn(x); torch.cuda.synchronize()  # 预热，避免把首次分配算进去
+        torch.cuda.set_sync_debug_mode('warn')
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always'); fn(x)
+        torch.cuda.set_sync_debug_mode('default')
+        return sum('sync' in str(r.message).lower() for r in w)
+
+    for e in (4, 8, 16):
+        mod = build(seed=3, num_experts=e).cuda().train()
+        x = torch.randn(2, 64, 64, device='cuda')
+        ref, now = syncs(lambda t: reference_masked_forward(mod, t), x), syncs(mod, x)
+        assert ref >= 3 * e, f'E={e}: 掩码写法应约 3E={3 * e} 次同步，实测 {ref}'
+        assert now == 1, f'E={e}: 排序写法应恒为 1 次同步，实测 {now}'
+
+
+def test_cuda_bit_identical():
+    """CUDA 上 fp32 与 bf16 autocast 同样逐位相同——排序改变了 kernel 的分块方式。"""
+    if not torch.cuda.is_available(): raise Skip('需要 CUDA')
+    for e, k in ((4, 1), (8, 2)):
+        mod = build(seed=3, num_experts=e, num_experts_per_tok=k).cuda().train()
+        torch.manual_seed(21)
+        x = torch.randn(4, 128, 64, device='cuda')
+        assert torch.equal(reference_masked_forward(mod, x), mod(x)), f'E={e} k={k}: CUDA fp32 不一致'
+        with torch.autocast('cuda', dtype=torch.bfloat16):
+            assert torch.equal(reference_masked_forward(mod, x), mod(x)), f'E={e} k={k}: CUDA bf16 不一致'
+
+
+TESTS = [test_bit_identical_to_masked_reference, test_edge_cases, test_zero_token_expert_still_gets_grad,
+         test_aux_loss_and_state_dict_unchanged, test_host_sync_count, test_cuda_bit_identical]
+
+
+def main():
+    print(f'torch {torch.__version__}  cuda={torch.cuda.is_available()}\n')
+    fails = skips = 0
+    for t in TESTS:
+        try:
+            t(); print(f'  PASS  {t.__name__}')
+        except Skip as s:
+            skips += 1; print(f'  SKIP  {t.__name__} ({s})')
+        except Exception:
+            fails += 1; print(f'  FAIL  {t.__name__}'); traceback.print_exc()
+    print(f'\n{len(TESTS) - fails - skips} passed, {fails} failed, {skips} skipped')
+    sys.exit(1 if fails else 0)
+
+
+if __name__ == '__main__':
+    main()

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')}

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,7 +23,7 @@ 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
 
 warnings.filterwarnings('ignore')
@@ -55,7 +54,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 +84,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"):
@@ -278,9 +272,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

@@ -95,9 +95,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

@@ -159,8 +159,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 +166,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 +191,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
@@ -174,4 +176,53 @@ class LMForRewardModel:
             {"role": "assistant", "content": response}
         ]
         score = self.model.get_score(self.tokenizer, eval_messages)
-        return max(min(score, 3.0), -3.0)
+        return max(min(score, 3.0), -3.0)
+
+
+# ===== 数学表达式安全求值（替代 eval，避免执行模型生成的任意代码） =====
+_MATH_FUNCS = {name: getattr(math, name) for name in (
+    'sqrt', 'pow', 'exp', 'log', 'log2', 'log10', 'sin', 'cos', 'tan', 'asin', 'acos', 'atan', 'atan2',
+    'sinh', 'cosh', 'tanh', 'floor', 'ceil', 'trunc', 'fabs', 'fmod', 'hypot', 'gcd', 'degrees', 'radians'
+)}
+_MATH_CONSTS = {'pi': math.pi, 'e': math.e, 'tau': math.tau}
+_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: operator.pow}
+_UNARY_OPS = {ast.UAdd: operator.pos, ast.USub: operator.neg}
+_MAX_EXPR_LEN, _MAX_EXP, _MAX_DIGITS = 512, 1e4, 100  # 表达式长度 / 指数 / 结果位数上限
+
+
+def _check_pow(base, exp):  # 防止 9**9**9 这类表达式把进程算死
+    if abs(exp) > _MAX_EXP or (abs(base) > 1 and abs(exp) * math.log10(abs(base)) > _MAX_DIGITS):
+        raise ValueError('幂运算结果过大')
+
+
+def _resolve_name(node):  # 同时兼容 sqrt(4) 与 math.sqrt(4) 两种写法
+    if isinstance(node, ast.Name): return node.id
+    if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == 'math': return node.attr
+    return None
+
+
+def _eval_ast(node):
+    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)) and not isinstance(node.value, bool):
+        return node.value
+    if isinstance(node, (ast.Name, ast.Attribute)) and _resolve_name(node) in _MATH_CONSTS:
+        return _MATH_CONSTS[_resolve_name(node)]
+    if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPS:
+        return _UNARY_OPS[type(node.op)](_eval_ast(node.operand))
+    if isinstance(node, ast.BinOp) and type(node.op) in _BIN_OPS:
+        left, right = _eval_ast(node.left), _eval_ast(node.right)
+        if isinstance(node.op, ast.Pow): _check_pow(left, right)
+        return _BIN_OPS[type(node.op)](left, right)
+    if isinstance(node, ast.Call) and _resolve_name(node.func) in _MATH_FUNCS and not node.keywords:
+        name, args = _resolve_name(node.func), [_eval_ast(arg) for arg in node.args]
+        if name == 'pow' and len(args) == 2: _check_pow(*args)
+        return _MATH_FUNCS[name](*args)
+    raise ValueError(f'不支持的表达式语法: {type(node).__name__}')
+
+
+def safe_math_eval(expression):
+    """对模型生成的数学表达式求值：只放行算术运算与 math 白名单函数，不执行任意代码。"""
+    expr = str(expression).replace('^', '**').replace('×', '*').replace('÷', '/').replace('−', '-')
+    expr = expr.replace('²', '**2').replace('³', '**3').replace('（', '(').replace('）', ')').strip()
+    if not expr or len(expr) > _MAX_EXPR_LEN: raise ValueError('表达式为空或过长')
+    return _eval_ast(ast.parse(expr, mode='eval').body)