はじめに

こんにちは。LeapMindで楽しく働いている前田です。LeapMindで開発中のEfficiera（えふぃしえら）という製品で使用しているメモリ割り当てのアルゴリズムを紹介します。このアルゴリズムの開発は前田が入社して最初にチャレンジしたタスクですが、従来アルゴリズムからの改善効果が大きく、また理論的にもエレガントなものにできました。

Efficieraとは

Efficieraは、LeapMindが開発しているディープラーニング推論向けアクセラレータIPです。極小量子化技術を利用することで、小さい回路規模・省電力・省メモリで高速な推論を実現します。

Efficieraの極小量子化は、2bitアクティベーション / 1bit重みを使います。これは、すでに実用化されている量子化推論で使われる8bitアクティベーション / 8bit重みと比較してさらに小さい表現です。このような極小ビット計算では、Intel OpenVINOやNVIDIA TensorRTなどの既存のフレームワークが提供する学習済みモデルを量子化モデルに変換する方法では精度が低下してしまいます。そのためEfficieraではモデル作成・学習から一貫して量子化ネットワークを扱うためのフレームワーク (Efficiera NDK / Compilation Tool) も開発・提供しています。

Efficieraとランタイム

Efficieraを動かすには、「ランタイム」と呼ばれるソフトウェアが必要です。これはEfficieraに接続したホストCPUで動作するライブラリで、Efficieraに演算内容を指示したり、Efficieraがサポートしない計算や、カメラデバイスなどのI/Oとの仲介を担当します(図1)。

runtime 図1: Efficieraを用いた推論におけるランタイムの役割

ランタイムにおけるメモリの使い方

ランタイムはさまざまな用途にメモリを使います。このうち最も重要なのが、計算したテンソルを保存しておく一時バッファです。DNNは多数の畳み込み演算 (Conv) がチェインする構造を持つので、これらのConvの出力を置いておく一時バッファも数多く必要です。ほとんどのConv出力は直後の演算でのみ使われるので、対応する一時バッファは推論処理全体のうち特定の期間だけ生存していれば十分です。この期間のことを以降では一時バッファの生存期間と呼びます(図2)。

lifetime 図2: ネットワークのグラフの一部(上段)とそれに対応する一時バッファの生存期間(下段)。グラフの各ノードの出力は、その全ての使用が終わるまで生存している必要がある。緑のノードの出力バッファは、黄色と紫の２つのノードが使用するので、他の出力バッファよりも生存期間が長い。

ランタイムは一時バッファに必要なメモリ領域（アリーナ）を初期化時に一括で確保します。アリーナを用いる理由はシステムの信頼性向上です。たとえば、推論中のメモリ確保が失敗したり、失敗せずとも推論の実行時間がばらつくことを避ける意味があります。

それぞれの一時バッファのアリーナでの配置を静的に決定するのは難しい問題です。直感的には、横軸を生存期間・縦軸をアリーナの領域とする二次元領域への矩形割り当て問題として考えることができます。上手な配置を行ってアリーナを密に配置できれば、推論に必要なメモリの使用量を削減できます(図3)。

alloc 図3: 図2のネットワークの一時バッファの配置例。左の配置では隙間が多く、アリーナのサイズが増大してしまう悪い配置である。右はアリーナのサイズが最小になる最も良い配置。

バッファ割り当て問題

バッファ割り当て問題を適切なアルゴリズムに落とすために定式化を行います。ここで重要なのは、推論の計算を有向非巡回グラフ (directed acyclic graph; DAG) で表現すること・計算の途中で一時的に使われる領域をDAG上の生存期間によって表現することです。

DAGはトポロジカルソートによって順序関係を保ったまま直列化できるので、以降グラフの各ノードは順序付けされているとします。すなわち、グラフが $n$ 個のノードから成る場合、各ノードの計算順序にしたがって、ノード $1, \ldots,$ ノード $n$ と順序付けられます。また、議論の簡単のために、各ノードの出力テンソルはひとつだけに限定し、ノードと出力テンソルを同一視します。複数の出力テンソルをもつノードは、ひとつだけの出力のノードの並置で常に置き換え可能なので、議論の一般性は損ないません。ノード $i(i=1, \ldots, n)$ の出力テンソルをテンソル $i$ と呼び、そのサイズを $\mathrm{size}[i]$ byteとします。

