PyTorch 프로파일링 (2부): nn.Linear에서 통합 MLP까지
Source: Hugging Face Blog
첫 번째 파트 “Profiling in PyTorch”에서는 torch.add(torch.matmul(x, w), b)를 사용해 PyTorch 프로파일러 트레이스를 읽는 방법을 배웠습니다. 또한 CPU 디스패치 체인, 런치 오버헤드, 오버헤드에 제한된 영역과 연산에 제한된 영역의 차이, torch.compile의 내부 동작 등 여러 주제를 다뤘습니다.
두 번째 iteration(이번 블로그 포스트)에서는 한 단계 위로 올라갑니다. 손으로 구현한 matmul‑add 쌍을 bias=True인 nn.Linear 로 교체합니다. 이는 모든 딥러닝 모델이 사용하는 기본 블록입니다. 그런 다음 예제에 맞게 세 개를 쌓고, 그 사이에 활성화 함수를 넣어 다층 퍼셉트론(MLP) 블록을 구성합니다.
이번 포스트에 사용된 스크립트는 여기에서 확인할 수 있습니다: 02_linear.py, 03_simple_mlp.py, 그리고 03_kernels_mlp.py. 이전과 마찬가지로 별도 탭에서 코드를 열어가며 읽는 것이 도움이 됩니다. 우리는 NVIDIA A100‑SXM4‑80GB GPU에서 스크립트를 실행합니다. Hugging Face 인프라에서 GPU를 손쉽게 설정하고 Dev Mode with Spaces로 실험해볼 수 있습니다. 또한 Hugging Face Jobs 파이프라인으로 스크립트를 실행할 수도 있습니다.
시작하기 전에, 앞으로 반복해서 사용할 두 가지 개념을 간단히 정리합니다:
- GPU 커널은 GPU의 수많은 스레드에서 병렬로 실행되는 프로그램입니다.
- CPU는 이러한 커널을 스케줄링하고 런치합니다. 프로파일러 트레이스에서 보이는 대부분의 PyTorch 오버헤드는 이 스케줄링 작업입니다.
From matmul-add to Linear
nn.Linear는 파트 1에서 이미 프로파일링한 행렬 곱셈과 덧셈을 감싸는 모듈 래퍼입니다. 차이점은 가중치와 편향을 파라미터로 소유하고, PyTorch 사용자에게 친숙한 forward 메서드를 제공한다는 점입니다.
# bias=True 로 하면 파트 1에서 본 곱셈·덧셈 연산을 정확히 흉내낼 수 있습니다.
linear_layer = nn.Linear(in_dim, out_dim, bias=True)
y = linear_layer(x)우리가 다루는 연산은 다음과 같이 표현됩니다:
y = x @ w.T + b여기서 x는 입력, w는 가중치, b는 편향입니다. 이제 02_linear.py를 실행하고 프로파일을 확인해 보겠습니다.
uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64
uvx trace-util traces -b tracestrace-util은 트레이스를 Hugging Face 버킷(Hugging Face storage)에 동기화하고, 터미널에 Perfetto URL을 출력해 주는 유틸리티입니다.
Figure 1. nn.Linear의 프로파일러 트레이스
Figure 1은 Linear 레이어의 forward 호출에 대한 프로파일 트레이스를 보여줍니다. 우리는 이전 트레이스와 동일한 schedule 설정(wait=1, warmup=1, active=3)으로 Linear 레이어의 forward를 트레이스했습니다. 그래서 CPU와 GPU 레인에 각각 세 개의 Profile Step이 나타납니다.
What is the transpose doing?
Figure 2. 전치 연산을 담당하는 CPU 행
Figure 2처럼 트레이스를 확대하면 aten::addmm(곱셈·덧셈) 앞에 aten::t(전치) 연산이 있는 것을 확인할 수 있습니다. 이는 nn.Linear가 가중치 파라미터를 전치한 뒤 입력과 곱한다는 뜻이며, aten::t 연산이 보이는 이유이기도 합니다.
중요한 점은 aten::t가 실제 데이터를 복사하거나 재배열하지 않는다는 것입니다. CPU에서 텐서 메타데이터(형태와 stride)만 바꿔 전치된 행렬을 표현합니다. GPU에서 커널을 런치하지도 않습니다. 이를 확인하는 방법은 두 가지입니다: 트레이스의 GPU 레인을 살펴보거나, 프로파일러 테이블에서 aten::t 행이 차지한 CUDA 시간을 확인하면 됩니다.
Why are there no separate mul and add kernels?
Figure 3. Linear 레이어 프로파일에 aten::add가 없는 이유
Figure 3에서 볼 수 있듯이 Linear 레이어의 디스패치 체인에는 aten::add(편향 추가) 연산이 없습니다. 이는 편향 추가가 에필로그(epilogue) 라는 기법으로 행렬 곱셈 커널에 접합(folded) 되었기 때문입니다.
에필로그란 GEMM(General Matrix Multiply) 커널이 결과를 HBM(High Bandwidth Memory, GPU 메인 메모리)에 기록하기 직전에 수행