深度學(xué)習(xí)相較傳統(tǒng)機(jī)器學(xué)習(xí)模型，對(duì)算力有更高的要求。尤其是隨著深度學(xué)習(xí)的飛速發(fā)展，模型體量也不斷增長(zhǎng)。于是，前幾年，我們看到了芯片行業(yè)的百家爭(zhēng)鳴和性能指標(biāo)的快速提升。正當(dāng)大家覺得算力問題已經(jīng)得到較大程度的緩解時(shí)，大語言模型（LLM, Large language model）的興起又帶來了前所未有的挑戰(zhàn)。當(dāng)網(wǎng)絡(luò)模型達(dá)到一定量級(jí)后（比如參數(shù)量達(dá)到10B級(jí)別），表現(xiàn)出In-context learning，Instruction following和Step-by-step reasoning等涌現(xiàn)能力（Emergent abilities）。這些能力是以往模型所不具備的，因此LLM引起了學(xué)術(shù)界與工業(yè)界的濃厚興趣以及廣泛關(guān)注。但是，得到這些能力的代價(jià)是模型的急劇膨脹。比較知名的如GPT-2（1.5B），Megatron-LM（8.3B），T5（11B），Turing-NLG（17B），LLaMA（65B），GPT-3（175B），Jurassic-1（178B），Gopher（280B），PaLM（540B），Switch Transformer（1.6T），參數(shù)量很快從B時(shí)代到達(dá)了T時(shí)代。如今，在這些“龐然大物”面前，那些歷史上曾經(jīng)的“巨無霸”，像GPT（110M）、BERT（340M）看起來都顯得有些迷你了。根據(jù)論文《Machine Learning Model Sizes and the Parameter Gap》中提到的，從1950到2018年，語言模型的模型大小增長(zhǎng)7個(gè)數(shù)量級(jí)，但從2018年到2022年，4年就暴漲5個(gè)數(shù)量級(jí)。網(wǎng)絡(luò)模型的陡然增大，隨之而來的是一系列數(shù)據(jù)集、訓(xùn)練方法、計(jì)算框架相關(guān)的問題。關(guān)于這些年大模型發(fā)展的回顧以及對(duì)這些問題業(yè)界的探索的闡述可參考論文《A Survey of Large Language Models》。LLM所涉及的東西非常多，這篇文章主要討論的限于對(duì)計(jì)算框架帶來的挑戰(zhàn)及如何用分布式的方法解決。

從網(wǎng)絡(luò)模型結(jié)構(gòu)方面，好消息是自從2017年Google在《Attention Is All You Need》中提出Transformer后，近幾年模型的結(jié)構(gòu)開始趨同于它。它不僅在于效果好，而且相比之前的如RNN結(jié)構(gòu)更便于計(jì)算。Transformer網(wǎng)絡(luò)中的主要負(fù)載是MHA和FF，而它們本質(zhì)上都是或大部分都是GEMM。而GEMM的計(jì)算是相對(duì)成熟的。這就將前沿的AI與成熟的計(jì)算理論與實(shí)踐結(jié)合起來了。另一方面Transformer是相同模塊的重復(fù)堆疊，這也給并行計(jì)算（如pipelining）提供了便利。因此，近幾年主要的LLM也都是基于Transformer。

那模型變大后對(duì)計(jì)算框架會(huì)帶來什么樣的影響呢？最直接的就是memory不夠用了。由于傳統(tǒng)CPU對(duì)神經(jīng)網(wǎng)絡(luò)這種調(diào)度并行的workload效率不高，因此目前主流的都是用GPU或者各種加速器。但是，這些芯片的memory比CPU的memory要小很多。

在模型的訓(xùn)練過程中，通常需要參數(shù)量數(shù)倍的memory。其中，消耗memory的大戶一般有這么幾個(gè)：

• Weight：模型的權(quán)重參數(shù)
• Activation：中間層的輸入與輸出
• Gradient & Optimizer state：梯度及優(yōu)化器狀態(tài)（如Adam中的momentum和variance）

其它還有些如temporary buffer或者fragmented memory等由于占得不多，因此優(yōu)化時(shí)放的精力也少。推理的話optimizer state就不用了，但weight與activation是訓(xùn)練與推理時(shí)都逃不了的。前面提到模型的size在短短數(shù)年增長(zhǎng)了幾個(gè)數(shù)量級(jí)，但相比之下，GPU的memory增長(zhǎng)大概只有幾倍。這讓本就不太夠用的GPU memory顯得更加捉襟見肘。論文《ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning》中分析了訓(xùn)練過程中的memory需求。它將weight，gradient和optimizer state歸為model state?；旌暇群褪褂肁dam優(yōu)化器情況下，單個(gè)參數(shù)應(yīng)對(duì)model state需要用到20 bytes。這樣，對(duì)于大模型的memory需求遠(yuǎn)遠(yuǎn)超過了單卡GPU（A100的話是80G）的承受范圍。對(duì)于100 B參數(shù)量的模型，光model state就需要64 GPU，如果是1T參數(shù)量的模型則需要512 GPU。論文《ZeRO: Memory Optimization Towards Training A Trillion Parameter Models》中提到參數(shù)量1.5 B的GPT-2光是model state就需要至少24 G。另外activation需要8G左右。這意味著一張V100基本就撐爆了。論文《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》中指出530 B的模型需要超過10 T的memory用于model state。Activation則會(huì)占用大約16.9 T（通過gradient accumulation優(yōu)化可減少到8.8 G）。此外，為了減少計(jì)算量使用KV cache的話，也會(huì)增加不少的memory開銷。這使得大模型計(jì)算的memory問題變得非常嚴(yán)重，很多時(shí)候它直接決定模型是否能跑得起來。

業(yè)界從多個(gè)角度提出方案來緩解memory的問題，如：

• Activation checkpointing / rematerization
• Heterogeneous Memory Management
• Compression
• Mixed precision training
• Memory efficient optimizer
• Live analysis & Reuse

本文主要關(guān)于分布式的方法，因此對(duì)上面這些點(diǎn)不作展開。減少內(nèi)存使用一定程度上解決“跑不起來”的問題。跑起來后，就需要考慮“跑得更快”的問題。畢竟大家都在折騰大模型，發(fā)展非常迅速，能早些得到結(jié)果就意味著領(lǐng)先。一個(gè)好的想法，如果實(shí)驗(yàn)它需要幾年，那等結(jié)果出來黃花菜都涼了。這就要求框架在計(jì)算大型AI模型時(shí)還需要提高硬件效率，最終提升訓(xùn)練效率。

進(jìn)入正題前，先過一些相關(guān)概念。

