feat: add device-agnostic training utilities
mengxin
parent
6fc918beb6
commit
3d4d313617
|
|
@ -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)
|
||||
return max(min(score, 3.0), -3.0)
|
||||
|
|
|
|||
Loading…
Reference in New Issue