撰文｜月踏
更新｜趙露陽
在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/
關鍵詞： 學習筆記
