使用 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):
...