OneFlow 如何做静态图的算子对齐任务

文章目录

1 前言
2 OneFlow 的 Graph 算子对齐概述
3 Graph 模式下自动测试实现原理
- 3.1 AutoTest 流程介绍
- 3.2 Graph 模式如何伴随 Eager 模式做算子对齐
- 3.3 Graph 模式的自动测试个性化
4 Graph 的 Debug 支持
5 总结
6 相关链接

1 前言

对于深度学习框架中模型的运行方式主要有两种，分别是动态图和静态图，动态图更易用，静态图性能更具优势，OneFlow 习惯将它们称为 Eager 模式和 Graph 模式。OneFlow 提供了 nn.Graph 模块，让用户可以用类似 Eager 模式的编程习惯，构建静态图训练测试。因此，需要同时保证 Eager 和 Graph 模式下算子行为和结果的正确性。

在之前的一篇文章中已经讲过深度学习框架如何优雅的做算子对齐任务，分析了 Eager Ops 的自动测试流程，包括如何产生随机数据测试用例和 AutoTest 核心代码实现，AutoTest 框架可以很轻易移植到其它深度学习框架使用。但是对于 Graph Ops 的的自动测试方法，还没有重点介绍，所以本文主要目的是介绍 OneFlow 如何完成 Graph 模式下算子的测试任务。到目前为止，OneFlow v0.7.0 已经新增所有 Op 在 nn.Graph 上做静态执行的单测支持，自动化单测功能完备。

文章中涉及到的代码位置：

https://github.com/Oneflow-Inc/oneflow/blob/master/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py
https://github.com/Oneflow-Inc/oneflow/blob/master/python/oneflow/test_utils/automated_test_util/generators.py

2 OneFlow 的 Graph 算子对齐概述

OneFlow 提供的 Eager 模式，用法与 PyTorch 对齐。所以在测试上，AutoTest 框架会随机出各种合法参数组合成的 Op ，并基于数值和类型完全相同的输入 Tensor（PyTorch 和 OneFlow 各有一份）分别运行 PyTorch 和 OneFlow 的代码，来完成算子的对齐工作。

此外，OneFlow 还提供了 Graph 模式，也基于面向对象的编程风格，让熟悉 Eager 开发的用户，只需要改很少的代码，就可以高效率的使用静态图。对比 Eager 模式，Graph 模式不易调试，但性能更好，易于优化和部署。那么，如何自动测试 Graph 模式下的 Ops 就是重点需要关注的问题。

在详细介绍 Graph 单侧之前，我们先看一下 AutoTest 框架里 Graph 打开方法，下面是一个测试 matmul 算子的例子。基于 random_pytorch_tensor 方法构造了两个随机的 tensor，它们的维度分别是 [n, k] 和 [k, m]，这些维度的值都是随机生成的，AutoTest 框架参数的随机性都是基于 generators.py 中的 generator 基类完成的。

    @autotest(check_graph = True)def test_flow_matmul_with_random_data(test_case):device = random_device()k = random(1, 6)x = random_tensor(ndim=2, dim1=k).to(device)y = random_tensor(ndim=2, dim0=k).to(device)z = torch.matmul(x, y)return z

通过调用 torch.matmul(x, y)，自动测试框架会分别运行 Torch 和 OneFlow 的 matmul 算子，会检查 Eager 模式下 OneFlow 和 PyTorch 算子的前向和反向结果是否一致。值得注意的是，代码中 @autotest 装饰器的 check_graph 开关为 True，表示此时会并行的做 Graph 的单测。

3 Graph 模式下自动测试实现原理

在了解了背景和使用方法之后，这里来讲解一下 Graph AutoTest 的实现思路。

3.1 AutoTest 流程介绍

在 Eager 的自动测试原理中，关于随机数据是如何产生的和 autotest() 装饰器的实现，不是本文的重点，在前置文章中有清晰的介绍。关于 AutoTest 框架核心流程实现，首先必须要关注的是 GetDualObject 函数，被用于 OneFlow 和 Pytorch 的算子对齐任务中。GetDualObject 这个函数会重写传入的原始 PyTorch 以及 OneFlow 对象的 __call__ 魔法函数，最后返回一个 DualObject 对象。这个过程中还包含跳过一些不需要关注的魔法函数；检查传入对象的属性是否合法；基于 nn.Module 和其它 API 默认参数的类型对 generator 继承类产生的随机数据绑定特定类型的工作（ get_args 函数中完成）；此外，在代码中还有对 tensor 方法的特判，因为 tensor 方法的调用方式（通过 getattr）和 nn.module 、 nn.functional 不同（通过 __call__）。

