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

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

全棧設(shè)計對于長上下文視覺語言模型至關(guān)重要。訓練大型模型通常是一項復雜而系統(tǒng)的工作，需要數(shù)據(jù)工程和系統(tǒng)軟件協(xié)同設(shè)計。與純文本 LLM 不同，VLM（例如 LLaVA）通常需要獨特的模型架構(gòu)和靈活的分布式訓練策略。

此外，長上下文建模不僅需要長上下文數(shù)據(jù)，還需要能夠支持內(nèi)存密集型長上下文訓練的基礎(chǔ)設(shè)施。因此，對于長上下文 VLM 來說，精心規(guī)劃的全棧設(shè)計（涵蓋系統(tǒng)、數(shù)據(jù)和 pipeline）是必不可少的。

本文，來自英偉達、MIT、UC 伯克利、得克薩斯大學奧斯汀分校的研究者引入了 LongVILA，這是一種用于訓練和部署長上下文視覺語言模型的全棧解決方案，包括系統(tǒng)設(shè)計、模型訓練策略和數(shù)據(jù)集構(gòu)建。

• 論文地址：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

對于訓練基礎(chǔ)設(shè)施，該研究建立了一個高效且用戶友好的框架，即多模態(tài)序列并行 (MM-SP)，它支持訓練記憶 - 密集型長上下文 VLM。

對于訓練 pipeline，研究者實施了一個五階段訓練流程，如圖 1 所示：即 (1) 多模態(tài)對齊，(2) 大規(guī)模預訓練，(3) 短監(jiān)督微調(diào)，(4) LLM 的上下文擴展，以及 (5) 長監(jiān)督微調(diào)。

對于推理，MM-SP 解決了 KV 緩存內(nèi)存使用率的挑戰(zhàn)， 這在處理非常長的序列時會成為瓶頸。

通過使用 LongVILA 增加視頻幀數(shù)，實驗結(jié)果表明該研究在 VideoMME 和長視頻字幕任務上的性能持續(xù)提高（圖 2）。在 1024 幀上訓練的 LongVILA 模型在 1400 幀的大海撈針實驗中實現(xiàn)了 99.5% 的準確率，相當于 274k 個 token 的上下文長度。此外， MM-SP 系統(tǒng)可以有效地將上下文長度擴展到 200 萬個 token 而無需梯度檢查點，與環(huán)形序列并行（ring sequence parallelism）相比實現(xiàn)了 2.1 倍至 5.7 倍的加速，與 Megatron 上下文并行 + 張量并行相比實現(xiàn)了 1.1 倍至 1.4 倍的加速。

下圖為 LongVILA 技術(shù)在處理長視頻字幕時的示例：在字幕開頭，8 幀的基線模型僅描述了靜態(tài)圖像和兩輛車。相比之下，256 幀的 LongVILA 描述了雪地上的汽車，包括車輛的前、后和側(cè)面視圖。在細節(jié)上，256 幀的 LongVILA 還描述了點火按鈕、變速桿和儀表盤的特寫，這些在 8 幀的基線模型中是缺失的。

多模態(tài)序列并行

訓練長上下文視覺語言模型（VLM）會產(chǎn)生大量內(nèi)存需求。例如下圖 1 中 Stage 5 的長視頻訓練，單個序列包含了產(chǎn)生 1024 個視頻幀的 200K tokens，這超出了單個 GPU 的內(nèi)存容量。

研究者開發(fā)了一個基于序列并行的定制系統(tǒng)。序列并行是當前基礎(chǔ)模型系統(tǒng)中常用的一種技術(shù)，用于優(yōu)化僅文本的 LLM 訓練。不過，研究者發(fā)現(xiàn)現(xiàn)有系統(tǒng)既不高效，擴展性也不足以處理長上下文 VLM 工作負載。

在確定現(xiàn)有系統(tǒng)的局限性之后，研究者得出結(jié)論，一個理想的多模態(tài)序列并行方法應該通過解決模態(tài)和網(wǎng)絡(luò)異構(gòu)性來優(yōu)先實現(xiàn)效率和可擴展性，并且擴展性不應受到注意力頭數(shù)量的限制。

MM-SP 工作流。為了應對模態(tài)異構(gòu)性的挑戰(zhàn)，研究者提出了一種兩階段式分片策略，以優(yōu)化圖像編碼和語言建模階段的計算工作負載。

