Cloud TPU v5e トレーニング [公開プレビュー]

プレビュー

このプロダクトまたは機能には、サービス固有の利用規約の「一般的なサービス規約」の「pre-GA サービス規約」が適用されます。 一般提供前のプロダクトと機能は「現状のまま」ご利用いただけますが、サポートが制限される場合があります。 詳しくは、リリース ステージの説明をご覧ください。

Cloud TPU v5e は、Google Cloud の最新の生成 AI アクセラレータです。 Pod あたり 256 チップのフットプリントが小さい v5e は、トランスフォーマー、text-to-image、および、畳み込みニューラル ネットワーク（CNN）のトレーニング、微調整、サービス提供に適した最適なプロダクトとなるように最適化されています。

Cloud TPU v5e のコンセプト

Cloud TPU を初めて使用する場合は、TPU ドキュメントのホームをご覧ください。

このセクションでは、v5e に関連するコンセプトを定義します。他の TPU 構成で使用される要素（スライス、ホスト、チップ、TensorCore など）については、TPU システム アーキテクチャ ページをご覧ください。

推論

推論は、トレーニングされたモデルを使用して新しいデータを予測するプロセスです。サービス提供プロセスで使用されます。

サービス提供

サービス提供とは、トレーニング済み機械学習モデルを本番環境にデプロイするプロセスであり、本番環境を使用して予測や決定を行うために使用できます。レイテンシとサービスレベルの可用性はサービス提供において重要です。

スライス

スライスは、高速相互接続相互接続（ICI）によって接続された同じ Pod 内に配置されたチップのコレクションを表します。v5e は 2D スライス形状です。サポートされているスライス形状については、アクセラレータ タイプのセクションの表をご覧ください。

チップ形状とチップ トポロジもスライス形状を指します。

単一ホストとマルチホスト

ホストは、VM を実行する物理コンピュータ（CPU）です。ホストは複数の VM を一度に実行できます。

8 個未満のチップを使用するスライスは、最大で 1 つのホストを使用します。8 チップを超えるスライスは、複数のホストにアクセスし、複数のホストを使用して分散トレーニングを実行できます。スライスとチップの詳細については、TPU システム アーキテクチャのページをご覧ください。

v5e は、マルチホストのトレーニングと単一ホストの推論をサポートしています。このプレビューは、マルチホストの推論またはサービス提供をサポートしていません。1 つのサービス提供ジョブで許可されるチップの最大数は 8 チップです。

TPU VM

基盤となる TPU にアクセスできる Linux を実行している仮想マシン。 v5e TPU の場合、各 TPU VM は、ユーザー指定のアクセラレータ タイプに応じて 1、4、8 チップに直接アクセスできます。TPU VM は、ワーカーとも呼ばれます。

ワーカー

TPU VM をご覧ください。

v5e システム アーキテクチャ

各 v5e チップには、1 つの TensorCore が含まれています。各 TensorCore には、4 つのマトリックス乗算ユニット（MXU）、ベクトル ユニット、スカラー ユニットがあります。次の表に、v5e の主要なチップの仕様と値を示します。Pod の仕様は、チップの仕様の次の表に示しています。

主なチップ仕様	v5e 値
チップあたりのピーク コンピューティング（bf16）	197 TFLOPs
チップあたりのピーク コンピューティング（Int8）	393 TFLOPs
HBM2 の容量と帯域幅	16 GB、819 GBps
インターチップ相互接続 BW	1600 Gbps

主な Pod の仕様	v5e 値
TPU Pod のサイズ	256 チップ
相互接続トポロジ	2D トーラス
Pod あたりのピーク コンピューティング	100 PetaOps（Int8）
Pod あたりの all-reduce 帯域幅	51.2 TB/秒
Pod あたりの二分割帯域幅	1.6 TB/秒
Pod あたりのデータセンターのネットワーク帯域幅	6.4 Tbps

TPU v5e チップ

次の図は、TPU v5 チップを示しています。

アクセラレータのタイプ

Cloud TPU v5e は、トレーニングと推論（サービス提供）を組み合わせたプロダクトです。AcceleratorType フラグは、トレーニング用にプロビジョニングされた TPU 環境と、サービス提供用にプロビジョニングされた環境を区別するために使用されます。