基于上述流程，通过执行样例代码中的 torch.matmul(x, y)，AutoTest 框架会通过调用 GetDualObject 函数生成 DualObject 对象，其中 torch 就可以理解为一个 DualObject 对象。最后执行 DualObject 对象，完成结果对比。自动测试流程中 Eager 算子对齐更多的细节也在前置文章中做了清晰介绍。

GetDualObject 函数实现在： https://github.com/Oneflow-Inc/oneflow/blob/7fe29cb3d24be41fa981c4ad6be3051dacc3b605/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py#L600 。
DualObject 类对象实现在：https://github.com/Oneflow-Inc/oneflow/blob/0826518cc49200dccada0f54d5c83accb9218c83/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py#L784

3.2 Graph 模式如何伴随 Eager 模式做算子对齐

从上面的分析中，可以大概总结出 AutoTest 的流程，生成随机数据、生成 DualObject 对象、执行 DualObject 对象和判断结果是否对齐。其中，AutoTest 框架在执行 DualObject 对象阶段，会并行的执行 OneFlow 算子的 Graph 版本，这样也就完成了 Graph 模式伴随 Eager 模式做算子对齐的任务。此外，本节也梳理了 GetDualObject 函数中应该如何识别需要静态（Graph）执行的对象。

在算子对齐任务中，存在 nn.module， nn.functional 和 tensor 方法三种类型。这里我们先以 nn.Module 类型为例，分析 Graph 模式伴随 Eager 模式测试的代码，其他三种类型处理方法基本一致。代码执行顺序如下图。

在 oneflow_eager_run_with_graph_check 中分别调用了 get_module_graph_test 和 get_oneflow_eager_res，得到 Graph 和 Eager 模式的两个结果，最后检查是否对齐。也就是说，对于一个测试 case，AutoTest 框架总共执行了三种代码，分别是 Pytorch、OneFlow Eager 模式和 OneFlow Graph 模式，来验证三种结果是否都对齐了。我们先来探究下 get_module_graph_test 这个接口，也就是如何得到 Graph 版本的计算结果。代码如下：

# NOTE(lixiang): When oneflow is of type nn.Module, build the following Graph for testing.
#   graph_train_oneflow: is a deepcopy of oneflow.
def get_module_graph_test(graph_train_oneflow, oneflow, *args):of_sgd = flow.optim.SGD(graph_train_oneflow.parameters(), lr=0.001, momentum=0.9,)graph_train_parameters_len = 0for param in oneflow._parameters.values():if param is not None:graph_train_parameters_len += 1class TestGraphOfModule(flow.nn.Graph):def __init__(self):super().__init__()self.test_module = graph_train_oneflowif global_backward and graph_train_parameters_len:self.add_optimizer(of_sgd)def build(self, *args):res = self.test_module(*args)forward_res = resif global_backward and graph_train_parameters_len:res = res.sum()res.backward()return forward_resreturn TestGraphOfModule()

其中 oneflow 是一个 nn.Module 对象，graph_train_oneflow 是它的深拷贝结果，主要为了防止在测试算子的 inplace 版本时，对相应的 DualObject 对象值进行了修改，造成 Graph 的输入和 Eager 不一致导致测试结果对不齐的情况。首先为了验证 Graph 的后向可以正常执行，构造了一个 Optimizer。在 __init__ 中复用 Eager 模式下的 nn.Module 对象后，在 build 中描述了 Graph 测试的计算过程，最终返回了 Graph 的实例。简单来说，就是构造一个适应所有算子的通用静态图模型。

在讨论如何构造静态执行代码计算 Graph 结果之后，识别需要静态执行的对象也是需要优先解决的问题。 oneflow_eager_run_with_graph_check 的完整代码如下：

