使用 XLA 作为编译后端加速模型训练¶
Note
XLA 作为 jit 后端目前属于实验特性，可能存在一些 bug
算子支持还不完备，可能存在部分算子不支持/算子 lowering 不正确的问题
相关接口不稳定， 后续版本可能会有 API 变动
XLA 编译后端¶
XLA 是 google 推出的可用于模型训练加速的编译器。 MegEngine 接入了 XLA 编译器用于提高训练性能。
在 MegEngine 中使用 XLA 编译器需要用户安装 mge_xlalib 包。
Note
仅支持在 MegEngine 1.13 及以上版本使用 XLA 编译后端。
安装的 mge_xlalib 包的 cuda 版本需要与用户使用的 MegEngine 包相同。
mge_xlalib 目前支持的 python 版本是 3.8 到 3.10， cuda 版本是 11.1 11.4 11.8 。
mge_xlalib cuda 11.8 加速效果最好， cuda 11.1 由于 cudnn 版本较低 XLA 加速效果较差。
mge_xlalib 安装命令如下：
# cuda 11.1
python3 -m pip install mge-xlalib==0.4.7+cuda11010.cudnn804 -f https://www.megengine.org.cn/whl/mge.html
# cuda 11.4
python3 -m pip install mge-xlalib==0.4.7+cuda11040.cudnn821 -f https://www.megengine.org.cn/whl/mge.html
# cuda 11.8
python3 -m pip install mge-xlalib==0.4.7+cuda11080.cudnn860 -f https://www.megengine.org.cn/whl/mge.html
XLA 编译器的使用方式与 MegEngine graph runtime 自带编译器类似， 需要用MegEngine提供的装饰器 (xla_trace) 对训练函数进行包装。 函数执行第一遍时会记录算子执行序列，以捕获静态图。 后续执行会把静态图用XLA编译， 并调用编译好的 XLA executable 加速训练过程。
由于 xla_trace 实现和 XLA 编译器自身的特性， 目前该功能的使用还有如下限制：
所有（非常量的）外部 Tensor /包含外部 Tensor 的对象（如 Module， Optimizer等）需要作为被装饰器包装的训练函数的参数传入。外部 Tensor 指所有不是由 Op 计算产生的 Tensor（对应静态图中的输入节点）。
由于 XLA 目前对动态 Shape 支持较差，需要保证所有外部 Tensor 的 Shape 是固定的。
需要用户保证被 trace 函数是静态的（函数不会在输入不同时执行不同的分支）。
具体使用示例见 使用 xla_trace 装饰器 和 使用 partial_trace 装饰器
使用 xla_trace 装饰器¶
您可以使用 jit.xla_trace 装饰器对需要加速的函数进行包装， 代码示例如下。
import numpy as np
import megengine as mge
import megengine.functional as F
from megengine.jit import partial_trace, xla_trace
def softmax(inp):
offset = inp.max(axis=-1, keepdims=True).detach()
cached = F.exp(inp - offset)
down = F.sum(cached, axis=-1, keepdims=True)
return cached / down
@xla_trace
def xla_fused_softmax(inp):
offset = inp.max(axis=-1, keepdims=True).detach()
cached = F.exp(inp - offset)
down = F.sum(cached, axis=-1, keepdims=True)
return cached / down
inp = mge.tensor(np.random.randn(256, 1000, 1000))
xla_fused_softmax(inp) # run in imperative runtime, trace op sequence
print (softmax(inp))
print (xla_fused_softmax(inp)) # run in xla
当模型训练迭代（Iteration）完全静态的情况下， 您也可以使用 jit.xla_trace 装饰器将训练迭代全部交由XLA执行。 需要将 optimizer， module 作为train_func 参数传入，同时 train_func 中需包含包含模型前向、 反向 、 参数更新等代码， 代码示例如下：
from megengine.jit import xla_trace
@xla_trace(capture_as_const=True) #capture_as_const为True时, 所有不在train_func 参数列表里的外部Tensor会被当成常量捕获
def train_func(data, label, *, opt, net):
gm = GradManager()
gm.attach(net.parameters())
with gm:
logits = net(data)
loss = F.loss.cross_entropy(logits, label)
gm.backward(loss)
opt.step().clear_grad()
return loss
for epoch in range(total_epochs):
total_loss = 0
for step, (batch_data, batch_label) in enumerate(dataloader):
data = mge.tensor(batch_data)
label = mge.tensor(batch_label)
loss = train_func(data, label, opt=optimizer, net=model)
total_loss += loss.numpy().item()
print("epoch: {}, loss {}".format(epoch, total_loss/len(dataloader)))
使用 partial_trace 装饰器¶
模型训练迭代中存在动态执行逻辑的情况下， 无法将整个计算交由 XLA 执行。 这种情况下可以使用 jit.patrial_trace 装饰器对其中静态的部分进行加速。 被 partial_trace 包装部分的前向/反向会使用 XLA 执行, 其他部分仍由 MegEngine 执行。 代码示例如下：
from megengine.jit import partial_trace
@partial_trace(backend="xla", capture_as_const=True)
def backbone(model, inp):
return model(inp)
for epoch in range(total_epochs):
total_loss = 0
gm = GradManager()
gm.attach(net.parameters())
for step, (batch_data, batch_label) in enumerate(dataloader):
data = mge.tensor(batch_data)
label = mge.tensor(batch_label)
with gm:
logits = backbone(net, data)
loss = F.loss.cross_entropy(logits, label)
gm.backward(loss)
opt.step().clear_grad()
total_loss += loss.numpy().item()
print("epoch: {}, loss {}".format(epoch, total_loss/len(dataloader)))