프로파일링(Profiling)할 수 없는 것은 최적화(Optimize)할 수 없습니다.

대규모 언어 모델 (LLM)에서 초당 더 많은 토큰을 뽑아내려 하든, 추론(Inference) 시간을 밀리초 단위로 단축하려 하든, 혹은 단순히 왜 학습 루프(Training loop)가 사양서(Spec sheet)에서 약속한 것보다 느리게 돌아가는지 이해하려 하든, 그 경로는 결국 프로파일링으로 이어집니다.

문제는 프로파일링의 진입 장벽이 매우 높다는 점입니다. 트레이스(Traces)는 색칠된 직사각형들이 빽빽하게 들어찬 벽과 같습니다. 이벤트(Events)들은 위협적인 이름을 가지고 있습니다. 대부분의 튜토리얼은 여러분이 이미 이를 읽을 수 있다고 가정합니다. 그래서 프로파일링을 해야 한다는 것을 알고 있음에도 불구하고, 트레이스를 여는 것은 나중으로 미루거나 (혹은 다른 사람에게 맡겨야 할) 귀찮은 일처럼 느껴질 수 있습니다. 이 포스트와 이 포스트가 시작하는 시리즈는 그 진입 장벽을 낮추기 위한 우리의 시도입니다.

이 글은 프로파일러 트레이스를 읽는 기술을 천천히 쌓아 올리고 이를 최적화로 연결하는 시리즈인 **PyTorch에서의 프로파일링 (Profiling in PyTorch)**의 첫 번째 포스트입니다. 계획은 다음과 같습니다:

Part 1 (본 포스트): 가장 단순한 연산인 행렬 곱셈(Matrix multiplication)과 그 뒤를 잇는 편향 더하기(Bias add)로 시작하여, 프로파일러가 반환하는 결과물을 읽는 법을 배웁니다.
Part 2: nn.Linear와 작은 MLP(Multi-Layer Perceptron)로 규모를 키우고, 트레이스를 사용하여 최적화의 동기를 부여하며

시작하기에 앞서, 아래의 모든 내용을 더 잘 이해하기 위한 두 가지 정의를 살펴보겠습니다.

• GPU **커널 (kernel)**은 GPU의 수많은 스레드에서 병렬로 실행되는 프로그램입니다.
• CPU는 이러한 커널을 **스케줄링하고 실행 (schedules and launches)**합니다.

일반적으로 GPU 커널을 직접 작성할 필요는 없습니다. PyTorch 연산을 사용하면, 해당 연산은 GPU에서 작업을 수행하는 하나 이상의 커널로 자동 변환됩니다.

이 두 가지 개념을 숙지했다면, 이제 질문을 던져봅시다.

이 포스트에서 사용하는 전체 스크립트는 다음과 같습니다:

01_matmul_add.py

이 스크립트를 별도의 탭에서 열고 코드를 단계별로 따라가 보시는 것을 권장합니다. 스크립트 실행에는 NVIDIA A100-SXM4-80GB GPU를 사용합니다.

Sara Hooker 박사가 정확하게 비유했듯이, 우리 몸이 주로 물로 구성되어 있는 것처럼, 심층 신경망 (Deep Neural Networks)은 주로 행렬 곱셈 (matrix multiplies)으로 구성되어 있습니다. 행렬 곱셈이 매우 근본적인 연산인 만큼, 프로파일링 여정을 다른 것이 아닌 이것으로 시작하는 것은 매우 의미 있는 일입니다.

def fn(x, w, b):
    return torch.add(torch.matmul(x, w), b)

행렬 곱셈과 함께 이루어지는 행렬 덧셈은 뉴런 내에서 가중치 (weights)와 편향 (biases)이 상호작용하는 방식을 모방합니다. 이 덧셈(addition, 중의적 표현)은 이 포스트의 후반부에서 컴파일 (compilation)을 위한 토대를 이해하는 데 도움이 될 것입니다.

