基于NVIDIA Megatron Core的MOE LLM實(shí)現(xiàn)和訓(xùn)練優(yōu)化

本文將分享阿里云人工智能平臺(tái) PAI 團(tuán)隊(duì)與 NVIDIA Megatron-Core 團(tuán)隊(duì)在 MoE (Mixture of Experts) 大語(yǔ)言模型（LLM）實(shí)現(xiàn)與訓(xùn)練優(yōu)化上的創(chuàng)新工作。分享內(nèi)容將按以下脈絡(luò)展開：

首先簡(jiǎn)短回顧 MoE 技術(shù)的發(fā)展歷程，提煉核心概念及其在實(shí)踐應(yīng)用中亟待解決的關(guān)鍵挑戰(zhàn)。

接著詳述雙方合作研發(fā)的 MoE 框架所具有的獨(dú)特能力和卓越性能，展示其在提升模型訓(xùn)練效率、資源利用以及模型表現(xiàn)等方面取得的驗(yàn)證結(jié)果。

最后，扼要介紹阿里云基于此合作成果所搭建的平臺(tái)工具及推薦的最佳實(shí)踐方案，賦能開發(fā)者高效運(yùn)用 MoE 技術(shù)，促進(jìn)大規(guī)模模型訓(xùn)練的深入探索與廣泛應(yīng)用。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIA Megatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

MoE 是一種模型結(jié)構(gòu)，它將稠密模型結(jié)構(gòu)拆分為多個(gè)子結(jié)構(gòu)，每個(gè)子結(jié)構(gòu)成為一個(gè)專家，通過(guò)在訓(xùn)練和推理過(guò)程中動(dòng)態(tài)選擇一組專家進(jìn)行計(jì)算，實(shí)現(xiàn)了模型參數(shù)的稀疏更新。

簡(jiǎn)而言之，MoE 模型將整體模型拆分為多個(gè)專業(yè)子模塊（專家），每次僅激活和更新少數(shù)與輸入相關(guān)的專家子結(jié)構(gòu)，而選擇哪些專家參與計(jì)算則是通過(guò)一個(gè)路由機(jī)制決定。在 GPT 等超大規(guī)模語(yǔ)言模型中，采用MoE 技術(shù)能夠顯著減少訓(xùn)練和推理時(shí)的計(jì)算負(fù)擔(dān)，因?yàn)椴皇撬袇?shù)都需要在每次操作時(shí)都更新，這極大提高了訓(xùn)練效率，并且在推理階段僅使用部分活躍的網(wǎng)絡(luò)參數(shù)，極大地削減了計(jì)算資源需求。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

2023 年末，Mistral AI 推出開源 MoE 架構(gòu)大語(yǔ)言模型 Mixtral 8x7B，憑借 46.7B 參數(shù)量，在多項(xiàng)下游任務(wù)榜單的效果勝過(guò)當(dāng)時(shí)的最佳開源稠密模型 Llama-2 70B。右側(cè)圖表揭示，在同等推理資源條件下，MoE 模型性能顯著優(yōu)于稠密模型。這一成就引發(fā)業(yè)界對(duì) MoE 模型的強(qiáng)烈關(guān)注，進(jìn)而推動(dòng)阿里云與 NVIDIA Megatron-Core 團(tuán)隊(duì)，共同深化在大模型領(lǐng)域的 MoE 技術(shù)合作與應(yīng)用。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

為了描述使用 MoE 結(jié)構(gòu)來(lái)實(shí)現(xiàn)工業(yè)級(jí)應(yīng)用需要解決的問(wèn)題，本文首先回顧 MoE 的發(fā)展歷史。90 年代初提出的多專家模型概念為 MoE 打下了理論基礎(chǔ)，即通過(guò)集合多個(gè)專家模型協(xié)同解決任務(wù)，并采用門控路由進(jìn)行專家選擇。

