profile.

scripts/analyze_profile.py

@@ -0,0 +1,197 @@
+"""
+Parse torch.profiler chrome trace, summarize top ops + backend detection.
+
+Usage:
+    # trace_out/ 是 train_pretrain.py --profile_out 指定的目录
+    python analyze_profile.py --trace_dir ../profile_out
+
+    # 或指定具体 trace 文件（一个 rank 一个）
+    python analyze_profile.py --trace_file ../profile_out/DEVvm54331_XXX.pt.trace.json
+"""
+import argparse
+import glob
+import gzip
+import json
+import os
+from collections import defaultdict
+
+
+# op 名 → bucket 分类规则（按前缀 / substring 匹配）
+BUCKET_RULES = [
+    ('matmul (GEMM)',        ['aten::mm', 'aten::bmm', 'aten::addmm', 'aten::linear']),
+    ('attention (SDPA/FA)',  ['flash', 'scaled_dot_product', 'sdpa', 'attention',
+                              'aten::_scaled_dot_product', 'FlashAttn', 'mem_efficient',
+                              'triton_flash', 'ampere_fp16_s16816gemm_bf16']),
+    ('elementwise (norm/act)', ['rms_norm', 'RMSNorm', 'aten::silu', 'aten::mul', 'aten::add',
+                                 'aten::div', 'aten::rsqrt', 'aten::sub', 'element_wise']),
+    ('softmax',              ['softmax', 'log_softmax', 'cross_entropy']),
+    ('rope',                 ['rope', 'rotary', 'cos', 'sin']),
+    ('embedding',            ['embedding', 'aten::embedding']),
+    ('dropout',              ['dropout']),
+    ('reduce (all-reduce/gather)', ['nccl', 'ncclAllReduce', 'ncclAllGather', 'ncclReduceScatter',
+                                     'allreduce', 'AllReduce', 'ProcessGroupNCCL']),
+    ('memory (copy/permute/view)', ['aten::to', 'aten::_to_copy', 'aten::contiguous',
+                                     'aten::permute', 'aten::view', 'aten::reshape',
+                                     'aten::transpose', 'aten::clone', 'Memcpy']),
+    ('optimizer',            ['adam', 'AdamW', 'optim', 'clip_grad']),
+]
+
+# attention backend 识别（chrome trace event name）
+SDPA_BACKENDS = {
+    'flash':      ['FlashAttn', 'flash_attn', 'triton_flash', 'ampere_fp16_s16816gemm'],
+    'mem_efficient': ['mem_efficient', 'MemEfficient', 'cutlassF'],
+    'math':       ['aten::baddbmm', 'aten::softmax'],  # naive fallback
+}
+
+
+def bucket_op(name: str) -> str:
+    ln = name.lower()
+    for bucket, patterns in BUCKET_RULES:
+        for p in patterns:
+            if p.lower() in ln:
+                return bucket
+    return 'other'
+
+
+def load_trace(path):
+    """支持 .json 和 .json.gz。"""
+    opener = gzip.open if path.endswith('.gz') else open
+    with opener(path, 'rt') as f:
+        data = json.load(f)
+    return data.get('traceEvents', data)
+
+
+def analyze(events):
+    # 只统计 CUDA kernel 时间（cat == 'kernel' 或 'cuda_runtime'）
+    # 但 torch.profiler 输出的 chrome trace 里 CUDA op 是 name+dur 的 complete event (ph=='X')
+    stats = defaultdict(lambda: {'time_us': 0, 'count': 0})
+    total_cuda_us = 0
+    sdpa_backend_hits = defaultdict(int)
+
+    for e in events:
+        if e.get('ph') != 'X':
+            continue
+        cat = e.get('cat', '')
+        name = e.get('name', '')
+        dur = e.get('dur', 0)
+        # 只关心 GPU kernel + collective + cudaLaunch，跳过 python-level
+        if cat not in ('kernel', 'cuda_runtime', 'user_annotation', 'gpu_memcpy', 'gpu_memset', 'cpu_op'):
+            continue
+        if cat == 'cpu_op' and 'ncclAllReduce' not in name and 'sync' not in name.lower():
+            # CPU op 大部分是 dispatch overhead，除了 NCCL / sync 都跳
+            continue
+        stats[name]['time_us'] += dur
+        stats[name]['count'] += 1
+        if cat in ('kernel', 'gpu_memcpy', 'gpu_memset'):
+            total_cuda_us += dur
+
+        # 探测 attention backend
+        for backend, patterns in SDPA_BACKENDS.items():
+            for p in patterns:
+                if p in name:
+                    sdpa_backend_hits[backend] += dur
+
+    return stats, total_cuda_us, sdpa_backend_hits
+
+
+def format_us(us):
+    if us > 1e6: return f'{us/1e6:.2f} s'
+    if us > 1e3: return f'{us/1e3:.2f} ms'
+    return f'{us:.0f} us'
+
+
+def report(stats, total_cuda_us, sdpa_backend_hits, top_n=25):
+    print(f'\n{"=" * 80}')
+    print(f'  Total CUDA time in trace: {format_us(total_cuda_us)}')
+    print(f'{"=" * 80}')
+
+    # bucket 汇总
+    bucket_time = defaultdict(int)
+    bucket_count = defaultdict(int)
+    for name, s in stats.items():
+        b = bucket_op(name)
+        bucket_time[b] += s['time_us']
+        bucket_count[b] += s['count']
+
+    print(f'\n  Time by bucket (% of total CUDA):')
+    print(f'  {"bucket":40s} {"time":>12s} {"% total":>8s} {"# calls":>8s}')
+    for bucket, t in sorted(bucket_time.items(), key=lambda kv: -kv[1]):
+        pct = t / max(total_cuda_us, 1) * 100
+        print(f'  {bucket:40s} {format_us(t):>12s} {pct:>7.1f}% {bucket_count[bucket]:>8d}')
+
+    # attention backend 探测
+    print(f'\n  Attention backend detection:')
+    if sdpa_backend_hits:
+        winner = max(sdpa_backend_hits.items(), key=lambda kv: kv[1])
+        for backend, t in sorted(sdpa_backend_hits.items(), key=lambda kv: -kv[1]):
+            marker = ' ← DOMINANT' if backend == winner[0] else ''
+            print(f'    {backend:20s}: {format_us(t):>12s}{marker}')
+        if winner[0] != 'flash':
+            print(f'  ⚠️  Not using FlashAttention. Expected FA2 on A100 for hidden 128/heads 16 → suboptimal.')
+            print(f'     Fix: 训练里用 `torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION)` context')
+    else:
+        print(f'    (未识别到 attention kernel——可能 kernel 名 encoded 或未走 SDPA)')
+
+    # top ops 明细
+    print(f'\n  Top {top_n} CUDA ops:')
+    print(f'  {"rank":>4s} {"time":>12s} {"% total":>8s} {"# calls":>8s}  {"name":s}')
+    sorted_ops = sorted(stats.items(), key=lambda kv: -kv[1]['time_us'])
+    for i, (name, s) in enumerate(sorted_ops[:top_n], 1):
+        pct = s['time_us'] / max(total_cuda_us, 1) * 100
+        # 截断长 kernel 名
+        display_name = name if len(name) < 80 else name[:77] + '...'
+        print(f'  {i:>4d} {format_us(s["time_us"]):>12s} {pct:>7.1f}% {s["count"]:>8d}  {display_name}')
+
+    # 分类建议
+    print(f'\n{"=" * 80}')
+    print(f'  优化提示（按 bucket 比例）:')
+    print(f'{"=" * 80}')
+    total = total_cuda_us
+    reduce_pct = bucket_time.get('reduce (all-reduce/gather)', 0) / max(total, 1) * 100
+    if reduce_pct > 15:
+        print(f'  ⚠️  reduce/collective 占 {reduce_pct:.1f}%（>15% 算高）→ 考虑：')
+        print(f'      - `no_sync()` 包 gradient accumulation 减少 3× 冗余 all-reduce')
+        print(f'      - FSDP: 用 ZeRO-2/3 拆 optimizer state 减少每次 sync 数据量')
+    elem_pct = bucket_time.get('elementwise (norm/act)', 0) / max(total, 1) * 100
+    if elem_pct > 15:
+        print(f'  ⚠️  elementwise 占 {elem_pct:.1f}%（>15% 算高）→ memory bandwidth bound')
+        print(f'      - Fused RMSNorm + SiLU + Add kernel（apex 或 triton）省 50%+')
+    mem_pct = bucket_time.get('memory (copy/permute/view)', 0) / max(total, 1) * 100
+    if mem_pct > 8:
+        print(f'  ⚠️  memory movement 占 {mem_pct:.1f}%（>8% 算高）→ 布局问题')
+        print(f'      - 检查 permute/transpose/contiguous 频次，可能重复布局转换')
+    attn_pct = bucket_time.get('attention (SDPA/FA)', 0) / max(total, 1) * 100
+    matmul_pct = bucket_time.get('matmul (GEMM)', 0) / max(total, 1) * 100
+    print(f'\n  Matmul + Attention (核心 compute): {matmul_pct + attn_pct:.1f}%')
+    print(f'  如果这个数 <50%，多数时间在做非核心 op → 有巨大优化空间')
+    print(f'  Llama 3 训练里这个数一般 60-75%')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--trace_file', type=str, default=None)
+    parser.add_argument('--trace_dir', type=str, default='../profile_out')
+    parser.add_argument('--top_n', type=int, default=25)
+    args = parser.parse_args()
+
+    if args.trace_file:
+        files = [args.trace_file]
+    else:
+        # torch.profiler 每 rank 一个文件；分析 rank 0 即可
+        files = sorted(glob.glob(os.path.join(args.trace_dir, '*.pt.trace.json*')))
+        if not files:
+            print(f'no trace file found in {args.trace_dir}')
+            return
+        # 只取第一个（一般 rank 0）
+        files = files[:1]
+
+    for f in files:
+        print(f'\n[loading] {f} ({os.path.getsize(f)/1e6:.1f} MB)')
+        events = load_trace(f)
+        print(f'[loaded] {len(events):,} events')
+        stats, total, sdpa_hits = analyze(events)
+        report(stats, total, sdpa_hits, top_n=args.top_n)
+
+
+if __name__ == "__main__":
+    main()

