Rajbhandari, Samyam, et al. "Zero: Memory optimizations toward training trillion parameter models." SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020.

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Samyam Rajbhandari*, Jeff Rasley*, Olatunji Ruwase, Yuxiong He<br>{samyamr, jerasley, olruwase, yuxhe}@microsoft.com

Abstract

대규모 딥러닝 모델은 상당한 정확도 향상을 제공하지만, 수십억에서 수조 개의 파라미터를 학습시키는 것은 어려운 일이다. 기존의 데이터 병렬화(data parallelism) 및 **모델 병렬화(model parallelism)**와 같은 솔루션들은 제한된 장치 메모리에 이러한 모델을 맞추는 데 근본적인 한계를 보이며, 동시에 계산, 통신 및 개발 효율성을 확보하는 데 어려움이 있다.

우리는 메모리를 최적화하고 학습 속도를 크게 향상시키면서 효율적으로 학습할 수 있는 모델 크기를 늘리는 **새로운 솔루션인 Zero Redundancy Optimizer (ZeRO)**를 개발했다. ZeRO는 데이터 병렬(data-parallel) 및 모델 병렬(model-parallel) 학습에서 발생하는 메모리 중복성을 제거하면서도, 낮은 통신량과 높은 계산 세분성을 유지하여, 장치 수에 비례하여 모델 크기를 확장할 수 있도록 하며 지속적으로 높은 효율성을 유지한다. 메모리 요구 사항 및 통신량에 대한 우리의 분석은 다음과 같은 점을 보여준다: ZeRO는 오늘날의 하드웨어를 사용하여 1조 개 이상의 파라미터로 확장될 잠재력을 가지고 있다.

우리는 ZeRO를 구현하고 평가했다: ZeRO는 400개의 GPU에서 100B개 이상의 파라미터를 가진 대규모 모델을 초선형(super-linear) 속도 향상으로 학습시키며, 15 Petaflops의 처리량을 달성한다. 이는 모델 크기에서 8배, 달성 가능한 성능에서 10배 증가한 것으로, state-of-the-art를 뛰어넘는 결과이다. 사용성 측면에서 ZeRO는 **최대 13B개의 파라미터(예: Megatron GPT 8.3B 및 T5 11B보다 큰 모델)**를 가진 대규모 모델을 모델 병렬화(model parallelism) 없이 학습시킬 수 있다. 이는 연구자들이 적용하기 더 어려운 방식이다. 마지막으로, 연구자들은 ZeRO의 시스템 혁신을 활용하여 **세계에서 가장 큰 언어 모델(17B 파라미터)**을 기록적인 정확도로 만들었다.

1 Extended Introduction

Deep Learning (DL) 모델은 점점 더 커지고 있으며, 모델 크기 증가는 상당한 정확도 향상을 제공한다. 자연어 처리(NLP) 분야에서 Transformer는 Bert-large (0.3B) [1], GPT-2 (1.5B) [2], Megatron-LM (8.3B) [3], T5 (11B) [4]와 같은 대규모 모델의 길을 열었다. 모델 크기 성장을 수십억 개에서 수조 개의 파라미터로 확장하기 위해서는 학습 과정에서 발생하는 문제들을 해결해야 한다. 이러한 모델들은 단일 장치(예: GPU 또는 TPU)의 메모리에 분명히 맞지 않으며, 단순히 장치를 추가하는 것만으로는 학습 확장에 도움이 되지 않는다.

기본적인 **Data Parallelism (DP)**은 장치당 메모리를 줄이지 못하며, 현재 세대 32GB 메모리 GPU에서는 1.4B 파라미터 이상의 모델에 대해 메모리 부족 현상이 발생한다. Pipeline Parallelism (PP), Model Parallelism (MP), CPU-Offloading 등 다른 기존 솔루션들은 기능성, 사용성, 그리고 메모리 및 연산/통신 효율성 사이에서 trade-off를 가지지만, 이 모든 요소는 빠르고 대규모의 학습에 필수적이다.

[^0]

대규모 모델 학습을 위한 다양한 기존 솔루션 중 MP가 가장 유망하다고 할 수 있다. 현재 문헌에서 가장 큰 모델인 11B T5 모델 [4]과 Megatron-LM 8.3B [3]은 모두 Model Parallelism에 의해 구동되었으며, 각각 Mesh-Tensorflow [5]와 Megatron-LM [3]에 구현되었다. 그러나 MP는 이러한 모델 크기를 훨씬 넘어 확장하기 어렵다. MP는 모델을 수직으로 분할하여 각 layer의 연산과 파라미터를 여러 장치에 분산시키므로, 각 layer 간에 상당한 통신이 필요하다. 결과적으로, GPU 간 통신 대역폭이 높은 단일 노드 내에서는 잘 작동하지만, 단일 노드를 넘어서면 효율성이 빠르게 저하된다 [3]. 우리는 Megatron-LM을 사용하여 두 개의 DGX-2 노드에 걸쳐 40B 파라미터 모델을 테스트했으며, V100 GPU당 약 5 Tflops(하드웨어 최고 성능의 5% 미만)를 관찰했다.

그렇다면, 기존 솔루션의 한계를 극복하고 대규모 모델을 더 효율적으로 학습시키는 방법은 무엇일까? 이 질문에 답하기 위해, 우리는 먼저 기존 시스템의 모델 학습 시 메모리 소비 스펙트럼 전체를 분석하고 이를 두 부분으로 분류한다:

1. 대규모 모델의 경우, 메모리의 대부분은 optimizer state (Adam [6]의 momentum 및 variance와 같은), gradient, parameter를 포함하는 model state가 차지한다.
2. 나머지 메모리는 activation, temporary buffer, 사용 불가능한 fragmented memory에 의해 소비되며, 이를 통칭하여 residual state라고 한다.

우리는 **ZeRO (Zero Redundancy Optimizer)**를 개발하여 높은 연산 및 통신 효율성을 유지하면서 이 두 부분 모두에서 메모리 효율성을 최적화한다. 이 두 부분이 서로 다른 문제에 직면하므로, 우리는 그에 상응하는 솔루션을 개발하고 논의한다.

Model State Memory 최적화
Model state는 학습 중에 가장 많은 메모리를 소비하는 경우가 많지만, DP 및 MP와 같은 기존 접근 방식은 만족스러운 솔루션을 제공하지 못한다. DP는 연산/통신 효율성은 좋지만 메모리 효율성이 떨어지는 반면, MP는 연산/통신 효율성이 떨어질 수 있다. 더 구체적으로, DP는 전체 model state를 모든 data parallel process에 복제하여 메모리 소비의 중복을 초래한다. 반면 MP는 이러한 state를 분할하여 높은 메모리 효율성을 얻지만, 종종 너무 세분화된 연산과 비싼 통신으로 이어져 확장 효율성이 떨어진다. 또한, 이러한 모든 접근 방식은 학습 중에 모든 model state가 항상 필요한 것은 아님에도 불구하고, 전체 학습 과정 동안 필요한 모든 model state를 정적으로 유지한다. 이러한 관찰을 바탕으로, 우리는 **ZeRO-DP (ZeRO-powered data parallelism)**를 개발하여 DP의 연산/통신 효율성을 달성하면서 MP의 메모리 효율성을 달성한다. ZeRO-DP는 model state를 복제하는 대신 분할하여 data-parallel process 간의 메모리 state 중복을 제거하며, 학습 중 동적 통신 스케줄을 사용하여 DP의 연산 granularity 및 통신 볼륨을 유지함으로써 연산/통신 효율성을 유지한다.

ZeRO-DP는 세 가지 주요 최적화 단계(Figure 1 참조)를 가지며, 이는 optimizer state, gradient, parameter의 분할에 해당한다. 누적적으로 활성화될 때:

1. Optimizer State Partitioning ($P_{os}$): 4배 메모리 감소, DP와 동일한 통신 볼륨;
2. Gradient Partitioning 추가 ($P_{os+g}$): 8배 메모리 감소, DP와 동일한 통신 볼륨;
3. Parameter Partitioning 추가 ($P_{os+g+p}$): 메모리 감소는 DP degree $N_d$에 비례한다. 예를 들어, 64개의 GPU($N_d=64$)에 분할하면 64배 메모리 감소를 가져온다. 통신 볼륨은 50% 소폭 증가한다.

ZeRO-DP는 메모리 중복을 제거하고 클러스터의 전체 집계 메모리 용량을 사용할 수 있도록 한다. 세 가지 단계가 모두 활성화되면, ZeRO는 단 1024개의 NVIDIA GPU에서 1조 개의 파라미터 모델을 학습할 수 있다. 16비트 정밀도로 Adam [6]과 같은 optimizer를 사용하는 1조 개의 파라미터 모델은 optimizer state, gradient, parameter를 저장하는 데 약 16테라바이트(TB)의 메모리가 필요하다. 16TB를 1024로 나누면 16GB가 되며, 이는 GPU(예: 32GB 온디바이스 메모리)에 대해 합리적인 범위 내에 있다.

Figure 1: ZeRO-DP 최적화의 세 가지 단계에 따른 장치당 model state 메모리 소비 비교. $\Psi$는 모델 크기(파라미터 수), $K$는 optimizer state의 메모리 승수, $N_d$는 DP degree를 나타낸다. 이 예시에서는 Adam optimizer를 사용한 mixed-precision 학습을 기반으로 모델 크기 $\Psi=7.5B$, DP degree $N_d=64$, $K=12$로 가정한다.

Residual State Memory 최적화
ZeRO-DP가 model state의 메모리 효율성을 높인 후, activation, temporary buffer, 사용 불가능한 메모리 조각에 의해 소비되는 나머지 메모리가 두 번째 메모리 병목 현상이 될 수 있다. 우리는 이 세 가지 요인에 의해 소비되는 residual memory를 각각 최적화하기 위해 ZeRO-R을 개발한다.

