MegaTrain: 단일 GPU에서 100B+ 파라미터 대형 언어 모델의 전정밀도 학습 구현

기존에는 수천억 개의 파라미터를 가진 모델을 학습하려면 수십 대의 GPU가 필요했으나, 이 논문은 단 한 대의 GPU와 충분한 CPU 메모리만으로도 대규모 모델 학습이 가능함을 입증했다. 이는 고가의 GPU 자원이 부족한 중소 연구소나 개인 개발자들도 거대 모델의 미세 조정 및 학습에 참여할 수 있는 길을 열어준다.

딥러닝 학습의 핵심인 Gradient Descent를 수행하려면 모델의 가중치(Weights)와 이를 업데이트하기 위한 옵티마이저 상태가 메모리에 상주해야 한다. 하지만 모델이 커질수록 이 데이터들이 GPU의 고속 메모리(HBM) 용량을 초과하게 되어 학습 자체가 불가능해지는 '메모리 벽' 문제가 발생한다.

MegaTrain은 이 문제를 해결하기 위해 '도서관 대출 시스템'과 같은 방식을 도입했다. 모든 책(파라미터)을 좁은 책상(GPU) 위에 다 올려두는 대신, 큰 창고(CPU 메모리)에 보관하고 지금 읽을 페이지(레이어)만 책상으로 가져와 읽은 뒤 즉시 돌려보내는 방식이다. 이때 책을 가져오는 시간 동안 흐름이 끊기지 않도록, 현재 페이지를 읽는 동안 다음 페이지를 미리 가져오는 이중 버퍼링 기술을 적용했다.

결과적으로 GPU는 오직 현재 계산 중인 레이어 하나만큼의 메모리만 사용하게 된다. 이를 통해 GPU 메모리 크기에 상관없이 CPU 메모리 용량이 허용하는 한 수천억 개의 파라미터를 가진 거대 모델도 단 한 대의 GPU에서 안정적으로 학습할 수 있게 된다.

MegaTrain은 CPU 메모리를 주 저장소로 활용하는 메모리 중심 아키텍처를 채택했다. 모든 파라미터와 옵티마이저 상태는 CPU 메모리에 유지되며, 연산 시에만 레이어 단위로 GPU에 스트리밍된다. 순전파 시에는 레이어 가중치를 GPU 버퍼로 불러와 연산 후 즉시 해제하며, 역전파 시에는 다시 가중치를 불러와 그래디언트를 계산한 뒤 이를 즉시 CPU로 전송한다.

데이터 전송 효율을 위해 Layer-Contiguous Tiling 기법을 사용한다. 각 레이어의 가중치, 그래디언트, 옵티마이저 상태를 하나의 연속된 메모리 블록으로 묶어 PCIe 대역폭을 최대한 활용하는 단일 버스트 DMA 전송을 수행한다. [레이어별 파라미터 묶음 → 단일 DMA 요청 → GPU 전송 → 전송 효율 극대화] 순으로 동작하여 작은 조각들을 여러 번 보낼 때 발생하는 오버헤드를 제거했다.

실행 엔진은 Compute, Weight-Transfer, Gradient-Transfer라는 3개의 CUDA 스트림을 병렬로 운영한다. 이중 버퍼링을 통해 n번째 레이어를 계산하는 동안 n+1번째 레이어 가중치를 미리 가져오고, n-1번째 레이어의 그래디언트를 CPU로 내보낸다. 또한 Stateless Layer Template을 도입하여 연산 로직과 데이터를 분리함으로써, 기존 PyTorch의 Autograd가 생성하는 거대한 연산 그래프 메타데이터를 GPU 메모리에서 제거했다.

단일 NVIDIA H200 GPU(141GB HBM)와 1.5TB CPU 메모리 환경에서 최대 120B 파라미터 모델의 전정밀도 학습을 성공적으로 수행했다. 32B 모델 학습 시 250 TFLOPS 이상의 성능을 유지했으며, 이는 기존의 오프로딩 기반 시스템들이 메모리 부족(OOM)으로 실패하는 영역에서도 안정적인 동작을 보여준 것이다.

효율성 측면에서 MegaTrain은 14B 모델 학습 시 DeepSpeed ZeRO-3(CPU 오프로딩 모드) 대비 1.84배 높은 처리량을 기록했다. 특히 모델의 깊이(Depth)가 깊어져도 성능 저하가 매우 완만하게 나타났는데, 28레이어에서 180레이어로 모델 규모가 3.95배 커지는 동안 처리량은 단 20.1%만 감소하여 뛰어난 확장성을 입증했다.

정확도 검증 결과, MegaTrain은 표준적인 전체 GPU 학습 방식 및 ZeRO 기반 방식과 동일한 수치적 정확도를 유지했다. 이는 제안된 블록 단위 재계산(Block-wise Recomputation)과 CPU 기반 옵티마이저 업데이트가 모델의 수렴 성능에 부정적인 영향을 주지 않음을 의미한다.

MegaTrain의 핵심은 GPU 메모리 점유율을 모델 전체 크기가 아닌 '단일 레이어의 크기'로 고정(Bound)시킨 것이다. 이를 위해 표준 Autograd의 전역 그래프 구조를 탈피하고, 레이어 단위로 실행을 제어하는 명시적 스케줄링 모델을 구현했다. 역전파 시에는 메모리 절약을 위해 블록 단위 재계산 전략을 사용하여 활성화 함수 값을 관리한다.

CPU 측면에서는 고성능 벡터 명령어(AVX-512 등)를 활용하여 Adam 옵티마이저 업데이트를 수행한다. 이는 PCIe 대역폭이 병목인 상황에서 GPU로 데이터를 주고받는 비용보다 CPU에서 직접 계산하는 것이 더 효율적이라는 판단에 근거한다. 또한 Pinned Slab Pool을 통해 호스트 측 메모리 고정(Pinning) 오버헤드를 일정하게 유지하며 시스템 안정성을 확보했다.

구현 측면에서는 PyTorch와 CUDA를 기반으로 하되, 성능에 민감한 파라미터 바인딩 및 메모리 관리 로직은 C++ 확장으로 작성하여 Python의 디스패치 오버헤드를 최소화했다. 특히 GH200과 같은 NVLink-C2C 아키텍처뿐만 아니라 일반적인 PCIe Gen4 기반 시스템(A100, RTX 3090 등)에서도 기존 기법들보다 월등한 성능 향상을 확인했다.

CPU 메모리 용량이 전체 모델 상태(파라미터, 그래디언트, 옵티마이저 상태)를 수용할 수 있을 만큼 충분해야 한다. 또한 CPU-GPU 간의 PCIe 대역폭이 극도로 낮은 환경에서는 연산과 전송의 오버랩이 완벽하지 않아 성능 저하가 발생할 수 있다.