用60%成本幹80%的事，DeepSeek分享沉澱多年的高效能深度學習架構
机器之心發表於2024-09-06
硬體發展速度跟不上 AI 需求，就需要精妙的架構和演算法。
根據摩爾定律，計算機的速度平均每兩年就會翻一倍，但深度學習的發展速度還要更快，如圖 1 和 2 所示。可以看到，AI 對算力的需求每年都以 10 倍幅度增長，而硬體速度每兩年增長 3 倍、DRAM 頻寬增長 1.6 倍、互連頻寬則僅有 1.4 倍。而大模型是大資料 + 大計算的產物，其引數量可達千億乃至萬億規模，需要成千上萬臺 GPU 才能有效完成訓練。這些實際情況提升了人們對高效能運算（HPC）的需求。為了獲得更多計算資源，人們不得不擴充套件更多計算節點。這就導致構建 AI 基礎設施的成本不斷激增。降低這些成本具有很大的好處，構建成本和能耗高效型計算機叢集也就自然成了一個熱門的研究方向。近日，DeepSeek（深度求索）釋出了一份基於硬體發展的實際情況及其多年實踐經驗的研究成果，其中提出了一些用於構建用於深度學習和 LLM 的 AI-HPC 系統的成本高效型策略。- 論文標題：Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning
- 論文地址：https://arxiv.org/pdf/2408.14158
具體來說，該團隊基於 Fire-Flyer AI-HPC 架構部署了一個包含 1 萬臺 PCIe A100 GPU 的計算叢集。下表比較了該叢集與英偉達的 DGX-A100 的硬體引數。Fire-Flyer 2：支援深度學習和早期 LLM 訓練如圖 3 所示，LLM 的記憶體需求量通常比較大。相較之下，其它模型的需求就小多了。ResNet、Mask-RCNN、BERT、MAE 等常用模型的引數量均少於 1B，這說明其記憶體需求較低。因此，在設計用於深度學習模型訓練的叢集時，使用 PCIe A100 GPU 可能就已經足夠了。Fire-Flyer 2：PCIe A100 GPU 架構基於該團隊的訓練工作負載，使用單個 200Gbps 的 NVIDIA Mellanox ConnectX-6 (CX6) InfiniBand (IB) 網路卡就能滿足 8 臺英偉達 PCIe A100 GPU 的儲存 IO 和計算通訊的頻寬需求。他們使用瞭如圖 4 所示的計算節點架構：之後，隨著 LLM 時代的到來，該團隊也在 PCIe 卡之間新增了 NVLink Bridge。網路拓撲：整合了儲存和計算的兩層 Fat-Tree他們選擇的拓撲結構是 Fat-Tree，原因是它具有極高的對分頻寬。網路連線解決方案則是 InfiniBand。具體來說，他們使用了 Mellanox QM8700 InfiniBand 交換機，其提供了 40 個速度 200 Gbps 的埠。整體而言，該叢集由 1 萬臺 A100 GPU 構成，包括約 1250 個 GPU 計算節點和近 200 個儲存伺服器，儘管雙層 Fat-Tree 最多可以容納 800 個節點（配置 20 個脊交換機和 40 個葉交換機）。為了降低成本，他們選擇了兩區網路配置而不是三層 Fat-Tree 解決方案，如圖 5 所示。每個計算區都包含一個 800 埠的 Fat-Tree，並連線到了大約 600 個 GPU 計算節點。每臺儲存伺服器配備兩個 IB 網路卡，分別連線到不同的區，因此所有 GPU 計算節點可以共享一組儲存服務。此外，這兩個區會透過有限數量的鏈路互連。他們的 HAI Platform 排程策略確保跨區計算任務最多限制為一個。無論是使用 NCCL 還是 DeepSeek 內部開發的通訊庫 HFReduce，都可以透過使用雙二叉樹演算法跨區執行。其排程器可確保在此拓撲中，只有一對節點跨區通訊因此，即使有任務需要用到所有節點，也能在 Fire-Flyer 2 AI-HPC 上高效執行。在 TF32 和 FP16 GEMM 基準上，相比於英偉達 DGX-A100 架構，DeepSeek 設計的這套架構的計算效能為前者的 83%。但是，其成本和能耗的下降幅度要大得多，僅為前者的 60%，如表 2 所示。DGX-A100 叢集使用了三層 Fat-Tree，其中包含 320 臺核心交換機、500 臺脊交換機和 500 臺葉交換機，總共 1320 臺交換機（如表 3 所示），而 DeepSeek 的這個架構只需要 122 臺交換機。這樣的設計具有更高的成本效益。此外，透過使用 800 個埠的 Frame 交換機，還能進一步降低光模組和線纜的成本。雖然由於 PCIe 卡規格和 SXM 之間的固有差異而存在效能差距，但 DeepSeek 的這一架構通常能以僅 60% 的成本實現 80% 的 DGX-A100 效能！此外，他們還將能耗降低了 40%，也由此降低了二氧化碳排放。從這些指標看，這一架構設計無疑是成功的。有了高效的硬體，也自然需要適配的軟體。該團隊開發了一個用於高效 allreduce 運算的軟體庫：HFReduce。HFReduce 的核心策略見圖 6，其包括節點內（演算法 1）和節點間（演算法 2）的 reduce。HFReduce 相較於 NCCL 的優勢有兩項：如圖 7a 所示，在 Fire-Flyer 2 AI-HPC 上執行資料大小為 186 MiB 的 allreduce 時，HFReduce 可以達到 6.3-8.1GB/s 的節點間頻寬，而 NCCL 的節點間頻寬僅為 1.6-4.8GB/s。另外，還能使用 NVLink 提升 HFReduce 的效能。透過安裝 NVLink Bridge，可透過速度 600 GB/s 的 NVLink 實現成對 GPU 間的高效通訊。為了緩解原 HFReduce 的記憶體限制問題，他們還實現了另一種 allreduce 模式，稱為 HFReduce with NVLink。其核心概念是先在透過 NVLink 互連的 GPU 之間執行 reduce 操作，再將梯度傳遞給 CPU。隨後，當 CPU 返回結果時，它會拆分結果資料並將它們分別返回給透過 NVLink 連線的配對的 GPU，然後透過 NVLink 執行 allgather。如圖 7b 所示，HFReduce with NVLink 實 現了超過 10 GB/s 的節點間頻寬。有關 HFReduce 的策略和瓶頸的更多深度分析請參閱原論文。HaiScale：針對深度學習模型訓練進行特別的最佳化HaiScale 分散式資料並行（DDP）是一種以 HFReduce 為通訊後端的訓練工具。這類似於 Python 的以 NCCL 為後端的 DDP。在反向傳播階段，HaiScale DDP 會對計算出的梯度執行非同步 allreduce 操作，允許此通訊與反向傳播中涉及的計算重疊。如圖 8a 所示，相較於使用 Torch DDP 的 NCCL 後端，使用 HFReduce 訓練 VGG16 模型所需的時間僅為前者的一半，當 GPU 數量從 32 增至 512 時可實現近 88% 的並行可擴充套件性。為了訓練大型語言模型（LLM），HaiScale 框架採用了多種並行策略，類似於 Megagron 和 DeepSpeed。他們針對 PCIe 架構在資料並行（DP）、管道並行（PP）、張量並行（TP）、專家並行（EP）等方面進行了特定的工程最佳化。1. 使用 NVLink Bridge 實現 PCIe GPU 之間的張量並行圖 8 和 9 展示了這些最佳化策略的一些實驗結果。可以看到，隨著 GPU 數量增長，這些策略能帶來非常好的可擴充套件性。此外，該團隊還在論文中分享了更高階的成本效率和聯合設計最佳化方法，其中包括一些降低計算 - 儲存整合網路中資訊擁堵的方法、高吞吐量分散式檔案系統 3FS 以及一個時間共享式排程平臺 HAI Platform。最後，他們驗證了這整套設計的穩定性和穩健性。下圖總結了他們在 2023-2024 年遇到的記憶體和網路故障趨勢。總體而言，Fire-Flyer 2 AI-HPC 在成本效能上表現優秀 —— 能以 60% 的能源消耗達到英偉達 DGX-A100 計算效能的 80%。當進行大規模訓練時，其能帶來的整體成本效益將非常可觀。如果你也打算構建自己的大規模訓練叢集，不妨考慮一下這套架構。