您的位置: 首页 > 技术杂谈 > 正文

推出全新分布式计算接口！OneFlow v0.7.0发布，LiBai代码库、Serving、MLIR一应俱全

2022-04-07 11:00 https://my.oschina.net/oneflow/blog/5510544 OneFlow深度学习框架次阅读条评论

今天是 OneFlow 开源的 610 天，OneFlow v0.7.0 正式发布。点击“阅读原文”，欢迎下载体验最新版本。

本次更新包含以下重点：

完善地提供了一种可以帮助用户轻松使用多机多卡执行的机制 ：Global Tensor是 OneFlow 为社区带来的分布式执行的易用方案，用它可以方便地实现各种分布式并行策略，极大提高分布式实现的灵活性和易用性。基于 Global Tensor，OneFlow已支持 ResNet50、Wide and Deep、GPT、Bert、Swin-Transformer、InsightFace 等模型的并行化。
持续完善 nn.Graph 功能，支持包括 ZeRO 、GradAcc、Checkpointing、Pipeline 相关的高级功能，丰富了 graph.debug 模式。新增支持任意 2D SBP 转换、支持 2D SBP 的半自动推导、支持断点续训等。新增 OneFlow Feature Stages 标识，并给 nn.Graph 所提供的每一个功能都增加该标识。就 nn.Graph 整体而言，基础功能进入 Beta Stage，可以支持对该功能的大部分需求；高级功能进入 Alpha Stage，可支持对该功能的标准需求。
深度优化 Eager 性能，在 V100 显卡上测试 Swin-Transformer 模型的单卡性能相比 v0.6.0 提升 3 倍。
算子相关进展：在单机单卡场景下，OneFlow 对 PyTorch 的兼容性进一步完善，OneFlow 已经支持的算子都保证和 PyTorch 的接口、语义、结果一致；另外设计了一套自动测试框架来验证一致性，常见网络可以做到import oneflow as torch 来完成迁移。相较于 v0.6.0， OneFlow 新增 16 个算子，优化 6 个算子的性能，修复 16 个算子存在的 bug。
支持 Einsum 算子和 View 机制。
OneFlow 正式接入 MLIR 编译器生态。
发布 OneFlow-Serving v0.1.0，提供了开箱即用的 Triton OneFlow backend 镜像（https://github.com/Oneflow-Inc/serving）。
发布 LiBai（李白） v0.1.0：这是一个针对 Transformer 的大规模分布式并行训练代码库，相比 Megatron-LM 等定制化代码库，基于模块化设计的 LiBai 为分布式训练提供了一系列模型和训练组件，让分布式下的模型训练像单卡一样方便（https://github.com/Oneflow-Inc/libai）。
发布 Flow-Vision v0.1.0：新增 DeiT、ConvNeXt、ReXNet 等模型，完善了使用教程和文档（https://github.com/Oneflow-Inc/vision）

以下为版本更新详情。

1 分布式

Global Tensor

Global Tensor 是OneFlow发布的一套全新的分布式计算接口，可以很方便地支持包括数据并行、模型并行和流水并行在内的任意并行方式。不同于普通 Tensor（下文叫 Local Tensor），Global Tensor 是一种全局视角下的 Tensor，它的数据以特定方式分布在集群中的一组计算节点上，每个节点存储了该 Tensor 的部分或全部数据。placement 和 SBP 是每个 Global Tensor 的基本属性，描述了其数据在集群中的分布方式。

Global Tensor 的数据分布方式

Global Tensor 支持三种不同的数据分布方式，我们将其统称为 SBP。

Split (dim)：数据以dim 维度平均切分并分布到每一个计算节点上。
Broadcast：数据在每一个计算节点间进行复制。
PartialSum：数据为每一个计算节点的 element-wise 加和。

统一的计算接口

Global Tensor 具有和 Local Tensor 基本一致的计算接口，支持以很少的改动就可以将一个单卡的代码转换成分布式方式执行。

支持 Local Tensor 与 Global Tensor 的转换

Local Tensor 可以使用 Tensor.to_global 接口创建一个 Global Tensor，并将该 Local Tensor 作为它在当前节点的本地分量。
Global Tensor 可以使用 Tensor.to_local 接口返回它在当前节点的本地分量。

支持 Global Tensor 在集群中重新分布

