TorchTPUとは何か — Google TPUでPyTorchを「ネイティブ」に走らせるハイブリッド Eager/Compiled アーキテクチャ

はじめに：TPUでPyTorchが動く時代

GoogleのTPU（Tensor Processing Unit）は、行列積に特化した演算ユニット「MXU」をシストリックアレイ状に並べたアクセラレータで、大規模言語モデルの訓練と推論において高い電力効率・コスト効率を発揮してきた。ただしソフトウェア面では長らくTensorFlowとJAXのエコシステムに属しており、PyTorchユーザーがTPUを使うには、torch_xlaというLazy Tensorベースのバックエンドを経由するか、JAX向けにコードを書き換えるという回り道が必要だった。

2026年、GoogleとMetaのPyTorchチームが共同で開発した「TorchTPU」が公開され、この前提が変わった。TorchTPUはPyTorchのディスパッチャに新しいC++バックエンドとして登録される仕組みで、TPUがGPUやCPUと並ぶPyTorchのネイティブデバイスになる。実務上の意味は明確で、既存コードのmodel.to("cuda")をmodel.to("tpu")に、分散のバックエンド指定をncclからtpu_distに変える程度で、モデルがそのまま動作する。

graph TD
    A["PyTorch<br/>Aten オペレータ / Dispatcher"] --> B["CUDA バックエンド<br/>NVIDIA GPU"]
    A --> C["CPU バックエンド<br/>x86 / ARM"]
    A --> D["TPU バックエンド（TorchTPU）<br/>PrivateUse1 として登録・autoload"]
    D --> E["eager 実行<br/>Strict / Debug / DeferAndFuse"]
    D --> F["compiled 実行<br/>torch.compile(backend='tpu') → StableHLO → XLA"]

設計の要点は二段構えにある。まず誰でも動かせることを優先し、通常のPyTorchコードをその場で実行するeagerモードを一級市民として扱う。そのうえで本番級の性能を狙う段階では、torch.compile(backend="tpu")によるグラフ全体のコンパイルや、Pallasなどで書くカスタムカーネルを組み合わせて詰めていく。Hugging Faceの人気モデル約5,000件を対象にした検証では、Tier-1の4,917モデル中98%（4,821件）がクラッシュなく動作したと報告されている。TPUが「JAX専用のハードウェア」ではなく、PyTorchのデバイス選択肢のひとつになった点が、この変化の本質である。

TorchTPUとは：一言でいうと「PyTorchの第一級TPUバックエンド」

TorchTPUを一言で表すなら、TPUをGPU・CPUと並ぶPyTorchの「第一級（first-class）デバイス」にするバックエンドである。従来のPyTorch/XLAやTorchAXが「PyTorchをJAXの作法に押し込む」アプローチだったのに対し、TorchTPUは逆方向、つまり「PyTorchの流儀のままTPUを動かす」ことを設計の出発点にしている。

何を解決するのか

第一の目的は、既存のPyTorchワークロードを最小限の変更でTPUに載せることだ。実務上の変更点はほぼ2箇所に集約される。

# 変更前（GPU）
model.to("cuda")
dist.init_process_group(backend="nccl")

# 変更後（TPU）
model.to("tpu")
dist.init_process_group(backend="tpu_dist")

GPUとTPUの両対応コードを書きたい場合は、torch.accelerator.current_accelerator()と引数なしのdist.init_process_group()を使う選択肢もある。いずれにせよ、モデル定義や学習ループの本体を書き換える必要はない。

技術的には、PyTorchのディスパッチャに新しいC++バックエンドをPrivateUse1として登録する「out-of-tree backend」方式を採っている。TPUマシンにTorchTPUのwheelが入っていれば、PyTorchをimportするだけでバックエンドが自動設定される（autoload）。PrivateUse1名はtpuに置き換えられ、サブクラスやラッパーを挟まずに、普通のPyTorchテンソルがそのままTPU上で扱える。テンソルの実体は、PyTorch標準のat::Tensor／c10::StorageImplの仕組み（c10::UniqueVoidPtr）を使ってTPU専用ストレージとして実装されている。

3つの設計原則

このバックエンドは、次の3本柱に沿って設計されている。

「PyTorch Citizenship」が示すように、TorchTPUはJAXに依存しない。これは単なる実装上の都合ではなく、PyTorchの命令型（eager）実行という思想そのものを守るための選択である。この原則が、後述するeager実行の3モードや、torch.compileとの二段構えの設計へとつながっていく。

なぜ重要か：PyTorch/XLAが抱えていた課題

TPUでPyTorchを動かす試みは2020年のPyTorch/XLAに始まる。これはLazy Tensor方式、つまり演算を呼ばれた瞬間には実行せず、グラフを組み立ててからまとめてコンパイルする方式を採用していた。CPUやGPUのPyTorchが1演算ずつその場で実行されるのに対し、TPU側では計算の実体が遅れて現れるため、次のような問題が積み重なっていた。

従来のPyTorch on TPU（PyTorch/XLA、TorchAX）の問題

• 遅延実行によるデバッグの困難さ：コードを書いた時点と実際にTPU上で計算が走る時点がずれるため、途中のテンソルを覗いたり、失敗した演算の位置を特定したりすることが難しい。
• 侵襲的なコード変更：どの範囲をコンパイルするかを明示するために、学習ループなどへ専用の呼び出しを挿入する必要があり、既存のPyTorchコードをそのまま持ち込めなかった。
• コンパイル遅延とeagerモードの欠如：グラフ全体のコンパイルに時間がかかり、対話的な試行錯誤や短いスクリプトの実行に向かない。PyTorch本来の命令型（eager）実行が存在しなかった。
• 分散プリミティブへの追随遅れ：DTensorやFSDPv2、SimpleFSDPといった新しい分散APIがPyTorch側で次々に登場する一方、バックエンドの対応が追いつかず、エコシステムの進化から離れてしまう。
• TorchAXの構造的な難しさ：PyTorchをJAXの作法に合わせて動かすアプローチであり、JAX側の仕様変更とPyTorch側の進化の両方に追随する必要があるため、整合を取り続けることが難しかった。

TorchTPUが変えたこと

TorchTPUは設計の出発点を反転させ、PyTorchの意味論を守ることを最優先に据えた。eager実行を一級市民として扱い、torch.compile(backend="tpu")と組み合わせることで本番級の性能に到達できるようにしている。さらに分散処理をPyTorchのProcess Group APIレベルで実装したため、DTensor・FSDPv2・SimpleFSDP・FairScaleといった既存の仕組みがそのまま動作する。

