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 的各种新特性，更多详情欢迎阅读《除了“烧钱”的硬件集群，阻止算法工程师训练GPT-3的原因是什么》。
今天是OneFlow开源的第320天。欢迎*载下**体验OneFlow v0.4.0： https://github.com/Oneflow-Inc/oneflow
1.Lazy 模式的功能更新
支持 2-D SBP
转为2维
withflow.scope.placement("gpu","0:0-3",(2,2)):
x=flow.hierarchical_parallel_cast(
x,parallel_distribution=["B","S(1)"]
)
转为1维
withflow.scope.placement("gpu","0:0-3",(4,)):
x=flow.hierarchical_parallel_cast(
x,parallel_distribution=["S(0)"]
)
支持流水并行的新接口
创建 pipeline_stage 的 scope
withflow.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 的新接口
withflow.experimental.scope.config(
checkpointing=True
):
欢迎阅读文章：《后向重计算在OneFlow中的实现：以时间换空间，大幅降低显存占用》
2.Eager 模式的功能更新
提供 oneflow.experimental 命名空间，部分对齐 torch.xxx 接口
新接口的使用方法
importoneflow.experimentalasflow
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 开始
4.下集预告
在下一个版本的 OneFlow 中，将包含更全面的 PyTorch 兼容，包括更丰富的接口支持以及更多 GPU 支持。同时，下个版本的 OneFlow 也将支持动静图转换功能。敬请期待。