Hugging Face と Dask による AI データ処理のスケーリング

Hugging Face と Dask は、ローカルメモリの上限を超える大規模 AI データセットを処理するためのスケーラブルなフレームワークを提供します。Dask の分散コンピューティング機能と Hugging Face の transformers と datasets を組み合わせることで、モデル推論やデータフィルタリングといった AI タスクを、ノートパソコン上の数行からマルチ GPU クラウドクラスター上の数億行へと拡張できます。

Dask を用いた分散データ処理

Dask はアウト・オブ・コア計算を可能にし、システムメモリに収まりきらないデータセットを管理しやすいチャンクに分割して処理できます。pandas に慣れ親しんだユーザーにとっては、Dask DataFrame が類似の API を提供するため、ローカルでのプロトタイピングから大規模本番環境への移行が容易になります。

Dask を AI データ処理に利用する主な利点は次のとおりです:

  • Efficient Loading: Dask は Hugging Face データセットのデフォルト形式である Parquet とネイティブに連携し、効率的な列指向フィルタリングと圧縮を実現します。
  • Parallel Execution: map_partitions 関数を使うことで、カスタム関数（例: モデル推論）を大きな Dask DataFrame 内の各 pandas DataFrame パーティションに対して並列に適用できます。
  • Distributed Writing: Dask は結果を Parquet 形式で並列に書き出すことをサポートし、Hugging Face データセットリポジトリと統合できます。

モデル推論のスケーリング: Pandas から Dask へ

スケーリングの実例として、Hugging Face は FineWeb データセット（英語ウェブデータ 15 兆トークン）と FineWeb‑Edu 分類器を使用し、教育価値の高いウェブページを特定しました。

Pandas でのローカルプロトタイピング

小規模（例: 100 行）では pandas を使って FineWeb-Edu 分類器を実行できます。GPU 搭載の M1 Mac ではこの処理に約 10 秒かかります。ワークフローは Hugging Face の pipeline を用いたテキスト分類で、ハードウェアデバイス（CUDA、MPS、CPU）は関数内部で動的に選択され、後でコードを分散実行する際の互換性が保たれます。

2.11 億行へのスケーリング

2.11 億行（ディスク上で 432 GB のクローラーの一部）にスケールすると、シリアル処理は実質的に遅くなります。そこで Dask DataFrame に切り替え、Hugging Face からデータを遅延ロードすることでタスクを並列化しました。

このスケールされたワークフローでは、compute_scores 関数を map_partitions 経由で適用します。パフォーマンス最適化のため、Hugging Face パイプラインの batch_size を（例: 768）に増やし、GPU ハードウェアをより有効に活用します。

クラウド上のマルチ GPU 並列推論

最大スループットを得るには、Dask をクラウドインフラ上にデプロイします。例では Coiled を使用し、AWS の g5.xlarge インスタンス（NVIDIA A10 Tensor Core GPU）を搭載した 100 ワーカーのクラスターをプロビジョニングしました。

インフラ自動化

Coiled は以下の重要なデプロイ手順を自動化します:

  • VM Provisioning: GPU 対応のクラウド VM を自動で起動。
  • Environment Setup: NVIDIA ドライバと CUDA ランタイムのインストールを処理。
  • Package Synchronization: ローカルの Python パッケージやファイルをクラウドワーカーに同期し、環境の一貫性を確保。

パフォーマンスと利用率

2.11 億行の処理には約 5 時間かかりました。モニタリング結果はハードウェアが効率的に使用されていることを示し、GPU の中央値利用率は 100%、メモリ使用率は 24 GB 中 21.5 GB（中央値）でした。

想定される AI 活用例

この分散処理パターンは、テキスト分類以外の大規模 AI タスクにも応用できます。例としては:

  • Genomic Data Filtering: 大規模ゲノムデータセットから関心遺伝子を抽出。
  • Structured Data Extraction: LLM を用いて非構造化テキストを構造化データに変換。
  • Web Data Cleaning: Common Crawl から取得した大規模スクレイプデータのクリーンアップとフィルタリング。
  • Multimodal Inference: マルチモーダルモデルで大規模音声、画像、動画データセットを解析。

Sources