Coop: 内存并非商品

Jianhao Zhang; Shihan Ma; Peihong Liu; Jinhui Yuan

Coop: 内存并非商品

机器学习 2023-11-02 v1 人工智能软件工程

作者: Jianhao Zhang , Shihan Ma , Peihong Liu , Jinhui Yuan

摘要

张量重材料化(tensor rematerialization)通过检查点模型并在需要时重新计算被驱逐的张量,使得在有限内存预算下训练深度神经网络(DNNs)成为可能。然而,现有的张量重材料化技术忽视了深度学习框架中的内存系统,并隐含地假设不同地址的空闲内存块是相同的。在这一有缺陷的假设下,不连续的张量被驱逐,其中一些并未用于分配新张量。这导致严重的内存碎片并增加了潜在重材料化的代价。为解决此问题,我们提出在滑动窗口内驱逐张量,以确保所有驱逐都是连续的并立即被使用。此外,我们提出了廉价的张量划分和可重算原地操作,通过优化张量分配进一步降低重材料化代价。我们将我们的方法命名为 Coop,因为它是张量分配与张量重材料化的协同优化。我们在八个代表性 DNN 上评估了 Coop。实验结果表明,与最先进的基线相比,Coop 实现了高达 $2\times$ 的内存节省,并大幅降低了计算开销、搜索延迟和内存碎片。

关键词

张量重材料化深度神经网络内存碎片张量分配滑动窗口协同优化检查点

引用

@article{arxiv.2311.00591,
  title  = {Coop: Memory is not a Commodity},
  author = {Jianhao Zhang and Shihan Ma and Peihong Liu and Jinhui Yuan},
  journal= {arXiv preprint arXiv:2311.00591},
  year   = {2023}
}

备注

NeurIPS 2023 spotlight

Coop: 内存并非商品

摘要

关键词

引用

备注

相关论文