為了更快得到結(jié)果，直觀的想法是會(huì)投入更多的資源來擴(kuò)展系統(tǒng)，而系統(tǒng)的擴(kuò)展可分為兩種方式：scale up和Scale out。通俗點(diǎn)說，前者是為單個(gè)計(jì)算節(jié)點(diǎn)添加更多的資源（如CPU，memory等）；后者是添加更多的計(jì)算節(jié)點(diǎn)。理想情況下，我們希望吞吐是隨著資源的投入線性增長(zhǎng)的，雖然現(xiàn)實(shí)情況中，由于同步、通信開銷等原因總會(huì)有些開銷。要達(dá)到近線性的加速，就需要算法與系統(tǒng)是易被并行的，而且硬件資源的使用平衡。

在分布式計(jì)算下，這些可并行的任務(wù)就可以放到多卡或多機(jī)上。這就涉及到通信問題。否則額外計(jì)算資源帶來的性能收益可能還補(bǔ)不上通信開銷。硬件上，可以通過NVLink，PCIe等（可參見論文《Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect》）。多機(jī)可以通過Ethernet，GDR（GPU Direct RDMA）等機(jī)制。為了使用更方便，Nvidia提供了NCCL通信庫(kù)利用多種通信機(jī)制向上提供諸如AllReduce、AllGather、ReduceScatter等常見的collective primitive。其它類似的通信庫(kù)還有OpenMPI、Gloo。

在講分布式并行前，先提下一些并行相關(guān)概念。首先workload得并行起來，才能分布到多卡或者多機(jī)上。并行是分布的基礎(chǔ)。關(guān)于在深度學(xué)習(xí)領(lǐng)域的并行可以參見2018年的綜述文章《Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis》。在此不展開。

關(guān)于并行帶來的加速影響，涉及兩個(gè)重要定律，即Amdahl和Gustafson定律。Amdahl定律可能是計(jì)算機(jī)領(lǐng)域最廣為人知的定律之一，它定量地給出了提升系統(tǒng)的部分性能對(duì)整個(gè)系統(tǒng)性能的影響。它有一個(gè)假設(shè)就是問題的規(guī)模是固定的。論文《Reevaluating Amdahl’s Law》中指出現(xiàn)實(shí)中，人們一般是為了解決更大規(guī)模問題而增加更多資源。Gustafson定律以此為假設(shè)給出了增加計(jì)算資源對(duì)系統(tǒng)性能的影響。這兩個(gè)定律從兩個(gè)不同角度闡述了加速比，處理器個(gè)數(shù)與串行比例之間的關(guān)系。前者強(qiáng)調(diào)了要降低程度的串行化比例，否則光摞計(jì)算資源性價(jià)比不高。后者強(qiáng)調(diào)當(dāng)并行化部分足夠大，加速比是可以隨著計(jì)算資源線性增長(zhǎng)的。

對(duì)于大模型的計(jì)算，目前流行的分布式并行方法可以大體分為幾類：

Data parallelism

AI模型在訓(xùn)練時(shí)一般都會(huì)將多個(gè)樣本數(shù)據(jù)組成mini-batch一起計(jì)算。DP利用了這一點(diǎn)，將mini-batch進(jìn)一步分成多個(gè)更小的mini-batch，再分給多個(gè)worker去處理。在這種模式下，每個(gè)worker都有完整的模型，包括weight，optimizer state等。每個(gè)worker都可以獨(dú)立完成整個(gè)模型的梯度計(jì)算，它們計(jì)算出來的梯度最終需要聚合到一起。

其實(shí)對(duì)于大模型的分布式計(jì)算框架的研究其實(shí)從很早就開始了。2012年Google的論文《Large Scale Distributed Deep Networks》中提到的DistBelief（也就是TensorFlow的前身）采用了DP和MP等并行技術(shù)，以及改進(jìn)的優(yōu)化器，可以用于在集群中訓(xùn)練參數(shù)超過1B的模型。只是由于當(dāng)時(shí)的局限性，它用的是CPU集群，而且是基于parameter server的中心式架構(gòu)。導(dǎo)致機(jī)器數(shù)一多后，PS成為通信瓶頸影響scalability。因此，后來越來越多使用AllReduce方法，因?yàn)榭梢猿浞掷盟衐evice的帶寬?？蓞⒁?#34;Massively Scale Your Deep Learning Training with NCCL 2.4"。

按同步方式還可以分為BSP（Bulk synchronous parallel）與ASP（Asynchronous parallel）兩種。前者會(huì)有同步開銷，后者沒有同步開銷，但無法保證收斂。因此，還有一種SSP（Stale synchronous parallel）方式是這兩者的折中。即每個(gè)worker允許使用stale的weight處理下一個(gè)mini-batch。這個(gè)staleness的閥值是可以定義的。

2017年論文《Integrated Model and Data Parallelism in Training Neural Networks》中將DP分為兩種：一種是Batch parallelism，即平時(shí)常見的那種；另一種是Domain parallelism。它將輸入圖片切分成多塊交給多個(gè)worker處理。這種方式中，每個(gè)worker與batch parallelism一樣也包含完整的模型參數(shù)，但只處理圖像的一部分。它在視覺網(wǎng)絡(luò)前幾層activation比較大的情況下有比較好的通信復(fù)雜度。

DP由于實(shí)現(xiàn)簡(jiǎn)單，在各種主流AI框架中基本都有支持。如PyTorch的DDP（DistributedDataParallel），可參見論文《PyTorch Distributed: Experiences on Accelerating Data Parallel Training》。TensorFlow中的tf.distribute.Strategy，可參見"Distributed training with TensorFlow"和"Multi-GPU and distributed training"。MXNet中可以參考"Run MXNet on Multiple CPU/GPUs with Data Parallelism"。

最基礎(chǔ)的DP的優(yōu)點(diǎn)除了實(shí)現(xiàn)簡(jiǎn)單外，它的計(jì)算與通信效率也很高。但是，它的缺點(diǎn)也很明顯：一是當(dāng)worker數(shù)增多時(shí)，mini-batch也會(huì)隨之變大，而過大的mini-batch會(huì)讓訓(xùn)練不穩(wěn)定甚至不收斂。此外，它最大的問題還在于，它無法解決單卡放不下大模型的問題。

為了解決DP無法減少memory footprint的問題，2019年的論文《ZeRO: Memory Optimization Towards Training A Trillion Parameter Models》提出了ZeRO（Zero Redundancy Optimizer）。ZeRO-DP用ZeRO擴(kuò)展了DP，在達(dá)到DP的計(jì)算與通信效率的同時(shí)還能達(dá)到MP的內(nèi)存效率。它通過將model state（parameter, gradient與optimizer state）在data parallel processes間進(jìn)行切分，從而消除memory redundancy，同時(shí)它使用dynamic communication schedule利用model state時(shí)間上的使用特性減少通信量。根據(jù)切分model state的程度不同，ZeRO可分為三級(jí)，或三個(gè)stage。即ZeRO-1切分optimizer state；ZeRO-2切分optimizer state與gradient；ZeRO-3切分optimizer state，gradient和parameter。第三級(jí)可以支持在1024張Nvidia GPU上訓(xùn)練T級(jí)別參數(shù)的模型。

