[리뷰] ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

ZeRO-Infinity는 DGX-2의 GPU & CPU 메모리와 NVMe 스토리지를 이용하여 1 Trillion 모델을 학습할 수 있는 솔루션이다.

GPU 메모리 Wall

1 Trillion 모델을 학습하기 위한 총 메모리양은 다음과 같이 계산할 수 있다.

• 모델 파라미터: 1x10¹² x 4 Bytes (FP32) = 1000G x 4 Bytes = 4 TB
• 학습에 필요한 최소 메모리양: 4 TB x 5 = 20 TB (학습에 필요한 최소 메모리양는 파라미터의 5배)

하지만 위 실제 학습에 필요한 최소 메모리양에 activation이 포함되지 않았으므로 이를 포함한다면 20TB + activation 만큼의 메모리가 필요하다.

1 Trillion 모델을 돌리기 위해 필요한 GPU의 개수는 DeepSpeed 3D parallelism기준 800개 V100 GPU (32GB ver.) 개가 필요하다.

DeepSpeed 3D parallelism으로 1T 모델을 학습하기 위한 GPU의 대수

Large Model 학습을 위한 저비용 학습 솔루션: ZeRO-Infinity의 새로운 기능

Microsoft는 3D parallelism 에서 더 나아가 Large Model 학습 문제를 해결하기 위한 새로운 학습 솔루션인 ZeRO-Infinity을 제안하였다. ZeRO-Infinity이 제공하는 새로운 기능은 다음과 같다.

1. 전례없는 모델 크기

• Infinity offload engine: CPU 메모리와 NVMe 스토리지를 동시에 활용하여 제한된 GPU 리소스로 Large Mode 학습을 지원함
• Memory-centric tiling: model parallelism이 필요없이 모델의 한 개의 개별 레이어가 너무 거대하여 GPU 메모리에 맞지 않는 경우에도 지원하는 GPU 메모리 최적화

Infinity offload engine와 memory-centric tiling을 이용하여 제한된 GPU 리소스를 가지고 Large Model 학습을 할 수 있다.

2.탁월한 학습 효율성

• Bandwidth-centric partitioning: 모든 디바이스의 총 메모리 대역폭을 활용하기 위한 새로운 데이터 파티셔닝 전략
• Overlap-centric design: compute와 communication을 중첩

Bandwidth-centric partitioning과 overlap-centric design으로 NVMe 스토리지와 CPU 메모리에 데이터를 offloading함에도 불구하고 제한된 대역폭에 방해받지 않고 탁월한 학습 효율성을 제공한다.

3.사용성

ZeRO-Infinity를 사용하면 3D parallelism과 같은 다중 병렬 처리를 모델에 적용할 필요가 없으며 수동으로 코드 리팩토링이 필요하지 않다.

메모리 요구사항

Memory for Residual States

Residual States는 주로 activation 메모리로 구성되어 있으며, activation 메모리는 모델 아키텍처, batch size(bsz), sequence length(seq)에 의존한다. Activation 메모리 요구량은 매우 크지만 activation checkpointing을 사용하면 상당히 줄일 수 있다. 다만 BWD 시 0.33x의 추가 recomputation 비용이 들어간다. T-NLG (17.2B)와 GPT-3(175B)와 같은 Large Model들은 activation checkpointing을 이용하여 학습된다. Activation checkpoints를 저장하기 위해 필요한 메모리는 다음과 같이 평가된다.

2*bsz*seq*hd*nl/ci (bytes)

(ci: 두 activiation checkpoints간 transformer 블록의 개수, bsz*seq*hd는 각 transformer의 입력 사이즈)

Note: 이 논문에서는 batch size=32, sequence length=1024, 1 activation per transformer block을 가정한다.

Activation checkpoints를 사용하면 activation의 full set보다 훨씬 더 작은 용량을 필요로 하나, trillion 파라미터를 초과하면 너무 커서 GPU 메모리에 맞지 않는다. (1 trillion: Act:7.13TB / Act ckpt:0.2TB, 10 trillion: Act:24.38TB/Act ckpt:0.76 TB)

Model State Working Memory(MSWM)

MSWM은 모든 model states가 CPU 또는 NVMe에 offload 된 후 모델에서 가장 큰 operator 상에 FWD/BWD를 수행할 때 요구되는 최소한의 GPU 메모리양을 의미한다. 대략 BWD 시 가장 큰 operator의 parameters와 gradients의 사이즈이며 BWD를 위한 parameters와 gradients을 담기 위한 최소한의 충분한 메모리가 있어야 하기 때문이다.

