Ray 2.55からGoogle TPUが正式サポート！GKEで分散学習を始める完全ガイド

なぜRay + TPUなのか？

GPUでRayを使っているなら、同じコードをGoogle TPUで実行できるようになりました。既存のRayライブラリ(Data, Train, Serve)をそのまま使いながら、TPUスライスをあたかも1つのGPUノードのように扱えます。

Ray 2.55からTPUは実験段階ではなくなり、公式プレビルドイメージとリリースパイプラインに含まれています。コミュニティサポートに依存する必要はありません。この記事では、Ray on TPUの基盤となる GKE連携 と スライス配置 の概念を中心に解説します。

TL;DR: TPUスライスを理解すれば、Rayの既存開発プロセスを維持したままTPUハードウェアを活用できます。

TPUスライスとは？

TPUチップは固定グループで接続されており、これを スライス と呼びます。複数のホストマシンのチップがICI(Inter-Chip Interconnect)という高速リンクで接続されています。

GPU環境に例えると、1つのスライスは NVLinkで接続されたマルチGPUボックス と同様です。ボックス外に出ると高速接続が切れるため、分散学習時のall-reduceなどの集合通信が完了せず、学習が停止する可能性があります。

TPUで重要なのは トポロジー です。スライスの形状を示し、例えば16チップスライスは 4x4 と表記します。RayとGKEがこのトポロジー情報を使用して、スライス全体をアトミックに予約します。

GKEでRay Operatorを設定する

RayをTPUで実行するには、GKEクラスタにRay Operatorアドオンを有効化する必要があります。AutopilotまたはStandardモードでコマンドが少し異なります。

Autopilot (完全マネージドノード)

gcloud container clusters create-auto CLUSTER \
--enable-ray-operator --location=LOCATION

Standard (ノードプールを直接管理)

gcloud container clusters create CLUSTER \
--addons=RayOperator --location=LOCATION &&
gcloud container node-pools create v6e-16-slice \
--cluster=CLUSTER \
--location=LOCATION \
--machine-type=ct6e-standard-4t \
--tpu-topology=4x4 \
--num-nodes=4

このコマンドはKubeRayとRay TPUウェブフックをインストールします。ウェブフックは各TPUホストに ray.io/tpu-slice-name などのラベルを付与し、Rayが同じスライスに属するマシンを識別できるようにします。

RayClusterマニフェストでTPUをリクエストする

RayClusterのworkerGroupSpecで nodeSelector を使用してTPU世代とトポロジーを指定し、リソースとしてチップ数をリクエストします。

# RayCluster workerGroupSpec内
nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice # TPU世代
  cloud.google.com/gke-tpu-topology: "4x4" # スライス形状
# ... google.com/tpuリソース制限でチップ数をリクエスト
numOfHosts: 4 # マルチホスト: スライスを構成するホストVM数

これでGKEがスライスをプロビジョニングし、ウェブフックがラベルを付け、Rayがラベルを読んでジョブを配置します。コードはそのままPythonで記述します。

Ray Coreのスライス配置グループ

Ray Coreは分散コンピューティングの基本レイヤーで、タスクとアクターをスケジューリングします。TPUサポートは ray.util.tpu APIに含まれており、主要関数は slice_placement_group() です。

この関数はスライス全体をアトミックに予約します。つまり、すべてのホストが成功するか、何も実行されないかのどちらかです。

from ray.util.tpu import slice_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

# v6e 4x4スライスを1つ予約 (4ホストにまたがる16チップ)
spg = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(spg.placement_group.ready(), timeout=600)

@ray.remote(resources={"TPU": 4})
def worker(rank, world):
    # 各ホストで実行される処理
    pass

tasks = [
    worker.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=spg.placement_group)
    ).remote(rank=i, world=spg.num_hosts)
    for i in range(spg.num_hosts)
]

実務ではRay AIライブラリ(Data, Train, Serve)が内部的にこの関数を呼び出すため、直接使う機会は多くありません。カスタム分散ワークロードを書く場合のみ直接使用します。

⚠️ このAPIは現在 alpha 安定性段階です。公開APIですが、リリース間で変更される可能性があります。

注意点と制限

1. スライス単位の配置: TPUスライスは1つの単位として予約されるため、スライスより小さい単位で実行できません。GPUのように柔軟なスケーリングは困難です。
2. アルファAPI: slice_placement_group() はまだアルファ段階なので、本番適用時はAPI変更の可能性を考慮する必要があります。
3. コスト: TPUはGPUより時間単価が低い場合がありますが、スライス全体を予約する必要があるため、アイドル時間が発生する可能性があります。

日本市場での適用コンテキスト

日本でもGPU調達難からTPU導入を検討する企業が増えています。特に大規模AIサービスを運営するスタートアップや大企業で、RayとGKEを併用するケースが増えると予想されます。ただし、国内クラウド環境でのTPU割り当て枠を確認することが優先です。

次の学習ステップ

この記事では基盤を扱いました。次にRay Data、Ray Train、Ray ServeをTPUで使用する方法を学んでみてください。特にvLLMを使用したLLMサービングとJaxTrainerによる学習が実務で役立ちます。

関連記事

• Azure Databricks、3年ROI 331%を獲得 — Forrester TEI分析まとめ
• JavaScript ShadowRealm API、ついにコード分離の時代が来る

まとめ

Ray on TPUは実験段階を超え、実務で活用できるレベルになりました。GKEがプロビジョニングとラベル付けを担当し、Ray Coreがスライス単位で予約するため、開発者はPythonコードに集中できます。

TPUスライスの概念と制約を理解すれば、既存のRayコードをほぼ修正せずにTPUへ移行できます。GPUで既にRayを使用しているなら、TPUへの拡張は思ったより簡単です。

詳細は根拠資料を参照してください。Happy building!