From eaf3a8cbec2d4dc8c3b9716d2bd8ad6529887077 Mon Sep 17 00:00:00 2001 From: baihaochen <1291308544@qq.com> Date: Sun, 7 Jun 2026 19:46:48 +0800 Subject: [PATCH] feat: add pretrain evaluation scripts --- eval_llm.py | 62 +++++++----- scripts/eval_pretrain_ppl.py | 177 +++++++++++++++++++++++++++++++++++ trainer/trainer_utils.py | 3 + 3 files changed, 220 insertions(+), 22 deletions(-) create mode 100644 scripts/eval_pretrain_ppl.py diff --git a/eval_llm.py b/eval_llm.py index f3a5c42..a8a7bb6 100755 --- a/eval_llm.py +++ b/eval_llm.py @@ -1,41 +1,51 @@ import time import argparse +import os import random import warnings import torch from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer -from model.model_minimind import MiniMindConfig, MiniMindForCausalLM -from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM +from model.model_minimind import MiniMindForCausalLM +from model.model_minimind_mla import MiniMindMLAForCausalLM from model.model_lora import * -from trainer.trainer_utils import setup_seed, get_model_params +from trainer.trainer_utils import ( + add_model_profile_args, + apply_model_profile, + build_lm_config, + get_model_params, + get_model_suffix, + resolve_attention_type, + setup_seed, +) warnings.filterwarnings('ignore') def init_model(args): tokenizer = AutoTokenizer.from_pretrained(args.load_from) if 'model' in args.load_from: - if args.use_mla: - config = MiniMindMLAConfig( - hidden_size=args.hidden_size, - num_hidden_layers=args.num_hidden_layers, - use_moe=bool(args.use_moe), - inference_rope_scaling=args.inference_rope_scaling, - kv_lora_rank=args.kv_lora_rank - ) + attention_type = resolve_attention_type(args) + config = build_lm_config( + hidden_size=args.hidden_size, + num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe), + attention_type=attention_type, + num_attention_heads=args.num_attention_heads, + num_key_value_heads=args.num_key_value_heads, + intermediate_size=args.intermediate_size, + kv_lora_rank=args.kv_lora_rank, + q_lora_rank=args.q_lora_rank, + rope_dim=args.rope_dim, + inference_rope_scaling=args.inference_rope_scaling, + ) + if attention_type == "mla": model = MiniMindMLAForCausalLM(config) else: - model = MiniMindForCausalLM(MiniMindConfig( - hidden_size=args.hidden_size, - num_hidden_layers=args.num_hidden_layers, - use_moe=bool(args.use_moe), - inference_rope_scaling=args.inference_rope_scaling - )) - model_suffix = '_moe' if args.use_moe else '' - if args.use_mla: model_suffix += '_mla' - ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{model_suffix}.pth' + model = MiniMindForCausalLM(config) + model_suffix = get_model_suffix(config) + ckp = os.path.join(args.save_dir, f'{args.weight}_{args.hidden_size}{model_suffix}.pth') model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True) if args.lora_weight != 'None': apply_lora(model) - load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}.pth') + load_lora(model, os.path.join(args.save_dir, f'{args.lora_weight}_{args.hidden_size}.pth')) else: model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) get_model_params(model, model.config) @@ -49,9 +59,15 @@ def main(): parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称（None表示不使用，可选：lora_identity, lora_medical）") parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度") parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") + parser.add_argument('--num_attention_heads', default=8, type=int, help="注意力头数量") + parser.add_argument('--num_key_value_heads', default=4, type=int, help="KV头数量") + parser.add_argument('--intermediate_size', default=None, type=int, help="FFN中间层维度") parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") + parser.add_argument('--attention_type', default='gqa', type=str, choices=['gqa', 'mha', 'mqa', 'mla'], help="注意力结构") parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）") parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）") + parser.add_argument('--q_lora_rank', default=256, type=int, help="MLA的Q压缩秩（仅use_mla=1时生效）") + parser.add_argument('--rope_dim', default=None, type=int, help="MLA解耦RoPE维度") parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）") parser.add_argument('--max_new_tokens', default=8192, type=int, help="最大生成长度（注意：并非模型实际长文本能力）") parser.add_argument('--temperature', default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）") @@ -60,7 +76,9 @@ def main(): parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）") parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）") parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") + add_model_profile_args(parser) args = parser.parse_args() + apply_model_profile(args) prompts = [ '你有什么特长？', @@ -108,4 +126,4 @@ def main(): print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n') if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/eval_pretrain_ppl.py b/scripts/eval_pretrain_ppl.py new file mode 100644 index 0000000..ae5aee5 --- /dev/null +++ b/scripts/eval_pretrain_ppl.py @@ -0,0 +1,177 @@ +import argparse +import json +import math +import os +import sys +import time + +import torch +from torch.utils.data import DataLoader, Dataset +from transformers import AutoTokenizer + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from trainer.trainer_utils import ( # noqa: E402 + add_model_profile_args, + apply_model_profile, + build_lm_config, + init_model, + resolve_attention_type, +) + + +class JsonlPretrainEvalDataset(Dataset): + def __init__(self, data_path, tokenizer, max_length=1024, max_samples=2000, start_index=0): + self.samples = [] + self.tokenizer = tokenizer + self.max_length = max_length + max_samples = int(max_samples) + start_index = int(start_index) + + with open(data_path, "r", encoding="utf-8") as f: + for line_idx, line in enumerate(f): + if line_idx < start_index: + continue + if max_samples > 0 and len(self.samples) >= max_samples: + break + line = line.strip() + if not line: + continue + item = json.loads(line) + text = str(item.get("text", "")) + if text: + self.samples.append(text) + + if not self.samples: + raise ValueError(f"No eval samples loaded from {data_path}") + + def __len__(self): + return len(self.samples) + + def __getitem__(self, index): + tokens = self.tokenizer( + self.samples[index], + add_special_tokens=False, + max_length=self.max_length - 2, + truncation=True, + ).input_ids + tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id] + input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens)) + input_ids = torch.tensor(input_ids, dtype=torch.long) + labels = input_ids.clone() + labels[input_ids == self.tokenizer.pad_token_id] = -100 + return input_ids, labels + + +def parse_args(): + parser = argparse.ArgumentParser(description="Evaluate pretrain checkpoint loss/PPL on JSONL text samples") + parser.add_argument("--data_path", required=True, help="JSONL file with a text field") + parser.add_argument("--save_dir", default="out", help="Directory containing native .pth weights") + parser.add_argument("--weight", required=True, help="Weight prefix, e.g. pretrain_searchlm_300m_v1") + parser.add_argument("--tokenizer_path", default="model") + parser.add_argument("--output", default="", help="Optional JSON output path") + parser.add_argument("--max_samples", type=int, default=2000, help="Number of eval samples; <=0 means all after start_index") + parser.add_argument("--start_index", type=int, default=0, help="Skip this many JSONL rows before evaluation") + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--max_seq_len", type=int, default=1024) + parser.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"]) + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--log_interval", type=int, default=20) + parser.add_argument("--hidden_size", default=768, type=int) + parser.add_argument("--num_hidden_layers", default=8, type=int) + parser.add_argument("--num_attention_heads", default=8, type=int) + parser.add_argument("--num_key_value_heads", default=4, type=int) + parser.add_argument("--intermediate_size", default=None, type=int) + parser.add_argument("--use_moe", default=0, type=int, choices=[0, 1]) + parser.add_argument("--attention_type", default="gqa", choices=["gqa", "mha", "mqa", "mla"]) + parser.add_argument("--use_mla", default=0, type=int, choices=[0, 1]) + parser.add_argument("--kv_lora_rank", default=128, type=int) + parser.add_argument("--q_lora_rank", default=256, type=int) + parser.add_argument("--rope_dim", default=None, type=int) + add_model_profile_args(parser) + return parser.parse_args() + + +def main(): + args = parse_args() + apply_model_profile(args) + attention_type = resolve_attention_type(args) + lm_config = build_lm_config( + hidden_size=args.hidden_size, + num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe), + attention_type=attention_type, + num_attention_heads=args.num_attention_heads, + num_key_value_heads=args.num_key_value_heads, + intermediate_size=args.intermediate_size, + kv_lora_rank=args.kv_lora_rank, + q_lora_rank=args.q_lora_rank, + rope_dim=args.rope_dim, + ) + + model, tokenizer = init_model( + lm_config, + from_weight=args.weight, + tokenizer_path=args.tokenizer_path, + save_dir=args.save_dir, + device=args.device, + ) + if args.device.startswith("cuda") and args.dtype != "float32": + dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + model = model.to(dtype=dtype) + model.eval() + + dataset = JsonlPretrainEvalDataset( + args.data_path, + tokenizer, + max_length=args.max_seq_len, + max_samples=args.max_samples, + start_index=args.start_index, + ) + loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False) + + total_nll = 0.0 + total_tokens = 0 + start_time = time.time() + with torch.inference_mode(): + for step, (input_ids, labels) in enumerate(loader, start=1): + input_ids = input_ids.to(args.device, non_blocking=True) + labels = labels.to(args.device, non_blocking=True) + outputs = model(input_ids, labels=labels) + valid_tokens = (labels[..., 1:] != -100).sum().item() + total_nll += float(outputs.loss.item()) * valid_tokens + total_tokens += valid_tokens + + if step % args.log_interval == 0 or step == len(loader): + loss = total_nll / max(total_tokens, 1) + ppl = math.exp(loss) if loss < 100 else float("inf") + print( + f"Eval step [{step}/{len(loader)}], " + f"loss: {loss:.4f}, ppl: {ppl:.2f}, tokens: {total_tokens}" + ) + + loss = total_nll / max(total_tokens, 1) + ppl = math.exp(loss) if loss < 100 else float("inf") + result = { + "weight": args.weight, + "model_profile": args.model_profile, + "data_path": args.data_path, + "start_index": args.start_index, + "num_samples": len(dataset), + "max_seq_len": args.max_seq_len, + "batch_size": args.batch_size, + "loss": loss, + "ppl": ppl, + "tokens": total_tokens, + "elapsed_sec": time.time() - start_time, + } + print(json.dumps(result, ensure_ascii=False, indent=2)) + + if args.output: + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + with open(args.output, "w", encoding="utf-8") as f: + json.dump(result, f, ensure_ascii=False, indent=2) + + +if __name__ == "__main__": + main() diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index d07393c..60d2e70 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -214,6 +214,7 @@ def _candidate_deepspeed_tags(lm_config, weight, save_dir="../checkpoints"): 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 = [] @@ -308,6 +309,7 @@ def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../chec loaded_tag = None for tag in _candidate_deepspeed_tags(lm_config, weight, save_dir): try: + #deepspeed会从save_dir中找tag对应的checkpoint load_path, client_state = model_engine.load_checkpoint(save_dir, tag=tag) except Exception as exc: errors.append(f"{tag}: {exc}") @@ -317,6 +319,7 @@ def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../chec continue loaded_tag = tag break + #这个的else的for else 意思是如果for循环没有被break打断就走else分支 else: if is_main_process(): Logger("DeepSpeed checkpoint not loaded" + (f": {'; '.join(errors)}" if errors else ""))