Compare commits
6 Commits
b7cf3e7521
...
f109aeaaa6
| Author | SHA1 | Date |
|---|---|---|
|
|
f109aeaaa6 | |
|
|
21ec325dba | |
|
|
cdf0a93e02 | |
|
|
c367572702 | |
|
|
9c088dd3a5 | |
|
|
0f6803e551 |
|
|
@ -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())},
|
||||
|
|
|
|||
|
|
@ -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')}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 死锁
|
||||
|
|
|
|||
|
|
@ -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,34 @@ 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 白名单(长度上限 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
|
||||
|
|
|
|||
Loading…
Reference in New Issue