이번 글에서는 PyTorch 프로파일링 2부에서 nn.Linear를 사용하여 행렬 곱셈과 덧셈을 분석하고, 이를 세 개 연결하여 MLP 블록을 구성하는 과정을 살펴봅니다.
nn.Linear는 가중치와 편향을 파라미터로 관리하며, 사용자가 익숙한 forward 메서드를 제공합니다. 프로파일링 결과, 가중치 전치 연산(aten::t)이 수행되고, 편향 덧셈은 GEMM 커널의 에필로그로 통합되어 메모리 접근 횟수를 줄이는 것을 확인할 수 있습니다.
또한 torch.compile을 사용하여 단일 Linear 연산의 성능을 개선하고, 튜닝된 커널 라이브러리를 활용하여 성능을 더욱 향상시키는 방법을 알아봅니다. 이 글은 PyTorch 내부 동작에 대한 이해를 높이고, 모델 최적화에 도움이 될 것입니다.
nn.Linear는 입력(x), 가중치(w), 편향(b)을 사용하여 y = x @ w.T + b 연산을 수행하며, 가중치 전치 연산(aten::t)은 텐서 메타데이터를 수정하는 작업입니다.
torch.compile은 커널 실행을 스케줄링하는 CPU의 오버헤드를 줄여 성능을 향상시키며, 퓨즈드 Triton 커널을 통해 더욱 효율적인 연산이 가능합니다.
튜닝된 커널 라이브러리를 사용하면 메모리 접근 횟수를 줄이고, 에필로그를 활용하여 성능을 최적화할 수 있습니다.