Distributed Training for Large Scale Transformer (1/6) - Overview of Parallelism, Data Parallelism (DP) and ZeRO

< 목차 >

• Motivation
• Parallelism
  • Data Parallelism (DP)
  • Model Parallelism (MP)
  • Pipeline (Model) Parallelism (PP)
  • Tensor (Model) Parallelism (TP)
  • Key Communication Operations For Distributed Training
  • Distributed Data Parallelism (DDP)
  • Summary of Parallelism
• Zero Redundancy Optimizer (ZeRO)
  • Where Did All The Memory Go ?
    • POV: Model States
    • POV: Residual Memory Consumption
  • ZeRO-DP
    • Overview of ZeRO-3
    • Quantative Analysis of ZeRO-DP
  • ZeRO-R
  • Analysis
    • Communication Analysis of ZeRO-DP
    • Communication Analysis of ZeRO-R (Focused on TP)
  • Evaluation
    • Per-device Memory Consumption of Different Optimization Settings
    • Comparison between ZeRO+TP vs TP)
    • Superlinear Scalability
    • Memory and Performance Analysis
• Other Techniques for Large Scale Modeling
  • Gradient Checkpoint (Activation Checkpoint)
  • Gradient Accumulation
• Library/Frameworks for Distributed Training
  • DeepSpeed and Pytorch Fully Sharded Data Parallel (FSDP)
  • Accelerate
• References

Deep Learning 분야의 model size가 점점 더 large scale로 가고 있다. 수십~수백 billion (B)단위의 large model을 학습하기 위해서는 다양한 technique들이 요구 된다. 이번 post에서는 병렬화 (prallelism) 부터 ZeRO, activation checkpointing 등 각종 large scale modeling에 관한 technique들에 대해 알아보려고 한다.

Motivation

분산 학습 (Distributed Training) 이란 하나의 main process만 이용해서 model을 학습하는 것이 아니라 여러개의 mini process를 사용해서 학습하는것을 말한다. 왜 분산 학습을 해야 하는 걸까? Deep learning model을 학습하기 위해 1,000,000개의 sample로 이루어진 dataset을 구축했다고 생각해 보자. 가장 간단한 idea는 모든 data를 보고 각각의 loss를 계산한 다음에 gradient를 구해서 model parameter를 update하는 것이다. 하지만 이는 불가능하기 때문에 우리는 dataset의 일부분만을 sampling하면서 update를 하며 결국 dataset을 다 보게 되는데, 이 때 얼만큼의 data sample을 한번에 쓸 것이냐?를 Batch Size라고 하며 dataset을 한 번 전부 순회 (iteration) 하는 단위를 epoch이라고 한다. 이렇게 학습하는 방식은 Batch learning이 아닌 Mini-batch learning 이라고 하거나 이런 알고리즘을 일반적으로 몇 개씩 무작위로 학습하기 때문에 sampling Stochastic Gradient Descent (SGD) 혹은 Online Gradient Descent라고도 부른다.

그런데 학습 하려는 model size가 너무 크면 어떻게 할까? GPU memory의 대부분을 model이 차지하게 될 것이고, 우리가 일반저그올 사용하는 Adam은 momentum 등을 같이 관리하기 때문에 이러한 optimizer state들까지 고려한다면 GPU memory가 부족한 상황이 발생한다. 당연하게도 batch size를 256, 128, … 8, 4 이렇게 줄여나감으로써 문제를 해결하게 될텐데, 일반적으로 large batch size를 사용하는것이 모델 성능 (pefromance) 면이나, 학습 시간 (training wall clock time) 면에서 모두 이득이기 때문에 다른 방법을 강구해야 한다. (Note that, 항상 large batch가 좋은 성능을 보장하진 않는다)

Parallelism

Data Parallelism (DP)

그 다른 방법은 별 게 없는데, 바로 GPU device를 더 쓰는 것이다. 어떤 workstation에 GPU가 8개 있다고 생각해 보자. 이 때 GPU 1개 당 batch size 4를 쓸 수 있었다고 치면, dataset 1,000,000개를 84로 나누어 31,250개 씩 각 GPU에 할당한다. (어떤 data가 어떤 GPU로 가는지는 일반적으로 매 epoch마다 random shuffle 된다) 만약 dataset이 1,000,135개 였다면 135개는 딱 나눠떨어지지 않기에 (not evenly divisible) 324=128개를 제외한 7개를 버리던지, 32*5=160개가 되도록 기존의 data들을 random하게 25개 정도 upsampling을 해서 나눠준다. 이런식으로 하면 GPU 1개당 4개의 gradient를 계산하므로 batch size를 32로 설정한 것과 다름이 없게 되며 모든 gradient는 main process로 모아서 처리하게 된다.

Fig. DP는 예를 들어 GPU 4개에 full size model이 copy되고 각 GPU가 서로 다른 batch를 처리한다. Source from link

만약 우리가 GPU 8대 workstation을 8개 가지고 있다면 우리는 64개 GPU로 256 batch size learning을 하는 효과를 누리게 되는데, 한 machine당 GPU 8대가 한계인데 machine을 하나만 쓰는 것을 multi-gpu, single-node상황 이라고 하며 machine을 8개 쓴 후자의 경우를 multi-node상황 이라고 한다.

Multi-node의 경우 각 machine의 각 GPU가 계산한 gradient를 main process로 모아서 parameter를 update하고 이를 다른 process들이 다시 가져가야 한다. 이 상황에서 node간 통신 비용 (communication cost)가 발생한다. 초당 data를 얼마나 보내는지를 결정하는 memory bandwidth는 GPU hardware마다 다르고 당연하게도 장비가 비쌀수록 좋아진다 (A100, H100이 비싼 이유). 이렇게 모든 GPU에 각각 동일한 model 을 띄우고 개별적으로 gradient를 계산하고 모으는 방식은 처리하고싶은 data가 만약 128개라면 이를 8개로 나눠 각 GPU (혹은 process; 아예 동의어는 아님)가 16개씩 나눠 처리해서 gradient를 계산하기 때문에 Data Parallelism (DP)라고 부른다.

DP방식은 모든 GPU에 같은 크기의 Model parameter를 전부 Copy하는 방식이기 때문에 Memory Efficiency는 떨어지는 방법이지만, model size가 그렇게 크지 않다면 선택할 수 있는 방법 중 가장 간단하고 시간적으로 큰 낭비가 없는 방법이다.

Model Parallelism (MP)

그런데, 만약 GPU 1대에 Model 하나가 아예 올라가지도 못 하는 상황이 발생하면 어떻게 해야 할까? Large Transformer 시대에 상대적으로 작은 규모인 10 Billion parameters의 model 의 경우 Floating Point 16 (FP16)을 쓸 경우 약 20GB의 memory가 필요한데, 학습을 하려면 mixed precision 등의 technique을 쓰더라도 paramter당 20 bytes로 총 200GB memory를 필요로 하게 된다. 이는 구매 가능한 현존 최고 spec인 A100 GPU의 Memory 가 80GB인 상황을 생각하더라도 1대에도 다 올릴 수 없는 양이다.

이런 경우는 model 의 각 layer를 쪼개서 각 GPU에 올리는 방식으로 해결을 할 수 있는데, 단순히 model을 쪼개서 올리는것 만으로는 부족하다. 즉 추가로 어떤 방식들들이 더 필요한데, 어떤 방식들에는 Zero Redundancy Optimizer (ZeRO) 같은 방법이 있는데 이는 후술하도록 하겠다.

Fig. multi GPU일 경우 model을 쪼개지 않아도 ZeRO를 쓰는 것만으로 4~5배 size의 model을 학습할 수 있다. Source from link

일단 model을 쪼개는 (partitioning) 방식만 생각해보자. 말 그대로 각 GPU에 부분 부분 layer를 올리는 방식으로 Model이 32개 layer를 갖고 있고 8 GPU 1 node 상황에서 각 GPU별로 딱 나눠 떠러지게 나눴다고 (evenly partitioning) 치면 각 GPU당 4개의 layer가 올라가 있는 상황이 된다.

Fig. MP vs DP. DP는 GPU 4개에 full size model이 copy되어있고, MP는 일부가 나눠져 올라가 있다. Source from link

문제는 Neural Network (NN)는 이전 layer의 output이 다음 layer의 output이 되어야 하기 때문에 어떤 GPU들은 대기상태에 들어가는 문제가 발생할 수 밖에 없습니다. 즉 병렬 처리 (Parallel Process)를 할 수 없게 되는것이다.

Fig.

Pipeline (Model) Parallelism (PP)

앞서 얘기한 것 처럼 MP가 NN의 sequential dependency 때문에 idle time이 생기기 때문에 효율적인 Pipelining 필요하게 되었다. Pipleline Parallelism (PP)는 Gpipe라는 Paper에서 최초로 제안되었는데, 이는 아래 figure에서 처럼 mini batch를 micro batch로 더 개념을 세분화해서 pipelining을 하게 된다.

Fig.

이렇게 하면 동시에 처리할 수 있는 job이 늘어나면서 MP의 idle time을 최소화 할 수 있게 되는 것이다.

Fig. Source from link

아래의 animation을 보면 micro batch를 세분화 할수록 Idle time (bubble time)이 줄어드는걸 볼 수 있다.

Fig. Source from link

그럼 micro batch로 bubble time을 얼마나 줄일 수 있는가? Gpipe paper에서는

• batch size: N
• micro batch size: M
• #GPU: K

라고 했을 때 bubble time을 아래와 같이 정량화 했다.

즉 M이 커질수록 이를 줄일 수 있게 되는 것인데, paper의 실험에서는 \(M \geq 4 \times K\) 일 때 bubble은 무시할 수준이 되었다고 한다. (backward, activation checkpoint를 모두 고려했을 때)

PP구현체들을 보면 GPU 8대에 예를 들어 model을 얼마의 단위로 등분할 것인가? (MP degree), 얼마나 data 병렬화를 할 것인가? (DP degree)라는 두 개의 개념이 있고, 이 둘을 곱한 수만큼의 GPU가 총 필요한 GPU 수가 된다.

Fig. Source from link

• PP configuration example)
  • world size: 2*8=16 (총 process 수)
    • node: 2 (machine 수)
    • #GPU per node: 8 (machine 당 GPU 수)
  • Degree
    • DP Degree 4 (total batch를 4등분)
    • MP Degree 2 (2GPU에 model을 등분)
  • total batch size per iteration: 256
    • mini batch size: 32 (2gpu 당 32)
    • micro batch size: 2 (32를 2씩, 즉 16개가 나옴)

한 편, 이런 pipeline 방식도 더 디테일하게 구현을 할 수 있는데, forwarding을 먼저 하고 backward를 순차적으로 하는 것이 일반적이지만,

Fig. Simple Pipeline. Pipeline schedule with forward passes (blue) for all microbatches (represented by numbers) followed by backward passes (green). The gray area represents the pipeline bubble time. Source from link

아래처럼 backward pass를 항상 최우선으로 처리할 수도 있다. 이를 중간중간 연산을 끼워넣는다고 하여 Interleaved Pipeline 혹은 1F1B라고 한다.

Fig. Interleaved Pipeline. Pipeline schedule with 1F1B schedule (initial warm-up followed by a forward plus a backward pass for some microbatch in steady state). Source from link

Tensor (Model) Parallelism (TP)

MP가 MOdel 너무 커서 layer 를 sequential 하게 partitioning 했다면 (inter-layer parallelism) Tensor Parallelism (TP)는 layer 내의 weight matrix tensor를 partitioning한 intra-layer parallelism 방법이다. (TP는 이번 post에서 핵심으로 다룰 Distributed Optimizer, 다른 말로는 Zero Redundancy Optimzier (ZeRO)라고 하는것과 양대 산맥을 이루는 parallelism technique이므로 distributed training series의 다음 post로 따로 빼서 작성했으니 이를 참고하길 바란다. 따라서 지금은 간략하게 개요만 설명하려고 한다)

이런 방식의 병렬화는 Mesh-TensorFlow: Deep Learning for Supercomputers와Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism라는 paper들에서 거의 처음 제안되었는데 핵심은 바로 아래 처럼 GEneral Matrix Multiplication (GEMM)을 병렬화 하자는 것이다.

Fig. Y = XA 를 column-wise, row-wise로 나눠서 해도 결국 같다.

이렇게 GEMM을 device별로 나눠 계산할 수 있는 이유는 Deep Learning에서 모두가 쓰는 Transformer는 matmul과 약간의 non-linear transformation으로 이뤄져 있기 때문에 병렬화 하기가 쉽다는 것이다. 주의할 점은 병렬화를 할 때 통신을 최소한으로 해야 하기 때문에 module을 잘 쪼개야 한다는 것이다.

Fig.

위의 (a) subfigure를 보면

라는 MLP block을 나눌 때

가 아니라

로 나눈 것을 볼 수가 있는데, 이는 \(Y=GELU(X_1 A_1 + X_2 A_2) \neq GELU(X_1 A_1) + GELU (X_2 A_2)\) 이기 때문에 중간에 sync를 한번 맞춰야 하기 때문이다. Subfigure a, b 에 있는 \(f,g\)는 해당 block의 출력이나 입력을 모든 process와 공유해서 같은 의미 (같은 position)의 값은 더한다는 All Reduce를 의미함으로 figure에서 보이는 MLP, Self-Attention은 forward backward 별로 한번씩의 All Reduce 통신만 발생한다.

Fig. Transformer Block 하나 당 4번의 통신이 필요.

즉 DP가 각 batch에서 통신을 딱 한번 (마지막에 gradient 모을 때) 하고 MP는 그보다 조금 더 (layer activation 전달할 때) 필요한 반면 TP는 layer를 forward, backward 할 때 마다 수십회의 통신을 하기 때문에 매우 통신 속도가 빨라야 하며, 최소 한 node 내에서만 통신을 하는것이 권장된다고 한다. (machine간 통신은 훨~씬 느리기 때문)

Fig. DP와 MP(MP)를 같이 쓸 수 있는데, 이런 경우 MP나 TP의 degree는 한 node의 GPU를 넘지 않아야 한다. Source from link

즉 2 node 4 GPU면 TP degree는 4인 것이 8인 것보다 훨씬 좋다는 것이다.

Fig. 얘기한거처럼 DP+복잡한 PP도 가능

Key Communication Operations For Distributed Training

이제 Large Scale Neural Network (NN)을 위한 거의 모든 종류의 Parallelism을 한 번씩 개괄적으로 살펴봤다.

그런데 앞서 몇 번 Gather, Reduce, All Reduce같은 용어들이 쓰였는데, 이들은 모드 collective communication을 위한 opeartion들이며, distributed training을 위해 꼭 필요한 개념들이다. 하지만 이를 처음 접하는 이들을 위해서 한 번 짚고 넘어가려고 한다.

이들은 각각 DP에서 device별로 계산한 batch의 gradient를 모으거나, TP에서 layer별 activation output을 한데 모으고 여러 device로 다시 나눠주는 데 쓰인다. 그리고 보통 pytorch에서 쓰는 이런 communication operator들은 NVIDIA Collective Communications Library (NCCL)라고 하는 NVIDIA에서 정의한 것들을 backend로 쓰게 된다. (당연히 우리가 쓰는 산업용 GPU는 NVIDIA의 것이므로…)

각각은 다음의 figure를 보면 된다.

Fig. Scatter는 [B, T]를 [B/N, T]로 나누는 것 where N is #GPU. Gather는 그 반대

Fig. Reduce는 각 GPU로부터 모든 값을 더해 master process로 가져오는 것. All Reduce는 이를 모든 GPU에 전부 할당하는 것

Fig. Broadcast는 Scatter와 대비되는데, [B, T]가 나눠져 각 GPU로 분배되지 않고 그대로 전파 되며 All-Gather는 더한 것이 아니라 모아서 각 GPU에 할당. (All-Reduce와 대비)

이를 사용해서 DP를 다시 아래처럼 나타낼 수 있겠다.

Fig. Forward and Backward passes with torch.nn.DataParallel. Source from link

한편 DP는 그럼 각 process에서 loss, gradient 까지 계산하고 이를 다 더한걸 모두에게 배포하는 all reduce를 해야만 하는데, 어떤식으로 이를 구현해야 할까? 가장 단순한 방법은 앞서 DP에서처럼 어떤 한 node에서 다 합치고 다시 분배를 하는 것이다. 하지만 torch DDP 에서는 simple all reduce 대신 Baidu의 Ring All Reduce을 썼다고 알려져 있다.

(in Pytorch official docs...)

The model is replicated on all the devices; 
each replica calculates gradients and simultaneously synchronizes 
with the others using the ring all-reduce algorithm.

이제 all reduce가 왜 효율적인지 까지만 알아보고 넘어 가도록 하자.

우리가 4개의 Process (P=4)가 각각 model output으로 length=4 (N=4)인 array를 얻었다고 치자. 우리가 원하는 것은 같은 위치의 각 element를 모두 더한 결과물을 모두가 가지게 되는 것이다.

Fig.

Ring All Reduce는 아래처럼 N크기의 array를 P개의 subarray로 나눈다. (지금은 마침 array length N도 4임)

Fig.

이를 chunk라고 부르고 우리가 하려는 operation은 sum이기 때문에, 동시에 서로 다른 p개 chunk를 다음 process에 넘기면서 순차적으로 정보를 누적시킨다.

Fig.

Fig.

이를 P-1 번 반복하면 마지막으로 방문한 process는 우리가 원하는 최종 결과물의 서로 다른 part (chunk)를 완성하게 된다.

Fig.

마지막으로 각 Part를 서로에게 공유해주면 완성이다.

Fig. All Reduce = (Reduce) (Scatter) (All Gather)의 합작이다. Source from link

앞서 간단하게 master node가 모두 모아서 더하고 다시 분배하는 방식은 (P-1)N 만큼의 data를 master가 받고 sum한 후 다시 이를 (P-1)N 만큼 처리해야 하기에 총 통신량이 P에 비례하나, Ring All Reduce는 각 process가 한번에 드는 통신 비용은 N/P 사이즈의 Chunk를 (P-1) 번 처리하며 마지막에 한번 더 같은 사이즈를 (P-1)번 교환하기에 2N(P-1)/P 가 되므로 실제로는 P에 independent해지면서 process가 늘어날수록 통신량이 선형적으로 증가하는 것을 막고 각 process에 계산과 통신을 균등하게 분배함으로 병목을 제거할 수 있었다고 한다.

Distributed training framework을 이해하고, 실제로 이를 구현하거나 사용하면서 문제가 생겼을 때 debugging하기 위해서는 반드시 이를 잘 이해하고 있어야 한다.

Distributed Data Parallelism (DDP)

마지막으로 ZeRO에 대해 알아보기 전에 잠시 다시 DP로 돌아가 Distributed DP (DDP)에 대해 짚고 넘어가려고 한다. Pytorch 같은 Open-source Framework을 보면 당연히 distributed training 을 위한 module들을 제공하는데, nn.DataParallel module class의 문서를 보면 아래와 같은 Warning 문구가 있다.

Fig. DP대신 DDP를 쓰라는 torch docs. Source from link

왜 DDP를 권하는 걸까?

사실 pytorch로 distributed training을 하는 그 누구도 naive DP를 쓰지 않는다. DP가 single node multi gpu 인 상황에서만 쓸 수 있는 반면 DDP는 Dnode가 여러개인 상황에서 distributed training을 가능하게 해줍니다. 하지만 문서에는 single node일 때도 DDP를 쓰는 것을 고려하라고 한다. 문서를 좀 더 찾아보면 DP의 구현상 이슈로 문제점이 몇 가지 있다고 얘기하는데,.

Fig. DP vs DDP 2. Source from link

물론 DP도 장점은 있다. 코드 1줄이면 DP를 쓸 수 있다는 것이다.

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
output = net(input_var)  # input_var can be on any device, including CPU

하지만 그외에 아래와 같은 문제점이 생기는 것이 더 크다고 할 수 있다.

• 매 iteration 마다 data, model copy 가 일어남.
• single process (master) 와 multi-thread로 구현되어있는데 python GIL때매 병렬 forward, backward가 느릴 수 있음.

이를 도식화 하면 아래와 같은데,

Fig. DP. Source from link

Master GPU 하나에서 data loading, loss계산, model update 등을 모두 모아서 처리하느라 빈번하게 scatter와 gather가 일어나는 걸 볼 수 있다. 이럴 경우 master GPU의 memory가 가중되는 문제까지 생겨서 batch size를 크게 못 가져갈 수도 있을 것 같다. DDP는 이런 비효율을 없앤 것이라 할 수 있는데, 하나의 process에 multi-thread가 아닌 여러 process를 띄우고 처음부터 각각 dataloader, model 을 띄운 뒤에 gradient 계산할 때만 all reduce를 하는 방식을 택한다.

Fig. DDP. Source from link

그래서 우리는 DDP를 써야 하는것이다.

Summary of Parallelism

지금까지의 Parallelism을 요약하자면 다음과 같다.

• (D)DP
  • 각 process가 dataloader, model을 갖고 있으며 gradient계산 후 all reduce만 하면 되는 편리함.
  • process 별 통신은 적지만 model parameter가 각 procee에 전부 올라가야 하므로 memory efficiency 하지는 않음.
• MP
  • model이 너무커서 layer를 각 GPU process에 나눠서 올림.
  • 예를 들어 layer0 의 activation을 layer1이 받아야 하므로 기다리는 동안 bubble time 생김
  • layer간 통신 필요
• PP
  • MP의 bubble을 최소화 하기 위해 micro batch 단위로 scheduling을 함.
  • classic MP보다 통신 더 필요
• TP
  • interlayer (MP의 방식)가 아닌 intralayer 병렬화
  • 통신 많이 필요
  • 같은 node (machine) 내에서 하길 권
• MP + PP + TP
  • 분할, scheduling을 잘 해야함.

이제 다음 stage로 넘어가서 ZeRO에 대해 알아보자.

Zero Redundancy Optimizer (ZeRO)

앞서 parallelism 에 대해 기본적으로 알아봤는데 사실 이것만으로는 large model을 학습하기엔 부족하다. 3B, 7B … 이상의 large model을 distributed training 하려면 Zero Redundancy Optimizer(ZeRO)를 반드시 알아야 한다.

Fig. Hardware (HW)의 발전속도는 Model Size의 발전 속도를 따라가지 못한다. 결국 ZeRO같은 algorithm을 사용해야만 HW의 한계를 이겨내고 학습할 수 있다. Source from here

ZeRO는 Microsoft에서 개발한 method로 모델의 크기가 수십 billion 을 넘어가는 large model을 상대적으로 적은 GPU를 가지고 학습하기 위해 제안되었다. 앞서 TP를 설명하며 언급드렸던 Megatron-LM이 수백 billion에 달하는 model을 학습하는 대표적인 algorithm이었으나 이를 학습하기 위해 천문학적인 GPU가 들어갔다고 알려져 있다. 여기에 ZeRO를 추가하거나 ZeRO만으로도 필요한 GPU수를 더 최적화 할 수 있다고 하는데, 본 post의 highlight인 ZeRO에 대해서 이제 본격적으로 알아보도록 하겠다.

Where Did All The Memory Go ?

ZeRO paper를 보면 다음과 같은 구절이 있다.

Let’s take a step back to examine the memory consumption of the current training system. 
For example, a 1.5B parameter GPT-2 model requires 3GB of memory 
for its weights (or parameters) in 16-bit precision, 
yet, it cannot be trained on a single GPU with 32GB memory using Tensorflow or PyTorch. 

One may wonder where all the memory goes. 

논문에서는 GPT-2를 예시로 드는데, GPT-2의 model size는 large size가 1.5 billion (B)정도 된다. 저자들은 model parameter를 부동 소수점 16자리로만 표현하는 Floating Point 16 (FP16)을 사용해서 학습할 때를 가정했다. FP16을 쓰면 weight element 하나당 2 bytes를 model을 다 upload하는 데 3GB 정도의 GPU memory를 쓰게 되는데,

그럼에도 불구하고 32GB (v100)의 GPU 1대로 학습이 불가하다고 한다. 그 이유는 FP16으로 모델 학습을 하는 것이 실제로는 Mixed Precision Training이라는 algorithm을 쓰는 것이고 이 mixed precision이 엄청난 Redundancy를 가지고 있기 때문이라고 저자들은 주장한다.

POV: Model States

어느 부분이 redundancy일까?

우선 Mixed Precision Training method 에 대해 짧게 생각해 보자. Deep Neural Network (DNN) model을 Single-Precision (FP32) 대신 Half-Precision (FP16)로 연산을 하면 속도가 수 배로 빨라질 수 있지만, 표현할 수 있는 수의 범위가 더 좁아지게 (narrower) 됨으로써 model accuracy 에 영향을 주게 된다. 그래서 fp16이나 custom fp16인 brain float 16 (bf16)등을 사용하여 model을 학습하고자 하는 욕구가 생겼는데, 이는 수를 표현하는 데 있어 정밀도 (precision)가 낮고 표현 가능한 범위가 좁은 문제를 가지고 있기 때문에 그냥 model weight을 전부 fp16, bf16으로 바꿔 학습하면 학습이 잘 되질 않는다. (bf16은 별 문제가 없을 수 있음)

Fig. FP32 vs FP16 vs BF16. fp16은 fp32나 bf16과 다르게 표현할 수 있는 실수의 범위가 더 적다. Source from link

그래서 mixed precision training paper에서 내놓은 idea가 바로 FP16과 FP32를 섞어 쓰자는 것이었다. 이는 크게 3가지 technique으로 구성되어 있다.

