跳到正文
IFAI

開源專案

NVIDIA 開源 TensorRT-LLM:H100 推理吞吐翻 8 倍

首次發布 傳播指數 48

NVIDIA 推出開源庫 TensorRT-LLM,專為加速大語言模型推理設計。它引入動態組批與 FP8 精度,在 H100 上相比 A100 帶來最高 8 倍吞吐提升,Llama 2 推理提速 4.6 倍。開發者通過 Python API 即可呼叫底層並行最佳化,大幅降低算力與電費支出。

NVIDIA 開源 TensorRT-LLM:H100 推理吞吐翻 8 倍
圖:developer.nvidia.com

TensorRT-LLM 是 NVIDIA 開源的大語言模型推理加速庫,專門解決大模型在 GPU 上跑起來又慢又貴的問題。它把底層複雜的平行計算和視訊記憶體最佳化打包起來,讓開發者用簡單的 Python API 就能在 Ampere、Ada Lovelace 和 Hopper 架構 GPU 上榨乾硬體效能。

它能做什麼

在實際部署中,大模型推理最耗費算力也最佔視訊記憶體。TensorRT-LLM 針對大模型推理常見的效能瓶頸做了深度最佳化,最直接的效果就是大幅拉高吞吐量並壓低成本。

在主流的 CNN/Daily Mail 文本摘要基準測試中,結合了 FP8 精度和動態組批(in-flight batching)的 H100 GPU,整體吞吐量達到了 A100 基礎 PyTorch 模式的 8 倍。在熱門的開源模型 Llama 2 上,推理速度也能提速 4.6 倍。

吞吐提升直接換來了真金白銀的成本下降。根據官方測試,在跑類似 GPT-J 6B 這樣的小引數模型時,相比 A100 基線,總體擁有成本(TCO)能降低 5.3 倍,電費能耗降低 5.6 倍。跑 70B 引數的 Llama 2 時,TCO 也能縮減 3 倍,能耗減少 3.2 倍。

它開箱支援當前幾乎所有主流模型架構,包括 Meta 的 Llama 2、OpenAI 的 GPT-2 和 GPT-3、Falcon、Mosaic MPT 以及 BLOOM 等數十個模型。

怎麼做到的

以前要把 70B 這樣的大模型跑起來,開發者往往得手動拆解模型權重、手寫底層程式碼來協調多張顯示卡。TensorRT-LLM 內建了張量並行機制,利用 NVLink 自動把權重矩陣拆分到多張 GPU 或多個節點上平行計算,完全不需要人工重寫模型結構。

另一個關鍵改進是被稱為 in-flight batching 的動態組批排程。大模型生成文本有長有短,傳統的批處理必須等最長的那條請求生成完才能開始下一批,很多算力在乾等。TensorRT-LLM 能在每個計算迭代中即時把已經完成的請求移出視訊記憶體,並立刻塞進新的請求,實測能讓 H100 上的吞吐直接翻倍。

在視訊記憶體最佳化上,它結合了 Hopper 架構的 Transformer Engine,開發者不需要改動模型程式碼,就能直接把權重轉成 FP8 格式執行。和傳統的 INT8 或 INT4 相比,FP8 在大幅節省視訊記憶體的同時幾乎不損失模型精度。此外,它還集成了 PagedAttention 以及 FlashAttention 等最新的底層運算元融合。

Databricks 工程副總裁 Naveen Rao 表示,MosaicML 已經在自家的推理服務裡無縫接入了這套工具,整合的過程非常輕鬆。

怎麼跑起來

硬體方面,TensorRT-LLM 支援 NVIDIA Ampere、Lovelace 和 Hopper 架構顯示卡。對於百億引數以上的大模型,建議搭配具備高速 NVLink 的多卡環境使用。

專案目前已經完全開源,程式碼託管在 GitHub 的 NVIDIA/TensorRT-LLM 倉庫中,同時也整合進了企業級的 NVIDIA NeMo 框架。開發者只要熟悉 Python,就可以通過模組化的 Python API 自定義模型結構、定義最佳化策略並啟動服務,不需要深入折騰底層的 C++ 或 CUDA 細節。

為什麼重要

我們認為,TensorRT-LLM 最大的價值在於大幅降低了企業自建大模型推理叢集的技術門檻與執行成本。過去只有少數頭部大廠有能力手寫底層運算元榨乾硬體效能,現在任何有部署需求的技術團隊都能直接拿來即用,省下了昂貴的高效能運算工程師人力支出。

對於行業來說,這不僅是 NVIDIA 在給自家新一代顯示卡帶貨,更是在大模型推理領域築深自家的 CUDA 護城河。不過需要注意的是,這套最佳化深度綁定了 NVIDIA 較新的 GPU 架構,特別是 FP8 相關紅利基本由 Hopper 獨享。如果你正在運維基於 A100 或 H100 的推理叢集,現在就非常值得把現有服務遷移上去試試。

來源 (1)

  1. [1]
    developer.nvidia.com(在新分頁開啟)原文developer.nvidia.com

每日 AI 內容精選