github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Compare commits

base: github:fe5ecbd7e5f0f84de71c6b8babbd4a7467aa09fd
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
...
compare: github:be0b8b4475f3406dd5eabc96169ab792e3774fa6
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

10 Commits
fe5ecbd7e5 ... be0b8b4475

Author SHA1 Message Date
Jiyang be0b8b4475
Merge 437eae89d6 into 7a9137d2e9 2026-09-15 19:34:48 +08:00
jingyaogong 7a9137d2e9 [update] trim verbose ddp comments 2026-09-15 14:53:11 +08:00
jingyaogong 7988bc81a3
Merge pull request #855 from basil-k-aji-dev/perf/ddp-followups 
[perf] skip the per-forward RoPE buffer broadcast in the RL trainers
 2026-09-15 14:51:37 +08:00
basil-k-aji-dev 37ac206e9a [perf] skip the per-forward RoPE buffer broadcast in the RL trainers 
Follow-up to the review on #854.

DDP re-broadcasts every registered buffer from rank 0 on each forward. The only
buffers here are freqs_cos / freqs_sin, which each rank computes deterministically
from config and never mutates from data, so the broadcast moves identical bytes
every step for nothing.

Verified with a two-rank CPU/Gloo run: with broadcast_buffers=False the
gradients stay identical across ranks and so do the buffers.

Also document the critic's tied-embedding assumption. CriticModel.forward never
calls lm_head, which is only safe under DDP because tie_word_embeddings (default
True) makes lm_head.weight the same Parameter as embed_tokens.weight. Untied, it
becomes a genuinely unused parameter and DDP raises.
 2026-09-14 22:22:00 +05:30
jingyaogong 21ec325dba [update] simplify safe math eval 2026-09-14 23:55:32 +08:00
jingyaogong cdf0a93e02
Merge pull request #847 from a0917-cell/fix/safe-math-eval 
[fix] 用 AST 求值器替换工具里的 eval()，并修复 Windows 上工具调用全部静默失败
 2026-09-14 23:34:28 +08:00
jingyaogong c367572702
Merge pull request #854 from basil-k-aji-dev/fix/rl-ddp-grad-sync 
[fix] route RL trainer forward through DDP so gradients all-reduce
 2026-09-14 22:28:03 +08:00
basil-k-aji-dev 9c088dd3a5 [fix] route RL trainer forward through DDP so gradients all-reduce 
train_grpo.py, train_ppo.py and train_agent.py wrap the policy in
DistributedDataParallel but run the loss-bearing forward through
.module, bypassing the wrapper. DDP arms its gradient all-reduce hooks
inside DistributedDataParallel.forward(), so calling .module directly
means backward() performs no all-reduce: every rank steps on purely
local gradients and the replicas silently diverge. Only rank 0 is
saved, so multi-GPU RL training quietly degrades to single-GPU
training on 1/N of the data with no error.

The supervised trainers are unaffected; they already call the wrapped
module. no_grad uses of .module (rollout generation, checkpoint
saving) are left untouched, where unwrapping is correct.

Single-GPU behaviour is unchanged: with no process group, model is the
bare module and these calls are identical to before.
 2026-09-14 14:54:07 +05:30
tkgcc 0f6803e551 [fix] replace tool eval() with an AST math evaluator 
execute_tool() relied on signal.SIGALRM, which does not exist on Windows.
The AttributeError was swallowed by a bare except, so every tool call
silently returned None and the model always observed
{"error": "tool not found"} (issue #846).

calculate_math also ran eval() on model-generated text. Two of the three
call sites (scripts/eval_toolcall.py, scripts/web_demo.py) passed no
globals at all, leaving full builtins reachable from sampled output.

safe_math_eval() in trainer_utils.py parses the expression with ast and
walks it, allowing only arithmetic operators and a math whitelist, with
caps on expression length and exponent size. It cannot loop forever, so
the signal-based timeout is no longer needed and is removed. Both sqrt(x)
and math.sqrt(x) spellings keep working.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
 2026-09-07 07:53:33 +08:00
Jiyang Liu 437eae89d6 [feat] add MPS support for Apple Silicon 
On macOS there's no CUDA, so these trainer scripts fall back to CPU. Detect MPS in the --device default and switch mixed precision to the device-agnostic torch.autocast / torch.amp.GradScaler so it works there.

CUDA and CPU paths are unchanged.
 2026-08-14 01:26:49 +08:00
11 changed files with 91 additions and 56 deletions
Show Stats Expand all files Collapse all files

4
scripts/eval_toolcall.py
View File

@ -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())},

