From 618b72dc2a5cc6193e06e09e6be8bff90b803945 Mon Sep 17 00:00:00 2001
From: mengxin <zhangxiaoshuai@xiaohongshu.com>
Date: Sun, 6 Sep 2026 02:41:29 +0800
Subject: [PATCH 1/4] feat: add device-agnostic training utilities

---
 trainer/trainer_utils.py | 117 +++++++++++++++++++++++++++++++++++----
 1 file changed, 107 insertions(+), 10 deletions(-)

diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py
index b171289..10ec4ad 100644
--- a/trainer/trainer_utils.py
+++ b/trainer/trainer_utils.py
@@ -9,6 +9,7 @@ import ast
 import operator
 import random
 import math
+from contextlib import nullcontext
 import numpy as np
 import torch
 import torch.distributed as dist
@@ -17,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))
@@ -43,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)
@@ -57,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)
@@ -71,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'
@@ -105,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')
@@ -118,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)
 
@@ -160,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()
@@ -177,8 +276,6 @@ 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 常量。"""

From 7b554a516b47de4c61f7c6d53993adbdc71f21e1 Mon Sep 17 00:00:00 2001
From: mengxin <zhangxiaoshuai@xiaohongshu.com>
Date: Sun, 6 Sep 2026 02:41:42 +0800
Subject: [PATCH 2/4] feat: enable MPS across training entrypoints

---
 eval_llm.py                   |  6 ++---
 trainer/train_agent.py        | 40 ++++++++++++++++++++------------
 trainer/train_distillation.py | 24 ++++++++++---------
 trainer/train_dpo.py          | 24 ++++++++++---------
 trainer/train_full_sft.py     | 24 ++++++++++---------
 trainer/train_grpo.py         | 37 ++++++++++++++++++------------
 trainer/train_lora.py         | 24 ++++++++++---------
 trainer/train_ppo.py          | 43 ++++++++++++++++++++++-------------
 trainer/train_pretrain.py     | 24 ++++++++++---------
 9 files changed, 143 insertions(+), 103 deletions(-)

diff --git a/eval_llm.py b/eval_llm.py
index 9359ca8..54a6dbc 100755
--- 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()
\ No newline at end of file
+    main()
diff --git a/trainer/train_agent.py b/trainer/train_agent.py
index 63a0782..0e815fd 100644
--- a/trainer/train_agent.py
+++ b/trainer/train_agent.py
@@ -15,7 +15,6 @@ 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
@@ -23,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, safe_math_eval
+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')
@@ -331,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
@@ -358,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
 
@@ -368,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__":
@@ -379,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="梯度裁剪阈值")
@@ -413,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))
 
@@ -422,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():
@@ -465,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)
 
@@ -482,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"))
diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py
index 085b2a7..4ebeecc 100644
--- 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()
\ No newline at end of file
+        dist.destroy_process_group()
diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py
index 479eb85..9bf58cf 100644
--- 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()
\ No newline at end of file
+        dist.destroy_process_group()
diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py
index c22eb80..5aa7bb7 100644
--- 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()
\ No newline at end of file
+        dist.destroy_process_group()
diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py
index ad9aa3b..3d34eb3 100755
--- 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')
@@ -144,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()
 
@@ -188,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
 
@@ -199,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()
 
@@ -212,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="梯度裁剪阈值")
@@ -246,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))
     
@@ -257,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
@@ -305,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)
     
@@ -324,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"))
@@ -334,4 +343,4 @@ if __name__ == "__main__":
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized():
         dist.barrier()
-        dist.destroy_process_group()
\ No newline at end of file
+        dist.destroy_process_group()
diff --git a/trainer/train_lora.py b/trainer/train_lora.py
index 19a1499..455fd0f 100644
--- 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()
\ No newline at end of file
+        dist.destroy_process_group()
diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py
index 4d762b6..7a97d8c 100644
--- 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')
@@ -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)
     
@@ -440,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)
@@ -450,4 +461,4 @@ if __name__ == "__main__":
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized():
         dist.barrier()
-        dist.destroy_process_group()
\ No newline at end of file
+        dist.destroy_process_group()
diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py
index 83a0382..62bca57 100644
--- 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()
\ No newline at end of file
+        dist.destroy_process_group()

From 791b5150cc76e7b2606318053cde700865f43b77 Mon Sep 17 00:00:00 2001
From: mengxin <zhangxiaoshuai@xiaohongshu.com>
Date: Sun, 6 Sep 2026 02:41:54 +0800
Subject: [PATCH 3/4] test: cover MPS device and training paths

---
 tests/benchmark_mps.py              | 68 ++++++++++++++++++++
 tests/fixtures/pretrain_smoke.jsonl |  8 +++
 tests/fixtures/sft_smoke.jsonl      |  4 ++
 tests/test_device_utils.py          | 98 +++++++++++++++++++++++++++++
 4 files changed, 178 insertions(+)
 create mode 100644 tests/benchmark_mps.py
 create mode 100644 tests/fixtures/pretrain_smoke.jsonl
 create mode 100644 tests/fixtures/sft_smoke.jsonl
 create mode 100644 tests/test_device_utils.py

diff --git a/tests/benchmark_mps.py b/tests/benchmark_mps.py
new file mode 100644
index 0000000..c9037af
--- /dev/null
+++ 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()
diff --git a/tests/fixtures/pretrain_smoke.jsonl b/tests/fixtures/pretrain_smoke.jsonl
new file mode 100644
index 0000000..743f55b
--- /dev/null
+++ b/tests/fixtures/pretrain_smoke.jsonl
@@ -0,0 +1,8 @@
+{"text":"苹果芯片可以通过 MPS 加速 PyTorch 模型训练。"}
+{"text":"MiniMind 是一个用于学习语言模型训练流程的小型项目。"}
+{"text":"梯度下降通过反向传播更新神经网络参数。"}
+{"text":"混合精度训练可以降低内存占用并提高吞吐量。"}
+{"text":"统一内存允许 CPU 和 GPU 共享同一块物理内存。"}
+{"text":"检查点可用于恢复模型参数和优化器状态。"}
+{"text":"注意力机制根据上下文计算不同词元之间的关系。"}
+{"text":"小规模冒烟测试可以快速发现训练环境问题。"}
diff --git a/tests/fixtures/sft_smoke.jsonl b/tests/fixtures/sft_smoke.jsonl
new file mode 100644
index 0000000..f39c67d
--- /dev/null
+++ 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":""}]}
diff --git a/tests/test_device_utils.py b/tests/test_device_utils.py
new file mode 100644
index 0000000..78f8184
--- /dev/null
+++ 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()

From 856a57cacb5483ea70d6c23273302e1ec90bc8b9 Mon Sep 17 00:00:00 2001
From: mengxin <zhangxiaoshuai@xiaohongshu.com>
Date: Sun, 6 Sep 2026 02:42:05 +0800
Subject: [PATCH 4/4] docs: document Apple Silicon MPS training

---
 .gitignore |  9 ++++++---
 README.md  | 13 +++++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 7956e1d..503b3b2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,9 @@
+.venv*/
 __pycache__/
-*.pyc
+*.py[cod]
 .DS_Store
-out
+out/
+checkpoints/
+.smoke-out/
 website/
-docs-minimind/
\ No newline at end of file
+docs-minimind/
diff --git a/README.md b/README.md
index dadc82c..f703be3 100644
--- 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' 下载数据