而后，論文《ZeRO-Offload: Democratizing Billion-Scale Model Training》中基于ZeRO提出ZeRO-Offload改進(jìn)ZeRO-DP。它將optimizer state放在CPU memory中，同時(shí)每個(gè)process將自己的那份梯度也放到CPU memory中。ADAM優(yōu)化器的更新在CPU上完成，更新完了再放回GPU。再使用GPU上的all-gather收集所有更新好的參數(shù)。論文《ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning》中提出ZeRO-Infinity利用了NVMe memory進(jìn)一步改進(jìn)了ZeRO-Offload。這些屬于異構(gòu)內(nèi)存管理的范疇，不是討論重點(diǎn)，就不展開了。

它有點(diǎn)類似于后面要提到的TP，因?yàn)樗鼤?huì)切tensor分到多個(gè)worker上。但它與TP的顯著差別是TP是切完后在worker內(nèi)部算完再聚合計(jì)算結(jié)果，因此它需要對(duì)代碼有修改。而ZeRO只是利用多個(gè)worker來分布式地存tensor，在算之前會(huì)先聚合它們，算完了再把除了自己要維護(hù)的部分丟掉，從而減少memory占用?？偟脕碚f，MP會(huì)傳輸activation，而ZeRO不會(huì)。MP處理的切分的輸入，而ZeRO處理的完整的輸入。這樣的話基本不需要修改代碼。有些地方也將它區(qū)別于DP和TP，單獨(dú)作為一類來介紹。這里為了乘法，同時(shí)因?yàn)檎撐闹蟹Q為ZeRO-DP，因此放到DP這一節(jié)來一并提了。

ZeRO實(shí)現(xiàn)于Microsoft自家的框架DeepSpeed中。由于它的有效性，其思想也被其它的框架所吸收。比如PyTorch中的FSDP（Fully Sharded Data Parallel）就是受ZeRO與FairScale的Fully Sharded Data Parallel的影響。詳細(xì)可以參見"Introducing PyTorch Fully Sharded Data Parallel (FSDP) API"和"GETTING STARTED WITH FULLY SHARDED DATA PARALLEL(FSDP)"。

在DP中，由于每個(gè)data-parallel process都會(huì)有完整的模型，它們需要一致地被更新。這個(gè)weight的更新操作需要在所有的模型replica上進(jìn)行。這是一個(gè)冗余操作。2020年的論文《Automatic Cross-Replica Sharding of Weight Updates in Data-Parallel Training》針對(duì)這個(gè)問題，提出了shareded DP。該方法在XLA框架上實(shí)現(xiàn)，用于Cloud TPU。它利用all-reduce操作可分為reduce-scatter與all-gather兩個(gè)階段的特點(diǎn)，使用靜態(tài)分析與變換，使多個(gè)replica共同分擔(dān)weight update計(jì)算。

關(guān)于DP的優(yōu)化也有不少相關(guān)工作，如論文《Preemptive All-reduce Scheduling for Expediting Distributed DNN Training》中的PACE是一個(gè)communication scheduler，用于搶占式調(diào)度all-reduce操作，從而最大化計(jì)算與通信的overlapping；論文《Supporting Very Large Models using Automatic Dataflow Graph Partitioning》中的ByteScheduler也是用于提升了DP中的overlapping ratio。

Model parallelism

MP是一類將model state在多個(gè)device上切分的方法。它的優(yōu)點(diǎn)是可以讓多卡分擔(dān)memory占用，從而解決模型放不下單卡的問題。但缺點(diǎn)是要使切分的工作量平衡不太容易，計(jì)算間有依賴，activation需要傳輸導(dǎo)致通信量增加。

對(duì)于一個(gè)模型，切法有很多。有些是比較通用的，有些適用于特定網(wǎng)絡(luò)。因此MP也可細(xì)分為好幾類：

Tensor parallelism

又稱為operator parallelism，operator-level parallelism，或horizontal MP。它將單個(gè)算子在多個(gè)device間切分。通常用于單個(gè)算子的weight就占用大量GPU memory的情況，如embedding layer。

最早，在論文《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》中使用了TP來并行化連續(xù)的FF層（Transformer中的self-attention與feed-forward）。它利用矩陣的分塊相乘性質(zhì)將weight切分，從而將相乘的計(jì)算量分到多個(gè)worker上。由于目前流行的Transformer網(wǎng)絡(luò)其主要計(jì)算量為矩陣相乘。TP應(yīng)用于GEMM就可以看作并行矩陣相乘（Parallel GEMM）問題。并行矩陣相乘按并行的維度可分為：

• 1D：上面提到的Megatron中將weight在單個(gè)維度上切分，是1D的。
• 1.5D：如論文《Communication-Avoiding Parallel Sparse-Dense Matrix-Matrix Multiplication》
• 2D：如經(jīng)典的Cannon’s algorithm和SUMMA（Scalable Universal Matrix Multiplication Algorithm）。還有論文《An Efficient 2D Method for Training Super-Large Deep Learning Models》中的Optimus使用SUMMA來提升內(nèi)存與通信效率。
• 2.5D：如論文《Tesseract: Parallelize the Tensor Parallelism Efficiently》是基于SUMMA與2.5D Matrix Multiplication。論文《Communication-Optimal Parallel 2.5D Matrix Multiplication and LU Factorization Algorithms》提升了Cannon’s algorithm的效率。
• 3D：原始3D并行矩陣乘法可參見論文《A three-dimensional approach to parallel matrix multiplication. IBM Journal of Research and Development》。論文《Maximizing Parallelism in Distributed Training for Huge Neural Networks》做了改進(jìn)，消除了原算法中廣播數(shù)據(jù)的冗余。論文《Communication-Optimal Parallel Recursive Rectangular Matrix Multiplication》提出CARMA是一種達(dá)到asymptotic lower bound的遞歸算法。論文《Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication》提出COSMA（Communication Optimal S-partition-based Matrix multiplication Algorithm）使用red-blue pebble游戲來建模數(shù)據(jù)依賴，并推出I/O最優(yōu)schedule。

Colossal-AI中集成了多種并行矩陣相乘方法，有興趣可以參考。

論文《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》中提出的GShard模塊包含一組輕量級(jí)的annotation API和XLA compiler的擴(kuò)展。它采用automatic sharding的方法，用于有600 B參數(shù)的帶Sparsely-Gated Mixture-of-Experts的transformer網(wǎng)絡(luò)。

基于GShard，論文《GSPMD: General and Scalable Parallelization for ML Computation Graphs》中的GSPMD使用了tensor sharding annotations來以一種統(tǒng)一的方式達(dá)到不同的并行（DP, in-layer MP和spatial partitioning等）。切分的方式很靈活，不同層可以用不同的方法切分，同一層可以使用不同方法切分。GSPMD泛化了其后端GShard，以partial tiling和manual subgroups擴(kuò)展了sharding representation，實(shí)現(xiàn)了像priority sharding propagation, recursive partitioning和pipepiling等特性。該方法可用于上千個(gè)TPUv3 cores上的視覺、語音和語言模型。

Pipeline parallelism

又稱為vertical MP。與TP切分單個(gè)算子不同，PP是將整個(gè)網(wǎng)絡(luò)在層間進(jìn)行切分。這樣整個(gè)計(jì)算圖就分為幾個(gè)子圖。對(duì)于有依賴的子圖計(jì)算時(shí)就像工廠的流水線一樣，沒有依賴的部分可以在多device上并行執(zhí)行。PP與TP相比，通信量更小。在TP中，每個(gè)Transformer layer需要兩次activation上的all-reduce操作。而對(duì)于PP，只需要點(diǎn)到點(diǎn)的通信。因此PP需要更少的通信，但它的問題是會(huì)有bubble和weight staleness。因此，兩者是正交互補(bǔ)的關(guān)系。

PP早期比較經(jīng)典的工作大致開始于2018年，包括Microsoft論文《PipeDream: Fast and Efficient Pipeline Parallel DNN Training》中提出的PipeDream和Google論文《GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism》中提出的GPipe。

PipeDream以一種通用和自動(dòng)的方式結(jié)合PP, MP和DP。PipeDream將DNN切成多個(gè)Stage。為了將這個(gè)過程自動(dòng)化，它通過profiler估計(jì)每個(gè)layer的計(jì)算時(shí)間與output size?；谶@些估計(jì)，optimizer對(duì)之進(jìn)行切分。切分算法的目標(biāo)是最小化總訓(xùn)練時(shí)間，等價(jià)于最小化最慢的stage的時(shí)間。該問題具有optimal sub-problem property（即給定機(jī)器數(shù)量最大化吞吐的pipeline，包含了最少機(jī)器數(shù)量時(shí)最大化吞吐的sub-pipeline），因此可以用DP求解。Work scheduling使用1F1B（one-forward-one-backward）來避免GPU的idle時(shí)間。將forward和backward的執(zhí)行交替，這樣就不需要等到所有的forward完成再執(zhí)行backward，從而提升硬件利用率。它采用異步backward update的pipeline中每個(gè)mini-batch的forward和backward pass使用不同的參數(shù)，即weight staleness問題，這可能導(dǎo)致訓(xùn)練無法收斂。另外，不同的stage可能面臨不同程度的staleness。這也會(huì)導(dǎo)致收斂性受影響。PipeDream使用了兩種技術(shù)：一是Weight Stashing維護(hù)多份weight，每個(gè)active mini-batch一份；二是Vertical Sync用于估計(jì)potential inconsistency across stages。

與PipeDream不同，GPipe使用的是同步SGD。因此，它會(huì)產(chǎn)生bubble。為了減少bubble的開銷，GPipe采用batch splitting技術(shù)，將mini-batch切分為更小的micro-batches，然后對(duì)其用pipeline。每個(gè)micro-batch稱為一個(gè)chunk。同步發(fā)生在mini-batch結(jié)束時(shí)。一個(gè)mini-batch中的micro-batches的梯度進(jìn)行累加后在mini-batch的結(jié)尾一次應(yīng)用。這個(gè)同步產(chǎn)生的bubble，當(dāng)device越多越嚴(yán)重。當(dāng)M >= 4 x K（M為micro-batches，K為accelerator數(shù)量）時(shí)bubble overhead可以忽略。另外，由于它需要將中間activation的結(jié)果存下來，因此內(nèi)存需求很大。

簡(jiǎn)單的pipelining會(huì)因?yàn)閕nconsistent weight的問題影響模型收斂性。為了解決這種問題需要在memory footprint與throughput間做trade-off。以前面的兩種方法為例，它維護(hù)一份weight，并周期性地進(jìn)行pipeline flush。而PipeDream采用的是weight stashing策略。它沒有周期性的pipeline flush，代價(jià)是需要維護(hù)多份weight。前者降低了throughput，后者增加了memory footprint。

2020年的論文《Memory-Efficient Pipeline-Parallel DNN Training》提出的PipeDream-2BW（2BW是double-buffered weight updates的縮寫）及使內(nèi)存更加高效。它使用weight gradient coalescing strategy和double buffering of weights技術(shù)。該方法利用了每個(gè)輸入產(chǎn)生的梯度不需要馬上應(yīng)用，因此可以被累加，從而減少需要維護(hù)的weight。它將1F1B與2BW scheme結(jié)合，在每個(gè)wokrer上最多stash兩份weight。避免pipeline flush同時(shí)降低memory footprint。梯度以micro-batches的粒度來計(jì)算。多個(gè)micro-batches的更新會(huì)累加，然后以batch為粒度來應(yīng)用梯度。PipeDream-2BW結(jié)合了PP與DP，它將模型分為多個(gè)stage放到多個(gè)worker中用PP，然后將每個(gè)stage復(fù)制多份用DP。PipeDream-2BW的planner基于cost model，通過搜索的方法決定如何在worker間切分模型。其變體PipeDream-Flush有比2BW更低的memory footprint，但吞吐也低一些。它維護(hù)一份weight，引入周期性的pipeline flush來保證consistent weight。

Megatron在"Scaling Language Model Training to a Trillion Parameters Using Megatron"一文中提到Interleaved schedule。它不限制將連續(xù)的layer分到一個(gè)device上，而是允許多個(gè)stage分到一個(gè)device上。在Interleaved 1F1B schedule中，device上的stage依次執(zhí)行。這樣的優(yōu)點(diǎn)是可以減少bubble，但缺點(diǎn)是由于stage數(shù)量加倍，導(dǎo)致通信量增加。

2019年論文《PipeMare: Asynchronous Pipeline Parallel DNN Training》提出的PipeMare是一種異步方法，文中它和前面提到的同步方法（PipeDream與GPipe）從delay, pipeline utilization和memory usage角度做了對(duì)比?，F(xiàn)有PP方法為了保證sync execution（即forward中的wegith和backward中用于計(jì)算gradient的是同一份），會(huì)引入了bubble，或者存micro-batch的多份weight。這些方法會(huì)降低硬件效率。PipeMare去除了sync execution的限制。為了解決由此引入的對(duì)收斂性的影響，它提出兩種技術(shù)：Learning rate rescheduling（step size與delay反比）和Discrepancy correction（調(diào)節(jié)backward中的weight）。

