Compare commits

...

10 Commits

Author SHA1 Message Date
AmirX0 c86e25bae7
Merge 56a21ed3c2 into 7a9137d2e9 2026-09-15 16:07:58 +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
AmirX0 56a21ed3c2 docs: expand bilingual dataset guide 2026-09-01 07:06:36 +03:30
9 changed files with 285 additions and 32 deletions

View File

@ -486,6 +486,8 @@ torchrun --nproc_per_node N train_xxx.py
MiniMind训练数据集下载地址: [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main) MiniMind训练数据集下载地址: [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main)
下载、文件用途、最终格式与自定义数据检查清单可参考[数据集双语速查指南](./dataset/dataset.md)。
> 无需全部clone,可单独下载所需的文件 > 无需全部clone,可单独下载所需的文件
将下载的数据集文件放到`./dataset/`目录下(✨为推荐的必须项) 将下载的数据集文件放到`./dataset/`目录下(✨为推荐的必须项)

View File

@ -485,6 +485,8 @@ Besides this, other RL data maintains the same format as SFT data, typically fil
MiniMind training dataset download links: [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main) MiniMind training dataset download links: [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main)
For downloads, file-to-trainer mapping, final schemas, and a custom-data checklist, see the [bilingual dataset quick reference](./dataset/dataset.md).
> No need to clone everything, you can download individual files as needed > No need to clone everything, you can download individual files as needed
Place the downloaded dataset files in the `./dataset/` directory (✨ indicates recommended essentials) Place the downloaded dataset files in the `./dataset/` directory (✨ indicates recommended essentials)

View File