図4: 図2のグラフの各ノードに順番付けしたグラフ。

テンソル $i ( i=1, \ldots, n)$ のアリーナへの割り当てを決定することは、アリーナ上に割り当てる領域の先頭の位置 $\mathrm{offset}[i]$ を決定することと等価です。すなわち、テンソル $i$ をアリーナ上の領域 $[\mathrm{offset}[i], \mathrm{offset}[i] + \mathrm{size}[i])$ に割り当てます。アリーナの大きさ $\mathrm{memory\\\_size}$ は、各テンソルが割り当てられた領域がはみ出さないように取れば良いので $\mathrm{memory\\\_size} = \max_{i=1, \ldots, n} \\\{\mathrm{offset}[i]+\mathrm{size}[i]\\\}$ となります。この値がより小さくなるような $\mathrm{offset}$ を決定するのが目的です。ただし、 $\mathrm{offset}$ を決定する際には各テンソルの生存期間を考慮する必要があります。

テンソルの生存期間

テンソルの $i (i=1, \ldots, n)$ の生存期間を $[\mathrm{begin}[i], \mathrm{end}[i] )$ と書くことにします。テンソル $i$ はノード $i$ の出力として生まれてから、他のノードの入力として使用し尽くされるまでの期間、アリーナ上に存在する必要があります。したがって、テンソル $i$ を入力として取るノードがノード $x_1, \ldots,$ ノード $x_l$ であるときに

\begin{cases} \mathrm{begin}[i] = i \\\\\\ \mathrm{end}[i] = \max \\\{i, x_1, x_2, \ldots, x_l \\\} + 1 \end{cases}

と表せます。

問題の定式化

バッファ割り当ての問題は、生存期間が重複する二つのテンソル i,j(i≠j) をアリーナ上で重ならないように配置する問題として、以下のように定式化できます。

入力: $\mathrm{begin}[i], \mathrm{end}[i], \mathrm{size}[i] (i= 1, \ldots, n)$ 出力: $\mathrm{offset}[i] (i=1, \ldots, n)$ 最小化: $\mathrm{memory\\\_size} = \max_{i=1,\ldots,n} \\\{\mathrm{offset}[i]+\mathrm{size}[i]\\\}$ 条件:

任意の $i, j(i \neq j)$ について、 $[\mathrm{begin}[i], \mathrm{end}[i])$ と $[\mathrm{begin}[j], \mathrm{end}[j])$ が重複するならば $[\mathrm{offset}[i], \mathrm{offset}[i]+\mathrm{size}[i])$ と $[\mathrm{offset}[j], \mathrm{offset}[j]+\mathrm{size}[j])$ は重複しない

アルゴリズム

従来のアルゴリズム

以前は次のようなアルゴリズムを用いていました:

arranged_idnex <- [1, ..., n]を、(end[i]-begin[i], size[i])の降順でソート

#arranged_indexの先頭から割り当てる
height <- [0, ..., 0] (全部0 長さn)
offset <- []
for k in arranged_index:
        max_height <- height[begin[k]: end[k]]の最大値
        offset[k] <- max_height
        height[begin[k]: end[k]] <- offset[k] + size[k]
  
return offset

上記のアルゴリズムは以下の2つのステップから成ります:

ステップ1: 入力を特定の順番で並べる ( $\mathrm{arranged\\\_index}$ ) ここでは、 $(\mathrm{end} - \mathrm{begin} , \mathrm{size})$ の降順でソートしています。
ステップ2: $\mathrm{arranged\\\_index}$ の順番に従って敷き詰めて割り当てる

図5: ステップ2において割り当てを行う様子の一例。アリーナのメモリについての軸(縦軸)を高さに見立てて、できるだけ低いところに割り当てるようにする。

ワーストケースの観察

このアルゴリズムの振る舞いにはよくない性質があります。例として図6に示す分岐のないネットワークを入力してみましょう。簡単のために全てのテンソルが同じサイズであるとします。

図6: 分岐のないネットワークのグラフ例

