From 618b72dc2a5cc6193e06e09e6be8bff90b803945 Mon Sep 17 00:00:00 2001 From: mengxin Date: Sun, 6 Sep 2026 02:41:29 +0800 Subject: [PATCH] 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 常量。"""