← all posts
AI 2026.05.03 · 13 min read Advanced

분산 학습의 네 가지 축 — 3D Parallelism, MoE, Checkpoint, Elastic

70B+ 모델 학습에서 DP×TP×PP 그룹 분해부터 MoE all-to-all, 비동기 체크포인트, Elastic 회복 효율까지 — 분산 학습 시스템의 설계 원칙을 추적한다.

📚 Distributed Training Deep Dive 7 / 7
  1. 01 분산 학습의 통신은 왜 전부 AllReduce로 귀결되는가
  2. 02 Data Parallelism의 수학 — AllReduce는 왜 정확한가
  3. 03 Tensor Parallelism은 왜 AllReduce가 정확히 2번인가
  4. 04 Pipeline Bubble은 어떻게 줄어드는가
  5. 05 ZeRO는 왜 단계적으로 분산하는가
  6. 06 Transformer 학습에서 Activation Memory는 왜 폭발하는가
  7. 07 분산 학습의 네 가지 축 — 3D Parallelism, MoE, Checkpoint, Elastic

70B 모델은 단일 GPU에 올라가지 않는다. 그래서 병렬화한다. 그런데 병렬화하면 장애가 나고, 장애가 나면 재시작해야 하고, 재시작하려면 체크포인트가 있어야 한다. 이 네 가지 문제 — 분할, 라우팅, 저장, 회복 — 는 서로 독립적으로 보이지만 하나의 설계 철학으로 연결되어 있다. 그 철학은 무엇인가?

world_size = DP × TP × PP

GPU 64개짜리 클러스터에서 LLaMA-70B를 학습시키는 표준 설정은 TP=8, PP=4, DP=2다. 이 세 축은 각자 다른 문제를 해결한다.

Tensor Parallelism은 단일 레이어의 가중치를 쪼갠다. QKV projection을 8-way로 나누면 각 GPU는 d/8 크기의 행렬만 보유한다. 대신 매 레이어마다 4번의 AllReduce가 필요하다 — 대역폭이 충분해야 한다. 그래서 TP는 NVLink로 연결된 노드 내부에 가둔다(~300 GB/s).

Pipeline Parallelism은 레이어를 시간 축으로 분할한다. 노드 0이 stage 0(레이어 0-19)을 처리하면 activation을 노드 1로 넘기고, 노드 1은 stage 1(레이어 20-39)을 처리한다. 통신은 레이어당 한 번의 P2P — 빈도는 낮지만 InfiniBand(~100 GB/s)를 쓴다.

Data Parallelism은 미니배치를 복제한다. 두 replica가 각자 forward/backward를 돌린 뒤 gradient를 AllReduce로 합산한다. backward 도중 overlap이 가능하다.

r=(rdp⋅TP⋅PP)+(rpp⋅TP)+rtpr = (r_{\text{dp}} \cdot \text{TP} \cdot \text{PP}) + (r_{\text{pp}} \cdot \text{TP}) + r_{\text{tp}}

world rank rr은 이 공식으로 (dp_id, pp_id, tp_id) 좌표로 분해된다. 같은 (dp_id, pp_id)를 공유하는 8개 rank가 TP group, 같은 (dp_id, tp_id)를 공유하는 4개 rank가 PP group이다.

✎ 통신 대역폭과 배치 원칙

TP는 매 레이어마다 실행되므로 높은 대역폭(NVLink)이 필수다. PP는 stage 경계마다 한 번이므로 중간 대역폭(InfiniBand)으로 충분하다. 이 비대칭이 “TP는 intra-node, PP는 inter-node” 원칙의 수학적 근거다.

MoE — active params를 줄이는 조건부 계산

Dense Transformer에서 FLOPs는 파라미터 수에 정비례한다. 1T 파라미터 모델은 7B 모델 대비 약 150배 연산이 필요하다. MoE(Mixture-of-Experts)는 이 제약을 깬다.

각 token은 전체 expert 중 top-kk개만 통과한다. Mixtral 8×7B는 총 47B 파라미터지만 토큰당 활성 파라미터는 13B(top-2)다 — dense 13B의 연산량으로 dense 47B의 표현력을 얻는다.

MoE(x)=∑i∈Tk(x)pi∑j∈Tk(x)pj⋅ei(x)\text{MoE}(x) = \sum_{i \in \mathcal{T}_k(x)} \frac{p_i}{\sum_{j \in \mathcal{T}_k(x)} p_j} \cdot e_i(x)

그러나 대가가 따른다. 라우팅이 불균등해지면 일부 expert만 과다하게 사용되고 나머지는 학습되지 않는다(dead expert). 이를 막기 위해 보조 손실을 추가한다.

Laux=α⋅E⋅∑i=1Efi⋅Pi\mathcal{L}_\text{aux} = \alpha \cdot E \cdot \sum_{i=1}^{E} f_i \cdot P_i

fif_i는 expert ii가 받은 토큰 비율, PiP_i는 평균 gating 확률이다. Cauchy-Schwarz 부등식으로 이 손실의 최솟값은 fi=Pi=k/Ef_i = P_i = k/E — 즉 완전 균등 분배 — 일 때 달성됨을 증명할 수 있다.

명제 1 · Load Balancing Loss의 minimizer

∑ifi=∑iPi=k\sum_i f_i = \sum_i P_i = k 제약 하에서 Laux=α⋅E⋅∑ifiPi\mathcal{L}_\text{aux} = \alpha \cdot E \cdot \sum_i f_i P_i의 최솟값은 fi=Pi=k/Ef_i = P_i = k/E (균등 분배)일 때 달성된다.