具體如下圖 4 所示，第一階段首先在序列并行進程組內(nèi)的設(shè)備之間均勻地分布圖像（比如視頻幀），從而在圖像編碼階段實現(xiàn)負載平衡。第二階段，研究者聚合全局視覺和文本輸入以進行 token 級分片。

2D 注意力并行。為了解決網(wǎng)絡(luò)異構(gòu)性并實現(xiàn)可擴展性，研究者結(jié)合環(huán)形（Ring）序列并行和 Ulysses 序列并行的優(yōu)勢。

具體來講，他們將跨序列維或注意力頭維的并行視為「1D SP」。該方法通過跨注意力頭和序列維的并行計算來實現(xiàn)擴展，將 1D SP 轉(zhuǎn)換為由獨立的 Ring（P2P）和 Ulysses（A2A）進程組組成的 2D 網(wǎng)格。

以下圖 3 左所示，為了實現(xiàn)跨 2 個節(jié)點的 8-degree 序列并行，研究者使用 2D-SP 構(gòu)建了一個 4×2 通信網(wǎng)格。

此外，在下圖 5 中，為了進一步解釋 ZIGZAG-RINGATTN 如何平衡計算以及 2D-Attention 機制如何運作，研究者解釋了使用不同方法的注意力計算計劃。

與 HuggingFace 的原生 pipeline 并行策略相比，本文的推理模式更加高效，原因在于所有設(shè)備同時參與計算，從而與機器數(shù)量呈正比地加速進程，具體如下圖 6 所示。同時，該推理模式是可擴展的，內(nèi)存均勻地分布給各個設(shè)備，以使用更多機器來支持更長的序列。

LongVILA 訓練流程

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

在 Stage 1，只有多模態(tài)映射器可以訓練，其他映射器被凍結(jié)。

在 Stage 2，研究者凍結(jié)了視覺編碼器，并訓練了 LLM 和多模態(tài)映射器。

在 Stage 3，研究者針對短數(shù)據(jù)指令遵循任務對模型全面進行微調(diào)，比如使用圖像和短視頻數(shù)據(jù)集。

在 Stage 4，研究者以持續(xù)預訓練的方式，使用僅文本的數(shù)據(jù)集來擴展 LLM 的上下文長度。

在 Stage 5，研究者通過長視頻監(jiān)督微調(diào)來增強指令遵循能力。值得注意的是，所有參數(shù)在該階段是可訓練的。

實驗結(jié)果

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

訓練與推理系統(tǒng)

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

表 2 顯示了吞吐量結(jié)果。與 ZIGZAG-RINGATTN 相比，本文系統(tǒng)實現(xiàn)了 2.1 倍至 5.7 倍的加速，性能與 DeepSpeed-Ulysses 相當。與 Megatron-LM CP 中更優(yōu)化的環(huán)形序列并行實現(xiàn)相比，實現(xiàn)了 3.1 倍至 4.3 倍的加速。

該研究通過逐步將序列長度從 1k 增加到 10k 來評估固定數(shù)量 GPU 支持的最大序列長度，直到發(fā)生內(nèi)存不足錯誤。結(jié)果總結(jié)在圖 9 中。

當擴展到 256 個 GPU 時，本文方法可以支持大約 8 倍的上下文長度。此外，所提系統(tǒng)實現(xiàn)了與 ZIGZAG-RINGATTN 類似的上下文長度擴展，在 256 個 GPU 上支持超過 200 萬的上下文長度。

表 3 比較了支持的最大序列長度，該研究提出的方法支持的序列比 HuggingFace Pipeline 支持的序列長 2.9 倍。

圖 11 展示了長視頻大海撈針實驗的結(jié)果。相比之下，LongVILA 模型（右）在一系列幀數(shù)和深度上都表現(xiàn)出了增強的性能。

表 5 列出了各種模型在 Video MME 基準上的表現(xiàn)，比較了它們在短視頻、中視頻和長視頻長度上的有效性以及整體性能。LongVILA-8B 采用 256 幀，總分為 50.5。

研究者還在表 6 對第 3 階段和第 4 階段的影響進行了消融研究。

表 7 顯示了在不同幀數(shù)（8、128 和 256）上訓練和評估的 LongVILA 模型的性能指標。隨著幀數(shù)的增加，模型的性能顯著提高。具體來說，平均分數(shù)從 2.00 上升到 3.26， 突顯了模型在更多幀數(shù)下生成準確豐富字幕的能力。