一文讀懂什么是AI Infra和大模型算法工程

什么是AI Infra和大模型算法工程

大模型的三要素是算法、算力和數據，我們需要海量的數據搭配上千卡甚至萬卡算力集群再加上優秀的算法才能訓練出一個優質的大模型。但把這三者連接起來并不是一件容易的事情。通常搞算法的同學如果想獨立地訓一個10B以上的模型，會有一段非常痛苦的調試期。在來我司之前，我在碩士期間是做的AI算法相關，當時實驗室拿到了8塊NVIDIA 4080非常興奮，覺得可以嘗試玩玩小的大模型了，結果找了個開源代碼倉花了幾天時間終于把代碼調通，發現立馬就OOM(out of memory)了，又上網找了一些省顯存的技巧，模型終于跑起來了，但巨慢無比，可謂是一波三折。

AI infra或者說大模型算法工程的角色，就是介于算法和算力芯片之間，搭起一道橋梁，能讓算法工程師更快速地把算法能在算力集群上高效地運行起來。具體來說，有下面幾件事要做：

維度

關鍵詞

一句話定義

數據

Feature Store / 數據治理

把「原始日志」轉成「模型可用特征」并持續保鮮。

訓練

分布式訓練 / 顯存優化/ 性能優化

讓 10B~100B 模型在千卡集群上高效收斂。

推理

Serving / KV Cache / 編譯器

讓大模型在 GPU/NPU上低延遲、高吞吐、低成本地跑起來。

平臺化

MLOps / 監控 / CI-CD

讓算法同學像寫 Python 腳本一樣把模型端到端上線。

小結：AI Infra = Data Infra + Training Infra + Inference Infra + MLOps Infra

說起來看似簡單，但從最底層的加速芯片，如NVIDIA的A100/H100和我們的昇騰910，一直到最上層的算法，有不止一層軟件棧。大致上可以分為：

1. 底層的硬件驅動和算子加速庫（Nvidia的CUDA，昇騰的CANN）
2. AI編程框架（Pytorch/Tensorflow/MindSpore）
3. 分布式訓練/推理加速庫（Megatron-Core/TransformerEngine/DeepSpeed/FSDP/MindSpeed/vLLM/sglang/MindIE）
4. 上層模型與應用套件（veRL/nemo/Megatron-LM）

行業標桿NVIDIA在這方面的積累很深厚，我們為了做好昇騰生態，近幾年也投入了大量的資源，要做好AI Infra其實是投入不菲的，這個“坑”巨大同時機會也很多。另外，這兩年大模型的發展非常迅速，從一開始的chatgpt到llama再到deepseek，模型的規模和參數量越來越大，模型架構也在不斷迭代升級，但硬件的發展速度沒有上層算法的發展快，這就會導致現有的infra不能滿足新的算法帶來的需求。因此，只要大模型算法和應用在不斷的發展，AI infra技術就需要不斷的演進來匹配新算法帶來的新需求，來更好的鏈接算法與算力。或者，從公司和商業角度來說，作為AI Infra人員，只要你的業務方算法能不斷地有新的算法或者新的業務形態，你就會有新的場景需要優化。基于此，目前大部分的頭部AI應用廠商例如阿里/字節/騰訊/百度以及算力廠商NVIDIA，昇騰等等都有AI Infra崗位的需求，仍然處于高速發展的階段。另外，DeepSeek更是證明了AI Infra的重要性，杰出的Infra和算法協同設計可以突破算力的限制，并大幅減少大模型serving的成本，打開了算法和Infra Co-Design的新時代。

當前AI Infra發展的主力在開源社區，例如業界知名的Megatron/vLLM兩個訓練和推理加速框架都是開源軟件，各個AI廠商都在持續貢獻開源社區，一起在推動著這個行業的發展。生態的開放性更是使得Infra行業參與的開發者團體十分巨大。所有在社區的貢獻都會被記錄和認可，開源代碼就是每個人的技術名片，這對于從業人員來說更是一項無形技術資產的積累。同時，大模型訓練和推理優化涉及到計算機網絡、內存，計算和負載均衡等方方面面的術問題，想要做好需要系統性地優化，并在多種約束條件下取得trade-off，通過設計一項好的加速特性，開發者可以體會到系統工程的魅力。

