MegaTrain: 단일 GPU로 100B+ 파라미터 LLM을 Full Precision으로 학습하기

MegaTrain: Full Precision Training of 100B+ Parameter LLMs on a Single GPU

Apr 8, 2026•chrsw•View Original

TL;DR Highlight

CPU 메모리를 주 저장소로 활용하고 GPU는 연산 엔진으로만 쓰는 방식으로, H200 GPU 단 한 장으로 120B 파라미터 모델을 풀 정밀도로 학습할 수 있는 시스템인 MegaTrain을 소개한다.

Who Should Read

GPU VRAM 부족으로 대형 모델 학습에 어려움을 겪고 있는 ML 엔지니어나 연구자, 특히 단일 GPU 환경에서 파인튜닝이나 실험을 돌려야 하는 개발자.

Core Mechanics

MegaTrain은 '메모리 중심(memory-centric)' 설계를 채택해서, 모델 파라미터와 옵티마이저 상태를 GPU VRAM이 아닌 CPU 호스트 메모리에 저장한다. GPU는 연산이 필요한 순간에만 데이터를 받아서 처리하는 '일시적 연산 엔진' 역할을 한다.
각 레이어를 처리할 때마다 파라미터를 CPU에서 GPU로 스트리밍해서 넣고, 계산된 그래디언트는 다시 CPU로 내보내는 방식을 쓴다. 이렇게 하면 GPU에 상주하는 데이터가 최소화되어 VRAM 사용량을 극적으로 줄일 수 있다.
CPU-GPU 간 대역폭 병목을 해결하기 위해 '파이프라인 더블 버퍼링(pipelined double-buffered) 실행 엔진'을 도입했다. 여러 CUDA 스트림을 사용해서 파라미터 프리페칭, 실제 연산, 그래디언트 오프로딩을 동시에 겹쳐서 실행하므로 GPU가 쉬지 않고 계속 돌아갈 수 있다.
기존의 autograd 그래프를 유지하는 대신 '무상태 레이어 템플릿(stateless layer templates)'을 도입해서, 파라미터가 스트리밍되어 들어올 때 동적으로 바인딩하는 방식을 쓴다. 이렇게 하면 그래프 메타데이터를 GPU에 계속 올려둘 필요가 없어지고 스케줄링 유연성도 높아진다.
H200 GPU 1장과 1.5TB CPU 메모리 조합으로 최대 120B 파라미터 모델 학습을 안정적으로 검증했다. 또한 14B 모델 기준으로 DeepSpeed ZeRO-3 + CPU 오프로딩 대비 1.84배 높은 학습 처리량(throughput)을 달성했다.
GH200(GPU-CPU 통합 아키텍처) 단일 기기에서 7B 모델을 512k 토큰 컨텍스트로 학습시킬 수 있다는 것도 확인했다. 이는 초장문 컨텍스트 학습을 단일 노드에서 가능하게 한다는 점에서 의미가 있다.

Evidence