各テンソルの生存期間は、 $[\mathrm{begin}[i], \mathrm{end}[i] ) = [i, i+2) \ (i=1,\ldots, n)$ となります。 $\mathrm{end}[i] - \mathrm{begin}[i] = 2 \ (i=1,\ldots,n)$ なので、上記のアルゴリズムのステップ1においてソートの基準となる値は全て等しくなります。すなわち、ソートの結果は用いるソートアルゴリズムのタイブレイクの仕様により決まります。例えば、マージソートなどの安定ソートを利用した場合、テンソルの並びは全く変更されません。

仮にノードの並びが変更されなかったとすると、ステップ2において以下のようにテンソルが階段上に配置されてしまい、メモリを効率的に利用できません。図7: テンソルが階段状に配置されてしまう様子

この結果を踏まえて、ソートの基準としてどのようなものが相応しいか検討しましょう。 $\mathrm{end}[i]-\mathrm{begin}[i]$ や $\mathrm{size}[i]$ はしばしば似たような値になるので、ソートの基準として用いるのは好ましくありません。また、 $\mathrm{begin}[i]$ や $\mathrm{end}[i]$ についてソートしても同様に階段上に配置されてしまいます。 $\mathrm{begin}[i], \mathrm{end}[i], \mathrm{size}[i]$ を組み合わせてソートの基準とするのは筋が良くないようです。

パス被覆による改善

特定の基準でソートするという考えから一旦離れて、さきほどのケースにおける最適な配置を考えます。明らかに次のように上下交互に配置するのが最適です。

図8: 最適な配置

このような良い配置を一般的なケースにおいて得るためには、以下のように処理を行うとよさそうです。

ステップ1: 各テンソルを、互いに生存期間が重複しないテンソルのグループに分割する。 $\mathrm{arranged\\\_index}=[\mbox{グループ1のテンソルたち}, \mbox{グループ2のテンソルたち}, \ldots]$ とする。
ステップ2: 旧アルゴリズムから据え置き。

このようなテンソルのグループ分けは無数に考えられますが、その中でグループ数が最小になるようなグループ分けをパス被覆という概念と対応させて考えることができます。グラフ $G$ を以下のように構成します。

ノード $i (i=1, \ldots, n)$ を頂点とする
任意の $i, j (i \neq j)$ について、 $\mathrm{end}[i] \leq \mathrm{begin}[j]$ であれば $\mbox{ノード}i \rightarrow \mbox{ノード}j$ に辺を張る

図9: グラフ $G$

このようにして構成したグラフ $G$ はDAGになります。グラフ $G$ を考える利点は、“互いに生存期間が重複しないのテンソルのグループ”とグラフ $G$ 上の有向パスが一対一に対応するところにあります。したがって、テンソルのグループ分けはグラフ $G$ のパス被覆(互いに交わらない有向パスによってグラフ $G$ の全てのノードを覆うこと) に対応します。特に、グループ数最小のグループ分けはグラフ $G$ の最小パス被覆に対応します

図10: パス被覆とグループ分けの対応

嬉しいことに、DAGに対する最小パス被覆は多項式時間で解くことができます。用いるアルゴリズムにもよりますが、一般に $O(n^3)$ または $O(n^{2.5})$ の時間計算量です。

グループ分けアルゴリズム

今回の問題では、次に示すようなシンプルな貪欲アルゴリズムによって、最小パス被覆問題を陽に解かずにグループ数最小のグループ分けを求めることができます。

テンソル $i (i=1, \ldots, n)$ を順番にグループに追加する
- テンソル $i$ を追加できるグループが存在するなら、任意に一つ選んで追加する
- そのようなグループが存在しないならテンソル $i$ のみからなる新しいグループを作成する

アルゴリズムの正当性の証明はAppendixに記します。

評価

グループ分けアルゴリズムによって得られる割り当て結果のアリーナのメモリ使用量 $\mathrm{memory\\\_size}$ は、以下のように評価できます: $\mathrm{memory\\\_size} \leq (2W) \cdot \mathrm{max\\\_size}$

ここで、 $\mathrm{max\\\_size}、W$ はそれぞれ以下のように定義されます。

