最近、Googleは、新しいCloud TPU仮想マシン（VM）を発表した。このVMによって、TPUホストマシンへ直接アクセスできるようになる。このVMを使って、同社はTensorFlow、PyTorch、およびJAX on Cloud TPUを開発・デプロイするためにユーザ体験を改善し、新たに提供する。
顧客はこれまでも、TPUチップセットを使ってGoogle Cloudに仮想インスタンスをセットアップできた。ただし、インスタンスが同じサーバ環境で実行されなかったため、これにはいくつかの欠点があった。TPUはネットワーク接続を介してチップセットにリモート接続され、アプリケーションはネットワークを介してデータをTPUに送信し、処理されたデータが返送されるのを待つ必要があるため、処理速度が低下していた。
Cloud TPU VMがプレビュー版となっているため、顧客はTPUチップセットをデプロイされたインスタンスに直接接続できる。これにより、TPUチップセットを使用する際の、さまざまなアプリケーションとGoogle Cloudインスタンスとの間のネットワーク遅延を防ぐことができる。Google AIのプロダクトマネージャーであるAlexander Spiridonov氏は、新しいCloud TPU VMに関するブログ投稿で次のように述べている。
この新しいCloud TPUシステムアーキテクチャは、よりシンプルで柔軟性があります。主なユーザビリティの利点に加えて、コードがTPUに到達するためにデータセンターネットワークを往復する必要がなくなるため、パフォーマンスの向上も実現できます。さらに、大幅なコスト削減も見られる場合があります。以前は、Cloud TPUポッドスライス内のリモートホストにデータを提供するために強力なCompute Engine VMのフリートが必要だったものが、今では、そのデータ処理をCloud TPUホストで直接実行し、Compute Engine VMを追加する必要がなくなります。
出典: https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms
Googleは、2つのバリエーションでCloud TPU VMを提供している。最初のバリエーションは、第2世代のTPUチップセットに基づくCloud TPU v2である。もう一つは、第3世代のTPUに基づく新しいCloud TPU v3バージョンである。Google Cloudによると、この2つの違いはパフォーマンスである。Cloud TPU v2は最大180テラフロップス、TPU v3は最大420テラフロップスで処理できる。
Cloud TPU VMのユースケースは、既存のCloud TPUポッドでアルゴリズムを開発することである。これは、TPUに基づくAIサーバの大規模なクラスターである。特に、このソリューションは、非常に複雑な機械学習モデルの実行に適している。たとえば、最速のクラスターは1秒あたり100ペタフロップスを超える性能を提供する。これによって、このクラスターでのアルゴリズムの構築が大幅に安価になる。顧客は、本番環境に移行するときに、ポッドの使用料とより強力なハードウェアへの移行コストを支払うだけで済む。さらに、Google Cloudは、量子コンピューティングの計画でCloud TPU VMを使うことを計画している。
AIコミュニティのTwitterアカウントであるHuggingfaceは、ツイートで次のように述べている。
JAX/Flaxと新しいCloud TPU V3-8のパワーにより、マスクされたLMをわずか18時間で事前トレーニングできるようになりました！
現在、プレビュー中のCloud TPU VMは現在、us-central1とeurope-west4リージョンで利用できる。これらのVMは、Googleのプリエンプティブサービスを備えたTPUホストマシンあたり1時間あたり1.35ドルから利用できる。詳細については、価格ページをご覧ください。最後に、顧客は、Cloud TPUとCloud TPUポッドを使用して、JAX、PyTorch、TensorFlowを使ったMLモデルのトレーニングをすばやく開始できる。その際、ドキュメントとJAX、PyTorch、TensorFlowのクイックスタートを活用できる。