NVSwitch와 TensorRT-LLM 멀티샷으로 3배 빠른 AllReduce 구현

Reading Time: 3 minutes

사용자 수가 수백 명에서 수십만 명까지 변동하고 각 요청의 입력 시퀀스 길이가 다른 프로덕션 환경에서 생성형 AI 워크로드를 배포하는 데는 여러 어려움이 있습니다. 이런 환경에서 짧은 지연 시간으로 추론하려면 GPU 세대나 메모리 용량에 관계없이 멀티 GPU 설정이 필수적입니다. 프로덕션 수준의 성능을 위해 NVIDIA NVLink Switch를 활용해 통신 속도를 최대 3배까지 높이는 새로운 멀티 GPU 통신 프로토콜, TensorRT-LLM 멀티샷을 소개합니다. 이 블로그에서는 이 기능이 기존 멀티 GPU 통신의 한계를 어떻게 해결하는지 간략히 설명합니다.

기존 AllReduce 알고리즘의 문제점

지연 시간이 짧은 추론을 위해서는 단일 GPU의 메모리 용량에 상관없이 멀티 GPU 설정이 필수적입니다. 하지만 동시성이 낮을 경우 GPU 간 데이터 교환 시간이 실제 연산 시간보다 길어질 수 있습니다. 최적의 성능을 위해서는, 각 GPU의 부분적인 결과를 결합하는 효율적인 AllReduce 연산이 중요합니다.

기존 방식에서는 링 기반 알고리즘을 사용해 부분 결과를 GPU 링에서 전달합니다. 각 GPU는 자신만의 값을 추가해 이웃 GPU에 전달하며, 이 과정이 2N-2회 반복됩니다(N은 작동하는 GPU의 수). 이 과정이 끝나면 모든 GPU가 동일한 합산 값을 얻게 됩니다. 마지막 GPU의 합산값을 다시 다른 GPU로 전파하려면 링을 두 번째로 돌아야 합니다.

이 링 방식은 통신 단계당 GPU 간 대역폭을 효율적으로 사용하지만, GPU 수가 증가하면 통신 단계도 늘어납니다. 따라서 모든 GPU가 각 단계마다 동기화되어야 하므로 지연 시간이 증가하고, 이로 인해 지연 시간 오버헤드가 커져 더 엄격한 지연 시간 목표를 달성하기 어려워질 수 있습니다.

Ring AllReduce 알고리즘은 아래에 설명되어 있습니다:

링 알고리즘: GPU-1 → GPU-2 → … → GPU-N → GPU-1 → GPU-2 → … → GPU-(N-1)
2N-2 단계, 각 단계마다 전체 텐서 송수신 포함
지연 시간: 2N-2 통신 단계. (N: GPU 개수)
트래픽: (4N-4)/N 텐서 바이트의 송신/수신

TensorRT-LLM 멀티샷으로 AllReduce 통신 문제 해결하기

TensorRT-LLM 멀티샷은 NVSwitch의 멀티캐스트 기능을 활용하여 기존 Ring AllReduce의 O(N) 지연 시간을 최대 3배까지 단축하는 새로운 알고리즘입니다. 멀티캐스트는 NVSwitch의 하드웨어 가속 기능을 통해 GPU가 데이터를 한 번 전송하면 모든 GPU에 동시에 전달할 수 있어 대역폭 효율을 유지하면서 통신 단계를 두 번의 동기화로 최소화할 수 있습니다. NVSwitch가 없을 경우 통신 대역폭은 N배로 증가하게 됩니다.

TensorRT-LLM 멀티샷은 AllReduce 연산을 ReduceScatter 연산과 AllGather 연산으로 나누어 수행합니다. (집합 연산에 관한 자세한 내용은 설명서를 참조하세요.)

각 GPU는 결과 텐서의 일부만을 누적 처리하도록 설계되어 있습니다.

첫 번째 단계(또는 “샷”)에서는 각 GPU가 텐서의 특정 조각을 담당하는 GPU로 해당 텐서 조각을 전송합니다.

로컬에서 누적된 후에는 각 GPU가 출력의 고유한 슬라이스에 대한 올바른 누적 합계를 갖추게 됩니다.

두 번째 단계(또는 “샷”)에서는 각 GPU가 NVSwitch 멀티캐스트 기능을 사용해 자신의 결과 슬라이스를 모든 GPU에 브로드캐스트합니다. NVSwitch가 데이터 증폭을 담당하므로 필요한 GPU당 대역폭이 최소화되며, 각 GPU는 전체 텐서의 1/N만 전송하고 동시에 전체 결과 텐서를 수신하게 됩니다.

이 전체 작업은 텐서 병렬 추론을 수행하는 GPU 수에 관계없이 두 번의 통신 단계만으로 완료됩니다.

TensorRT-LLM 멀티샷 알고리즘: GPU_N이 슬라이스를 전송하고, 슬라이스 합계를 계산하고, 결과를 브로드캐스트하여 단일 멀티캐스트 연산으로 처리합니다.
지연 시간: 2단계 통신(GPU 수와 무관)
트래픽: 2 텐서 바이트의 송수신(GPU 수에 관계없음)

알고리즘의 중요성

이 알고리즘은 2N-2(여기서 N은 GPU 수) 대신 단 두 번의 통신 단계만 거치므로, MultiShot은 기존 Ring AllReduce보다 최대 3배 빠른 성능을 제공합니다. 특히, 작은 메시지 크기와 높은 병렬 처리가 요구되는 상황에서 지연 시간을 최소화해 사용자 경험을 크게 개선하는 데 효과적입니다.

또한, 이 알고리즘은 지연 시간을 줄이거나 동일한 지연 시간에서 처리량을 증가시키는 데 유리합니다. 특히, 지연 시간 요구가 까다로운 환경에서는 GPU 수에 따라 초선형 성능 향상이 가능해집니다.