隨著 2017 年深度學(xué)習(xí)突飛猛進(jìn)，參數(shù)量對(duì)模型性能的重要性日益顯現(xiàn)，傳統(tǒng)稠密結(jié)構(gòu)深度學(xué)習(xí)框架面臨容量瓶頸。谷歌率先將 MoE 與深度學(xué)習(xí)結(jié)合，首次在 RNN 中驗(yàn)證了 MoE 的可行性。至 2020 年，Transformer 結(jié)構(gòu)在語(yǔ)言模型領(lǐng)域展現(xiàn)出卓越效果和優(yōu)秀的可擴(kuò)展性。此后，Google Gshard 項(xiàng)目將 MoE 融入 Transformer，通過(guò)深度設(shè)計(jì)與實(shí)驗(yàn)顯著提升效果和性能，并在 Switch Transformers 中將參數(shù)總量推向萬(wàn)億級(jí)別，奠定了 MoE 架構(gòu)在大模型領(lǐng)域的基礎(chǔ)。

后續(xù)研究如 ST-MoE、Tutel MoE、FasterMOE 和 MegaBlocks 等，針對(duì)速度和穩(wěn)定性進(jìn)行深入探索。2023 年，GPT-4 的驚艷表現(xiàn)引發(fā)了業(yè)內(nèi)對(duì)其是否采用 MoE 結(jié)構(gòu)的熱議。同年，Mixtral 的實(shí)現(xiàn)似乎驗(yàn)證了這一猜想，表明 MoE 在大模型架構(gòu)中的應(yīng)用正逐步走向成熟和廣泛認(rèn)可。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

后續(xù)我們將簡(jiǎn)要剖析幾個(gè)標(biāo)志性工作，首先是將 MoE 與深度學(xué)習(xí)融為一體的 Sparsely-Gated MoE。該研究將 RNN 網(wǎng)絡(luò)劃分為多個(gè)專家子網(wǎng)絡(luò)，并采用 Softmax Gating 機(jī)制來(lái)控制 token 的路由分配。然而，原始 Softmax Gating 邏輯可能導(dǎo)致負(fù)載不均衡，隨著訓(xùn)練推進(jìn)，部分專家網(wǎng)絡(luò)可能過(guò)于活躍，而其他專家的參數(shù)訓(xùn)練不足。

為解決此問(wèn)題，研究者在計(jì)算 Softmax 之前，對(duì)門控矩陣參數(shù)和輸入特征進(jìn)行噪聲 (noise) 注入，以實(shí)現(xiàn)更為均衡的專家選擇。此外，在門控網(wǎng)絡(luò)損失函數(shù) (loss function) 中增設(shè)輔助損失，引導(dǎo)模型實(shí)現(xiàn)更佳的負(fù)載分配。同時(shí)，路由策略上嘗試每次挑選 top-2 至 top-4 的專家，這些優(yōu)化舉措最終助力模型參數(shù)容量首次突破千億門檻。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

而谷歌的 Gshard 項(xiàng)目進(jìn)一步將 MoE 應(yīng)用于 Transformer 架構(gòu)中，其主要?jiǎng)?chuàng)新包括：

設(shè)計(jì)了 MoE Transformer 的基本結(jié)構(gòu)和并行模式，每間隔一層前向?qū)硬捎?MoE 層替代，專家分布在不同設(shè)備上獨(dú)立計(jì)算，其他層的參數(shù)共享。

針對(duì)專家負(fù)載均衡問(wèn)題，提出了“專家容量（expert capacity）”概念，限制每個(gè)專家處理的 token 數(shù)量，并采用殘差連接策略繞過(guò)已滿負(fù)荷的專家。

引入輔助損失與隨機(jī) top-k 路由策略，以優(yōu)化專家選擇過(guò)程。

在通信效率上，創(chuàng)新提出“l(fā)ocal group dispatching”方案，通過(guò)門控網(wǎng)絡(luò)預(yù)篩選后再按專家 ID 分組傳輸數(shù)據(jù)，有效提升通信效率。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

實(shí)驗(yàn)結(jié)果顯示，Gshard 成功將 Transformer 參數(shù)量從 150B 提升至 600B，而計(jì)算開銷僅增加 50%，首次驗(yàn)證了 MoE 在模型擴(kuò)展性上的顯著優(yōu)勢(shì)。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