1. Activation (backward pass를 수행하기 위해 forward pass에서 저장됨)의 경우, checkpointing [7]이 도움이 되지만 대규모 모델에는 충분하지 않다는 것을 확인했다. 따라서 ZeRO-R은 activation partitioning을 통해 기존 MP 접근 방식에서 activation 복제를 식별하고 제거하여 activation memory를 최적화한다. 또한 적절한 경우 activation을 CPU로 offload한다.
2. ZeRO-R은 메모리 및 연산 효율성의 균형을 맞추기 위해 temporary buffer의 적절한 크기를 정의한다.
3. 우리는 학습 중에 다양한 tensor의 수명 변화로 인해 fragmented memory가 발생하는 것을 관찰한다. 단편화로 인한 연속 메모리 부족은 충분한 여유 메모리가 있더라도 메모리 할당 실패를 유발할 수 있다. ZeRO-R은 tensor의 다른 수명을 기반으로 메모리를 사전에 관리하여 메모리 단편화를 방지한다.

ZeRO-DP와 ZeRO-R이 결합되어 DL 학습을 위한 강력한 메모리 최적화 시스템을 형성하며, 이를 통칭하여 ZeRO라고 한다.

ZeRO와 MP: ZeRO가 DP의 메모리 비효율성을 제거하므로, "MP가 여전히 필요한가, 그리고 언제 필요한가? ZeRO는 MP와 어떻게 작동하는가?"라는 질문이 자연스럽게 제기된다. ZeRO를 사용하면 MP는 대규모 모델을 단독으로 수용하는 목적에 있어서는 덜 매력적인 옵션이 된다. ZeRO-DP는 장치당 메모리 footprint를 줄이는 데 MP만큼 효과적이거나, MP가 모델을 균등하게 분할할 수 없을 때는 더 효과적이다. 또한, 비교할 만하거나 더 나은 확장 효율성을 가진다. 더욱이, Data Parallelism은 사용하기 매우 쉬워 다양한 워크로드에 널리 적용될 수 있는 반면,

Figure 2: 다양한 모델 크기에 대한 ZeRO 학습 처리량 및 SOTA baseline 대비 속도 향상. ZeRO의 경우 MP는 항상 노드에 맞았지만, baseline의 경우 40B보다 큰 모델은 노드 간 MP가 필요했다.

오늘날의 MP 접근 방식은 모델 개발자가 모델을 수정하고, 시스템 개발자가 분산 연산자를 개발하는 데 상당한 노력이 필요하며, Megatron-LM과 같은 기존 작업은 제한된 연산자 및 모델만 지원한다.

그럼에도 불구하고, MP를 활용하고 싶은 경우가 여전히 있다: i) ZeRO-R과 함께 사용될 때, MP는 매우 큰 모델의 activation memory footprint를 줄일 수 있다. ii) activation memory가 문제가 되지 않는 작은 모델의 경우, DP만 사용하여 집계 batch size가 너무 커서 좋은 수렴을 얻기 어려울 때 MP가 이점을 가질 수 있다. 이러한 경우, ZeRO와 MP를 결합하여 허용 가능한 집계 batch size로 모델을 수용할 수 있다.

우리는 ZeRO가 MP와 결합될 수 있음을 보여주며, 이는 DP degree $N_d$와 MP degree $N_m$을 사용하여 각 장치에서 최대 $N_d \times N_m$배의 이론적인 메모리 감소를 가져온다. 이를 통해 1024개의 GPU에서 16-way Model Parallelism(각 DGX2 노드 내)과 노드 간 64-way Data Parallelism을 사용하여 1조 개의 파라미터 모델을 수용하고, 적당한 batch size로 효율적으로 실행할 수 있다!

구현 및 평가
ZeRO의 전체 최적화 세트는 오늘날의 하이엔드 하드웨어 클러스터(예: 1K V100 GPU)에서 1조 개의 파라미터 모델을 실행할 수 있게 하지만, 하드웨어 연산 용량은 여전히 너무 제한적이며 학습 시간이 비현실적으로 길 수 있다(>1년). 따라서 이 구현의 초점은 현재 하드웨어의 연산 능력 범위 내에서 SOTA보다 10배 많은 파라미터(약 100B 파라미터)를 가진 모델을 효율적으로 지원하는 것이다. 우리는 이 목표를 달성하기 위해 ZeRO의 최적화 중 일부인 ZeRO-100B (ZeRO-DP의 $P_{os+g}$와 ZeRO-R)를 구현하고 평가한다. 결과는 다음과 같다:

모델 크기
MP와 결합된 ZeRO-100B는 170B 파라미터 모델을 효율적으로 실행하는 반면, Figure 2에서 보듯이 Megatron 단독 사용과 같은 기존 시스템은 40B 파라미터 이상으로 효율적으로 확장할 수 없다. 이는 SOTA 대비 모델 크기가 8배 이상 증가한 것이다.

속도
향상된 메모리 효율성은 더 높은 처리량과 더 빠른 학습을 가능하게 한다. Figure 2에서 보듯이, ZeRO는 400개의 Nvidia V100 GPU 클러스터에서 100B 파라미터 모델을 GPU당 38 TFlops 이상으로 실행하며, 총 성능은 15 Petaflops 이상이다. 이는 동일한 모델 크기에 대해 SOTA 대비 학습 속도가 10배 이상 향상된 것이다.

[^1] Figure 3: ZeRO-100B를 사용한 60B 파라미터 모델의 초선형 확장성 및 GPU당 학습 처리량.

확장성
우리는 64-400 GPU 범위에서 초선형 속도 향상을 관찰했으며, GPU 수를 두 배로 늘리면 성능이 두 배 이상 증가한다. 이는 DP degree를 늘릴수록 model state의 메모리 footprint를 줄여 GPU당 더 큰 batch size를 수용할 수 있게 하여 더 나은 성능을 가져오는 ZeRO-DP의 특성이다. 우리는 GPU 수를 400개 이상으로 늘려도 이러한 동작이 계속될 것으로 예상한다.

대규모 모델 학습의 민주화
ZeRO-100B는 데이터 과학자가 모델 재구성을 요구하는 MP나 PP 없이도 최대 13B 파라미터 모델을 학습할 수 있도록 지원하며, 13B는 문헌에서 가장 큰 모델(T5의 11B 파라미터)보다 더 많은 파라미터이다. 따라서 데이터 과학자는 병렬화에 대한 걱정 없이 대규모 모델을 자유롭게 실험할 수 있다. 이에 비해 기존 시스템(예: PyTorch Distributed Data Parallel)은 1.4B 파라미터 모델에서도 메모리 부족 현상이 발생한다.

새로운 SOTA 모델
ZeRO는 17B 파라미터의 가장 큰 언어 모델과 기록적인 정확도를 가진 Turing-NLG를 구동한다.

우리는 ZeRO를 DeepSpeed라고 불리는 오픈 소스 DL 학습 최적화 라이브러리의 일부로 공유한다. 우리는 2020년 5월 말까지 이 논문에서 설명된 모든 구현을 공개하고, ZeRO-DP 3단계 파라미터 분할($P_{os+g+p}$)을 활성화하여 1조 개의 파라미터까지 지원하도록 확장할 계획이다. 우리는 ZeRO를 DL 커뮤니티에 완전히 공개하여 대규모 모델 학습의 진화와 민주화를 촉진할 계획이다.

[^2]

2 Related Work

2.1 Data, Model and Pipeline Parallelism

대규모 모델을 학습할 때 병렬화(parallelization)는 핵심적인 전략이다. 학습을 위해 장치 메모리에 맞는 모델의 경우, **데이터 병렬화(Data Parallelism, DP)**가 여러 장치로 학습을 확장하는 데 사용된다. DP에서는 모델 파라미터가 각 장치에 복제된다. 각 단계에서 미니 배치(mini-batch)는 모든 데이터 병렬 프로세스에 균등하게 분할되어, 각 프로세스는 데이터 샘플의 다른 하위 집합에서 forward 및 backward propagation을 실행하고, 프로세스 간 평균화된 gradient를 사용하여 모델을 로컬에서 업데이트한다.

모델이 장치 메모리에 맞지 않을 때는 모델 병렬화(Model Parallelism, MP) [5, 3]와 파이프라인 병렬화(Pipeline Parallelism, PP) [10, 11]가 각각 수직적(vertical) 및 수평적(horizontal) 방식으로 모델을 프로세스 간에 분할한다. Section 1에서는 ZeRO가 DP 및 MP와 어떻게 관련되는지 논의했다. 이제 PP와 그것이 메모리 소비 감소와 어떻게 관련되는지 논의한다.

PP는 모델을 layer별로 수평적으로 분할하여 각 파티션을 다른 장치에서 실행하고, micro-batching을 사용하여 파이프라인 버블(pipeline bubble)을 숨긴다 [10, 11]. tied-weights 및 batch-normalization과 같은 모델 기능은 수평적 분할 및 micro-batching으로 인해 구현하기 어렵다. G-pipe [10]와 같은 인기 있는 PP 구현은 모델 파라미터와 전체 activation을 모두 분할하지만, 파이프라인 버블을 숨기기 위해 파이프라인 파티션 수에 비례하는 배치 크기를 요구한다. 큰 배치 크기는 수렴 속도에 영향을 미칠 수 있으며, activation을 저장하는 데 상당한 메모리를 요구하기도 한다. PipeDream [12]의 다른 PP 구현은 배치 크기를 크게 늘리지 않고 파이프라인 버블을 숨기기 위해 여러 개의 오래된(stale) 파라미터 복사본을 유지하여 메모리 효율성이 떨어진다. 또한, 이 구현은 표준 DL 학습과 동일하지 않으며 학습 수렴에 영향을 미친다. 이와 대조적으로, ZeRO는 PP의 기능, 성능 및 수렴 관련 제약 없이 PP와 동일하거나 더 나은 메모리 효율성을 달성한다.