scripts/launch_1b_pretrain.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+# 1B pretrain launcher — 4 卡 DDP，SkyPile stage 1（真 raw web pretrain）
+#
+# Config 基于 plan L84 + dry-run 学到的教训：
+#   - hidden=2048, layers=20, heads=16 (kv=4 GQA) → ~1.02B params
+#   - seq=1024 (SkyPile 长度足够)
+#   - global batch = 16 per GPU × 4 GPU × accum 4 = 256 → ~256*1024=262k tokens/step
+#   - lr 3e-4 (1B from-scratch 标准) + warmup 1000 (前 500 步发散风险 >50%)
+#   - eval 每 500 step，用 minimind pretrain 做 proxy holdout（1B 未见过）
+#   - save 每 2000 step（2GB × 264 次约 500GB 磁盘，可承受；1B 训 ~150k step 总量）
+#   - grad_checkpointing 0（1B seq=1024 在 A100 80GB 用不上，省 30% 时间）
+#
+# Stage 2 (minimind curated Q&A mid-training) 是另开 run，从 pretrain_1b ckpt 起点继续
+#
+# 用法:
+#   scripts/launch_1b_pretrain.sh                # 默认 4 卡
+#   scripts/launch_1b_pretrain.sh --nproc 2       # 覆盖为 2 卡
+#   scripts/launch_1b_pretrain.sh -- --learning_rate 5e-4  # 传给 python
+set -e
+
+# 默认参数
+NPROC=4
+# 分离 launcher 参数和 python 参数
+PY_ARGS=()
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --nproc) NPROC="$2"; shift 2 ;;
+        --) shift; PY_ARGS+=("$@"); break ;;
+        *) PY_ARGS+=("$1"); shift ;;
+    esac
+done
+
+REPO_ROOT=$(cd "$(dirname "$0")/.." && pwd)
+cd "$REPO_ROOT/trainer"
+
+# 数据：SkyPile 是 stage 1 主战场
+DATA_PATH="../dataset/skypile_test.jsonl"
+# eval holdout：用 minimind pretrain 作 proxy（1B 训练时未见过）
+EVAL_DATA="../dataset/pretrain_t2t.jsonl"
+
+echo "[launch_1b_pretrain] NPROC=$NPROC data=$DATA_PATH"
+echo "[launch_1b_pretrain] extra args: ${PY_ARGS[*]}"
+
+torchrun --nproc_per_node=$NPROC --standalone train_pretrain.py \
+    --hidden_size 2048 \
+    --num_hidden_layers 20 \
+    --max_seq_len 1024 \
+    --batch_size 24 \
+    --accumulation_steps 4 \
+    --learning_rate 3e-4 \
+    --warmup_steps 1000 \
+    --grad_checkpointing 0 \
+    --dtype bfloat16 \
+    --num_workers 4 \
+    --use_compile 1 \
+    --data_path "$DATA_PATH" \
+    --eval_interval 500 \
+    --eval_data_path "$EVAL_DATA" \
+    --eval_max_batches 20 \
+    --save_interval 500 \
+    --log_interval 20 \
+    --save_weight pretrain_1b \
+    --save_dir ../out \
+    --wandb_project MiniMind-1B-Pretrain \
+    "${PY_ARGS[@]}"

