OneFlow v0.4.0 正式发布

OneFlow深度学习框架

于 2021-06-15 19:46:04 发布

阅读量392

点赞数 2

分类专栏： CHANGELOG 文章标签： tensorflow pytorch 自然语言处理深度学习神经网络

本文链接：https://blog.csdn.net/OneFlow_Official/article/details/117931681

版权

CHANGELOG 专栏收录该内容

3 篇文章 0 订阅

订阅专栏

OneFlow v0.4.0 正式发布，这是自 OneFlow 开源以来的最大版本更新。本次更新新增了多项功能，主要有：增加了 2-D SBP、流水并行，Checkpointing 的新接口以及大量对齐 PyTorch 的接口，还支持 CUDA 11.2。

此前，OneFlow 开源了 GPT 源码（https://github.com/Oneflow-Inc/OneFlow-Benchmark/tree/master/LanguageModeling/GPT），其中大量使用了 v0.4.0 的各种新特性，更多详情欢迎阅读《OneFlow —— 让每一位算法工程师都有能力训练 GPT》。

今天是OneFlow开源的第320天。欢迎下载体验OneFlow v0.4.0：https://github.com/Oneflow-Inc/oneflow

1.Lazy 模式的功能更新

支持 2-D SBP

转为2维

with flow.scope.placement("gpu", "0:0-3", (2, 2)):
    x = flow.hierarchical_parallel_cast(
        x, parallel_distribution=["B", "S(1)"]
    )

转为1维

with flow.scope.placement("gpu", "0:0-3", (4,)):
    x = flow.hierarchical_parallel_cast(
        x, parallel_distribution=["S(0)"]
    )

支持流水并行的新接口

创建 pipeline_stage 的 scope

with flow.experimental.scope.config(
        pipeline_stage_id_hint=dist_util.get_layer_stage(layer_idx)
    ):
    ...

为了使流水并行能更好的工作，必须使用梯度累加，可以使用有限内存跑更大 batch。通过 config 设置梯度累加的步数：

func_cfg = flow.FunctionConfig()
...
func_cfg.train.num_gradient_accumulation_steps(args.num_accumulation_steps)
@flow.global_function(..., function_config=func_cfg)

支持 ZeRO 优化

开启方式：

func_cfg = flow.FunctionConfig()
...
func_cfg.optimizer_placement_optimization_mode(mode) # mode  = "non_distributed" or "distributed_split"
@flow.global_function(..., function_config=func_cfg)

示例代码请参考这个测试用例：https://tinyurl.com/yzhd8u7v

mode = "distributed_split" 对应 DeepSpeed ZeRO 优化的 stage 2

支持 Checkpointing 的新接口

with flow.experimental.scope.config(
    checkpointing=True
):

欢迎阅读文章：《后向重计算在OneFlow中的实现：以时间换空间，大幅降低显存占用》

2.Eager 模式的功能更新

提供 oneflow.experimental 命名空间，部分对齐 torch.xxx 接口

新接口的使用方法

import oneflow.experimental as flow
flow.enable_eager_execution() # 启用 eager

目前部分对齐的功能

  flow.nn.Conv2d  <->  torch.nn.Conv2d  flow.nn.BatchNorm2d  <->  torch.nn.BatchNorm2d  flow.nn.ReLU  <->  torch.nn.ReLU  flow.nn.MaxPool2d  <->  torch.nn.MaxPool2d  flow.nn.AvgPool2d  <->  torch.nn.AvgPool2d  flow.nn.Linear  <->  torch.nn.Linear  flow.nn.CrossEntropyLoss  <->  torch.nn.CrossEntropyLoss  flow.nn.Sequential  <->  torch.nn.Sequential    flow.nn.Module.to  <->  torch.nn.Module.to  flow.nn.Module.state_dict  <->  torch.nn.Module.state_dict  flow.nn.Module.load_state_dict  <->  torch.nn.Module.load_state_dict    flow.save  <->  torch.save  flow.load  <->  torch.load    flow.Tensor  <->  torch.Tensor  flow.tensor  <->  torch.tensor  flow.tensor.to  <->  torch.tensor.to  flow.tensor.numpy  <->  torch.tensor.numpy  flow.tensor 加减乘除  <->  torch.tensor 加减乘除  flow.tensor.flatten  <->  torch.tensor.flatten  flow.tensor.softmax  <->  torch.tensor.softmax    flow.optim.SGD  <->  torch.optim.SGD

基于上述模块已经可以轻松搭建常用网络，如：ResNet、BERT、MobileNetV3 等。后续版本将对齐/支持更多接口，届时可将大多数基于 PyTorch 搭建的网络，轻松切换到 OneFlow。

快速上手例子 lenet:

https://github.com/Oneflow-Inc/models/blob/main/quick_start_demo_lenet/lenet.py
新接口文档链接：

https://oneflow.readthedocs.io/en/master/experimental.html
对齐 torch vision 的 resnet50 示例代码：

https://github.com/Oneflow-Inc/models/tree/main/resnet50
接下来的版本会增加更多对齐 PyTorch 的接口
experimental 下对齐的接口在 0.6.0 版本更新时会被移动到 OneFlow 的命名空间下，届时会完全对齐 PyTorch，OneFlow 0.6.0 会将 eager 作为默认的执行方式
eager 模式目前只支持单 GPU 运行，在 0.5.0 会支持多 GPU 运行

3.其他更新

新的 Python Pip 包名和版本号规则

此前 OneFlow 的一个版本采取了“不同包名，相同版本名”的规则，如 oneflow_cu102==0.3.4，从 0.4.0 之后将采取“相同包名，不同版本名”的规则，如oneflow==0.4.0+cu102，最新安装方式请参考 README （https://github.com/Oneflow-Inc/oneflow#install-with-pip-package）

支持 CUDA 11.2

stable 版本和 nightly 版本的 OneFlow 都支持 CUDA 11.2 平台（cu112）

ONNX 模块独立仓库

ONNX 模块目前在新仓库

https://github.com/Oneflow-Inc/oneflow_convert_tools 中维护，OneFlow 仓库中的代码将在下个版本移除。具体细节可看《深度学习框架OneFlow是如何和ONNX交互的？》一文。oneflow_convert_tools 是针对 OneFlow 的 lazy 模式开发，目前最新版本号为v0.3.2，后续针对 eager 模式的 oneflow_convert_tools 版本号将从 v0.4.0 开始