中文

Checkmate:以最优张量重材料化突破内存墙

机器学习 2020-05-15 v3 计算机视觉与模式识别 分布式、并行与集群计算 机器学习

摘要

我们将深度神经网络(DNN)训练时间与内存需求之间的权衡问题形式化为张量重材料化优化问题,这是对先前检查点策略的推广。我们介绍了 Checkmate,这是一个利用现成的 MILP(混合整数线性规划)求解器在合理时间(一小时以内)内求解最优重材料化调度、或使用近似算法获得近最优调度,然后利用这些调度加速数百万次训练迭代的系统。我们的方法可扩展到复杂、真实的架构,并通过使用特定加速器、基于性能剖析的成本模型实现硬件感知。除了降低训练成本外,Checkmate 还能使真实网络以最多 5.1 倍更大的输入尺寸进行训练。Checkmate 是一个开源项目,可在 https://github.com/parasj/checkmate 获取。

关键词

引用

@article{arxiv.1910.02653,
  title  = {Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization},
  author = {Paras Jain and Ajay Jain and Aniruddha Nrusimha and Amir Gholami and Pieter Abbeel and Kurt Keutzer and Ion Stoica and Joseph E. Gonzalez},
  journal= {arXiv preprint arXiv:1910.02653},
  year   = {2020}
}

备注

In Proceedings of 3rd Conference Machine Learning and Systems 2020 (MLSys 2020)