From 21d501f385f8c3df31a654d62a27e8a788c2ff61 Mon Sep 17 00:00:00 2001 From: Longfang Zhao Date: Wed, 1 Jul 2026 10:53:44 -0700 Subject: [PATCH] profile. --- scripts/analyze_profile.py | 197 ++++++++++++++++++++++++++++++++++ scripts/launch_1b_pretrain.sh | 65 +++++++++++ trainer/train_pretrain.py | 84 ++++++++++++--- 3 files changed, 331 insertions(+), 15 deletions(-) create mode 100644 scripts/analyze_profile.py create mode 100755 scripts/launch_1b_pretrain.sh diff --git a/scripts/analyze_profile.py b/scripts/analyze_profile.py new file mode 100644 index 0000000..2d08a6c --- /dev/null +++ b/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() diff --git a/scripts/launch_1b_pretrain.sh b/scripts/launch_1b_pretrain.sh new file mode 100755 index 0000000..8f39a47 --- /dev/null +++ b/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[@]}" diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index cf654b6..031a4cd 100644 --- a/trainer/train_pretrain.py +++ b/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,12 +70,21 @@ 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 - with autocast_ctx: - res = model(input_ids, labels=labels) - loss = res.loss + res.aux_loss - loss = loss / args.accumulation_steps + # 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() - scaler.scale(loss).backward() + 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: scaler.unscale_(optimizer) @@ -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():