graph LR
    subgraph OLD["従来のPyTorch on TPU"]
        A1["PyTorch / Aten"] --> A2["Lazy Tensor<br/>遅延グラフ蓄積"] --> A3["XLA コンパイル"]
    end
    subgraph NEW["TorchTPU"]
        B1["PyTorch / Aten<br/>Dispatcher"] --> B2["eager実行<br/>Strict / Debug / DeferAndFuse"]
        B1 --> B3["torch.compile<br/>Torch Dynamo"]
        B2 --> B4["StableHLO"]
        B3 --> B4
        B4 --> B5["XLA コンパイル"]
    end

こうした設計変更により、TPUは「JAXやTensorFlowのためのハードウェア」から、PyTorchのデバイス選択肢のひとつへと位置づけを変えた。

TPUというハードウェアの基礎知識

TPU（Tensor Processing Unit）は、深層学習で支配的な巨大テンソル演算のために設計されたASIC（特定用途向け集積回路）である。GPUが汎用の並列プロセッサであるのに対し、TPUは行列積とその周辺のデータ移動に機能を絞り込んでいる。この違いを押さえておくと、TorchTPUがなぜ独自の実行モデルやメモリ管理を必要とするのかが理解しやすくなる。

チップの中身

計算の中心はTensorCoreで、行列積ユニット（MXU）、ベクタユニット（VPU）、スカラユニットから構成される。MXUは256×256の積和演算器をシストリックアレイ状に並べたもので、世代ごとに規模が異なる。入力データが演算器の間を流れながら積和が進むため、大きな行列積を高い電力効率で処理できる。これに加えてSparseCoreという専用ユニットがあり、gather/scatterや埋め込みテーブルの参照、集合通信（collective）を担当する。SparseCoreはTensorCoreと並行して動かせるため、密な行列積とスパースな処理、あるいは計算と通信を重ね合わせられる。

graph LR
    TC["TensorCore<br/>MXU（256×256 シストリックアレイ）<br/>VPU / スカラユニット"]
    SC["SparseCore<br/>gather・scatter・埋め込み<br/>集合通信"]
    HBM["HBM<br/>TPU 8i で 288GB / 8.6TB/s"]
    ICI["ICI（チップ間）<br/>19.2TBps"]
    DCN["DCN（スライス間）<br/>TPU 8t で 400GBps"]
    TC <--> HBM
    SC <--> HBM
    TC <--> ICI
    SC <--> ICI
    ICI <--> DCN

接続経路によって帯域が大きく異なる点も、TPUの性能を左右する。

TPU Podはスライスを論理的にまとめた単位で、チップ同士は世代に応じて3DトーラスやBoardflyと呼ばれる構成で接続される。階層ごとに帯域が違うため、並列化の戦略を物理トポロジに合わせることが重要になる。

プログラミングモデルの要点

TPUのプログラミングモデルはGPUとは発想が異なる。TensorCoreは事実上シングルスレッドの演算器であり、性能を引き出すには行列の形状を十分に太くしてMXUを使い切る必要がある。GPUのように小さなカーネルを大量に並列起動する戦略は通用しない。また、1つのTPUデバイスには1つのユーザープロセスからしかアクセスできず、TorchTPUはJAXの「1ホスト1プロセス」ではなくPyTorch流の「1デバイス1プロセス」を採用している。

ソフトウェア側の制約も設計に影響する。集合通信はNCCLのような別ライブラリではなくStableHLOの演算として表現され、XLAコンパイラが処理する。StableHLOはSSA（静的単一代入）に従うため値をin-placeで書き換えられず、PyTorchの「view」やミューテーションの意味論をそのまま表現できない。加えてXLAはコンパイル時にテンソル形状が確定していることを前提とするため、実行中にバッチサイズや系列長が変わる動的形状と衝突する。これらの差分を、TorchTPUはviewの再現やBounded Dynamismといった仕組みで埋めている。

実行モデル：EagerとCompiledをどう両立させているか

TorchTPUの実行モデルは「eager実行を一級市民として残し、必要な範囲だけをまとめてコンパイルする」二段構えだ。eager実行とは、Atenオペレータを呼び出し順にその場でディスパッチする命令的な実行方式で、GPU上の通常のPyTorchと同じ感覚でコードが動く。コンパイルはその上に載る高速化オプションという位置づけになる。

3つのEagerモード

DeferAndFuseはStrict比で50%以上速くなる場合があり、融合で中間バッファが減るためピークHBM使用量も下がる。学習スクリプトの段階的最適化では、素のeager 11.49秒がDeferAndFuseで9.95秒（1.15倍）になった例がある。

遅延実行の仕組み

ディスパッチャから呼ばれると、TorchTPUは計算せず「Deferred Op」を生成する。内部にはStableHLOビルダーと、キャッシュキー計算に使う静的パラメータの対応表を持つ。.to("cpu")や.item()のような同期点（マテリアライゼーション・ポイント）に到達すると、そこまでの遅延領域を解析し、生存解析で実体化すべき演算を決める。集合通信のような副作用のある演算は必ず順序ごと保存され、Pythonの参照カウントによりレプリカ間でも生存判定が決定論的に一致する。同じ演算列が繰り返される場合は、Repeated Ops Heuristicで早めに実体化を起こす。

sequenceDiagram
    participant U as ユーザーコード
    participant D as Atenディスパッチャ
    participant B as 遅延バッファ
    participant C as コンパイルキャッシュ
    U->>D: オペレータ呼び出し
    D->>B: Deferred Opを生成
    U->>D: .item() / .to("cpu") などの同期点
    D->>B: 生存解析で実体化対象を決定
    D->>C: フィンガープリントを照会
    alt キャッシュヒット
        C->>U: 既存バイナリを起動して値を返す
    else キャッシュミス
        C->>U: 非同期にコンパイルしてから起動
    end

コンパイルキャッシュとCompiledモード

キャッシュは3階層に分かれる。Tier-1はプロセス内メモリ（常時有効）、Tier-2はホスト内共有メモリ（分散実行で既定有効、単一デバイスでは無効）、Tier-3はCloud Storageなどの永続キャッシュ（既定無効、明示指定で有効）で、ホストやジョブの再起動をまたいで再利用できる。キーはStableHLOグラフそのものではなく、構造・テンソルサイズ・コンパイラオプションとバージョンの3つの64ビット整数（計192ビット）のフィンガープリントだ。

