단일 GPU로 120B 모델 학습하는 법 - MegaTrain

들어가며

100B급 모델을 학습한다고 생각해보겠습니다. 보통은 GPU 수십 장을 묶은 분산 학습을 떠올립니다. 가중치와 옵티마이저 상태만으로 GPU 한 장의 메모리를 크게 초과하기 때문입니다. 그런데 2026년에 공개된 MegaTrain은 H200 1장과 host RAM 1.5TB로 최대 120B 모델 학습을 주장합니다. 이 글에서는 학습에 필요한 메모리가 어떻게 구성되는지 계산해보고 MegaTrain이 단일 GPU로 이를 감당하는 방법을 알아보겠습니다.

 

1. 학습 메모리의 회계

학습에는 추론보다 훨씬 많은 메모리가 필요합니다. mixed-precision 학습 기준으로 파라미터 1개당 유지해야 하는 상태는 다음과 같습니다.

항목 정밀도 크기

weight	fp16	2 B
gradient	fp16	2 B
master weight	fp32	4 B
(Adam 모멘텀 2종)	fp32	8 B

모멘텀을 제외해도 파라미터당 8 B입니다. Adam까지 포함하면 16 B입니다. 120B 모델에 적용하면 weight만 240 GB이고 persistent state 전체는 약 887 GB입니다. Adam 포함 시 약 1.85 TB까지 늘어납니다. H200의 HBM은 141 GB이므로 단순 계산으로 GPU 7~13장이 필요합니다.

여기에 activation이 추가됩니다. forward 과정의 중간 결과는 backward에서 다시 쓰이므로 어딘가에 들고 있어야 합니다.

 

2. 메모리 계층과 offload의 기본 발상

GPU에 다 들어가지 않는 데이터를 다루는 공통 전략이 있습니다. 데이터를 메모리 계층에 나눠 두고 연산 직전에만 GPU로 끌어올리는 offload입니다. ZeRO-3 Offload, ZeRO-Infinity, FSDP, MegaTrain 모두 같은 발상에서 출발합니다.

계층 매체 용량 대역폭 역할

L1	GPU SRAM	수십 MB	~80 TB/s	연산 직전 캐시
L2	HBM	수십~수백 GB	TB/s 급	GPU 주력 작업 메모리
L3	host DDR5/LPDDR5X	TB 단위	수백 GB/s (HBM의 약 1/10)	대용량 보관
L4	NVMe	수십 TB	한 자리 GB/s	persistent storage

이 전략이 성립하는 근거는 학습 연산의 구조에 있습니다. forward와 backward는 레이어 단위로 순차 진행됩니다. 그래서 HBM에는 지금 연산 중인 레이어 1개분(120B 기준 약 10 GB)만 있으면 됩니다. 나머지 weight는 host RAM에 둡니다.

 

3. 기존 offload의 두 가지 한계

같은 발상의 ZeRO 계열이 이미 있는데 MegaTrain이 새로 필요한 이유는 두 가지입니다.

전송이 연산을 막습니다. host와 GPU 사이의 PCIe 대역폭은 HBM보다 훨씬 좁습니다. ZeRO는 전송을 하나의 레인이 모두 담당하는 구조라 backward에서 GPU가 데이터를 기다리며 유휴 상태로 머무는 시간이 커집니다. 게다가 작은 덩어리 다수로 나눠 보내기 때문에 PCIe 대역폭도 충분히 활용하지 못합니다.

activation은 레이어 단위로 쪼갤 수 없습니다. weight는 연산 중인 레이어 것만 있으면 되지만 activation은 다릅니다. backward가 끝날 때까지 전체 레이어의 중간 결과를 유지해야 합니다. ZeRO는 이 activation을 GPU에 남겨 둡니다. 그래서 모델이 32B를 넘으면 activation만으로 OOM이 발생합니다.

 

4. 해법 1: 3개 stream과 double-buffer