2.2 Non-parallelism based approach to reduce memory

MP와 PP 외에도, DL 학습의 메모리 오버헤드를 줄이는 것을 목표로 하는 여러 연구들이 있다.

2.2.1 Reducing Activation Memory

활성화(activation)의 메모리 사용량을 줄이기 위한 여러 연구가 압축(compression) [13], activation checkpointing [7, 14], 또는 실시간 분석(live analysis) [15]에 초점을 맞춰 진행되었다. 이러한 노력들은 상호 보완적이며 ZeRO와 함께 작동할 수 있다. 실제로 ZeRO-R의 activation memory reduction 기능은 activation checkpointing과 병렬적으로 작동한다.

2.2.2 CPU Offload

16, 17]은 각각 알고리즘 설계 또는 가상화된 메모리를 통해 모델 상태를 CPU 메모리로 오프로드하여 오늘날 컴퓨팅 노드의 이기종 특성을 활용한다. 학습 시간의 최대 50%가 GPU-CPU-GPU 전송에 소요될 수 있다 [16]. ZeRO는 PCI-E로 인해 대역폭이 심각하게 제한되는 CPU 메모리에 모델 상태를 저장하지 않고도 메모리 소비를 크게 줄인다는 점에서 다르다. 드문 경우지만, ZeRO-R은 성능 향상을 위해 매우 큰 모델의 activation checkpoint만 오프로드할 수 있다 (자세한 내용은 6.1절 참조).

2.2.3 Memory Efficient Optimizer

18, 19는 모델 파라미터와 gradient의 더 거친(coarser-grained) 통계를 유지함으로써 adaptive optimization method의 메모리 소비를 줄이는 데 초점을 맞추고 있으며, 이는 모델 수렴 보장에 잠재적인 영향을 미칠 수 있다. ZeRO는 이러한 노력들과 **직교(orthogonal)**하며, 그 최적화는 모델 최적화 방법을 변경하거나 모델 수렴에 영향을 미치지 않는다. 대신, 디바이스당 optimizer state와 gradient의 메모리 사용량을 효과적으로 줄여준다.

2.3 Training Optimizers

Adaptive optimization method [20, 6, 21, 22]는 대규모 모델의 효과적인 학습을 통해 SOTA 성능과 정확도를 달성하는 데 필수적이다. 이 방법들은 SGD와 비교하여, 각 모델 파라미터와 gradient에 대한 세밀한 1차 및 2차 통계량을 유지하지만, 이는 상당한 메모리 사용량을 초래한다. ZeRO는 이러한 optimizer들의 메모리 사용량을 획기적으로 줄여주어, 제한된 장치 메모리를 가진 하드웨어에서도 대규모 모델 학습에 이러한 정교한 최적화 방법들을 실용적으로 사용할 수 있게 한다. 또한, 이는 더 나은 수렴 성능을 가질 수 있는, 더욱 복잡하고 메모리를 많이 사용하는 optimizer들을 개발하고 활용하는 것을 가능하게 한다.

3 Where Did All the Memory Go?

현재 학습 시스템의 메모리 소비를 자세히 살펴보자. 예를 들어, 1.5B 파라미터 GPT-2 모델은 16비트 정밀도에서 가중치(파라미터)에 3GB의 메모리를 필요로 하지만, 32GB 메모리를 가진 단일 GPU에서 Tensorflow나 PyTorch를 사용하여 학습될 수 없다. 모든 메모리가 어디로 가는지 궁금할 수 있다. 모델 학습 중 대부분의 메모리는 모델 상태(model states), 즉 옵티마이저 상태(optimizer states), gradient, 파라미터로 구성된 텐서에 의해 소비된다. 이러한 모델 상태 외에 나머지 메모리는 **activation, 임시 버퍼, 단편화된 메모리(fragmented memory)**에 의해 소비되는데, 이를 **잔여 상태(residual states)**라고 부른다. 우리는 이 두 가지 측면에서 메모리 소비를 자세히 살펴본다.

3.1 Model States: Optimizer States, Gradients and Parameters

학습 중에는 모델 상태(model states)가 대부분의 장치 메모리를 소비한다. 예를 들어, DL 학습에서 가장 널리 사용되는 optimizer 중 하나인 **Adam [6]**을 고려해보자. Adam은 업데이트를 계산하기 위해 두 가지 optimizer state를 저장해야 한다: (i) 시간 평균 momentum과 (ii) gradient의 분산(variance). 따라서 Adam으로 모델을 학습시키려면, momentum과 gradient의 variance 사본을 모두 저장할 수 있는 충분한 메모리가 있어야 한다. 또한, gradient와 가중치(weights) 자체를 저장할 수 있는 메모리도 필요하다. 이 세 가지 유형의 파라미터 관련 텐서 중에서, 특히 mixed-precision training이 적용될 때 optimizer state가 가장 많은 메모리를 소비한다.

Mixed-Precision Training
현재 세대의 NVIDIA GPU에서 대규모 모델을 학습시키는 최신 접근 방식은 mixed precision (fp16/32) training이다 [23]. 이 방식에서는 파라미터와 activation이 fp16으로 저장되어, GPU의 고처리량 tensor core unit [24]을 활용할 수 있게 한다. mixed-precision training 중에는 forward 및 backward propagation이 모두 fp16 가중치와 activation을 사용하여 수행된다. 그러나 backward propagation이 끝난 후 업데이트를 효과적으로 계산하고 적용하기 위해, mixed-precision optimizer는 파라미터의 fp32 사본과 다른 모든 optimizer state의 fp32 사본을 유지한다.

Adam을 구체적인 예시로 들어보자. Adam을 사용하여 $\Psi$ 파라미터를 가진 모델을 mixed precision training하려면, 파라미터와 gradient의 fp16 사본을 저장할 수 있는 충분한 메모리가 필요하며, 각각 $2\Psi$ 바이트와 $2\Psi$ 바이트의 메모리를 요구한다. 또한, optimizer state를 저장해야 한다: 파라미터, momentum, variance의 fp32 사본이 필요하며, 각각 $4\Psi$, $4\Psi$, $4\Psi$ 바이트의 메모리를 요구한다.
optimizer state의 메모리 승수(memory multiplier)를 $K$라고 하자. 즉, optimizer state를 저장하는 데 필요한 추가 메모리는 $K\Psi$ 바이트이다. **Mixed-precision Adam의 경우 $K=12$**이다.
총 메모리 요구량은 $2\Psi + 2\Psi + K\Psi = 16\Psi$ 바이트가 된다.
15억 개의 파라미터를 가진 GPT-2와 같은 모델의 경우, 이는 최소 24GB의 메모리를 요구하며, 이는 fp16 파라미터만을 저장하는 데 필요한 3GB의 메모리보다 훨씬 높은 수치이다.

3.2 Residual Memory Consumption

Activations는 학습 중에 상당한 양의 메모리를 차지할 수 있다 [7]. 구체적인 예시로, 1.5B 파라미터 GPT-2 모델을 시퀀스 길이 1K, 배치 크기 32로 학습할 경우 약 60GB의 메모리가 필요하다. Activation checkpointing (또는 activation recomputation)은 전체 activation 메모리를 약 제곱근 수준으로 줄이는 일반적인 접근 방식이며, 이는 33%의 재계산 오버헤드를 감수한다. 이 방식을 사용하면 해당 모델의 activation 메모리 사용량을 약 8GB로 줄일 수 있다.

이러한 상당한 감소에도 불구하고, 더 큰 모델의 경우 activation checkpointing을 사용하더라도 activation 메모리는 상당히 커질 수 있다. 예를 들어, 1,000억 개의 파라미터를 가진 GPT-유사 모델은 activation checkpointing을 사용하더라도 배치 크기 32에서 약 60GB의 메모리를 필요로 한다.

중간 결과를 저장하는 데 사용되는 임시 버퍼는 대규모 모델에서 상당한 양의 메모리를 소비한다. Gradient all-reduce 또는 gradient norm 계산과 같은 연산은 처리량(throughput)을 개선하기 위해 연산을 적용하기 전에 모든 gradient를 단일의 평탄화된(flattened) 버퍼로 융합(fuse)하는 경향이 있다. 예를 들어, 장치 간 all-reduce의 대역폭은 메시지 크기가 클수록 향상된다. Gradient 자체는 일반적으로 fp16 텐서로 저장되지만, 융합된 버퍼는 연산에 따라 fp32 텐서일 수 있다. 모델의 크기가 클 때, 이러한 임시 버퍼의 크기는 무시할 수 없다. 예를 들어, 1.5B 파라미터를 가진 모델의 경우, 평탄화된 fp32 버퍼는 6GB의 메모리를 필요로 한다.

메모리 단편화 (Memory Fragmentation): 지금까지는 학습 중 실제 메모리 소비에 대해 논의했다. 추가적으로, 사용 가능한 메모리가 충분함에도 불구하고 사용 가능한 메모리가 부족해지는 상황이 발생할 수 있다. 이는 메모리 단편화로 인해 발생한다. 요청된 메모리보다 총 사용 가능한 메모리가 더 많더라도, 요청을 충족할 만큼 충분한 연속적인 메모리가 없으면 메모리 요청은 실패한다. 우리는 매우 큰 모델을 학습할 때 상당한 메모리 단편화를 관찰했으며, 이는 일부 극단적인 경우 30% 이상의 메모리가 여전히 사용 가능함에도 불구하고 메모리 부족(out of memory) 문제를 야기했다.

4 ZeRO: Insights and Overview

ZeRO는 두 가지 최적화 세트를 포함한다: i) ZeRO-DP: 모델 상태(model states)의 메모리 사용량(memory footprint)을 줄이는 데 중점을 둔다. ii) ZeRO-R: 잔여 메모리 소비량(residual memory consumption)을 줄이는 데 중점을 둔다.

