構建支持大模型訓練的計算機系統需要考慮的4個問題

鄭緯民

清華大學計算機科學與技術系，北京 100084

1 三大算力系統

目前，有3種支持大模型訓練的算力系統，分別為基于英偉達公司GPU的系統、基于國產AI芯片的系統，以及基于超級計算機的系統，3種算力系統各有優劣。

基于英偉達公司GPU的系統的硬件性能好、編程生態好，但由于中美博弈引發的禁售風險，目前我國無法購買相關高端芯片。此外，英偉達GPU產品價格高，價格暴漲的同時，一卡難求。

當然，這也給國產AI芯片帶來了一個難得的發展機會?；趪aAI芯片的系統已經取得了很多成果，但是其生態系統不好。要改變生態系統不好的局面，需要在編程框架、并行加速、通信庫、算子庫、AI編譯器、編程語言、調度器、內存分配系統、容錯系統、存儲系統10個關鍵軟件上下功夫。

除了GPU這種算力形態，超級計算機也可以用于支持大模型訓練。部分超算中心的算力利用率并不飽和，可將空閑算力用于大模型訓練?；诔売嬎銠C的系統更要注重軟硬件協同設計，從而更好地服務于大模型訓練。

2 國產AI芯片的破局之路

2.1 芯片的生態系統

芯片的生態系統是指與該芯片相關的一系列硬件、軟件、工具和支持服務，這些共同構成了一個完整的生態系統。如何評價芯片生態系統的優劣？以下是一些常見的評價標準。

● 軟硬件支持：豐富的軟硬件支持能夠提高芯片的可用性和靈活性。

● 兼容性和標準：具備廣泛的兼容性，支持多種硬件和軟件標準，有助于降低開發成本，提高整個系統的集成性。

● 健壯性：具備強大的適應性，能夠及時響應市場需求和技術變革，保持長期的競爭力。

● 安全性：具備成熟的安全性措施，能夠防范各類威脅，特別是在涉及敏感數據或關鍵基礎設施的情況下。

● 供應鏈：具備健全、可靠的供應鏈，在市場需求波動時能夠保持穩定的供應。

● 開發者社區：能夠提供技術支持、經驗分享和創新性的解決方案，促使生態系統不斷發展和優化。

芯片的生態系統涉及多個方面，將上述評價標準落實到具體場景中，評價起來也不難。例如，在基于英偉達公司GPU的系統上訓練的模型，若能便捷地移植到基于國產AI芯片的系統上，那就可以說其生態系統好；如果移植過程中碰到了很多問題，需要耗費一年、兩年，甚至更長的時間，那就是生態系統不好。

2.2 國產AI芯片需要做好的10個軟件

目前國內已經有30多家公司推出了國產芯片，并取得了很多優秀的成果，進步很大。但用戶不喜歡用國產AI芯片，最主要的原因是國產AI芯片的生態系統不好。要改變國產AI芯片生態系統不好的局面，需要做好10個關鍵軟件，分別是：編程框架、并行加速、通信庫、算子庫、AI編譯器、編程語言、調度器、內存分配系統、容錯系統、存儲系統。

（1）編程框架

編程框架需要支持常見的大模型，如TensorFlow和PyTorch，使開發者能夠無縫遷移其現有的代碼。

（2）并行加速

其需要提供高效的并行計算庫和工具，充分利用硬件的并行處理能力，優化系統性能。

（3）通信庫

通信庫要支持在多個芯片或多機之間進行數據傳輸和協同計算，提高數據傳輸效率。

（4）算子庫

算子庫需要提供常見的神經網絡層的各種運算支持，簡化開發工作，提高開發效率。

（5）AI編譯器

AI編譯器需要提高對主流深度學習框架的支持和優化，在異構處理器上對人工智能程序、算子庫不能提供的操作等生成高效的目標代碼。XLA、TVM在機器學習編譯方面表現突出。

（6）編程語言

編程語言需要滿足以下3個要求：一是能夠支持在異構處理器上方便編寫并行程序；二是能夠描述底層硬件功能，與底層硬件緊密配合，充分發揮硬件性能；三是能夠編寫人工智能模型的基本算子（operator），支持AI開發。英偉達的CUDA和Intel的oneAPI在這一方面提供了很好的學習范本。