トレーニング ジョブはスループットと可用性に対して最適化されていますが、サービス提供ジョブはレイテンシに対して最適化されています。 したがって、サービス提供用にプロビジョニングされた TPU でのトレーニング ジョブはエクスペリエンスが低下する可能性があります（低可用性）。同様に、トレーニング用にプロビジョニングされた TPU で実行されるサービス提供ジョブは、エクスペリエンスが低下する可能性があります（高レイテンシ）。

AcceleratorType は、ノード内の TensorCore の数を使用してスライスのサイズを記述します。AcceleratorType は、v$VERSION_NUMBER-$CORES_COUNT という形式の文字列です。

アクセラレータ タイプの詳細については、サポートされている TPU 構成をご覧ください。

v5e にはチップごとに 1 つのコアがあります。v5 では、次の 2D スライス形状がサポートされています。

トポロジ	TPU チップ数	ホストの数
1x1	1	1/8
2x2	4	1/2
2x4	8	1
4x4	16	2
4x8	32	4
8x8	64	8
8x16	128	16
16x16	256	32

トポロジごとの VM 数については、VM のタイプをご覧ください。

Cloud TPU v5e のタイプ

トレーニングは最大 256 個のチップに対応しています。

v5e トレーニング ジョブに TPU をプロビジョニングするには、CLI または API TPU 作成リクエストで次のいずれかのアクセラレータ タイプを使用します。

• v5e-1
• v5e-4
• v5e-8
• v5e-16
• v5e-32
• v5e-64
• v5e-128
• v5e-256

推論のプロビジョニングについては、v5e 推論の概要をご覧ください。

TPU VM のタイプ

TPU スライスの各 TPU VM には 1、4、8 個のチップが含まれています。4 チップ以下のスライスには、同じ Non Uniform Memory Access（NUMA）ノードがあります。NUMA ノードの詳細については、システム アーキテクチャ ドキュメントをご覧ください。8 チップ TPU VM の場合、NUMA パーティション内では CPU-TPU 通信の方が効率的です。たとえば、次の図では、CPU0-Chip0 通信は CPU0-Chip4 通信よりも高速です。

始める

容量の確保

Cloud TPU v5e は公開プレビュー版になりました。AI ワークロードに対して Cloud TPU v5e の使用を開始するには、Cloud セールスまでお問い合わせください。

Google Cloud プロジェクトを準備する

1. Google アカウントにログインします。 Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。
2. Google Cloud コンソールで、[プロジェクト セレクタ] ページから Cloud プロジェクトを選択するか作成します。

3. Google Cloud の使用には請求の設定が必要なため、プロジェクトで課金が有効になっていることを確認します。

   公開プレビューの使用に対する課金は、Cloud TPU 料金ページに記載されている標準のリージョン料金に従います。

4. gcloud alpha components をインストールします。

5. TPU ユーザーが既存の gcloud alpha コンポーネントを再利用する場合は、関連するコマンドとフラグがサポートされるように更新してください。

   gcloud components update
   

6. Cloud Shell で、次の gcloud コマンドを使用して TPU API を有効にします。（Google Cloud コンソールから有効にすることもできます）。

   gcloud services enable tpu.googleapis.com
   

7. TPU サービス アカウントを有効にします。

   サービス アカウントにより、Cloud TPU サービスが他の Google Cloud サービスにアクセスできるようになります。ユーザー管理のサービス アカウントは、Google Cloud で推奨される方法です。ガイドに沿って作成し、ロールを付与します。次のロールが必要となります。

   • TPU 管理者
   • ストレージ管理者: Cloud Storage へのアクセスに必要
   • ログ書き込み: Logging API でログを書き込むために必要
   • モニタリング指標の書き込み: Cloud Monitoring に指標を書き込むために必要

8. プロジェクトとゾーンを構成します。

   プロジェクト ID は、Cloud コンソールに表示されるプロジェクトの名前です。

   export PROJECT_ID=your_project_ID
   export ZONE=us-west4-a
   
   gcloud alpha compute tpus tpu-vm service-identity create --zone=${ZONE}
   
   gcloud auth login
   gcloud config set project ${PROJECT_ID}
   gcloud config set compute/zone ${ZONE}
   

