Hugging Face 高效多模態資料管線

Hugging Face 開發了一套高效的多模態資料管線，以解決視覺語言模型（VLM）訓練中的關鍵瓶頸，特別是 GPU 利用率不足與過度填充。透過將批次處理視為打包問題，團隊實作了平衡背包策略，最大化每個批次的資料密度，並確保 GPU 之間的影像分配均衡。

問題：GPU 利用率不足與填充浪費

低效的資料管線常會導致「GPU 閒置」，硬體在等待資料時未被充分利用。主要原因是「填充地獄」，即批次被填入大量無用的填充 token，以匹配批次中最長序列的長度。在 nanoVLM 專案的初步測試中，Hugging Face 發現天真的填充方式會浪費約 60% 的批次空間於空白 token，造成大量計算資源浪費與成本上升。

資料管線策略的演變

Hugging Face 透過五個迭代階段優化其管線，從天真方法逐步過渡到精密的打包系統。

階段 1 與 2：可視化與天真填充

最初的工作聚焦於可視化資料集（影像、文字提示與回應）並套用天真填充。在此方法中，批次中的每個序列都會被填充至與最長序列相同長度。這會導致大量浪費，因為 GPU 必須處理大量空白 token。

階段 3：受限填充

為了減少浪費，實作了全域最大長度。超過此長度的樣本會被剔除。雖然此舉減少了一部分填充，但仍需將所有序列填充至固定長度，無論實際內容為何，且因過濾而導致訓練樣本的流失。

階段 4：文字的背包打包

為了消除填充，Hugging Face 將「背包問題」的概念套用於批次處理。目標是在不超過最大 token 限制（max_length）的前提下，將盡可能多的序列（項目）打包進一個批次（背包）中。

使用玩具資料測試了兩種策略：

  • 貪婪打包：序列依序加入，直至批次滿為止。此方法快速，但常在後續批次留下空隙。
  • 箱子打包（先適配遞減）：先依長度排序（由長到短），再放入第一個有空間的包裹。此方式可產生更緊密的批次，最小化浪費空間。

為支援此動態批次，團隊將資料集從 map‑style 轉換為 IterableDataset，並使用 Python 隊列實作生產者‑消費者模式，以確保打包過程不會成為 GPU 的瓶頸。

階段 5：多模態資料的平衡背包

最終階段將這些概念套用於多模態資料，結合 token 限制與影像預算。此舉確保每個樣本的影像在 GPU 之間均衡分配，避免單一 GPU 處理過多影像資料。

ConstantLengthDataset 類別透過以下方式管理此流程：

  1. 讀取影像與文字。
  2. 過濾超出 token 或影像限制的樣本。
  3. 使用平衡貪婪背包策略打包樣本。
  4. 僅對最終批次進行固定長度的填充，從而使整體填充量最小化。

打包策略的技術比較

概念 階段 4（玩具資料） 階段 5（多模態資料）
項目 整數（序列長度） 完整樣本（影像、提示、回應）
權重 整數本身 Token 數量（len(input_ids)）
背包 整數批次 $≤$ max_length 樣本批次 $≤$ seq_length 且受影像限制
打包策略 貪婪或箱子打包 在 token 與影像限制下的貪婪打包
輸出 整數列表 字典，包含 input_ids、labels、attention_mask、images

結論與來源

資料管線的平衡背包策略來源於 NVIDIA 的研究論文《Eagle 2：從頭構建前沿視覺語言模型的後訓練資料策略》。透過從天真填充轉向受限且平衡的打包方法，開發者能確保 GPU 完全被利用，且訓練成本更具效益。

Sources