高效的多模態資料管道

你已經準備好了一切——資料、模型、強大的 GPU 配置。你點選“執行”，然後……等待。再等一會兒。你的 GPU 幾乎沒有出什麼力，而你的錢包卻在每小時變輕。

聽起來很熟悉？我們也經歷過。在對我們的 nanoVLM 專案進行了一些調查後，我們發現真正的罪魁禍首不是我們的模型或硬體，而是我們的資料管道效率極其低下。

以下是我們的發現

1. GPU 空閒：我們的模型實際上在等待資料送達
2. 填充地獄：每個批次都塞滿了無用的填充令牌，這些令牌對訓練毫無貢獻

在這篇文章中，我們將分五個階段構建一個高效的管道。在每個階段，我們都會在前一步的基礎上進行增刪，並評論哪些地方做得對，哪些地方做得不好。

目錄：

• 第 0 階段：準備工作
• 第 1 階段：視覺化資料集
• 第 2 階段：樸素填充
• 第 3 階段：約束填充
• 第 4 階段：使用揹包演算法進行更智慧的打包
• 第 5 階段：將揹包演算法應用於多模態資料
• 結論

【第 0 階段】準備工作

為了更容易地跟進資料準備任務，我們建立了一個專門關注資料管道的獨立倉庫。我們希望這比在與 nanoVLM 倉庫整合後閱讀程式碼更容易理解。此外，這對於啟動其他資料管道也可能很有用！

倉庫：https://github.com/ariG23498/mmdp

要跟上進度，您只需克隆該倉庫即可。它包含了最終的資料準備任務，但其設計旨在展示過程中的每一步。

$ git clone https://github.com/ariG23498/mmdp.git

【第 1 階段】視覺化資料集

在最佳化任何東西之前，我們需要了解我們正在處理什麼。我們的多模態資料集包含影像、文字提示和響應。

$ uv run 01_check_dataset.py

熟悉你的訓練資料對於成功至關重要。前面的指令碼每次執行時都會顯示一個隨機樣本；你可能需要將程式碼片段複製到筆記本中並多次執行，以便對資料有一個感覺。

【第 2 階段】樸素填充

我們第一次訓練嘗試採用了顯而易見（且非常常見）的方法

• 對所有內容進行分詞
• 找到每個批次中最長的序列
• 將其他所有內容填充以匹配長度

$ uv run 02_naive_pad_dataloader.py

結果令人痛苦。看看這個視覺化圖

看到那些灰色的部分了嗎？那就是填充。那是 GPU 在你支付計算時間的同時，卻在處理絕對空無一物的東西。我們大約浪費了 60% 的批次在空令牌上。

【第 3 階段】約束填充

我們的下一步很簡單。設定一個全域性最大長度並遵守它。如果一個樣本太長，我們就直接丟棄它。

你可能已經注意到，批次現在少了一個樣本。這是由於過濾過程。這有所幫助，但我們仍然將所有內容填充到相同的固定長度，而不管實際內容如何。比以前好，但仍然很浪費。

【第 4 階段】：使用揹包演算法進行更智慧的打包

現在我們準備完全重新思考批處理。填充是敵人，我們需要一個策略來最小化它，同時最大化每個批次中可以容納的資料量。進入揹包問題，這是計算機科學中的一個經典問題，非常適合這個場景。

想象一下你正在為一次徒步旅行打包揹包。它的承重量有限，而你希望儘可能多地塞進有用的物品。在我們的情況下

• 揹包是一個具有最大令牌限制（max_length）的訓練批次。
• 每個物品是一個序列（一個分詞後的提示-響應對），其重量是令牌的數量。
• 我們的目標是在不超過令牌限制的情況下，將盡可能多的序列打包到批次中，從而最大限度地減少浪費的空間。

為了測試這個想法，我們從一個玩具資料集開始：只是一個從 1 到 25 的數字列表，每個數字代表一個序列長度。這讓我們可以在不涉及影像和文字複雜性的情況下進行實驗。

切換到可迭代資料集

大多數 PyTorch 資料集是 map-style 的（你用 dataset[i] 訪問它們）。但對於動態批處理，我們需要更靈活的東西。因此，我們透過子類化 torch.utils.data.IterableDataset 構建了一個 iterable-style 資料集。這讓我們能夠動態生成批次，並處理像在多個工作程序之間分片資料這樣的技巧。

def _get_data_range(self):
    worker_info = get_worker_info()
    if worker_info is None:  # single worker, return the entire dataset
        return self.start, self.end
    else:  # multiple workers, split the data load
        per_worker = int(
            math.ceil((self.end - self.start) / worker_info.num_workers)
        )
        worker_id = worker_info.id
        iter_start = self.start + worker_id * per_worker
        iter_end = min(iter_start + per_worker, self.end)
        return iter_start, iter_end

生產者-消費者模式的魔力

打包序列可能會很慢，尤其是在我們進行排序或打亂時。為了保持流暢，我們使用生產者-消費者模式，藉助 Python 佇列

def _producer(self, data_iter, queue, stop_signal):
    if self.strategy == "greedy":
        for pack in self._greedy_packing(data_iter):
            queue.put(pack)
    elif self.strategy == "binpack":
        while True:
            buffer = list(itertools.islice(data_iter, self.buffer_size))
            if not buffer:
                break
            knapsacks = self._bin_packing(buffer)
            for pack in knapsacks:
                queue.put(pack)
    queue.put(stop_signal)

生產者執行緒打包批次並將其放入佇列，而主執行緒則根據需要從中取出。這種重疊使得管道能夠順暢流動。

貪婪打包

首先，我們嘗試一種簡單的貪婪打包策略

def _greedy_packing(self, iterator):
    pack, pack_sum = [], 0
    for item in iterator:
        if item > self.max_length:
            continue
        if pack_sum + item <= self.max_length:
            pack.append(item)
            pack_sum += item
        else:
            yield pack
            pack = [item]
            pack_sum = item
    if pack:
        yield pack

這種方法按順序遍歷資料，將專案新增到一個包中直到滿，然後開始一個新的包。它速度快但不完美。以下是批次的樣子

=== Strategy: GREEDY ===
[tensor([1]), tensor([2]), tensor([3]), tensor([4]), tensor([5]), tensor([6]), tensor([7]), tensor([8]), tensor([9]), tensor([10]), tensor([11]), tensor([12]), tensor([13])]
[tensor([14]), tensor([15]), tensor([16]), tensor([17]), tensor([18]), tensor([19])]
[tensor([20]), tensor([21]), tensor([22]), tensor([23])]
[tensor([24])]

注意到後面的批次變得稀疏了嗎？我們留下了空隙。

使用裝箱演算法實現更緊密的擬合

讓我們嘗試一種更智慧的方法：裝箱演算法（具體來說是首次適應遞減法）

def _bin_packing(self, buffer: List[int]):
    buffer = sorted(buffer, reverse=True)
    knapsacks = []
    for item in buffer:
        for pack in knapsacks:
            if sum(pack) + item <= self.max_length:
                pack.append(item)
                break
        else:
            knapsacks.append([item])

這種方法按長度對序列進行排序（最長的在前），並嘗試將每個序列裝入第一個有空間的包中。如果都裝不下，就開啟一個新的包。結果如何？

=== Strategy: BINPACK ===
[tensor([24]), tensor([23]), tensor([22]), tensor([21]), tensor([10])]
[tensor([20]), tensor([19]), tensor([18]), tensor([17]), tensor([16]), tensor([9]), tensor([1])]
[tensor([15]), tensor([14]), tensor([13]), tensor([12]), tensor([11]), tensor([8]), tensor([7]), tensor([6]), tensor([5]), tensor([4]), tensor([3]), tensor([2])]

這些批次要緊密得多，浪費的空間更少。這就像玩俄羅斯方塊一樣，把資料塊緊密地拼接在一起。

【第 5 階段】將揹包演算法應用於多模態資料

現在是動真格的時候了，將揹包打包演算法應用到我們的多模態資料集中。

我們又回到了影像、提示和響應，我們需要在遵守令牌限制和影像預算的同時高效地打包它們。影像預算的設定是為了平衡每個樣本的影像數量。我們希望避免一個 GPU 需要處理比另一個 GPU 多得多的影像的情況。

我們新的 ConstantLengthDataset 類負責繁重的工作。以下是它與第 4 階段相比的工作方式

ConstantLengthDataset 包攬了一切

• 讀取樣本（影像和文字）。
• 過濾掉太長或影像太多的樣本。
• 使用貪婪揹包策略將樣本打包成批次，平衡令牌數和影像數。
• 將最終批次填充到固定長度，但填充量比以前少得多。

這是結果

看那個！灰色（填充）部分非常少，批次裡充滿了有用的資料。這就像打包行李箱一樣，裝得那麼好，你不用坐在上面就能拉上拉鍊。

這張圖乍一看可能不直觀，但讓我們把它和約束填充的圖並排比較一下。

在這裡你會注意到，揹包演算法中的樣本分佈更均勻。我們也不會因為過濾而導致批次中的樣本數量減少。

結論

最初只是一個簡單的“為什麼訓練這麼慢？”的調查，最終導致我們徹底重新思考了處理多模態資料的方式。

用於資料管道的平衡揹包策略來自 NVIDIA 的論文 Eagle 2: Building Post-Training Data Strategies from Scratch for Frontier Vision-Language Models。

關鍵經驗

• 將所有內容填充到最長序列是一個不錯的初步方法（但很浪費）
• 將批處理視為一個打包問題
• 考慮所有約束（文字長度、影像記憶體等）
• 首先用玩具資料測試以驗證你的方法

想深入瞭解嗎？請檢視

• 我們的 nanoVLM 部落格文章
• nanoVLM GitHub
• 這個管道的程式碼

祝你訓練愉快（願你的 GPU 保持忙碌）！