PyTorch / XLA が Cloud TPU で一般提供開始

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

PyTorch 機械学習（ML）フレームワークは、その柔軟性と使いやすさで ML コミュニティで人気があり、これを Google Cloud でサポートできることをうれしく思います。このたび、Cloud TPU の PyTorch / XLA サポートが一般提供されるようなったことをお知らせします。具体的には、PyTorch ユーザーは、十分にサポートされ安定した PyTorch 統合を使用して、大規模で低コストの Cloud TPU ハードウェア アクセラレータにアクセスできます。

PyTorch / XLA は、PyTorch の直感的な API と XLA 線形代数コンパイラの長所を組み合わせたもので、CPU、GPU、Cloud TPU Pod を含む Cloud TPU をターゲットにすることができます。PyTorch / XLA は、最小限の変更でほとんどの標準的な PyTorch プログラムを実行し、CPU にフォールバックして TPU で現時点では未対応の操作を実行します。PyTorch / XLA が生成する詳細なレポートを活用して、PyTorch ユーザーはボトルネックを発見し、Cloud TPU でより効率的に実行するようにプログラムを適応させることができます。

マサチューセッツ工科大学（MIT）の博士号候補者である Jonathan Frankle 氏は、次のように述べています。「PyTorch / XLA により、PyTorch ワークフローにほとんど変更を加えることなく Cloud TPU で数千ものテストを実施できました。PyTorch / XLA は、PyTorch の使いやすさ、TPU のスピードと費用対効果という、それぞれの長所を兼ね備えています。」同氏は、ICLR 2019 で Best Paper Award を受賞した画期的な論文「宝くじ仮説: 訓練できるまばらなニューラル ネットワークを見つけること」に関連する最新の研究を PyTorch / XLA を使ってスケールアップしています。

アレン人工知能研究所（AI2）では最近、複数のプロジェクトにわたって Cloud TPU で PyTorch / XLA を使用しています。AI2 の研究者である Matthew Peters 氏は現在、PyTorch / XLA を使用して、言語理解能力を向上させるために最先端の言語モデルに視覚的なコンポーネントを追加する方法を研究しています。同氏は次のように述べています。「PyTorch / XLA はまだ新しいテクノロジーですが、すでに PyTorch に投資して機械学習モデルをトレーニングしている組織に前途多望な新しいプラットフォームを提供します。」

PyTorch / XLA のご利用開始にあたって、Google Cloud は、広く使用されているディープ ラーニング モデルと関連チュートリアルで構成され、ますます増え続けるオープンソースの実装をサポートしています。ここでは、ResNet-50、Fairseq Transformer、Fairseq RoBERTa、そして、DLRM のチュートリアルをご紹介します。また、ML モデルの継続的なテストを容易にするオープンソース ツールを開発しており、PyTorch Lightning チームや Hugging Face チームがこのフレームワークを使用して Cloud TPU で独自のテストを実施できるように支援してきました（こちらは PyTorch Lightning チームの関連するブログ投稿です）。

上記リンク先のチュートリアルを確認して、Colab を介してブラウザで PyTorch / XLA を試し、PyTorch / XLA の GitHub リポジトリに問題や pull リクエストを投稿してください。また、PyTorch 1.6 とともに PyTorch / XLA がプリインストールされた新しいディープ ラーニング VM（DLVM）イメージもリリースしました。ここでは、この新しい DLVM イメージをすぐに使い始める手順について説明します。

サンプルコードを含む PyTorch / XLA の技術情報については、PyTorch Medium の公式サイトに掲載されている関連記事をぜひお読みください。

-Cloud AI プロダクト管理ディレクター Craig Wiley