Transformer 기반 모델에서 가장 큰 operator는 hidden states를 hd부터 4hd로 변환하기 위한 linear 레이어이다. 이 linear 레이어의 parameters와 gradients 사이즈는 다음과 같다.

4*hd*4hd Bytes

0.1 Trillion 파라미터를 초과하는 모델에 대해선 MSWM이 급격히 커지고 있어, (10 Trillions 모델의 경우 무려 64GB의 working memory 필요) 충분한 contiguous memory 부족때문에 학습 시 OOM(Out-of-Memory)의 원인이 된다. 이와 같은 문제는 “Memory-centric tiling for working memory” 방법을 사용하여 model parallelism 없이 MSWM 문제를 해결할 수 있다.

Activation Working Memory(AWM)

AWM은 실제 BWD을 실행하기 전에 activation을 재계산하기 위해 BWD시 필요한 메모리이다. 이것은 두 연속적인 activation checkpoints간 activation 사이즈이다. Transformer 블록당 1개의 activation checkpoint를 생성한다면, transformer 블록당 총 activation 사이즈는 대략 다음과 같다.

bsz * seq * ci * (16*hd + 2*attn_heads*seq) Bytes

싱글 parameters & gradients로 구성되어 있는 MSWM와 달리, AWM은 수십개의 activation으로 구성되어 있으며 전체 AWM이 GPU 메모리에 맞기만 한다면 contiguous memory 부족 때문에 메모리 이슈를 발생시키지 않는다. (즉, activation은 수십개로 분할된 상태이므로 contiguous memory가 아니어도 된다라는 의미)

Memory requirements for massive models

Available memory and achievable bandwidth on NVIDIA V100 DGX-2 Cluster

Bandwidth 요구사항

딥러닝 학습 시 AIT 정량화

iteration당 total Computation: Transformer는 linear 레이어의 연산이 지배적인 연산이다.

• (1) FWD 시 computation: 2*bsz*seq*parameters
• (2) BWD 시 computation: 2*2*bsz*seq*parameters
• (3) Activiation checkpointing: 2*bsz*seq*parameters
• computation_per_iter ((1)+(2)+(3)) = 2*4*bsz*seq*parameters = 2*4*12*bsz*seq*nl*hd²

parameters와 gradients에 대한 AIT: FWD와 BWD 시, 소스 위치로부터 GPU register까지 model parameters가 최소한 두번 load 되어야 한다.(i. FWD , ii. BWD) Activation checkpointing이 사용되므로 BWD 중에 recomputation을 위한 추가 시간과 또다른 1*parameters의 데이터 이동이 추가된다. 게다가 gradients는 GPU register로부터 최종 위치에 최소한 한번 저장되어야 한다. 따라서 parameters와 gradients동일한 최종 위치에 저장된다고 가정하면 총 데이터 이동은 parameter는 3번, gradient는 1번이다.

• total data movement :4*parameters
• total data movement(bytes): 2*4*parameters
• ait w.r.t parameter & gradients: computation per iteration/data movement = (2*4*bsz*seq*parameters)/(2*4*parameters) = bsz*parameters

Optimizer States에 대한 AIT: optimizer states는 optimizer step 중에 최소한 한번씩 read/write 해야 한다.

• total data movement: 2*optimizer_state
• total data movement(bytes): 2*16*optimizere_state
• ait w.r.t optimizer states: computation per iteration/data movement =(2*4*bsz*seq*parameters)/(2*16*parameters)=bsz*seq/4

Activation Checkpoints에 대한 AIT: FWD 동안 activation checkpoints는 최종 위치에 저장되었다가 BWD 동안 복원되어야 한다.

• total data movement(bytes): 2*total_activation_checkpints_bytes=2*2*bsz*hd*nl/ci=4*bsz*seq*hd*nl/ci=(2*4*12*bsz*seq*nl*hd²)/(4*bsz*seq*hd*nl/ci)=24*hd*ci

parameter 와 gradient의 bandwidth 요구사항

• parameter와 gradient를 위해 70 GB/s 이상의 bandwidth를 갖는다면 가장 작은 batch size에서도 50% 이상의 효율성을 달성가능하다.
• 70GB/s 이상의 bandwidth를 갖는다면, 이론적으로 데이터 이동이 완벽하게 연산과 중첩되어 100% 효율성을 달성할 수 있다.

optimizer states의 bandwidth 요구사항