torch.compile(model, backend="tpu")では、Torch DynamoがFXグラフとサンプル入力を渡し、TorchTPUは全オペレータを実体化せずにトレースして、グラフ全体を1つのStableHLOプログラムとしてコンパイルする。TorchInductorを経由せずXLAを主バックエンドに選んだのは、eager向けに構築したコンパイル基盤を再利用でき、XLA/TPUコンパイラがGoogleで長年実運用され安定しているためだ。なお現時点では動的形状（SymInt）対応が既定で有効でないため、torch.compile時にdynamic=Falseが自動設定される。これはBounded Dynamismの本番化で解消予定の一時的な制約である。

動的形状への対処：Bounded Dynamism

XLA（Accelerated Linear Algebra）は、TPU向けの計算グラフをコンパイルするコンパイラであり、コンパイル時に形状が確定していることを前提に、最適なハードウェアレイアウトや融合カーネルを生成する。この前提は、バッチサイズや系列長が実行中に変わる推論・学習と相性が悪い。TorchTPUは「bounded dynamism（上限付きの動的形状）」という仕組みで、この衝突を吸収している。

問題：静的シェイプ前提のコンパイラ

静的シェイプを前提とするコンパイラでは、実行中に入力形状が変わると、素朴な実装では形状ごとにグラフ全体を再コンパイルすることになり、実行性能が大きく劣化する。

よく知られた回避策は、あらかじめ大きめのテンソルを確保し、手作業でパディングとスライスを挟む方法である。しかしこれには3つの問題がある。第一に、モデル全体に侵襲的なコード変更を要求する。第二に、パディング値は演算ごとに異なる（torch.sumなら0、torch.mulなら1など）ため、演算をまたいで一様に扱えない。第三に、パディングとスライスを行う関数自体が可変形状の入力を受け取るため、結局その関数の再コンパイルが発生してしまう。

解決：上限付きの動的形状

TorchTPUは、次元の最大上限だけをコンパイル時に与え、その範囲内の任意の形状を1つのコンパイル済みバイナリで処理する。コンパイルの流れは次のとおりである。

flowchart TD
    A["torch._dynamo.mark_dynamic で動的次元と上限を指定"] --> B["Dynamo が SymInt としてシンボリックトレース"]
    B --> C["バックエンドが SymInt をスカラテンソルへ置換し<br/>_pad_real_to_static と _set_dimension_logical_size を挿入"]
    C --> D["境界付き動的テンソル i32[2, ≤16] を生成"]
    D --> E["StableHLO へ lowering してコンパイル"]
    E --> F["実行時: PadRealToStatic が末尾をパディングし<br/>DynamicPadder がパディング領域を単位元へリセット"]

実行時の振る舞いにも工夫がある。動的サイズのテンソルはそのままコンパイル済みモジュールへ渡され、カスタムopのPadRealToStaticが動的次元の末尾をパディングして上限付きの静的バッファを作る。演算はこのバッファ全体に対して一様に実行されるが、XLAのDynamicPadderパスがパディング領域の値を演算の単位元（加算なら0、乗算なら1）へ自動的にリセットするため、結果には影響しない。中間バッファは論理サイズの情報を保持し、動的テンソルが出力になるときはパディングが末尾に位置するよう構成されているため、テンソルサイズのメタデータを実サイズへ更新するだけで同じバッファを返せる。物理的なスライスは不要である。

検証では、因果LLM 839モデルのうち834モデル（99%）で数値が正しいことが確認されている。一方、torch.compileでは現時点でdynamic=Falseが自動設定される。SymInt対応に必要なbounded dynamismが既定で有効になっていないためで、本番化によって解消される予定である。上限を超える形状や上限の再設定が起きた場合は従来どおり再コンパイルが発生する点は、運用時に押さえておきたい。

数値精度の扱い：BF16と3パス／6パス

TPUのネイティブ形式

TPUの行列積ユニット（MXU）は、bfloat16（BF16）をネイティブな入力形式として設計されている。BF16は符号1ビット・指数8ビット・仮数7ビットの16ビット浮動小数点数で、指数部をFP32と同じ幅に保つことで値のレンジを維持しつつ、仮数部を削って演算器の面積と帯域を節約する。MXUはこのBF16入力を受け取り、内部ではFP32で積和を蓄積する。つまり「入力はBF16、蓄積はFP32」という組み合わせが、TPUの行列積の基本形である。

一方、行列積以外の非収縮演算（要素ごとの加算、活性化関数、LayerNorm、Softmaxなど）は、ベクタユニット（VPU）がIEEE-754準拠のFP32で実行する。行列積だけがBF16に落ち、それ以外はFP32のままという分担が、TPUの精度設計の土台になっている。

PyTorchの精度モードとの対応

TorchTPUはPyTorchの精度モード（medium／high／highest）を、TPUのハードウェア特性に合わせて次のように対応付けている。

mediumは乗算をネイティブのBF16に変換し、蓄積だけをFP32で行う。BF16の仮数部は7ビットしかないため単純な丸めでは精度が落ちるが、蓄積をFP32に保つことで、実用上はほぼすべての深層学習ワークロードで十分な精度が得られる。これがTorchTPUの既定になっている。

highとhighestは、BF16では表現しきれない仮数部の情報を「パスを増やす」ことで回復する方式である。highは入力を上位・下位の2つに分割し、その積を3つのクロス項として計算する3パス方式で、FP32精度の大半を取り戻す。highestはすべてのクロス項を計算する6パス方式で、IEEE FP32に近い精度になる。ただし、その代償としてMXUの計算量はhighで約3倍、highestで約6倍に増える。

graph LR
    A["FP32 入力"] --> B["上位ビット（BF16）"]
    A --> C["下位ビット（BF16）"]
    B --> D["上位×上位"]
    B --> E["上位×下位"]
    C --> F["下位×上位"]
    C --> G["下位×下位"]
    D --> H["FP32 蓄積"]
    E --> H
    F --> H
    G --> H

