github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: add pretrain evaluation scripts

pull/791/head
baihaochen 2026-06-07 19:46:48 +08:00
parent 4dcda45bfa
commit eaf3a8cbec
3 changed files with 220 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
eval_llm.py
View File

@ -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()
main()

177
scripts/eval_pretrain_ppl.py Normal file
View File

@ -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()

3
trainer/trainer_utils.py
View File

@ -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 ""))