Memory Optimization

ZeRO: 1조개 파라미터 모델 학습을 위한 메모리 최적화

ZeRO (Zero Redundancy Optimizer)는 수십억에서 수조 개의 파라미터를 가진 대규모 딥러닝 모델 학습의 메모리 한계를 극복하기 위해 개발된 최적화 기술입니다. 기존의 Data Parallelism (DP)이 모델 상태(optimizer states, gradients, parameters)를 모든 장치에 복제하여 메모리 비효율성을 야기하는 문제를 해결합니다. ZeRO는 이 모델 상태들을 여러 데이터 병렬 프로세스에 걸쳐 분할(partitioning)하여 메모리 중복을 제거합니다. 이 최적화는 세 단계로 구성됩니다: 1) Optimizer State Partitioning (Pos), 2) Gradient Partitioning (Pos+g), 3) Parameter Partitioning (Pos+g+p). 이 단계들을 통해 메모리 사용량을 크게 줄이면서도, DP의 높은 연산 및 통신 효율성은 유지합니다. 결과적으로 ZeRO는 기존 하드웨어에서 훨씬 더 큰 모델을 효율적으로 훈련할 수 있게 하며, 1000억개 파라미터 모델에서 SOTA 대비 10배의 성능 향상을 보여주었고, 세계에서 가장 큰 언어 모델인 Turing-NLG (17B)를 탄생시켰습니다. 논문 제목: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models