ML Infra 一大戰場 - 分散式訓練

Graph from Nvidia Blog. The graph illustrates that the growth in computing demand for transformer-based models is outpacing that of traditional machine learning models.

5 年前在做 ML System 研究時，我們看見 ML Job 所需要的算力需求，已經超過硬體成長的 Moore’s Law。有鑑於此，如何有效的分散 ML Workload，便是當年的最大課題。

5 年後的現在，ML 進入 LLM 時代。傳統 ML 的算力需求成長，是每兩年 8 倍。LLM 的算力需求成長，是每兩年 275 倍。因此，研究這個問題的急迫程度，可說是提高了數十倍 XD

上次我們已經聊過 On-Device AI ，這次我們將轉向 Server Side，聊聊我們該怎麼分散化 ML Jobs。

分散式的各種切法

From ChainerMN.

講到分散式訓練 ( Distributed Training )，直覺就有該怎麼分、以及該怎麼合。

廣義上來說，分法有兩種：Data Parallelism 與 Model Parallelism。

Data Parallelism 是將 Dataset 本身切開，一部分 data 給 device 1 跑、另一部分給 device 2 跑。

跑完以後，各自在將 gradient 互相分享，更新參數，完成一次 batch run。

Model Parallelism 則會將 Model 本身拆開。比如說，一個 6 layer 的 model ，前 3 給 device 1 跑、後 3 給 device 2 跑。

單純切還是太笨了

Data Parallelism 的合併問題

然而 Data Parallelism 切完以後，就有如何分享成果、 Gradient Aggregation 的問題。

基本上，有分為 Synchronize 的方法及 Asynchronize 的方法。

Synchronize 亦指，等全部 device train 完以後，再一起 sync。

當然，只要任一個 device 跑得特別慢，全部就要停下來等他。

From Tensorflow

Asynchronize 意指，有一個共享的 Parameter Server，負責收集大家的 gradient，並和大家溝通。

各 Device 將不再互等，只要跟 Parameter Server 一直 update 大家的成果即可。

當然，每個 Device 跑速不同，收集到的 gradient 可能不精確，準確度不一定比較好。

Model Parallelism 的相依問題

Model Parallelism 也有相依性的問題。簡單的說，如果 Device 1 還沒 train 玩，Device 2 只能乾等。

因此，近年來的 Pipeline Parallelism，便為了解決這問題。

《Chimera: Efficient Training Large-Scale Neural Networks with Bidirectional Pipelines》便提出，如果我們能將各 Model、各 batch run，以相互交叉的方式排序，便可以極大化提升 device 使用率。

From paper: Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines

跨越節點

上述平行化方法，我們都只討論 Model Training 都只在一個機器上發生。

現今 LLM 的世代，模型在大多數情況，都擠不進一台機器，必須多台機器、甚至 multiple clusters。

Topology Scheduling 便成為最大課題。當 communication 已經 over network，就要考慮 network 可能斷、network 有 hotspot 等因素。

實務上，學界也會針對 Intra-node（一個節點內）、Inter-node（多個節點）分開討論、甚至一起討論。

還有更多

上述大多只提及 Model Training 的平行化方法。實際上，還有 Serving、甚至是多重 workloads 的角度可以切入，也會在這個系列中持續探索，非常歡迎大家追蹤這個沙龍，就不會錯過任何文章！