trainer/train_pretrain.py

@@ -51,7 +51,7 @@ def eval_holdout(model, eval_loader, max_batches, autocast_ctx, device):
     return avg_loss, math.exp(avg_loss)
 
 
-def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None):
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None, profiler=None):
     start_time = time.time()
     last_log_time, last_log_step = start_time, start_step
     last_step = start_step
@@ -70,11 +70,20 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
+        # DDP gradient accumulation optimization：
+        # 默认 DDP 每次 .backward() 都做全参数 all-reduce（4GB grad × 85 bucket ≈ 30ms）。
+        # 但 accum=N 时前 N-1 次 all-reduce 是浪费（grad 还要继续累积）。用 model.no_sync()
+        # 让 DDP 在非最后一个 micro-step 里 skip all-reduce，只在最后一个 micro-step 才同步。
+        # profile 显示 all-reduce 占 73% 时间 → 修完预期总时间省 ~55%，MFU 17%→~35%。
+        is_last_accum_step = (step % args.accumulation_steps == 0)
+        need_sync = is_last_accum_step or not isinstance(model, DistributedDataParallel)
+        sync_ctx = nullcontext() if need_sync else model.no_sync()
+
+        with sync_ctx:
             with autocast_ctx:
                 res = model(input_ids, labels=labels)
                 loss = res.loss + res.aux_loss
                 loss = loss / args.accumulation_steps
