#!/usr/bin/env python
# coding: utf-8
# # Zero Redundancy Optimization (ZeRO)
#
# 이번 세션에는 Microsoft의 뉴럴넷 학습 최적화 솔루션인 ZeRO에 대해서 알아보도록 하겠습니다.
# ## 1. Mixed Precision
#
# 최신 GPU들이 Lower precision에 대한 계산을 지원하면서 현대의 뉴럴넷 학습은 대부분 FP16(half)과 FP32(single)을 함께 사용하는 Mixed precision 방식을 사용합니다. V100 기준으로 FP32에서 속도가 14TFLOPS 정도라면, FP16에서는 100TFLOPS의 속도로 모델을 학습할 수 있습니다. 또한 FP16을 사용하면 모델의 사이즈가 줄기 때문에 학습 뿐만 아니라 배포시에도 장점이 있죠.
#
#
#
# ![](../images/mixed_precision_1.png)
#
#
#
# ### 그런데 왜 Mixed?
# 그런데 여기에서 의문이 듭니다. FP16으로만 모델을 학습시키면 되지, 굳이 FP32와 FP16을 같이 쓸 필요가 있을까요? 결과부터 말하자면 FP16만으로 학습시 Loss가 심하게 발산하여 학습이 거의 불가능합니다. Gradient를 FP16로 유지하면 대부분의 소수점을 버리는 것이기 때문에 정밀한 학습이 불가능해집니다. 따라서 속도가 빠른 FP16과 정확도가 높은 FP32를 모두 사용해서 두 방식의 장점만을 취하려고 하는 것이죠.
#
# ![](../images/ddp_analysis_3.png)
#
# Computation cost가 큰 Forward와 Backward는 FP16 모델로 하고, 계산된 Gradient를 정밀도가 높은 FP32 모델에 복사해서 weight를 업데이트 합니다. 그런데 여기서 궁금한 점이 생깁니다. FP16의 Gradient를 FP32에 적용하려면 어떻게 해야할까요? 연구진이 실험한 결과, FP16으로 계산된 Loss를 Backward 하면 크기가 크기가 작았던 일부 값들(그림에서 왼쪽)은 계산이 되면서 0으로 변해버렸다고 합니다.
#
# ![](../images/mixed_precision_4.png)
#
#
#
# ### Loss Scaling
# 이러한 문제를 어떻게 해결할 수 있을까요? 매우 심플한 아이디어로, Loss Gradient에 매우 큰 값을 곱해줘서 분포를 오른쪽으로 밀어주면 됩니다. 이러한 기술의 이름을 Loss scaling이라고 합니다. FP16의 Loss에 매우 큰 값을 곱하면, FP32에 적용 했을 때 사라져 버릴 수 있는 값들도 잘 살려낼 수 있죠.
#
# ![](../images/mixed_precision_5.png)
#
# In[ ]:
"""
참고: apex/apex/amp/opt.py
"""
import contextlib
@contextlib.contextmanager
def scale_loss(self, loss):
if not self._amp_handle.is_active():
yield loss
return
# When there are multiple losses per-optimizer, we need
# to save out current grad accumulation, since we won't be
# able to unscale this particulare loss once the grads are
# all mixed together.
cached_grads = []
if self._loss_idx > 0:
for p in master_params(self._optimizer):
if p.grad is not None:
cached_grads.append(p.grad.data.detach().clone())
else:
cached_grads.append(None)
self._optimizer.zero_grad()
loss_scale = self._cur_loss_scaler().loss_scale()
yield loss * loss_scale
# In[ ]:
"""
참고: apex/tests/L0/run_amp/test_fused_sgd.py
"""
with amp.scale_loss(loss0, optimizer, loss_id=loss_ids[0]) as scaled_loss:
scaled_loss.backward()
if i == inject_inf and which_backward == 0:
if inject_inf_loc == "fp32":
model0.weight0.grad[0] = float('inf')
elif inject_inf_loc == "fp16":
model0.weight1.grad[0] = float('inf')
# 실제로 아래 그림처럼 Loss에 큰 값을 곱해주면 발산하지 않고 학습이 잘 되었다고 합니다. 회색 그래프는 scaling을 하지 않았을때, 녹색은 scaling 했을때의 성능입니다. 놀랍게도 FP32와 성능이 거의 흡사하죠.
#
# ![](../images/mixed_precision_2.png)
#
# 이러한 이유로 FP16과 FP32를 함께 사용하는 Mixed precision은 현대 뉴럴넷 학습에 거의 필수가 되었습니다. FP16 정도의 저장 용량으로 FP32의 커버리지를 커버하는 bfloat16 (Google TPU) 방식이 지금보다 더 다양한 GPU에서 지원되고 대중화 되기 전까지는 FP16 + 32의 Mixed precision training은 뉴럴넷 학습에 필수적으로 쓰이는 기술일 것입니다.
#
#
#
# ### Mixed Precision의 동작방식
#
# 다음은 Mixed Precision의 동작 방식을 나타낸 그림입니다. 코드와 수식을 이용해 진행 과정을 자세히 살펴봅시다.
#
#
#
# ![](../images/mixed_precision_33.png)
# ### 0) 모델과 옵티마이저 생성
#
# 2개의 레이어를 가진 뉴럴넷을 정의합니다.
# In[1]:
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super().__init__()
self.w1 = nn.Linear(512, 512, bias=False)
self.w2 = nn.Linear(512, 1, bias=False)
def forward(self, x):
z1 = self.w1(x)
z2 = self.w2(z1)
return z2
# 학습할 뉴럴넷과, 옵티마이저 생성합니다.
# In[2]:
from torch.optim import SGD
fp32_model= Net().to("cuda")
optimizer = SGD(fp32_model.parameters(), lr=1e-2)
# In[3]:
f"GPU = {torch.cuda.memory_allocated(0) / (1024 ** 2)} GiB"
#
#
# ### 1) Float2Half
#
# 이 과정은 단순히 `0.524796132`와 같은 파라미터를 `0.5247`과 같이 잘라내는 작업입니다.
#
# 보시다시피 용량도 FP32 모델의 절반정도 사이즈를 가집니다. (1.0 GB + 0.5 GB)
# In[4]:
fp16_model = Net().half().to("cuda")
fp16_model.load_state_dict(fp32_model.state_dict())
# In[5]:
f"GPU = {torch.cuda.memory_allocated(0) / (1024 ** 2)} GiB"
#
#
# ### 2) Forward
#
# fp16으로 복사된 모델을 이용하여 forward pass를 수행합니다.
#
# $z_1 = w_1 \cdot x \; $ (FWD: layer1)
#
# $z_2 = w_2 \cdot z_1 \; $ (FWD: layer2)
# In[6]:
import torch
# example input sizes
batch_size, hidden_size = 4, 512
# create dummy data (bsz=4, hid=256)
x = torch.randn(batch_size,hidden_size, dtype=torch.half, device="cuda")
# do forward
z2 = fp16_model(x)
# check dtypr of output logits
f"logits type = {z2.dtype}"
# 계산된 FP16의 출력값을 이용하여 Loss를 계산합니다.
#
# $L = \frac{(y - z_2)^2}{2} \; $ (Loss computation)
# In[7]:
# craete dummy data (bsz=4)
y = torch.tensor([[1.9], [9.5], [0.9], [1.2]], dtype=torch.half, device="cuda")
# compute mean square error loss
L = torch.nn.functional.mse_loss(z2, y)
# check dtype of loss
f"loss type = {L.dtype}"
#
#
# ### 3) Backward
#
# 이제 $w_n := w_n - lr \cdot \frac{dL}{dw_n}$와 같은 Gradient Descent Rule로 모델의 파라미터를 업데이트 해야 합니다.
#
# 따라서 $\frac{dL}{dw_1}$과 $\frac{dL}{dw_2}$와 같은 Gradient를 구해야 하는데요. 이들은 대략 아래와 같습니다. (chain rule에 의해서 원하는 결과를 얻을 수 있습니다.)
#
# $\frac{dL}{dw_2} = \frac{dL}{dz_2} \cdot \frac{dz_2}{dw_2}$
#
# $\frac{dL}{dw_1} = \frac{dL}{dz_2} \cdot \frac{dz_2}{dz_1} \cdot \frac{dz_1}{dw_1}$
#
#
#
#
#
# 구체적으로는 아래와 같습니다.
#
# $\frac{dL}{dz_2} = y - z_2 \; $ (BWD-activation: layer2)
#
# $\frac{dz_2}{dw_2} = z_1 \;$ (BWD-weight: layer2)
#
# $\frac{dz_2}{dz_1} = w_2 \;$ (BWD-activation: layer1)
#
# $\frac{dz_1}{dw_1} = x \; $ (BWD-weight: layer1)
#
#
#
# $\frac{dL}{dw_2} = (y - z_2) \cdot z_1$
#
# $\frac{dL}{dw_1} = (y - z_2) \cdot w_2 \cdot x$
#
# In[8]:
# loss scaling
L *= 1024
# do backward
L.backward()
#
#
# ### 4) Update Weight
#
# 마지막으로 파라미터를 업데이트하기 위해 `optimizer.step()`를 수행합니다.
#
# $w_1 := w_1 - lr \cdot \frac{dL}{dw_1} \; $ (Weight Update)
#
# $w_2 := w_2 - lr \cdot \frac{dL}{dw_2} \; $ (Weight Update)
# In[9]:
print(f'before: {fp32_model.w1.weight}\n')
optimizer.step()
print(f'after: {fp32_model.w1.weight}\n')
# 생각해보면, FP32 모델은 forward & backward를 수행한적이 없었죠. 따라서 gradient 텐서를 갖고있지 않습니다. 그래서 `optimizer.step()`을 수행 해도 값이 변하지 않았습니다. 따라서 `optimizer.step()`을 수행하기 전에, `backward()`를 거친 FP16모델의 gradient를 복사해야 합니다.
#
# 참고로 PyTorch는 파라미터(`nn.Parameter`) 중 `requires_grad=True`로 설정된 파라미터들은 모두 `grad`라는 애트리뷰트를 가지고 있습니다. 모델이 출력한 텐서의 `backward`가 호출되면 graph를 타고 뒤로 돌아오면서 미분 계산을 수행하고 결과 값을 `grad`라는 공간에 저장합니다. `grad`는 해당 텐서와 동일한 사이즈이기 때문에 모델의 용량이 10GB라면 gradient도 10GB 만큼 필요합니다. 우리가 인퍼런스 할 때 보다 학습할때 메모리가 훨씬 많이 필요한 이유 중 하나입니다. 따라서 학습에 사용될 텐서가 아니라면 반드시 `requires_grad`를 `False`로 설정해야 불필요한 메모리 소모를 막을 수 있습니다.
#
# In[10]:
# copy gradient to FP32 model
fp32_model.w1.weight.grad = fp16_model.w1.weight.grad.float()
fp32_model.w2.weight.grad = fp16_model.w2.weight.grad.float()
# In[ ]:
print(f'before: {fp32_model.w1.weight}\n')
optimizer.step()
print(f'after: {fp32_model.w1.weight}\n')
# ### Pytorch에서 Mixed precision training 수행하기
#
# Pytorch에서는 다음과 같이 손쉽게 Mixed precision training을 수행할 수 있습니다.
# In[ ]:
# 참고: https://pytorch.org/blog/accelerating-training-on-nvidia-gpus-with-pytorch-automatic-mixed-precision/
import torch
# Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler()
for data, label in data_iter:
optimizer.zero_grad()
# Casts operations to mixed precision
with torch.cuda.amp.autocast():
loss = model(data)
# Scales the loss, and calls backward()
# to create scaled gradients
scaler.scale(loss).backward()
# Unscales gradients and calls
# or skips optimizer.step()
scaler.step(optimizer)
# Updates the scale for next iteration
scaler.update()
#
#
# ### Dynamic Loss Scaling
#
# Loss Scaling은 Mixed Precision 학습을 매우 효과적으로 만들어줬습니다. 그러나 scale 수치를 몇으로 설정하는 것이 가장 좋을지 알기가 매우 어렵습니다. 따라서 몇몇 오픈소스에는 이러한 문제를 해결하기 위해 Dynamic Loss Scaling 기법을 제안합니다. 이는 NVIDIA의 `amp`나 MS의 `deepspeed`에도 구현되어 있습니다.
#
# Dynamic Loss Scaling의 아이디어는 매우 간단합니다. **목표는 Gradient의 소수점들이 Overflow 되지 않는 선에서 scale값을 최대로 유지하는 것**입니다. Gradient 값을 키우면 키울수록 좋지만 너무 커지면 Overflow가 발생하기 때문에 Overflow가 되지 않는 선에서 최대로 키워주는 것이죠.
#
# 따라서 학습 초반에 매우 큰 값을 scale 값으로 설정합니다. `deepspeed`의 경우 기본 값이 $2^{32}$로 설정되어 있습니다. 이 값으로 Loss를 backward 해보고 만약 Gradient가 Overflow 되었다면 scale 값을 2배 줄입니다. 이 과정을 여러번 반복하면서 Overflow가 발생하지 않는 최대의 scale값을 찾아내는 것이 바로 Dynamic Loss Scaling입니다.
#
#
# ### AMP (Apex Mixed Precision)
#
# `apex`는 NVIDIA에서 개발한 라이브러리로, Mixed Precision 라이브러리 중에서 가장 유명한 인지도를 가지고 있습니다. 요즘에는 `torch`자체에 mixed precision 기능이 내장되기도 하고 DeepSpeed, Pytorch-Lightning 등의 도구가 많이 나오게 돼서 `apex`를 예전만큼은 자주 사용하지 않지만 그래도 여전히 많이 사용되고 있는 라이브러리입니다. 사용법은 아래와 같이 매우 간단합니다.
# In[ ]:
import torch
from apex import amp
# Declare model and optimizer as usual, with default (FP32) precision
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
# Allow Amp to perform casts as required by the opt_level
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
# loss.backward() becomes:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
# 위 코드를 보면 `opt_level`이라는 파라미터가 보입니다. `apex`에는 mixed precision의 level을 설정할 수 있는 기능이 있는데 이를 알아두면 추후에 `apex`를 사용할 일이 생길때 매우 유용할 것입니다. (참고로 알파벳 O + 숫자 0,1,2,3입니다.)
#
# ![](../images/apex.png)
#
# - `O0`: FP32 학습
# - `O1`: FP16을 잘 지원하는 Tensor Core 연산들은 FP16 / 나머지는 FP32
# - `O2`: Normalization의 weight를 제외한 모든 파라미터를 FP16으로 설정
# - `O3`: FP16 학습
#
#
# ## 2. Zero Redundancy Optimization
#
# FP16과 FP32를 함께 사용하게 됨으로써 학습 속도는 매우 빨라지게 되었지만 단점이 생겼습니다. 바로 메모리인데요. FP32의 master weight과 FP16 파라미터, Gradient를 모두 GPU에 올려둔 상태이기 때문에 메모리가 기존보다 더 많이 필요해집니다.
#
# ![](../images/zero_1.png)
#
# 그리고 모델 파라미터가 FP16로 존재한다고 해도, Optimization은 FP32에서 일어나기 때문에 AdaGrad, Adam 등의 Adaptive optimizer 들이 필요로 하는 Variance 및 Momentum과 같은 텐서들은 여전히 FP32로 보관되어야 합니다.
#
# ![](../images/adam.png)
#
# In[ ]:
"""
참고: pytorch/torch/optim/adam.py
"""
@torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.
Args:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
params_with_grad = []
grads = []
exp_avgs = []
exp_avg_sqs = []
max_exp_avg_sqs = []
state_steps = []
beta1, beta2 = group['betas']
for p in group['params']:
if p.grad is not None:
params_with_grad.append(p)
if p.grad.is_sparse:
raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
grads.append(p.grad)
state = self.state[p]
# Lazy state initialization
# 모든 파라미터에 대해서 동일 사이즈로 `exp_avg`와 `exp_avg_sq`로 가지고 있음
# 이 때문에 Adam 기반의 optimizer를 사용하면 모델 2개에 해당하는 GPU 메모리가 더 필요해짐.
if len(state) == 0:
state['step'] = 0
# Exponential moving average of gradient values
state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# Exponential moving average of squared gradient values
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
if group['amsgrad']:
# Maintains max of all exp. moving avg. of sq. grad. values
state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
exp_avgs.append(state['exp_avg'])
exp_avg_sqs.append(state['exp_avg_sq'])
if group['amsgrad']:
max_exp_avg_sqs.append(state['max_exp_avg_sq'])
# update the steps for each param group update
state['step'] += 1
# record the step after step update
state_steps.append(state['step'])
F.adam(params_with_grad,
grads,
exp_avgs,
exp_avg_sqs,
max_exp_avg_sqs,
state_steps,
amsgrad=group['amsgrad'],
beta1=beta1,
beta2=beta2,
lr=group['lr'],
weight_decay=group['weight_decay'],
eps=group['eps'])
return loss
# 지금까지 FP16 parameter, gradient, FP32 parameter, gradient, momentum, variance 등 우리가 모델을 학습 할 때 메모리에 할당되는 텐서들의 종류에 대해서 조사했습니다. 놀라운 것은 진짜 모델이 차지하는 영역은 얼마 안된다는 것이죠. 이렇게 학습시에는 모델 외에도 **부가적으로 어마어마한 양의 텐서가 GPU 메모리에 할당됩니다.**
#
#
#
# ![](../images/memory.png)
#
#
#
# 추가로 **Data 텐서**와 **Activation 텐서**도 메모리에 할당됩니다. Data 텐서는 모델에 입력되기 전의 토큰 상태의 텐서를 의미하며, Activation 텐서는 Forward & Bacward 과정에서 연산되는 Hidden states 등의 텐서를 의미합니다. 추가로 분산처리를 수행하면 **통신 중에 텐서들을 담아둘 Bucket 공간** 등도 필요합니다. 버킷에 대해서는 이미 Data Parallelism 세션에서 Gradient Bucketing 등으로 다루었던 적이 있죠. 따라서 **모델과 데이터만 병렬화 할 것이 아니라 이러한 Optimizer States(분산, 모멘텀), Data & Activation Memory 등도 관리할 필요**가 있습니다.
#
#
#
# Zero Redundancy Optimization (이하 ZeRO)는 이러한 부분들을 매우 효율적으로 관리 할 수 있도록 도와주는 **메모리 최적화 기술의 집합체**입니다. 크게 **ZeRO-DP** (ZeRO Data Parallelism)과 **ZeRO-R** (ZeRO Residual States) 등의 솔루션이 존재합니다. 이제부터 차근 차근 알아봅시다.
#
#
# ## 3. ZeRO Data Parallelism
#
# 가장 먼저 메모리 상태를 조사해보면, 위 그림에서 왼편 (FP16, 32, model & optimizer & gradient)가 가장 큰 공간을 차지합니다. 따라서 이들을 효율적으로 쪼개서 관리해야 합니다. ZeRO-DP는 Data Parallel과 함께 이러한 텐서들을 디바이스마다 쪼개서 관리 할 수 있도록 도와줍니다.
#
# ![](../images/zero_2.png)
#
# ZeRO-DP는 4개의 stage로 나누어서 제공되고 있으며 `DeepSpeed` 라이브러리를 통해 선택적으로 적용 할 수 있습니다.
#
# - **Stage 0**:
# - No Partitioning
# - ZeRO-DP를 적용하지 않습니다.
# - **Stage 1**:
# - Optimizer States Partitioning
# - Optimizer Stages(모멘텀, 분산) 텐서를 여러 GPU로 분할합니다.
# - 메모리 소비량 4배 감소
# - 기존과 비슷한 양의 Communication Cost
# - **Stage 2**:
# - Stage 1 + Gradient partitioning
# - Gradient(기울기) 텐서를 여러 GPU로 분할합니다.
# - 메모리 소비량 2배 더 감소
# - 기존과 비슷한 양의 Communication Cost
# - **Stage 3**:
# - Parameter partitioning
# - Parameter(모델) 텐서를 여러 GPU로 분할합니다.
# - 메모리 소비량 분할 수준에 따라 선형적 감소
# - 기존보다 1.5배 많은 Communication Cost
#
# ZeRO-DP의 동작은 매우 복잡하기 때문에 영상으로 확인하겠습니다.
#
# https://www.microsoft.com/en-us/research/uploads/prod/2020/02/Turing-Animation.mp4?_=1
# In[15]:
from IPython.display import HTML
HTML("""