# NOTE(lixiang): Check if the results of eager and graph are equal when oneflow is of type nn.Module or functional.
def oneflow_eager_run_with_graph_check(oneflow, oneflow_args, oneflow_kwargs, testing_graph, verbose, *args
):if testing_graph:graph_args, graph_kwargs = get_args_copy(oneflow_args, oneflow_kwargs)if isinstance(oneflow, flow.nn.Module):graph_train_oneflow = copy.deepcopy(oneflow)if not is_global():arg_device_type = "cpu"for arg in oneflow_args:if flow.is_tensor(arg):arg_device_type = arg.device.typegraph_train_oneflow = graph_train_oneflow.to(arg_device_type)else:graph_functional_oneflow = copy.deepcopy(oneflow)oneflow_res = get_oneflow_eager_res(oneflow, oneflow_args, oneflow_kwargs, verbose)if testing_graph:if verbose:print("After running eager module or functional: ", repr(oneflow),)find_check_module_func = Trueignore_apis_list = ["tensor", "train"]test_g_res = []if isinstance(oneflow, flow.nn.Module):test_g = get_module_graph_test(graph_train_oneflow, oneflow, *args)if verbose:print("Run graph of module: ", repr(oneflow))test_g.debug(3)# When testing module methods, kwargs are not considered.test_g_res = test_g(*graph_args)if verbose:print("The result after running graph module: ", test_g_res,)elif oneflow.__name__ in ignore_apis_list:find_check_module_func = False# 1. "oneflow.nn.modules" not in oneflow.__module__: For avoid run nn.Module branch graph test, like fold op call Fold Module actually.# 2. inspect.isfunction(oneflow): Compared with the ordinary flow.xxx, oneflow.nn.modules.math_ops series op exist an extra layer of python wrapper.# 3. inspect.ismethod(oneflow) and "oneflow.nn.modules" in oneflow.__module__:  For op that only has Tensor.xxx method, and call oneflow.xxx actually, like masked_fill.elif (("oneflow.nn.modules" not in oneflow.__module__)or inspect.isfunction(oneflow)or (inspect.ismethod(oneflow) and "oneflow.nn.modules" in oneflow.__module__)):test_g_res = get_functional_graph_res(graph_functional_oneflow,oneflow,oneflow_res,oneflow_args,oneflow_kwargs,verbose,*graph_args,**graph_kwargs,)if find_check_module_func:if isinstance(test_g_res, tuple):for _, g_res in enumerate(test_g_res):check_eager_graph_tensor(oneflow_res, g_res)else:check_eager_graph_tensor(oneflow_res, test_g_res)return oneflow_res

在 oneflow_eager_run_with_graph_check 中，需要判断哪些对象需要静态执行测试。因为 OneFlow 设定部分代码需要静态化，比如有些 Eager 模式下的方法，在 Graph 模式下没有定义。上面的代码中首先通过 if testing_graph: 判断是否打开了 Graph 开关，既是否需要并行的做 Graph 的单测；再对 oneflow 对象的类型做 isinstance 判断，当为 nn.Module 时才需要静态执行，调用 get_module_graph_test。否则调用 get_functional_graph_res 等处理，在测试框架中其他需要类似判断的地方也同理。

    if testing_graph:······if isinstance(oneflow, flow.nn.Module):···test_g = get_module_graph_test(graph_train_oneflow, oneflow, *args)···elif:······

3.3 Graph 模式的自动测试个性化

在 3.2 介绍了 Graph 如何伴随 Eager 模式做算子对齐的任务之后，本节主要分析 Graph 模式自动测试的个性化内容。

在 Graph 模式下，需要处理 nn.module、nn.functional 和 tensor 三个类别的方法，AutoTest 框架采用先判断后构图的方式。首先，GetDualObject 函数中，相关的接口包括： get_pytorch_oneflow_res、 get_pytorch_oneflow_tensor_res、 oneflow_eager_run_with_graph_check、 oneflow_tensor_eager_run_with_graph_check、 get_oneflow_eager_res、 get_tensor_graph_res、 get_functional_graph_res 和 get_module_graph_test 。看一下每个接口的功能，如下表。

Function	Info
`get_module_graph_test`	返回 `nn.module` 模块中算子的 Graph 实例。
`get_functional_graph_res`	返回 `nn.functional` 模块中算子的 Graph 计算结果。
`get_tensor_graph_res`	返回 `tensor` 模块中算子的 Graph 计算结果。
`get_pytorch_oneflow_res`	分别获得 OneFlow 和 Torch 的计算结果。
`get_pytorch_oneflow_tensor_res`	tensor Ops 的特化，同上。
`oneflow_eager_run_with_graph_check`	测试 eager 模式是否对齐，伴随Graph检查。
`oneflow_tensor_eager_run_with_graph_check`	tensor Ops 的特化，同上。
`get_oneflow_eager_res`	获得 OneFlow Eager 模式算子的计算结果。

大概了解了每个函数功能之后，再来看一下调用链，如下流程图所示，图中包含了 AutoTest 框架中对于 Graph 模式存在 nn.module、nn.functional 和 tensor 三个类别的方法如何处理，对应图中的三个灰色框。

在分析了 nn.module、nn.functional 和 tensor 三个类别的处理方法之后，其中，自动测试 Graph 时也存在一个反向的梯度测试，但是并没有取出 tensor 对应的梯度，也就是说，可以保证后向执行是正常的，没有检查 grad 值。对于使用方法，当 @autotest() 打开 auto_backward=True 时（默认就是打开的），不仅会跑 Eager 的 Backward 测试（这里会对梯度结果做比较），还会跑对应 Graph 的 Backward测试（这里不做梯度比较）。

对应上述描述的代码，可以在文章 3.2 部分的代码中找到：

if (global_backwardand graph_train_parameters_len
):self.add_optimizer(of_sgd)
···
···
···
if (global_backwardand graph_train_parameters_len
):res = res.sum()res.backward()

此外，对于一些算子 inplace 版本的 Graph 检查，需要对输入做深拷贝，来保证 Graph 和 Eager 的 input 始终一致。如下代码中，get_args_copy（在 torch_flow_dual_object.py 中）分别对普通参数和关键字参数做了 deepcopy。类似的，在 Graph 单侧中，存在 oneflow 深拷贝为 graph_train_oneflow 的行为，主要为了防止在测试一些算子时，Eager 的值被 Eager Inplace 修改，造成 Graph 的输入和 Eager 不一致导致测试出错的情况。

# NOTE(lixiang): Deepcopy the input parameters in order to correctly test the inplace version of the op.
def get_args_copy(args, kwargs):copy_args = []for arg in args:if flow.is_tensor(arg):copy_arg = arg.clone().detach()else:copy_arg = copy.deepcopy(arg)copy_args.append(copy_arg)copy_kwargs = {}for key, value in kwargs.items():if flow.is_tensor(value):copy_kwargs[key] = value.clone().detach()else:copy_kwargs[key] = copy.deepcopy(value)return copy_args, copy_kwargs

最后，为了保证 tensor deepcopy 的正确性，在 OneFlow 中，copy.deepcopy 会调用 tensor 的 getState 和 setState 方法，tensor 的 state 需要同时包括 data、dtype 和 device 信息，缺一不可。具体代码见： https://github.com/Oneflow-Inc/oneflow/blob/e00ba51364ff87e39edc409be395e5ed493a4ac0/python/oneflow/framework/check_point_v2.py#L159 。

4 Graph 的 Debug 支持

在 3.2 的代码中，可以发现存在 if verbose: 的判断，当 verbose = True 时，会输出 Graph 的 debug 信息（如算子运行 Graph 模式后的计算结果等），当然也包括 eager 下的其他需要的调试信息。当测试出现问题时，可以通过该功能拿到错误样例，构造最小复现代码。开启方法通过环境变量控制：ONEFLOW_TEST_VERBOSE = 1。AutoTest 框架里这个功能更多针对开发者，OneFlow 的 Graph 针对用户也提供了调试功能。

Graph 模式支持了学习率的调试输出，开启方法和 Eager 相同。

optimizer = flow.optim.SGD(model.parameters(), lr=1e-3)
# Set verbose=True
scheduler = flow.optim.lr_scheduler.CosineDecayLR(optimizer, decay_steps=100, alpha=0.98, verbose=True)

此外，调用 Graph 对象的 debug 方法，就开启了 Graph 的调试模式。

graph.debug(v_level = 1) # 可以简写为：graph.debug(1)

v_level=0 时，只输出最基础的警告和构图阶段信息，如构图时间。
v_level=1 时，将额外打印每个 nn.Module 的构图信息。
v_level=2 时，在构图阶段，将额外打印每个 Op 的创建信息，包括名称、输入内容、设备和 SBP 信息等。
v_level=3 时，将额外打印每个 Op 更详细的信息，如与代码位置有关的信息，方便定位代码问题。

这部分更详细的内容可以在 https://docs.oneflow.org/master/basics/08_nn_graph.html#graph_3 中发现。

5 总结

AutoTest 框架的灵活性和易用性都比较强，本文主要介绍了 Graph 模式如何伴随 Eager 模式做算子对齐和 Graph 的自动测试个性化内容。Eager 到 Graph 的 Local ops 执行测试覆盖也已经在 OneFlow v0.7.0 中完成，在 0.8 版本中将会保证 Graph Global ops 单测的正确性。此外，静态图的 debug 和其他功能等将更完备。欢迎大家学习或者使用。

6 相关链接

https://github.com/Oneflow-Inc/oneflow
https://github.com/pytorch/pytorch