使用 DTR 进行显存优化¶

MegEngine 通过引入 DTR 1 技术来进行动态图下的显存优化，同时也支持在静态图下开启。

1: Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic tensor rematerialization. In International Conference on Learning Representations. 2021. URL: https://openreview.net/forum?id=Vfs_2RnOD0H.

DTR 使用与配置方式¶

在训练代码之前添加一行代码，即可启用动态图的 DTR 显存优化：

>>> megengine.dtr.enable()

1.5 新版功能: 用户现在可以直接开启 DTR 优化，不再需要设置一个显存阈值 eviction_threshold 作为触发条件。 MegEngine 默认会在当前空闲的显存无法满足一次申请时尝试进行优化，根据 DTR 策略找出最优的 Tensor 并释放其显存，直到该次显存申请成功。

而在 1.4 版本中，必须提前设置好显存阈值，才能开启 DTR 显存优化：

>>> megengine.dtr.eviction_threshold = "5GB"
>>> megengine.dtr.enable()

显存阈值的设置技巧

一般情况下，显存阈值设得越小，显存峰值就越低，训练耗时也会越大；显存阈值设得越大，显存峰值就越高，训练耗时也会越小。但值得注意的是，当显存阈值接近显卡容量时，容易引发碎片问题。因为 DTR 是根据活跃的显存大小来执行释放操作的，释放掉的 Tensor 在显卡上的物理地址很可能不连续。例如：释放了两个物理位置不相邻的 100MB 的 Tensor, 仍然无法满足一次 200MB 显存的申请。此时就会自动触发碎片整理操作，对性能造成巨大影响。

结合分布式训练

在分布式情景下，我们通常会使用 launcher 将一个函数包装成一个多进程运行的函数，此时如果想要开启 DTR 显存优化，需要在被包装的函数中定义 DTR 的参数：

@dist.launcher
def main():

    megengine.dtr.enable()

如果你还不清楚相关概念，可参考分布式训练（Distributed Training）页面了解细节。

参见

还有一些其它的接口如 evictee_minimum_size, enable_sqrt_sampling … 可以对 DTR 策略进行自定义，更多配置说明请参考 dtr 模块的 API 文档页面。

在静态图下开启 DTR¶

用户在编译静态图时使用 DTRConfig 设置 trace 的参数 dtr_config, 就可以打开 DTR 优化：

from megengine.jit import trace, DTRConfig

config = DTRConfig(eviction_threshold=8*1024**3)

@trace(symbolic=True, dtr_config=config)
def train_func(data, label, * , net, optimizer, gm):
    ...

通过重计算节省显存（Recomputation）

使用 Sublinear 进行显存优化