• optimizer states는 50% 이상의 효율성을 달성하기 위해 parameter & gradient 가 요구하는 bandwidth에 비해 거의 4x 더 높은 bandwidth를 요구한다.
• GPU 당 batch size=2로 90% 효율성을 달성하기 위해 필요 유효 bandwidth는 1.5TB/s이다.

Activation의 bandwidth 요구사항

• activation checkpointing이 활성화되었을 때, hidden size가 2K에서 조차 2GB/s이면 50% 이상의 efficiency를 유지할 수 있다.
• hidden size가 8K 이상 커지면 요구 bandwidth는 1GB/s 이하로 내려간다.

ZeRO-Infinity 디자인 개요

1.모델 사이즈 스케일링

ZeRO-Infinity는 이기종 메모리 시스템을 활용하여 3D parallelism보다 50배 더 증가한 NVIDIA DGX-2 노드당 1T 모델을 지원한다. 다음은 모델 사이즈를 스케일링하기 위한 ZeRO-Infinity에 적용된 기술이다.

Infinity Offload Engine

ZeRO-Infinity는 모든 model states를 분할하여 메모리 중복성을 제거하는 ZeRO-3을 기반으로 모든 분할 model states를 CPU 또는 NVMe 스토리지로 offloading한다.

Activation을 위한 CPU offload

Model states 외에 Activation 메모리를 CPU 메모리로 offload할 수 있다. 10 T 모델의 activation checkpoint(0.76 TB)는 DGX-2의 1.5TB CPU 메모리에 모두 들어갈 수 있는 크기이다.

Working 메모리를 위한 Memory-centric tiling

ZeRO-Infinity는 ZeRO-3의 데이터 fetch와 release 패턴을 활용하여 큰 연산자를 순차적으로 실행할 수 있는 작은 tile로 분할하여 working 메모리의 요구사항을 줄인다. 예를들어 ZeRO-Infinity는 OP(OPerator)를 원래 OP의 파라미터 tile로 구성된 수학적으로 동등한 더 작은 선형 OP 시퀀스로 나타내고 그것들을 순차적으로 실행한다. ZeRO-3와 결합하면 tile 수에 비례하여 working 메모리가 줄어든다. 따라서 ZeRO-Infinity는 model parallelism에 의존하지 임의 크기의 OP를 지원할 수 있다.

2. 탁월한 학습 효율성을 위한 설계

ZeRO-Infinity가 모든 model states 및 activation을 CPU 메모리 또는 NVMe 스토리지에 offload하는 것은 고효율을 달성해야 실용적이다. 왜냐하면 GPU-to-GPU이 1x이라고 한다면 CPU-to-GPU이 1/10x, NVMe-to-GPU이 1/100x 느리다.

Parameter과 Gradients에 대해 효율성

Parameter와 Gradient에 대한 데이터 이동 대역폭은 DGX-2의 GPU-to-GPU 대역폭에 가까운 70GB/s 이상이어야 한다. FWD 및 BWD에 Owner GPU로부터 나머지 GPU로 parameter가 전파되는 ZeRO-3와 같은 병렬 학습 솔루션은 통신이 중첩되기만 한다면 효율적으로 동작될 수 있다.

반면, GPU-to-CPU Memory(100GB/s), GPU-to-NVMe Storage(25GB/s)의 대역폭은 12GB/s의 PCIe 대역폭의 제한을 받기 때문에 이기종 학습을 지원하는데 충분하지 않다. 따라서 ZeRO-Offload와 같은 이기종 솔루션은 owner GPU가 나머지 GPU로 parameter를 전파하기 전에 parameter를 CPU로부터 owner GPU로 이동시켜야 하므로 GPU 당 충분히 큰 배치 사이즈가 필요하다.

이는 두가지 문제를 야기시킨다.

• 1) Massive 모델의 경우, activation 메모리가 너무 커져 CPU 메모리에 맞지 않는다.
• 2) 효율적인 수렴을 위해 수백~수천개 GPU로 스케일할 때 효과적인 배치 크기가 너무 커진다. (note: 지나치게 batch size가 크면 학습 수렴에 문제가 있다.)

ZeRO-Infinity는 bandwidth-centric partioning과 overlap centric design을 이용 위의 2가지 문제를 해결한다.

Optimizer States에 대해 효율성

