NVLink Switch가 탑재된 NVIDIA HGX H200의 Medusa로 최대 1.9배 향상된 Llama 3.1 성능

Reading Time: 3 minutes

거대 언어 모델(LLM)의 크기와 복잡성이 계속 증가함에 따라, 실시간 생성형 AI 애플리케이션이 요구하는 짧은 지연 시간과 높은 처리량을 제공하기 위해서는 멀티 GPU 컴퓨팅이 필수입니다.

성능은 초고속 GPU 간 통신을 통해 결합된 GPU가 ‘하나의 강력한 GPU’로서 요청을 처리하는 능력과 여러 GPU를 최대한 활용할 수 있는 고급 소프트웨어에 따라 달라집니다. 추측 디코딩과 같은 고급 알고리즘과 함께 텐서 병렬 처리라는 기술을 사용하여 각 모델 계층의 계산을 사용 가능한 GPU에 분할함으로써 토큰 생성 지연 시간을 줄여 대화형 사용자 경험을 제공할 수 있습니다.

지연 시간이 매우 짧은 Llama 3.1 서비스를 위해 클라우드 서비스는 각각 8개의 H200 텐서 코어 GPU와 4개의 올투올 NVLink Switch 칩이 통합된 전체 NVIDIA HGX H200 서버를 사용할 수 있습니다. 서버 내의 각 GPU는 NVLink Switch를 통해 다른 GPU와 최대 900GB/s의 대역폭으로 통신할 수 있습니다. 대화형 사용 사례에서 멀티 GPU 통신이 병목 현상이 발생하지 않도록 하려면 높은 GPU 간 패브릭 대역폭이 필요합니다.

*그림 1. 4개의 NVLink Switch 칩이 장착된 NVIDIA HGX H200.*

NVIDIA H200 HGX 시스템에서 최적화 알고리즘을 효율적으로 구현하기 위해 NVIDIA TensorRT-LLM이 사용됩니다. TensorRT-LLM은 텐서 병렬 처리 및 추측 디코딩을 비롯한 다양한 기술을 사용하여 최신 LLM에서 최첨단 추론 성능을 제공하는 오픈 소스 TensorRT 라이브러리입니다.

추론적 디코딩 알고리즘인 Medusa의 개선을 포함한 향후 TensorRT-LLM 최적화는 HGX H200에서 각각 초당 268 토큰과 초당 108 토큰의 뛰어난 저지연 성능을 제공하는 Llama 3.1 70B 및 Llama 3.1 405B에서 제공됩니다.

Medusa, NVIDIA HGX H200에서 토큰 생성량 최대 1.9배 증가

트랜스포머 기반 LLM은 자동 회귀적이기 때문에 토큰을 순차적으로 생성해야 하므로 생성 단계당 처리량이 단 하나의 토큰으로 제한됩니다. 일반적으로 LLM 추론 중에 단일 토큰이 생성되는 속도는 모델 가중치가 메모리에 얼마나 빨리 로드되는지에 따라 달라집니다. 즉, 워크로드로 인해 H200 GPU의 상당한 텐서 코어 기능을 제대로 활용하지 못할 수 있습니다.

투기적 디코딩은 “드래프트 모델”을 사용하여 다음 토큰 이후의 여러 후속 토큰을 예측함으로써 토큰 생성 단계당 토큰 생성 처리량을 증가시키는 기술입니다. 그런 다음 대상 LLM은 예측 후보를 ‘일괄 처리’하고 다음 토큰과 병렬로 검증하여 사용 가능한 병렬 GPU 컴퓨팅 리소스를 보다 효과적으로 활용합니다. 후보 시퀀스가 원본 LLM에 의해 승인되면 생성 단계에서 여러 개의 토큰이 생성되므로 토큰 생성 속도가 빨라집니다.

이 백서에서 설명하는 Medusa는 원본 모델을 초안 모델로 사용하는 투기적 디코딩 알고리즘으로, 별도의 초안 모델을 사용할 때 발생하는 시스템 복잡성과 분배 불일치를 피합니다. 이 기법은 Medusa 헤드라고 하는 추가 디코딩 “헤드”를 사용하여 다음 토큰을 넘어 후보 토큰을 예측합니다. 각 Medusa 헤드는 이전 토큰을 넘어서는 토큰의 분포를 생성합니다. 그런 다음 트리 기반 주의 메커니즘이 원본 모델의 유효성을 검사하기 위해 일부 후보 시퀀스를 샘플링합니다. 병렬 후보 시퀀스의 수를 초안 길이라고 하며, 생성 단계당 수락되는 토큰의 평균 수를 수락률이라고 합니다. 수락률이 높을수록 전체 토큰 생성 처리량이 증가합니다.

*그림 2. Medusa가 있는 경우와 없는 경우의 Llama 3.1 모델에서 NVLink Switch가 있는 NVIDIA HGX H200 성능.*

*추가 각주: MTBench 코딩 데이터세트, 단일 턴, 정확도는 기본 모델과 비트 단위로 동일, TensorRT-LLM의 내부 빌드, HGX H200 TP8, FP8, BS=1*

Medusa를 사용하면 HGX H200은 사용자당 초당 268개의 토큰을 생산할 수 있으며, Llama 3.1 70B의 경우 108개, Llama 3.1 405B의 경우 108개를 생산할 수 있습니다. 이는 Medusa를 사용하지 않을 때보다 Llama 3.1 70B에서는 1.5배 이상, Llama 3.1 405B에서는 1.9배 이상 빠른 속도입니다. 헤드의 미세 튜닝 방식에 따라 작업마다 Medusa의 수용률에 차이가 있지만, 전반적인 성능은 다양한 작업에서 일반화됩니다.

Llama 3.1 70B 및 Llama 3.1 405B용 Medusa 헤드는 NVIDIA NeMo 프레임워크와 통합된 NVIDIA TensorRT 모델 옵티마이저를 사용하여 훈련되었습니다. 메두사 헤드 트레이닝은 프로즌 백본을 사용하여 Medusa를 사용해도 기본 모델과 동일한 정확도를 얻을 수 있도록 했습니다.

멈추지 않는 NVIDIA 풀스택 혁신

NVLink Switch 및 TensorRT-LLM이 탑재된 NVIDIA HGX H200은 이미 인기 있고 가장 까다로운 커뮤니티 모델에서 뛰어난 실시간 추론 성능을 제공하고 있습니다. 사용자 경험을 지속적으로 개선하고 추론 비용을 절감하기 위해 칩, 시스템, 소프트웨어 라이브러리, 알고리즘 등 기술 스택의 모든 계층에서 끊임없이 혁신하고 있습니다.

플랫폼과 LLM 생태계가 발전함에 따라 향후 지연 시간이 짧은 추론 성능에 대한 업데이트를 공유할 수 있기를 기대합니다.