HN 日本語サマリー

← 一覧へ戻る
AI・機械学習

GoogleがTPUネイティブPyTorchバックエンド「TorchTPU」をオープンソース化

Google open-sources TorchTPU, a native PyTorch back end for TPUs (github.com)

370 pointsby ilreb79 コメント

要約

Googleは、TPU（Tensor Processing Unit）上でPyTorchモデルを実行するためのネイティブバックエンドであるTorchTPUをオープンソースで公開しました。これにより、ユーザーはPyTorchの演算をTPUデバイスで直接実行できるようになります。このプロジェクトは、TPU向けのカスタムATenカーネル、XLA計算の最適化、分散学習やモデル推論の例を提供します。

全文翻訳

［画像: TorchTPU Logo］

説明

このプロジェクトは、GoogleのTensor Processing Units（TPU）向けのPyTorchバックエンドを提供し、ユーザーがTPUデバイス上でPyTorchモデルや操作を実行できるようにします。これには、さまざまなPyTorch操作のためのカスタムATenカーネル、実行を最適化するためのコンパイルキャッシュ、分散トレーニングとモデル推論の例が含まれています。

特長

  • TPU向けのPyTorch ATenカーネル実装。
  • XLA計算の最適化された実行のためのコンパイルキャッシュ。
  • 分散トレーニング（データ並列、テンソル並列）のサポート。
  • さまざまなモデル（Llama、Qwen、ResNet、DLRM、minGPT、UNet）の例。
  • グラフモード実行のためのtorch.compile()との統合。
  • デバッグとベンチマークのためのユーティリティ。

ユーザーガイドとインタラクティブノートブック

ユーザーガイドはhttps://torchtpu.readthedocs.io/で入手できます。これはインタラクティブノートブックの静的バージョンです。ノートブックはdocsフォルダにあります。

ノートブックはMarimo（Jupyterに似ています）です。marimo edit file.pyで実行でき、その後ブラウザでアクセスできます。TPU VMから実行する必要があります。v6e GCEインスタンスが最も簡単です。

インストール

同じ手順が上記のユーザーガイドで利用可能です。

TorchTPUのPythonバージョンサポートポリシーは、PyTorchのリリース互換性マトリックスと一致しています。現在、Python 3.11+がサポートされています。詳細については、Pythonバージョンサポートポリシーを参照してください。

まずPython仮想環境を作成します：

python -m venv .venv
source .venv/bin/activate

または、uvを使用する場合：

uv venv .venv
source .venv/bin/activate
pipによるインストール

最新の公開リリース（PyPI上の0.2.0）をインストールします（uvの場合はすべてのpipコマンドの前にuvを付けます）：

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
ソースからビルド

ホイールは以下のようにビルドできます：

bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel

すべてのOSS Bazelビルドは、.bazelrcを介してデフォルトで最適化モード（-c opt）になります。非最適化ビルド（-c fastbuild、-c dbg）はパフォーマンスを低下させる可能性があり、ホイールビルドには使用されません。C++コードの開発またはデバッグを行っており、明示的に非最適化ホイールが必要な場合は、--//:allow_unoptimized_wheel=Trueを渡してください：

bazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel \
    --//:allow_unoptimized_wheel=True

--config=wheel_commonは必須であり、これがないとビルド時に失敗します。この設定は--//:wheel_build=Trueを設定し、共有XLA/MLIRバックエンドをtorch_version-
リセット遷移にルーティングして、pywrapがそれを単一のlibxla_base.soにファクタリングします。このフラグなしでビルドすると、バージョンごとの共通バックエンドが独自のコピーをシップし、結果のホイールは重複した静的登録によりimport torchでアボートします。

wheel_commonはリモートビルド実行（RBE）のリモートキャッシュとリモートエグゼキュータもオンにします。RBE認証情報がない場合は、それらのフラグを削除するために--config=no_rbeを追加してください：

bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe

順序は重要です：--config=no_rbeは--config=wheel_commonの後である必要があります。なぜなら、後者の設定がリモートキャッシュとリモートエグゼキュータフラグの両方で優先されるからです。

ホイールは以下のようにインストールします：

cd <path_to_repo>
python -m venv .venv; source .venv/bin/activate
# --extra-index-urlは、デフォルトのCUDAビルドではなくCPUバージョンのtorchを提供します。
pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpu

--extra-index-url https://download.pytorch.org/whl/cpuなしでインストールするとCUDAバージョンのtorchがインストールされます。上記のように--extra-index-urlを渡すか、CPUビルドを別途インストールしてください。

開始方法

開始に役立つポインタ：

依存関係

このプロジェクトは依存関係管理にBazelを使用しています。主要なPython依存関係は次のとおりです。

TBD

Bazelセットアップ

このプロジェクトは、ビルドと依存関係管理にBazelを使用しています。

Bazelをインストールするには、実際にはBazel自体のダウンロードとインストールを処理する透過的なラッパーであるBazeliskをインストールします。

Bazeliskをインストールする方法はいくつかあります。一般的な方法を以下に示します：

  • Goがインストールされている場合：

    go install github.com/bazelbuild/bazelisk@latest

  • バイナリを直接ダウンロード：

    wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64

インストール方法にかかわらず、その場所をPATHに追加してください。

インストールを確認するには、リポジトリのルートディレクトリでbazel infoを実行します：

bazel info

エラーなしでリポジトリに関する情報が表示されるはずです。

VSCode clangdセットアップ

C++コードナビゲーションのために、VSCodeとclangd拡張機能を使用することが推奨されます。Bazelアクショングラフから生成されたコンパイルデータベースに依存することで、clangdはコード補完、定義への移動などのさまざまなスマート機能でエディタを強化できます。

  1. UIまたは次のコマンドを使用してVSCode拡張機能をインストールします：

    code --install-extension llvm-vs-code-extensions.vscode-clangd
# 干渉を避けるためにMicrosoft C++拡張機能をアンインストールします。
code --uninstall-extension ms-vscode.cpptools

  2. 初期compile_commands.jsonを生成します：

    tools/setup_clangd.py

初回実行で完全なビルドを行うには時間がかかる場合があります（30分以上）。再実行は増分的で高速です。最近bazel buildを実行した場合、生成されたファイルはすでに存在するため、重いビルドを待機するのをスキップするために--no-buildを渡すことができます（新しいワークスペースでは、ビルドをスキップするとclangdが生成されたファイルを解決できない可能性があります）。スクリプトのファイルレベルコメントで詳細なオプションを確認してください。

または、noxを介してスクリプトを呼び出すこともできます（--の後の引数は転送されます）：

nox -s refresh_compile_commands -- --no-build
nox -s refresh_compile_commands -- //src/torch_tpu/ops/...

貢献

PyTorch on TPUプロジェクトへの貢献を歓迎します！貢献方法に関するガイドラインについては、CONTRIBUTING.mdを参照してください。これには、貢献者ライセンス契約（CLA）への署名やコードレビュープロセスが含まれます。

スター履歴

［画像: Star History Chart］

ライセンス

このプロジェクトはApache 2.0ライセンスの下でライセンスされています。詳細については、LICENSEファイルを参照してください。