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

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

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

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

維度

關(guān)鍵詞

一句話定義

數(shù)據(jù)

Feature Store / 數(shù)據(jù)治理

把「原始日志」轉(zhuǎn)成「模型可用特征」并持續(xù)保鮮。

訓(xùn)練

分布式訓(xùn)練 / 顯存優(yōu)化/ 性能優(yōu)化

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

推理

Serving / KV Cache / 編譯器

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

平臺化

MLOps / 監(jiān)控 / CI-CD

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

小結(jié)：AI Infra = Data Infra + Training Infra + Inference Infra + MLOps Infra

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

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

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

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

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

底層算子

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

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

AI框架

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

Pytorch框架支持動態(tài)圖和即時算子下發(fā)模式（Eager Mode），寫法上與普通Python代碼無異，每執(zhí)行一行python代碼就下發(fā)一個算子，計算圖可以不斷動態(tài)拓展。這給深度學(xué)習(xí)算法開發(fā)帶來了極大地便利性，寫法與普通 Python 無異，打印中間變量、斷點(diǎn)調(diào)試、即時修改網(wǎng)絡(luò)結(jié)構(gòu)都更直觀，科研迭代極快。由于極佳的易用性，Pytorch已經(jīng)成為了業(yè)界最主流的AI框架，已經(jīng)成為了事實(shí)上的標(biāo)準(zhǔn)。

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

我個人認(rèn)為，AI Infra/大模型算法工程最好的學(xué)習(xí)切入點(diǎn)是Pytorch框架，它是上層大部分訓(xùn)練加速庫以及AI軟件的框架基礎(chǔ)，同時也起到了算子下發(fā)連接異構(gòu)加速硬件的作用，可以說是AI軟件屆的"linux操作系統(tǒng)"，通過學(xué)習(xí)torch框架的基本用法與原理至關(guān)重要。

分布式訓(xùn)練加速庫

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

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

業(yè)界主流加速庫：Megatron/MindSpeed/FSDP/DeepSpeed

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

Megatron

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

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

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

MindSpeed

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

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

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

DeepSpeed

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

FSDP

Fully-Sharded Data Parallel(FSDP) 是pytorch官方推出的大模型分布式訓(xùn)練方案，可基于torch原生API支持，無需依賴上層加速庫，在易用性方面?zhèn)涫芡瞥纾艿搅藦V大算法人員的喜愛。FSDP通過在DP域內(nèi)對模型也同時進(jìn)行分片，緩解了純DP并行模型參數(shù)冗余的問題，但同時也引入了高昂的通信。總體而言，F(xiàn)SDP目前在易用性方面極佳，但在極致性能以及集群線性度方面和Megatron對比尚有差距，不過后續(xù)隨著社區(qū)的不斷優(yōu)化，F(xiàn)SDP的性能也會得到不斷提升，逐漸逼近Megatron，因此需要對FSDP保持關(guān)注。

TP/PP/SP/DP/CP/EP

由于目前Megatron是當(dāng)前性能最好的分布式訓(xùn)練加速框架，頭部廠商的內(nèi)部訓(xùn)練框架大多基于Megatron進(jìn)行二次開發(fā)，且Megatron在性能優(yōu)化方面積累了很多優(yōu)秀的特性，值得學(xué)習(xí)與分析。

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