Compare commits
13 Commits
744f16d802
...
856a57cacb
| Author | SHA1 | Date |
|---|---|---|
mengxin
|
856a57cacb | |
|
mengxin
|
791b5150cc | |
|
mengxin
|
7b554a516b | |
|
mengxin
|
618b72dc2a | |
jingyaogong
|
7a9137d2e9 | |
jingyaogong
|
7988bc81a3 | |
basil-k-aji-dev
|
37ac206e9a | |
|
jingyaogong
|
21ec325dba | |
|
jingyaogong
|
cdf0a93e02 | |
jingyaogong
|
c367572702 | |
|
basil-k-aji-dev
|
9c088dd3a5 | |
|
jingyaogong
|
a3c7b01cc0 | |
tkgcc
|
0f6803e551 |
|
|
@ -1,6 +1,9 @@
|
|||
.venv*/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.py[cod]
|
||||
.DS_Store
|
||||
out
|
||||
out/
|
||||
checkpoints/
|
||||
.smoke-out/
|
||||
website/
|
||||
docs-minimind/
|
||||
13
README.md
13
README.md
|
|
@ -284,6 +284,19 @@ print(torch.cuda.is_available())
|
|||
若 `cuda` 不可用，也仍可根据自身设备选择 `CPU` 或 `MPS` 运行，但训练速度与兼容性会有非常大的差异。
|
||||
如需安装或更换 PyTorch 版本，可参考 [torch_stable](https://download.pytorch.org/whl/torch_stable.html) 与[链接](https://blog.csdn.net/weixin_45456738/article/details/141029610?ops_request_misc=&request_id=&biz_id=102&utm_term=%E5%AE%89%E8%A3%85torch&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-2-141029610.nonecase&spm=1018.2226.3001.4187)
|
||||
|
||||
#### Apple Silicon / MPS
|
||||
|
||||
在支持 Metal 的 Apple Silicon Mac 上，训练脚本会按 `CUDA → MPS → CPU` 的顺序自动选择设备，也可显式传入 `--device mps`。MPS 支持 `float32`、`float16` 和 `bfloat16`；建议较新的 Apple Silicon 优先尝试默认的 `bfloat16`，遇到算子兼容问题时改用 `float32`。
|
||||
|
||||
```bash
|
||||
cd trainer
|
||||
PYTORCH_ENABLE_MPS_FALLBACK=1 python train_pretrain.py \
|
||||
--device mps --dtype bfloat16 --num_workers 0 \
|
||||
--batch_size 8 --accumulation_steps 32
|
||||
```
|
||||
|
||||
`PYTORCH_ENABLE_MPS_FALLBACK=1` 会让尚未支持的 MPS 算子回退到 CPU，兼容性更好但可能降低速度。确认训练全程没有回退后可移除此环境变量。Mac 只有一个 MPS 设备，请直接使用 `python train_xxx.py`，不要使用面向 CUDA/NCCL 多卡训练的 `torchrun`。MPS 不使用 CUDA 的 pinned memory，保存检查点后也会自动释放 MPS 缓存。
|
||||
|
||||
</details>
|
||||
|
||||
### 1' 下载数据
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import torch
|
|||
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from model.model_lora import *
|
||||
from trainer.trainer_utils import setup_seed, get_model_params
|
||||
from trainer.trainer_utils import setup_seed, get_model_params, get_default_device
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
def init_model(args):
|
||||
|
|
@ -45,7 +45,7 @@ def main():
|
|||
parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）")
|
||||
parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）")
|
||||
parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）")
|
||||
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
|
||||
parser.add_argument('--device', default=get_default_device(), type=str, help="运行设备 (cuda:0/mps/cpu)")
|
||||
args = parser.parse_args()
|
||||
|
||||
prompts = [
|
||||
|
|
|
|||
|
|
@ -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')}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,68 @@
|
|||
"""Short MiniMind MPS training benchmark; it does not save checkpoints."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from trainer.trainer_utils import get_autocast_context, get_grad_scaler
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--batch-size', type=int, default=8)
|
||||
parser.add_argument('--seq-len', type=int, default=340)
|
||||
parser.add_argument('--hidden-size', type=int, default=768)
|
||||
parser.add_argument('--num-hidden-layers', type=int, default=8)
|
||||
parser.add_argument('--dtype', choices=['float32', 'float16', 'bfloat16'], default='bfloat16')
|
||||
parser.add_argument('--warmup-steps', type=int, default=1)
|
||||
parser.add_argument('--steps', type=int, default=3)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not torch.backends.mps.is_available():
|
||||
raise RuntimeError('MPS is not available')
|
||||
|
||||
config = MiniMindConfig(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
max_position_embeddings=max(args.seq_len, 512),
|
||||
)
|
||||
model = MiniMindForCausalLM(config).to('mps').train()
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4)
|
||||
scaler = get_grad_scaler('mps', args.dtype)
|
||||
input_ids = torch.randint(
|
||||
0, config.vocab_size, (args.batch_size, args.seq_len), device='mps'
|
||||
)
|
||||
|
||||
durations = []
|
||||
total_steps = args.warmup_steps + args.steps
|
||||
for step in range(total_steps):
|
||||
torch.mps.synchronize()
|
||||
started = time.perf_counter()
|
||||
with get_autocast_context('mps', args.dtype):
|
||||
loss = model(input_ids, labels=input_ids).loss
|
||||
scaler.scale(loss).backward()
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
torch.mps.synchronize()
|
||||
elapsed = time.perf_counter() - started
|
||||
if step >= args.warmup_steps:
|
||||
durations.append(elapsed)
|
||||
|
||||
mean_seconds = sum(durations) / len(durations)
|
||||
tokens_per_second = args.batch_size * args.seq_len / mean_seconds
|
||||
print(f'mean_step_seconds={mean_seconds:.3f}')
|
||||
print(f'tokens_per_second={tokens_per_second:.1f}')
|
||||
print(f'loss={loss.detach().float().item():.4f}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
{"text":"苹果芯片可以通过 MPS 加速 PyTorch 模型训练。"}
|
||||
{"text":"MiniMind 是一个用于学习语言模型训练流程的小型项目。"}
|
||||
{"text":"梯度下降通过反向传播更新神经网络参数。"}
|
||||
{"text":"混合精度训练可以降低内存占用并提高吞吐量。"}
|
||||
{"text":"统一内存允许 CPU 和 GPU 共享同一块物理内存。"}
|
||||
{"text":"检查点可用于恢复模型参数和优化器状态。"}
|
||||
{"text":"注意力机制根据上下文计算不同词元之间的关系。"}
|
||||
{"text":"小规模冒烟测试可以快速发现训练环境问题。"}
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
{"conversations":[{"role":"user","content":"MPS 是什么？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"MPS 是 Apple 平台上的 GPU 加速框架。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
{"conversations":[{"role":"user","content":"MiniMind 有什么用途？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"它用于学习和实验小型语言模型训练流程。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
{"conversations":[{"role":"user","content":"什么是梯度？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"梯度描述损失函数相对模型参数的变化方向。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
{"conversations":[{"role":"user","content":"为什么使用混合精度？","reasoning_content":"","tools":"","tool_calls":""},{"role":"assistant","content":"它通常能减少内存占用并提高训练速度。","reasoning_content":"","tools":"","tool_calls":""}]}
|
||||
|
|
@ -0,0 +1,98 @@
|
|||
import os
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from trainer.trainer_utils import (
|
||||
empty_device_cache,
|
||||
get_autocast_context,
|
||||
get_default_device,
|
||||
get_grad_scaler,
|
||||
init_model,
|
||||
init_distributed_mode,
|
||||
LMForRewardModel,
|
||||
should_pin_memory,
|
||||
)
|
||||
|
||||
|
||||
class DeviceUtilsTest(unittest.TestCase):
|
||||
def test_default_device_prefers_mps_over_cpu(self):
|
||||
with mock.patch('torch.cuda.is_available', return_value=False), \
|
||||
mock.patch('torch.backends.mps.is_available', return_value=True):
|
||||
self.assertEqual(get_default_device(), 'mps')
|
||||
|
||||
def test_pin_memory_is_cuda_only(self):
|
||||
self.assertTrue(should_pin_memory('cuda:0'))
|
||||
self.assertFalse(should_pin_memory('mps'))
|
||||
self.assertFalse(should_pin_memory('cpu'))
|
||||
|
||||
def test_distributed_mps_fails_with_actionable_message(self):
|
||||
with mock.patch.dict(os.environ, {'RANK': '0'}):
|
||||
with self.assertRaisesRegex(RuntimeError, 'single process'):
|
||||
init_distributed_mode('mps')
|
||||
|
||||
@unittest.skipUnless(torch.backends.mps.is_available(), 'MPS is not available')
|
||||
def test_init_model_uses_default_device_when_device_is_omitted(self):
|
||||
config = MiniMindConfig(
|
||||
hidden_size=64,
|
||||
num_hidden_layers=2,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
model, _ = init_model(
|
||||
config,
|
||||
from_weight='none',
|
||||
tokenizer_path='model',
|
||||
save_dir='out',
|
||||
device=None,
|
||||
)
|
||||
|
||||
self.assertEqual(next(model.parameters()).device.type, 'mps')
|
||||
|
||||
def test_reward_model_uses_default_device_when_device_is_omitted(self):
|
||||
fake_model = mock.Mock()
|
||||
fake_model.to.return_value = fake_model
|
||||
fake_model.eval.return_value = fake_model
|
||||
|
||||
with mock.patch(
|
||||
'trainer.trainer_utils.get_default_device', return_value='mps'
|
||||
), mock.patch(
|
||||
'trainer.trainer_utils.AutoTokenizer.from_pretrained'
|
||||
), mock.patch(
|
||||
'trainer.trainer_utils.AutoModel.from_pretrained', return_value=fake_model
|
||||
):
|
||||
LMForRewardModel('unused', device=None, dtype=torch.float32)
|
||||
|
||||
fake_model.to.assert_called_once_with('mps')
|
||||
|
||||
@unittest.skipUnless(torch.backends.mps.is_available(), 'MPS is not available')
|
||||
def test_minimind_optimizer_step_on_mps_for_all_dtypes(self):
|
||||
for dtype in ('float32', 'float16', 'bfloat16'):
|
||||
with self.subTest(dtype=dtype):
|
||||
config = MiniMindConfig(
|
||||
hidden_size=64,
|
||||
num_hidden_layers=2,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
model = MiniMindForCausalLM(config).to('mps').train()
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
||||
scaler = get_grad_scaler('mps', dtype)
|
||||
input_ids = torch.randint(0, config.vocab_size, (2, 32), device='mps')
|
||||
|
||||
with get_autocast_context('mps', dtype):
|
||||
loss = model(input_ids, labels=input_ids).loss
|
||||
|
||||
self.assertTrue(torch.isfinite(loss).item())
|
||||
scaler.scale(loss).backward()
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
torch.mps.synchronize()
|
||||
empty_device_cache('mps')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
@ -84,7 +84,7 @@ class TorchRolloutEngine(RolloutEngine):
|
|||
).clone() # [B*num_gen, P+R]
|
||||
prompt_len = prompt_ids.size(1)
|
||||
completion_ids = output_ids[:, prompt_len:] # [B*num_gen, R]
|
||||
full_mask = (output_ids != self.tokenizer.pad_token_id).long()
|
||||
full_mask = torch.cat([attention_mask.repeat_interleave(num_generations, dim=0), attention_mask.new_ones(output_ids.size(0), completion_ids.size(1))], dim=1)
|
||||
per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask)
|
||||
completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
|
||||
return RolloutResult(output_ids, completion_ids, per_token_logps, completions,
|
||||
|
|
|
|||
|
|
@ -10,13 +10,11 @@ import gc
|
|||
import json
|
||||
import math
|
||||
import random
|
||||
import signal
|
||||
import argparse
|
||||
import warnings
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
|
|
@ -24,7 +22,12 @@ 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,
|
||||
get_default_device, get_autocast_context,
|
||||
get_grad_scaler, should_pin_memory,
|
||||
safe_math_eval)
|
||||
from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
|
@ -55,7 +58,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 +88,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"):
|
||||
|
|
@ -106,34 +104,33 @@ def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_
|
|||
open_thinking = random.random() < thinking_ratio
|
||||
for turn in range(max_turns):
|
||||
context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=open_thinking)
|
||||
inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False).to(device)
|
||||
context_ids = inputs["input_ids"][0].tolist()
|
||||
if prompt_ids is None:
|
||||
prompt_ids = context_ids
|
||||
prompt_ids = tokenizer(context, add_special_tokens=False)["input_ids"]
|
||||
input_ids = torch.tensor([prompt_ids + response_ids], device=device)
|
||||
rollout_result = rollout_engine.rollout(
|
||||
prompt_ids=inputs["input_ids"],
|
||||
attention_mask=inputs["attention_mask"],
|
||||
prompt_ids=input_ids,
|
||||
attention_mask=torch.ones_like(input_ids),
|
||||
num_generations=1,
|
||||
max_new_tokens=max_new_tokens,
|
||||
temperature=0.8,
|
||||
)
|
||||
new_ids = rollout_result.completion_ids[0].tolist()
|
||||
new_logps = rollout_result.per_token_logps[0].tolist()
|
||||
if len(new_ids) != len(new_logps): Logger(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
|
||||
pairs = [(t, lp) for t, lp in zip(new_ids, new_logps) if t != tokenizer.pad_token_id and t != tokenizer.eos_token_id]
|
||||
new_ids = [t for t, _ in pairs]
|
||||
new_logps = [lp for _, lp in pairs]
|
||||
valid_len = int(rollout_result.completion_mask[0].sum().item())
|
||||
new_ids = rollout_result.completion_ids[0, :valid_len].tolist()
|
||||
new_logps = rollout_result.per_token_logps[0, :valid_len].tolist()
|
||||
if len(new_ids) != len(new_logps):
|
||||
raise RuntimeError(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
|
||||
new_text = rollout_result.completions[0]
|
||||
all_outputs.append(new_text)
|
||||
response_ids.extend(new_ids)
|
||||
response_mask.extend([1] * len(new_ids))
|
||||
response_mask.extend([int(t != tokenizer.eos_token_id) for t in new_ids])
|
||||
response_old_logps.extend(new_logps)
|
||||
final_context = context + new_text
|
||||
calls = parse_tool_calls(new_text)
|
||||
if not calls:
|
||||
break
|
||||
unfinished = turn == max_turns - 1
|
||||
messages.append({"role": "assistant", "content": new_text})
|
||||
assistant_message = {"role": "assistant", "content": new_text}
|
||||
messages.append(assistant_message)
|
||||
for call in calls:
|
||||
name, raw = call.get("name", ""), call.get("arguments", {})
|
||||
if isinstance(raw, str):
|
||||
|
|
@ -143,10 +140,17 @@ def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_
|
|||
result_str = (json.dumps(result, ensure_ascii=False) if result else '{"error": "tool not found"}')[:2048] # 防止天文数字撑爆tokenizer
|
||||
messages.append({"role": "tool", "content": result_str})
|
||||
|
||||
marker = f"<|agent_observation_{id(messages)}_{len(response_ids)}|>"
|
||||
assistant_message["content"] += marker
|
||||
marked_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
|
||||
assistant_message["content"] = new_text
|
||||
_, found, observation = marked_context.partition(marker)
|
||||
if not found:
|
||||
raise RuntimeError("chat template did not preserve the assistant content boundary")
|
||||
obs_delta = tokenizer(observation, add_special_tokens=False)["input_ids"]
|
||||
if new_ids and new_ids[-1] == tokenizer.eos_token_id and obs_delta[:1] == [tokenizer.eos_token_id]:
|
||||
obs_delta = obs_delta[1:]
|
||||
observe_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
|
||||
observe_ids = tokenizer(observe_context, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
|
||||
current_len = len(prompt_ids) + len(response_ids)
|
||||
obs_delta = observe_ids[current_len:]
|
||||
response_ids.extend(obs_delta)
|
||||
response_mask.extend([0] * len(obs_delta))
|
||||
response_old_logps.extend([0.0] * len(obs_delta))
|
||||
|
|
@ -268,13 +272,14 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
|
|||
prompt_lens = torch.tensor([prompt_len for _, _, prompt_len, _ in packed_samples], device=args.device)
|
||||
full_response_masks = torch.tensor([mask + [0] * (max_len - len(mask)) for _, mask, _, _ in packed_samples], device=args.device, dtype=torch.float32)
|
||||
old_per_token_logps = torch.tensor([old_logps + [0.0] * ((max_len - 1) - len(old_logps)) for _, _, _, old_logps in packed_samples], device=args.device, dtype=torch.float32)
|
||||
full_mask = (input_ids != tokenizer.pad_token_id).long()
|
||||
full_mask = (torch.arange(max_len, device=args.device).unsqueeze(0) < seq_lens.unsqueeze(1)).long()
|
||||
|
||||
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)
|
||||
|
|
@ -330,11 +335,13 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
|
|||
policy_loss = (((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean()
|
||||
if valid_rows.any() else per_token_loss.sum() * 0.0)
|
||||
loss = (policy_loss + aux_loss) / args.accumulation_steps
|
||||
loss.backward()
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
if step % args.accumulation_steps == 0:
|
||||
if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step(); scheduler.step(); optimizer.zero_grad()
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad()
|
||||
|
||||
if step % args.log_interval == 0 or step == iters:
|
||||
pl = loss.item() * args.accumulation_steps
|
||||
|
|
@ -357,7 +364,8 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
|
|||
state_dict = raw_model.state_dict()
|
||||
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
|
||||
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
|
||||
scheduler=scheduler, scaler=scaler)
|
||||
model.train()
|
||||
del state_dict
|
||||
|
||||
|
|
@ -367,8 +375,10 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
|
|||
del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
|
||||
|
||||
if last_step > start_step and last_step % args.accumulation_steps != 0:
|
||||
if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step(); scheduler.step(); optimizer.zero_grad()
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
@ -378,8 +388,8 @@ 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("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="训练数据类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -412,7 +422,7 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agent", help="SGLang共享存储路径")
|
||||
args = parser.parse_args()
|
||||
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -421,9 +431,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
wandb = None
|
||||
if args.use_wandb and is_main_process():
|
||||
|
|
@ -464,6 +473,7 @@ if __name__ == "__main__":
|
|||
model.load_state_dict(ckp_data['model'])
|
||||
optimizer.load_state_dict(ckp_data['optimizer'])
|
||||
scheduler.load_state_dict(ckp_data['scheduler'])
|
||||
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
|
||||
start_epoch = ckp_data['epoch']
|
||||
start_step = ckp_data.get('step', 0)
|
||||
|
||||
|
|
@ -472,7 +482,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):
|
||||
|
|
@ -480,7 +491,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device), collate_fn=collate_fn)
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch+1}/{args.epochs}]: skip {start_step} steps')
|
||||
rl_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
|
||||
|
|
|
|||
|
|
@ -11,13 +11,16 @@ import warnings
|
|||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -151,8 +154,8 @@ 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("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -177,7 +180,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -188,9 +191,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -210,7 +212,6 @@ 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'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -235,7 +236,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, teacher_model, lm_config_student, start_step, wandb, args.alpha, args.temperature)
|
||||
|
|
|
|||
|
|
@ -11,13 +11,16 @@ import warnings
|
|||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import DPODataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -135,8 +138,8 @@ 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("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -156,7 +159,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -166,9 +169,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -190,7 +192,6 @@ 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'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -215,7 +216,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, ref_model, lm_config, start_step, wandb, args.beta)
|
||||
|
|
|
|||
|
|
@ -10,13 +10,16 @@ import time
|
|||
import warnings
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -88,8 +91,8 @@ 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("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -109,7 +112,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -119,9 +122,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -136,7 +138,6 @@ 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'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -161,7 +162,8 @@ if __name__ == "__main__":
|
|||
setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
|
||||
|
|
|
|||
|
|
@ -14,7 +14,6 @@ import torch
|
|||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from transformers import AutoTokenizer
|
||||
from contextlib import nullcontext
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
|
|
@ -22,7 +21,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
|
|||
from transformers import AutoModel
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import RLAIFDataset
|
||||
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
|
||||
from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed,
|
||||
SkipBatchSampler, init_model, LMForRewardModel,
|
||||
get_default_device, get_autocast_context,
|
||||
get_grad_scaler, should_pin_memory)
|
||||
from trainer.rollout_engine import create_rollout_engine
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
|
@ -91,12 +94,14 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
prompt_lens = rollout_result.prompt_lens.to(args.device)
|
||||
full_mask = (outputs != tokenizer.pad_token_id).long()
|
||||
logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0)
|
||||
full_mask.scatter_(1, logp_pos + 1, rollout_result.completion_mask.to(args.device, dtype=full_mask.dtype))
|
||||
|
||||
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)
|
||||
|
||||
|
|
@ -142,12 +147,14 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
|
||||
policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean()
|
||||
loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar
|
||||
loss.backward()
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
if step % args.accumulation_steps == 0:
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
|
|
@ -186,7 +193,8 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
state_dict = raw_model.state_dict()
|
||||
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
|
||||
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
|
||||
scheduler=scheduler, scaler=scaler)
|
||||
model.train()
|
||||
del state_dict
|
||||
|
||||
|
|
@ -197,8 +205,10 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
|
|||
|
||||
if step > start_step and step % args.accumulation_steps != 0:
|
||||
if args.grad_clip > 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
|
|
@ -210,8 +220,8 @@ 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("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -244,7 +254,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -255,9 +265,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -303,6 +312,7 @@ if __name__ == "__main__":
|
|||
model.load_state_dict(ckp_data['model'])
|
||||
optimizer.load_state_dict(ckp_data['optimizer'])
|
||||
scheduler.load_state_dict(ckp_data['scheduler'])
|
||||
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
|
||||
start_epoch = ckp_data['epoch']
|
||||
start_step = ckp_data.get('step', 0)
|
||||
|
||||
|
|
@ -312,7 +322,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. 开始训练 ==========
|
||||
|
|
@ -321,7 +332,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
|
||||
|
|
|
|||
|
|
@ -10,14 +10,17 @@ import time
|
|||
import warnings
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from model.model_lora import save_lora, apply_lora
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -82,8 +85,8 @@ 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("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -102,7 +105,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -112,9 +115,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -148,7 +150,6 @@ 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'))
|
||||
optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
|
||||
|
||||
# ========== 7. 从ckp恢复状态 ==========
|
||||
|
|
@ -173,7 +174,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, lora_params, start_step, wandb)
|
||||
|
|
|
|||
|
|
@ -12,7 +12,6 @@ import warnings
|
|||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
|
|
@ -20,7 +19,11 @@ from torch.nn.utils import clip_grad_norm_
|
|||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import RLAIFDataset
|
||||
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
|
||||
from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed,
|
||||
SkipBatchSampler, init_model, LMForRewardModel,
|
||||
get_default_device, get_autocast_context,
|
||||
get_grad_scaler, should_pin_memory)
|
||||
from trainer.rollout_engine import create_rollout_engine
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
|
@ -36,7 +39,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):
|
||||
|
|
@ -120,6 +123,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
|
|||
resp_idx = torch.arange(resp_labels.size(1), device=gen_out.device).unsqueeze(0)
|
||||
logp_pos = prompt_lens.unsqueeze(1) - 1 + resp_idx
|
||||
resp_pad_mask = rollout_result.completion_mask.to(args.device).bool()
|
||||
full_mask.scatter_(1, logp_pos + 1, resp_pad_mask.to(full_mask.dtype))
|
||||
resp_lengths = resp_pad_mask.sum(dim=1); valid_resp = resp_lengths > 0; eos_mask = resp_labels.eq(tokenizer.eos_token_id) & resp_pad_mask
|
||||
has_eos = eos_mask.any(dim=1); eos_pos = torch.argmax(eos_mask.int(), dim=1)
|
||||
resp_lengths = torch.where(has_eos, eos_pos + 1, resp_lengths).long().clamp(min=1)
|
||||
|
|
@ -158,8 +162,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
|
||||
|
|
@ -167,11 +169,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
|
||||
# 计算造成额外数值偏差。
|
||||
|
|
@ -190,7 +194,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 死锁
|
||||
|
|
@ -224,7 +228,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
|
|||
else:
|
||||
loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps
|
||||
|
||||
loss.backward()
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
policy_loss_sum += policy_loss.item()
|
||||
value_loss_sum += value_loss.item()
|
||||
|
|
@ -237,20 +241,26 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
|
|||
grad_accum_step += 1
|
||||
|
||||
if grad_accum_step % args.accumulation_steps == 0:
|
||||
scaler.unscale_(actor_optimizer)
|
||||
scaler.unscale_(critic_optimizer)
|
||||
clip_grad_norm_(actor_model.parameters(), args.grad_clip)
|
||||
clip_grad_norm_(critic_model.parameters(), args.grad_clip)
|
||||
actor_optimizer.step()
|
||||
critic_optimizer.step()
|
||||
scaler.step(actor_optimizer)
|
||||
scaler.step(critic_optimizer)
|
||||
scaler.update()
|
||||
actor_scheduler.step()
|
||||
critic_scheduler.step()
|
||||
actor_optimizer.zero_grad()
|
||||
critic_optimizer.zero_grad()
|
||||
|
||||
if grad_accum_step % args.accumulation_steps != 0:
|
||||
scaler.unscale_(actor_optimizer)
|
||||
scaler.unscale_(critic_optimizer)
|
||||
clip_grad_norm_(actor_model.parameters(), args.grad_clip)
|
||||
clip_grad_norm_(critic_model.parameters(), args.grad_clip)
|
||||
actor_optimizer.step()
|
||||
critic_optimizer.step()
|
||||
scaler.step(actor_optimizer)
|
||||
scaler.step(critic_optimizer)
|
||||
scaler.update()
|
||||
actor_scheduler.step()
|
||||
critic_scheduler.step()
|
||||
actor_optimizer.zero_grad()
|
||||
|
|
@ -297,7 +307,8 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
|
|||
lm_checkpoint(lm_config, weight=args.save_weight, model=actor_model, optimizer=actor_optimizer,
|
||||
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
|
||||
scheduler=actor_scheduler, critic_model=critic_model,
|
||||
critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler)
|
||||
critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler,
|
||||
scaler=scaler)
|
||||
actor_model.train()
|
||||
del actor_state
|
||||
|
||||
|
|
@ -314,8 +325,8 @@ 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("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -353,7 +364,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -363,9 +374,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -419,6 +429,7 @@ if __name__ == "__main__":
|
|||
critic_optimizer.load_state_dict(ckp_data['critic_optimizer'])
|
||||
actor_scheduler.load_state_dict(ckp_data['scheduler'])
|
||||
critic_scheduler.load_state_dict(ckp_data['critic_scheduler'])
|
||||
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
|
||||
start_epoch = ckp_data['epoch']
|
||||
start_step = ckp_data.get('step', 0)
|
||||
|
||||
|
|
@ -428,8 +439,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. 开始训练 ==========
|
||||
|
|
@ -438,7 +450,8 @@ if __name__ == "__main__":
|
|||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb)
|
||||
|
|
|
|||
|
|
@ -10,13 +10,16 @@ import time
|
|||
import warnings
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from model.model_minimind import MiniMindConfig
|
||||
from dataset.lm_dataset import PretrainDataset
|
||||
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
|
||||
from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
|
||||
init_distributed_mode, setup_seed, init_model,
|
||||
SkipBatchSampler, get_default_device,
|
||||
get_autocast_context, get_grad_scaler,
|
||||
should_pin_memory)
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -87,8 +90,8 @@ 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("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||
parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||
|
|
@ -108,7 +111,7 @@ if __name__ == "__main__":
|
|||
args = parser.parse_args()
|
||||
|
||||
# ========== 1. 初始化环境和随机种子 ==========
|
||||
local_rank = init_distributed_mode()
|
||||
local_rank = init_distributed_mode(args.device)
|
||||
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||
setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
|
||||
|
||||
|
|
@ -118,9 +121,8 @@ 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"
|
||||
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 = get_autocast_context(args.device, args.dtype)
|
||||
scaler = get_grad_scaler(args.device, args.dtype)
|
||||
|
||||
# ========== 4. 配wandb ==========
|
||||
wandb = None
|
||||
|
|
@ -135,7 +137,6 @@ 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'))
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
# ========== 6. 从ckp恢复状态 ==========
|
||||
|
|
@ -160,7 +161,8 @@ if __name__ == "__main__":
|
|||
setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
|
||||
pin_memory=should_pin_memory(args.device))
|
||||
if skip > 0:
|
||||
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
|
||||
|
|
|
|||
|
|
@ -5,8 +5,11 @@ 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
|
||||
from contextlib import nullcontext
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
|
@ -15,6 +18,91 @@ from torch.utils.data import Sampler
|
|||
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
|
||||
from model.model_minimind import MiniMindForCausalLM
|
||||
|
||||
|
||||
SUPPORTED_DTYPES = {
|
||||
'float32': torch.float32,
|
||||
'float16': torch.float16,
|
||||
'bfloat16': torch.bfloat16,
|
||||
}
|
||||
|
||||
|
||||
def get_device_type(device):
|
||||
"""Return the accelerator type for strings such as cuda:0, mps, and cpu."""
|
||||
return torch.device(device).type
|
||||
|
||||
|
||||
def get_default_device():
|
||||
"""Prefer CUDA, then Apple Silicon MPS, and finally CPU."""
|
||||
if torch.cuda.is_available():
|
||||
return 'cuda:0'
|
||||
if torch.backends.mps.is_available():
|
||||
return 'mps'
|
||||
return 'cpu'
|
||||
|
||||
|
||||
def validate_device(device):
|
||||
device_type = get_device_type(device)
|
||||
if device_type == 'cuda' and not torch.cuda.is_available():
|
||||
raise RuntimeError('CUDA was requested but is not available.')
|
||||
if device_type == 'mps' and not torch.backends.mps.is_available():
|
||||
reason = ('this PyTorch build has no MPS support'
|
||||
if not torch.backends.mps.is_built()
|
||||
else 'the current macOS device does not expose MPS')
|
||||
raise RuntimeError(f'MPS was requested but is not available: {reason}.')
|
||||
if device_type not in {'cuda', 'mps', 'cpu'}:
|
||||
raise ValueError(f'Unsupported training device: {device}')
|
||||
return device_type
|
||||
|
||||
|
||||
def get_torch_dtype(dtype):
|
||||
try:
|
||||
return SUPPORTED_DTYPES[dtype]
|
||||
except KeyError as exc:
|
||||
choices = ', '.join(SUPPORTED_DTYPES)
|
||||
raise ValueError(f'Unsupported dtype {dtype!r}; choose one of: {choices}') from exc
|
||||
|
||||
|
||||
def get_autocast_context(device, dtype):
|
||||
"""Create the appropriate AMP context for CUDA or Apple MPS."""
|
||||
device_type = validate_device(device)
|
||||
torch_dtype = get_torch_dtype(dtype)
|
||||
if device_type == 'cpu' or torch_dtype == torch.float32:
|
||||
return nullcontext()
|
||||
return torch.autocast(device_type=device_type, dtype=torch_dtype)
|
||||
|
||||
|
||||
def get_grad_scaler(device, dtype):
|
||||
"""Enable gradient scaling for fp16 CUDA/MPS and use a no-op otherwise."""
|
||||
device_type = validate_device(device)
|
||||
enabled = dtype == 'float16' and device_type in {'cuda', 'mps'}
|
||||
try:
|
||||
return torch.amp.GradScaler(device_type, enabled=enabled)
|
||||
except (TypeError, RuntimeError):
|
||||
# Older PyTorch releases only expose a CUDA GradScaler. Keep CUDA
|
||||
# compatibility and fail clearly for MPS instead of silently training
|
||||
# fp16 without scaling.
|
||||
if device_type == 'cuda':
|
||||
return torch.cuda.amp.GradScaler(enabled=enabled)
|
||||
if enabled:
|
||||
raise RuntimeError(
|
||||
'This PyTorch version cannot scale MPS float16 gradients. '
|
||||
'Upgrade PyTorch or use --dtype bfloat16/float32.'
|
||||
)
|
||||
return torch.cuda.amp.GradScaler(enabled=False)
|
||||
|
||||
|
||||
def should_pin_memory(device):
|
||||
"""Pinned host memory accelerates CUDA copies but does not benefit MPS."""
|
||||
return get_device_type(device) == 'cuda'
|
||||
|
||||
|
||||
def empty_device_cache(device):
|
||||
device_type = get_device_type(device)
|
||||
if device_type == 'cuda' and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
elif device_type == 'mps' and torch.backends.mps.is_available():
|
||||
torch.mps.empty_cache()
|
||||
|
||||
def get_model_params(model, config):
|
||||
total = sum(p.numel() for p in model.parameters()) / 1e6
|
||||
n_routed = getattr(config, 'n_routed_experts', getattr(config, 'num_experts', 0))
|
||||
|
|
@ -41,10 +129,17 @@ def get_lr(current_step, total_steps, lr):
|
|||
return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
def init_distributed_mode(device=None):
|
||||
if int(os.environ.get("RANK", -1)) == -1:
|
||||
return 0 # 非DDP模式
|
||||
|
||||
device = device or get_default_device()
|
||||
if get_device_type(device) != 'cuda':
|
||||
raise RuntimeError(
|
||||
'Distributed training in MiniMind currently requires CUDA/NCCL. '
|
||||
'Run MPS training as a single process with `python train_xxx.py --device mps`.'
|
||||
)
|
||||
|
||||
dist.init_process_group(backend="nccl")
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.cuda.set_device(local_rank)
|
||||
|
|
@ -55,10 +150,13 @@ def setup_seed(seed: int):
|
|||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
if torch.backends.mps.is_available():
|
||||
torch.mps.manual_seed(seed)
|
||||
|
||||
def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
|
@ -69,6 +167,7 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
|
|||
if model is not None:
|
||||
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
|
||||
raw_model = getattr(raw_model, '_orig_mod', raw_model)
|
||||
model_device = next(raw_model.parameters()).device
|
||||
state_dict = raw_model.state_dict()
|
||||
state_dict = {k: v.half().cpu() for k, v in state_dict.items()}
|
||||
ckp_tmp = ckp_path + '.tmp'
|
||||
|
|
@ -103,7 +202,7 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
|
|||
torch.save(resume_data, resume_tmp)
|
||||
os.replace(resume_tmp, resume_path)
|
||||
del state_dict, resume_data
|
||||
torch.cuda.empty_cache()
|
||||
empty_device_cache(model_device)
|
||||
else: # 加载模式
|
||||
if os.path.exists(resume_path):
|
||||
ckp_data = torch.load(resume_path, map_location='cpu')
|
||||
|
|
@ -116,7 +215,8 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
|
|||
return None
|
||||
|
||||
|
||||
def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'):
|
||||
def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device=None):
|
||||
device = device or get_default_device()
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
|
||||
|
|
@ -158,7 +258,8 @@ class SkipBatchSampler(Sampler):
|
|||
|
||||
|
||||
class LMForRewardModel:
|
||||
def __init__(self, model_path, device="cuda", dtype=torch.float16):
|
||||
def __init__(self, model_path, device=None, dtype=torch.float16):
|
||||
device = device or get_default_device()
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
self.model = AutoModel.from_pretrained(model_path, torch_dtype=dtype, trust_remote_code=True)
|
||||
self.model = self.model.to(device).eval()
|
||||
|
|
@ -175,3 +276,31 @@ class LMForRewardModel:
|
|||
]
|
||||
score = self.model.get_score(self.tokenizer, eval_messages)
|
||||
return max(min(score, 3.0), -3.0)
|
||||
# ===== 数学表达式安全求值：替代 eval，只放行算术运算与 math 白名单（长度上限 512） =====
|
||||
def safe_math_eval(expression):
|
||||
"""对模型生成的数学表达式求值：支持 + - * / // % **、math 白名单函数与 pi/e/tau 常量。"""
|
||||
def pow_guard(base, exp): # 幂运算统一走这里，拦截 9**9**9 / 10**99999 这类把进程算死的输入
|
||||
if abs(exp) > 1e4 or (abs(base) > 1 and abs(exp) * math.log10(abs(base)) > 100): raise ValueError('幂运算结果过大')
|
||||
return base ** exp
|
||||
def resolve(node): # 同时兼容 sqrt(4) 与 math.sqrt(4) 两种写法
|
||||
if isinstance(node, ast.Name): return node.id
|
||||
if isinstance(node, ast.Attribute) and getattr(node.value, 'id', '') == 'math': return node.attr
|
||||
def walk(node):
|
||||
if isinstance(node, ast.Constant) and type(node.value) in (int, float): return node.value
|
||||
name = resolve(node)
|
||||
if name in consts: return consts[name]
|
||||
if isinstance(node, ast.UnaryOp): return unary_ops[type(node.op)](walk(node.operand))
|
||||
if isinstance(node, ast.BinOp): return bin_ops[type(node.op)](walk(node.left), walk(node.right))
|
||||
if isinstance(node, ast.Call) and not node.keywords: return funcs[resolve(node.func)](*map(walk, node.args))
|
||||
raise ValueError(f'不支持的表达式语法: {type(node).__name__}')
|
||||
consts = {'pi': math.pi, 'e': math.e, 'tau': math.tau}
|
||||
funcs = {'pow': pow_guard, **{n: getattr(math, n) for n in 'sqrt exp log log2 log10 sin cos tan asin acos atan atan2 sinh cosh tanh floor ceil trunc fabs fmod hypot gcd degrees radians'.split()}}
|
||||
bin_ops = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod, ast.Pow: pow_guard}
|
||||
unary_ops = {ast.UAdd: operator.pos, ast.USub: operator.neg}
|
||||
chars = str.maketrans({'^': '**', '×': '*', '÷': '/', '−': '-', '²': '**2', '³': '**3', '（': '(', '）': ')'})
|
||||
expr = str(expression).translate(chars).strip()
|
||||
if not expr or len(expr) > 512: raise ValueError('表达式为空或过长')
|
||||
try:
|
||||
return walk(ast.parse(expr, mode='eval').body)
|
||||
except (KeyError, SyntaxError, TypeError):
|
||||
raise ValueError('不支持的表达式语法') from None
|
||||
|
|
|
|||
Loading…
Reference in New Issue