過(guò)去幾年，Transformer雖穩(wěn)坐AI架構(gòu)「鐵王座」，但其二次方復(fù)雜度帶來(lái)的算力消耗和長(zhǎng)序列處理瓶頸，限制了大模型在推理階段處理長(zhǎng)文本。

Mamba憑借「線性復(fù)雜度」異軍突起，非常適合長(zhǎng)序列任務(wù)，有望成為Transformer架構(gòu)的替代品，但在處理全局關(guān)系上偏弱。

Mamba+Transformer混合架構(gòu)可以將二者的優(yōu)勢(shì)互補(bǔ)，實(shí)現(xiàn)「效率」和「性能」的雙豐收。

最近英偉達(dá)發(fā)布了Nemotron-H系列模型，模型尺寸為8B和56B（蒸餾版本47B），用Mamba-2層替換了Transformer中的自注意力層，關(guān)鍵創(chuàng)新在于對(duì)Transformer和Mamba的平衡，實(shí)現(xiàn)了高效處理長(zhǎng)上下文的同時(shí)，還不犧牲模型性能，顯著提高了推理速度，并且內(nèi)存占用更少。

論文鏈接：https://arxiv.org/pdf/2504.03624

實(shí)驗(yàn)結(jié)果表明，Nemotron-H模型在準(zhǔn)確度上優(yōu)于同尺寸的開(kāi)源Transformer模型（例如Qwen-2.5-7B/72B和Llama-3.1-8B/70B），同時(shí)在推理速度上提速3倍。

為了進(jìn)一步提高推理速度并降低推理時(shí)所需的內(nèi)存量，研究人員使用一種剪枝和蒸餾壓縮技術(shù)MiniPuzzle，將56B尺寸的模型蒸餾為NemotronH-47B-Base，在基準(zhǔn)的準(zhǔn)確率上與56B模型相當(dāng)，同時(shí)推理速度提升20%

論文中還提出了一種基于FP8的訓(xùn)練方案，使56B模型實(shí)現(xiàn)了與BF16訓(xùn)練相當(dāng)?shù)男阅堋?/span>

Nemotron-H架構(gòu)

Nemotron-H模型由Mamba-2、自注意力層和前饋神經(jīng)網(wǎng)絡(luò)（FFN）層組成，其中總層數(shù)的8%為自注意力層，均勻分布在模型中。

即，Nemotron-H-8B模型包含52層，其中4層為注意力層；Nemotron-H-56B模型包含118層，其中10層為注意力層，其余層平均分配為FFN層和Mamba-2層。

為了與標(biāo)準(zhǔn)Transformer模塊的結(jié)構(gòu)一致，研究人員提出三條設(shè)計(jì)準(zhǔn)則：模型的第一層必須是Mamba-2層，最后一層必須是FFN層，并且自注意力層總是位于FFN層之前。

模型的具體參數(shù)

需要注意的是，8B和56B模型在FFN隱藏維度、注意力查詢頭和狀態(tài)維度設(shè)置上存在差異；

• 對(duì)于Mamba-2層，保持默認(rèn)的頭維度（64）、擴(kuò)展因子（2）和卷機(jī)窗口（4）；
  
• 使用RMSNorm進(jìn)行歸一化；
• 不適用位置嵌入；
• 模型的嵌入層和輸出層使用了獨(dú)立的權(quán)重，沒(méi)有使用線性層的偏置權(quán)重，也沒(méi)有使用dropout；
• 在每個(gè)Mamba-2層、自注意力層和FFN層周圍都加入了殘差跳躍連接。

訓(xùn)練過(guò)程

數(shù)據(jù)準(zhǔn)備

訓(xùn)練數(shù)據(jù)從來(lái)源上大體上可以分為多語(yǔ)言、網(wǎng)絡(luò)爬取、學(xué)術(shù)、代碼、維基百科和數(shù)學(xué)數(shù)據(jù)，這種數(shù)據(jù)組合可以全面覆蓋通用知識(shí)，同時(shí)在編程和數(shù)學(xué)等領(lǐng)域培養(yǎng)強(qiáng)大的專業(yè)能力。

其中多語(yǔ)言數(shù)據(jù)涵蓋了九種語(yǔ)言：德語(yǔ)、西班牙語(yǔ)、法語(yǔ)、意大利語(yǔ)、葡萄牙語(yǔ)、中文、日語(yǔ)、韓語(yǔ)和俄語(yǔ)。

研究人員設(shè)計(jì)數(shù)據(jù)組合時(shí)，確保所有相同質(zhì)量的數(shù)據(jù)源權(quán)重相似，而高質(zhì)量的數(shù)據(jù)源權(quán)重會(huì)高于低質(zhì)量的數(shù)據(jù)源。

在訓(xùn)練56B尺寸的模型時(shí)，使用了大約20萬(wàn)億個(gè)token的數(shù)據(jù)，其中，網(wǎng)頁(yè)爬取數(shù)據(jù)占比最大，達(dá)到了59%，其次是代碼數(shù)據(jù)，占20%，學(xué)術(shù)內(nèi)容占8.8%

