作者|鄭建華 更新|趙露陽
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)。
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。
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.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ù)組。
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]
為了提高性能,從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)過程中用到。
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é)果。
實(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。
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。
之前都是做各種準(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ù)量。
推導(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)呢?
NdSbp 的推導(dǎo)主要包括3步
調(diào)用 GetValidNdSbpSignatureList 獲取有效的簽名
剔除不能包含 nd_sbp_constraints 的簽名
貪心搜索較優(yōu)的簽名
重點(diǎn)看一下有效簽名的獲取。主要是兩步:
GetNdSbpSignatureList: 獲取全部簽名
FilterNdSbpSignatureListByLogicalShape: 過濾不合適的簽名
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?的含義。
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)鍵詞: