AI 시스템 성능 엔지니어링 (Korean Edition)

제14장. PyTorch 컴파일러, OpenAI Triton 및 XLA 백엔드

이 작품은 AI를 사용하여 번역되었습니다. 여러분의 피드백과 의견을 환영합니다: [e-mail]

제13장에서는 PyTorch 기반 훈련 및 추론 워크로드를 최적화하고 튜닝하는 여러 방법을 논의했습니다. PyTorch 컴파일러와 코드 변경을 최소화하면서도 성능을 향상시키기 위해 커널 융합 및 기타 커널 수준 기법을 자동화하는 방식에 대해 간략히 살펴보았습니다.

이 장에서는 TorchDynamo, 사전 컴파일 자동 미분(AOT Autograd), PrimTorch 중간 표현(IR)(일명 Prims 또는 Prims IR)과 같은 구성 요소와 TorchInductor, 가속 선형 대수(XLA), OpenAI의 Triton 생태계와 같은 컴파일러 백엔드를 포함하여 동적 PyTorch 컴파일 스택을 더 깊이 파고듭니다. PyTorch 컴파일러 스택은 그림 14-1에 표시되어 있습니다.

또한 컴파일 파이프라인 디버깅 도구와 다중 GPU 및 다중 노드 클러스터에서 PyTorch를 확장하는 라이브러리도 다룹니다. 이후 동적 형상( torch.compile )이 내부적으로 어떻게 작동하는지, 그리고 동적 형상과 가변 시퀀스 길이를 효율적으로 처리하는 방법을 살펴보겠습니다.

또한 PyTorch 컴파일러와 OpenAI Triton 생태계의 통합도 살펴볼 것입니다. 우리의 목표는 PyTorch의 유연한 즉각 실행 개발 환경을 희생하지 않으면서 PyTorch 모델과 애플리케이션의 가속화와 확장성을 확보하는 것입니다.

그림 14-1. PyTorch 컴파일러 스택개요

PyTorch 컴파일러 심층 분석

13장에서 설명한 바와 같이, PyTorch의 자동 사후 컴파일러( torch.compile )는 PyTorch 코드(및 모델)를 컴파일하여 상당한 속도 향상을 제공합니다( ). 대부분의 경우, 다음 예시처럼 단 한 줄의 코드로 이를 수행할 수 있습니다. 다양한 옵션에 대해서는 진행하면서 설명하겠습니다:

compiled_model = torch.compile(model, 
  mode="max-autotune",
#  ...
) 

이 섹션에서는 TorchDynamo의 그래프 캡처, AOT Autograd의 전진/후진 그래프 통합 최적화, PrimTorch IR, TorchInductor의 코드 생성을 포함한 의 PyTorch 컴파일 파이프라인 단계를 상세히 설명합니다. 이 파이프라인은 대상 GPU 하드웨어용 최적화된 커널을 생성하며, 그림 14-2에 표시되어 있습니다.

그림 14-2. PyTorch ...