（7）調度器

調度器需要設計高效的調度算法，合理分配芯片上的計算資源，進一步提高集群資源的利用率，如Kubernetes（K8S）、華為ModelArts。調度器應該在大規模系統上高效調度人工智能任務。

（8）內存分配系統

針對人工智能應用的特點，需要提供高效的內存分配策略，充分利用芯片的內存層次結構，提高數據讀寫效率，減少數據傳輸延遲。

（9）容錯系統

引入容錯機制，確保在模型訓練過程中面臨硬件故障或其他異常情況時，系統能夠快速恢復模型訓練，保持穩定運行。這對可靠性要求較高的應用領域非常重要。

（10）存儲系統

存儲系統需要支持訓練過程中高效的數據讀寫，如檢查點、訓練數據等?？梢钥紤]采用先進的存儲技術，如快速閃存、持久內存等，以提高整體的存儲效率。

目前，上述10個軟件都有國產的，但仍存在不足，比如種類做得不夠齊全、性能不夠好等。這也導致用戶不喜歡用國產AI芯片。針對大多數任務，用戶不會因為芯片性能只有60%而有明顯感知，用戶感覺不好用還是生態系統存在問題。如果國產AI芯片硬件性能達到國外芯片性能的60%，只要把軟件和生態做好，用戶也會滿意的。

3 做好軟硬件協同設計，超算也可以支持大模型訓練

超級計算機是計算機中功能最強、運算速度最快、存儲容量最大的一類計算機，多用于國家高科技領域和尖端技術研究。目前，我國超算水平已經處于國際第一梯隊，有14個國家一級超算中心，另外還有不少由地方和行業建設運營的超算中心。這些超算中心在科學計算上做得很不錯，發揮了很大的作用。但是部分超算中心的算力利用率并不飽和，可將空閑算力用于大模型訓練。

3.1 新一代神威超級計算機

新一代神威超級計算機采用全新的申威SW26010-Pro處理器，性能比上一代提高了4倍。它配備了96 000個處理器，共計37 440 000個核，半精度性能達到5.308 EFLOPS。此外，它的內存空間超過9 PB，聚合帶寬高達23 Pbit/s。

神威高性能計算機體系結構如圖1所示。最左邊的結構是核心組，每個核心組包含8×8個核，再加上1個主核，即65個核；每個核心組都有自己的內存控制器，配備16 GB內存，理論帶寬達 307.2 Gbit/s。而中間的結構是SW26010-Pro處理器，它包含6個核心組，即配備390個核、96 GB內存，相比于SW26010有了顯著提升。超節點由256個CPU組成，超節點內部的CPU之間均有直接通路，因此超節點內部的通信速度非?？?；超節點之間采用裁剪網絡連接，因此超節點間的通信性能比超節點內部低。

圖1 神威高性能計算機體系結構

3.2 利用超級計算機做大模型訓練面臨的3個問題

大規模算力為擴展預訓練模型提供了一個絕佳的機會。利用超級計算機做大模型訓練，仍面臨以下3個問題。

（1）新型硬件層出不窮

隨著應用程序對算力需求的日益提升和摩爾定律的逐漸放緩，新型計算硬件，特別是異構加速處理器，成為高性能計算發展的主流。為了打破存儲墻對應用程序計算性能的桎梏，新型存儲器件（如固態存儲設備SSD等）也被廣泛應用于高性能計算機。過去，一臺機器的基本配備是“CPU+存儲器+硬盤”，如今，除了CPU，還涌現了GPU、TPU（張量處理單元）、NPU（神經網絡處理器）、SSD等新硬件。

由于近幾年美國的芯片限令，國內正在積極探索硬件的國產化之路。在高性能計算機方面，新一代的神威、天河等超級計算機均采用了國產的處理器和加速器，如圖2、圖3所示。

圖2 基于申威系列處理器的神威超級計算機

圖3 基于飛騰處理器和邁創系列加速器的天河超級計算機

（2）新型應用程序快速發展