$\mathrm{max\\\_size} := \max_{i=1,\ldots, n} \mathrm{size}[i]$
$W$ : 推論ネットワークのグラフに対して定まる量。グラフ上の有向パスの集合であって、グラフの各辺がパス集合内の少なくとも1つのパスに含まれるようなものを考える。通常のパス被覆と違って、パス集合が頂点を覆う必要はない。また、2つ以上のパスに含まれる頂点や辺が存在してもよいこのようなパス集合の中で、集合のサイズが最小のものを考え、そのサイズを $W$ とする。

証明はAppendixに記します。

ニューラルネットワークで用いるグラフでは、 $W$ の値は小さくなる傾向にあります。例えば、resnetではグラフの全ての辺を2本のパスで覆うことができます(図11)。1本のパスで全ての辺を覆うことは不可能なので、 $W=2$ となります。

図11: 赤い矢印で表現された2つのパスによって、resnetのグラフの各辺を覆う様子。

多くのモデルにおいて、 $W$ の値はネットワークの深さに依存せず定数となります。したがって、扱うネットワークの層が深くなってもメモリ使用量は漸近的に増加しません。

性能改善

例として、開発中の100GOP程度のネットワークに適用した場合、従来アルゴリズムに比べてメモリ使用量を40%削減できました。社内開発している複数のモデルにおいても同様の有効性が確認できています。

おわりに

LeapMindでは、競技プログラミングで培ったアルゴリズムの経験を生かして楽しく仕事をしています。ところで、LeapMindではエンジニアを募集しています。競技プログラミングで培ったアルゴリズムの経験を生かして楽しく仕事したい人も、そうでない人も、興味がある方はぜひ一度LeapMindの採用ページから応募してみてください。

Appendix

グループ分けアルゴリズムの正当性の証明

アルゴリズムの正当性を示します。明らかにアルゴリズムが出力するグループ分けは正しいグループ分けです(各グループに属するテンソルの生存期間は互いに重複しない)。したがって、グループ数が最小であることを示せば十分です。これは次のように証明できます。

\mathrm{overlap\\\_max} :=\displaystyle \max_{1 \leq x \leq n} \\\# \\\{i \ (1 \leq i \leq n) | \mathrm{begin}[i] \leq x < \mathrm{end}[i] \\\}

と定義すると、明らかに $\mathrm{overlap\\\_max}$ はグループ数に対する下界となっています。上記のアルゴリズムが出力するグループ分けのグループ数が、 $\mathrm{overlap\\\_max}$ ちょうどであることを示せば十分です。 $\mathrm{overlap\\\_max}+1$ 個以上のグループからなるグループ分けが出力されたと仮定して矛盾を導きましょう。

テンソル $k$ を追加する際に $\mathrm{overlap\\\_max}+1$ 個目のグループが作成されたとします。仮定より、既に作成された $\mathrm{overlap\\\_max}$ 個のグループグループ $1, …,$ グループ $\mathrm{overlap\\\_max}$ それぞれについてテンソル $k$ を追加することができなかったということになります。すなわち,以下の条件が成り立つような $\mathrm{overlap\\\_max}$ 個のテンソル、テンソル $g_1, \ldots,$ テンソル $g_{\mathrm{overlap\\\_max}}$ が存在します:

テンソル $g_i$ はグループiに属する( $i=1, …, \mathrm{overlap\\\_max}$ )
テンソル $g_i$ とテンソル $k$ の生存期間は重複している

今、テンソルは $i = \mathrm{begin}[i]$ の小さい方から追加しているので、 $\mathrm{begin}[g_i] \leq \mathrm{begin}[k] (i=1, \ldots, \mathrm{overlap\\\_max})$ が成り立ちます。テンソル $g_i (i =1, \ldots, \mathrm{overlap\\\_max})$ とテンソル $k$ の生存期間が重複しているという仮定と合わせると $\mathrm{begin}[k] < \mathrm{end}[g_i] ( i=1, \ldots, \mathrm{overlap\\\_max})$ が従います。ところが、

