Hugging Face 串流資料集更新

Hugging Face 已顯著增強 load_dataset('dataset', streaming=True) 功能，使用戶能夠在不在本地下載的情況下訓練多 TB 資料集。這些優化消除了常見問題，例如「磁碟空間不足」錯誤和 429 速率限制回應，提供的效能可匹配高計算環境中的本地 SSD 速度。

大規模效能提升

對串流後端的改進為大規模訓練帶來了顯著的效率提升。在使用 256 個工作進程的 64xH100 GPU 上進行訓練時，Hugging Face 報告了以下指標：

  • 啟動請求： 最多提高 100 倍效率（請求更少）。
  • 資料檔案解析時間： 提升 10 倍速度。
  • 串流速度： 最多提升 2 倍。
  • 進行中的請求： 最多提高 2 倍效率。
  • 穩定性： 在 256 個併發工作進程時零工作進程崩潰。

技術優化

效能提升分為兩個主要階段：啟動和串流。

啟動優化

為防止「請求風暴」——每個 DataLoader 工作進程獨立初始化資料集，Hugging Face 推出了：

  • 持續資料檔案快取： 第一個工作進程從 Hub 解析檔案清單，所有後續工作進程從本地快取讀取，幾乎消除了重複的啟動請求。
  • 優化解析邏輯： 獲取檔案清單所需的 API 呼叫現在被更有效地捆綁，以降低延遲。

串流吞吐量

為確保 GPU 不會因等待資料而閒置，新增了以下功能：

  • Parquet 預取： 庫現在會在背景中預取下一個資料塊，同時模型處理當前塊。
  • 可設定的緩衝區： 進階使用者現在可以調整緩衝區的區塊大小和預取體積。例如，可以使用 pyarrow.dataset.ParquetFragmentScanOptions 將最小請求大小從預設的 32MiB 增加到 128MiB。

資料傳輸與存儲基礎設施

Hugging Face 使用 Xet，一種基於去重的儲存系統，來加速上傳和下載。透過利用 Parquet 內容定義分塊 (CDC)，重複的資料僅傳輸一次，使上傳到 Hugging Face 的速度快於傳統遠端儲存。這進一步由 pyspark_huggingface 套件支援，該套件提供了一個 Spark 資料來源，用於讀取和寫入具有 Xet 和 Parquet CDC 支援的資料集。

自訂串流管線

對於 datasets 函式庫不原生支援的格式，或需要細粒度控制的使用者，Hugging Face 改進了 huggingface_hub 函式庫中的 HfFileSystem。這使得可以使用 .read()、.readline() 和 .seek() 高效遠端讀取資料庫儲存庫中的檔案。

當 HfFileSystem 傳遞給 PyTorch 的 DataLoader 時，它會重複使用來自 .ls() 和 .glob() 的快取結果，在列出資料檔案時無需額外請求。此方法已在 LeRobot 函式庫中用於影像幀採樣，以及在 WebDataset 函式庫中用於串流 TAR 檔案。

實際應用：nanoVLM

這些增強功能在開發 nanoVLM 以迎接 SmolVLMs 下一代時經過實戰檢驗。Hugging Face 發現串流現在與讀取本地 SSD 同樣快速，消除了之前將資料傳輸到本地 SSD 會導致訓練啟動延遲三小時的瓶頸。

Sources