计算机系统应用教程网站

网站首页 > 技术文章 正文

OneFlow v0.4.0 正式发布

btikc 2024-09-05 12:40:43 技术文章 9 ℃ 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 的各种新特性,更多详情欢迎阅读《除了“烧钱”的硬件集群,阻止算法工程师训练GPT-3的原因是什么》。


今天是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 开始


4.下集预告


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


欢迎下载体验OneFlow新一代开源深度学习框架:https://github.com/Oneflow-Inc/oneflow

Tags:

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表