Cloud TPU 環境をプロビジョニングする

Cloud TPU v5 は、キューに格納されたリソースとして管理されます。queued-resource create コマンドを使用すると、容量をプロビジョニングできます。

環境変数を作成する

TPU の作成に必要な環境変数を設定します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

変数	説明
PROJECT_ID	Google Cloud プロジェクト名
ACCELERATOR_TYPE	サポートされているスライス形状については、アクセラレータ タイプのセクションをご覧ください。
ZONE	すべての公開プレビュー容量は us-west4-a にあります。
RUNTIME_VERSION	公開プレビューには v2-alpha-tpuv5-lite を使用します。
SERVICE_ACCOUNT	これは、Google Cloud コンソール -> IAM] -> サービス アカウント で確認できるサービス アカウントのアドレスです。 例: [e-mail]
TPU_NAME	キューに格納されたリソース リクエストの割り当て時に作成される TPU のユーザー割り当て ID。
QUEUED_RESOURCE_ID	キューに格納されたリソース リクエストのユーザー割り当て ID。キューに格納されたリソースについては、キューに格納されたリソースのドキュメントをご覧ください。

TPU リソースの作成

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID}
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved

注: --reserved フラグは、オンデマンド リソースではなく、予約済みリソースに対してのみ必要です。

キューに格納されたリソースが正常に作成されると、response フィールド内の状態は ACCEPTED または、FAILEDのいずれかになります。 キューに格納されたリソースが ACCEPTED 状態の場合は、キューに格納されたリソースが予備検証に合格し、容量を待機していることを意味します。容量が利用可能になると、リクエストは PROVISIONING に移行します。ACCEPTED の状態であることは、割り当てが保証されることを意味するわけではなく、ACCEPTED ステータスから ACTIVE ステータスに変更されるまでには時間がかかることがあります。 キューに格納されたリソースが FAILED 状態の場合、失敗の理由が出力に表示されます。

QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

list コマンドまたは describe コマンドを使用して、キューに格納されたリソースのステータスをクエリします。

    gcloud alpha compute tpus queued-resources describe {QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE}
   

状態は、キューに格納されたリソースのステータスを表します。状態は次のように定義されます。

都道府県	Description
ACCEPTED	キューに格納されたリソース create コマンドを受信しました。容量が使用可能になり次第、プロビジョニングを開始します。
PROVISIONING	TPU スライスがプロビジョニング中です。
ACTIVE	すべての TPU がプロビジョニングされ、使用できる状態になりました。起動スクリプトを指定した場合、キューに格納されたリソース状態が ACTIVE に遷移すると、すべての TPU でスクリプトの実行が開始されます。
FAILED	スライスをプロビジョニングできませんでした。
SUSPENDING	1 つ以上のスライスを削除中です。
SUSPENDED	基盤となるスライスはすべて削除されますが、キューに格納されたリソースは明示的に削除されるまで残ります。現時点では、一時停止されたキューに登録されたリソースを再開することはできません。削除する必要があります。
DELETING	キューに格納されたリソースを削除中です。

重要: ACCEPTED の状態は非推奨になりました。WAITING_FOR_RESOURCES に置き換えられます。

TPU VM に SSH 接続する

次のセクションでは、TPU スライスの各 TPU VM にバイナリをインストールしてコードを実行する方法について説明します。このコンテキストで、TPU VM はワーカーとも呼ばれます。

スライスに含まれる VM の数を計算するには、VM のタイプのセクションをご覧ください。

バイナリをインストールするかコードを実行するには、tpu-vm ssh command を使用して TPU VM に SSH 接続します。

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}

特定の TPU VM に SSH 接続するには、0 ベースのインデックスに続けて --worker フラグを使用します。

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

スライス形状が 8 チップを超える場合、1 つのスライスに複数の VM があります。この場合、--worker=all フラグを使用すると、1 つのステップですべての TPU VM をインストールできます。VM ごとに個別に SSH 接続する必要はありません。次に例を示します。

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} /
--zone  ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

管理

TPU VM の管理に使用できるすべてのコマンドについては、TPU の管理をご覧ください。

フレームワークの設定

このセクションでは、TPU v5e で JAX または PyTorch を使用したカスタムモデル トレーニングの一般的な設定プロセスについて説明します。TensorFlow のサポートはまもなく有効になります。

推論の設定手順については、v5e 推論の概要をご覧ください。

JAX の設定

スライス形状が 8 チップを超える場合、1 つのスライスに複数の VM があります。この場合、SSH を使用して個別にログインすることなく、--worker=all フラグを使用してすべての TPU VM に 1 つのステップでインストールを実行する必要があります。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone  ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

次のコマンドを実行して、デバイスの数を確認できます（ここに示した出力は、v5litepod-16 スライスで生成されました）。このコードは、JAX から Cloud TPU TensorCore が検出され、基本的なオペレーションを実行できることを確認することによって、すべてが正しくインストールされていることをテストします。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

出力は次のようになります。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() は、指定されたスライスのチップの合計数を示します。jax.local_device_count() は、このスライス内の 1 つの VM によってアクセス可能なチップの数を示します。

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

出力は次のようになります。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

JAX を使用して v5e トレーニングを開始するには、このドキュメントの JAX チュートリアルをお試しください。

PyTorch の設定

v5e は PJRT ランタイムのみをサポートしていることに注意してください。PyTorch 2.1 以降では、すべての TPU バージョンのデフォルト ランタイムとして PJRT が使用されます。

このセクションでは、すべてのワーカー用のコマンドで PyTorch/XLA を使用して v5e 上で PJRT を使用する方法について説明します。

依存関係をインストールする

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='
   sudo apt-get update -y
   sudo apt-get install libomp5 -y
   pip3 install mkl mkl-include
   pip3 install tf-nightly tb-nightly tbp-nightly
   pip3 install numpy
   sudo apt-get install numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
'

torch/torch_xla/torchvision のホイールをインストールできず、pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222 のようなエラーが表示された場合は、次のコマンドを使用してバージョンをダウングレードしてください。

pip3 install setuptools==62.1.0

PJRT を使用して単純なスクリプトを実行します。

大きい割り当てが頻繁に発生するモデルの場合、tcmalloc を使用すると、デフォルトの malloc 実装よりもトレーニング時間が大幅に短縮されるため、TPU VM でのデフォルトの malloc は tcmalloc です。ただし、ワークロードによっては（たとえば、埋め込みテーブルへの割り当てが非常に大きい DLRM の場合）、tcmalloc では速度が低下する可能性があります。その場合は、代わりにデフォルトの malloc を使用して次の変数の設定を解除してみてください。

unset LD_PRELOAD

Python スクリプトを使用して v5e VM で簡単な計算を行う例を次に示します。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker all --command='
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/$USER/.local/lib/
export PJRT_DEVICE=TPU_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
unset LD_PRELOAD
export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so

python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"
'

これにより、次のような出力が生成されます。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

PyTorch を使用して v5e トレーニングを開始するには、このドキュメントの PyTorch チュートリアルをお試しください。

セッションの終了時に TPU とキューに格納されたリソースを削除します。 キューに格納されたリソースを削除するには、次の 2 つの手順でスライスを削除してから、キュー内のリソースを削除します。

gcloud alpha compute tpus tpu-vm delete $TPU_NAME \
--project=$PROJECT_ID \
--zone=$ZONE \
--quiet

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--quiet

これら 2 つのステップは、FAILED 状態にある、キューに格納されたリソース リクエストを削除するためにも使用できます。

モニタリングとプロファイル

Cloud TPU v5e は、前世代の Cloud TPU と同じ方法でモニタリングとプロファイリングをサポートします。プロファイリングの詳細については、Cloud TPU ツールでモデルをプロファイリングするをご覧ください。モニタリングの詳細については、Cloud TPU VM をモニタリングするをご覧ください。

JAX/FLAX の例

v5e で ImageNet をトレーニングする

このチュートリアルでは、架空の入力データを使用して ImageNet を v5e でトレーニングする方法について説明します。実際のデータを使用する場合は、GitHub の README ファイルをご覧ください。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved
   

   注: --reserved フラグは、オンデマンド リソースではなく、予約済みリソースに対してのみ必要です。
   注: キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
   

   When the QueuedResource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. 最新バージョンの JAX と jaxlib をインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   注: jax>=0.4.8 の使用を提案するエラーが表示された場合は、メッセージを無視してかまいません。

