Hot Chips 29 - FPGAを活用するMicrosoftの「Project BrainWave」

しかし、MicrosoftはGoogleとは違って、モデルパラメータがすべてFPGA内部のBRAMに格納できるサイズにニューラルネットを分割するというアプローチを取った。このようにすれば、処理の開始時にはすべてのモデルパラメータをFPGA内部のメモリに転送する必要があるが、その後は、FPGA内部のBRAMにモデルパラメータは残っており(Persistentであるので)BRAMを読めばよい。この構成では、外部メモリがなく、BRAMのバンド幅は非常に大きいのでメモリバンド幅は性能の制約にならない。

この構成であれば、1つのリクエストが入力されたときには、FPGAチップのすべての計算リソースをその処理につぎ込めばよく、応答が速いというメリットがあるという。

ただし、BrainWaveでは、1つのFPGAにモデルパラメータが収まりきらない場合は、複数のFPGAにニューラルネットを分割してFPGAのBRAMに格納できるようにしている。

次の図は大きなLSTM(Long Short-Term Memory:時系列データに対するニューラルネット)を分割して、パイプラインで処理する場合の図であるが、それぞれのFPGAに実装されたLSTMのステージが、データセンターネットワークを使って右隣のFPGAにデータを送る。

FPGAはCPUの介在なしに直接ネットワークを使ってデータを送るので、ホップ当たりの遅延は2μs以下とのことである。

BrainWaveのソフトDPUでは、RTLを書く必要はなく、シングルスレッドのCプログラムを書けばよい。これで行列積や非線形アクティべーションなどの専用命令を含むISAが作られる。

そして、ニューロンの計算はパラメータ化された低精度で計算が行われる。マイクロアーキテクチャもパラメータ化ができ、最大1M ALM(Adaptive Logic Module)のFPGAまでスケールができる。

次の図は、FPGAで作られるソフトDPUのブロック図であるが、命令を順に解釈して処理を実行するというれっきとしたプロセサで、基本的なアーキテクチャをベースとして、パラメータに応じて、演算器の個数や演算精度などをカストマイズしてDPUコアを生成しているのではないかと思われる。

マトリクスとベクタの積の計算は、入力ベクタをまとめてバッチで計算する方が演算性能は上がるが、Webサービスの場合、問い合わせをまとめてバッチにすると応答のレーテンシが長くなってしまう。このため、Microsoftは、マトリクス×1つのベクタの計算の速度を速くするようマトリクス、ベクタユニットの設計を最適化している。

マトリクス、ベクタユニットの入出力は16ビット長データのFP16形式であるが、形式変換を行って、ユニット内部の計算はMicrosoft独自のms-fp8とかms-fp9といった低精度で計算を行っている。

次の図は低精度の計算の効果を示すもので、左は225MHz動作のStratix Vと500MHz動作のStratix10 FPGAを使った場合にTops値(Tera Ops)を精度ごとにプロットしたもので、Stratix 10の場合、INT16では12Topsであるが、ms-fp9では65Tops、ms-fp8では90Topsと大幅に計算性能が上がっている。

右の図は3つのモデルをFP32、ms-fp9とms-fp9で再学習をさせた場合の推論の精度を比較した棒グラフである。FP32を1.0とすると、ms-fp9での学習の場合は、1%程度精度が下がっているが、再学習を行えばFP32と同じか若干上回る精度になっている。

つまり、ms-fp9を使えば、精度はFP32並みかそれ以上で、INT16と比較しても5倍あまりの計算性能が得られるというわけである。

結論であるが、FPGAを使うMicrosoft BrainWaveは、適応可能な低精度計算を行うことができ、低バッチサイズの処理に優れた性能を発揮する。そして、FPGAは再構成が可能であり、将来、精度要求やアルゴリズムが変わっても対応できる。

Microsoftのサーバで動作するBrainWaveは、AIクラウドの強力なプラットフォームであり、より深く、大きなCNNを実行して認識精度を改善したり、より高次のRNNを処理して自然言語の理解を高めたり、音声認識の精度を改善したりして、クラウドAIで何ができるかの限界を押し広げる。

なお、BrainWaveは、現状は、社外には公開されていないが、Stay tunedと書かれており、何らかの形で社外でも使えるようにすることを考えていると思われる。