A100 买不到了，只有小显卡怎么训大模型

为了达到更好的训练效果，通常炼丹师们会使用更大的模型和更大的 Batch size，但因此带来的大显存占用，却成为不可避免的硬伤。

尤其是如今 GPU 越来越贵，甚至还可能买不到......

MegEngine v1.5 及以上版本，支持动态图和静态图的显存优化，显存占用可降至 1/4。

先上对比效果

背后的逻辑很简单：计算换显存 - MegEngine 使用重计算策略，实现了用时间换空间。

使用方法

动态图下开启：

在训练代码之前只需添加一行代码：

megengine.dtr.enable()

静态图下开启：

在编译静态图时使用 DTRConfig 设置 trace 的参数 dtr_config ：

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

更多使用技巧，见官方文档

使用 DTR 进行显存优化 - MegEngine 1.11 文档

延展阅读：

To 新朋友的特别提示：

DTR 作为 MegEngine 原生特性，无法脱离 MegEngie 独立使用。

好在 MegEngine 顶层 API 基于 Python，采取了类似于 PyTorch 的风格，已被众多用户反馈易上手，入门简单。

还有丰富的用户指南文档：

所以，不用顾虑太多，可以勇敢尝试哦~

开源地址：

（欢迎 star~ :D