6

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

 1 year ago
source link: https://www.v2ex.com/t/913802
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

V2EX  ›  程序员

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

  MegEngineBot · 1 小时 22 分钟前 · 662 次点击

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

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

image.png

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

先上对比效果

up-893b768c757e3cf316b0635e694acf9e05a.jpg

背后的逻辑很简单:计算换显存 - 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~  :D


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK