桥接PyTorch和TVM
桥接PyTorch和TVM
人工智能最引人入胜的一些应用是自然语言处理。像BERT或GPT-2之类的模型及其变体，可以获住足够多的文本信息。
这些模型属于称为Transformers的神经网络类体系结构。 HuggingFace transformers library是实现最受欢迎的库之一。
与已经高度优化的实现的卷积模型或LSTM相比，对于Transformers而言，情况并非如此。本文探索TVM如何填补空白。分两个步骤进行操作：
- 首先，在TVM上，使用BERT inference推理和调优。
- 其次，进行一些更基本的探索，以了解如何在PyTorch中使用TVM进行训练。考虑到实验性质，将重点更多地放在可行性上，而不是性能上。
- 使用TVM优化BERT推理
如何将BERT从Transformers库传输到TVM？
有用的是，transformers支持使用PyTorch JIT跟踪其模型，直到建立跟踪模型。
使用示例输入在AMD Radeon VII上进行100次运行后，PyTorch跟踪模型大约需要0.65-0.7秒，意味着每次运行6.5-7ms。看看是否可以使用TVM更快。将模型转换为TVM轻而易举：
shape_list = [(i.debugName().split('.')[0], i.type().sizes()) for i in list(traced_model.graph.inputs())[1:]]
mod_bert, params_bert = tvm.relay.frontend.pytorch.from_pytorch(traced_model,
shape_list, default_dtype="float32")
找不到dtype信息将有一些警告，但是一切顺利！构建并运行，构建遵循标准的TVM配置。将PyTorch（cpu）张量转换为TVM数组。
target = 'rocm -model=gfx906' # use what matches your GPU
target_host = 'llvm'
ctx = tvm.context(target)
tt_a = tvm.nd.array(tokens_tensor.numpy(), ctx)
st_a = tvm.nd.array(segments_tensors.numpy(), ctx)
tvm.relay.backend.compile_engine.get().clear() # just to be sure, see https://github.com/apache/incubator-tvm/pull/5724
with tvm.transform.PassContext(opt_level=3):
graph, lib, params = tvm.relay.build(mod_bert,
target=target,
target_host=target_host,
params=params_bert)
module = tvm.contrib.graph_runtime.create(graph, lib, ctx)
这会警告几次：
WARNING:autotvm:Cannot find config for ... batch_matmul.cuda .... A fallback configuration is used, which may bring great performance regression.
可能会带来很大的性能下降。
但是首先运行模型，看看输出是否匹配：
(8.583069e-06, 8.493662e-07)
看起来不错。记住，在float32中进行计算，因此$ 10 ^ {-6-6 $$ ish是一个很好的结果。
建立模型并设置参数后，按以下方式对模型进行计时：
def x():
for i in range(100):
module.run()
ctx.sync()
x()
%timeit x()
该模型每100次运行需要6.65s，或者模型运行67ms。确实很慢，警告说找不到（调整的）配置。调整调度。
调整确实需要半天左右的时间（在TVM调整，进行有关使用autotvm进行ResNet调整的介绍。）
此后，可以再次使用新配置来构建模型。这次应该看不到有关缺少配置的任何评论。现在，每次运行的时间约为6.5-7毫秒，类似于PyTorch。从最基本的算子优化中得到的结果。不过，可以将其进一步推进。
要了解如何操作，深入研究BERT建模和TVM。
如果不想获取完整的详细信息，请跳过下一部分并向下滚动到Results。补充一点，希望调优部分会在某种意义上过时，因为在不久的将来，可以立即获得改进，或者至少在进行一些初始调优后，速度会更好。因此，如果看不到这里和Results之间的加速，那是因为在提交补丁时做了功课。
BERT模型
让仔细看看BERT中的情况。
像许多深度学习模型一样，BERT带有一些前言（词汇嵌入）和结语（池化），并且大部分被组织成相似的块，这里有12个BertLayer模块。该attention_mask防止BERT在与问题打交道时看答案。
因此，放大并详细查看BertLayer，这最终是需要快速完成的工作。正如在网络图中看到的那样，BertLayer模块的主要部分是一个子模块BertSelfAttention。
现在，BertSelfAttention捕捉到了著名的自我关注机制，这是Transformers模型的标志。（不能推荐Sascha Rush的带注释的Transformers来作为详细的演练。）
BertLayer细节
如果要详细介绍，应该单独运行一个BertLayer。获取BertLayer的输入，然后BertLayer像对整个模型一样，将单个转换为TVM。
为了查看TVM模块，定义了一个小的可视化帮助器（基于TVM PR＃4370）。
import graphviz
def visualize(expr, collapse_small=True, node_attr_dict = {}):
def collect_ops(node):
ops = set()
def visitor(e):
if isinstance(e, tvm.ir.Op):
ops.add(e.name)
tvm.relay.analysis.post_order_visit(node, visitor)
return ops
# node_dict maps a Relay node to an index (node ID)
def _traverse_expr(node, node_dict):
if node in node_dict:
return
node_dict[node] = len(node_dict)
node_dict = {}
tvm.relay.analysis.post_order_visit(expr, lambda x: _traverse_expr(x, node_dict))
relayviz_nodes = []
dot = graphviz.Digraph(format='svg', )
dot.attr('node', shape = 'box')
def to_str(node):
if isinstance(node, tvm.relay.Constant):
return repr(node).lstrip('Constant(')[:-1]
else:
raise NotImplementedError("to_str:" + repr(node))
def is_small_const(c):
if not (collapse_small and isinstance(c, tvm.relay.Constant)):
return False
if isinstance(c.data, tvm.runtime.ndarray.NDArray):
return numpy.prod(c.data.shape) < 10
return True
# Sort by node ID
for node, node_id in sorted(node_dict.items(), key=lambda x: x[1]):
if isinstance(node, tvm.relay.Function):
dot.node(str(node_id), 'Function', **node_attr_dict.get(node, {}))
dot.edge(str(node_dict[node.body]), str(node_id))
elif isinstance(node, tvm.relay.Var):
if node.type_annotation is not None:
if hasattr(node.type_annotation, 'shape'):
shape = tuple([int(x) for x in node.type_annotation.shape])
dtype = node.type_annotation.dtype
typstr = 'Tensor[{}, {}]'.format(shape, dtype)
else:
typstr = str(node.type_annotation)
else:
typstr = '?'
d = dict(shape = 'ellipse')
d.update(node_attr_dict.get(node, {}))
dot.node(str(node_id),
'{}: {}'.format(
node.name_hint, typstr
), **d)
elif isinstance(node, tvm.relay.Tuple):
dot.node(str(node_id), 'Tuple[...])', **node_attr_dict.get(node, {}))
for field in node.fields:
dot.edge(str(node_dict[field]), str(node_id))
elif isinstance(node, tvm.relay.Constant):
if not is_small_const(node): # small consts are shown in ops
dot.node(str(node_id), 'Constant({}, {})'.format(node.data.shape, node.data.dtype),
**node_attr_dict.get(node, {}))
elif isinstance(node, tvm.relay.Call):
args_with_edge = []
arg_str_list = []
for arg in node.args:
if is_small_const(arg):
arg_str_list.append(to_str(arg))
else:
arg_str_list.append('·')
args_with_edge.append(arg)
arg_str = ', '.join(arg_str_list)
if isinstance(node.op, tvm.ir.Op):
name = node.op.name
attrs = {k:getattr(node.attrs, k) for k in node.attrs.keys()} if hasattr(node.attrs, 'keys') else {}
#attrs = inspect.getmembers(node.attrs)
attr_str_list = [k+'='+(str(v) if len(str(v))<20 else "...") for k, v in attrs.items()]
if attr_str_list:
attr_str = '| '+ ', '.join(attr_str_list)
else:
attr_str = ''
else:
ops = collect_ops(node)
if ops:
name = '_'.join(ops)
else:
name = '...'
attr_str = ''
s = f'{name}({arg_str}{attr_str})'
dot.node(str(node_id), s, **node_attr_dict.get(node, {}))
for arg in args_with_edge:
dot.edge(str(node_dict[arg]), str(node_id))
elif isinstance(node, tvm.ir.Op):
# dot.node(str(node_id), 'Op {}'.format(node.name))
pass # covered in call
elif isinstance(node, tvm.relay.TupleGetItem):
dot.node(str(node_id), 'TupleGetItem(idx={})'.format(node.index), **node_attr_dict.get(node, {}))
dot.edge(str(node_dict[node.tuple_value]), str(node_id))
elif isinstance(node, tvm.relay.Let):
dot.node(str(node_id), 'Let(XX)', **node_attr_dict.get(node, {}))
dot.edge(str(node_dict[node.value]), str(node_id))
dot.edge(str(node_id), str(node_dict[node.var]))
else:
raise RuntimeError(
'Unknown node type. node_id: {}, node: {}'.format(node_id, type(node)))
return dot
在主要功能上运行。由于某种原因（可能是完全笼统），PyTorchtransformers会将Linear图层转换为batch_matmul，而不是dense。由于TVMbatch_matmul在两个算子上都有scale（与PyTorch不同），也有很多转置算子。
visualize(mod['main'])
除了命名输入外，还看到许多未命名（编号）的变量。这些是神经网络参数。
编译模型
就像完整模型一样，可以在检查子模块计算出相同数量后运行并计时。
100次运行需要20.2毫秒。包络计算的背后是，BertLayer在PyTorch中，在这一层上花费了约0.2ms，在12层上花费了约2.4ms-不是大多数，而是整个6-7ms runtime中的一部分。与TVM进行比较。（一个好的规则是永远不要进行优化，也不进行测量。）
TVM的时钟runtime为18.2ms，可运行100次。再次与PyTorch大致相同。
从图片中看到的一件事是输入被重塑了3次。有一个TVM优化阶段调用，称为“公共子表达式消除”（CSE），结合了三种重塑形式。（前一阵子没有成功，具有不同的形状参数， TVM开发人员在动态到静态转换过程中解决了这个问题。），模型参数被重塑和转置了。也可以摆脱吗？是的。为此，将首先绑定参数，即将其放入模型中。然后，参数已变为常量，而不是输入节点。通过该Foldconstant过程，可以通过transposes和reshapes传播常数，以使更靠近matmuls。
在这三个之后（当编译中继模型时，TVM将执行此算子），模型如下所示：
将具有相同输入的三个批处理matmul合并为一个更有效batch_matmul。在TVM PR 5791中实施了此操作。还有另一个恒定折叠的传递。
new_mod = tvm.relay.transform.CombineParallelBatchMatmul()(new_mod)
new_mod = tvm.relay.transform.FoldConstant()(new_mod)
visualize(new_mod["main"])
经过检查，仍然得到相同的结果。可以再次计时：100次运行需要25.2毫秒。再次有点慢，需要调整新的形状。调整后，在100次运行中处于12.6ms，从大约0.2ms变为大约0.13-0.15ms，这是一个不错的加速。通过手工计算，这应该从总运行时间中减少0.6-0.8ms，或者说介于5％-10％之间，检查。
优化后对整个BERT模型的结果
定义一个函数，结合上面的优化过程，然后在整个BERT模型上运行。进行与上述相同的训练。
可以达到624毫秒进行100次运行。在PyTorch中从6.5-7ms缩短到了TVM中的6.2ms。这是5％-10％的加速。仅采用了特定的形状，而不是很大的形状。更多的分析，将考虑更多的问题形式。
可能会更进一步-例如，在批处理matmul之后，通过处理重塑来融合添加的内容，现在暂时将其保留。此外，还将受益于TVM的进一步改进，因此，随着时间的推移，基准会如何提高也会很有趣。特别地，即将到来的Ansor调整机制，似乎很有希望。
比较模型的实现
一直将PyTorch与TVM输出进行比较，查看是否良好。另外，当研究某个内层时，获取了该内层的输入，转换并输入到TVM模型中。这是一种非常有效的技术。
有时很难评估结果之间的偏差是由于数值精度，还是由于某处的误差。最初转换模型时，SelfAttentionTVM模型，将子模块输出复制到大约1e-6。但是，BertLayer转换的内容类似于1-e3。不确定是由于累积的数字误差，还是某些地方的材料偏差引起的。（原来是GELU激活，已转换为FastGELU。）
在这种情况下，想做的一件事，跳转到双精度并检查。数值误差应该变得更小，而其它偏差将保持相同的数量级。使用PyTorch前端，如果传递default_dtype="float64"给转换函数，可以在PyTorch端跟踪转换为float64的模型。
运行模块并与PyTorch进行比较，应该有1e-14左右的偏差。
TVM的改进
必须弥合一些差距（git checkout包括所有差距）：
- TVM PyTorchtransformers不支持fp32以外的输入。实施了改进的转换，也包括在TVM升级中。
- TVM调度（即main操作的计算组织，batch_matmul）是固定的，非常慢（类似于现在没有经过调整的调度运行）。因此，实施了可调度的调度表。
- PyTorchtransformers生成批量matmul操作（也可以将其更改为生成密集层）。较大的速度优势之一，将查询key和value线性层组合，实现了对批处理matmul操作的融合。
- 当比较计算结果时，注意到GELU函数已转换为其FastGELU变体。（TVM中有一个快速的数学优化过程，可以对误差函数进行一些替换，尽管没有检查，是否对用误差函数表示的GELU产生FastGELU。）
- TVM最初（并且在某种程度上仍然）专注于静态形式。尝试动态算子，动态重塑（以目标形式作为参数）是这些实验的早期阶段，由于通用子表达式消除过程未检测到可以合并相同的输入重塑，阻止了批处理的融合，情况有所改善。
- 使用TVM计算训练Pytorch模型
看看在PyTorch中训练BERT时，是否可以使用TVM。当需要处理自动分化时，这将打开全新的worms。从上面保留 theme，并以BertLayer示例为例，但方法通常代表着 non-trivial 模块。希望将训练期间的计算转移到TVM。
用户可以采用一个（可跟踪的）模块并执行
add_tvm_dispatch(module, sample_input)
如果用与sample_input形状相同的输入来调用模块，将获得TVM计算的输出（当然是PyTorch张量），否则，将使用常规正向。
如何实现这些任务。仍将无法获得很大的提速。
通过运行BertLayer从TransformersBert模型到的跟踪，得到了relay模型tvm.relay.frontend.from_pytorch。
在这之间，要做的一件事是将PyTorch中的模块化接口（带有命名参数）转移到功能接口（TVM可以为做的事情）。按照可以使用的顺序排列函数参数-即首先以直接的方式输入模块，然后与PyTorch使用相同的顺序来传递参数。完成此操作后，BertLayer 在TVM中的性能如下：
就像在BERT推理中一样，运行一些优化过程。
进行一些新的转换：
- Autodifferentiation的一个特殊之处，将使用大量..._like算子来广播或“取消广播”（总和是广播带有autodifferentiation的对偶）。现在有两个张量参数，后者实际上并不需要梯度。ZappLike将这些算子替换为带有shape参数的相应函数。
- 另一件事是派生类。TVM生成一个张量，所有张量都具有与函数的返回值相同的形状，以此作为链法则的起点。将这些乘以算子的派生产品。相乘并没有多大作用，类似地，TVM将变量（输入）的梯度初始化为相同形式的零。如果未使用，则渐变将为零，但如果使用，则“真实渐变”将添加到零。但是添加零也可以消除。这些由ZeroZapp和OneZapp负责。
- TVM没有针对LayerNorm（BatchNorm或其它）的训练变体。实现了通过以阐明计算的过程。
- TVM也没有训练停止。由于TVM当前没有随机数，在某种程度上很难解决。取而代之的是，用一个采用随机bernoulli抽取（值为0/1的值）的构造替换掉落对象，并以此模拟掉落对象。将使用PyTorch为生成版本。这样做还有一个好处，即（如果以与PyTorch相同的顺序生成滤除版本），将获得完全相同的结果。
正如上面所暗示的，TVM的梯度获取，假定是计算中的最后一个元素（上面讨论的“张量”）。这与PyTorch的模块化视图不太吻合，因为PyTorch希望grad_out每个输出都给出一个。令人高兴的是，这在计算上等效于乘以渐进求和，以此来修改函数。希望具有灵活性，允许两个函数都返回一个张量，而两个函数都返回一个张量元组。
应用这些修改后，模型如下所示：
随着let节点数量的增加，使用ToGraphNormalForm传递，将其恢复为正常形式。TVM的渐变获取返回一个函数，该函数具有与原始函数相同的参数（在本例中，使用grad_out和修改），然后返回原始返回值的元组和包含所有输入的梯度的元组。要做的第一件事就是放下所有的渐变色grad_out和dropout不需要。然后，运行简化流程。
因此，这是向前和向后的图表：
在PyTorch中，首先计算前向，然后计算后向，必须取出并拆分图形。困难的问题之一，如何处理为正向和反向计算的事物。这是一个与MinCut问题有关的难题。
极端选择可能是：
- 只能保留输入并根据需要重新计算所有内容。
- 如果有一个Salal输出，可以计算梯度并与后面的后一层的导数相乘。（损失函数可能会这样做。）但是，这不适用于非标量张量输出。
执行以下算子：通常计算正向，但保留所有将在向后使用。这很可能是没有看到端到端加速的原因。下面讨论一些潜在的启发式方法。
在这里使用一种颜色。首先，将正向计算的所有节点都涂成红色。然后，遍历梯度计算，然后从向后的蓝色为所需的节点着色。在可视化中展示属性支持。
一点（PyTorch）术语：当有一个功能Layer：x y，然后是一些Loss：y l∈ℝ时，向后是BackwardOfLayer：grad _out↦grad _in，grad _out = dl / dy和* grad _in = dl / dx`。
拆分功能，收集蓝色节点，进行捕获-但是常量将被复制，并且输入（Var节点）需要分别处理。现在可以拆分向后，用变量替换所有蓝色节点。
接下来，进行转发并进行修改，以返回所需的中间体。前向看起来像这样：
TVM无法返回嵌套元组，将函数中的输出展平。再次，区分张量值函数和元组值函数（那些可能返回多个张量的函数）。
最后，可以让TVM发挥其魔力并编译功能，对gr_only_compiled_module 和fw_and_cap_compiled_module，定义了便利函数，可以在PyTorch和TVM之间移动张量，并以TVM字典的形式获取模型参数。
def tensor_to_tvm(t):
return tvm.nd.from_dlpack(torch.utils.dlpack.to_dlpack(t))
def tensor_from_tvm(a):
return(torch.utils.dlpack.from_dlpack(a.to_dlpack()))
model_params_tvm = {k: tensor_to_tvm(v) for k, v in pytorch_model.state_dict().items()}
同样，在PyTorch和TVM中的GPU上获得输入。
记录的三个随机抽取的发生顺序与模型中顺序相同。在计算图上进行了深度优先搜索，如果dropout的值连接在图中而不是位于独立的分支上，也是PyTorch绘制矩阵的顺序。
torch.manual_seed(12345)
drop_c = {}
for k in dropout_info.keys(): # we don't know the order
p, typ = dropout_info[k]
drop_c[k] = torch.nn.functional.dropout(torch.ones([int(i) for i in typ.shape],
dtype=getattr(torch, typ.dtype), device="cuda"), p=p)*(1-p)
drop_tvm = {n: tensor_to_tvm(t) for n, t in drop_c.items()}
现在可以前进了。
fw_and_cap_compiled_module.set_input('input', inp_tvm[0])
fw_and_cap_compiled_module.set_input('attention_mask', inp_tvm[1])
fw_and_cap_compiled_module.set_input(**model_params_tvm)
fw_and_cap_compiled_module.set_input(**drop_tvm)
fw_and_cap_compiled_module.run()
可以将输出与PyTorch的输出进行比较：
torch.manual_seed(12345)
pytorch_model.train()
res = pytorch_model(*inp_c)[0]
numpy.abs(fw_and_cap_compiled_module.get_output(0).asnumpy()-res.detach().cpu().numpy()).max()
这给了2.1457672e-06。
非常好。让也尝试向后。生成一个grad_out，设置所有变量，并运行向后模型
gr_out_c = torch.randn(res.shape, device="cuda", dtype=res.dtype)
num_captures = len(capture_vars)
num_regular_outputs = len(fw_and_cap_fn_flattened.body.fields) - num_captures
captured_values = {v.name_hint: fw_and_cap_compiled_module.get_output(num_regular_outputs + i) for i, v in enumerate(capture_vars)}
gr_only_compiled_module.set_input(**drop_tvm)
gr_only_compiled_module.set_input(**model_params_tvm)
gr_only_compiled_module.set_input(**captured_values)
gr_only_compiled_module.set_input('gr:out:0', tensor_to_tvm(gr_out_c))
gr_only_compiled_module.run()
在PyTorch方面，最简单的方法是重新运行前向（记住要重置随机种子），并获取证书。
torch.manual_seed(12345)
pytorch_model.train()
inp_c_rq = [i.requires_grad_() for i in inp_c]
for p in pytorch_model.parameters():
p.requires_grad_()
res = pytorch_model(*inp_c_rq)[0]
grads_pt = torch.autograd.grad(res, inp_c_rq + list(pytorch_model.parameters()), gr_out_c, allow_unused=True)
奏效了吗？看来是这样的：
for i, g_pt in enumerate(grads_pt):
print(numpy.abs(gr_only_compiled_module.get_output(i).asnumpy() - g_pt.cpu().numpy()).max())
给列出了1e-5ish范围内的数字。
但是想在PyTorch中运行一些东西，对吗？
遵循PyTorch的工作方式，首先定义一个autograd.Function刚刚手动完成的操作：
在forward：
- 产生随机值，
- 向前运行，
- 记录向后所需的catch，输入和丢失值。
在中backward，向后运行并返回结果（作为PyTorch张量）。
这样，得到了一个PyTorch autograd.Function，调用TVM（需要一个小的package）。
需要做的add_tvm_dispatch(module, sample_inputs)就是跟踪模块，从中创建基于TVM的autograd函数，然后替换调用该函数的正向调用（使用参数）（如果适用）或退回常规方法。向前，Python的无限动力使这种相对容易。由于所有这些都不是真正与TVM相关的，提供了一些保留。
性能
就性能而言，并不是一个最终想要达到的目标。调整任务后（以及来自HuggingFace BERT + PyTorch JIT不太实际的推理示例），向前和向后运行了100次启用TVM的BertLayer迭代，类似于为推理所做的方式。一次迭代通过TVM花费6.2毫秒，而在PyTorch上花费1.3毫秒。
可以通过TVM运行模型。但这还不如通常的方法快。
更严重的是，有两条直接的途径来提高性能：
- 查找一组更好的捕获节点。
- 在TVM图上找到优化。
就前者的启发式（记住它很可能是NP困难的，没有得出正式的认证），人们会想重新进行廉价的计算，最主要的是逐点计算（或者除了matmul之外，什么都可以？）。
