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。

分散式的各種切法

講到分散式訓練 ( 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 跑得特別慢，全部就要停下來等他。

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 使用率。

跨越節點

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

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

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

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

還有更多

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