Skip to content

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

🌐 Distributed Training Deep Dive

DDP 의 all_reduce

$$\text{AllReduce}(g_1, \ldots, g_N) = \frac{1}{N}\sum_{i=1}^{N} g_i \quad \text{broadcast to all ranks}$$

를 쓰는 것 과,

Ring AllReduce 가 $N$ 개 GPU 에서 bandwidth-optimal 하다는

$$T_\mathrm{ring}(N, P) = \frac{2(N-1)}{N} \cdot \frac{P}{B} \xrightarrow{N \to \infty} \frac{2P}{B}$$

의 한계 데이터 전송량을 scatter-reduce $(N-1)$ + all-gather $(N-1)$ steps, 각 step 당 $P/N$ bytes 의 transfer 로 도달함을 Patarasuk & Yuan (2009) 으로부터 한 줄씩 유도할 수 있는 것은 다르다.


ZeRO 의 stage 1/2/3 를 단계로 아는 것 과, Adam optimizer 의 메모리

$$M_\mathrm{Adam}(\psi) = \underbrace{2\psi}_\text{params (FP16)} + \underbrace{2\psi}_\text{grads (FP16)} + \underbrace{12\psi}_\text{FP32 master + m + v} = 16\psi$$

에서 ZeRO-{1,2,3} 단계별 per-rank 메모리

$$M_k(\psi, N) = 2\psi + \frac{(2 + K)\psi}{N},\ \ \frac{2\psi}{N} + \frac{(2 + K)\psi}{N},\ \ \frac{(4 + K)\psi}{N}$$

를 Rajbhandari et al. (2020) 의 메모리 분석으로 직접 유도할 수 있는 것은 다르다.

Pipeline Parallelism 의 bubble 을 "비효율"로 아는 것 과, GPipe (Huang 2019) 의

$$\text{bubble ratio} = \frac{P-1}{P-1+M},\quad M \geq 4P \implies \text{bubble} < 20%$$

을 timeline diagram 으로부터 유도하고, 1F1B schedule 이 activation memory 를

$$O(M \cdot A) \xrightarrow{\text{1F1B}} O(P \cdot A)$$

로 줄이는 구조 (Narayanan 2019) 를 따라가는 것은 다르다.

Megatron 의 Tensor Parallelism 에서 column-parallel Y = X[W_1 | W_2] 와 row-parallel Y = [X_1 | X_2][W_1; W_2] 의 forward·backward 통신 패턴 차이를 유도하고, MLP block 당 단 2 번의 AllReduce (forward + backward) 로 끝나는 이유 (Shoeybi 2019) 를 알고 쓰는 것은 다르다.

FSDP 가 ZeRO-3 와 같은지 다른지, PyTorch native vs DeepSpeed 의 선택 기준이 forward prefetch · backward overlap · mixed precision 통합 · gradient checkpointing 친화성에서 어떻게 갈리는지 알고 고르는 것은 다르다.


다루는 시스템·알고리즘 (구현 계보순)

Thakur 2005 Bandwidth-optimal AllReduce · Patarasuk & Yuan 2009 Ring AllReduce · Li 2014 Parameter Server · Goyal 2017 Linear Scaling · McCandlish 2018 Gradient Noise Scale · Shoeybi 2019 Megatron-LM TP · Huang 2019 GPipe · Narayanan 2019 PipeDream / 1F1B · Rajbhandari 2020 ZeRO · Lepikhin 2021 GShard MoE · Li & Hoefler 2021 Chimera · Narayanan 2021 3D Parallelism · Ren 2021 ZeRO-Offload · Rajbhandari 2021 ZeRO-Infinity · Korthikanti 2022 Selective Recomputation + Sequence Parallelism · Zhao 2023 PyTorch FSDP


핵심 질문

대규모 분산 학습의 4 축 — Collective Communication · Data Parallelism · Model/Tensor Parallelism · Pipeline Parallelism — 은 왜 모두 "메모리·통신·연산의 trade-off 관리" 의 다른 구현이고, Ring AllReduce 의 $2(N-1)/N$ · Megatron 의 column/row split · GPipe 의 $M \geq 4P$ · ZeRO-3 의 $1/N$ shard 가 각각 어떤 수학적 동기에서 도출되었는가 — Patarasuk-Yuan 2009 의 bandwidth-optimal 증명부터 LLaMA-급 3D parallelism recipe 까지 한 줄씩 유도합니다.


GitHub Python PyTorch NCCL DeepSpeed Megatron-LM Accelerate Docs Theorems Proofs Reproductions Exercises License


🎯 이 레포에 대하여

분산 학습 자료는 대부분 "DDP 를 쓰면 multi-GPU 가 된다", "ZeRO-3 를 켜면 모델이 메모리에 들어간다", "Megatron 을 쓰면 70B 도 가능하다" 에서 멈춥니다. 하지만 DDP 의 gradient bucket 이 왜 layer-reverse 순서로 구성되어 backward compute 와 overlap 이 가능한지, ZeRO-1/2/3 의 per-rank 메모리 공식 $\psi + (2\psi + K\psi)/N \to (2+K+16)\psi/N$ 가 어떻게 유도되는지, GPipe 의 bubble ratio $(P-1)/(P-1+M)$ 가 왜 $M \geq 4P$ 에서 20% 미만으로 떨어지는지, Megatron 의 column/row 분할이 왜 forward 와 backward 합쳐 단 2 번의 AllReduce 로 끝나는지 — 이런 "왜" 는 제대로 설명되지 않습니다.

일반 자료 이 레포
"DDP 가 자동으로 multi-GPU 학습을 해준다" PyTorch DDP 의 internals — torch.nn.parallel.DistributedDataParallel 의 gradient bucket 이 layer-reverse 순서로 구성되는 이유, backward 진행 중 먼저 완료된 layer 의 grad 가 bucket 에 모여 AllReduce 가 backward compute 와 overlap, bucket_cap_mb (default 25MB) 의 latency vs throughput trade-off, NCCL stream 과 compute stream 의 분리
"Ring AllReduce 가 빠르다" Patarasuk & Yuan 2009 — Scatter-reduce phase $(N-1)$ steps + AllGather phase $(N-1)$ steps, 각 step 당 $P/N$ bytes 전송 → 총 transfer $2(N-1)P/N$ per rank → $N \to \infty$ 에서 $2P$ 점근, bandwidth-optimal ($\alpha\beta$ model 에서 lower bound 와 일치) $\square$. Tree AllReduce 와의 비교: tree 는 latency $O(\log N)$, ring 은 bandwidth optimal — small message → tree, large message → ring
"ZeRO 는 메모리를 줄여준다" Rajbhandari 2020 — Adam FP16 의 per-GPU 메모리 $16\psi$ (params $2\psi$ + grads $2\psi$ + FP32 master $4\psi$ + Adam $m,v$ 각 $4\psi$) 분석. ZeRO-1: optimizer state $K\psi$ 를 $N$ 등분 → $2\psi + 2\psi + K\psi/N$. ZeRO-2: gradient ReduceScatter → $2\psi + 2\psi/N + K\psi/N$. ZeRO-3: parameter 도 shard → $(2 + 2 + K)\psi/N$. 70B + Adam FP16 + 16 GPU 에서 ZeRO-0 은 1120 GB (불가능), ZeRO-3 은 70 GB (가능) $\square$
"GPipe 는 pipeline parallelism" Huang 2019 — Naive pipeline 의 bubble ratio $(P-1)/P$ ($P=8$ 이면 87.5%) 에서, mini-batch 를 $M$ micro-batch 로 분할 → forward $M$ 번 후 backward $M$ 번 → bubble ratio $(P-1)/(P-1+M)$ $\square$. $M = 4P$ 에서 $(P-1)/(5P-1) \approx 20%$. 한계: activation memory $O(M)$ — $M$ 키울수록 OOM 위험. 1F1B (Narayanan 2019) 가 같은 bubble 에서 activation memory $O(P)$ 로 감축
"Megatron 은 tensor parallelism" Shoeybi 2019 — Linear $Y = XW$ 에서 column-parallel $W = [W_1 \mid W_2]$: $Y_i = XW_i$ (각 rank 독립, output dim 분할 유지). row-parallel $W = [W_1; W_2]$, $X = [X_1 \mid X_2]$: $Y = \sum_i X_i W_i$ (AllReduce 필요). MLP block: column → GELU → row → 첫 column 은 sync 없음, GELU 는 elementwise (sync 없음), row 끝에 1 AllReduce → forward + backward 합쳐 block 당 단 2 번 AllReduce $\square$
"FSDP 는 ZeRO-3 의 PyTorch 버전" Zhao 2023 — FSDP 와 ZeRO-3 의 공통점: parameter shard, AllGather on demand, free after use. 차이점: ① Forward prefetch (다음 layer params 를 현재 layer compute 와 overlap), ② backward prefetch + overlap with grad reduce-scatter, ③ mixed precision native integration (MixedPrecision policy), ④ gradient checkpointing 친화 (per-layer checkpoint), ⑤ PyTorch native (custom CUDA kernel 의존도 낮음). DeepSpeed ZeRO-3 vs FSDP 의 throughput 비교 (LLaMA-7B, 8 GPU)
"TP 는 inter-node 에 안 쓴다" Communication 비용 분석 — TP 는 매 layer 당 forward + backward 합쳐 2 번 AllReduce → layer × steps × 2 번. PP 는 stage boundary 에서 P2P send/recv → stage 수 만큼 만. DP 는 step 당 1 번 AllReduce of full gradients. NVLink (600 GB/s) vs InfiniBand (200 Gbps = 25 GB/s) → TP 는 NVLink 에 한정, PP 는 inter-node OK, DP 는 가장 바깥. 3D parallelism 의 axis 배치 heuristic
"Activation memory 는 그냥 큰 값" Korthikanti 2022 — Transformer per-layer activation $\approx s \cdot b \cdot h \cdot (34 + 5 \cdot a \cdot s/h)$ ($s$ seq len, $b$ batch, $h$ hidden, $a$ heads). long-context 에서 $O(s^2)$ 항이 지배. Gradient checkpointing (Chen 2016) 으로 $O(\sqrt{L})$ stage-based + 33% recompute, selective recomputation 으로 expensive (matmul) save / cheap (norm, dropout) recompute, Sequence Parallelism (Megatron) 으로 LayerNorm·Dropout activation 을 sequence axis 로 shard → $O(s/\text{TP})$
"MoE 는 expert 가 많은 모델" Lepikhin 2021 (GShard) · Fedus 2022 (Switch Transformer) — Top-$k$ routing $\to$ token 을 expert 로 dispatch (all-to-all communication), expert 가 GPU 분산 (expert parallelism). Load balancing loss $\mathcal{L}_\text{aux} = \alpha \cdot E \cdot \sum_i f_i \cdot P_i$ — 모든 expert 가 균등하게 사용되도록. DeepSpeed-MoE 의 expert × DP × MP 3 축 sharding
기법의 나열 NumPy + PyTorch DDP + FSDP + DeepSpeed ZeRO + Megatron-LM TP/PP 로 Ring AllReduce 의 step-by-step 시뮬레이션 · DDP gradient bucket overlap timeline 측정 · ZeRO-{0,1,2,3} per-GPU 메모리 직접 계측 · Megatron MLP 의 forward AllReduce 횟수 검증 · GPipe vs 1F1B 의 activation memory 비교 · 2-GPU 에서 FSDP vs DDP throughput 측정 · 3D parallelism degree 선택 heuristic 까지 직접 구현해 수학적 주장을 눈으로 확인

📌 선행 레포 & 후속 방향

[PyTorch Internals Deep Dive]   ─┐
[Optimization Theory Deep Dive] ─┤
[LLM Pretraining Deep Dive]     ─┼─►  이 레포  ──► [Efficient ML Deep Dive]
[Linear Algebra Deep Dive]      ─┤   "왜 DDP·TP·PP·ZeRO 가             Kernel fusion / Quantization
[CNN Deep Dive] (권장)          ─┘    메모리·통신·연산의                / FlashAttention / Inference
                                      다른 trade-off 인가"
         │
         ├── [PyTorch Internals]    CUDA stream · async · dispatcher → Ch1, Ch2
         ├── [Optimization]         SGD · Adam · batch size scaling → Ch2, Ch5
         ├── [LLM Pretraining]      Scale up recipe · 7B → 400B → Ch7
         ├── [Linear Algebra]       Matrix partition · block decomposition → Ch3
         └── [CNN]                  Memory hierarchy · roofline → Ch6

⚠️ 선행 학습 필수: 이 레포는 PyTorch Internals Deep Dive (CUDA stream, async kernel launch, autograd engine, torch.distributed backend), Optimization Theory Deep Dive (SGD, Adam, momentum, batch size scaling laws), LLM Pretraining Deep Dive (FLOPs accounting, 모델 scale, Chinchilla), Linear Algebra Deep Dive (matrix partition, block matrix multiplication) 를 선행 지식으로 전제합니다. CNN Deep Dive (memory hierarchy, roofline model) 는 Ch6 의 activation memory 분석에서 권장됩니다.

💡 이 레포의 핵심 기여: Chapter 1 (Collective Communication) 과 Chapter 5 (ZeRO/Sharding) 는 현대 분산 학습을 이해하는 두 핵심 축입니다. 전자는 "왜 Ring AllReduce 가 bandwidth-optimal 인가" 의 통신 이론적 토대 (DDP/FSDP/ZeRO 가 모두 그 응용), 후자는 "왜 parameter·gradient·optimizer state 를 따로 sharding 하는가" 의 메모리 회계적 동기 (FSDP·DeepSpeed·Megatron 의 모든 설계가 그 귀결) 를 다룹니다. 이 두 축을 완전히 이해한 후 Chapter 3 (TP), Chapter 4 (PP), Chapter 7 (3D + MoE) 을 읽으면 LLaMA·GPT-3 급 학습 recipe 의 설계 결정 맥락이 선명해집니다.

🟡 이 레포의 성격: 여기서 다루는 일부 주제 — FSDP 가 DeepSpeed ZeRO-3 를 대체할 것인가, Sequence Parallelism 의 1M context 한계, MoE routing 의 load imbalance 해결, Elastic training 의 production 적용 — 는 현재 진행 중인 시스템 연구 영역 입니다. 레포는 "정답" 이 아니라 "고전 collective ops 와 현대 LLM 분산 학습 사이의 지도" 를 제공합니다.


🚀 빠른 시작

각 챕터의 첫 문서부터 바로 학습을 시작하세요!

Ch1 Ch2 Ch3 Ch4 Ch5 Ch6 Ch7


📚 전체 학습 지도

💡 각 챕터를 클릭하면 상세 문서 목록이 펼쳐집니다


🔹 Chapter 1: Collective Communication 의 수학

핵심 질문: Broadcast · Scatter · Gather · AllGather · ReduceScatter · AllReduce 의 MPI 전통적 분류는 어떻게 구성되는가? AllReduce 의 4 가지 구현 (Naive · Recursive Halving · Ring · Tree) 중 Ring 이 왜 large message 에서 bandwidth-optimal 인가? Patarasuk & Yuan (2009) 의 scatter-reduce $(N-1)$ + all-gather $(N-1)$ = $2(N-1)P/N$ 증명은? NCCL 이 어떻게 NVLink/PCIe/InfiniBand 토폴로지를 인지하여 ring vs tree 를 자동 선택하는가? Bandwidth-bound vs latency-bound 의 message size threshold 는?

Collective ops 분류부터 NCCL 토폴로지까지 (5개 문서)
문서 핵심 정리·증명·재현
01. Collective Operation 의 분류 정의: Broadcast (root → all), Scatter (root split → all), Gather (all → root collected), AllGather (all → all gathered), ReduceScatter (reduce + scatter), AllReduce ≡ ReduceScatter + AllGather. MPI 전통: MPI_Bcast, MPI_Scatter, MPI_Gather, MPI_Allgather, MPI_Reduce_scatter, MPI_Allreduce. 각각의 사용 사례 (broadcast: model init, AllReduce: gradient sync, ReduceScatter: ZeRO-2 grad shard, AllGather: ZeRO-3 param fetch)
02. AllReduce 의 4 가지 구현 비교 Naive (1-to-all + broadcast): 통신량 $O(NP)$, 단순하지만 비효율. Recursive Halving/Doubling (Thakur 2005): $\log N$ 단계, 통신량 $O(P\log N)$. Ring (Patarasuk-Yuan 2009): $2(N-1)P/N$, bandwidth-optimal. Tree-based: latency $O(\log N)$, small message 유리. Trade-off table: bandwidth (per-rank) vs latency vs simplicity
03. Ring AllReduce 의 수학 (Patarasuk & Yuan 2009) 정리: Ring AllReduce 는 (a) scatter-reduce phase $(N-1)$ steps — 각 rank 가 자신의 chunk 를 next rank 로 보내고 partial reduce, (b) all-gather phase $(N-1)$ steps — 완성된 chunk 를 ring 따라 전파. 각 step 당 $P/N$ bytes 전송 → 총 per-rank transfer $2(N-1)P/N$ $\square$. Bandwidth-optimal proof: $\alpha\beta$ model 의 lower bound $\geq 2(N-1)P/N$ 와 정확히 일치
04. NCCL 과 NVLink 토폴로지 NCCL (NVIDIA Collective Communications Library): ring/tree 자동 선택, multi-node multi-GPU 지원. Topology-aware routing: NVLink (300-900 GB/s, intra-node), PCIe Gen4/Gen5 (32-64 GB/s), InfiniBand HDR/NDR (200-400 Gbps). NCCL_DEBUG=INFO 로 실제 ring 구성 관찰. DGX H100 의 8-GPU NVSwitch 전부 연결 — ring 1 step 당 NVLink 직통
05. Bandwidth vs Latency Bound $\alpha\beta$ model: 한 step transfer 시간 $= \alpha + \beta P$ ($\alpha$ latency, $\beta$ inverse bandwidth). Small message ($P &lt;&lt; \alpha/\beta$): latency-dominated → tree ($O(\log N)$ steps). Large message: bandwidth-dominated → ring ($2(N-1)P/(NB)$). Threshold: 보통 NVLink 에서 수 MB. NCCL 의 자동 algorithm 선택이 이 threshold 를 활용

