초대형 AI 훈련, 메모리 걱정 뚝: ZeRO 3단계 완전 정복 가이드

초거대 AI 훈련, 메모리 병목 현상 돌파구: ZeRO 3단계 완전 정복

AI 모델의 크기가 기하급수적으로 커지면서, 이를 훈련하기 위한 GPU 메모리 요구량 또한 천정부지로 치솟고 있습니다. 이는 마치 도시 전체를 레고로 만드는 초대형 프로젝트에서, 각 조립가(GPU)가 자신의 작은 작업대(GPU 메모리) 위에 도시 전체에 필요한 모든 종류의 블록(모델 상태 정보)을 전부 올려놓고 작업하려는 것과 같습니다. 당연히 작업 공간은 순식간에 가득 차 버리겠죠. 이러한 메모리 병목 현상을 해결하기 위해 ZeRO(Zero Redundancy Optimizer)라는 혁신적인 기술이 등장했습니다. ZeRO는 여러 GPU가 협력하여 모델을 훈련할 때, 각 GPU가 중복으로 가지고 있어야 했던 정보들을 효율적으로 나누어 관리함으로써, GPU당 메모리 사용량을 획기적으로 줄여줍니다. 마치 각 조립가가 자신이 맡은 구역의 블록만 작업대에 두고, 필요할 때 다른 조립가와 블록을 공유하거나 중앙 창고에서 가져오는 것과 유사합니다. 이제 ZeRO의 3단계 최적화 과정을 통해 어떻게 메모리 부담을 덜 수 있는지 살펴보겠습니다.

ZeRO 1단계: 옵티마이저 상태 분할

개념: 첫 단계는 모델 학습 전략의 핵심인 '옵티마이저 상태(Optimizer States)' 정보의 중복을 제거하는 것입니다. 기존 데이터 병렬(DP) 방식에서는 모든 GPU가 이 옵티마이저 상태 정보 전체를 각자 복사하여 보유해야 했습니다. 모델이 클수록 이 정보의 양도 방대해져 메모리 부담이 컸습니다. ZeRO 1단계에서는 이 옵티마이저 상태를 여러 GPU에 걸쳐 분할하여, 각 GPU는 자신이 담당하는 모델 파라미터 일부에 대한 옵티마이저 상태만 저장합니다.

비유: 레고 도시 조립 프로젝트에서, 각 조립가(GPU)는 더 이상 도시 전체의 모든 건물(모델 파라미터)에 대한 상세 조립 계획(옵티마이저 상태) 전체를 가질 필요가 없습니다. 대신, 자신이 맡은 구역의 건물들에 대한 조립 계획 부분만 갖게 됩니다.

메모리 절감: 이로 인해 각 GPU는 옵티마이저 상태 저장을 위한 메모리 공간을 크게 아낄 수 있습니다. 'Baseline'에서 녹색 부분(Optimizer States)이 'Pos' 단계에서 크게 줄어든 것을 확인할 수 있습니다. 예시에서 GPU당 120GB가 필요했던 메모리가 31.4GB로 감소했습니다. GPU 간 통신량은 기존과 동일하게 유지됩니다.

ZeRO 2단계: 그래디언트 분할 추가

개념: 2단계에서는 1단계의 옵티마이저 상태 분할에 더해, 모델 학습 과정에서 계산되는 '그래디언트(Gradients)' 정보까지 분할합니다. 그래디언트는 모델 파라미터를 올바른 방향으로 업데이트하기 위한 중요한 지표인데, 이 역시 기존에는 모든 GPU에 중복 저장되었습니다. ZeRO 2단계는 이 그래디언트 또한 각 GPU가 담당하는 파라미터 부분에 대해서만 저장하도록 합니다.

비유: 레고 도시 조립에서, 이제 각 조립가는 자신이 만든 건물 구역에 대한 품질 검토팀의 '수정 지시사항(그래디언트)'만 선택적으로 받습니다. 이전에는 모든 조립가가 도시 전체의 모든 수정 지시를 다 받고 그중 필요한 것을 골라야 했다면, 이제는 훨씬 효율적으로 필요한 피드백만 받는 것입니다.

메모리 절감: 옵티마이저 상태에 이어 그래디언트까지 분할되므로 메모리가 더욱 절약됩니다. 이미지의 'Pos+g' 단계에서 주황색 부분(Gradients)이 추가로 줄어든 것을 볼 수 있습니다. 예시에서 메모리 사용량은 31.4GB에서 16.6GB로 더욱 감소했습니다. 통신량은 여전히 기존과 동일합니다.

ZeRO 3단계: 모델 파라미터 분할 추가

개념: 가장 강력한 최적화 단계인 3단계는 옵티마이저 상태와 그래디언트뿐만 아니라, '모델 파라미터(Parameters)' 자체까지 분할합니다. 이는 모델을 구성하는 핵심 데이터까지 여러 GPU가 나누어 갖는다는 의미입니다. 각 GPU는 대부분의 시간 동안 자신이 담당하는 파라미터의 일부만 메모리에 유지합니다. 특정 계층의 계산이 필요할 때만 해당 파라미터들을 모든 관련 GPU로부터 일시적으로 불러와 사용하고, 계산 후 즉시 해제합니다.

비유: 레고 도시 조립의 최종 단계입니다. 이제 각 조립가는 자신이 맡은 구역의 실제 레고 블록(모델 파라미터) 더미만 자신의 작업대에 둡니다. 도시 전체의 모습을 보거나(순방향/역방향 계산) 다른 구역과 연결 작업을 할 때만, 필요한 블록들을 다른 조립가들로부터 잠시 가져와 사용하고 돌려줍니다. 이렇게 하면 각 조립가는 아주 작은 작업 공간만으로도 거대한 레고 도시를 완성할 수 있게 됩니다.

메모리 절감: 모델의 모든 주요 구성 요소가 분할되므로 메모리 절감 효과가 극대화됩니다. 이미지의 'Pos+g+p' 단계에서 파란색 부분(Parameters)마저 극적으로 줄어든 것을 볼 수 있습니다. 예시에서 메모리 사용량은 16.6GB에서 단 1.9GB로 획기적으로 감소했습니다. 이론적으로 GPU 수에 비례하여 메모리 절감 효과가 커져, 엄청난 크기의 모델 훈련도 가능해집니다. 다만, 파라미터를 자주 주고받아야 하므로 GPU 간 통신량은 기존 대비 약 1.5배로 다소 증가하는 트레이드오프가 있습니다.

마무리하며

ZeRO는 단계별 최적화를 통해 거대 모델 훈련 시 발생하는 극심한 메모리 부족 문제를 해결하는 효과적인 접근법을 제시합니다. 각 단계는 트레이드오프를 고려하여 선택적으로 적용될 수 있으며, 이를 통해 제한된 하드웨어 자원에서도 이전에는 불가능했던 규모의 AI 모델 연구와 개발을 가능하게 합니다.

---

https://arxiv.org/abs/1910.02054

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models