6
scripts/web_demo.py
View File

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

26
trainer/train_agent.py
View File

@ -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)
@ -384,7 +379,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=2, help="批次大小")
parser.add_argument("--learning_rate", type=float, default=3e-7, help="学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.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="梯度累积步数")
@ -427,9 +422,9 @@ if __name__ == "__main__":
max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
wandb = None
if args.use_wandb and is_main_process():
@ -478,7 +473,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled')
rollout_engine.update_policy(model)
if dist.is_initialized():
model = DistributedDataParallel(model, device_ids=[local_rank])
# 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费
model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
rollout_engine.update_policy(model)
for epoch in range(start_epoch, args.epochs):

8
trainer/train_distillation.py
View File

@ -151,7 +151,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=6, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.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="梯度累积步数")
@ -188,9 +188,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
@ -210,7 +210,7 @@ if __name__ == "__main__":
Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M')
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'))
scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ==========

8
trainer/train_dpo.py
View File

@ -135,7 +135,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=4, help="batch size")
parser.add_argument("--learning_rate", type=float, default=4e-8, help="初始学习率（建议<=5e-8避免遗忘）")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.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="梯度累积步数")
@ -166,9 +166,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
@ -190,7 +190,7 @@ if __name__ == "__main__":
train_ds = DPODataset(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'))
scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ==========

8
trainer/train_full_sft.py
View File

@ -88,7 +88,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=16, help="batch size")
parser.add_argument("--learning_rate", type=float, default=1e-5, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.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="梯度累积步数")
@ -119,9 +119,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
@ -136,7 +136,7 @@ if __name__ == "__main__":
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
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'))
scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ==========

14
trainer/train_grpo.py
View File

@ -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)
@ -211,7 +212,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=2, help="batch size")
parser.add_argument("--learning_rate", type=float, default=3e-7, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.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="梯度累积步数")
@ -256,9 +257,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
@ -313,7 +314,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled')
rollout_engine.update_policy(model)
if dist.is_initialized():
model = DistributedDataParallel(model, device_ids=[local_rank])
# 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费
model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
rollout_engine.update_policy(model)
# ========== 8. 开始训练 ==========

8
trainer/train_lora.py
View File

@ -82,7 +82,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=10, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.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="梯度累积步数")
@ -112,9 +112,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
@ -148,7 +148,7 @@ if __name__ == "__main__":
# ========== 6. 定义数据和优化器 ==========
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'))
scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
# ========== 7. 从ckp恢复状态 ==========

23
trainer/train_ppo.py
View File

@ -36,7 +36,7 @@ def rep_penalty(text, n=3, cap=0.5):
class CriticModel(MiniMindForCausalLM):
def __init__(self, params):
super().__init__(params)
# 替换lm_head为输出单一价值的线性层
# lm_head 不参与 forward，仅靠 tie_word_embeddings 与 embed_tokens 共享权重，解绑后 DDP 会报未使用参数
self.value_head = nn.Linear(params.hidden_size, 1)
def forward(self, input_ids=None, attention_mask=None, **kwargs):
@ -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 死锁
@ -315,7 +315,7 @@ if __name__ == "__main__":
parser.add_argument("--batch_size", type=int, default=2, help="batch size")
parser.add_argument("--learning_rate", type=float, default=3e-7, help="Actor学习率")
parser.add_argument("--critic_learning_rate", type=float, default=5e-7, help="Critic学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.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="梯度累积步数")
@ -364,9 +364,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
@ -429,8 +429,9 @@ if __name__ == "__main__":
Logger('torch.compile enabled')
rollout_engine.update_policy(actor_model)
if dist.is_initialized():
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
# freqs_cos/freqs_sin 各 rank 由 config 确定性算出，默认每步广播一次纯属浪费
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank], broadcast_buffers=False)
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank], broadcast_buffers=False)
rollout_engine.update_policy(actor_model)
# ========== 8. 开始训练 ==========

8
trainer/train_pretrain.py
View File

@ -87,7 +87,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
@ -118,9 +118,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
@ -135,7 +135,7 @@ if __name__ == "__main__":
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
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'))
scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ==========

34
trainer/trainer_utils.py
View File

@ -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