之后的 Switch Transformer 工作對(duì) MoE Transformer 的訓(xùn)練流程進(jìn)行了深度優(yōu)化，以實(shí)現(xiàn)模型參數(shù)規(guī)模的極限拓展。首先，路由策略選擇了更為激進(jìn)的 top-1 方案，即每次僅選擇一個(gè)專家進(jìn)行參數(shù)更新。此舉不僅能顯著降低路由的計(jì)算負(fù)擔(dān)，與 top-2 策略相比計(jì)算量近乎減半，同時(shí)每個(gè)專家處理的 batch size 也因此減少，進(jìn)一步減輕了通信成本。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

其次，針對(duì)先前存在的專家負(fù)載均衡問(wèn)題，工作引入了輔助損失函數(shù)。該函數(shù)中，N 代表專家總數(shù)，f 表示第 i 個(gè)專家分配到的 token 比例，P 代表路由器分配給第i個(gè)專家的概率。優(yōu)化目標(biāo)旨在促使 token 均勻分配至每個(gè)專家。通過(guò)最小化損失函數(shù)，使得 f 和 P 趨向于 N 分之一，此時(shí)損失函數(shù)中的調(diào)節(jié)參數(shù) α 通常取值 0.01 左右，能夠?qū)崿F(xiàn)較好的負(fù)載均衡效果。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

在分布式訓(xùn)練方面，除了先前提及的數(shù)據(jù)并行和專家并行外，該工作對(duì)非專家結(jié)構(gòu)部分實(shí)施了模型并行策略，來(lái)進(jìn)一步降低單卡顯存需求，使得模型參數(shù)量得以顯著提升。

針對(duì)路由負(fù)載均衡引起的模型訓(xùn)練隨機(jī)性與不穩(wěn)定問(wèn)題，該工作提出了一系列穩(wěn)定訓(xùn)練的策略。其中包括在本地計(jì)算時(shí)采用較高的 FP32 精度，而在通信階段使用較低的 BF16 精度的隨機(jī)精度策略；采用具有較小縮放系數(shù)的正態(tài)分布初始化權(quán)重；引入專家 dropout 即專家內(nèi)部的 dropout 技術(shù)以減少過(guò)擬合現(xiàn)象。

通過(guò)這些優(yōu)化措施，MoE transformer 的訓(xùn)練收斂速度相較于稠密模型在迭代次數(shù)和時(shí)間上均有數(shù)倍提升。其中，最大的 Switch-C 模型參數(shù)量高達(dá) 1.5 萬(wàn)億。后續(xù)的 ST-MoE 工作則更深入地探究了如何進(jìn)一步改善 MoE 模型的訓(xùn)練穩(wěn)定性和提升模型性能。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

在預(yù)訓(xùn)練穩(wěn)定性方面，為解決專家路由隨機(jī)性和數(shù)值精度造成的 roundoff 誤差，即路由網(wǎng)絡(luò)計(jì)算 logits 值過(guò)大導(dǎo)致的訓(xùn)練不收斂問(wèn)題，工作提出了 z-loss 函數(shù)。該函數(shù)中，B 表示一個(gè) batch 的 token 數(shù)量，N 代表專家數(shù)量，X 則對(duì)應(yīng)輸入路由網(wǎng)絡(luò)的 logits 的維度。通過(guò)引入 z-loss 有效抑制了由精度等因素導(dǎo)致的logits 值過(guò)高現(xiàn)象，從而增強(qiáng)了模型訓(xùn)練的穩(wěn)定性，z-loss 同樣適用于稠密 LLM 訓(xùn)練的穩(wěn)定性提升。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

此外，該工作還深入探究了影響 MoE 訓(xùn)練穩(wěn)定性和效果均衡的多種因素，并通過(guò)細(xì)致實(shí)驗(yàn)歸納出一系列最佳實(shí)踐。

譬如，向路由網(wǎng)絡(luò)引入適量噪聲 (noise)，例如 dropout 操作，能夠增強(qiáng)訓(xùn)練穩(wěn)定性，但可能模型效果受損。在調(diào)優(yōu) (Fine-tuning) 階段，ST-MoE 的研究進(jìn)一步發(fā)現(xiàn)，增大 dropout 參數(shù)對(duì)稠密 MoE 模型和非稠密 MoE 模型的影響存在顯著差異，對(duì)于稠密模型影響更大。同時(shí)，與稠密模型不同，小型 batch size 和較大學(xué)習(xí)率 (learning rate) 對(duì) MoE 模型訓(xùn)練效果的提升更為有利。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

