단일 GPU로 1000억 개 LLM 학습? MegaTrain의 비밀

MegaTrain은 기존 GPU 중심 시스템과 달리, 매개변수와 옵티마이저 상태를 호스트 메모리(Host Memory), 즉 CPU 메모리에 저장한다. GPU는 일시적인 연산 엔진으로 활용되며, 각 레이어(Layer)마다 매개변수를 스트리밍(Streaming)하고 기울기를 계산하여 장치 상태를 최소화한다. 이러한 접근 방식은 VRAM 제약(VRAM Constraint)을 극복하고 대규모 모델 학습을 가능하게 한다.

MegaTrain은 CPU-GPU 간의 대역폭 병목 현상(Bandwidth Bottleneck)을 해결하기 위해 두 가지 주요 최적화를 적용했다. 첫째, 파이프라인 이중 버퍼링(Pipelined Double-buffered Execution) 엔진을 도입하여 매개변수 프리페칭(Prefetching), 계산, 기울기 오프로딩(Offloading)을 여러 CUDA 스트림(CUDA Streams)에서 겹치도록 하여 GPU의 지속적인 실행을 가능하게 한다. 둘째, 영구적인 자동 구문 분석 그래프(Autograd Graph)를 상태 비저장 레이어 템플릿(Stateless Layer Templates)으로 대체하여 동적으로 가중치를 바인딩하고, 영구적인 그래프 메타데이터를 제거한다.

MegaTrain은 140억 개 매개변수 모델 학습 시 DeepSpeed ZeRO-3(CPU 오프로딩)보다 1.84배 높은 처리량을 달성했다. 하지만, 댓글에서는 단일 3090 GPU에서 140억 개 모델 학습 시 341 tok/s의 속도를 보인다는 점을 지적하며, 다른 최적화 기법(예: Muon, 양자화)을 활용하면 더 높은 성능을 얻을 수 있다고 주장한다. 또한, H200 GPU에서 1200억 개 모델 학습에 걸리는 시간과 8개의 GPU를 사용할 경우의 성능에 대한 질문도 제기되었다.

일부 사용자는 Apple M 시리즈와 같은 통합 메모리 아키텍처(Unified Memory Architecture)에서의 MegaTrain 작동 방식에 대한 궁금증을 나타냈다. 이러한 아키텍처는 GPU와 CPU가 메모리를 공유하므로, 데이터 전송 오버헤드를 줄일 수 있다. 따라서, MegaTrain의 이중 버퍼링(Double-buffering)과 같은 기법이 통합 메모리 환경에서 어떤 이점을 가질 수 있는지에 대한 추가적인 분석이 필요하다.