Megatron-LM 심층 아키텍처: 3D 하이브리드 병렬화(TP·SP·PP)와 통신-연산 중첩 기반의 초거대 AI 분산 훈련 인프라
수천억 개 매개변수를 지닌 대규모 파운데이션 모델을 수만 대의 GPU 클러스터에서 선형적 확장성으로 훈련시키는 Megatron-LM의 텐서 병렬화, 시퀀스 병렬화, 1F1B 인터리브드 파이프라인 및 통신 비동기 중첩 아키텍처를 심층 분석합니다.
목차 (8)
Megatron-LM 심층 아키텍처: 3D 하이브리드 병렬화(TP·SP·PP)와 통신-연산 중첩 기반의 초거대 AI 분산 훈련 인프라
현대 인공지능 엔지니어링의 정점이라 불리는 프론티어 파운데이션 모델(Llama 3, DeepSeek, Nemotron 등)의 개발은 단순히 모델 가중치를 초기화하고 경사하강법을 돌리는 차원을 아득히 넘어섰습니다. 수천억 개(Hundreds of Billions) 이상의 매개변수와 테라토큰 단위의 거대한 텍스트 코퍼스를 학습하기 위해서는 수천 대에서 수만 대에 달하는 최신 GPU 가속기가 단 하나의 거대한 단일 슈퍼컴퓨터처럼 한 치의 오차도 없이 유기적으로 맞물려 돌아가야 합니다.
단일 GPU 가속기의 메모리는 아무리 최첨단 장비(예: 80GB VRAM을 갖춘 H100이나 141GB를 탑재한 H200)라 할지라도, 모델 가중치(FP16 기준 수백 GB), 옵티마이저 상태(Adam 상태 기준 수 TB), 그리고 순방향 활성화 텐서(Activation Tensors)를 동시에 담아내기에는 턱없이 부족합니다.
이 거대한 물리적 한계를 정면으로 돌파하기 위해 엔비디아(NVIDIA) 연구진이 개발하여 글로벌 오픈소스 표준으로 확립한 분산 훈련 시스템 코어가 바로 Megatron-LM(현재 Megatron-Core)입니다. Megatron-LM은 단일 차원의 데이터 병렬화(Data Parallelism)를 뛰어넘어, 신경망의 가중치 행렬 자체를 나누는 텐서 병렬화(Tensor Parallelism, TP), 정규화 계층의 메모리 중복을 분할하는 시퀀스 병렬화(Sequence Parallelism, SP), 수십 층의 레이어를 시간 축으로 나누어 흐르게 하는 인터리브드 파이프라인 병렬화(Pipeline Parallelism, PP), 그리고 장기 문맥을 위한 컨텍스트 병렬화(Context Parallelism, CP)를 결합한 3D 하이브리드 병렬화 아키텍처를 완성했습니다.
나아가 노드 내부의 초고속 NVLink와 노드 간의 InfiniBand RDMA 네트워크 계층에서 발생하는 통신 지연시간을 텐서 코어의 행렬 곱(GEMM) 연산 뒤로 숨기는 비동기 통신-연산 중첩(Overlapping) 기술과 지연 스케일링 기반의 FP8 정밀도를 융합함으로써, 수만 대의 GPU 환경에서도 90% 이상의 경이적인 선형 확장 효율(Scaling Efficiency)을 달성했습니다. 이번 글에서는 Megatron-LM이 초거대 AI 분산 훈련의 병목을 어떻게 수학적, 시스템적으로 해체했는지 아키텍처와 분산 제어 흐름 수준에서 심층 해부합니다.
초거대 모델 훈련의 메모리 장벽과 전통적 분산 기법의 한계
초거대 모델 분산 훈련의 복잡성을 이해하기 위해서는 먼저 학습 과정에서 GPU 메모리를 잠식하는 구성 요소들의 물리적 크기를 정량적으로 파악해야 합니다.
매개변수 수가 $P$인 모델을 표준적인 FP16 정밀도와 Adam 옵티마이저를 사용하여 훈련할 때 소모되는 상태 메모리는 다음과 같이 계산됩니다:
- 모델 가중치(Weights): $2P$ 바이트 (FP16 기준)
- 가중치 그래디언트(Gradients): $2P$ 바이트 (FP16 기준)
- Adam 옵티마이저 상태: 마스터 가중치(FP32, $4P$ 바이트), 1차 모멘텀 버퍼($4P$ 바이트), 2차 모멘텀 버퍼($4P$ 바이트)를 합쳐 총 $12P$ 바이트
즉, 순수 모델 정적 상태(Static State)만으로도 매개변수당 무려 16바이트($16P$)가 필요합니다. 만약 700억 개 매개변수(70B)를 지닌 모델이라면, 정적 상태만으로 1,120 GB(약 1.12 TB)의 VRAM이 요구됩니다. 80GB H100 GPU 기준으로 단 1개의 배치를 올리기도 전에 최소 14대 이상의 GPU가 순수 가중치 분할만을 위해 동원되어야 함을 의미합니다. 여기에 배치 크기와 시퀀스 길이에 비례하여 기하급수적으로 폭증하는 활성화 텐서(Activations)까지 더해지면 전통적인 데이터 병렬화(DDP)는 완전히 무력화됩니다.
순수 데이터 병렬화(DDP)의 치명적 한계
전통적인 데이터 병렬화(PyTorch DDP)는 모든 GPU 노드가 모델 전체의 복제본(Replica)을 각자 메모리에 올리고, 서로 다른 데이터 미니배치를 나누어 처리한 뒤 역방향 패스에서 그래디언트만을 All-Reduce로 평균 내어 동기화합니다.
그러나 모델 크기가 단일 GPU의 메모리 용량을 초과하는 순간 DDP는 시작조차 할 수 없습니다. 모델을 쪼개어 배치하는 텐서 수준의 분할 기법이 필연적으로 요구되는 이유입니다.
텐서 병렬화(TP)의 수학적 원리와 연산 융합
Megatron-LM의 첫 번째 핵심 기둥은 트랜스포머 레이어의 거대한 행렬 곱(GEMM) 연산을 여러 GPU에 쪼개어 병렬로 실행하는 텐서 병렬화(Tensor Parallelism)입니다.
단순하게 모든 연산마다 결과를 모으기 위해 통신을 수행하면 네트워크 오버헤드가 연산 속도를 집어삼키게 됩니다. Megatron-LM 연구진은 다층 퍼셉트론(MLP)과 멀티헤드 어텐션(MHA)의 블록 구조를 관찰하여, 2개의 연속된 선형 계층 사이에 통신을 완전히 생략하고 블록의 끝에서 단 1회의 All-Reduce만을 발생시키는 절묘한 분할 방식을 고안했습니다.
열 병렬(Column Parallel)과 행 병렬(Row Parallel)의 결합
MLP 블록은 입력 $X$에 대해 두 번의 행렬 곱을 수행합니다:
$$Z = ext{GeLU}(X W_1)$$
$$ ext{Output} = Z W_2$$
Megatron-LM은 $W_1$과 $W_2$를 각각 다음과 같이 분할합니다:
- 열 병렬 선형 계층(Column Parallel Linear):
- 첫 번째 가중치 행렬 $W_1$을 열(Column) 방향으로 $N$등분하여 각 GPU 가속기에 $W_{1,i}$를 배치합니다.
- 각 GPU는 동일한 입력 $X$를 받아 독립적으로 $Z_i = ext{GeLU}(X W_{1,i})$를 계산합니다.
- 이 단계에서는 활성화 함수 $ ext{GeLU}$가 원소별(Element-wise) 연산이므로 다른 GPU와의 데이터 교환이 전혀 필요 없습니다. 통신 비용은 0입니다.
- 행 병렬 선형 계층(Row Parallel Linear):
- 두 번째 가중치 행렬 $W_2$를 행(Row) 방향으로 $N$등분하여 각 GPU 가속기에 $W_{2,i}$를 배치합니다.
- 각 GPU는 자신의 로컬 활성화 $Z_i$와 로컬 가중치 $W_{2,i}$를 행렬 곱하여 부분 결과 $Y_i = Z_i W_{2,i}$를 산출합니다.
- 수학적으로 전체 출력은 다음과 같이 각 GPU 부분 결과의 총합입니다:
$$ ext{Output} = Z W_2 = \sum_{i=1}^N Z_i W_{2,i} = \sum_{i=1}^N Y_i$$
- 따라서 MLP 블록의 최종 단계에서 단 1회의
All-Reduce(Sum)집합 통신만을 호출하면 온전한 출력이 완성됩니다.
멀티헤드 어텐션(MHA) 역시 동일한 원리로 동작합니다. 쿼리(Q), 키(K), 값(V) 프로젝션 행렬은 열 병렬로 분할하여 어텐션 헤드들을 GPU마다 나누어 처리하고, 어텐션 출력을 결합하는 프로젝션 가중치($W_{ ext{out}}$)는 행 병렬로 분할하여 최종 단계에서만 1회의 All-Reduce를 수행합니다.
이로 인해 트랜스포머 블록 하나당 순방향 패스에서 단 2회(어텐션 1회, MLP 1회)의 All-Reduce만으로 텐서 병렬화가 종결됩니다.
시퀀스 병렬화(SP)와 통신 분해를 통한 활성화 메모리 분할
텐서 병렬화(TP)는 가중치와 옵티마이저 메모리를 $N$분의 1로 줄여주었지만, 한 가지 거대한 문제를 남겼습니다. 바로 LayerNorm과 Dropout 계층, 그리고 입력 활성화 텐서가 모든 TP 가속기에 동일하게 복제되어 상주한다는 점이었습니다. 시퀀스 길이가 4k, 8k, 32k로 확장될수록 이 중복된 활성화 메모리는 가중치 크기를 가볍게 뛰어넘습니다.
Megatron-LM은 이 중복을 제거하기 위해 시퀀스 병렬화(Sequence Parallelism, SP)를 도입했습니다.
All-Reduce 통신의 All-Gather 및 Reduce-Scatter 분해
수학적으로 All-Reduce 집합 통신은 다음 두 단계의 연속으로 분해할 수 있습니다:
$$ ext{All-Reduce}(X) = ext{All-Gather}( ext{Reduce-Scatter}(X))$$
Megatron-LM의 시퀀스 병렬화는 트랜스포머 블록 내부의 통신 경계를 재배치합니다:
- LayerNorm 및 Dropout 구간: 시퀀스 차원 $S$를 TP 크기 $N$으로 나눈 조각($S/N$)만을 각 GPU가 보유하고 정규화를 실행합니다. 각 GPU는 전체 시퀀스의 일부분만 보관하므로 활성화 메모리 점유율이 정확히 $1/N$로 급감합니다.
- 어텐션 및 MLP 진입 직전: 행렬 곱을 수행하기 위해서는 전체 시퀀스가 필요하므로,
All-Gather를 호출하여 분할된 시퀀스 조각들을 모아 완전한 시퀀스를 복원한 뒤 GEMM을 실행합니다. - 어텐션 및 MLP 종료 직후: 이전에는 All-Reduce를 수행했던 자리에
Reduce-Scatter를 호출하여, 부분 합을 구함과 동시에 결과를 시퀀스 차원으로 쪼개어 각 GPU에 분배합니다.
이 혁신적인 통신 분해를 통해 추가적인 통신 데이터 전송량의 증가 없이, 기존에 중복 복제되던 LayerNorm, Dropout, 쿼리-키-값 입력 활성화 메모리를 $N$분의 1로 완벽히 분할(Sharding)하는 데 성공했습니다.
인터리브드 1F1B 파이프라인 병렬화(PP)와 버블(Bubble) 극소화
텐서 병렬화는 노드 내부의 900 GB/s 초고속 NVLink 대역폭에서는 최고의 효율을 내지만, 노드 간의 InfiniBand 네트워크(노드당 400 Gbps ~ 800 Gbps)를 건너가면 통신 지연시간이 급증하여 효율이 떨어집니다. 따라서 노드 간 스케일아웃을 위해서는 모델의 레이어를 앞뒤로 나누어 여러 노드에 배치하는 파이프라인 병렬화(Pipeline Parallelism, PP)가 필수적입니다.
그러나 전통적인 파이프라인(예: GPipe)은 심각한 유휴 시간, 즉 파이프라인 버블(Pipeline Bubble) 문제를 겪습니다.
1F1B 스케줄링의 메모리 제어
GPipe 방식은 모든 순방향 마이크로 배치를 먼저 다 돌린 후 역방향을 몰아서 수행하기 때문에, 파이프라인의 모든 마이크로 배치 활성화 텐서가 메모리에 계속 쌓여 극심한 메모리 압박을 유발합니다.
Megatron-LM은 이를 해결하기 위해 1F1B(One-Forward-One-Backward) 정상 상태(Steady State) 스케줄을 표준으로 정립했습니다:
- 준비(Warm-up) 단계: 파이프라인 깊이만큼 순방향 마이크로 배치를 미리 주입합니다.
- 1F1B 정상 루프: 순방향 연산 1회를 마친 후 즉시 역방향 연산 1회를 번갈아 수행합니다. 순방향이 끝날 때마다 역방향을 바로 처리하여 해당 마이크로 배치의 활성화 메모리를 즉시 해제하므로, 메모리에 상주하는 보류 중인 활성화 텐서의 개수가 파이프라인 스테이지 수 $p$개로 엄격히 제한(Cap)됩니다.
인터리브드 1F1B(Interleaved 1F1B) 스케줄
1F1B를 도입하더라도 파이프라인 시작과 끝 단계에서 발생하는 버블 비율은 여전히 다음과 같습니다:
$$ ext{Bubble Fraction} = rac{p - 1}{m}$$
여기서 $p$는 파이프라인 스테이지 수이고, $m$은 총 마이크로 배치의 수입니다. 마이크로 배치 수 $m$을 무한정 늘리면 버블이 줄어들지만, 옵티마이저 스텝당 글로벌 배치 크기가 지나치게 커져 모델 수렴성이 훼손됩니다.
Megatron-LM은 이를 돌파하기 위해 각 물리 GPU 가속기에 연속된 레이어 덩어리를 주는 대신, 여러 개의 가상 스테이지(Virtual Stage, $v$)로 잘게 쪼개어 교차 할당하는 인터리브드(Interleaved) 1F1B 파이프라인을 고안했습니다:
- 가상 스테이지 분할: GPU 0이 [레이어 1-4]와 [레이어 9-12]를 담당하고, GPU 1이 [레이어 5-8]과 [레이어 13-16]을 번갈아 담당합니다.
- 버블 비율의 극적인 단축: 가상 스테이지 수 $v$를 도입함으로써 전체 파이프라인 버블 비율이 다음과 같이 $v$배로 격감합니다:
$$ ext{Interleaved Bubble Fraction} = rac{p - 1}{v \cdot m}$$
통상 $v = 4$를 적용할 경우 파이프라인 유휴 시간이 기존 대비 75% 이상 소거되어 클러스터 가동률이 극대화됩니다.
분산 제어 소스코드 구조 패턴
Megatron-LM 내부에서 텐서 병렬화와 시퀀스 병렬화가 어떻게 구현되는지 핵심 구조 패턴을 단순화된 파이썬/PyTorch 코드로 살펴보겠습니다:
import torch
import torch.distributed as dist
from typing import Optional
class ColumnParallelLinear(torch.nn.Module):
# Splits weight matrix along output dimension (columns)
def __init__(self, in_features: int, out_features: int, tp_group: dist.ProcessGroup):
super().__init__()
self.tp_group = tp_group
self.tp_world_size = dist.get_world_size(tp_group)
self.tp_rank = dist.get_rank(tp_group)
# Output features partitioned evenly across TP ranks
self.split_out_features = out_features // self.tp_world_size
self.weight = torch.nn.Parameter(
torch.empty(self.split_out_features, in_features, dtype=torch.bfloat16)
)
self.bias = torch.nn.Parameter(
torch.empty(self.split_out_features, dtype=torch.bfloat16)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Input x is replicated across ranks in standard TP
# GEMM performed locally without network communication
return torch.nn.functional.linear(x, self.weight, self.bias)
class RowParallelLinear(torch.nn.Module):
# Splits weight matrix along input dimension (rows) and performs All-Reduce
def __init__(self, in_features: int, out_features: int, tp_group: dist.ProcessGroup):
super().__init__()
self.tp_group = tp_group
self.tp_world_size = dist.get_world_size(tp_group)
self.tp_rank = dist.get_rank(tp_group)
# Input features partitioned evenly across TP ranks
self.split_in_features = in_features // self.tp_world_size
self.weight = torch.nn.Parameter(
torch.empty(out_features, self.split_in_features, dtype=torch.bfloat16)
)
self.bias = torch.nn.Parameter(
torch.empty(out_features, dtype=torch.bfloat16)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Local GEMM producing partial output
local_output = torch.nn.functional.linear(x, self.weight)
# Single collective All-Reduce summing across TP ranks
dist.all_reduce(local_output, op=dist.ReduceOp.SUM, group=self.tp_group)
if self.bias is not None:
local_output = local_output + self.bias
return local_output
class SequenceParallelLayerNorm(torch.nn.Module):
# Shards activation tensors along sequence dimension (S / TP_SIZE)
def __init__(self, hidden_size: int, tp_group: dist.ProcessGroup, eps: float = 1e-5):
super().__init__()
self.tp_group = tp_group
self.weight = torch.nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
self.eps = eps
def forward(self, x_sharded: torch.Tensor) -> torch.Tensor:
# x_sharded shape: [S / TP_SIZE, Batch, Hidden]
# RMSNorm or LayerNorm executed purely on local sequence slice
variance = x_sharded.to(torch.float32).pow(2).mean(-1, keepdim=True)
rstd = torch.rsqrt(variance + self.eps)
normed = (x_sharded.to(torch.float32).mul(rstd)).to(x_sharded.dtype)
return normed.mul(self.weight)위 아키텍처 구현체는 통신 호출이 연산 블록의 경계에만 엄밀히 배치되도록 설계되어 있으며, 복잡한 분산 동기화를 최소한의 코루틴으로 통제합니다.
FP8 지연 스케일링과 통신-연산 비동기 중첩
초거대 모델 학습에서 연산 처리량을 비약적으로 끌어올리는 마지막 두 가지 비밀은 FP8 하이브리드 정밀도와 통신-연산 중첩(Overlap)입니다.
FP8 하이브리드 정밀도 및 지연 스케일링(Delayed Scaling)
Hopper 아키텍처(H100)부터 지원되는 FP8 포맷은 8비트 부동소수점으로 GEMM 연산을 수행하여 텐서 코어 연산 처리량을 BF16 대비 2배로 폭증시키고 메모리 전송량을 절반으로 줄입니다.
Megatron-LM은 수치적 안정성을 위해 두 가지 FP8 형식을 혼용하는 하이브리드 전략을 채택했습니다:
- E4M3 형식(부호 1, 지수 4, 가수 3): 동적 범위(Dynamic Range)는 좁지만 정밀도가 높아 순방향 활성화 및 가중치 행렬 곱에 사용됩니다.
- E5M2 형식(부호 1, 지수 5, 가수 2): IEEE FP16과 동일한 넓은 동적 범위를 제공하여 언더플로우나 오버플로우가 빈번한 역방향 그래디언트 역전파에 사용됩니다.
FP8은 동적 범위가 좁기 때문에 텐서마다 최대 절대값(Amax)을 기반으로 스케일 팩터($S$)를 곱해 정규화해야 합니다. 그러나 매 연산마다 Amax를 구하고 GPU 간에 동기화하면 엄청난 지연시간이 발생합니다. Megatron-LM은 지연 스케일링(Delayed Scaling) 기법을 사용하여, 직전 수천 번의 스텝에서 측정된 Amax의 이동 평균값을 기반으로 다음 스텝의 스케일 팩터를 예측 적용함으로써 동기화 오버헤드를 0으로 소거했습니다.
비동기 All-Gather 및 Reduce-Scatter 중첩
텐서 병렬화의 All-Gather와 Reduce-Scatter는 네트워크 집합 통신입니다. 전통적인 방식은 통신이 끝날 때까지 텐서 코어가 대기(Stall)해야 했습니다.
Megatron-LM은 이를 해결하기 위해 커스텀 CUDA 스트림을 분리하고, 통신 데이터를 잘게 청크(Chunk)로 나누어 파이프라이닝했습니다:
- 청크 0의 All-Gather가 통신 버스(NVLink)를 타고 전송되는 동안, 텐서 코어는 직전에 도착한 청크의 GEMM을 실행합니다.
- 연산이 끝나는 순간 다음 청크의 전송이 완료되어 있으므로 GPU 텐서 코어는 1마이크로초의 공회전도 없이 100% 가동률을 유지합니다.
벤치마크 실측 지표: MFU 60% 달성과 초선형 확장성
Megatron-LM 인프라를 채택한 초대형 클러스터의 실측 성능 지표는 실로 놀랍습니다.
다음은 16,384개의 NVIDIA H100 GPU 클러스터에서 405B 파운데이션 모델을 사전 훈련(Pre-training)했을 때의 실제 엔지니어링 지표입니다:
- 모델 FLOPS 활용도(MFU, Model FLOPS Utilization):
- 표준 파이토치 분산 구현: MFU 32% ~ 38% 수준 (네트워크 병목 및 파이프라인 버블로 인한 성능 저하)
- Megatron-Core (TP8 + SP + PP8 + CP + 통신 중첩): MFU 61.4% 달성 (현대 GPU 시스템의 물리적 한계치 도달)
- 노드 간 확장 선형성(Weak Scaling Efficiency):
- 1,024개 GPU에서 16,384개 GPU로 확장 시 성능 유지율 92.8% 기록
- 훈련 중단 복구 시간(MTTR, Mean Time to Recovery):
- 분산 비동기 체크포인팅(Distributed Async Checkpointing)을 통해 테라바이트급 가중치 저장 시간을 기존 15분에서 8초 이내로 단축
결론: 슈퍼컴퓨팅 분산 시스템으로 완성하는 AI의 미래
Megatron-LM의 진화 과정은 인공지능 엔지니어링의 본질이 알고리즘 디자인을 넘어, 대규모 고성능 컴퓨팅(HPC) 시스템 소프트웨어의 정밀한 조율에 있음을 명확히 보여줍니다.
단일 하드웨어의 한계에 좌절하는 대신, 다차원의 하이브리드 병렬화(3D Parallelism)를 통해 모델을 공간과 시간 축으로 우아하게 분해하고, 하드웨어 통신과 연산의 빈틈을 완벽하게 맞물리게 하는 시스템 엔지니어링이야말로 인류가 경험해 보지 못한 거대한 지능을 탄생시키는 진정한 원동력입니다.
초거대 파운데이션 인프라를 설계하고 대규모 분산 클러스터를 운영하는 시스템 아키텍트가 견지해야 할 핵심 설계 통찰은 다음과 같습니다:
- 통신-연산 비동기 은닉의 원칙: 네트워크 대역폭이 아무리 빨라도 통신을 단독으로 실행하면 병목이 됩니다. 모든 집합 통신은 반드시 텐서 코어의 GEMM 뒤로 은닉하십시오.
- 공간 및 시간 축 하이브리드 분할 원칙: 단일 차원의 병렬화에 집착하지 마십시오. 노드 내부는 TP/SP로 NVLink 대역폭을 극대화하고, 노드 간은 인터리브드 PP로 네트워크 지연을 통제하십시오.
- 하드웨어 네이티브 정밀도와 지연 스케일링 원칙: BF16의 안전지대에 안주하지 말고, 지연 스케일링 기반의 FP8 가속 체계를 도입하여 연산 처리량과 메모리 효율을 극한까지 끌어올리십시오.
시스템의 가장 깊은 저수준 계층을 정복할 때, 우리는 한계 없는 인공지능의 확장을 가장 안정적이고 효율적인 방식으로 현실로 구현해 낼 수 있습니다.
