Hugging Face ストリーミングデータセットのアップデート

Hugging Face は load_dataset('dataset', streaming=True) 機能を大幅に強化し、ユーザーがローカルにダウンロードせずにマルチテラバイト規模のデータセットで学習できるようにしました。これらの最適化により、"disk out of space" エラーや 429 のレートリミット応答といった一般的な問題が解消され、ハイコンピュート環境でローカル SSD の速度に匹敵するパフォーマンスを実現します。

大規模でのパフォーマンス向上

ストリーミングバックエンドの改善により、大規模学習で大幅な効率向上が実現しました。64 台の H100 GPU と 256 のワーカーで学習した際、Hugging Face は以下の指標を報告しています。

  • Startup Requests: 最大 100 倍効率的（リクエストが減少）。
  • Data Files Resolution Time: 10 倍高速。
  • Streaming Speed: 最大 2 倍高速。
  • In-flight Requests: 最大 2 倍効率的。
  • Stability: 256 の同時ワーカーでワーカーのクラッシュはゼロ。

技術的最適化

パフォーマンス向上は主に 2 つのフェーズ、起動とストリーミングに分かれています。

起動最適化

各 DataLoader ワーカーが独立してデータセットを初期化することで発生する「リクエストストーム」を防ぐため、Hugging Face は以下を導入しました。

  • Persistent Data Files Cache: 最初のワーカーが Hub からファイルリストを取得し、以降のワーカーはローカルキャッシュから読み込むことで、実質的に冗長な起動リクエストを排除します。
  • Optimized Resolution Logic: ファイルリスト取得に必要な API 呼び出しがより効率的にバンドルされ、レイテンシが低減されました。

ストリーミングスループット

GPU がデータ待ちになることを防ぐため、以下の機能が追加されました。

  • Prefetching for Parquet: ライブラリは現在、モデルが現在のチャンクを処理している間に次のデータチャンクをバックグラウンドで取得します。
  • Configurable Buffering: 上級ユーザーはバッファのブロックサイズやプリフェッチ量を調整できるようになりました。例えば、pyarrow.dataset.ParquetFragmentScanOptions を使用して、デフォルトの 32MiB から 128MiB に最小リクエストサイズを増やすことができます。

データ転送とストレージインフラストラクチャ

Hugging Face は重複排除ベースのストレージシステム Xet を利用してアップロードとダウンロードを高速化しています。Parquet Content Defined Chunking (CDC) を活用することで、重複データは一度だけ転送され、従来のリモートストレージよりも Hugging Face へのアップロードが高速になります。これに加えて、pyspark_huggingface パッケージが Xet と Parquet CDC をサポートした Spark データソースを提供し、データセットの読み書きを可能にしています。

カスタムストリーミングパイプライン

datasets ライブラリでネイティブにサポートされていない形式や、細かな制御が必要なユーザー向けに、Hugging Face は huggingface_hub ライブラリの HfFileSystem を改良しました。これにより、.read()、.readline()、.seek() を使用してデータセットリポジトリからファイルを効率的にリモート読み取りできます。

HfFileSystem を PyTorch の DataLoader に渡すと、.ls() と .glob() のキャッシュ結果を再利用し、データファイルの一覧取得時に追加リクエストが不要になります。この手法は、LeRobot ライブラリでのビデオフレームサンプリングや WebDataset ライブラリでの TAR アーカイブのストリーミングに活用されています。

実際の適用例: nanoVLM

これらの強化は、次世代 SmolVLM 用 nanoVLM の開発中に実戦でテストされました。Hugging Face は、ストリーミングがローカル SSD からの読み取りと同等の速度になり、データをローカル SSD に転送することでトレーニング開始が 3 時間遅れるという以前のボトルネックが解消されたことを確認しました。

Sources