Connector for PyTorch

Cloud Storage の Connector for PyTorch は、Google がサポートするオープンソース プロダクトで、Cloud Storage と PyTorch を直接統合します。

概要

Connector for PyTorch には、トレーニング中のデータの読み込みと、チェックポインティングとモデルの読み込みの両方にメリットがあります。

トレーニング中のデータ読み込みでの Connector for PyTorch の利点は次のとおりです。

  • Connector for PyTorch の最適化により、主に 1 MB 未満のファイルで構成されるデータセットで、デフォルトの PyTorch よりも最大で 3 倍速いトレーニングを実現できます。
  • Connector for PyTorch は、Cloud Storage バケットからトレーニング データを効率的に読み込むために PyTorch のデータセット プリミティブを実装しています。
  • ランダム データアクセス パターンに対して map-style データセットを、ストリーミング データアクセス パターンに対して iterable-style データセットをそれぞれサポートしています。
  • ダウンロードされた未加工のバイトデータを任意の形式に変換できるため、PyTorch DataLoader で NumPy 配列や PyTorch テンソルを柔軟に処理できます。

チェックポインティングとモデルの読み込みでの Connector for PyTorch の利点は次のとおりです。

  • モデルのチェックポイントを Cloud Storage バケットに直接保存し、モデルのチェックポイントをバケットから読み込むための便利なインターフェースが用意されています。
  • Connector for PyTorch は、PyTorch Lightning の CheckpointIO の DatafluxLightningCheckpoint 実装を使用し、PyTorch Lightning のチェックポインティングをサポートしています。
  • Connector for PyTorch は、PyTorch 分散チェックポインティング用に StorageWriter と StorageReader の実装を提供します。Connector for PyTorch デモライブラリには、PyTorch Lightning FSDP ワークロードでこの機能を使用するサンプルコードが含まれています。
  • コネクタ チェックポインティングでは、Lightning とベースの PyTorch の両方で非同期チェックポイントの保存がサポートされています。

詳細については、Connector for PyTorch の GitHub ランディング ページをご覧ください。

フレームワーク

Connector for PyTorch は、次のフレームワーク バージョンでサポートされています。

  • Python 3.8 以降
  • PyTorch Lightning 2.0 以降
  • PyTorch 2.3.1 以降

ご利用にあたって

Connector for PyTorch を使用するには、次のものが必要です。

  • 操作するデータを含む Cloud Storage バケット。
  • バケットに保存されているデータを操作するための次の権限。
    • storage.objects.create
    • storage.objects.list
    • storage.objects.get
    • storage.objects.delete（作成されたダウンロードを使用する場合）

これらの権限は、Storage オブジェクト ユーザーなどの IAM ロールを使用して、Connector for PyTorch が認証に使用するアカウントに付与する必要があります。

インストール

Connector for PyTorch をインストールするには、次のコマンドを使用します。

pip install gcs-torch-dataflux

構成

Connector for PyTorch のアプリケーションのデフォルト認証情報を使用するには、次のいずれかの方法で認証情報を指定する必要があります。

  • Compute Engine VM で Connector for PyTorch を実行する場合、デフォルトでは、アプリケーションのデフォルト認証情報として VM に接続されているサービス アカウントが自動的に使用されます。詳細については、ワークロードの認証方法を選択するをご覧ください。
  • アプリケーションのデフォルト認証情報は手動で構成することもできます。Google Cloud CLI を使用して直接ログインできます。
gcloud auth application-default login

例

Connector for PyTorch の完全なサンプルセットが Connector for PyTorch GitHub リポジトリの demo ディレクトリにあります。次のようなものがあります。

パフォーマンス

Connector for PyTorch には ML ワークロード用に設計された特別な最適化機能があり、Cloud Storage に対する直接 API 呼び出しよりも非常に優れたパフォーマンスを実現できます。

  • リスティングのパフォーマンスを最適化するため、Connector for PyTorch は高速リスティング アルゴリズムを使用し、並列化されたオブジェクト リスティング プロセス間でリスティング ワークロードの負荷を分散します。
  • 小さなファイルをダウンロードする際のパフォーマンスを最適化するために、Connector for PyTorch は作成オペレーションを実行し、一連の小さなオブジェクトを連結して 1 つの大きなオブジェクトとして処理します。これらの新しい複合オブジェクトは、ソース オブジェクトと同じバケットに保存されます（名前の先頭に dataflux-composed-objects/ が付いています）。
  • チェックポイント書き込みに対するマルチパート アップロードにより、標準のチェックポイント アップロードよりもパフォーマンスを最大で 10 倍向上させることができます。