Mixtral 模型在實(shí)現(xiàn)上借鑒了 Megablocks 框架的獨(dú)特思路，該框架的主要特點(diǎn)是將 MoE 層的計(jì)算表述為塊稀疏運(yùn)算。面對(duì)前面提及的動(dòng)態(tài)路由問(wèn)題，即不同專家處理的 token 數(shù)量各異，導(dǎo)致用戶在計(jì)算過(guò)程中需抉擇使用 drop token 或 padding，前者關(guān)乎模型性能，后者則會(huì)增加訓(xùn)練成本。

Megablocks 創(chuàng)新性地將多個(gè)矩陣乘的操作統(tǒng)一定義為一個(gè)大型塊稀疏矩陣，即將多位專家的計(jì)算視作一個(gè)固定尺寸的大矩陣，其中的計(jì)算任務(wù)則細(xì)分為多個(gè)小矩陣塊的 GEMM（General Matrix Multiply）操作。

為優(yōu)化這種塊稀疏矩陣的處理，Megablocks 框架利用了 Block Compressed Sparse Row (BCSR) 數(shù)據(jù)結(jié)構(gòu)，對(duì)矩陣的行和列訪問(wèn)速度以及轉(zhuǎn)換操作進(jìn)行了優(yōu)化。通過(guò)這種方法實(shí)現(xiàn)了路由矩陣的高效操作。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

上述 MoE 發(fā)展史詳細(xì)說(shuō)明了其開發(fā)需要解決的問(wèn)題，下文將簡(jiǎn)介阿里云為解決這些問(wèn)題，與 NVIDIA Megatron-Core 團(tuán)隊(duì)的技術(shù)合作及產(chǎn)出。Megatron-Core 是一個(gè)成熟且輕量化的大規(guī)模 LLM 訓(xùn)練框架，集成了訓(xùn)練大規(guī)模 LLM 的核心技術(shù)，比如多元化的模型并行支持、算子優(yōu)化、通信優(yōu)化、顯存優(yōu)化，以及低精度訓(xùn)練（如 FP8）等先進(jìn)技術(shù)。

Megatron-Core 沿襲了 Megatron-LM 的優(yōu)秀能力，并在代碼質(zhì)量、穩(wěn)定性、功能完備度及測(cè)試覆蓋范圍等維度全面提升。尤為關(guān)鍵的是，該框架設(shè)計(jì)上更注重解耦和模塊化，開發(fā)者在做二次開發(fā)或探索新模型架構(gòu)時(shí)享有高度靈活性。因此我們選擇了與 Megatron-Core 團(tuán)隊(duì)開展合作。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

接下來(lái)探討 Megatron-Core 對(duì) MoE 架構(gòu)的主要支持特性。在并行化方面，Megatron-Core MoE 不僅支持專家并行，還支持 3D 并行，包括數(shù)據(jù)并行、張量并行、流水并行以及序列并行等。對(duì)于超大規(guī)模 MoE 模型，它能夠靈活地將專家并行與其他并行策略有機(jī)結(jié)合。

在 token 分發(fā)機(jī)制上，Megatron-Core MoE 采用了 dropless MoE 操作，即不丟棄任何 token。在路由和負(fù)載均衡優(yōu)化層面，它支持多種路由策略，如通用的 top-k，并在負(fù)載均衡算法上支持 Sinkhorn 算法、z-loss 以及 load balancing loss 等多種方案。

此外，為解決多個(gè)專家接收變長(zhǎng)輸入問(wèn)題，Megatron-Core MoE 引入了 GroupedGMM 技術(shù)，并優(yōu)化效率較低的操作，將其替換為優(yōu)化的 CUDA kernel。

同時(shí)，在模型遷移與適配上，Megatron-Core MoE 提供了豐富的模型 checkpoint 轉(zhuǎn)換功能，允許用戶導(dǎo)入HuggingFace 模型，并自由調(diào)整 TP（tensor parallelism）、PP（pipeline parallelism）和 EP（expert parallelism）等結(jié)構(gòu)，隨后利用 Megatron-Core 高效啟動(dòng)模型訓(xùn)練任務(wù)。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

阿里云人工智能平臺(tái) PAI 團(tuán)隊(duì)基于 Megatron-Core，開發(fā)了一套簡(jiǎn)易、高效的大模型訓(xùn)練工具。依托阿里云靈駿計(jì)算服務(wù)，團(tuán)隊(duì)推出了 PAI-Megatron-Patch 工具庫(kù)，可實(shí)現(xiàn)從十余種主流開源大模型的模型格式到Megatron-LM 和 Megatron-Core 的無(wú)縫轉(zhuǎn)換。用戶通過(guò) PAI DSW 和 DLC 等產(chǎn)品，能輕松啟動(dòng) Megatron-LM 及 Megatron-Core 進(jìn)行大規(guī)模稠密模型和 MoE 模型訓(xùn)練。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

針對(duì)開源大模型多采用 HuggingFace 格式，與 Megatron 訓(xùn)練框架存在一定差異的問(wèn)題，PAI-Megatron-patch 的主要功能在于提供從 HuggingFace 模型到 Megatron 框架的權(quán)重轉(zhuǎn)換服務(wù)。該轉(zhuǎn)換實(shí)質(zhì)上是對(duì)模型命名空間進(jìn)行映射，涵蓋了 layernorm 層、attention 層以及 MLP 層等核心組件的定義轉(zhuǎn)換。PAI-Megatron-patch 內(nèi)置簡(jiǎn)潔的權(quán)重轉(zhuǎn)換腳本，使得用戶能夠便捷地執(zhí)行模型格式轉(zhuǎn)換操作，極大地簡(jiǎn)化了遷移流程。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

在進(jìn)行 MoE 模型轉(zhuǎn)換時(shí)，需確保轉(zhuǎn)換后的文件能直接用于分布式訓(xùn)練，為此我們也做了大量驗(yàn)證工作。以Mixtral 模型為例，在擁有 8 個(gè)專家和 16 張卡的場(chǎng)景下，采取 TP=4 和 EP=4 的切分策略。每個(gè) TP 分區(qū)內(nèi)的文件夾中包含四個(gè)文件，存儲(chǔ)全部八個(gè)專家的 FFN 權(quán)重，即每個(gè)文件承載兩位專家的 FFN 權(quán)重信息。

經(jīng)過(guò)轉(zhuǎn)換，模型在零樣本（Zero-shot）損失精度方面的表現(xiàn)如途中下表所示，數(shù)據(jù)顯示轉(zhuǎn)換前后模型的精度差異非常微小，確保了模型轉(zhuǎn)換的有效性和準(zhǔn)確性。

我們的工作不僅關(guān)注模型的轉(zhuǎn)換環(huán)節(jié)，同樣嚴(yán)謹(jǐn)?shù)仳?yàn)證了整個(gè)訓(xùn)練流程的穩(wěn)健性。仍以 Mixtral 8x7B 模型為例，我們?cè)谌齻€(gè)訓(xùn)練階段——從頭預(yù)訓(xùn)練、基于 checkpoint 的續(xù)訓(xùn)及指令微調(diào) (Finetune)，均進(jìn)行了細(xì)致測(cè)試。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

在從頭預(yù)訓(xùn)練階段，我們?cè)O(shè)定的參數(shù)包括：global_size=256，LR (learning rate) =1e-4，seq_len=2048，TP=4。經(jīng)過(guò) 24 小時(shí)至 2.4K 個(gè)訓(xùn)練步驟后，損失 loss 成功收斂至約 1.9。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

在繼續(xù)預(yù)訓(xùn)練階段，設(shè)定 global_size = 256，LR=5e-5，seq_len=2048，TP=4。在 18 小時(shí)達(dá)到 2,000 個(gè)訓(xùn)練步驟后，loss 亦表現(xiàn)出正常的收斂行為。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

針對(duì) Finetune 階段，我們做了三組測(cè)試:

橙色曲線代表 finetune base 模型，設(shè)定 global_size=128，LR=1e-5，seq_len=2048，TP=4。

藍(lán)色曲線代表 finetune instruct 模型，其參數(shù)配置與橙色曲線一致。

紅色曲線代表僅計(jì)算 answer loss 的 finetune instruct 模型，其 global_size=32，LR 依舊為1e-5，seq_len 調(diào)整為 128，TP 依舊為 4。

