ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
host memoryやNVMeを使ったout-of-coreの実装を真面目にクソ頑張ることで、より大きなtransformerを高い実行効率を維持しながら学習できるようにする話。
手法
基本的には各種データ（param, optimizer stats, activationなど）のアクセスパターンを考える。次に、それに応じて要求される読み込みbandwidthを計算する。実行効率を維持したいので、「計算とoverlapさせられるかどうか」で要求bandwidthが決まる。あとは、それに応じてデータの配置をする。
これを気合を入れて実装する。
感想
何かびっくりするようなアイディアが有る訳ではない。それもあってか論文の書かれ方は結果が凄いことを何度もゴリ押しして強調してくるからだるい。
こうやってちゃんと必要な性能を計算して設計するアプローチは当たり前っちゃ当たり前だがすごく大事と思う。実装めちゃくちゃ大変そうなので凄い。
