作者｜鄭建華 更新｜趙露陽

OneFlow 的 Global Tensor 有兩個(gè)必要屬性：

(相關(guān)資料圖)

Placement：決定了 tensor 數(shù)據(jù)分布在哪些設(shè)備上。

SBP：決定了 tensor 數(shù)據(jù)在這些設(shè)備上的分布方式。例如：

split：將切分后的不同部分放到不同設(shè)備；同時(shí)指定切分的 axis。

broadcast：將數(shù)據(jù)復(fù)制到各個(gè)設(shè)備。

如果參與運(yùn)算的 tensor 的 SBP 不一樣，結(jié)果 tensor 的 SBP?是什么呢？例如下面的代碼：

# export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=0 LOCAL_RANK=0# export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=1 LOCAL_RANK=1import oneflow as flowP0 = flow.placement("cpu", ranks=[0, 1])t1 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.split(0))# t1 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.broadcast)t2 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.split(1))t3 = t1 + t2# oneflow.placement(type="cpu", ranks=[0, 1])print(t3.placement)# (oneflow.sbp.split(dim=0),)print(t3.sbp)

t1和t2是分布在相同設(shè)備上的兩個(gè) tensor。t1.sbp是S(0)，在行上切分；t2.sbp是S(1)，在列上切分。

計(jì)算結(jié)果t3的 SBP?不需要用戶手動(dòng)指定，系統(tǒng)可以自動(dòng)推導(dǎo)出t3.sbp為S(0)。這個(gè)過程中的一個(gè)核心步驟，就是 SBP Signature 的推導(dǎo)。

1?

SBP相關(guān)概念

1.1 SBP

SBP是OneFlow中獨(dú)有的概念，其描述了張量邏輯上的數(shù)據(jù)與張量在真實(shí)物理設(shè)備集群上存放的數(shù)據(jù)之間的一種映射關(guān)系。以下內(nèi)容參考SBP官方文檔(https://docs.oneflow.org/master/parallelism/02_sbp.html#sbp)：

詳細(xì)而言：

split表示物理設(shè)備上的 Tensor，是將全局視角的 Tensor 切分得到的。切分時(shí)，需要指定切分的維度。物理設(shè)備上的 Tensor ，經(jīng)過拼接，可以還原得到全局視角的 Tensor 。

broadcast表示全局視角下的 Tensor，會(huì)復(fù)制并廣播到所有的物理設(shè)備上。

partial 表示全局視角下的 Tensor 與物理設(shè)備上的 Tensor 的 形狀相同，但是物理設(shè)備上的值，只是全局視角下 Tensor 的 一部分。以 partial sum 為例，如果我們將集群中所有設(shè)備的張量按位置相加，那么就可以還原得到全局視角的 Tensor。除了 sum 外，min、max 等操作也適用于 partial。

下圖中分別展示了 SBP 的情況，分別是 split(0)、split(1)、broadcast 和 partial sum。

1.2 SBP Signature

SBP Signature即SBP簽名，是OneFlow中獨(dú)創(chuàng)且很重要的概念。本節(jié)以下文字摘自SBP Signature的官方文檔：

對(duì)于一個(gè)孤立的 Tensor，我們可以隨意設(shè)置它的 SBP 屬性。但是，對(duì)于一個(gè)有輸入、輸出數(shù)據(jù)的算子，我們卻不可以隨意設(shè)置它的輸入、輸出的 SBP 屬性。這是因?yàn)殡S意設(shè)置一個(gè)算子輸入輸出的 SBP 屬性，可能不符合全局視角下算子的運(yùn)算法則。

對(duì)于某個(gè)算子，其輸入輸出的一個(gè)特定的、合法的 SBP 屬性組合，稱為這個(gè)算子的一個(gè) SBP Signature。

算子作者根據(jù)算子的運(yùn)算法則，在開發(fā)算子時(shí)，就已經(jīng)羅列并預(yù)設(shè)好該算子所有可能的 SBP Signature。

