File size: 1,946 Bytes
f31f59b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Where does one training micro-step go? Arm B, micro 8 x 1024, bf16. CUDA time by operator group.
python profile_step.py [--arm B] [--attn sdpa]
"""
import argparse, collections
import numpy as np
import torch
from torch.profiler import profile, ProfilerActivity
import train as T
ap = argparse.ArgumentParser()
ap.add_argument("--arm", default="B")
ap.add_argument("--attn", default="sdpa")
a = ap.parse_args()
counts = np.load("data_smoke/unigram_counts.npy")
m = T.Model(a.arm, counts, a.attn).cuda()
opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad], lr=1e-4)
data = T.Data("data_smoke", 0)
def step():
x = data.batch(8).cuda()
with torch.autocast("cuda", dtype=torch.bfloat16):
loss = m.loss(x)
loss.backward()
opt.step()
opt.zero_grad(set_to_none=True)
for _ in range(3):
step()
torch.cuda.synchronize()
with profile(activities=[ProfilerActivity.CUDA]) as prof:
for _ in range(3):
step()
torch.cuda.synchronize()
groups = collections.Counter()
for e in prof.key_averages():
t = e.self_device_time_total if hasattr(e, "self_device_time_total") else e.self_cuda_time_total
k = e.key.lower()
if any(s in k for s in ("flash", "attention", "fmha", "sdpa", "efficient_attention", "softmax_backward_data")):
g = "attention kernels"
elif any(s in k for s in ("gemm", "matmul", "mm", "cutlass", "sm90", "sm100", "sm120", "ampere", "xmma")):
g = "matrix multiplies (all linears incl. readout)"
elif "log_softmax" in k or "nll" in k or "cross_entropy" in k:
g = "cross-entropy"
elif "adam" in k or "foreach" in k:
g = "optimizer"
else:
g = "elementwise / norms / other"
groups[g] += t
tot = sum(groups.values())
for g, t in groups.most_common():
print(f"{g:48s} {t / 1e3 / 3:8.1f} ms/step {t / tot * 100:5.1f}%")
print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=12))
|