4. ImageNet モデルのクローンを作成し、対応する要件をインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='git clone -b throughput_testing https://github.com/gkroiz/flax.git && cd flax && pip install -e . && cd examples/imagenet && pip install -r requirements.txt && pip install --upgrade git+https://github.com/google/flax.git'
   

5. 架空のデータを生成するには、データセットのサイズに関する情報が必要です。これは、ImageNet データセットのメタデータから収集できます。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
   

モデルのトレーニング

前のすべての手順が完了したら、モデルをトレーニングできます。GitHub リポジトリには、使用可能な構成ファイルがいくつかあります。別の方法として、コマンドラインでオーバーライドして、構成ファイルをカスタマイズすることもできます。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='cd flax/examples/imagenet && JAX_USE_PJRT_C_API_ON_TPU=1 JAX_PLATFORMS=tpu TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0 TPU_MIN_LOG_LEVEL=0 python3 imagenet_fake_data_benchmark.py --config=configs/fake_data_16_chips.py'
    

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

Hugging Face FLAX モデル

FLAX で実装された Hugging Face モデルを Cloud TPU v5e ですぐに使用できます。このセクションでは、人気のあるモデルを実行する手順について説明します。

Imagenette で ViT をトレーニングする

このチュートリアルでは、Cloud TPU v5e で Fast AI imagenette データセットを使用して、HugggFace から Vision Transformer（ViT）モデルをトレーニングする方法について説明します。

ViT モデルは、畳み込みネットワークと比較して優れた結果を持ち、ImageNet で Transformer エンコーダを正常にトレーニングした最初のモデルでした。詳しくは、次のリソースをご覧ください。

• ViT の概要
• ViT コードリソース

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT}
   

   キューに登録されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project ${PROJECT_ID} --zone ${ZONE}
   

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. JAX とそのライブラリをインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Hugging Face リポジトリをダウンロードし、要件をインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/vision/requirements.txt'
   

5. Imagenette データセットをダウンロードします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=${ZONE} --worker=all --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
   

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd transformers && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'
  

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --quiet
   

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

ViT ベンチマークの結果

トレーニング スクリプトは v5e-4、v5e-16、v5e-64 で実行されました。次の表に、異なるアクセラレータ タイプでのスループットを示します。

アクセラレータ タイプ	v5e-4	v5e-16	v5e-64
Epoch	3	3	3
グローバル バッチサイズ	32	128	512
スループット（例/秒）	263.40	429.34	470.71

Pokémon 上で Diffusion をトレーニングする

このチュートリアルでは、Cloud TPU v5e で Pokémon データセットを使用して、HuggingFace から Stable Diffusion モデルをトレーニングする方法について説明します。

Stable Diffusion モデルは、テキスト入力からフォトリアリスティックな画像を生成する、潜在的 text-to-image モデルです。詳しくは、次のリソースをご覧ください。

• Stable Diffusion の概要
• Stable Dufffusion のコードソース

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved
   

   注: --reserved フラグは、オンデマンド リソースではなく、予約済みリソースに対してのみ必要です。
   注: キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
   

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. JAX とそのライブラリをインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. HuggingFace リポジトリをダウンロードし、要件をインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
   

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --=$PROJECT=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd diffusers/examples/text_to_image && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --fake_data=1 --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'
   

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --quiet
   

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

分散のベンチマーク結果

トレーニング スクリプトは、v5e-4、v5e-16、v5e-64 で実行されました。次の表は、スループットを示しています。

アクセラレータ タイプ	v5e-4	v5e-16	v5e-64
トレーニング ステップ	1500	1500	1500
グローバル バッチサイズ	32	64	128
スループット（例/秒）	36.53	43.71	49.36

OSCAR データセットで GPT2 をトレーニングする

このチュートリアルでは、Cloud TPU v5e で OSCAR データセットを使用して HuggingFace から GPT2 モデルをトレーニングする方法を説明します。

GPT2 は、ヒューマン ラベリングなしで、未加工のテキストで事前にトレーニングされた変形モデルです。センテンス内の次の単語を予測するようにトレーニングされています。 詳しくは、次のリソースをご覧ください。

• GPT2 の概要
• GPT2 のコードソース

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved
   

   注: --reserved フラグは、オンデマンド リソースではなく、予約済みリソースに対してのみ必要です。
   注: キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
   

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. JAX とそのライブラリをインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. HuggingFace のリポジトリをダウンロードし、要件をインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/language-modeling/requirements.txt'
   

5. 構成ファイルをダウンロードしてモデルをトレーニングします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-v2-images-artifacts/config/v5litepod-preview/jax/gpt .'
   

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'
   

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

GPT2 のベンチマーク結果

トレーニング スクリプトは、v5e-4、v5e-16、v5e-64 で実行されました。次の表は、スループットを示しています。

	v5e-4	v5e-16	v5e-64
Epoch	3	3	3
グローバル バッチサイズ	64	64	64
スループット（例/秒）	74.60	72.97	72.62

PyTorch/XLA

PJRT ランタイムを使用して ResNet をトレーニングする

PyTorch/XLA は、PyTorch 2.0 以降、XRT から PjRt に移行しています。PyTorch/XLA トレーニング ワークロード用に v5e を設定する手順が更新されています。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved
   

   注: --reserved フラグは、オンデマンド リソースではなく、予約済みリソースに対してのみ必要です。
   注: QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
   

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. PyTorch/XLA のクローンを作成する

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command='cd /usr/share/
   sudo git clone --recursive https://github.com/pytorch/pytorch
   cd pytorch/
   sudo git clone --recursive https://github.com/pytorch/xla.git
   cd xla/
   sudo git checkout 3b10663
   sudo pip3 install numpy'
   

4. Torch、Torch/xla、Torchvision をインストールする

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} --project ${PROJECT_ID} --worker=all --command='
   pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torchvision-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
   sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
   sudo pip3 install torch_xla[tpuvm]
   '
   

5. Torch/XLA 固有の依存関係をインストールする

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} --zone  ${ZONE} --worker all --command='
   sudo apt-get update -y
   sudo apt-get install libomp5 -y
   pip3 install mkl mkl-include
   pip3 install tf-nightly tb-nightly tbp-nightly
   pip3 install numpy
   sudo apt-get install numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
   '
   

6. libtpu をダウンロードする

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
   sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/libtpu.so /home/$USER/.local/lib/python3.8/site-packages/libtpu/libtpu.so
   "
   

7. v5e 構成で tpu.py を更新する

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
   sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/tpu.py /home/$USER/.local/lib/python3.8/site-packages/torch_xla/experimental/tpu.py
   "
   

ResNet モデルをトレーニングする

      gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet

ベンチマークの結果

次の表に、ベンチマークのスループットを示します。

アクセラレータ タイプ	スループット（例/秒）
v5e-4	4240 ex/s
v5e-16	10,810 ex/s
v5e-64	46,154 ex/s

v5e で GPT2 をトレーニングする

このチュートリアルでは、wikitext データセットを使用して PyTorch/XLA で HuggingFace のリポジトリを使用して、v5e で GPT2 を実行する方法を説明します。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT}
   

   注: QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
   

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. torch/xla 依存関係をインストールします。

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} --zone  ${ZONE} --worker all --command='
   sudo apt-get update -y
   sudo apt-get install libomp5 -y
   pip3 install mkl mkl-include
   pip3 install tf-nightly tb-nightly tbp-nightly
   pip3 install numpy
   sudo apt-get install numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+fc838ad-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
   '
   

4. HuggingFace のリポジトリをダウンロードし、要件をインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command="
   rm -rf transformers/
   git clone https://github.com/huggingface/transformers.git; \
   cd transformers; \
   sudo pip3 install -e .; \
   git checkout -q ebdb185befaa821304d461ed6aa20a17e4dc3aa2; \
   pip3 install datasets; \
   pip3 install evaluate; \
   pip3 install scikit-learn;"
   

5. 事前トレーニング済みモデルの構成ファイルをダウンロードします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command="
   cd transformers &&
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/my_config_2.json examples/pytorch/language-modeling/
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/fsdp_config.json examples/pytorch/language-modeling/
   "
   

モデルのトレーニング