この分解の考え方は、GPUのTensor CoreがTF32（符号1・指数8・仮数10の19ビット）という専用形式でFP32近似を行うのと対照的である。TPUのMXUはBF16（16ビット）しか受け付けないため、高精度が必要な場合は「専用ユニットに切り替える」のではなく「同じMXUでパスを増やす」というアルゴリズム的な分解で対応する。CUDAではhigh／highestでもFP32蓄積のままほぼ基準速度〜1〜2倍程度のコストに収まるのに対し、TPUでは3倍／6倍と跳ね上がる点が、実運用上の大きな違いになる。

TorchTPUは大多数のワークロードにmedium（既定）を推奨しており、high／highestを明示的に選ぶとTORCH_WARN_ONCEによる性能警告が表示される。これらは精度に起因する収束問題を具体的に確認した場合にのみ使うべきものと位置づけられている。さらに、グローバルな設定に加えてスレッドローカルのtorch.tpu.precisionコンテキストマネージャが用意されており、数値的に敏感な層だけをhighに上げ、残りはmediumのままにするといった細かな制御が可能である。

PyTorchの意味論をどう守っているか

Viewとin-place書き換え

PyTorchのテンソルは、線形なバイト配列へのポインタとストライド（各次元の歩幅）の組として表現される。複数のテンソルが同じバイト配列を共有でき、あるテンソルへの書き込みが別のテンソルからも見える——このエイリアシングと書き換えの可視性が意味論の根幹にあり、viewやas_strided、分散学習でのパラメータ配列の連結といったパターンがこれに依存している。

一方、TorchTPUが中間表現に用いるStableHLOは純粋関数的で、すべての値は不変（イミュータブル）である。ストライドを直接操作する命令もin-place書き換えも存在せず、部分更新はdynamic_update_sliceで「元のテンソルの一部を新しい値に差し替えた新しいテンソル」として表すしかない。

TorchTPUはこの差を「view分解」「view反転」「遅延オペレータ」の組み合わせで埋める。viewの生成時には、そのストライドパターンをreshape（隣接次元の平坦化・非平坦化）、transpose、broadcast、slice、bitcastといったStableHLOのプリミティブ列に分解し、分解結果を計算する遅延オペレータとして保持する。viewへの書き込みが起きた場合は、その書き込みと元のview操作の逆操作を合成的に組み合わせた新しい遅延オペレータを基底テンソル側に作り、ストレージのポインタを差し替える。物理的には同じバッファを共有しないまま、両方のテンソルが正しい値を返すようになる。

具体例として、t1 = torch.arange(4)からt2 = t1.as_strided(size=(2,), stride=(2,))を作り、t2.fill_(5)を実行すると、t2は[5., 5.]、t1は[5., 1., 5., 1.]になる。StableHLOにas_stridedは存在しないが、この書き換えの波及は再現される。ブロードキャスト次元やオーバーラップした次元へのin-place書き込みはPyTorch自体が未定義動作として扱っており、TorchTPUもサポートしない。この仕組みはeager・compiled・ハイブリッドのいずれのモードでも、ユーザーの介入なしに機能する。

空テンソル

torch.empty()は「未初期化のメモリを確保して返す」関数だが、値が不変な中間表現には確保と未初期化という概念をそのまま持ち込めない。そこでTorchTPUはHBMを確保せず、NaNで埋めたテンソルを計算する遅延オペレータを返す。返り値がそのまま書き込み先として使われる場合、NaNが読み出されることはないため実体化されず、初期化のコストも発生しない。値が実際に読まれたときにだけNaNが現れるため、未初期化メモリを読むPyTorchの挙動と整合する。

カーネル拡張：Pallas・Tokamax・Helion

TorchTPUが「既存コードをそのまま動かす」ことを第一に設計されている一方で、TPUのピーク性能を引き出すには、モデル固有の部分をデバイスに合わせて書き下すカスタムカーネルが欠かせない。TorchTPUはこの領域に、成熟度の異なる3つの経路を用意している。

まず、PyTorchフレームワーク自身が提供するカーネル（torch.nn.functional.scaled_dot_product_attentionなど）については、当初Pallasで書いたTPUカーネルをMosaicへシンボリック形状でエクスポートし、呼び出し時に形状を特化する方式を実装した。しかし想定より柔軟性が低かったため、動的形状を直接扱えるMosaicカーネルを書く新しい方式を実験中だ。

ユーザーが自前のカーネルを書く場合の選択肢は次の3つである。

• Pallas（本番利用可）：JAXのカーネルDSLであるPallasでTPUカーネルを書き、torch.tpu.pallas.jax_opでラップして通常のPyTorch関数として呼び出す。低レベルなハードウェア命令まで制御でき、device="tpu"のテンソルにそのまま使える。
• Tokamax（本番利用可）：OpenXLAの高性能TPUカーネルライブラリ。MoEのビルディングブロック（グループ行列積、トークンのgather/scatter）、attention変種（paged attention、Gated DeltaNet、Kimi delta attention、multiheaded latent attention）、TopK計算など、高度に最適化された実装を提供する。オートチューニングやXProfベースのベンチマーク、回帰テストといった運用基盤も備え、カーネル自体はJAX/Pallasで書かれているが、torch.tpu.pallas.jax_op APIが相互運用を透過的に管理する。
• Helion（実験的）：PyTorchベースのDSLで書いたカーネルを、同じコードのまま複数のアクセラレータ向けにコンパイルできる。@helion.kernel(backend="pallas")で定義し、hl.tileで自動分割されたタイルをループする。

graph TD
    A["Pallas（JAX DSL）<br/>本番利用可"] --> D["torch.tpu.pallas.jax_op"]
    B["Tokamax（OpenXLA ライブラリ）<br/>本番利用可"] --> D
    C["Helion（PyTorch DSL）<br/>実験的"] --> E["複数アクセラレータ向けコンパイル"]
    D --> F["PyTorch テンソル<br/>device='tpu'"]
    E --> F

実務上の使い分けは明快だ。まずはTokamaxが提供する既製カーネルで性能を引き上げ、それでも足りない部分をPallasで書き下す。Helionは「GPUとTPUの両方で同じカーネルを動かしたい」という将来ニーズに向けた選択肢であり、現時点では実験的な位置づけにとどまる。いずれの経路も最終的には通常のPyTorchテンソルに着地するため、モデル本体のコードをカーネル実装から切り離したまま保てる点が、TorchTPUのカーネル拡張の設計上の特徴である。

分散実行：PyTorch流の集合通信をTPUに載せる

