Coop: Memory is not a Commodity

要約

テンソルの再具体化により、必要に応じてモデルにチェックポイントを設定し、削除されたテンソルを再計算することで、限られたメモリ バジェットの下でディープ ニューラル ネットワーク (DNN) のトレーニングが可能になります。
ただし、既存のテンソル再具体化技術は、深層学習フレームワークのメモリ システムを見落としており、異なるアドレスにある空きメモリ ブロックが同一であると暗黙的に想定しています。
この欠陥のある仮定の下では、不連続なテンソルが排除され、その一部は新しいテンソルの割り当てに使用されません。
これにより、深刻なメモリの断片化が発生し、潜在的な再具体化のコストが増加します。
この問題に対処するために、すべてのエビクションが連続してすぐに使用されるようにするために、スライディング ウィンドウ内でテンソルをエビクトすることを提案します。
さらに、テンソル割り当てを最適化することで再実体化コストをさらに削減するために、安価なテンソル分割とその場で再計算可能な方法を提案しました。
テンソル割り当てとテンソル再具体化の同時最適化であるため、このメソッドを Coop と名付けました。
Coop を 8 つの代表的な DNN で評価しました。
実験結果は、Coop が最先端のベースラインと比較して最大 $2\times$ のメモリ節約を実現し、計算オーバーヘッド、検索レイテンシ、メモリの断片化を大幅に削減することを示しています。

要約(オリジナル)

Tensor rematerialization allows the training of deep neural networks (DNNs) under limited memory budgets by checkpointing the models and recomputing the evicted tensors as needed. However, the existing tensor rematerialization techniques overlook the memory system in deep learning frameworks and implicitly assume that free memory blocks at different addresses are identical. Under this flawed assumption, discontiguous tensors are evicted, among which some are not used to allocate the new tensor. This leads to severe memory fragmentation and increases the cost of potential rematerializations. To address this issue, we propose to evict tensors within a sliding window to ensure all evictions are contiguous and are immediately used. Furthermore, we proposed cheap tensor partitioning and recomputable in-place to further reduce the rematerialization cost by optimizing the tensor allocation. We named our method Coop as it is a co-optimization of tensor allocation and tensor rematerialization. We evaluated Coop on eight representative DNNs. The experimental results demonstrate that Coop achieves up to $2\times$ memory saving and hugely reduces compute overhead, search latency, and memory fragmentation compared to the state-of-the-art baselines.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google