Megatron Core 병렬화 전략

Megatron Core는 여러 병렬화 전략을 지원하고, 이를 조합해서 수십억에서 수조 파라미터의 모델을 수천 개 GPU에서 효율적으로 학습할 수 있어요. 이 페이지에서는 각 병렬화 전략이 무엇을 최적화하는지, 언제 쓰는 게 좋은지, 어떻게 조합하는지 설명해 드릴게요.

출처: Megatron Core Parallelism Strategies Guide (공식)

병렬화 전략 개요

전략	병렬화 목표	추천 용도
데이터 병렬 (DP)	배치 차원	데이터 확장, 표준 학습
텐서 병렬 (TP)	개별 레이어	큰 레이어·활성화, GPU 메모리 제약
파이프라인 병렬 (PP)	모델 깊이	매우 깊은 모델
컨텍스트 병렬 (CP)	시퀀스 길이	긴 시퀀스 (8K+ 토큰)
엑스퍼트 병렬 (EP)	MoE 엑스퍼트	Mixture-of-Experts 모델
완전 샤딩 데이터 병렬 (Megatron-FSDP)	모델 상태	매우 큰 모델, DP 교체

데이터 병렬 (DP)

표준 DDP는 모델을 GPU 전체에 복제하고 배치를 나눠요.

torchrun --nproc_per_node=8 pretrain_gpt.py \
    --data-parallel-sharding-strategy no_shard

각 GPU는 모델의 전체 복사본을 갖고 배치의 일부를 처리해요.

Megatron-FSDP는 모델 파라미터, 그래디언트, 옵티마이저 상태를 GPU에 샤딩해 메모리 사용을 줄여요.

--use-megatron-fsdp
--data-parallel-sharding-strategy optim_grads_params
--ckpt-format fsdp_dtensor
--init-model-with-meta-device

--data-parallel-sharding-strategy 옵션은 ZeRO 단계와 연결돼요.

optim — 옵티마이저 상태만 샤딩 (ZeRO-1)
optim_grads — 그래디언트 + 옵티마이저 샤딩 (ZeRO-2)
optim_grads_params — 파라미터 + 그래디언트 + 옵티마이저 샤딩 (ZeRO-3)

텐서 병렬 (TP)

개별 모델 레이어를 GPU 여러 개에 나눠 담는 전략이에요. 큰 은닉 차원에 권장돼요.

--tensor-model-parallel-size 4   # 4-way 텐서 병렬화
--sequence-parallel              # 시퀀스 병렬화 (권장)

레이어가 단일 GPU에 안 들어가거나, 은닉 차원이 크거나(4096+), DP·PP와 함께 쓸 때 유용해요.

파이프라인 병렬 (PP)

모델 레이어를 깊이(수직) 기준으로 GPU에 나눠 담아요.

--pipeline-model-parallel-size 8              # 8 파이프라인 스테이지
--num-layers-per-virtual-pipeline-stage 4     # 로드밸런싱용 가상 파이프라인

매우 깊은 모델(50+ 레이어), 큰 모델에 TP와 함께, GPU 간 메모리 분산에 유용해요.

컨텍스트 병렬 (CP)

긴 시퀀스를 GPU에 나눠 담아 장문 컨텍스트 학습을 효율화해요.

--context-parallel-size 2           # 2-way 컨텍스트 병렬화
--cp-comm-type p2p                  # 통신 유형

8K+ 토큰의 긴 시퀀스, 활성화 메모리 절감에 유용하고, TP·PP·DP와 조합할 수 있어요.

엑스퍼트 병렬 (EP)

Mixture-of-Experts 모델에서 엑스퍼트를 GPU에 분산시켜요.

--expert-model-parallel-size 8   # 8-way 엑스퍼트 병렬화
--num-experts 64                 # MoE 레이어당 64 엑스퍼트
--moe-grouped-gemm               # 엑스퍼트 계산 최적화

EP와 TP를 함께 쓸 때는 반드시 Sequence Parallelism을 켜야 해요.

--tensor-model-parallel-size 4
--expert-model-parallel-size 8
--sequence-parallel   # TP + EP 사용 시 필수

전략 조합하기

총 GPU 수는 이렇게 계산돼요.

Total GPUs = TP × PP × CP × EP × DP

예를 들어 LLaMA-3 70B를 64 GPU에서 학습한다면 TP=4, PP=4, CP=2, DP=2 → 4 × 4 × 2 × 2 = 64 GPU예요.

torchrun --nproc_per_node=8 pretrain_gpt.py \
    --tensor-model-parallel-size 4 \
    --pipeline-model-parallel-size 4 \
    --context-parallel-size 2 \
    --num-layers 80 \
    --hidden-size 8192 \
    --num-attention-heads 64 \
    --seq-length 8192 \
    --micro-batch-size 1 \
    --global-batch-size 512 \
    --bf16

전략 고르기

간단하게 시작하려면 데이터 병렬(DP)만 쓰고, 모델이 안 들어가면 TP를, 그래도 부족하면 PP를, 긴 시퀀스면 CP를 추가하는 흐름이 좋아요. 메모리 제약에는 FSDP로 모델 상태를, TP로 큰 레이어를, PP로 모델 깊이를 나누고, 극단적인 경우 활성화 체크포인팅·오프로딩을 켜요.

Megatron Core 병렬화 전략