ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

읽게 된 이유 TL;DR 공홈 홍보글 Data Parallel & Model Parallel & Pipeline Parallel Abstract 현존하는 거대한 Language Model들은 학습하는 것에 문제가 있음.우리는 vram 효율도 좋고 속도도 빠르고 엄청 큰 모델도 올릴 수 있는 것을 만들었다!Extended Introduction LM이 너무 커지고 있다.Model Parallel vs Pipeline Parallel ZeRO에서 제안하는 방향 Optimize Model state Memory ZeRO-DP: ZeRO-powered Data Parallel ZeRO-DP의 3가지 단계 Optimize Residual State Memory ZeRO-R: Residual 메모리 최적화 그렇다면 MP는 이제 필요가 없나?1-Bit Adam Optimizer ZeRO의 장점 ZeRO-Offload ZeRO-Offload의 구성 ZeRO-Offload는 어떻게 10B처럼 큰 모델을 단일 GPU로 올리나?DeepSpeed Sparse Attention Reference

읽게 된 이유

N모사에서 NVIDIA Superpod으로 175B짜리 한국어 모델을 굽는 중 (GPT-3)

하지만 저거는 1024대 이상의 GPU(A100)을 쓰는데.. MS에서 1대 만으로도 10B 모델을 구울 수 있다고?

이건 해봐야 해...!

TL;DR

ZeRO1, 2와 CPU OffLoad까지 사용하면 10B 모델을 1대의 V100 GPU(32G vram)에서 학습할 수 있다.

이때 Model Parallel, Pipeline Parallel 등 모델 파티셔닝을 극한으로 쪼개고, 여기에 CPU 메모리까지 추가로 사용해서 Vram의 사용 효율을 극대화 한다.

v1은 ZeRO1, Oct 2019

v2는 ZeRO2와 ZeRO-Offload, May 2020

(주의: 1GPU 사용시에는 MP/DP 없이 CPU Offload만 사용함)

DeepSpeed & ZeRO2 설명 영상

공홈 홍보글

BERT-Large 학습에 한시간!

Data Parallel & Model Parallel & Pipeline Parallel

딥러닝 뉴럴넷을 쪼개서 Parallel하게 학습하는 3가지 방법

1) Data Parallel

PyTorch의 ddp 쓰면 곧바로 사용 가능

Batch를 나눠서 할당

모델은 한 GPU에 다 올라가는 셈

2) Model Parallel

Layer 한 층을 여러개로 쪼개서 여러 GPU로 계산

GPU간에 통신이 자주 필요함

3) Pipeline Parallel (제일 위 그림에서는 workload partitioning이라고 표현)

Layer를 층별로 나눔

앞선 Layer의 계산이 끝날때까지 기다렸다가 결과를 받아 다음 GPU로 넘겨줌

위로 인한 공백 시간이 발생하는 것을 막기 위해, Micro batch in Mini Batch 기법을 사용함

Batch Size가 128이라면 이걸 4로 나눠서 32씩 계산하고 다음 GPU로 넘기고 하는 방법
GPU Utilization 높이기 (= 연산 전체 속도 상승)

Abstract

현존하는 거대한 Language Model들은 학습하는 것에 문제가 있음.

학습을 빠르게 하고싶다 = 많은 GPU + 큰 Batch size → Data Parallel

모델을 복사해야 해서 Redundency 발생
한 GPU 보다 큰 모델은 학습 못함

엄청 큰 모델을 학습하고 싶다 → 모델을 쪼개자 → Model Parallel

Transformer의 MultiHead를 각 head별로 쪼개서 GPU에서 각각 계산하자!
커뮤니케이션이 결국 필요함
개발도 어려움

우리는 vram 효율도 좋고 속도도 빠르고 엄청 큰 모델도 올릴 수 있는 것을 만들었다!

→ "ZeRO: Zero Redundancy Optimizer"

DP/MP에서 발생하는 Memory Redundancy를 제거

Communication cost를 최소화

Computational Power를 최대한 유지(각 GPU별로)

이 방법으로 최대 1Trillion Params를 가진 모델을 현존하는 장비(아마도 V100 1024대)로 학습할 수 있다.

그리고 이거 오픈소스(DeepSpeed)로 공개함 → 100B params GPT를 400GPU에서 학습 가능했음

Model Parallel 없이도 13B 모델을 단일 GPU에 올릴 수 있음

Extended Introduction

기존에 ZeRO1에서 버전 올리면서 한 페이퍼에 분량이 늘어남. 사실상 이 부분만 봐도 논문 다 보는 셈

