機械学習のための機械学習技術の進展

 (blog.research.google)
1 ポイント 投稿者 GN⁺ 2023-12-17 | 1件のコメント | WhatsAppで共有
  • MLモデルが大規模になるほど、実行効率はハードウェアだけでなく コンパイラ最適化 に大きく左右され、Google ResearchとGoogle DeepMindはこの領域自体をMLで改善するアプローチを公開した
  • TensorFlow、JAX、PyTorchのようなフレームワークの下で下される コンパイラの決定 は、同じモデルであっても実行時間やリソース使用量を大きく変えうる
  • TpuGraphs は、TPU向けMLプログラムの計算グラフ、コンパイル設定、実行時間を収めた、学習ベースのコストモデル研究に利用できるデータセットである
  • 大規模な計算グラフを一度に学習しにくい制約を減らすために、Graph Segment Training はグラフをセグメントに分割してメモリ使用量を下げ、学習時間を3倍短縮する
  • Kaggle大会では、66か国から616チーム792人が参加し、グラフ圧縮、パディング値の調整、ノード特徴の追加、設定間attentionなど、実際のコスト予測モデル改善手法が検証された

MLコンパイラが実行性能を左右する理由

  • 最新のMLモデルは、自然言語理解、対話、画像生成、動画生成のような作業を実行し、TensorFlow、JAX、PyTorchのような MLプログラミングフレームワーク で記述・学習される
  • フレームワークは、行列積や畳み込みのような線形代数演算と、2D畳み込みレイヤーやTransformerレイヤーのようなニューラルネットワーク層を提供する
  • ユーザーは、モデルをハードウェア上で効率よく実行する細かな方法を直接扱う必要がなく、フレームワークの下にある コンパイラ がモデルを自動最適化する
  • ただしコンパイラは、複雑な最適化問題をヒューリスティックで解くことが多く、常に最適な性能を出せるとは限らない

計算グラフと二段階最適化

  • MLコンパイラは、ユーザーが記述した数学的命令を、実際のハードウェアで実行可能な命令へ変換する
  • MLプログラムは 計算グラフ として表現できる
    • ノードはmatrix multiplicationのようなテンソル演算を表す
    • エッジは、あるノードから別のノードへ流れるテンソルを表す
  • コンパイラ最適化は大きく2種類に分かれる
    • グラフレベル最適化: グラフ全体の文脈を考慮して決定を下し、グラフ全体を変換する
    • カーネルレベル最適化: fused subgraphである1つのカーネルを、他のカーネルとは独立に変換する

メモリレイアウトの性能トレードオフ

  • 行列のような2Dテンソルは、メモリ上で [A B C a b c] または [A a B b C c] の形で保存されうり、これはそれぞれrow-majorとcolumn-majorレイアウトに対応する
  • MLコンパイラの重要な最適化の1つは、プログラム内のすべての中間テンソルに メモリレイアウト を割り当てる作業である
  • 特定のレイアウトが個々の演算では最も効率的でも、addとconvolutionの間でレイアウトが一致しなければ、コンパイラは追加の copy演算 を挿入しなければならない
  • 逆に、個々の演算性能はやや低くても、レイアウト変換が不要な構成のほうが、全体の実行では有利な場合がある
  • XLA benchmark suiteでは、デフォルトのコンパイラ設定の代わりに最適なレイアウト構成を選ぶと、最大 32%の高速化 が観測された

TpuGraphsデータセット

  • TpuGraphs は、Googleのカスタム TPU 上で実行されるプログラムのための 学習ベースのコストモデル データセットである
  • 目標は、入力プログラムとコンパイラ設定を受け取り、プログラムの実行時間を予測するコストモデルを学習することにある
  • データセットは2種類のXLAコンパイラ設定を対象にしている
    • layout: 行列のrow-major・column-majorという概念を高次元テンソルへ一般化した設定
    • tiling: タイルサイズ設定
  • 各サンプルは、MLワークロードの計算グラフ、コンパイル設定、その設定でコンパイルしたときの実行時間を含む
  • グラフはオープンソースのMLプログラムから収集され、ResNet、EfficientNet、Mask R-CNN、Transformerのようなモデルアーキテクチャを含む
  • ダウンロード方法と開始用コードは TpuGraphs GitHub で提供されている
  • TpuGraphsは、同程度のグラフサイズを持つ既存最大のgraph property predictionデータセットよりグラフ数が 25倍 多く、既存のMLプログラム性能予測データセットより平均グラフサイズが 770倍 大きい

ベースラインのコストモデルとGNN構造

  • TpuGraphsはベースラインの学習ベースコストモデルも提供しており、入力プログラムがグラフとして表現されるため GNN を用いる
  • ノード特徴は2つの部分で構成される
    • opcode id: テンソル演算の種類を示す、最も重要なノード情報
    • その他のノード特徴
  • ベースラインモデルは、opcode idをembedding lookup tableによって opcode embedding に変換する
  • opcode embeddingと残りのノード特徴を結合し、GNN入力として使用する
  • GNNが生成したノードembeddingは、sumやmeanのような単純なgraph pooling reductionによって固定長のグラフembeddingへ結合される
  • 最終的なグラフembeddingはfeedforward layerを経て、1つのscalar出力へ変換される