某一層算子只要有輸入的 SBP 屬性，OneFlow 就可以根據(jù) SBP Signature 推導(dǎo)出該層算子輸出的 SBP 屬性。

所謂的 SBP Signature 自動(dòng)推導(dǎo)，指的是：在給定所有算子的所有合法的 SBP Signature 的前提下，OneFlow 有一套算法，會(huì)基于傳輸代價(jià)為每種合法的 SBP Signature 進(jìn)行打分，并選擇傳輸代價(jià)最小的那個(gè) SBP Signature。這樣使得系統(tǒng)的吞吐效率最高。

如果 OneFlow 自動(dòng)選擇的 SBP Signature，上一層算子的輸出與下一層算子的輸入的 SBP 屬性不匹配時(shí)，那怎么辦呢？OneFlow 會(huì)檢測(cè)到這種不一致，并且在上游的輸出和下游的輸入間插入一類算子，做相關(guān)的轉(zhuǎn)換工作。這類自動(dòng)加入做轉(zhuǎn)換的算子，就稱為 Boxing 算子。

總結(jié)一下，SBP Signature 的要點(diǎn)如下：

每個(gè)算子都需要設(shè)置相應(yīng)的SBP簽名，用于描述數(shù)據(jù)(Tensor)的分布方式。

SBP簽名包括算子的全部輸入、輸出的SBP。缺少（部分）輸入，或（部分）輸出，不能構(gòu)成簽名。

所以SbpSignature.bn_in_op2sbp_parallel是一個(gè)map結(jié)構(gòu)，key就是各個(gè)input和output的標(biāo)識(shí)。

輸入與輸出的SBP簽名組合，在算子的運(yùn)算法則下必須是合法的，算子的作者需要列出合法SBP簽名的候選集。

如果輸入數(shù)據(jù)(input tensor)的SBP與該算子合法的SBP簽名不一致，則為了得到該算子正確計(jì)算所需要的數(shù)據(jù)（tensor），OneFlow 會(huì)在上游的輸出和下游的輸入間插入boxing算子（可能包含nccl等集合通信操作），做自動(dòng)轉(zhuǎn)換工作，這類自動(dòng)轉(zhuǎn)換的過程，就稱為 Boxing。例如，eager global模式下的interpreter在GetBoxingOutput方法中完成Boxing過程。

1.3 NdSbp 及 NdSbpSignature

在上面1.1小節(jié)中，我們了解到SBP用于描述一個(gè)邏輯張量（Tensor），與其對(duì)應(yīng)物理設(shè)備上的映射關(guān)系，那OneFlow中的2D甚至ND SBP又是什么意思呢？

簡(jiǎn)單理解就是，普通的SBP（1D/1維 SBP）只能比較粗粒度地對(duì)張量進(jìn)行切分，譬如split(0)就表示，沿著張量第0維進(jìn)行切分，如果在此基礎(chǔ)上，想進(jìn)行更細(xì)粒度的切分，譬如繼續(xù)沿著第1維再“切一刀”，那么普通的1D SBP就無法做到了，于是需要2D或者ND SBP。

以下文字主要參考官方文檔2D SBP。

我們可以通過ranks=[0, 1, 2, 3]指定tensor的數(shù)據(jù)分布在這4個(gè)設(shè)備上。這4個(gè)設(shè)備組成了一個(gè)一維的設(shè)備矩陣。對(duì)應(yīng)的 SBP 如split(1)，是單個(gè)值，即 1D SBP。

Tensor 數(shù)據(jù)的分布也可以指定為ranks=[[0, 1], [2, 3]]。四個(gè)計(jì)算設(shè)備被劃分為2x2的設(shè)備矩陣。這時(shí)，SBP 也必須與之對(duì)應(yīng)，是一個(gè)長(zhǎng)度為 2 的數(shù)組。對(duì)應(yīng)的NdSbp.sbp_parallel的類型就是數(shù)組。

例如sbp = (broadcast, split(0))。這個(gè) 2D SBP 的含義是：