우리는 이러한 최적화 방식과 그 배경이 되는 통찰력을 개괄적으로 설명하며, 이를 통해 ZeRO가 효율성을 유지하면서도 메모리 사용량을 줄일 수 있는 이유를 제시한다. 여기서 효율성이 핵심 제약 조건이라는 점에 유의해야 한다. 이 제약이 없다면, 모든 파라미터 상태를 CPU 메모리로 옮기거나 MP(Model Parallelism) 정도를 임의로 늘리는 것과 같은 trivial한 해결책으로도 메모리 사용량을 줄일 수 있을 것이다.

4.1 Insights and Overview: ZeRO-DP

ZeRO 기반의 DP는 세 가지 핵심 통찰력에 기반한다: a) DP는 MP보다 확장 효율성(scaling efficiency)이 더 좋다. 그 이유는 MP가 계산의 granularity를 줄이는 동시에 통신 오버헤드를 증가시키기 때문이다. 특정 지점을 넘어서면, 낮은 계산 granularity는 GPU당 효율성을 감소시키고, 증가된 통신 오버헤드는 GPU 간 확장성(특히 노드 경계를 넘을 때)을 저해한다. 반대로 DP는 더 높은 계산 granularity와 더 낮은 통신량을 가지므로, 훨씬 더 높은 효율성을 허용한다. b) DP는 모델 상태(model states)가 모든 데이터 병렬 프로세스에 걸쳐 중복으로 저장되기 때문에 메모리 비효율적이다. 반대로 MP는 모델 상태를 분할하여 메모리 효율성을 확보한다. c) DP와 MP 모두 전체 학습 과정에 필요한 모든 모델 상태를 유지하지만, 모든 것이 항상 필요한 것은 아니다. 예를 들어, 각 layer에 해당하는 파라미터는 해당 layer의 forward propagation과 backward propagation 동안에만 필요하다.

이러한 통찰력을 바탕으로, ZeRO-DP는 DP의 학습 효율성을 유지하면서 MP의 메모리 효율성을 달성한다. ZeRO-DP는 모델 상태를 복제하는 대신 분할하고(Section 5), 모델 상태의 본질적인 시간적 특성을 활용하면서 통신량을 최소화하는 동적 통신 스케줄을 사용한다(Section 7). 이를 통해 ZeRO-DP는 모델의 장치당 메모리 사용량(per-device memory footprint)을 DP degree에 비례하여 선형적으로 감소시키면서, 통신량은 기본 DP와 유사하게 유지하여 효율성을 보존한다.

4.2 Insights and Overview: ZeRO-R

4.2.1 Reducing Activation Memory

두 가지 핵심 통찰은 다음과 같다: a) **MP(Model Parallelism)**는 모델 상태를 분할하지만, 종종 activation memory의 복제를 요구한다. 예를 들어, 선형 레이어의 파라미터를 수직으로 분할하여 두 개의 GPU에서 병렬로 계산할 경우, 각 GPU는 해당 파티션을 계산하기 위해 전체 activation을 필요로 한다. b) GPT-2 또는 그 이상의 대규모 모델의 경우, 산술 강도(arithmetic intensity)(반복당 계산량 대 반복당 activation checkpoint량의 비율)가 매우 크고($\geq 10K$), hidden dimension에 비례하여 증가한다. 이는 대역폭이 낮더라도 activation checkpoint를 위한 데이터 이동 비용을 숨길 수 있게 한다.

ZeRO는 activation checkpoint를 GPU 간에 분할함으로써 MP의 메모리 중복성을 제거하고, 필요할 때 allgather를 사용하여 이를 재구성한다. 이로 인해 activation memory footprint는 MP degree에 비례하여 감소한다. 매우 큰 모델의 경우, ZeRO는 activation 파티션을 CPU 메모리로 이동시키는 것을 선택할 수도 있으며, 이러한 모델의 높은 산술 강도 덕분에 여전히 좋은 효율성을 달성할 수 있다.

4.2.2 Managing Temporary buffers

ZeRO-R은 모델 크기가 증가함에 따라 임시 버퍼가 과도하게 커지는 것을 방지하기 위해 고정된 크기의 버퍼를 사용하며, 동시에 효율성을 유지할 수 있을 만큼 충분히 큰 크기로 버퍼를 설정한다.

4.2.3 Managing fragmented Memory

메모리 단편화(memory fragmentation)는 수명이 짧은(short lived) 메모리 객체와 수명이 긴(long lived) 메모리 객체가 교차(interleaving)하여 할당될 때 발생한다. forward propagation 중에는 activation checkpoint가 수명이 길지만, 재계산되는 activation은 수명이 짧다. 마찬가지로, backward computation 중에는 activation gradient가 수명이 짧은 반면, parameter gradient는 수명이 길다. 이러한 통찰을 바탕으로,

ZeRO는 activation checkpoint와 gradient를 미리 할당된 연속적인 메모리 버퍼로 이동시켜 실시간(on-the-fly) 메모리 조각 모음(defragmentation)을 수행한다. 이는 메모리 가용성을 높일 뿐만 아니라, 메모리 할당자가 사용 가능한 연속적인 메모리를 찾는 데 걸리는 시간을 줄여 효율성을 향상시킨다.

5 Deep Dive into ZeRO-DP

기존 DP(Data Parallel) 방식은 각 장치에 모델 상태를 복제하여 상당한 메모리 오버헤드를 발생시키는 반면, ZeRO-DP는 optimizer state, gradient, parameter를 데이터 병렬 프로세스에 걸쳐 분할함으로써 이러한 메모리 중복성을 제거한다. Figure 1은 ZeRO-DP를 적용했을 때와 적용하지 않았을 때의 메모리 요구량을 정량화하여 시각적으로 보여준다. 이 그림은 (1) optimizer state, (2) gradient, (3) parameter 중복성을 누적적으로 분할한 후의 메모리 사용량을 나타낸다. 우리는 이를 ZeRO-DP의 세 가지 최적화 단계인 **$P_{os}, P_{g}, P_{p}$**라고 부르며, 이에 대해서는 아래에서 자세히 설명한다.

5.1 $\mathbf{P}_{o s}$ : Optimizer State Partitioning

DP(Data Parallel) 정도가 $N_d$일 때, 우리는 optimizer state를 $N_d$개의 동일한 파티션으로 그룹화하여, $i$번째 data parallel 프로세스가 $i$번째 파티션에 해당하는 optimizer state만 업데이트하도록 한다. 따라서 각 data parallel 프로세스는 전체 optimizer state의 $\frac{1}{N_d}$만 저장하고 업데이트하면 되며, 이후 파라미터의 $\frac{1}{N_d}$만 업데이트한다. 각 학습 단계의 끝에는 모든 data parallel 프로세스에 걸쳐 all-gather를 수행하여 완전히 업데이트된 파라미터를 모든 data parallel 프로세스에 걸쳐 얻는다.

메모리 절약 (Memory Savings): Figure 1에서 볼 수 있듯이, optimizer state 파티션 이후의 메모리 소비는 $4\Psi + K\Psi$에서 $4\Psi + \frac{K\Psi}{N_d}$로 감소한다. Figure 1에 제시된 구체적인 예시처럼, 7.5B 파라미터 모델은 64-way DP($N_d=64$)를 사용하는 $P_{os}$ 방식으로는 31.4GB의 메모리를 필요로 하는 반면, 표준 DP 방식으로는 120GB를 필요로 한다. 더욱이, $N_d$가 클 때 모델 상태에 대한 메모리 요구량은 $4\Psi + 12\Psi = 16\Psi$ 바이트에서 $4\Psi + \frac{12\Psi}{N_d} \approx 4\Psi$ 바이트로 감소하여, 4배의 감소 효과를 가져온다.

5.2 $\mathbf{P}_{g}$ : Gradient Partitioning

각 데이터 병렬(data parallel) 프로세스는 해당 파라미터 파티션만 업데이트하므로, 해당 파라미터에 대한 reduced gradient만 필요하다. 따라서 역전파(backward propagation) 중 각 레이어의 gradient가 사용 가능해지면, 해당 파라미터를 업데이트하는 데이터 병렬 프로세스에서만 gradient를 reduce한다. reduction 후에는 gradient가 더 이상 필요 없으므로 메모리를 해제할 수 있다. 이는 gradient를 유지하는 데 필요한 메모리 공간을 $2 \Psi$ 바이트에서 $\frac{2 \Psi}{N_{d}}$ 바이트로 줄여준다.

이는 효과적으로 Reduce-Scatter 연산과 같다. 여기서 서로 다른 파라미터에 해당하는 gradient는 서로 다른 프로세스로 reduce된다. 실제 환경에서 이를 더 효율적으로 만들기 위해, 우리는 bucketization 전략을 사용한다. 이 전략에서는 특정 파티션에 해당하는 모든 gradient를 bucket으로 묶어 한 번에 reduction을 수행한다. 이는 NVIDIA의 AMP 25 optimizer가 all-reduce gradient 계산을 bucket으로 묶어 통신과 계산을 중첩시키는 방식과 유사하다. 우리의 경우, 메모리 공간을 줄이고 계산과 통신을 중첩시키기 위해 파티션 경계에서 all-reduce 대신 reduction을 수행한다.

메모리 절약: gradient 및 optimizer state 중복을 제거함으로써, 메모리 공간을 $2 \Psi+\frac{14 \Psi}{N_{d}} \approx 2 \Psi$ 바이트로 더욱 줄일 수 있다. 이는 약 8배의 감소를 의미한다. Figure 1의 예시처럼, 7.5B 파라미터 모델은 64-way DP($N_{d}=64$)를 사용하는 $P_{os+g}$ 방식에서 16.6GB의 메모리만 필요하지만, 표준 DP에서는 120GB가 필요하다. $N_{d}$가 클 때, 모델 상태의 메모리 요구량은 $2 \Psi+14 \Psi=16 \Psi$ 바이트에서 $2 \Psi+\frac{14 \Psi}{N_{d}} \approx 2 \Psi$ 바이트로 줄어든다.

