前言
目前比較常見的並行訓練是資料並行，這是基於模型能夠在一個GPU上儲存的前提，而當這個前提無法滿足時，則需要將模型放在多個GPU上。現有的一些模型並行方案仍存在許多問題，本文提出了一種名為PatrickStar的異構訓練系統。PatrickStar通過以細粒度方式管理模型資料來更有效地使用異構記憶體，從而克服了這些缺點。
本文附上了PatrickStar的使用示例。PatrickStar與模型定義無關，在PyTorch指令碼上新增幾行程式碼可以帶來端到端的加速。
本文來自公眾號CV技術指南的
關注公眾號CV技術指南 ，專注於計算機視覺的技術總結、最新技術跟蹤、經典論文解讀。
論文：PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
程式碼：https://github.com/Tencent/PatrickStar
使用示例
先介紹一個大家比較關心的問題--復不復雜，好不好實現。
PatrickStar與模型定義無關，在PyTorch指令碼上新增幾行程式碼可以帶來端到端的加速。
Background
現在人工智慧的共識是採用PTMs(Pre-Trained Models)作為任務的骨幹，而不是在與任務相關的資料集上從頭開始訓練模型。PTMs的高效能伴隨著眾多的引數，這對計算和儲存資源提出了巨大的要求。
由於大型的模型資料在單個GPU的記憶體無法儲存，所以最常用的資料並行技術不適用於PTM（比較常見的例子是2012年的AlexNet）。通過利用並行訓練在多個GPU記憶體之間分配模型資料，例如ZeRO、模型並行和流水線並行，最新進展使得PTM規模的增長成為可能。SOTA解決方案是將它們組合成3D並行，它可以在數千個GPU上將PTM擴充套件到數萬億個引數。
創新思路
作者觀察到在PTM訓練期間必須管理的兩類訓練資料：模型資料由引數、梯度和優化器狀態組成，它們的足跡與模型結構定義相關；非模型資料由operator生成的中間張量組成。非模型資料根據訓練任務的配置(如批次大小)動態變化。模型資料和非模型資料相互競爭GPU記憶體。
現有的解決方案在不考慮非模型資料的情況下，在CPU和GPU記憶體之間靜態劃分模型資料，並且它們的記憶體佈局對於不同的訓練配置是恆定的。這種靜態的分割槽策略導致了幾個問題。
首先，當GPU記憶體或CPU記憶體不足以滿足其相應的模型資料要求時，即使當時其他裝置上仍有記憶體可用，系統也會崩潰。其次，當資料在張量粒度的不同記憶體空間之間傳輸時，通訊效率低下，當可以將模型資料提前放置到目標計算裝置上時，CPU-GPU通訊量是不必要的。
針對上述問題，本文提出了一種名為PatrickStar的異構訓練系統。PatrickStar通過以細粒度方式管理模型資料來更有效地使用異構記憶體，從而克服了這些缺點。
作者將模型資料張量組織成塊，即大小相同的連續記憶體塊。在訓練期間，根據組塊的張量狀態來動態編排組塊在異構儲存空間中的分佈。通過重用不共存的塊，PatrickStar還比SOTA解決方案進一步降低了模型資料的記憶體佔用。
作者使用預熱迭代(warm-up iteration)在執行時收集模型資料的可用GPU記憶體的統計資訊。設計了一種高效的塊回收策略和基於統計資料的裝置感知操作符放置策略，以減少CPU-GPU的資料移動量。使用零冗餘優化器，基於塊的記憶體管理可以通過塊的GPU內通訊(intra-GPU communication)有效地與資料並行共生。
Contributions
1.從頭開始構建了一個基於塊的記憶體管理DNN訓練系統，名為PatrickStar。與現有的異構訓練方法相比，該方法通過縮小非模型記憶體空間、提高記憶體效率和降低CPU-GPU通訊量，支援更大的模型規模和更高的計算效率。
2.基於塊的管理與零冗餘優化器資料並行自然是共生的。基於塊的集體通訊模式降低了GPU內頻寬需求和提高了頻寬利用率。
3.作者在擁有8x V100 GPU、240 GB和120 GB DRAM記憶體CPU的雲端計算節點上對系統進行了評估。在240 GB記憶體的情況下，PatrickStar訓練了一個120億引數的類似GPT的模型，這是DeepSpeed最大模型規模的1.5倍。在120 GB記憶體的情況下，PatrickStar的模型比例是DeepSpeed的4倍。
4.PatrickStar的計算效率比DeepSpeed更高，在8x GPU上實現了超線性伸縮。
Design Overview
論文設計了一個並行PTM訓練系統PatrickStar，它的工作原理如下圖所示。
PatrickStar通過將模型資料分塊管理並儲存在異構空間中，提高了現有異構訓練的模型規模和效率。圖中的圓圈表示引數的元素，它們在記憶體中以塊的形式排列。當操作符的計算被觸發時(右圖的彩色部分)，PatrickStar將引數所在的塊安排到所需的計算裝置上。
論文提出了在異構記憶體空間中高效編排塊的優化方案。此外，它還可以與零冗餘優化器相結合，擴充套件到多個GPU。必要時，在訓練過程中將組塊移動到所需裝置。
System design on a single GPU
PatrickStar作為PyTorch和異構記憶體之間的中介軟體，如上圖所示。
系統由在預處理階段工作的靜態模組和在訓練階段工作的執行時模組組成。PatrickStar的靜態模組在訓練前進行處理。
基於神經網路結構，構造了張量與組塊之間的對映模式。
在訓練過程中，執行時模組通過將張量重定向到所管理的基於塊的儲存空間來接管PyTorch的記憶體訪問，並使用塊管理器來智慧地管理異構儲存空間中的塊。
Preprocessing Stage
在訓練開始之前，在預處理階段為模型資料的每個張量分配一塊塊空間，並生成塊-張量對映模式。一個有效的對映模式需要具備以下三個特徵：1.增加張量訪問的區域性性。2.降低峰值記憶體消耗。3.對並行友好。
論文提出了一個高效的對映模式。
簡單來說，根據模型資料中張量的型別，將組塊分為四種型別，即引數FP16列表、引數FP32列表、動量列表和方差列表，共14M位元組(M為引數數)。組塊具有相同的大小，因此不同的組塊可以重用相同的記憶體，並且有利於並行訓練中的集體通訊。特別的是，PatrickStar沒有分配分級FP16列表。梯度FP16張量可以重用引數FP16列表的塊空間，消除了梯度FP16張量對引數FP16張量的依賴。
與最小模型資料記憶體佔用為16M位元組的ZeRO-Offload相比，PatrickStar減少了記憶體佔用。此外，它還分配了額外的GPU記憶體，用於儲存要移動到CPU的漸變，而PatrickStar則消除了這一開銷。因此，與其他PTM訓練方案相比，PatrickStar的模型資料儲存空間最小。
Training Stage
在訓練過程中，PatrickStar需要正確高效地編排異構儲存空間中的資料塊。論文介紹了導塊正確移動的機制，並介紹了一種優化策略以提高其效率。此機制比較複雜，論文有非常具體的介紹，這裡不予解讀。
Scaling to Multiple GPUS
PatrickStar使用多程式在多個GPU之間執行資料並行化。假設程式數為nproc。每個程式負責單個GPU，而所有程式共享CPU。程式的異構記憶體空間由GPU的全部記憶體和1/nproc的CPU記憶體空間組成。該程式以與零冗餘優化器相同的方式管理位於其本地異構儲存空間中的總塊的1/nproc。本地空間中的塊稱為本地塊，而不在本地空間中的塊稱為遠端塊。
如圖所示，通訊組由塊列表的nproc連續塊組成，其中每個塊屬於不同的程式。根據前面的塊張量對映模式，我們可以設計一種以最小的程式間通訊量實現資料並行的通訊方案。程式只需要在FWD和BWD階段傳送引數FP16和梯度FP16。需要最大量資料(包括動量、方差、引數FP32和梯度FP16)的ADAM階段在本地執行。
論文中有非常詳細的過程介紹與演算法，這裡不予解讀。
與相關工作相比，PatrickStar獲得了更低的GPU內頻寬需求和更高的頻寬利用率。根據成本模型，PatrickStar的頻寬要求為2(p−1)/p×2m(all-gathers)+(p−1)/p×2m(reduce-scatter)=6(p−1)/p×M，其中p為並行度，M為引數個數。在ZeRO-Offload和ZeRO-DP中，每層的引數由單個程式擁有並廣播給其餘程式。與all-gathers相比，廣播將資料傳輸集中在單個GPU上，並未充分利用聚合頻寬。
基於廣播的頻寬要求為4(p−1)/p×2M(廣播)+(p−1)/p×2M=10(p−1)/p×M，比PatrickStar提高了2/3。雖然Zero-Infinity也採用了all-gathers的方式，但PatrickStar的頻寬利用率仍然較高。
已經證明，將序列張量佈置為大塊(bucket)來傳輸的分組化策略可以導致更高的頻寬利用率，因為它將在每次通訊中傳輸更多的資料。PatrickStar中基於塊的方法自然是分塊的，而ZeRO-Infinity的傳輸單元是張量。PatrickStar進一步避免了資料複製開銷，以提高效能。
Optimization
將模型資料託管在GPU記憶體中相比，異構訓練方法引入了額外的CPU-GPU資料移動開銷。對基於塊的記憶體管理進行了優化，使其更加高效和強大。
首先，它可以以細粒度的方式佈局運算子，使記憶體密集型運算子不在他們的首選裝置上，但它可以減少資料移動量，提高系統的整體效率。為此，PatrickStar提出了一種Device-aware Opetator Placement優化。
其次，當資料不能永久駐留在運營商的計算裝置上時，如果不使用，資料塊就會被逐出。除了在更大的模型規模上進行改進外，PatrickStar還通過塊逐出策略將塊逐出量降至最低。
為了實現上述優化，我們必須在訓練過程中很好地瞭解每臺裝置上有多少GPU記憶體空間可以分配給塊，這在本節中稱為可分塊記憶體(Chunkable Memory)。因此，PatrickStar提供了在執行時收集可分塊記憶體統計資訊的方法。
這些方案細節均在論文有詳細介紹。
Conclusion
PatrickStar通過縮小非模型記憶體空間、提高記憶體效率和降低CPU-GPU通訊量，支援更大的模型規模和更高的計算效率。
PyTorch、DeepSpeed和PatrickStar在同一GPU上的訓練吞吐量。OOM表示超記憶體。
在多個GPU上使用DP訓練PyTorch、DeepSpeed和PatrickStar的吞吐量。
論文提出了一種創新的異構訓練系統，稱為PatrickStar。它將模型資料組織成塊，從而在異構儲存空間中更靈活地編排它們。該系統與零冗餘優化器資料並行是共生的。PatrickStar成功降低了PTM培訓的硬體要求，並更高效地擴充套件到多個GPU。
在雲端計算平臺的8xGPU 240 GB CPU節點上，它在最大模型規模上比SOTA提高了2倍，速度也比SOTA快。
未來的一些工作可以在PatrickStar上進行研究。
首先，通過對非模型資料檢查點和解除安裝策略的聯合優化，可以得到更好的塊清除方法。其次，基於組塊的異構方法可以與其他並行訓練方法相結合進行多節點伸縮。
關注公眾號CV技術指南 ，專注於計算機視覺的技術總結、最新技術跟蹤、經典論文解讀。
在公眾號中回覆關鍵字 “入門指南“可獲取計算機視覺入門所有必備資料。
其它文章