NVIDIA H100 및 B200에서 PyTorch 컴파일을 통한 정규화 커널 성능 최적화

섹션별 상세

기존 torch.compile은 H100 및 B200 하드웨어의 높은 메모리 대역폭을 충분히 활용하지 못하는 기본 자동 튜닝 설정을 가지고 있었습니다. 이를 해결하기 위해 Inner reduction의 RBLOCK 크기를 확장하고, 영구적 리덕션 시 XBLOCK 크기를 조정하여 메모리 대역폭 포화도를 높였습니다. 이러한 휴리스틱 개선을 통해 순전파 성능이 Quack 라이브러리와 동등한 수준으로 향상되었습니다.

정규화의 역전파는 입력에 대한 그래디언트(dX)와 가중치에 대한 그래디언트(dW)를 모두 계산해야 하므로 복잡한 리덕션 구조를 가집니다. 기존에는 이를 별도의 커널로 처리하여 동일한 데이터를 두 번 읽는 비효율이 발생했으나, 두 리덕션을 단일 커널로 통합하는 MixOrderReduction을 도입했습니다. 이 기법은 공유 메모리를 활용해 데이터 재사용성을 높이고 메모리 읽기 비용을 절반으로 줄입니다.

메모리 대역폭 한계에 도달하기 위해 GEMM과 같은 연산 집약적 작업에서 주로 쓰이던 소프트웨어 파이프라이닝(Prefetching)을 리덕션 커널에 적용했습니다. Inductor의 자동 튜닝 파라미터에 num_stages를 추가하여 데이터 로드와 연산을 병렬화함으로써 대규모 배치 크기에서 성능을 추가로 20% 개선했습니다. 이는 특히 Blackwell 아키텍처처럼 메모리 대역폭이 매우 높은 환경에서 중요한 역할을 합니다.

벤치마크 결과 B200 GPU에서 MixOrderReduction을 적용한 torch.compile은 Eager 모드 대비 17.07배, 기존 컴파일 방식 대비 약 1.7배 빠른 성능을 보였습니다. 또한 수동 최적화된 Liger Kernel보다 1.45배, Quack보다 1.34배 빠른 속도를 기록하며 SOTA 성능을 입증했습니다. 개별 커널의 최적화뿐만 아니라 주변 연산과의 자동 융합이 가능하다는 점이 컴파일러 기반 접근법의 핵심 장점입니다.

근거

MixOrderReduction을 적용한 torch.compile은 Eager 모드보다 17.07배 빠르며, 이전 컴파일 방식보다 1.7배 이상 향상되었습니다. — Backwards 벤치마크 결과 섹션 및 B200 벤치마크 차트
B200 환경에서 torch.compile이 Liger보다 1.45배, Quack보다 1.34배 빠른 성능을 기록했습니다. — Benchmark Results 섹션의 수치 비교

기술

PyTorch
Inductor
Triton
Quack
Liger Kernel
NVIDIA H100
NVIDIA B200

활용 사례

LLM 학습 가속화
정규화 레이어 최적화
컴파일러 기반 커널 자동 생성

언급된 리소스

GitHubQuack GitHub Repository

논문LayerNorm Paper

논문RMSNorm Paper