diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index b6211d5..e3f617b 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -7,6 +7,7 @@ __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import random import math +from contextlib import nullcontext import numpy as np import torch import torch.distributed as dist @@ -15,6 +16,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 +127,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 +148,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 +165,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 +200,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 +213,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 +256,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 +273,4 @@ class LMForRewardModel: {"role": "assistant", "content": response} ] score = self.model.get_score(self.tokenizer, eval_messages) - return max(min(score, 3.0), -3.0) \ No newline at end of file + return max(min(score, 3.0), -3.0)