AI推理与Compiler

AI推理与Compiler

AI芯片编译器能加深对AI的理解， AI芯片编译器不光涉及编译器知识，还涉及AI芯片架构和并行计算如OpenCL/Cuda等。如果从深度学习平台获得IR输入，还需要了解深度学习平台如Tensorflow、TVM等。

编译器领域的知识本身就非常艰深，和AI模型本身的关系也不是特别紧密，很难将AI建模作为发展方向，可以多关注GPGPU Architecture。即使AI芯片过气了，GPGPU还是会长盛不衰。

OneFlow是有其独特的设计理念和技术路线的。目前市面上已有的众多开源框架，用户最多的是PyTorch和TensorFlow，除此之外还有各大公司的自研框架：PaddlePaddle、MXNet、MindSpore、MegEngine等等。其中TensorFlow突出的特点是最完备，推理、serving、XLA、tensorboard等一应俱全，工业部署的主流还是TensorFlow；PyTorch的特点是易用性最好，eager执行、动态图、跟python的交互方式等等。

本文以 Batch Norm 为例介绍了推理计算图的具体实现，以及 MatMul 在 CPU 上的优化细节。作为 CPU 推理优化的基石，最优的推理计算图是实现高性能 CPU 推理的前提条件，极致性能的 MatMul 计算基础算子将为实现卷积计算中的 Im2col 和 Winograd 提供性能保障。
概述

天元（MegEngine）深度学习框架凭借「训练与推理一体化」的独特范式，能够极大程度上（90%）节省模型从研发到部署的整体成本，降低转换难度，真正实现小时级转化；同时，天元（MegEngine）在 CPU 推理方面所做的大量优化工作，也使得开发者在推理时能够发挥出处理器的最佳性能。

CPU 推理

CPU 推理的性能优化至关重要，经过优化的推理性能可以较未经优化的原始性能提升数十倍。

 

 

 在进行模型推理阶段的优化时，首先需要对模型的计算图进行优化，以避免冗余的计算与访存，确保计算图在推理时是最优的；其次，在大多 CV 相关模型中，卷积的计算比重最高，达到模型总计算量 90% 以上，因此对模型推理的优化主要聚焦在对卷积计算的优化上。

卷积计算的实现方式有 3 种：direct 卷积，Im2col 卷积以及 Winograd 卷积。

Ÿ direct 卷积：根据卷积的计算公式直接对 FeatureMap 上进行滑窗计算；

Ÿ Im2col 卷积：根据卷积计算需要在输入通道上进行 reduce sum 的特点，将卷积运行转化为 MatMul 计算；

Ÿ winograd：在保证计算无误的前提下，使用加法替代乘法，达到优化卷积乘法计算量的目的，在中间过程需要使用 MatMul 进行计算。

以 Im2col 或 Winograd 方法进行的卷积计算，频繁使用到 MatMul，对 MatMul基础算子的优化尤为重要。

本文将首先介绍模型优化中的图优化，然后介绍基础算子 MatMul 在 CPU 上的优化方法。

推理计算图优化

在训练阶段定义模型的计算图，主要是为了满足模型参数的训练需求。当训练结束，模型参数固定后，对计算图的进一步优化能够帮助模型推理更加高效。

推理和训练的计算图是一张有向无环图(DAG)，在天元中，开发者能够以类似 LLVM 的方式对 DAG 计算图定义许多优化方法，这里简称 OptPass。OptPass 可以根据用户的配置有选择性的加入到图优化的 OptPass 列表中，从而帮助用户灵活地为图优化定义 OptPass。 

计算图优化

天元定义了多个为推理进行计算图优化的 OptPass，开发者使用这些 OptPass 之后，将得到一张用于推理的最优计算图。下面以 MobileNetV1 中， Convolution+Batch Norm+Relu 这样的典型结构经过计算图优化之后 Fuse 为 ConvBias 的过程为例，介绍天元的计算图优化过程。

由于 Batch Norm 在推理阶段除了输入 Tensor 外都是常数，可以简化为多个 elemwise 的组合，天元实现了一个 ConvertBatchNormToElemwisePass 的 Optpass，这个 OptPass 将模型中的所有 Batch Norm 转化为 Elemwise，具体转化原理如下：

 

 

 紧接这一过程，天元将运行 OptPass ParamRedistributePass，该 OptPass 会将上述 Batch Norm 转化而来的 Elemwise 中的 scale 融合到 convolution 的权重中，具体实现原理如下：

 

 

 天元将运行 OptPass FuseConvBiasNonlinPass，该 OptPass 会将计算图中的 Convolution+Elemwise 转化为天元内部实现的 ConvBias Op 中，还会设置 ConvBias Op 中 NonelineaMode 参数。