本文后續將會按照底層算子、AI框架、分布式訓練框架、分布式推理框架、強化學習框架的順序進行介紹，在此之前建議先閱讀大模型的一些前置知識

底層算子

上層軟件定義的計算邏輯最終實在AI Core硬件單元上執行的。在底層硬件上有非常多的硬件執行單元，為了保證計算的正確性和高效性，需要對這些單元進行精細地指令控制，是一個非常復雜的過程。算子層的作用就是將一系列計算動作，例如矩陣乘，rms norm等封裝成可供上層軟件直接調用的操作函數（OP），屏蔽上層用戶對硬件底層實現的感知，把復雜性留給算子本身，讓上層用戶可以通過調用一系列封裝好的OP執行自己需要的計算流程。

目前我大部分的工作僅停留在分布式訓練/推理加速庫，并未深入到算子本身的優化，因此本章節只能點到為止，更多的信息大家可以學習昇騰的CANN或者NVIDIA的cuda，二者在底層算子實現上有很多的積累和實踐。

AI框架

在有了底層的算子實現能夠讓AI計算在GPU/NPU上高效地執行起來以后，上層開發者仍然需要一套框架能夠方便地將多個OP組合成算法業務邏輯。這個過程需要同時具備易用性和高效性，盡量能讓開發者聚焦在算法業務邏輯的開發，盡量屏蔽復雜的底層算子實現。AI編程框架對開發者提供高階API，把矩陣運算、自動求導、分布式通信等底層細節封裝起來，研究者只需寫幾十行代碼就能訓練大規模神經網絡，當前最主流的AI編程框架是Pytorch和Tensorflow，其中Pytorch更是占據了大約90%以上的使用場景。

Pytorch框架支持動態圖和即時算子下發模式（Eager Mode），寫法上與普通Python代碼無異，每執行一行python代碼就下發一個算子，計算圖可以不斷動態拓展。這給深度學習算法開發帶來了極大地便利性，寫法與普通 Python 無異，打印中間變量、斷點調試、即時修改網絡結構都更直觀，科研迭代極快。由于極佳的易用性，Pytorch已經成為了業界最主流的AI框架，已經成為了事實上的標準。

同時，Pytorch可支持多種硬件后端，為GPU/TPU/XPU/NPU等多種芯片提供了注冊接口，例如torch-npu項目實現了插件化的npu適配，能支持大部分pytorch原生特性，同時無縫接入了NPU底層算子，從而能在NPU上支持torch生態項目。

我個人認為，AI Infra/大模型算法工程最好的學習切入點是Pytorch框架，它是上層大部分訓練加速庫以及AI軟件的框架基礎，同時也起到了算子下發連接異構加速硬件的作用，可以說是AI軟件屆的"linux操作系統"，通過學習torch框架的基本用法與原理至關重要。

分布式訓練加速庫

隨著大模型的參數量與日俱增，當前參數量已經來到了萬億級別，早期的單機單卡訓練模式已經無法滿足百億/千億/萬億大模型的訓練。當訓練拓展到多機多卡時，就需要卡間和機間的通信來確保計算結果正確，分布式并行算法就成為了必需品。同時，大模型訓練消耗海量的計算資源，通常需要千卡集群運行數天甚至數月，每一秒鐘都是真金白銀，優化計算效率和訓練吞吐對于模型成本和迭代效率至關重要。

因此，在AI Infra層需要沉淀出一個大模型專用加速庫，給大模型訓練提供多種分布式并行算法，例如數據并行（DP），張量并行（TP），流水并行（PP）等，讓大模型能在千卡集群上運行起來。同時也要在計算，通信，顯存等多維度沉淀大模型加速能力。

業界主流加速庫：Megatron/MindSpeed/FSDP/DeepSpeed