• (기본) 대부분의 forward/backward 연산은 fp16으로 한다.
• fp32의 weight을 copy한 master copy를 따로 관리하며 매 training step마다 fp16 weight으로 forwarding해서 구한 gradient는 fp32 weight에 더해 update하고 fp16은 이를 copy해서 쓴다. (매 번 copy가 일어남)
• gradient value가 0이될 경우의 수를 줄이기 위해 loss를 scaling해서 gradient를 계산한 뒤, 나중에 다시 unscaling한다.
• 특정 operation들은 fp16이 아닌 fp32로 한다.

(더 관심있는 이들은 Lowered Precision Training에 대한 post를 따로 작성했으니 참고하면 좋을 것 같다.)

이 중 ZeRO의 motivation은 바로 첫 번째 FP32 copy를 두고 이를 통해 update한다라는 technique 이다. 이는 backward pass를 통해 계산 된 weight gradients가 그보다 더 작은 값인 learning rate와 곱해지면서 더 값이 작아지고, 결국 상대적으로 큰 수인 weight과 update가 더해질 때 아무런 변화가 생기지 않는 것을 방지한다.

Fig.

저자들은 가장 보편적인 Adam optimizer을 쓸 경우 model parameter update를 위해서 아래 세가지 optimizer state를 항상 저장하고 있어야 한다고 했는데,

• Time averaged momentum
• Variance of the gradients
• Gradients and Weights themselves

얘기한 것 처럼 mixed precision을 쓸 경우 model parameter와 각 layer의 activation들은 모두 fp16으로 저장되고, forward pass, backward pass 모두 fp16으로 수행되기 때문에 v100이상의 장비에서 엄청난 throughput (초당 sample 처리량)을 자랑하나 효과적인 parameter update를 위해 결국 backward pass의 마지막에는 optimizer가 model parameter의 fp32 version copy를 가지고 있어야 하며, optimizer state 또한 fp32로 가지고 있어야 하는 문제가 발생한다.

Fig. state_dict of Adam. 실제로 torch optimizer는 optimizer state를 관리하는 dictionary를 가지고 있는데, momentum 등은 CUDA tensor형태이다. 즉 model parameter와 같은 크기의 momentum이 또 memory를 잡아먹는 것. Source from link

이를 ZeRO에서는 정략적으로 다음과 같이 나타냈다.

• the number of model parameter: \(\Psi\)
  • for default
    • for model param: \(2 * \Psi\) (fp16)
    • for gradients: \(2 * \Psi\) (fp16)
  • for optimizer
    • for fp32 copy of model param: \(4*\Psi\) (fp32)
    • momentum: \(4*\Psi\) (fp32)
    • variance: \(4*\Psi\) (fp32)

Optimizer 를 위해 더 필요한 메모리, 즉 memory multiplier 를 \(K\)라 하면 Adam은 언제나 model parameter의 수 \(\Psi\)에 대해 \(2\Psi + 2\Psi + K\Psi = 16\Psi\) 씩이나 필요한 것이다.

Fig.

즉 1.5B GPT-2가 실제로는 학습에 필요한 Memory가 3GB가 아니라 기본적으로 8배나 큰 24GB나 필요했던 것이다.

Fig. 2 device로 DP를 한다면 이렇게 heavy한 요소들을 모든 device가 들고 있어야 한다. 파란색은 model parameter, \(2\psi\)를 의미하고 노란색은 gradient, \(2\psi\)를 그리고 나머지는 optimizer state \(K \psi\)를 의미한다.

그러니까 mixed precision은 사실 학습 속도를 높혀주지만 memory save는 장담이 안되는 것이다.

Fig. link에는 efficient training을 위한 tool들이 소개되어 있는데, mixed precision은 memory saving이 된다고는 볼 수는 없다.

물론 model parameter와 더불어 layer들의 output activation 값들이 다 fp16일 것이기 때문에 model size가 작을 때는 memory가 save가 돼서 batch를 좀 늘려볼 수 있지만 model size가 커지면 문제가 생길 수 있다.

Fig. “누가 memory save 해준대?” 언제나 Huggingface의 Engineer Sgugger는 매운 답변을 달아준다. Source from link

POV: Residual Memory Consumption

앞서 model parameter, optimizer state 관점에서의 redundancy에 대해 알아봤다. 이번에는 residual memory관점에서 redundancy를 분석하고 이에 대한 해결책들에 대해 얘기해보려고 한다. (이들은 보이지 않는 부분에 대한 최적화로 사소해 보일 수 있으나 결코 사소한 것들이 아니다) paper에서는 아래 세 가지를 주로 얘기한다.

• Activations
• Temporary buffers
• Memory Fragmentation

먼저 첫 번째 Activation이다. Model이 GPT-2 large로 1.5B크기를 가지고 있고 주어진 sequence input의 batch size가 32, 길이가 1K일 때 필요한 memory량은 얼마일까? 60GB나 된다고 한다. 여기서 activation이란 matmul이나 addition이 발생하는 모든 부분에서의 연산 결과를 의미한다. Transformer의 경우 activation에 필요한 memory는 다음에 비례하는데,.

(the number of transformer layer) * (hidden dim) * (sequence length) * (batch size) 

GPT-2의 경우 coefficient 12배가 곱해지면 60GB만큼의 memory가 필요하게 된다. 그러니까 우리는 GPT-2 XLarge를 (당시 최신 spec이었던) V100에서 조차 학습이 불가했던 것이다.

Fig. Source from Low-Memory Neural Network Training: A Technical Report

Model size나 batch size가 커지면 당연히 부담은 더 커질텐데, 이를 해결하기 위해서 MP를 할 경우 더 문제가 된다고 한다. 왜냐하면 MP중 하나인 TP (vertical MP)를 한다고 할 경우 (보통 communication cost때문에 vectical MP인 TP를 씀), layer output을 각 gpu device로 서로 copy해줘야 다음 layer의 forward computation을 할 수 있기 때문에 문제가 된다고 한다.

그래서 이를 효과적으로 해결해주는 방법이 제안되었는데, 바로 activation (gradient) checkpoint이다. ZeRO paper에서는 분석한 결과 (아마도 GPT-2에 대해) activation checkpointing을 쓰면 memory 소비를 8GB까지 줄일 수 있었지만 속도가 33%정도 느려졌다고 한다 (re-computation overhead 때문에). 8GB까지로 줄어든 이유는 activation checkpointing이 보통 원래 memory cunsumption을 \(x\)라 할 때, \(\sqrt{x}\)까지 줄여주기 때문이다.

Fig. Memory used while training a ResNet model with large batch size, using the regular tf.gradients function and using memory-optimized gradient implementation

Activation checkpointing에 대해서 간단히만 얘기하자면, key idea는 layer들의 activation output들을 다 저장하지 않고 일부만 띄엄 띄엄 저장한 다음에 forward가 다 끝나면 loss를 계산하고 backpropagation을 할 때 빈 activation output들을 다시 계산 (re-computation)하자는 것이다. 즉 Training time을 20% 추가로 쓰는 대신 10배 더 큰 NN 모델을 gpu memory에 넣을 수 있는 기법으로, parallelism이나 CPU-offloading등과 함께 orthogonal 하게 쓸 수 있어 Large Model을 학습하는 데 있어 중요한 기술 중 하나이다.

아래의 animation을 보면 vanilla backpropagation을 할 경우, forward activation을 다 저장해 둔 다음에 끝에가서 loss를 계산하고 차례 차례 memory를 release하는 걸 확인할 수 있다.

Fig. Vanilla Backprop

만약 GPU에 memory를 model 올리는 데 이미 거의 다 썼다면 아래와 같이 정말 비효율적으로 다시 계산하는 방법을 택할 수 도 있을것이다.

Fig. Memory Poor Backprop

Activation Checkpointing은 이 둘의 절충안으로 checkpoint지점들을 두고 그 지점부터만 현재 node까지 빈 activation을 다시 계산하면서 backward를 하는 것이다.

Fig. Checkpoint of Activation Checkpointing method

Fig. Checkpointed Backprop

하지만 ZeRO가 지적하는 점은 model size가 커지면 activation checkpointing으로도 감당이 안된다고 한다. 실제로 ZeRO는 activation checkpointing은 이미 적용을 했음에도 불구하고 학습이 불가능한 case에 대해 얘기한다. 가령 100B정도 되는 model size가 되면 32 batch size를 쓸 때 checkpointing을 해도 60GB에 달하는 memory가 필요하기 때문에 이미 model을 올리는 순간부터 문제가 발생한다. 이를 해결하는 것이 바로 ZeRO-DP이다.