2020年論文《DAPPLE: A Pipelined Data Parallel Approach for Training Large Models》中的DAPPLE是一個(gè)結(jié)合了DP和PP的同步的分布式訓(xùn)練框架。其中的Profiler以DNN模型為輸入，然后對(duì)每一層的執(zhí)行時(shí)間，activation size和parameter size進(jìn)行profile?；谶@些profile結(jié)果，以及硬件配置等信息，并行策略的planner基于動(dòng)態(tài)規(guī)劃算法解決partition與placement（device mapping）問題，從而得到執(zhí)行時(shí)間最優(yōu)的DP與PP混合策略。另外它還提出一種運(yùn)行時(shí)調(diào)度算法來減少device memory的占用。它的特點(diǎn)是可以確定和交叉的方式來執(zhí)行forward與backward，從而讓每個(gè)pipeline任務(wù)的內(nèi)存可以盡早釋放。

2020年論文《Pipelined Backpropagation at Scale: Training Large Models without Batches》改進(jìn)了PB（Pipelined Backpropagation）方法。PP一般會(huì)將mini-batch切成micro-batch，并組成pipeline。同步方式中，需要等pipeline空了才開始下一輪mini-batch的參數(shù)更新。但這個(gè)pipeline的fill和drain的過程會(huì)嚴(yán)重降低硬件利用率，尤其是當(dāng)batch size相對(duì)pipeline stage數(shù)較小時(shí)。而PB是一種異步訓(xùn)練方法，它不需要等pipeline清空，因此避免了fill和drain帶來的開銷。但是它會(huì)帶來inconsistent weight和stale gradient的問題，導(dǎo)致準(zhǔn)確率下降和訓(xùn)練不穩(wěn)定。另一方面，fine-grained pipelining將每一層作為一個(gè)pipeline stage（也就是每個(gè)process處理一層），被發(fā)現(xiàn)在一些場(chǎng)景下有很好的性能。但它會(huì)使得pipeline stage變多，從而加重weight inconsistency問題。另外，small micro-batch size能減少fine-grained PB中的gradient delay與weight inconsistency問題。針對(duì)fine-grained, small batch場(chǎng)景下PB的準(zhǔn)確率退化問題，本文引入SC（Spike Compensation）和LWP（Linear Weight Prediction）兩個(gè)方法來解決。

2020年論文《HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism》結(jié)合PP與DP，應(yīng)用在異構(gòu)GPU cluster上。這些cluster中有不同型號(hào)的GPU，有一些可能是幾代前的，單卡無法進(jìn)行訓(xùn)練。HetPipe（Heterogeneous Pipeline）中，一組GPU，稱為virtual worker，以pipeline方式處理minibatch。多個(gè)virtual worker采用DP。另外，還提出Wave Synchronous Parallel（WSP）。所謂wave指在一個(gè)virtual worker中可以并行處理的一串minibatch。每個(gè)virtual worker一次性更新wave的所有minibatch，而不是每個(gè)minibatch更新一次到PS。它會(huì)通過profile得到perf model。然后將partitioning問題建模成linear programming問題，最后通過CPLEX求解。

2020年論文《GEMS: GPU-Enabled Memory-Aware Model-Parallelism System for Distributed DNN Training》中的GEMS是GPU上的Memory-Aware MP系統(tǒng)。它結(jié)合了MP和DP。為了克服現(xiàn)有系統(tǒng)資源利用率低和實(shí)現(xiàn)復(fù)雜的缺點(diǎn)，它提出了兩種設(shè)計(jì)：MAST（Memory Aware Synchronize Training）和MASTER（Memory Aware Synchronized Training with Enhanced Replications）。通過訓(xùn)練額外一份模型replica，來填滿做完一次forward與backward后空閑的硬件資源。在1024張V100上來訓(xùn)練1K層的ResNet，并達(dá)到相當(dāng)高的scaling-efficiency。

2021年論文《Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines》提出的Chimera也是一種PP方法。它結(jié)合了bidirectional pipeline方法。其主要思想是結(jié)合不同方向的兩個(gè)pipeline。與最新的同步pipeline方法相比，它減少最多50%的bubble。它是一種同步方法，因此不影響accuracy和convergence。與SOTA方法具有相同的peak activation memory用量，且memory消耗更加平衡。

2021年論文《PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers》中的automated elastic piplining提升了Transformer分布式訓(xùn)練的效率。調(diào)節(jié)PP與DP的程度。Freeze training的研究表明神經(jīng)網(wǎng)絡(luò)中的參數(shù)常常是由下而上收斂的，因此就可以動(dòng)態(tài)分配資源使它們聚焦于某一些active layer。在訓(xùn)練過程中，它漸進(jìn)地識(shí)別并凍結(jié)一些layer，將資源用于訓(xùn)練其它的active layer。PipeTransformer由幾部分組成：1) Freeze Algorithm從training loop中采樣，并作出freezing的決策。2) AutoPipe是一個(gè)elastic pipeline module將frozen layer剔除，剩下的active layer打包塞入GPU。3) AutoDP得到AutoPipe傳來的pipeline length信息，它會(huì)spawn更多的pipeline replica，盡可能提升DP寬度。4) AutoCache是跨pipeline的caching module。它有專門的模塊在shared memory中維護(hù)cache，這樣所有的訓(xùn)練進(jìn)程可以實(shí)時(shí)訪問。當(dāng)activation太大無法放入CPU memory，它還會(huì)將之swap到disk，并自動(dòng)做pre-fetching。

可以看到，以上方法大致可分為兩步和異步兩種。同步方法，像GPipe, GEMS, DAPPLE, Chimera通過periodic pipeline flush保證weight的一致性，從而保證了收斂性。因此更c(diǎn)onvergence-friendly，但有bubble。而異步方法(如PipeDream, PipeMare)沒有flush，雖然吞吐會(huì)更高，但會(huì)導(dǎo)致forward與backward的weight版本mismatch（如AMPNet和PipeMare），或者雖然weight版本一致，但是會(huì)有weight staleness問題（如PipeDream）。這樣雖然能減少或者消除bubble，但有stale weight的問題。為了使得weight版本一致，就需要保留多份weight，從而使memory footprint增大。

Sequence parallelism

相比前兩種，SP是一種更加針對(duì)特定模型的方法。語言類的模型一般要處理或者生成token序列，SP在這個(gè)維度上做并行。

