行列の積は何回の計算? AIチップの演算量と演算強度

身近なところでは

AI用のチップの仕様表には、たいてい2種類の数字が並びます。1秒に何回計算できるかと、メモリから1秒に何バイト運べるかです。例えばNVIDIAのH100(2022年発表、SXM版)は、16ビットの行列の計算が1秒に約1,000兆回です。運べる量は1秒に3.35TBです。

では、手元のAIの速さは、どちらの数字で決まるのでしょうか。答えは「計算の中身による」です。この記事では、行列の積の仕事の量と運ぶデータの量を数え、どちらが速さを決めるかを見分けます。

GPUが行列の計算を大勢の計算回路で同時に進めることは、「なぜAIにGPUが使われるの?」で説明しました。この記事は、その先の「数える」ことに絞ります。

しくみ1: ベクトルと行列は「数の表」

数を1列に並べたものをベクトルと呼びます。ある生徒の国語・数学・英語の点(80, 60, 70)は、3個の数のベクトルです。文章のAIでは、1つの語を数千個の数のベクトルで表します。

数を縦横の表に並べたものが行列です。縦にm行、横にk列ある行列を「m×kの行列」と書きます。AIの重み(どの入力をどれだけ重く見るかを表す数)も、行列の形でしまわれています。

行列にベクトルを掛けると、表の各行について「入力と掛けて足す」を1回ずつ行い、行の数だけ答えが出ます。これが行列×ベクトルの積です。1人の生徒の点から、理系向けと文系向けの2つの点を出す計算がこれにあたります。

100人分の点をまとめて計算すると、入力も表(3×100)になり、行列×行列の積になります。同じ重みの表を、100人分に使い回すのが特徴です。チャットのAIが答えを1語ずつ作る計算は、1人分だけなら前者に、多くの人の分をまとめれば後者に近くなります。

しくみ2: 計算の回数を数える(2mnk)

m×kの行列Aと、k×nの行列Bを掛けると、答えCはm×nの行列になります。Aの横の数とBの縦の数は、同じkでないと掛けられません。答えのマス1つは、Aの1行とBの1列から、k組の数を掛けて足して作ります。

掛け算1回と足し算1回の組を、積和演算と呼びます。英語の頭文字でMACとも書きます。答えのマスはm×n個あり、1マスにk回の積和が要るので、全部でm×n×k回です。

計算機の世界では、掛け算と足し算を別々に1回と数えることが多く、その場合は2×m×n×k回になります。この数え方の1回をFLOP(フロップ。小数の計算1回)と呼びます。NVIDIAの性能の手引きも、行列の積をこの2mnkで数えています。

小さな例で確かめます。2×3と3×2の行列の積は、答えが4マスで、1マスに3回の積和なので、12回の積和、つまり24 FLOPです。Pythonで3重のくり返し(ループ)を書いて数えると、掛け算は12回、足し算は8回でした。足し算が少ないのは、各マスの最初の積は、何かに足す必要がないからです。

この差は答えのマスの数(m×n)だけです。kが数千ある実際の計算では1,000分の1以下なので、2mnkで数えるのが標準になっています。縦横4,096の正方形の行列どうしなら、2×4,096×4,096×4,096で約1,370億FLOPです。

FLOPを1秒に何回できるかがFLOPS(フロップス)で、チップの速さを表します。兆を表す「テラ」をつけて、TFLOPSと書きます。整数の計算では「演算」と数え、1秒に何兆回かをTOPSと書きます。

しくみ3: 運ぶデータの量を数える

計算回路が働くには、AとBの数をメモリから読み、答えCをメモリに書く必要があります。読み書きする数は、m×k+k×n+m×n個です。これに、数1個のバイト数を掛けると、運ぶデータの量になります。

数1個の大きさは、数の形式で決まります。AIでよく使う16ビットの小数(FP16)は2バイト、8ビットの整数(INT8)は1バイトです。形式の違いは「FP32・FP16・BF16・INT8の違い」にまとめています。

縦横4,096の例をFP16で数えると、3つの行列で約5,000万個の数、約1億バイト(約100MB)です。ここでは、どの数も一度だけ運べば済む、理想の場合を数えています。後で見るように、実際に運ぶ量はプログラムの書き方で何倍にもなります。

しくみ4: 演算強度は「1バイトで何回計算するか」

計算の回数を、運ぶバイト数で割った値を演算強度と呼びます。運んできた1バイトで何回計算できるかを表し、単位はFLOP/バイトです。

行列×ベクトルの積では、演算強度は約1しかありません。FP16の重み1個(2バイト)は、1回の積和(2 FLOP)に使うだけで、もう出番がないからです。NVIDIAの手引きも、行列×ベクトルの積は演算強度が1未満で、常にメモリの速さで決まると説明しています。

行列×行列では事情が変わります。縦横4,096の例では、約1,370億FLOPを約1億バイトで割って、演算強度は約1,365です。運んだ数の1個1個が、答えの1行や1列の分だけ、何千回も使われるからです。演算強度は「運んだ数を何回使い回すか」の目安だと考えられます。

しくみ5: ルーフライン: 2つの上限の低いほう

チップの速さには、2つの上限があります。1つは、計算回路が1秒にできる計算の回数です。もう1つは、1秒に運べるバイト数(メモリ帯域)に、演算強度を掛けたものです。数が届かなければ計算できないので、実際の速さは2つの低いほうを超えられません。

この関係を、横軸に演算強度、縦軸に計算の速さをとって描いた図を、ルーフライン(屋根の線)と呼びます。2009年にWilliamsらが、性能を見積もる図として提案しました。左の斜めの線が運ぶ速さの上限、右の平らな線が計算の上限です。

2本の線が交わる点の演算強度は、計算の上限をメモリ帯域で割ると求まります。H100では、約989兆÷3.35兆で約300です。演算強度がこれより小さい仕事は、計算回路がメモリを待つ時間が長くなります。この状態をメモリ律速と呼び、反対に計算回路が休まず働く状態を演算律速と呼びます。

1人分ずつの行列×ベクトル(演算強度 約1)は、交点よりずっと左にあり、H100の計算の力の約0.3%しか使えません。大きな行列の積(約1,365)は右側にあり、計算の上限で走れます。

しくみ6: バッチ: 同じ重みを何人分にも使う

演算強度を上げるいちばん素直な方法は、入力をまとめることです。まとめて処理する入力の数をバッチサイズと呼びます。重みの表を一度読めば、まとめた全員分の計算に使えます。

縦横4,096の重みの表(FP16で約3,400万バイト)なら、バッチをBにすると、Bが小さいうちは演算強度もおよそBです。計算の回数はB倍になるのに、運ぶ量はほぼ重みの分のままだからです。H100の交点の約300に届くには、数百人分をまとめる必要があります(後の「数字で見る」の表)。

しくみ7: タイリング: チップの中の小さなメモリで使い回す

ここまでは、どの数も一度だけ運ぶと仮定してきました。しかし、計算回路のそばにある速いメモリ(オンチップメモリ)は小さめです。HBMなどの外のメモリよりずっと小さく、大きな行列は丸ごとは入りません。

何も工夫しないと、答えのマスを1つ作るたびに、Aの1行とBの1列を外のメモリから読み直すことになります。すると、運ぶ量は積和の回数と同じくらいに増え、演算強度は行列×ベクトル並みの1前後まで下がります。

そこで、行列を小さな正方形の区画に切り分けます。Aの区画とBの区画をチップの中に読み込み、その組でできる計算をすべて済ませてから、次の区画に進みます。この方法をタイリングと呼びます。区画の1辺をb個にすると、外のメモリから運ぶ量は、工夫しない場合のおよそb分の1に減ります。

NVIDIAの手引きによると、GPUは答えの行列を区画に分けて計算しています。テンソルコアのような専用回路も、小さな区画どうしの積を、1つの命令でまとめて計算します。

しくみ8: 画像の畳み込みも行列の積にできる

画像のAI(CNN)でよく使う畳み込みは、小さな窓(例えば3×3画素)をずらしながら、窓の中の数に重みを掛けて足す計算です。窓の中身を1行に並べ直し、窓の位置の数だけ縦に積むと、大きな行列ができます。これに重みの行列を掛ければ、畳み込みが行列の積になります。この並べ直しを、im2col(イムツーコル)と呼びます。

ただし、窓は重なり合うので、並べ直すと同じ数が何度も複製されます。3×3の窓を1画素ずつずらすなら、1つの数が最大9回現れ、その分メモリを多く使います。そこで、並べ直した行列をチップの中だけで少しずつ作る方法も使われます(cuDNNの2014年の論文)。層ごとの積和の回数は、CNNの出力サイズの計算ツールで試せます。

しくみ9: 量子化: 数を短くして運ぶ量を減らす

小数を、少ないビットの整数などに置き換えて、数1個の大きさを減らすことを、AIでは量子化と呼びます。FP16をINT8にすれば運ぶバイト数は半分、32ビットの小数(FP32)からなら4分の1です。

メモリ律速の仕事では、運ぶ量が半分になれば、時間もほぼ半分になります。1人分の行列×ベクトルでは、INT8にすると演算強度が約2に上がります。演算律速の仕事でも得をします。H100の仕様では、INT8の計算の上限は16ビットの2倍です。ただし、数を粗くしすぎると答えがずれるので、精度を確かめながら使います。

しくみ10: シストリックアレイ: 使い回しを形にした回路

積和の回路を碁盤の目のように並べ、数を隣の回路へ順に手渡しながら行列の積を計算する作りを、シストリックアレイと呼びます。1970年代の終わりに、H. T. KungとC. E. Leisersonが提案しました。データはメモリから出て、回路の間を規則正しく流れ、またメモリへ戻ります。Kungは1982年の解説で、この流れを心臓を出入りする血液にたとえています。

GoogleのTPU(初代)は、この方式で行列の積を計算します。重みを各回路に置いたまま、入力の数を左から右へ流し、途中の和を上から下へ流します。メモリから読んだ数1個が、横に並ぶ256個の回路で使われるので、メモリを読む回数が大きく減ります。

初代TPUは、256×256=65,536個の8ビットの積和の回路で、1秒に最大92兆回の演算ができました(2017年の論文)。演算強度を上げる工夫を、回路の並べ方そのものに組み込んだ設計です。

よくある誤解

数字で見る: バッチを変えると何が起きるか

縦横4,096の重みの表(FP16)に、バッチの数だけ入力を掛ける計算を、H100の仕様の値で見積もりました。計算と運ぶことは同時に進むとし、長いほうの時間がかかるとした理想の値です。「計算」「運ぶ」はそれぞれにかかる時間(µsは100万分の1秒)、「稼働」は計算回路が働く時間の割合です。

バッチ1から256までは、かかる時間がほぼ10〜11 µsのままで、まとめた分だけ得をします。512では計算の時間が運ぶ時間を上回り、その先はバッチを2倍にすると時間も2倍です。1人あたりの処理の速さはもう上がらず、1回の待ち時間だけが延びます。

もっと詳しく学ぶ

• 用語の詳しい解説: 演算強度 / Rooflineモデル / シストリックアレイ / 学習後量子化 / 混合精度
• 表で違いを見る: 学習と推論の違い / CPU・GPU・NPUの違い / DDR・LPDDR・GDDR・HBMの違い
• 数字で試す: メモリ帯域の計算 / CNNの出力サイズ・積和の回数
• 図で見る: 図解の一覧(ルーフラインモデル、GPUのスレッドの階層)
• 土台から読む: 「なぜAIにGPUが使われるの?」 / 「CPUとは?」