MegEngine 正式支持 XLA 啦!
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 中已经支...