Global Tensor 使用 Tensor.to_global 接口支持在集群中进行数据的重新分布，既可以选择分布到另外一组节点上，也可以改变它在这组节点上的分布方式（即改变 SBP ）。重新分布通常会发生跨进程的数据通信，Tensor.to_global 这个接口很好地屏蔽了复杂的底层通信逻辑。

>>> import oneflow as flow>>> x = flow.tensor([1.0, 2.0], placement=flow.placement("cuda", ranks=[0, 1]), sbp=flow.sbp.split(0))>>> y = x.to_global(placement=flow.placement("cuda", ranks=[2, 3]), sbp=flow.sbp.broadcast)

OneFlow 中每一种计算接口都定义了一套其所能支持的输入和输出的 SBP 组合，Global Tensor 支持自动重新分布，以满足执行某个计算接口对其 SBP 的要求。比如下面的代码：

   
   
   
    
    
    >>> import oneflow as flow>>> x = flow.randn(4, 4,             placement=flow.placement("cuda", ranks=[0, 1]),             sbp=flow.sbp.split(0))>>> y = flow.randn(4, 4,             placement=flow.placement("cuda", ranks=[0, 1]),             sbp=flow.sbp.split(1))>>> z = x + y

当执行 x + y 时由于 x 是按第 0 维切分，y 是按第 1 维切分，它们在每个节点上的分量无法直接完成相加，那么它就会自动将 x 的 SBP 转换成flow.sbp.split(1) 或者将 y 自动转换成flow.sbp.split(0) ，计算得到的结果 z 的 SBP 为 flow.sbp.split(1) 或 flow.sbp.split(0) 。

注意

Global Tensor 目前不支持和 DDP 接口混合使用；
Global Tensor 的代码要求所有节点一起执行，有分支的代码可能会因为执行路径分散而导致进程死锁，我们会持续改进这里的用户体验。

2 持续完善 nn.Graph 的功能
新增 OneFlow Feature Stages 标识

OneFlow Feature Stages 标识 OneFlow 功能的成熟度等级依次为 Pre-alpha Stage、Alpah Stage、Beta Stage、Release candidate (RC) Stage、 Stable Stage 。它给用户提供功能的状态说明，以了解该功能下所提供的保证，如功能完备性、API 稳定性、文档等；它还为开发者提供完善功能的标准，并据此推进对应功能走向成熟。
nn.Graph v0.7.0 进展概述

基础功能进入 Beta Stage，可以支持对该功能的大部分需求；
高级功能进入 Alpha Stage，可支持对该功能的标准需求；
已经支持了 ResNet50、Wide and Deep、GPT、Bert、Swin-Transformer、InsightFace 等模型。

nn.Graph 静态图下 Feature

