📖 약 90분

Chapter 2: 분산 학습 기초

Neural Network의 학습 원리부터 대규모 분산 학습 패러다임까지, Checkpointless Training을 이해하기 위한 기초 지식을 교과서 수준으로 다룹니다.

1. Neural Network Training 기초

1.1 Forward Pass

입력 데이터 X가 네트워크의 각 레이어를 통과하며 예측값을 계산하는 과정입니다. 각 레이어 l에서는 다음 연산이 수행됩니다:

최종 출력 Ŷ를 얻고, 실제 정답 Y와의 차이를 Loss Function L(Ŷ, Y)로 계산합니다. 대표적인 Loss Function으로는 회귀의 MSE(Mean Squared Error), 분류의 Cross-Entropy Loss가 있습니다.

1.2 Backward Pass (Backpropagation)

오차를 최소화하기 위해 Chain Rule(연쇄 법칙)을 사용하여 각 가중치에 대한 Loss의 기울기(Gradient)를 계산합니다:

이 Gradient는 해당 가중치를 어느 방향으로 얼마나 수정해야 Loss가 줄어드는지를 나타내는 벡터입니다. Neural Network의 출력(Loss)은 하나이고 파라미터는 수십억 개이므로, 출력에서 입력 방향으로 진행하는 Reverse-mode differentiation(역방향 미분)이 효율적입니다 — 한 번의 backward pass로 모든 파라미터의 gradient를 계산할 수 있습니다.

1.3 Training Step 전체 흐름

2. Optimizers 상세

2.1 SGD (Stochastic Gradient Descent)

가장 기본적인 최적화 알고리즘으로, Learning Rate η만큼 Gradient의 반대 방향으로 파라미터를 이동시킵니다:

SGD with Momentum: 이전 Gradient의 이동 방향을 기억하여 진동을 줄이고 수렴을 가속합니다:

β는 보통 0.9로 설정합니다. 관성(momentum)을 활용하여 local minima 탈출 가능성을 높입니다.

2.2 Adam (Adaptive Moment Estimation)

Momentum과 RMSprop을 결합한 적응적 학습률 최적화 알고리즘입니다. 각 파라미터마다 학습률을 개별 조절합니다:

1차 모멘트 (Momentum):

2차 모멘트 (Variance, RMSprop 역할):

Bias Correction (초기화가 0이므로 초반 편향 보정):

파라미터 업데이트:

Parameter	Default	의미
β1	0.9	1차 모멘트 decay (Momentum)
β2	0.999	2차 모멘트 decay (Variance)
ε	10-8	Numerical stability (0으로 나누기 방지)
η	10-3	Learning rate

2.3 AdamW (Decoupled Weight Decay)

표준 Adam에서 L2 Regularization은 gradient에 weight decay term을 추가하므로, Adam의 adaptive learning rate 메커니즘에 의해 의도치 않은 상호작용이 발생합니다. AdamW는 weight decay를 gradient update와 분리(decouple)합니다:

여기서 λ는 weight decay coefficient입니다.

비교 항목	Adam + L2	AdamW
Weight decay 적용	Gradient에 포함 (∇L + λw)	Parameter에 직접 적용
Adaptive scaling 영향	Weight decay도 적응적 스케일링 적용받음	Weight decay는 별도 적용
효과	큰 gradient 파라미터의 decay가 약해짐	모든 파라미터에 균일한 decay
일반화 성능	상대적으로 낮음	더 우수 (LLM 학습 표준)

3. Mixed Precision Training

메모리 사용량을 줄이고 연산 속도를 높이기 위해 다양한 정밀도의 부동소수점 데이터 타입을 혼합하여 사용합니다.

Format	Bits	Exponent	Mantissa	Dynamic Range	Use Case
FP32	32	8	23	±3.4×1038	Master weights, accumulation
FP16	16	5	10	±65,504	Forward/backward (Loss Scaling 필수)
BF16	16	8	7	±3.4×1038	LLM 학습 표준 (Loss Scaling 불필요)
FP8 (E4M3)	8	4	3	±448	Forward pass (정밀도 우선)
FP8 (E5M2)	8	5	2	±57,344	Backward pass (범위 우선)

Mixed Precision 학습 방식

Forward/Backward 연산은 BF16/FP16으로 빠르게 수행하고, 파라미터 업데이트 시 누적 오차를 방지하기 위해 Master Weights를 FP32로 유지하여 업데이트합니다:

# PyTorch AMP (Automatic Mixed Precision) 예시
from torch.amp import autocast, GradScaler

scaler = GradScaler()  # FP16 사용 시만 필요 (BF16은 불필요)

for input, target in data:
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.bfloat16):
        output = model(input)       # BF16으로 연산
        loss = loss_fn(output, target)

    loss.backward()                 # Gradient도 BF16
    optimizer.step()                # FP32 master weights 업데이트

4. Learning Rate Schedulers

4.1 Linear Warmup

학습 초기에 불안정한 가중치에 큰 gradient가 곱해져 발산(divergence)하는 것을 방지합니다. 매우 작은 학습률에서 시작하여 목표 학습률까지 선형적으로 증가시킵니다:

4.2 Cosine Decay

Warmup 이후 코사인 곡선을 따라 학습률을 0에 가깝게 서서히 감소시킵니다. 학습 후반부에 미세한 가중치 조정(fine-tuning)이 가능해져 더 좋은 local minima를 찾습니다:

5. Data Parallelism (DP / DDP)

5.1 DataParallel (DP) — 단일 머신

GPU 0이 모든 것을 관장하는 방식입니다:

1. GPU 0에서 batch를 읽고 각 GPU에 mini-batch 분배
2. GPU 0에서 최신 모델을 다른 GPU에 복제
3. 각 GPU에서 forward pass → 출력을 GPU 0으로 전송 → loss 계산
4. GPU 0에서 loss를 각 GPU에 분배 → backward pass
5. 모든 gradient를 GPU 0으로 수집하여 평균

5.2 DistributedDataParallel (DDP) — 멀티 머신

각 GPU가 완전히 독립적으로 동작하며, backward pass 중에 AllReduce로 gradient를 동기화합니다:

1. 메인 프로세스가 GPU 0의 모델을 모든 GPU에 복제 (초기 1회)
2. 각 GPU가 독립적으로 자신의 mini-batch 처리 (forward + backward)
3. Backward pass 중에 AllReduce로 gradient 동기화 (autograd hook)
4. 각 GPU가 동기화된 gradient로 독립적으로 optimizer step

import torch.distributed as dist

# 분산 환경 초기화
dist.init_process_group(
    backend='nccl',      # GPU용 (CPU는 'gloo')
    rank=rank,             # 현재 프로세스 번호
    world_size=world_size  # 전체 GPU 수
)

# DDP 래핑
model = DistributedDataParallel(model, device_ids=[local_rank])

# backward() 반환 시 param.grad는 이미 동기화된 gradient를 포함

6. Tensor Parallelism (TP)

하나의 거대한 행렬 곱셈 연산을 여러 GPU로 쪼개어 분산 연산합니다. Megatron-LM에 의해 널리 알려졌습니다.

6.1 Column-Parallel Linear Layer

행렬 A를 column 방향으로 분할합니다:

각 GPU가 A의 일부 column을 담당합니다. Forward에서 각 GPU가 독립적으로 연산하고, 결과를 column-wise로 concat합니다.

6.2 Row-Parallel Linear Layer

행렬 A를 row 방향으로 분할합니다:

Forward 시 부분합을 더하기 위해 AllReduce가 필요합니다. 일반적으로 Column-parallel 다음에 배치하여 통신을 최소화합니다.

6.3 Transformer에서의 Tensor Parallelism

MLP Block에서의 TP

Column-Parallel Linear → GeLU → Row-Parallel Linear

GeLU는 element-wise 연산이므로 통신이 불필요합니다. Row-parallel 출력 시에만 AllReduce가 발생합니다.

Self-Attention Block에서의 TP

Q, K, V projection: Column-parallel (각 GPU가 attention head의 일부를 담당)
Output projection: Row-parallel
제약: Attention head 수가 TP degree의 배수여야 합니다.

통신 위치	횟수 (per layer)	연산
Forward pass	2번	AllReduce (MLP 1번 + Attention 1번)
Backward pass	2번	AllReduce (MLP 1번 + Attention 1번)

7. Pipeline Parallelism (PP)

모델의 레이어를 순차적으로 여러 GPU에 배치합니다:

GPU 0: Layers 0-7    (Stage 0)
GPU 1: Layers 8-15   (Stage 1)
GPU 2: Layers 16-23  (Stage 2)
GPU 3: Layers 24-31  (Stage 3)

7.1 Pipeline Bubble 문제

Naive한 구현에서는 대부분의 GPU가 대기 상태(idle)입니다:

Time →
GPU 0: [F0][  ][  ][  ][B0][  ][  ][  ]
GPU 1: [  ][F0][  ][  ][  ][B0][  ][  ]
GPU 2: [  ][  ][F0][  ][  ][  ][B0][  ]
GPU 3: [  ][  ][  ][F0][  ][  ][  ][B0]
         ↑ Bubble (유휴 시간) ↑

7.2 Micro-batch를 통한 해결

Mini-batch를 더 작은 Micro-batch로 분할하여 파이프라인에 연속으로 밀어 넣습니다:

Time → (M=4 micro-batches)
GPU 0: [F0][F1][F2][F3][B0][B1][B2][B3]
GPU 1: [  ][F0][F1][F2][F3][B0][B1][B2]
GPU 2: [  ][  ][F0][F1][F2][F3][B0][B1]
GPU 3: [  ][  ][  ][F0][F1][F2][F3][B0]

Bubble 비율 공식:

P = Pipeline stages (GPU 수), M = Micro-batches 수. M이 클수록 bubble 비율이 감소합니다. 예: P=4, M=32이면 bubble은 9.4%에 불과합니다.

7.3 1F1B (One Forward One Backward) Schedule

Forward와 Backward를 interleave하여 activation memory 사용량을 최적화합니다. 모든 micro-batch의 forward를 먼저 수행하는 대신, forward 하나 → backward 하나를 교대로 실행합니다.

8. Sequence / Context Parallelism (SP/CP)

최근 100K~1M 이상의 긴 Context를 처리하기 위해 Sequence 차원을 여러 GPU에 분할합니다.

왜 필요한가?

Self-Attention의 메모리 복잡도는 O(sequence_length²)입니다. 시퀀스 길이가 증가하면 activation memory가 급격히 증가하여 단일 GPU로 처리할 수 없습니다.

Ring Attention

시퀀스를 GPU 수만큼 분할하고, 각 GPU가 자신의 Query chunk에 대해 K, V를 Ring 구조로 순환시키며 Attention을 계산합니다. 통신과 연산을 오버랩하여 효율을 높입니다.

DeepSpeed Ulysses

AlltoAll 통신으로 시퀀스 차원의 분할과 head 차원의 분할을 교환합니다. Flash Attention과 함께 사용하여 효율성을 극대화합니다.

설정	최대 시퀀스 길이
단일 H100 GPU (80GB)	~500K tokens
단일 노드 (8 GPU)	~3.7M tokens
4 노드 (32 GPU)	~15M tokens

9. Expert Parallelism (EP)

Mixture of Experts (MoE) 모델에서 각 Expert를 다른 GPU에 분산 배치합니다.

MoE 구조

Router가 각 토큰을 적절한 Expert에 배정(Top-k routing)합니다. 각 토큰은 일부 Expert만 활성화(sparse activation)되므로, 파라미터 수에 비해 연산량이 적습니다.

All-to-All 통신

1. Token dispatch: 각 GPU의 토큰을 담당 Expert가 있는 GPU로 전송
2. Expert computation: 각 GPU가 자신의 Expert 연산 수행
3. Token combine: 결과를 원래 GPU로 반환

Load Balancing

Expert에 토큰이 불균형하게 배정되면 일부 GPU에 병목이 발생합니다. 이를 방지하기 위한 Auxiliary Loss:

f_i: Expert i에 배정된 토큰 비율, P_i: Router가 Expert i에 배정할 확률, α: balancing coefficient. 이 loss를 main loss에 추가하여 균등 분배를 유도합니다.

10. 3D Parallelism (TP + PP + DP)

대규모 모델 학습에서는 세 가지 parallelism을 계층적으로 조합합니다:

Parallelism	적용 범위	이유	통신 패턴
Tensor Parallelism	노드 내 (NVLink)	고대역폭 필수 (900 GB/s)	AllReduce, AllGather
Pipeline Parallelism	노드 간 (InfiniBand/EFA)	낮은 대역폭 OK	Point-to-point Send/Recv
Data Parallelism	전체 클러스터	가장 효율적 확장	AllReduce (gradient sync)

1 Trillion Parameter 모델 예시