在 ranks 的第一維度執(zhí)行廣播，將數(shù)據(jù)分別拷貝到group 0(rank [0, 1])和group 1(rank [2, 3])。

在 ranks 的第二維度分別執(zhí)行split(0)。

例如，對(duì)于group 0，將上一步中分配給它的數(shù)據(jù)按行拆分成(1，2)和(3，4)分別給device 0和device 1。

示意圖如下：

如果 Tensor 的數(shù)據(jù)分布形式是多維的，如[[0, 1], [2, 3]]，算子對(duì)應(yīng)的 SBP Signature 也是多維的，所以NdSbpSignature中，每個(gè) input/output 對(duì)應(yīng)的 sbp_parallel 都是數(shù)組。

2?

placement.hierarchy

placement 對(duì)應(yīng)的 C++ 類型是ParallelDesc。構(gòu)造 placement 的 ranks 可以是多維數(shù)組，表示設(shè)備的多維分布矩陣。

placement.hierarchy表示了placement上ranks的層次信息。簡(jiǎn)單理解，hierarchy就是用于描述ranks分布的形狀（類似于shape可用于描述tensor數(shù)據(jù)分布的形狀），hierarchy存儲(chǔ)了 ranks 在各個(gè)維度的 size 信息。

hierarchy 數(shù)組的長(zhǎng)度是 ranks 的維數(shù)。

hierarchy 數(shù)組的元素值，是 ranks 對(duì)應(yīng)維度的 size。

構(gòu)造 hierarchy 的 C++ 代碼可參考GetRanksShape。

運(yùn)行下面的代碼可以觀察 hierarchy 的值。

import oneflow as flowplacements = [ flow.placement("cpu", ranks=[ 0, 1, 2, 3, 4, 5]), flow.placement("cpu", ranks=[[0, 1, 2], [3, 4, 5]]),]for p in placements: print(p.hierarchy)# outputs:# [6]# [2, 3]

3?

tensor add 是哪個(gè)算子？

為了提高性能，從v0.8.0開始，Tensor 的接口基本都通過 C API 提供給Python。

PyTensorObject_methods中定義了很多 Tensor 方法。不過，add 方法是通過 Python C API 的 number protocol 實(shí)現(xiàn)的，指定 PyTensorObject_nb_add 實(shí)現(xiàn)加法操作，實(shí)際由functional::add實(shí)現(xiàn)。

functional::add的定義在functional_api.yaml.pybind.cpp中，這是一個(gè)在構(gòu)建期自動(dòng)生成的文件。順著這個(gè)找，容易發(fā)現(xiàn)示例代碼對(duì)應(yīng)的是AddFunctor。Op的名字是"add_n"，自動(dòng)生成的文件op_generated.cpp中定義了add_n對(duì)應(yīng)的Op是AddNOp。add_n_op.cpp中定義的 AddNOp 的幾個(gè)方法，會(huì)在 SBP Signature 推導(dǎo)過程中用到。

4?

一維 SBP 的推導(dǎo)過程

SBP Signature 推導(dǎo)相關(guān)的類關(guān)系如下：

示例代碼中的 tensor add 操作（t1 + t2），執(zhí)行到 Interpreter的中調(diào)用GetOrInfer時(shí)，會(huì)進(jìn)行 SBP Signature 的推導(dǎo)。在GlobalTensorInferCache::GetOrInfer中，會(huì)以GlobalTensorMetaInferArgs作為 key 把推導(dǎo)結(jié)果存起來，不需要每次都進(jìn)行推導(dǎo)。

GlobalTensorMetaInferArgs的 hash 函數(shù)主要依賴輸入 tensor 的如下信息：

shape

dtype

nd_sbp

placement

consumer_nd_sbp_constraint

不同的 tensor 對(duì)象，只要這些元信息相同，就可以復(fù)用同一個(gè)推導(dǎo)結(jié)果。

UserOpExpr通過GlobalTensorInferCache持有所有推導(dǎo)過的結(jié)果。

4.1 GlobalTensorInferCache 中的推導(dǎo)準(zhǔn)備

