DeepSpeed ZeRO-3

개요

DeepSpeed ZeRO-3(Zero Redundancy Optimizer Stage 3)는 초대규모 언어 모델 학습 시 발생하는 메모리 병목을 극복하기 위한 Microsoft의 DeepSpeed 프레임워크 구성 요소입니다. 모델 파라미터, 옵티마이저 상태, 그래디언트를 모두 분산 저장하여 GPU 메모리 사용을 획기적으로 줄이고, 100B+ 파라미터 모델의 학습을 단일 또는 소수 GPU 환경에서도 가능하게 만듭니다.

이 글에서는 ZeRO-3의 구조, 작동 방식, 핵심 기술, 장점 및 실무 적용 전략을 살펴봅니다.

---

1. 개념 및 정의

항목	설명
정의	ZeRO-3는 학습 시 모든 파라미터 관련 메타정보를 GPU 간 분산 저장하여 메모리 중복을 제거하는 분산 최적화 기법입니다.
목적	대규모 모델의 학습 가능성을 확보하면서도 효율적인 자원 활용을 실현
필요성	기존 DDP, ZeRO-1, 2 등은 전체 파라미터의 일부만 샤딩 → 한계 존재

ZeRO-3는 모델 병렬화 없이도 모든 메모리 병목을 제거할 수 있는 강력한 전략입니다.

---

2. 특징

특징	설명	기존 방식과 비교
전체 파라미터 샤딩	파라미터, 옵티마이저 상태, 그래디언트를 모두 분산 저장	ZeRO-1/2는 일부만 샤딩
Offload 지원	CPU/NVMe로 메모리 이동 가능	GPU 사용량 급감, 느린 속도는 trade-off
Parameter Partitioning	각 GPU는 해당 shard만 로컬 보유	메모리 중복 완전 제거

ZeRO-3는 기존 메모리 최적화 방식의 정점을 구현한 Stage입니다.

---

3. 아키텍처 및 작동 방식

단계	설명	핵심 처리
초기화	모델 파라미터를 각 GPU에 나눠서 저장	shard 단위 분배
Forward	필요한 shard만 All-Gather 후 순전파 계산	나머지는 로컬에서 로딩 안 함
Backward	그래디언트 계산 후 Reduce-Scatter로 나눠 저장	통신 효율적으로 수행
Optimizer Step	옵티마이저 상태도 분산된 shard 기준 업데이트	offload 시 CPU에서 수행 가능

통신과 연산 흐름을 파이프라인화하여 성능을 최적화합니다.

---

4. 기술 요소

기술 요소	설명	효과
NVMe Offload	옵티마이저 상태를 디스크로 이전	GPU 메모리 사용 최소화
Activation Checkpointing	중간 활성값을 재계산하여 저장 공간 절약	메모리 사용량 감소 vs 연산량 증가
ZeRO-Inference 연계	학습 모델을 추론 최적화 구조로 변환	통합 모델 추론 배포 가능
CPU Offload	일부 모듈을 CPU로 이동하여 GPU 사용량 절감	CPU I/O 성능 의존성 고려 필요

DeepSpeed는 PyTorch와 완전히 호환되며, 다양한 전략을 함께 사용할 수 있습니다.

---

5. 장점 및 이점

장점	설명	기대 효과
학습 가능한 모델 크기 확대	기존 GPU 대비 10배 이상 큰 모델도 학습 가능	예: GPT-175B를 64 GPU로 학습 가능
메모리 중복 제거	전체 파라미터 중복 없음 → 메모리 효율 극대화	GPU 비용 절감
통합 프레임워크	DeepSpeed 내에서 ZeRO, Offload, Mixed Precision 등 통합 운영 가능	설정 일관성 확보

ZeRO-3는 초대형 모델을 위한 효율성과 실용성을 모두 갖춘 전략입니다.

---

6. 활용 사례 및 고려사항

활용 사례	설명	고려사항
GPT-3 학습	Megatron-LM과 함께 대규모 모델 구성 시 활용	통신 병목, FP16 안정성 검토 필요
사내 모델 미세조정	제한된 GPU 자원에서 LLaMA 등 모델 튜닝	오프로딩 비율과 속도 트레이드오프 고려
저비용 연구 환경	연구소, 학교 등 GPU 수 제한 환경에서 10B+ 모델 실험	Checkpointing 조합 전략 중요

학습 속도 ↔ 메모리 절약 간의 균형 설정이 최적화 성패를 좌우합니다.

---

7. 결론

DeepSpeed ZeRO-3는 초대규모 모델 학습을 실현하는 데 필수적인 기술로, 메모리 최적화, 통신 효율, 오프로드 기능 등 다양한 전략을 결합해 자원 효율과 확장성을 극대화합니다. 특히 100B+ 모델을 학습 또는 튜닝하고자 하는 조직이나 연구소에는 GPU 자원 효율화를 위한 최적의 해법입니다.

대규모 모델 학습 비용을 낮추고 GPU 한계를 극복하고자 한다면, ZeRO-3는 반드시 고려해야 할 핵심 기술입니다.