LM이 너무 커지고 있다.

기존의 큰 Language Model에는

BERT-Large 0.3B

분명히 KcBERT 학습할 때는 엄청 커보였는데 이제는 너무 작고 귀여워 보인다..

GPT-2 1.3B

Meagtron-LM 8.3B

T5 11B

GPT-3 175B

하지만 MultiGPU를 사용하기 위해 가장 쉽게 사용하는 Data Parallel(DP, DDP) 방식은 한 GPU에 모든 Model param이 올라가야 한다는 점으로 인해, 32G V100 기준 1.4B params 이상은 올릴 수가 없다.

다른 솔루션으로는 Pipeline Parallel 방식으로, 각 layer(층)별로 모델을 쪼개거나 혹은 Model Parallel로 Multihead attention에서 head를 쪼개는 등의 방식을 사용할 수도 있다.

또다른 방식으로는 CPU Off loading(CPU Ram 활용) 등을 쓸 수도 있다.

하지만 여전히 이러한 방식을 써도 속도 이슈 등이 발생한다.

Model Parallel vs Pipeline Parallel

위 3가지 방법 중 가장 쉽게 적용 가능한 것이 Model Parallel로, T5와 Megatron-LM은 MP를 적용해서 학습함.

하지만 MP는 모델 크기에 따라 최대로 늘릴 수 있는 것에 제약이 있다.

MultiHead att 등에서 HEAD를 쪼개는 것이기 때문에 최대로 늘릴 수 있는 경우가 Head 수로 제약이 됨

좀더 General하게 이야기 하면, 각 layer에 있는 params를 쪼개서 여러 GPU로 나눠 계산하는 것.

각 layer 단에서 계산된 값을 공유하는 것으로 인해 communication cost가 무척 높음

따라서, 단일 Node(multiGPU인 한 장비)에서 쓰는 것은 괜찮지만, Multi Node로 가는 순간 속도 급락

40B 모델로 테스트한 경우(2대 DGX-2) → V100 GPU당 겨우 5Tflops 나옴. (각 GPU 최대 성능의 5%)

ZeRO에서 제안하는 방향

"어디에서 Memory comsumption이 발생하는가?"

대부분은 model states = optimizer states + gradients + params
나머지는 Residual states = activation + tmp buffer + fragmented memory

ZeRO에서 제안하는 방향

Memory Efficiency
High compute & less communication

Optimize Model state Memory

model state는 학습 중에 가장 많은 GPU 메모리를 차지한다.

하지만 현존하는 DP와 MP는 좋은 솔루션이 아니다.

DP: Good compute & communication efficiency, Poor memory efficiency

모델 전체를 copy 해서 모든 GPU에 할당
각 Batch size를 N개(GPU)로 나눠서 계산

MP: Poor compute & comm efficiency.

위 두가지 방법 모두, 모든 Model state가 항상 필요한 것이 아님에도 GPU 메모리를 차지함

ZeRO-DP: ZeRO-powered Data Parallel

DP급의 computation & communication efficiency 유지하지만 MP급의 메모리 여유

Memory state redundancy 제거

모델을 copy하는 대신 Partitioning

Dynamic communication schedule 통해서 communication volume을 DP급으로 줄임

ZeRO-DP의 3가지 단계

사실 이 Figure 하나가 논문 전체를 대변해 주는 것과 동일할 정도.

Optimizer State Partitioning (P_os)

1/4로 메모리 감소, communication cost 동일

P_os + Gradient Partitioning (P_os+g)

1/8로 메모리 감소, comm cost 동일

P_os+g + Parameter Partitioning

N개의 GPU 쓸 경우 1/N으로 메모리 감소, comm cost는 50% 증가

위 방법 통해서 1T params 모델을 1024 V100 GPU로 학습시킬 수 있다.

만약 Adam Optimizer + fp16으로 학습한다면 1GPU당 약 16G 메모리 사용함 (충분히 가능!)

Optimize Residual State Memory

model state뿐만 아니라, activation, temp buffers, unusable memory fragments 이슈가 있음

ZeRO-R: Residual 메모리 최적화

Activation 최적화

Checkpointing은 약간 도움 되지만 큰 모델에는 그닥 도움 x
MP에 있는 중복 Activation을 제거 (activation partitioning)
만약 가능한 경우 CPU memory로 Activation을 Offload

적당한 Temp buffer 사이즈 책정

Tensor의 Lifetime 사용으로 fragmented memory 없도록 만들기

그렇다면 MP는 이제 필요가 없나?

NO, 여전히 쓸만할 경우가 있음

ZeRO-R을 사용할 경우 MP를 사용하면 Actvation memory footprint를 줄일 수 있음 (엄청 큰 모델의 경우)

DP 혼자 쓰는 것보다 MP 함께 쓰면 큰 Batch size로 aggreagate하는 효과를 누릴 수 있음

좀더 좋은 Convergence 가능함

1-Bit Adam Optimizer

Adam Optimizer를 1bit로 전송함

Gradient Momentum 나온 값을 UP / DOWN 두가지 중 하나만을 선택해서 정해진 alpha값만큼 증가 혹은 감소하도록 처리

모델이 워낙 크고 작은 lr을 사용하니 가능한 방식

Same Converge 보장, 1/5 정도의 communication cost 소모

초반에는 Gradient의 Variance update가 굉장히 차이가 커서 정보량이 크지만, 일정 수준(Warmup) epochs 이후에는 Variance update 값이 매우 작아지기 때문에 고정된 값을 사용해도 된다. 대신 +인지 - 인지 방향만 전달해주면 OK.

연구에서는 약 15-20%정도만 warmup해주면 OK.

ZeRO의 장점

Model Size

ZeRO-100B로 170B짜리 학습 가능
현존하는 Megatron-LM 코드로는 40B 모델이 최대

Speed

400대의 V100으로 학습할 경우 100B 모델 학습 시간 기준 기존 대비 10배 속도
한 GPU당 학습 속도가 약 38TFlops, 총합 15PFlops

Scalability

64→400GPU까지 올리는데 큰 지장이 없음. 그 이상도 가능

→ MegatronLM vs ZeRO

ZeRO-100B 사용하면 PP/MP 없이도 한 GPU에 13B 모델 올릴 수 있음

기존 PyTorch DDP 사용하면 1.4B 이상시 OOM 발생함

이 외에도 Sparse transformer 지원을 통해 엄청 긴 Sequence를 학습할 수 있음 (Sparse attention)

위 모든 것 구현된 곳 → DeepSpeed

DeepSpeed

DeepSpeed is a deep learning optimization library that makes distributed training easy, efficient, and effective. 10x Larger Models 10x Faster Training Minimal Code Change DeepSpeed delivers extreme-scale model training for everyone, from data scientists training on massive supercomputers to those training on low-end clusters or even on a single GPU: Extreme scale: Using current generation of GPU clusters with hundreds of devices, 3D parallelism of DeepSpeed can efficiently train deep learning models with trillions of parameters.

https://www.deepspeed.ai/

ZeRO-Offload

단일 GPU에서 10B 모델을!

ZeRO-Offload의 구성

ZeRO-Offload는 ZeRO-2 위에서 동작한다.

Offload의 핵심은 Compute efficiency loss를 최소화하는 방향이라서, GPU를 최대한 많이 쓰는 방향으로 동작한다.

ZeRO-Offload는 어떻게 10B처럼 큰 모델을 단일 GPU로 올리나?

가장 기본적으로, CPU Offload를 통해 Optimizer State와 Gradients를 CPU 메모리로 올린다.

ZeRO-2에서는 위 두 가지를 각각 다른 GPU에서 사용할 수 있도록 Partitioning을 하지만, 그 대신 CPU ram에 올리는 셈.

전체 학습 과정에서 Optimizer states는 항상 CPU memory에 올라가 있고, Gradients는 학습에 사용하는 GPU들에서 Backward 과정 중 Reduce-scatter 한 뒤, 각각의 GPU(DataParallel)에서 나온 Gradients들의 Avg로 다시 CPU로 Offload한다. (각각 GPU 학습에 사용하는 부분 말고는 버림)

Reduce-Scatter?

ReduceScatter는 기본적으로 Reduce와 동일하지만 결과물이 각각 다른 Partition 으로 나간다는 차이

위 과정으로 CPU에 Gradient가 올라오고 나면 CPU상에서 DataParallel 각각에서 Optimizer State partition이 업데이트(p update)되고, 업데이트 된 파라미터들은 All-gather 연산으로 합쳐진 뒤 각각의 GPU로 다시 내려간다.

이때, g offload 와 g swap 사이의 간격을 최소화 하기 위해 각 DataParallel별로 끝나는 시점에 곧장 CPU로 올려서 연산을 시작하고, 연산이 끝나는 즉시 GPU로 내리는 순서를 진행한다.

(이로 인해 각각 다른 CUDA stream을 사용한다고 함)