在預(yù)訓(xùn)練Nemotron-H基礎(chǔ)模型時(shí)，研究人員采用了分階段的數(shù)據(jù)混合方法：

第一階段，使用一種促進(jìn)數(shù)據(jù)多樣性的數(shù)據(jù)組合；在第二和第三階段，主要使用高質(zhì)量的數(shù)據(jù)集（例如維基百科），其中第二階段為訓(xùn)練進(jìn)度達(dá)到60%時(shí)，第三階段為訓(xùn)練進(jìn)度達(dá)到80%時(shí)；第四階段則使用最后3800億個(gè)訓(xùn)練token

在后訓(xùn)練階段，研究人員調(diào)整了數(shù)據(jù)的分布，更加注重有監(jiān)督的微調(diào)（SFT）樣本。

FP8訓(xùn)練策略

Nemotron-H訓(xùn)練的一個(gè)創(chuàng)新在于使用8位浮點(diǎn)數(shù)（FP8），在降低內(nèi)存需求和計(jì)算成本的同時(shí)，還能保持模型的質(zhì)量，主要包括以下關(guān)鍵點(diǎn)：

采用逐張量（per-tensor）的當(dāng)前縮放技術(shù)，以提高訓(xùn)練的穩(wěn)定性。

將模型中最初的四個(gè)和最后四個(gè)矩陣乘法（GEMM）操作保持在BF16精度，以確保關(guān)鍵部分的高精度處理；

在訓(xùn)練過(guò)程中，F(xiàn)P8訓(xùn)練逐漸與BF16訓(xùn)練收斂，最終達(dá)到類似的性能水平。

實(shí)驗(yàn)結(jié)果表明，F(xiàn)P8訓(xùn)練在多種基準(zhǔn)測(cè)試中能夠達(dá)到或超過(guò)BF16訓(xùn)練的性能，不僅提高了訓(xùn)練效率，還通過(guò)在MMLU、常識(shí)理解、代碼生成和GSM8K等基準(zhǔn)測(cè)試中的表現(xiàn)，證明了其在保持或提升模型質(zhì)量方面的有效性

模型壓縮

為了進(jìn)一步提高模型部署的效率，研究人員開(kāi)發(fā)了一種新型壓縮框架MiniPuzzle的，結(jié)合了剪枝、神經(jīng)架構(gòu)搜索和知識(shí)蒸餾技術(shù)。

MiniPuzzle壓縮框架的工作流程如下圖所示，展示了從預(yù)訓(xùn)練模型到壓縮模型的轉(zhuǎn)變過(guò)程，包括重要性評(píng)估、神經(jīng)架構(gòu)搜索和蒸餾等步驟。

MiniPuzzle方法包含五個(gè)階段：

1. 重要度估計(jì)：分析每一層對(duì)模型性能的貢獻(xiàn)。

def importance_estimation(model, dataset):# Compute per-layer importance scores    scores = []for layer in model.layers:# Zero out layer outputs and measure impact on loss        scores.append(measure_impact_on_loss(model, layer, dataset))return scores

def importance_estimation(model, dataset):
# Compute per-layer importance scores
    scores = []for layer in model.layers:
# Zero out layer outputs and measure impact on loss
        scores.append(measure_impact_on_loss(model, layer, dataset))return scores

2. 層重要度分析：研究人員需要分析出哪些層對(duì)模型性能的貢獻(xiàn)最大。

3. 條件神經(jīng)架構(gòu)搜索：探索不同壓縮架構(gòu)方案，在每個(gè)候選壓縮模型中保留不同的層。

4. 內(nèi)存與性能權(quán)衡：根據(jù)內(nèi)存使用量和準(zhǔn)確度對(duì)模型進(jìn)行評(píng)估，對(duì)候選架構(gòu)的內(nèi)存負(fù)載與基準(zhǔn)性能進(jìn)行權(quán)衡。

5. 知識(shí)蒸餾：通過(guò)訓(xùn)練，使壓縮后的模型能夠匹配甚至超越原始模型的能力。

通過(guò)這一過(guò)程，Nemotron-H-56B模型成功被壓縮為Nemotron-H-47B模型，參數(shù)減少了16%，同時(shí)保持了相當(dāng)?shù)臏?zhǔn)確度，并將推理吞吐量提高了20%。

實(shí)驗(yàn)結(jié)果

Nemotron-H模型在性能和效率方面相較于其他基于Transformer的模型取得了顯著進(jìn)步。

推理吞吐量

混合架構(gòu)使得推理速度大幅提升，尤其是在處理長(zhǎng)序列時(shí)：

• Nemotron-H-56B的推理吞吐量比Qwen-2.5-72B和Llama-3.1-70B高出多達(dá)3倍。
• Nemotron-H-8B在類似準(zhǔn)確度水平下，比Qwen-2.5-7B的吞吐量高出1.8倍。
  

效率提升在處理長(zhǎng)序列（例如65,536個(gè)token）時(shí)尤為明顯，突顯了Mamba層在輸出token時(shí)計(jì)算復(fù)雜度固定的優(yōu)勢(shì)。

