Zenn
🍎

【U22 プロコン】僕がやってる深層学習フレームワークの設計などについて

2023/11/21に公開
AI
Common Lisp
tech

前置き

つい先日(2023/11/18~20)に東京へ行き, U22 プログラミングコンテストに参加してきました。自分の作品はありがたいことに最終審査まで残り, しかもテクノロジー部門で最高賞である経済産業大臣賞を頂きました。本当にありがとうございます。

他の入選者の方々におかれましても本当に最終審査会お疲れ様でした！自分の近い年齢の方でこんなにすごいものを作れる方がこんなにたくさんいるのかと, 終始感銘を受けていた1日でした。

最終審査会では事前に提出した作品説明の資料に加えて, 自分で10分間程度のプレゼンテーションを行い, 審査員の方々はそれらを総合的に評価し点数を付けて受賞作品を決定します。(多分)

今回提出した作品は事前にいろんな場所のLT会でも発表したし, 当日は落ち着いてやれば大丈夫だろう...と見積もっていたのですが, いざ他の方プレゼンが始まるとどの作品もあまりに圧倒的で(一次審査に残る時点で相当厳選されているのでそれはそう), しかもこういう時に限って僕の発表順が一番最後で, 緊張に完全に支配されてマイクを持つ手が震えながら超早口で発表を終えてしまいました...

今回の受賞結果は恐らく事前に提出した資料も加味していただいたものだと思うのですが, それでも他の方に自分が作ったものがうまく伝わらないまま終えてしまったのが心残りなので, 一度文字に起こしてわかりやすく伝えようと思いこの記事を執筆しました。

とりあえず今回はバックエンドだけに焦点を絞って話します。

深層学習フレームワークを作った

cl-waffe2という深層学習フレームワークを開発し提出しました。

https://github.com/hikettei/cl-waffe2

深層学習フレームワークというのは深層学習モデルを学習, 推論させるための基盤となるプログラムです。機械学習や深層学習に触れたことがある方ですと, 恐らくTensorflowやPyTorchあたりに馴染みがある方も多いと思います。

どの規模でフレームワークを開発するかはいろいろありますが, cl-waffe2は主に三つのレイヤーに分かれています：

  • バックエンド → コンパイラ, モデル最適化の機能など
  • フロントエンド → Numpy-Likeな行列演算ライブラリ, モデルの表現など
  • その他ユーティリティ → 確率分布の高速生成だとか, NNモデルの標準実装など

大前提: 深層学習コンパイラってなんぞや？？

多分これを先に話さなかったのが一番の失敗点です, ずっと科学計算の勉強を一人でやってて, どこまでが一般的な概念でどこからが専門的な話なのか考慮する努力を完全に欠いていました。

平たく言えば, 深層学習モデルのパラメーターは行列で表現できて, プログラム上ではN次元配列として表すことができます：

// 疑似コード
for i=0..10 {
    for j=0..10 {
        A[i * 10 + j] = B[i + 10 * j];
    }
}

速度の面で言えば上の疑似コードは完全ではありません: CPUならL1/L2キャッシュの失敗を防ぐためにTilingするだとか, メモリ局所性を最適化するだとか, 並列化, 命令融合など... ある計算資源の能力を最大限活用するために考慮することがたくさんあります。それら低レイヤーの知識を有した方がC/C++/CUDAなどを直接書くのも可能ですが, それら最適化を施したコードを生成するコンパイラを作ることも可能です。

そういった中で Petalisp というCommon Lispに埋め込まれた ドメイン固有言語(Embedded DSL) の先行研究が昔からあって有名で, cl-waffe2はこのプロジェクトにとても影響を受けています。

(なぜCommon Lispなのかも作者の方が解説しています)
https://github.com/marcoheisig/Petalisp/tree/master

(古い資料なので注意...)

https://european-lisp-symposium.org/static/2018/heisig.pdf

僕の理解が間違っていなければ, Petalispは:

    1. map, reduce, transform, range, fuse ... など8つくらいの基本的操作にλ関数を組み合わせて線形代数の操作を表現する
    1. (1.)に対して, Petalispは中間表現の最適化, Polyhedral Compilerによる並列化の自動Scheduling, JIT Compiler, などを適用することで高速なコードを生成できる

といったことも可能で, 例えば行列積みたいな基本的な操作も更に小さな操作に分割して表現することができて:

;; Quote from https://github.com/marcoheisig/Petalisp/blob/master/README.org, Example2
(defun matrix-multiplication (A B)
  (lazy-reduce #'+
   (lazy #'*
    (lazy-reshape A (transform m n to n m 1))
    (lazy-reshape B (transform n k to n 1 k)))))

僕の手元だとOpenBLASのgemmを呼び出す時と同じくらいの速度で動作します。更にPetalispにはCUDAのバックエンドも存在するので一度Petalispで欠いちゃえばCUDAでも(有志の方が実装すれば!!)Apple Siliconでも動作します。

