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는 반드시 고려해야 할 핵심 기술입니다.