PyTorchの分散実行は、torch.distributedのProcess Group APIを起点に、各プロセスが集合通信（all_reduce、all_gather、reduce_scatterなど、全プロセスが参加する協調的なデータ交換）を呼び出す形で組み立てられる。TorchTPUはこれを別ライブラリに逃がさず、c10d::Backendを継承したバックエンドとして実装する。バックエンド名はtpu_distで、eager実行でもtorch.compileによるコンパイル実行でも集合通信を捕捉できるため、DTensor、FSDPv2、SimpleFSDP、FairScaleといった分散プリミティブがそのまま動く。プロセス構成もGPUと同じ「1デバイス1プロセス」を採用する。

ただしTPUにはGPUと異なる難しさがある。GPUでは1プロセスがホスト上の全デバイスを見えるため、プロセスグループ初期化の前にモデルをcuda:{local_rank}へ移せるが、TPUでは1プロセスが1デバイスしか見えず、ホストに何台接続されているかも分からない。そのため分散実行の意図、ワールドサイズ、TPU種別（トポロジに影響する）を初期化の前に知る必要がある。さらにTPUの集合通信はStableHLOの演算として表現され、stablehlo.all_reduceは自分が属さないサブグループも含めた全サブグループの情報（replica_groups）を要求する。

最大の難所はMPMD（Multiple Program, Multiple Data）対応だ。TPUスタックは伝統的にSPMDを前提に構築されてきたが、PyTorchエコシステムではrankに条件付けた分岐や、テンソル次元がデバイス数で割り切れないときのDTensorの不均等分割により、プロセスごとに実行グラフが異なりうる。集合通信を含むモジュールがランク間で一致しないと、同期フラグIDやバッファアドレスがずれ、予測不能なデッドロックやメモリ破壊に至る。

この問題に対し、TorchTPUは段階的に3つの方式を用意している。

起動面では、通常のtorchrunが設定方法を知らないTPU固有の変数が2つ必要になる。TORCH_TPU_TOPOLOGYはTPUスライスの物理メッシュ形状（2x2x1なら2,2,1）、TORCH_TPU_SLICEBUILDER_ADDRESSESはチップ間インターコネクト（ICI）のメッシュ構築に使うポート一覧だ。torch.tpu.distributed.set_tpu_launch_env()を使えばPythonから設定できる。ノードの既定プロセスグループは、そのノードのスライスに対応する。

物理トポロジへの配置も性能を左右する。TPUは階層型ネットワークを持ち、スライス内のチップを直接つなぐICIは数十〜数百GB/sと高速だが、スライス間をつなぐDCNはEthernetベースで帯域が大幅に低く遅延も大きい。原則は「通信強度の低い次元から高い次元へメッシュを並べる」ことで、毎層all_reduce／reduce_scatterを伴うTensor Parallelを最速のICIリングへ、all_to_allを使うContext/Sequence ParallelとExpert Parallelをスライス内のローカルICI軸へ、FSDPをローカルICI軸へ、勾配のall_reduceだけを行うレプリカData ParallelをICIまたはDCNへ割り当てる。torch.tpu.topology_aware_mesh()が物理ICIリングに沿った最適なrank配置を計算し、それをDeviceMeshに渡して構築する。

graph TB
    subgraph Slice["TPUスライス（高速ICIメッシュ）"]
        A["dp=0 / tp=0"] --- B["dp=0 / tp=1"]
        C["dp=1 / tp=0"] --- D["dp=1 / tp=1"]
        A --- C
        B --- D
    end
    E["別スライス"] -. "DCN（低速・高遅延）" .- Slice

さらに、ホストCPUを介さずICI上のハードウェアRDMAでチップ間通信を直接行う、分散Pallasカーネルも書ける。Ring Attention、カスタムall-to-all、パイプライン並列のシフト、ハロー交換などに向く。pltpu.make_async_remote_copyで非同期リモートコピーを発行し、jax.shard_mapとDMAセマフォでラップし、torch.tpu.pallas.jax_opでPyTorchのテンソル上から呼び出せる。以前は生のハードウェアデバイスIDで通信先を決めていたため、IDとrankの順序がずれると意図しない物理デバイスへ送ってしまっていた。現在はPallasのパーティションをPyTorchのrankへ直接マッピングし、DeviceIdType.LOGICALと正規のデバイス割り当てを強制することで、全ワールド構成でも孤立サブグループ構成でも一貫したルーティングを保証している。

StreamとSparseCore：計算・通信・転送の重ね合わせ

Stream（ストリーム）は、PyTorchが非同期実行の順序を管理するためのコマンドキューの抽象である。同じStreamに投入された演算は順序どおりに実行され、異なるStreamの演算は互いに独立して進む。GPUではこの仕組みを使い、SMで走る計算カーネル、NCCLによる集合通信、ホストとデバイス間のデータ転送という3種類の処理を別々のキューに振り分けて重ね合わせる。TorchTPUもこのPyTorchのStream APIを実装しているが、TPUでは前提が異なる。TPUのTensorCoreは本質的にシングルスレッドの演算器であり、1つのTensorCore上で複数の計算を同時にディスパッチすることはできない。

その代わり、TensorCoreにはDMAエンジンとSparseCoreエンジンが付随する。DMAエンジンはホスト・デバイス間の転送をTensorCoreの演算と並行して実行できる。SparseCoreはgather/scatterや埋め込みテーブルの参照のようなスパース演算に最適化されたユニットで、集合通信もサポートする。原理的にはStableHLOのあらゆる演算を実行しうるが、スパース計算ではTensorCoreより速く、密な行列積ではMXUを備えたTensorCoreに劣る。TPUにおける並行実行とは、同じ計算ユニットを増やすことではなく、計算・通信・転送を別々のエンジンに割り当てて同時に動かすことだと言える。

StreamをSparseCoreに固定するには torch.tpu.Stream(pinned="sparse") を使う。このStreamで集合通信を実行し、既定（TensorCoreで動く）Streamでは行列積を流し、record_event と wait_event で依存関係を張れば、通信と計算を重ねられる。

sparse_stream = torch.tpu.Stream(pinned="sparse")
with torch.tpu.stream(sparse_stream):
    dist.all_gather_into_tensor(out_tensor, in_tensor)
    ev = sparse_stream.record_event()

with torch.tpu.stream(torch.tpu.default_stream()):
    y = torch.matmul(a, b)          # TensorCoreで密な計算
    torch.tpu.default_stream().wait_event(ev)  # 集合通信の完了を待つ