近年來，新型的應用程序也在快速發展，如基于張量計算的人工智能應用、基于圖數據的圖計算應用、面向大規模數據的大數據應用等，如圖4所示。新型硬件的使用對軟件系統的設計也是一個巨大的挑戰。如何在新型硬件系統上設計與優化新型應用程序，是一個亟待解決的關鍵問題。

圖4 新型應用與新型硬件

（3）整體系統工程化的挑戰

軟硬件都在快速發展，整體系統工程化仍面臨一定的挑戰。如圖5所示，在硬件層面，新型異構高性能計算機的體系結構在計算、網絡、存儲等方面存在硬件復雜性；在軟件層面，不規則應用程序導致節點間負載不均衡、并行擴展難。因此，面臨的主要挑戰是如何有效適配應用程序到復雜的硬件系統。

圖5 整體系統工程化的主要挑戰

3.3 軟硬件協同設計與優化

“八卦爐”是清華大學聯合北京智源人工智能研究院、阿里達摩院等開發的超大規模預訓練模型，模型參數接近于人腦中的突觸數量（人腦有大約1 000萬億個突觸），比GPT3多1 000倍。要訓練如此多的模型參數，需要一臺性能更強的機器提供算力支持?！鞍素誀t”的模型訓練是在青島超算中心完成的。本節以“八卦爐”大模型的訓練為例，介紹大規模預訓練模型系統如何進行軟硬件協同設計與優化。

（1）拓撲感知的混合并行策略

大模型的訓練包括單節點訓練和并行式訓練，其中單節點訓練受限于計算性能與內存，模型難以擴展，而分布式訓練通過不同的并行模式，能夠擴展模型規模與吞吐量。分布式訓練能夠加速深度學習模型的訓練過程，是大模型常用的訓練方式。

當前分布式訓練并行策略各異，各自具有特定的通信模式。如圖6所示，數據并行采用All-Reduce通信，能夠擴展吞吐量；模型并行采用All-Gather等通信，能夠擴展模型規模；專家并行采用All-to-All通信，能夠同時擴展模型規模和吞吐量。簡單使用單一并行策略，可能導致嚴重的性能降級，比如網絡裁剪的影響。

圖6 3 種典型的大模型并行訓練方式

如何在新一代神威的網絡上高效訓練模型？拓撲感知的混合并行策略是一個好的解決方法，它面向網絡拓撲結構選擇合適的并行模式，將通信需求更高的任務分配到高帶寬網絡。我們采用了一種混合數據并行與專家并行的模式。由于數據并行的通信量大、帶寬高，將其分配到超節點內部；由于專家并行的通信量小、帶寬低，將其分配到超節點之間。相比對稱策略，拓撲感知的混合并行策略的性能提升可達到1.6倍。

（2）體系結構感知的訪存性能優化

如圖7所示，申威26010 Pro異構處理器包含6個核組，每個核組有64個計算核心和1個計算控制核心。申威26010 Pro異構處理器配備了6個存儲控制器，每個存儲控制器管理16 GB的內存空間。核組和存儲控制器通過片上環形網絡互連，訪存操作經環形網絡到存儲控制器進行處理。環形網絡和存儲控制器是影響訪存操作性能的主要瓶頸。

圖7 申威26010 Pro 異構處理器

環形網絡的一個潛在問題是網絡擁塞。如圖8所示，當CPU的390個核同時訪問內存時，大量請求被提交到環形網絡，環形網絡無法及時處理請求。從而導致網絡擁塞，吞吐量大大降低，應用程序的訪存性能會明顯下降。

圖8 環形網絡的網絡擁塞

如何優化訪存性能，可以參考以下方法：一是選擇合適的拓撲結構，優化環網上的路由算法，使訪存請求能夠以最短路徑到達目標節點，降低節點之間的通信時延；二是采用高效的請求調度算法，縮短訪存請求排隊和等待的時間；三是設計有效的流量控制策略，合理分配網絡帶寬；四是通過合并請求、壓縮數據等方式減小通信負擔。

（3）大規模檢查點存儲性能優化

大模型所需的訓練時間較長，可能花費1個月、1年，甚至更長時間。機器在執行大規模訓練任務時負載重，發生錯誤的概率高。訓練過程中，經常會發生一些硬件錯誤和軟件錯誤。針對訓練時間較長的任務，若每3小時出錯一次，出錯后又重新訓練，那么訓練任務根本無法完成。

