博客

MegEngine 正式支持 XLA 啦!
作者:MegEngine 发布日期:2023/12/13

XLA(Accelerated Linear Algebra)是 Google 提出的一个神经网络编译器,可以用于加速 AI 模型的训练和推理。MegEngine 1.13.1 中也已经支持了 XLA,在训练模型时可以选择开启此项功能,不同的模型可以获得 10%~80% 不等的速度提升。

主要的目标场景

MegEngine 现在是动态执行的,即 python 中每一个 mge.functional 的调用都对应着底层 gpu 上的一次 kernel 执行。这种模式的好处在于实际的执行方式与代码逻辑一致,所见即所得,非常的灵活;不过其问题是难以优化,性能可能不是最优。

而 XLA 采取静态执行的方式,会将模型计算过程表达成一张静态计算图,称为 “HLO” (High-Level Optimized)。HLO 中包含计算图的相关操作,张量的数据流程和形状等信息。XLA 随后会对 HLO 进行一系列的优化,并最终生成一个更优的计算图,从而更快的完成计算。而 XLA 的局限性就在于不够灵活,对于 Tensor Shape 改变或者控制流等信息无法很好的表达。

现在 MegEngine 中已经支持了 XLA,模型训练中一些比较静态的场景,我们可以使用 XLA 来进行加速,从而缩短整个训练过程的时间。

使用方法与效果

在使用 MegEngine 进行训练时,可以通过对原来的训练函数增加 xla_trace/partial_trace 装饰器来启用 XLA 编译优化。

当整个模型是完全静态时,我们可以使用 xla_trace 将整张网络表达成一张静态图,然后交由 XLA 做后续的优化编译,后续的执行过程将执行这张优化后的计算图提升速度。

而如果我们模型中有一些动态性,比如训练过程中一些 Tensor Shape 会发生变化,亦或者是存在控制流,我们可以使用 partial_trace,将网络中静态的部分 trace 成一些子图并分别交给 XLA 进行编译优化,而网络中其他部分仍然保持动态执行,同时保证性能与灵活性。

下面展示了在 MegEngine 中,XLA 功能开启前后,主流的神经网络模型性能变化。其中蓝色为 XLA 开启之前的训练速度,橙色为 XLA 开启之后的训练速度。在开启 XLA 后,大部分模型的性能可以获得 10%~40% 的提升,最多可以超过 80%。
1.png

关于 xla 的更多信息及具体的使用方法可以参考 https://www.megengine.org.cn/doc/stable/zh/user-guide/model-development/jit/xla.html

更多 MegEngine 信息获取,您可以:查看文档、和 GitHub 项目,或加入 MegEngine 用户交流 QQ 群:1029741705。欢迎参与 MegEngine 社区贡献,成为 Awesome MegEngineer,荣誉证书、定制礼品享不停。

查看更多

上一篇: 分享

编译器上手指南,算子开发及开源项目指导手册,直播报名通道限时开启!

下一篇: 分享

MegEngine dataloader 新工具帮助定位性能瓶颈,快来体验吧!

相关推荐
  • MegEngine 版本最新发布!新增支持寒武纪思元系列 AI 芯片训练和推理

    2024/01/02

  • 一个深度学习框架的年度报告

    2022/12/30

  • MegEngine 11-12 双月报:新版本发布,开发者福利课程,MegEngine 使用技巧,精彩不容错过!

    2024/01/05

  • MegEngine dataloader 新工具帮助定位性能瓶颈,快来体验吧!

    2023/12/19

  • 旷视天元Beta版核心技术升级全面解读

    2020/07/24