프로파일링을 위해 우리는 torch.profiler 모듈을 사용할 것입니다. 포함된 단계는 다음과 같습니다:

• 프로파일링할 코드를 준비합니다 (여기서는 행렬 곱셈과 행렬 덧셈을 감싸는 def fn).
• 알고리즘에 주석을 답니다 (Annotate). 이는 완전히 선택 사항이지만, 수행하는 것을 권장합니다. record_function은 우리의 함수를 matmul_add로 주석 처리하며, 이는 (나중에 언급하겠지만) 트레이스 (traces)에서 탐색하기 쉽게 만들어 줍니다.

def step():
    with torch.profiler.record_function("matmul_add"):
        return fn(x, w, b)

• 코드를 torch.profiler.profile 컨텍스트 매니저 (context manager)로 감쌉니다.

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU, # CPU 활동
        ...

• 프로파일을 내보냅니다 (Export).

프로파일러 테이블 (profiler table)

prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)

프로파일러 트레이스 (profiler trace)

...

프로파일러는 두 가지 별개의 결과물 (artifacts)을 내보냅니다:

• 프로파일러 테이블 (The profiler table): 알고리즘의 통계적 요약을 제공합니다. "무엇이 가장 많은 시간을 차지하는가"에 대한 답을 줍니다. 이는 핫스팟 (hotspots)을 파악하는 데 매우 유용합니다. 핫스팟이란 가장 많은 시간을 차지하는 이벤트, 파이프라인의 병목 (bottleneck)이 될 수 있는 이벤트, 또는 매우 빈번하게 트리거되는 이벤트를 의미합니다.
• 프로파일러 트레이스 (The profiler trace): 시간적 실행 뷰 (temporal execution view)를 제공합니다. "언제, 왜 특정 연산이 발생했는가"에 대한 답을 주며, CPU와 GPU에서 발생하는 활동들을 묘사합니다. 이는 실행된 커널 (kernel)들, 커널 실행 지연, CPU와 GPU 활동 간의 중첩 (overlap) 등을 조사하고자 할 때 유용합니다.

첫 번째 실행을 통해 이 두 가지를 실제로 살펴보겠습니다. (여기 전체 01_matmul_add.py 스크립트가 있습니다)

이 스크립트는 GPU가 있는 머신에서 실행하는 것을 권장합니다.

uv run 01_matmul_add.py --size 64

위 스크립트를 (GPU 머신에서) 실행하면, 다음 두 가지 결과물이 포함된 traces/01_matmul_add 폴더를 찾을 수 있습니다:

64_bf16_cold_eager.json
64_bf16_cold_eager.txt

.txt 파일은 프로파일러 테이블을 담고 있습니다. 그림 1에 표시된 것처럼 파일을 열면, 첫 번째 열이 프로파일링 범위 (scope) 내에서 트리거된 이벤트들로 구성된 커다란 테이블을 마주하게 됩니다.

다른 열들은 해당 이벤트가 CPU, GPU 또는 torch.profiler.profile 내의 activities에 지정된 기타 장치 (device)에서 소요한 시간과 관련이 있습니다. 어떤 이벤트가 가장 많은 시간을 차지하는지 살펴보고, 실제로 해당 이벤트가 그만큼의 시간이 걸리는 것이 타당한지 직관적으로 이해해 보십시오. 또한 이벤트가 몇 번 트리거되었는지를 나타내는 "# of Calls" 열을 확인하는 것도 중요합니다.

이왕 하는 김에, "Self CPU/CUDA"와 "CPU/CUDA total"의 차이에 대해서도 이야기해 보겠습니다. "Self" 열은 자식 이벤트(children)를 제외하고 해당 이벤트 자체 내부에서만 소비된 시간을 측정합니다. 반면 "total" 열은 해당 이벤트와 그에 속한 모든 자식 이벤트를 합산한 시간을 포함합니다. 따라서 만약 matmul_add의 "CPU total"을 확인한다면, [IMG:1] 이는 self 시간과 해당 이벤트가 트리거한 자식 이벤트들의 시간을 합친 것으로 구성됩니다. 이는 유의해야 할 중요한 뉘앙스입니다.