RTX 3080(VRAM 10GB)을 가진 한 사용자는 40M~50M 파라미터 이상 모델 학습 시 OOM이 나서 고생하고 있었는데, 이 방식을 쓰면 CPU RAM이 많은 자신의 PC에서 훨씬 큰 모델을 로컬로 학습할 수 있겠다며 긍정적으로 반응했다.
한 댓글러는 이 아이디어 자체는 새롭지 않다고 지적하면서, 실제 속도 측면에서 한계가 있다고 비판했다. 논문에서 14B 모델을 단일 3090으로 341 tok/s 달성했다고 했는데, 자신은 비슷한 방식으로 4090에서 약 1,000 tok/s를 냈다고 주장하면서, 그래도 실용적 pretraining에 쓰기엔 너무 느리다는 점을 강조했다.
같은 댓글러는 논문에서 쓰지 않은 추가 최적화 기법들도 있다고 언급했다. 예를 들어 그래디언트를 별도로 오프로딩하지 않고 옵티마이저 상태에 바로 누적하는 방법, Adam 대신 VRAM을 절반만 쓰는 Muon 옵티마이저 사용, 파라미터와 옵티마이저 상태 모두 4-bit 양자화 적용 등이다.
H200 + 1.5TB 호스트 메모리라는 조건에 대해 '단일 GPU라고는 하지만 결코 가벼운 세팅이 아니다'라는 냉소적인 반응도 있었다. 일반 개발자가 쉽게 접근할 수 있는 환경이 아니라는 점을 지적한 것이다.
PyTorch의 FSDP(Fully Sharded Data Parallel) 기능과 유사해 보인다는 의견도 있었다. 구체적으로 `torch.distributed.fsdp` 프리미티브만으로 이 방식을 어느 정도 재현할 수 있을지 궁금하다는 질문이 제기됐고, Apple M 시리즈처럼 GPU-CPU 메모리가 통합된 아키텍처에서는 이 기법이 얼마나 효과적일지에 대한 의문도 나왔다.
이 기법은 실용적으로는 소규모 파인튜닝 작업에나 쓸 법하고, 대규모 pretraining에 쓰기엔 속도가 너무 느리다는 의견이 있었다. DeepSpeed와 유사하다는 반응도 있었다.

How to Apply

VRAM이 부족한 단일 GPU(예: RTX 3090, RTX 4090) 환경에서 14B~30B 규모 모델을 파인튜닝해야 하는 경우, DeepSpeed ZeRO-3 + CPU 오프로딩보다 MegaTrain 방식을 적용하면 동일 하드웨어에서 최대 1.84배 빠른 학습 처리량을 기대할 수 있다.
H200 같은 고사양 단일 GPU와 1TB 이상의 CPU 메모리가 있는 서버 환경에서 100B 이상 모델을 멀티 GPU 없이 실험해보고 싶다면, MegaTrain의 레이어 단위 스트리밍 방식이 다중 노드 세팅 없이도 120B 모델 학습을 가능하게 해준다.
초장문 컨텍스트(512k 토큰 이상)가 필요한 7B 모델 학습을 단일 장비에서 돌려야 하는 경우, GH200처럼 GPU-CPU 메모리가 고속으로 연결된 아키텍처와 MegaTrain을 조합하면 멀티 GPU 없이도 이를 처리할 수 있다.
직접 구현하기 어렵다면, PyTorch의 `torch.distributed.fsdp` (Fully Sharded Data Parallel) 기능이 유사한 CPU 오프로딩을 지원하므로 이를 먼저 검토한 뒤 성능이 부족할 때 MegaTrain 논문의 더블 버퍼링 기법을 추가로 적용하는 단계적 접근을 고려할 수 있다.

Terminology

Full Precision Training모델 파라미터를 FP32(32비트 부동소수점) 그대로 학습하는 방식. 반대는 FP16이나 BF16 같은 혼합 정밀도(mixed precision) 학습이다. 정밀도가 높을수록 메모리를 더 많이 쓴다.

CPU OffloadingGPU VRAM에 다 담기 어려운 데이터(파라미터, 옵티마이저 상태 등)를 CPU RAM으로 옮겨두는 기법. GPU는 필요할 때만 가져다 쓴다. 속도는 느려지지만 훨씬 큰 모델을 다룰 수 있다.

ZeRO-3DeepSpeed의 메모리 최적화 전략 중 가장 공격적인 단계. 파라미터, 그래디언트, 옵티마이저 상태를 여러 GPU에 분산 저장하고, 필요할 때만 수집해서 쓴다.

Double Buffering다음 작업에 필요한 데이터를 미리 준비해두는 기법. A를 처리하는 동안 B를 미리 로드해서, A가 끝나면 기다림 없이 바로 B를 처리할 수 있게 한다.

Autograd GraphPyTorch 등이 역전파(backpropagation)를 위해 내부적으로 유지하는 연산 흐름 그래프. 학습 중 메모리를 상당히 차지한다.

Throughput (처리량)단위 시간당 처리할 수 있는 토큰 수(tok/s). 학습 속도를 나타내는 지표로, 높을수록 같은 시간에 더 많은 데이터를 학습할 수 있다.