2022年論文《Sequence Parallelism: Long Sequence Training from System Perspective》中討論了sequence parallelism。Transformer模型中self-attention的memory消耗與sequence length是平方關(guān)系，因此大的sequence length常常會(huì)限制大模型的運(yùn)行?，F(xiàn)有的并行方法都要求整個(gè)sequence放在一個(gè)device上。SP將activation沿著sequence維度進(jìn)行切分，然后分到多個(gè)device上，讓每個(gè)device只維護(hù)一部分sequence。為了跨device計(jì)算attention score，SP使用了ring-style communication，提出Ring Self-Attention（RSA）。其實(shí)現(xiàn)基于PyTorch，且與DP, PP和TP兼容，因此稱為4D parallelism。該方法和DP類似，需要參數(shù)和optimizer state在所有device上拷貝，因此對(duì)于大模型仍然不是很友好。

2021年論文《TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models》中的TeraPipe利用了Transformer-based model的auto-regressive性質(zhì)，用于single training sequence內(nèi)的PP。

2022年的論文《Reducing Activation Recomputation in Large Transformer Models》提出兩種技術(shù)：sequence parallelism和selective activation recomputation。兩者都可以將內(nèi)存需求減少約一半。其中的SP結(jié)合了TP。

Expert Parallelism

在不顯著增加計(jì)算量的情況下增加模型capacity一個(gè)方法是利用網(wǎng)絡(luò)結(jié)構(gòu)上的sparsity（注意不要與諸如矩陣計(jì)算中的sparsity混淆）。近幾年比較火的一種網(wǎng)絡(luò)范式就是Mix of Expert(MoE)。這種結(jié)構(gòu)是由2017年論文《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》提出。它擁有一個(gè)很好的性質(zhì)，就是可以在計(jì)算代碼不增加的情況下增加參數(shù)量。方法是添加sparsely activated expert block，并添加一個(gè)parameterized routing function在它們前面。

2021年論文《FastMoE: A Fast Mixture-of-Expert Training System》中的FastMoE是一個(gè)開源的基于PyTorch的分布式的MoE訓(xùn)練系統(tǒng)。FastMoE支持將expoert分布到多個(gè)計(jì)算節(jié)點(diǎn)的多個(gè)worker上。

2023年論文《A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training》提出了DeepSpeed-TED。它結(jié)合了ZeRO的DP，MegatronLM的TP和DeepSpeed-MoE的expert parallelism。其中expert Parallelism利用了expert block可以并行計(jì)算的特點(diǎn)，將它們放到不同的GPU上。

Hybrid parallelism

2022年論文《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》中從memory和compute efficiency角度比較了幾種并行方法的優(yōu)劣。

• DP的compute efficiency高，易于實(shí)現(xiàn)，但是隨著worker增加，batch size也需要對(duì)應(yīng)地增加，但增加batch size會(huì)影響收斂。Memory efficiency由于復(fù)制了model和optimizer，因此不高。梯度需要在worker間通信與聚合，而梯度的通信開銷隨著模型增大而增大。這會(huì)讓通信帶寬限制compute efficiency。
• TP的memory footprint隨著worker數(shù)量增加同比減少。每個(gè)worker的計(jì)算量減少，會(huì)影響compute efficiency。
• PP的memory footprint隨著pipeline stage同比減少，但不減少activation的memory footprint。PP的通信開銷最小，它要求切分盡可能均衡。

實(shí)際工程中，幾種并行方式是可以同時(shí)使用的，這樣可以結(jié)合各種并行方法的優(yōu)勢(shì)。前面提到的很多論文中的工作其實(shí)都結(jié)合了多種并行方式，比如PipeDream就結(jié)合了DP和MP。另外像論文《Integrated Model and Data Parallelism in Training Neural Networks》中也結(jié)合了DP與MP。

結(jié)合多種并行方法的一般原則是節(jié)點(diǎn)內(nèi)使用TP，節(jié)點(diǎn)間使用PP。再用DP復(fù)制到更多節(jié)點(diǎn)。比如論文《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》提到Megatron-Turing NLG 530B在280 A100集群中在node內(nèi)部使用8-way TP，在node間使用35-way PP。然后用DP擴(kuò)展到千卡。

基礎(chǔ)設(shè)施

為了支持靈活的分布式并行，不能啥都從頭手?jǐn)]，還需要一些軟件基礎(chǔ)設(shè)施還提高效率。業(yè)界在這方面也有一些相關(guān)探索。

2018年論文《Mesh-TensorFlow: Deep Learning for Supercomputers》中的Mesh-TensorFlow支持SPMD式的per-operator partitioning。Mesh-TF引入了一種tensor partition的表達(dá)。它將processors視作多維mesh。張量可以沿著mesh維度復(fù)制或切分。Mesh-TensorFlow引入了一種可以指定分布式TP計(jì)算的語言，可以編譯計(jì)算圖，加上collective primitive。使用的案例比如Switch Transformers。

2021年論文《CoCoNet: Co-Optimizing Computation and Communication for Distributed Machine Learning》提到目前機(jī)器學(xué)習(xí)中的計(jì)算與通信是獨(dú)立的抽象，由不同的庫(kù)實(shí)現(xiàn)。而打破這個(gè)邊界可以解鎖更多優(yōu)化。比如 Interface（模型參數(shù)放在非連續(xù)buffer中，而要調(diào)用collective communication（如AllReduce）前需要拷貝到一個(gè)單獨(dú)的buffer中），Fusion（產(chǎn)生執(zhí)行多個(gè)通信與計(jì)算操作的kernel，從而減少memory bandwidth），Reorder（將計(jì)算移到通信前或后，這樣可以分散計(jì)算或者使能fusion），Overlapping（多個(gè)計(jì)算和通信操作以細(xì)粒度編排，充分利用計(jì)算與網(wǎng)絡(luò)資源）。

2021年論文《CoCoNet: Co-Optimizing Computation and Communication for Distributed Machine Learning》中的CoCoNET（Communication and Computation optimization for neural networks）使用了嵌入于C++的DSL表達(dá)包含計(jì)算與通信操作的kernel。受Halide啟發(fā)，CoCoNET包含scheduling語言，用一組變換指定執(zhí)行調(diào)度。然后由AutoTuner自動(dòng)應(yīng)用這些變換優(yōu)化程序。最后Code generator根據(jù)程序和schedule生成高性能kernel。

2022年論文《MSCCL: Microsoft Collective Communication Library》
GC3提供面向數(shù)據(jù)的DSL用于寫自定義的collective communication algorithm，optimizing compiler將之lower成可執(zhí)行形式。最后基于interpreter runtime來執(zhí)行。它可用于AllReduce和AllToAll操作。GC3是一種內(nèi)嵌于Python的聲明式語言，采用chunk-oriented programming style（受dataflow編譯語言啟發(fā)）。它可以為特定的topology和input size，又或者是自定義collective（如AllToNext）產(chǎn)生MPI抽象的新算法。

2022年論文《OneFlow: Redesign the Distributed Deep Learning Framework from Scratch》中的OneFlow提出了SBP（split, broadcast, partial-value）抽象，包含S，B和P（分別對(duì)應(yīng)split, broadcast, partial-value）。它用于指定global tensor和對(duì)應(yīng)local tensor的對(duì)應(yīng)關(guān)系。這個(gè)對(duì)應(yīng)關(guān)系稱為SBP signature。通過對(duì)應(yīng)的API，用戶就無需接觸low-level的通信原語，只需指定placement和SBP signature?；谶@些，compiler可以自動(dòng)生成DP, MP和PP的執(zhí)行計(jì)劃。

自動(dòng)化策略

上面講了分布式并行的方法，以及相應(yīng)的基礎(chǔ)設(shè)施，那剩下的核心問題就是切分策略了。

最簡(jiǎn)單就是方式當(dāng)然是手工制定切分策略。如2021年論文《GSPMD: General and Scalable Parallelization for ML Computation Graphs》中提出的GSPMD需要用戶通過annotation的方式指定sharding strategy。也就是說，這種方式中，如何切分就需要使用者來制定。

隨著網(wǎng)絡(luò)模型和硬件越來越大，越來越復(fù)雜，手工配置的缺點(diǎn)就越來越明顯了。精力花費(fèi)不說，模型或者運(yùn)行環(huán)境一變可能就得重來。

更自動(dòng)化一些的方式是針對(duì)特定的網(wǎng)絡(luò)結(jié)構(gòu)與硬件特性，根據(jù)經(jīng)驗(yàn)制定切分的規(guī)則。如2014年的論文《One weird trick for parallelizing convolutional neural networks》中討論針對(duì)CNN的Conv與FC制定的并行策略。Conv由于計(jì)算量大但參數(shù)少，適合DP，而FC相比之下正好相反，適合MP。2022年的論文《Efficiently Scaling Transformer Inference》討論了TPU v4（3D torus topology）上運(yùn)行PaLM模型的一些方法。其中就針對(duì)網(wǎng)絡(luò)中的Attention與FF，結(jié)合硬件特性提出了切分的策略。

根據(jù)給定網(wǎng)絡(luò)模型與硬件配置找最優(yōu)的并行策略本質(zhì)上是一個(gè)組合優(yōu)化問題。對(duì)于這樣一類問題，更加理想一些的方式是自動(dòng)化地找出最優(yōu)解。常見的方式比如整數(shù)線性規(guī)劃，啟發(fā)式搜索或者強(qiáng)化學(xué)習(xí)。

對(duì)于MP的策略搜索其實(shí)也可以看作device placement問題。對(duì)于該問題，幾年前就有不少相關(guān)工作。雖然大部分是針對(duì)單機(jī)多卡，或者CPU與GPU的異構(gòu)計(jì)算等小規(guī)模場(chǎng)景，但原理是類似的。其中很多是基于強(qiáng)化學(xué)習(xí)方法。如2017年的論文《Device Placement Optimization with Reinforcement Learning》中的ColocRL使用強(qiáng)化學(xué)習(xí)讓RNN策略網(wǎng)絡(luò)來預(yù)測(cè)計(jì)算圖中操作的placement，從而來解決CPU與多GPU環(huán)境中的device placement問題。2018年的論文《A Hierachical Model for Device Placement》是對(duì)它的改進(jìn)，引入了更加靈活的方式，免除了需要人工對(duì)算子分組的繁瑣。還有2018年的論文《Spotlight: Optimizing Device Placement for Training Deep Neural Networks》中的Spotlight基于強(qiáng)化學(xué)習(xí)中的PPO算法。2018年的論文《Post: Device Placement with Cross-entropy Minimization and Proximal Policy Optimization》結(jié)合了cross-entropy minimization和PPO算法。前者用于batch間，后者用于batch內(nèi)，從而提高學(xué)習(xí)效率。2018年的論文《Placeto: Efficient Progressive Device Placement Optimization》中的Placeto也是基于強(qiáng)化學(xué)習(xí)。它通過迭代式生成placement方案和使用graph embedding提高了模型的泛化能力。2019年的論文《REGAL: Transfer Learning For Fast Optimization of Computation Graphs》中提出的REGAL（REINFORCE-based Genetic Algorithm Learning）則是結(jié)合了強(qiáng)化學(xué)習(xí)與遺傳算法（更具體地，BRKGA）。該算法由強(qiáng)化學(xué)習(xí)（REINFORCE算法）訓(xùn)練GNN表示的策略網(wǎng)絡(luò)，負(fù)責(zé)輸出分布，用于遺傳算法中的mutation和crossover。

2018年的論文《Beyond Data and Model Parallelism for Deep Neural Networks》中的FlexFlow使用自動(dòng)搜索的方法來找圖中operation在考慮DP與MP情況下的最優(yōu)切分。FlexFlow是一個(gè)高用編排TP計(jì)算的DP框架。它使用Markov Chain Monte Carlo算法來進(jìn)行策略的搜索。其中會(huì)考慮sharding strategy和placement。

2019年的論文《Optimizing Multi-GPU Parallelization Strategies for Deep Learning Training》提出MP與DP結(jié)合可以有效解決DP的scaling與statistical效率降低的問題。為了找到混合DP與MP下的最優(yōu)策略，它使用了基于整數(shù)線性規(guī)劃的工具-DLPlacer用于求解device placement問題，最大化并行度，減少數(shù)據(jù)通信開銷。它以最小化單步的訓(xùn)練時(shí)間為目標(biāo)，給出operation在硬件資源上的分配方案，schedule與activation等數(shù)據(jù)的通信路徑。

2020年的論文《Efficient Algorithms for Device Placement of DNN Graph Operators》基于cost model自動(dòng)切分模型到多個(gè)device。Device placement問題中的核心是組合優(yōu)化問題。文中使用基于Integer Programming（IP）與Dynamic Programming（DP）來解決。并且支持non-contiguous fragment of graph。類似的DP也在PipeDream中使用，區(qū)別是PipeDream只支持線性的layer-granularity graphs。

2020年的論文《Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training》是一種在DP中對(duì)weight update做自動(dòng)切分的方法。在DP中，所有的replica都會(huì)做weight update計(jì)算（因?yàn)閣eight沒有batch維）。一種方案是weight update計(jì)算sharding。但簡(jiǎn)單的做法會(huì)導(dǎo)致data formatting和replica間通信開銷極大增加。該文自動(dòng)將weight update計(jì)算在replica間做sharding。它基于XLA，基于static analysis（為了正確性與性能）和graph transformation。不需要修改模型代碼。