寫容錯檢查點是一個很好的解決方法，在大規模訓練中至關重要。容錯檢查點是如何發揮作用的？它在故障發生前將當前的軟硬件狀態存儲到硬件中，存儲完成后繼續訓練。當機器發生故障時，可直接取出硬盤中存儲的軟硬件狀態，從這個狀態繼續訓練。但是大模型的訓練參數非常龐大，檢查點需要存儲大量的數據，會耗費很長時間。未經優化時，寫一次檢查點可能需要耗費多個小時。

如何優化檢查點存儲，是大模型訓練面臨的重要問題。一個方法是使用高性能的存儲硬件，如SSD、高速磁盤；另一方法是優化文件系統算法，利用并行I/O技術，將檢查點數據劃分為多個塊，并通過多個存儲節點并行寫入。經過優化，我們在青島的超級計算機上寫一次容錯檢查點只需20 min。

大規模預訓練模型系統將以上優化方法整合到神威平臺，并且做了以下調整：在算子庫方面，完善了swTensor，支持混合精度算子；在深度學習框架方面，深度優化了swPyTorch，支持混合并行模式；在預訓練模型方面，實現分層混合精度策略，支持負載均衡方法。

“八卦爐”大模型的訓練被擴展到新一代神威超級計算機的全機規模，這是首次在國產超算平臺上支持完整的預訓練流程，首次支持高達百萬億規模的模型訓練。通過軟硬件協同設計與優化，依托超算算力資源進行大模型訓練完全可以達到英偉達GPU平臺的性能。除此之外，我們團隊目前已經在青島超算中心成功適配了LLaMA大模型和國內的百川大模型，半精度和全精度的訓練效果可以與英偉達GPU平臺完全對齊，而成本只需要英偉達平臺的六分之一。

4 大模型基礎設施的4個問題

大模型基礎設施通常是指用于訓練和部署大規模深度學習模型的龐大計算和存儲系統。這些基礎設施在訓練復雜模型、處理大量數據和提供高性能推理服務方面發揮了關鍵作用。在搭建大模型基礎設施時，需要考慮多個方面的問題，以確保系統的性能、可靠性和可擴展性。

（1）半精度運算性能與雙精度運算性能的平衡

設計中不僅要考慮半精度運算性能，還要考慮雙精度運算能力。根據團隊經驗，雙精度與半精度運算性能之比為1:100～1:50比較合適。根據科學智能（AI for science）和大模型訓練的發展趨勢，我們提出了變精度平衡設計的思想，為適應科學計算和更廣泛的AI算法和應用提供了保障。

（2）網絡平衡設計

網絡設計不能只針對CNN類算法，還需考慮大規模預訓練模型對系統的需求。大規模預訓練模型需要高帶寬、低時延的網絡，這個網絡要支持數據并行、模型并行和專家平行等模式。

假設你買了1萬塊卡，這些卡如何連接？最好的辦法是每塊卡和其他9 999塊卡都建設一條直連通路。如此便需要大量的連接卡，連接卡的花費甚至高于AI卡，并且每塊卡無法提供如此多的插槽用于連接。一種解決方法是每128臺機器保持兩兩互連，以此支持數據并行、模型并行和專家平行模式。不同的并行方式需要的連接方式也不一樣。

（3）內存平衡設計

大量訪問內存的請求會導致網絡擁塞，吞吐量大大降低，應用程序的訪存性能會顯著下降。多個訪問內存的請求可能訪問同一存控對應的內存空間，負載不均，存控需要順序處理訪存請求。因此，要做好體系結構感知的內存平衡設計。

（4）I/O子系統平衡設計

系統的本地NVMe SSD僅通過本地文件系統訪問，這限制了其應用范圍，可將每臺服務器上的快速本地NVMe整合成應用可見的全局分布式文件系統。此外，I/O子系統可以通過增加SSD來支持檢查點。

英偉達GPU的利用率為50%左右，而國產卡很多只能達到20%。為什么只有20%？就是這4個問題沒解決好。大模型基礎設施的4點平衡設計做得好，別人要用1萬塊卡，我們可能用9 000塊卡就可以了。