支援1024幀、準確率近100％，英偉達「LongVILA」開始發力長影片

机器之心發表於2024-08-21

現在，長上下文視覺語言模型（VLM）有了新的全棧解決方案 ——LongVILA，它集系統、模型訓練與資料集開發於一體。


現階段，將模型的多模態理解與長上下文能力相結合是非常重要的，支援更多模態的基礎模型可以接受更靈活的輸入訊號，以便人們可以以更多樣化的方式與模型互動。而更長的上下文使模型處理的資訊更多，例如長文件、長影片，這種能力同樣為更多現實世界的應用程式提供了所需的功能。

然而，目前面臨的問題是一些工作已經啟用了長上下文視覺語言模型（VLM），但通常是採用簡化的方法，而不是提供一個全面的解決方案。

全棧設計對於長上下文視覺語言模型至關重要。訓練大型模型通常是一項複雜而系統的工作，需要資料工程和系統軟體協同設計。與純文字 LLM 不同，VLM（例如 LLaVA）通常需要獨特的模型架構和靈活的分散式訓練策略。

此外，長上下文建模不僅需要長上下文資料，還需要能夠支援記憶體密集型長上下文訓練的基礎設施。因此，對於長上下文 VLM 來說，精心規劃的全棧設計（涵蓋系統、資料和 pipeline）是必不可少的。

本文，來自英偉達、MIT、UC 伯克利、得克薩斯大學奧斯汀分校的研究者引入了 LongVILA，這是一種用於訓練和部署長上下文視覺語言模型的全棧解決方案，包括系統設計、模型訓練策略和資料集構建。
圖片
  • 論文地址：https://arxiv.org/pdf/2408.10188
  • 程式碼地址：https://github.com/NVlabs/VILA/blob/main/LongVILA.md
  • 論文標題：LONGVILA: SCALING LONG-CONTEXT VISUAL LANGUAGE MODELS FOR LONG VIDEOS

對於訓練基礎設施，該研究建立了一個高效且使用者友好的框架，即多模態序列並行 (MM-SP)，它支援訓練記憶 - 密集型長上下文 VLM。

對於訓練 pipeline，研究者實施了一個五階段訓練流程，如圖 1 所示：即 (1) 多模態對齊，(2) 大規模預訓練，(3) 短監督微調，(4) LLM 的上下文擴充套件，以及 (5) 長監督微調。

對於推理，MM-SP 解決了 KV 快取記憶體使用率的挑戰， 這在處理非常長的序列時會成為瓶頸。

透過使用 LongVILA 增加影片幀數，實驗結果表明該研究在 VideoMME 和長影片字幕任務上的效能持續提高（圖 2）。在 1024 幀上訓練的 LongVILA 模型在 1400 幀的大海撈針實驗中實現了 99.5% 的準確率，相當於 274k 個 token 的上下文長度。此外， MM-SP 系統可以有效地將上下文長度擴充套件到 200 萬個 token 而無需梯度檢查點，與環形序列並行（ring sequence parallelism）相比實現了 2.1 倍至 5.7 倍的加速，與 Megatron 上下文並行 + 張量並行相比實現了 1.1 倍至 1.4 倍的加速。
圖片
下圖為 LongVILA 技術在處理長影片字幕時的示例：在字幕開頭，8 幀的基線模型僅描述了靜態影像和兩輛車。相比之下，256 幀的 LongVILA 描述了雪地上的汽車，包括車輛的前、後和側面檢視。在細節上，256 幀的 LongVILA 還描述了點火按鈕、變速桿和儀表盤的特寫，這些在 8 幀的基線模型中是缺失的。
圖片
多模態序列並行

訓練長上下文視覺語言模型（VLM）會產生大量記憶體需求。例如下圖 1 中 Stage 5 的長影片訓練，單個序列包含了產生 1024 個影片幀的 200K tokens，這超出了單個 GPU 的記憶體容量。

研究者開發了一個基於序列並行的定製系統。序列並行是當前基礎模型系統中常用的一種技術，用於最佳化僅文字的 LLM 訓練。不過，研究者發現現有系統既不高效，擴充套件性也不足以處理長上下文 VLM 工作負載。
圖片
在確定現有系統的侷限性之後，研究者得出結論，一個理想的多模態序列並行方法應該透過解決模態和網路異構性來優先實現效率和可擴充套件性，並且擴充套件性不應受到注意力頭數量的限制。

MM-SP 工作流。為了應對模態異構性的挑戰，研究者提出了一種兩階段式分片策略，以最佳化影像編碼和語言建模階段的計算工作負載。

具體如下圖 4 所示，第一階段首先在序列並行程序組內的裝置之間均勻地分佈影像（比如影片幀），從而在影像編碼階段實現負載平衡。第二階段，研究者聚合全域性視覺和文字輸入以進行 token 級分片。
圖片
2D 注意力並行。為了解決網路異構性並實現可擴充套件性，研究者結合環形（Ring）序列並行和 Ulysses 序列並行的優勢。

具體來講，他們將跨序列維或注意力頭維的並行視為「1D SP」。該方法透過跨注意力頭和序列維的平行計算來實現擴充套件，將 1D SP 轉換為由獨立的 Ring（P2P）和 Ulysses（A2A）程序組組成的 2D 網格。

以下圖 3 左所示，為了實現跨 2 個節點的 8-degree 序列並行，研究者使用 2D-SP 構建了一個 4×2 通訊網格。
圖片
此外，在下圖 5 中，為了進一步解釋 ZIGZAG-RINGATTN 如何平衡計算以及 2D-Attention 機制如何運作，研究者解釋了使用不同方法的注意力計算計劃。
圖片
與 HuggingFace 的原生 pipeline 並行策略相比，本文的推理模式更加高效，原因在於所有裝置同時參與計算，從而與機器數量呈正比地加速程序，具體如下圖 6 所示。同時，該推理模式是可擴充套件的，記憶體均勻地分佈給各個裝置，以使用更多機器來支援更長的序列。
圖片
LongVILA 訓練流程

上文提到，LongVILA 的訓練流程分為 5 個階段完成。各個階段的主要任務分別如下：

在 Stage 1，只有多模態對映器可以訓練，其他對映器被凍結。

在 Stage 2，研究者凍結了視覺編碼器，並訓練了 LLM 和多模態對映器。

在 Stage 3，研究者針對短資料指令遵循任務對模型全面進行微調，比如使用影像和短影片資料集。

在 Stage 4，研究者以持續預訓練的方式，使用僅文字的資料集來擴充套件 LLM 的上下文長度。

在 Stage 5，研究者透過長影片監督微調來增強指令遵循能力。值得注意的是，所有引數在該階段是可訓練的。

實驗結果

研究者從系統和建模兩個方面對本文全棧解決方案進行評估。他們首先展示了訓練和推理結果，從而說明了可支援長上下文訓練和推理的系統實現了效率和可擴充套件性。接著評估了長上下文模型在字幕和指令遵循任務上的表現。

訓練與推理系統

該研究對訓練系統的吞吐量、推理系統的延遲以及支援的最大序列長度進行了定量評估。

表 2 顯示了吞吐量結果。與 ZIGZAG-RINGATTN 相比，本文系統實現了 2.1 倍至 5.7 倍的加速，效能與 DeepSpeed-Ulysses 相當。與 Megatron-LM CP 中更最佳化的環形序列並行實現相比，實現了 3.1 倍至 4.3 倍的加速。
圖片
該研究透過逐步將序列長度從 1k 增加到 10k 來評估固定數量 GPU 支援的最大序列長度，直到發生記憶體不足錯誤。結果總結在圖 9 中。

當擴充套件到 256 個 GPU 時，本文方法可以支援大約 8 倍的上下文長度。此外，所提系統實現了與 ZIGZAG-RINGATTN 類似的上下文長度擴充套件，在 256 個 GPU 上支援超過 200 萬的上下文長度。
圖片
表 3 比較了支援的最大序列長度，該研究提出的方法支援的序列比 HuggingFace Pipeline 支援的序列長 2.9 倍。
圖片
圖 11 展示了長影片大海撈針實驗的結果。相比之下，LongVILA 模型（右）在一系列幀數和深度上都表現出了增強的效能。
圖片
表 5 列出了各種模型在 Video MME 基準上的表現，比較了它們在短影片、中影片和長影片長度上的有效性以及整體效能。LongVILA-8B 採用 256 幀，總分為 50.5。
圖片
研究者還在表 6 對第 3 階段和第 4 階段的影響進行了消融研究。
圖片
表 7 顯示了在不同幀數（8、128 和 256）上訓練和評估的 LongVILA 模型的效能指標。隨著幀數的增加，模型的效能顯著提高。具體來說，平均分數從 2.00 上升到 3.26， 突顯了模型在更多幀數下生成準確豐富字幕的能力。

相關文章