阿里安全使用 NVIDIA NeMo 框架和 TensorRT-LLM 的大模型工程化落地實踐

前言

隨著 ChatGPT 的一夜爆火，大模型如今越來越廣泛的應用到各種業務領域中，阿里安全的業務領域對大模型技術的應用也已經 2 年有余。本文對阿里安全在大模型工程領域積累的實踐經驗做出總結和分享。

在大模型實際應用實踐的過程中，阿里安全采用 NVIDIA NeMo^TM 框架和 TensorRT-LLM 大語言模型推理加速庫，顯著優化了模型訓練與推理性能。其中 NeMo 在多卡環境可實現 2-3 倍的訓練加速，TensorRT-LLM 結合 SmoothQuant Int8 可實現領先的推理加速比，動態批處理策略 (Dynamic Batch) 將計算步驟減少 30%，實際 QPS 增益 2-3 倍。Prompt 優化策略在特定業務中提升吞吐高達 10 倍。整體優化成果顯著增強了模型性能與業務效率。

掃描下方二維碼、點擊文章底部閱讀原文，或復制鏈接（https://www.nvidia.cn/webinars/sessions/?session_id=240516-32936&ncid=so-wech-741757-vt04），登陸后即可觀看相關演講，并下載講義。

大模型訓練

首先，我們從工程的角度，分析并總結了一些大模型訓練的相關經驗，針對過去 2 年阿里安全工程團隊對大模型訓練加速這一部分做出分享。

NVIDIA NeMo 框架 (Megatron-LM) 是 NVIDIA 提供的一個端到端的云原生框架，無論是在本地還是在云上，用戶可以靈活地構建、定制和部署生成式 AI 模型。它包含但不限于預訓練模型、數據管護工具、模型對齊工具、訓練和推理框架、檢索增強工具和護欄工具包，為用戶使用生成式 AI 提供了一種既方便、又經濟的方法，同時，NeMo 也支持多模態模型的訓練，包括但不限于 Stable Diffusion，Vision Transformer 等。本文關注焦點在于大模型訓練框架的速度對比，因此只聚焦 NVIDIA Megatron Core 部分。在使用 NeMo 進行大模型訓練過程中，影響訓練速度比較大的 feature 主要如下：

• 模型并行，包括張量并行和流水線并行。當然目前還有比較新的 feature，比如長序列下的 context parallel。
• Layer & Kernel Fusion，Megatron Core 會將多個算子的計算融合在一起，放在一個 kernel 中計算，提升訓練速度。
• Distributed Optimizer， Megatron Core 會把 OptimizerStates 分配到各個 device 中，減少顯存的占用。還有比較新的FSDP，支持了梯度和參數的進一步分塊，降低顯存占用。

當然，NeMo 還有其他的一些 feature，比如計算和通信做 overlap，調用基于 CuDNN 實現的 FlashAttention 等等。

訓練性能對比

在 Megatron-LM 的公開論文中可以看到（如圖 1 所示），Megatron Core 可以保證在 GPU 水平擴展的時候，單卡的 FLOPs 基本能保持不變，而像 DeepSpeed 框架則有比較大的衰減。175B 模型在 1,536 卡的規模上，Megatron-LM 的性能是 DeepSpeed 的 3 倍多，530B 模型在 2,240 卡規模上Megatron-LM 也是 DeepSpeed 的 3 倍多。

圖 1：Megatron-LM 論文中性能對比數據

我們團隊在 Llama2-13B 的模型做了類似的實驗，得出的結論也是 NeMo 比 DeepSpeed 性能高，具體的數據如下標所示：

無論在單機 8 卡，還是雙機 16 卡的規模上，NeMo 的性能都是 DeepSpeed 的 2 倍。

NeMo 性能評測小結

• 由于訓練過程中，模型并行以及梯度交換等需要在多個節點中通信，帶寬會成為一個比較重要的瓶頸，在訓練大模型的時候，網卡的帶寬很容易成為瓶頸，建議使用 RDMA+ 多網卡進行大模型訓練。
• 為了降低多節點之間的模型的 OptimizerStates 的通信量，同時也保證訓練的精度不會太大的影響，可以將 Distributed Optimizer dataType 設置為 FP32，同時將 grad_sync_dtype 設置為 BF16。

NeMo 使用問題總結

• 影響訓練性能的三個重要參數：megatron_amp_O2x0；混合精度 O2 優化選項，設為 False 會降低訓練性能；optimx0；.grad_sync_dtype 梯度同步精度選項，跟顯存占用相關，使用 FP16 可以減少顯存占用，optim.optimizer_dtype optimizer states 參數類型，設為 FP16 或者 BF16 會降低模型精度。
• SFT Chat 類模型時需要將 NeMo 默認 prompt 配置修改成原來模型的 prompt，否則會造成訓練 loss 增大，影響模型最終效果。
• 關于 TP 和 PP 在實踐過程中參數如何配置的問題，我們通過實踐發現，一般情況下，對于 13B 左右的模型，當卡的規模小于 32，如果有高效的 RDMA 網絡，一般只需要開啟 TP（TP>1），PP=1 的配置；如果模型更大一些例如 70B 模型，需要同時開啟 PP（PP>1）。

大模型推理

大模型快速部署流程

針對大模型推理，NVIDIA 推出了 TensorRT-LLM，實現了業界領先的性能。經多方比較，TensorRT-LLM 被我們選為構建阿里安全大模型高性能推理的基石。在阿里安全的業務中，由于大模型服務比較多，為了讓算法同學可以快速部署大模型，工程團隊開發了一系列功能讓算法同學可以快速、高效、平穩的部署大模型。部署的具體流程圖如下：

圖 2：大模型部署流程（圖片來源于阿里安全）

• 模型校驗和標準化導出階段：對用戶提供的模型文件和包含相關參數的配置文件進行校驗，并生成標準的數據格式，方便后續的模型編譯工作。
• 模型編譯階段：使用 TensorTR-LLM 將模型編譯成 engine 格式。
• 服務 DAG 編排階段：將服務的各個模塊定義為 Op，通過 DAG 的方式將其展示在可視化界面中，使得算法同學可以自主編排自己的服務邏輯。
• DAG 調試和構建階段：用戶在 DAG 中可以自己構造服務的輸入數據，完成對整個 DAG 進行調試。同時調試是實時的，用戶可以在不部署的服務的情況下調試模型服務，調試結束后，可以將 DAG 固化，最后構建成真實的服務配置。
• K8s 服務部署階段：我們的模型服務都是 K8s 進行部署，便于快速部署，快速復制擴容等。

大模型服務在線推理架構

大模型推理場景有什么特點

• 每條請求推理耗時的分布是非常不均勻，耗時短的請求可能在幾百毫秒能返回，而耗時長的請求可能需要幾十秒才能返回，這種情況在非 LLM 的模型上一般是不存在。
• 目前的大模型是生成式的，在實際應用過程中，會依賴上文的信息，即之前的推理結果需要在后續的推理過程中作為輸入再次傳進來。
• 部署的服務要求能做到實時監控內部的健康狀態，因為大模型服務在使用 GPU 過程中，GPU 內部的錯誤（例如，數組越界等）會導致后續所有的計算都會失敗，因此推理框架需要能及時感知這種錯誤，并且讓框架做到快速重啟。

針對這種服務特點，我們采用異步服務進行部署，即上游調用先提交一次請求，然后可以通過輪詢或者回調的方式獲取服務計算的結果，具體的架構圖如下所示：

圖 3：模型服務架構圖（圖片來源于阿里安全）

大模型服務性能優化篇

TensorRT-LLM 在 In-Flight Batching 的基礎上，新增了調度機制，稱之為動態批處理 （Dynamic Batch）。為了進一步提升模型推理的速度，我們仔細分析了大模型在推理過程中一些計算，從上述公式中我們可以看到，要想讓推理過程中的計算密度變大，只能調大 batch_size；而 generate 過程中，每個樣本的輸入長度和輸出長度都不太一樣，必然會導致算力有浪費的情況，具體的推理過程如下圖所示：

圖 4：大模型 generate 生成過程示意圖（圖片來源于阿里安全）

從上述的案例中可以看到，一共有 4 條樣本需要進行推理，而若 GPU 最多只能一次處理 3 條樣本，則共需要 9 步完成這三條樣本的推理，而第 1 條樣本在 step1 的時候就推理結束，input2 在 step5 的時候推理結束，因此在這 9 步推理過程中，出現很多 EOS 的 token（我們把這種叫做氣泡），氣泡越多，算力浪費越嚴重；剩余 input4 只能按 batch_size=1 進行推理，算力浪費較為嚴重。

在實際的推理服務中，如果一個 batch 中的其中一條輸出的 output_len 很大，就可能會導致該 batch 的氣泡比例可能超過 80%，就造成了算力的嚴重浪費，為了解決這種問題，我們提出新的推理算法，算法的邏輯如下：

1. 準備輸入的數據（n = max_batch_size*10），形成一個候選集合 S。
2. 如果集合 S 為空，退出；否則，對集合 S 的文本序列排序。
3. 從集合 S 中取出 m 個文本進行推理，其中 m = find_max_batch_size(current_seq_len)，推理到一定 step（一般取 100），或者本次 batch 的 70% 的樣本已推理完并且未推理完成的樣本數大于 1/max_batch_size，推理中斷退出。
4. 將上次未推理完成的樣本拼接上已生成的部分文本作為一個新文本，插入到候選集合 S 中，然后重復第 2 步。

算法的核心思想是依據當前是 seq_len 設置每次 batch 的大小（為了讓 GPU 內存占滿，不造成算力的浪費)，并且每次推理過程中會動態檢測是否需要提前終止，終止的條件是大部分樣本已完成了推理或者推理的 step 到達一定限度。針對圖九中的案例，我們使用新的調度思想之后的效果如下圖所示：

圖 5：啟用新的調度策略之后的推理（圖片來源于阿里安全）

采用新的調度邏輯后，這個案例的 step 數從原來的 20 步減少到 14 步；在實際的業務生產過程中，這個 step 的減少數遠遠大于案例中的 case。我們通過實踐發現，在我們的業務中采用上述算法，服務的每秒查詢率（QPS）一般有 2-3 倍以上的提升。

優化 Prompt，提升服務性能

一般情況下，優化 prompt 可以提升模型的效果，案例如下：

輸入：請判斷一下，這段文字是否正確，如果錯誤請說明理由。「古代帝王把帝位讓給賢能的人稱為禪讓，是原始社會末期部落首領的選舉制度」
輸出：對

輸入：請判斷一下，這段文字是否正確，如果錯誤請說明理由。「禮部為六部之一，掌管禮儀、祭祀、土地、戶籍等職事，部長官稱為禮部尚書」
輸出：錯，"禮部"，在明代有尚書一人，左、右侍郎各一人，下設四司，分別為：儀制清吏司，掌嘉禮、軍禮以及管理學務、科舉考試事；祠祭清吏司，掌吉禮、兇禮事物；主客清吏司，掌賓禮及接待外賓事物；精膳清吏司，掌筵饗等事務。

上述案例是讓大模型判斷一段文字的正確性，如果錯誤則說明理由。在實際服務的生產過程中，對于回答是正確的 case，輸出的 token 數是 1，對于回答錯誤的 case，輸出的 token 數可能很大，如 Dynamic Batch 一節分析的那樣，當輸出的長度差別很大的時候，推理過程中產生的氣泡會很大，這種情況將造成嚴重的算力浪費。為了解決該問題，我們可以將問題進行分解，例如對大模型提問兩次，第一次讓大模型判斷問題的正確性，第二次針對事實錯誤的 case 問大模型原因。

輸入：請判斷這段文字是否正確，「古代帝王把帝位讓給賢能的人稱為禪讓，是原始社會末期部落首領的選舉制度」
輸出：對
輸入：請判斷這段文字是否正確，「禮部為六部 之一，掌管禮儀、祭祀、土地、戶籍等職事，部長官稱為禮部尚書」
輸出：錯

輸入：這段有事實錯誤，請指出原因，「禮部為六部之一，掌管禮儀、祭祀、土地、戶籍等職事，部長官稱為禮部尚書」
輸出："禮部"，在明代有尚書一人，左、右侍郎各一人，下設四司，分別為：儀制清吏司，掌嘉禮、軍禮以及管理學務、科舉考試事；祠祭清吏司，掌吉禮、兇禮事物；主客清吏司，掌賓禮及接待外賓事物；精膳清吏司，掌筵饗等事務。

在實際業務實踐中，對于這類模式的問題采用上述策略，部分業務的吞吐提升 10 倍以上。

模型推理量化

我們的模型量化大部分都是基于 NVIDIA TensorRT Model Optimizer（簡稱 ModelOpt，原名AMMO）做的。 ModelOpt 提供了簡明易用的接口，可以對各種第三方模型進行訓練后量化 （PTQ），并跟 TensorRT-LLM 實現良好銜接。目前 TensorRT-LLM 已經完美的支持了 常見的各種量化方法，如 INT8 weight only（W8A16），AWQ/GPTQ（W4A16 groupwise），FP8 (W8A8) 等。

模型量化后的性能表現

我們分別在 Baichuan2-13B 和 Qwen-14B 的模型上使用了上述各種量化方法進行實驗，量化后的性能結果如下：

圖 6 （圖片來源于阿里安全）

模型推理我們基于 TensorRT-LLM 進行實現的，從實驗中看，在 batch_size 和 seq_len 都相同的條件下，sq_int8 的推理速度是最快的。

模型量化后的業務效果

如下是在以實際落地的兩個業務說明量化的效果：

圖 7（圖片來源于阿里安全）

我們得出的結論是，業務 1 上 smoothQuant int8 表現最好，業務 2 上 int8 weight only 表現更好。在我們的實踐中發現，并沒有哪一種方法始終是最好的，但是相對來說，smoothQuant 在大部分業務場景表現都比較穩健。某些資料表示，smoothQuant int8 對模型精度有一定影響；但在我們的應用場景下，ModelOpt 量化出來的 smoothQuant int8 精度令人滿意。

模型量化經驗總結

• 前期是不同量化方法的選擇問題，我們會在標準的業務數據集上測試各種量化的效果，然后選取比較穩定的量化方法應用到實際業務中。
• 量化后期針對模型在實際業務中存在精度損失問題： 理論上，量化后的模型肯定做不到和未量化的模型的精度相同，但我們在實際業務中能將量化的損失控制在 1% 以內，如果精度損失過大，一般可以調整量化過程中的校準數據集，而校準樣本一般 2,000-8,000 條，而對校準數據集的要求是分布盡可能和實際業務的樣本的分布一致。

大模型工程落地的一些思考

模型性能優化總結

• Attention 是 Transformers 的計算瓶頸，從Transformers 的 FLOPs 分析的最后計算公式中可以看出，Transformer 模型的主要計算是集中在 attention 計算上，過去兩年，業內主要的針對 Transformer 模型的優化也集中這塊，比如 FlashAttention，FlashAttention-2 等方法都是針對 attention 計算的優化。而TensorRT-LLM 中針對不同模型的 attention，基于 flash 的思想，實現了更加豐富的功能及性能的支持。
• 大模型推理解決第 2-n 個推理的加速問題成為最迫切的問題，過去 1 年業內提出的 flash-decoding++ 就是解決這塊計算加速的問題，在 TensorRT-LLM 中是對該部分也有特定的性能優化，在此基礎上，通過調大 batch_size 的方式進一步提升其吞吐量。
• 優化 prompt 減少推理過程中氣泡也是性能有重要手段，具體問題具體分析，總體原則是，在一個 batch 推理的過程中要避免 token 無效計算。

模型效果優化總結

• Prompt 工程是目前大模型實際應用中非常重要的一環，prompt 設計的好壞，不僅僅影響模型的業務效果，同時也是極大的影響模型服務的吞吐量。prompt 是一個經驗工程，需要開發者在實踐中不斷嘗試和總結。
• 另外，prompt 的 few-shot 中的 example 的順序也會影響模型推理的結果。
• 為了避免大模型沒法輸出預期的結果，建議開發者設計多套 prompt，逐步引導大模型正確輸出，并在 prompt 中設計結束符，避免大模型出現幻覺。
• 對于結構化的 prompt，不推薦直接使用換行或者空格直接把各個部分分開，建議使用 markdown 語法或者 xml 語法，這個可能是大模型在預訓練的時候，使用大量的 xml 和 markdown 語料訓練。
• Prompt 的設計需要考慮一些邊界情況，例如 prompt 中可以添加這樣的語句：

if you cannot fetch any information, output `no answer`

未來計劃

• 為了進一步提升大模型的 generate 的速度，后續將會嘗試 Medusa Decode 等解碼手段。
• 后續會考慮嘗試使用 Attention Sink 相關 skills 在長文本中嘗試。
• 過去 1 年一直沒有在 FP8 上進行嘗試，后續將會用 FP8 在實際業務中實踐。
• RAG 的應用，目前大模型能力有一定限制，后續將會嘗試使用 RAG（檢索增強）的方式提升大模型在業務中效果。
• 持續關注 TensorRT-LLM 最新的進展和 feature，TensorRT-LLM 這個開源項目更新迭代的頻率還是挺快的，基本 2-3 個月就會有一個大版本出來。

感謝

本文最后感謝我的主管和團隊，以及 NVIDIA 解決方案技術團隊和 GPU 計算專家團隊對本文的指導。