그런데 사실 paper의 contribution은 ZeRO-DP뿐만 아니라 ZeRO-R이라는 것도 있다. 이것이 해결하는 문제는 일시적인 버퍼 (Temporary Buffers)이나 메모리 조각화 (Memory Fragmentation)같은 문제이다.

먼저 Temporary buffers는 학습시 발생하는 중간 결과 (intermediate results)를 저장하는데 쓰이는 buffer들인데, 가령 (바로 아래에서 배우게 될) ZeRO-DP를 할 경우 각 device별로 흩어져있는 gradient를 한데 모으거나 gradient norm을 계산하는 행위를 효율적으로 하려면 보통 여러 gradient를 하나로 모아서 처리 (통신) 한다고 한다. 그 이유는 all-reduce같은 연산이 large message size에 대해서 더 효율적 (throughput이 잘나옴) 이기 때문이라고 한다. 이런 fused buffer의 memory overhead는 model size에 비례하게 되므로 너무 작아질수도 커질수도 있어서 문제가 발생할 수 있다.

그 다음으로 Memory Fragmentation는 짧은 시간 동안만 존재하는 (short-lived) tensor들과 긴 시간 존재하는 (long-lived) tensor들이 번갈아 나타나는 (interleaving) 현상 때문에 발생하는데, activation checkpointing 과 gradient를 계산할 때 발생한다고 한다. 생각해보면 forward시에 어떤 activation들은 checkpoint지점이기 때문에 backward가 계산될 때까지 남고, checkpoint지점이 아닌 activation tensor들은 지워지며 backward에서도 parameter gradient는 마지막에 parameter update를 할 때까지 살아남지만 gradient를 계산하는데 쓰이는 나머지들은 바로 필요없어지죠. 여기서 문제가 발생하는데 예를 들어 activation들이 일부가 지워졌다고 치자. 이렇게 지워진 공간들을 다 합치면 충분히 memory가 많이 saving이 될테지만, 이 freed memory의 위치들이 연속적이지 않아서 (contiguous하지 않아서) 나중에 어떤 큰 memory를 할당하려고 하면 OOM이 발생할 수 있는것이다. 둘째로 contiguous piece가 있다 하더라도 memory allocator가 이를 찾는 데 시간을 많이 써서 비효율이 발생할 수 있게 된다. 그래서 우리는 이 Memory Fragmentation을 줄여야 하는것이다.

ZeRO-DP

그렇다면 어떻게 ZeRO가 memory 관리를 해주길래 V100 GPU한장에서 13B model까지 학습이 가능하다고 광고하는 것일까? 한 번 알아보도록 하자.

Overview of ZeRO-3

먼저 model state 관점에서의 최적화만 알아보자. ZeRO-DP는 이름에서도 알수 있듯 DP를 하되 Redundancy를 제거한 것이다. Paper에서는 이를 3가지 단계로 구분했다.

Fig.

ZeRO-1에서 ZeRO-3로 점점 더 많이 쪼갠 것이라고 생각하시면 되는데, 먼저 직관적으로 이해하기 위해 DeepSpeed team의 video를 animation으로 만들었다. 이하는 ZeRO-3가 어떻게 학습과정에서 쓰이는지를 시각화 한 것이다. 왜 ZeRO-1, 2가 아닌 ZeRO-3를 먼저 보는지에 대해서는 ZeRO-3 영상밖에 없기도 하지만 stage 3가 1, 2를 포함하기 때문에 그렇다.

Fig. 우리가 학습할 model이 16층 짜리 transformer라고 가정.

Fig. 먼저 거대한 dataset이 있다.

Fig. 이를 4개의 GPU device에 나눠서 학습하려고 한다. Data Prallel (DP)를 하는 것이다.

Fig. batch size를 micro batch size 로 4등분 했다. 이 animation에서는 모든 memory optimization technique을 다 사용한 Stage-3를 설명할 것.

Fig.

Fig. 파란색 작은 block은 model parameter를, 노란색은 gradient를, 그리고 초록색은 optimizer state를 의미한다. 그리고 맨 상단에 있는 가장 큰 block은 transformer layer의 진행 과정을 시각화함.

Fig. 파란색 block은 fp16 model parameter

Fig. 노란색 block은 fp16 gradient

Fig. 초록색 block은 optimizer state를 나타내는데, 이는 fp32 parameter, fp32 variance, fp32 momentum, fp32 gradient를 의미한다. 이부분이 딱 보기에도 가장 많은 memory를 차지한다.

Fig. 추가로, 최상단에 있는 파란색 block은 각 layer들의 activation들을 말한다. batch size, model size가 커질수록 잡아먹는 용량이 더 커진다.

Fig. 이제 모든준비가 끝났습니다. 각 process가 다른 Data를 처리하는 DP이긴 한데, 각 process가 서로 다른 model state들을 나눠서 들고있는 형태가 ZeRO-DP의 기본적인 형태이다.

Fig. (Animation) model forward를 해야하는데, GPU1,2,3은 초반부 layer에 대한 model parameter가 없는 상태이기 때문에 GPU0 에서 나머지 device로 나눠준다 (broadcast).

Fig. 이제 모든 device가 첫 4개 Layer의 model parameter를 들고 있다.

Fig. 이제 forward pass 연산을 한다. 최상단의 activation을 보면 띄엄 띄엄 값이 저장되고 있는 걸 알 수 있는데, 이를 activation checkpointing이라고 한다. (후술)

Fig. (Animation) forward pass 과정을 시각화한 것

Fig. (Animation) 이제 GPU1,2,3 들은 activation 값들을 얻었으니 필요가 없어진 layer parameter를 다 지운다.

Fig. 그 다음 또 (4~7) 4개 layer를 broadcast를 해준다.

Fig. (Animation) 또 forward를 진행해줍니다. (계속 activation checkpointing)

Fig. (Animation) 또 필요없어진 parameter들은 지운다.

Fig. 이를 반복하면서 마지막까지 전부 forward를 한다.

Fig. 최종적으로 모든 process에서 각각 Loss를 구한다.

Fig. 이제 parameter를 update하기 위해서 각 activation, parameter를 활용해 gradient를 구해야 한다. 마지막 layer의 경우 forward를 위해서 broadcast한 parameter를 지우지 않고 그대로 가지고 있는걸 볼 수 있는데, 이는 어차피 gradient를 구하기 위해서 다시 broadcast를 해야 하기 때문이다.

Fig. backward를 해서 fp16 gradient를 구한다.

Fig. (Animation) 그 다음 reduce를 통해 모든 gradient를 한 곳으로 모으는데, 마지막 4개 layer는 GPU3이 관리하므로 GPU3으로 모아준 뒤, 나머지 Device들에서는 gradient, parameter들을 모두 지운다.

Fig. (Animation) (나머지 Device들에서는 gradient, parameter들을 모두 지움)

Fig. 그 다음 layer들도 backward, reduce, 지우기 반복.

Fig. 이렇게 모든 device가 자신이 관리하는 layer들의 fp16 gradient를 가지게 되었다.

Fig. 이제 갖고있는 fp16 gradient들로 parameter를 각각 update 해야한다. 각 device에서 optimizer state를 병럴적으로 돌린다.

Fig. (Animation) optimizer run 중

Fig. mixed precision을 하는 중이므로 fp32 model parameter가 momentum, variance, graident를 사용해서 update 된다.

Fig. 이제 이를 다시 fp16 으로 변환해주고,

Fig. 각 device들에게 broadcasting을 해주면 한 iteration이 끝이 납니다.

아마 느끼셨겠지만 device 수가 늘어날수록 ZeRO-3의 경우 기하급수적으로 이득을 볼 수 있게 된다.

Quantative Analysis of ZeRO-DP

이제 정량적인 수치에 대해서 각 stage를 적용했을 때 얼마나 save 할 수 있는지를 알아보도록 하겠다.

Fig.

앞서 얘기한 것 처럼 model size가 \(\psi\)이고 fp16 mixed precision을 썼으며 adam optimizer를 쓸 때를 가정해 vanilla DDP를 하면 \((2(\text{fp16 param})+2(\text{gradient})+K(\text{optimizer states}))\cdot \psi\) 가 든다. Model size가 현존하는 open source LLM중 가장 큰 llama2에서 가장 큰 model인 70B 이라고 가정하고 얘기해 보자.

정말 말이 안되는 수치라고 할 수 있다. 현존하는 GPU중 가장 좋은 A100-80GB 의 hardware도 이를 감당할 수는 없다. 반면 70B model을 각 stage를 사용했을 때 총 필요한 gpu memory는 다음과 같다.