gantt
    title TensorCore・SparseCore・DMAの重ね合わせ
    dateFormat X
    axisFormat %s
    section TensorCore
        torch.matmul :a1, 0, 6
        torch.matmul（後続） :a2, 6, 5
    section SparseCore
        all_gather :b1, 0, 4
        gather / scatter :b2, 5, 5
    section DMA
        ホスト→デバイス転送 :c1, 0, 2
        損失のデバイス→ホスト転送 :c2, 11, 2

TorchTPUは「1プロセス1TensorCore」のモデルを採るため、複数のStreamをTensorCoreに割り当てても計算そのものは速くならない。またSparseCoreに固定したStreamは、演算サポートの網羅性や性能特性が既定のStreamと同一ではない点に注意が必要だ。転送の非同期化にもStreamは有効で、訓練サンプルのホスト→デバイス転送を1つのStreamで開始し、その完了を待って訓練ステップを実行し、計算終了後に損失のデバイス→ホスト転送を別Streamで開始する構成を組める。CPU側はホスト→デバイス転送のStreamだけを同期すれば、ホストメモリを解放してよい時点を把握しつつ、計算を止めずに済む。

この重ね合わせは、集合通信を独立した小さなStableHLOモジュールに分けて安全側に倒すSkinny Collectivesの性能損失を補う手段としても使われる。Skinny CollectivesはXLAが計算と通信を重ねられないためエンドツーエンドで30〜40%の性能低下を招きうるが、集合通信をSparseCore固定のストリームで走らせ、TensorCoreの密な演算と並行させれば損失の大半を回復できる見込みで、改善が進行中である。

実力検証：6つのケーススタディ

設計思想だけでなく、実際にどこまで動くのか。論文は6つのケーススタディでTorchTPUの実力を裏付けている。ここでは数値を中心に整理する。

ケース1：Hugging Faceモデルの網羅検証

NLP・Vision・Audio・マルチモーダル・拡散生成にわたるHugging Faceモデルを継続的に検証する自動テスト基盤を構築し、TPU v5e以降の世代でeager・compiled両モードを対象にしている。人気重視のTier-1では4,917モデル中98%（4,821件）がクラッシュなしで動作し、アーキテクチャ多様性重視のTier-2では657モデル中638件（97%）がGPU比で数値一致した。判定は因果LLMがperplexityの相対差0.10以内、Vision/Audio/Encoderが相対MSE 0.10以下（または絶対MSE 0.05未満）という二重閾値で行われる。

ケース2：素のPyTorchコードでの性能

nn.LinearやF.scaled_dot_product_attentionなど標準モジュールだけで書いたMuse Glimmer 30B（モデル固有カーネルは意図的に不使用）で検証した。単一TPU v7チップの推論（prefill）は系列長5,120で約25.7k TPS・MFU 59%、4チップのFSDP学習はTPUあたり約5.6k TPS・MFU 38%。さらにGPU/CUDA向けに書かれたTorchTitan版を、未対応のFlexAttentionを標準APIに置き換える程度の最小改修で載せても、4チップ学習でTPUあたり約5.3k TPS・MFU 36%を出した。カスタムカーネルなしでも実用的な性能が出る点が、導入のしやすさを支えている。

ケース3：Pod規模へのスケーリング

Qwen3 MoE 235BをTorchTitanで最大1つのTPU v7 Podまで検証した。64チップではPyTorch/TorchTitanが1,828 TPSと、JAX/MaxTextの1,937 TPSと同水準。スケールを上げてもMFUがほぼ横ばいの弱スケーリングが確認されている。

マルチスライスHSDPでは、逆伝播中の層t+1のall-reduce（DCN経由）を層tの計算と完全に重ね合わせることで、スライス数を増やしてもTPS/TensorCoreがほぼ一定に保たれる。

ケース4：vLLMによるLLM推論（SemiAnalysis InferenceX）

第三者評価であるSemiAnalysisのInferenceXで、TPU v7とvLLM-TorchTPUスタックがMoEワークロードで高いコスト効率を示した。InferenceX 8K/1Kでは100万トークンあたり2.01セント（9月初出時は3.6セント）、エージェント的ワークロードを測るAgentXでは0.83セントで応答性も最高水準。エキスパート並列MoE層の単一SPMDプログラム化や、FP8 KVキャッシュのパディングを排除するSequence-on-Laneレイアウト、Gated DeltaNetのPallasカーネル融合といった最適化が効いている。

ケース5：TorchTitanによるファインチューニング

FLUX.1-dev 12Bで段階的最適化の効果を定量化した。素のeagerの11.49秒から、DeferAndFuse（9.95秒）→SimpleFSDP（6.52秒）→Pallas SplashAttention（5.55秒）→torch.compile(backend="tpu")（3.02秒）→ピープホール書き換え（2.78秒）と、合計3.80倍の高速化を達成している。大規模側ではLlama 3 405Bを1,024チップ・8Kコンテキスト・BF16で動かし、500トークン/秒/チップ・MFU 50%超を記録した。

ケース6：推薦モデルとSparseCore（TorchRec）

高基数の埋め込みテーブルがボトルネックになる推薦モデルでは、TorchRecにTPUバックエンドを上流統合し、ホスト側のボトルネックを迂回してSparseCoreを直接使う。勾配をHBMに実体化しないIn-Place Optimizer Fusionで勾配メモリを削減し、DLRM v2を32チップのTPU v7xで評価したところ、JAXの最適化ベースラインと同一の18.3msステップ遅延（100%の性能一致）と同一の損失軌跡を確認した。

これら6つの結果は、TorchTPUが「動く」だけでなく、推論・訓練・推薦という主要ワークロードで実運用に耐える性能を持つことを示している。

エコシステムとの統合とオープンソース化

2026年10月のPyTorch Conferenceで、TorchTPUとその周辺の4つの基盤リポジトリ（TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU、TorchTitan）がGitHubでオープンソース化された。ライセンスはApache 2.0、対応するPythonは3.11以降で、PyTorchのリリース互換マトリクスに合わせている。インストールはpip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpuのように行う。extra indexの指定はCPUビルドのtorchを取得するために必要で、これがないと巨大なCUDAビルドが引き込まれる。

統合は、サービング・訓練／強化学習・モデル配布・運用支援の4つのレイヤに整理できる。