次の処理に関するパフォーマンス データが GitHub にあります。

  • Lightning テキストベース トレーニング
  • Lightning 画像トレーニング
  • 単一ノードのチェックポインティング
  • マルチノードのチェックポインティング

考慮事項

ワークロードごとに次の点を考慮する必要があります。

高速なリスティング オペレーション

Connector for PyTorch の高速リスティング アルゴリズムにより、Connector for PyTorch は通常のシーケンシャル リスティングよりも多くのリスト オペレーションを使用します。リスト オペレーションは、クラス A オペレーションとして課金されます。

複合オブジェクトの使用

一時的な複合オブジェクトの操作でストレージの過剰な使用料金や早期削除の料金が発生しないようにするには、バケットで次の設定を使用する必要があります。

Connector for PyTorch によって作成された複合オブジェクトは通常、トレーニング ループの終了時に自動的に削除されますが、まれに削除されないことがあります。オブジェクトがバケットから確実に削除されるようにするには、次のコマンドを実行します。

gcloud storage rm gs://<my-bucket>/dataflux-composed-objects/ --recursive

複合オブジェクトの使用を無効にするには、作成するデータセットの config 部分に disable_compose=True または max_composite_object_size=0 を含めます。ただし、この動作をオフにすると、トレーニング ループの時間が大幅に長くなる可能性があります（特に小さなファイルで作業する場合）。

複合オブジェクトを使用すると、Cloud Storage は、ファイルを直接ダウンロードする場合よりも低いスケールで QPS とスループットの上限に達します。複合オブジェクトを使用しない場合でもプロジェクトの QPS またはスループットの上限に達するようなマルチノード スケールで実行する場合は、複合オブジェクトの使用を無効にする必要があります。

429 エラーとパフォーマンスの低下

Connector for PyTorch の使用中に、429 エラーが発生したり、実行時間が想定よりも長くなることがあります。これには、次のような一般的な理由があります。

  • 多くの ML プロジェクトでは、PyTorch Lightning や Ray などのツールを活用した高度な分散トレーニング モデルが使用されています。これらのモデルは Connector for PyTorch と互換性がありますが、多くの場合、Cloud Storage のレート制限がトリガーされます。
  • 429 エラーが発生し、「このワークロードが Cloud Storage から送信帯域幅を過剰に使用している」、「このワークロードが Cloud Storage の下り（外向き）帯域幅の上限をトリガーした」などのメッセージが返された場合は、ワークロードのデータ スループット レートが Google Cloud プロジェクトの最大容量を超えていることを示しています。この問題の対応策としては、次の対処法をお試しください。
  • QPS の上限により、TooManyRequests を示すメッセージとともに 429 エラーがトリガーされることがありますが、通常は、実行時間が想定よりも遅くなることでこの現象を確認できます。多くの場合、小さなファイルを大量に処理すると QPS のボトルネックが発生します。バケット QPS の上限は時間の経過とともに自然にスケーリングされるため、ウォームアップ期間を設定することでパフォーマンスが向上することがあります。ターゲット バケットのパフォーマンスの詳細を確認するには、Google Cloud コンソールでバケットを表示し、[オブザーバビリティ] タブを確認します。
  • エラー メッセージに dataflux-composed-objects というキーワードを含む TooManyRequests エラーでワークロードが失敗する場合は、作成されたオブジェクトの使用を無効にすることをおすすめします。これにより、作成オペレーションが大規模に実行された場合の QPS の負荷を軽減できます。

メモリ消費

アップロードとダウンロードのパフォーマンスを最適化するため、推論用の最終モデルを含むチェックポイントの書き込みと読み込みはメモリに完全にステージングされます。これらのパフォーマンス向上を活用できるように、各マシンにはチェックポイントをメモリにステージングするのに十分な空き RAM が必要です。

サポートを利用する

公式の Google Cloud サポート チャネルのいずれかで、サポートの利用、一般的な質問の提出、新機能のリクエストを行うことができます。GitHub で問題を報告してサポートを受けることもできます。

PyTorch、PyTorch ロゴ、および関連するマークは、The Linux Foundation の商標です。