隨著大語言模型（LLMs）的規模不斷擴大，提供易于使用且高效的部署路徑變得越來越重要，因為為此類 LLMs 提供服務的成本越來越高。一種降低這一成本的方法是應用訓練后量化（PTQ），其中包括減少為經過訓練的模型提供服務所需的計算和內存需求的技術。
在本文中，我們概述了如何在 NVIDIA NeMo 中實施 PTQ。這是通過使用 NVIDIA TensorRT Model Optimizer 庫實現的，該庫可量化和壓縮深度學習模型，以在 GPU 上優化推理。它還使用 NVIDIA TensorRT-LLM，這是一個開源庫，用于優化大語言模型（LLM）推理。我們展示了量化模型的準確性和性能結果。在整個示例中，我們使用了 Llama 3 模型。
PTQ 是 NeMo LLM 構建和自定義功能的自然擴展，使用 NVIDIA TensorRT 模型優化器和 NVIDIA TensorRT-LLM 實現無縫、高效的部署路徑。例如，NVIDIA NIM 受益于 NeMo 中的 PTQ 工作流程。
從技術角度來看，量化具有以下優勢：
- 它減小了模型大小，因此適合使用較少的GPU和較低的總設備可用內存進行部署。
- 它通過使用更少的位數據類型來減輕內存帶寬壓力
- 例如，它可顯著加快 NVIDIA 架構上的矩陣乘法（GEMM）運算，與微基準測試中的 FP16/BF16 數據類型相比，FP8 的運算速度最高可提升 2 倍。
NeMo 功能概述
NVIDIA NeMo 是一個端到端平臺，可隨時隨地開發自定義生成式 AI。它包含用于訓練、微調、檢索增強生成、護欄和工具包的工具、數據管護工具和預訓練模型，為企業提供了一種簡單、經濟高效且快速的生成式 AI 采用方式。
在使用 NeMo 工具包提供的各種選項構建模型后，可以使用 NeMo 導出和部署工具應用 PTQ 方法，并為優化模型提供服務。
最近發布的 NeMo 容器版本是一個獨立工具包，包含應用 PTQ 和部署量化 LLM 所需的所有依賴項。
NeMo 和 TensorRT 模型優化器提供了適用于量化的各種模型，包括以下系列：
PTQ 支持還提供多節點支持，用于使用適當的張量和流水線并行設置校準最大的 LLMs。
量化和部署 NeMo 模型
總的來說，PTQ 工作流程包括以下步驟：
- 加載模型。
- 校準模型以獲取較低精度的 GEMMs 縮放系數，并將量化模型導出到 TensorRT-LLM 檢查點。
- 構建 TensorRT-LLM 引擎。
- 部署模型 (例如，使用 PyTriton)。
加載模型
典型的 PTQ 用例始于以高精度格式（例如 FP16 或 BF16）訓練的模型，該模型應以較低精度的數據類型（例如 FP8）提供。輸入模型可以是通過之前的工作流步驟獲得的基礎 LLM 或指令調優 LLM。
NeMo 還為各種模型提供社區模型轉換器，這些模型可用于生成相應的 NeMo 檢查點。
校準和導出量化模型
在 PTQ 中，校準是一種獲取模型層中執行的矩陣乘法運算的縮放系數的過程，以便使用低于訓練精度的格式進行計算。
您可以直接從 NeMo 容器（例如使用 torchrun）輕松啟動此步驟，也可以在多節點用例的 Slurm 集群上使用 NeMo 框架啟動器。
簡言之，量化代碼可以歸結為以下代碼示例：
from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModelfrom nemo.export.quantize import Quantizer# Set quantization and export configs appropriately, see https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/conf/megatron_gpt_ptq.yamlquantizer = Quantizer(quantization_config, export_config)model = MegatronGPTModel.restore_from(...)dataloader = ... # A dataloader that yields lists of stringsdef forward_loop(model): # Model forward pass for collecting activation statistics for calibration ...model = quantizer.quantize(model, forward_loop)quantizer.export(model) |
完整的腳本 megatron_gpt_ptq.py 是校準工作流程的入口點。重要的量化參數在 megatron_gpt_ptq.yaml 配置文件中指定，建議使用默認設置。最重要的是，可用的低精度格式和量化算法是 FP8、INT4 AWQ 和 INT8 SQ。
通常，數據集的選擇不會對準確性產生重大影響。然而，對于高度特定領域的應用程序（例如 StarCoder2 等代碼完成模型），建議使用代碼數據集以準確估計校正統計數據。
校正步驟的最后一步是以適合在下一步構建TensorRT-LLM引擎的TensorRT-LLM檢查點格式保存模型。
總體而言，使用 NVIDIA DGX H100 GPU 節點對具有 70B 參數的中等大小模型使用 tensor 并行，校正過程只需幾分鐘。
構建 TensorRT-LLM 引擎
在運行 TensorRT-LLM 之前，您需要編譯一組二進制文件，以構建推理引擎，這些二進制文件考慮了對特定 GPU 硬件、模型架構和推理設置的優化。
使用與常規 NeMo 模型相同的 API，為在校準步驟中獲得的量化檢查點構建引擎。若要構建 FP8 引擎，必須使用具有必要 FP8 支持的計算資源（例如 NVIDIA H100 Hopper 或 NVIDIA L40 Ada Lovelace GPU）運行此步驟。
以下 Python 命令展示了如何構建 TensorRT-LLM 引擎，以及如何在模型中傳遞示例提示。
from nemo.export.tensorrt_llm import TensorRTLLMtrt_llm_exporter = TensorRTLLM(model_dir=”path/to/trt_llm_engine”)trt_llm_exporter.export( nemo_checkpoint_path=”path/to/model_qnemo”, max_batch_size=8, max_input_len=2048, max_output_len=512,)trt_llm_exporter.forward(["How does PTQ work?"]) |
導出命令通常需要幾分鐘才能完成構建或導出 TensorRT-LLM 引擎，并將其保存到 model_dir 參數中。
部署模型
可以使用 PyTriton 輕松部署給定的 TensorRT-LLM 引擎。
from nemo.deploy import DeployPyTritonfrom nemo.export.tensorrt_llm import TensorRTLLMtrt_llm_exporter = TensorRTLLM(model_dir="path/to/trt_llm_engine")nm = DeployPyTriton( model=trt_llm_exporter, triton_model_name="llama3_70b_fp8", port=8000,)nm.deploy()nm.serve() |
最后，在客戶端上，NeMo Framework 提供了一個專用類，用于向服務器發送查詢。以下代碼示例展示了如何使用它。
from nemo.deploy.nlp import NemoQueryLLMnq = NemoQueryLLM( url="localhost:8000", model_name="llama3_70b_fp8",)nq.query_llm( prompts=["How does PTQ work?"], top_k=1,) |
Llama 3 PTQ 示例和結果
出于演示目的，我們展示了兩個預訓練 Llama 3 模型變體（8B 和 70B）的 Llama 3 PTQ 吞吐量和準確性結果。我們分別使用 benchmark.py 和 mmlu.py 腳本評估了 TensorRT-LLM 引擎的性能和準確性。
以下結果是使用 TensorRT-LLM 0.12.0 和 TensorRT Model Optimizer 0.17.0 的 NVIDIA H100 80GB GPU 的結果。軟件棧隨附最新的 NeMo 框架容器，為您提供完整的環境。
準確性結果
圖 2 顯示了跨不同量化方法的兩個 Llama 3 模型大小的 MMLU 結果，以及基準 FP16 結果。
值得注意的是，FP8 量化在很大程度上保持了準確度。在 INT8 SQ 和 Llama 3 模型大小的情況下，我們發現 SmoothQuant alpha 參數可以提高準確性。此參數用于管理量化焦點，從權重僅到激活僅，可以在量化配置中方便地進行設置。在兩個 Llama 3 模型大小的情況下，alpha=0.8 的中間值都會產生最佳的 MMLU 結果。
在表 1 中，括號中的百分比數字是基準 FP16 分數的一小部分，用于測量給定場景保持準確性的程度。
| ? | FP16 | FP8 | INT8 SQ | INT4 AWQ |
| Llama 3 8B | 0.654 | 0.649 (99.2%) | 0.629 (96.2%) | 0.629 (96.2%) |
| Llama 3 70B | 0.790 | 0.787 (99.6%) | 0.772 (97.7%) | 0.777 (98.4%) |
性能結果
圖 3 顯示了定義為不同量化方法和兩個 Llama 3 模型大小的量化模型與 FP16 基準的吞吐量比的推理加速。本文稍后將詳細介紹實現的確切吞吐量結果。
在所有實驗中，我們分別使用 2048 和 512 的輸入和輸出長度來構建 TensorRT-LLM 引擎并收集性能數據。這些值可視為文本摘要場景的代表性參數。
下表顯示了用于為給定量化格式構建引擎的GPU數量，以及兩種批量大小（32和1）的FP16基準結果。吞吐量相對于使用的GPU數量進行了標準化。
| 模型 | 格式 | GPU | 吞吐量【令牌/秒】 | 加速 |
| Llama 3 8B | FP16 | 1 | 2293.08 | – |
| FP8 | 1 | 3330.85 | 1.45 | |
| INT8 SQ | 1 | 3203.50 | 1.40 | |
| INT4 AWQ | 1 | 2475.98 | 1.08 | |
| Llama 3 70B | FP16 | 4 | 256.10 | – |
| FP8 | 2 | 464.22 | 1.81 | |
| INT8 SQ | 2 | 463.25 | 1.81 | |
| INT4 AWQ | 2 | 360.57 | 1.41 |
我們觀察到，在較小的 Llama 3 變體中，FP8、INT8 SQ 和 INT4 AWQ 的速度分別提升了 1.45 倍、1.40 倍和 1.08 倍。在較大的模型中，FP8 和 INT8 SQ 的速度提升高達 1.81 倍。
INT4 AWQ 是一種僅權重的量化方法，推薦用于小批量大小。它主要提高了內存帶寬，但對于更大的批量大小，會受到計算限制。
我們展示批量大小為 1 的結果以供比較。在本例中，對于 Llama 3 8B 和 Llama 3 70B 模型，我們分別獲得了比 FP16 基準高 1.56 倍和 2.66 倍的性能優勢。所有 Llama 3 70B 模型的量化變體都可以使用僅一個 NVIDIA H100 GPU 提供服務，而基準 FP16 精度至少需要兩個 GPUs。
| 模型 | QFORMAT | GPU | 吞吐量【令牌/秒】 | 加速 |
| Llama 3 8B | FP16 | 1 | 135.79 | – |
| FP8 | 1 | 170.75 | 1.26 | |
| INT8 SQ | 1 | 158.90 | 1.17 | |
| INT4 AWQ | 1 | 211.50 | 1.56 | |
| Llama 3 70B | FP16 | 2 | 17.75 | – |
| FP8 | 1 | 32.64 | 1.84 | |
| INT8 SQ | 1 | 32.18 | 1.81 | |
| INT4 AWQ | 1 | 47.14 | 2.66 |
報告的吞吐量數據不應視為峰值性能，因為可以使用 TensorRT-LLM 的其他功能，例如在飛批處理，以進一步改進這些數據。
我們還使用 TensorRT-LLM 的 gptManagerBenchmark 工具檢查了性能統計數據，重點關注 batch size = 32 的 FP16 基準和 FP8 量化引擎。
在 Llama 3 8B 模型中，第一個令牌（TTFT）的時間得到改善，令牌間延遲（ITL）的加速大致相當于本文前面報告的基于吞吐量的加速。
對于更大的 Llama 3 70B 模型，量化引擎在少 2 倍的 GPUs 上運行所實現的 TTFT 和 ITL 結果與基準 FP16 結果相似。這直接轉化為 2 倍的資源節省。借助 PTQ，可以使用更少的 GPUs 更高效地服務模型。
總結
本文向您展示了如何在 NeMo 中使用 PTQ 為 LLM 部署構建高效的 TensorRT-LLM 引擎。在未來迭代中，用于模型的位數將大幅減少，因為新一代 NVIDIA B100 Blackwell 架構將支持 FP4。
還值得一提的是，對于某些應用程序，PTQ 可能已經足夠，而其他應用程序可能需要量化感知訓練（QAT）技術來微調量化權重以保持模型準確性。NeMo 中也提供了 QAT 來滿足這些需求。
有關更多信息，請參閱訓練后量化（Post-Training Quantization，PTQ）。PTQ 的切入點是 megatron_gpt_ptq.py 腳本。你可能還發現 NeMo 框架訓練后量化（PTQ）手冊很有用。它使用兩個示例模型（Llama 3 和 Nemotron-340b）指導你完成整個部署過程。
至于 QAT，入門指南是 megatron_gpt_qat.py 腳本，相應的手冊是適用于 Llama2 SFT 模型的 NeMo 框架量化感知訓練（QAT）。有關更多信息，請參閱調整 TensorRT-LLM 性能的最佳實踐。
致謝
非常感謝 NVIDIA 各個團隊的許多敬業工程師的幫助，他們為成功集成 NeMo 和 TensorRT 模型優化器做出了貢獻，包括 Asma Kuriparambil Thekkumpate、Keval Morabia、Wei-Ming Chen、Huizi Mao、Ao Tang、Dong Hyuk Chang、Alexandros Koumparoulis、Enwei Zhu 和 Simon Layton。
?