Petalispの開発者様は世間がChainerを使っていた頃からPetalispの開発をされていたと思うので, ようやく時代が追いついてた感を感じていますが, 近年のLLMブームあたりから深層学習向けのコンパイラの研究も更に盛んになってると感じます:

The Deep Learning Compiler: A Comprehensive Survey (深層学習コンパイラのSurvey論文)

https://zenn.dev/acd1034/articles/230325-dl-compiler-overview

深層学習コンパイラになって, 複数バックエンドへの移植を考慮, 線形代数固有の最適化を追加するなど新しく考えることもあると思いますが, 個人的にはPetalispのアーキテクチャが綺麗だと思っています。

(多分)ここらへんの文章を読んでる時にcl-waffe2のアイデアが浮かんだんだと思います。

cl-waffe2のアイデア

結局cl-waffe2のバックエンドでは何がやりたかったのかを箇条書きでまとめます：

  • Petalispの設計を深層学習向けに最適化したかった

    • 命令はλ関数じゃなくてAbstractNodeクラスを用いる
    • 数学関数とかもSLEEFのSIMD拡張だとか使えるようにしたい
    • オペレーターが多すぎるとデバッグが辛いのでネットワーク可視化に関する機能をバックエンドレベルで追加したい
    • どうせmacroletするからElement-Wiseな命令もmapだけじゃなくて別々のクラスとして表現したい (そうじゃないとlog1p fusionとかが難しい)

  • 深層学習フレームワークをもっとコンパクトにしたい

    • 小さい命令を組み合わせる → 複雑な計算を表現 の仕組みを全面的に取り入れたい

      • Numpy-LikeなAPIではどうしても最適化できない冗長なメモリ読み書き問題を最適化したい

      • CPU/CUDA以外に移植を書くのが超簡単

      • (P.S.: 深層学習の低レイヤーを全て自動生成するみたいな表現が適切だと思います)

ただ後者のアイデアに関しては, 実は同じくらいの時期に始まったとあるプロジェクトとアイデアが被っちゃってました:

https://github.com/tinygrad/tinygrad/tree/master

こっちは人材も資本力もたくさんあるし, そもそも実装がcl-waffe2の何億倍も良い。言語選定以外は!! コンパイラをPythonで書くな!! ましてプロジェクトを主導してる方がかの有名なGeorge Hotzなので将来性しかない。絶対流行る。

(閑話休題) そんな感じのバックエンドを実装した上で:

  • デバッグとProfilingの容易ささえ確保できてるなら, 例えAPIが使いにくくても, 将来的に美しいフロントエンドでWrapするから問題ない
  • Primitiveなバックエンドと美しいフロントエンド, その両方を一つのパッケージで完結させたら最強じゃね？→cl-waffe2

みたいなのが僕の脳内でした。

フロントエンド

こっちもいろいろ調べてたんですけど, またいつか機会があったら紹介します:

最後に

改めて経済産業大臣賞に選んでいただきありがとうございました。

受賞後各方面からたくさんの反応を受け取っていますが, とりあえずcl-waffe2は引き続き僕の趣味として細々とメンテナンスします。というのも今回紹介した技術も正直背伸びをしながら勉強しすぎたというか, あまり焦って開発してもいつかボロが出そうなのでとりあえず大学行ってゆっくり勉強します。加えてもう既にtinygradがたくさんの資本力と人材を動員して素晴らしいフレームワークを実現しようとしていますので, もしcl-waffe2に興味を持っていただけた方がいましたらこちらも参照していただけると楽しめると思います。

また, cl-waffe2で失敗したなと思うこともいくつかあって:

  • Frontend/Backendが思ってたより規模が大きくなってしまい, それぞれ別のリポジトリで開発すべきだった。(モジュラー性の大切さを学べました)
  • もうちょっと慎重に実装すべきだった。クソコードが多すぎる
  • 大きいプロジェクトになればなるほど名前空間の設計は慎重になるべきだった
  • 一人で開発するのはむずかしい(小並感)

これらの反省も今後の開発に活かしていきたいです。

失敗したことも多かったですが, 振り返ってみると高校生活最後の一年間を自分の技術力の限界を試すために費やしたのは間違った選択ではなかったと感じています。某自称進学校(注: 現在所属している学校では断じてありません！！！)のように 高校生は黙って学校のことを聞いて勉強してろやって言う大人も存在しますが, 何かで突き抜ければ必ずどこかで評価してくれる人間は見つけられることを行動で証明できたので満足です。

References

あとがき

もっといい説明思いついた (いつか推敲しておく)

Tinygrad/Petalisp = RISC (CNNとかRNNも自動生成)
cl-waffe2.                = RISCとCISCの中間 (自動生成もできるけどC++手動で書くのもいける)
TVM/PyTorch.        = CISC (CNNとかRNNをC++で手動で書く)

Discussion