• stage 1 (\(P_{os}\)): optimizer state (fp32 param, momentum 등)만 각 device에 나눠서 올리는 것
  • 총 Memory = \(2\psi + 2\psi + \frac{K}{N_d}\psi = \color{red}{293.125 \text{ GB}}, \text{ where } K=12 \text{ and } N_d=64\)
  • time: DP와 동일한 communication cost, memory: 4배 save
• stage 2 (\(P_{os}+P_{g}\)): stage 1에 추가로 gradient 까지 sharding
  • 총 Memory = \(2\psi + \frac{(2+K)}{N_d}\psi = \color{red}{155.3125 \text{ GB}}, \text{ where } K=12 \text{ and } N_d=64\)
  • time: DP와 동일한 communication cost, memory: 8배 save
• stage 3 (\(P_{os}+P_{g}+P_{p}\)): stage 2에 추가로 model parameter 까지 sharding
  • 총 Memory = \(\frac{(2+2+K)}{N_d}\psi = \color{red}{17.5 \text{ GB}}, \text{ where } K=12 \text{ and } N_d=64\)
  • time: parameter를 통신해야 하므로 DDP보다 50%증가, memory: 64배 save

수식을 보면 알겠지만 각 stage들은 device의 수가 늘어날수록 stage 1은 \(4\psi\)로 수렴하고, stage 2는 \(3\psi\), 마지막으로 stage 3는 \(\psi\)까지 개선이 될 수 있다는 걸 알 수 있다. (계산에 사용한 code는 아래와 같다)

bytes = {
    "fp32": 4,
    "fp16": 2,
    "bf16": 2,
}
def get_memory(model_size, dp_degree, precision="fp16", K=12, stage=3):
    assert precision in ["fp32", "fp16", "bf16"]
    assert stage in [1, 2, 3]

    if stage == 1:
        return (bytes[precision] + bytes[precision] + (K)/dp_degree) * model_size
    elif stage == 2:
        return (bytes[precision] + (bytes[precision] + K)/dp_degree) * model_size
    elif stage == 3:
        return ((bytes[precision] + bytes[precision] + K)/dp_degree) * model_size

ZeRO-R

ZeRO-DP에 대해서 알아봤으니 이번에는 ZeRO-R에 대해서 알아보자. ZeRO-R의 R은 앞서 간단히 얘기한 것 처럼 Residual State Memory를 의미한다.

• Activations
• Temporary buffers
• Memory Fragmentation

Paper에서 언급하는 Residual State Memory를 해결하기 위한 전략은 크게 세 가지가 있다.

• Partitioned Activation Checkpointing: \(P_a\)
  • CPU offloading: \(P_{a+cpu}\)
• Constant Size Buffers: \(C_B\)
• Memory Defragmentation: \(M_D\)

먼저 activation 을 checkpointing하는 것은 기본이다. 그런데 이 checkpointing 이 된걸 device별로 나누기까지 하고, 이를 Activation Partitioning, \(P_a\)이라 부른다. (사실 이는 TP를 위한 technique이라 할 수 있는데 지금은 곧 따로 언급하겠다)

그래도 부족하면 CPU에 Offloading까지 했다가 필요할 때 (Backward 계산 시) 꺼내 쓴다, ,\(P_{a+cpu}\). CPU Offloading은 예를 들어 model parameter, optimizer state, activation 등을 CPU로 내렸다가 필요할 때 꺼내서 사용하는 기술로, model state를 CPU에 offloading 하는것은 학습 시간의 최대 50%를 offloading에 쓴다거나 하는 비효율이 발생하기 때문에 ZeRO-R 에서는 activation 만 offloading하는 방법을 택했다고 한다. 이렇게 CPU offloading을 막 해도 된는 (?) 이유는 matmul연산이 매우 크기 때문에 그동안 data movement를 함으로써 bandwidth가 낮더라도 movement cost를 숨기는 것이 가능하기 때문이라고 한다. (NN training시 연산 결과물들을 CPU나 다른 Disk에 cashing하는 것은 ZeRO가 나오기 전에도 이미 있던 technique이었겠으나 이를 민주화 해서 opensource package로 제공까지 한 것은 ZeRO가 처음일 것이다)

그 다음은 all-reduce 같은 연산을 할 시 Constant size Buffer를 쓰는 것이다. 기본적으로 large all-reduce operation이 small인 것 보다 훨씬 높은 bandwidth를 갖는다고 한다. 그래서 NVIDIA Apex나 Megatron 같은 library들은 all-reduce를 하기 전에 parameter 들을 모두 하나의 buffer에 넣어서 쓴다고 되어있는데, 이러면 parameter size에 비례해서 3B 일때 12GB가 필요하는 등 (32-bit)의 문제가 생겨 너무 작지도 않지만 너무 커지지도 않게 하는 고정된 buffer를 쓰는 전략을 취한다는 것이다.

마지막으로 Memory Defragmentation는 forward시 activation을 지우는 행위 등에 따라 남는 memory들이 조각화 되지 않도록 잘 관리해주는 것으로, 미리 연속적인 (contiguous) memory chunk들을 할당해 두고 checkpointed activation이나 gradient 등을 미리 할당된 memory로 옮김으로써 문제를 해결할 수 있다.

Analysis

이번 section에서는 ZeRO-DP와 ZeRO-R을 baseline과 비교해 communcation volumne이 얼마나 증가했는지? 그리고 이것이 training speed (or throughput)에는 영향이 없는지? 등에 대해서 알아볼 것이다. ZeRO-DP같은 것이 optimizer state, param, gradient등을 device별로 partitioning함으로써 memory reduction을 이뤄낸 만큼 communication cost를 더 지불하기 때문에 당연히 할 수 궁금할 부분일 것이다.

이 부분은 사실 매우 중요하다고 할 수 있는데, 우리가 DeepSpeed와 같은 distributed training framework를 따로 만들거나 사용할 때 현재 우리의 model size와 distributed training plan에서 어느 부분이 bottleneck인지 이론적으로 알 수 있어야 하기 때문이다. 물론 직접 해보고 profiling해보는 과정도 필요하겠으나 어느정도 알고있어야 debugging도 쉬울 것이니 말이다.

Communication Analysis of ZeRO-DP

먼저 ZeRO-DP와 baseline-DP을 비교해보자. 분석 결과 이에 대한 답은 다음과 같다고 한다.

• ZeRO-1, 2에 대해서는 additional comm cost가 없이 memory를 8x 가까이 줄일 수 있다.
• ZeRO-3을 쓸 경우 (param partitioning 추가) 1.5배 정도 느려지지만 device수에 비례해서 훨씬 memory를 줄일 수 있다.

왜 그런지에 대해서 분석하기 위해 먼저 vanilla DP에 대해 생각해보자. Vanilla DP은 process별로 mini-batch forwarding을 하고 error backpropagation을 통해 gradient를 구할 때 까지는 communication이 없다. 그리고 backprop이 끝나고 나서야 gradient를 all-reduce한다. 보통 larger model에 대해서 all-reduce는 communication bandwidth bound이기 때문에 paper에서는 이에만 집중했다고 한다 (즉 얼마나 data를 보내는지).

Fig. All-Reduce.

All Reduce는 distributed process별 output들을 모아서 다 더하고 process별로 똑같이 나눠주는 것인데, 이는 output들을 다 모아서 더한 뒤 하나의 process에만 저장하는 reduce의 확장판이다.

Fig. Reduce

그런데 사실 all-reduce는 reduce-scatter와 all-ather를 연속적으로 수행하는 것이라고 할 수 있다.

Fig. Reduce-Scatter

Fig. All-Gather

그리고 이 때 \(\Psi\) elemetns의 data를 위해 각각의 communication operation이 필요로하는 data movement 양은 \(\Psi\)로, all-reduce는 \(2\Psi\)만큼의 data movement를 필요로 한다고 한다. (pipelined approach를 사용할 경우라고 하는듯, 아마 Ring all-reduce같은 것들인 것 같은데, 원래는 \(\Psi \times \frac{N_d-1}{N_d}\)인데 \(N_d \right \infty\)이면 \(2\Psi\))

이제 ZeRO-2, \(P_{os+g}\)에 대해 생각해보자. 이는 optimizer state, gradient가 partition되어 있기 때문에 all-reduce를 할 필요는 없다. 가령 12th layer에 대해 backprop을 통해 gradient를 계산했다고 치자. 12th layer를 담당하는 optimizer state는 cuda:3이라고 치면, 다른 device들은 reduced gradient를 가지고 있을 필요가 없으므로 cuda:3만 reduce를 하면 된다.

Fig. gradient는 cuda:3만 가지고 있으면 된다.