Table 1: DP 정도에 따른 $Z e R O-\mathrm{DP}$의 다양한 최적화 방식별 장치당 메모리 소비량. 굵게 표시된 텍스트는 모델이 32GB V100 GPU 클러스터에 적합한 조합을 나타낸다.

$5.3 \mathbf{P}_{p}$ : Parameter Partitioning

Optimizer state 및 gradient와 마찬가지로, 각 프로세스는 자신의 파티션에 해당하는 파라미터만 저장한다. forward 및 backward propagation에 자신의 파티션 외의 파라미터가 필요한 경우, 해당 파라미터는 적절한 data parallel 프로세스로부터 broadcast를 통해 수신된다. 언뜻 보기에 상당한 통신 오버헤드가 발생할 것처럼 보이지만, 우리는 이 방식이 baseline DP 시스템의 총 통신량을 1.5배만 증가시키면서도, $N_d$에 비례하는 메모리 절감을 가능하게 함을 보여준다.

메모리 절감 (Memory Savings): 파라미터 파티셔닝을 통해, 우리는 $\Psi$ 파라미터 모델의 메모리 소비를 $16\Psi$에서 $\frac{16\Psi}{N_d}$로 줄인다. Figure 1의 예시처럼, 7.5B 파라미터 모델은 64-way DP($N_d=64$)를 사용하는 $P_{os+p+g}$ 방식에서 1.9GB의 모델 상태 메모리를 필요로 하는 반면, 표준 DP에서는 120GB를 필요로 한다. 이는 ZeRO가 모델 상태를 공유할 충분한 수의 장치만 있다면, DP가 임의의 크기 모델을 수용할 수 있도록 지원한다는 중요한 의미를 가진다.

5.4 Implication on Model Size

$P_{os}$, $P_{os+g}$, $P_{os+g+p}$의 세 가지 파티셔닝 단계는 모델 상태(model states)에 대한 각 데이터 병렬(data parallel) 프로세스의 메모리 소비를 각각 최대 4배, 8배, $N_d$배까지 줄여준다.
Table 1은 다양한 DP(Data Parallel) 정도에 대해 ZeRO-DP 최적화의 3단계에서 몇 가지 예시 모델의 모델 상태 메모리 소비를 분석한다. ZeRO가 없을 경우, 메모리 소비는 DP 정도와 관계없이 테이블의 첫 번째 행과 동일하다.
주목할 점은, $N_d=64$일 때, ZeRO는 $P_{os}$, $P_{os+g}$, $P_{os+g+p}$를 각각 사용하여 최대 7.5B, 14B, 128B 파라미터를 가진 모델을 학습할 수 있다는 것이다.
$N_d=1024$일 때, **모든 최적화가 활성화된 ZeRO ($P_{os+g+p}$)**는 1조(Trillion) 파라미터를 가진 모델을 학습할 수 있다! 또는 잠재적으로 **임의의 크기(Arbitrary size)**의 모델도 가능하다!
ZeRO가 없을 경우, DP만으로 실행할 수 있는 가장 큰 모델은 1.5B 파라미터 미만이다.

6 Deep Dive into ZeRO-R

6.1 $P_{a}$ : Partitioned Activation Checkpointing

Section 4.2에서 논의했듯이, **MP(Model Parallelism)**는 설계상 **activation의 복제(replication)**를 필요로 하여, 모델 병렬 GPU 전반에 걸쳐 activation의 중복 복사본을 생성한다. ZeRO는 **activation을 분할(partitioning)**함으로써 이러한 중복성을 제거하며, **activation이 계산에 사용되기 직전, 한 번에 하나의 activation layer만 복제된 형태로 구체화(materialize)**한다.
더 구체적으로, 모델의 한 layer에 대한 forward propagation이 계산되면, 입력 activation은 모든 model parallel 프로세스에 걸쳐 분할되며, 이는 backpropagation 중에 다시 필요할 때까지 유지된다. 이때 ZeRO는 all-gather 연산을 사용하여 activation의 복제된 복사본을 다시 구체화한다. 우리는 이 최적화를 **$P_a$**라고 부른다. 이는 activation checkpointing [7]과 함께 작동하며, 복제된 복사본 대신 분할된 activation checkpoint만 저장한다.
또한, 매우 큰 모델과 매우 제한된 장치 메모리의 경우, 이러한 분할된 activation checkpoint는 CPU로 오프로드될 수도 있어, 추가적인 통신 비용으로 activation 메모리 오버헤드를 거의 0으로 줄일 수 있다. 이에 대해서는 Section 7에서 논의할 것이다. 우리는 이를 **$P_{a+cpu}$**라고 부른다.

메모리 절약 (Memory Saving)
분할된 activation checkpointing을 통해 ZeRO는 activation footprint를 MP degree에 비례하는 요소만큼 줄인다.
Table 4에 제시된 100B 모델을 batch size 32, sequence length 1024, MP degree 16으로 학습하는 경우를 고려해보자. 각 Transformer layer에 대해 단일 activation을 checkpointing한다면, activation checkpoint를 저장하는 데에만 GPU당 약 33GB의 메모리가 필요할 것이다. 그러나 ZeRO의 $P_a$를 사용하면 GPU당 약 2GB로 줄일 수 있다.
더 나아가, 이 2GB는 CPU로 오프로드되어 activation을 위한 메모리 footprint를 거의 0으로 줄일 수 있다.

6.2 $\mathrm{C}_{B}$ : Constant Size Buffers

ZeRO는 메모리 및 연산 효율성의 균형을 맞추기 위해 temporal-data buffer의 크기를 신중하게 선택한다. 학습 과정에서 일부 연산의 계산 효율성은 입력 크기에 크게 의존하며, 입력이 클수록 더 높은 효율성을 달성한다. 예를 들어, 큰 all-reduce 연산은 작은 연산보다 훨씬 높은 대역폭을 얻을 수 있다. 따라서 더 나은 효율성을 얻기 위해 NVIDIA Apex 또는 Megatron과 같은 고성능 라이브러리는 이러한 연산을 적용하기 전에 **모든 파라미터를 단일 버퍼로 융합(fuse)**한다. 그러나 융합된 버퍼의 메모리 오버헤드는 모델 크기에 비례하며, 이는 제한적인 요소가 될 수 있다. 예를 들어, 3B 파라미터 모델의 경우, 32비트 융합 버퍼는 12GB의 메모리를 필요로 한다. 이 문제를 해결하기 위해, 우리는 모델이 너무 커질 때 성능 효율적인 고정 크기의 융합 버퍼를 사용한다. 이렇게 함으로써 버퍼 크기가 모델 크기에 의존하지 않게 되며, 버퍼 크기를 충분히 크게 유지함으로써 여전히 좋은 효율성을 달성할 수 있다.

6.3 $\mathrm{M}_{D}$ : Memory Defragmentation

모델 학습 시 **메모리 단편화(memory fragmentation)**는 activation checkpointing과 gradient 계산으로 인해 발생한다. activation checkpointing을 사용하는 forward propagation 과정에서, back propagation을 위해 선택된 activation만 저장되고, 대부분의 activation은 back propagation 중에 다시 계산될 수 있으므로 버려진다. 이로 인해 수명이 짧은 메모리(버려진 activation)와 수명이 긴 메모리(checkpoint된 activation)가 교차되어 메모리 단편화가 발생한다. 마찬가지로, backward propagation 중에는 parameter gradient는 수명이 길지만, activation gradient 및 parameter gradient 계산에 필요한 다른 버퍼들은 수명이 짧다. 다시 한번, 수명이 짧은 메모리와 수명이 긴 메모리의 이러한 교차는 메모리 단편화를 유발한다.

Table 2: 메모리 분석을 통한 최대 모델 크기 (왼쪽) 및 $Z e R O-O S$ 실행 시 측정된 모델 크기 (오른쪽). $P_{o s}$로 측정된 모델 크기는 이론적 최대값과 일치하며, 이는 우리의 메모리 분석이 모델 크기에 대한 현실적인 상한선을 제공함을 보여준다.

메모리 여유가 충분할 때는 제한적인 메모리 단편화는 일반적으로 문제가 되지 않지만, 제한된 메모리로 대규모 모델을 학습할 때 메모리 단편화는 두 가지 문제를 야기한다: i) 충분한 가용 메모리가 있음에도 불구하고 연속적인 메모리 부족으로 인한 OOM(Out Of Memory), ii) 메모리 할당자가 메모리 요청을 충족시키기 위해 연속적인 메모리 조각을 찾는 데 상당한 시간을 소비하여 발생하는 낮은 효율성.

ZeRO는 activation checkpoint와 gradient를 위해 연속적인 메모리 청크를 미리 할당하고, 이들이 생성될 때 미리 할당된 메모리로 복사함으로써 실시간으로 메모리 단편화를 제거한다. **$\mathrm{M}_{D}$**는 ZeRO가 더 큰 배치 크기로 더 큰 모델을 학습할 수 있도록 할 뿐만 아니라, 제한된 메모리로 학습할 때 효율성도 향상시킨다.

7 Communication Analysis of ZeRO-DP

$Z e R O$가 메모리 중복성을 제거하여 모델 크기를 확장하는 방식이므로, 메모리 효율성을 위해 통신량을 희생하는 것은 아닌지 질문하는 것은 당연하다. 즉, $Z e R O$ 기반 DP(Data Parallelism) 접근 방식의 통신량은 baseline DP 접근 방식과 비교하여 어떤가?

