コンピューティング

TPU 上の ML モデルをトレーニングしやすくする新しい Cloud TPU VM

2021年6月15日

Google Cloud Japan Team

※この投稿は米国時間 2021 年 5 月 29 日に、Google Cloud blog に投稿されたものの抄訳です。

Google は本日、新しい Cloud TPU VM を公開しました。これを使用すると、TPU ホストマシンに直接アクセスが可能になり、Google の業界最先端の TPU ハードウェアを従来より簡単に使えるようになります。さらに、新しい改善されたユーザーエクスペリエンスにより、Cloud TPU 上で TensorFlow、PyTorch、JAX を開発してデプロイできるようになります。ネットワーク上からリモートで Cloud TPU にアクセスする代わりに、Cloud TPU VM で各 TPU ホストマシンに独自のインタラクティブな開発環境をセットアップできます。

単一の TPU VM を使って行単位で ML モデルを作成してデバッグし、後から Cloud TPU Pod スライス上でスケールアップし、超高速な TPU 相互接続を活用できます。作成するすべての TPU VM に対して root アクセス権があるため、TPU アクセラレータとの密なループで、任意のコードをインストールして実行できます。ローカルストレージの使用や、入力パイプラインでのカスタムコードの実行が可能で、調査や本番環境のワークフローに Cloud TPU を簡単に統合できます。

Cloud TPU と TensorFlow、PyTorch、JAX との統合に加えて、VM 上の新しい libtpu 共有ライブラリにより独自の統合を作成することもできます。

「TPU VM への直接アクセスにより、TPU 上に構築可能なものが完全に変化し、デベロッパーエクスペリエンスとモデルのパフォーマンスが劇的に向上しました。」 - Cohere 共同創設者兼 CEO Aidan Gomez

新しい Cloud TPU アーキテクチャについての詳細

これまで、Cloud TPU へのアクセスはリモートに限られていました。通常は 1 つ以上の VM を作成し、その VM が gRPC を使用してネットワーク上で Cloud TPU ホストマシンと通信を行うという方式でした。

https://storage.googleapis.com/gweb-cloudblog-publish/images/Previous_network-attached_Cloud_TPU_system.max-1300x1300.jpg

これに対して、Cloud TPU VM は次の図に示すように、TPU アクセラレータに直接接続されている TPU ホストマシンで実行されます。

https://storage.googleapis.com/gweb-cloudblog-publish/images/New_Cloud_TPU_VM_system_architecture.max-1400x1400.jpg

この新しい Cloud TPU のシステムアーキテクチャはよりシンプルで、柔軟性があります。ユーザビリティの大きな利点に加えて、コードがデータセンターのネットワークを往復して TPU とやり取りを行う必要がなくなるため、パフォーマンスも向上します。さらに、大幅なコスト削減も見込めます。これまで Cloud TPU Pod のスライス内でリモートホストにデータを供給するために多数の強力な Compute Engine VM を必要としていた場合、Cloud TPU ホストで直接データ処理を行えるようになるため、追加の Compute Engine VM は不要になります。

お客様の声

早期アクセスのお客様は昨年の 10 月から Cloud TPU VM を使用しており、複数の研究者やエンジニアのチームが、それ以来集中的な使用テストを行っています。お客さまからは、次のようなコメントが寄せられています。

Alex Barron 氏は Gridspace のリード機械学習エンジニアです。Gridspace は、音声通話をリアルタイムで 100% モニタリング、分析、自動化できる、設定不要のプロダクトを供給しています。同社のソフトウェアは、USAA、Bloomberg、Square をはじめとする大手企業の音声案内に使用されています。

「Gridspace では、膨大な音声と言語のモデルをトレーニングするため、JAX と Cloud TPU VM を使用していました。これらのモデルにより、当社で最大のお客様お問い合わせセンターで、高度な分析と自動化の機能を使用できます。以前の Cloud TPU 製品と比べて、同じサイズの TPU でのトレーニングの実行が 2 倍も高速になり、コードの変更なしに 32 のホストを持つ v3-256 にスケーリングできました。当社は Cloud TPU VM の性能と使いやすさに非常に満足しており、今後も使い続ける予定です。」 - Alex Barron 氏

James Townsend 氏はロンドンの UCL Queen Square Institute of Neurology の研究者です。氏のチームは、医療イメージングにディープラーニングを応用するため、Cloud TPU VM で JAX を使用してきました。