-
             scaler.scale(loss).backward()
 
         if step % args.accumulation_steps == 0:
@@ -86,16 +95,24 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
 
             optimizer.zero_grad(set_to_none=True)
 
+        # profiler step 打点（每 step 一次；schedule 决定何时真的记录）
+        if profiler is not None:
+            profiler.step()
+
         if step % args.log_interval == 0 or step == iters:
             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
-            # MFU = 实际 FLOPs/s / GPU 峰值；按 window 内的 token 吞吐算（避开 warmup / ckpt save 污染）
-            # tokens 用 per-rank batch × seq_len × world_size × window_steps（DDP 每卡同时处理 batch）
+            # MFU = per-card 实际 FLOPs/s / per-card 峰值。分子分母都要一致口径：
+            # - 分子: 4 卡总 tokens × flops/token = 总 flops；÷ world_size 得 per-card flops
+            # - 分母: gpu_peak_tflops (单卡峰值，312 TFLOPS on A100)
+            # 早期版本 bug：分子用全部 tokens 但分母只除 1 卡峰值 → 报数是真实 MFU 的 world_size 倍。
             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
+            total_flops_per_s = flops_per_token * tokens_window / window_dt
+            per_card_flops_per_s = total_flops_per_s / world_size
+            mfu = per_card_flops_per_s / (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
@@ -107,11 +124,23 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
 
         # ---- eval loop（可选） ----
         if args.eval_interval > 0 and step % args.eval_interval == 0 and eval_loader is not None:
+            # ⚠️ 全 rank 都跑 eval：单 rank eval (即使 unwrap raw model) 也会导致 DDP guard 失效 /
+            # torch.compile recompile → rank 0 长 stall → 下一步 all-reduce seq 不同步 → NCCL 10min timeout。
+            # 解决：所有 rank 用同一份 eval_loader (顺序一致) 跑同样的 batch，然后 all-reduce mean loss。
+            # 每 rank 都工作，无 rank 掉队问题。max_batches=20 * 5 rank = 20 batch × ~0.3s ≈ 6s 开销。
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            val_loss_local, _ = eval_holdout(raw_model, eval_loader, args.eval_max_batches, autocast_ctx, args.device)
+            if dist.is_initialized():
+                loss_t = torch.tensor(val_loss_local, device=args.device)
+                dist.all_reduce(loss_t, op=dist.ReduceOp.AVG)
+                val_loss = loss_t.item()
+            else:
+                val_loss = val_loss_local
+            val_ppl = math.exp(val_loss)
             if is_main_process():
-                val_loss, val_ppl = eval_holdout(model, eval_loader, args.eval_max_batches, autocast_ctx, args.device)
                 Logger(f'[eval] step={step} val_loss={val_loss:.4f} val_ppl={val_ppl:.2f}')
                 if wandb: wandb.log({"val_loss": val_loss, "val_ppl": val_ppl, "step": step})
-            if dist.is_initialized(): dist.barrier()
 
         if step % args.save_interval == 0 or step == iters:
             if is_main_process():
@@ -181,6 +210,11 @@ if __name__ == "__main__":
                         help="holdout jsonl 路径；不设则关闭 eval")
     parser.add_argument("--eval_max_batches", type=int, default=20,
                         help="每次 eval 最多跑几个 batch（20*batch_size 通常够 ~1s 内出数）")
