了解SOK的原理

在上期文章中，我們對(duì) HugeCTR Sparse Operation Kit （以下簡(jiǎn)稱(chēng)SOK） 的基本功能，性能，以及 API 用法做了初步的介紹，相信大家對(duì)如何使用 SOK 已經(jīng)有了基本的了解。在這期文章中，我們將從在 TensorFlow 上使用 SOK 時(shí)常見(jiàn)的“數(shù)據(jù)并行-模型并行-數(shù)據(jù)并行”流程入手，帶大家詳細(xì)了解 SOK 的原理。

圖 1：SOK 訓(xùn)練的數(shù)據(jù)并行-模型并行-數(shù)據(jù)并行流程

1. Input Dispatcher

Input Dispatcher 的職責(zé)是將數(shù)據(jù)以并行的形式的輸入，分配到各個(gè) GPU 上。總共分為以下幾個(gè)步驟：

第一步：對(duì)每個(gè) GPU 接收到的數(shù)據(jù)并行的 category key，按照 key 求余 GPU 的數(shù)量計(jì)算出其對(duì)應(yīng)的 GPU ID，并分成和 GPU 數(shù)量相同的組；同時(shí)計(jì)算出每組內(nèi)有多少 key。例如圖 2 中，GPU 的總數(shù)為 2，GPU 0 獲取的輸入為［0， 1， 2， 3］，根據(jù)前面所講的規(guī)則，它將會(huì)被分成［0， 2］， ［1， 3］兩組。注意，在這一步，我們還會(huì)為每個(gè)分組產(chǎn)生一個(gè) order 信息，用于 output dispacher 的重排序。

第二步：通過(guò) NCCL 交換各個(gè) GPU 上每組 key 的數(shù)量。由于每個(gè) GPU 獲取的輸入，按照 key 求余 GPU 數(shù)量不一定能夠均分，如圖 3 所示，提前在各個(gè) GPU 上交換 key 的總數(shù)，可以在后面交換 key 的時(shí)候減少通信量。

第三步：使用 NCCL，在各個(gè) GPU 間按照 GPU ID 交換前面分好的各組 key，如圖 4 所示。

圖 4：GPU 間交換 Input key

Step4：對(duì)交換后的所有 key 除以 GPU 總數(shù)，這一步是為了讓每個(gè) GPU 上的 key的數(shù)值范圍都小于 embedding table size 整除 GPU 的數(shù)量，保證后續(xù)在每個(gè) worker 上執(zhí)行 lookup 時(shí)不會(huì)越界，結(jié)果如圖 5 所示。

總而言之，經(jīng)過(guò)上面 4 個(gè)步驟，我們將數(shù)據(jù)并行地輸入，按照其求余 GPU 數(shù)量的結(jié)果，分配到了不同對(duì)應(yīng)的 GPU 上，完成了 input key 從數(shù)據(jù)并行到模型并行的轉(zhuǎn)化。雖然用戶(hù)往每個(gè) GPU 上輸入的都可以是 embedding table 里的任何一個(gè) key，但是經(jīng)過(guò)上述的轉(zhuǎn)化過(guò)程后，每個(gè) GPU 上則只需要處理 embedding table 里 1/GPU_NUMBER 的 lookup。

圖 5：整除 input key

2. Lookup

Lookup 的功能比較簡(jiǎn)單，和單機(jī)的 lookup 的行為相同，就是用 input dispatcher 輸出的 key，在本地的 embedding table 里查詢(xún)出對(duì)應(yīng)的 embedding vector，我們同樣用一個(gè)簡(jiǎn)單的圖來(lái)舉例。注意下圖中 Global Index 代表每個(gè) embedding vector 在實(shí)際的 embedding table 中對(duì)應(yīng)的 key，而 Index 則是當(dāng)前 GPU 的“部分”embedding table 中的 key。

圖 6：使用 Embedding Table 進(jìn)行 Lookup

3. Output Dispatcher

和 input dispatcher 的功能對(duì)應(yīng)，output dispatcher 是將 embedding vector 按照和 input dispatcher 相同的路徑、相反的方向?qū)?embedding vector 返回給各個(gè) GPU，讓模型并行的 lookup 結(jié)果重新變成數(shù)據(jù)并行。

第一步：復(fù)用 input dispatcher 中的分組信息，將 embedding vector 進(jìn)行分組，如圖 7 所示。

圖 7：Embedding vector 的分組

第二步：通過(guò) NCCL 將 embedding vector 按 input dispatcher 的路徑返還，如圖 8 所示。

圖 8：Embedding vector 的返還

第三步：復(fù)用 input dispatcher 第一步驟的結(jié)果，將 embedding vector 進(jìn)行重排序，讓其和輸入的 key 順序保持一致，如圖 9 所示。

圖 9：Embedding vector 的重排序

可以看到， GPU 0 上輸入的［0， 1， 3， 5］，最終被轉(zhuǎn)化為了［0.0， …］， ［0.1， …］， ［0.3， …］， ［0.5， …］ 四個(gè) embedding vector，雖然其中有 3 個(gè) embedding vector 被存儲(chǔ)在 GPU 1 上，但是以一種對(duì)用戶(hù)透明的方式，在 GPU 0 上拿到了對(duì)應(yīng)的 vector。在用戶(hù)看來(lái)，就好像整個(gè) embedding table 都存在 GPU 0 上一樣。

4. Backward

在 backward 中，每個(gè) GPU 會(huì)得到和 input 的 key 所對(duì)應(yīng)的梯度，也就是數(shù)據(jù)并行的梯度。此時(shí)的梯度對(duì)應(yīng)的 embedding vector 可能并不在當(dāng)前 GPU 上，所以還需要做一步梯度的交換。這個(gè)步驟和 output dispatcher 的第三步驟中的工作流程的路徑完全相同，只是方向相反。 仍然以前面的例子舉例，GPU 0 獲取了 key ［0， 1， 3， 5］的梯度，我們把它們分別叫做 grad0， grad1， grad3， grad5；由于 grad1，grad3，grad5 對(duì)應(yīng)的 embedding vector 在 GPU 1 上，所以我們把它們和 GPU 1 上的 grad4， grad6 進(jìn)行交換，最終在得到了 GPU 0 上的梯度為［grad0， grad4， grad6］，GPU 1 上的梯度為［grad1， grad3， grad5， grad5， gard7］。

結(jié)語(yǔ)

以上就是 SOK 將數(shù)據(jù)并行轉(zhuǎn)化為模型并行再轉(zhuǎn)回?cái)?shù)據(jù)并行的過(guò)程，這整個(gè)流程都被封裝在了 SOK 的 Embedding Layer 中，用戶(hù)可以直接調(diào)用相關(guān)的 Python API 即可輕松完成訓(xùn)練。

審核編輯：郭婷