그리고 최종적으로 gradient를 통해 예를 들어 Adam의 EMA update rule등을 이용해 각 device가 맡은 부분의 parameter를 update하고 모든 process에 대해서 all-gather를 해주면 되므로, ZeRO-2의 경우도 \(2\Psi\)의 cost만이 들 것이다.

마지막으로 ZeRO-3, \(P_{os+g+p}\)는 forward pass에서 partitioned param을 process간 broadcast해줘야 하는 issue가 있다. 하지만 저자들은 이 overhead를 pipelining을 통해 hiding할 수 있다고 얘기한다.

Fig. Broadcast

Forward시 parameter를 broadcast한 뒤 activation을 만들고 parameter의 owner process가 아닌 경우는 다시 이를 지우게 된다. 하지만 backward시 checkpointed activation을 다시 recomputation하는 과정이 또 있기 때문에 한번 더 broadcast가 필요하다. 그리고 만들어진 process별 gradient를 owner process에게 전달하면 되는데 (reduce), 앞선 vanilla DP나 ZeRO-2와 다르게 ZeRO-3는 gradient로 param update를 한 뒤 이를 전체 process로 all-gather를 해줄 필요가 없다. 이는 당연히 parameter별로 owner process만 updated param을 가지고 있으면 되기 때문인데, reduced gradient를 가지고 있는 시점에서 해당 parameter의 optimizer state도 이미 가지고 있기 때문에 더 이상의 comm은 필요 없기 때문이다. 이렇게 해서 ZeRO-3는 \(3\Psi\)만큼의 cost가 들게 되며, 이것이 Trillion scale model을 학습할 수 있는 대신 지불하는 1.5배 느려지는 이유이다.

Communication Analysis of ZeRO-R (Focused on TP)

이번에는 ZeRO-R에 대해서 얘기해보자. 사실 ZeRO-R에서 memory defragment등의 technique이 있었지만 이것들은 분석할 필요까진 없어 보이고, TP에 partitioned activation을 적용했을 때 얼마나 efficient해지는지를 분석할 것이다.

사실 partitioned activation은 TP를 쓰지 않는 한 필요 없는 technique일 것이다. 원래 Megatron TP의 경우 activation checkpointing을 한다고 가정할 때 전체 model에서 forward-backward에 필요한 communication volume는 아래와 같다.

• total: \(12 \times B \times T \times d_{\text{model}} \times n_{\text{layers}}\)
  • forward: 2 all-reduce for \(B \times T \times d_{\text{model}}\)
  • re-computation (forward) for backward: 2 all-reduce for \(B \times T \times d_{\text{model}}\)
  • backward: 2 all-reduce for \(B \times T \times d_{\text{model}}\)

왜냐하면 all-reduce는 message_size의 2배에 달하는 communication이 필요하며, 모든 layer에서 이 연산이 일어나기 때문이다.

앞서 model size, \(\Psi\)에 대해서 ZeRO-3가 고작 \(3\Psi\)라고 했던 것에 비하면 \(3\Psi\)는 다음과 같기 때문에

(where factor of 12 is due to 4x for QKVO proj of self_attn and 8x for intermediate linear layers) TP와 ZeRO-DP의 comm volume ratio는 아래와 같이 계산된다.

이는 \(B,T\)에 비례해서 커지는 만큼 ZeRO-DP와 비교해서 매우 큰 연산량이라고 할 수 있다.

그런데 여기에 \(P_a\)를 적용할 경우 MP process group 별로 activation을 나눠 가지고 있기 때문에 backprop시 recomputation을 위해서 additional all-gather를 매 transformer block별로 한 번씩 해줘야 하는데 (보통 layer별로 하니까) all-gather는 comm cost가 message_size, \(B \times T \times d_{\text{model}} \times n_{\text{layers}}\)만큼 필요하기 때문에 original comm cost의 \(10%\)도 안되는 추가 cost로 memory를 매우 save할 수 있는 것이다.

그러니까 일반적으로 2 node, 16 GPU devices로 TP를 했다고 치면 \(P_a\)를 쓸 경우 activation에 의한 VRAM memory를 16배 줄일 수 있는 것이다. 다시 말하지만 \(P_a\)의 장점은 TP group의 process가 full size의 activation을 들고있을 필요가 없기 때문에 TP degree에 비례하게 memory를 줄일 수 있고, 그 결과 batch size를 그만큼 늘릴 수 있는게 TP+ZeRO의 throughput 증가의 핵심이다. 즉 TP degree가 16이라면 batch size를 16배 증가시킬 수 있는 것이다. (사실 paper에 DP communication volume이 batch size에 반비례한다고 얘기하는데 뭔소린지 잘 모르겠다. DP communication volume자체는 model size가 정해지면 변하지 않는 것 같은데 DP communication이 전체 wall clock time에서 차지하는 비율이 batch size에 반비례 한다고 얘기하고 싶은 것 같다.)

마지막으로 cpu offloading까지 한 \(P_{a+cpu}\)에 대해서, 이를 적용할 경우 per device activation memory requirement는 거의 0에 가까워지는데 그 댓가로 CPU <-> GPU간의 data movement가 발생하기 때문에 \(P_a\)의 memory movement와 비교해서 2배의 비용을 지불해야 하게 된다. 저자들은 small batch size를 쓰는 경우 처럼 DP communication volume이 bottleneck일 경우 CPU <-> GPU data transfer가 overhead가 되지 않을 때 까지 batch size를 늘리는 것이 좋을 것이라고 얘기한다.

아래는 ZeRO paper에 나오는 TP+ZeRO에 대한 paragraph들이다.

Fig. Paragraph for ZeRO powered TP (1). activation checkpointing을 위한 data-movement cost는 low bandwidth더라도 hiding이 가능하다.

Fig. Paragraph for ZeRO powered TP (2). 33GB 수준이 partitioning을 하면 2GB수준이 되고, CPU offloading을 하면 거의 0이 된다.

Evaluation

끝으로 여태까지 소개한 technique들을 적용해 일반적인 DP 대비 얼마나 많은 memory save가 가능하며, 그에 대한 반동으로 througput을 잃게 되는지? 아니면 분석처럼 hiding할 수 있는 수준이라 linear한 throughput 증가량을 보일 수 있을지 알아보도록 하자.

Per-device Memory Consumption of Different Optimization Settings

먼저 ZeRO Stage에 따라 얼마나 큰 model을 device에 올릴 수 있는지를 나타낸 table을 보도록 하자.

Fig. 128B model을 OOM없이 V100에 올리기 위해선 64장이 필요하다.

이 수치는 V100-32GB 장비를 가정하고 만든 Table이기 때문에 VRAM memory는 32GB가 limit이다. Table을 보면 \(N_d=1024\) (num. GPUs 1024, num. nodes 128) 일 경우 ZeRO-3를 쓰면 (\(P_{os+g+p}\)) 1 Trillion (T) 크기의 model까지 올릴 수 있다는 것을 확인할 수 있다. 하지만 꼭 device에 optimizer, gradient, parameter가 다 올라간다고 해서 학습이 가능한 것은 아니라는 점에 주의해야 한다. 왜냐하면 forward+backward 과정에서 생성되는 activation 같은 것들이 있기 때문에 실제로는 더 많은 memory가 필요하기 때문인데, 일단 이런건 무시하고 생각했을 때 그렇다는 것 같다.

Comparison between ZeRO+TP vs TP)

그 다음은 Megatron TP (vertical MP)와 ZeRO+TP의 GPU당 token 처리량 (Throughput)을 비교한 것이다.

보면 Baseline-MP와 Baseline w. internode MP가 나눠져 있는 것을 볼 수 있다. 일단 baseline-MP라는 것은 TP와 naive DP를 같이 적용한 것인데, 이를 나눈 이유는 TP에는 왠만하면 node를 넘어선 sharding은 느리니까 하면 느려지기 때문에 지양하자라는 생각이 기본이기 때문이다. 그러나 32GB-V100에서 model size가 40B만 넘어가도 node간의 (inter-node) sharding를 할 수밖에 없어진다. 그런데 이러면 TP는 위 figure에서 처럼 througpuht이 확 떨어지는 문제가 생긴다. 왜냐하면 앞서 얘기한 것 처럼 TP는 매 transformer block을 통과할 때 마다 activation을 all-to-all communication을 해야 하기 때문에 communication이 빈번하게 일어나는데, 이 때 model size가 커져서 inter-node partition을 하게 되면 high communication volume을 요구함과 동시에, 같은 node 내 (intra-node)의 300GB/sec per link (NVSwitch) 수준을 사용하다가 node 간 (inter-node)의 12.5 GB/sec per link (Infiniband EDR) 수준으로 communication bandwidth가 확 떨어지게 되기 때문이다.

