엔비디아(NVIDIA)는 JAX 기반 대규모 언어모델(LLM) 훈련 과정에서 발생하는 고대역폭 메모리(HBM) 병목 현상을 해결하기 위한 새로운 접근 방식인 '호스트 오프로딩(Host Offloading)' 기술을 공개했습니다. 이 기술은 GPU의 제한적인 HBM 용량을 효율적으로 사용하게 함으로써, 연구자와 개발자들이 더 큰 규모의 LLM을 훈련하거나 기존 훈련의 효율성을 크게 향상시킬 수 있도록 돕습니다.
호스트 오프로딩은 GPU의 HBM에 모든 모델 파라미터를 저장하는 대신, 일부 파라미터를 CPU의 메인 메모리(호스트 메모리)로 옮겨 저장하는 방식입니다. 필요할 때만 GPU로 데이터를 전송하여 HBM 사용량을 최적화합니다. 엔비디아는 이 기술을 통해 HBM 사용량을 최대 50%까지 줄일 수 있으며, 이는 모델 크기를 두 배로 늘리거나 동일한 HBM 용량으로 더 많은 모델을 처리할 수 있음을 의미합니다. 특히, JAX와 같은 프레임워크에서 분산 훈련 시 메모리 관리의 유연성을 높여줍니다.
이 기술은 LLM 개발 비용을 절감하고 접근성을 높이는 데 중요한 역할을 할 수 있습니다. GPU HBM은 고가이며 용량 확장이 어렵기 때문에, HBM 사용량을 줄이는 것은 곧 더 적은 자원으로 더 큰 모델을 훈련할 수 있다는 뜻입니다. 이는 스타트업이나 연구 기관 등 자원 제약이 있는 곳에서도 대규모 AI 모델 개발에 참여할 수 있는 기회를 확대하며, 궁극적으로 AI 혁신의 속도를 가속화할 것으로 기대됩니다.