MegaTrain의 첫 번째 해법은 전송 시간을 연산 시간 뒤에 숨기는 것입니다. GPU가 데이터를 기다리며 유휴 상태로 머무는 시간을 제거합니다.

• 3개 CUDA stream: 연산(S_comp), 다음 레이어 weight 올리기(S_H2D), gradient 내리기(S_D2H)를 별도 stream으로 분리합니다. 올리기와 내리기가 서로를 막지 않습니다.
• Double-buffer: 레이어 크기의 버퍼 2개를 둡니다. 연산 stream이 버퍼 0을 읽는 동안 전송 stream이 버퍼 1을 채웁니다. 두 버퍼를 번갈아 사용합니다.
• Large-burst tiling: 레이어 하나의 상태를 큰 덩어리 한 번에 전송합니다. ZeRO의 잦은 소량 전송과 달리 PCIe 대역폭을 거의 포화시킵니다.

여기에 activation까지 host로 offload합니다. ZeRO가 GPU에 남겨 두던 activation을 같은 겹치기 방식으로 host에 내리고 올립니다. 32B에서 막히던 한계가 120B까지 늘어난 핵심 차이입니다.

 

5. 해법 2: Stateless Layer Template

두 번째 해법은 GPU 메모리 점유를 레이어 1개분으로 고정하는 장치입니다. 레이어의 데이터 스키마만 GPU에 상주시킵니다. streaming buffer로 도착한 데이터를 스키마에 동적으로 bind하고 연산한 뒤 unbind합니다. 레이어마다 별도의 메모리를 잡지 않으므로 모델이 커져도 GPU 메모리 사용량은 늘지 않습니다.

 

6. 성능과 적용 범위

논문이 제시하는 결과는 다음과 같습니다.

환경 결과

H200 1장 + host 1.5 TB	최대 120B 모델 학습
GH200 1장, 7B 모델	512K token context 학습
GH200 1장, 14B 모델	ZeRO-3 Offload 대비 1.84× throughput (264 vs 약 143 TFLOPS)

512K context 실험은 activation 병목을 보여주는 사례입니다. 7B 모델의 weight는 14 GB에 불과합니다. 그런데 context가 4,096에서 524,288 토큰으로 128배 늘면 activation이 폭증합니다. attention score는 시퀀스 길이의 제곱에 비례하기 때문입니다. weight가 아니라 activation이 병목인 상황도 offload로 해결할 수 있음을 보여줍니다.

적용 범위는 분명합니다. MegaTrain은 multi-GPU 분산 학습에는 아직 부적합합니다. 단일 GPU에 큰 모델을 올려야 하는 상황에 특화된 기법입니다. ZeRO-3가 OOM으로 막히는 30B 이상을 단일 GPU로 학습해야 한다면 현재로서는 유일한 선택지입니다.

 

정리

항목 내용

문제	120B 학습의 persistent state는 약 887 GB로 GPU 1장을 한참 초과
기본 전략	weight·activation을 host RAM에 두고 연산 직전에만 GPU로 올리는 offload
ZeRO의 한계	전송이 연산을 막고 activation을 GPU에 남겨 32B+에서 OOM
해법 1	3개 CUDA stream + double-buffer + large-burst로 전송을 연산에 겹침
해법 2	Stateless Layer Template로 GPU 점유를 레이어 1개분으로 고정
적용 범위	단일 GPU에 큰 모델. multi-GPU 분산에는 부적합

GPU 1장의 HBM 용량이 모델 크기를 결정하지 않는다는 점이 이 논문의 핵심입니다. 전송을 연산 뒤에 완전히 숨길 수 있다면 host RAM까지가 사실상의 학습 메모리가 됩니다.

 

출처

• "MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU", arXiv:2604.05091 — https://arxiv.org/abs/2604.05091
• 코드: https://github.com/DLYuanGod/MegaTrain
• 사내 세미나 발표 정리 노트 (2026-05-20, 발표자: 진희솜)