XEngine：异构环境下神经网络的最优张量重物化
机器学习
2022-12-22 v1 编程语言
摘要
在资源受限的设备上训练深度学习网络时，内存效率至关重要。在反向传播过程中，前向张量被用于计算梯度。尽管可以选择将这些依赖项保留在内存中直到它们在反向传播中被重用，但某些前向张量可以被丢弃，并在之后从保存的张量（即所谓的检查点）重新计算。这尤其允许资源受限的异构环境利用所有可用的计算设备。不幸的是，这些检查点的定义是一个非平凡问题，并对程序员构成挑战——不当或过度的重计算会抵消检查点带来的好处。在本文中，我们提出了XEngine，一种通过确定张量的检查点和重计算，将网络算子调度到低内存环境下的异构设备上的方法。我们的方法为每个时间步和算子选择合适的资源，并在考虑每个设备内存限制的情况下优化神经网络的端到端时间。为此，我们将深度学习网络算子在异构系统上的调度问题形式化为一个混合整数二次规划（MIQP）。我们将我们的MIQP求解器XEngine与Checkmate进行了比较，后者是一种解决单设备上重计算问题的混合整数线性规划（MILP）方法。我们的求解器找到的解决方案比仅在单个设备上计算网络的最快Checkmate调度快达22.5%。如果内存限制不允许仅将网络调度到图形处理单元，我们还能为同时利用中央处理单元和图形处理单元的网络找到有效调度。
引用
@article{arxiv.2212.09290,
title = {XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environments},
author = {Manuela Schuler and Richard Membarth and Philipp Slusallek},
journal= {arXiv preprint arXiv:2212.09290},
year = {2022}
}