Changelog v0.4.0

Highlights

在这个版本，我们为 OneFlow 新增了大量的功能，0.4.0 是 OneFlow 自开源以来最大的更新。在这个版本中，我们增加了 2-D SBP、流水并行，Checkpoint 的新的接口，以及大量对齐 pytorch 的接口，还支持了 CUDA 11.2。在之前，我们已经开源了 OneFlow 的 GPT 源码，其中大量使用了这个版本的各种新特性，同时也欢迎移步阅读《OneFlow —— 让每一位算法工程师都有能力训练 GPT》这篇文章。

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)

• 示例代码请参考这个测试用例
• mode = "distributed_split" 对应 DeepSpeed ZeRO 优化的 stage 2

支持 Checkpointing 的新接口

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

欢迎阅读相关文章：亚线性内存优化—activation checkpointing在oneflow中的实现

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 运行

其他更新

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

之前一个 OneFlow 的版本采取的是“不同包名，相同版本名”的规则，如 oneflow_cu102==0.3.4，从 0.4.0 之后将采取“相同包名，不同版本名”的规则，如oneflow==0.4.0+cu102，最新安装方式请参考 README 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 主仓库中 的 ONNX 相关的代码将在下个版本移除，具体细节可以看《深度学习框架OneFlow是如何和ONNX交互的？》 一文。oneflow_convert_tools 目前是针对 OneFlow 的 lazy 模式开发，目前最新版本号为 v0.3.2，后面针对 eager 模式的 oneflow_convert_tools 版本号将从 0.4.0 开始

"下集预告"

在下一个版本的 OneFlow 中，将包含更全面的 PyTorch 兼容，包括更多更丰富的接口支持以及多 GPU 支持。同时，下个版本的 OneFlow 也将支持动静图转换的功能。敬请期待！