표의 마지막 두 줄을 살펴보면 프로파일러가 다음과 같이 알려주는 것을 확인할 수 있습니다.

Self CPU time total: 2.314ms
Self CUDA time total: 23.104us

CPU 시간은 ms 단위인 반면, GPU 시간은 us 단위입니다. [IMG:2] 상황을 객관적으로 비교해 보자면, GPU에서 소비된 시간(커널 ampere_bf16_s16816gemm...)은 CPU에서 소비된 시간(matmul_add 연산)의 1% 미만입니다. GPU가 대부분의 시간 동안 유휴(idle) 상태로 머물러 있는데, 이는 즉각적인 위험 신호(red flag)입니다. 이런 현상이 발생하는 이유는 GPU가 작은 행렬 곱셈(matmul)을 매우 빠르게 계산할 수 있기 때문입니다. 따라서 우리의 코드는 커널을 준비하고, 이를 GPU에 실행(launch)하며, 곱셈을 위한 데이터를 전송하고, 결과를 수집하는 데 대부분의 시간을 소비하게 됩니다. 이 개념을 오버헤드 제한 (overhead-bound) 알고리즘이라고 합니다.

이러한 영역에서 벗어나는 가장 쉬운 방법은 더 큰 행렬 곱셈을 사용하는 것입니다.

uv run 01_matmul_add.py --size 4096

그림 2의 마지막 두 줄은 다음과 같습니다.

Self CPU time total: 4.908ms
Self CUDA time total: 4.495ms

두 시간 모두 ms 단위이며, 이는 행렬 곱셈의 크기를 키우는 것만으로도 더 많은 GPU 시간을 확보했음을 의미합니다. 그림 2를 보면 이제 가장 많은 CUDA 시간이 이를 실행한 CPU 연산(matmul_add)이 아니라 GPU 커널(ampere_bf16_s16816gemm_..)에 의해 점유되고 있음을 알 수 있습니다. 이는 우리가 실제로 오버헤드 제한(overhead bound) 상태에서 연산 제한(compute bound) 상태로 이동했음을 의미합니다.

이제 .json 파일 내부에 존재하는 디스패치 체인(dispatch chain)을 시각화하는 단계로 넘어가겠습니다.

아티팩트(artifacts)입니다. 이를 Perfetto UI에 업로드하여 트레이스(traces)를 확인하거나, uvx trace-util traces -b traces를 사용하여 Perfetto 링크를 직접 생성할 수 있습니다.

그림 3에서는 행렬 곱셈(matrix multiplication)과 덧셈(addition)에 대한 프로파일러 트레이스(profiler trace)를 볼 수 있습니다. 여기서 막대의 너비는 이벤트(event)의 지속 시간을 나타내며, 수직적인 중첩은 호출 계층 구조(call hierarchy)를 나타냅니다. CPU 레인(lane)은 CPU에서 발생하는 이벤트를 나타내고, GPU 레인은 실제 커널 실행(kernel executions)을 보여줍니다. 또한 대기 시간(waiting) 또는 유휴 시간(idle time)을 나타내는 빈 공간도 확인할 수 있습니다.

스크립트는 다음과 같은 기본 설정(default configurations)으로 실행되었습니다:

• size 64: 입력(inputs), 가중치(weights), 편향(biases)의 크기가 (64, 64)입니다.
• dtype bf16: 데이터 타입(data type)은 bfloat16입니다.
• no compile: torch 연산을 컴파일(compiled)하지 않았습니다.
• no warmup: 프로파일링(profiling) 전에 GPU를 웜업(warmed up)하지 않았습니다.

Perfetto를 사용할 때는 트레이스에 더 빠르게 접근하기 위해 키보드 사용을 권장합니다.

그림 6에 표시된 약 228 µs는 record_function("matmul_add")에 진입한 시점과