實(shí)際的推導(dǎo)在GlobalTensorInferCache::Infer中進(jìn)行。

4.1.1 推導(dǎo) output 的 shape 和 dtype

user_op_expr.InferLogicalTensorDesc的作用主要是推導(dǎo) output 的 shape 和 data_type，結(jié)果保存到 output_mut_metas。這里涉及到 UserOpExpr 和 Op 兩個(gè)模塊之間的交互關(guān)系。后面會(huì)總結(jié)一下幾個(gè)模塊之間的部分交互接口。

user_op_expr.InferLogicalTensorDesc中用到的兩個(gè)函數(shù)對(duì)象，在Op中定義，并注冊(cè)到OpRegistry中。OpRegistryResult 的函數(shù)對(duì)象來自 Op 注冊(cè)。示例代碼中 tensor add 對(duì)應(yīng)的 Op 是 AddNOp。

AddNOp 場(chǎng)景的實(shí)際調(diào)用順序示例如下：

user_op_expr.InferLogicalTensorDesc

logical_tensor_desc_infer_fn_->AddNOp::InferLogicalTensorDesc

out.shape = in[0].shape

dtype_infer_fn_->AddNOp::InferDataType

out.data_type = in[0].data_type

4.1.2 構(gòu)造 UserOp

MakeOp(user_op_expr...)返回一個(gè)Operator，具體類型是UserOp（參考之前靜態(tài)圖的討論）。這個(gè)對(duì)象負(fù)責(zé)執(zhí)行具體的推導(dǎo)。

CheckInputParallelDescIdentical要求所有 inputs 的 placement 是一致的。因?yàn)檫@里是針對(duì)UserOp做的推導(dǎo)，例如 tensor add、matmul 等操作，操作數(shù)都在相同的設(shè)備時(shí)，這些操作才能直接計(jì)算，否則，就需要通過系統(tǒng) Op 將數(shù)據(jù)搬運(yùn)到一起，再進(jìn)行計(jì)算。

既然所有 inputs 的 placement 都是一樣的，那就用第一個(gè)作為代表，并賦值給 UserOp 保存。

op->InferParallelSignatureIf()的作用是將 placement 填充到op.bn2parallel_desc_。

對(duì)于 AddNOp 來說，key是in_0, in_1, out_0，value 是 inputs[0].placement。

infer_args.MakeInputBlobDescs操作用偽碼表示如下：

# for each input index iblob_descs[i].shape = inputs[i].shapeblob_descs[i].stride = inputs[i].strideblob_descs[i].data_type = inputs[i].data_type

infer_args.MakeNdSbpInferHints操作用偽碼表示如下：

# for each input index ihints[i].parallel_desc = inputs[i].parallel_deschints[i].blob_desc = blob_descs[i]hints[i].nd_sbp = inputs[i].nd_sbp

blob_descs的作用是為了構(gòu)造pd_infer_hints，pd_infer_hints是為了構(gòu)造NdSbpInferHint4Ibn，將相關(guān)信息封裝到這個(gè)函數(shù)對(duì)象中。這個(gè)函數(shù)對(duì)象被傳遞給UserOp進(jìn)行推導(dǎo)。在UserOp中，通過這個(gè)函數(shù)對(duì)象，根據(jù)input/output的標(biāo)識(shí)bn（blob name），獲取NdSbpInferHint，從而可以得到上述元信息。

UserOp推導(dǎo)完畢后，GlobalTensorInferCache會(huì)將 inputs/outputs 的元信息，連同推導(dǎo)得到的 NdSbp ，一起保存到GlobalensorInferResult。

4.2 Operator 中的推導(dǎo)準(zhǔn)備

Operator::InferNdSbpSignatureIf中，調(diào)用InferNdSbpSignature進(jìn)行實(shí)際的推導(dǎo)，然后調(diào)用FillNdSbpSignature保存推導(dǎo)結(jié)果。

InferNdSbpSignature是一個(gè)虛函數(shù)。UserOp會(huì)先檢查Op有沒有定義自己的 SBP Signature 推導(dǎo)函數(shù)，AddNOp 沒有這方面的函數(shù)，就調(diào)用 Operator::InferNdSbpSignature。

