Lightmatterの光AI推論アクセラレータ「Mars」を読み解く - Hot Chips 32

Marsの電力はどこで消費されるのか？

重みはMZIに近い場所に電荷として蓄えられており、記憶のための消費電力はほぼゼロである。また、64×64のフォトニック乗算器の消費電力もほぼゼロである。

重みのアップデートのためのダイナミックな電力消費は必要であるが、この電力は電子的な回路の場合と同程度の消費電力である。電子回路と同様に、入力を纏めて処理するバッチ化を行えば重みのアップデートの電力は少なくて済むようになる。

したがって、Marsの場合の電力消費の大部分はデータ変換に伴うものであり、Marsは計算のためには電力を使っていない。

Marsでは重みの記憶はMZIの近くに電荷として記憶しており、記憶のための電力はほぼゼロである。また、フォトニックな計算の電力もほぼゼロであり、電力消費の大部分はデータ変換のために使っている

次の図はMarsの電力消費の内訳を示すものである。そして、 ResNet-50で99%の精度を実現したケースで、Marsの全体の消費電力はレーザー光源の電力を含めて3W以下で、その大部分がデータ移動のための電力である。

電力消費の内訳は、48.7％がデジタル部のデータ移動などのダイナミックな消費電力、ベクタのDACが11.3％、重みのDACが10.3％、レーザーが8.8％、ADCが12.7％といったところが大口の電力消費者である。

Marsの電力消費の内訳。デジタル部分のダイナミックな動作に全体の48.7％の電力が使われている。また、DACとADCで、合計34.4％の電力を使っている

そして、Marsチップは、次の写真のように、90nmプロセスで作られたフォトニックチップの上に、14nmプロセスのASICチップを搭載した構造になっている。フォトニックチップの電力が小さいので、チップを重ねても放熱の問題は生じない。そして、チップのスタッキングでデータ移動の距離が減り、エネルギーも減るという効果が出ている。

Marsはオプティカルチップの上にデジタルのSoCを積み重ねて搭載している。このスタッキングでデータの移動距離を縮めて、さらに消費電力を減らしている

MarsアクセラレータはPCI Expressカードの形態で作られており、ホストCPUに接続される。マシンラーニングのフレームワークやコンパイラはホストCPUで実行される。

フレームワークとしては、PyTorch、TensorFlowとONNXをサポートしている。Lightmatterのソフトウェア環境で、モデルのシミュレーションを行い、コンパイルや実行を行うことができる。そして、実行中のモデルの内部状態をモニタできるデバッガなども提供している。また、実行中のプロファイルを取り、チューニングを行えるツールも提供している。