TPU VM での AI/ML ワークロードのストレージに関するベスト プラクティス

TPU VM で AI/ML ワークロードのパフォーマンスと費用対効果を最大化するには、ワークロードに適したストレージ ソリューションを選択して構成します。I/O ボトルネックを解消することで、TPU アクセラレータのアイドル時間を短縮し、トレーニング時間とコストを削減できます。

このドキュメントでは、TPU VM でのトレーニング、チェックポイント設定、サービング、キャッシュ保存に関するワークロード固有のストレージの推奨事項と最適化のベスト プラクティスについて説明します。これらのプラクティスを適用する前に、利用可能な TPU データのストレージ オプションを確認してください。このドキュメントは、TPU VM に精通しており、Cloud Storage リソースのプロビジョニングの基本的な経験があることを前提としています。

ワークロード固有のガイダンス

次の表に、ワークロードごとのストレージの推奨事項を優先順位の高い順に示します。

ワークロード 推奨事項 最適化とツール(該当する場合)
トレーニング データセット(データ準備を含む)
  1. Cloud Storage リージョン バケットRapid Cache およびクライアントサイド チューニング(先読みなど)を組み合わせて使用すると、コストを最小限に抑えることができます。
  2. 最適なパフォーマンスとスケーリングを実現するには、Cloud Storage Rapid Bucket を使用します。
  3. メタデータの同時実行性が高い場合や、ファイルが小さい場合(1 MB 未満)に Lustre(ファイル システム ストレージ)を標準化する場合は、Managed Lustre を使用します。
  • リージョン Cloud Storage バケットの場合は、階層型名前空間を有効にします。
  • リージョン Cloud Storage バケットまたは Rapid Bucket の場合は、Cloud Storage FUSE と gRPC を使用します。GKE の場合は、gcsfusecsi-training プロファイルを使用します。
  • Managed Lustre の場合は、動的ティアを使用して費用を削減し、パフォーマンスを自動的に最適化することを検討してください。
チェックポイントと強化学習の重み
  1. Managed Lustre は、同期チェックポイントや強化学習の高速重み伝播など、低レイテンシ(1 ミリ秒未満)の要件を満たすために大規模なチューニングを行わずに高パフォーマンスを実現します。
  2. 非同期およびマルチティアのチェックポイント処理と、ストレージへのチェックポイント オペレーションの並列化で最高のスループットを実現するには、Cloud Storage Rapid Bucket を使用します。Rapid Bucket は、これらのワークロードに必要な高スループットを提供します。
Rapid Bucket の場合は、gcsfusecsi-checkpointing プロファイルを使用します。
モデルの保存とダウンロード
  1. 小規模なモデルのキャッシュ保存と最大 2,500 個のインスタンスには、Hyperdisk ML を使用します。
  2. 単一のストレージ ソリューションを標準化する場合は、トレーニング データセットに使用するストレージ ソリューション(Managed Lustre、Rapid Bucket、リージョン バケット)と同じものを使用します。
モデルをダウンロードするには、GKE Run:ai Model Streamer または Cloud Storage FUSE を使用し、gcsfusecsi-serving プロファイルを使用して別のマウント ポイントを使用します。
Key-Value(KV)キャッシュのオフロード
  1. 常にホスト CPU メモリ(RAM)をプライマリ ティアとして使用し、レイテンシを大幅に増やすことなく VRAM を解放します。
  2. CPU メモリが不足している場合は、Managed Lustre をキャッシュ階層の並列低レイテンシの第 2 階層として使用して、高スループット(最大 10 TB/秒)での KV キャッシュ オフロードの超低レイテンシ(1 ミリ秒未満)のニーズを満たします。

Cloud Storage の最適化

以降のセクションでは、TPU VM で Cloud Storage を使用する際のパフォーマンスを最適化するためのベスト プラクティスについて説明します。

メタデータの最適化のために階層型名前空間を有効にする

メタデータのパフォーマンスを向上させるには、AI/ML ワークロード用のリージョン バケットを作成するときに、階層型名前空間を有効にします。メタデータのパフォーマンスとは、ファイルの内容の読み取りや書き込みではなく、オブジェクト パスとフォルダの検索、一覧表示、変更に関連するオペレーションを Cloud Storage が処理する速度を指します。

