Merge 856a57cacb into 7a9137d2e9

docs: document Apple Silicon MPS training
test: cover MPS device and training paths
2026-09-16 21:10:35 +08:00 · 2026-09-16 21:08:40 +08:00 · 2026-09-16 21:08:40 +08:00 · 2026-09-16 21:08:40 +08:00 · 2026-09-16 21:08:19 +08:00 · 2026-09-15 14:53:11 +08:00
18 changed files with 506 additions and 140 deletions
--- a/.gitignore
+++ b/.gitignore
@ -1,6 +1,9 @@
+.venv*/
 __pycache__/
-*.pyc
+*.py[cod]
 .DS_Store
-out
+out/
+checkpoints/
+.smoke-out/
 website/
-docs-minimind/
+docs-minimind/
--- a/README.md
+++ b/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' 下载数据
--- a/eval_llm.py
+++ b/eval_llm.py
@ -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 = [
@ -91,4 +91,4 @@ def main():
        print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n')

 if __name__ == "__main__":
-    main()
+    main()
--- a/scripts/eval_toolcall.py
+++ b/scripts/eval_toolcall.py
@ -12,7 +12,7 @@ from datetime import datetime
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from openai import OpenAI
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
-from trainer.trainer_utils import setup_seed, get_model_params
+from trainer.trainer_utils import setup_seed, get_model_params, safe_math_eval
 warnings.filterwarnings('ignore')

 TOOLS = [
@ -27,7 +27,7 @@ TOOLS = [
 ]

 MOCK_RESULTS = {
-    "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("²", "**2").replace("³", "**3").replace("（", "(").replace("）", ")")))},
+    "calculate_math": lambda args: {"result": str(safe_math_eval(args.get("expression", "0")))},
    "get_current_time": lambda args: {"datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "timezone": args.get("timezone", "Asia/Shanghai")},
    "random_number": lambda args: {"result": random.randint(int(args.get("min", 0)), int(args.get("max", 100)))},
    "text_length": lambda args: {"characters": len(args.get("text", "")), "words": len(args.get("text", "").split())},
--- a/scripts/web_demo.py
+++ b/scripts/web_demo.py
@ -2,6 +2,7 @@ import random
 import re
 import json
 import os
+import sys
 from threading import Thread

 import torch
@ -9,6 +10,9 @@ import numpy as np
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from trainer.trainer_utils import safe_math_eval
+
 st.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed")

 st.markdown("""
@ -125,7 +129,7 @@ def execute_tool(tool_name, args):
    import datetime
    try:
        if tool_name == 'calculate_math':
-            return {"result": eval(args.get('expression', '0'))}
+            return {"result": safe_math_eval(args.get('expression', '0'))}
        elif tool_name == 'get_current_time':
            tz = args.get('timezone', 'Asia/Shanghai')
            return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
--- a/tests/benchmark_mps.py
+++ b/tests/benchmark_mps.py
@ -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()
--- a/tests/fixtures/pretrain_smoke.jsonl
+++ b/tests/fixtures/pretrain_smoke.jsonl
@ -0,0 +1,8 @@
+{"text":"苹果芯片可以通过 MPS 加速 PyTorch 模型训练。"}
+{"text":"MiniMind 是一个用于学习语言模型训练流程的小型项目。"}
+{"text":"梯度下降通过反向传播更新神经网络参数。"}
+{"text":"混合精度训练可以降低内存占用并提高吞吐量。"}
+{"text":"统一内存允许 CPU 和 GPU 共享同一块物理内存。"}
+{"text":"检查点可用于恢复模型参数和优化器状态。"}
+{"text":"注意力机制根据上下文计算不同词元之间的关系。"}
+{"text":"小规模冒烟测试可以快速发现训练环境问题。"}
--- a/tests/fixtures/sft_smoke.jsonl
+++ b/tests/fixtures/sft_smoke.jsonl
@ -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":""}]}
--- a/tests/test_device_utils.py
+++ b/tests/test_device_utils.py
@ -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()
--- a/trainer/train_agent.py
+++ b/trainer/train_agent.py
@ -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"):
@ -278,9 +276,10 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model

        rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch)

-        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
        with autocast_ctx:
-            res = model_unwrapped(input_ids, attention_mask=full_mask)
+            # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
+            # DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
+            res = model(input_ids, attention_mask=full_mask)
            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
            logits = res.logits[:, :-1, :]
            per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
@ -336,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
@ -363,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

@ -373,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__":
@ -384,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="梯度裁剪阈值")
@ -418,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))

@ -427,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():
@ -470,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)

@ -478,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):
@ -486,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"))
--- a/trainer/train_distillation.py
+++ b/trainer/train_distillation.py
@ -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)
@ -245,4 +247,4 @@ if __name__ == "__main__":
    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()
--- a/trainer/train_dpo.py
+++ b/trainer/train_dpo.py
@ -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)
@ -225,4 +227,4 @@ if __name__ == "__main__":
    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()
--- a/trainer/train_full_sft.py
+++ b/trainer/train_full_sft.py
@ -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)
@ -171,4 +173,4 @@ if __name__ == "__main__":
    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()
--- a/trainer/train_grpo.py
+++ b/trainer/train_grpo.py
@ -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')
@ -95,9 +98,10 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod

        rewards = calculate_rewards(prompts, completions, reward_model).to(args.device)  # [B*num_gen]

-        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
        with autocast_ctx:
-            res = model_unwrapped(outputs, attention_mask=full_mask)
+            # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
+            # DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
+            res = model(outputs, attention_mask=full_mask)
            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
            per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)

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

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

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

@ -211,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="梯度裁剪阈值")
@ -245,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))
    
@ -256,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
@ -304,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)
    
@ -313,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. 开始训练 ==========
@ -322,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"))
@ -332,4 +343,4 @@ if __name__ == "__main__":
    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()
--- a/trainer/train_lora.py
+++ b/trainer/train_lora.py
@ -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)
@ -183,4 +185,4 @@ if __name__ == "__main__":
    # ========== 10. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()
--- a/trainer/train_ppo.py
+++ b/trainer/train_ppo.py
@ -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):
@ -159,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
@ -168,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
                    # 计算造成额外数值偏差。
@ -191,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 死锁
@ -225,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()
@ -238,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()
@ -298,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

@ -315,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="梯度裁剪阈值")
@ -354,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))
    
@ -364,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
@ -420,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)
    
@ -429,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. 开始训练 ==========
@ -439,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)
@ -449,4 +461,4 @@ if __name__ == "__main__":
    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()
--- a/trainer/train_pretrain.py
+++ b/trainer/train_pretrain.py
@ -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)
@ -170,4 +172,4 @@ if __name__ == "__main__":
    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()
--- a/trainer/trainer_utils.py
+++ b/trainer/trainer_utils.py
@ -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()
@ -174,4 +275,32 @@ class LMForRewardModel:
            {"role": "assistant", "content": response}
        ]
        score = self.model.get_score(self.tokenizer, eval_messages)
-        return max(min(score, 3.0), -3.0)
+        return max(min(score, 3.0), -3.0)
+# ===== 数学表达式安全求值：替代 eval，只放行算术运算与 math 白名单（长度上限 512） =====
+def safe_math_eval(expression):
+    """对模型生成的数学表达式求值：支持 + - * / // % **、math 白名单函数与 pi/e/tau 常量。"""
+    def pow_guard(base, exp):  # 幂运算统一走这里，拦截 9**9**9 / 10**99999 这类把进程算死的输入
+        if abs(exp) > 1e4 or (abs(base) > 1 and abs(exp) * math.log10(abs(base)) > 100): raise ValueError('幂运算结果过大')
+        return base ** exp
+    def resolve(node):  # 同时兼容 sqrt(4) 与 math.sqrt(4) 两种写法
+        if isinstance(node, ast.Name): return node.id
+        if isinstance(node, ast.Attribute) and getattr(node.value, 'id', '') == 'math': return node.attr
+    def walk(node):
+        if isinstance(node, ast.Constant) and type(node.value) in (int, float): return node.value
+        name = resolve(node)
+        if name in consts: return consts[name]
+        if isinstance(node, ast.UnaryOp): return unary_ops[type(node.op)](walk(node.operand))
+        if isinstance(node, ast.BinOp): return bin_ops[type(node.op)](walk(node.left), walk(node.right))
+        if isinstance(node, ast.Call) and not node.keywords: return funcs[resolve(node.func)](*map(walk, node.args))
+        raise ValueError(f'不支持的表达式语法: {type(node).__name__}')
+    consts = {'pi': math.pi, 'e': math.e, 'tau': math.tau}
+    funcs = {'pow': pow_guard, **{n: getattr(math, n) for n in 'sqrt exp log log2 log10 sin cos tan asin acos atan atan2 sinh cosh tanh floor ceil trunc fabs fmod hypot gcd degrees radians'.split()}}
+    bin_ops = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod, ast.Pow: pow_guard}
+    unary_ops = {ast.UAdd: operator.pos, ast.USub: operator.neg}
+    chars = str.maketrans({'^': '**', '×': '*', '÷': '/', '−': '-', '²': '**2', '³': '**3', '（': '(', '）': ')'})
+    expr = str(expression).translate(chars).strip()
+    if not expr or len(expr) > 512: raise ValueError('表达式为空或过长')
+    try:
+        return walk(ast.parse(expr, mode='eval').body)
+    except (KeyError, SyntaxError, TypeError):
+        raise ValueError('不支持的表达式语法') from None
Author	SHA1	Message	Date
CharlesCheung996	86472c9900	Merge `856a57cacb` into `7a9137d2e9`	2026-09-16 21:10:35 +08:00
mengxin	856a57cacb	docs: document Apple Silicon MPS training	2026-09-16 21:08:40 +08:00
mengxin	791b5150cc	test: cover MPS device and training paths	2026-09-16 21:08:40 +08:00
mengxin	7b554a516b	feat: enable MPS across training entrypoints	2026-09-16 21:08:40 +08:00
mengxin	618b72dc2a	feat: add device-agnostic training utilities	2026-09-16 21:08:19 +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