Graph Segment Trainingで大規模グラフを学習する

  • Graph Segment Training は、メモリ容量が限られたデバイス上で大規模グラフを扱うための GNN学習スケーリング手法 である
  • この手法は、予測対象がノードやエッジではなくグラフ全体であるgraph-level predictionの状況を対象としている
  • 計算グラフは数十万個のノードを含みうるため、グラフ全体を一度に使うFull Graph Trainingは計算上不可能な場合がある
  • GSTは大規模グラフを小さなセグメントに分割し、セグメントのランダムな部分集合だけを選んでモデルを更新する
  • 残りのセグメントは中間activationを保存せずにembeddingを生成し、メモリ使用量を削減する
  • すべてのセグメントembeddingを結合して元の大規模グラフのembeddingを作り、これを予測に使用する
  • historical embedding tableとsegment dropoutも導入し、historical embeddingのstalenessを緩和する
  • この手法全体により、end-to-endの学習時間が 3倍 短縮された

Kaggle大会で検証された改善手法

  • Fast or Slow? Predict AI Model Runtime Kaggle大会はTpuGraphsデータセットを基に開催され、66か国から616チーム792人が参加した
  • 提出数は 10,507件 で、153人にとっては初めてのKaggle大会であり、そのうち47人は上位100人に入った
  • 参加チームはさまざまな手法を実験した
    • グラフ pruning・compression: GSTの代わりに大規模グラフを圧縮する方法が試され、設定可能なノードとその直近の近傍を含むsubgraphだけを保持する方式が使われた
    • パディング値の変更: デフォルトのpadding値0は有効なfeature値と衝突するため、-1を使うとモデル精度を大きく改善できた
    • ノード特徴の追加・エンコーディング変更: dot generalのcontracting dimensionsのような追加ノード特徴が重要であり、ノード特徴のエンコーディング方式も結果に影響しうる
    • cross-configuration attention: 優勝チームは、モデルが設定同士を明示的に比較できる単純なlayerを設計し、各設定を個別に推論させる方式よりはるかに良い結果を出した
  • 大会結果と優勝ソリューションは、2023年12月16日にNeurIPSの ML for Systems workshop 大会セッションで扱われる予定である

NeurIPS Expo関連セッション

  • 構造化データと人工知能研究に関心のある読者向けに、NeurIPS Expoパネル Graph Learning Meets Artificial Intelligence が2023年12月9日に開催される
  • このパネルでは、学習ベースのコストモデルの進展などが取り上げられる

関連記事

