インフラ

TPUとは

Tensor Processing Unit

Googleが開発した、ニューラルネットワークの行列演算に特化した専用プロセッサ

ひとことで言うと

Googleが作った、AI計算専用のプロセッサ。

概要

TPU(Tensor Processing Unit)とは、Googleが機械学習のワークロード、特にニューラルネットワークで多用される行列積演算(テンソル演算)の高速処理のために独自開発した専用プロセッサ(ASIC)。 GPUが元々グラフィックス処理向けの汎用的な並列計算チップを機械学習に転用したものであるのに対し、TPUは設計段階から行列演算に特化しており、多数の積和演算器を格子状に配置したシストリックアレイと呼ばれる構造により、大規模な行列積を高い電力効率で実行できる。 Google CloudのAPIやサービスを通じて外部にも提供されており、Google自身の大規模モデルの学習・推論基盤としても利用されている。

背景

機械学習、特に深層学習の計算の大部分は行列積演算が占める。 GPUはこうした並列計算に適した汎用アーキテクチャを持つが、機械学習に用途を絞ることでさらなる効率化が可能だと考えられ、Googleは自社のワークロード向けに特化した専用チップとしてTPUを開発した。

歴史

2016年: GoogleがTPU(第1世代)を発表し、自社のデータセンターで機械学習の推論処理に利用していることを明らかにする。 2017年: 学習にも対応した第2世代TPUを発表。 2018年以降: Google CloudのTPUとして外部の開発者にも提供が拡大し、世代を重ねるごとに性能・スケールが向上。

アーキテクチャ

TPUの中核は、積和演算(乗算と加算)を行う演算器を格子状に多数配置し、データを演算器間で直接受け渡しながら処理するシストリックアレイと呼ばれる構造。汎用性を抑え行列積へ特化した設計とすることで、同等の演算性能をより少ない電力・チップ面積で実現している。複数のTPUチップを高速な専用ネットワークで接続した「Pod」と呼ばれる構成により、大規模モデルの分散学習にも対応する。

ワークフロー

モデルの計算グラフをXLA等のコンパイラでTPU向けに最適化 → 行列演算をシストリックアレイへ展開 → 大規模モデルでは複数チップからなるPodへ処理を分散 → 結果をホスト側へ返す。

コード例

JAXでTPUデバイスを確認し行列演算を実行する

import jax
import jax.numpy as jnp

print(jax.devices())  # TPUデバイスの一覧を確認


@jax.jit
def matmul(a, b):
    return jnp.dot(a, b)


result = matmul(jnp.ones((1024, 1024)), jnp.ones((1024, 1024)))

利点

行列積演算に特化した設計により、機械学習ワークロードでの電力効率・処理性能に優れる
複数チップを接続したPod構成により、大規模な分散学習にスケールしやすい
Google製の機械学習フレームワーク(JAX等)との親和性が高く、大規模学習での実績が豊富

欠点

機械学習以外の汎用的な計算には向かず、用途が限定される
Google Cloud以外の環境では利用しにくく、GPUと比べてエコシステムやツールの選択肢が限られる
PyTorchを中心とするエコシステムと比べ、対応ツール・ドキュメントの層がGPUよりやや薄い

比較

GPU — GPUは元々グラフィックス処理向けの汎用チップを機械学習に転用したものであるのに対し、TPUは機械学習向けに専用設計されたチップ
CUDA — CUDAはNVIDIA GPU向けの並列計算プラットフォームであり、TPUはGoogleが提供する別系統の専用ハードウェア
NPU — NPUがエッジデバイス向けの小型な推論用プロセッサであるのに対し、TPUは主にデータセンター向けの学習・推論用プロセッサ

よくある質問

TPUは個人でも使える？

Google CloudのサービスとしてTPUをレンタルする形で利用でき、Google Colab等でも一部無料枠が提供されている。

TPUとGPU、どちらを選ぶべき？

Google Cloud中心の環境で大規模なJAX/TensorFlowベースの学習をするならTPU、より汎用的なエコシステムやマルチクラウド対応を重視するならGPUが選ばれやすい。

参考文献

DocumentationGoogle Cloud: Cloud TPU
Research PaperIn-Datacenter Performance Analysis of a Tensor Processing Unit