如此便完成了从 Convolution+Batch Norm+Relu 到 ConvBias 的转换，整体转换过程如下图：

 

 

 实验验证图优化之后的性能

在推理之前，如果要对完成训练的模型进行图优化，则需要在模型 dump 的期间进行，当然，也可以在 SDK 运行模型之前进行。下面是在模型 dump 时进行图优化的代码：

from megengine.jit import trace@trace(symbolic=True)def pred_fun(data, *, net):    net.eval()    pred = net(data)    pred_normalized = F.softmax(pred)    return pred_normalized    # 使用 trace 类的 trace 接口无需运行直接编译pred_fun.trace(data, net=xor_net)# 使用 trace 类的 dump 接口进行部署pred_fun.dump("xornet_deploy.mge", arg_names=["data"], optimize_for_inference=True, enable_fuse_conv_bias_nonlinearity=True)

复制代码

上面的 optimize_for_inference=True 将在 dump 模型时候针对 inference 进行优化， enable_fuse_conv_bias_nonlinearity=True 将在模型中进行 Op Fuse 的优化，此外天元还支持其优化参数，具体可见天元的文档。

下表展示的，是在实验过程中对模型进行图优化前、后，模型运行性能的测试对比。

 

 

 图优化对模型性能的提升效果显著，具体提升比例由模型自身决定。

MatMul 优化

MatMul 作为卷积运算的基础算子，会频繁地被 Im2col、Winograd 以及 FullyConnect 使用。在天元中，MatMul 既被封装为单独的 Op，也可以作为单独的 algo 供卷积的实现使用。

MatMul 也是计算最为密集的算子，因此天元进行了极致的优化。

优化

MatMul 是线性代数中的矩阵乘，假设矩阵 A 大小为 MK，矩阵 B 大小为 KN，则得到矩阵 C 大小为 M*N，其中 C 的每个元素的计算公式如下：

 

 

 在 MatMul 的计算中乘法和加法的计算量为 2MNK (计算 C 中每个元素时，加法和乘法计算量分别为 K，C 的总元素个数为 MN)，访存量为 2MNK (计算每个 C 中元素需要 2K 访存)+ 2MN(整个 C 矩阵读一次和写一次)。由于计算量固定(排除 Strassen)，所以只能优化访存，使得乘法和加法运算达到处理器的极限性能，从而实现 MatMul 的最佳性能。

MatMul 分块

关于减少 MatMul 计算时的访存量，最有效的方法是对 MatMul 的计算进行分块，并将分块之后的数据保存在 CPU 的 Cache 中。

如下图所示，将 A 按照 mr 进行行分块，将 B 按照 nr 进行列分块，计算时将需要用到的分块保存在 CPU 的各级 Cache 中，从而极大的减少了内存的读写。

 

 

 进一步细化上面的分块计算过程如下图所示，A 中的一个行块都要重复地和 B 中的每一个列块进行小分块的 MatMul，写入到 C 的小块中，为了使得分块尽量的大 (如上所述，能够减少内存访存量)，Cache miss 率尽量低，因此需要根据 CPU 的 Cache 结构特点(速度：L1D>L2>L3，容量 L1D<L2<L3) 来分配 Cache 和数据块之间的存放关系，天元中进行的分配如下：

· 将下图中红色 (访问重复次数最多的 A 的行块，计算时需要的 B 的一个列块以及计算结果的 C 的小块) 部分都保存在 L1 中。

· 由于计算完每一个 C 的行块，都需要重复遍历一次整个 B 矩阵，因此将 B 存放在 L2 中使得每次读取 B 的一个列块都是从 L2 中读取，代价最小。

· 将重复访问率最高的 C 的累加中间结果保存在速度最快的 CPU 处理器的寄存器中。

 

 

 通过上面的分配策略，并结合 CPU 中资源(寄存器数量，L1D 和 L2 的大小)，便可以确定最佳的 MatMul 计算中的 Nr，Kr：

 · 可以根据 CPU 处理器的寄存器数量得到 mr 和 nr 的具体大小，寄存器容量 > mr*nr

· 根据 L1D Cache 的大小结合 mr 和 nr 计算出 Kr，Kr=L1D/(mr+nr)

· 再根据 L2 的大小计算出 B 矩阵中的 Nr，Nr=(L2-L1D)/Kr

在上面计算 N 时，用 L2-L1D 的原因是，由于当前 CPU 使用的 Cache 是 Inclusive 的，且 L2 是指令缓存和数据缓存合并的，另外上面 M 没有明确限制。

在得到上面最佳 Nr、Kr、mr 和 nr 之后，进一步便可以首先对 MatMul 计算中的 N、K 进行 Nr 和 Kr 分块，然后在 Nr、Kr 的基础上再进行 mr 和 nr 分块。如此，MatMul 计算中的计算访存比达到最高，且 CPU 处理器的资源也得到充分利用。

经过分块之后，由于在最内层的计算 Kernel 为 A(mrKr)B(Krnr)=C(mrnr)的分块矩阵乘，决定了整个 MatMul 的计算性能，因此需要极致地优化。

Kernel 计算

最核心的计算 Kernel 进行的是尺寸为 (mrKr)x(Krnr)–>(mr*nr) 的小尺寸 MatMul，其计算示意图如下：

 

 

 如上图所示，Kernel 在计算时会读取 A 中一列， B 中一行，进行矩阵乘，得到 大小为 mr*nr 的 C，然后和原来 C 中的值相加，如此循环 Kr 次，完成该 Kernel 的计算。

在该 Kernel 的计算过程中乘法和加法的计算量为 2mrnrKr，访存量为 (mr+nr)Kr+2mrnr，可以根据处理器来判断该计算能否隐藏数据的访存。下面以 ARM cortex A76 为例进行分析，根据 A76 的数据手册得到：

· FP32 SIMD Load throughput=2，即单周期可以 load 8 个 float 数据

· FP32 SIMD FMLA throughput=2，即单周期可以进行 16 个乘加运算

当 Kernel 的尺寸 mr=8、nr=12、Kr=256，计算量为 49152 次乘加运算，访存量为 5312 个 float 数据时，该计算访存量为 9.25，大于处理器的计算访存比 2。因此可以得出结论，如果 A 和 B 均在 L1 中，则该 Kernel 的计算不会因为数据的 Load 被阻塞，所以计算单元能够发挥出处理器的最佳性能。

虽然在对 MatMul 进行分块时，已经计划将 A 和 B 这样的小矩阵置于 L1D Cache 中，但是数据在真正运行时却不一定都在 L1D 中，原因在于，B 矩阵的列块在原来的大矩阵中内存并不连续，其次 A 中的一列由于内存不连续，也不能使用 SIMD 进行 load，为此需要对 A 和 B 进行数据 PACK。

MatMul 数据 PACK

上文 kernel 在计算过程中，需要同时读取 A 矩阵 mr 行数据，而每行数据之间内存不连续，因此这种访问对 Cache 不友好，同样，在读取矩阵 B 中 nr 列的时候也存在数据读取不连续的问题，加之 A 的所有行块和 B 中的所有列块将被读取多次，因此可以通过对 A 和 B 提前进行数据 PACK， 实现在后续计算中频繁多次访问数据时是连续的（只在第一次 PACK 时对 Cache 不友好），进而获得巨大收益。

 

 

 对矩阵 A 进行数据 PACK 是将 A 中 mr 行数据的相同列拷贝到一起，如上图中将 A PACK 到 A’ 的步骤。重复完所有 A 中的行块便完成了 A 矩阵的数据 PACK。B 矩阵的 PACK 操作是，将 nr 列数据拷贝到连续的内存地址中，对应上图 B PACK 到 B’ 的过程。

实 验

按照文介绍方式方式，天元在 X86 和 ARM 上分别对 MatMul 进行了优化。下表展示了 ARM64 上的性能测试结果，实验平台为 kirin 980。

首先，对该处理器进行分析可以看到，其主频为 2.6 GHz，每个周期能够进行 16 次乘加计算，因此其理论计算峰值为 16*2.6=41.6 Gflops。

 

 

 可以看到，经天元优化的 MatMul 计算，发挥出了该处理器 90% 以上的计算性能。