階層型名前空間が有効になっていないバケットでは、フォルダは実際のリソースとして存在せず、スラッシュ(/)で区切られたオブジェクト名の接頭辞で表されるシミュレートされたフォルダです。このため、システムはその接頭辞を持つすべてのオブジェクトをスキャンする必要があるため、ディレクトリの内容の一覧表示やディレクトリの名前変更などのオペレーションが遅くなります。階層名前空間は、AI/ML ワークロードにとって重要な真のファイル システム構造を提供します。その理由は次のとおりです。

  • アトミック ディレクトリ名の変更: ML フレームワークは、ディレクトリ名の変更を使用してチェックポイントを確定します。階層型名前空間はアトミックな名前変更をサポートしているため、チェックポイントを迅速に確定できます。
  • 初期 QPS の増加: 階層型名前空間では、階層型名前空間が有効になっていないバケットと比較して、読み取りと書き込みの初期秒間クエリ数(QPS)が最大 8 倍になります。これにより、多くの TPU が同時にストレージにアクセスする際のボトルネックを防ぐことができます。
  • 効率的なフォルダレベルのオペレーション: 特定のディレクトリ内のファイルの検索と一覧表示が大幅に高速化され、トレーニングとデータ読み込み中のレスポンス タイムが短縮されます。

Rapid Bucket で提供されるゾーンバケットは、デフォルトで階層型名前空間を使用します。詳細については、階層型名前空間の概要をご覧ください。

適切なプロファイルで Cloud Storage FUSE を使用する

Cloud Storage FUSE は、バケットをローカル ファイル システムとしてマウントできる FUSE アダプタです。Google Kubernetes Engine を使用する場合は、Cloud Storage FUSE CSI ドライバCloud Storage FUSE プロファイルを使用して、パフォーマンス チューニングを自動化することをおすすめします。

Cloud Storage FUSE の使用に関するベスト プラクティスの詳細については、パフォーマンス調整のベスト プラクティスをご覧ください。

TPU VM ブートディスクをカスタマイズする

起動スクリプトを使用するか、カスタム イメージを作成することで、TPU VM のゲスト OS 環境をカスタマイズできます。ブートディスクのカスタマイズは、次のようなシナリオで役立ちます。

  • ソフトウェアとライブラリのプリロード: 特定の ML フレームワーク、依存関係、カスタム ソフトウェアをインストールして、VM の起動時間を短縮し、一貫した環境を確保します。
  • 標準以外の OS ディストリビューションを使用する: Google 管理リストに含まれていない OS ディストリビューションまたはバージョンを使用します。
  • セキュリティとモニタリングの構成を適用する: カスタム セキュリティ設定の適用、モニタリング エージェントのインストール、環境変数の設定を行います。

ただし、TPU VM のブートディスクの復元は制限されています。オフライン修復のためにブートディスクを切り離したり、スナップショットを作成したりすることはできません。そのため、ブートプロセスに影響する変更を行う際は注意してください。これらのベスト プラクティスに従うことで、TPU VM 環境をカスタマイズする際の起動失敗のリスクを軽減できます。

ブートディスクをカスタマイズする際は、次の主要原則に留意してください。

  • ブートディスクの変更を最小限に抑える: 可能な限り、ブートディスクを大幅に変更するのではなく、Persistent Disk ボリュームまたは Hyperdisk ボリュームにアプリケーションをインストールしてデータを保存します。

  • マウントに UUID を使用する: /etc/fstab ファイルにエントリを追加する場合は、/dev/sdb1 などのデバイス名ではなく、常に UUID を使用してディスクとパーティション(UUID=...)を識別します。自動生成されたデバイス名は、再起動間で安定しているとは限りません。