답변은 두 부분으로 나뉜다: i) $Z e R O$-DP는 $P_{o s}$와 $P_{g}$를 사용할 때 추가적인 통신량을 발생시키지 않으면서 최대 8배의 메모리 절감을 가능하게 한다. ii) $Z e R O$-DP는 $P_{o s}$와 $P_{g}$ 외에 $P_{p}$를 사용할 때 최대 1.5배의 통신량을 발생시키면서 메모리 사용량을 $N_{d}$배 더 줄인다.

이 섹션에서는 이에 대한 분석을 제시한다. 먼저 표준 DP의 통신량에 대한 간략한 개요부터 시작한다.

7.1 Data Parallel Communication Volume

데이터 병렬 학습 중에는 backward propagation이 끝난 후, 다음 단계의 업데이트를 계산하기 전에 모든 데이터 병렬 프로세스 간의 gradient가 평균화된다. 이 평균화는 **all-reduce 통신 집합(collective)**을 사용하여 수행된다. 모델 크기가 클 경우, all-reduce 통신은 전적으로 통신 대역폭에 의해 제한되므로, 우리는 각 데이터 병렬 프로세스로부터 송수신되는 **총 통신량(communication volume)**에 분석을 한정한다.

All-reduce의 최신 구현은 두 단계 접근 방식을 사용한다. 첫 번째 단계는 reduce-scatter 연산으로, 서로 다른 프로세스에서 데이터의 다른 부분을 축소(reduce)한다. 다음 단계는 all-gather 연산으로, 각 프로세스가 모든 프로세스에서 축소된 데이터를 수집(gather)한다. 이 두 단계의 결과가 all-reduce이다. Reduce-scatter와 all-gather는 모두 파이프라인 방식으로 구현되어, 각각 $\Psi$ 요소(데이터가 $\Psi$ 요소를 가질 때)에 대해 총 $\Psi$ 요소의 데이터 이동을 발생시킨다. 따라서 표준 **DP(Data Parallel)**는 각 학습 단계에서 $2\Psi$의 데이터 이동을 발생시킨다.

7.2 ZeRO-DP Communication Volume

7.2.1 Communication Volume with $P_{o s+g}$

Gradient partitioning을 사용하면, 각 프로세스는 자신이 담당하는 파라미터 파티션을 업데이트하는 데 필요한 gradient 부분만 저장한다. 따라서 $Z e R O$는 all-reduce 대신 scatter-reduce 연산만을 필요로 하며, 이로 인해 발생하는 통신량은 $\Psi$이다. 각 프로세스가 자신이 담당하는 파라미터 파티션을 업데이트한 후에는, 모든 데이터 병렬 프로세스로부터 업데이트된 모든 파라미터를 수집하기 위해 allgather 연산이 수행된다. 이 또한 $\Psi$의 통신량을 발생시킨다. 따라서 각 학습 단계당 총 통신량은 $\Psi+\Psi=2 \Psi$이며, 이는 baseline DP와 정확히 동일하다.

7.2.2 Communication Volume with $P_{o s+g+p}$

파라미터 분할(parameter partitioning) 후, 각 데이터 병렬(data parallel) 프로세스는 자신이 업데이트하는 파라미터만 저장한다. 따라서 forward propagation 중에는 다른 모든 파티션의 파라미터를 받아야 한다. 하지만 이는 메모리 오버헤드를 피하기 위해 파이프라인화될 수 있다.

특정 파티션에 해당하는 모델 부분에서 forward propagation을 계산하기 전에, 해당 파티션을 담당하는 데이터 병렬 프로세스는 가중치를 모든 데이터 병렬 프로세스에 브로드캐스트할 수 있다. 해당 파티션에 대한 forward propagation이 완료되면, 파라미터는 폐기될 수 있다. 따라서 **총 통신량은 $\frac{\Psi \times N_{d}}{N_{d}}=\Psi$**이다. 즉, 파라미터 all-gather를 전체 forward propagation에 걸쳐 분산시키고, 사용된 파라미터는 즉시 폐기함으로써 재조정한다. 그러나 이 all-gather는 backward propagation을 위해 역순으로 다시 한 번 발생해야 한다는 점에 유의해야 한다.

따라서 총 통신량은 이러한 all-gather로 인해 발생하는 통신량과 gradient의 reducescatter로 인해 발생하는 통신량의 합이다. 총 통신량은 $3\Psi$로, 이는 baseline 대비 1.5배이다. gradient 및 parameter partitioning은 모든 gradient 및 parameter 상태가 항상 필요한 것은 아니라는 통찰을 활용하여, 상태를 신중하게 통신함으로써 메모리를 최적화한다.

8 Communication Analysis of $Z e R O-\mathrm{R}$

우리는 ZeRO-R의 partitioned activation checkpointing ($P_a$) 통신량과 baseline MP를 비교하여, $P_a$가 baseline MP 통신량의 10분의 1 미만으로 통신량 증가를 유발함을 보인다. 나아가, $P_a$가 더 큰 batch size를 허용하고 DP 통신량을 줄임으로써 효율성을 향상시키는 시나리오를 파악하기 위해 $P_a$의 통신 오버헤드와 DP 통신량 간의 관계를 분석한다. 우리는 이러한 분석을 활용하여 $P_a$ 및 $P_{a+cpu}$를 적용할지 여부와 시점을 결정한다.

activation checkpoint를 분할하는 통신량 trade-off는 모델 크기, checkpointing 전략, MP 전략에 따라 달라진다. 구체적인 통찰력을 공유하기 위해, 우리는 SOTA MP 접근 방식인 Megatron-LM을 사용하여 구현된 Transformer 기반 모델의 맥락에서 분석을 수행한다.

activation checkpointing이 적용된 Megatron-LM에서 각 Transformer block은 forward propagation에서 batch $\times$ seq_length $\times$ hidden_dim 크기의 all-reduce 연산을 두 번 수행하고, forward re-computation을 위해 두 번, backward propagation에서 두 번 더 수행한다. all-reduce의 통신량은 $2 \times$ message_size이므로, 블록당 총 통신량은 $12 \times$ seq_length $\times$ hidden_dim이다.

ZeRO-R이 activation checkpoint를 분할할 때, 각 activation checkpoint의 backward propagation의 forward recomputation 전에 추가적인 all-gather 연산이 필요하다. 일반적으로 우리는 각 Transformer block의 입력 activation을 checkpointing하며, 이는 Transformer block당 하나의 all-gather를 요구한다. all-gather의 통신량은 message_size이므로, $P_a$의 통신 오버헤드는 seq_length $\times$ hidden_dim이다. 따라서 $P_a$의 총 통신 오버헤드는 model parallelism의 원래 통신량의 10% 미만이다.

MP가 DP와 함께 사용될 때, $P_a$는 model-parallel 통신량의 10% 증가를 대가로 data-parallel 통신량을 한 자릿수(order of magnitude) 줄이는 데 사용될 수 있으며, data-parallel 통신이 성능 병목일 때 효율성을 크게 향상시킨다. $P_a$는 MP degree만큼 activation 메모리 소비를 줄여 batch size를 비례적으로 증가시킨다는 점에 주목해야 한다. 대규모 모델의 경우, MP는 16(DGX-2 노드의 GPU 수)만큼 커질 수 있어 batch size를 최대 16배까지 증가시킬 수 있다. data-parallel 학습의 통신량은 batch size에 반비례한다. 따라서 $P_a$로 인한 batch size의 한 자릿수 증가는 data-parallel 통신량의 한 자릿수 감소로 이어질 수 있다.

마지막으로 $P_{a+cpu}$가 적용되면, 분할된 activation checkpoint는 CPU로 offload되어, $P_a$에 비해 CPU 메모리로의 데이터 이동이 2배 추가되는 대신 activation 메모리 요구 사항이 거의 0으로 줄어든다. $P_a$를 사용하더라도 작은 batch size로 인해 DP 통신량이 주요 병목인 극단적인 경우, CPU 데이터 전송 오버헤드가 DP 통신량 오버헤드보다 작다면(이는 일반적으로 작은 batch size에서 해당됨) $P_{a+cpu}$는 batch size를 증가시켜 효율성을 향상시킬 수 있다.

모델 및 하드웨어 특성을 고려하여, 우리는 위의 분석을 활용하여 $P_a$ 및 $P_{a+cpu}$를 적용할지 여부와 시점을 결정한다.

9 Step Towards 1 Trillion Parameters

오늘날 발표된 가장 큰 모델들은 100억 개 파라미터 범위에 있으며, 이는 이미 학습하기에 도전적인 수준이다. 1조 개 파라미터, 즉 3자릿수 더 큰 규모에 도달하는 것은 필연적으로 일어날 일이지만, 그 길은 수많은 난관과 놀라움, 그리고 혁신으로 가득할 것이다. 우리는 이 모든 것을 알고 있거나 해결한다고 주장하지는 않지만, ZeRO는 시스템 관점에서 가장 근본적인 도전 과제 중 하나를 해결한다: 현재 하드웨어에서 이 규모의 모델을 수용하면서도 우수한 시스템 확장성으로 학습할 수 있도록 하는 능력이다.

State-of-Art를 뛰어넘는 도약
현재 state-of-art 프레임워크인 Megatron이 허용 가능한 처리량으로 학습할 수 있는 가장 큰 모델은 DGX-2 시스템에서 16-20B 파라미터 모델이다. 여러 DGX 노드에 걸쳐 model parallelism을 적용하여 추가로 확장하면, 제한된 노드 간 대역폭으로 인해 효율성이 크게 저하된다.
ZeRO는 효율적으로 실행 가능한 모델 크기를 크게 증가시킨다. 이는 현재 세대의 하드웨어가 노드 경계를 넘어가는 fine-grained model parallelism 없이도 훨씬 더 큰 모델을 실행할 수 있도록 한다. Table 1에서 보여주듯이, ZeRO는 모든 최적화가 켜진 상태($\mathrm{P}_{os+g+p}$)에서 DP만 사용하여 1024개의 GPU에 1조 개 이상의 파라미터를 수용할 수 있었다. 또는 model parallelism과 결합했을 때(Table 2 참조), ZeRO는 1024개의 GPU에서 16-way model parallelism(각 DGX2 노드 내)과 64-way data parallelism(노드 간)으로 1조 개 이상의 파라미터를 수용할 수 있었다. 1조 개 파라미터 모델을 효율적으로 실행하는 것은 더 이상 불가능한 일이 아니다!

연산 능력 격차
그러나 1조 개 파라미터 모델을 허용 가능한 시간 내에 end-to-end로 학습시키려면 여전히 상당한 양의 연산 능력이 필요하며, 이는 오늘날의 AI 클러스터에는 부족한 부분이다.

자원 요구 사항을 이해하기 위해 Bert-Large와 간략하게 비교해 본다. Bert-Large는 1024 GPU DGX-2H 클러스터에서 67분 만에 학습될 수 있다 [26]. 1조 개 파라미터 모델은 데이터 샘플당 Bert-Large 모델보다 쉽게 3000 x (1조 / 3억 3천만) 배 더 많은 연산을 포함할 수 있다.

Figure 4: ZeRO-DP를 사용한 최대 모델 처리량.

Figure 5: ZeRO로 구현된 SOTA Turing-NLG.

Figure 6: 최대 모델 크기.

Figure 7: 할당된 최대 캐시.

Figure 8: GPU당 처리량.

동일한 시퀀스 길이와 모델 학습에 필요한 총 샘플 수가 동일하다고 가정하더라도, 1T 모델을 학습하는 데는 동일한 하드웨어와 유사한 연산 효율성을 가정할 때 140일이 소요될 것이다. 실제로는 데이터 샘플과 시퀀스 길이 모두 모델 크기 증가에 따라 늘어날 가능성이 높으므로, 학습에 1년 이상이 소요될 수 있다. 1T 파라미터 모델을 합리적인 시간 내에 학습시키려면 엑사플롭스(exa-flop) 시스템이 필요할 것이다. 그러나 그러한 연산 능력이 가용해질 때, ZeRO가 1T 모델을 효율적으로 실행할 수 있는 시스템 기술을 제공하기를 바란다.

10 Implementation and Evaluation

우리는 현재 공개된 가장 큰 모델(예: T511B [4])보다 약 10배 큰, 약 100B 파라미터를 가진 모델의 효율적인 학습을 지원하는 데 중점을 두었다. 이러한 모델들은 현재 하드웨어(예: 1K V100 GPU)에서 합리적인 시간 내에 학습될 수 있다. 우리는 이러한 목표를 달성하기 위해 ZeRO-DP와 ZeRO-R을 포함하는 ZeRO-Pos+g의 최적화된 부분 집합을 구현하고 평가하였다. 이 구현을 ZeRO-100B라고 부를 것이다.
우리의 결과는 ZeRO-100B가 최대 170B 파라미터를 가진 모델을 효율적으로 학습할 수 있음을 보여준다. 이는 SOTA 모델보다 8배 크고, 최대 10배 빠르며, 사용 편의성도 향상되었다. ZeRO-100B는 새로운 SOTA 정확도를 달성한 세계 최대 규모의 공개 모델인 Turing-NLG의 기반 기술이다.

10.1 Implementation and Methodology

구현 (Implementation)
우리는 ZeRO-100B를 PyTorch로 구현했으며, 여기에는 Pos+g 및 ZeRO-R의 모든 최적화 기능이 포함되어 있다. 이 인터페이스는 torch.nn.module로 구현된 모든 모델과 호환된다. 사용자는 이 인터페이스를 사용하여 모델을 래핑하기만 하면, 기존의 classic DP를 사용하듯이 ZeRO-powered DP를 활용할 수 있다. 모델을 수정할 필요가 없다. ZeRO-powered DP는 Megatron-LM을 포함한 모든 형태의 MP와 결합될 수 있다.

하드웨어 (Hardware)
우리는 **400개의 V100 GPU (25개의 DGX-2 노드)**와 800Gbps의 노드 간 통신 대역폭을 갖춘 클러스터에서 실험을 수행했다.

Baseline
MP를 사용하지 않는 실험의 경우, **torch의 distributed data parallel (DDP)**을 baseline으로 사용한다. MP를 사용하는 실험의 경우, 우리가 아는 한 state-of-the-art인 Megatron-LM을 사용한다. 우리는 NVIDIA에서 2019년 9월에 공개된 Megatron-LM의 오픈소스 버전을 사용한다. 가장 최근의 Megatron-LM 결과는 32개의 DGX-2 노드 (총 512개의 32GB V100 GPU)를 사용하여 16B 파라미터 모델까지 확장할 수 있는 능력을 보고한다 [3].

ZeRO
MP를 사용하지 않는 실험에서는 ZeRO-100B에 구현된 ZeRO-powered DP를 사용한다. MP를 사용하는 실험에서는 ZeRO-powered DP를 Megatron-LM의 MP와 결합한다.

모델 구성 (Model Configurations)
이 섹션에서 제시된 모델들은 GPT-2 [2]와 유사한 Transformer 기반 모델이다. 우리는 hidden dimension과 layer 수를 변경하여 다양한 파라미터 수를 가진 모델을 얻는다. Table 4는 우리 실험에 사용된 구성 파라미터를 보여주며, 추가 세부 사항은 AE Appendix에 있다.

10.2 Speed and Model Size

ZeRO-100B는 최대 170B 파라미터 모델을 400개의 GPU에서 효율적으로 실행할 수 있으며, 이는 Megatron-LM보다 8배 이상 큰 규모이다. Figure 2는 ZeRO-100B와 MP를 함께 사용했을 때와 Megatron MP만 사용했을 때의 모델 크기별 GPU당 처리량을 보여준다. ZeRO-100B는 **8B에서 100B 파라미터 모델에 대해 평균적으로 15 PetaFlops의 지속적인 처리량(최대 성능의 30% 이상)**을 달성한다. 이에 비해, baseline MP 성능은 모델 크기가 증가함에 따라 빠르게 저하된다. MP는 GPU 간 높은 통신량을 발생시키며, 더 큰 모델을 수용하기 위해 단일 노드를 넘어설 경우 통신 대역폭이 링크당 300GB/sec (NVSwitch)에서 링크당 12.5GB/sec (Infiniband EDR)로 감소하여 상당한 성능 저하를 초래한다. ZeRO-100B는 baseline 대비 최대 10배의 속도 향상을 달성하며, 대규모 모델에서 훨씬 뛰어난 성능을 보인다.

ZeRO-100B의 경우, 100B를 초과하는 모델에서 성능이 약간 감소하는 것은 더 큰 batch size를 실행할 충분한 메모리가 부족하기 때문이다. 다음 섹션에서 논의하겠지만, ZeRO-100B의 초선형(super-linear) 속도 향상으로 인해 GPU 수를 늘리면 성능이 향상될 것으로 예상한다.

10.3 Super-Linear Scalability

ZeRO-100B는 매우 큰 모델 크기에서 초선형(super-linear) 확장성을 보여준다. Figure 3은 60B 파라미터 모델이 64개에서 400개의 GPU로 확장될 때의 확장성 결과를 보여주며, 우리는 이러한 추세가 더 많은 GPU에서도 계속될 것으로 예상한다. **$P_{os+g}$**는 DP(Data Parallel) 정도가 증가함에 따라 GPU당 메모리 소비를 줄여주어, ZeRO-100B가 GPU당 더 큰 batch size를 수용할 수 있게 한다. 이는 결과적으로 연산 강도(arithmetic intensity)를 증가시켜 처리량(throughput)을 향상시킨다.

| | | $Z e R O-\mathrm{DP}$ | $Z e R O-\mathrm{R}$ | | :---: | :---: | :---: | | 1 | $\mathrm{P}_{o s}$ | $\mathrm{C}_{B}+\mathrm{M}_{D}$ | | 2 | $\mathrm{P}_{o s}$ | $\mathrm{C}_{B}+\mathrm{M}_{D}+\mathrm{P}_{a}$ | | 3 | $\mathrm{P}_{o s+g}$ | $\mathrm{C}_{B}+\mathrm{M}_{D}$ | | 4 | $\mathrm{P}_{o s+g}$ | $\mathrm{C}_{B}+\mathrm{M}_{D}+\mathrm{P}_{a}$ | | 5 | $\mathrm{P}_{os+g}$ | $\mathrm{C}_{B}+\mathrm{M}_{D}+\mathrm{P}_{a+c p u}$ |

Table 3: ZeRO 구성

10.4 Democratizing Large Model Training

MP와 PP를 사용하는 것은 많은 데이터 과학자들에게 어려운 일이며, 이는 대규모 모델을 학습시키는 데 잘 알려진 난관이다. ZeRO는 모델 자체에 어떤 변경도 요구하지 않으며, baseline DP만큼 간단하게 사용할 수 있으면서도 모델 크기와 속도를 크게 향상시킨다. Figure 4는 ZeRO-100B가 128개의 GPU에서 MP 없이 최대 13B 파라미터를 가진 모델을 학습할 수 있으며, 평균적으로 GPU당 40 TFlops 이상의 처리량을 달성함을 보여준다. 이에 비해, ZeRO가 없는 경우 DP만으로는 가장 큰 학습 가능한 모델이 1.4B 파라미터에 불과하며, GPU당 처리량은 20 TFlops 미만이다. 더욱이, MP에서 발생하는 통신 오버헤드가 없기 때문에, 이러한 모델들은 NVLINK 또는 NVSwitch와 같은 매우 빠른 노드 내 상호 연결(interconnect) 없이도 학습될 수 있다. 이러한 고속 상호 연결은 MP에서 좋은 효율성을 달성하기 위해 필수적이다.