OneFlow

1. OneFlow的设计思路，独特性
2. OneFlow的特色一：Actor机制
3. OneFlow的特色二：SBP机制
4. 对人工智能/深度学习未来的看法
5. 对开源的感想和总结

OneFlow是独特的

可以说完备性和易用性这两极分别被TF和PyTorch占据，OneFlow作为一个小团队打造的框架，如果单纯的模仿别的框架，跟TF比完备性，跟PyTorch比易用性，那是绝对没戏的。

但深度学习框架还有第三极：性能。OneFlow的特点就是追求极致的性能，而且是分布式多机多卡环境下的横向扩展性。OneFlow的核心设计理念就是从分布式的性能角度出发，打造一个多机多卡体验就像使用单卡一样容易，而且速度最快的深度学习框架。

深度学习是吞没算力的巨兽。多机多卡的理想很丰满，现实很骨感，用户会遇到多机多卡加速比不高，得不偿失，用户会遇到参数梁巨大时，现有框架不支持模型并行而无法训练。为解决此类问题，业界不仅仅改进深度学习框架自身，还研发了多种第三方插件，譬如NCCL, Horovod, BytePS，HugeCTR，Mesh-tensorflow, Gpipe等等。但是，仍只满足一小部分需求。

分别介绍两点OneFlow相较于其它框架的独特设计，来体现OneFlow是如何看待分布式场景下的深度学习训练的。

Actor——用一套简洁的机制解决所有令人头秃的技术难题

（关键特性：去中心化调度流水线数据搬运是一等公民传输被计算掩盖控制逻辑被执行逻辑掩盖）

在OneFlow的设计中，分为Compiler和Runtime两个时期，Compiler把用户定义的网络、分布式环境信息等编译成一个静态图的中间表示（称之为Plan）；Runtime时期，各个机器根据Plan里的Actor描述信息真实地创建属于机器的众多Actor，整个深度学习训练期间，OneFlow的执行基本单元就是Actor，Actor之间通过消息机制通信，静态执行图上的节点就是Actor，节点之间的连边是Register，Register存储了Actor之间的生产者消费者信息，以及生产者Actor产生的数据。

1、Actor机制实现去中心化调度

OneFlow的运行时去中心化调度就是用Actor机制实现的。在整个由Actor构成的静态图中，没有一个中心的调度节点，每个Actor都只需要关心消费的那些Actor，和消费生产出的数据的那些Actor即可。这样在超大规模的分布式训练场景下，完全的去中心化调度可以避免中心调度的单点性能瓶颈问题。

每个Actor内部都有一个状态机，Actor之间的收发消息和Actor内部执行都会改变状态。当Actor收到了执行所需要读取的Register，且当前有空闲的Register可以生产的时候，这个Actor就执行（Act）一次，生产出一个Register。生产完以后，该Actor就向需要消费这个Register的那些消费者Actor们发消息，表示可以来读取产生的数据了；同时该Actor还需要把消费完的那些Register还给这些Regsiter的生产者Actor们，表示用完了数据，可以回收了。Actor内部的状态机如图4所示。

 

2、Actor机制实现流水线

上面介绍了Actor的内部状态机，Actor之间的消息传递和数据传递是依赖Register实现的。一个Actor是否能执行，只跟两个条件相关：1）消费的那些Register是否可读；2）生产的那些Register是否有空闲块可写。对于一个Register，如果运行时分配多个空闲块，那么相邻的两个Actor就可以同时工作，重叠起来，这样就实现了各个Actor之间的流水线。理想状态下整个静态执行图的执行时间就是整个系统中是性能瓶颈的那个Actor运行的总时间，其余Actor的执行时间就被流水线掩盖起来了。

举一个例子来解释Actor机制下的流水线是如何运转起来的。图2是一个由3个Actor（a,b,c）组成的静态图的时序图。其中深绿色的Regst方块表示正在被使用的Register块，白色的Regst方块表示同一个Register的备用空闲块。

1）在Time0时刻，Actor a产出了一个Regst_a_0，Actor b 和 Actor c由于没有可读的Register，所以处在等待状态。假设每个Actor的执行时间都是单位时间。

2）到Time1时刻，Actor a给Actor b发消息说可以来读产出的Regst_a_0了，Actor b收到了消息，并检查生产的Register b是否有空闲Regst块可用，发现有可用的Regst_b_0，于是Time1时刻Actor b执行，读取Regst_a_0，写Regst_b_0；同时Actor a还会去看是否有空闲块可写，发现有，Time1时刻Actor a也在执行，写Regst_a_1。（这里需要说明的是，Regst_a_0和Regst_a_1逻辑上是属于同一个Register，只是空间上分成了不同的空闲块备份而已。在深度学习训练任务中，Regst_a_0和Regst_a_1里存放的是同一个operator产出的不同batch的数据。）于是Actor a和Actor b就并行工作起来了。Actor c由于没有数据可读，仍在等待。

3）到Time2时刻，Actor b 生产出了Regst_b_0，于是给下游的消费者Actor c发消息说可以来读生产的Regst_b_0，同时给上游的生产者Actor a发消息说用完了的Regst_a_0。此时Actor a 已经把刚刚生产的Regst_a_1又发给了Actor b，Actor b检查仍有Regst_b_1空闲，于是Actor b开始读Regst_a_1，写Regst_b_1；Actor c收到Regst_b_0，发现有Regst_c_0空闲，于是Actor c开始执行，读Regst_b_0, 写Regst_c_0；Actor a收到了Actor b用完还回来的Regst_a_0，检查Regst_a_0所有的消费者都用完了，于是将Regst_a_0回收，标记为空闲块，同时Actor a还可以继续执行，写Regst_a_2。

在上面的例子中，到了Time2时刻，其实Actor a、b、c都在工作，在深度学习训练任务中，Time2时刻Regst_b_0、Regst_c_0存放的是Batch 0的数据，Regst_a_1、Regst_b_1存放的是Batch 1的数据，Regst_a_2存放的是Batch 2的数据。通过一个Register有多个空闲块的设计，Actor机制就实现了流水并行。

 

3、数据搬运

在多机多卡的分布式环境中，各个机器和各个设备之间的数据传输往往是影响系统横向扩展性的最重要因素，如果传输开销可以被计算开销掩盖，那么分布式深度学习训练就可以达到理想的线性加速比。相较于其框架，OneFlow把数据搬运视为跟数据计算同等地位的操作，提出“数据搬运是一等公民”的思想。

已有框架在编译期的关注焦点是数据计算，数据搬运是背后隐式发生的，在静态分析计算图时略过计算和搬运的重叠编排，OneFlow在计算图中显式表达了数据搬运，而且在静态分析时同等对待数据搬运和数据计算，以最大化重叠搬运和计算。

在OneFlow最终的执行图中，数据搬运操作也是一个个Actor。除了在设备上做数据计算用的Actor以外，还有计算机内存到GPU显存之间的数据拷贝Actor，机器之间做网络通信的网络Actor，负责数据的切分、合并、复制的Actor，负责读取磁盘数据的Actor，负责加载保存模型的Actor等等。很多其框架都把数据加载、多卡模型梯度的同步、网络、模型加载更新等分别做成一个单独的模块，而OneFlow的设计是所有的功能都在一张由Actor组成的静态执行图里实现了。OneFlow这样的设计不仅简洁、优雅，还非常高效。

 

图3表示了没有GPU-Direct的情况下，在OneFlow的Runtime阶段，一个设备上的计算节点如果消费了另一个设备的计算节点，数据是如何搬运过去的。

4、尽可能并行

在OneFlow的设计中，所有的出发点都是希望可以尽可能并行，从而达到最优的分布式性能。比如考虑到分布式训练模型梯度同步时，显存到内存的传输带宽高于机器之间的网络传输带宽，OneFlow会做两级的scatter和gather操作（本机的和各个机器之间的），用于增加locality，提高整体性能；又比如在异步启动深度学习训练时，python端用户的控制逻辑跟OneFlow运行时的执行图是并行执行的，同时OneFlow有一套互斥临界区的设计保证执行的高效性和正确性；数据加载部分无论是从磁盘读数据还是从python端喂数据，OneFlow都能保证尽可能并行，使得计算设备不会因为要等数据而导致性能下降。

已有框架要尽可能重叠数据搬运和计算，一般借助多层回调（callback）函数，当嵌套层次过多时，会遇到所谓的callback hell麻烦，正确性和可读性都可能下降。但在OneFlow中，以上的这些并行并发特性，都是在这一套简洁的Actor机制下实现的，解决了令人头秃的callback hell问题。

