On Hands Profiling

import torch.cuda.nvtx as nvtx

for i, batch in enumerate(loader):
    nvtx.range_push(f"step_{i}")

    nvtx.range_push("forward")
    out = model(batch)
    loss = criterion(out, target)
    nvtx.range_pop()

    nvtx.range_push("backward")
    loss.backward()
    nvtx.range_pop()

    nvtx.range_push("optimizer")
    optimizer.step()
    optimizer.zero_grad()
    nvtx.range_pop()

    nvtx.range_pop()  # step

nsys 로 launch:

nsys profile -t cuda,nvtx,osrt \
  --capture-range=cudaProfilerApi \
  -o trace python train.py

결과 — timeline 위에 “step_0”, “forward”, “backward”, “optimizer” 가 nested range 로 표시. iteration 안에서 어디가 비는지가 명확해진다.

PyTorch 가 자동으로 추가하는 NVTX 마커도 있다 — autograd 의 op 이름들. 하지만 사용자가 자기 모듈 단위로 한 단계 추가하는 게 가독성을 비약적으로 높인다.

# 예 — fp16 GEMM kernel 한 자리
void cutlass::Kernel<
  cutlass_80_tensorop_f16_s16816gemm_f16_128x128_32x4_tn_align8
>(...);

각 부분의 의미.

• cutlass — CUTLASS template 위에서 만들어진 kernel.
• 80 — Ampere (SM 8.0) 용. 이게 자기 GPU 와 맞아야 한다. H100 위에서 80 prefix 면 fallback — 차선.
• tensorop — Tensor Core 사용. (없으면 CUDA core 만 — 큰 손해)
• f16_s16816 — fp16 input, MMA 모양 16×16×16.
• 128x128_32x4 — CTA tile 128×128, K 축 32, num_stages 4.
• tn — A 가 transposed (T), B 가 normal (N).
• align8 — 8-element aligned load (8×16-bit = 128-bit vector).

이 한 줄에서 즉시 진단 가능한 것.

• “H100 인데 sm80 fallback” → driver/build target 확인.
• “tensorop 이 안 보임” → fp32 input 일 수도 있음. dtype cast 점검.
• “align4” 처럼 작은 align → input pointer 가 align 안 됨. contiguous() 추가.
• “128x128 인데 작은 batch” → tail effect. 더 작은 tile 의 kernel 이 나았을 수도.

cost 1 — torch.compile 의 recompile. compile 의 cache 가 shape 별로 잡혀 있어서, 새 shape 가 들어오면 recompile. 1-2 초 stall. 처음 N iteration 만 그렇고 그 다음엔 cache hit — 하지만 production 의 처음 N 번이 critical.

해결 — torch.compile(dynamic=True) 또는 mark_dynamic 으로 일부 dim 을 dynamic 으로 표시. 컴파일러가 단일 generic 커널 생성. 속도는 약간 떨어지지만 recompile 안 일어남.

cost 2 — cudaMalloc 의 비용. 매 iteration 마다 새 size 의 buffer 가 필요. CUDA caching allocator (PyTorch 의 default) 가 이전 size 의 buffer 를 재사용 못 하면 새 cudaMalloc — 0.1-1 ms stall.

해결 — pre-allocated pool, torch.cuda.set_per_process_memory_fraction, 또는 caching_allocator_warmup. PyTorch profiler 의 memory section 으로 확인.

import torch.profiler as profiler
from torch.profiler import profile, record_function

# memory profile 켜기
with profile(profile_memory=True,
             record_shapes=True) as prof:
    for _ in range(5):
        out = model(batch)
        loss = criterion(out, label)
        loss.backward()
        optimizer.step()

# peak vs allocated 분리
print(prof.key_averages().table(
    sort_by="self_cuda_memory_usage",
    row_limit=10))

Taylor 가 강조한 디테일.

• activation memory 는 backward 시점에 peak — forward 끝까지 모든 layer 의 activation 이 남아 있어야 backward 가 가능.
• optimizer state 는 학습 시작 시 peak — Adam 은 m, v 의 두 부가 state 가 weight 와 같은 크기.
• gradient + 작은 임시 buffer — 보통 작지만 큰 batch 에서 늘어남.

처방. activation 이 큰 자리 → gradient checkpointing (recompute). optimizer 가 큰 자리 → ZeRO 또는 8-bit Adam. fragment 가 큰 자리 → allocator 의 fragmentation 분석.

flame chart — 함수 호출 stack 을 위로 쌓는다. 가장 많이 시간 쓰는 함수가 가장 넓은 막대로. CPU profiler 의 표준.

icicle plot — 같은 stack 을 위에서 아래로. CPU 코드와 GPU 코드를 같은 화면에 놓을 때 GPU 의 “이 kernel 이 어떤 CPU op 가 launch 한 결과인가” 의 부모-자식 관계를 표현.

chrome://tracing 의 PyTorch profiler 가 icicle 을 쓴다. 위쪽 lane 이 CPU, 아래쪽이 GPU. flow event 화살표가 CPU op 와 그것이 launch 한 GPU kernel 을 잇는다.