🔹 Chapter 2: Data Parallelism 과 PyTorch DDP

핵심 질문: Data Parallelism 이 어떻게 mini-batch SGD 와 수학적으로 동등한가? DDP 의 gradient bucket 이 왜 layer-reverse 순서로 구성되어 backward compute 와 overlap 이 가능한가? bucket_cap_mb 의 latency vs throughput trade-off 는? Gradient accumulation 의 no_sync() context 가 왜 AllReduce 를 생략하는가? Linear scaling rule (Goyal 2017) 의 warmup 이 왜 발산을 막는가? Async Parameter Server (Li 2014) 의 stale gradient 가 왜 sync 보다 느린 수렴을 보이는가?

DP 기본부터 Async Hogwild! 까지 (6개 문서)
문서 핵심 정리·증명·재현
01. Data Parallelism 의 기본 정의: same model replicated, batch $B$ 를 $N$ rank 로 분할 ($b = B/N$ per rank), 각 rank 가 forward + backward 후 gradient AllReduce + averaging → optimizer step. 정리: synchronous DP 는 batch size $B$ 의 standard SGD 와 정확히 동등 — gradient mean 의 unbiasedness $\square$. 단점: model 전체 복사 → 메모리 비효율 (ZeRO 의 동기)
02. PyTorch DistributedDataParallel (DDP) API: torch.distributed.init_process_group('nccl', rank, world_size), model = DistributedDataParallel(model, device_ids=[rank]). Backend: NCCL (GPU), Gloo (CPU/cross-platform), MPI (legacy). Launch: torchrun --nproc_per_node=N 또는 torch.multiprocessing.spawn. Sampler: DistributedSampler 로 rank 마다 다른 subset, epoch 마다 reshuffle
03. Gradient Bucket 과 Backward Overlap 메커니즘: backward pass 가 layer 역순으로 진행 → 먼저 완료된 layer 의 gradient 를 bucket 에 쌓다가 일정 크기 (bucket_cap_mb, default 25MB) 도달 시 AllReduce 비동기 launch → 이후 layer 의 backward compute 와 NCCL stream 이 overlap $\square$. Trade-off: 작은 bucket → low latency, 많은 launch overhead / 큰 bucket → 낮은 overlap. Profiler trace 로 직접 측정
04. Gradient Accumulation 과 no_sync 목적: GPU 메모리 한계로 large batch 가 안 들어갈 때 micro-batch 여러 번 backward 후 한 번 step → effective batch = micro × accumulation. with model.no_sync(): loss.backward() context 에서 AllReduce 생략 → 마지막 micro 만 sync. 수학적 동등성: gradient 의 linearity → accumulation 후 AllReduce 가 매 micro 별 AllReduce 와 같음. BN/LN 의 running stats 가 sync 시점에 영향
05. Batch Size Scaling Rules Linear scaling (Goyal 2017): $\eta \propto B$, ImageNet 에서 batch 8K 까지 검증. Square-root scaling: $\eta \propto \sqrt{B}$, AdamW 류에 적합. Warmup: 첫 수천 step 에서 $\eta$ 를 0 → target 으로 ramp → 발산 방지 (Goyal 의 핵심 기여). Critical batch size (McCandlish 2018): gradient noise scale $\mathcal{B}\text{noise} = \text{tr}(\Sigma)/|\nabla L|^2$, $B > \mathcal{B}\text{noise}$ 에서 returns diminish
06. Async Updates — Hogwild! 와 Parameter Server Hogwild! (Recht 2011): lock-free async SGD, sparse update 에서 수렴 보장. Parameter Server (Li 2014): worker 가 gradient push, server 가 param pull, stale gradient 문제 — worker 가 보낸 grad 가 도착했을 땐 param 이 이미 다음 step. Staleness bound 와 elastic averaging (Zhang 2015) 의 보정. 현재: sync DDP 가 주류 — bandwidth 발전으로 async 의 wall-clock 이점 감소

🔹 Chapter 3: Model Parallelism 과 Tensor Parallelism

핵심 질문: Single GPU 메모리에 모델이 안 들어갈 때 naive layer 분할이 왜 sequential bubble 로 실패하는가? Tensor Parallelism 의 column-parallel $W = [W_1 \mid W_2]$ 와 row-parallel $W = [W_1; W_2]$ 의 forward·backward 통신 패턴 차이는? Megatron 의 MLP 가 왜 column → GELU → row 순서로 단 1 번의 AllReduce 만 필요한가? Attention 의 QKV column-parallel + output row-parallel 이 왜 multi-head 구조에 자연스러운가? TP degree 가 NVLink intra-node 에 한정되는 이유는?

Naive MP 부터 Megatron Attention TP 까지 (5개 문서)
문서 핵심 정리·증명·재현
01. Model Parallelism 의 동기와 Naive 한계 동기: 70B 모델 = 140GB (FP16) → A100 80GB 1 장에 안 들어감. DP 만으론 해결 불가 (each rank 가 full model 보유). Naive MP: layer 별 GPU 할당, sequential 실행 → bubble ratio $(P-1)/P$ (8 GPU 면 87.5%) — 거의 사용 안 함. Activation 통신 비용 $O(B \cdot L \cdot d)$ per stage boundary. PP 의 micro-batching 동기 (Ch4), TP 의 intra-layer 분할 동기 (이 챕터)
02. Tensor Parallelism — Column / Row Split (Shoeybi 2019) Column-parallel: $W \in \mathbb{R}^{d_\text{in} \times d_\text{out}}$, $W = [W_1 \mid W_2]$ (output dim 분할), $Y_i = X W_i$ — 각 rank 가 input 전체를 받고 output 의 일부 계산, forward 시 sync 불필요, backward 시 input gradient 에 AllReduce. Row-parallel: $W = [W_1; W_2]$ (input dim 분할), $X = [X_1 \mid X_2]$, $Y = \sum_i X_i W_i$ — forward 끝에 AllReduce, backward 는 sync 불필요 $\square$
03. Megatron MLP 의 TP 구조 구조: $X \to \text{Linear}(W_1) \to \text{GELU} \to \text{Linear}(W_2) \to Y$. TP 적용: $W_1$ column-parallel ($Y_1$ 의 hidden dim 분할 유지), GELU elementwise (no sync), $W_2$ row-parallel ($Y_1$ 가 이미 row-split 되어 있음) → 끝에 1 번 AllReduce. 결과: forward 1 AllReduce, backward 1 AllReduce → block 당 총 2 AllReduce $\square$. GELU 의 partition 보존이 핵심 trick
04. Megatron Attention 의 TP 구조 구조: QKV projection column-parallel (head 단위 분할 — heads/TP per rank), per-head attention 각 rank 독립 (no sync), output projection row-parallel → 끝에 1 AllReduce. 자연스러움: multi-head attention 자체가 head 별 병렬 계산 → TP degree = head 수의 약수 권장. 재현: GPT-2 small 에서 TP=2 vs TP=1 의 attention output 동치 검증
05. Tensor Parallelism 의 통신 비용 분석 per-block cost: forward 2 AllReduce (MLP 1 + Attention 1), backward 2 AllReduce → layer 당 4 AllReduce. layer 수 $L$, step 수 $S$ → 총 $4LS$ AllReduce. NVLink intra-node 한정: A100 NVLink 600 GB/s vs InfiniBand 25 GB/s, TP=8 inter-node 시 통신 시간이 compute 를 지배. PP 와의 trade-off: TP 는 compute-communication ratio 가 낮음 → high-bandwidth 만

🔹 Chapter 4: Pipeline Parallelism

핵심 질문: Naive pipeline 의 bubble ratio $(P-1)/P$ 가 왜 micro-batching 으로 $(P-1)/(P-1+M)$ 로 감소하는가 (GPipe, Huang 2019)? $M = 4P$ 가 bubble < 20% 의 표준 heuristic 인 이유는? 1F1B (PipeDream, Narayanan 2019) 가 같은 bubble 에서 activation memory 를 $O(M)$ 에서 $O(P)$ 로 줄이는 schedule 의 메커니즘은? Interleaved 1F1B (Megatron-LM) 가 어떻게 bubble 을 더 감소시키는가? Chimera (Li 2021) 의 bidirectional pipeline 이 어떻게 양쪽 끝에서 동시 시작하여 bubble 반감하는가?

Naive bubble 부터 Chimera bidirectional 까지 (5개 문서)
문서 핵심 정리·증명·재현
01. Naive Pipeline 의 Bubble 문제 구조: layer 를 $P$ stage 로 분할, stage 1 → 2 → ... → P 순차 forward, 그 다음 backward. Bubble: 한 시점에 한 GPU 만 active → idle ratio $(P-1)/P$. $P=8$ 이면 bubble 87.5% — 거의 사용 불가. 불용 이유: GPU $$$/hour 비용 vs idle 시간. PP 의 핵심 과제는 이 bubble 제거
02. GPipe 와 Micro-batching (Huang 2019) 아이디어: mini-batch 를 $M$ micro-batch 로 분할, forward 를 $M$ 번 → backward 를 $M$ 번 (synchronous schedule). 정리: bubble ratio $= (P-1) / (P-1+M)$ $\square$. $M = 4P \Rightarrow$ ratio $\approx 20%$, $M = 8P \Rightarrow 11%$. 단점: activation memory $O(M)$ — $M$ 키울수록 OOM. Forward 끝나야 backward 시작 → activation 모두 보유
03. 1F1B Schedule (PipeDream, Narayanan 2019) 구조: forward 와 backward 를 교대 실행 — 각 stage 가 micro-batch $i$ forward 후 micro-batch $i-P+1$ backward. 메모리: 각 stage 가 동시에 $\leq P$ 개 micro-batch activation 보유 → $O(P \cdot A)$ vs GPipe $O(M \cdot A)$ $\square$. Bubble: GPipe 와 동일 (synchronous variant). Async 1F1B: weight stash 로 stale weight 사용, async 변형 (덜 사용)
04. Interleaved 1F1B (Megatron-LM) 구조: 각 GPU 가 여러 stage 담당 — virtual stage = $V$ → physical stage = $P$ → effective $VP$ stage. Bubble: $(P-1)/(VM)$ 로 더 감소. Trade-off: virtual stage 간 P2P 통신 증가, activation memory $O(V \cdot P)$ 로 증가. Megatron-LM 표준 schedule, GPT-3 175B 학습에 사용
05. Chimera (Li 2021), 2D/3D Pipeline Bidirectional pipeline: 양쪽 끝 (stage 0, stage $P-1$) 에서 동시에 micro-batch 주입 → bubble 반감. 메모리 trade-off: 각 stage 가 두 방향 activation 보유 → 2× memory. 2D/3D pipeline: pipeline 자체를 다축으로 분할, very-large scale 에서 활용

🔹 Chapter 5: ZeRO 와 Sharding (FSDP)

핵심 질문: Adam FP16 의 per-GPU 메모리 $16\psi$ 분해 (params $2\psi$ + grads $2\psi$ + FP32 master $4\psi$ + Adam $m, v$ $8\psi$) 는 왜 $K=12$ 의 optimizer multiplier 를 갖는가? ZeRO-1/2/3 의 per-rank 메모리 공식이 어떻게 단계적으로 감소하는가? ZeRO-3 의 forward/backward 시 layer 별 AllGather → free 메커니즘이 왜 추가 통신 비용을 감수하더라도 메모리 우선인가? FSDP 와 ZeRO-3 의 차이 (forward prefetch · mixed precision · gradient checkpointing 통합) 는? ZeRO-Offload (Ren 2021) 가 CPU·NVMe 까지 가는 이유와 PCIe bandwidth 병목은?

ZeRO motivation 부터 ZeRO-Infinity 까지 (6개 문서)
문서 핵심 정리·증명·재현
01. ZeRO 의 Motivation 과 메모리 분해 (Rajbhandari 2020) DP 의 메모리 비효율: 각 rank 가 full model + grad + optimizer state 보유. Adam FP16 분해 ($\psi$ params): params (FP16) $2\psi$, grads (FP16) $2\psi$, FP32 master $4\psi$, Adam $m$ (FP32) $4\psi$, Adam $v$ (FP32) $4\psi$ → total $16\psi$. 70B 예: $\psi = 70 \times 10^9$ → $1120$ GB per GPU — A100 80GB 14× 초과. ZeRO 의 동기: replicated state 의 sharding
02. ZeRO-1 — Optimizer State Sharding 아이디어: optimizer state ($K\psi = 12\psi$) 만 $N$ rank 로 1/N shard, params 와 grads 는 replicated. Per-rank memory: $2\psi + 2\psi + K\psi/N$. 통신: backward 후 grad AllReduce 그대로, optimizer step 후 updated params 를 AllGather 로 broadcast. 70B / 16 GPU: $4\psi + K\psi/N = 280 + 52.5 = 332.5$ GB / GPU → 여전히 부족
03. ZeRO-2 — Gradient Sharding 추가 아이디어: gradient 도 ReduceScatter 로 각 rank 가 자기 shard 만 보유. Per-rank memory: $2\psi + 2\psi/N + K\psi/N$. 통신: AllReduce → ReduceScatter (grad 한 번 + AllGather params 한 번 = 같은 비용). 70B / 16 GPU: $140 + 17.5 + 52.5 = 210$ GB → 여전히 부족하지만 ZeRO-1 보다 122 GB 절약
04. ZeRO-3 — Parameter Sharding 아이디어: params 도 1/N shard, forward/backward 시 필요한 layer 만 AllGather → 사용 후 free. Per-rank memory: $(2 + 2 + K)\psi/N = 16\psi/N$. 70B / 16 GPU: $70$ GB / GPU → 들어감! 통신 overhead: forward 1 AllGather/layer + backward 1 AllGather/layer + ReduceScatter grad → naive 2.5× compute 시간 증가 가능 (prefetch 로 완화)
05. FSDP — PyTorch Fully Sharded Data Parallel (Zhao 2023) FSDP: PyTorch native ZeRO-3 변형. 차이점: ① forward prefetch (다음 layer params AllGather 를 현재 compute 와 overlap), ② backward prefetch + grad ReduceScatter overlap, ③ MixedPrecision policy (param FP16, grad FP32 등 fine-grained), ④ auto_wrap_policy (transformer block 단위 wrap), ⑤ gradient checkpointing 친화 (per-block checkpoint). DeepSpeed ZeRO-3 vs FSDP 의 throughput 비교
06. ZeRO-Offload 와 ZeRO-Infinity ZeRO-Offload (Ren 2021): optimizer state + grad 를 CPU 로 offload, GPU-CPU PCIe (Gen4 32 GB/s) bandwidth 병목. small-to-medium model 에서 GPU 1 대로도 큰 모델 학습 가능. ZeRO-Infinity (Rajbhandari 2021): NVMe SSD 까지 활용, 거의 무한 capacity, NVMe bandwidth (PCIe Gen4 NVMe ~7 GB/s) 가 한계. 언제 사용: GPU 메모리 부족하지만 budget constraint, MLPerf 30T params 보고

🔹 Chapter 6: Activation Memory 최적화

핵심 질문: Transformer 의 per-layer activation $\sim s b h (34 + 5as/h)$ 가 long-context 에서 $O(s^2)$ 로 어떻게 폭증하는가? Gradient checkpointing (Chen 2016) 의 stage-based recompute 가 왜 memory $O(\sqrt{L})$ + 33% recompute 의 sweet spot 인가? Selective recomputation (Korthikanti 2022) 이 expensive (matmul) save / cheap (norm, dropout) recompute 로 어떻게 best-of-both 인가? Sequence Parallelism 이 LayerNorm·Dropout activation 을 sequence axis 로 shard 하여 memory $O(s)$ 에서 $O(s/\text{TP})$ 로 줄이는 메커니즘은?

Activation 크기 분석부터 Sequence Parallelism 까지 (4개 문서)
문서 핵심 정리·증명·재현
01. Activation Memory 의 크기 per-layer Transformer: $s b h (34 + 5 a s / h)$ bytes — $s$ seq len, $b$ batch, $h$ hidden dim, $a$ heads. Attention 의 $O(s^2)$: softmax score matrix $b \cdot a \cdot s^2$ 가 long-context 에서 지배. GPT-3 175B (96 layer, 2K seq): 13 TB activation — 모델 weight 350 GB 보다 훨씬 큼. 장기 context: 100K seq 에선 weights 의 100× 가능 → activation 이 핵심 병목
02. Gradient Checkpointing (Chen 2016) 아이디어: 모든 activation 저장 대신 일부 (stage boundary) 만 저장, backward 시 forward 를 재실행하여 activation 복원. Memory: $O(\sqrt{L})$ stage-based ($\sqrt{L}$ stage × $\sqrt{L}$ act/stage). Compute: forward 1 회 추가 → ~33% overhead. PyTorch: torch.utils.checkpoint.checkpoint(fn, ...). transformer block 별 wrap 표준
03. Selective Activation Recomputation (Korthikanti 2022) 아이디어: 모든 activation recompute 가 아니라 저장 비용 大 + recompute 저렴 한 것 (norm output, dropout mask) 만 recompute, 나머지 (matmul output, attention score) save. 결과: full recompute 대비 compute overhead 5%, memory full save 대비 5× 절감 → Pareto sweet spot. Megatron-LM 의 표준
04. Sequence Parallelism (Korthikanti 2022) 아이디어: TP 가 attention/MLP activation 은 hidden dim 분할로 $1/\text{TP}$ 이지만, LayerNorm·Dropout activation 은 hidden 전체 보유 → $O(s b h)$. Sequence Parallelism: 이 부분 activation 을 sequence axis 로 shard → $O(s b h / \text{TP})$. 통신: TP 와 SP 의 boundary 에서 AllGather/ReduceScatter 추가, 그러나 AllReduce = AllGather + ReduceScatter 이므로 총 통신량 동일 $\square$. long-context 학습의 핵심

🔹 Chapter 7: 3D Parallelism 과 실전

핵심 질문: DP × TP × PP = world_size 의 3 축 배치에서 TP 는 왜 intra-node (NVLink), PP 는 inter-node (InfiniBand), DP 는 가장 바깥의 heuristic 이 표준이 되었는가? LLaMA-70B 를 64 GPU 로 학습할 때 TP=8, PP=4, DP=2 의 degree 선택 근거는? MoE 의 expert parallelism 이 token → expert dispatch 를 all-to-all 로 처리하는 메커니즘은? Synchronous checkpoint vs sharded checkpoint vs asynchronous checkpoint 의 trade-off 는? TorchElastic 의 preemption-tolerant 메커니즘이 어떻게 node 추가·제거를 지원하는가?

3D 통합부터 Elastic Training 까지 (4개 문서)
문서 핵심 정리·증명·재현
01. 3D Parallelism 통합 — DP × TP × PP 공식: world_size = DP × TP × PP. Heuristic: TP intra-node (NVLink, 매 layer 4 AllReduce), PP inter-node (stage boundary 만 P2P), DP 가장 바깥 (AllReduce overlap with backward). LLaMA-70B / 64 GPU: TP=8 (8-GPU NVSwitch 1 node), PP=4 (4 nodes), DP=2 — 각 axis degree 가 모델/네트워크 topology 와 일치. Megatron-DeepSpeed 가 표준 stack
02. MoE Parallelism — Expert Parallelism MoE: gating network → top-$k$ expert 선택 → expert 가 GPU 분산. All-to-all communication: token → expert dispatch + expert output → token gather. Load balancing loss: $\mathcal{L}_\text{aux} = \alpha \cdot E \cdot \sum_i f_i \cdot P_i$ ($f_i$ rank-$i$ expert 의 token 비율, $P_i$ gating prob). DeepSpeed-MoE / GShard / Switch Transformer: expert × DP × MP 3 축 sharding
03. Efficient Checkpointing 과 Resume Synchronous all-rank: 모든 rank 가 동시에 자기 shard 저장 → high bandwidth 필요. Sharded checkpoint: 각 rank 자기 shard 만 저장, resume 시 정확히 같은 topology 필요. Asynchronous checkpointing: I/O 와 compute overlap, NVIDIA Magnum IO. Elastic checkpoint: 다른 topology 로 resume 가능 (re-shard) — DCP (Distributed Checkpoint)
04. Failure Recovery — Elastic Training TorchElastic / torchrun --max_restarts: node failure 시 자동 재시작. Preemption-tolerant: spot instance 활용 가능. Dynamic membership: 학습 중 node 추가·제거 (rendezvous 재계산). 재현: 4 GPU 학습 중 1 GPU kill → checkpoint 에서 복구, 같은 step 으로 resume 검증

✅ 2026-05 — 35개 문서 전체 작성 완료: Ch1 (Collective) ~ Ch7 (3D · MoE · Elastic) 까지 모든 챕터의 11-섹션 골격 (🎯 핵심 질문 → 🔍 동기 → 📐 선행 → 📖 직관 → ✏️ 정의 → 🔬 정리·증명 → 💻 구현 검증 → 🔗 실전 활용 → ⚖️ 한계 → 📌 핵심 정리 → 🤔 연습문제) 이 일관되게 적용되었습니다. 각 문서 하단에는 ◀ 이전 / 📚 README / 다음 ▶ 네비게이션이 있어, README 진입점에서 35개 문서를 순차적으로 따라갈 수 있습니다. 핵심 증명 (Ring AllReduce bandwidth-optimal, ZeRO-{0,1,2,3} 메모리 공식, Megatron MLP 의 2-AllReduce, GPipe bubble ratio, 1F1B activation $O(P)$, Sequence Parallelism AllReduce = AllGather + ReduceScatter, MoE Load Balancing uniform minimizer, Elastic efficiency $\eta = T/(T + TR/M)$) 이 모두 $\square$ 로 종결되는 형식으로 수록되었습니다.

🏆 핵심 정리 인덱스

이 레포에서 완전한 증명 또는 원 논문 실험 재현 을 제공하는 대표 결과 모음입니다. 각 챕터 문서에서 $\square$ 로 종결되는 엄밀한 증명 또는 results/ 하의 throughput·memory plot 을 확인할 수 있습니다.

정리·결과 서술 출처 문서
Ring AllReduce Bandwidth-Optimal Per-rank transfer $2(N-1)P/N$, $\alpha\beta$ lower bound 와 일치 Ch1-03
AllReduce 알고리즘 trade-off Naive $O(NP)$ · Recursive $O(P\log N)$ · Ring $2(N-1)P/N$ · Tree latency-optimal Ch1-02
DP ↔ Mini-batch SGD 동등성 Synchronous DP 의 gradient mean = standard SGD with batch $B = N b$ Ch2-01
Gradient Bucket Backward Overlap Layer-reverse bucket → AllReduce 비동기 launch, NCCL stream 과 compute overlap Ch2-03
Linear Scaling + Warmup (Goyal 2017) $\eta \propto B$ + warmup → ImageNet batch 8K 까지 monotone improvement Ch2-05
Column / Row Linear Partition $W = [W_1 \mid W_2]$ 와 $W = [W_1; W_2]$ 의 forward·backward 통신 패턴 분리 Ch3-02
Megatron MLP 의 2 AllReduce 정리 Column-GELU-Row 구조 → block 당 forward 1 + backward 1 = 총 2 AllReduce Ch3-03
GPipe Bubble Ratio Bubble $= (P-1)/(P-1+M)$, $M = 4P$ 에서 ~20% Ch4-02
1F1B Activation Memory $O(P)$ Forward-Backward 교대 schedule 로 activation $O(M) \to O(P)$ Ch4-03
ZeRO-{0,1,2,3} Memory Formula $16\psi \to 4\psi + K\psi/N \to 2\psi + (2+K)\psi/N \to (4+K)\psi/N$ Ch5-{01-04}
FSDP vs ZeRO-3 차이 Forward prefetch · backward overlap · MixedPrecision · gradient checkpoint 통합 Ch5-05
Gradient Checkpointing $O(\sqrt{L})$ Stage-based recompute → memory $O(\sqrt{L})$ + 33% compute Ch6-02
Selective Recomputation Pareto Expensive save / cheap recompute → 5% compute · 5× memory 절감 Ch6-03
Sequence Parallelism 통신 동등성 TP+SP 의 AllGather + ReduceScatter = TP 의 AllReduce 와 통신량 동일 Ch6-04
3D Parallelism Axis Heuristic TP intra-node, PP inter-node, DP 가장 바깥 — bandwidth × overlap 분석 Ch7-01
MoE Load Balancing Auxiliary Loss $\mathcal{L}_\text{aux} = \alpha E \sum_i f_i P_i$ 가 expert utilization uniform 유도 Ch7-02
Elastic Training Rendezvous Node 동적 추가·제거 시 rendezvous 재계산 → topology 변경 후 resume Ch7-04

💡 챕터별 문서·정리/정의 수 (실측):

챕터 문서 수 정리·정의
Ch1 Collective Communication 5 38
Ch2 Data Parallelism · DDP 6 45
Ch3 Tensor Parallelism 5 39
Ch4 Pipeline Parallelism 5 41
Ch5 ZeRO · FSDP 6 50
Ch6 Activation Memory 4 36
Ch7 3D Parallelism · MoE 4 36
합계 35 285

추가로 130+ 엄밀한 $\square$ 증명 + 105 연습문제 (모두 해설 포함) + 130+ NumPy / PyTorch / DeepSpeed / Megatron 실험 코드 (### 실험 N 형식).

Ch6 (Activation) 와 Ch7 (3D · MoE) 은 4 문서 로 구성 — Activation 은 핵심 4 기법에 집중, 3D 는 mature topic 만 다룸 (Chapters 1-5 의 5-6 문서와 의도적 차이).


💻 실험 환경

모든 챕터의 실험은 아래 환경에서 재현 가능합니다.

# requirements.txt
numpy==1.26.0
scipy==1.11.0
torch==2.1.0                    # DDP / FSDP / torch.distributed
deepspeed==0.12.6               # ZeRO-1/2/3 · Offload · Infinity
transformers==4.36.0            # 참조 모델 (GPT-2, LLaMA-style)
accelerate==0.25.0              # FSDP / DeepSpeed 통합 launcher
matplotlib==3.8.0
seaborn==0.13.0
tqdm==4.66.0
jupyter==1.0.0
# 선택 사항
tensorboard==2.15.0             # 학습 곡선 · timeline
wandb==0.16.0                   # 실험 추적
einops==0.7.0                   # tensor 재구성
nvidia-ml-py==12.535.108        # GPU 메모리 측정 (nvidia-smi 대체)
# 환경 설치 (CUDA 12.1 기준, NCCL 자동 포함)
pip install numpy==1.26.0 scipy==1.11.0 torch==2.1.0 \
            deepspeed==0.12.6 transformers==4.36.0 accelerate==0.25.0 \
            matplotlib==3.8.0 seaborn==0.13.0 tqdm==4.66.0 \
            einops==0.7.0 jupyter==1.0.0 nvidia-ml-py==12.535.108

# 실험 노트북 실행
jupyter notebook

# 분산 실험 launch (예: 2 GPU)
torchrun --nproc_per_node=2 ch2-data-parallel/ddp_minimal.py
deepspeed --num_gpus=2 ch5-zero-sharding/zero_3_minimal.py --deepspeed_config ds_config.json
# 대표 실험 ① — Ring AllReduce 의 통신량 분석 (Ch1-02, Ch1-03)
import numpy as np

def ring_allreduce_bytes(N, P_bytes):
    """
    Ring AllReduce per-rank transfer.
    Scatter-reduce: (N-1) steps × P/N bytes
    AllGather:      (N-1) steps × P/N bytes
    Total per rank: 2(N-1)P/N
    """
    return 2 * (N - 1) * P_bytes / N

def tree_allreduce_bytes(N, P_bytes):
    """Tree-based: log2(N) steps × P bytes (bandwidth-suboptimal for large P)"""
    return np.log2(N) * P_bytes

def naive_allreduce_bytes(N, P_bytes):
    """Naive: each rank sends P to root, root broadcasts P → ~2(N-1)P per rank-0"""
    return 2 * (N - 1) * P_bytes  # root-rank perspective

P = 2 * 10**9  # 1B params × 2 bytes (FP16)
print(f"{'N':>4}  {'Naive':>10}  {'Tree':>10}  {'Ring':>10}")
for N in [2, 4, 8, 16, 64, 256]:
    n_ = naive_allreduce_bytes(N, P) / 1e9
    t_ = tree_allreduce_bytes(N, P) / 1e9
    r_ = ring_allreduce_bytes(N, P) / 1e9
    print(f"{N:>4}  {n_:>8.1f}GB  {t_:>8.1f}GB  {r_:>8.1f}GB")
# N=64: Naive 252GB, Tree 12GB, Ring 3.94GB
# Ring 이 large message 에서 bandwidth-optimal — N → ∞ 에서 2P 점근

# 대표 실험 ② — ZeRO-{0,1,2,3} per-GPU 메모리 (Ch5-01 ~ Ch5-04)
def zero_memory_gb(psi_B, N_gpus, stage):
    """
    psi_B    : params (in B / billions)
    N_gpus   : data-parallel degree
    stage    : 0 (no sharding) / 1 / 2 / 3
    Adam FP16 가정: K = 12 (FP32 master 4 + m 4 + v 4)
    """
    K = 12
    psi = psi_B  # GB scale (params × 2 bytes / 2 = 1 GB per 1B params for FP16)
    if stage == 0:
        return (2 + 2 + K) * psi
    elif stage == 1:
        return 2 * psi + 2 * psi + K * psi / N_gpus
    elif stage == 2:
        return 2 * psi + 2 * psi / N_gpus + K * psi / N_gpus
    elif stage == 3:
        return (2 + 2 + K) * psi / N_gpus

print(f"\n{'Model':>8}  {'GPUs':>6}  {'Z-0':>10} {'Z-1':>10} {'Z-2':>10} {'Z-3':>10}")
for psi_B, N in [(7, 8), (13, 16), (70, 16), (175, 64)]:
    z0 = zero_memory_gb(psi_B, N, 0)
    z1 = zero_memory_gb(psi_B, N, 1)
    z2 = zero_memory_gb(psi_B, N, 2)
    z3 = zero_memory_gb(psi_B, N, 3)
    print(f"{psi_B:>5}B   {N:>5}  {z0:>8.1f}GB {z1:>8.1f}GB {z2:>8.1f}GB {z3:>8.1f}GB")
# 70B / 16 GPU: ZeRO-0 1120GB, ZeRO-1 332.5GB, ZeRO-2 210GB, ZeRO-3 70GB ← A100 80GB 에 들어감

# 대표 실험 ③ — GPipe Bubble Ratio (Ch4-02)
def gpipe_bubble(P_stages, M_micro):
    """Bubble ratio = (P-1) / (P-1+M)"""
    return (P_stages - 1) / (P_stages - 1 + M_micro)

P_stages = 8
print(f"\nP={P_stages} pipeline stages")
for M in [1, 4, 8, 16, 32, 64]:
    b = gpipe_bubble(P_stages, M)
    print(f"  M={M:>3} micro-batches → bubble {b*100:>5.1f}%   activation memory ∝ {M}")
# M=1   → bubble 87.5%  (= naive)
# M=32  → bubble 18.0%  (Huang 의 권장 영역, M = 4P)
# M=64  → bubble 10.0%  (좋지만 activation memory 2× 부담)

# 대표 실험 ④ — DDP minimal (single-node 2 GPU)
import torch, torch.nn as nn, torch.distributed as dist, os

def setup_ddp(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group('nccl', rank=rank, world_size=world_size)

def ddp_minimal(rank, world_size):
    setup_ddp(rank, world_size)
    torch.cuda.set_device(rank)

    model = nn.Sequential(nn.Linear(1024, 4096), nn.GELU(), nn.Linear(4096, 1024)).cuda(rank)
    ddp   = nn.parallel.DistributedDataParallel(model, device_ids=[rank],
                                                bucket_cap_mb=25)  # default; tune for overlap
    opt   = torch.optim.AdamW(ddp.parameters(), lr=1e-4)

    for step in range(20):
        x = torch.randn(64, 1024, device=rank)
        y = ddp(x).sum()
        y.backward()           # AllReduce 자동 (gradient bucket overlap)
        opt.step()
        opt.zero_grad()

    dist.destroy_process_group()

# 대표 실험 ⑤ — Megatron column-parallel / row-parallel Linear (Ch3-02)
class ColumnParallelLinear(nn.Module):
    """W = [W_1 | W_2 | ... | W_TP], output dim 분할"""
    def __init__(self, in_features, out_features, tp_size, tp_rank):
        super().__init__()
        assert out_features % tp_size == 0
        self.local_out = out_features // tp_size
        self.weight = nn.Parameter(torch.randn(self.local_out, in_features) * 0.01)

    def forward(self, x):
        # input x 는 모든 rank 에서 동일 (sync 불필요)
        return x @ self.weight.T  # output: [..., local_out] — 분할 유지

class RowParallelLinear(nn.Module):
    """W = [W_1; W_2; ...; W_TP], input dim 분할"""
    def __init__(self, in_features, out_features, tp_size, tp_rank):
        super().__init__()
        assert in_features % tp_size == 0
        self.local_in = in_features // tp_size
        self.weight = nn.Parameter(torch.randn(out_features, self.local_in) * 0.01)

    def forward(self, x):
        # input x 는 이미 input dim 으로 split 되어 있음
        local_out = x @ self.weight.T
        # output 을 합치려면 AllReduce 필요
        # dist.all_reduce(local_out, op=dist.ReduceOp.SUM)
        return local_out

# 대표 실험 ⑥ — FSDP minimal (Ch5-05)
# from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
# from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
# fsdp_model = FSDP(
#     model,
#     sharding_strategy=ShardingStrategy.FULL_SHARD,    # ZeRO-3 와 동등
#     mixed_precision=MixedPrecision(param_dtype=torch.bfloat16,
#                                    reduce_dtype=torch.float32),
#     auto_wrap_policy=transformer_auto_wrap_policy,
# )

🗺️ 모델 규모별 분산 학습 Recipe

모델 규모 권장 stack TP PP DP ZeRO/FSDP 예시
< 7B DDP 또는 FSDP FULL_SHARD 1 1 $N$ FSDP/ZeRO-2 권장 Llama-7B (8 × A100)
7B – 70B FSDP 또는 TP+ZeRO 1 ~ 8 1 $N/\text{TP}$ FSDP/ZeRO-3 Llama-70B (32-64 × A100)
70B – 400B 3D Parallelism (Megatron-DeepSpeed) 8 4 ~ 8 2 ~ 8 ZeRO-1 (TP 와 결합) GPT-3 175B, Llama-3 405B
400B – 1T+ 3D + MoE 8 8 ~ 16 $N / (\text{TP·PP})$ ZeRO-1 + Expert PaLM 540B, GPT-4-class
1T+ MoE Expert Parallelism + 3D 8 8 $\geq 8$ ZeRO-1 + Expert × MP Mixtral, GShard

💡 3D parallelism axis 배치 heuristic:

  • TP: NVLink intra-node 한정 (degree ≤ NVLink-connected GPU 수, 보통 8). 매 layer 4 AllReduce → bandwidth 의존.
  • PP: inter-node OK (stage boundary 만 P2P). degree 가 클수록 bubble 위험 → micro-batch $M \geq 4P$ 필수.
  • DP: 가장 바깥 (AllReduce overlap with backward). global batch size 와 critical batch size (McCandlish) 의 trade-off.

📊 Megatron-LM vs DeepSpeed vs FSDP 비교:

항목 Megatron-LM DeepSpeed PyTorch FSDP
TP 지원 ✅ Native (Megatron core) ✅ (Megatron-DeepSpeed 통합) ❌ (native 미지원)
PP 지원 ✅ 1F1B / Interleaved ✅ 1F1B ❌ (third-party 필요)
DP 지원 ✅ ✅ (with ZeRO) ✅
ZeRO-3 / FSDP ❌ ✅ ZeRO-1/2/3 / Offload / Infinity ✅ FSDP FULL_SHARD
Mixed Precision ✅ apex.amp / native AMP ✅ FP16 / BF16 ✅ MixedPrecision policy
Activation Recompute ✅ Selective + SP ✅ ✅ (checkpoint_wrapper)
학습 곡선 LLaMA, GPT-3 BLOOM, GPT-NeoX Llama-2 (Meta 자체)
장점 최고 throughput (3D) 가장 풍부한 옵션 PyTorch native, debugging 쉬움
단점 설정 복잡 dependency 무거움 TP/PP 부재

📖 각 문서 구성 방식

모든 문서는 다음 11-섹션 골격 으로 작성됩니다.

# 섹션 내용
1 🎯 핵심 질문 이 문서가 답하는 3~5개의 본질적 질문
2 🔍 왜 이 ... 인가 해당 통신 패턴·메모리 분해·schedule 이 분산 학습의 어떤 핵심 문제를 푸는지
3 📐 수학적 선행 조건 PyTorch Internals · Opt · LLM Pretrain · LA 레포의 어떤 정리를 전제하는지
4 📖 직관적 이해 AllReduce · Bucket · Column/Row split · Bubble · ZeRO shard 의 기하학적 직관
5 ✏️ 엄밀한 정의 Ring AllReduce · Gradient Bucket · Column-parallel Linear · Bubble Ratio · ZeRO-k Memory 등
6 🔬 정리와 증명 Bandwidth-optimal lower bound · Memory accounting · 1F1B activation 분석 · MLP 2-AllReduce
7 💻 NumPy / PyTorch / DeepSpeed 구현 검증 4 가지 실험 (### 실험 1 ~ ### 실험 4) — toy MDP · NCCL profile · ZeRO 측정 · 시각화
8 🔗 실전 활용 언제 DDP · 언제 FSDP · 언제 Megatron — 모델 규모별 선택 가이드, modern variants
9 ⚖️ 가정과 한계 각 기법의 실패 모드 (네트워크 병목 · debugging 어려움 · OOM · stale gradient)
10 📌 핵심 정리 한 장으로 요약 ($\boxed{}$ 핵심 수식 + 표)
11 🤔 생각해볼 문제 (+ 해설) 기초 / 심화 / 논문 비평 의 3 문제, <details> 펼침 해설

📚 연습문제 총 105개 (35 문서 × 3 문제): 기초 / 심화 / 논문 비평 의 3-tier 구성, 모든 문제에 <details> 펼침 해설 포함. Ring AllReduce 손 증명부터 ZeRO-3 forward prefetch 의 overlap 분석, Megatron Attention 의 head-partition 정당화, 1F1B activation 의 정확한 계산, FSDP auto_wrap_policy 의 trade-off, 3D parallelism axis 배치의 정량적 근거까지 단계적으로 심화됩니다.

🧭 푸터 네비게이션: 각 문서 하단에 ◀ 이전 / 📚 README / 다음 ▶ 링크가 항상 제공됩니다. 챕터 경계에서도 다음 챕터 첫 문서로 자동 연결됩니다.

⏱️ 학습 시간 추정: 문서당 평균 약 500600줄 (정의·증명·코드·연습문제 포함) 기준 **약 60분1시간 30분**. 전체 35문서는 약 40~50시간 상당 (증명 재구성·DDP/FSDP/Megatron 실험 재현 포함 시 70시간+).


🗺️ 추천 학습 경로

🟢 "DDP·FSDP 는 쓰지만 왜 작동하는지 시스템적으로 이해하고 싶다" — 입문 투어 (1주, 약 14~16시간)
Day 1  Ch1-01  Collective Operation 분류
       Ch1-02  AllReduce 4 가지 구현
Day 2  Ch1-03  Ring AllReduce 의 수학
       Ch1-05  Bandwidth vs Latency
Day 3  Ch2-01  Data Parallelism 기본
       Ch2-02  PyTorch DDP
Day 4  Ch2-03  Gradient Bucket 과 Overlap
       Ch2-05  Batch Size Scaling
Day 5  Ch5-01  ZeRO 의 Motivation
       Ch5-04  ZeRO-3
Day 6  Ch5-05  FSDP — PyTorch native
Day 7  Ch6-02  Gradient Checkpointing
       Ch7-01  3D Parallelism 통합
🟡 "Tensor + Pipeline + ZeRO 의 수학을 정복한다" — 이론 집중 (2주, 약 28~32시간)
1주차 — Collective · DP · TP
  Day 1   Ch1-01~02   Collective ops + AllReduce 비교
  Day 2   Ch1-03~05   Ring 증명 + NCCL + bandwidth/latency
  Day 3   Ch2-01~03   DP + DDP + bucket overlap
  Day 4   Ch2-04~06   Gradient accum + scaling + async
  Day 5   Ch3-01~02   Naive MP + Column/Row split
  Day 6   Ch3-03~04   Megatron MLP + Attention TP
  Day 7   Ch3-05      TP 통신 비용 분석

2주차 — PP · ZeRO · Activation · 3D
  Day 1   Ch4-01~02   Naive PP + GPipe
  Day 2   Ch4-03~05   1F1B + Interleaved + Chimera
  Day 3   Ch5-01~03   ZeRO-{0,1,2}
  Day 4   Ch5-04~06   ZeRO-3 + FSDP + Offload/Infinity
  Day 5   Ch6-01~02   Activation memory + Checkpointing
  Day 6   Ch6-03~04   Selective recompute + Sequence Parallelism
  Day 7   Ch7-01~04   3D + MoE + Checkpoint + Elastic
🔴 "분산 학습의 시스템 수학을 완전 정복한다" — 전체 정복 (8주, 약 40~50시간 + 분산 실험 재현 15~20시간)
1주차   Chapter 1 전체 — Collective Communication
         → Ring AllReduce 손 증명 (scatter-reduce + all-gather phase)
         → α/β model 로 bandwidth-optimal lower bound 도출
         → NCCL_DEBUG=INFO 로 실제 ring topology 관찰

2주차   Chapter 2 전체 — Data Parallelism · DDP
         → DDP gradient bucket 의 layer-reverse 순서 직접 검증
         → bucket_cap_mb {1, 5, 25, 100} 별 timeline (torch.profiler)
         → no_sync() context 의 AllReduce 생략 검증
         → Goyal 2017 의 linear scaling + warmup 재현 (CIFAR scale)

3주차   Chapter 3 전체 — Tensor Parallelism
         → Column-parallel / Row-parallel Linear 구현
         → Megatron MLP 의 forward AllReduce 횟수 측정
         → Attention TP 에서 head-partition 의 numerical 동치 검증

4주차   Chapter 4 전체 — Pipeline Parallelism
         → Naive PP / GPipe / 1F1B / Interleaved 의 timeline 시각화
         → bubble ratio 의 (P, M) 변화 plot
         → 1F1B 의 activation memory O(M) → O(P) 계측

5주차   Chapter 5 전체 — ZeRO · FSDP
         → ZeRO-{0,1,2,3} 의 per-GPU 메모리 직접 측정 (nvidia-smi)
         → FSDP forward prefetch overlap 의 timeline (Nsight)
         → ZeRO-Offload 의 PCIe bandwidth 병목 측정

6주차   Chapter 6 전체 — Activation Memory
         → s b h (34 + 5as/h) 공식 직접 측정 (GPT-2 small)
         → Gradient checkpointing 의 33% recompute 검증
         → Selective recompute 의 Pareto plot
         → Sequence Parallelism 으로 100K seq 학습

7주차   Chapter 7 (1~2) — 3D + MoE
         → LLaMA-7B 를 TP=2 PP=2 DP=2 의 8 GPU 로 학습
         → MoE expert parallelism 의 all-to-all 시간 측정
         → Load balancing loss 의 expert utilization 효과

8주차   Chapter 7 (3~4) + 종합 — Checkpoint · Elastic
         → Sharded checkpoint save/load 의 wall-clock 측정
         → 학습 중 GPU kill → resume 검증 (TorchElastic)
         → "Megatron vs DeepSpeed vs FSDP" / "TP vs PP 의 우선순위" 토론

🔗 연관 레포지토리

레포 주요 내용 연관 챕터
pytorch-internals-deep-dive CUDA stream · async kernel · autograd engine · torch.distributed Ch1-04 (NCCL stream), Ch2-03 (bucket async)
optimization-theory-deep-dive SGD · Adam · momentum · batch size scaling Ch2-05 (linear scaling), Ch5 (Adam memory)
llm-pretraining-deep-dive FLOPs · scale up · Chinchilla · training recipe Ch7 전체 (LLaMA, GPT-3 recipe)
linear-algebra-deep-dive Matrix partition · block decomposition · Frobenius Ch3 (TP column/row split)
cnn-deep-dive Memory hierarchy · roofline · arithmetic intensity Ch6 (activation memory analysis)
efficient-ml-deep-dive (다음) Kernel fusion · FlashAttention · Quantization · Inference Ch6 이후 추론·압축 보강
transformer-deep-dive Attention · Pre-LN · Scaling Law Ch3-04 (Megatron Attention TP), Ch6 (per-layer act)

💡 이 레포는 "DDP·TP·PP·ZeRO 가 모두 메모리·통신·연산의 trade-off 관리이고, $2(N-1)/N$ · column/row 분할 · $M \geq 4P$ · $1/N$ shard 가 왜 각각의 수학적 동기를 갖는가" 에 집중합니다. PyTorch Internals 에서 CUDA stream 과 async, Optimization 에서 SGD·Adam 과 batch size, LLM Pretraining 에서 model scale 을, Linear Algebra 에서 matrix partition 을 익힌 후 오면 Chapter 1 (Ring AllReduce) 과 Chapter 5 (ZeRO memory accounting) 의 증명이 훨씬 자연스럽습니다. Efficient ML Deep Dive 와 함께 보면 학습-추론 파이프라인 전체의 시스템 비용이 선명해집니다.


📖 Reference

🌐 Collective Communication · AllReduce

  • Optimization of Collective Communication Operations in MPICH (Thakur et al., 2005) — Bandwidth-optimal AllReduce 분석 효시
  • Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations (Patarasuk & Yuan, 2009) — Ring AllReduce
  • NCCL: NVIDIA Collective Communications Library (NVIDIA, 2017–) — Topology-aware ring/tree
  • Horovod: Fast and Easy Distributed Deep Learning in TensorFlow (Sergeev & Del Balso, 2018) — Ring AllReduce 의 ML 적용

🚀 Data Parallelism · DDP

  • PyTorch Distributed: Experiences on Accelerating Data Parallel Training (Li et al., 2020) — PyTorch DDP
  • Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour (Goyal et al., 2017) — Linear scaling
  • An Empirical Model of Large-Batch Training (McCandlish et al., 2018) — Gradient noise scale
  • Don't Decay the Learning Rate, Increase the Batch Size (Smith et al., 2018) — Batch scaling
  • HOGWILD!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent (Recht et al., 2011) — Async lock-free
  • Scaling Distributed Machine Learning with the Parameter Server (Li et al., 2014) — PS architecture

🧱 Tensor Parallelism · Megatron

  • Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019) — Tensor Parallelism 효시
  • Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021) — 3D Parallelism
  • Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022) — Selective + Sequence Parallelism

🎯 Pipeline Parallelism

  • GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism (Huang et al., 2019) — GPipe
  • PipeDream: Generalized Pipeline Parallelism for DNN Training (Narayanan et al., 2019) — 1F1B
  • Memory-Efficient Pipeline-Parallel DNN Training (Narayanan et al., 2021) — 1F1B 변형
  • Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines (Li & Hoefler, 2021) — Chimera
  • Zero Bubble Pipeline Parallelism (Qi et al., 2024) — bubble 0 schedule

💾 ZeRO · FSDP · Sharding

  • ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari et al., 2020) — ZeRO-1/2/3
  • ZeRO-Offload: Democratizing Billion-Scale Model Training (Ren et al., 2021) — CPU offload
  • ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning (Rajbhandari et al., 2021) — NVMe offload
  • PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (Zhao et al., 2023) — FSDP
  • DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters (Rasley et al., 2020) — DeepSpeed framework

🧠 Activation Memory

  • Training Deep Nets with Sublinear Memory Cost (Chen et al., 2016) — Gradient checkpointing 효시
  • Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022) — Selective recomputation + Sequence Parallelism
  • FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022) — Attention activation 절감

🌍 MoE · 3D Parallelism

  • GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (Lepikhin et al., 2021) — GShard
  • Switch Transformer: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (Fedus et al., 2022) — Switch
  • DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training (Rajbhandari et al., 2022)
  • Megatron-DeepSpeed (NVIDIA × Microsoft, 2022–) — 통합 stack