此外，在多机的网络通信部分，OneFlow底层的网络通信库原生支持RDMA的高性能通信，也有一套基于epoll的高效通信设计。而相比于最流行的Pytorch，多机还需要通过RPC来做数据同步。

Placement+SBP——OneFlow如何做到分布式最易用

OneFlow是目前分布式场景中支持数据并行、模型并行、流水并行等最易用的深度学习框架。用户只需要像单卡一样去搭建网络模型，并告诉OneFlow有哪些机器哪些卡，OneFlow就会用最高效的方式把这些机器和设备使用起来。

这源于OneFlow的一套独特的设计：ConsistentView（一致性视角）。对于多机多卡，OneFlow会抽象成一个超级大的设备，称之为逻辑上的设备，这个逻辑设备的显存是实际多个物理设备的显存之和，这个逻辑设备的算力也是实际多个物理设备的算力之和。用户只需要定义在这个逻辑上的超级设备里，深度学习模型是如何构建的，其余的便不需要用户来操作，由OneFlow来完成逻辑上的设备到物理上的设备的映射。

这里先明确两个概念：“逻辑上的”和“物理上的”。“逻辑上的”表示OneFlow把分布式集群抽象成一个超级计算机之后的计算和数据，“物理上的”表示那些真实的部署到各个机器和设备上的计算和数据。深度学习网络是由Op构成的计算图，Op之间生产消费Tensor数据。在多机多卡的环境下，一个逻辑上的Op会对应多个真实的物理上的Op，每个物理上的Op实际执行的计算都是这个逻辑Op计算的一部分，一个逻辑上的Tensor也会对应多个物理上的Tensor，每个物理上的Tensor都是逻辑Tensor的一部分。

对于其框架定义的分布式训练，每张卡是一个“world”，多卡之间根据暴露出来的接口来同步模型梯度；而对于OneFlow而言，多机多卡也都是一个“world”，使用一套Placement+SBP的方式做全局的统筹管理。

1、Placement

在OneFlow的计算图搭建过程中，每个计算Op都有一个属性叫做Placement，表示了该逻辑上的Op，是要部署到哪些机器哪些设备上的。对于常见的数据并行，就是所有的Op都部署到所有的设备上。但OneFlow也支持用户指定Op的Placement，比如当网络过大单卡根本放不下的时候，在OneFlow可以让网络的前一部分在一张卡上，后一部分在另一张卡上，用一种“接力”的方式工作，实现流水并行。

图4展示了一种可能的Placement例子：

 

2、SBP

SBP是OneFlow独有的概念，他是三个单词的首字母组合：Split、Broadcast、PartialSum（以PartialSum为例，实际上还可以是PartialMin, PartialMax等reduce操作），全称叫SbpParallel，表示一种逻辑上的Tensor跟物理上的多个Tensor的映射关系。

其中Split表示物理上的Tensor是逻辑Tensor按照某一维度切分后得到的， Split有个参数axis，表示切分的维度，如果把多个物理上的Tensor按照Split的维度进行拼接，就能还原出逻辑Tensor；Broadcast表示物理上的Tensor是跟逻辑上的Tensor完全相同的；PartialSum表示物理上的Tensor虽然跟逻辑上的Tensor形状一致，但是物理上的Tensor里的值是逻辑Tensor里对应位置的一部分，如果把物理上的多个Tensor按照对应位置相加，即可还原出逻辑上的Tensor。图5展示了SBP的简单示例。

 

 

 SbpSignature是一个SbpParallel的集合，在OneFlow的设计里是Op的属性，描绘了一个逻辑上的Op被映射成各个设备上的多个物理上的Op以后，这些物理上的Op是如何看待输入输出Tensor在逻辑上和物理上的映射关系的。一个Op会有多个合法的SbpSignature，一个最简单的合法signature就是输入输出都是Broadcast，这表示了这个Op需要整个逻辑上的Tensor数据。

当用户构建的逻辑上的计算图确定以后，OneFlow在Compiler生成分布式的物理上的执行图时，会考虑每个Op的Placement和该Op允许的合法SbpSignature列表，在其中选择一个传输开销最小的SbpSignature作为本次训练的SbpSignature，用于指导Compiler生成最高效的执行图。

