기술 요약: 멀티 노드 LLM 추론에서의 통신 성능 이해 및 개선

문제 제기

대규모 언어 모델 (LLM) 의 규모가 커짐에 따라 메모리 사용량이 단일 GPU 의 용량을 초과하는 경우가 많아, 멀티 노드 간 분산 추론이 필요해지고 있습니다. 텐서 병렬화 (TP) 와 파이프라인 병렬화 (PP) 와 같은 모델 병렬화 전략은 단일 노드 추론에 대해서는 잘 연구되어 왔으나, 멀티 노드 환경에서의 성능은 아직 충분히 탐구되지 않았습니다. 멀티 노드 추론은 노드 간 연결 (예: NVLink) 에 비해 노드 간 지연 시간이 훨씬 크기 때문에 상당한 도전을 안겨줍니다.

저자들은 TP 가 메모리 병목 및 디코드 (decode) 중심 워크로드에서는 종종 우수하지만, 빈번한 all-reduce 연산으로 인해 상당한 통신 오버헤드를 겪는다고 지적합니다. 반면, TP 와 PP 를 결합한 하이브리드 병렬화 (HP) 는 통신 오버헤드를 줄이지만, 작은 행렬 곱셈의 특성상 디코드 중심 환경에서는 계산 시간을 효과적으로 단축하지 못합니다. 여기서 중요한 병목 현상으로 지목된 것은 LLM 디코드 단계에서 전형적인 작은 메시지 크기 (128 KB~2 MB) 에 대한 표준 NCCL all-reduce 구현의 비최적 성능이며, 특히 노드 간 확장 시 두드러집니다.

방법론

저자들은 기존 최첨단 추론 엔진과 맞춤형 연구 프로토타입을 결합하여 체계적인 성능 연구를 수행했습니다:

1. 실험 설정:

   • 하드웨어: 실험은 Perlmutter 시스템 (NVIDIA A100 GPU, HPE Slingshot-11 인터커넥트) 과 Vista 시스템 (NVIDIA GH200 GPU, InfiniBand 인터커넥트) 에서 수행되었습니다.
   • 모델: bf16 정밀도의 Llama 3.1 70B 및 405B(Instruct) 모델.
   • 워크로드: 두 가지 주요 환경이 테스트되었습니다: "Prefill-heavy"(대규모 프롬프트 길이) 와 "Decode-heavy"(긴 생성 시퀀스).
   • 엔진: 연구는 vLLM과 SGLang을 활용했으며, 멀티 노드 HPC 환경에서 통제된 실험을 용이하게 하기 위해 저자들이 개발한 커스텀 오픈소스 추론 엔진인 YALIS도 사용되었습니다. YALIS 는 통합 모델 정의 계층, Torch Compile 기반 실행, 그리고 AxoNN 기반 텐서 병렬화 기능을 제공합니다.

2. 분석 접근법:

   • 저자들은 단일 노드에서 32 노드 구성으로 GPU 수를 증가시키며 엔드 투 엔드 배치 지연 시간을 측정하는 강력한 확장 (strong-scaling) 실험을 수행했습니다.
   • 병목 현상을 식별하기 위해 실행 시간을 계산 (Matmul, 기타 계산), 통신, 유휴 시간으로 분해했습니다.
   • 마이크로 배치 (PP) 와 텐서 분할 (TP) 이 행렬 곱셈 커널에 미치는 영향을 분리하기 위해 합성 GEMM 벤치마크를 사용했습니다.

3. 제안된 솔루션 (NVRAR):

   • 통신 병목 현상을 해결하기 위해 저자들은 NVSHMEM 을 기반으로 한 재귀적 더블링 (recursive doubling) 방식의 계층적 all-reduce 알고리즘인 NVRAR을 개발했습니다.
   • 설계: NVRAR 은 세 단계로 작동합니다:
     1. 노드 내 Reduce-Scatter: NCCL 을 사용하여 노드 내부에서 데이터를 축소합니다.
     2. 노드 간 재귀적 더블링: NVSHMEM 의 비차단 RMA 원시 연산 (put_nbi) 을 사용하여 재귀적 더블링 토폴로지를 통해 노드 간 all-reduce 를 수행합니다.
     3. 노드 내 All-Gather: NCCL 을 사용하여 축소된 데이터를 노드 내 모든 GPU 로 다시 수집합니다.
   • 최적화:
     • 청크 기반 비차단 통신: 데이터를 청크 단위로 처리하여 계산과 통신을 중첩합니다.
     • 퓨즈드 페이로드: Slingshot/InfiniBand 에서 비싼 명시적 신호 원시 연산을 피하기 위해 데이터와 동기화 플래그를 단일 8 바이트 페이로드로 결합합니다.
     • 시퀀스 번호 동기화: 전역 장벽 (fence/quiet) 의 지연 오버헤드를 피하면서도 정확성을 보장하기 위해 고유한 시퀀스 번호를 사용하여 전역 동기화를 수행합니다.

