""" 训练工具函数集合 """ import os import sys __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import re import json import random import math import shutil 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 MiniMindConfig, MiniMindForCausalLM from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM DEFAULT_DEEPSPEED_CONFIG = os.path.join(os.path.dirname(__file__), "ds_config_zero2.json") def _resolve_repo_path(path): if os.path.isabs(path) and os.path.exists(path): return path if os.path.exists(path): return path repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) candidate = os.path.join(repo_root, path) return candidate def add_model_profile_args(parser): parser.add_argument("--model_profile", type=str, default="", help="SearchLM模型规模配置名，如 SearchLM-100M/SearchLM-300M") parser.add_argument("--model_profiles_path", type=str, default="configs/searchlm_profiles.json", help="模型规模配置文件") def add_deepspeed_args(parser, default_config=DEFAULT_DEEPSPEED_CONFIG): parser.add_argument("--deepspeed_config", type=str, default=default_config, help="DeepSpeed ZeRO配置JSON") parser.add_argument("--local_rank", type=int, default=-1, help="DeepSpeed launcher自动注入") def apply_model_profile(args): """将 SearchLM profile 写入 argparse args；显式传空则不修改。""" profile_name = getattr(args, "model_profile", "") if not profile_name: return None profiles_path = _resolve_repo_path(getattr(args, "model_profiles_path", "configs/searchlm_profiles.json")) with open(profiles_path, "r", encoding="utf-8") as f: profiles = json.load(f) if profile_name not in profiles: raise ValueError(f"Unknown model_profile={profile_name}. Available: {', '.join(profiles)}") profile = profiles[profile_name] if "hidden_size" not in profile: raise ValueError(f"Profile {profile_name} is not a trainable from-scratch profile") fields = [ "hidden_size", "num_hidden_layers", "num_attention_heads", "num_key_value_heads", "attention_type", "intermediate_size", "kv_lora_rank", "q_lora_rank", "rope_dim", "max_seq_len" ] for field in fields: if field in profile and hasattr(args, field): setattr(args, field, profile[field]) return profile def get_model_suffix(lm_config): """根据模型配置返回 checkpoint 文件名后缀""" suffix = '' if getattr(lm_config, 'use_moe', False): suffix += '_moe' attention_type = getattr(lm_config, "attention_type", "gqa") if isinstance(lm_config, MiniMindMLAConfig) or attention_type == "mla": suffix += '_mla' elif attention_type in {"mha", "mqa"}: suffix += f'_{attention_type}' return suffix def unwrap_model(model): """兼容 DDP / DeepSpeed / torch.compile 的模型解包。""" raw_model = getattr(model, "module", model) return getattr(raw_model, "_orig_mod", raw_model) def get_model_weight_path(lm_config, save_dir="../out", weight="model", epoch=None, step=None): model_suffix = get_model_suffix(lm_config) step_suffix = f"_epoch{epoch}_step{step}" if epoch is not None and step is not None else "" return f"{save_dir}/{weight}_{lm_config.hidden_size}{model_suffix}{step_suffix}.pth" def save_model_weights(lm_config, model, save_dir="../out", weight="model", dtype=torch.float16): """保存一份轻量推理权重；DeepSpeed完整训练状态由 save_checkpoint 单独保存。""" os.makedirs(save_dir, exist_ok=True) ckp_path = get_model_weight_path(lm_config, save_dir=save_dir, weight=weight) raw_model = unwrap_model(model) state_dict = raw_model.state_dict() state_dict = {k: v.detach().to(dtype).cpu() for k, v in state_dict.items()} tmp_path = ckp_path + ".tmp" torch.save(state_dict, tmp_path) os.replace(tmp_path, ckp_path) del state_dict if torch.cuda.is_available(): torch.cuda.empty_cache() return ckp_path def load_deepspeed_config(args, total_steps): """加载并补全 DeepSpeed 配置，保持 batch / lr / precision 与 CLI 参数一致。""" ds_path = _resolve_repo_path(getattr(args, "deepspeed_config", DEFAULT_DEEPSPEED_CONFIG)) with open(ds_path, "r", encoding="utf-8") as f: cfg = json.load(f) world_size = dist.get_world_size() if dist.is_initialized() else int(os.environ.get("WORLD_SIZE", "1")) if cfg.get("train_micro_batch_size_per_gpu") == "auto": cfg["train_micro_batch_size_per_gpu"] = args.batch_size if cfg.get("gradient_accumulation_steps") == "auto": cfg["gradient_accumulation_steps"] = args.accumulation_steps if cfg.get("train_batch_size") == "auto": cfg["train_batch_size"] = args.batch_size * args.accumulation_steps * world_size if "gradient_clipping" in cfg: cfg["gradient_clipping"] = args.grad_clip optimizer_cfg = cfg.setdefault("optimizer", {"type": "AdamW", "params": {}}) optimizer_cfg.setdefault("params", {})["lr"] = args.learning_rate scheduler_cfg = cfg.get("scheduler") if scheduler_cfg and scheduler_cfg.get("type") == "WarmupDecayLR": sched_params = scheduler_cfg.setdefault("params", {}) if sched_params.get("warmup_max_lr") == "auto": sched_params["warmup_max_lr"] = args.learning_rate if sched_params.get("warmup_num_steps") == "auto": sched_params["warmup_num_steps"] = max(1, total_steps // 20) if sched_params.get("total_num_steps") == "auto": sched_params["total_num_steps"] = max(1, total_steps) dtype = getattr(args, "dtype", "float16") cfg.setdefault("fp16", {})["enabled"] = dtype == "float16" cfg.setdefault("bf16", {})["enabled"] = dtype == "bfloat16" return cfg def get_wandb_id(wandb=None): if not wandb: return None if hasattr(wandb, "get_run"): run = wandb.get_run() return getattr(run, "id", None) if run else None return getattr(wandb, "id", None) def get_cuda_peak_memory_gb(): if not torch.cuda.is_available(): return None return torch.cuda.max_memory_allocated(torch.cuda.current_device()) / (1024 ** 3) def reduce_metrics(metrics, average=True): """聚合所有 rank 的标量指标，避免日志只反映 rank0 局部 batch。""" if not dist.is_initialized(): return metrics numeric_items = [(key, value) for key, value in metrics.items() if isinstance(value, (int, float))] if not numeric_items: return metrics device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") values = torch.tensor([float(value) for _, value in numeric_items], device=device) dist.all_reduce(values, op=dist.ReduceOp.SUM) if average: values /= dist.get_world_size() reduced = dict(metrics) for (key, _), value in zip(numeric_items, values.tolist()): reduced[key] = value return reduced def get_deepspeed_tag(lm_config, weight): return f"{weight}_{lm_config.hidden_size}{get_model_suffix(lm_config)}" def get_deepspeed_step_tag(lm_config, weight, epoch, step): return f"{get_deepspeed_tag(lm_config, weight)}_epoch{epoch}_step{step}" def get_deepspeed_latest_file(lm_config, weight, save_dir="../checkpoints"): return os.path.join(save_dir, f"{get_deepspeed_tag(lm_config, weight)}_latest") def _read_deepspeed_latest_tag(lm_config, weight, save_dir="../checkpoints"): latest_file = get_deepspeed_latest_file(lm_config, weight, save_dir) if not os.path.exists(latest_file): return None with open(latest_file, "r", encoding="utf-8") as f: tag = f.read().strip() return tag or None def _write_deepspeed_latest_tag(lm_config, weight, tag, save_dir="../checkpoints"): latest_file = get_deepspeed_latest_file(lm_config, weight, save_dir) tmp_file = latest_file + ".tmp" with open(tmp_file, "w", encoding="utf-8") as f: f.write(tag) os.replace(tmp_file, latest_file) def _candidate_deepspeed_tags(lm_config, weight, save_dir="../checkpoints"): base_tag = get_deepspeed_tag(lm_config, weight) candidates = [] latest_tag = _read_deepspeed_latest_tag(lm_config, weight, save_dir) if latest_tag: candidates.append(latest_tag) candidates.append(base_tag) if os.path.exists(save_dir): step_tags = [] prefix = f"{base_tag}_epoch" for name in os.listdir(save_dir): path = os.path.join(save_dir, name) if name.startswith(prefix) and os.path.isdir(path): step_tags.append((os.path.getmtime(path), name)) candidates.extend(name for _, name in sorted(step_tags, reverse=True)) dedup = [] seen = set() for tag in candidates: if tag not in seen: seen.add(tag) dedup.append(tag) return dedup def _parse_deepspeed_tag_metadata(tag): match = re.search(r"_epoch(\d+)_step(\d+)$", tag) if not match: return {} return { "epoch": int(match.group(1)), "step": int(match.group(2)), } def _prune_old_deepspeed_step_tags(lm_config, weight, keep_tag, save_dir="../checkpoints"): if not os.path.exists(save_dir): return base_tag = get_deepspeed_tag(lm_config, weight) prefix = f"{base_tag}_epoch" for name in os.listdir(save_dir): if name == keep_tag or not name.startswith(prefix): continue path = os.path.join(save_dir, name) if os.path.isdir(path): shutil.rmtree(path, ignore_errors=True) def assert_no_training_output_collision(lm_config, weight, save_dir="../out", checkpoint_dir="../checkpoints", from_resume=0, overwrite=0): if int(from_resume) == 1 or int(overwrite) == 1: return base_tag = get_deepspeed_tag(lm_config, weight) latest_weight = get_model_weight_path(lm_config, save_dir=save_dir, weight=weight) latest_file = get_deepspeed_latest_file(lm_config, weight, checkpoint_dir) existing = [] if os.path.exists(latest_weight): existing.append(latest_weight) if os.path.exists(latest_file): existing.append(latest_file) if os.path.exists(os.path.join(checkpoint_dir, base_tag)): existing.append(os.path.join(checkpoint_dir, base_tag)) if os.path.exists(checkpoint_dir): prefix = f"{base_tag}_epoch" for name in os.listdir(checkpoint_dir): if name.startswith(prefix): existing.append(os.path.join(checkpoint_dir, name)) break if not existing: return raise RuntimeError( "检测到同名训练产物，已阻止覆盖。" f" save_weight={weight}, hidden_size={lm_config.hidden_size}, existing={existing}. " "如果是继续训练，请加 --from_resume 1；如果是新实验，请换一个 --save_weight；" "如果确认要覆盖/复用旧目录，请显式加 --overwrite_output 1。" ) def save_deepspeed_checkpoint(model_engine, lm_config, weight, epoch=0, step=0, wandb=None, save_dir="../checkpoints", **kwargs): """DeepSpeed checkpoint 必须所有 rank 都调用，不能只在 rank0 保存。""" os.makedirs(save_dir, exist_ok=True) wandb_id = get_wandb_id(wandb) if is_main_process() else None if dist.is_initialized(): wandb_box = [wandb_id] dist.broadcast_object_list(wandb_box, src=0) wandb_id = wandb_box[0] client_state = { "epoch": epoch, "step": step, "world_size": dist.get_world_size() if dist.is_initialized() else 1, "wandb_id": wandb_id, } client_state.update({k: v for k, v in kwargs.items() if v is not None}) tag = get_deepspeed_step_tag(lm_config, weight, epoch, step) model_engine.save_checkpoint(save_dir, tag=tag, client_state=client_state) if is_main_process(): _write_deepspeed_latest_tag(lm_config, weight, tag, save_dir) _prune_old_deepspeed_step_tags(lm_config, weight, keep_tag=tag, save_dir=save_dir) Logger(f"DeepSpeed checkpoint saved: {save_dir}/{tag}") if dist.is_initialized(): dist.barrier() return tag def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../checkpoints"): errors = [] loaded_tag = None for tag in _candidate_deepspeed_tags(lm_config, weight, save_dir): try: load_path, client_state = model_engine.load_checkpoint(save_dir, tag=tag) except Exception as exc: errors.append(f"{tag}: {exc}") continue if load_path is None: errors.append(f"{tag}: load_path=None") continue loaded_tag = tag break else: if is_main_process(): Logger("DeepSpeed checkpoint not loaded" + (f": {'; '.join(errors)}" if errors else "")) return None client_state = client_state or {} tag_state = _parse_deepspeed_tag_metadata(loaded_tag or "") for key, value in tag_state.items(): client_state.setdefault(key, value) client_state.setdefault("world_size", 1) saved_ws = client_state.get("world_size", 1) current_ws = dist.get_world_size() if dist.is_initialized() else 1 if saved_ws != current_ws and is_main_process(): Logger(f"DeepSpeed checkpoint world_size changed: {saved_ws} -> {current_ws}") if is_main_process(): Logger( f"DeepSpeed checkpoint loaded: {load_path}, " f"resume epoch={client_state.get('epoch')}, step={client_state.get('step')}" ) return client_state def build_lm_config(hidden_size=768, num_hidden_layers=8, use_moe=False, attention_type="gqa", kv_lora_rank=128, **kwargs): """统一构建不同注意力结构的 MiniMind 配置。""" if attention_type == "mla": return MiniMindMLAConfig( hidden_size=hidden_size, num_hidden_layers=num_hidden_layers, use_moe=use_moe, kv_lora_rank=kv_lora_rank, **kwargs ) if attention_type not in {"gqa", "mha", "mqa"}: raise ValueError(f"Unsupported attention_type: {attention_type}") return MiniMindConfig( hidden_size=hidden_size, num_hidden_layers=num_hidden_layers, use_moe=use_moe, attention_type=attention_type, **kwargs ) def resolve_attention_type(args): """兼容旧参数 --use_mla，同时支持新的 --attention_type。""" if getattr(args, "use_mla", 0): return "mla" return getattr(args, "attention_type", "gqa") def get_model_params(model, config): total = sum(p.numel() for p in model.parameters()) / 1e6 #getattr用于动态获取对象的属性或方法 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 log_training_setup(args, lm_config, stage="train", dataset_len=None, iters=None, tokens_per_sample=None, extra=None): """打印可直接沉淀到实验报告里的 DDP/训练配置信息。""" world_size = dist.get_world_size() if dist.is_initialized() else 1 rank = dist.get_rank() if dist.is_initialized() else 0 per_gpu_batch = getattr(args, "batch_size", 1) accumulation = getattr(args, "accumulation_steps", 1) effective_batch = per_gpu_batch * world_size * accumulation max_seq_len = tokens_per_sample or getattr(args, "max_seq_len", None) tokens_per_step = effective_batch * max_seq_len if max_seq_len else None attention_type = getattr(lm_config, "attention_type", "gqa") device = getattr(args, "device", "unknown") gpu_name = None if torch.cuda.is_available() and "cuda" in str(device): try: gpu_name = torch.cuda.get_device_name(torch.cuda.current_device()) except Exception: gpu_name = None if rank == 0: Logger("=" * 80) Logger(f"[Training Setup] stage={stage}") backend = getattr(args, "training_backend", None) if backend is None: backend = "DDP" if world_size > 1 else "single-process" Logger( f" distributed={backend}, world_size={world_size}, " f"device={device}" + (f", gpu={gpu_name}" if gpu_name else "") ) Logger( f" model: attention_type={attention_type}, hidden_size={lm_config.hidden_size}, " f"num_layers={lm_config.num_hidden_layers}, use_moe={getattr(lm_config, 'use_moe', False)}" ) if attention_type == "mla" or isinstance(lm_config, MiniMindMLAConfig): Logger( f" mla: kv_lora_rank={lm_config.kv_lora_rank}, q_lora_rank={lm_config.q_lora_rank}, " f"rope_dim={lm_config.rope_dim}" ) Logger( f" batch: per_gpu={per_gpu_batch}, accumulation_steps={accumulation}, " f"effective_batch={effective_batch}" ) if tokens_per_step: Logger(f" sequence: max_seq_len={max_seq_len}, tokens_per_optimizer_step~{tokens_per_step:,}") if dataset_len is not None: Logger(f" data: samples={dataset_len:,}" + (f", batches_per_epoch={iters:,}" if iters is not None else "")) Logger( f" precision={getattr(args, 'dtype', 'unknown')}, lr={getattr(args, 'learning_rate', 'unknown')}, " f"epochs={getattr(args, 'epochs', 'unknown')}, compile={getattr(args, 'use_compile', 0)}" ) if extra: for key, value in extra.items(): Logger(f" {key}: {value}") Logger("=" * 80) 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))) 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): random.seed(seed)#设置python的随机种子 np.random.seed(seed)#设置numpy的随机种子 torch.manual_seed(seed)#设置cpu随机种子 torch.cuda.manual_seed(seed)#设置pytorch在当前GPU上的随机种子 torch.cuda.manual_seed_all(seed)#设置pytorch在所有GPU上的随机种子 torch.backends.cudnn.deterministic = True#使用确定性算法 torch.backends.cudnn.benchmark = False#禁用cudnn的自动优化 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 = get_model_suffix(lm_config) 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:#存储模型模式 #如果模型被DDP包了，真正的模型在model.module里 raw_model = unwrap_model(model) state_dict = raw_model.state_dict() #half表示转成float16，然后移到cpu上 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用于将一个文件替换成另一个文件 os.replace(src, dst) 将 src 文件替换为 dst 文件，如果 dst 已经存在，则会被替换掉。这是一个原子操作，确保在替换过程中不会出现文件损坏或丢失的情况。 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, #训练时用了几张GPU '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 = unwrap_model(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) #删除临时变量，然后清理cuda缓存 del state_dict, resume_data torch.cuda.empty_cache() else: # 加载模式 if os.path.exists(resume_path): #加载到gpu 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) if isinstance(lm_config, MiniMindMLAConfig): model = MiniMindMLAForCausalLM(lm_config) else: model = MiniMindForCausalLM(lm_config) if from_weight!= 'none': model_suffix = get_model_suffix(lm_config) weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{model_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 #对sampler的重写，主要重写iter方法和len方法，sampler决定取哪些数据，给dataloaer拿到索引后，通过getitem方法取数据 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 #self.sampler每次吐出一个样本的索引 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) #clip到[-3,3] return max(min(score, 3.0), -3.0)