关于Op的合法SbpSignature的列表，举一个矩阵乘法（matmul）的Op的例子。定义: Y = matmul(A,B) , A, B, Y都是Tensor，表示Y = AB。那么至少存在两种合法的SbpSignature：

• Y: Split(0), A:Split(0), B:Broadcast
• Y: Split(1), A:Broadcast, B: Split(1)

两种合法的signature在两个设备上的示意图如图6所示。假设逻辑上的MatMul的输入输出Tensor的形状是：

A(64, 10) X B(10, 50) -> Y(64, 50)

且该Op分布在两个设备上。在第一种SbpSignature下，0号设备上的A是逻辑上A的前一半，1号设备上的A是逻辑A的后一半（按照第0维切分），而两个设备上的B跟逻辑上的B完全一致，两个设备输出的Y分别是逻辑上的Y的前一半和后一半。同样可以分析第二种SbpSignature。

值得一提的是，当A是数据，B是模型的时候，第一种SbpSignature就是数据并行，第二种SbpSignature就是模型并行。如果两个相邻的MatMul op，前一个使用第一种SbpSignature，后一个使用第二种SbpSignature，整个网络就实现了混合并行。图7是一个混合并行的示例，定义了Y0 = MatMul_0(A0, B0) , Y1 = MatMul_1(Y0, B1)这样一个由两个op组成的计算图，其中A0, Y0, Y1是数据Tensor，B0, B1是模型Tensor。

 

 

 

 

 在图7中MatMul_0产出的Y0被MatMul_1消费，但是这两个op对同一个Tensor的SBP看待方式是不一样的，MatMul_0认为Y0是Split(axis=0)切分，但是MatMul_1需要一个Broadcast的Y0输入。这时候OneFlow会自动插入一个“万能”的Boxing Op做必要的数据裁剪、拼接、搬运和求和等等操作，使得所有的Op都可以在分布式环境下拿到想要的数据。

另外在数据并行的时候，训练的前向模型Tensor的是Broadcast，对应反向传播的梯度就是PartialSum，当Optimizer需要全部的梯度来更新模型时，就会触发OneFlow的Boxing机制进行高效的梯度同步工作。

3、最易用的分布式并行框架

OneFlow的这套Placement + SBP + Boxing的机制，可以使得用户定义的计算图中的Op、Tensor以任意的方式分布在各个机器和各个设备上，无论是数据并行、模型并行还是流水并行，对于OneFlow而言，都只是一个特定Placement下的特定SbpSignature的组合而已，用户可以方便的配置，也可以交给OneFlow来做自动的处理。

另外，早在微软推出ZeRO-2框架之前，OneFlow就已经支持了类似的功能，多机多卡情况下，每个模型Tensor都只保存在其中一个设备上，降低梯度计算中的内存占用。

综上，在编译期，OneFlow通过在一套数学上严谨的形式系统来表示所有合法的并行模式，并支持编译器较方便地自动搜索最优并行方案；在运行期，则通过Actor系统最优地、灵活地支持并行、并发执行。OneFlow的内核具有简洁、高效和高扩展性的优点。

浅谈人工智能和深度学习的未来

OneFlow区别于其它深度学习框架的核心就是对于分布式训练的思考和设计，OneFlow致力于分布式训练最快。

随着深度学习的发展，模型会越来越大，训练深度学习模型所需的算力会越来越高，同时模型增大的速度要大于GPU单卡显存扩容的速度；训练对算力的增长要求要大于GPU单卡算力增长的速度。所以分布式深度学习训练会越来越常见、越来越重要。BERT、GPT-3等模型的出现印证了判断。单个芯片的能力总是受到物理约束，相信算力增长的解决之道在于横向扩展，而这必须从系统软件的角度去求解，一旦解决，即使是把性能“一般”的芯片互联起来，只要系统软件可以让协同工作好，就可以满足任何算力需求。

在分布式深度学习训练中，性能至关重要。性能和横向扩展性是深度学习框架的核心竞争力。

如果深度学习的未来是单机单卡就能搞得定的规模的话，那么深度学习也就只能做做目前已知的一些简单任务，深度学习的浪潮也会褪去。

开源，OneFlow还有很长的路要走

OneFlow团队是一个小团队，OneFlow作为一个开源框架还有很多不足。框架的易用性和完善性不如Pytorch和Tensorflow。团队还在努力补上OneFlow的短板，同时也非常需要开源社区的支持。