1件のコメント

 
GN⁺ 2023-12-17
Hacker News のコメント

  • ML コンパイラは過大評価されている。従来のコンパイラと同じトレードオフだ。専門の性能プログラマを雇うよりははるかに多くのスループットを得られるが、後者のほうがたいていずっと速く、場合によっては数桁の差で上回ることもある。
    いくつもの層で不足している。アルゴリズムのレベルでは、ネットワークをより高速に変えるためのコツを人間へフィードバックできず、ごく基本的なシグナルを与える程度にとどまる。意図も失われる。ML ネットワークの設計者は Python で構造を指定するが、何段階もの下位変換を経ると、まったく見当違いの結果になることがある。最近、あるコンパイラが slice update を行う際に、配列の取り得る全インデックス範囲を作り、それを切り出して更新対象のインデックスを得てから scatter しているのを見たので、単一の memcpy 呼び出しに置き換えた。カーネルも非効率だ。こうしたコンパイラ出力が熟練したアセンブリプログラマと対決するたびにコンパイラは負け、たいてい 30%以上 の差が出る。簡単に解決できそうに見えるが、この50年間誰もまともに解けていないのなら、言うほど単純ではないのは明らかだ。

    • チェスエンジンの Stockfish を見ると、盤面評価で何年も人間が書いてきたヒューリスティックを捨て、小さなニューラルネットワークに置き換えたところ、より良くなった。
      コンパイラにもインライン化、ループ展開、ベクトル化のようなヒューリスティックが多いので、ニューラルネットワークが役立つ可能性があり、人間が書いた多数のヒューリスティックより保守もしやすいかもしれない。
    • 従来のコンパイラと同じトレードオフだと言っているが、その従来のコンパイラがものすごく有用だったという点が面白い。
    • あまりに断定的で、閉じた態度に見える。
    • その通り。手でアセンブリを書く人を雇えばいいのに、なぜ誰かが gcc/clang を使うのだろう？
    • その スループット こそが重要だ。すべての ML 作業に性能専門家を付けることはできない。
      こうした最適化がまったくないよりは、今でもはるかにましだ。

  • これをもう少し現実的に説明してもらえる？ 現在の ML コンパイラ の実際の状況がどうで、近い将来に何を期待できるのか知りたい。

    • もっとも手軽なアプローチの一つは torch.compile だ。PyTorch コンパイラの最新の反復版で、以前の方式には TorchScript と FX Tracing があった。
      単に model = torch.compile(model) と書けばよい。「これら163個のオープンソースモデル全体で、torch.compile は93%のケースで動作し、NVIDIA A100 GPU での学習が43%速くなった。Float32 精度では平均21%、AMP 精度では平均51%速くなった。」[1] Google は、この種の研究開発にもっと多くの人が参加できるようにしようとしているように見える。
      [1] https://pytorch.org/get-started/pytorch-2.0/
    • 近い将来に期待できるのは、モデルが開発されたフレームワークに対して明示的なベンダーサポートがなくても、AMD、CUDA、TPU、CPU などを使えるようになることだ。
      現実は複雑なのでかなり単純化して言うと、計算グラフを何らかの中間表現にコンパイルし、それに合うバックエンドを実装する方式だ。関連プロジェクトとしては stableHLO、IREE、openXLA を見ればよい。Jax の jit コンパイラも、こうしたコンパイラの一形態と見なせる。トレースされた演算を XLA に下ろし、XLA がさらにバックエンドで動くようにいろいろな魔法を使う。結局、下に降りるほどずっと 変換と抽象化 だ。
    • torch.compile を見ればよい。

  • 要約すると、計算グラフの実行時間性能予測を グラフニューラルネットワーク（GNN） で改善する取り組みだ。各ノードの opcode に対する埋め込み辞書と、shape、bits、window size のような他のノード特徴量を一緒に使う（[1]）。
    さまざまな XLA コンパイル設定と TPU での結果性能を含む大規模なグラフデータセットを [2] で公開しており、グラフを分割する方式（METIS graph partition は初めて見た）と複数の学習手法によって、[3] で以前より大きなグラフの予測を改善している。これは与えられたグラフの性能を予測する話であって、同等な新しいグラフを改善・提案・修正する話ではない。FunSearch のように予測力がそこそこあるモデルは、進化的探索と組み合わせて使える。
    [1] https://github.com/google-research-datasets/tpu_graphs#featu...
    [2] TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs https://arxiv.org/abs/2308.13490
    [3] Learning Large Graph Property Prediction via Graph Segment Training https://arxiv.org/abs/2305.12322

  • あのグラフで 畳み込み がどう動くのか説明してもらえる？ shape が [2,4,16] のテンソルに shape が [4,16,8] のカーネルを畳み込んだら [2,8] のテンソルが出てくるなんて、どうして可能なのか？

    • 役に立つかわからないが、入力テンソル [2,4,16] の 2 はバッチサイズ、4 は入力特徴次元、16 は入力チャネル次元と見なせる。
      カーネル [4,16,8] の 4 はフィルタ窓のサイズ、16 は入力チャネル次元に対応する値、8 は出力チャネル次元だ。出力 [2,8] の 2 はバッチサイズとして維持され、8 はカーネルの出力チャネル次元に対応する。一見すると次元が合わないように見えるが、グラフ上の畳み込みは 近傍構造 を利用する。カーネルがグラフ上を移動しながら、現在のノードと一定半径内の近傍特徴に重みを適用し、その重み付き和を集めて各出力チャネルの新しい特徴を作る。グラフ構造やエッジ重み、パディングやストライドといった実装上の詳細も出力 shape に影響し得る。

  • Gemini はどんな状況？

    • GPT-4 が引き続き優勢だという点が興味深い: https://twitter.com/lmsysorg/status/1735729398672716114
      思い浮かぶだけでも Llama、Claude、Gemini、Falcon、Mistral のような基盤モデルが少なくとも5つあり、互いに抜きつ抜かれつしているが、GPT はいまだに一段上にいて、1年経ってもそうだ。Transformer ベースの大規模言語モデルは十分に単純で、GPU 時間に100万ドルほど使えるなら誰でも作れることが明らかになったが、それでも OpenAI に完全には追いつけていない。彼らの特別な秘訣は何なのだろう？

  • Transformer 自体はどうなのか？ どんな意味で最適だという手がかりはあるのか？

  • 最初の段落で核心を埋もれさせてしまっている感じはするが、残りは素晴らしい内容だ。

  • 今の ML の進歩速度 は驚くべきものだ。シンギュラリティは信じていないが、ソフトウェアと社会を誰にも予測できない形で変えつつある。

    • これと FunSearch を見ると、シンギュラリティが差し迫っているように見える。
      https://deepmind.google/discover/blog/funsearch-making-new-d...
    • 私には、ドットコム、モバイル、クラウド、VR の次に来た、もう一つの ゴールドラッシュ に見える。
    • 5年後には、人々は今のようにはプログラミングしていない気がする。
    • まずは治療が難しい病気の 治療法 を出してくれるところを見たい。シンギュラリティそのものは、人間に利益がなければ意味がなく、その利益は主に健康の改善と苦痛の軽減にあるべきだ。