AI & 機械学習

PyTorch の AI トレーニングを促進する、Cloud Storage のための Dataflux Dataset のご紹介

2024年5月14日

Google Cloud Japan Team

※この投稿は米国時間 2024 年 5 月 3 日に、Google Cloud blog に投稿されたものの抄訳です。

はじめに

機械学習（ML）モデルは巨大なデータセット上で育つものですので、費用対効果の高い ML トレーニングには高速でのデータ読み込みが鍵となります。Google は先日、Google の Cloud Storage からのデータの読み込みを高速化する PyTorch Dataset 抽象化である Dataflux Dataset をリリースしました。Dataflux を使用すると、小さいファイルであれば、fsspec と比較して最大 3.5 倍高速でのトレーニングが可能になります。

今回のリリースの構築の基盤をなしているのは、TensorFlow や JAX、TFX、MLIR、KubeFlow、Kubernetes といった OSS への 20 年以上にわたる貢献にずっと関わってきた Google のオープン標準への Google の取り組みと、Project Jupyter や NumFOCUS などの重要な OSS データサイエンスイニシアティブへのスポンサーシップです。

また Google ではディープラーニング IO（DLIO）ベンチマーク上で Dataflux Dataset を検証し、大規模なファイルでも同様のパフォーマンスが得られることを確認しました。これだけ大きくパフォーマンスが向上していますので、トレーニングワークフローのための他のライブラリや Cloud Storage API の直接の呼び出しにも Dataflux Dataset を使用することを Google では推奨します。

Dataflux Dataset の主な機能には次のようなものがあります。

Cloud Storage の直接統合: 最初にデータをローカルにダウンロードする必要がなくなります。
パフォーマンスの最適化: 特に小さいファイルでトレーニング時間が最大 3.5 倍早くなります。
PyTorch Dataset の基本ロール: 使い慣れた PyTorch のコンセプトでシームレスに作業できます。
チェックポイントのサポート: モデルのチェックポイントを Cloud Storageに直接保存したり読み込んだりできます。

Dataflux Datasets を使用する

前提条件: Python 3.8 以上
インストール: $ pip install gcs-torch-dataflux
認証: Google Cloud を使用 application-default authentication

例: トレーニング用の画像の読み込み

Dataflux Dataset を有効化するために必要な変更はほんの少しです。PyTorch をお使いで Cloud Storage にデータがあれば、すでに独自のデータセット実装は作成済みだと思われます。下記のスニペットで、Dataflux Dataset の作成がどれだけ容易かを説明します。詳細は、Google の GitHub ページをご覧ください。

読み込んでいます...

仕組み

Dataflux の大幅なパフォーマンス上昇を実現するため、Google は ML トレーニングのワークフローに生じるデータロード時のパフォーマンスのボトルネックに対処しました。トレーニングの実行中、データはストレージから一括で読み込まれ、いくつかの処理の後に、CPU から ML トレーニングのコンピューティングを行う GPU へ送られます。バッチの読み取りと構築にかかる時間が GPU のコンピューティングより長いと、GPU の機能が事実上停止して使用率が下がるため、トレーニングにかかる時間が長くなります。

クラウドベースのオブジェクトストレージシステム（Google の Cloud Storage など）にデータがある場合、特に小さなオブジェクトにデータが入っていると、ローカルディスクの場合よりもデータのフェッチに時間がかかります。これは Time to First Byte レイテンシによるものです。しかし、オブジェクトがいったん「開く」と、クラウドストレージプラットフォームのスループットが高くなります。Dataflux では、オブジェクトの作成という Cloud Storage 機能を採用しています。これは、多数の小さなオブジェクトをより大きなオブジェクトに動的に組み込める機能です。そうすると、たとえば 1024 個の小さなオブジェクト（バッチサイズ）をフェッチするのではなく、30 個の大きなオブジェクトだけをフェッチしてメモリにダウンロードすればよくなります。その後その大きなオブジェクトが個々の小さなオブジェクトに分解され、データセットサンプルとして返されます。このプロセスで一時的に作成されたオブジェクトも消去されます。

Dataflux Datasets で採用されているもうひとつの最適化機能が、データセットに必要な最初のメタデータの取得を高速化する高スループットの同時リスティングです。Dataflux では、リスティングを大幅に高速化する、ワークスティーリングという高度なアルゴリズムを使用しています。それによって、初めての AI トレーニング（「エポック」）でも、同時リスティングがない Dataflux Datasets と比較して処理が速くなります。オブジェクトの数が数千万に及ぶデータセットでも同様です。

加えて、高速リスティングと動的作成により、Dataflux を使用した ML トレーニングでの GPU の停止が最小限に抑えられ、トレーニングにかかる時間が大幅に削減されると同時に、アクセラレータの使用率が上昇します。

高速リスティングと動的作成は Dataflux クライアントライブラリの一部で、GitHub から入手可能です。Dataflux Dataset ではこれらのクライアントライブラリを内部で使用しています。

Dataflux は今すぐご利用可能

Dataflux Dataset for PyTorch（自社で ML トレーニング用のデータセットコードを作成される場合は Dataflux Python クライアントライブラリ）をぜひお試しいただき、ワークフローがどのくらい加速されたか Google までお知らせください。

この機能やその他の Google のストレージ AI 関連機能の詳細については、Google Cloud Next ’24 の録画セッション「How to define a storage infrastructure for AI and analytical workloads（AI および分析ワークロード向けのストレージインフラストラクチャを定義する方法）」でご確認いただけます。