所有模型在訓(xùn)練過(guò)程中均呈現(xiàn)出正常收斂態(tài)勢(shì)，其中 finetune base 模型導(dǎo)致 loss 下降幅度較大，這一現(xiàn)象符合預(yù)期。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

為了進(jìn)一步驗(yàn)證 Finetune 的效果，我們選取了一個(gè)代碼生成任務(wù)作為評(píng)估基準(zhǔn)。我們利用大約 80K 條公開的代碼訓(xùn)練樣本，測(cè)試集在 HumanEvol 平臺(tái)上進(jìn)行評(píng)估，此處的參數(shù)配置沿襲了上述 SFT 設(shè)定。經(jīng)過(guò) 2,500 步訓(xùn)練后，模型在 HumanEvol 上的性能指標(biāo)從最初的 45.73% 顯著提升至 53.05%，有力證明了訓(xùn)練流程的合理性和有效性。此外，在速度對(duì)比方面，該方法優(yōu)于同等資源條件下運(yùn)行的 Megablocks。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

我們通過(guò)阿里云人工智能平臺(tái) PAI 提供了一套官方最佳實(shí)踐指南，指南是對(duì)我們從數(shù)據(jù)處理、訓(xùn)練、推理、評(píng)估、直至在線服務(wù)部署的完整 AI 開發(fā)流程的測(cè)試成果總結(jié)：

https://mp.weixin.qq.com/s?__biz=Mzg4MzgxNDk2OA==&mid=2247492431&idx=1&sn=1e135a9e61ac65e88a156969d66ab5be&chksm=cf430418f8348d0e1ed97db178536e73d3ae7093e428c6e3e25bc0dc84a5754f49e3c20ce98e&cur_album_id=2918365856378880004&scene=189#wechat_redirect

該指南涵蓋了從數(shù)據(jù)準(zhǔn)備階段起，指導(dǎo)用戶如何從 OSS （對(duì)象存儲(chǔ)服務(wù)）和 NAS （網(wǎng)絡(luò)附加存儲(chǔ)）讀取原始數(shù)據(jù)，并在 PAI DSW 環(huán)境中執(zhí)行高效的數(shù)據(jù)預(yù)處理操作。預(yù)處理后的數(shù)據(jù)能夠方便地回存至 NAS 或 OSS，以為后續(xù)的模型訓(xùn)練做好準(zhǔn)備。

這種方式的模型訓(xùn)練支持在 PAI DLC 進(jìn)行大規(guī)模分布式訓(xùn)練，同時(shí)也兼容 DSW 提供的單機(jī)訓(xùn)練環(huán)境，確保用戶可根據(jù)實(shí)際需求靈活選擇訓(xùn)練方式。訓(xùn)練完成后，模型的 checkpoint 可以直接導(dǎo)出至 OSS 或 NAS 存儲(chǔ)系統(tǒng)。

完成訓(xùn)練的 checkpoint 可先進(jìn)行離線推理和模型性能評(píng)估。一旦推理和評(píng)估驗(yàn)證無(wú)誤，開發(fā)者只需一鍵即可部署至 EAS 提供的模型在線服務(wù)。依托此服務(wù)的接口，開發(fā)者可以輕松構(gòu)建各類 APP 和業(yè)務(wù)場(chǎng)景。

圖片來(lái)源于 GTC 2024 大會(huì) China AI Day 線上專場(chǎng)的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實(shí)現(xiàn)和訓(xùn)練優(yōu)化》

這份最佳實(shí)踐指南提供了詳細(xì)的步驟說(shuō)明，確保用戶能清晰掌握每一步驟的操作方法。

未來(lái)，阿里云人工智能平臺(tái) PAI 團(tuán)隊(duì)將繼續(xù)深化與 NVIDIA Megatron-Core 團(tuán)隊(duì)的合作，致力于在密集型和稀疏型模型的訓(xùn)練表現(xiàn)和效率上取得更大突破，為推進(jìn) AGI（通用人工智能）技術(shù)的發(fā)展貢獻(xiàn)力量。我們熱忱歡迎全球開發(fā)者共同參與到開源社區(qū)項(xiàng)目以及阿里云的建設(shè)之中，攜手共進(jìn)，共創(chuàng)智能未來(lái)。

審核編輯：劉清