目前業界主流的分布式訓練加速庫主要有Megatron，pytorch社區原生的FSDP和DeepSpeed，另外針對昇騰設備上的大模型訓練加速，昇騰推出了MindSpeed支持插件化適配Megatron，支持原生Megatron能力的同時提供親和昇騰NPU的訓練加速能力。

Megatron

Megatron是NVIDIA推出的基于Pytorch的大模型加速庫，起初Megatron的核心主打特性是張量并行（tensor parallelism），通過模型參數的張量切分實現了模型并行。相比于之前工作聚焦于數據并行，Megatron-LM首次將模型并行工程化落地，也成為了后續3D并行范式的起點。

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

Reducing Activation Recomputation in Large Transformer Models

后續Megatron又陸續推出了流水線并行（Pipeline Parallel）、序列并行（sequence parallelism）和傳統的數據并行（DP）一起構成了3D并行范式，即TP-PP-DP。3D并行能夠以90%以上的集群線性度將任意規模的大模型訓練擴展到千卡甚至萬卡集群，同時Megatron后續引入了一系列計算、通信和顯存優化，極致優化大模型訓練效率，目前是開源大模型訓練加速庫中性能最好的。

MindSpeed

MindSpeed是專為昇騰設備設計的大模型加速解決方案，通過插件化適配的方式，使原生Megatron-LM框架支持昇騰軟硬件環境。用戶只需進行簡單的適配，即可在昇騰設備上使用Megatron提供的大部分加速特性和并行。

同時，MindSpeed針對昇騰硬件特點，額外提供了昇騰親和的加速算法，用戶可以使能MindSpeed提供的各類加速、優化技術，從而獲得在昇騰上大模型訓練的極致性能，性能提升可達30%以上。

另外，對于使用自研的類Megatron-LM分布式訓練框架的用戶，MindSpeed同樣提供了靈活的集成方案。通過集成MindSpeed中的關鍵特性，用戶可以實現加速效果。MindSpeed中的各類特性相互獨立、相互兼容，通過封裝成獨立的模塊或API，為用戶提供了方便的集成和優化途徑。此外，MindSpeed采用開源開放策略，用戶可直接獲取MindSpeed源碼，并根據自身需求進行源碼修改和特性集成。MindSpeed的很多特性在客戶界面得到了應用和借鑒，同時部分優秀特性也被Megatron吸收例如選擇性重計算框架，是昇騰上訓練大模型兼具易用性和高性能的解決方案。

DeepSpeed

DeepSpeed是一個由微軟開發的開源深度學習優化庫，旨在提高大規模模型訓練的效率和可擴展性。它通過多種技術手段來加速訓練，包括模型并行化、梯度累積、動態精度縮放、本地模式混合精度等。DeepSpeed還提供了一些輔助工具，如分布式訓練管理、內存優化和模型壓縮等，以幫助開發者更好地管理和優化大規模深度學習訓練任務。此外，deepspeed基于pytorch構建，只需要簡單修改即可遷移。

FSDP

Fully-Sharded Data Parallel(FSDP) 是pytorch官方推出的大模型分布式訓練方案，可基于torch原生API支持，無需依賴上層加速庫，在易用性方面備受推崇，受到了廣大算法人員的喜愛。FSDP通過在DP域內對模型也同時進行分片，緩解了純DP并行模型參數冗余的問題，但同時也引入了高昂的通信。總體而言，FSDP目前在易用性方面極佳，但在極致性能以及集群線性度方面和Megatron對比尚有差距，不過后續隨著社區的不斷優化，FSDP的性能也會得到不斷提升，逐漸逼近Megatron，因此需要對FSDP保持關注。

TP/PP/SP/DP/CP/EP

由于目前Megatron是當前性能最好的分布式訓練加速框架，頭部廠商的內部訓練框架大多基于Megatron進行二次開發，且Megatron在性能優化方面積累了很多優秀的特性，值得學習與分析。

一切大模型訓練性能優化的起點是分布式并行策略，即模型并行和數據并行。模型并行層面主要包括張量并行（TP）和流水線并行（PP），分別在層內和層間對模型進行切分。數據并行層面主要包括上下文并行（CP）和數據并行（DP），分別在序列維度和批次維度對數據進行切分。