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