撰文|月踏
更新|趙露陽
在OneFlow中,Global View也被稱作一致性視角,用來把一個物理集群抽象成一個邏輯設備,并使用Placement和SBP來實現這種抽象。本文從基本概念、數據結構、接口實現等方面對其進行學習和總結。
1
Placement
1.1 使用示例
Placement用來描述設備信息,包括設備類型、設備分布信息,先看一個具體的使用示例,然后根據這個示例來做分析:
import?oneflow?as?flowx?=?flow.placement(type="cuda",?ranks=[[0,?1,?2,?3],?[4,?5,?6,?7]])
type(x)的輸出為:
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可以找到下面的定義:
ONEFLOW_API_PYBIND11_MODULE("",?m)?{py::class_
通過上面的多個def接口可以看到,通過調用PlacementSymbolExportUtil::CreateParallelDescSymbol()來構造Placement對象,這個函數是個重載函數,定義在同一個文件中,多個重載版本只是參數有區別,其中的創建Placement的邏輯基本一致,下面列一個重載版本作為示例:
// create Symbol
由這個函數的返回值可見,Placement在C++中對應的數據結構是ParallelDesc,這個數據結構后面再說,現在先繼續看創建邏輯,這里繼續調用了CreateParallelDesc函數,同樣定義在PlacementSymbolExportUtil這個類中:
static Maybe
這里最重要的是調用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的主要實現:
Maybe
再繼續看下GetParallelDescSymbol是怎么根據cfg::ParallelConf的對象得到ParallelDesc類型對象的,GetParallelDescSymbol定義在oneflow/core/framework/instructions_builder.cpp+230:
Maybe
大概過程就是在一個全局表里面去查有沒有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,我們只需要關注這個數據結構就行:
class ParallelDesc final { ... ... Optional
這里面的數據結構看起來很復雜,我也不完全明白所有成員的含義,但歸根結底這里數據成員的值還都是根據cfg::ParallelConf中的內容來的,在前面調用GetParallelDescSymbol時,如果全局表中沒有找到,就會根據cfg::ParallelConf類型對象創建一個ParallelConf類型對象,再根據這個ParallelConf類型對象創建一個ParallelDesc類型對象,在ParallelDesc的構造函數中會調用類內的MaybeInit函數,位于oneflow/core/job/parallel_desc.cpp+112,這里面會完成ParallelDesc數據成員的賦值:
Maybe
以上就是在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環境做下面這個簡單的示例:
import oneflow as flows=flow.sbp.split(1)b=flow.sbp.broadcastp=flow.sbp.partial_sum
type(s)、type(b)、type(p)的輸出如下:
print(s)、print(b)、print(p)的輸出如下:
oneflow.sbp.split(axis=1)oneflow.sbp.broadcastoneflow.sbp.partial_sum
2.3 追蹤代碼
先找入口,在python/oneflow/__init__.py+196:
from?.?import?sbp
這用到了同目錄下的這個module文件:python/oneflow/sbp.py,內容如下:
import?oneflowfrom oneflow.framework.distribute import split_sbp as splitimport?oneflow._oneflow_internalsbp = oneflow._oneflow_internal.sbp.sbpbroadcast = oneflow._oneflow_internal.sbp.broadcast()partial_sum?=?oneflow._oneflow_internal.sbp.partial_sum()#?其中split_sbp的定義如下def split_sbp(axis: int) -> oneflow._oneflow_internal.sbp.sbp: assert type(axis) is int return?oneflow._oneflow_internal.sbp.split(axis)
可見split、broadcast和partial_sum都是定義在pybind定義的_oneflow_internal這個module的子module sbp的內部,在oneflow/api/python/symbol/sbp_symbol.cpp+85可以找到下面定義:
ONEFLOW_API_PYBIND11_MODULE("sbp", m) { m.attr("max_split_axis") = kMaxSplitAxis; py::class_
可以看到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
message SplitParallel { required int64 axis = 1; }message BroadcastParallel { }message PartialSumParallel { }message SbpParallel { oneof parallel_type { SplitParallel split_parallel = 1; BroadcastParallel broadcast_parallel = 2; PartialSumParallel partial_sum_parallel = 3; }}message SbpSignature { map
?
繼續看前面定義SBP的Python接口時所調用的GetSplitSbpParallel、GetBroadcastSbpParallel、GetPartialSumSbpParallel這三個C++函數,位于oneflow/api/python/symbol/sbp_symbol.cpp+41:
Maybe
它們各自又分別調用了MakeSplitSbpParallel、MakeBroadcastSbpParallel、MakePartialSumSbpParallel這三個函數,位于oneflow/core/job/sbp_parallel.cpp+68:
Maybe
SymbolOf背后用到的是Symbol這個OneFlow的基本組件,它的實現就不在這里展開了,總體來講,它是把創建的對象維護到下面這個全局的SymbolMap中:
std::unordered_map
這樣以后再用到的話,如果已經存在就不需要重新創建了,直接返回就好。
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,下面是主要的數據成員:
class TensorStorage { ... size_t blob_bytes_; std::unique_ptr
繼續看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
其他人都在看
資源依賴的“詛咒”?
“遠見者”特斯拉AI主管Karpathy
TVM:成為深度學習領域的“Linux”
對抗軟件系統復雜性:恰當分層,不多不少
解讀 Pathways (二):向前一步是 OneFlow
OneFlow v0.7.0發布:全新分布式接口,LiBai、Serving等一應俱全
歡迎下載體驗OneFlow v0.7.0最新版本:https://github.com/Oneflow-Inc/oneflow/https://github.com/Oneflow-Inc/oneflow/
關鍵詞: 學習筆記