

Global View的概念和实现|OneFlow学习笔记
source link: https://blog.csdn.net/OneFlow_Official/article/details/124138777
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

Global View的概念和实现|OneFlow学习笔记
撰文|月踏
更新|赵露阳
在OneFlow中,Global View也被称作一致性视角,用来把一个物理集群抽象成一个逻辑设备,并使用Placement和SBP来实现这种抽象。本文从基本概念、数据结构、接口实现等方面对其进行学习和总结。
1
Placement
1.1 使用示例
Placement用来描述设备信息,包括设备类型、设备分布信息,先看一个具体的使用示例,然后根据这个示例来做分析:
type(x)的输出为:
<class 'oneflow._oneflow_internal.placement'>
print(x)的输出为:
oneflow.placement(type="cuda", ranks=[[0, 1, 2, 3], [4, 5, 6, 7]])
可见Placement有下面两个属性:
-
type:表示设备类型,目前只支持CPU和CUDA
-
ranks:一个Python list,用于表示device的排布信息,ranks可以是一维至多维的,其shape表示了设备的排布信息(hierarchy)。上述ranks表示Tensor存放在集群中的2个节点中,其中节点1中使用设备0~3,节点2中使用设备4~7。
1.2 追踪代码
先看Python端的接口,在python/oneflow/__init__.py+27可以看到下面语句:
placement = oneflow._oneflow_internal.placement
可见Placement是在前文《OneFlow学习笔记:python到C++调用过程分析》
讲的一个pybind定义的_oneflow_internal这个module的子module,在oneflow/api/python/symbol/placement_symbol.cpp+226可以找到下面的定义:
通过上面的多个def接口可以看到,通过调用PlacementSymbolExportUtil::CreateParallelDescSymbol()来构造Placement对象,这个函数是个重载函数,定义在同一个文件中,多个重载版本只是参数有区别,其中的创建Placement的逻辑基本一致,下面列一个重载版本作为示例:
由这个函数的返回值可见,Placement在C++中对应的数据结构是ParallelDesc,这个数据结构后面再说,现在先继续看创建逻辑,这里继续调用了CreateParallelDesc函数,同样定义在PlacementSymbolExportUtil这个类中:
这里最重要的是调用MakeParallelConf这个函数,位于oneflow/core/framework/parallel_conf_util.cpp+38,它根据传入的device type、machine device ids、hierarchy shape信息创建了一个cfg::ParallelConf类型的对象parallel_conf。
这里需要注意的是,hierarchy shape表示设备的排放层次序列,即上面通过const auto& shape = JUST(GetRanksShape(obj));得到的由ranks参数所表示的list shape。创建完parallel_conf之后,通过后面的GetParallelDescSymbol接口来得到需要返回的ParallelDesc类型对象,下面是MakeParallelConf的主要实现:

再继续看下GetParallelDescSymbol是怎么根据cfg::ParallelConf的对象得到ParallelDesc类型对象的,GetParallelDescSymbol定义在oneflow/core/framework/instructions_builder.cpp+230:
大概过程就是在一个全局表里面去查有没有cfg::ParallelConf对应的已经创建好的ParallelDesc的对象,有的话直接返回,没有的话就创建出来放到全局表中去,至此就得到了前面展示的pybind接口中需要的ParallelDesc对象。
下面继续看下相关的数据结构,主要是cfg::ParallelConf和ParallelDesc,它们都和下面这个proto文件相关:
-
oneflow/core/job/placement.proto
这个proto文件是所有placement相关数据结构的源头,根据它会先自动生成下面三个文件:
-
build/oneflow/core/job/placement.pb.h
-
build/oneflow/core/job/placement.pb.cc
-
build/of_cfg_proto_python/oneflow/core/job/placement_pb2.py
前两者的接口主要是为了对placement数据做序列化,但是这些接口不适合对接python,所以使用tools/cfg中的工具对第三个文件做了处理,生成了下面三个方便给python端提供接口的文件:
-
build/oneflow/core/job/placement.cfg.h
-
build/oneflow/core/job/placement.cfg.cpp
-
build/oneflow/core/job/placement.cfg.pybind.cpp
cfg::ParallelConf这个数据结构就定义在build/oneflow/core/job/placement.cfg.h这个自动生成的文件中,再看ParallelDesc,它其实可以看作是cfg::ParallelConf的一层wrapper,主要是用在c++代码中来表示placement的数据结构,位于oneflow/core/job/parallel_desc.h+46,我们只需要关注这个数据结构就行:

这里面的数据结构看起来很复杂,我也不完全明白所有成员的含义,但归根结底这里数据成员的值还都是根据cfg::ParallelConf中的内容来的,在前面调用GetParallelDescSymbol时,如果全局表中没有找到,就会根据cfg::ParallelConf类型对象创建一个ParallelConf类型对象,再根据这个ParallelConf类型对象创建一个ParallelDesc类型对象,在ParallelDesc的构造函数中会调用类内的MaybeInit函数,位于oneflow/core/job/parallel_desc.cpp+112,这里面会完成ParallelDesc数据成员的赋值:

以上就是在python端使用placement时从上到下的大概过程和placement相关的数据结构。
2
SBP
2.1 基本概念
SBP是OneFlow发明的概念,在OneFlow的官方文档和论文中都有详细的说明(具体链接都在文末Reference中列出),这里只做简单介绍,它是下面三个单词的缩写:
-
Split:表示把数据按照指定的维度进行切分,被切分出的数据块会被分发到前面Placement指定的各个物理设备中去
-
Broadcast:表示把整份数据广播到前面Placement指定的各个物理设备中去
-
Partial:表示前面Placement指定的各个物理设备中所存的数据不是最终的运算结果,需要对各个物理设备上的数据进行Elementwise的add/min/max等操作,才能得到最终的结果
SBP描述了一致性视角下的数据与物理设备上的数据的映射关系,计算的时候,数据会根据自己的SBP属性被分发到各个物理设备进行计算,下面贴一张OneFlow官方文档的截图来直观的展示SBP的三种情况:
图1
2.2 使用示例
在Python环境做下面这个简单的示例:
type(s)、type(b)、type(p)的输出如下:
print(s)、print(b)、print(p)的输出如下:
2.3 追踪代码
先找入口,在python/oneflow/__init__.py+196:
from . import sbp
这用到了同目录下的这个module文件:python/oneflow/sbp.py,内容如下:
可见split、broadcast和partial_sum都是定义在pybind定义的_oneflow_internal这个module的子module sbp的内部,在oneflow/api/python/symbol/sbp_symbol.cpp+85可以找到下面定义:
可以看到sbp对接python接口时用的是cfg::SbpParallel这个数据结构,这里它和placement中的cfg::ParallelConf一样,同样下面这个proto文件自动生成出来:
-
oneflow/core/job/sbp_parallel.proto
编译的时候protoc会先根据这个proto文件生成下面三个文件:
-
build/oneflow/core/job/sbp_parallel.pb.h
-
build/oneflow/core/job/sbp_parallel.pb.cc
-
build/of_cfg_proto_python/oneflow/core/job/sbp_parallel_pb2.py
其中前两个文件提供接口用于对SBP数据做序列化,接口都属于oneflow namespace,第三个文件结合tools/cfg中的工具用于生成下面三个文件给Python端来用,文件中的接口属于cfg namespace:
-
build/oneflow/core/job/sbp_parallel.cfg.h
-
build/oneflow/core/job/sbp_parallel.cfg.cpp
-
build/oneflow/core/job/sbp_parallel.cfg.pybind.cpp
OneFlow的内部c++代码中用的是cfg::NdSbp这个数据结构,它其实可以看作是vector<cfg::SbpParallel>,这些数据结构之间的关系,直接看proto文件最为直接:
继续看前面定义SBP的Python接口时所调用的GetSplitSbpParallel、GetBroadcastSbpParallel、GetPartialSumSbpParallel这三个C++函数,位于oneflow/api/python/symbol/sbp_symbol.cpp+41:

它们各自又分别调用了MakeSplitSbpParallel、MakeBroadcastSbpParallel、MakePartialSumSbpParallel这三个函数,位于oneflow/core/job/sbp_parallel.cpp+68:

SymbolOf背后用到的是Symbol这个OneFlow的基本组件,它的实现就不在这里展开了,总体来讲,它是把创建的对象维护到下面这个全局的SymbolMap中:
std::unordered_map<HashEqTraitPtr<const T>, std::shared_ptr<const T>>;
这样以后再用到的话,如果已经存在就不需要重新创建了,直接返回就好。
3
Global Tensor
Tensor的基本概念就不用说了,相信没有人不知道,在OneFlow的设计中,Global Tensor就是为了能够满足Global View所需抽象的一种Tensor,里面需要有前面讲的Placement和SBP相关的属性,下面把OneFlow所有Tensor一并总结列出:
图2
OneFlow的Tensor设计采用了bridge design pattern,把接口和实现做了分离,在Global View的情况下,用到的是上图中的ConsistentTensor(0.7版统称GlobalTensor),可以看到它持有一个指向ConsistentTensorImpl的指针,真正的实现就在ConsistentTensorImpl这个类中,下面是TensorImpl系列类的hierarchy图示:
图3
先看这个图里的基类部分,橙色部分是它包含的数据成员,总体来讲这个基类维护了一些用于反向求导的信息。
再看EagerConsistentTensorImpl,前面讲过,Global View实际上是一个逻辑视角,对应的global tensor实际上也是个逻辑tensor,那么它实际的数据存在于集群的每台机器的每张卡对应的tensor中,即图2中的MirroredTensor中,EagerConsistentTensorImpl持有指向MirroredTensor的指针,MirroredTensor持有指向MirroredTensorImpl的指针,MirroredTensorImpl的子类EagerMirroredTensorImpl中则持有指向TensorStorage的指针,tensor中的数据最终是存在于TensorStorage对象中,它定义在oneflow/core/eager/eager_blob_object.h+32,下面是主要的数据成员:
继续看ConsistentTensorImpl,它持有一个指向ConsistentTensorMeta的指针,TensorMeta这个系列类维护了Tensor的一些元信息,如shape、data_type、device等,如果要是ConsistentTensor的话,还会持有placement和SBP的信息,下面是TensorMeta系列类的hierarchy图示:
图4
可以看到在ConsistentTensorMeta中,维护了Placement和SBP的信息。
本文大概梳理了一下Global View的基本概念和部分具体实现,主要的参考资料是OneFlow的官方代码、官方文档和论文,以下是具体链接:
1.https://github.com/Oneflow-Inc/oneflow
2.https://arxiv.org/abs/2110.15032
3.https://docs.oneflow.org/master/parallelism/02_sbp.html
4.https://docs.oneflow.org/master/parallelism/03_consistent_tensor.html
其他人都在看
欢迎下载体验OneFlow v0.7.0最新版本:https://github.com/Oneflow-Inc/oneflow/https://github.com/Oneflow-Inc/oneflow/
Recommend
-
12
Softmax操作是深度学习模型中最常用的操作之一。在深度学习的分类任务中,网络最后的分类器往往是Softmax + CrossEntropy的组合:
-
14
【深度学习】分布式训练常用技术总结 - OneFlow深度学习框架的个人空间 - OSCHINA - 中文开源技术交流社区
-
8
撰文 | BBuf 原文首发于公众号GiantPandaCV 这篇文章基于自己为OneFl...
-
4
OneFlow v0.5.0正式上线:四大特性实现轻快上手,高效、易用从此兼得 9月27日,在中关村论坛国际技术交易大会上...
-
9
作者 | BBuf 原文首发于公众号GiantPandaCV 0x0.总...
-
10
OneFlow深度学习框架的博客_CSDN博客-前沿技术,业界观点,社区之星领域博主 OneFlow深度学习框架 个人简介:不止于成为世界上最快的开源深度学习框架。https://github.com/Oneflow-Inc...
-
7
OneFlow学习笔记:从Python到C++调用过程分析
-
5
Autograd解析|OneFlow学习笔记
-
5
OneFlow源码解析:Global Tensor ...
-
5
机器学习编译器的前世今生_OneFlow深度学习框架的博客-CSDN博客_webgpu
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK