목차

핵심 한 줄 정리

Large-scale distributed training은 거대한 model과 batch를 여러 accelerator에 나누는 것만으로 끝나지 않으며, GPU와 cluster의 계층적인 memory bandwidth에 맞춰 computation과 communication을 겹치고 data, parameter, sequence, layer 축의 parallelism을 조합해 Model FLOPs Utilization을 높이는 문제이다.

반드시 기억할 개념

Distributed training의 두 가지 작업

Computer가 수행하는 작업은 크게 computation과 communication으로 나눌 수 있다.

• Computation은 input bit로부터 새로운 output bit를 계산하는 작업이다.
• Communication은 data를 한 memory 또는 device에서 다른 곳으로 옮기는 작업이다.

Accelerator 수를 늘리면 이론적인 compute는 거의 비례해서 증가한다. 반면 device 사이의 communication bandwidth는 GPU 내부 compute만큼 빠르게 증가하지 않는다. 따라서 distributed training의 핵심은 다음 두 조건을 동시에 만족하는 것이다.

1. 각 accelerator에 독립적으로 실행할 충분한 computation을 배정한다.
2. 필요한 communication을 computation과 overlap하여 accelerator가 기다리는 시간을 줄인다.

전체 step time은 단순히 compute time만으로 결정되지 않는다.

$$t_{\text{step}} \approx t_{\text{compute}} + t_{\text{communication}} + t_{\text{input}} + t_{\text{idle}}$$

Communication을 compute와 완전히 overlap할 수 있다면 두 비용의 합보다 큰 쪽에 가까워질 수 있다.

$$t_{\text{step}} \approx \max \left( t_{\text{compute}}, t_{\text{communication}} \right)$$

실제 목표는 theoretical compute를 늘리는 것보다 $t_{\text{idle}}$을 줄이고 communication을 숨겨 device를 계속 유용한 연산에 사용하는 것이다.

GPU의 memory hierarchy

GPU에서는 compute core에 가까운 memory일수록 작지만 빠르고, 멀수록 크지만 느리다. NVIDIA H100을 예로 들면 대략 다음과 같은 hierarchy를 가진다.

Level	대략적인 크기	역할
Register file와 L1 cache	SM당 약 $256\text{ KB}$	현재 kernel의 operand와 자주 쓰는 data
L2 cache	GPU 전체 약 $50\text{ MB}$	여러 SM이 공유하는 cache
HBM	약 $80\text{ GB}$	Model parameter, activation, optimizer state 저장

HBM은 High-Bandwidth Memory이며, H100의 HBM과 compute core 사이 bandwidth는 약 $3\text{ TB/s}$이다. 매우 큰 값처럼 보이지만 tensor core의 peak compute를 계속 채우려면 같은 data를 cache와 register에서 여러 번 재사용해야 한다.

Matrix multiplication에서 input을 HBM에서 매번 다시 읽으면 compute core보다 memory transfer가 먼저 한계에 도달한다. 그래서 고성능 kernel은 tiling을 사용해 matrix block을 cache 또는 register에 올린 뒤 여러 연산에서 재사용한다.

이 관점은 distributed training에도 그대로 확장된다. GPU 내부 HBM, 같은 server 안의 GPU, 같은 pod의 GPU, 다른 pod의 GPU 순으로 멀어질수록 일반적으로 bandwidth는 낮아지고 latency는 커진다.

Streaming Multiprocessor와 tensor core

H100에는 132개의 active Streaming Multiprocessor, SM이 있다. 각 SM에는 일반적인 FP32 arithmetic core와 matrix multiplication에 특화된 tensor core가 함께 들어 있다.

FP32 core 하나가 한 cycle에 fused multiply-add를 계산한다고 하자.

$$y = ax+b$$

Multiply와 add를 각각 하나의 FLOP으로 세면 fused multiply-add는 2 FLOPs이다. SM의 FP32 core 128개는 한 cycle에 총 256 FLOPs를 처리할 수 있다.

Tensor core는 작은 matrix tile에 대한 matrix multiply-accumulate를 전용 회로로 계산한다.

$$D = AB+C$$

H100 tensor-core tile의 한 예는 $A \in \mathbb{R}^{16 \times 4}$, $B \in \mathbb{R}^{4 \times 8}$, $C,D \in \mathbb{R}^{16 \times 8}$이다. 이 연산은 tensor core 하나에서 1,024 FLOPs이며, SM당 tensor core 4개를 사용하면 cycle당 4,096 FLOPs에 해당한다.

$$16 \times 4 \times 8 \times 2 = 1024$$

따라서 deep learning workload가 peak throughput에 가까워지려면 일반 FP32 core보다 tensor core가 처리할 수 있는 큰 matrix multiplication으로 연산을 구성해야 한다.

Mixed precision

Tensor core는 보통 FP16이나 BF16 같은 낮은 precision의 input으로 multiplication을 수행하고, accumulation은 FP32처럼 더 높은 precision으로 수행한다.

$$\operatorname{FP16/BF16}(A) \times \operatorname{FP16/BF16}(B) \xrightarrow{\text{accumulate}} \operatorname{FP32}(D)$$

낮은 precision은 memory traffic을 줄이고 tensor core throughput을 활용하게 한다. 높은 precision의 accumulation은 긴 dot product에서 rounding error가 누적되는 문제를 줄인다.

Model을 무조건 FP32로 실행하면 tensor core 대신 상대적으로 느린 arithmetic path를 사용할 수 있다. Framework에서는 automatic mixed precision을 사용하되 다음을 확인해야 한다.

• 어떤 operator가 FP16 또는 BF16으로 실행되는가?
• Reduction과 optimizer state는 어떤 precision으로 유지되는가?
• FP16의 작은 gradient가 underflow되지 않도록 loss scaling이 필요한가?
• BF16을 지원하는 hardware에서 FP16 대신 BF16이 더 안정적인가?

Hardware spec을 비교하면 2013년 K40은 약 $5\text{ TFLOPs}$의 FP32 peak를 가졌고, B200의 발표 spec은 약 $83.3\text{ TFLOPs}$ FP32와 약 $5{,}000\text{ TFLOPs}$ mixed-precision peak를 제시하였다. 이 값은 실제 application throughput이 아니라 특정 precision에서의 theoretical peak이다.

GPU cluster도 하나의 memory hierarchy이다

같은 H100 server에는 일반적으로 GPU 8개가 있고, GPU 사이에는 NVLink와 NVSwitch 같은 고속 interconnect가 사용된다. Llama3 training cluster를 바탕으로 한 hierarchy example은 다음과 같다.

범위	구성	GPU 간 bandwidth 예시
GPU 내부	HBM과 compute core	약 $3\text{ TB/s}$
한 server	GPU 8개	약 $900\text{ GB/s}$
한 pod	GPU 3,072개	임의 GPU 쌍 기준 약 $50\text{ GB/s}$
Pod 사이	전체 cluster	$50\text{ GB/s}$보다 낮음

Llama3 training cluster의 공개된 구성에서는 다음과 같이 scale-out하였다.

• 한 rack에 server 2대와 GPU 16개를 배치한다.
• Rack 192개를 묶어 GPU 3,072개의 pod를 만든다.
• Pod 8개를 묶어 GPU 24,576개의 cluster를 만든다.

GPU 24,576개가 각각 $80\text{ GB}$ HBM을 가진다면 총 HBM은 약 $1.97\text{ PB}$이다.

$$24{,}576 \times 80\text{ GB} = 1{,}966{,}080\text{ GB} \approx 1.97\text{ PB}$$

FP32 core는 약 4억 1,500만 개, tensor core는 약 1,300만 개이며 mixed-precision peak compute는 약 $24\text{ EFLOPs}$이다. 중요한 점은 이 숫자를 더하는 것이 아니라, bandwidth가 다른 수만 개의 device를 하나의 computer처럼 효율적으로 사용하는 것이다.

통신량이 큰 parallelism은 같은 server처럼 빠른 link 위에 배치하고, 통신량이 작은 parallelism은 server나 pod를 넘는 느린 link 위에 배치해야 한다. Logical parallelism topology를 physical network topology와 맞추는 것이 핵심이다.

다른 training accelerator

Distributed training의 원리는 NVIDIA GPU에만 한정되지 않는다.

• Google TPU는 matrix operation에 특화된 accelerator이며 v5p는 최대 8,960개 chip의 pod로 구성할 수 있다.
• AMD는 Instinct 계열 accelerator를 제공한다.
• AWS Trainium은 cloud에서 large-scale training에 사용하는 전용 accelerator이다.

각 hardware의 memory hierarchy, collective library, supported precision, compiler stack은 다르지만 computation을 나누고 communication을 숨기는 원리는 같다.

%% title: GPU Cluster의 계층적 Memory·Interconnect
%% caption: Compute에 가까울수록 빠르고 용량이 작으며, 멀어질수록 용량은 커지지만 bandwidth와 latency가 불리해진다.
flowchart TB
    compute["Tensor Cores / CUDA Cores"] --> regs["Registers<br/>가장 빠름 · 가장 작음"]
    regs --> shared["Shared Memory / L1"]
    shared --> l2["L2 Cache"]
    l2 --> hbm["GPU HBM"]
    hbm --> nvlink["NVLink / NVSwitch<br/>Server 내 GPU"]
    nvlink --> network["RDMA Network<br/>Server 간 GPU"]
    network --> host["Host Memory / Storage<br/>가장 큰 용량 · 가장 느림"]

Transformer에서 나눌 수 있는 축

Transformer는 $L$개 layer를 쌓고, 각 layer가 다음 3차원 activation을 처리한다고 볼 수 있다.

$$X \in \mathbb{R}^{B \times S \times D}$$

$B$는 batch size, $S$는 sequence length, $D$는 model dimension이다. 여기에 layer 축까지 포함하면 주요 parallelism은 다음 네 축에 대응한다.

나누는 축	Parallelism	각 device가 담당하는 것
Batch $B$	Data parallelism	서로 다른 sample
Sequence $S$	Context parallelism	한 sequence의 서로 다른 token 구간
Model dimension $D$	Tensor parallelism	한 matrix operation의 서로 다른 shard
Layer $L$	Pipeline parallelism	서로 다른 layer 구간

Model이 작으면 data parallelism만으로 충분하다. Model, sequence, device 수가 커질수록 여러 축을 동시에 사용하는 multidimensional parallelism이 필요하다.

Collective communication

여러 device가 협력할 때 반복해서 사용하는 통신 pattern을 collective operation이라고 한다.

Operation	Input	Output	대표 용도
Broadcast	한 rank에 tensor 하나	모든 rank에 같은 tensor	Parameter 배포
All-reduce	모든 rank에 tensor	모든 rank에 합 또는 평균	DDP gradient 동기화
All-gather	모든 rank에 서로 다른 shard	모든 rank에 전체 tensor	Sharded parameter 복원
Reduce-scatter	모든 rank에 tensor	Reduce 결과의 서로 다른 shard	Sharded gradient 집계
All-to-all	모든 rank의 여러 shard	목적 rank별로 재배치	Context, expert parallelism

Naive하게 한 device로 data를 모두 모으면 해당 device와 link가 bottleneck이 된다. 실제 collective library는 ring이나 tree algorithm을 사용해 traffic을 여러 link에 분산한다.

Distributed algorithm을 이해할 때는 계산 결과만 볼 것이 아니라 다음도 함께 봐야 한다.

• 한 step에서 parameter 크기의 몇 배를 통신하는가?
• 모든 device가 결과 전체를 가지는가, shard만 가지는가?
• Collective를 어느 network group에서 실행하는가?
• Communication을 어느 compute kernel과 overlap할 수 있는가?

Data parallelism의 수학

$M$개의 GPU가 있고 각 GPU가 local batch $N$개를 처리한다고 하자. Global batch size는 다음과 같다.

$$B_{\text{global}} = MN$$

GPU $i$의 sample $j$에 대한 loss를 $\ell(x_{i,j};W)$라고 하면 global loss는 다음과 같다.

$$L(W) = \frac{1}{MN} \sum_{i=1}^{M} \sum_{j=1}^{N} \ell(x_{i,j};W)$$

Gradient의 linearity를 사용하면 다음과 같이 device별 local gradient와 device 사이 평균으로 나눌 수 있다.

$$\nabla_W L = \frac{1}{M} \sum_{i=1}^{M} \left( \frac{1}{N} \sum_{j=1}^{N} \nabla_W \ell(x_{i,j};W) \right)$$

괄호 안은 GPU $i$가 독립적으로 계산할 수 있는 local gradient이다. 바깥 평균만 device 사이 collective communication이 필요하다. 따라서 synchronous data parallelism은 같은 global batch를 한 개의 큰 GPU에서 계산한 것과 수학적으로 같은 gradient를 만든다.

Distributed Data Parallel

Distributed Data Parallel, DDP의 한 iteration은 다음 순서로 진행된다.

1. 모든 GPU가 같은 model parameter와 optimizer state의 replica를 가진다.
2. 각 GPU는 서로 다른 local minibatch를 읽는다.
3. 각 GPU가 독립적으로 forward와 backward를 수행해 local gradient를 계산한다.
4. Gradient를 all-reduce하여 모든 GPU에 같은 global gradient를 만든다.
5. 모든 GPU가 같은 optimizer update를 적용한다.

모든 replica가 같은 parameter에서 시작하고 같은 gradient로 update하므로 다음 step에서도 같은 parameter를 유지한다.

서로 다른 GPU가 실수로 같은 minibatch를 읽으면 global batch가 커지지 않고 같은 sample을 중복 계산한다. Distributed sampler가 rank마다 서로 다른 data shard를 제공하는지 반드시 확인해야 한다.

%% title: Distributed Data Parallel의 한 Iteration
%% caption: 각 GPU는 같은 model replica로 다른 minibatch를 처리하고, local gradient를 all-reduce한 뒤 모두 같은 optimizer update를 적용한다.
flowchart LR
    d1["Local Batch 1"] --> g1["GPU 1<br/>Model Replica<br/>Forward + Backward"]
    d2["Local Batch 2"] --> g2["GPU 2<br/>Model Replica<br/>Forward + Backward"]
    d3["Local Batch M"] --> g3["GPU M<br/>Model Replica<br/>Forward + Backward"]
    g1 --> reduce["Gradient All-Reduce<br/>sum / average"]
    g2 --> reduce
    g3 --> reduce
    reduce --> u1["Same Optimizer Update<br/>GPU 1"]
    reduce --> u2["Same Optimizer Update<br/>GPU 2"]
    reduce --> u3["Same Optimizer Update<br/>GPU M"]

Backward와 all-reduce overlap

Backward는 마지막 layer에서 첫 layer 방향으로 gradient를 만든다. 전체 backward가 끝난 뒤 한꺼번에 all-reduce할 필요는 없다.

• Layer $l+1$의 gradient가 준비되면 즉시 all-reduce를 시작한다.
• 동시에 compute core는 layer $l$의 backward를 계산한다.
• Gradient 여러 개를 bucket으로 묶어 collective 호출 횟수를 줄인다.

이상적으로는 다음과 같은 pipeline이 만들어진다.

$$\text{Backward}(l) \quad\Vert\quad \text{AllReduce}\left(\nabla W_{l+1}\right)$$

$\Vert$는 두 작업이 동시에 진행됨을 뜻한다. 마지막 backward가 끝날 때 모든 gradient communication도 끝나 있으면 optimizer가 기다리지 않고 update를 시작할 수 있다.

PyTorch의 DistributedDataParallel은 gradient hook과 bucket을 이용해 이 overlap을 자동화한다. 하지만 bucket size, layer별 compute, network 속도에 따라 communication이 완전히 숨겨지는지는 달라지므로 profiling이 필요하다.

Synchronous와 asynchronous SGD

Synchronous SGD는 매 step 모든 worker의 gradient를 모은 뒤 같은 parameter update를 수행한다. 결과를 이해하고 재현하기 쉽지만 가장 느린 worker를 기다려야 한다.

Asynchronous SGD는 worker가 서로 다른 시점의 parameter로 독립적인 update를 계산하고 parameter server 또는 다른 worker에 전달할 수 있다. Synchronization barrier가 적지만 stale gradient가 생긴다.

$$W_{t+1} = W_t - \eta g\left(W_{t-\tau}\right)$$

$\tau$는 gradient가 계산된 뒤 적용될 때까지의 staleness이다. $\tau$가 크면 현재 parameter와 맞지 않는 방향으로 update할 수 있어 training이 불안정하고 재현하기 어려워진다. 가능한 경우에는 synchronous training이 일반적으로 더 단순하고 안정적이다.

DDP의 model-state memory bottleneck

DDP는 모든 GPU에 parameter, gradient, optimizer state를 복제한다. Adam을 단순화하면 parameter 하나마다 다음 네 값을 저장한다.

1. Parameter $W$
2. Gradient $\nabla W$
3. First moment $m$
4. Second moment $v$

모두 16-bit라고 가정한 lower bound는 parameter당 8 bytes이다.

$$4 \times 2\text{ bytes} = 8\text{ bytes/parameter}$$

따라서 parameter 10억 개는 최소 약 $8\text{ GB}$가 필요하다.

$$10^9 \times 8\text{ bytes} = 8\text{ GB}$$

실제로는 optimizer state나 master parameter를 FP32로 유지할 수 있어 parameter당 memory가 더 크다. 여기에 activation, temporary buffer, allocator overhead도 추가된다. HBM이 $80\text{ GB}$라고 해도 DDP만으로 학습할 수 있는 model 크기는 빠르게 제한된다.

Fully Sharded Data Parallel

Fully Sharded Data Parallel, FSDP는 data parallelism을 유지하면서 parameter, gradient, optimizer state를 GPU 사이에 shard한다. Persistent model-state memory는 이상적으로 GPU 수 $M$에 반비례한다.

$$\text{Model-state memory per GPU} \approx \frac{P}{M}$$

$P$는 전체 model state의 memory이다. 각 GPU는 여전히 서로 다른 local batch의 full forward와 backward를 수행하지만, 현재 layer를 계산할 때만 필요한 parameter shard를 모은다.

Forward에서 layer $l$을 처리하는 과정은 다음과 같다.

1. 모든 rank의 $W_l$ shard를 all-gather하여 일시적으로 full parameter를 만든다.
2. 각 rank가 자신의 local activation에 대해 layer $l$을 계산한다.
3. 계산이 끝나면 owner shard 이외의 full parameter buffer를 해제한다.
4. Layer $l$을 계산하는 동안 $W_{l+1}$을 prefetch한다.

Backward에서도 parameter를 다시 all-gather해 local gradient를 계산한다. 이후 reduce-scatter로 모든 rank의 local gradient를 합치고, 각 rank에는 자신이 소유한 gradient shard만 남긴다. 각 rank는 자신의 optimizer-state shard에 대해서만 update한다.

$$\text{AllGather}(W_l) \rightarrow \text{Backward}(l) \rightarrow \text{ReduceScatter}(\nabla W_l)$$

Steady state에서는 다음 세 작업을 서로 다른 인접 layer에 겹칠 수 있다.

• Layer $l-1$의 parameter prefetch
• Layer $l$의 backward compute
• Layer $l+1$의 gradient reduce-scatter와 optimizer update

FSDP는 DDP보다 memory를 크게 줄이는 대신 forward와 backward 중 parameter communication이 추가된다. Payload를 parameter 크기 단위로 단순화하면 한 iteration에서 forward parameter, backward parameter, gradient까지 약 $3P$에 해당하는 communication이 필요하다고 볼 수 있다.

%% title: FSDP의 Layer별 Parameter·Gradient 통신
%% caption: 평소에는 model state shard만 유지하고, 현재 layer를 계산할 때만 parameter를 all-gather한다. Backward 후에는 reduce-scatter로 gradient shard만 남긴다.
flowchart LR
    shards["Persistent Wₗ Shards<br/>GPU 1 … M"] --> gather["All-Gather Wₗ"]
    gather --> full["Temporary Full Wₗ"]
    full --> forward["Layer ℓ Forward / Backward"]
    forward --> localgrad["Local Full Gradient"]
    localgrad --> scatter["Reduce-Scatter ∇Wₗ"]
    scatter --> gradshards["Gradient Shards<br/>GPU 1 … M"]
    gradshards --> update["Local Optimizer-State<br/>Shard Update"]
    forward --> free["Release Full Parameter Buffer"]

Hybrid Sharded Data Parallel

Hybrid Sharded Data Parallel, HSDP는 GPU를 2차원 grid로 보고 두 축에 서로 다른 전략을 적용한다.

• Fast-link group 안에서는 FSDP로 model state를 shard한다.
• Group 사이에서는 DDP처럼 full model replica를 유지한다.

예를 들어 server마다 GPU가 8개라면 같은 server의 8개 GPU를 FSDP group으로 묶고, 여러 server 사이에는 data-parallel replica group을 만들 수 있다.

$$G_{\text{total}} = G_{\text{shard}} \times G_{\text{replica}}$$

FSDP group 내부에서는 parameter all-gather와 gradient reduce-scatter가 빈번하므로 빠른 NVLink를 사용한다. Server 사이에서는 iteration당 gradient synchronization 위주로 통신하므로 상대적으로 느린 network를 사용할 수 있다.

HSDP는 algorithm의 communication pattern을 cluster topology에 맞춘 대표적인 예이다. Shard group과 replica group의 크기는 model size, local batch, network bandwidth를 profiling해 결정해야 한다.

Activation memory bottleneck

Model state를 shard해도 backward에 필요한 activation이 memory를 차지한다. Llama3-405B example은 transformer layer 126개, model dimension 약 16,000, sequence length 4,096를 사용한다. Batch와 sequence가 커지면 layer마다 저장하는 hidden state가 parameter보다 큰 bottleneck이 될 수 있다.

일반적인 $N$-layer network는 forward에서 각 layer의 activation을 저장하므로 activation memory가 $O(N)$이다. Backward에서는 저장한 activation을 사용해 gradient를 계산한다.

아무 activation도 저장하지 않고 각 backward 단계마다 처음부터 forward를 다시 계산하면 memory는 매우 작아지지만 computation은 다음처럼 커질 수 있다.

$$N + (N-1) + \cdots +1 = O(N^2)$$

Activation checkpointing

Activation checkpointing은 일부 layer의 activation만 저장하고 나머지는 backward 중 다시 계산한다. Segment 길이를 $C$라고 하면 checkpoint activation은 약 $N/C$개이고, 현재 segment를 재계산할 때 최대 $C$개의 activation이 필요하다.

$$M_{\text{activation}} = O\left( \frac{N}{C} +C \right)$$

$C \approx \sqrt{N}$으로 선택하면 activation memory를 다음 수준으로 줄일 수 있다.

$$M_{\text{activation}} = O\left(\sqrt{N}\right)$$

Segment 단위 checkpointing은 각 segment의 forward를 backward 직전에 한 번 더 수행하므로 asymptotic compute는 여전히 $O(N)$이지만 forward recomputation만큼 constant factor가 증가한다. 정확한 memory와 compute trade-off는 checkpoint schedule에 따라 달라진다.

Activation checkpointing은 memory를 compute로 바꾸는 방법이다. GPU가 memory 때문에 더 큰 batch나 model을 처리하지 못할 때 유용하지만, recomputation이 MFU와 step time에 미치는 영향을 함께 측정해야 한다.

%% title: Activation Checkpointing의 저장과 재계산
%% caption: Forward에서 segment 경계 activation만 저장하고, backward가 필요한 segment를 지날 때 그 구간의 forward activation을 다시 계산한다.
flowchart LR
    x["Input"] --> s1["Segment 1"]
    s1 --> c1["Checkpoint 1<br/>saved"]
    c1 --> s2["Segment 2"]
    s2 --> c2["Checkpoint 2<br/>saved"]
    c2 --> s3["Segment 3"]
    s3 --> loss["Loss"]
    loss -. "backward" .-> r3["Recompute Segment 3"]
    c2 -.-> r3
    r3 -.-> r2["Recompute Segment 2"]
    c1 -.-> r2
    r2 -.-> r1["Recompute Segment 1"]
    x -.-> r1

Parallelism을 늘리는 실용적인 순서

Model과 cluster마다 경계는 다르지만 대략적인 scaling recipe는 다음과 같다.

1. Model이 약 10억 parameter 이하이고 GPU가 약 128개 이하라면 DDP부터 시작한다.
2. Model state가 GPU memory를 압박하면 FSDP로 전환한다.
3. Activation이 bottleneck이면 activation checkpointing을 적용한다.
4. GPU가 수백 개로 늘어 FSDP의 wide collective가 비싸지면 HSDP를 고려한다.
5. 약 500억 parameter 이상, GPU 1,000개 이상, sequence length 10,000 이상이라면 context, pipeline, tensor parallelism을 조합한다.

이는 고정된 법칙이 아니다. GPU memory, interconnect, model architecture, global batch 제한에 따라 전환점이 달라진다. Local batch size는 가능한 한 GPU memory를 충분히 사용하도록 키우되 optimization 성질과 global batch size를 함께 확인해야 한다.

Hardware FLOPs Utilization

Hardware FLOPs Utilization, HFU는 실제 hardware가 수행한 FLOPs가 theoretical peak의 어느 비율인지 나타낸다.

$$\operatorname{HFU} = \frac{F_{\text{executed}}} {G \cdot F_{\text{peak}} \cdot t_{\text{step}}}$$

$G$는 accelerator 수, $F_{\text{peak}}$는 accelerator 하나의 초당 peak FLOPs, $F_{\text{executed}}$는 step 동안 실제 실행된 전체 hardware operation이다.

H100의 tensor-core peak를 약 $989.4\text{ TFLOPs/s}$라고 할 때, 크기가 약 $8{,}000 \times 8{,}000$인 dense matrix multiplication benchmark는 약 $80\%$ HFU를 얻을 수 있었다. 큰 matrix가 tensor core를 잘 채우고 launch overhead를 상쇄하기 때문이다.

HFU는 recomputation이나 model 이외의 auxiliary compute도 hardware work로 셀 수 있다. 따라서 hardware가 바쁘다는 사실은 알려 주지만, 유용한 model forward와 backward에 얼마나 사용했는지는 직접 말해 주지 않는다.

Model FLOPs Utilization

Model FLOPs Utilization, MFU는 model의 유용한 forward와 backward computation이 theoretical peak의 어느 비율을 차지했는지 나타낸다.

$$\operatorname{MFU} = \frac{F_{\text{model per step}}} {G \cdot F_{\text{peak}} \cdot t_{\text{step}}}$$

$F_{\text{model per step}}$은 architecture와 global batch를 기준으로 계산한 useful model FLOPs이다. Data loading, communication, idle time, activation recomputation이 step을 늦추면 분모의 시간만 커져 MFU가 낮아진다.

Activation checkpointing을 사용하면 recomputation은 HFU에는 포함될 수 있지만 useful model FLOPs에는 포함하지 않으므로 MFU가 더 낮아질 수 있다.

대략적인 해석은 다음과 같다.

• MFU가 $30\%$보다 훨씬 낮으면 큰 bottleneck이 있을 가능성이 높다.
• $30\%$ 이상이면 실용적으로 괜찮은 수준이다.
• $40\%$ 이상이면 large-scale training에서 매우 좋은 수준이다.

Llama3-405B의 마지막 training phase는 GPU 약 8,000~16,000개에서 high-$30\%$부터 low-$40\%$ 수준의 MFU를 보고하였다. MFU는 parallelism degree, local batch, microbatch 수, checkpoint 범위 등을 선택할 때 중심 metric이 된다.

더 최신 GPU가 항상 더 높은 MFU를 얻는 것은 아니다. A100에서 H100으로 compute peak는 약 3배 증가했지만 memory bandwidth는 약 2배 증가해 compute와 communication의 격차가 더 커졌다. Device가 빨라질수록 communication을 숨기기 어려워질 수 있다.

Context parallelism

Context parallelism은 한 sequence의 token을 여러 GPU에 나눈다.

$$X = \operatorname{Concat} \left( X^{(1)}, X^{(2)}, \ldots, X^{(C)} \right)$$

LayerNorm, residual connection, token-wise MLP는 token마다 독립적으로 계산하므로 각 shard에서 바로 실행할 수 있다. 어려운 부분은 모든 token pair가 상호작용하는 self-attention이다.

$$A = \operatorname{softmax} \left( \frac{QK^T}{\sqrt{d}} \right)$$

Query shard 하나도 전체 $K,V$와 상호작용해야 정확한 attention output을 얻는다.

• Ring attention은 $K,V$ block을 ring 형태로 device 사이에 순환시키며 blockwise attention과 online softmax를 계산한다.
• Ulysses는 attention head와 sequence shard를 all-to-all로 재배치해 core attention을 head 축으로 병렬화한다.

Llama3 pretraining example에서는 sequence length 약 8,000인 초기 phase에는 context parallelism을 사용하지 않았고, 약 130,000 token으로 늘린 phase에는 16-way context parallelism을 사용하였다. 이 경우 sequence 하나를 GPU 16개가 함께 처리한다.

Context parallelism은 sequence가 길어 activation과 attention memory가 커질 때 특히 중요하다. Sequence가 짧으면 collective overhead가 이득보다 클 수 있다.

Pipeline parallelism

Pipeline parallelism은 $L$개 layer를 $P$개의 stage로 나누고 각 stage를 다른 GPU 또는 GPU group에 배치한다.

$$\text{Stage } p = \left\{ l_p, \ldots, l_{p+1} - 1 \right\}$$

Forward에서는 activation이 stage 1부터 $P$까지 순서대로 이동한다. Backward에서는 upstream gradient가 반대 방향으로 이동한다. Batch 하나만 처리하면 한 stage가 계산하는 동안 나머지 stage 대부분은 idle 상태가 된다.

Pipeline이 채워지고 비워지는 idle 구간을 pipeline bubble이라고 한다. Microbatch 하나만 사용하면 $P$개 stage의 utilization은 대략 $1/P$까지 떨어질 수 있다.

Batch를 $M$개의 microbatch로 나누어 서로 다른 stage에서 동시에 처리하면 bubble 비율을 줄일 수 있다. 단순한 forward pipeline의 이상적인 utilization은 다음과 같이 볼 수 있다.

$$U_{\text{pipeline}} \approx \frac{M} {M+P-1}$$

$P=4$, $M=4$이면 다음과 같다.

$$U_{\text{pipeline}} = \frac{4}{4+4-1} = \frac{4}{7} \approx 57\%$$

Microbatch를 늘리면 utilization은 좋아지지만 동시에 살아 있는 activation이 많아져 memory가 증가한다. Pipeline schedule, microbatch 수, activation checkpointing을 함께 조정해야 한다.

Stage별 compute가 불균형하면 가장 느린 stage가 전체 pipeline throughput을 제한한다. Layer 수뿐 아니라 attention, MLP, embedding, output head의 실제 실행 시간을 기준으로 stage를 나누어야 한다.

Tensor parallelism

Tensor parallelism은 하나의 weight matrix와 matrix multiplication을 여러 GPU에 나눈다. FSDP가 계산 직전에 full layer parameter를 모아 각 GPU가 local batch의 complete operation을 수행한다면, tensor parallelism은 parameter shard를 유지한 채 같은 operation의 일부를 각 GPU가 계산한다.

Matrix multiplication을 생각해 보자.

$$Y = XW$$

$W$를 column 방향으로 $P$개 shard로 나누면 각 GPU가 output feature의 일부를 계산한다.

$$W = \left[ W^{(1)} \;W^{(2)} \;\cdots \;W^{(P)} \right]$$ $$Y^{(p)} = XW^{(p)}$$

전체 $Y$가 다음 layer에 필요하면 $Y^{(p)}$를 all-gather해야 한다. 하지만 transformer MLP의 연속된 두 linear layer는 communication을 줄이도록 짝지을 수 있다.

첫 weight를 column-parallel로 나누어 hidden activation의 shard를 만든다.

$$H^{(p)} = \phi \left( XW_1^{(p)} \right)$$

두 번째 weight를 대응하는 row-parallel shard로 나누면 각 GPU가 output의 partial sum을 계산할 수 있다.

$$Z^{(p)} = H^{(p)}W_2^{(p)}$$

마지막에 한 번만 all-reduce하면 full output을 얻는다.

$$Z = \sum_{p=1}^{P}Z^{(p)}$$

이 방식은 두 linear layer 사이의 activation all-gather를 피한다. Transformer FFN이 두 linear layer로 구성되므로 tensor parallelism과 잘 맞는다.

Tensor parallelism은 layer 안에서 자주 collective를 수행하므로 latency와 bandwidth 요구가 높다. 보통 같은 server의 NVLink-connected GPU처럼 가장 빠른 group에 배치한다.

Multidimensional parallelism

아주 큰 training run은 한 전략만 선택하지 않고 여러 축을 곱해 사용한다.

$$G_{\text{total}} = G_{\text{DP}} \times G_{\text{CP}} \times G_{\text{PP}} \times G_{\text{TP}}$$

Llama3-405B의 큰 run example은 다음 degree를 동시에 사용하였다.

• 8-way tensor parallelism
• 16-way context parallelism
• 16-way pipeline parallelism
• 8-way data parallelism

$$8 \times 16 \times 16 \times 8 = 16{,}384$$

가장 빈번하게 통신하는 tensor-parallel group은 같은 server에 놓고, 비교적 낮은 빈도로 gradient를 동기화하는 data-parallel group은 더 먼 server에 걸쳐 배치하는 식으로 logical axis를 physical topology에 대응시킨다.

Parallelism degree를 키우는 것 자체가 목표는 아니다. 다음 조건을 만족하는 조합을 찾아야 한다.

• Model state와 activation이 각 GPU memory에 들어간다.
• Global batch size가 optimization에 적절하다.
• Matrix shape가 tensor core를 충분히 활용한다.
• Collective가 compute와 overlap된다.
• Pipeline bubble과 data-loading idle time이 작다.
• 최종적으로 MFU가 높다.

%% title: Transformer의 다차원 Parallelism 축
%% caption: 하나의 training run에서 batch, sequence, model dimension, layer 축을 독립적으로 나누고, 통신 빈도에 맞춰 물리 cluster topology에 배치한다.
flowchart TB
    tensor["Transformer Work<br/>B × S × D × L"]
    tensor --> dp["Data Parallel<br/>Batch B"]
    tensor --> cp["Context Parallel<br/>Sequence S"]
    tensor --> tp["Tensor Parallel<br/>Dimension D"]
    tensor --> pp["Pipeline Parallel<br/>Layers L"]
    tp --> fast["Fastest Links<br/>same NVLink domain"]
    cp --> fast
    pp --> medium["Stage-to-Stage Links"]
    dp --> wide["Wider Cluster<br/>lower-frequency sync"]
    fast --> product["DP × CP × TP × PP<br/>Total GPU Count"]
    medium --> product
    wide --> product

Parallelism 전략 비교

전략	나누는 대상	줄어드는 memory	주요 communication	주된 한계
DDP	Batch	줄지 않음	Gradient all-reduce	Model state가 모든 GPU에 복제됨
FSDP	Batch + model state	Parameter, gradient, optimizer state	Parameter all-gather, gradient reduce-scatter	Layer마다 communication 증가
HSDP	Shard group + replica group	Group 내 model state	Group 내 FSDP, group 간 all-reduce	Group 크기와 topology tuning 필요
Context parallelism	Sequence	Long-context activation	Attention용 ring 또는 all-to-all	Attention의 global interaction
Pipeline parallelism	Layer	Layer별 parameter와 activation	Stage 간 activation과 gradient	Pipeline bubble
Tensor parallelism	Weight matrix	Layer parameter와 activation shard	Layer 내부 all-gather/all-reduce	매우 빠른 interconnect 필요

과제에서 확인할 것

Mixed precision과 tensor core

• Model parameter, input, accumulation, gradient, optimizer state의 dtype을 각각 확인한다.
• FP32 실행과 BF16 또는 FP16 autocast 실행의 throughput을 비교한다.
• Matrix dimension이 tensor-core tile에 잘 맞지 않을 때 utilization이 어떻게 달라지는지 확인한다.
• FP16에서 loss scaling이 없을 때 gradient underflow가 발생하는지 확인한다.

DDP correctness

• Rank마다 다른 minibatch가 들어가는지 sampler index를 확인한다.
• Global batch size가 local_batch_size × world_size × gradient_accumulation_steps와 일치하는지 확인한다.
• Gradient를 sum하는지 average하는지에 따라 learning rate와 loss normalization이 어떻게 달라지는지 확인한다.
• Single-GPU global-batch gradient와 DDP all-reduced gradient가 수치적으로 일치하는지 비교한다.

Communication overlap

• Profiler에서 backward compute와 gradient all-reduce가 실제로 겹치는지 확인한다.
• Gradient bucket size를 바꿨을 때 collective 호출 횟수와 overlap이 어떻게 달라지는지 확인한다.
• DataLoader가 다음 batch를 비동기로 준비하지 못해 GPU가 기다리는 구간이 있는지 확인한다.
• Intra-node와 inter-node collective bandwidth를 각각 측정한다.

FSDP와 HSDP

• Parameter, gradient, optimizer state가 rank별로 실제 shard되는지 memory snapshot으로 확인한다.
• Forward와 backward에서 parameter all-gather가 몇 번 발생하는지 확인한다.
• FSDP group을 server 내부로 제한했을 때와 server 여러 대에 펼쳤을 때 throughput을 비교한다.
• Prefetch와 reduce-scatter가 compute에 숨겨지는지 timeline으로 확인한다.

Activation checkpointing

• Checkpoint를 적용하기 전후 peak memory와 step time을 함께 측정한다.
• Segment 길이 $C$를 바꾸며 memory와 recomputation의 trade-off를 확인한다.
• Checkpoint된 forward가 backward 중 다시 실행되므로 random operation의 재현성이 유지되는지 확인한다.
• 모든 layer를 무작정 checkpoint하는 것이 MFU에 미치는 영향을 확인한다.

Utilization metric

• Model architecture와 global batch로 step당 useful FLOPs를 계산한다.
• Device peak FLOPs와 실제 step time으로 MFU를 계산한다.
• HFU와 MFU가 다르게 나오는 이유를 activation recomputation과 auxiliary compute 관점에서 설명한다.
• Parallelism degree와 microbatch를 바꿀 때 memory만 보지 말고 MFU도 함께 비교한다.

Advanced parallelism

• Context parallelism에서 local query가 전체 key와 value를 보도록 어떤 collective가 필요한지 확인한다.
• Pipeline stage 수와 microbatch 수로 theoretical utilization을 계산하고 profiler 결과와 비교한다.
• Tensor-parallel MLP에서 첫 linear layer를 column-wise, 두 번째 layer를 row-wise로 나누는 이유를 block matrix multiplication으로 확인한다.
• Tensor, context, pipeline, data parallel group이 physical cluster topology의 어느 범위에 배치되는지 확인한다.