주요 기여

1. 체계적인 성능 연구: 이 논문은 멀티 노드 환경에서 TP 와 HP 의 확장 행동을 상세히 분석하여, TP 가 일반적으로 디코드 중심 워크로드에 우수하지만 통신 오버헤드에 의해 제한받고, HP 는 이러한 환경에서 계산 시간 단축에 어려움을 겪음을 규명했습니다.
2. YALIS 엔진: 멀티 노드 HPC 환경에서 계측 가능성과 통제된 실험을 위해 설계된 연구 중심 추론 엔진인 YALIS 의 개발.
3. NVRAR 알고리즘: LLM 디코드 단계의 특징인 작은 메시지 환경에 최적화된 커스텀 all-reduce 알고리즘인 NVRAR 의 제안 및 구현. 이는 Slingshot 및 InfiniBand 네트워크에서 표준 NCCL 구현보다 낮은 지연 시간을 달성하기 위해 NVSHMEM 을 활용합니다.
4. 통합 및 평가: NVRAR 을 YALIS 와 vLLM 에 통합하여 대규모 모델 (Llama 3.1 405B) 과 전문가 혼합 (MoE) 모델 (Qwen3-235B) 에서 엔드 투 엔드 성능 향상을 입증했습니다.

결과

• 확장 행동: TP 와 HP 모두 멀티 노드 환경에서 강력한 확장 (strong scaling) 이 부진합니다. HP 는 계산 병목 (prefill-heavy) 환경에서 TP 보다 우수하지만, TP 는 메모리 병목 (decode-heavy) 환경에서 HP 보다 우수합니다.
• 통신 병목: NCCL all-reduce 는 작은 메시지 (128 KB~2 MB) 에 대해 노드 간 확장성이 낮아, 이러한 특정 환경에서는 종종 MPI 기반 솔루션보다 성능이 떨어집니다.
• NVRAR 성능:
  • 마이크로 벤치마크: HPE Slingshot 및 InfiniBand 인터커넥트에서 128 KB2 MB 크기의 메시지에 대해 NVRAR 은 NCCL 대비 **1.9 배3.6 배 낮은 지연 시간**을 달성합니다.
  • 엔드 투 엔드 추론: YALIS 와 vLLM 에 통합된 NVRAR 은 텐서 병렬화를 사용하는 디코드 중심 워크로드에서 Llama 3.1 405B 모델의 엔드 투 엔드 배치 지연 시간을 최대 1.72 배 단축합니다.
  • MoE 모델: NVRAR 은 병렬화 전략의 TP 구성 요소를 가속화하여 MoE 모델 (예: Qwen3-235B) 의 처리량도 향상시켜, NCCL 기반 구성 대비 최대 1.14 배 높은 처리량을 달성합니다.

의의

이 논문은 LLM 이 더 커지고 추론이 멀티 노드 클러스터로 이동함에 따라 통신 효율성이 성능의 주요 결정 요인이 된다고 주장합니다. 저자들은 표준 통신 라이브러리 (NCCL) 가 LLM 디코드 단계의 특정 메시지 크기와 지연 시간 제약에 최적화되어 있지 않음을 입증했습니다. NVRAR 을 도입함으로써 그들은 통신 오버헤드를 크게 줄이는 표적 솔루션을 제공하여 대규모 모델에 대한 텐서 병렬화의 더 효율적인 확장을 가능하게 합니다. 이 연구는 차세대 분산 추론 시스템을 위한 커스텀 통신 원시 연산의 필요성을 강조하며, 특정 HPC 인터커넥트 환경에서 표준 집단 통신 라이브러리보다 계층적 NVSHMEM 기반 접근법의 유용성을 검증합니다.