PyTorch 프로파일링 (Part 2): nn.Linear에서 Fused MLP까지

본 시리즈의 첫 번째 파트인 "PyTorch 프로파일링 (Profiling in PyTorch)"에서는 PyTorch 프로파일러 (profiler) 트레이스 (trace)를 읽는 법을 배우기 위해 torch.add(torch.matmul(x, w), b)를 사용했습니다.

또한 우리는 CPU 디스패치 체인 (CPU dispatch chain), 런치 오버헤드 (launch overhead), 오버헤드 제한 (overhead-bound) 영역과 연산 제한 (compute-bound) 영역의 차이, 그리고 torch.compile의 내부 동작 등 우리가 마주했던 여러 다른 주제들에 대해서도 논의했습니다.

두 번째 단계(본 블로그 포스트)에서는 한 단계 더 올라갑니다. 수동으로 작성한 matmul-add 쌍을 nn.Linear (bias=True 포함)로 교체합니다. 이것은 모든 딥러닝 모델이 사용하는 빌딩 블록 (building block)입니다. 그런 다음 그 사이에 활성화 함수 (activation)를 넣어 세 개의 레이어를 쌓아(본 예제에 특화됨) 다층 퍼셉트론 (MLP, Multilayer Perceptron) 블록을 형성합니다.

본 블로그 포스트를 위한 스크립트는 여기에 있습니다:

02_linear.py, 03_simple_mlp.py, 그리고 03_kernels_mlp.py.

이전과 마찬가지로, 스크립트를 별도의 탭에서 열고 코드를 읽으면서 따라가는 것이 도움이 됩니다. 우리는 스크립트를 실행하기 위해 NVIDIA A100-SXM4-80GB GPU를 사용합니다. Hugging Face 인프라에서 GPU를 설정하고 Spaces의 Dev Mode를 사용하여 스크립트를 실험하는 것은 매우 쉽습니다. 또한 Hugging Face Jobs 파이프라인을 통해 스크립트를 실행할 수도 있습니다.

시작하기 전에, 우리가 반복해서 의존하게 될 두 가지 아이디어를 빠르게 요약해 보겠습니다:

• GPU **커널 (kernel)**은 GPU의 많은 스레드 (threads)에서 병렬로 실행되는 프로그램입니다.
• CPU는 이러한 커널을 **스케줄링하고 실행 (schedules and launches)**합니다. 프로파일러 트레이스에서 보이는 PyTorch 오버헤드의 대부분은 바로 이 스케줄링 작업입니다.

nn.Linear는 우리가 파트 1에서 이미 프로파일링했던 동일한 행렬 곱셈 (matrix multiplication) 및 덧셈을 감싸고 있는 모듈 래퍼 (module wrapper)입니다. 유일한 차이점은 가중치 (weight)와 편향 (bias)을 파라미터 (parameters)로 소유하며, PyTorch 사용자들이 익숙해진 forward 메서드를 노출한다는 점입니다.

# bias=True는 시리즈의 파트 1에서 보았던
# 곱셈 및 덧셈 연산을 실제로 에뮬레이션합니다
linear_layer = nn.Linear(in_dim, out_dim, bias=True)
...

현재 다루고 있는 연산은 다음과 같이 작성될 수 있습니다:

y = x @ w.T + b

여기서 x

는 입력(input)이고, w는 가중치(weight), b는 편향(bias)입니다. 02_linear.py를 실행하여 프로파일(profile)을 확인해 보겠습니다.

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64
uvx trace-util traces -b traces

trace-util은 트레이스(trace)를 Hugging Face 버킷으로 동기화한 후 터미널에 Preffeto URL을 제공하는 유틸리티입니다.

그림 1은 선형 레이어(linear layer)의 포워드(forward) 호출에 대한 프로파일러 트레이스를 보여줍니다. 이전 트레이스와 유사하게 wait=1, warmup=1, active=3과 같은 schedule 설정을 사용하여 선형 레이어의 forward 호출을 트레이스했습니다. CPU와 GPU 레인(lane)에서 세 개의 프로파일 스텝(Profile Steps)이 보이는 이유가 바로 이것입니다.

그림 2에서처럼 프로파일러 트레이스를 확대해 보면, aten::addmm (행렬 곱셈 및 덧셈) 연산 이전에 aten::t (전치, transpose) 연산이 있음을 알 수 있습니다. 이를 통해 nn.Linear가 가중치 파라미터를 전치한 다음 입력과 곱한다는 것을 이미 파악할 수 있습니다. 이것이 aten::t 연산이 보이는 이유입니다.

주의 깊게 봐야 할 중요한 점은 aten::t가 실제로 데이터를 복사하거나 재구성하지 않는다는 것입니다. 이는 단지 전치된 행렬을 나타내기 위해 CPU에서 텐서 메타데이터(shape 및 stride)를 다시 작성할 뿐입니다. GPU에서는 커널(kernel)을 실행하지 않습니다. 이는 두 가지 방법으로 확인할 수 있습니다. 트레이스의 GPU 레인을 살펴보거나, 프로파일러 테이블의 aten::t 행과 CUDA에서 소요된 시간을 확인하는 것입니다.

그림 3에서 볼 수 있듯이, 선형 레이어의 디스패치 체인(dispatch chain)에는 aten::add (편향 덧셈)가 없습니다. 이는 편향 덧셈이 **에필로그 (epilogue)**라고 불리는 기법을 사용하여 행렬 곱셈 커널에 폴딩 (folded) 되었기 때문입니다.

**에필로그 (epilogue)**는 GEMM (General Matrix Multiply) 커널이 결과를 HBM (High Bandwidth Memory, GPU의 메인 메모리)에 다시 쓰기 직전, 맨 마지막에 수행하는 작은 연산입니다. 편향을 더하거나, 활성화 함수(activation)를 적용하거나, 상수를 곱해 스케일링하는 것 모두 전형적인 에필로그입니다. 에필로그를 사용하는 목적은 메모리 트래픽이 연산 비용을 높이기 때문에, HBM에 두 번 로드하거나 쓰는 것을 방지하기 위함입니다.

nn.Linear는 torch.nn.functional.linear를 호출하며, 이는 다시 aten::linear를 호출합니다. aten::linear는 입력을 확인하여 bias(편향)가 전달되었음을 인지하고, 행렬 곱셈 (matmul)과 덧셈 (add)을 별도로 수행하는 대신 aten::addmm(bias, x, weight)를 디스패치 (dispatch)합니다. addmm은 다음과 같이 계산합니다:

out = x @ weight.T + bias

GPU에서 실행되는 cuBLAS GEMM 커널에는 bias-add 변형이 내장되어 있으며, aten::addmm은 바로 그 커널을 선택합니다. 덧셈은 별도의 커널로 나타나지 않는데, 이는 덧셈이 matmul 커널의 writeback (쓰기 작업)의 일부이기 때문이며, 이것이 바로 에필로그 (epilogue)가 하는 역할입니다.

여기서 미묘한 점 하나를 주목할 필요가 있습니다. Part 1의 --compile 옵션 아래에서 보았던 커널(addmm)은 eager 모드의 nn.Linear가 이미 사용하고 있는 커널입니다. 따라서 torch.compile이 여기서 더 융합 (fuse)할 수 있는 것은 남아있지 않으며, 이는 다음 단계에서 검증할 내용입니다.

forward 호출을 컴파일하고 프로파일러 트레이스 (profiler trace)를 살펴보겠습니다. (프로파일러 트레이스는 다음 섹션에서 시각화됩니다.)

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile
uvx trace-util traces -b traces

단일 nn.Linear의 forward에 대해 eager 트레이스와 컴파일된 트레이스를 비교해 보면 다음과 같은 사실을 알 수 있습니다:

• GPU 상에서 동일한 cuBLAS GEMM 커널이 실행됩니다.
• CPU 상에서 동일한 aten::addmm 연산 (op)이 실행됩니다.
• CPU 레인 (lane)에 compile 특유의 몇몇 추가 행이 나타납니다.

이 점을 내재화할 가치가 있습니다. 모델이 느리다고 느껴질 때마다 습관적으로 torch.compile을 사용하는 경우가 많습니다. 하지만 bias가 포함된 단일 GEMM의 경우, compile이 할 수 있는 일은 거의 없습니다. 이는 버그가 아니라, compile이 융합 (fusing)을 수행하려면 하나 이상의 연산이 필요하기 때문입니다. MLP를 살펴보며 이를 증명해 보겠습니다.

두 트레이스 (eager vs compile)를 주의 깊게 읽는 독자라면, eager 모드의 CPU 디스패치 체인 (dispatch chain)이 컴파일된 것보다 더 많은 내용을 포함하고 있음을 눈치챌 것입니다.

Figure 4: aten::linear가 aten::t (transpose)를 거쳐 aten::addmm으로 이어지는 eager 디스패치 체인

aten::linear 내부의 eager CPU 디스패치 체인은 aten::t 다음에 aten::addmm이 이어지는 구조입니다.

(Figure 4). aten::t가 실제로 무엇을 하는지 이해하려면, *스트라이드 (strides)*와 *뷰 (views)*에 대해 잠시 살펴볼 필요가 있습니다.

텐서는 메모리 내에서 하나의 평탄하고 연속적인(contiguous) 숫자 배열로 데이터를 저장합니다. shape (형태)와 stride (스트라이드)는 해당 배열 위에 놓인 메타데이터로, PyTorch가 데이터를 어떻게 탐색할지 알려줍니다. 예를 들어 (s0, s1) 스트라이드는 "한 행을 이동하려면 s0개의 요소를 건너뛰고, 한 열을 이동하려면 s1개의 요소를 건너뛴다"는 의미입니다. 메타데이터를 변경하면 데이터의 복사 없이도 동일한 원시 데이터에 대한 다른 *뷰 (view)*를 얻을 수 있습니다:

>>> M = torch.tensor([[0, 1],
... [2, 3],
... [4, 5]])
...

M.t()는 단 하나의 숫자도 이동시키지 않았습니다. 대신 스트라이드가 서로 바뀐 새로운 뷰를 반환했을 뿐입니다. 따라서 이제 행 단위로 읽으면 원래의 버퍼인 0, 1, 2, 3, 4, 5를 전치(transposed)된 순서로 읽게 됩니다. 근본적인 데이터는 동일하며, 오직 메타데이터만 다릅니다.

이것이 바로 선형 레이어(linear layer) 내부에서 aten::t가 수행하는 작업입니다. 새로운 텐서를 할당하거나 데이터를 복사하지 않고, 스트라이드가 재작성된 가중치(weight)의 *뷰 (view)*를 생성합니다.

Figure 5에서 볼 수 있듯이, compile은 GPU 커널을 제거한 것이 아니라, 해당 뷰를 디스패치(dispatching)하는 과정에서 발생하는 *CPU 오버헤드 (CPU overhead)*를 제거한 것입니다. Inductor는 컴파일 타임에 뷰 체인을 추적하여 결과적인 스트라이드를 한 번 계산하고, 해당 스트라이드가 하드코딩된 직접적인 aten::addmm 호출을 생성했습니다. GPU는 동일한 연산을 수행하는 동안, 몇 마이크로초(microseconds)의 CPU 작업이 사라지게 됩니다.

예상할 수 있듯이, 입력 데이터가 컴파일러에 의해 미리 계산된 스트라이드를 위반할 경우 에러가 발생할 것입니다.

두 트레이스(trace)의 GPU 레인을 살펴보면, 포워드(forward)당 정확히 하나의 커널이 존재하며, 두 번 모두 동일한 커널입니다:

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8

만약 전치(transpose) 커널이 실행되지 않았다면, 누가 GEMM에게 가중치 행렬을 전치된 순서로 읽도록 가르쳤을까요? 답은 커널의 이름에 있습니다. 접미사(suffix)를 확인해 보세요:

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8
^^

저 tn

is 레이아웃 기술자 (layout descriptor)입니다. cuBLAS와 CUTLASS는 입력 레이아웃의 각 조합에 대해 *별도의 커널 바이너리 (separate kernel binary)*를 사전 컴파일합니다.

n (non-transposed, 비전치)과 t (transposed, 전치)는 커널이 내부 루프 (inner loop) 동안 입력을 어떻게 순회하는지를 설명합니다. 디스패처 (dispatcher)의 역할은 입력 스트라이드 (strides)를 확인하여 어떤 접미사 (suffix) 조합이 일치하는지 결정하고, 적절한 사전 컴파일된 커널을 선택하는 것입니다.

프로파일러 트레이스 (profiler trace)에서의 커널 이름은 커널의 정체성을 해시 덤프 (hash dump)한 것입니다. 만약 두 번의 실행에서 동일한 커널 이름이 나타난다면, GPU는 동일한 작업을 수행하고 있는 것입니다. 만약 이름이 다르다면 (예: _tn_ vs _nn_, bf16 vs fp16, 또는 s16816gemm vs s161616gemm), GPU는 서로 다른 작업을 수행하고 있으며 디스패처가 다른 분기 (branch)를 선택한 것입니다. 이 이름을 읽는 법을 배우는 것은 트레이스를 비교할 때 가장 유용한 습관 중 하나입니다.

이 섹션에서는 다층 퍼셉트론 (MLP, Multilayer Perceptron)을 프로파일링할 것입니다. 이를 더 흥미롭게 만들기 위해, (실무에서 매우 빈번하게 사용되는) GeGLU 활성화 함수 변형을 사용하는 피드포워드 네트워크 (feed-forward network)를 프로파일링하겠습니다. 이는 딥러닝 연구 역사상 가장 위대한 코드 라인 중 하나에 경의를 표하는 방식이기도 합니다 (그림 6).

class SimpleGeGLUMLP(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        ...

전체 스크립트는 여기에서 찾을 수 있습니다: 03_simple_mlp.py.

다음과 같이 실행하세요:

uv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072
uvx trace-util traces -b traces

트레이스를 열기 전에, 우리가 무엇을 보게 될지 함께 생각해 봅시다. forward 함수는 상당한 양의 연산을 수행하지만, 그 대부분은 우리에게 이미 익숙한 것입니다.

우리는 각 nn.Linear 레이어에 대해 하나씩, 총 세 번의 aten::linear 디스패치 (dispatch)가 발생할 것을 예상해야 합니다. 또한 GeLU를 위한 하나와 곱셈을 위한 하나, 총 두 번의 포인트와이즈 커널 실행 (pointwise kernel launch)도 예상해야 합니다. 확인하기 전에 이러한 예상을 세우는 것은 프로파일링 여정에서 가장 유용한 단 하나의 습관입니다. 즉, 트레이스를 읽는 목적은 처음부터 추측을 만들기 위함이 아니라, 추측을 확인하거나 깨뜨리기 위함입니다.

그림 7을 통해 우리의 직관이 옳았음을 확인하며 스스로를 칭찬할 수 있습니다. 순전파(forward pass) 한 번(mlp_fwd [IMG:1])당 GPU는 정확히 5개의 커널(kernel)을 실행합니다. 그림 8은 선형 투영(linear projection) 레이어의 CPU 레인에서 볼 수 있는 "점유율 쿼리(occupancy query)"를 강조하여 보여줍니다.

연산 (Op)	CPU 연산 (CPU op)	GPU 커널 (GPU kernel)	실행 (launches)
gate_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x128_...	occupancy query + cudaLaunchKernel
up_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x128_...	occupancy query + cudaLaunchKernel
gelu	aten::gelu	vectorized_elementwise_kernel<4, GeluCUDAKernelImpl...>	cudaLaunchKernel
h * u	aten::mul	vectorized_elementwise_kernel<4, ...MulFunctor...>	cudaLaunchKernel
down_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x256_...	occupancy query + cudaLaunchKernel

세 개의 GEMM(General Matrix Multiply)은 각각 실행 전에 추가적인 cudaOccupancyMaxActiveBlocksPerMultiprocessor [IMG:2] 호출을 수행합니다. 이에 대해서는 Part 1에서 별도의 섹션으로 다루었으며, [여기]에서 확인하실 수 있습니다. 이는 cuBLAS가 그리드(grid)의 크기를 결정하는 과정입니다. 포인트와이즈 연산(pointwise ops, GeLU 및 mul)은 점유율 쿼리 없이 직접 실행됩니다. 따라서 "선형 연산(a linear)"은 실제로는 쿼리(query) + 실행(launch)인 반면, "포인트와이즈 연산(a pointwise op)"은 단순히 실행(launch)뿐입니다.

aten::t [IMG:3], aten::transpose [IMG:4], aten::reshape [IMG:5], aten::view [IMG:6], aten::as_strided [IMG:7], 그리고 aten::_unsafe_view [IMG:8] 연산들은 커널을 전혀 실행하지 않습니다. 이들은 표(그림 9)에서 CUDA 시간이 0.000us [IMG:9]로 표시되는데, 그 이유는 이 연산들이 CPU에서 텐서 메타데이터(shape 및 stride)만 재작성하기 때문입니다. 표를 훑어보는 독자는 선형 연산당 약 6개의 연산 이름을 보게 되지만, 그중 오직 하나(mm [IMG:10])만이 실제로 GPU에 도달합니다.

MLP는 행렬 곱셈(matmul)을 위해 [batch, seq, dim] [IMG:11]을 [batch * seq, dim] [IMG:12]로 평탄화(flatten)합니다. 명령줄 호출에서 batch [IMG:13]를 64로, seq [IMG:14]를 128로 설정했으므로, 아래의 8192 (batch * seq = 64 * 128 [IMG:15])라는 수치가 도출됩니다.

트레이스(trace) 결과: [IMG:16]

Linear	aten::mm input dims	M·K·N	cuBLAS kernel	avg CUDA
gate_proj	[8192,768] x [768,3072]	8192·768·3072	…128x128…stages_32x5_tn	0.19ms
up_proj	[8192,768] x [768,3072]	8192·768·3072	…128x128…stages_32x5_tn	0.19ms
down_proj	[8192,3072] x [3072,768]	8192·3072·768	…128x256…stages_64x3_tn	0.17ms

세 GEMM(General Matrix Multiplication) 모두 동일한 FLOP 카운트인 2·8192·768·3072 ≈ 38.7 GFLOP를 가지고 있지만, down_proj가 약 10% 더 빠릅니다. 같은 작업이지만 다른 형태(N=768 대신 3072)이기 때문에 cuBLAS는 해당 형태에 대해 더 나은 재사용성을 갖는 다른 타일(128×256, 더 깊은 stages_64x3 파이프라인)을 선택합니다.