Fig. Paragraph for ZeRO powered TP (3)

그러나 ZeRO를 TP와 같이 쓰면 100B model size에서 baseline보다 10배 빠르게 학습할 수 있었다고 한다. 그 이유는 당연히 앞서 분석한 것 처럼 partitioned activation를 통해 VRAM memory를 줄여 batch size를 늘릴 수 있었기 때문으로 보인다. (여기에 추가적으로 optimizer state만 partitioning하는 ZeRO-1도 같이 쓸 수 있는 것으로 알고있다)

Superlinear Scalability

그 다음은 60B model size에 대해서 GPU수가 늘어남에 따라 Total Throughput (TFLOPs)가 linear하게 증가하는가?에 대한 실험 결과이다.

결과는 super linear하다. 여기서 ZeRO-100B는 저자들이 정의한 setting으로 100B 미만에 대해서는 ZeRO-2, \(P_{os+g}\)와 ZeRO-R을 적용한 configuration을 쓸 것인데 그 이름을 ZeRO-100B로 지은 것이다. 당연히 이 결과는 memory save에 따른 batch size증가에 의한 것이다.

Memory and Performance Analysis

마지막으로 각종 ZeRO configuration에 따라

• Max model size
• Max cached allocated
• Throughput per GPU

등을 비교한 것인데, 이는 별로 해석에 어려움이 없을 것이 없기 때문에 넘어가도록 하겠다.

Fig

Fig

Fig

하나만 언급하고 넘어가자면 Figure 8에 나와있는 것 처럼 일반적으로는 C1->C4로 갈수록 memory save에 따른 batch size증가로 achieved throughput이 증가하게 되는데, C4->C5는 그렇지 않다. 이는 activation을 CPU offloading하면서 발생한 비용이 batch size를 키운 것 보다 훨씬 비효율적이어서 그런 것이다. 그러므로 이를 사용할 때는 꼭 throughput을 비교하는 주의를 기울여야 할 것이라는 점을 언급하고 넘어간다.

Other Techniques for Large Scale Modeling

Gradient Checkpoint (Activation Checkpoint)

Check these pages for more details. (i’m out of energy ;D)

• Fitting larger networks into memory. from Yaroslav Bulatov
• How Activation Checkpointing enables scaling up training deep learning models

Gradient Accumulation

Gradient Accumulation는 간단하게 말해 여러개의 GPU를 쓰는 분산 환경에서 model forward, backward를 통해서 계산한 gradient를 바로 model parameter에 반영해 업데이트 하지 않고 몇 step 누적시킨 후 계산하는 기법이다.
원래는 large batch를 쓰고 싶을 때 memory issue가 있는 경우 small batch size로 여러 번 forward-backward해서 gradient를 averaging하는 method이다.

하지만 accumulation을 잘 해주면 distributed training에서 장점이 있을 수 있는데, 예를 들어 아래의 figure 처럼 4개의 GPU를 쓸 때 gradient를 각 GPU에서 계산하고 나면 결국 model parameter를 업데이트 하기 위해서 rank 0의 machine에게 gradient를 전달해줘야 (통신) 하는데 이 때 모든 GPU들은 마지막 연산이 끝날 때 까지 기다리게 되는 문제가 생기는 것을 최소화 할 수 있다. (이는 아무래도 batch 구성을 잘해줄 경우 그럴 것이다)

Fig.

Library/Frameworks for Distributed Training

DeepSpeed and Pytorch Fully Sharded Data Parallel (FSDP)

어떻게 하면 ZeRO같은 기술을 쓸 수 있을까?

크게는 Microsoft’s DeepSpeed와 Pytorch’s Fully Sharded Data Parallel(FSDP)를 쓰는 방법이 있겠다.

• Microsoft’s DeepSpeed
• Pytorch’s Fully Sharded Data Parallel(FSDP)
  • Blog
  • Tutorial
  • lessw2020/fsdp_review

그 밖에도 많은 opensource가 있을 것이고 TensorFlow나 jax유저라면 또 다른 option이 있을 것이고 쌩으로 구현을 해도 된다. 하지만 깡으로 구현 하는 것이 distributed trainign에 대한 해상력을 높히고 opensource대비 abstraction을 많이 쓰지 않을 것이기에 최적화를 하기도 쉽겠으나 빠르고 정확한 framework을 만들기는 쉽지 않을 수 있다.

그렇기에 가장 먼저 추천한느 것은 아무래도 DeepSpeed를 쓰는 것인데, 그 이유는 ZeRO라는 기술이 microsoft에서 제안된 것이기 때문이다.

FSDP도 사실 ZeRO와 거의 기술적인 차이는 없다. 하지만 Pytorch를 개발한 Meta에서 관리하는 native ZeRO인 만큼 장기적으로 봤을 때 최적화가 더 잘 될것으로 보인다. 실제로 이 둘을 비교하여 뭐가 좋다는 결론을 낸 사람은 없는 것으로 알고있다. (아무리 찾아도 안나온다)

사실 FSDP는 같은 Meta의 분산 학습 library인 Fairscale에서 나온것이다. 보통 Meta가 개발하는 방식이 이런데, memory efficient attention, flash attention의 구현체가 xformers에 먼저 도입되었다가 latest Pytorch에 편입되듯 어떤 FSDP에 관한 feature를 가장 먼저 만나보려면 Fairscale을 보는 것도 좋을 것이다.

Accelerate

혹은 Huggingface’s Accelerate를 쓰는 것도 방법이다. 이는 huggignface model family에 ZeRO를 적용하는데 유용한데, 사실 이 library는 뭘 구현했다기 보다는 DeepSpeed나 FSDP같은 기술을 user들이 code 몇 줄 추가하는것으로 쉽게 사용할 수 있도록 해주는 wrapper이다.

Fig. Accelerate을 쓰는 방법은 매우 쉽다.

Fig. 231007 latest version docs에 PP, TP는 지원 안한다고 한다.

Accelerate가 하는일은 예를 들어 DeepSpeed를 사용할 경우 맨 처음 hugginface model class를 init을 할 때 ZeRO.Init을 하도록 하는 context manager 역할을 자동으로 한다거나, model init이 끝난 후 학습이 시작되기 전 deepspeed.initialize를 통해 model을 ZeRO에 맞게 준비시킨다거나 하는 것 등이고, 실제로 forward시 all-gather, backward시 all-reduce등을 하는 건 전부 deepspeed나 FSDP의 내부 logic에 따라 처리된다.

References

• Microsoft DeepSpeed Team
  • Full Paper List from Deepspeed Team
    • ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
      • older version
    • ZeRO-Offload: Democratizing Billion-Scale Model Training
    • ZeRO++: Extremely Efficient Collective Communication for Giant Model Training
    • ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
    • DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales
  • Blog Posts
    • ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters
    • ZeRO-2 & DeepSpeed: Shattering barriers of deep learning speed & scale
    • DeepSpeed: Extreme-scale model training for everyone
    • ZeRO-Infinity and DeepSpeed: Unlocking unprecedented model scale for deep learning training
    • DeepSpeed ZeRO++: A leap in speed for LLM and chat model training with 4X less communication
  • Github
  • Docs
  • ZeRO-Offload slides from USENIX
  • ds config json
• Google
  • Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training
• Meta
  • Pytorch Docs > FullyShardedDataParallel
  • Fairscale Docs > fairscale.nn.data_parallel.fully_sharded_data_parallel
  • Introducing PyTorch Fully Sharded Data Parallel (FSDP) API
• Huggingface
  • Methods and tools for efficient training on a single GPU
  • Efficient Training on Multiple GPUs from Huggingface Blog
  • How to use deepspeed from accelerate
• CS 886: Recent Advances on Foundation Models from Wenhu Chen
  • Efficient LLM Training
• Deep Learning Systems Course
  • Lecture 15 - Training Large Models
• Others
  • parallelism
    • How to Train Really Large Models on Many GPUs?
    • Parallelism in Distributed Deep Learning Insu Jang
    • Large-scale language modeling tutorials with PyTorch from Hyun Woong Ko
    • Large Language Models (in 2023) at SNU from Hyung Won Chung
      • slide
    • 초고성능 딥러닝 클러스터 구축하기 1편 from HyungMo Kim
  • gradient checkpointing
    • Fitting larger networks into memory. from Yaroslav Bulatov
    • How Activation Checkpointing enables scaling up training deep learning models
    • Saving memory using gradient-checkpointing
    • Efficient Training on Multiple GPUs
    • Methods and tools for efficient training on a single GPU
    • transformers issue for Supporting Selective Activation Checkpointing and CPU Offloading Option