桥接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之外，什么都可以？）。