Static Graph下的 Op 动静转换功能，从 Alpha Stage 到 Beta Stage
- 新增所有合法 Op 在 nn.Graph 做静态执行的单测，自动化单测功能完备；
- 新增支持更为灵活的输入输出，包括 List/Tuple/Dict 以及它们的嵌套，修复返回大小为 1 的 Tuple 问题；
- 新增后向的自动测试。
Static Graph 下的 Optimizer 和 LR Scheduler, 从 Alpha Stage 进步到 Beta Stage
- 添加更多的内置 LR scheduler，例如 WarmupLR, CosineAnnealingWarmRestarts 等常见的 scheduler ，同时提供 SequentialLR 和 ChainedScheduler 来为 scheduler 提供不同的组合能力；
- 重构了 scheduler 的 get_lr 函数，将其改造成纯函数的实现，目的是为了把 lr 的计算由迭代解切换到解析解，为 scheduler 的组合使用提供支撑；
- add_optimizer 接口新增参数 is_sparse。用以支持 graph 模式下的稀疏更新，支持稀疏更新的 optimizer 有 Adam 和 SGD。Eager 模式下的 optimizer 还未支持稀疏更新策略，后续版本会同稀疏张量一起支持。功能状态为 Pre-alpha Stage；
- 新增 LR 和 Step 的 Debug 打印功能，打开 LR Scheduler 的verbose 开关即可。
Static Graph 下新增state_dict 和load_state_dict ，支持断点续训，功能状态为 Beta Stage
Static Graph 下的 Debug，从 Alpha Stage 进入 Beta Stage
- 新增 debug(2) 、 debug(3) ，可以分 nn.Module 去定位问题，可定位 C++ 层 Op 对应的 Python 代码，可定位 Op 的前向图创建和推理；
- 新增显示内存开销。
Static Graph 下新增 ZeRO-DP 的支持，在数据并行下缩减和 Optimizer 关联的显存开销，功能状态为 Alpha Stage
Static Graph 下的 Global Tensor，多种并行执行，整体状态为 Alpha 和 Beta 之间
- 已在 LiBai 等多个模型库中使用；
- 已经在 OneFlow 模型库中广泛使用，单测的覆盖在进行中；
- 1D Global Tensor支持只定义 Source Tensor 的 SBP，下游可以自动推导，而且效果良好，Beta Stage；
- 新增 2D Global Tensor 支持只定义 Source Tensor 的 SBP ，下游可以自动推导，而且效果良好，Alpha Stage；
- 新增支持 1D to ND 与 ND to 1D 的转换, Alpha Stage；
- 新增支持任意 2D SBP 的转换, Alpha Stage；
- 1D&2D 单 Op 的测试在覆盖中，Pre-alpha Stage；
- 支持选择半自动推导 SBP的挑选策略，Pre-alpha Stage。
Static Graph 下的梯度累积(Gradient Accumulation)，重构和修复 Reshape 的支持，新增 API 文档，当前接口为mini-batch 的输入，下个版本将更新为体验更好的micro-batch 的输入，功能状态从 Pre-Alpha 到 Alpha；
Static Graph 下的流水并行，完善了教程，在 Libai 等多个模型库进入使用，功能状态为 Beta；
Static Graph 下的自动混合精度 AMP，新增 API 文档，功能状态 Pre-Alpha 到 Alpha；
Static Graph 下的 Activation Checkpointing，新增 API 文档，功能状态从 Pre-Alpha 到 Alpha；
Static Graph 下的多种 Op Fuse 优化，新增 API 文档，功能状态 Pre-Alpha 到 Alpha；
Static Graph 下的 XLA/TensorRT/OpenVINO 执行，新增 API 文档，功能状态 Pre-Alpha 到 Alpha。

教程

En https://docs.oneflow.org/en/master/basics/08_nn_graph.html
中 https://docs.oneflow.org/master/basics/08_nn_graph.html

API文档

En https://oneflow.readthedocs.io/en/master/graph.html
中 https://start.oneflow.org/oneflow-api-cn/graph.html

流水并行的教程

En https://docs.oneflow.org/en/master/parallelism/06_pipeline.html
中 https://docs.oneflow.org/master/parallelism/06_pipeline.html

nn.Graph 静态图下的模型支持

支持ResNet50单机单卡和单机多卡（ https://github.com/Oneflow-Inc/models/tree/main/Vision/classification/image/resnet50）
支持了Wide and Deep模型（ https://github.com/Oneflow-Inc/models/tree/main/RecommenderSystems/wide_and_deep）
支持了Libai中的GPT、Bert、Swin Transformer（ https://github.com/Oneflow-Inc/libai）
修复了以上多种模型的支持中遇到的功能问题

3 深度优化 Eager 性能

深度优化 Eager 性能，OneFlow 在 V100 显卡上测试 Swin-Transformer 模型性能，单卡比 PyTorch 快25%， 8卡 DDP 比 PyTorch 快10%
优化 DDP 中的 NCCL 通信调度逻辑
DDP 支持 AllReduce fuse 优化，减少碎片化的 AllReduce 引起的额外开销，在 ResNet50 上测试有约 5% 的性能提升
VM 支持指令融合优化，大幅节省零碎小 Kernel 的调度开销
优化了 CPU 负载较高时的额外内存开销
Eager DataLoader 支持进程间内存共享优化
深度优化 Clip Grad 性能

4 算子相关进展

