Merge b59f54cfc5 into 512eed0b65

scripts/download_skypile.py

@@ -0,0 +1,163 @@
+"""
+Streaming 下载 SkyPile-150B 子采样，转成 minimind PretrainDataset 兼容格式（{"text": ...}）。
+
+特性：
+- Streaming：不全下，边读边写，磁盘只占目标量
+- Token 计数：精确控制下载量（按 budget 而非样本数停）
+- 断点续传：进程被杀也能从最后写到的行接着下
+- 进度日志：每 N 个样本打印进度
+- 简易过滤：太短 / 重复字符多的样本跳过
+
+Usage:
+    # 默认下 25B tokens（1B 模型 Chinchilla 最优）
+    python download_skypile.py
+
+    # 改下载量
+    python download_skypile.py --target_tokens 10_000_000_000
+
+    # 改输出路径
+    python download_skypile.py --output ../dataset/skypile_25b.jsonl
+
+需要 fwdproxy（HuggingFace 在 allowlist 但 OD 无直接 DNS）：
+    export https_proxy=http://fwdproxy:8080
+    export http_proxy=http://fwdproxy:8080
+"""
+import os
+import sys
+
+__package__ = "scripts"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import argparse
+import json
+import time
+import warnings
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+warnings.filterwarnings('ignore')
+
+
+def count_tokens_fast(text: str) -> int:
+    """快速估算 token 数：中文按字符 ≈ token，英文按词 × 1.3。
+    用于在线计数（精确 tokenize 太慢）。"""
+    n_chars = len(text)
+    # 中文 token ≈ 1 char；英文 1 token ≈ 4 char
+    # 简单估算：取个折中
+    return int(n_chars / 1.8)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Stream-download SkyPile-150B subset")
+    parser.add_argument('--repo', default='Skywork/SkyPile-150B', type=str)
+    parser.add_argument('--split', default='train', type=str)
+    parser.add_argument('--output', default='../dataset/skypile_25b.jsonl', type=str)
+    parser.add_argument('--target_tokens', type=int, default=25_000_000_000,
+                        help='目标 token 数（estimate），到达后停止')
+    parser.add_argument('--min_length', type=int, default=200,
+                        help='样本最小字符数，太短跳过')
+    parser.add_argument('--log_every', type=int, default=5000)
+    parser.add_argument('--tokenizer_path', type=str, default='../model',
+                        help='用 minimind tokenizer 精确计算 token 数（每 N 个样本核对一次）')
+    parser.add_argument('--exact_count_every', type=int, default=50000,
+                        help='每 N 个样本精确 tokenize 一次校准估算误差')
+    args = parser.parse_args()
+
+    output_path = os.path.abspath(args.output)
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    # 断点续传：count 已有文件的行数
+    n_written, est_tokens_written = 0, 0
+    if os.path.exists(output_path):
+        print(f'[resume] found existing {output_path}, counting...')
+        with open(output_path) as f:
+            for line in f:
+                n_written += 1
+                try:
+                    est_tokens_written += count_tokens_fast(json.loads(line)['text'])
+                except Exception:
+                    pass
+        print(f'[resume] {n_written:,} samples already written, ~{est_tokens_written/1e9:.2f}B tokens')
+        if est_tokens_written >= args.target_tokens:
+            print('[resume] target already reached, exit')
+            return
+
+    # 校准 tokenizer（可选；下载完后跑 eval_perplexity 时会用到）
+    tokenizer = None
+    if args.tokenizer_path:
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
+            print(f'[tokenizer] loaded for calibration: vocab={tokenizer.vocab_size}')
+        except Exception as e:
+            print(f'[tokenizer] failed to load ({e}), will skip calibration')
+
+    print(f'[start] streaming {args.repo} split={args.split} → {output_path}')
+    print(f'[target] {args.target_tokens/1e9:.1f}B tokens (estimated)')
+
+    ds = load_dataset(args.repo, split=args.split, streaming=True)
+
+    skipped_short = 0
+    n_seen = 0
+    t0 = time.time()
+    calibration_ratio = 1.0  # actual_tokens / estimated_tokens，每次校准更新
+
+    with open(output_path, 'a') as fout:
+        for sample in ds:
+            n_seen += 1
+            # 跳过已写过的（断点续传）
+            if n_seen <= n_written:
+                continue
+
+            text = sample.get('text', '')
+            if len(text) < args.min_length:
+                skipped_short += 1
+                continue
+
+            est_tok = count_tokens_fast(text)
+            actual_est = int(est_tok * calibration_ratio)
+
+            # 写入 minimind 格式
+            fout.write(json.dumps({'text': text}, ensure_ascii=False) + '\n')
+            n_written += 1
+            est_tokens_written += actual_est
+
+            # 进度日志
+            if n_written % args.log_every == 0:
+                elapsed = time.time() - t0
+                rate = (n_written / elapsed) if elapsed > 0 else 0
+                progress = est_tokens_written / args.target_tokens * 100
+                eta_h = (args.target_tokens - est_tokens_written) / max(est_tokens_written / elapsed, 1) / 3600
+                print(f'[{n_written:>9,} samples] {est_tokens_written/1e9:>5.2f}B tok '
+                      f'({progress:>5.1f}%) | {rate:.0f} samples/s | ETA {eta_h:.1f}h | '
+                      f'skipped_short={skipped_short:,}')
+                fout.flush()
+
+            # 周期性精确 tokenize 几条校准估算
+            if tokenizer and n_written % args.exact_count_every == 0:
+                actual = sum(len(tokenizer(t, add_special_tokens=False).input_ids)
+                             for t in [text])
+                estimate = est_tok
+                if estimate > 0:
+                    new_ratio = actual / estimate
+                    calibration_ratio = 0.9 * calibration_ratio + 0.1 * new_ratio
+                    print(f'[calibration] est_ratio updated → {calibration_ratio:.3f} '
+                          f'(this sample: est={estimate}, actual={actual})')
+
+            if est_tokens_written >= args.target_tokens:
+                print(f'\n[done] target reached: {est_tokens_written/1e9:.2f}B est tokens')
+                break
+
+    elapsed = time.time() - t0
+    print(f'\n{"=" * 60}')
+    print(f'Total: {n_written:,} samples, ~{est_tokens_written/1e9:.2f}B est tokens')
+    print(f'Elapsed: {elapsed/60:.1f} min ({elapsed/3600:.1f}h)')
+    print(f'Skipped (too short): {skipped_short:,}')
+    print(f'Output: {output_path} ({os.path.getsize(output_path)/1e9:.2f} GB)')
+    print(f'{"=" * 60}')
+    if tokenizer:
+        print(f'⚠️  估算 token 数有 ~{abs(1-calibration_ratio)*100:.1f}% 误差。')
+        print(f'   精确数：跑 eval_perplexity.py 时会有 actual count。')
+
+
+if __name__ == "__main__":
+    main()

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

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

trainer/train_full_sft.py

@@ -16,17 +16,23 @@ from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from model.model_minimind import MiniMindConfig
 from dataset.lm_dataset import SFTDataset
-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')
 
 
 def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
     start_time = time.time()
+    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 仅在 --enable_optimizations=1 + pin_memory=True 时有效
+        input_ids = input_ids.to(args.device, non_blocking=args.enable_optimizations == 1)
+        labels = labels.to(args.device, non_blocking=args.enable_optimizations == 1)
         last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
@@ -49,14 +55,21 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             optimizer.zero_grad(set_to_none=True)
 
         if step % args.log_interval == 0 or step == iters:
-            spend_time = time.time() - start_time
+            now = time.time()
+            spend_time = now - start_time
+            window_dt = max(now - last_log_time, 1e-6)
+            window_steps = step - last_log_step
+            window_its = window_steps / window_dt
+            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'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), 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})
+            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, "it/s": window_its, "MFU": mfu})
 
         if (step % args.save_interval == 0 or step == iters) and is_main_process():
             model.eval()
@@ -104,7 +117,13 @@ if __name__ == "__main__":
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT", help="wandb项目名")
-    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1],
+                        help="是否使用 torch.compile 加速（0=否，1=是，A100 上推荐 1）")
+    parser.add_argument("--enable_optimizations", default=1, type=int, choices=[0, 1],
+                        help="开关：A/B 对比优化效果。1=启用 fused AdamW + non_blocking + persistent_workers（默认）；"
+                             "0=baseline（关掉这些，仅保留 log 测量）。compile 由 --use_compile 独立控制。")
+    parser.add_argument("--gpu_peak_tflops", type=float, default=0,
+                        help="GPU bf16 理论峰值 TFLOPS（用于 MFU 计算）；默认 0 表示按 GPU 名自动检测")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
@@ -136,7 +155,19 @@ if __name__ == "__main__":
     train_ds = SFTDataset(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
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate,
+                            fused=(args.enable_optimizations == 1))
+
+    # ========== 5b. MFU 测量准备（始终启用，无 perf 影响）==========
+    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}')
+    Logger(f'[OPT] enable_optimizations={args.enable_optimizations} '
+           f'(non_blocking={args.enable_optimizations}, fused_adam={args.enable_optimizations}, '
+           f'persistent_workers={args.enable_optimizations}), use_compile={args.use_compile}')
 
     # ========== 6. 从ckp恢复状态 ==========
     start_epoch, start_step = 0, 0
@@ -149,8 +180,10 @@ if __name__ == "__main__":
     
     # ========== 7. 编译和分布式包装 ==========
     if args.use_compile == 1:
-        model = torch.compile(model)
-        Logger('torch.compile enabled')
+        # enable_optimizations=1 时用 max-autotune（找更优 kernel），0 时用 default mode
+        compile_mode = "max-autotune" if args.enable_optimizations == 1 else "default"
+        model = torch.compile(model, mode=compile_mode)
+        Logger(f'torch.compile enabled (mode={compile_mode})')
     if dist.is_initialized():
         model = DistributedDataParallel(model, device_ids=[local_rank])
     
@@ -160,7 +193,9 @@ 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.enable_optimizations == 1 and 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)

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,7 +152,15 @@ 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
@@ -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)

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模式