\begin{align*} \mathrm{overlap\\\_max}&=\displaystyle \max_{1 \leq x \leq n} \\\# \\\{i \ (1 \leq i \leq n) | \mathrm{begin}[i] \leq x < \mathrm{end}[i] \\\} \\\\\\ &\geq \\\# \\\{i \ (1 \leq i \leq n) | \mathrm{begin}[i] \leq \mathrm{begin}[k] < \mathrm{end}[i] \\\} \\\\\\ &\geq \\\# \\\{k, g_1,\ldots,g_{\mathrm{overlap\\\_max}} \\\} \\\\\\ &=\mathrm{overlap\\\_max}+1 \end{align*}

となるので矛盾します。(証明終)

$\mathrm{memory\\\_size} \leq (2W) \cdot \mathrm{max\\\_size}$ の証明

$\mathrm{memory\\\_size} \leq \mathrm{overlap\\\_max} \cdot \mathrm{max\\\_size}$ が成り立つので、 $\mathrm{overlap\\\_max} \leq 2W$ を示せば十分です。辺を覆うパス被覆を構成する $W$ 本のパスを考えます。パス $p (p=1,\ldots, W)$ について、

パス}p\mbox{に含まれる}$$ となるように$\mathrm{path\\\_group}[p]$を定義します。まずは次の補題を準備します。 ### 補題1 任意の$p(1 \leq p \leq W)$と任意の$i, j (i < j) \in \mathrm{path\\\_group}[p]$について、以下の条件が成り立ちます: $$\mathrm{end}[i] \leq \mathrm{begin}[j] + 1$$ ### 補題1の証明 仮定より、$\mbox{ノード}i, \mbox{ノード}(\mathrm{end}[i]-1)$はパス$p$に含まれます。すなわち、ノード$(\mathrm{end}[i]-1)$はパス$p$から誘導されるノード$i \rightarrow $ノード$j$ パス上に含まれます。したがって、 $$\mathrm{end}[i]-1 \leq \mathrm{begin}[j]$$ が成り立ちます。(補題1の証明終) 次に、以下の補題を準備します。 ### 補題2 任意の$p(1 \leq p \leq W)$について、以下の条件が成り立ちます: $$\displaystyle \max_{1 \leq x \leq n} \\\# \\\{i \in \mathrm{path\\\_group}[p] \mid \mathrm{begin}[i] \leq x < \mathrm{end}[i] \\\} \leq 2$$ ### 補題2の証明 任意の$i,j,k(i < j < k) \in \mathrm{path\\\_group}[p]$について、 $$\mathrm{end}[i] \leq \mathrm{begin}[k]$$ を示せば十分です。$i, j$と$j, k$に対して補題1を適用すると、以下の2つの不等式が得られます。 - $\mathrm{end}[i] \leq \mathrm{begin}[j]+1$ - $\mathrm{end}[j] \leq \mathrm{begin}[k]+1$ また定義より $$\mathrm{end}[j]-1 \geq j+1 = \mathrm{begin}[j]+1$$ が成り立ちます。したがって、 $$\mathrm{end}[i] \leq \mathrm{begin}[j]+1 \leq \mathrm{end}[j]-1 \leq \mathrm{begin}[k]$$ が成り立ちます。(補題2の証明終) 議論を単純にするため、グラフには孤立点が含まれないと仮定します。孤立点が含まれるグラフは計算グラフとしては不適なので、議論の一般性を損ないません。$\mathrm{path\\\_group}[p](p=1,\ldots, W)$の定義より、各ノードは一つ以上の$\mathrm{path\\\_group}$に含まれます。したがって、

\begin{align*} \mathrm{overlap\_max}&=\displaystyle \max_{1 \leq x \leq n} \# \{i \ (1 \leq i \leq n) | \mathrm{begin}[i] \leq x < \mathrm{end}[i] \} \\\ &\leq \displaystyle \max_{1 \leq x \leq n} \sum_{1 \leq p \leq W} \# \{i \in \mathrm{path\_group}[p] | \mathrm{begin}[i] \leq x < \mathrm{end}[i] \} \\\ &\leq \displaystyle \sum_{1 \leq p \leq W} \max_{1 \leq x \leq n} \# \{i \in \mathrm{path\_group}[p] | \mathrm{begin}[i] \leq x < \mathrm{end}[i] \} \\\ &\leq 2W \end{align*}

が成り立ちます。(証明終)

Tech Blog