ZeRO 옵티마이저

ZeRO 옵티마이저 (Zero Redundancy Optimizer)

모델이 커지면 GPU 메모리에 한 번에 안 들어가는 경우가 생겨요. 그래서 옵티마이저가 관리하는 상태를 여러 GPU에 나눠 담자는 아이디어가 나왔는데, 그게 바로 ZeRO예요. "Zero Redundancy"라는 이름처럼, 원래는 모든 GPU가 중복해서 들고 있던 상태들을 제로 리던던시(zero redundancy)로 만들어요.

출처: DeepSpeed ZeRO Optimizer 튜토리얼 (공식)

ZeRO가 해결하는 문제

학습 중에 GPU가 들고 있어야 할 것은 크게 세 가지예요. 옵티마이저 상태, 그래디언트, 그리고 모델 파라미터. 분산 학습에서 각 GPU가 이걸 전부 복제해 들고 있으면, 모델 하나를 올리기 위해 필요한 메모리가 그대로 N배로 늘어나요. ZeRO는 이 상태들을 데이터 병렬 그룹 안에서 논리적으로 분할해서, 각 GPU가 일부만 들고 있게 만들어요.

핵심 요점은, 코드 변경 없이 설정 JSON만 바꾸면 ZeRO를 켤 수 있다는 거예요. 그래서 기존 학습 코드를 그대로 두고도 훨씬 큰 모델을 학습할 수 있어요.

ZeRO 활성화하기

DeepSpeed 설정 JSON에 zero_optimization 키를 추가하면 ZeRO 최적화가 활성화돼요. 단계(stage)에 따라 어느 상태까지 분할할지 달라져요. ZeRO-2는 옵티마이저 상태와 그래디언트를, ZeRO-3는 파라미터까지 분할해요.

GPT-2 10B 모델을 ZeRO-2로 학습하는 예시 설정은 이런 모양이에요.

--model-parallel-size 1 \
--num-layers 50 \
--hidden-size 4096 \
--num-attention-heads 32 \
--batch-size 1 \
--deepspeed_config ds_zero_stage_2.config \
--checkpoint-activations

활성화 체크포인팅(--checkpoint-activations)까지 켜면 중간 활성화를 저장해 두었다가 역전파 때 다시 계산해서, 메모리를 더 아낄 수 있어요.

ZeRO-Infinity로 트릴리언 스케일

ZeRO-Infinity는 모델 상태 전체를 CPU나 NVMe 메모리로 오프로딩해, GPU 메모리 한계를 넘어서는 모델을 학습할 수 있게 해줘요. 아래처럼 설정에서 오프로딩을 켜면 돼요.

from deepspeed import zero

with deepspeed.zero.Init(
    data_parallel_group=mpu.get_data_parallel_group(),
    remote_device=get_args().remote_device,
    enabled=get_args().zero_stage == 3
):
    model = GPT2Model(num_tokentypes=0, parallel_output=True)

ZeRO-Offload: ZeRO-2에서 옵티마이저·그래디언트 상태를 CPU 메모리로 오프로딩하는 초기 시스템이에요.
ZeRO-Infinity: ZeRO-3에서 쓸 수 있는 차세대 오프로딩 기능이에요. CPU와 NVMe까지 활용해 더 큰 모델을 다뤄요.

ZeRO 옵티마이저