2020年的論文《Optimizing Distributed Training Deployment in Heterogeneous GPU Clusters》提出的HeteroG是一個(gè)用于異構(gòu)GPU集群中加速深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練的模塊。HeteroG使用operation-level的hybrid parallelism，communication architecture selection和execution scheduling。它基于GNN-based learning和組合優(yōu)化。

2021年的論文《Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training》中介紹了Amazon的SageMaker。它采用module server架構(gòu)，支持PP，TP和DP多種并行方法，并可以對(duì)模型進(jìn)行自動(dòng)的切分與分布。切分算法的目標(biāo)是盡可能使device上的計(jì)算與訪存平衡，并減少通信次數(shù)。它還為PP提供專用的通信后端。它管理node內(nèi)與node間的device-to-device通信，而無需依賴NCLL。

2021年的論文《Varuna: Scalable, Low-cost Training of Massive Deep Learning Models》中的Varuna自動(dòng)配置分布式(MP與DP)的訓(xùn)練，它可以利用spot VMs而非dedicated VMs，使成本降低5倍。它解決三個(gè)挑戰(zhàn)：慢及不穩(wěn)定的網(wǎng)絡(luò)，資源可能被搶占，對(duì)用戶透明。對(duì)于網(wǎng)絡(luò)問題，它的micro-batch調(diào)度算法可能更加高效應(yīng)對(duì)network jitter，另一方面stage的切分還考慮了網(wǎng)絡(luò)帶寬。對(duì)于搶占問題，它采用了一種dynamic, semantics-preserving reconfiguration方法，稱為Job morphing。它可以在訓(xùn)練過程中配置pipeline depth等參數(shù)，從而適配資源。為了找到更優(yōu)的配置，它有提供了micro-benchmark-driven simulation機(jī)制。對(duì)于易用性問題，它引入了自動(dòng)找模型中cut-points的機(jī)制。

2022年的論文《Accelerating Large-Scale Distributed Neural Network Training with SPMD Parallelism》根據(jù)DNN model與cluster specification自動(dòng)找sharding strategy。它通過overlap通信與計(jì)算來減少通信開銷。兩個(gè)device上計(jì)算兩個(gè)micro-batch可以并行，一個(gè)通信時(shí)另一個(gè)計(jì)算。構(gòu)建training time cost model，并設(shè)計(jì)基于DP的搜索算法。基本思想是將每個(gè)device上分配的計(jì)算圖復(fù)制一份。每個(gè)用一半的輸入數(shù)據(jù)訓(xùn)練。這樣，兩個(gè)micro-batch就能形成pipeline，從而將計(jì)算與通信overlap。 基于PyTorch實(shí)現(xiàn)系統(tǒng)HiDup。它的輸入為模型和集群的 spec。輸出是可用于運(yùn)行在多個(gè)GPU上的模型。

2022年論文《Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning》中的Alpa是一個(gè)用于GPU cluster的DL編譯系統(tǒng)。它構(gòu)建于Jax之上，支持自動(dòng)地將JAX函數(shù)并行化，分析計(jì)算圖并根據(jù)計(jì)算圖與cluster產(chǎn)生并行執(zhí)行計(jì)劃，從而結(jié)合DP，TP和PP將DL運(yùn)行在分布式集群上。Alpa將并行分為兩個(gè)層次：inter-operator與intra-operator并行，將它們分別看作子問題來求解。后者由于通信量需求大，一般分配在由高帶寬連接的設(shè)備中。它將device cluster切成若干device mesh。首先，Inter-op pass使用DP將stage分配給mesh。然后，Intra-op pass使用ILP優(yōu)化給定stage與mesh上的并行執(zhí)行計(jì)劃，再將cost反饋給inter-op pass。這些信息幫助inter-op pass使用DP最小化inter-op的執(zhí)行時(shí)間，從而獲得更優(yōu)的stage與mesh劃分策略。此外，runtime orchestration pass用來滿足相鄰stage的通信需求。

2023年的論文《Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models》是一個(gè)基于PyTorch的自動(dòng)化的intra-op （不考慮inter-op）并行系統(tǒng)。可以對(duì)分布式并行和activation checkpoint進(jìn)行聯(lián)合優(yōu)化。它使用anayzer模塊收集硬件性能和計(jì)算圖開銷信息，采用采用two-stage的求解器找到執(zhí)行計(jì)劃，最后將PyTorch模型編譯成分布式的模型。Analyzer包含三個(gè)部分：symbolic profiler收集由PyTroch FX模塊生成的靜態(tài)計(jì)算圖的計(jì)算和訪存開銷。Cluster detector收集硬件特性和集群的拓?fù)浣Y(jié)構(gòu)信息。Tensor layout manager用于確定tensor layout的conversion策略。并行執(zhí)行計(jì)劃的搜索分為intra-op與activation checkpoint兩層。前者采用的是Alpa中用于處理intra-op的ILP求解器；后者采用Rotor算法。Generator將生成的執(zhí)行計(jì)劃應(yīng)用于原圖，然后將優(yōu)化過的圖編譯成PyTorch代碼。其中的pass會(huì)插入通信節(jié)點(diǎn)，對(duì)參數(shù)做sharding和對(duì)參數(shù)做reshape conversion。

2023年的論文《OSDP: Optimal Sharded Data Parallel for Distributed Deep Learning》中的OSDP(Optimal Sharded Data Parallel)是結(jié)合DP與MP的自動(dòng)并行訓(xùn)練系統(tǒng)。它在內(nèi)存使用與硬件利用率間做trade-off，自動(dòng)產(chǎn)生最大化吞吐率的分布式的計(jì)算圖。它指出ZeRO系統(tǒng)的兩個(gè)缺點(diǎn)：與vanilla DP比有額外50%的通信開銷，中間的巨大張量會(huì)使peak memory usage超限。對(duì)于第一個(gè)問題，它對(duì)于每個(gè)operator確定是否要做parameter sharding。對(duì)于第二個(gè)問題，它使用operator splitting技術(shù)。OSDP包含三個(gè)modules：Profiler, Search Engine和Scheduler。Profiler負(fù)責(zé)根據(jù)Search Engine給的execution plan與batch size，基于cost model估計(jì)memory和time cost。Search Engine基于Profiler的cost估計(jì)，用DFS遍歷所有operator，生成memory不超過限制下的最優(yōu)execution plan。Scheduler迭代地收集Search Engine給出的plan與throughput作為candidate，同時(shí)增加batch size直到所需memory超過device限制。