【ITニュース解説】PyTorch on TPU: Hands On With Colab

この解説文は、大規模言語モデル「Gemma 4 E2B」をGoogleのTPU（Tensor Processing Unit）上でPyTorchを使って効率的に動作させるためのさまざまな試みとその結果について説明する。TPUはAI計算に特化した高性能なハードウェアであり、PyTorchはその上でAIモデルを構築・実行するための主要なフレームワークである。しかし、必ずしもPyTorchがTPUの性能を最大限に引き出せるわけではなく、そのギャップをどのように埋めるかが課題となる。この記事では、既存のJAXエンジンとの比較を通じて、PyTorchの最適化手法を段階的に検証している。

まず、PyTorchのコードがTPU上でどのように実行されるかの基本的なモードが検証された。大きく分けて3つのモードがある。一つは「Eager（即時実行）」モードで、Pythonコードが書かれた通りに即座に実行される。これは開発がしやすい反面、最適化が効きにくく速度は出にくい。二つ目は「Defer-and-fuse（遅延結合）」モードで、いくつかの操作をまとめてから実行することで、効率を上げる。そして三つ目が「torch.compile（コンパイル）」モードで、指定された関数全体を事前にコンパイルし、TPUに最適化されたプログラムとして実行する。このコンパイルモードが最も高速な実行を可能にするが、入力データの形状が変わると再コンパイルが必要になるという特性を持つ。例えば、ある操作のスタックでは、コンパイルモードはEagerモードの98倍以上高速に動作した。この実験から、コンパイルがいかに重要であるかがわかる。

次に、Gemma 4 E2BモデルをTPUにロードする方法と、その際に発生する課題が検証された。このモデルは「4-bit量子化」という技術で軽量化されており、モデルの重みを通常の半分のサイズに圧縮している。しかし、Hugging Face Transformersライブラリを使ってモデルをロードし、TPUに送るだけでは、モデルの重みがTPU上で自動的に「bf16（bfloat16）」という、より精度は高いがメモリを消費する形式に展開されてしまうことが判明した。これは、4-bit量子化のメリットであるメモリ削減効果を失わせるだけでなく、計算速度にも影響を与える。また、モデルを正しくTPUに配置するためには、モデルの各要素（テンソル）の名前を適切にマッピングし、「device_map='tpu'」という設定を使って明示的にTPUへ送る必要がある。

モデルをロードした後の、テキスト生成の初期性能も測定された。Hugging Faceのgenerateメソッドをそのまま使用すると、約5トークン/秒程度の速度しか出ず、JAXエンジンが同じモデルとチップで達成する130トークン/秒以上という速度には遠く及ばないことが明らかになった。これは、generateがトークンを一つずつ生成するたびに、PyTorchの実行コンテキストが頻繁に切り替わることや、コンパイルのメリットを十分に活かせないことに起因する。特にDefer-and-fuseモードでは、シーケンスの長さが伸びるたびに新しい形状として扱われ、コンパイルオーバーヘッドが大きくなるため、最初の呼び出しに非常に時間がかかることがわかった。

この性能差を縮めるために、いくつかの最適化手法が適用された。まず、「コンパイル済みデコードステップ」の導入である。これは、一つのトークンを生成する処理（デコードステップ）全体をtorch.compileで最適化するというものだ。しかし、compressed-tensorsというライブラリがモデルのモジュールに「ラッパー」と呼ばれる追加の処理を付加しており、これがtorch.compileによるプログラム全体の追跡とコンパイルを阻害していることが判明した。このラッパーを解除することで、デコードステップ全体を一度だけコンパイルし、その後は入力形状が同じであれば高速に再利用できるようになり、生成速度が大幅に向上した。これにより、約5トークン/秒だった速度が31〜33トークン/秒まで改善された。

さらに、最も大きい埋め込み層の重みを4-bitのままTPU上に保持するカスタムコードを導入することで、さらなる性能向上が見られた。通常はbf16に展開される重みを4-bitで保持することで、約3.4 GBのメモリを節約でき、デコード速度も50トークン/秒まで向上した。これは、モデルが元々4-bit量子化で学習されていた（QAT: Quantization-Aware Trained）ため、4-bitで計算しても精度が保たれることを示している。ただし、生成されたトークンを一つずつホスト（CPU側）に読み出すようなストリーミング処理では、TPU上での計算が速くなっても、ホストとのデータ転送がボトルネックとなり、速度向上のメリットが相殺されてしまうこともわかった。

最後に、長文のコンテキストに対応するための「リングキャッシュ」が導入された。Hugging Face Transformersのデフォルトのキャッシュ機構では、スライディングウィンドウサイズが512トークンに固定されており、それ以上の長文コンテキストを効率的に扱えない。これを解決するため、Gemmaモデルの特定の層が使用するキーとバリューのキャッシュを、リングバッファ構造を持つカスタム実装に置き換えた。このリングキャッシュは、モデルの一部の層でフルアテンション（全トークンを見る）を維持しつつ、他の層ではスライディングウィンドウ（直近のトークンを見る）を適用することで、最大131,328トークンもの長大なコンテキストをわずか0.82 GBのキャッシュメモリで保持できるようになった。これにより、長文プロンプトの読み込みも高速化され、デコード速度も38〜45トークン/秒で安定した。

これらの最適化の結果、PyTorchのGemma 4 E2Bモデルは、初期のHugging Facegenerateの約5トークン/秒から、最も最適化された場合で約50トークン/秒（短文）または約38トークン/秒（長文）まで性能が向上した。これは、同じチップ上のJAXエンジンが達成する130トークン/秒にはまだ及ばないものの、当初の14.8倍の差から2.6倍の差まで縮めることに成功したことを意味する。すべての検証で、生成されるテキストの内容は同じであり、推論精度が犠牲になっていないことも確認された。

結論として、短いチャットのような用途では、コンパイル済みデコードステップと4-bitテーブル維持の最適化が有効であり、長文対応にはリングキャッシュが推奨される。最高の性能を求めるならJAXエンジンが有利であるが、PyTorchでも工夫次第でTPUのポテンシャルを大幅に引き出せることを示した。ただし、コンパイルには一定のメモリ（ピーク時16.6〜16.9 GB）と時間が必要であり、TPU v5e-1チップの16 GB HBMメモリ容量に迫るため、より大きなモデルではメモリがボトルネックになる可能性もある。