PyTorch가 실제로 aten::matmul을 디스패치 (dispatch) 하는 시점 사이의 "데드 윈도우 (dead window)"입니다.

이는 워크스페이스 할당 (workspace allocations), cuBLAS (NVIDIA의 기본 선형 대수 연산을 수행하기 위한 독점적 GPU 가속 라이브러리) 휴리스틱 (heuristics), 또는 지연 모듈 로딩 (lazy module loading)을 포함한 여러 가지 이유로 발생할 수 있습니다. 우리는 이를 무시하거나, 프로파일링 (profiling)을 하기 전에 몇 번 더 웜업 (warmup) 단계를 실행할 수 있습니다 (이것이 표준적인 방법입니다).

프로파일링 관점에서 웜업 (warmup)이란 실제로 프로파일링을 수행하기 전에 이벤트를 몇 번 실행하는 것을 의미합니다. GPU에 의해 수행되는 사전 작업(위에서 언급한 사항들 포함)은 일회성 작업이며, 우리는 이를 프로파일링하고 싶지 않습니다. 우리의 예제에는 두 가지 웜업 단계가 있습니다. 하나는 프로파일러에 진입하기 전에 실제로 함수를 루프 (loop) 돌리는 단계이고, 다른 하나는 warmup 인자를 통해 프로파일러 내부에서 수행되는 단계입니다. 이 섹션에서는 스케줄 (schedule)과 함께 실제 반복 (iterations)을 활성화했습니다.

uv run 01_matmul_add.py --warmup

그림 7을 보면 각 프로파일 단계가 유사한 시간을 소요하는 것을 볼 수 있지만, 이것이 일회성 오버헤드 (overheads)를 최적화할 수 있었다는 의미는 아닙니다. 우리는 오버헤드가 프로파일링되지 않도록 실행을 웜업 (warmup)한 것입니다. 이 섹션을 해결책에 대한 힌트 없이 갑자기 끝내는 것은 독자에게 불공평하다고 생각되므로, 런치 오버헤드 (launch overheads)를 추가로 최적화하는 방법에 대해 읽어볼 수 있는 링크를 제공합니다.

그림 8에서는 CPU와 GPU 레인 (lanes) 사이에 약 2.5 ms의 오프셋 (offset)이 있는 것을 볼 수 있습니다. 이는 CPU가 CUDA 커널 (kernels)을 제출한 후 실제로 실행을 시작할 때까지의 지연 시간입니다. 웜업 (warmup) 단계와 스케줄의 wait 및 warmup을 결합하면 GPU를 계속 바쁘게 유지하여 오프셋을 줄일 수 있을 것이라고 생각할 수도 있습니다.

실제로 어떤 일이 일어나고 있는지 밝혀내기 위해, 스케줄을 약간 변경해 보겠습니다.

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=3, repeat=1)

그림 9는 Activity Buffer Request가 있음을 보여줍니다.

어떠한 연산이 수행되기 전 GPU 레인(lane)에 존재합니다. 조금 더 자세히 살펴보겠습니다.

GPU 트레이스(trace)를 확대해 보면, ProfileStep#0에 대한 행렬 곱셈(matmul) 및 덧셈(add) 커널(kernel)이

(그림에서는 CPU 트레이스가 보이지 않음) 연속해서 발생하는 반면, ProfileStep#1에 대한 커널들 사이에는

간격(window)이 있음을 알 수 있습니다. 이에 대한 가장 적절한 설명은 버퍼 오버플로(overflow)가 발생하여, 커널 실행 중에 또 다른 버퍼 요청(GPU VRAM에 메모리를 할당하라는 요청)이 발행되었다는 것입니다.

다른 가능성을 배제하는 가장 좋은 방법은 더 많은 반복 횟수로 프로파일링을 수행하여 트레이스의 다른 부분에서도 유사한 간격이 나타나는지 확인하는 것입니다. 이를 위해 active=20으로 실행합니다.

그림 11에서 볼 수 있듯이, ProfileStep#1에서 유사한 경향이 나타납니다.