TPU VM での AI/ML ワークロードのストレージに関するベスト プラクティス
TPU VM で AI/ML ワークロードのパフォーマンスと費用対効果を最大化するには、ワークロードに適したストレージ ソリューションを選択して構成します。I/O ボトルネックを解消することで、TPU アクセラレータのアイドル時間を短縮し、トレーニング時間とコストを削減できます。
このドキュメントでは、TPU VM でのトレーニング、チェックポイント設定、サービング、キャッシュ保存に関するワークロード固有のストレージの推奨事項と最適化のベスト プラクティスについて説明します。これらのプラクティスを適用する前に、利用可能な TPU データのストレージ オプションを確認してください。このドキュメントは、TPU VM に精通しており、Cloud Storage リソースのプロビジョニングの基本的な経験があることを前提としています。
ワークロード固有のガイダンス
次の表に、ワークロードごとのストレージの推奨事項を優先順位の高い順に示します。
| ワークロード | 推奨事項 | 最適化とツール(該当する場合) |
|---|---|---|
| トレーニング データセット(データ準備を含む) |
|
|
| チェックポイントと強化学習の重み |
|
Rapid Bucket の場合は、gcsfusecsi-checkpointing プロファイルを使用します。 |
| モデルの保存とダウンロード |
|
モデルをダウンロードするには、GKE Run:ai Model Streamer または Cloud Storage FUSE を使用し、gcsfusecsi-serving プロファイルを使用して別のマウント ポイントを使用します。 |
| Key-Value(KV)キャッシュのオフロード |
|
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 |
チェックポイントの帯域幅を推定する
次の式を使用して、トレーニング ワークロードに必要な最小チェックポイント帯域幅を見積もることができます。データ読み取り、複数のトレーニング実行、トレーニングと推論の場合は、推定帯域幅要件を比例して増やします。
- チェックポイント サイズ: パラメータ数 × パラメータあたりのバイト数(FP16 + オプティマイザーの状態の場合、パラメータあたり約 12 ~ 16 バイト)。オプティマイザーの状態とさまざまな精度に対してバッファ(約 3 倍)を追加します。
- チェックポイント間隔: チェックポイントを保存する頻度(15 分ごとなど)。
- 必要な帯域幅: チェックポイント サイズ ÷ チェックポイント間隔。
次の例は、Qwen3-72B の最小チェックポイント帯域幅を推定する方法を示しています。
- チェックポイント サイズ: 72B パラメータ × 12 バイト ≈ チェックポイントあたり 864 GB。バッファあり: 3 × 864 GB ≒ 2.5 TB。
- チェックポイント間隔: 2 分 = 120 秒。
- 必要な帯域幅: 2.5 TB ÷ 120 秒 ≈ 20 Gbps。
参照レシピ
特定のハードウェアとワークロードのストレージ構成の例については、次のレシピをご覧ください。
- TPU7x 推論:
- TPU7x トレーニング:
割り当てと帯域幅の上限
Cloud Storage と Compute Engine の帯域幅は、デフォルトの割り当てによって制限されます。割り当てを超えると、入力リクエストと出力リクエストがスロットリングされることがあります。
Cloud Storage の割り当てと増加をリクエストする方法については、Cloud Storage ドキュメントの割り当てと上限をご覧ください。Hyperdisk と Persistent Disk の Compute Engine の割り当てについては、ディスクの割り当てをご覧ください。
次のステップ
- TPU データのストレージ オプション
- TPU VM を Cloud Storage バケットに接続する
- 耐久性のあるブロック ストレージを TPU VM にアタッチする
- Cloud Storage FUSE のパフォーマンス調整のベスト プラクティス