ZeRO 의 역할과 최적화 단계

ZeRO 는 데이터 병렬 (data parallelism) 처리의 계산 및 메모리 자원을 활용하여 모델 학습 시 각 장치 (GPU) 의 메모리와 계산 요구 사항을 줄여줍니다. 이를 위해 ZeRO 는 모델 학습에 필요한 다양한 상태 (가중치, 그래디언트, 옵티마이저 상태) 를 분산된 학습 하드웨어 (GPU 및 CPU) 에 걸쳐 분할하여 각 GPU 의 메모리 사용량을 감소시킵니다.

구체적으로, ZeRO 는 최적화를 단계별로 점진적으로 구현하고 있으며, 초기 단계에서 적용된 최적화는 이후 단계에서도 계속 사용할 수 있습니다. ZeRO 에 대한 더 자세한 내용은 논문 을 참고하세요.

ZeRO-DP(Data Parallelism) 는 기존의 Data Parallelism 과 Model Parallelism 이 가진 메모리 문제를 해결하기 위한 솔루션입니다. 일반적인 Data Parallelism 에서는 모델 상태 (model states) 가 각 디바이스에 복제되면서 메모리 낭비가 발생하는데, ZeRO-DP 는 이러한 model states 를 분할하여 메모리 중복을 제거합니다.

ZeRO-DP 는 디바이스의 메모리가 충분히 크다면 매우 큰 모델도 처리할 수 있습니다. ZeRO 는 세 가지 주요 최적화 단계를 통해 이를 실현합니다.

B) 주요 최적화 단계

Optimizer States Partitioning (Stage-1): 4 배의 메모리 절감
Optimizer States Partitioning + Gradient Partitioning (Stage-2): 8 배의 메모리 절감
Optimizer States Partitioning + Gradient Partitioning + Parameter Partitioning (Stage-3): GPU 개수에 따라 최대 64 배까지 메모리 절감 가능 (예: 64 개의 GPU 사용 시). 다만, 통신량 (Comm. Volume) 은 약 50% 증가합니다.

C) ZeRO-1

ZeRO-1 은 Optimizer States Partitioning 기법을 적용한 첫 번째 단계로, 주로 optimizer 상태를 분할하여 메모리를 절약하는 방식입니다.

D) ZeRO-2

ZeRO-2 는 기존 ZeRO-1 보다 더 큰 모델을 학습할 수 있도록 설계되어 있으며, 학습 효율성을 두 배로 향상시킵니다. 특히 Residual States 문제를 해결하고 gradient, activation 및 fragmented memory(메모리 단편화) 와 같은 다양한 학습 중 발생하는 메모리 소비 문제를 개선합니다.

D.1) 세부 최적화 내용

Model State Memory: ZeRO-1 에서 사용된 Optimizer States Partitioning(Stage-1) 에 추가적으로 Gradient Partitioning(Stage-2) 을 적용하여 디바이스당 소비되는 메모리를 두 배 더 줄입니다.
Activation Memory: Activation partitioning 기법을 통해 model parallelism 에서 발생하는 activation 복제를 제거합니다. 만약 모델 크기가 매우 크거나 GPU 메모리가 부족한 경우에는 activation 데이터를 CPU 로 옮겨서 처리할 수도 있습니다.
Fragmented Memory: 학습 중 tensor 들의 생명 주기가 달라지면서 발생하는 메모리 단편화를 해결합니다. 이로 인해 가용한 연속적인 공간이 부족해지는 상황에서도 효율적으로 공간을 활용할 수 있게 됩니다.
Constant Size Buffer: All-reduce 과정에서 사용하는 Bucket 사이즈를 일정하게 유지함으로써 성능을 최적화합니다. Bucket 사이즈가 어느 정도 크게 유지되기만 해도 충분히 좋은 효율성을 얻을 수 있습니다.

E) With Megatron

Megatron-LM 은 tensor-slicing model parallelism 를 사용한다. Megatron-LM 의 tensor-slicing model parallelism 에 DeepSpeed ZeRO-2 의 data parallelism 을 추가하면 DeepSpeed 가 Megatron-LM 보다 10x 더 빠르다.

F) ZeRO

G) Zero2

zero2-bf16.yaml

H) DeepSpeed ZeRO: 메모리 최적화 기법

ZeRO(Zero Redundancy Optimizer)는 데이터 병렬 처리(Data Parallelism) 상황에서 발생하는 메모리 중복을 제거하여 각 GPU의 메모리 부담을 획기적으로 줄이는 기술입니다. 훈련에 필요한 구성요소(모델 파라미터, 그래디언트, 옵티마이저 상태)를 여러 GPU에 분산하여 저장하는 방식입니다.

