From 17f67f989787bdb3903e2686f5d28e7a0bcc3f29 Mon Sep 17 00:00:00 2001 From: Longfang Zhao Date: Tue, 30 Jun 2026 09:30:04 -0700 Subject: [PATCH] Add MFU cal utils. --- scripts/eval_completion.py | 147 +++++++++++++++++++++++++++++++++++++ scripts/eval_perplexity.py | 147 +++++++++++++++++++++++++++++++++++++ trainer/train_pretrain.py | 44 ++++++++--- trainer/trainer_utils.py | 47 ++++++++++++ 4 files changed, 375 insertions(+), 10 deletions(-) create mode 100644 scripts/eval_completion.py create mode 100644 scripts/eval_perplexity.py diff --git a/scripts/eval_completion.py b/scripts/eval_completion.py new file mode 100644 index 0000000..35cdb95 --- /dev/null +++ b/scripts/eval_completion.py @@ -0,0 +1,147 @@ +""" +Qualitative eval of a MiniMind pretrain (or post-pretrain) checkpoint by sampling +text continuations from a fixed prompt set spanning multiple domains. + +For pretrain (base) models — does NOT use chat template. Use eval_llm.py for SFT models. + +Usage: + # Default: 64M pretrain ckpt + python eval_completion.py + + # 1B ckpt + python eval_completion.py --hidden_size 2048 --num_hidden_layers 20 --weight pretrain + + # Custom prompt + python eval_completion.py --prompt "今天天气不错，我打算" + + # Save outputs to markdown for design doc + python eval_completion.py --output_md eval_results.md + +What to look for: + ✓ Grammatically coherent Chinese, no repetition loops + ✓ Stays roughly on topic of the prefix + ✓ Variety across runs (different temperature should give different outputs) + ✗ Repetitive: "的的的的" / 同一句重复 → 训练或解码有问题 + ✗ 全部生成都一模一样 → 模型坍缩 / temperature 太低 + ✗ 乱码 / 非中文片段（如果是纯中文 prefix）→ tokenizer 问题 +""" +import os +import sys + +__package__ = "scripts" +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import argparse +import warnings +import time +import torch +from transformers import AutoTokenizer, TextStreamer +from model.model_minimind import MiniMindConfig, MiniMindForCausalLM + +warnings.filterwarnings('ignore') + + +DEFAULT_PROMPTS = [ + ("叙事", "在很久很久以前，有一个小村庄，"), + ("日常", "今天天气不错，我打算"), + ("技术", "Python 是一种编程语言，它的特点是"), + ("事实", "中华人民共和国成立于"), + ("代码", "def fibonacci(n):\n if n < 2:"), + ("科学", "牛顿第一定律说的是"), + ("文学", "李白是唐代著名的"), + ("数学", "二的十次方等于"), + ("常识", "水的沸点在标准大气压下是"), + ("地理", "中国的首都是"), +] + + +def load_model(args): + cfg = MiniMindConfig( + hidden_size=args.hidden_size, + num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe), + ) + model = MiniMindForCausalLM(cfg) + moe_suffix = '_moe' if args.use_moe else '' + ckp = f'{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' + state_dict = torch.load(ckp, map_location='cpu', weights_only=True) + model.load_state_dict(state_dict, strict=False) + model = model.half().eval().to(args.device) + return model, cfg + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser(description="MiniMind Pretrain Completion Eval") + parser.add_argument('--hidden_size', type=int, default=768) + parser.add_argument('--num_hidden_layers', type=int, default=8) + parser.add_argument('--use_moe', type=int, default=0, choices=[0, 1]) + parser.add_argument('--save_dir', type=str, default='../out') + parser.add_argument('--weight', type=str, default='pretrain') + parser.add_argument('--tokenizer_path', type=str, default='../model') + parser.add_argument('--prompt', type=str, default=None, + help='单条自定义 prompt；不指定则跑 DEFAULT_PROMPTS 全集') + parser.add_argument('--max_new_tokens', type=int, default=80) + parser.add_argument('--temperature', type=float, default=0.7) + parser.add_argument('--top_p', type=float, default=0.9) + parser.add_argument('--top_k', type=int, default=50) + parser.add_argument('--repetition_penalty', type=float, default=1.1) + parser.add_argument('--num_samples_per_prompt', type=int, default=1, + help='每个 prompt 采样几次（验证多样性）') + parser.add_argument('--output_md', type=str, default=None, + help='输出 markdown 报告路径（可选）') + parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu') + parser.add_argument('--seed', type=int, default=42) + args = parser.parse_args() + + torch.manual_seed(args.seed) + + print(f'[load] ckpt={args.save_dir}/{args.weight}_{args.hidden_size}.pth') + model, cfg = load_model(args) + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) + + prompts = [('custom', args.prompt)] if args.prompt else DEFAULT_PROMPTS + + md_lines = [] + md_lines.append(f'# Completion Eval: `{args.weight}_{args.hidden_size}.pth`') + md_lines.append(f'temperature={args.temperature}, top_p={args.top_p}, top_k={args.top_k}, ' + f'rep_penalty={args.repetition_penalty}, max_new={args.max_new_tokens}\n') + + for label, prompt in prompts: + print(f'\n{"=" * 60}') + print(f'[{label}] Prompt: {prompt!r}') + md_lines.append(f'## [{label}] `{prompt}`\n') + for sample_i in range(args.num_samples_per_prompt): + print(f'\n--- Sample {sample_i + 1} ---') + # 必须加 BOS：训练时所有样本以 [BOS] 开头，eval 时不加 BOS 模型会产生 degenerate + # 输出（"的的的"/"是是是" 重复），因为它从没见过"无 BOS 开头"的序列 + body = tokenizer(prompt, add_special_tokens=False).input_ids + ids = torch.tensor([[tokenizer.bos_token_id] + body], device=args.device) + t0 = time.time() + output_ids = model.generate( + ids, + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + top_p=args.top_p, + top_k=args.top_k, + repetition_penalty=args.repetition_penalty, + do_sample=True, + eos_token_id=tokenizer.eos_token_id, + ) + elapsed = time.time() - t0 + text = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True) + continuation = text[len(prompt):] + print(continuation) + new_tokens = output_ids.shape[1] - ids.shape[1] + print(f' [{new_tokens} new tokens in {elapsed:.2f}s = {new_tokens/elapsed:.1f} tok/s]') + md_lines.append(f'**Sample {sample_i + 1}**: {continuation.strip()}\n') + md_lines.append('') + + if args.output_md: + with open(args.output_md, 'w') as f: + f.write('\n'.join(md_lines)) + print(f'\n[saved] {args.output_md}') + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_perplexity.py b/scripts/eval_perplexity.py new file mode 100644 index 0000000..99061f8 --- /dev/null +++ b/scripts/eval_perplexity.py @@ -0,0 +1,147 @@ +""" +Compute perplexity of a MiniMind checkpoint on a holdout text dataset. + +Usage: + # Default: 64M pretrain ckpt, hash-sampled 1% of pretrain_t2t.jsonl + python eval_perplexity.py + + # 1B ckpt with explicit holdout + python eval_perplexity.py --hidden_size 2048 --num_hidden_layers 20 \\ + --weight pretrain --holdout_file ../dataset/holdout.jsonl + + # Compare two ckpts (e.g. before/after QAT) + python eval_perplexity.py --weight pretrain + python eval_perplexity.py --weight qat + +Reference PPL ranges (rough): + 64M from-scratch on Chinese pretrain holdout: 8-15 + 200M: 6-10 + 1B: 5-8 + Anything > 30 → training is broken (data leak, lr too high, etc.) + Anything < 5 → suspect data leak (test set seen during training) +""" +import os +import sys + +__package__ = "scripts" +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import argparse +import json +import math +import warnings +import time +import torch +from transformers import AutoTokenizer +from model.model_minimind import MiniMindConfig, MiniMindForCausalLM + +warnings.filterwarnings('ignore') + + +def load_model(args): + cfg = MiniMindConfig( + hidden_size=args.hidden_size, + num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe), + ) + model = MiniMindForCausalLM(cfg) + moe_suffix = '_moe' if args.use_moe else '' + ckp = f'{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' + state_dict = torch.load(ckp, map_location='cpu', weights_only=True) + # strict=False — tied embeddings have duplicate keys; load handles via shared storage + model.load_state_dict(state_dict, strict=False) + model = model.half().eval().to(args.device) + return model, cfg + + +def iter_holdout(args): + """Yield text strings from holdout. Either a separate file or hash-sampled from pretrain.""" + if args.holdout_file: + with open(args.holdout_file) as f: + for line in f: + d = json.loads(line) + yield d.get('text') or d.get('content') or '' + return + # Hash-based holdout: take 1/sampling_mod of training file (deterministic, repeatable) + # NOTE: only meaningful as proxy — model HAS seen these during training. + # True holdout requires reserving samples BEFORE training (Week 1 trainer fix). + with open(args.train_file) as f: + for line in f: + if hash(line) % args.sampling_mod == 0: + d = json.loads(line) + yield d.get('text') or d.get('content') or '' + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser(description="MiniMind Perplexity Eval") + parser.add_argument('--hidden_size', type=int, default=768) + parser.add_argument('--num_hidden_layers', type=int, default=8) + parser.add_argument('--use_moe', type=int, default=0, choices=[0, 1]) + parser.add_argument('--save_dir', type=str, default='../out') + parser.add_argument('--weight', type=str, default='pretrain', + help='ckpt 前缀，如 pretrain / full_sft / qat') + parser.add_argument('--tokenizer_path', type=str, default='../model') + parser.add_argument('--holdout_file', type=str, default=None, + help='独立 holdout jsonl 文件；若为 None 则从 train_file 按 hash 抽样') + parser.add_argument('--train_file', type=str, default='../dataset/pretrain_t2t.jsonl', + help='没 holdout 时从这里 hash 抽样作 proxy holdout') + parser.add_argument('--sampling_mod', type=int, default=1000, + help='hash 抽样模数：1000=取 0.1%, 100=取 1%') + parser.add_argument('--max_samples', type=int, default=500) + parser.add_argument('--max_seq_len', type=int, default=512) + parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu') + args = parser.parse_args() + + print(f'[load] ckpt={args.save_dir}/{args.weight}_{args.hidden_size}.pth') + model, cfg = load_model(args) + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) + pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 + + total_loss, total_tokens, n_samples = 0.0, 0, 0 + t0 = time.time() + # ⚠️ 必须加 BOS/EOS：minimind PretrainDataset 训练时格式是 [BOS] + tokens + [EOS] + pad。 + # eval 时不加 BOS，第一个 token 的 loss 会异常高（模型从来没见过"无 BOS 开头"）， + # 整体 PPL 会虚高 ~4x（实测从 4.7 → 20.2）。 + bos_id = tokenizer.bos_token_id + eos_id = tokenizer.eos_token_id + for text in iter_holdout(args): + if not text or len(text) < 10: + continue + body = tokenizer(text, max_length=args.max_seq_len - 2, + truncation=True, add_special_tokens=False).input_ids + tokens = [bos_id] + body + [eos_id] + if len(tokens) < 2: + continue + ids = torch.tensor([tokens], device=args.device) + labels = ids.clone() + labels[labels == pad_id] = -100 + out = model(input_ids=ids, labels=labels) + n_tokens = (labels != -100).sum().item() - 1 # 减 1：last token 不算 loss + if n_tokens <= 0: + continue + total_loss += out.loss.item() * n_tokens + total_tokens += n_tokens + n_samples += 1 + if n_samples >= args.max_samples: + break + + avg_loss = total_loss / max(total_tokens, 1) + ppl = math.exp(avg_loss) + elapsed = time.time() - t0 + print(f'\n{"=" * 60}') + print(f'Samples evaluated: {n_samples}') + print(f'Total tokens: {total_tokens:,}') + print(f'Avg loss: {avg_loss:.4f}') + print(f'Perplexity (PPL): {ppl:.2f}') + print(f'Elapsed: {elapsed:.1f}s ({total_tokens/elapsed:.0f} tokens/s)') + print(f'{"=" * 60}') + + if args.holdout_file is None: + print(f'⚠️ Used hash-sampled "holdout" from train_file (model HAS seen these).') + print(f' For real eval, reserve a holdout BEFORE training.') + + +if __name__ == "__main__": + main() diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index b3f71af..c2f64d7 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -16,7 +16,11 @@ from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import PretrainDataset -from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler +from trainer.trainer_utils import ( + get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, + setup_seed, init_model, SkipBatchSampler, + detect_gpu_peak_tflops_bf16, compute_model_flops_per_token, +) warnings.filterwarnings('ignore') @@ -26,8 +30,9 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): last_log_time, last_log_step = start_time, start_step last_step = start_step for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): - input_ids = input_ids.to(args.device) - labels = labels.to(args.device) + # non_blocking=True overlaps host->device copy with compute (pin_memory=True 已在 DataLoader) + input_ids = input_ids.to(args.device, non_blocking=True) + labels = labels.to(args.device, non_blocking=True) last_step = step lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: @@ -52,14 +57,20 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): if step % args.log_interval == 0 or step == iters: now = time.time() spend_time = now - start_time - window_its = (step - last_log_step) / max(now - last_log_time, 1e-6) + window_dt = max(now - last_log_time, 1e-6) + window_steps = step - last_log_step + window_its = window_steps / window_dt + # MFU = 实际 FLOPs/s / GPU 峰值；按 window 内的 token 吞吐算（避开 warmup / ckpt save 污染） + # tokens 用 per-rank batch × seq_len × world_size × window_steps（DDP 每卡同时处理 batch） + tokens_window = window_steps * args.batch_size * args.max_seq_len * world_size + mfu = (flops_per_token * tokens_window / window_dt) / (gpu_peak_tflops * 1e12) * 100 last_log_time, last_log_step = now, step current_loss = loss.item() * args.accumulation_steps current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 current_logits_loss = current_loss - current_aux_loss current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 - Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') + Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, MFU: {mfu:.1f}%, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters) and is_main_process(): @@ -108,6 +119,8 @@ if __name__ == "__main__": parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名") parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是，A100 上推荐默认 1）") + parser.add_argument("--gpu_peak_tflops", type=float, default=0, + help="GPU bf16 理论峰值 TFLOPS，用于 MFU 计算；默认 0 表示按 GPU 名自动检测（A100=312）") args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== @@ -139,8 +152,16 @@ if __name__ == "__main__": train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) - optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) - + # fused=True 把所有 param 的 Adam 更新融成一个 CUDA kernel（A100/3090 + CUDA tensor 支持） + optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=True) + + # ========== 5b. MFU 测量准备 ========== + flops_per_token = compute_model_flops_per_token(model, lm_config) + gpu_peak_tflops = detect_gpu_peak_tflops_bf16(override=args.gpu_peak_tflops) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + Logger(f'[MFU] params_flops/token={flops_per_token/1e9:.3f} GFLOPS, ' + f'GPU peak={gpu_peak_tflops:.0f} TFLOPS bf16, world_size={world_size}') + # ========== 6. 从ckp恢复状态 ========== start_epoch, start_step = 0, 0 if ckp_data: @@ -152,8 +173,10 @@ if __name__ == "__main__": # ========== 7. 编译和分布式包装 ========== if args.use_compile == 1: - model = torch.compile(model) - Logger('torch.compile enabled') + # max-autotune 在 Triton 层为每个 fused kernel 试多套 tile/block 配置，挑最快的。 + # 编译时间多 ~30-60s，稳态可比 default 快 5-15%。 + model = torch.compile(model, mode="max-autotune") + Logger('torch.compile enabled (mode=max-autotune)') if dist.is_initialized(): model = DistributedDataParallel(model, device_ids=[local_rank]) @@ -163,7 +186,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=True, persistent_workers=args.num_workers > 0) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') train_epoch(epoch, loader, len(loader) + skip, start_step, wandb) diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 96ad3dd..9347d06 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -41,6 +41,53 @@ def get_lr(current_step, total_steps, lr): return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps))) +# A100/H100/3090 等常见卡的 bf16 理论峰值 TFLOPS。MoE 等价 FLOPs 是 active params 不是 total， +# active params 算出来的 MFU 才有意义；MoE 的 active 比例由 num_experts_per_tok / num_experts 决定。 +_GPU_PEAK_BF16_TFLOPS = { + 'A100': 312, 'A100-SXM4-80GB': 312, 'A100-SXM4-40GB': 312, 'A100-PCIE-80GB': 312, 'A100-PCIE-40GB': 312, + 'H100': 989, 'H100-SXM5-80GB': 989, 'H100-PCIE-80GB': 756, 'H100 80GB HBM3': 989, + 'H200': 989, + 'RTX 3090': 71, 'GeForce RTX 3090': 71, 'RTX 3090 Ti': 80, + 'RTX 4090': 165, 'GeForce RTX 4090': 165, + 'V100': 125, 'V100-SXM2-32GB': 125, + 'L40': 181, 'L40S': 362, +} + + +def detect_gpu_peak_tflops_bf16(device_index: int = 0, override: float = None) -> float: + """Look up the bf16 TensorCore peak for the current GPU; fall back to A100 (312).""" + if override is not None and override > 0: + return float(override) + if not torch.cuda.is_available(): + return 312.0 + name = torch.cuda.get_device_name(device_index) + for key, peak in _GPU_PEAK_BF16_TFLOPS.items(): + if key in name: + return float(peak) + Logger(f'[MFU] Unknown GPU "{name}", falling back to A100 (312 TFLOPS). Pass --gpu_peak_tflops to override.') + return 312.0 + + +def compute_model_flops_per_token(model, config=None) -> float: + """Compute the "model FLOPs" per training token (forward + backward). + + Standard approximation for dense transformers: 6 × P, where P is the total + trainable params count. For MoE we use active params = base + experts_per_tok × per_expert. + """ + raw = model.module if isinstance(model, DistributedDataParallel) else model + raw = getattr(raw, '_orig_mod', raw) + total = sum(p.numel() for p in raw.parameters()) + if config is not None and getattr(config, 'use_moe', False): + n_routed = getattr(config, 'num_experts', 0) + n_active = getattr(config, 'num_experts_per_tok', 1) + if n_routed > 0: + expert = sum(p.numel() for n, p in raw.named_parameters() if 'mlp.experts.0.' in n) + base = total - expert * n_routed + active = base + expert * n_active + return 6.0 * active + return 6.0 * total + + def init_distributed_mode(): if int(os.environ.get("RANK", -1)) == -1: return 0 # 非DDP模式