▷ 증명

f,Pf, P의 합이 각각 kk로 고정된 상태에서 Lagrangian을 세우면: ∂L/∂fi=EPi−λ=0⇒Pi=const\partial \mathcal{L}/\partial f_i = E P_i - \lambda = 0 \Rightarrow P_i = \text{const}, ∂L/∂Pi=Efi−μ=0⇒fi=const\partial \mathcal{L}/\partial P_i = E f_i - \mu = 0 \Rightarrow f_i = \text{const}. 따라서 최솟값은 fi=Pi=k/Ef_i = P_i = k/E이고 이때 Laux=αk2\mathcal{L}_\text{aux} = \alpha k^2이다. Hessian이 양반정치(positive semi-definite)임을 확인해 최솟값임을 보인다. □\square

∎

MoE의 또 다른 비용은 통신이다. EP(Expert Parallelism)에서 각 GPU는 일부 expert만 보유하므로, 토큰을 owning GPU로 보냈다가 결과를 받아오는 2회 all-to-all이 매 레이어마다 발생한다. 이 때문에 EP는 NVLink 내부로 제한하는 것이 권장된다.

체크포인트 — 저장 방식이 재시작 유연성을 결정한다

70B 모델 + Adam optimizer의 체크포인트 크기는 약 1.1 TB다. 이것을 어떻게 저장하느냐가 재시작의 유연성과 속도를 결정한다.

Single-rank gather: rank 0이 전체를 모아 한 파일로 저장. 어떤 topology로도 resume 가능하지만, 단일 NIC 병목으로 10분 이상 걸린다.

Sharded sync: 각 rank가 자기 shard만 저장. 16-rank 학습이라면 per-rank 70 GB를 동시에 쓴다. 속도는 NN배 빠르지만 같은 topology로만 resume 가능하다.

Async sharded: GPU → 고정 CPU 메모리로 snapshot 후 즉시 다음 step 시작. background thread가 disk에 쓰는 동안 GPU는 쉬지 않는다. stall 시간이 ~0에 수렴한다.

Tsharded async≈0,Tsync sharded=ΨBagg,Tsingle=ΨBsingleT_\text{sharded async} \approx 0, \quad T_\text{sync sharded} = \frac{\Psi}{B_\text{agg}}, \quad T_\text{single} = \frac{\Psi}{B_\text{single}}

**PyTorch DCP(Distributed Checkpoint)**는 각 shard에 chunk_offset, chunk_size, total_size 메타데이터를 포함시켜 다른 topology로의 resume을 가능하게 한다. 새 rank는 자신에게 필요한 byte range를 여러 shard 파일에서 partial read해 재조합한다.

Elastic Training — 회복 효율의 수학

700-GPU 클러스터의 MTBF(Mean Time Between Failures)는 수 시간에 불과하다. Meta OPT-175B 학습은 175일 동안 100회 이상 재시작했다. Elastic Training 없이는 사실상 불가능한 규모다.

TorchElastic의 핵심은 rendezvous — 모든 worker가 합의하는 (world_size, rank, master_addr) 튜플이다. 노드 장애 시 rendezvous가 재계산되고, 살아있는 노드들이 새 world_size로 init_process_group을 호출해 학습을 재개한다.

ηelastic=TT+(T/M)⋅R\eta_\text{elastic} = \frac{T}{T + (T/M) \cdot R}

학습 시간 T=100T=100시간, MTBF M=10M=10시간, resume 비용 R=5R=5분이면 η≈99.2%\eta \approx 99.2\%다. RR을 DCP async로 1분으로 줄이면 η≈99.7%\eta \approx 99.7\% — 3%p 차이가 100시간 학습에서 3시간 compute를 구한다.

⚠ Elastic range와 유효 배치 크기

nnodes=MIN:MAX 설정에서 노드 수가 변하면 유효 배치 크기가 달라진다. Gradient accumulation steps KK를 노드 수에 반비례로 동적 조정해 B_micro × N_nodes × K = const를 유지해야 학습 곡선이 연속된다.

정리

네 개의 챕터는 서로 다른 문제를 다루지만 하나의 질문으로 수렴한다 — 100+ GPU에서 수 주짜리 학습을 어떻게 안정적으로 완료할 것인가.

  • 3D Parallelism: world_size = DP × TP × PP. TP는 NVLink(intra-node), PP는 InfiniBand(inter-node). 통신 빈도와 대역폭의 매칭이 핵심이다.
  • MoE: active params를 k/Ek/E로 줄이되, load balancing auxiliary loss로 expert collapse를 막는다. EP는 가능하면 NVLink 내부로 제한한다.
  • Checkpointing: async sharded가 stall을 0에 가깝게 만든다. topology 변경이 필요하면 DCP re-shard — 단 TP/PP/MoE 축 변경은 별도 처리가 필요하다.
  • Elastic Training: η≈1−R/M\eta \approx 1 - R/M. resume 비용 RR을 줄이는 것이 장애 내성을 높이는 직접적인 수단이다.

이 네 가지 설계는 독립적으로 선택할 수 없다 — 체크포인트 방식이 elastic 회복의 유연성을 결정하고, parallelism topology가 체크포인트 re-shard 가능성을 제약하며, MoE의 expert 수가 EP degree를 강제한다.

REF
REF
Jiang et al. · 2024 · Mixtral of Experts · arXiv
# distributed-training # 3d-parallelism # moe # checkpointing # elastic-training