""" 训练工具函数集合 """ import os import sys __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import random import math import numpy as np import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel from torch.utils.data import Sampler from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification from model.model_minimind import MiniMindForCausalLM 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)) n_active = getattr(config, 'num_experts_per_tok', 0) n_shared = getattr(config, 'n_shared_experts', 0) expert = sum(p.numel() for n, p in model.named_parameters() if 'mlp.experts.0.' in n) / 1e6 shared_expert = sum(p.numel() for n, p in model.named_parameters() if 'mlp.shared_experts.0.' in n) / 1e6 base = total - (expert * n_routed) - (shared_expert * n_shared) active = base + (expert * n_active) + (shared_expert * n_shared) if active < total: Logger(f'Model Params: {total:.2f}M-A{active:.2f}M') else: Logger(f'Model Params: {total:.2f}M') def is_main_process(): return not dist.is_initialized() or dist.get_rank() == 0 def Logger(content): if is_main_process(): print(content) def get_lr(current_step, total_steps, lr): return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps))) # A100/H100/3090 等常见卡的 bf16 理论峰值 TFLOPS。MoE 等价 FLOPs 是 active params 不是 total， # active params 算出来的 MFU 才有意义；MoE 的 active 比例由 num_experts_per_tok / num_experts 决定。 _GPU_PEAK_BF16_TFLOPS = { 'A100': 312, 'A100-SXM4-80GB': 312, 'A100-SXM4-40GB': 312, 'A100-PCIE-80GB': 312, 'A100-PCIE-40GB': 312, 'H100': 989, 'H100-SXM5-80GB': 989, 'H100-PCIE-80GB': 756, 'H100 80GB HBM3': 989, 'H200': 989, 'RTX 3090': 71, 'GeForce RTX 3090': 71, 'RTX 3090 Ti': 80, 'RTX 4090': 165, 'GeForce RTX 4090': 165, 'V100': 125, 'V100-SXM2-32GB': 125, 'L40': 181, 'L40S': 362, } def detect_gpu_peak_tflops_bf16(device_index: int = 0, override: float = None) -> float: """Look up the bf16 TensorCore peak for the current GPU; fall back to A100 (312).""" if override is not None and override > 0: return float(override) if not torch.cuda.is_available(): return 312.0 name = torch.cuda.get_device_name(device_index) for key, peak in _GPU_PEAK_BF16_TFLOPS.items(): if key in name: return float(peak) Logger(f'[MFU] Unknown GPU "{name}", falling back to A100 (312 TFLOPS). Pass --gpu_peak_tflops to override.') return 312.0 def compute_model_flops_per_token(model, config=None) -> float: """Compute the "model FLOPs" per training token (forward + backward). Standard approximation for dense transformers: 6 × P, where P is the total trainable params count. For MoE we use active params = base + experts_per_tok × per_expert. """ raw = model.module if isinstance(model, DistributedDataParallel) else model raw = getattr(raw, '_orig_mod', raw) total = sum(p.numel() for p in raw.parameters()) if config is not None and getattr(config, 'use_moe', False): n_routed = getattr(config, 'num_experts', 0) n_active = getattr(config, 'num_experts_per_tok', 1) if n_routed > 0: expert = sum(p.numel() for n, p in raw.named_parameters() if 'mlp.experts.0.' in n) base = total - expert * n_routed active = base + expert * n_active return 6.0 * active return 6.0 * total def init_distributed_mode(): if int(os.environ.get("RANK", -1)) == -1: return 0 # 非DDP模式 dist.init_process_group(backend="nccl") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) return local_rank def setup_seed(seed: int, deterministic: bool = False): # cudnn 决定性 + 禁 benchmark 是预训练吞吐杀手（A100 上 ~3-5x 慢）。 # 默认走 benchmark=True 让 cudnn 选最快 kernel；要严格复现传 deterministic=True。 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 = deterministic torch.backends.cudnn.benchmark = not deterministic # A100 TF32 矩阵乘加速（autocast bf16 路径之外的 fp32 matmul 才生效，但无副作用） torch.set_float32_matmul_precision('high') 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) moe_path = '_moe' if lm_config.use_moe else '' ckp_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}.pth' resume_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}_resume.pth' if model is not None: raw_model = model.module if isinstance(model, DistributedDataParallel) else model raw_model = getattr(raw_model, '_orig_mod', raw_model) state_dict = raw_model.state_dict() state_dict = {k: v.half().cpu() for k, v in state_dict.items()} ckp_tmp = ckp_path + '.tmp' torch.save(state_dict, ckp_tmp) os.replace(ckp_tmp, ckp_path) wandb_id = None if wandb: if hasattr(wandb, 'get_run'): run = wandb.get_run() wandb_id = getattr(run, 'id', None) if run else None else: wandb_id = getattr(wandb, 'id', None) resume_data = { 'model': state_dict, 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'step': step, 'world_size': dist.get_world_size() if dist.is_initialized() else 1, 'wandb_id': wandb_id } for key, value in kwargs.items(): if value is not None: if hasattr(value, 'state_dict'): raw_value = value.module if isinstance(value, DistributedDataParallel) else value raw_value = getattr(raw_value, '_orig_mod', raw_value) resume_data[key] = raw_value.state_dict() else: resume_data[key] = value resume_tmp = resume_path + '.tmp' torch.save(resume_data, resume_tmp) os.replace(resume_tmp, resume_path) del state_dict, resume_data torch.cuda.empty_cache() else: # 加载模式 if os.path.exists(resume_path): ckp_data = torch.load(resume_path, map_location='cpu') saved_ws = ckp_data.get('world_size', 1) current_ws = dist.get_world_size() if dist.is_initialized() else 1 if saved_ws != current_ws: ckp_data['step'] = ckp_data['step'] * saved_ws // current_ws Logger(f'GPU数量变化({saved_ws}→{current_ws})，step已自动转换为{ckp_data["step"]}') return ckp_data return None def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'): tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) model = MiniMindForCausalLM(lm_config) if from_weight!= 'none': moe_suffix = '_moe' if lm_config.use_moe else '' weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{moe_suffix}.pth' weights = torch.load(weight_path, map_location=device) model.load_state_dict(weights, strict=False) get_model_params(model, lm_config) Logger(f'Trainable Params: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f}M') return model.to(device), tokenizer class SkipBatchSampler(Sampler): def __init__(self, sampler, batch_size, skip_batches=0): self.sampler = sampler self.batch_size = batch_size self.skip_batches = skip_batches def __iter__(self): batch = [] skipped = 0 for idx in self.sampler: batch.append(idx) if len(batch) == self.batch_size: if skipped < self.skip_batches: skipped += 1 batch = [] continue yield batch batch = [] if len(batch) > 0 and skipped >= self.skip_batches: yield batch def __len__(self): total_batches = (len(self.sampler) + self.batch_size - 1) // self.batch_size return max(0, total_batches - self.skip_batches) class LMForRewardModel: def __init__(self, model_path, device="cuda", dtype=torch.float16): 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() self.device = device @torch.no_grad() def get_score(self, messages, response): history_text = "\n".join([f"{m['role']}: {m['content']}" for m in messages[:-1]]) last_query = messages[-1]['content'] if messages else "" message_context = f"{history_text}\n以上是对话历史。我的新问题是：\n{last_query}" if history_text else last_query eval_messages = [ {"role": "user", "content": message_context}, {"role": "assistant", "content": response} ] score = self.model.get_score(self.tokenizer, eval_messages) return max(min(score, 3.0), -3.0)