NVIDIA는 호스트 메모리 오프로딩 기법을 사용하여 JAX 기반 LLM 학습의 고대역폭 메모리 병목을 줄이는 연구를 발표했습니다.
대규모 언어 모델 훈련 워크로드는 계산이 완전히 활용되기 전에 점점 더 GPU 메모리 제한에 직면합니다. 모델 가중치, 그래디언트, 최적화기 상태, 통신 버퍼, 중간 활성화가 모두 GPU 고대역폭 메모리(HBM)를 두고 경쟁합니다. 모델 크기, 시퀀스 길이, 배치 크기가 증가함에 따라 HBM 용량이 종종 주요 확장 병목이 됩니다.
오픈소스 JAX 라이브러리의 호스트 오프로딩은 정방향 패스 중 선택된 활성화를 고정된 호스트 메모리로 이동시키고 역방향 패스에서 필요할 때 다시 스트리밍함으로써 HBM 압력을 줄입니다. 이 접근 방식은 선택된 활성화를 호스트 메모리에서 다시 로드하는 대신 재계산해야 하는 활성화 재연산에 대한 대안을 제공합니다.
호스트 오프로딩은 NVIDIA Grace Blackwell 시스템에서 특히 유리합니다. NVIDIA Grace CPU와 NVIDIA Blackwell GPU는 NVLink-C2C를 통해 연결되어 양방향으로 900 GB/s의 대역폭을 제공하므로 고정된 호스트 메모리가 선택된 활성화를 위한 실용적인 임시 저장소가 됩니다. Vera CPU와 Rubin GPU는 양방향 속도를 1.8 TB/s의 일관된 대역폭으로 두 배로 높여 성능을 더욱 개선합니다. 그러나 고대역폭 CPU-GPU 연결성만으로는 충분하지 않으며, 성능을 개선하려면 활성화 전송이 유용한 GPU 작업과 겹쳐야 합니다.
MaxText(가속화된 선형 대수(XLA) 컴파일러를 활용하는 JAX LLM 훈련 프레임워크)를 사용하여 실험을 수행했습니다. 모든 결과는 128개의 GPU를 사용하는 NVIDIA GB200 NVL72 시스템에서 두 가지 워크로드로 측정되었습니다: 고정 배치 크기에서 대상 쿼리, 키, 값(QKV) 활성화 오프로딩을 연구하는 데 사용되는 밀집 디코더 전용 트랜스포머 모델인 Llama 3.1 405B, 그리고 처리량 및 메모리 용량 효과를 연구하는 데 사용되는 다중 헤드 잠재 주의(MLA)를 갖춘 희소 전문가 혼합 모델인 DeepSeek-V3 671B.
DeepSeek-V3 671B은 61개의 디코더 레이어를 포함합니다. 처음 3개는 밀집 다층 퍼셉트론(MLP) 블록을 사용하고, 나머지 레이어는 MoE 블록을 사용합니다. 스택을 지배하는 반복적인 MoE 디코더 레이어에 대한 활성화 오프로딩 정책은 선택된 MLA 쿼리와 키/값 프로젝션 중간값, 그리고 선택된 MoE 상향 프로젝션 중간값을 오프로드합니다. 이러한 활성화는 더 큰 배치 구성이 메모리에 맞을 수 있는지를 결정하기에 충분할 정도로 큽니다.
오프로딩, 지연 시간 숨김 스케줄러(LHS), 파이프라인된 전송이 활성화되면 DeepSeek-V3 671B은 초당 디바이스당 908.2 TFLOPs를 달성했습니다. 이는 동일한 배치 구성에서 활성화 재연산보다 57% 빠르고, LHS 또는 파이프라인 없는 오프로딩보다 67.7% 빠릅니다. 밀집 Llama 워크로드에서는 LHS만으로도 지연 시간을 숨기기에 충분하지만, DeepSeek-V3 MoE 및 MLA 레이어의 대규모 활성화 풋프린트는 파이프라인된 전송이 전체 처리량에 뚜렷한 긍정적 영향을 미친다는 의미입니다. 이 성능 이점은 NVIDIA의 소프트웨어와 하드웨어의 긴밀한 공동 설계를 반영합니다. Blackwell 시스템에서 XLA 사용자 정의 스케줄링 플래그는 데이터가 비동기적으로 이동되도록 보장하는 전용 복사 스트림과 함께 작동합니다. 이 통합을 통해 플랫폼은 컴파일러-상호연결 통합이 부족한 아키텍처에 대해 접근할 수 없는 대규모 배치 구성을 잠금 해제할 수 있습니다.
용량 비교를 위해 디바이스에 선택된 활성화를 저장하면 미니 배치 2와 전역 배치 256이 가능했으며, 최적화된 호스트 오프로딩은 미니 배치 8과 전역 배치 1024를 가능하게 했습니다. 오프로딩이나 재연산이 없으면 디바이스는 미니 배치 8과 전역 배치 1024에서 메모리 부족 오류를 만납니다. 호스트 오프로딩은 선택된 활성화 저장소를 GPU 메모리 외부로 이동함으로써 이 구성을 가능하게 했으며, 모델 상태, 통신 버퍼, 런타임 작업 공간, 활성 계산을 위해 더 많은 HBM을 사용할 수 있게 했습니다. LHS 및 파이프라인된 전송이 활성화되면 오프로드 구성은 165.2 GiB의 GPU 메모리를 사용하며, 이러한 최적화가 없으면 145.6 GiB입니다. 이 증가는 전송과 계산을 겹치기 위해 더 많은 복사 버퍼와 미리 가져온 활성화를 GPU 메모리에 유지하여 일부 메모리 용량을 더 나은 겹침과 더 높은 처리량으로 거래하는 것을 반영합니다.
Llama 3.1 405B 실험은 배치 크기 2, 시퀀스 길이 8,192, 완전히 분할된 데이터 병렬화 128, bfloat16 활성화, NVFP4 4비트 가중치 양자화를 사용하여 합성 데이터에서 10단계를 실행했습니다. LHS를 사용한 QKV 활성화 오프로딩은 처리량을 2,669에서 2,746 TFLOPs/s/device로 개선했습니다. 이는 오프로딩 없는 기준선보다 2.9% 증가입니다. LHS를 비활성화하면 QKV 오프로드 처리량이 2,569 TFLOPs/s/device로 감소하며, 호스트 오프로딩이 다른 GPU 작업과의 효과적인 겹침에 달려 있음을 보여줍니다. 이 구성의 경우 LHS는 파이프라인 없이 2,746 TFLOPs/s/device에서 최고 처리량을 제공했으며, 파이프라인을 사용하면 2,718 TFLOPs/s/device였습니다. LHS가 이미 계산과 통신 뒤에서 대부분의 전송 지연 시간을 숨기기 때문입니다.
70.9 GiB 호스트 메모리 값은 한 순간에 절약되는 GPU 메모리의 양이 아니라 모든 126개 레이어에 걸친 전체 QKV 활성화 저장소를 나타냅니다. 배치 크기 2 및 시퀀스 길이 8,192에서 한 레이어의 bfloat16 QKV 활성화는 약 576 MiB를 필요로 합니다: 쿼리의 경우 512 MiB, 키와 값의 경우 각각 32 MiB. 레이어에 대해 스캔 루프가 활성화되면 역방향 패스는 한 번에 한 레이어를 처리하며, GPU에서 한 번에 하나의 레이어의 QKV 활성화만 필요합니다. 이 워크로드에서 QKV 오프로딩은 주로 역방향 패스 QKV 재연산을 계산 및 통신과 겹치는 전송으로 바꾸어 성능을 최적화하며, GPU 피크 메모리는 모델 상태, 통신 버퍼, 런타임 작업 공간이 지배합니다. 밀집 Llama 3.1 405B 모델은 희소 DeepSeek V3 671B보다 작은 이득을 보여주지만, 동일한 기본 메커니즘을 드러냅니다. 대상 지정된 QKV 오프로딩은 역방향 패스 재연산을 계산과 통신과 겹치는 전송으로 바꿉니다.