10.5 Memory and Performance Analysis

우리는 최대 모델 크기, 메모리 소비, 성능에 대한 다양한 최적화 기법의 이점과 영향을 살펴본다. 이러한 최적화 기법들은 Table 3에서 Config 1부터 5 (C1-C5)로 지칭된다.

최대 모델 크기 (Maximum Model Size)
Figure 6은 고정된 batch size와 MP(Model Parallelism) 16에서 다양한 ZeRO 최적화를 활성화했을 때 학습 가능한 가장 큰 모델 크기를 보여준다. C1과 C2를 비교하면, C2에서 $P_a$를 사용하여 activation memory를 16배(MP degree) 줄였기 때문에 모델 크기가 40B에서 60B로 증가한다. 반면, C4에서 140B로 크게 증가하는 것은 $P_{os+g}$를 활성화했기 때문인데, 이는 C2의 $P_{os}$에 비해 모델 상태(model states)에 필요한 메모리 요구량을 절반으로 줄여준다. C5에서 150B로 증가하는 것은 오직 파티션된 activation checkpoint를 CPU 메모리로 오프로딩하여 activation memory를 추가로 줄였기 때문이다.

최대 캐시 메모리 (Max Cached Memory)
Figure 7은 40B 및 100B 파라미터 모델에 대해 각 학습 iteration 동안 PyTorch가 캐시한 최대 메모리를 보여준다. C1에서 C2로 갈수록 캐시 메모리 크기가 예상대로 감소한다. C2와 C3 간의 메모리 소비 차이는 activation memory 대비 모델 상태의 크기에 따라 달라지며, activation memory가 더 클 때는 증가하고, 모델 상태가 더 클 때는 감소할 수 있다. 주목할 점은 40B 모델의 경우 C4에서 C5로 갈 때 캐시 메모리가 감소하지 않지만, 100B 모델의 경우에는 감소한다는 것이다. 이는 단순히 100B 모델의 activation memory가 훨씬 커서 감소 효과가 눈에 띄게 나타나기 때문이다. 이러한 특성으로 인해 $P_{a+cpu}$는 매우 큰 모델에서 더 큰 batch size를 사용할 수 있게 해주는 유용한 도구가 된다. Figure 8에서 170B 모델은 메모리 부족 없이 실행되기 위해 $P_{a+cpu}$가 필요하다.

[^5]| Figure 2 | | | Figures 3, 4] | | | | :--- | :--- | :---: | :--- | :--- | :---: | | | Layers | HD | | Layers | HD | | 1.5 B | 48 | 1600 | $1.16 \mathrm{~B}-2.5 \mathrm{~B}$ | $24,34,54$ | 1920 | | 8 B | 72 | 3072 | 4 B | 64 | 2304 | | $40 \mathrm{~B}-60 \mathrm{~B}$ | 88,132 | 4096 | $6 \mathrm{~B}-8 \mathrm{~B}$ | 52,72 | 3072 | | $80 \mathrm{~B}-170 \mathrm{~B}$ | $100,125,150$ | 8192 | $10 \mathrm{~B}-13 \mathrm{~B}$ | $50,54,58,62$ | 4096 | | $140 \mathrm{~B}-170 \mathrm{~B}$ | 175,212 | 8192 | 60 B | 75 | 8192 |

Table 4: Figures 2, 3, 4에 걸쳐 다양한 모델 크기, 레이어 수, hidden dimension (HD)에 대한 구성.

최대 달성 가능 성능 (Max Achievable Performance)
Figure 8은 다양한 최적화 설정에서 달성 가능한 최상의 성능을 보여준다. 성능 향상이 최적화 간의 메모리 소비 감소와 일치함을 주목하라. 앞서 언급했듯이, 메모리 소비가 낮아지면 더 큰 batch size를 사용할 수 있게 되어 성능이 향상된다. 유일한 예외는 60B 파라미터 모델에서 C4와 C5 사이의 성능 하락이다. C5는 메모리 소비가 더 낮음에도 불구하고, activation을 CPU로 이동시키고 다시 가져오는 오버헤드가 발생한다. 이는 대부분의 경우 성능 저하로 이어진다. 다만, 모델이 너무 커서 C5 없이는 실행할 수 없거나, C5 없이 실행 가능한 batch size가 매우 작은 경우(예: Figure 8의 170B 파라미터 모델)에는 예외이다. 학습 중에는 $P_{a+cpu}$가 유익할 때만 활성화된다.

10.6 Turing-NLG, the SOTA language model with 17B parameters

2020년 5월 12일 기준으로, Turing-NLG는 170억 개 이상의 파라미터를 가진 세계에서 가장 큰 모델이다. 이 모델은 Webtext-103 perplexity 10.21을 달성하며 language model 분야에서 새로운 SOTA를 기록했다. Turing-NLG는 ZeRO-100B를 사용하여 end-to-end로 학습되었으며, Figure 5는 이전 SOTA인 Megatron-LM 8.3B 파라미터 모델과 비교하여 30만 iteration에 걸친 validation perplexity를 보여준다. ZeRO-100B는 이 모델 학습에서 GPU당 41.4 TFlops의 지속적인 처리량을 달성한다.

11 Concluding Remarks

HPC 및 시스템 관점에서, 우리는 ZeRO가 대규모 모델 학습 환경에 혁명적인 변화를 가져올 것이라고 믿는다. 우리의 구현인 ZeRO-100B는 모델 크기를 8배 증가시키고, 처리량(throughput)을 10배 이상 향상시키며, 최신 GPU 클러스터에서 초선형(super-linear) 속도 향상을 달성하고, 세계에서 가장 큰 모델을 학습시켰지만, 이는 아직 빙산의 일각에 불과하다. ZeRO는 전체적으로 모델 크기를 한 자릿수 더 늘릴 잠재력을 가지고 있으며, 미래의 조(trillion) 단위 파라미터 모델 학습을 가능하게 할 것이다.

아마도 우리가 ZeRO에 대해 가장 낙관적으로 생각하는 점은 데이터 과학자들에게 어떠한 장애물도 부과하지 않는다는 것이다. **MP(Model Parallelism)**나 **PP(Pipeline Parallelism)**와 같은 기존 접근 방식과 달리, 모델 재구성(refactoring)이 전혀 필요 없으며, 표준 DP(Data Parallelism)만큼 사용하기 쉽다. 이러한 점은 ZeRO를 대규모 모델 학습에 대한 향후 연구의 최적의 후보로 만든다. 우리는 오픈 소싱과 커뮤니티 피드백을 통해 ZeRO를 DL 커뮤니티에 완전히 공개하여, 대규모 모델 학습의 발전과 민주화를 촉진할 계획이다.

Acknowledgement

Junhua Wang의 귀중한 지원과 조언에 감사드립니다. Minjia Zhang, Elton Zheng, Shaden Smith, Reza Yazdani Aminabadi, Arash Ashari, Niranjan Uma Naresh에게 작업 평가에 대한 훌륭한 피드백과 도움에 감사드립니다. Brandon Norick, Corby Rossett, Gopi Kumar, Jack Zhang, Jing Zhao, Payal Bajaj, Rangan Majumder, Saksham Singhal, Saurabh Tiwary, Xia Song에게 유익한 토론과 제안에 감사드립니다.

Table 5: ZeRO 처리량과 baseline 비교와 관련된 Figure 2의 모델 구성.

Table 6: 초선형 확장성(superlinear scalability)과 관련된 Figure 3의 모델 구성.

Table 7: 다양한 ZeRO 구성에서의 최대 모델 크기와 관련된 Figure 4의 모델 구성.

Table 8: 다양한 ZeRO 구성에서의 할당된 메모리와 관련된 Figure 5의 모델 구성.

Table 9: 다양한 ZeRO 구성에서의 처리량과 관련된 Figure 6의 모델 구성.

Table 10: 데이터 병렬화만 사용했을 때 최대 모델 크기 대 처리량 평가와 관련된 Figure 7의 모델 구성.

이 표들은 본 논문에서 제시된 실험에 사용된 모든 모델 구성과 배치 크기를 포함한다. Figure 2에서, 일부 baseline 실험의 총 GPU 수가 ZeRO의 400개에 비해 384개 또는 256개인 점에 유의하라. 이는 총 GPU 수가 MP의 배수여야 하며, 우리가 총 400개의 GPU에만 접근할 수 있었기 때문이다. 모델 구성 값에는 hidden size가 attention head로 나누어 떨어져야 하고, hidden size가 MP로 나누어 떨어져야 하며, attention head가 MP로 나누어 떨어져야 하는 등 몇 가지 추가적인 제약이 존재한다. baseline의 경우, 모델에 적합한 2의 거듭제곱 중 가장 낮은 GPU 수를 사용했다. 예를 들어, 170B 파라미터 모델의 경우 baseline은 256개였다. 우리가 400개의 GPU만 가지고 있었기 때문에, baseline은 256개의 GPU로만 실행할 수 있었다.

이것이 baseline에 ZeRO보다 유리하게 작용한다는 점을 지적하고 싶다. 왜냐하면 GPU 수가 적을수록 baseline의 통신 처리량이 더 좋기 때문이다. 예를 들어, 170B 파라미터 모델의 경우 baseline의 DP는 1이므로 DP에 대한 통신 비용이 발생하지 않는다. 본 논문에서 제시된 결과는 baseline의 이러한 이점에도 불구하고 얻어진 것이다.

또한, 우리는 집계된 성능이 아닌 GPU당 성능을 비교하고 있으며, 따라서 결과는 여전히 공정한 비교이지만 baseline에 약간의 이점을 제공한다는 점을 지적하고 싶다.