システム変更時の起動障害のリスクを軽減するには、次の推奨事項に従ってください。

  • エラー処理: スクリプトに堅牢なエラーチェックとグレースフル フェイルオーバー モードを実装します。デバッグに役立つように、詳細なメッセージをシリアル コンソールと Cloud Logging に記録します。

  • 重要な依存関係: /etc/fstab ファイル、ネットワーク構成、ブートローダー設定など、起動に不可欠なファイルを変更する際は、細心の注意を払ってください。構文エラーや誤ったエントリがあると、VM が起動できなくなる可能性があります。

  • セカンダリ ディスク: スクリプトがセカンダリ ディスクに依存している場合は、ディスクが存在しない場合や、接続に予想よりも時間がかかる場合を処理するようにします。絶対に必要な場合を除き、ブートプロセスがセカンダリ ディスクのマウントに大きく依存しないようにします。

    セカンダリ ディスクのマウントで推奨される /etc/fstab エントリと推奨されない /etc/fstab エントリの例を次に示します。

    • 推奨: UUID=a1b2c3d4-e5f6-7890-1234-567890abcdef /mnt/mydata ext4 defaults,nofail 0 2
    • 非推奨: /dev/sdb1 /mnt/mydata ext4 defaults 0 2

    nofail オプションを使用すると、ディスクが見つからない場合にシステムが停止するのを防ぐことができますが、マウント ポイントが使用できない場合でもアプリケーションが処理できるようにしてください。

  • パッケージ管理: サードパーティ リポジトリを追加する際は注意してください。信頼できるものであり、ベース OS イメージと互換性があることを確認します。インストールするパッケージの依存関係と、システム ライブラリに与える可能性のある影響を把握します。

  • ディスク容量: ブートディスクの使用量をモニタリングします。広範なロギングや大規模なソフトウェアのインストールにより、ブートディスクが満杯になり、VM が起動できなくなることがあります。

  • ロギング: TPU VM での起動に関する問題を診断するための主なツールであるシリアル コンソールに詳細なログを記録するように、アプリケーションとスクリプトを構成します。

ストレージ容量を計画する

アクセラレータを最大限に活用するために、ワークロードに必要なストレージ容量を計画することが重要です。これには、ストレージ容量とチェックポイント帯域幅が含まれます。

ストレージを見積もる

次の見積もりをストレージ要件の出発点として使用できます。

ワークロード タイプ データセット ストレージ チェックポイント ストレージ
LLM の事前トレーニング TPU あたり 2 TB TPU あたり 200 GB
マルチモーダル トレーニング TPU あたり 12 TB TPU あたり 1 TB
推論 TPU あたり 1 TB TPU あたり 1 GB

チェックポイントの帯域幅を推定する

次の式を使用して、トレーニング ワークロードに必要な最小チェックポイント帯域幅を見積もることができます。データ読み取り、複数のトレーニング実行、トレーニングと推論の場合は、推定帯域幅要件を比例して増やします。

  1. チェックポイント サイズ: パラメータ数 × パラメータあたりのバイト数(FP16 + オプティマイザーの状態の場合、パラメータあたり約 12 ~ 16 バイト)。オプティマイザーの状態とさまざまな精度に対してバッファ(約 3 倍)を追加します。
  2. チェックポイント間隔: チェックポイントを保存する頻度(15 分ごとなど)。
  3. 必要な帯域幅: チェックポイント サイズ ÷ チェックポイント間隔

次の例は、Qwen3-72B の最小チェックポイント帯域幅を推定する方法を示しています。

  1. チェックポイント サイズ: 72B パラメータ × 12 バイト ≈ チェックポイントあたり 864 GB。バッファあり: 3 × 864 GB ≒ 2.5 TB。
  2. チェックポイント間隔: 2 分 = 120 秒。
  3. 必要な帯域幅: 2.5 TB ÷ 120 秒 ≈ 20 Gbps。

参照レシピ

特定のハードウェアとワークロードのストレージ構成の例については、次のレシピをご覧ください。

割り当てと帯域幅の上限

Cloud Storage と Compute Engine の帯域幅は、デフォルトの割り当てによって制限されます。割り当てを超えると、入力リクエストと出力リクエストがスロットリングされることがあります。

Cloud Storage の割り当てと増加をリクエストする方法については、Cloud Storage ドキュメントの割り当てと上限をご覧ください。Hyperdisk と Persistent Disk の Compute Engine の割り当てについては、ディスクの割り当てをご覧ください。

次のステップ