FWD와 BWD 동안 순차적으로 소비되고 생산되는 Parameters와 gradients와 달리 optimizer states는 병렬로 동시에 업데이트될 수 있다. 결과적으로 aggregate GPU 또는 CPU 메모리 BW는 GPU와 CPU의 개수가 증가함에 따라 요구 BW 1.5TB/s보다 더 높아질 수 있다.

ZeRO-Infinity는 ZeRO-3를 기반으로 하기 때문에 optimization 단계를 위해 optimizer states를 CPU 메모리로 offload할 때 aggregated CPU compute를 활용할 수 있다. 하지만 NVMe-to-CPU Memory로 offload 시 한번에 한 chunk씩 최적화 단계를 수행하기 위해 CPU Memory에 들어갈 수 있는 chunk로 다시 가져와야 한다. 따라서 최적화 단계는 NVMe-CPU Memory 대역폭에 의해 제한된다.

Infinity offload engine은 거의 peak NVMe 대역폭을 달성할 수 있을 뿐만이 아니라 ZeRO-Infinity가 NVMe-CPU 읽기와 CPU-NVMe 쓰기를 중첩할 수 있도록 허용한다. 동시에 데이터 이동을 위해 임시 버퍼를 주의깊게 재사용하여 메모리 단편화를 최소화한다.

Activation에 대한 효율성

DGX-2 노드의 각 GPU는 PCIe를 통해 3GB/s 정도의 GPU-to-CPU 읽기/쓰기가 병렬로 가능하다. Hidden size가 8K 또는 그이상에서 80% 효율 이상으로 재학습하는 동안 activation 체크포인트는 CPU 메모리로 offload가 가능하다. 더 작은 hidden size에서 높은 효율성을 위해 ZeRO-Infinity는 activation 체크포인트의 통신을 효율적으로 중첩할 뿐만 아니라 activation 체크포인트의 빈도를 줄일 수 있다.

효율적인 최적화

Bandwidth-centric Partitioning

ZeRO-Offload가 broadcast를 기반으로 개별 파라미터를 분배하는 반면, ZeRO-Infinity는 all-gather를 사용한다. 데이터가 GPU가 있는 경우, broadcast와 all-gather 모두 동일한 통신 cost 를 갖는다.

ZeRO-Offload가 Single PCIe만 활성화되고 나머지 PCIe 링크는 모두 idle인 상태로 aggregated PCIe bandwidth가 낮은 편이다. ZeRO-Infinity는 분할된 파라미터와 all-gather기반 통신을 사용하여 모든 PCIe 링크가 병렬로 활성화된다. 결과적으로 NVMe-to-GPU or CPU-to-GPU의 유효 통신 bandwidth는 선형적을 증가한다.

예를들어 broadcast기반 방식을 사용하면 DGX-2 PCIe Gen3기준으로 CPU/NVMe-to-GPU bandwidth는 12GB/s으로 일정하게 유지된다. 하지만 all-gather 방식을 사용하면 유효 bandwidth가 약 48/25GB/s (GPU당 3.0/1.6GB/s)로 증가하며 DGX-2 당 최대 aggregated PCIe bandwidth와 max NVMe bandwidth에 의해 제한된다.

Overlap-centric Design

ZeRO-Infinity는 GPU 연산과 GPU-GPU 통신을 중첩할 뿐만 아니라 NVMe-CPU, CPU-GPU 통신을 중첩하는 엔진이 있다. 중첩 엔진에는 1) FWD와 BWD pass에서 소비되기 전에 파라미터를 재구성하기 위한 데이터 이동을 중첩하는 동적 프리페처 2) BWD 계산과 병렬로 gradient에 필요한 데이터 이동을 실행하기 위한 통신 및 Offload 중첩 메커니즘과 2개의 컴포넌트가 있다.

ZeRO-Infinity 평가

Methodology

• Hardware: 32개 DGX-2 (512 V100 GPUs), 800 Gbps (Inter-node communication bandwidth)
• Baseline: DDP(Pytorch, Distributed Data Parallelism), Megatron-LM, 3D parallelism, ZeRO, ZeRO-Offload
• Model configuration: GPT-like Transformer 모델 사용, sequence length=1024(fixed), hidden dimension(4K, 8K, 18K, 25K, 48K, 64K, 88K), # of layers (50, 62, 124, 125, 128, 174, 200, 205), # of params (10B, 50B, 100B, 0.5T, 1T, 5T, 10T, 20T)

Model Size & Speed

Model Size

• ZeRO-Infinity (32T) vs 3D parallelism (650B) → 50x 이상 더 큰 모델을 학습 가능하다.

Model Speed

• 3D parallelism: 500B에서는 동일한 throughput, 그 이상 모델 사이즈에 대해서는 Out of Memory이 발생한다.
• ZeRO-Infinity: 5T(49TFlops) → 10T(43 TFlops) → 20T(34 TFlops)와 같이 수십 T 모델을 학습할 수 있다. 모델 사이즈가 커질수록 성능 저하가 일어나는 이유는 NVMe의 bandwidth이 원인이 아니라 20T로 스케일 시 activation checkpoint이 저장될 CPU 메모리 용량 부족으로 인해 극도로 작은 batch size를 사용하기 때문에 나타나는 결과이다. 이는 CPU 메모리를 증가하던지 activation checkpoint를 NVMe에 offload하여 해결 가능하다. (NVMe에 offload하는 것은 추후 구현 예정임)

Super-linear Scalability

• 1T 모델을 학습할 때, 4~32 nodes까지 Super-linear Scalability를 달성한다.
• 노드당 배치 크기를 일정하게 유지하고 노드 수가 증가함에 따라 전체 배치 크기를 증가시키는 Weak Scaling 결과이다.
• ZeRO-Infinity는 벌써 4 nodes으로 2.8 pflops 이상을 달성하였다. (44 Tflops/GPU)
• Aggregate NVMe bandwidth는 보통 좋은 효율성을 달성하는데 충분하다.

Large Model Training의 민주화

• 10T~1T 모델에 Model Parallelism없이 ZeRO-Infinity 만으로 finetuning이 가능하다. 반면 3D parallelism은 20B 이상 모델으로 스케일할 수 없다.
• ZeRO-Infinity는 Single DGX-2로 1T 모델을 fine-tuning 가능하며, Model Parallelism 없이 model code 리팩토링없이 모델을 쉽게 확장할 수 있다.

모델 규모상의 시스템 기능의 영향

• Max. Model Size: ZeRO-2와 ZeRO-Offload를 사용하였을 때 Single Node에서 Data parallel보다 9배인 13B 파라미터까지 scale-up이 가능하다. ZeRO-infinity를 CPU Memory로 확대하였을 때 50배 더 큰 모델을 학습할 수 있으며, NVMe로 확대하였을 때 700배 더 큰 모델을 학습할 수 있다.

Max model size w.r.t. ZeRO strategies

• Max. Hidden Size: memory-centric tiling 없이 학습할 수있는 가장 큰 hidden size는 8K이며 16 memory-centric tiling factor를 사용하여 64K의 hidden size까지 학습할 수 있다. Memory-centric triling을 이용하여 model parallelism 사용을 피할 수 있으며 큰 hidden size를 학습할 수 있다.

성능상의 시스템 기능의 영향

ZeRO-Infinity vs ZeRO-Offload: ZeRO-Infinity와 ZeRO-Offload의 8B 파라미터 모델을 Gradient를 CPU로 offload하였을 때 back-propagation 시간을 측정한 비교한 결과이다. 64 GPU기준 ZeRO-Infinity가 ZeRO-Offload에 비해 거의 2배의 속도 향상이 보인다. 이는 ZeRO-Offload가 Single PCIe bandwidth만을 사용하기 때문이다.

Prefetching & Overlapping: 64 GPU & 8B 파라미터 모델 환경에서 overlapping & prefetching을 on/off 했을 때 상대적인 결과로 GPU 당 작은 배치 사이즈일 때 좋은 성능을 달성한다.

Activation checkpoint offload: hidden Dimension가 클수록 CPU 메모리에 activation checkpoint를 offloading하였을 때 recomputation에 의한 overhead이 최소화된다.

결론 & Future Implications

ZeRO-Infinity를 사용하면 가속 디바이스의 메모리가 더 이상 모델 규모 또는 학습 효율성의 제한으로 작용하지 않는다.

하지만 10~100 T 모델을 합리적인 시간에 학습하는 것은 여전히 compute에 대규모 도약이 요구되고, 이와 같은 미래의 장치에서 효율적으로 실행하려면 디바이스-디바이스 대역폭 또한 compute와 비례하여 도약이 필요하다.

V100 GPU에 비해 가속기의 compute가 10배 증가한다면, 각 가속기와 느린 메모리(e.g., NVMe)간 bandwidth가 30 GB/s의 대역폭만 유지하면 된다.

레퍼런스

[1] DeepSpeed: Extreme-scale model training for everyone

DeepSpeed: Extreme-scale model training for everyone - Microsoft Research

[2] ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning