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

MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU

April 6, 2026

저자: Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye

cs.AI

초록

본 논문에서는 단일 GPU에서 완전 정밀도(full precision)로 100B+ 파라미터 규모의 대규모 언어 모델을 효율적으로 학습시키는 메모리 중심 시스템인 MegaTrain을 제안한다. 기존의 GPU 중심 시스템과 달리 MegaTrain은 매개변수와 옵티마이저 상태를 호스트 메모리(CPU 메모리)에 저장하고 GPU는 일시적인 연산 엔진으로 취급한다. 각 레이어마다 매개변수를 스트리밍 방식으로 읽어오고 그래디언트를 계산하여 출력함으로써 디바이스에 지속적으로 유지되는 상태를 최소화한다. CPU-GPU 대역폭 병목 현상을 해결하기 위해 두 가지 핵심 최적화를 적용했다. 1) 여러 CUDA 스트림 간에 매개변수 프리페치, 연산, 그래디언트 오프로딩을 중첩시키는 파이프라인 방식의 이중 버퍼 실행 엔진을 도입하여 GPU가 지속적으로 연산을 수행할 수 있게 했다. 2) 지속적인 autograd 그래프 대신 상태 비저장 레이어 템플릿을 사용하며, 매개변수가 스트리밍되어 들어올 때 동적으로 가중치를 바인딩하여 지속적인 그래프 메타데이터를 제거함과 동시에 스케줄링 유연성을 제공한다. 1.5TB 호스트 메모리를 장착한 단일 H200 GPU에서 MegaTrain은 120B 파라미터 규모의 모델을 안정적으로 학습시킬 수 있다. 또한 14B 모델 학습 시 DeepSpeed ZeRO-3(CPU 오프로딩) 대비 1.84배의 학습 처리량을 달성했다. MegaTrain은 단일 GH200으로 512k 토큰 컨텍스트 길이의 7B 모델 학습도 가능하게 한다.

English

We present MegaTrain, a memory-centric system that efficiently trains 100B+ parameter large language models at full precision on a single GPU. Unlike traditional GPU-centric systems, MegaTrain stores parameters and optimizer states in host memory (CPU memory) and treats GPUs as transient compute engines. For each layer, we stream parameters in and compute gradients out, minimizing persistent device state. To battle the CPU-GPU bandwidth bottleneck, we adopt two key optimizations. 1) We introduce a pipelined double-buffered execution engine that overlaps parameter prefetching, computation, and gradient offloading across multiple CUDA streams, enabling continuous GPU execution. 2) We replace persistent autograd graphs with stateless layer templates, binding weights dynamically as they stream in, eliminating persistent graph metadata while providing flexibility in scheduling. On a single H200 GPU with 1.5TB host memory, MegaTrain reliably trains models up to 120B parameters. It also achieves 1.84times the training throughput of DeepSpeed ZeRO-3 with CPU offloading when training 14B models. MegaTrain also enables 7B model training with 512k token context on a single GH200.