178 lines
6.9 KiB
Python
178 lines
6.9 KiB
Python
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()
|