@ -1,5 +1,219 @@
# MiniMind Datasets # MiniMind 数据集指南 / Dataset Guide
将所有下载的数据集文件放置到当前目录. [中文](#user-content-chinese) | [English](#user-content-english)
Place the downloaded dataset file in the current directory. <a id="chinese"></a>
## 中文
本目录存放 MiniMind 各训练阶段使用的最终 JSONL 文件。官方发布的数据已经完成收集、蒸馏、清洗、去重、长度控制与格式统一,可直接由 `dataset/lm_dataset.py` 中的数据加载器读取。
### 下载与最小复现
数据集可从 [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) 或 [Hugging Face](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main) 按文件下载,无需克隆全部数据。
将文件放在当前目录:
```text
dataset/
├── pretrain_t2t_mini.jsonl
├── sft_t2t_mini.jsonl
├── pretrain_t2t.jsonl
├── sft_t2t.jsonl
├── dpo.jsonl
├── rlaif.jsonl
├── agent_rl.jsonl
└── agent_rl_math.jsonl
```
只想快速跑通从预训练到对话模型的流程时,下载以下两个文件即可:
- `pretrain_t2t_mini.jsonl`
- `sft_t2t_mini.jsonl`
然后从仓库根目录运行:
```bash
cd trainer
python train_pretrain.py
python train_full_sft.py
```
### 文件与训练脚本对应关系
| 文件 | 用途 | 默认使用者 |
| --- | --- | --- |
| `pretrain_t2t_mini.jsonl` | 快速预训练 | `train_pretrain.py` |
| `pretrain_t2t.jsonl` | 完整 MiniMind-3 预训练 | 通过 `train_pretrain.py --data_path ...` 指定 |
| `sft_t2t_mini.jsonl` | 快速 SFT、蒸馏及 LoRA 自定义数据格式参考 | `train_full_sft.py`、`train_distillation.py` |
| `sft_t2t.jsonl` | 完整 MiniMind-3 SFT | 通过 `train_full_sft.py --data_path ...` 指定 |
| `dpo.jsonl` | 偏好优化 | `train_dpo.py` |
| `rlaif.jsonl` | PPO、GRPO 与 CISPO | `train_ppo.py`、`train_grpo.py` |
| `agent_rl.jsonl` | 多轮 Tool-Use Agentic RL | `train_agent.py` |
| `agent_rl_math.jsonl` | 带可验证最终答案的数学 Agentic RL | 通过 `train_agent.py --data_path ...` 指定 |
上表中的训练脚本都支持 `--data_path`,因此也可直接指向自己的 JSONL 文件。
### 最终数据格式
预训练数据每行包含一个 `text` 字段:
```jsonl
{"text": "Transformer 通过自注意力机制建模上下文关系。"}
```
SFT、RLAIF 与常规对话数据使用 `conversations` 数组:
```json
{
"conversations": [
{"role": "user", "content": "你好"},
{"role": "assistant", "content": "你好!"}
]
}
```
DPO 数据分别提供偏好回复与非偏好回复:
```json
{
"chosen": [
{"role": "user", "content": "Q"},
{"role": "assistant", "content": "good answer"}
],
"rejected": [
{"role": "user", "content": "Q"},
{"role": "assistant", "content": "bad answer"}
]
}
```
Agentic RL 数据还可在 system 消息中包含序列化的 `tools`,并用顶层 `gt` 字段提供可验证目标。具体解析逻辑以 [`lm_dataset.py`](./lm_dataset.py) 为准。
### 来源、处理与自定义数据
当前公开版本使用的主要社区来源包括:
- [匠数大模型数据集](https://www.modelscope.cn/datasets/deepctrl/deepctrl-sft-data)
- [Magpie-Align](https://www.modelscope.cn/organization/Magpie-Align)
- [R1-Distill-SFT](https://www.modelscope.cn/datasets/AI-ModelScope/R1-Distill-SFT)
- [COIG](https://huggingface.co/datasets/BAAI/COIG)
- [Step-3.5-Flash-SFT](https://huggingface.co/datasets/stepfun-ai/Step-3.5-Flash-SFT)
- [DPO-En-Zh-20k](https://huggingface.co/datasets/llamafactory/DPO-En-Zh-20k)
本仓库目前提供训练所需的最终数据、格式说明和加载器;原始数据的完整收集、蒸馏与清洗流水线并未作为可执行脚本包含在本目录中。准备自定义数据时应:
1. 确认每个来源允许当前用途与再分发方式,并保留来源和许可证记录。
2. 完成去重、质量过滤、隐私信息清理与长度控制。
3. 转换为上面的最终 JSONL 格式。
4. 先用小样本运行对应训练脚本,确认 schema 和 chat template 正确,再开始完整训练。
数据来源、推荐组合、Tool Calling 格式与 token 长度说明以仓库根目录的[中文 README](../README.md#-数据介绍)为准。
<a id="english"></a>
## English
This directory stores the final JSONL files consumed by each MiniMind training stage. The published datasets have already gone through collection, distillation, cleaning, deduplication, length control, and format normalization, and can be read directly by the loaders in `dataset/lm_dataset.py`.
### Download and minimal reproduction
Download individual files from [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) or [Hugging Face](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main). You do not need to clone the full dataset repository.
Place the downloaded files in this directory:
```text
dataset/
├── pretrain_t2t_mini.jsonl
├── sft_t2t_mini.jsonl
├── pretrain_t2t.jsonl
├── sft_t2t.jsonl
├── dpo.jsonl
├── rlaif.jsonl
├── agent_rl.jsonl
└── agent_rl_math.jsonl
```
For the quickest pretraining-to-chat-model reproduction, only these two files are required:
- `pretrain_t2t_mini.jsonl`
- `sft_t2t_mini.jsonl`
Then run from the repository root:
```bash
cd trainer
python train_pretrain.py
python train_full_sft.py
```
### Dataset-to-trainer map
| File | Purpose | Default consumer |
| --- | --- | --- |
| `pretrain_t2t_mini.jsonl` | Quick pretraining | `train_pretrain.py` |
| `pretrain_t2t.jsonl` | Full MiniMind-3 pretraining | Pass with `train_pretrain.py --data_path ...` |
| `sft_t2t_mini.jsonl` | Quick SFT, distillation, and the reference schema for custom LoRA data | `train_full_sft.py`, `train_distillation.py` |
| `sft_t2t.jsonl` | Full MiniMind-3 SFT | Pass with `train_full_sft.py --data_path ...` |
| `dpo.jsonl` | Preference optimization | `train_dpo.py` |
| `rlaif.jsonl` | PPO, GRPO, and CISPO | `train_ppo.py`, `train_grpo.py` |
| `agent_rl.jsonl` | Multi-turn Tool-Use Agentic RL | `train_agent.py` |
| `agent_rl_math.jsonl` | Math Agentic RL with verifiable final answers | Pass with `train_agent.py --data_path ...` |
Every trainer listed above accepts `--data_path`, so it can also read a custom JSONL file directly.
### Final data schemas
Each pretraining line contains a `text` field:
```jsonl
{"text": "Transformers model context with self-attention."}
```
SFT, RLAIF, and regular conversation data use a `conversations` array:
```json
{
"conversations": [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hello!"}
]
}
```
DPO data provides preferred and rejected conversations separately:
```json
{
"chosen": [
{"role": "user", "content": "Q"},
{"role": "assistant", "content": "good answer"}
],
"rejected": [
{"role": "user", "content": "Q"},
{"role": "assistant", "content": "bad answer"}
]
}
```
Agentic RL records may also contain serialized `tools` on the system message and a top-level `gt` field for verifiable targets. See [`lm_dataset.py`](./lm_dataset.py) for the authoritative parsing behavior.
### Sources, processing, and custom data
Major community sources used by the current published datasets include:
- [Craftsman LLM Dataset](https://www.modelscope.cn/datasets/deepctrl/deepctrl-sft-data)
- [Magpie-Align](https://www.modelscope.cn/organization/Magpie-Align)
- [R1-Distill-SFT](https://www.modelscope.cn/datasets/AI-ModelScope/R1-Distill-SFT)
- [COIG](https://huggingface.co/datasets/BAAI/COIG)
- [Step-3.5-Flash-SFT](https://huggingface.co/datasets/stepfun-ai/Step-3.5-Flash-SFT)
- [DPO-En-Zh-20k](https://huggingface.co/datasets/llamafactory/DPO-En-Zh-20k)
This repository currently provides the final training datasets, their schemas, and the loaders. The complete raw-data collection, distillation, and cleaning pipeline is not included here as an executable script. When preparing custom data:
1. Confirm that each source permits the intended use and redistribution, and retain provenance and license records.
2. Deduplicate, quality-filter, remove private information, and control sample lengths.
3. Convert records to the final JSONL schemas above.
4. Run the matching trainer on a small sample to validate the schema and chat template before starting a full training run.
For source details, recommended dataset combinations, Tool Calling schemas, and token-length guidance, see the [English README](../README_en.md#-data-introduction).

View File

@ -12,7 +12,7 @@ from datetime import datetime
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from openai import OpenAI from openai import OpenAI
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM 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') warnings.filterwarnings('ignore')
TOOLS = [ TOOLS = [
@ -27,7 +27,7 @@ TOOLS = [
] ]
MOCK_RESULTS = { 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")}, "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)))}, "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())}, "text_length": lambda args: {"characters": len(args.get("text", "")), "words": len(args.get("text", "").split())},

View File

@ -2,6 +2,7 @@ import random
import re import re
import json import json
import os import os
import sys
from threading import Thread from threading import Thread
import torch import torch
@ -9,6 +10,9 @@ import numpy as np
import streamlit as st import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer 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.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed")
st.markdown(""" st.markdown("""
@ -125,7 +129,7 @@ def execute_tool(tool_name, args):
import datetime import datetime
try: try:
if tool_name == 'calculate_math': 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': elif tool_name == 'get_current_time':
tz = args.get('timezone', 'Asia/Shanghai') tz = args.get('timezone', 'Asia/Shanghai')
return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

View File

@ -10,7 +10,6 @@ import gc
import json import json
import math import math
import random import random
import signal
import argparse import argparse
import warnings import warnings
import torch import torch
@ -24,7 +23,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
from transformers import AutoTokenizer from transformers import AutoTokenizer
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
from dataset.lm_dataset import AgentRLDataset 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.rollout_engine import create_rollout_engine, compute_per_token_logps
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -55,7 +54,7 @@ UNIT_DATA = {"km_miles": 0.621371, "miles_km": 1.60934, "kg_pounds": 2.20462, "p
# ======== 模拟执行 ======== # ======== 模拟执行 ========
MOCK_RESULTS = { 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)}, "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_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_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) fn = MOCK_RESULTS.get(name)
if not fn: return None if not fn: return None
try: try:
signal.signal(signal.SIGALRM, lambda *_: (_ for _ in ()).throw(TimeoutError()))
signal.alarm(1)
return fn(args) return fn(args)
except: except Exception:
return None return None
finally:
try: signal.alarm(0)
except: pass
# ======== 多轮 Rollout ======== # ======== 多轮 Rollout ========
def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"): 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) 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: 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) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
logits = res.logits[:, :-1, :] logits = res.logits[:, :-1, :]
per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1) per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
@ -478,7 +473,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(model) rollout_engine.update_policy(model)
if dist.is_initialized(): 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) rollout_engine.update_policy(model)
for epoch in range(start_epoch, args.epochs): for epoch in range(start_epoch, args.epochs):

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] 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: 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) 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) per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
@ -313,7 +314,8 @@ if __name__ == "__main__":
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(model) rollout_engine.update_policy(model)
if dist.is_initialized(): 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) rollout_engine.update_policy(model)
# ========== 8. 开始训练 ========== # ========== 8. 开始训练 ==========

View File

@ -36,7 +36,7 @@ def rep_penalty(text, n=3, cap=0.5):
class CriticModel(MiniMindForCausalLM): class CriticModel(MiniMindForCausalLM):
def __init__(self, params): def __init__(self, params):
super().__init__(params) super().__init__(params)
# 替换lm_head为输出单一价值的线性层 # lm_head 不参与 forward,仅靠 tie_word_embeddings 与 embed_tokens 共享权重,解绑后 DDP 会报未使用参数
self.value_head = nn.Linear(params.hidden_size, 1) self.value_head = nn.Linear(params.hidden_size, 1)
def forward(self, input_ids=None, attention_mask=None, **kwargs): 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 clipfrac_sum = 0.0
aux_loss_sum = 0.0 aux_loss_sum = 0.0
log_count = 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): for ppo_epoch in range(args.ppo_update_iters):
if stop_ppo: if stop_ppo:
break 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): for i in range(0, B, mb_size):
inds = b_inds[i:i + 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]) mb_resp_values = mb_values_seq.gather(1, logp_pos[inds])
with autocast_ctx: 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) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
# 在 autocast 内计算 log_softmax,避免直接对 fp16/bf16 logits # 在 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_max={torch.exp(_lrv).max().item():.6f} "
f"ratio_min={torch.exp(_lrv).min().item():.6f} " f"ratio_min={torch.exp(_lrv).min().item():.6f} "
f"dropout={getattr(lm_config, 'dropout', None)} " 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 = (0.5 * (log_ratio ** 2) * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
# 同步各卡的 approx_kl,防止某卡 break 而其它卡继续导致 DDP 死锁 # 同步各卡的 approx_kl,防止某卡 break 而其它卡继续导致 DDP 死锁
@ -429,8 +429,9 @@ if __name__ == "__main__":
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(actor_model) rollout_engine.update_policy(actor_model)
if dist.is_initialized(): if dist.is_initialized():
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank]) # freqs_cos/freqs_sin 各 rank 由 config 确定性算出,默认每步广播一次纯属浪费
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank]) 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) rollout_engine.update_policy(actor_model)
# ========== 8. 开始训练 ========== # ========== 8. 开始训练 ==========

View File

@ -5,6 +5,8 @@ import os
import sys import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import ast
import operator
import random import random
import math import math
import numpy as np import numpy as np
@ -174,4 +176,34 @@ class LMForRewardModel:
{"role": "assistant", "content": response} {"role": "assistant", "content": response}
] ]
score = self.model.get_score(self.tokenizer, eval_messages) 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