File size: 8,593 Bytes
45cf443 c6bc767 45cf443 92c7321 c6bc767 45cf443 17c552f 45cf443 17c552f 45cf443 d58698c 64d62ee c6bc767 45cf443 d58698c 45cf443 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | import os
import sys
import random
import math
import numpy as np
import yaml
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 models import LMForCausalLM
def _load_yaml_config(config_path):
with open(config_path, 'r', encoding='utf-8') as f:
raw = yaml.safe_load(f) or {}
defaults = {}
for section in ('model', 'train', 'paths'):
for key, value in (raw.get(section) or {}).items():
defaults[key] = value
return defaults
def apply_config(parser, default_config=None):
"""让 YAML 配置成为 argparse 默认值，CLI 显式传参仍可覆盖。
用法：parser 需已定义 --config 参数。本函数会先 parse_known_args 读取 --config；
若未传 --config 且给定了 default_config，则回退到该默认路径。命中 YAML 后扁平化
（model/train/paths 三个 section）作为 argparse 默认值注入，再做最终 parse。
返回最终的 args。
"""
pre, _ = parser.parse_known_args()
config_path = getattr(pre, 'config', None) or default_config
if config_path and os.path.exists(config_path):
defaults = _load_yaml_config(config_path)
parser.set_defaults(**defaults)
return parser.parse_args()
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
_log_file = None
def init_logger(save_dir='../checkpoint', name='train'):
global _log_file
if not is_main_process():
return
os.makedirs(save_dir, exist_ok=True)
log_path = os.path.join(save_dir, f'{name}.log')
_log_file = open(log_path, 'a', encoding='utf-8')
Logger(f'日志写入: {os.path.abspath(log_path)}')
def Logger(content):
if not is_main_process():
return
print(content)
if _log_file is not None:
_log_file.write(str(content) + '\n')
_log_file.flush()
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
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)
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
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', save_dir='../checkpoint', tokenizer_dir=None, device='cuda', model_dir=None):
tokenizer_dir = tokenizer_dir or os.path.join(save_dir, 'tokenizer')
tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
model = LMForCausalLM(lm_config)
if from_weight != 'none':
moe_suffix = '_moe' if lm_config.use_moe else ''
weight_dir = model_dir or save_dir
weight_path = f'{weight_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)
|