InferNdSbpSignature 中會(huì)根據(jù) parallel_desc.hierarchy() 判斷是 1D SBP，還是 ND SBP。

先只看 1D SBP 的情況。調(diào)用傳入的 NdSbpInferHint4Ibn 函數(shù)對(duì)象，查到 GlobalTensorInferCache 中創(chuàng)建的 NdSbpInferHint，轉(zhuǎn)為 NdSbpInferHint 并存到 map 中。因?yàn)槭且痪S的，所以只需要取 sbp_parallel 的第一個(gè)元素。然后調(diào)用 InferSbpSignature（名字中少了 Nd），將推導(dǎo)結(jié)果寫到 SbpSignature。

無論是一維還是多維，結(jié)果的類型都是 NdSbpSignature。所以要將 SbpSignature 轉(zhuǎn)為 NdSbpSignature。

Operator::InferSbpSignature的作用主要是構(gòu)造兩個(gè)函數(shù)對(duì)象，SbpInferHint4Ibn 和 CalcOrderValue4SbpSig，然后調(diào)用子類 override 的、同名重載的虛函數(shù) InferSbpSignature。

SbpInferHint4Ibn?是將傳入的 map 數(shù)據(jù)封裝到函數(shù)對(duì)象中，用于查詢輸入輸出的元信息。

CalcOrderValue4SbpSig給每個(gè) SbpSignature 計(jì)算一個(gè)序值，用于對(duì)簽名進(jìn)行排序。

InferSbpSignature 也是一個(gè)虛函數(shù)。因?yàn)?AddNOp 沒有定義簽名推導(dǎo)函數(shù)，會(huì)調(diào)用 Operator::InferSbpSignature。

4.3 SbpSignature 的推導(dǎo)

之前都是做各種準(zhǔn)備，Operator::InferSbpSignature里才進(jìn)行真正的推導(dǎo)。簡(jiǎn)單講就3步：

獲取候選集

過濾不合適的簽名

排序

4.3.1 SbpSignature 的候選集

調(diào)用 GetValidNdSbpSignatureList會(huì)獲取 SbpSignature 的候選集。在這個(gè)函數(shù)中，先調(diào)用 GetNdSbpSignatureList獲取初步的候選集，再通過FilterNdSbpSignatureListByLogicalShape過濾得到正確可用的候選集。候選集都保存到sbp_sig_list。

GetNdSbpSignatureList是一個(gè)虛函數(shù)，UserOp 實(shí)現(xiàn)了自己的版本。這個(gè)函數(shù)中最核心的操作就是val_->get_nd_sbp_list_fn，實(shí)際調(diào)用AddNOp::GetSbp。UserOpSbpContext是 UserOp 與 AddNOp 等類之間的協(xié)議接口的一部分。

如前所述，提供 SBP Signature 的候選集，是算子的責(zé)任。AddNOp這個(gè)算子比較簡(jiǎn)單，只給出兩類簽名：

對(duì)輸入 tensor 的 shape 的每個(gè) axis i，所有的 input/output 都創(chuàng)建一個(gè) split(i)。

對(duì)于 tensor add 來說，input/output 的 shape 一樣才能直接計(jì)算，所以 split 的 axis 也都一樣。

所有的 input/output 都創(chuàng)建一個(gè) partialsum。

broadcast?的情況會(huì)在 Operator 中默認(rèn)設(shè)置，因?yàn)槔碚撋纤衖nputs/outputs都應(yīng)該支持以broadcast的方式進(jìn)行運(yùn)算。

候選集數(shù)據(jù)示例如下：

{"sbp_signature":[{"bn_in_op2sbp_parallel":{"in_0":{"split_parallel":{"axis":"0"}},"in_1":{"split_parallel":{"axis":"0"}},"out_0":{"split_parallel":{"axis":"0"}}}},{"bn_in_op2sbp_parallel":{"in_0":{"split_parallel":{"axis":"1"}},"in_1":{"split_parallel":{"axis":"1"}},"out_0":{"split_parallel":{"axis":"1"}}}},{"bn_in_op2sbp_parallel":{"in_0":{"partial_sum_parallel":{}},"in_1":{"partial_sum_parallel":{}},"out_0":{"partial_sum_parallel":{}}}},{"bn_in_op2sbp_parallel":{"in_0":{"broadcast_parallel":{}},"in_1":{"broadcast_parallel":{}},"out_0":{"broadcast_parallel":{}}}}]}

4.3.2 過濾不合適的簽名

分兩步過濾不合適的簽名

FilterAndCheckValidSbpSignatureListByLogicalShape中，對(duì)于每個(gè)輸入tensor ibn，簽名中 ibn 的 split axis，必須小于 tensor ibn 的 shape axes 數(shù)量。換句話說，如果 tensor 是二維的，就無法接受split(2)，只能是split(0)或split(1)。

FilterSbpSignatureList的作用是檢驗(yàn)sbp_sig_conf約束，也就是從GlobalTensorInferCache一路傳過來的參數(shù)nd_sbp_constraints。這個(gè)過濾規(guī)則要求，符合條件的簽名，其內(nèi)容必須包含sbp_sig_conf。

4.3.3 簽名排序

SortSbpSignatureListByCopyCost對(duì)候選簽名進(jìn)行排序。

優(yōu)先按 OrderValue 比較

OrderValue 相等時(shí)，按 CopyCost 比較 二者都是較小的值優(yōu)先。

OrderValue4SbpSig是對(duì)CalcOrderValue4SbpSig的封裝，預(yù)先計(jì)算所有簽名的 OrderValue 存到 map 中，便于 sort 函數(shù)查找。IbnCopyCost4SbpSig也是同理。

回過頭來看CalcOrderValue4SbpSig的定義。因?yàn)锳ddNOp是有輸入的，對(duì)于每個(gè)輸入 tensor ibn 會(huì)加上一個(gè)權(quán)重，當(dāng) ibn 的 sbp 與 簽名中對(duì)應(yīng)的 sbp 相同時(shí)，權(quán)重值為-10，即增加了選中的機(jī)會(huì)，因?yàn)?sbp 一致通常就不需要數(shù)據(jù)搬運(yùn)。而parallel_num的條件判斷在UserOp下應(yīng)該是都成立的。

當(dāng) sbp_sig_conf 不空時(shí)，CalcOrderValue4SbpSig 直接返回0。因?yàn)槿绻灻话?sbp_sig_conf，即使 SBP 都一致，簽名也不一定符合要求，所以直接返回0。

簽名成本由ComputeIbnCopyCost4SbpSig計(jì)算。主要是根據(jù)輸入和簽名的 sbp 計(jì)算 cost：

如果 sbp 一致，cost 為0

partial_sum 和 broadcast 的 cost 都是一個(gè)超大的數(shù)字。

否則 cost 等于 input tensor 的數(shù)據(jù)傳輸字節(jié)數(shù)量。

4.4 推導(dǎo)結(jié)果

推導(dǎo)得到的nd_sbp_signature如下：

{"bn_in_op2nd_sbp":{"in_0":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]},"in_1":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]},"out_0":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]}}}

示例代碼中，如果一個(gè)輸入是split，另一個(gè)是broadcast，推導(dǎo)的簽名結(jié)果都是broadcast。如果推斷的sbp簽名是split，是否能減少數(shù)據(jù)搬運(yùn)呢？

5?

NdSbp 的推導(dǎo)過程

NdSbp 的推導(dǎo)主要包括3步

調(diào)用 GetValidNdSbpSignatureList 獲取有效的簽名

剔除不能包含 nd_sbp_constraints 的簽名

貪心搜索較優(yōu)的簽名

重點(diǎn)看一下有效簽名的獲取。主要是兩步：

GetNdSbpSignatureList: 獲取全部簽名

FilterNdSbpSignatureListByLogicalShape: 過濾不合適的簽名

5.1 NdSbp 簽名的候選集

GetNdSbpSignatureList 核心是兩步：

GetSbpSignaturesIf: 得到一維的簽名（和 1D SBP 的情況相同）

DfsGetNdSbpSignature: 根據(jù)一維簽名拓展到多維

這個(gè)過程，如果深入到數(shù)據(jù)細(xì)節(jié)去看，會(huì)涉及 input/output、ranks、NdSbp 等多個(gè)維度，有點(diǎn)抽象復(fù)雜。如果從官方文檔 2D SBP中說明的 ranks 和 NdSbp 的物理含義出發(fā)，會(huì)更容易理解。

以ranks=[[0, 1, 2], [3, 4, 5]]為例（ranks=[r1, r2]）

這是一個(gè)二維的設(shè)備矩陣/陣列。算子的每個(gè)輸入、輸出也都有兩個(gè) sbp，NdSbpSignature 中的 value 是二維的，有兩個(gè)槽位。假設(shè) Op 的 1D Sbp 有 n 個(gè)簽名。

從形式上看，NdSbpSignature 是先按 bn 組織數(shù)據(jù)。但是從數(shù)據(jù)分布的過程看，是先按SbpSignature組織數(shù)據(jù)。一個(gè) NdSbpSignature 等價(jià)于 SbpSignature 數(shù)組。NdSbp中的每個(gè)槽位，都表示一個(gè) 1D Sbp 的數(shù)據(jù)分布（所有的 input/output一起分布）。

比如第 0 個(gè)槽位，就是在r1和r2這兩個(gè) sub group 之間分布數(shù)據(jù)，這個(gè)分布必須是一個(gè)有效的 1D SbpSignature（所有的 input/output一起分布）。

第 1 個(gè)槽位，對(duì)于r1，就是將分配給它的數(shù)據(jù)子集，再根據(jù)一個(gè) SbpSignature 進(jìn)行分布（所有的 input/output一起分布）。

所以，只需要按 SbpSignature整體 填滿兩個(gè)槽位就行。每個(gè)槽位各有 n 種可能，一共有 n*n 個(gè)候選簽名。這樣生成的候選集是完整的，不會(huì)漏掉候選項(xiàng)。這應(yīng)該就是 direct product of 1D sbp signatures?的含義。

6

模塊間協(xié)作關(guān)系

SbpSignature 推導(dǎo)的實(shí)現(xiàn)用了大量 functional 的代碼。應(yīng)該是為了不同模塊間的信息屏蔽，或者父類、子類之間的邏輯復(fù)用、信息傳遞等目的，很多信息都封裝到 function 中，需要時(shí)再檢索、轉(zhuǎn)換。

下圖展示了不同模塊之間的部分關(guān)系：

參考資料

oneflow v0.9.1（https://github.com/Oneflow-Inc/oneflow/tree/0ea44f45b360cd21f455c7b5fa8303269f7867f8/oneflow）

SBP Signature（https://docs.oneflow.org/master/parallelism/02_sbp.html#sbp-signature）

2D SBP（https://docs.oneflow.org/master/parallelism/04_2d-sbp.html）

placement api（https://oneflow.readthedocs.io/en/master/tensor_attributes.html?highlight=placement#oneflow-placement）

https://segmentfault.com/a/1190000042625900

其他人都在看

ChatGPT背后的經(jīng)濟(jì)賬

OneFlow v0.9.0正式發(fā)布

開源ChatGPT要來了；軟件2.0智能革命

比快更快，開源Stable Diffusion刷新作圖速度

OneEmbedding:單卡訓(xùn)練TB級(jí)推薦模型不是夢(mèng)

GLM訓(xùn)練加速：性能最高提升3倍，顯存節(jié)省1/3

“一鍵”模型遷移，性能翻倍，多語言AltDiffusion推理速度超快

歡迎Star、試用OneFlow最新版本：https://github.com/Oneflow-Inc/oneflow/https://github.com/Oneflow-Inc/oneflow/???

關(guān)鍵詞：