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

然而，目前面臨的問題是一些工作已經啟用了長上下文視覺語言模型（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， 突顯了模型在更多幀數下生成準確豐富字幕的能力。