多基準(zhǔn)測(cè)試中的準(zhǔn)確度

盡管架構(gòu)發(fā)生了變化，但Nemotron-H模型在廣泛的基準(zhǔn)測(cè)試中仍保持了強(qiáng)勁的性能表現(xiàn)。

在評(píng)估的17項(xiàng)任務(wù)中，Nemotron-H-56B在16項(xiàng)任務(wù)中的表現(xiàn)優(yōu)于Llama-3.1-70B，在數(shù)學(xué)推理任務(wù)上表現(xiàn)尤為出色。

應(yīng)用與多功能性

Nemotron-H模型可以進(jìn)行擴(kuò)展，以適應(yīng)各種應(yīng)用場(chǎng)景。

• 視覺(jué)-語(yǔ)言能力：基礎(chǔ)模型通過(guò)NVLM-D架構(gòu)擴(kuò)展，創(chuàng)建了視覺(jué)-語(yǔ)言模型（VLM），在VQAv2、GQA和VizWiz等基準(zhǔn)測(cè)試中表現(xiàn)出色，顯示出混合架構(gòu)對(duì)多模態(tài)任務(wù)的適應(yīng)性。
  
  
• 代碼生成：模型在與代碼相關(guān)的任務(wù)上表現(xiàn)尤為出色。訓(xùn)練數(shù)據(jù)中包含大量代碼數(shù)據(jù)（占比20%），使得模型能夠理解和生成多種編程語(yǔ)言的高質(zhì)量代碼。
  
  
• 長(zhǎng)文本處理：混合架構(gòu)的一個(gè)顯著優(yōu)勢(shì)是能夠高效處理長(zhǎng)文本。Nemotron-H-8B模型經(jīng)過(guò)專門的長(zhǎng)文本處理能力微調(diào)，在RULER基準(zhǔn)測(cè)試和其他長(zhǎng)文本評(píng)估任務(wù)中表現(xiàn)出色。
  
  
• 針對(duì)不同能力的數(shù)據(jù)分布：研究人員針對(duì)不同的訓(xùn)練階段精心調(diào)整了數(shù)據(jù)分布，以培養(yǎng)特定的能力，通過(guò)調(diào)整不同數(shù)據(jù)類型（網(wǎng)頁(yè)爬取、代碼、數(shù)學(xué)、學(xué)術(shù)等）的比例，可以在不需要架構(gòu)變更的情況下增強(qiáng)模型的特定能力。比如針對(duì)STEM能力優(yōu)化時(shí)，訓(xùn)練數(shù)據(jù)增加了數(shù)學(xué)和代碼內(nèi)容的比重。

Mamba架構(gòu)簡(jiǎn)介

原版Mamba架構(gòu)

Mamba是一種新型的序列建模架構(gòu)，通過(guò)選擇性狀態(tài)空間模型（Selective State Space Model, SSM）和硬件優(yōu)化算法，將計(jì)算復(fù)雜度降低到線性級(jí)別O(L），能夠高效處理長(zhǎng)達(dá)百萬(wàn)級(jí)的序列長(zhǎng)度，推理速度比Transformer快5倍，在短序列任務(wù)中也實(shí)現(xiàn)了超越Transformer的性能。

選擇性SSM的思路是，通過(guò)讓模型參數(shù)依賴于輸入內(nèi)容，實(shí)現(xiàn)對(duì)信息的選擇性傳播和遺忘，過(guò)濾無(wú)關(guān)信息，從而提高對(duì)密集模態(tài)（如語(yǔ)言和基因組）的建模能力。

Mamba中的硬件感知并行算法，可以避免顯式存儲(chǔ)擴(kuò)展?fàn)顟B(tài)，利用GPU的內(nèi)存層次結(jié)構(gòu)優(yōu)化計(jì)算過(guò)程，實(shí)現(xiàn)線性時(shí)間復(fù)雜度，并顯著提升推理速度。

論文鏈接：https://arxiv.org/pdf/2312.00752

Mamba-2架構(gòu)

Mamba-2結(jié)合了狀態(tài)空間模型（SSMs）和注意力機(jī)制，基于State Space Duality (SSD) 框架，通過(guò)結(jié)構(gòu)化矩陣的分解和優(yōu)化算法，實(shí)現(xiàn)了線性擴(kuò)展的訓(xùn)練效率。

與Mamba相比，Mamba-2對(duì)核心層進(jìn)行了優(yōu)化，簡(jiǎn)化了狀態(tài)轉(zhuǎn)移矩陣的結(jié)構(gòu)，并引入了更大的頭維度，從而顯著提高了訓(xùn)練效率，速度提高了2-8倍；還引入了多頭結(jié)構(gòu)和張量并行等技術(shù)，進(jìn)一步增強(qiáng)了模型的表達(dá)能力和并行計(jì)算效率，在大規(guī)模訓(xùn)練和推理中更加高效。

在實(shí)驗(yàn)中，Mamba-2在語(yǔ)言建模和多查詢關(guān)聯(lián)回憶任務(wù)上均優(yōu)于Mamba和基于注意力的模型。

論文鏈接：https://arxiv.org/abs/2405.21060