サービング層では、vLLMが連続バッチング、量子化、マルチチップ並列をTPU上でサポートする。上流のワーカー・スケジューラ・APIサーバをそのまま保つ設計のため、既存の推論スタックを変更せずに載せられる。SGLangはRadixAttentionでリクエスト間のKVキャッシュを再利用し、繰り返し現れるプロンプトプレフィックスの再計算を丸ごと省く。転送を分離したバックエンドも備える。llm-d 1.0はCNCF Sandboxプロジェクトで、計算の重いprefillとメモリの重いdecodeを分離し、プレフィックスキャッシュの有無を意識してリクエストを文脈を持つレプリカへ振り分ける。GKEではGPUとTPUを共通のエンドポイントで扱える。

訓練と強化学習では、TorchTitanがLlama 3の70B／405B、Qwen 3の8B・32B・235B MoE、MuseGlimmer 30B、FLUX.1 12Bといったレシピを提供し、最大8,000チップのIronwood TPU（TPU v7x）で検証されている。VeRLにはGRPO向けのTPUレシピが上流統合され、サンプリングと訓練を分離した構成を取る。RayはCPU・GPU・TPUをまたぐデータ処理・訓練・強化学習・サービングを単一のワークフローで編成できる。

モデル配布とエッジでは、Hugging Face Transformers v5がdevice_map="tpu"の指定だけでモデルをTPUへ配備でき、Diffusersはeager実行、torch.compile、テンソル並列をサポートする。TPU v6e-4でQwen-Image-2.1の1024×1024・40ステップを生成した例では、コンパイル後3.5秒に対しeagerは8.0秒だった。LiteRTはlitert-torch.convert()でPyTorchモデルを変換し、Android、iOS、macOS、Windows、Linux、Web、IoTデバイスへ展開する。TPU agent suiteはプレビュー版として、オンボーディングや診断、推論性能のチューニングを自動化する。

graph TD
    T["TorchTPU<br/>PyTorch ネイティブ TPU バックエンド"]
    T --> S["サービング<br/>vLLM / SGLang / llm-d 1.0"]
    T --> R["訓練・強化学習<br/>TorchTitan / VeRL / Ray"]
    T --> M["モデル配布・エッジ<br/>Transformers v5 / Diffusers / LiteRT"]
    T --> O["運用支援<br/>TPU agent suite"]

制約と注意点

TorchTPUは既存のPyTorchコードをほぼそのままTPUへ載せられるが、すべてのワークロードや構成が等しく快適に動くわけではない。導入を検討する際に把握しておくべき制約を整理する。

とくに誤解されやすいのは、「動く」と「速い」の間に大きな隔たりがある点である。行列積ユニット（MXU）はシストリックアレイであり、行列の形状が細いと演算器を使い切れない。Attentionのヘッド次元が64のままでは現行世代のTPUでピークの行列積効率が出にくく、128や256といった太い形状にするほうが有利になる。まずは既存コードをeagerモードで正しく動かし、計測しながら段階的に最適化を積み上げる進め方が現実的である。

graph TD
    A["既存のPyTorchコードを<br/>eagerモードで実行"] --> B{"性能は<br/>十分か"}
    B -->|いいえ| C["DeferAndFuse で<br/>オペレータを融合"]
    C --> D{"まだ足りない"}
    D -->|いいえ| E["SimpleFSDP などで<br/>分散構成を見直す"]
    E --> F["torch.compile<br/>backend='tpu'"]
    F --> G["Pallas カーネルや<br/>専用スタックを導入"]
    G --> H["形状・シャーディングを<br/>TPU向けに再設計"]
    B -->|はい| I["そのまま運用"]
    D -->|いいえ| I

段階的な最適化の効果は大きく、拡散モデルの学習では素のeager実行からDeferAndFuse、SimpleFSDP、Pallasカーネル、torch.compile、ピープホール書き換えと進めることで、ステップ時間が11.49秒から2.78秒へ短縮された例が報告されている。逆に、GPU向けのままで最高性能を狙うと、性能差がそのままコスト差として現れる点には注意が必要である。

今後のロードマップ

TorchTPUはすでに公開リポジトリとして配布されており、今後6〜12ヶ月の開発は「基盤の堅牢化」「Pod規模の分散訓練」「新興アーキテクチャとサービング」「高効率通信と診断」の4本柱で進められる。方向性を一言でまとめれば、eager実行の体験をさらに磨きつつ、大規模学習と低遅延推論の両端でTPUの性能を取り切るための土台を整えるというものだ。

graph LR
    A["基盤の堅牢化"] --> B["分散訓練の拡大"]
    B --> C["新興アーキテクチャ対応"]
    C --> D["通信と診断の高度化"]
    D -.->|フィードバック| A

注目したいのは、優先順位が単なる性能向上ではなく「予測可能な体験」の確立に置かれている点である。事前コンパイル済みの動的カーネルライブラリが揃えば、初回実行時のコンパイル待ちが減り、Bounded Dynamismの本番化は動的なバッチサイズや系列長で再コンパイルが起きない状態につながる。分散訓練の側では、Qwen3 MoE 235Bを8,192チップまで検証した実績を足がかりに、対応するモデルとシャーディング手法を広げ、訓練とvLLM/SGLangによるサンプリングを密結合した強化学習ループをRay Core上で編成する構想が示されている。

診断基盤の強化も見逃せない。WebGPUベースのTrace Viewer 2.0への移行と常時プロファイリング（XProfz）は、Pod規模の実行で断続的にしか現れないボトルネックを捉えるための仕組みである。モデルアーキテクチャの優先順位づけについて外部からのフィードバックも求められており、公開リポジトリを通じて開発の方向性に影響を与えられる段階に入っている。

よくある質問

導入や移行を検討する際につまずきやすい論点を、質問形式でまとめる。

TorchTPUを使うにはコードを書き換える必要がありますか？

基本的には不要です。model.to("cuda")をmodel.to("tpu")やtorch.device("tpu")に置き換え、分散実行のバックエンド指定をncclからtpu_distに変える程度で、既存のPyTorchコードはそのまま動作します。TPUマシンにTorchTPUのwheelが入っていれば、PyTorchをimportするだけでバックエンドが自動的に有効化されます。ただし最高性能を狙う段階では、形状やシャーディングの見直し、torch.compile(backend="tpu")、Pallasカーネルの導入といった追加作業が効いてきます。

PyTorch/XLAと何が違うのですか？

最大の違いは、eager実行を一級市民として扱っている点です。PyTorch/XLAはLazy Tensorベースで、コンパイル対象を明示するための侵襲的なコード変更、コンパイル遅延の大きさ、eagerモードの不在が課題でした。TorchTPUはTorch Dynamoベースのtorch.compile、TPUの集合通信をProcess Group APIレベルで実装した分散対応、PyTorchのviewとin-place書き換えの意味論の維持を特徴とします。

GPU向けに書いたPyTorchモデルはそのままTPUで動きますか？

動きます。GPU/CUDA向けに書かれたTorchTitan版Muse Glimmer 30Bも、未対応のFlexAttentionを標準のF.scaled_dot_product_attentionに置き換える程度の最小改修で、TPU上で実用的な性能を出しています。ただしTPUのMXUは行列の形状が太いほど効率が上がるため、GPU向けに最適化された形状やシャーディングのままではピーク性能には届きません。

動的なバッチサイズや系列長で再コンパイルは起きませんか？

Bounded Dynamismで解消を進めています。次元の最大上限だけを与えてグラフを一度コンパイルし、実行時は上限内の形状であれば既存のバイナリを再利用する仕組みです。因果LLM 839モデルでの検証では、834モデル（99%）で数値の正しさが確認されています。ただし現時点のtorch.compileではdynamic=Falseが自動設定されるという一時的な制約があります。

大規模な分散学習には対応していますか？

対応しています。DTensor・FSDPv2・SimpleFSDP・FairScaleなどをそのまま使え、Qwen3 MoE 235Bを最大8,192チップまで検証して、MFUがほぼ横ばいになる弱スケーリングを確認しています。rankごとに実行内容が異なるMPMDのコードには、Skinny Collectives、@spmd_safe、Compilation Handshakeという3つの仕組みで対処します。なお、サンプリングと訓練を同一チップに同居させるRLのコロケーションは「1デバイス1プロセス」制約のため、現状はそのままではサポートされません。

まとめ

TorchTPUは、GoogleのTPUをGPUやCPUと並ぶPyTorchのネイティブデバイスにするために、GoogleとMetaのPyTorchチームが共同で開発したオープンソースのバックエンドである。2026年のPyTorch Conferenceで公開され、model.to("cuda")をmodel.to("tpu")へ、分散のバックエンド指定をncclからtpu_distへ置き換える程度で、既存のPyTorchコードがそのまま動く。TPUがJAX専用のハードウェアではなく、PyTorchのデバイス選択肢のひとつになったことが最大の変化である。

設計の核心は二段構えにある。日常的な開発ではAtenオペレータを逐次ディスパッチするeager実行を一級市民として扱い、本番性能を狙う段階では複数のオペレータを遅延させて融合するDeferAndFuseと、torch.compile(backend="tpu")によるグラフ全体のコンパイルを重ねてXLAへ渡す。この両立を、3階層のコンパイルキャッシュ、上限付きの動的形状（Bounded Dynamism）、StableHLO上でのviewとin-place書き換えの意味論の再現が支えている。

TPU固有の制約にも個別の対処が用意されている。シングルスレッドのMXU、SPMDを前提とするXLAコンパイラ、1デバイス1プロセスという条件に対して、StreamとSparseCoreによる計算・通信・転送の重ね合わせ、ICIとDCNの階層に合わせたデバイスメッシュ配置、Process Group APIレベルでの集合通信実装、MPMD向けの3つの仕組みで応える。

実運用を裏付ける結果も揃っている。Hugging Faceのモデル4,917件のうち98%がクラッシュなく動作し、Qwen3 MoE 235Bは8,192チップまでMFUがほぼ横ばいの弱スケーリングを示した。FLUX.1-devのファインチューニングは段階的な最適化で11.49秒から2.78秒へ短縮され、推薦モデルのDLRM v2はJAXの最適化ベースラインと同一の18.3msステップ遅延を記録している。

一方でピーク性能が自動的に得られるわけではない。GPU向けの形状やシャーディングのままではMXUを使い切れず、high・highest精度は計算コストがそれぞれ約3倍・6倍に増える。まず既存コードをeagerで動かして正しさを確認し、計測しながら段階的に詰めていく進め方が現実的である。

graph LR
    A["既存コードを eager で実行<br/>Strict / Debug"] --> B["DeferAndFuse で<br/>オペレータを融合"]
    B --> C["SimpleFSDP / DTensor で<br/>分散を最適化"]
    C --> D["torch.compile(backend='tpu')<br/>グラフ全体を XLA へ"]
    D --> E["Pallas / Tokamax の<br/>カスタムカーネル"]

参考資料

• TorchTPU - A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU（arXiv:2610.24161v1） https://arxiv.org/abs/2610.24161
• Introducing TorchTPU（PyTorch Blog） https://pytorch.org/blog/introducing-torchtpu/
• TorchTPU: Running PyTorch Natively on TPUs at Google Scale（Google for Developers Blog） https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/
• TorchTPU is now open source: Run native PyTorch on TPUs（Google Cloud Blog） https://cloud.google.com/blog/topics/ai-infrastructure/torchtpu-is-now-open-source-run-native-pytorch-on-tpus
• google-pytorch/torch_tpu（GitHubリポジトリ） https://github.com/google-pytorch/torch_tpu
• Hugging Face ❤️ TorchTPU: TPUs, Natively（Hugging Face Blog） https://huggingface.co/blog/hf-torchtpu
• Google Open-Sources TorchTPU, Putting Native PyTorch on Its TPUs（Unite.AI） https://www.unite.ai/google-open-sources-torchtpu-putting-native-pytorch-on-its-tpus/
• TPU Inference Externalization Full Steam Ahead — InferenceX（SemiAnalysis） https://newsletter.semianalysis.com/p/tpu-inferencex-full-steam
• StableHLO仕様 https://openxla.org/stablehlo/spec
• Tokamax https://github.com/openxla/tokamax
• Pallas（JAX） https://docs.jax.dev/en/latest/pallas/index.html
• Helion https://github.com/pytorch/helion
• SimpleFSDP 論文（arXiv:2411.00284） https://arxiv.org/abs/2411.00284
• TPU system architecture https://cloud.google.com/tpu/docs/system-architecture-tpu-vm
• TorchTPUドキュメント https://torchtpu.readthedocs.io/

関連記事

• Google、TPU衛星を軌道投入　宇宙データセンター構想の現実味と課題