구분	ZeRO-1	ZeRO-2	ZeRO-3
분할 대상	옵티마이저 상태 (Optimizer States)	ZeRO-1 대상 + 그래디언트 (Gradients)	ZeRO-2 대상 + 모델 파라미터 (Parameters)
메모리 절약	약 4배	약 8배	GPU 수(N)에 비례하여 크게 증가
핵심 아이디어	옵티마이저 상태(예: Adam의 모멘텀, 분산)는 훈련 중 메모리를 많이 차지하지만, 업데이트 시에만 필요합니다. 이를 GPU들에 분할 저장하여 메모리 중복을 없앱니다.	그래디언트 또한 파라미터 업데이트 시에만 전체가 필요하므로, 이 역시 분할하여 메모리 사용량을 더욱 줄입니다.	모델의 파라미터(가중치) 자체를 분할합니다. 각 GPU는 전체 모델의 일부 파라미터만 보유하고, 순전파/역전파 시 필요한 파라미터를 다른 GPU로부터 가져와 사용합니다. 가장 높은 수준의 메모리 절약을 제공합니다.

I) 모델 병렬 처리 기법

모델 병렬 처리는 모델 자체가 너무 커서 단일 GPU에 올라가지 않을 때, 모델을 여러 GPU로 분할하는 기술입니다.

구분	데이터 병렬 (DP) Data Parallelism	텐서 병렬 (TP) Tensor Parallelism	파이프라인 병렬 (PP) Pipeline Parallelism	시퀀스 병렬 (SP) Sequence Parallelism
개념	동일한 모델을 여러 GPU에 복제하고, 데이터를 분할하여 처리	모델의 단일 레이어(연산)를 여러 GPU로 분할하여 병렬 처리	모델을 여러 레이어 청크로 나누어 각기 다른 GPU에 할당	Transformer의 입력 시퀀스를 분할하여 병렬 처리
분할 단위	미니배치 데이터	행렬 연산 (텐서)	모델 레이어 그룹	입력 데이터의 시퀀스 차원
장점	구현이 가장 간단하고 일반적입니다.	매우 큰 레이어를 가진 모델(Wide Model)에 효과적입니다. GPU 간 통신량이 비교적 적습니다.	매우 깊은 모델(Deep Model)에 효과적입니다. 활성화(Activation) 메모리를 절약할 수 있습니다.	긴 시퀀스(Long Context)에서 급증하는 Attention 계산량과 활성화 메모리를 크게 줄여줍니다.
단점	각 GPU가 모델 전체를 복제해야 하므로 메모리 비효율성이 발생합니다.	구현이 복잡하고, 일부 레이어에만 적용 가능합니다.	파이프라인 버블(bubble)로 인한 GPU 유휴 시간(idle time)이 발생할 수 있습니다.	Transformer 아키텍처에 특화되어 있으며, 텐서 병렬(TP)과 함께 사용될 때 가장 효과적입니다.
주요 적용 대상	대부분의 딥러닝 훈련	거대한 은닉층을 가진 모델 (예: GPT의 MLP, Attention)	수백 개의 레이어를 가진 초거대 모델 (예: GPT-3)	문서 요약, 긴 대화 등 긴 컨텍스트를 처리하는 모델

*사용자께서 질문하신 컨텍스트 병렬 처리(Context Parallelism) 는 일반적으로 시퀀스 병렬 처리(Sequence Parallelism) 와 동일한 개념으로 사용됩니다. 특히 Transformer 모델에서 입력의 길이를 ‘Context’라고 부르기 때문입니다.

I.1.1) 기술들의 조합: 어떻게 함께 사용되는가?

현대의 초거대 모델 훈련은 위 기술들을 복합적으로 사용합니다.

기본은 데이터 병렬(DP): 훈련 속도 향상을 위해 여러 GPU로 데이터를 분산 처리하는 것이 기본입니다.
메모리 부족 해결사, ZeRO: DP 환경에서 각 GPU의 메모리가 부족하면 ZeRO를 적용하여 메모리 사용량을 줄입니다. (예: ZeRO-DP)
모델이 너무 클 때, PP와 TP: ZeRO를 사용해도 모델이 단일 GPU 메모리에 맞지 않으면, 모델 자체를 분할합니다.
- 파이프라인 병렬(PP) 로 모델을 여러 GPU에 걸쳐 분산시키고,
- 텐서 병렬(TP) 로 개별 레이어의 크기를 줄입니다.
컨텍스트가 너무 길 때, SP: 긴 문서나 대화를 처리해야 해서 입력 시퀀스가 길어지면, 시퀀스 병렬(SP) 을 추가하여 활성화 메모리 부담을 줄입니다.

결론적으로, 초거대 모델은 데이터 병렬(DP) + 파이프라인 병렬(PP) + 텐서 병렬(TP) + 시퀀스 병렬(SP) 을 모두 조합하고, 여기에 ZeRO를 더해 각 GPU의 메모리 효율을 극대화하는 방식으로 훈련됩니다. 이를 3D Parallelism (+ZeRO) 이라고도 부릅니다.

I.2) ZeRO-3 Vs 텐서 병렬(TP) 비교

ZeRO-3와 TP는 같지 않습니다.

