ZeRO (딥 러닝)

Zero Redundancy Optimizer, ZeRO; ZeRO
데이터 병렬 학습에서 GPU마다 똑같이 복제하던 옵티마이저 상태·기울기·가중치를 나눠 갖게 해, 통신량은 거의 그대로 두고 GPU당 메모리를 GPU 수에 반비례하게 줄이는 기법

ZeRO(Zero Redundancy Optimizer, "제로")는 마이크로소프트가 2019년 발표한 대규모 신경망 학습용 메모리 최적화 기법이다.^[1] 마이크로소프트의 오픈 소스 학습 라이브러리 DeepSpeed의 핵심 기능으로 구현되었고, 같은 생각을 파이토치가 자체 구현한 것이 FSDP(Fully Sharded Data Parallel)다.

일반적인 데이터 병렬 학습(신경망 분산 학습 참고)은 GPU마다 모델 전체를 복제한다. GPU를 64장 써도 각 GPU가 가중치, 기울기, 옵티마이저 상태를 모두 들고 있으므로 64벌이 중복된다. ZeRO는 이 중복을 없앤다. 각 GPU는 전체의 $1 / N_{d}$ 만 저장하고, 계산에 필요한 순간에만 다른 GPU에서 받아 왔다가 버린다. 계산 방식은 데이터 병렬 그대로라 모델 코드를 고칠 필요가 없다는 점이 텐서 병렬화나 파이프라인 병렬화와 다르다.

학습 메모리는 어디에 쓰이나

논문은 학습 메모리를 모델 상태(model states)와 잔여 상태(residual states)로 나눈다.^[1]

모델 상태: 가중치, 기울기, 옵티마이저 상태(Adam의 1차·2차 모멘트 등)
잔여 상태: 활성값, 임시 버퍼, 메모리 단편화로 못 쓰는 공간

혼합 정밀도와 Adam으로 매개변수 $Ψ$ 개짜리 모델을 학습하면 모델 상태는 다음과 같다.

항목	자료형	바이트
가중치	FP16/BF16	$2 Ψ$
기울기	FP16/BF16	$2 Ψ$
마스터 가중치 사본	FP32	$4 Ψ$
1차 모멘트(momentum)	FP32	$4 Ψ$
2차 모멘트(variance)	FP32	$4 Ψ$
합계		$(2 + 2 + K) Ψ = 16 Ψ$ , $K = 12$

15억 매개변수 GPT-2만 해도 모델 상태가 최소 24GB로, FP16 가중치만의 3GB보다 훨씬 크다. 논문은 32GB GPU에서 일반 데이터 병렬로는 약 14억 매개변수를 넘는 모델이 메모리 부족에 걸린다고 했다.^[1]

ZeRO-DP의 세 단계

$N_{d}$ 를 데이터 병렬 GPU 수라 하면, 단계가 올라갈수록 나누는 대상이 늘어난다.

단계	나누는 대상	GPU당 모델 상태	7.5B 모델, 64 GPU	통신량(기준 대비)
0 (일반 데이터 병렬)	없음	$(2 + 2 + K) Ψ$	120GB	$2 Ψ$ (1배)
1 ( $P_{o s}$ )	옵티마이저 상태	$4 Ψ + \frac{K Ψ}{N_{d}}$	31.4GB	$2 Ψ$ (1배)
2 ( $P_{o s + g}$ )	+ 기울기	$2 Ψ + \frac{(2 + K) Ψ}{N_{d}}$	16.6GB	$2 Ψ$ (1배)
3 ( $P_{o s + g + p}$ )	+ 가중치	$\frac{(2 + 2 + K) Ψ}{N_{d}}$	1.9GB	$3 Ψ$ (1.5배)

$N_{d}$ 가 클 때 1단계는 약 4배, 2단계는 약 8배 메모리를 줄이고, 3단계는 GPU 수에 비례해 줄인다(64장이면 64배). 표의 7.5B 예시 값은 논문 그림 1과 같다.^[1] 아래 함수로 직접 계산해 볼 수 있다.

def zero_model_state_gb(n_params, n_gpus, stage, k=12):
    """혼합 정밀도 Adam 기준 GPU 한 장의 모델 상태 메모리(GB)"""
    p, g, o = 2 * n_params, 2 * n_params, k * n_params
    if stage >= 1: o /= n_gpus     # 옵티마이저 상태 분할
    if stage >= 2: g /= n_gpus     # 기울기 분할
    if stage >= 3: p /= n_gpus     # 가중치 분할
    return (p + g + o) / 1e9

for s in range(4):
    print(s, round(zero_model_state_gb(7.5e9, 64, s), 1))   # 120.0, 31.4, 16.6, 1.9

이 값은 모델 상태만이다. 활성값은 따로 들며, 이는 그래디언트 체크포인팅이나 텐서 병렬화로 줄인다.

왜 통신량이 늘지 않는가

일반 데이터 병렬은 역전파 뒤 기울기를 올 리듀스로 평균 낸다. 대역폭 최적 올 리듀스(링 올 리듀스 등)는 리듀스-스캐터(reduce-scatter)와 올-개더(all-gather) 두 단계로 이루어지고, 각 단계에서 GPU 하나가 주고받는 양이 약 $Ψ$ 이므로 합계 $2 Ψ$ 다.

1·2단계: 각 GPU는 자기가 맡은 $1 / N_{d}$ 조각의 기울기만 있으면 되므로 올 리듀스 대신 리듀스-스캐터( $Ψ$ )만 한다. 그 조각으로 자기 몫의 가중치를 갱신한 뒤 올-개더( $Ψ$ )로 갱신된 가중치를 모두에게 나눈다. 합계 $2 Ψ$ 로 일반 데이터 병렬과 같다. 올 리듀스를 두 단계로 쪼개 그 사이에 갱신을 끼워 넣은 셈이다.
3단계: 가중치도 조각으로만 있으므로 순전파 때 층마다 올-개더( $Ψ$ ), 역전파 때 다시 올-개더( $Ψ$ ), 기울기 리듀스-스캐터( $Ψ$ )로 합계 $3 Ψ$ , 기준의 1.5배다.^[1]

3단계의 한 학습 단계는 다음 순서로 진행된다.

순전파: 층 $l$ 을 계산하기 직전에 그 층의 가중치 조각을 올-개더해 전체 가중치를 만든다. 계산이 끝나면 다른 GPU 몫은 버린다. 다음 층의 올-개더를 미리 시작해(prefetch) 통신과 계산을 겹친다.
역전파: 층 $l$ 의 가중치를 다시 올-개더해 기울기를 계산하고, 그 기울기를 리듀스-스캐터해 각 GPU가 자기 조각의 기울기 합만 받는다.
갱신: 각 GPU가 자기 조각의 FP32 마스터 가중치와 옵티마이저 상태로 갱신한다.

ZeRO-R: 잔여 메모리 최적화

같은 논문은 모델 상태 밖의 메모리도 다뤘다.^[1]

활성값 분할: 텐서 병렬(메가트론 방식)에서 GPU마다 복제되던 활성값 체크포인트를 나눠 저장하고, 필요하면 CPU로 내린다.
고정 크기 버퍼: 기울기 통신용 임시 버퍼 크기를 모델 크기와 무관하게 고정한다.
단편화 방지: 수명이 긴 텐서(체크포인트, 기울기)와 짧은 텐서를 구분해 연속된 메모리에 미리 배치한다.

논문은 이를 합쳐 1000억 매개변수가 넘는 모델을 GPU 400장에서 초선형(super-linear) 속도 향상으로 학습해 15PFLOPS를 냈고, 모델 병렬 없이 130억 매개변수까지 학습할 수 있다고 보고했다. 마이크로소프트의 170억 매개변수 언어 모델 Turing-NLG가 ZeRO로 학습되었다.^[1]

후속 기법

이름	발표	핵심 아이디어	보고된 결과
ZeRO-Offload	2021	옵티마이저 상태와 갱신 계산을 CPU 메모리·CPU로 넘김. GPU는 순전파·역전파만. 최적화한 CPU Adam과 한 단계 지연 갱신(one-step delayed update)으로 CPU 병목 완화	V100 한 장으로 130억 매개변수 학습, 100억 매개변수 모델에서 GPU당 40TFLOPS^[2]
ZeRO-Infinity	2021	3단계 위에서 가중치·기울기·옵티마이저 상태를 CPU 메모리와 NVMe SSD까지 내림	DGX-2 한 대로 1조 매개변수 미세 조정, V100 512장에서 25PFLOPS 이상^[3]
ZeRO++	2023	가중치 올-개더를 블록 양자화(qwZ), 노드 안에 가중치 2차 사본을 둬 역전파 올-개더를 노드 내부로 한정(hpZ), 기울기를 양자화해 올-투-올로 평균(qgZ)	ZeRO 통신량 4배 감소, GPU 384장에서 처리량 최대 2.16배^[4]

ZeRO-Offload에서 옵티마이저 오프로드는 1~3단계 모두에서, 가중치 오프로드는 3단계에서만 쓸 수 있다.^[5]

DeepSpeed에서 쓰기

설정 파일의 zero_optimization.stage로 단계를 고른다(0은 끔).^[5]

{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "bf16": { "enabled": true },
  "gradient_clipping": 1.0,
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8,
    "stage3_prefetch_bucket_size": 5e8,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": { "device": "none" }
  }
}

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(), config="ds_config.json")

for batch in loader:
    loss = model_engine(batch["input"], batch["labels"])
    model_engine.backward(loss)   # 기울기 분할·통신은 엔진이 처리
    model_engine.step()           # 자기 조각만 갱신하고 가중치를 다시 모음

deepspeed --num_gpus 8 train.py처럼 전용 실행기로 띄운다. 허깅 페이스 Transformers의 Trainer와 Accelerate도 DeepSpeed 설정 파일을 받아 같은 기능을 쓴다.

FSDP와의 관계

PyTorch FSDP 논문은 FSDP가 DeepSpeed의 ZeRO에서 출발했지만 PyTorch의 다른 구성 요소와 맞게 설계와 구현을 새로 했다고 밝힌다. FSDP는 모델을 단위(unit)로 나누고 각 단위 안의 매개변수를 평탄화해 조각낸 뒤, 계산 직전에 모으고 직후에 버린다.^[6]

ZeRO	FSDP1(`ShardingStrategy`)	FSDP2(`fully_shard`)
0단계(일반 DDP)	`NO_SHARD`	해당 없음(DDP 사용)
2단계	`SHARD_GRAD_OP`	`reshard_after_forward=False`
3단계	`FULL_SHARD`	`reshard_after_forward=True`(루트 아닌 모듈의 기본값)
노드 안에서만 3단계, 노드 사이는 복제	`HYBRID_SHARD`	2차원 DeviceMesh, 또는 `reshard_after_forward`에 정수(노드 크기)
ZeRO-Offload	`CPUOffload`	`CPUOffloadPolicy`

FSDP2는 FSDP1의 평탄화(flat parameter) 대신 매개변수마다 0번 차원으로 나눈 DTensor로 표현해 다루기 쉽게 했다.^[7]

import torch
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

# torchrun --nproc_per_node 8 train.py 로 실행, 프로세스 그룹은 미리 초기화
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
for block in model.layers:          # 트랜스포머 블록마다 하나의 통신 단위
    fully_shard(block, mp_policy=mp)
fully_shard(model, mp_policy=mp)    # 나머지(임베딩, 출력층)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)  # DTensor 매개변수로 생성
for x, y in loader:
    loss = loss_fn(model(x), y)
    loss.backward()                 # 블록 단위 all-gather / reduce-scatter
    optimizer.step(); optimizer.zero_grad()

블록 단위로 감싸는 것이 중요하다. 모델 전체를 한 덩어리로 감싸면 순전파 시작 때 모든 가중치를 한꺼번에 모으게 되어 메모리 절감 효과가 사라진다.

다른 병렬화와 비교

방식	나누는 것	모델 코드 변경	통신	잘 맞는 상황
데이터 병렬(DDP)	데이터만	없음	기울기 올 리듀스	모델이 GPU 한 장에 들어갈 때
ZeRO 1·2	데이터 + 옵티마이저 상태·기울기	없음	DDP와 같음	옵티마이저 상태 때문에 메모리가 모자랄 때
ZeRO 3 / FSDP	데이터 + 모델 상태 전부	거의 없음(감쌀 단위 지정)	DDP의 1.5배, 층마다 올-개더	가중치 자체가 한 장에 안 들어갈 때
텐서 병렬화	층 안의 행렬	층 구조에 맞춘 분할 필요	층마다 올 리듀스(활성값)	노드 안 고대역폭 연결, 활성값까지 줄여야 할 때
파이프라인 병렬(GPipe, 파이프드림)	층 묶음	단계 분할 필요	단계 경계의 활성값 전송	노드 사이 대역폭이 낮을 때

ZeRO 3·FSDP는 GPU당 배치가 작을수록 계산에 비해 가중치 올-개더 비율이 커져 효율이 떨어진다. 그래서 수천 장 규모에서는 노드 안은 텐서 병렬, 노드 사이는 FSDP·ZeRO 또는 파이프라인 병렬로 섞는 다차원 병렬화를 쓴다.

같이 보기

각주

↑ ^1.0 ^1.1 ^1.2 ^1.3 ^1.4 ^1.5 ^1.6 Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He, 〈ZeRO: Memory Optimizations Toward Training Trillion Parameter Models〉, SC20. arXiv:1910.02054
↑ Jie Ren 외, 〈ZeRO-Offload: Democratizing Billion-Scale Model Training〉, USENIX ATC 2021. arXiv:2101.06840
↑ Samyam Rajbhandari 외, 〈ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning〉, SC21. arXiv:2104.07857
↑ Guanhua Wang 외, 〈ZeRO++: Extremely Efficient Collective Communication for Giant Model Training〉, 2023. arXiv:2306.10209
↑ ^5.0 ^5.1 DeepSpeed, 〈DeepSpeed Configuration JSON〉. deepspeed.ai
↑ Yanli Zhao 외, 〈PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel〉, VLDB 2023. arXiv:2304.11277
↑ PyTorch, 〈torch.distributed.fsdp.fully_shard〉, 2.14 문서. docs.pytorch.org

참고 문헌

[zero-1] 1.0 ^1.1 ^1.2 ^1.3 ^1.4 ^1.5 ^1.6 Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He, 〈ZeRO: Memory Optimizations Toward Training Trillion Parameter Models〉, SC20. arXiv:1910.02054

[offload-2] Jie Ren 외, 〈ZeRO-Offload: Democratizing Billion-Scale Model Training〉, USENIX ATC 2021. arXiv:2101.06840

[infinity-3] Samyam Rajbhandari 외, 〈ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning〉, SC21. arXiv:2104.07857

[zeropp-4] Guanhua Wang 외, 〈ZeRO++: Extremely Efficient Collective Communication for Giant Model Training〉, 2023. arXiv:2306.10209

[dscfg-5] 5.0 ^5.1 DeepSpeed, 〈DeepSpeed Configuration JSON〉. deepspeed.ai

[fsdp-6] Yanli Zhao 외, 〈PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel〉, VLDB 2023. arXiv:2304.11277

[fsdp2-7] PyTorch, 〈torch.distributed.fsdp.fully_shard〉, 2.14 문서. docs.pytorch.org

[1]

[2]

[3]

[4]

[5]

[6]

[7]