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

MicrosoftはHot Chips 29において、「Project BrainWave」について発表を行った。Project BrainWaveはFPGA利用するスケーラブルなDNN(Deep Neural Network)サービス用のプラットフォームを作るというプロジェクトである。

Microsoftは昔からFPGAのパワーに着目し、「Catapult」というプロジェクトでFPGAを搭載したアクセラレータを開発し、Bingサーチの高度化や翻訳の精度向上などに利用してきた。

FPGAを使うディープラーニングは、バッチサイズが小さくても性能が高く、レーテンシも短い。そして、1つのDNNサービスを多数のFPGAで処理する場合のスケーラビリティが高い。

さらに、FPGAの柔軟性は、急速に進歩するマシンラーニングに適応していくのに最適である。推論に必要な精度の演算器が容易に作れるし、疎な行列を圧縮する処理もやりやすい。

Project BrainWaveは、CNTK(Microsoft Cognitive Tool Kit)やCaffeなどのフレームワークで記述された学習前のDNNモデルを多数のFPGAでアクセラレートされたサーバに乗せて学習や推論をスケーラブルに動作させるというプロジェクトである。BrainWaveでは、FPGAを使ってソフトDPU(Deep Learning Processing Unit)を作っている。

BrainWaveのスタックは、フレームワークの下にコンパイラとランタイムの層があり、その下に、低精度で計算を行うDNN命令セットや流動的なAIアルゴリズムの変更に追随できるアーキテクチャ層、さらにソフトDPUを実現するマイクロアーキテクチャ層がある。

そして、巨大モデルで必要となる大量のパラメータを多数のFPGAに分割して記憶し、多数のFPGAによる大規模モデルのスケーラブルな実現を可能にするPersistency at Scale層がある。そして、その下にFPGAを搭載したサーバによるマイクロサービス層があるというスタックになっている。

次の図はサーバの構成を示すものであるが、FPGAはネットワークに直結されており、CPUはNIC(Network Interface Card)、FPGAを経由してネットワークにつながる構造となっている。右上の図がFPGAを搭載したCatapult V2ボードで、メザニンカードに実装されている。そして、右下の図のようにNICカードと一緒にサーバの後方に搭載されている。

次の図はFPGAがどのように使われているかを示す図で、FPGA同士はルータで接続され、CPUとは別の処理プレーンを構成する。前の図のように、FPGAがデータセンターネットワークに直接つながる構造になっているので、このような構造が採れる。

このため、ワークロードに対するCPUの割り当てとFPGAの割り当ては独立に行うこともでき、例えば次の図のように、FPGA群は一部がDNN処理、一部はSQL処理、一部はWebサーチのランキング、一部はSDN(Software Defined Network)の処理のオフロードというような使い方ができる。

なお、この図ではNICが一方のCPU、FPGAは他方のCPUに接続されており、前の図と接続が違っている。物理的に考えて、多分、こちらの図の方が正しいと思われる。

Caffeモデル、CNTKモデル、Tensorflowモデルなどはグラフ分割と最適化が行われ、それぞれのコンパイラで運用パッケージが作られる。その過程で、右の図のように、グラフが大きい場合は、複数のFPGAにグラフを分割する。

通常は、モデルが大きくFPGAのBRAM(Brock RAM)の中にはすべてのモデルパラメータ(重み)を格納することはできないので、FPGAに接続したDRAMに格納するという方法が採られる。しかし、この場合、パラメータの読み出しが計算量に比較して大きい処理では、DRAMメモリの読み出しが性能ネックになってしまう。

この問題を軽減するため、Googleの2代目のTPUでは、モデルパラメータを格納するメモリを通常のDDR4からHBM2に変更してメモリバンド幅を拡大して性能を上げている。