最新深度学习框架——OneFlow：新分布式训练（附源代码）

计算机视觉研究院专栏

作者：Edison_G

OneFlow 是什么？

OneFlow是开源的、采用全新架构设计，世界领先的工业级通用深度学习框架。

为什么是OneFlow？

• 分布式训练全新体验，多机多卡如单机单卡一样简单

• 完美契合一站式平台(k8s + docker)

• 原生支持超大模型

• 近零运行时开销、线性加速比

• 灵活支持多种深度学习编译器

• 自动混合精度

• 中立开放，合作面广

• 持续完善的算子集、模型库

1、背景

随着深度学习的发展，用户越来越依赖 GPU 或者其他加速器进行大规模运算。人工智能（Artificial Intelligence）需要更优秀的软件来释放硬件的能量已成业界共识。一方面，各种框架需要进一步降低编写深度学习分布式训练程序的门槛；另一方面，用户期待系统可以支持不同的深度学习网络模型，并实现线性加速。各知名深度学习框架正在朝这方面努力，但用户在使用这些框架时仍会遇到横向扩展性的难题，或者是投入很多计算资源但没有看到效率收益，或者是问题规模超过 GPU 显存限制而无法求解。

先说下分布式训练：

分布式训练

在 OneFlow 中您只需要简单的几行配置，OneFlow 框架内部会自动处理任务调度、资源并行等问题，因此，您并不需要特别改动网络结构和业务逻辑代码，就可以方便地使用分布式训练。

OneFlow 的分布式训练能力独树一帜，是 OneFlow 区别于其它框架的 最重要特性。

将介绍：

• 如何将单机程序修改为分布式程序

• OneFlow 中节点概念及分工

OneFlow 分布式优势

• 采用去中心化的流水架构，而非 master 与 worker 架构，最大程度优化节点网络通信效率

• 提供 consistent view ，整个节点网络中只需要逻辑上唯一的输入与输出

• 提供兼容其它框架的mirrored view，熟悉其它框架分布式训练的用户可直接上手

• 极简配置，由单一节点的训练程序转变为分布式训练程序，只需要几行配置代码

配置分布式训练网络

通过 OneFlow 提供的分布式配置的接口，您只需要简单的几行配置(指定分布式计算的节点 ip 以及每个节点使用 gpu 的数量)即可实现分布式的训练网络。

换句话说，这使得单机训练程序与分布式训练程序几乎是一样的，作为 OneFlow 用户，只需要专注于程序的 业务逻辑 及 模型结构本身 ，而不用操心分布式执行问题。OneFlow 框架会自动帮您处理复杂的任务调度、资源并行等问题。

下面，我会介绍一个例子：将单机版的训练作业，通过添加几行配置代码后将其改造为分布式训练作业。

单机训练程序

以下是单机训练程序的框架，因为各个函数的代码会在下文分布式程序中呈现，在此就未详细列出。

import numpy as np
import oneflow as flow
import oneflow.typing as tp


BATCH_SIZE = 100


def mlp(data):
  #构建网络...




@flow.global_function(type="train")
def train_job(
    images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
    labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
  #作业函数实现...
  #配置训练优化方法和参数




if __name__ == '__main__':
  #调用作业函数，开始训练...
      loss = train_job(images, labels)
  #...

上一期“计算机视觉研究院”已经将OneFlow的背景及基础知识详细分析过，今天我们主要分析OneFlow深度学习框架最核心的部分——Actor 机制和SBP（Split、Broadcast、PartialSum） 机制。

2、OneFlow 的设计目标

OneFlow 的设计目标是追求极致的性能，特别是分布式多机多卡环境下的横向扩展性，希望能让用户使用多机多卡就像使用单机单卡一样容易，且享受线性加速的运行效率。

OneFlow还出品了OneBrain产品！

OneBrain产品优势

• 端到端的智能数据平台

包含数据上传、数据预处理、特征工程、模型训练、模型评估和模型发布至离线或在线环境。

• 可视化开发

提供可视化的代码开发、工作流设计器页面，简单拖拽和开发，即可完成复杂的数据分析任务。

• 智能数据构建与管理

提供数据集成、数据开发、数据质量和数据服务等专业高效、安全可靠的一站式大数据智能平台。

• 开放灵活的算法支持

内置了企业常用的统计分析、机器学习、深度学习算法，面向数据分析应用提供基础算法支持。

基于Serverless云函数实现算子的灵活封装、集成，并支持发布到已有算法库中。

OneBrain 产品功能

• 数据集成

支持多种数据源接入配置

支持实时、历史数据的批量、增量同步

支持任意结构化、非结构化、半结构化的数据传输

• 数据开发

SQL智能编辑器

提供SQL格式化、智能补齐、关键字高亮、错误提示、SQL内部结构等人性化功能，带来更顺滑的SQL开发体验

多引擎混编工作流，串联跨引擎数据节点开发，每个类型的引擎数据节点都有对应的环境进行开发

• 知识图谱

知识图谱数据治理将各类数据（结构化数据、半结构化数据、文本数据、音视频图像等非结构化数据）进行识别，将其中的各类数据串连、组织、融合汇总成为人、地、事、物、组织等实体为节点，属性、时空、语义、特征等联系为边的一张关系网中，形成知识图谱库，再现真实世界对象之间的错综复杂的关系，利用AI技术让计算机仿照人类的思维方式去“思考”从而精准推断问题的答案

• 数据标注

覆盖各类语音、视频、文本、图片和点云工具，每个控件和快捷键都经过多年大量任务的打磨

• 模型训练

支持OneFlow、TensorFlow、PyTorch、Caffe和MXNet等多种深度学习和机器学习框架

用户可以基于预置的环境，定制自己的开发环境

支持分布式训练模型

支持模型训练的参数、代码和数据的跟踪

• 模型发布

可以轻松在生产环境中一键式部署您的受训模型，以便您开始针对实时或批量数据生成预测 支持弹性多副本部署，实现工业级的高可用部署

3、Actor 机制（一套简洁的机制解决几乎所有技术难题）

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

Actor 内部状态机

在 Actor 启动之后，会根据与其他 Actor 之间收发消息来切换自己的两个状态：等待状态 和 执行状态 。

一个 Actor 收到的消息一般分为几个类型：

• 上游的生产者 Actor 发来消息说：你可以来读我生产的数据了；

• 下游的消费者 Actor 发来消息说：我用完你生产的数据了。

当这个数据被所有消费者都用完以后，就可以回收成为空闲块等待下一次被该 Actor 重新生产一份新的数据。

一个 Actor 收到消息以后都会去尝试判断当前是否满足执行条件，执行条件一般有两个：

• 需要读取的数据是否都到齐了；

• 是否有空闲块可以拿来被生产。当满足执行状态以后 Actor 就开始调用自己内部的 Kernel 真实的去读写数据。

执行完毕后 Actor 会向上下游发消息：

• 向下游的消费者 Actor 发消息说：我刚生产了一块数据，你可以来读了；

• 向上游的生产者 Actor 发消息说：我刚用完了你之前发给我的数据了。

Actor 只需要关心上下游的消息就能判断自己能不能执行。每个 Actor 都通过自己内部的状态机和收发消息机制实现了 完全去中心化 的分布式协同工作。

3.2、Actor 机制实现流水线

上面我们介绍了 Actor 的内部状态机，Actor 之间的消息传递和数据传递是依赖 Register 实现的。一个 Actor 是否能执行，只跟两个条件相关：

• 自己消费的那些 Register 是否可读；

• 自己生产的那些 Register 是否有空闲块可写。

对于一个 Register，如果我们运行时给它分配多个空闲块，那么相邻的两个 Actor 就可以同时工作，工作时间重叠起来，这样就实现了各个 Actor 之间的流水线。理想状态下整个静态执行图的执行时间就是整个系统中是性能瓶颈的那个 Actor 运行的总时间，其余 Actor 的执行时间都被流水线掩盖起来了。

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

Actor 生产消费关系和执行时序图

• 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 机制就实现了流水并行。

在这里我们抛出一个更进一步深入的问题：整个数据流的执行像一个网络，数据在网络中的流动就完成了计算，如何避免生产者生产太快，消费者消费不及，以及如何避免生产者生产太慢，消费者感到饥饿的问题，这涉及到对计算、内存、传输带宽的规划，尽可能使系统的瓶颈之处最宽，需要解决流控（flow control）的问题以及资源分配问题（如每个 Actor 的 Register 到底分配几个内存块配额），这非常关键，也是 OneFlow 系统已解决的问题。

3.3、数据搬运是一等公民

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

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

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

数据是如何从一个设备搬运到另一个设备上的

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

4.4、尽可能并行

在 OneFlow 的设计中，所有的出发点都是希望可以尽可能并行，从而达到最优的分布式性能。比如考虑到分布式训练模型梯度同步时，显存到内存的传输带宽高于机器之间的网络传输带宽，OneFlow 会做两级的 scatter 和 gather 操作（本机的和各个机器之间的），用于增加 locality，提高整体性能。

又比如在异步启动深度学习训练时，Python 端用户的控制逻辑跟 OneFlow 运行时的执行图是并行执行的，同时 OneFlow 有一套互斥临界区的设计保证执行的高效性和正确性。

数据加载部分无论是从磁盘读数据还是从 Python 端喂数据，OneFlow 都能保证尽可能并行，使得计算设备不会因为要等数据而导致性能下降。

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

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

4、SBP（Split、Broadcast、PartialSum） 机制

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 的方式做全局的统筹管理。

Placement

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

下图展示了一种可能的 Placement 例子。用户定义了一个由3个 Op 组成的网络：Op_0 -> Op_1 -> Op_2。

其中 Op_0 和 Op_1 的 Placement 是 Device 0，Op_2 的 Placement 是 Device 1，这就是一个流水并行的例子，Oneflow 会自动在 Op_1 和 Op_2 之间插入需要的数据搬运的 Copy Op。

一个流水并行的Placement示例图

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。

下图展示了 SBP 的简单示例。

几种 SbpParallel 的简单情形

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：

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

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

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

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

MatMul的两种合法SbpSignature

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

值得一提的是，当 A 是数据，B 是模型的时候，第一种 SbpSignature 就是 数据并行 ，第二种 SbpSignature 就是 模型并行 。如果两个相邻的 MatMul op，前一个使用第一种 SbpSignature，后一个使用第二种 SbpSignature，整个网络就实现了 混合并行 。

下图是一个混合并行的示例，定义了 Y0 = MatMul_0(A0, B0) , Y1 = MatMul_1(Y0, B1) 这样一个由两个op组成的计算图，其中A0, Y0, Y1是数据Tensor，B0, B1 是模型Tensor。

混合并行示例

在上图中 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 机制进行高效的梯度同步工作。

最易用的分布式并行框架

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

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

5、总结

综上，在编译期，OneFlow 通过设计一套数学上严谨的形式系统来表示所有合法的并行模式，并支持编译器较方便地自动搜索最优并行方案。

在运行期，OneFlow 通过 Actor 系统最优地、灵活地支持并行、并发执行。OneFlow 的内核具有简洁、高效和高扩展性的优点。

基于此设计，OneFlow 使得分布式训练的性能达到极致，且分布式训练跟单卡一样简单易用。

---

下一期我们给大家距离说明“OneFlow的并行特色”！

✄-------------------------------------

我们开创一段时间的“计算机视觉协会”知识星球，也得到很多同学的认可，我们定时会推送实践型内容与大家分享，在星球里的同学可以随时提问，随时提需求，我们都会及时给予回复及给出对应的答复。

如果想加入我们“计算机视觉研究院”，请扫二维码加入我们。我们会按照你的需求将你拉入对应的学习群！

计算机视觉研究院主要涉及深度学习领域，主要致力于人脸检测、人脸识别，多目标检测、目标跟踪、图像分割等研究方向。研究院接下来会不断分享最新的论文算法新框架，我们这次改革不同点就是，我们要着重”研究“。之后我们会针对相应领域分享实践过程，让大家真正体会摆脱理论的真实场景，培养爱动手编程爱动脑思考的习惯！

扫码关注我们

公众号 : 计算机视觉研究院

扫码回复“OneFlow”获取源代码