Chapter 3: FSDP, ZeRO, NCCL

분산 학습의 메모리 효율화 기술(ZeRO, FSDP)과 GPU 간 통신의 핵심인 NCCL을 심층적으로 이해합니다. 이 기술들은 Checkpointless Training의 기반이 됩니다.

1. ZeRO 개요

DDP의 메모리 중복 문제

표준 Data Parallelism (DDP)은 각 GPU가 동일한 모델 복사본(Model Replica)을 가지고, 데이터만 분산하여 처리합니다. 각 GPU는 Forward/Backward pass를 독립적으로 수행한 뒤 그래디언트를 All-Reduce로 동기화합니다.

학습 상태(Training State) 구성

Mixed Precision Training (BF16/FP16 + FP32 Master Weights)과 Adam 옵티마이저를 사용할 때, 단일 파라미터당 필요한 메모리는 다음과 같습니다:

이 1.12TB를 단일 GPU(80GB VRAM)에 올리는 것은 불가능합니다. ZeRO는 이 문제를 해결합니다.

2. ZeRO Stage 1: Optimizer State Partitioning

동작 원리

ZeRO Stage 1 ($P_{os}$)은 가장 메모리를 많이 차지하는 옵티마이저 상태(Optimizer States)만 GPU들에 분산(Sharding)합니다.

• 각 GPU는 전체 파라미터 중 자신이 맡은 부분의 옵티마이저 상태만 유지
• 모델 파라미터와 그래디언트는 모든 GPU에 복제 (DDP와 동일)
• 옵티마이저 스텝 후, 업데이트된 파라미터를 All-Gather로 동기화

메모리 절감 계산

{
  "zero_optimization": {
    "stage": 1,
    "reduce_bucket_size": 5e8,
    "allgather_bucket_size": 5e8
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 1e-4,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  }
}

3. ZeRO Stage 2: + Gradient Partitioning

동작 원리

ZeRO Stage 2 ($P_{os+g}$)는 Stage 1에 추가로 그래디언트(Gradients)도 분산합니다.

• Backward pass 후, 전체 그래디언트를 Reduce-Scatter 연산으로 합산 + 분배
• 각 GPU는 자신이 담당하는 파라미터의 그래디언트만 유지
• 옵티마이저 스텝은 각 GPU가 자신의 파라미터 조각에 대해서만 수행

Reduce-Scatter 연산

Reduce-Scatter는 All-Reduce를 두 단계로 분리한 것 중 첫 번째입니다:

메모리 절감

{
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8,
    "allgather_bucket_size": 5e8
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16
  },
  "gradient_clipping": 1.0,
  "train_batch_size": 32,
  "train_micro_batch_size_per_gpu": 4
}

4. ZeRO Stage 3: + Parameter Partitioning

동작 원리

ZeRO Stage 3 ($P_{os+g+p}$)는 모델 파라미터(Parameters)까지 모두 분산합니다. 이것이 PyTorch FSDP와 기술적으로 거의 동일한 구현입니다.

• 각 GPU는 모델의 1/N 조각만 메모리에 상주
• 연산이 필요할 때만 다른 GPU로부터 파라미터를 All-Gather로 가져옴
• 연산 후 즉시 해제하여 메모리 확보

All-Gather on Demand

Forward pass 시 특정 레이어 연산 직전에 해당 레이어의 파라미터를 모든 GPU로부터 수집합니다:

메모리 절감 (N-fold)

{
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8,
    "stage3_prefetch_bucket_size": 5e8,
    "stage3_param_persistence_threshold": 1e6,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "bf16": {
    "enabled": true
  },
  "gradient_clipping": 1.0,
  "train_batch_size": 256,
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 32
}

5. ZeRO-Infinity: NVMe Offloading

개념

ZeRO-Infinity는 ZeRO Stage 3에 NVMe SSD 오프로딩을 추가하여 GPU 메모리를 넘어 시스템 메모리와 NVMe 스토리지까지 활용합니다.

Memory Pool Hierarchy

• Offload Optimizer: 옵티마이저 상태를 CPU 메모리로 오프로드
• Offload Param: 파라미터까지 CPU/NVMe로 오프로드
• NVMe Offload: CPU 메모리도 부족할 때 NVMe SSD 활용
• 비동기 I/O(aio)를 통한 prefetch로 성능 저하 최소화

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 5,
      "fast_init": false
    },
    "offload_param": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 5,
      "buffer_size": 1e8,
      "max_in_cpu": 1e9
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "aio": {
    "block_size": 1048576,
    "queue_depth": 8,
    "thread_count": 1,
    "single_submit": false,
    "overlap_events": true
  },
  "bf16": {
    "enabled": true
  },
  "train_batch_size": 512,
  "gradient_accumulation_steps": 64
}

6. FSDP 동작 원리

FSDP란?

FSDP (Fully Sharded Data Parallel)는 PyTorch의 ZeRO Stage 3 네이티브 구현입니다. 파라미터, 그래디언트, 옵티마이저 상태를 모든 GPU에 분산(Shard)합니다.

Forward Pass

# FSDP Forward Pass 동작 과정

# 1. At Rest: 각 GPU는 1/N 파라미터만 보유
GPU_0: [Shard_0]  GPU_1: [Shard_1]  GPU_2: [Shard_2]  GPU_3: [Shard_3]

# 2. Before Forward: All-Gather로 전체 파라미터 재구성
All-Gather() → 모든 GPU: [Full Parameters]

# 3. Forward Compute: 전체 파라미터로 연산 수행
output = layer(input)  # with full parameters

# 4. After Forward: 사용한 파라미터 해제 (메모리 확보)
del full_parameters  # keep only local shard

Backward Pass

# FSDP Backward Pass 동작 과정

# 1. Before Backward: 다시 All-Gather로 파라미터 재구성
All-Gather() → 모든 GPU: [Full Parameters]

# 2. Backward Compute: 그래디언트 계산
gradients = backward(loss)

# 3. After Backward: Reduce-Scatter로 그래디언트 합산 + 분배
Reduce-Scatter(gradients) → 각 GPU: [1/N Gradient Shard]

# 4. Optimizer Step: 각 GPU가 자신의 샤드만 업데이트
optimizer.step(local_shard)  # only 1/N of parameters

전체 흐름 다이어그램

7. FSDP Sharding Strategies

8. FSDP2 vs FSDP1

주요 차이점

DTensor란?

DTensor (Distributed Tensor)는 PyTorch 2.0에서 도입된 분산 텐서 추상화입니다. 텐서가 여러 디바이스에 어떻게 분산되어 있는지를 메타데이터로 관리합니다.

FSDP2 코드 예시

from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.device_mesh import init_device_mesh

# Device Mesh 초기화 (2D: DP x TP)
mesh = init_device_mesh("cuda", (dp_size, tp_size), mesh_dim_names=("dp", "tp"))

# Mixed Precision 정책
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
)

# FSDP2: 개별 모듈에 fully_shard 적용
for layer in model.transformer.layers:
    fully_shard(layer, mesh=mesh["dp"], mp_policy=mp_policy)

# 최상위 모듈에도 적용
fully_shard(model, mesh=mesh["dp"], mp_policy=mp_policy)

# 이제 model은 FSDP2로 샤딩됨
output = model(input_ids)

9. 메모리 계산 예시

70B 모델 시나리오별 메모리

Activation Memory 공식

Forward pass 시 Backward를 위해 중간 연산 결과(Activation)를 저장해야 합니다:

10. NCCL Collective Operations

NCCL이란?

NCCL (NVIDIA Collective Communications Library)은 분산 GPU 학습에서 GPU 간 통신을 담당하는 고성능 라이브러리입니다. "네트워크 스택"처럼 GPU들이 서로 대화하는 방법을 정의합니다.

주요 Collective Operations

All-Reduce

모든 GPU가 기여하고, 모든 GPU가 합산 결과를 받습니다. DDP에서 그래디언트 동기화에 필수.

# All-Reduce: 모든 GPU의 그래디언트 합산 → 모든 GPU에 동일한 결과
GPU 0: [1, 2, 3]     GPU 0: [10, 20, 30]
GPU 1: [2, 4, 6]  →  GPU 1: [10, 20, 30]  (sum)
GPU 2: [3, 6, 9]     GPU 2: [10, 20, 30]
GPU 3: [4, 8, 12]    GPU 3: [10, 20, 30]

All-Gather

각 GPU의 조각을 모아 전체 텐서를 구성, 모든 GPU가 동일한 전체 텐서를 받습니다. FSDP에서 파라미터 재구성에 사용.

# All-Gather: 각 GPU의 샤드 수집 → 전체 텐서 재구성
GPU 0: [A]           GPU 0: [A, B, C, D]
GPU 1: [B]        →  GPU 1: [A, B, C, D]  (concatenate)
GPU 2: [C]           GPU 2: [A, B, C, D]
GPU 3: [D]           GPU 3: [A, B, C, D]

Reduce-Scatter

모든 GPU가 기여하고, 결과를 N등분하여 각 GPU가 다른 조각을 받습니다. FSDP Backward에서 그래디언트 분배에 사용.

# Reduce-Scatter: 합산 + 분배
GPU 0: [1,2,3,4]     GPU 0: [10] (chunk 0의 합)
GPU 1: [2,4,6,8]  →  GPU 1: [20] (chunk 1의 합)
GPU 2: [3,6,9,12]    GPU 2: [30] (chunk 2의 합)
GPU 3: [4,8,12,16]   GPU 3: [40] (chunk 3의 합)

Broadcast

하나의 GPU(root)가 데이터를 보내고, 모든 GPU가 받습니다. 초기 가중치 분배에 사용.

# Broadcast: Root GPU의 데이터를 모든 GPU로 복사
GPU 0: [W]  (root)   GPU 0: [W]
GPU 1: [?]        →  GPU 1: [W]  (copied from root)
GPU 2: [?]           GPU 2: [W]
GPU 3: [?]           GPU 3: [W]

All-to-All

각 GPU가 다른 모든 GPU에게 서로 다른 데이터를 보냅니다. MoE에서 Expert Parallelism에 사용.

# All-to-All: 각 GPU가 다른 GPU들에게 각각 다른 데이터 전송
GPU 0: [A0,A1,A2,A3]     GPU 0: [A0,B0,C0,D0]
GPU 1: [B0,B1,B2,B3]  →  GPU 1: [A1,B1,C1,D1]
GPU 2: [C0,C1,C2,C3]     GPU 2: [A2,B2,C2,D2]
GPU 3: [D0,D1,D2,D3]     GPU 3: [A3,B3,C3,D3]

11. Ring vs Tree Algorithm

Ring Algorithm

GPU들을 논리적인 링(Ring) 형태로 연결하여 데이터를 순환시킵니다.

# Ring All-Reduce 동작 (4 GPU 예시)

   GPU 0 ←→ GPU 1
     ↑         ↓
   GPU 3 ←→ GPU 2

# 단계 1: 각 GPU가 이웃에게 청크 전송
# 단계 2: 받은 데이터와 로컬 데이터 합산
# 단계 3: N-1번 반복하면 모든 GPU가 전체 합산 결과 보유

# 복잡도
대역폭: O(데이터 크기)  # 데이터 크기에 비례, GPU 수 무관
지연시간: O(N)           # GPU 수에 비례 (단점)

Tree Algorithm (Double Binary Tree)

GPU들을 트리(Tree) 구조로 연결하여 계층적으로 데이터를 집계합니다.

# Tree All-Reduce 동작

        Root (GPU 0)
       /           \
    GPU 1         GPU 2
   /     \       /     \
GPU 3   GPU 4  GPU 5   GPU 6

# Reduce Phase: 리프 → 루트 (합산)
# Broadcast Phase: 루트 → 리프 (분배)

# 복잡도
대역폭: O(데이터 크기)     # Ring과 동일
지연시간: O(log N)         # GPU 수의 로그에 비례 (장점)

비교표

12. NCCL 초기화

TCPStore Rendezvous 과정

분산 학습을 시작하려면 수백~수천 개의 프로세스(GPU)가 서로의 존재와 위치를 알아야 합니다.

# 전통적인 NCCL 초기화 과정

# 1. Master Node (Rank 0)가 TCPStore 서버 시작
Rank 0: TCPStore 서버 오픈 (IP:PORT)

# 2. 모든 Worker가 Master에 연결하여 자신의 정보 등록
Rank 1 → Rank 0: "내 주소는 192.168.1.2:29501"
Rank 2 → Rank 0: "내 주소는 192.168.1.3:29501"
...
Rank N → Rank 0: "내 주소는 ..."

# 3. 모든 Worker가 등록 완료되면 NCCL Unique ID 생성
Rank 0: NCCL Unique ID 생성 및 브로드캐스트

# 4. 각 Rank가 Communicator 형성
ncclCommInitRank(comm, nranks, uniqueId, rank)

init_process_group 코드

import torch.distributed as dist
import os

# 환경변수에서 분산 설정 읽기
rank = int(os.environ['RANK'])
world_size = int(os.environ['WORLD_SIZE'])
master_addr = os.environ['MASTER_ADDR']
master_port = os.environ['MASTER_PORT']

# Process Group 초기화 (TCPStore 기반 rendezvous)
dist.init_process_group(
    backend='nccl',           # GPU 통신용
    init_method=f'tcp://{master_addr}:{master_port}',
    rank=rank,
    world_size=world_size,
)

# 이제 collective operations 사용 가능
tensor = torch.ones(10).cuda()
dist.all_reduce(tensor)  # 모든 GPU의 tensor 합산

13. Topology Discovery

NCCL의 자동 토폴로지 탐색

NCCL은 시작 시 hwloc 라이브러리를 사용하여 하드웨어 구성을 자동으로 탐색합니다.

연결 계층 (P2P Level)

NCCL_TOPO_FILE

복잡한 토폴로지에서는 XML 파일로 직접 토폴로지를 지정할 수 있습니다:

# 토폴로지 파일 지정
export NCCL_TOPO_FILE=/path/to/topology.xml

# 토폴로지 탐색 결과 덤프
export NCCL_TOPO_DUMP_FILE=/tmp/nccl_topo.txt

14. NCCL 환경변수

중요 환경변수 테이블

# AWS EFA (Elastic Fabric Adapter) 최적화 설정

# 기본 EFA 설정
export FI_PROVIDER=efa
export FI_EFA_USE_DEVICE_RDMA=1
export FI_EFA_FORK_SAFE=1

# NCCL EFA 플러그인
export NCCL_NET=aws-ofi-nccl
export NCCL_DEBUG=INFO

# P5 인스턴스 (H100 x8) 최적화
export NCCL_NVLS_ENABLE=1  # NVLink SHARP
export NCCL_IB_TIMEOUT=22
export NCCL_MIN_NCHANNELS=4

# 대역폭 최적화
export NCCL_BUFFSIZE=8388608  # 8MB
export NCCL_P2P_NET_CHUNKSIZE=524288  # 512KB

# 문제 발생 시 상세 로깅
export NCCL_DEBUG=TRACE
export NCCL_DEBUG_SUBSYS=INIT,COLL,P2P,NET
export NCCL_DEBUG_FILE=/tmp/nccl_debug_%h_%p.log

# 토폴로지 확인
export NCCL_TOPO_DUMP_FILE=/tmp/nccl_topo.xml

# Hang 감지 (30초 타임아웃)
export NCCL_TIMEOUT=30
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1

# 디버그 정보 출력 예시
[Rank 0] NCCL INFO Bootstrap: Using eth0:192.168.1.10<6379>
[Rank 0] NCCL INFO Trees [0] -1/-1/-1->0->1 [1] -1/-1/-1->0->1
[Rank 0] NCCL INFO Channel 00 : 0 1 2 3

요약