バッチサイズ 16 を使用して 2B モデルをトレーニングします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='
   export PJRT_DEVICE=TPU_C_API
   export PT_XLA_DEBUG=0
   export USE_TORCH=ON
   export XLA_USE_BF16=1
   export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
   export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so
   cd transformers
   numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py --num_cores 4
    examples/pytorch/language-modeling/run_clm.py --num_train_epochs 3
    --dataset_name wikitext
    --dataset_config_name wikitext-2-raw-v1
    --per_device_train_batch_size 16
    --per_device_eval_batch_size 16
    --do_train
    --do_eval
    --output_dir /tmp/test-clm
    --overwrite_output_dir
    --config_name examples/pytorch/language-modeling/my_config_2.json --cache_dir /tmp
    --tokenizer_name gpt2
    --block_size 1024
    --optim adafactor
    --adafactor true
    --save_strategy no
    --logging_strategy no
    --fsdp "full_shard"
    --fsdp_config examples/pytorch/language-modeling/fsdp_config.json'
   

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

ベンチマークの結果

トレーニング スクリプトは、v5e-4、v5e-16、v5e-64 で実行されました。次の表は、さまざまなアクセラレータ タイプのベンチマークのスループットを示しています。

	v5e-4	v5e-16	v5e-64
Epoch	3	3	3
構成	600M	20 億の	16B
グローバル バッチサイズ	64	128	256
スループット（例/秒）	66	77	31

v5e で ViT をトレーニングする

このチュートリアルでは、cifar10 データセット上の PyTorch/XLA で HuggingFace のリポジトリを使用して、v5e で VIT を実行する方法について説明します。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued-resource-id
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved
   

   注: --reserved フラグは、オンデマンド リソースではなく、予約済みリソースに対してのみ必要です。
   注: QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
   

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. torch/xla 依存関係をインストールする

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} --zone  ${ZONE} --worker=all --command='
   sudo apt-get update -y
   sudo apt-get install libomp5 -y
   pip3 install mkl mkl-include
   pip3 install tf-nightly tb-nightly tbp-nightly
   pip3 install numpy
   sudo apt-get install numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly+20230726-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly+20230726-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
   '
   

4. HuggingFace のリポジトリをダウンロードし、要件をインストールします。

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="
     git clone https://github.com/suexu1025/transformers.git vittransformers; \
     cd vittransformers; \
     pip3 install .; \
     pip3 install datasets; \
     wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
   

モデルのトレーニング

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME 

   --project=$PROJECT_ID 

   --zone=$ZONE 

   --worker=all 

   --command="
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers;
      numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10
      --remove_unused_columns=False
      --label_names=pixel_values
      --mask_ratio=0.75
      --norm_pix_loss=True
      --do_train=true
      --do_eval=true
      --base_learning_rate=1.5e-4
      --lr_scheduler_type=cosine
      --weight_decay=0.05
      --num_train_epochs=3
      --warmup_ratio=0.05
      --per_device_train_batch_size=8
      --per_device_eval_batch_size=8
      --logging_strategy=steps
      --logging_steps=30
      --evaluation_strategy=epoch
      --save_strategy=epoch
      --load_best_model_at_end=True
      --save_total_limit=3
      --seed=1337
      --output_dir=MAE
      --overwrite_output_dir=true
      --logging_dir=./tensorboard-metrics
      --tpu_metrics_debug=true"
   

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME 

   --project=${PROJECT_ID} 

   --zone=$ZONE 

   --quiet
   

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} 

   --project ${PROJECT_ID} 

   --zone ${ZONE} 

   --quiet
   

ベンチマークの結果

次の表は、さまざまなアクセラレータ タイプのベンチマークのスループットを示しています。

	v5e-4	v5e-16	v5e-64
Epoch	3	3	3
グローバル バッチサイズ	32	128	512
スループット（例/秒）	201	657	2,844

サポートとフィードバック

フィードバックをぜひお寄せください。フィードバックをお寄せいただける、またはサポートを必要とされる場合は、こちらにお問い合わせいただくか、[e-mail] までメールをお送りください。

利用規約

この公開プレビューに関して、Google から提供されるすべての情報は Google の機密情報であり、Google Cloud Platform 利用規約（または Google Cloud Platform の使用に適用されるその他の契約）の機密保持条項の対象となります。