ZeRO-3: 데이터 병렬(Data Parallelism)의 메모리 최적화 기법입니다.
TP (Tensor Parallelism): 모델 병렬(Model Parallelism)의 한 종류로, 연산 자체를 분할하는 기법입니다.

아래 비교를 통해 명확하게 이해하실 수 있습니다.


구분	ZeRO-3 (메모리 최적화)	텐서 병렬 (TP) (모델 병렬)
주요 목표	메모리 절약. 데이터 병렬 처리 시 각 GPU가 모델 전체 파라미터를 복제해 갖는 메모리 중복을 제거하는 것이 목표입니다.	모델 분할 및 연산 병렬화. 단일 GPU에 담을 수 없는 거대한 레이어(예: 20GB짜리 Weight 행렬)를 여러 GPU로 쪼개어 동시에 연산하는 것이 목표입니다.
동작 방식	1. 평소에는 각 GPU가 전체 파라미터의 일부만 저장합니다(1/N). 2. 순전파(Forward) 시, 특정 레이어를 계산할 때 필요한 전체 파라미터를 다른 GPU로부터 All-Gather로 가져옵니다. 3. 가져온 전체 파라미터로 표준 연산을 수행합니다. 4. 연산이 끝나면 필요 없어진 파라미터를 즉시 해제(Discard)합니다.	1. 시작부터 거대한 파라미터(Weight 행렬)를 물리적으로 조각내어(Shard) 각 GPU가 나눠 갖습니다. 2. 순전파(Forward) 시, 입력 텐서(Activation)도 쪼개거나 복제하여 각 GPU가 자신이 가진 조각난 파라미터로 연산을 수행합니다. 3. 각 GPU에서 계산된 부분 결과를 All-Reduce 같은 통신으로 합쳐 최종 결과를 만듭니다.
기반 철학	”파라미터는 연산하는 그 순간에만 전체가 필요하다. 평소에는 분산 저장하여 메모리를 아끼자."	"하나의 행렬 곱셈이 너무 크다. 여러 GPU가 동시에 힘을 합쳐 이 계산을 해치우자.”
통신 오버헤드	순전파/역전파 시 각 레이어마다 전체 파라미터를 모으고 흩뿌리는 통신이 발생합니다.	연산 중간에 부분 결과를 합치는 통신이 발생합니다.
어떤 관계인가?	데이터 병렬(DP)의 확장판입니다. ZeRO는 DP를 전제로 동작합니다.	데이터 병렬(DP)과 직교(orthogonal) 관계입니다. DP와 함께 사용될 수 있습니다.

I.2.1) 비유를 통한 이해

ZeRO-3: 공용 창고와 개인 작업대
- 여러 명의 기술자(GPU)가 있습니다. 거대한 엔진(모델)을 조립해야 하는데, 각자의 작업대(VRAM)는 좁습니다.
- 그래서 모든 공구(파라미터)를 중앙 공용 창고(여러 GPU의 메모리 공간에 분산 저장)에 둡니다.
- 기술자 A가 특정 부품(레이어)을 조립할 차례가 되면, 창고에 가서 필요한 모든 공구(해당 레이어의 전체 파라미터)를 빌려와 작업을 마친 뒤, 즉시 반납합니다.
- 핵심은 ‘빌려와서 혼자 작업하고 반납’하는 것입니다.
텐서 병렬(TP): 거대한 부품을 동시에 조립하는 여러 기술자
- 하나의 부품(하나의 큰 행렬)이 너무 커서 기술자(GPU) 한 명이 들거나 작업할 수 없습니다.
- 그래서 이 부품을 네 조각으로 잘라 네 명의 기술자가 각자 한 조각씩 맡아 동시에 작업합니다.
- 작업이 끝나면 각자의 결과물을 합쳐서 완전한 결과를 만듭니다.
- 핵심은 ‘하나의 일을 동시에 나눠서 처리’하는 것입니다.

I.2.2) 결론: 어떻게 함께 사용되는가?

현대의 초거대 모델 훈련에서는 이 둘을 함께 사용합니다.

모델의 특정 레이어가 너무 커서 한 GPU의 VRAM에 올릴 수 없을 때, TP를 사용해 레이어 자체를 여러 GPU에 분할합니다.
이렇게 TP로 분할된 모델조차도 전체 파라미터의 양이 엄청나므로, 데이터 병렬 학습을 위해 여러 GPU에 모델을 복제하면 메모리가 부족해집니다.
이때 ZeRO-3를 적용하여, TP로 분할된 파라미터를 포함한 모든 파라미터, 그래디언트, 옵티마이저 상태를 여러 데이터 병렬 그룹(Data Parallel Ranks)에 걸쳐 분산 저장합니다.

즉, TP로 모델의 ‘너비’를 줄이고, ZeRO-3로 ‘메모리 사용량’을 최적화하는, 상호 보완적인 관계로 함께 활용됩니다.

J) References

The Technology Behind BLOOM Training

Zzong's Notes

탐색기

deepspeed