Rockmate：PyTorch 中一种高效、快速、自动且通用的重计算工具
机器学习
2023-07-06 v1 编程语言
摘要
我们提出 Rockmate 以控制 PyTorch DNN 模型训练时的内存需求。Rockmate 是一种自动工具，从模型代码出发，生成等价模型，使用预定义大小的内存用于激活值，代价是少量重计算。Rockmate 自动检测计算与数据依赖的结构，并将初始模型重写为复杂块序列。我们表明这种结构广泛存在，可在文献中许多模型（基于 Transformer 的模型、ResNet、RegNets 等）中找到。该结构使我们能够以快速高效的方式解决问题，在单个块层面使用 Checkmate（在整个模型上太慢但通用）的适配版本，在序列本身层面使用 Rotor（快速但限于序列模型）的适配版本。我们在许多模型上的实验表明，Rockmate 与 Rotor 一样快、与 Checkmate 一样高效，并且在许多情况下能以相当可忽略的开销（约 10% 到 20%）显著降低了激活值的内存消耗（降低 2 到 5 倍）。Rockmate 是开源的，可在 https://github.com/topal-team/rockmate 获取。
引用
@article{arxiv.2307.01236,
title = {Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch},
author = {Xunyi Zhao and Théotime Le Hellard and Lionel Eyraud and Julia Gusak and Olivier Beaumont},
journal= {arXiv preprint arXiv:2307.01236},
year = {2023}
}