cDLRM: Look Ahead Caching for Scalable Training of Recommendation Models
2021/01/13
著者: https://gyazo.com/c9fc2ee5e864ba490a1c89405079862a
選んだ理由
現実上での課題として、つよつよモデルを作っても学習時間の問題でプロダクトに組み込むことが難しいというものがある。
時間よりもマシンリソースの話題だった
特にディープを使ったものに関して顕著で、これに対するアプローチの手段として、参考にならないかと考えたため。
どんなもの?
近年の Deep learning recommendation models (DLRMs) は下記の 2 つから構成させる
スパースなカテゴリ変数を扱うための embed 層
multi-layer perceptrons (MLPs) のようなデンスな値を扱うための層
既存の問題
embed 層は巨大になりがちで、これらすべてを GPU に載せる必要があった
このために複数の GPU が必要で、学習にかかる計算資源へのコストが非常に大きかった
一方で 1 度のバッチで参照される embed 層対応部分は全体に対して小さく、GPU の資源を効率よく活用できているわけではない
問題への対処(cDLRM の提案)
そこで embed 層のデータは CPU に載せて置く
学習前のバッチの情報を先読みしておき、バッチの内で必要な embed 層のデータだけ GPU に配置する
先行研究と比べてどこがすごい?
embed 層の hash を使ったメモリ削減方法が先行研究にあるが、これは精度を毀損する課題がある
We would like to stress that the objective of cDLRM is not to be the fastest recommendation model training system, but to democratize recommendation model training to the point where it does not require hundreds of thousands of dollars of hardware and GPUs to be able to just fit the model.
一番言いたいのはこれ感
技術や手法のキモはどこ?
達成したいことは学習時のバッチに対応する必要な embed 層の records のみを GPU に載せるということである。
cDLRM は次のようなフローで実行される
1. 値に対応する embed 層の index を queue $ Q_{lookahead} に追加する
2. $ Q_{lookahead} からデータを取得して GPU 上のメモリに乗せる
3. forward & backward propagation を実行する
4. GPU 上で書き換えられた embed を元の CPU 上の embed に上書きする(eviction 処理)
注意として、GPU に先に乗った embed 層のパラメータが、他の処理によって更新された場合は古いパラメータを参照して学習するという課題がある。しかし、これはレアケースなので無視する(本当か...?)
https://gyazo.com/7aebb23e1dcf8d1e51ac1929d960a32b
cDLRM は単一 GPU で動作することを目的にしたものだが、スピードアップのため複数の GPU でも採用可能
一貫性のためにキャッシュの取得は 1 つの GPU が行い、それを他の GPU に broadcast する
back propagation の際には 2 つの方法をとる
MLP aggregation
gradient を全 GPU で平均する
Embedding Table Cache aggregation
更新後の embed 層のパラメータを全 GPU で平均する
どうやって有効だと検証した?
Facebook が提供する DLRM モデルをベースに cDLRM モデルを構築した
https://gyazo.com/e131613a5d27389ee62022910010effd
精度毀損がないことを検証した
https://gyazo.com/db57141a6ff8717d6dcb3bca4ca36ab7
実行時間についてはコンポネントを 2 つに分割して検証した
batch computation time: バッチの実行時間
amortized caching overhead: キャッシュに載せるまでのオーバーヘッド
https://gyazo.com/d66b6fef3f11b209f072184bfe8a118c
(a), (b) Cache Size が十分大きいと先読みする window size (Lookahead window size) の増加に伴いキャッシュに乗せるの実行時間が減る
(c) Cache Size と Lookahead window size はトレードオフの関係にあり、Lookahead window size を大きくしすぎると、cache の衝突が起こりパフォーマンスが下がる
複数の GPU を採用した場合
https://gyazo.com/133328d33e6748ff9d029b146bbeadb0
(a) 複数 GPU を採用することによって、巨大なバッチについても速く処理できる事がわかる
(b) GPU を増やすとキャッシュを載せるためのオーバーヘッドが全体の時間にしめる割合が高くなる
(c) 更新後の embed 層を平均するスパンを大きくすると、精度をより強く毀損することがわかる
議論はある?
GPU もりもり環境の時間との比較があればよりよかった
所感
実装が見たい気持ち
パラメーター間の比較だけではなく、既存の学習方法との比較もみたかった。