Hugging Face、PyTorch / XLA TPUのサポートを強化

Hugging Faceは、Facebook PyTorchチームとGoogle TPUチームの協力により、2019年のPyTorch Developer Conferenceで発表されたPyTorch-TPUプロジェクトを基盤に、PyTorch / XLAを用いたCloud TPU上でのトレーニングをサポートする統合を発表しました。この統合により、PyTorchユーザーはHugging Faceのトレーナーインターフェースをそのまま利用して、Cloud TPU上でモデルを実行・スケールさせることが可能になります。PyTorch / XLAはPyTorchに新しいxlaデバイスタイプを追加し、既存のPyTorchインターフェースと互換性を保ちつつ、TPUなどのXLAデバイスでの実行を可能にします。また、TPUコアを複数使用する際の勾配交換を効率化する`xm.optimizer_step()`や、チェックポイント保存時のテンソルデバイス管理を容易にする`xm.save()`といったAPIも提供されます。PyTorch / XLAは、操作をグラフとして遅延実行し、XLAコンパイラによって最適化することで、CPUとTPUのアイドル時間を削減します。これにより、モデルのトレーニング効率が向上し、特にHugging Faceライブラリのような静的な形状を持つモデルでは、コンパイル時間のオーバーヘッドを最小限に抑え、高いパフォーマンスを発揮します。記事では、`bert-large-uncased`モデルをv3-8 Cloud TPUでトレーニングした際のパフォーマンスベンチマークも示されており、BF16精度でグローバルバッチサイズ128の場合、106.4分でトレーニングが完了することが報告されています。