本文章主要介紹 Microsoft 所開發的 DeepSpeed, deep learning optimization library 是透過何種優化來使得我們的 deep learning training 可以更加的快速有效率，並同時使用更少的 GPU memory !
因此也同時會介紹平時我們在訓練時究竟 GPU 記憶體都去哪裡了，也會介紹現在大家對於多 GPU 主要使用的 parallelism 方法的優點及缺點！
Introduction
現在的 deep learning 套件 (e.g., tensorflow, and pytorch) 非常的方便且容易上手使用，只需要一行程式碼就可以輕鬆的做完整個 forward 的流程，也只需要一行程式碼就可以輕鬆的透過 chain rule 來計算出所有 parameter 的 gradient。
但是當我們想要訓練一個非常龐大的模型 (e.g., GPT-2) 的時候，現在的 deep learning 套件可能就會面臨到一個瓶頸：OOM(out of memory)。造成 OOM 的原因主要是在於模型的 parameters 過多，導致我們將模型放入 GPU 之後，我們沒有多餘的記憶體空間來存放 forward/backward 所需要存放的額外資訊 (e.g., activations 和 gradients)。因此，在本篇文章中，首先會介紹我們平常在訓練模型時，究竟 GPU 記憶體都是被哪些東西所佔據了，接著我們會介紹現在大家在訓練模型時最常使用的 MP (model parallelism) 和 DP (data parallelism) 是什麼且他們分別有什麼優點和缺點，同時我們也會介紹由 Microsoft 所開發的 DeepSpeed, deep learning optimization library，是透過什麼樣的方式來對 GPU 記憶體甚至是訓練速度來進行優化！
Where Did All the Memory Go?
在這個章節我們會介紹究竟我們在訓練一個模型時，都是誰在佔據我們的 GPU 記憶體，而我們會分成兩個項目作介紹，分別是 Model States 以及 Residual Memory Consumption。
Model States
在我們訓練一個模型的時候，大部分的 GPU 記憶體都會被 Model States 所佔據，而 model states 包括了optimizer states, gradients 和 parameters。
舉例來說，當我們使用 adam optimizer 作為訓練時的 optimizer 時，我們便需要儲存每個 parameter 本身的 variance 以及 momentum，因此可以想像的是當我的模型 parameter 量有 10MB 的時候，我便需要額外儲存 20MB 的 optimizer state。
Residual Memory Consumption
而除了 model states 之外，其他需要佔據 GPU 記憶體的我們皆稱之為 residual memory consumption，其中包括了 activations, temporary buffers 以及 fragmented memory，由於這三個名字可能比較不常聽到，因此接下來會介紹以下這三個分別代表了什麼意思。
Activations：
在訓練模型時，除了獲得最後的 output 之外，我們同時會將 forward 過程中所有的 activations (e.g., 經過一個 conv layer 的 output，經過一個 relu layer 的 output，甚至是經過一個加法運算的 output) 都儲存在記憶體當中，原因在於我們計算模型梯度時是透過 chain rule，因此我們需要將過程中走過的所有路 (activations) 都儲存下來，才能沿著走過來的路透過 chain rule 的方式，一路算出最一開始的 gradients。
因此可以想像的是當我們將 input 餵入一個模型當中，我們會經過非常多的運算，而我們又需要把每個運算結果都儲存下來，這樣所需要的 GPU 記憶體會有多麼的龐大！
Temporary buffers：
在現今訓練模型的過程當中，我們多會選擇使用 ”多 GPUs” 來訓練模型，原因在於這樣可以增大我們的 batch size，使得我們在 batch normalization 甚至是更新的過程更加穩定，甚至也可以減少訓練模型的時間。而在使用多 GPUs 訓練模型的時候，我們會需要考慮到很多的東西，舉例來說，我們是要把一個模型切分成很多子等分然後放在不同的 GPUs 當中 (model parallelism) 還是我們要把一個模型複製在每一個 GPUs 當中 (data parallelism)。
因此假如我們選擇將一個模型複製在每一個 GPUs 當中的方法 (data parallelism) 時，我們要如何更新模型的 parameters。因為每個 GPU 都會吃到不同的 input，所以對於“每份複製出來的模型的 parameters”我們都會計算出不同的梯度，因此在這個情況下，我們要做的事情就是將 ”每個複製出來的模型所計算出的 parameter gradient” 彙整起來，並且用彙整後的梯度統一更新每個 GPU 上的模型 parameters，而這樣彙整的方式，我們稱之為 gradients all-reduce。而對於 all-reduce 這樣的操作，我們通常會在彙整前先將他們放入一個 buffer 當中，藉此改進 throughput，提升整體效率。因此我們想像得到當我們模型的 parameters 越多時，我們所需要的 buffer size 就會越大，這也會使得我們需要更多的 GPU 記憶體來存放！
Fragmented memory：
而在訓練的過程當中，也有可能發生 GPU 記憶體明明夠，但卻顯示 OOM 的情況。而這樣的情況通常是因為 memory fragmentation，也就是說記憶體當中沒有連續的記憶體區段可以分配。而會造成這個問題的其中一個是因為在訓練過程中所產生的 short lived memory objects 和 long lived memory objects 所導致。
DeepSpeed
DeepSpeed 是由 Microsoft 所開發的 deep learning optimization library，主要是希望可以 optimize GPU 記憶體的使用效率，同時也提高運算和資訊傳輸的效率。為了達到這個目的，DeepSpeed 提出了 ZeRO — Zero Redundancy Optimizer，而在 ZeRO 當中又可以分成 ZeRO-DP 以及 ZeRO-R 來針對 model state 以及 residual memory consumption 來進行優化。
因此在這個章節當中，我們會先了解何謂 MP (model parallelism) 以及 DP (data parallelism)，並知道目前我們所常用的 deep learning package 對於這兩個 parallelism 方式有何缺點。接下來我們首先會介紹 ZeRO-DP 是透過何種方式來優化 model state 的記憶體需求和效能，再來會介紹 ZeRO-R 是透過何種方式來優化 residual memory consumption。
Data Parallelism and Model Parallelism
在訓練一個龐大的模型或是訓練在一個龐大的資料集的時候，我們會選擇使用多張 GPUs 來進行訓練，而使用的方式多是 DP (data parallelism) 或是 MP (model parallelism)，接下來我們會了解什麼是 DP 以及什麼是 MP。
Data Parallelism：
Data parallelism 又稱作資料平行，所以我們可以從他的名字當中猜到，DP 希望將很多很多的資料同時平行在多張 GPUs 當中，要達到這個目的 DP 首先將模型複製了很多份，並且在每張 GPU 當中都放了一個複製的版本，如此以來當我今天有很多很多的資料要餵入，我便可以將資料拆成很多很多小等分，並且各自餵入不同的 GPUs 當中，最後再將多個不同的 outputs 合併起來。我們可以用下圖了解 DP 的運作：
而在最後計算梯度時，每個 GPU 當中的 ”複製模型“ 皆會計算出各自的梯度，接著 DP 會將每個 GPUs 不同的梯度計算出平均梯度，並且用平均梯度來更新模型 parameters。
因此 DP 本身的缺點便在於其 “複製了很多份的模型”，這也導致我們在訓練的過程會同時有很多份的 model state memory (e.g., optimizer state, model parameters, 和 parameters gradients)，理論上 N 張 GPUs 便會需要 N 倍的 GPU 記憶體需求，這樣的方式對於 GPU memory 非常不友善。舉例來說，當我們有一個模型非常龐大，龐大到整張 GPU 都放不下這個模型的 parameters，在這個情況下，就算我們有 100 萬張 GPUs，都沒有辦法透過 DP 來訓練這個模型。
但 DP 的缺點卻也正是他的優點，由於我們是透過 “複製模型” 的方式來訓練模型，因此在一次 forward 或是 backward 的過程當中，我們可以確保所有的運算是在 “相同的 GPU 當中”，且避免資訊在不同的 GPUs 之間傳遞，這使 DP 具有良好的 compute/communication efficiency。
Model Parallelism：
Model Parallelism 又稱作模型平行，顧名思義便是希望模型本身可以平行化，也就是説我會將一個模型拆成很多的子等分，並且每張 GPU 上皆會放上模型的一小等分，因此當我們有一份 input data 的時候，我們會先在第一張 GPU 進行運算，並且將第一張 GPU 的 output 再放入第二張 GPU 作為 input ，直到整個模型的所有子等分都運算完，並且得到最後的 output。那假如我們的模型是具有 “branch” 的話，我們便可以將相同的 input 放入不同的 GPU 當中進行不同的運算，並且將各自得到不同的 output 再餵入之後的 GPU 裏面做相加或是相乘的運算。我們可以用下圖了解 MP 的運作：
因此 MP 本身的缺點便在於其額外的 communication 成本，也就是説在一次的 forward/backward 過程中，其需要不斷的將 input/output 傳遞於不同的 GPUs 之間。
但是 MP 的優點便是其不需要將模型複製多份，這樣的方式可以節省很多的 GPU memory，尤其是當一個模型的 parameters 放不進一張 GPU 的時候，MP 便可以很有效的將模型切個成多個等分，並且各自的放在不同的 GPU 當中。
ZeRO-DP
DeepSpeed 透過 ZeRO-DP 來優化 model state (e.g., parameters, gradients, 和 optimizer state) 的記憶體需求。在剛剛我們了解了 MP 以及 DP 的優點和缺點，因此我們可以知道對於訓練速率來說，DP 最為理想，但是其需要非常大量的 GPU 記憶體卻也讓人望而生怯，因此 ZeRO-DP 便是希望針對 DP 本身的這個缺點進行優化，也就是説 ZeRO-DP 可以保有 MP 的 GPU 記憶體使用量同時也保持著 DP 的 communication efficiency。
要達到這樣的目的，ZeRO-DP 將原先複製多個 model state 的方式改成為分割 model state 成多個小部分，當需要時使用到某個 model state 時，才會透過 dynamic communication schedule 的方式將所需要的 model state 複製到各個 GPUs 當中，這樣的方式既可以避免 GPU 記憶體隨著 GPU 數量線性成長的瓶頸，同時也可以保持著接近於 DP 本身的 communication efficiency。
因此對於 ZeRO-DP 來說，總共可以分成三個部分，分別是
- Optimizer State Partitioning
- Gradient Partitioning
- Parameter Partitioning
接下來，我們針對每個部分來進行更詳細的介紹：
Optimizer State Partitioning
假設我們的 DP 具有 N 個 degree (data processes)，optimizer state partitioning 顧名思義就是根據這 N 個 degree 來將 optimizer state 切割成 N 個等分，並且將各個等分存放在各個 degree 當中。但是跟原先的 DP 不同的地方在於，對於存放在不同等分的 data processes，其除了只需要存放一小等分的 optimizer state 之外，同時在更新 parameters 的時候其也只需要更新部分的 parameters 和部分的 optimizer state，因此對於 optimizer state paritioning 並不會產生相較於原先 DP 多餘的 communication 負擔，且可以避免隨著 GPU 數量線性成長的 optimizer state GPU 記憶體需求。
而在每個 data process 更新完其所需要的 parameters 和 optimizer state 之後，ZeRO-DP 會再透過 all-gather 來將所有的 parameters 進行更新且同步。
Gradient Partitioning
由於在 optimizer state partitioning 當中，每個 data process 都只會需要更新部分的 parameters，因此對於每個 data process 來說，存放整個模型的 gradient 也因此顯得很冗餘 (因為就算存了整個模型的 gradient，也只會用其中部分的 gradient 來更新部分的 parameters 而已)。因此 gradient partitioning 便如同 optimizer state partitioning 一般，將整個模型的 gradient 切割成 N 個等分，並且存放在不同的 data processes 當中，如此一來，也可以避免隨著 GPU 數量線性成長的 gradient GPU 記憶體需求。
而要有效率的將 gradients 拆成多個等分並且存放在不同的 data process 當中時，DeepSpeed 有使用了 Reduce-Scatter operation，同時也使用了 bucketization strategy，也就是 bucketize 所有應該要存放在一起的 gradient，並統一使用 reduction operation 使得這樣的過程更加的有效率。
Parameter Partitioning
這邊的話大家可能也可以猜出來，parameter partitioning 便是將原先複製整個模型當中的 parameters 到每個 data process 的過程，變成是在每一個 data process 當中只存放部分的 parameter 而已。
但是這邊會出現一個問題，就是在 forward/backward 的過程中，理論上應該是每一個 data process 都會需要用到模型當中所有的 parameters，那要如何進行 parameter partitioning 呢？因此在 DeepSpeed 當中，當需要某個 data process 需要用到的 parameters 並不存在於他本身，該 data process 便會以 broadcast 的方式從存放需要的 parameters 的 data process 那邊接收需要的 parameters 過來。
看到這邊，可能大家又會有一個問題，這樣一天到晚傳來傳去難道不會增加很多的 communication 成本並讓整個訓練過程需要更多的時間嗎？因此在 DeepSpeed 當中有根據這個問題作解釋：This approach only increases the total communication volume of a baseline DP system to 1.5x, while enabling memory reduction proportional to N。也就是説雖然這個方式會增加額外的 communication 成本，但是相比於節省了 N 倍 (N 為 GPU 數量) 的 GPU 記憶體來說還是非常划算的！
ZeRO-R
介紹完了 ZeRO-DP，我們了解到 DeepSpeed 是透過什麼方式來減少 model state (e.g., parameters, gradients, 和 optimizer state) 的 GPU 記憶體需求，接下來我們將了解 DeepSpeed 是如何透過 ZeRO-R 來減少 residual memory consumption (e.g., activations, and temporary buffers) 的 GPU 記憶體需求。
對於 ZeRO-R 來說，總共可以分成三個部分，分別是：
- Partitioned Activation Checkpointing
- Constant Size Buffers
- Memory Defragmentation
接下來，我們針對每個部分來進行更詳細的介紹：
Partitioned Activation Checkpointing
在上面我們理解到 MP (model parallelism) 將 model 切割成多個子部分，並且放置在不同的 GPUs 當中，然而，MP 在設計時，便會將整個 activations 複製到每一個 GPU 當中，這使得 MP 雖然節省了複製 model 的 GPU 記憶體，但是 activations 的複製卻依舊需要蠻大量的 GPU 記憶體。
因此，ZeRO-R 以將 activations 分割成多個子等分，直到需要使用該 activations 後 (e.g., backpropagation) 才複製至各個 GPUs 的方式取代了一開始就將所有 activations 複製於所有 GPUs 的方式減少了大量的 GPU 記憶體需求。因此，當需要使用 activations 時，ZeRO-R 便透過 all-gather 的方式將 activations 複製到各個 GPUs 當中。
而將 activation 進行 partition 的方式，是可以跟 activations checkpointing 同時使用，使得 GPU 記憶體需求更加減少。
Activations checkpointing: 只儲存少部分的 activations，而在 backward 時，當需要某個特定卻沒有存下來的 activations 時，便會重新對於該 block 進行 forward，來得到所需要的 activations，由於不需要存下所有的 activations，因此可以節省大量的 GPU 記憶體，但是因為需要重新計算部分的 forward，因此會使得整體的運算時間變得較長，可以想像這是一個以時間換取空間的方法！
然而，雖然上面的方法看似可以使得 GPU 記憶體更加節省，但是對於某些非常巨大的模型來說，這樣的方法依舊使得 GPU 記憶體的需求非常龐大，因此 ZeRO-R 同時也提供將部分的 activations 從 GPU 記憶體 offload 到 CPU 記憶體當中，使得 GPU 記憶體的需求可以再更加的減少！
Constant Size Buffers
在訓練一個模型的時候，部分 libraries 會在進行某些運算時 (e.g., all-reduce)，先將所需要的 parameters 放入一個 buffers，使得整體的運算效率可以更加的提升。然而，這樣的方式雖然可以提升效率，但是也會使得 buffers 的大小隨著 input data 的大小或是模型的大小，而跟著變大，而造成額外的 GPU 記憶體的負擔。
在 DeepSpeed 當中，同樣也有使用 buffer 來增進整體訓練的效率，但是相較於隨著模型和 input data 無限增大的 buffer 來說，DeepSpeed 使用了 constant size buffers，也就是說當模型超出一定的大小之後，便不會使得 buffer 的大小跟著一直變大，而會維持在一個 constant，這樣的方式使得 DeepSpeed 可以保持一定的效率，卻不會對 GPU 記憶體造成龐大的負擔。
Memory Defragmentation
造成 memory fragementation 的主要是在訓練一個模型時所產生的 short lived memory 和 long lived memory 不斷交錯產生所造成的。舉例來說，在計算 gradients 時，parameter gradients 往往會存在很久直到我們透過 optimizer 更新參數，而用來進行 chain rules 的 activation gradients 卻往往只會存在一下下，這便是 long lived memory (parameter gradients) 和 short lived memory (activation gradients)。
在我們有非常大的記憶體的時候，memory fragementation 往往不會產生什麼問題，但是當我們的模型非常龐大且能使用的記憶體非常有限的時候，便很容易產生以下兩個問題：
- 具有足夠的記憶體，但由於都不連續，而產生 OOM (out of memory)
- Memory allocator 時常耗費非常大量的時間在找尋連續的記憶體，使得整體訓練效能降低
因此，ZeRO-R 解決 memory fragmentation 的方式便是針對 long lived memory 和 short lived memory 預先 allocate 一整個的連續的記憶體區塊。當 long lived memory 或是 short lived memory 產生時，便會直接將之複製到預先分配好的記憶體區塊當中，這樣的方式除了可以使得我們可以在有限的記憶體當中訓練更大的模型或是使用更大的 batch size，還可以提升整體訓練的效能。
Conclusions
在這篇文章當中，主要介紹了平時我們在訓練模型時，GPU 記憶體都被什麼東西佔用了，同時也稍微介紹了當今大家最常使用的 MP (model parallelism) 和 DP (data parallelism) 的方法和他們的優缺點，最後我們介紹了 DeepSpeed, a deep learning optimization library，是如何透過 ZeRO-DP 和 ZeRO-R 的方式來減少整體訓練所需要的 GPU 記憶體需求，並且增進整體訓練的效能。
最後，希望就算對於 DeepSpeed 沒有興趣的讀者，在閱讀這篇文章之後也可以對 ”訓練” 這件事情更加的了解，同時也非常推薦希望了解更多的讀者可以去看看 DeepSpeed 本身的論文，一定可以收穫良多的！