「Google Cloud TPU VM により、実装の複雑性を最小限に抑えながら、研究の大幅なスケールアップが可能になりました。単一の TPU デバイスでモデルの実装とデバッグを行う状態から、マルチデバイスとマルチホスト（ポッドスケール）でのトレーニングに、最小限の手間で移行できます。この規模でこの使いやすさはなかなかありません。研究の可能性が大きく広がりました。この作業により、大きな発展を期待しています。」 - James Townsend 氏

Patrick von Platen 氏は Hugging Face の研究エンジニアです。Hugging Face は自然言語処理（NLP）テクノロジーをオープンソースで提供しており、広く使われている Transformers ライブラリの創立者です。研究者やエンジニアは Hugging Face を使用して、わずか数行のコードで最先端の NLP モデルを活用できます。

「Hugging Face は最近、自社の Transformers ライブラリに、TensorFlow や PyTorch とともに JAX を統合しました。これにより、NLP コミュニティは BERT などの一般的な NLP モデルを Cloud TPU VM 上で効率的にトレーニングできるようになりました。単一の v3-8 により、最大 2,048 のバッチサイズを使用して、基本サイズの BERT モデルを 1 日以内で事前トレーニングできるようになりました。Hugging Face は、Cloud TPU VM へのアクセスを簡単にすることで、小規模な新設企業や教育機関も含めて、より広範な NLP コミュニティで、大規模な言語モデルの事前トレーニングが可能になると確信しています。」 - Patrick von Platen 氏

Ben Wang 氏は個人研究者で、言語とマルチモーダルアプリケーション用の Transformer ベースのモデルに取り組んでいます。氏は、Cloud TPU VM 上で大規模な Transformer をトレーニング、およびRay を使用して複数の Cloud TPU VM 上でのトレーニングをオーケストレーションするためのオープンソースコードを公開しました。

「Cloud TPU VM 上の JAX により、TPU に高パフォーマンスで直接アクセス可能になり、Ray を使用するプリエンプティブ TPU Pod スライス間のパイプライン並列トレーニングなど、自由なトレーニング設定を柔軟に構築できるようになりました。」 - Ben Wang 氏

Keno Fischer 氏は Julia プログラミング言語の中核デベロッパーで、Julia Computing の共同創設者でもあり、科学モデルの作成とシミュレーションに機械学習を応用するためのチームを率いています。氏は Julia のオリジナルの TPU バックエンドも含めて、Julia コンパイラの多くの部分を作成しました。

「新しい TPU VM プロダクトは、クラウド上での TPU のユーザビリティという点において飛躍的に進歩しています。TPU ハードウェアの利点を直接的に活用できるため、中間にあるネットワーク接続の帯域幅やレイテンシに制限されることはなくなりました。これは、機械学習モデルが多くの場合に、ホストマシンで実行される科学的シミュレーションと直接結合されている、当社の作業環境では非常に重要です。」 - Keno Fischer 氏

Julia チームは、新しい libtpu 共有ライブラリを使用して、第 2 世代の Cloud TPU 統合に取り組んでいます。最新情報を受け取るには、こちらからご登録ください。

Shrestha Basu Mallick は Sandbox@Alphabet チームのプロダクトマネージャーです。このチームは量子コンピュータの従来型シミュレーションと、大規模な量子化学計算の実行のため、TPU を適応することに成功しました。

「Google Cloud TPU VM と、TPU コアを 1 個から 2,048 個までシームレスにスケーリングできる能力により、私たちのチームは最も強力な量子回路の従来型シミュレータの一つを構築できました。このシミュレータは 40 キュビットの波動関数を展開できます。これには 1 兆個の複素振幅の操作が必要です。また、当チームが最大 500,000 個の電子軌道を持つ巨大な分子の量子化学計算を行うためには、TPU のスケーラビリティは重要なものでした。Cloud TPU に大きな期待を寄せています。」 - Shrestha Basu Mallick

料金と可用性

Cloud TPU VM は現在、us-central1 と europe-west4 リージョンでプレビューとして利用できます。単一の Cloud TPU デバイスと、Cloud TPU Pod スライスの両方を使用でき、TPU v2 または TPU v3 アクセラレータハードウェアを選択できます。Cloud TPU VM は、Google のプリエンプティブサービスとして、TPU ホストマシンごとに 1 時間につき $1.35 から利用できます。価格情報について詳しくは、こちらをご覧ください。