0
  • 聊天消息
  • 系統(tǒng)消息
  • 評論與回復(fù)
登錄后你可以
  • 下載海量資料
  • 學(xué)習(xí)在線課程
  • 觀看技術(shù)視頻
  • 寫文章/發(fā)帖/加入社區(qū)
會員中心
創(chuàng)作中心

完善資料讓更多小伙伴認(rèn)識你,還能領(lǐng)取20積分哦,立即完善>

3天內(nèi)不再提示

了解SOK的原理

星星科技指導(dǎo)員 ? 來源:NVIDIA ? 作者:Kunlun Lee ? 2022-08-15 15:22 ? 次閱讀

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

poYBAGL59JiARZz4AACrm9j_bTw807.png

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

1. Input Dispatcher

Input Dispatcher 的職責(zé)是將數(shù)據(jù)以并行的形式的輸入,分配到各個(gè) GPU 上??偣卜譃橐韵聨讉€(gè)步驟:

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

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

pYYBAGL59J-ANIPxAABP8FWKNo8475.png

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

poYBAGL59KWAChITAADEtKC5X1Q738.png

圖 4:GPU 間交換 Input key

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

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

pYYBAGL59KyAZ2_9AABYzcg3JBk527.png

圖 5:整除 input key

2. Lookup

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

poYBAGL59LKAPnZrAADtkpy9jkU142.png

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

3. Output Dispatcher

和 input dispatcher 的功能對應(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 所示。

pYYBAGL59LuAWMvcAADkv5anG2M175.png

圖 7:Embedding vector 的分組

第二步:通過 NCCL 將 embedding vector 按 input dispatcher 的路徑返還,如圖 8 所示。

poYBAGL59NuAervyAADkv5anG2M680.png

圖 8:Embedding vector 的返還

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

poYBAGL59NCAaZ8HAADMql64990015.png

圖 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 被存儲在 GPU 1 上,但是以一種對用戶透明的方式,在 GPU 0 上拿到了對應(yīng)的 vector。在用戶看來,就好像整個(gè) embedding table 都存在 GPU 0 上一樣。

4. Backward

在 backward 中,每個(gè) GPU 會得到和 input 的 key 所對應(yīng)的梯度,也就是數(shù)據(jù)并行的梯度。此時(shí)的梯度對應(yīng)的 embedding vector 可能并不在當(dāng)前 GPU 上,所以還需要做一步梯度的交換。這個(gè)步驟和 output dispatcher 的第三步驟中的工作流程的路徑完全相同,只是方向相反。 仍然以前面的例子舉例,GPU 0 獲取了 key [0, 1, 3, 5]的梯度,我們把它們分別叫做 grad0, grad1, grad3, grad5;由于 grad1,grad3,grad5 對應(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é)語

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

審核編輯:郭婷

聲明:本文內(nèi)容及配圖由入駐作者撰寫或者入駐合作網(wǎng)站授權(quán)轉(zhuǎn)載。文章觀點(diǎn)僅代表作者本人,不代表電子發(fā)燒友網(wǎng)立場。文章及其配圖僅供工程師學(xué)習(xí)之用,如有內(nèi)容侵權(quán)或者其他違規(guī)問題,請聯(lián)系本站處理。 舉報(bào)投訴
  • gpu
    gpu
    +關(guān)注

    關(guān)注

    28

    文章

    4762

    瀏覽量

    129145
  • API
    API
    +關(guān)注

    關(guān)注

    2

    文章

    1508

    瀏覽量

    62225
  • python
    +關(guān)注

    關(guān)注

    56

    文章

    4802

    瀏覽量

    84889
收藏 人收藏

    評論

    相關(guān)推薦

    了解超聲AFE設(shè)備的CW模式

    電子發(fā)燒友網(wǎng)站提供《了解超聲AFE設(shè)備的CW模式.pdf》資料免費(fèi)下載
    發(fā)表于 10-25 09:23 ?0次下載
    <b class='flag-5'>了解</b>超聲AFE設(shè)備的CW模式

    了解MSP430閃存數(shù)據(jù)保留

    電子發(fā)燒友網(wǎng)站提供《了解MSP430閃存數(shù)據(jù)保留.pdf》資料免費(fèi)下載
    發(fā)表于 10-21 09:34 ?0次下載
    <b class='flag-5'>了解</b>MSP430閃存數(shù)據(jù)保留

    了解Davinci預(yù)覽引擎

    電子發(fā)燒友網(wǎng)站提供《了解Davinci預(yù)覽引擎.pdf》資料免費(fèi)下載
    發(fā)表于 10-15 09:56 ?0次下載
    <b class='flag-5'>了解</b>Davinci預(yù)覽引擎

    了解Davinci調(diào)整器

    電子發(fā)燒友網(wǎng)站提供《了解Davinci調(diào)整器.pdf》資料免費(fèi)下載
    發(fā)表于 10-15 09:32 ?0次下載
    <b class='flag-5'>了解</b>Davinci調(diào)整器

    了解I2C總線

    電子發(fā)燒友網(wǎng)站提供《了解I2C總線.pdf》資料免費(fèi)下載
    發(fā)表于 10-08 11:13 ?2次下載
    <b class='flag-5'>了解</b>I2C總線

    了解電源模塊運(yùn)行限制

    電子發(fā)燒友網(wǎng)站提供《了解電源模塊運(yùn)行限制.pdf》資料免費(fèi)下載
    發(fā)表于 09-24 11:13 ?0次下載
    <b class='flag-5'>了解</b>電源模塊運(yùn)行限制

    了解智能柵極驅(qū)動器

    電子發(fā)燒友網(wǎng)站提供《了解智能柵極驅(qū)動器.pdf》資料免費(fèi)下載
    發(fā)表于 09-21 09:03 ?0次下載
    <b class='flag-5'>了解</b>智能柵極驅(qū)動器

    了解電源器件中的欠壓鎖定

    電子發(fā)燒友網(wǎng)站提供《了解電源器件中的欠壓鎖定.pdf》資料免費(fèi)下載
    發(fā)表于 09-04 10:10 ?0次下載
    <b class='flag-5'>了解</b>電源器件中的欠壓鎖定

    了解TPS56837的軟啟動

    電子發(fā)燒友網(wǎng)站提供《了解TPS56837的軟啟動.pdf》資料免費(fèi)下載
    發(fā)表于 08-29 14:15 ?0次下載
    <b class='flag-5'>了解</b>TPS56837的軟啟動

    了解功能隔離

    電子發(fā)燒友網(wǎng)站提供《了解功能隔離.pdf》資料免費(fèi)下載
    發(fā)表于 08-28 10:37 ?0次下載
    <b class='flag-5'>了解</b>功能隔離

    你認(rèn)識貼片電阻嗎,你對他了解多少?

    你認(rèn)識貼片電阻嗎,你對他了解多少?
    的頭像 發(fā)表于 08-27 15:49 ?544次閱讀
    你認(rèn)識貼片電阻嗎,你對他<b class='flag-5'>了解</b>多少?

    想要了解下大模型知識

    工作需要,想要了解一下大模型算力建設(shè)知識。
    發(fā)表于 08-20 15:31

    SOK在手機(jī)行業(yè)的應(yīng)用案例

    通過封裝 NVIDIA Merlin HugeCTR,Sparse Operation Kit(以下簡稱 SOK)使得 TensorFlow 用戶可以借助 HugeCTR 的一些相關(guān)特性和優(yōu)化加速 GPU 上的分布式 Embedding 訓(xùn)練。
    的頭像 發(fā)表于 07-25 10:01 ?378次閱讀
    <b class='flag-5'>SOK</b>在手機(jī)行業(yè)的應(yīng)用案例

    帶你了解電阻封裝

    了解和掌握電阻的封裝類型、尺寸和性能對于確保電路設(shè)計(jì)的正確性和穩(wěn)定性具有重要意義。希望以上信息能夠幫助您更好地理解電阻封裝并為您在實(shí)際應(yīng)用中提供幫助。
    的頭像 發(fā)表于 05-15 13:28 ?2395次閱讀

    了解EMI濾波器:從基礎(chǔ)到應(yīng)用?

    了解EMI濾波器:從基礎(chǔ)到應(yīng)用?|深圳比創(chuàng)達(dá)電子
    的頭像 發(fā)表于 01-23 10:26 ?604次閱讀
    <b class='flag-5'>了解</b>EMI濾波器:從基礎(chǔ)到應(yīng)用?