OneFlow 成功适配 oneDNN 用于 CPU 算子加速，unary 和 binary element-wise 等 CPU 算子的性能提升 4 倍，Swin-Transformer 的 dataloader 速度提升 2.5 倍。
DataLoader 新增进程间内存共享功能，大幅提升 DataLoader 在 DDP 情况下的性能。
新增 Bool 类型 Tensor。
新增 To_contiguous 算子服务于 view 机制。
新增 Scalar div 算子。
新增 Lamb 优化器。
新增 Polynomial Learning Rate Scheduler。
新增 Tensor_split，As_strided 算子。
新增 Cumprod 算子。
新增 Tensor.T() 和 oneflow.t() 算子。
新增 Normalize 算子。
新增 div 和 sub 算子的 inplace 版本。
新增 Module.zero_grad 功能。
新增 Scalar Tensor 作为索引来做 list indexing 的功能。
新增 Leaky ReLU 算子的 half 类型支持。
新增 Mask Select 算子支持。
新增 Bool 类型的 Broadcast 及 Allgather 等非 reduce 通信操作。
基于自动测试框架开发支持 eager global 的自动测试。
优化 ReduceSum CUDA Kernel 的性能。
优化 Gather 算子的 CUDA Kernel 的性能。
优化 NCHW 情况下的 MaxPool 和 AvgPool 算子的 CUDA Kernel性能。
优化 PReLU 算子的后向计算部分，一般情况下可以节省较多显存。
优化 LayerNorm 后向 Kernel，进一步节省显存。
Conv1D/2D/3D 和 DeConv1D/2D/3D Kernel，stride 和 dilation 参数支持单个 int 传参，新增 Tensor.zero_() 接口，对齐 PyTorch tensor.norm，torch.max，torch.min 用法，flow.nn.functional.dropout 支持 inplace。
修复 BatchNorm 模块在 affine 参数为 False 运行报错的 bug。
修复 Maximum，Mimimum 反向的 bug。
修复 Var 算子在某些情况下结果不符合预期的 bug。
修复 Tensor deepcopy 时行为不正确的 bug。
修复 Slice 算子输入 index 是 scalar tensor 时的 bug。
修复 BinaryCrossEntropy 在 half 情况下可能产生 nan 的 bug。
修复 Pow 算子底数和指数分别为实数和 Tensor 类型时报错的 bug。
修复 Stack 算子后向的 bug。
修复 Clip grad 在默认配置下并在 CUDA 上执行时 CPU 同步导致的效率过低问题。
修复 Batch Gather 和 Unsorted Batch Segment Sum 算子的 sbp 推导，global 单测通过。
修复 Affine Grid 算子的 Physical Shape 推导，并修复某些 SBP 情况下计算结果不符合预期的 bug ，global 单测通过。
修复 Arange 算子不支持产生 0 size tensor 的问题，global 单测通过。
修复 Flip 算子 SBP 推导不正确的问题，global 单测通过。
修复 Advanced Indexing 和 ZerosLike 算子 SBP 的 bug。
修复 Eager global inplace 可能不成功的 bug。

5 支持 Einsum & View 机制

新增 einsum 算子，einsum 提供了一套既简洁又优雅的规则，可实现包括但不限于内积、外积、张量乘法、张量转置和张量收缩（tensor contraction）等张量操作，熟练运用 einsum 可以很方便实现各种复杂的张量操作且不容易出错。
新增 view 机制。通过 view 机制，一些常用算子可以实现 Tensor 的内存复用/共享，这样就能省去 Kernel Launch/Compute 的过程，并达到节省显存的效果。目前，新增了 reshape, view, squeeze, unsqueeze 等不会改变 tensor.is_contiguous() 属性的 view 算子，后续会增加更多 view 算子（如 transpose, permute, narrow, expand, unfold 等）。

6 编译器相关进展

OneFlow 正式接入 MLIR 生态，OneFlow Dialect 组件已经完备。成功完成了 OneFlow Job（OneFlow nn.Graph 的计算图）和 MLIR 的 RoundTrip，并对 OneFlow 所有的算子在 CI 流程中进行 RoundTrip 测试。
基于 MLIR DRR 实现了一系列自动 Fused 算子的静态图优化，加速 OneFlow 模型训练和推理。

7 OneFlow Serving

OneFlow Serving 发布 v0.1.0 版本，特性如下：

提供用于推理的 OneFlow C++ API，支持加载模型和静态图推理。
模型训练者可以在 Python 中执行 flow.save(graph) 来同时保存模型权重和 MLIR 格式的计算图，用于在 C++ API 中加载并推理（暂不支持在 Python API 中加载计算图）。
支持自动使用 TensorRT 和 OpenVINO 推理 OneFlow 模型，无需模型转换（基于 OneFlow XRT 模块），在 NVIDIA GPU 和 Intel CPU 上可以取得更好的加速效果。
实现 Triton OneFlow backend
- 提供开箱即用的 Docker 镜像
- 支持 auto configuration，部署时只需要给出模型路径，不需要写 Triton 配置文件
在 OF 智能云上线了一个使用 Triton OneFlow backend 进行部署的项目（ https://oneflow.cloud/drill/#/project/public/code?id=7fc904d8dbe0069820da5d6d32a764fe ），欢迎试玩。