Abstract
현존하는 거대한 Language Model들은 학습하는 것에 문제가 있음.
- 학습을 빠르게 하고싶다 = 많은 GPU + 큰 Batch size → Data Parallelism
- 엄청 큰 모델을 학습하고 싶다 → 모델을 쪼개자 → Model Parallelism
위 두 가지와 모두 호환 가능하면서 GPU 메모리 부족 문제로 어려운 초대형 모델 training을 가능하게 함.
즉, GPU 메모리 사용량을 줄여줌.
- DP/MP에서 발생하는 Memory Redundancy를 제거
- Communication cost를 최소화
ZeRO Overview
ZeRO는 메모리 소비 스펙트럼을 분석하여 Model States와 Residual States로 나누어 각각에 대한 솔루션을 제안한다.
메모리의 대부분을 Model States가 점유하고 나머지 Residual States으로 구성된다.
1) ZeRO-1: ZeRO-DP (powered Data Parallelism)
- Model States 문제를 해결하기 위한 솔루션
2) ZeRO-2: ZeRO-R (powered Residual States Solutions)
- Residual States 문제를 해결하기 위한 솔루션
3) ZeRO-3: ZeRO-offload (CPU offloading)
4) ZeRO-Inifinity
아래 4가지 특성을 장점으로 내세운다.
Model scale
ZeRO-1은 100B 모델까지 학습가능하며, ZeRO-2는 200B 모델까지 학습가능하다. ZeRO-3(ZeRO-Offload)는 1T 모델까지 학습할 수 있다.
Compatible performance
Megatron-LM은 tensor-slicing model parallelism를 사용한다. Megatron-LM의 tensor-slicing model parallelism에 DeepSpeed ZeRO-2의 data parallelism을 추가하면 DeepSpeed가 Megatron-LM보다 10x 더 빠르다.
Scalability
디바이스 개수의 증가에 따라 superlinear speedup을 달성할 수 있다. 녹색 그래프는 perfect linear scalability로 실제로 디바이스 개수에 따른 산술적인 성능이라면, DeepSpeed를 사용하면 perfect linear scalability보다 더 높은 성능을 보인다.
예를 들어 GPU을 2배로 증가시키면 시스템 throughput이 2배 이상 증가한다. 이는 GPU당 더 큰 배치 크기를 맞출 수 있도록 DP 정도를 증가시킴에 따라 모델 상태의 메모리 공간을 줄여 더 나은 성능을 제공한다.
=> GPU가 많아질 수록 GPU끼리 서로 효과를 보아 GPU 메모리 사용량이 더 줄어들 요소가 있나?
Usability
DeepSpeed는 Pytorch 기반이며 몇 줄의 코드를 추가하면 DeepSpeed를 사용가능하다. 또한 1.4B 파라미터 이상의 Pytorch기반 모델은 Data Parallelism만을 사용할 때 OOM이 발생하지만 ZeRO-2는 13B까지 Model parallelism 없이 Data Parallelism을 사용할 수 있다.
Parameter 정의
* Model States: Parameters, Gradients, Optimizer states
* Residual States: Activation, Temporary Buffer, Fragmented Memory 등
- Transformer Model states 메모리 구성
- Parameter: FP16 Parameters(weight)
- Gradient: FP16 Gradient
- Optimizer states: FP32 Parameters, FP32 Gradient, FP32 Variance, FP32 Momentum <- back-prop후 업데이트되는 파라미터
- Activation: 모델의 각 레이어가 입력 데이터를 처리한 후 나오는 중간 결과
Optimizer states는 학습 중에 옵티마이저가 관리하는 파라미터 상태들을 의미합니다.
- 구성:
파라미터 (weights): 모델의 가중치 값.
모멘텀 (Momentum): optimizer가 사용하는 과거의 gradient 정보.
예를 들어, Adam 옵티마이저는 1차 및 2차 모멘텀(moving averages of gradients and squared gradients)을 유지합니다.
이러한 정보는 weight 업데이트 시 gradient를 더 잘 반영하게 돕습니다.
- m (1차 모멘텀): 그래디언트의 지수 이동 평균(velocity).
- v (2차 모멘텀): 그래디언트의 제곱의 지수 이동 평균(acceleration)
이 두 값이 각각의 파라미터마다 저장됩니다.
왜 Activation이 Residual State로 분류되나요?
- Activation은 모델의 forward pass 중에 한 번 계산된 후, backward pass에서 다시 사용됩니다. 그래서 GPU 메모리에 유지되어야 합니다.
그러나 그 자체로는 모델 파라미터처럼 업데이트되지 않으며, 일정 시간 동안만 필요하고 이후에는 재사용되지 않습니다.
이처럼 일시적으로 사용되기 때문에 "Residual State"라는 범주에 속하게 됩니다.
특히 Residual Connection이 있는 Transformer 구조에서, activation 값이 backward pass에 매우 중요한 역할을 하기 때문에 이를 "Residual" 상태로 부르는 경향이 있습니다.
Activation memory usage
- Total activations =12×hidden dimension×batch size×sequence length×number of layers
- ex) GPT-2 1.5B parameter model
- Hidden dimension: 1600
- Batch size: 32
- Sequence length: 1000 tokens
- Number of layers: 48 layers
- Total activations = 12×1600×32×1000×48 = 29 G (12: multiple attention heads)
- Memory for activations = 29G×2bytes = 58Gbytes = 58 GB (@FP16)
ZeRO 목적
ZeRO에서 제안하는 방향
- Memory Efficiency
- High compute & less communication
=> model layer를 gpu 갯수만큼 n등분하여 나누면서 optimizer state, gradient, parameter도 세 부분으로 나눠 필요한 순간에만 각 데이터를 communication하여 communication cost를 줄이고 model memory 사용도도 줄임.
ZeRO 동작 원리
- Optimizer State와 Gradient의 메모리 사용량을 DP 그룹 내에서 줄임
- DP와 함께 진행됨
ZeRO 상세 설명
1) ZeRO-1: ZeRO-DP
[1910.02054] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arxiv.org)
- DeepSpeed library에서 제안한 방법론
- DP(data parallelism)를 기반으로 하여 Memory redundancy를 줄이고자 함.
- Data parallelism에서, 모든 parameter들을 각 GPU에 full set으로 가지고 있는 대신 fraction data만 들고 있다가, 필요할 때 다른 GPU로부터 받아서 사용한다
- 세 가지 종류의 partitioning에 대해 다룸
* Optimizer state partitioning
* Gradient partitioning
* Parameter partitioning
- 앞의 두 partitioning의 경우 communication overhead X
- gradient는 only backward에서 전달받은 parameter를 가지고 gradient 계산을 한 후에 해당 gradient를 담당하는 GPU에 reduce 진행함.
- optimizer state는 train 1 loop가 끝나면 각 GPU안에서 생성된 gradient 값에 따라 variance, momentum로 weight 값 업데이트함.
- Parameter partitioning의 경우 50% 가량의 communication overhead가 있지만 GPU의 수에 비례한 메모리 사용을 줄일 수 있다는 결과
- parameter는 forward/backward 모두 partitioning broadcast를 수행함.
Parallelism 개념
- Data Parallel은 데이터 셋을 쪼개는 것이기에 총 train 시간은 줄일 수 있어도 model param은 모두 올려야 해서 deep model size가 커지면 운용 못 함.
- Pipeline Parallel 방식으로 각 layer(층)별로 모델을 쪼개거나
- Model Parallel로 Multihead attention에서 head를 쪼개는 등의 방식
=> (단점) MultiHead att 등에서 HEAD를 쪼개는 것이기 때문에 최대로 늘릴 수 있는 경우가 Head 수로 제약이 됨
=> 각 layer 단에서 계산된 값을 공유하는 것으로 인해 communication cost가 무척 높음
=> 위 특성으로 인해 DP 보다 scaling efficiency가 떨어짐. 그래서 ZeRO는 DP 기반으로 아이디어를 제시함.
ZeRO에서의 해결책
- (기존) DP: Good compute & communication efficiency, Poor memory efficiency
- 모델 전체를 copy 해서 모든 GPU에 할당
- 각 Batch size를 N개(GPU)로 나눠서 계산
=> (솔루션) 모델을 copy하는 대신 Partitioning을 통해 memory consumption을 줄임
- (기존) MP: High communication cost
=> (술루션) Dynamic communication schedule(optimizer state, gradient, parameter 3개로 나눔) 통해서 communication volume을 DP급으로 줄임
2) ZeRO-2: ZeRO-R (powered Residual States Solutions)
ZeRO-1과 비교하여 ZeRO-2는 DeepSpeed로 학습할 수 있는 모델의 크기를 2x 늘려 학습 효율성을 크게 향상시킨다. ZeRO-2는 Residual States 문제를 해결하기 위한 솔루션으로 gradient, activation, fragmented 메모리를 절감하기 위한 기술을 소개한다.
ZeRO-2는 model states, activation, fragmented 메모리를 포함한 학습 중 메모리 소비의 전체 스펙트럼을 최적화한다.
- Activation Memory: Activation checkpointing (or activation recomputation)을 통해 모든 레이어가 아닌 일부 레이어의 activation만 메모리에 저장하고 나머지는 나중에 다시 계산한다. MP를 사용하면 각 GPU가 연산하는 레이어의 activation만 해당 GPU에 저장하고 back-propagation에서 필요할 때 all-gather를 통해 계산하여 사용한다. 따라서 activation memory는 기존에서 MP 값 만큼 나눈 양이 소모 된다. 이렇게 되면 activation 메모리 양이 100B model의 경우 2GB 까지 줄 수 있어 이런 경우 CPU 메모리로 다 offload해서 사용할 수도 있다.
- Fragmented Memory: tensor들의 생명 주기가 달라 학습시 메모리 단편화가 발생한다. 메모리 단편화로 인해 연속적인 메모리의 부족으로 가용 메모리가 남아 있었음에도 불구하고 메모리 할당에 실패한다. ZeRO-2는 메모리 단편화를 제거하여 공간을 효율적으로 활용한다.
- Constant Size Buffer: All-reduce 연산 같은 경우 여러 레이어의 gradient를 fused 시켜 1개 flatten buffer에 넣을 수 있다. 이에 대한 사이즈를 모델에 의존하여 어느정도 크게 유지하기만 해도 충분히 좋은 효율성을 얻을 수 있다
Operating ex.
- OPT-13B model -> 1대 A100 80G GPU로 가능
=> 기본 16bit로 잡으면 13B * (2 + 2 + 12) bytes = 208 GB (적어도 A100 3대 필요)
- 100B model -> 400대 V100 GPU로 가능 (38 TFlops/GPU)
3) ZeRO-3: ZeRO-offload (CPU offloading - DRAM 메모리 사용)
- 목표: GPU 연산의 효율성은 최대한 유지하면서 CPU를 사용해 GPU 메모리의 부담을 줄이는 것
- GPU는 빠른 연산을 담당하고,
- CPU는 메모리 관리를 담당하여 GPU의 메모리 한계를 보완합니다.
ZeRO-Offload라고 불리우는 ZeRO-3는 gradients, optimizer states와 optimizer 연산을 CPU로 offloading한다. Parameters와 forward & backward 연산은 GPU에서 유지한다.
아래는 ZeRO-offload 스케쥴링
- 전체 학습 과정에서 Optimizer states는 항상 CPU memory에 올라가 있고,
- Gradients는 학습에 사용하는 GPU들에서 Backward 과정 중 Reduce-scatter 한 뒤, 각각의 GPU(DataParallel)에서 나온 Gradients들의 Avg로 다시 CPU로 Offload한다.
- CPU에 gradient가 업데이트되면 이를 optimizer states에 업데이트되고(p update) all-gather로 합친 뒤 각 GPU로 다시 내린다.
Operating ex.
- 10B model -> 1 V100 GPU (40 TFlops/GPU)
- 70B model -> 1 DGX-2(16 V100) box (with MP)
4) ZeRO-Infinity
Heterogeneous Memory Utilization: ZeRO-Infinity allows deep learning models to utilize multiple layers of memory within a system:
- GPU Memory: The fastest memory closest to the compute core.
- CPU Memory (DRAM): Larger but slower than GPU memory.
- NVMe Storage: Even larger but significantly slower than both GPU and CPU memory.
Reference
paper
- https://arxiv.org/abs/1910.02054
- https://arxiv.org/abs/2101.06840
- https://arxiv.org/abs/2104.07857
official tutorial & github
- https://www.deepspeed.ai/tutorials/zero/
- https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-chat
- https://www.microsoft.com/en-us/research/blog/zero-2-deepspeed-shattering-barriers-of-deep-learning-speed-scale/
blog
- https://junbuml.ee/zero-memory-optimizations-toward-training-trillion-parameter-models
youtube
- https://www.youtube.com/watch?v=by86GQhKgKQ&list=PL2HUvXaMsPQlkLv8c749I9qHyQr7-m9wG&index=3
- https://www.youtube.com/watch?v=CaseqC45DNc&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=30
댓글