写点什么

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

作者:MegEngineBot
  • 2023-02-06
    北京
  • 本文字数:652 字

    阅读完需:约 2 分钟

为了达到更好的训练效果,通常炼丹师们会使用更大的模型和更大的 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):


...


更多使用技巧,见官方文档  https://www.megengine.org.cn/doc/stable/zh/user-guide/model-development/dtr/


延展阅读:


https://zhuanlan.zhihu.com/p/375642263


https://www.bilibili.com/video/BV1Bg411c7cf?spm_id_from=333.999.0.0


To 新朋友的特别提示:


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

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

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

https://www.megengine.org.cn/doc/stable/zh/user-guide/index.html

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


开源地址: https://github.com/MegEngine/MegEngine   (欢迎 star~  :

用户头像

MegEngineBot

关注

工业级研发的开源深度学习框架-MegEngine​ 2022-10-21 加入

官网:https://www.megengine.org.cn/ MegEngine 技术交流 QQ 群:1029741705 框架使用相关交流或反馈,欢迎访问论坛:https://discuss.megengine.org.cn/ GitHub :http://github.com/MegEngine/MegEngine

评论

发布
暂无评论
A100 买不到了,只有小显卡怎么训大模型_深度学习_MegEngineBot_InfoQ写作社区