mosaicml/streaming

A Data Streaming Library for Efficient Neural Network Training

What it solves

訓練大規模 AI 模型通常需要數據集太大，無法放入本地磁碟，這會導致從雲端存儲加載數據時出現瓶頸。本專案提供了一個高性能的流式傳輸庫，允許模型即時獲取數據，從而減少對大規模本地存儲的需求，並消除在恢復訓練運行時的漫長等待時間。

How it works

StreamingDataset 可作為 PyTorch 的 IterableDataset 的直接替換方案。它將原始數據轉換為稱為 Mosaic Data Shard (MDS) 的專用格式，然後上傳到雲端存儲（AWS, GCS, Azure, 等）。在訓練期間，該庫將數據分片 (shards) 流式傳輸到本地緩存，僅下載當前步驟所需的數據。它使用確定性洗牌算法來確保無論在集群中使用了多少個 GPU 或節點，數據交付的順序都保持一致。

Who it’s for

需要在多節點、分布式集群上訓練大模型（LLMs, Diffusion models, 等）的機器學習工程師和研究人員，他們需要高效地處理存儲在雲端的超大規模數據集。

Highlights

  • True Determinism: 確保在不同硬件配置下具有相同的樣本排序，使調試和復現更加容易。
  • Instant Mid-Epoch Resumption: 允許在故障後幾秒鐘內恢復訓練，避免了大規模數據集常見的漫長數據加載等待時間。
  • Seamless Data Mixing: 支持即時以特定比例或絕對樣本數量進行多個數據集的混合。
  • High Throughput: 優化的 MDS 格式提供比傳統加載器更低的樣本延遲和更高的吞吐量。
  • Random Access: 支持通過索引訪問任何特定樣本，即使它尚未被下載。
  • Disk Usage Limits: 包含一個 cache_limit 功能，通過刪除最近最少使用 (LRU) 的分片來自動管理本地磁碟空間。