+    parser.add_argument("--profile_steps", type=int, default=0,
+                        help="torch.profiler 抓 N 步（0=关）；建议 20。会在 wait=10 warmup=5 active=N 后自动退出。"
+                             "trace 存到 --profile_out 目录，chrome://tracing 打开或 tensorboard --logdir profile_out")
+    parser.add_argument("--profile_out", type=str, default="../profile_out",
+                        help="profile trace 输出目录")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
@@ -247,16 +281,34 @@ if __name__ == "__main__":
         model = DistributedDataParallel(model, device_ids=[local_rank])
 
     # ========== 7b. eval loader（可选） ==========
+    # 全 rank 都建 eval loader（不能只 rank 0，否则 all-reduce 时 rank 1/2/3 无 eval loss 参与聚合）
     eval_loader = None
-    if args.eval_interval > 0 and args.eval_data_path and is_main_process():
-        # rank-0 only：只 main rank 需要 eval loader（其他 rank barrier 等待）
+    if args.eval_interval > 0 and args.eval_data_path:
         eval_ds = PretrainDataset(args.eval_data_path, tokenizer, max_length=args.max_seq_len)
+        # 所有 rank 用同一份数据 + shuffle=False → 保证跑相同 batch → all-reduce 有意义
         eval_loader = DataLoader(eval_ds, batch_size=args.batch_size,
                                  shuffle=False, num_workers=2, pin_memory=True)
         Logger(f'[eval] loader ready: {len(eval_ds)} samples from {args.eval_data_path}, '
                f'eval_interval={args.eval_interval}, max_batches={args.eval_max_batches}')
 
+    # ========== 7c. profiler（可选） ==========
+    profiler_ctx = None
+    if args.profile_steps > 0 and is_main_process():
+        # schedule: 前 10 步 warmup（skip compile 编译期），5 步观察，profile_steps 步真记录
+        from torch.profiler import profile, schedule, ProfilerActivity, tensorboard_trace_handler
+        os.makedirs(args.profile_out, exist_ok=True)
+        profiler_ctx = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            schedule=schedule(wait=10, warmup=5, active=args.profile_steps, repeat=1),
+            on_trace_ready=tensorboard_trace_handler(args.profile_out),
+            record_shapes=True, with_stack=False, with_flops=False,
+        )
+        Logger(f'[profile] enabled: warmup 15 step → capture {args.profile_steps} step → trace to {args.profile_out}')
+        Logger('[profile] 训练完 profile 步后不自动退出——手动 Ctrl+C 停即可。trace 已 flush 到磁盘。')
+
     # ========== 8. 开始训练 ==========
+    if profiler_ctx is not None:
+        profiler_ctx.__enter__()
     for epoch in range(start_epoch, args.epochs):
         train_sampler and train_sampler.set_epoch(epoch)
         setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
@@ -266,9 +318,11 @@ if __name__ == "__main__":
                             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, eval_loader)
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb, eval_loader, profiler_ctx)
         else:
-            train_epoch(epoch, loader, len(loader), 0, wandb, eval_loader)
+            train_epoch(epoch, loader, len(loader), 0, wandb, eval_loader, profiler_ctx)
+    if profiler_ctx is not None:
+        profiler_ctx.__exit__(None, None, None)
 
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized():