DeepSpeed로 대규모 모델 학습 가속화

Stage 1: 옵티마이저 상태(Optimizer States)를 GPU 간에 분할
Stage 2: 옵티마이저 상태와 그래디언트(Gradients)를 분할
Stage 3: 옵티마이저 상태, 그래디언트, 모델 파라미터를 모두 분할
Offload: 그래디언트와 옵티마이저 상태를 CPU/디스크로 오프로드하여 메모리 사용량 최소화

Hugging Face의 Accelerate 라이브러리를 통해 DeepSpeed ZeRO 기술을 적용하여 대규모 모델 학습 시 메모리 문제를 해결하는 방법을 소개한다.

대규모 모델 학습 시 발생하는 OOM(Out of Memory) 문제를 해결하기 위해 DeepSpeed의 ZeRO(Zero Redundancy Optimizer) 기술을 활용하는 방법을 다룬다.

ZeRO 기술은 다음과 같은 단계로 메모리 효율을 높인다:

Hugging Face의 Accelerate 라이브러리를 사용하면 별도의 코드 수정 없이 accelerate config 설정을 통해 DeepSpeed ZeRO를 즉시 적용할 수 있다.

실제 벤치마크 결과, DeepSpeed ZeRO Stage-2를 적용했을 때 기존 DDP(Distributed Data Parallel) 방식 대비 GPU당 배치 사이즈를 5배 더 크게 가져갈 수 있어 학습 효율이 대폭 개선됨을 확인했다.