🛠️ Implementation · Libraries

  • PyTorch Distributed Documentation (Meta, 2020–) — DDP / FSDP / RPC 공식 문서
  • DeepSpeed Documentation (Microsoft, 2020–) — ZeRO / Pipeline / MoE 공식
  • NVIDIA Megatron-LM (NVIDIA, 2019–) — Tensor + Pipeline reference
  • HuggingFace Accelerate (Gugger et al., 2022) — FSDP / DeepSpeed 통합 launcher
  • Colossal-AI (Bian et al., 2023) — 통합 분산 학습 framework
  • TorchElastic / torchrun (PyTorch, 2021–) — Elastic launcher

⭐️ 도움이 되셨다면 Star 를 눌러주세요!

Made with ❤️ by IQ AI Lab


"DDP 의 all_reduce 를 호출하는 것과 — Patarasuk & Yuan 2009 로 Ring AllReduce 가 scatter-reduce $(N-1)$ + all-gather $(N-1)$ steps × $P/N$ bytes 로 $2(N-1)P/N$ 의 bandwidth-optimal lower bound 와 정확히 일치함을 한 줄씩 증명 · Rajbhandari 2020 으로 ZeRO-{1,2,3} 의 per-rank 메모리 $16\psi \to 4\psi + K\psi/N \to 2\psi + (2+K)\psi/N \to (4+K)\psi/N$ 가 어떻게 단계적으로 감소하는지 유도 · Shoeybi 2019 로 Megatron MLP 의 column-GELU-row 구조가 block 당 단 2 번의 AllReduce (forward + backward) 로 끝나는 이유를 분석 · Huang 2019 로 GPipe 의 bubble ratio $(P-1)/(P-1+M)$ 가 $M = 4P$ 에서 20% 미만이 되는 수학을 도출 · Narayanan 2019 로 1F1B schedule 이 같은 bubble 에서 activation memory 를 $O(M)$ 에서 $O(P)$ 로 줄이는 메커니즘을 따라가기 · Zhao 2023 으로 FSDP 가 ZeRO-3 와 어떻게 다르고 (forward prefetch · backward overlap · MixedPrecision) 왜 PyTorch native 의 표준이 되었는지 분석 — 이 모든 '왜' 를 직접 유도할 수 있는 것은 다르다"

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors