大約一年以前，我們?展示[1]?了如何在第三代?英特爾至強(qiáng)可擴(kuò)展[2]?CPU (即 Ice Lake) 集群上分布式訓(xùn)練 Hugging Face transformers 模型。最近，英特爾發(fā)布了代號(hào)為 Sapphire Rapids 的第四代至強(qiáng)可擴(kuò)展 CPU，該 CPU 包含了令人興奮的深度學(xué)習(xí)加速新指令。

通過本文，你將會(huì)學(xué)到如何在一個(gè) AWS Sapphire Rapids 集群上加速一個(gè) PyTorch 訓(xùn)練任務(wù)。我們會(huì)使用?英特爾 oneAPI 集合通信庫[3]?(oneAPI Collective Communications Library, oneCCL) 來分布式化訓(xùn)練任務(wù)，并使用?英特爾 PyTorch 擴(kuò)展庫[4]?(Intel Extension for PyTorch，IPEX) 來自動(dòng)使用新指令進(jìn)行性能優(yōu)化。因?yàn)檫@兩個(gè)庫均已集成入 Hugging Face transformers 庫，我們能夠做到在不修改一行代碼的前提下開箱即用地運(yùn)行我們的示例代碼。

在隨后的另一篇文章里，我們還會(huì)探討如何使用 Sapphire Rapids CPU 進(jìn)行推理及其性能提升。

為何你應(yīng)該考慮在 CPU 上訓(xùn)練

在英特爾至強(qiáng) CPU 上訓(xùn)練一個(gè)深度學(xué)習(xí)模型是一個(gè)性價(jià)比高且可擴(kuò)展的方案，在使用分布式訓(xùn)練或者在小數(shù)據(jù)集或中等數(shù)據(jù)集上微調(diào)模型時(shí)尤其如此。

至強(qiáng) CPU 支持一些先進(jìn)的特性，如 512 位先進(jìn)矢量擴(kuò)展 (Advanced Vector Extensions，AVX-512) 以及超線程 (Hyper-Threading) ，這些特性提高了深度學(xué)習(xí)模型的并行性和效率，使得我們可以在得到更好的硬件資源使用率的同時(shí)訓(xùn)練得更快。

另外，一般而言，相比用于訓(xùn)練大型深度學(xué)習(xí)模型的專門硬件如 GPU 等而言，至強(qiáng) CPU 更便宜和易得。至強(qiáng) CPU 還更容易用于其他生產(chǎn)任務(wù)，從網(wǎng)絡(luò)服務(wù)到數(shù)據(jù)庫不一而足，這使得它們成為 IT 基礎(chǔ)設(shè)施的一個(gè)萬用且靈活的選擇。

最后，云用戶還可以通過使用 spot 實(shí)例的方式進(jìn)一步降低在至強(qiáng) CPU 上的訓(xùn)練成本。Spot 實(shí)例使用空閑計(jì)算資源，因此以折扣價(jià)售賣。與按需實(shí)例相比，spot 實(shí)例提供了高至 90% 的顯著的成本節(jié)約。最后同樣重要的是，CPU spot 實(shí)例一般來講比 GPU 實(shí)例更容易獲得。

現(xiàn)在，讓我們看一下 Sapphire Rapids 架構(gòu)引入的新指令。

先進(jìn)矩陣擴(kuò)展 (AMX)：深度學(xué)習(xí)新指令

Sapphire Rapids 架構(gòu)引入了英特爾先進(jìn)矩陣擴(kuò)展 (Advanced Matrix Extensions, AMX) 用于加速深度學(xué)習(xí)工作負(fù)載。用戶只需安裝最新版本的 IPEX 即可受益于新指令，無需更改任何 Hugging Face 代碼。

AMX 指令用于加速矩陣乘法，該操作是深度學(xué)習(xí)批量訓(xùn)練的核心操作。AMX 指令支持 Brain 浮點(diǎn) (BF16) 和 8 比特整型 (INT8) 數(shù)據(jù)類型，覆蓋不同訓(xùn)練場(chǎng)景的加速需求。

AMX 指令引入了新的 2 維 CPU 寄存器，稱作 tile 寄存器。因?yàn)檫@些寄存器在上下文切換時(shí)需要保存和恢復(fù)，所以需要內(nèi)核相關(guān)支持。在 Linux 上，內(nèi)核版本需要在 v5.16 及以上方可支持。

現(xiàn)在，讓我們看看怎樣構(gòu)建一個(gè) Sapphire Rapids CPU 集群用于分布式訓(xùn)練。

構(gòu)建一個(gè) Sapphire Rapids CPU 集群

截至本文撰寫之時(shí)，使用 Sapphire Rapids 服務(wù)器的最簡(jiǎn)單的方式是使用新的亞馬遜 EC2?R7iz[5]?實(shí)例家族。由于它尚在預(yù)覽期，你必須?登記注冊(cè)[6]?以獲得訪問權(quán)限。另外，虛擬機(jī)尚未支持 AMX，因此，我們將使用裸金屬實(shí)例 (r7iz.metal-16xl, 64 vCPU, 512GB RAM) 。

為避免手動(dòng)設(shè)置集群中的每個(gè)節(jié)點(diǎn)，我們首先建立一個(gè)主節(jié)點(diǎn)并依此創(chuàng)建一個(gè)新的亞馬遜機(jī)器鏡像 (Amazon Machine Image，AMI[7]) ，然后，我們用這個(gè) AMI 啟動(dòng)其他節(jié)點(diǎn)。

從網(wǎng)絡(luò)的角度，我們需要如下設(shè)置：

• 打開 22 端口，用于所有實(shí)例上的 ssh 訪問創(chuàng)建和調(diào)試

• 配置從主實(shí)例 (你啟動(dòng)訓(xùn)練的那個(gè)實(shí)例) 到所有其他實(shí)例 (包含主實(shí)例本身) 的免密 SSH 訪問。換句話說，主節(jié)點(diǎn)的 ssh 公鑰必須在所有階段上被授權(quán)

• 允許集群內(nèi)的所有網(wǎng)絡(luò)通信，使得分布式訓(xùn)練可以不受阻礙地運(yùn)行。AWS 提供了安全組這一安全便捷的方式支持這個(gè)功能。我們只需創(chuàng)建一個(gè)安全組，確保所有集群內(nèi)的實(shí)例屬于同一安全組，并允許同一安全組內(nèi)的所有網(wǎng)絡(luò)通信即可，以下是我使用的設(shè)置:

讓我們開始創(chuàng)建集群的主節(jié)點(diǎn)。

設(shè)置主節(jié)點(diǎn)

我們首先啟動(dòng)一個(gè)安裝了 Ubuntu 20.04 AMI (ami-07cd3e6c4915b2d18) 并加入了我們之前創(chuàng)建的安全組的?r7iz.metal-16xl?實(shí)例，用于創(chuàng)建主節(jié)點(diǎn)。該 AMI 雖然只包含了 Linux v5.15.0，但是幸運(yùn)的是英特爾和 AWS 已經(jīng)為這個(gè)內(nèi)核版本打上了 AMX 支持的補(bǔ)丁。因此，我們不需要升級(jí)內(nèi)核至 v5.16。

一旦實(shí)例運(yùn)行起來后，我們 ssh 登錄上它并通過?lscpu?命令檢查 AMX 是否確實(shí)已被支持。你應(yīng)該會(huì)在 flags 部分看到如下內(nèi)容：

然后，我們開始安裝本地依賴以及 Python 依賴。

接著，我們使用?ssh-keygen?創(chuàng)建一個(gè)新的 ssh 密鑰對(duì)，命名為?cluster，并保存在缺省位置 (~/.ssh)。

最后，我們用該實(shí)例創(chuàng)建一個(gè)新的 AMI。

設(shè)置集群

一旦 AMI 準(zhǔn)備就緒，我們用它啟動(dòng)另外 3 個(gè)?r7iz.16xlarge-metal?實(shí)例，不要忘了把他們加入之前創(chuàng)建的安全組中。

當(dāng)這些實(shí)例啟動(dòng)的時(shí)候，我們 ssh 登錄進(jìn)主節(jié)點(diǎn)并完成網(wǎng)絡(luò)設(shè)置。首先，我們編輯位于?~/.ssh/config?的 ssh 配置文件，使其支持從主節(jié)點(diǎn)到其他節(jié)點(diǎn)的免密連接，這里我們只需使用它們各自的私有 IP 及之前創(chuàng)建的密鑰對(duì)即可。以下是我的配置文件。

到此為止，我們可以使用?ssh node [1-3]?去免密連接任何節(jié)點(diǎn)。

在主節(jié)點(diǎn)側(cè)，我們創(chuàng)建一個(gè)?~/hosts?文件，并填入集群中所有節(jié)點(diǎn)的名稱，這些名稱已在上面的 ssh 配置文件中定義。我們用?localhost?代表主節(jié)點(diǎn)，因?yàn)槲覀儠?huì)在該節(jié)點(diǎn)啟動(dòng)訓(xùn)練腳本。我的文件如下所示。

集群現(xiàn)已準(zhǔn)備就緒，讓我們開始訓(xùn)練吧！

啟動(dòng)一個(gè)分布式訓(xùn)練任務(wù)

在本例中，我們將在?SQUAD[8]?數(shù)據(jù)集上微調(diào)一個(gè)用于問答的?DistilBERT[9]?模型。如果你想試試別的示例的話，盡管去做吧。

我們首先冒個(gè)煙，啟動(dòng)一個(gè)單實(shí)例訓(xùn)練任務(wù)。請(qǐng)注意如下幾個(gè)重要的標(biāo)志變量：

• no_cuda?確保使用 CPU 進(jìn)行訓(xùn)練，忽略 GPU

• use_ipex?使能 IPEX 庫，確保 AMX 和 AVX 指令的使用

• bf16?使能 BF16 訓(xùn)練

不必等到任務(wù)完成，我們只運(yùn)行 1 分鐘用于確保所有的依賴已被正常安裝。同時(shí)，這也給了我們一個(gè)單實(shí)例訓(xùn)練的基線性能：1 個(gè) epoch 花費(fèi)大約?26 分鐘。供參考，我們測(cè)量了同樣的任務(wù)在一個(gè)相當(dāng)?shù)?Ice Lake 實(shí)例 (c6i.16xlarge) 上的性能，基于相同的軟件設(shè)置，每個(gè) epoch 需要?3 小時(shí) 30 分鐘。加速比達(dá)到?8 倍。我們已經(jīng)能看到新指令帶來的好處！

現(xiàn)在，讓我們把訓(xùn)練任務(wù)分布式部署到 4 個(gè)實(shí)例上。一個(gè)?r7iz.16xlarge?實(shí)例有 32 個(gè)物理 CPU 核，我們傾向于直接使用物理核而不是虛擬核 (vCPUs) (KMP_HW_SUBSET=1T)。我們決定分配 24 個(gè)核用于訓(xùn)練 (OMP_NUM_THREADS) ，2 個(gè)核用于集合通信 (CCL_WORKER_COUNT) ，剩下的 6 個(gè)核給內(nèi)核和其他進(jìn)程使用。這 24 個(gè)訓(xùn)練線程分配給 2 個(gè) Python 進(jìn)程使用 (NUM_PROCESSES_PER_NODE) 。因此，一個(gè) 4 節(jié)點(diǎn)的集群上共有 8 (NUM_PROCESSES) 個(gè) Python 進(jìn)程。

現(xiàn)在，我們啟動(dòng)分布式訓(xùn)練任務(wù)。

現(xiàn)在，一個(gè) epoch 僅需?7 分 30 秒。

任務(wù)如下圖所示。圖的上半部分是主節(jié)點(diǎn)，同時(shí)你也可以看到其他 3 個(gè)節(jié)點(diǎn)每個(gè)均有 2 個(gè)訓(xùn)練進(jìn)程在運(yùn)行。

4 節(jié)點(diǎn)的完美線性擴(kuò)展需要 6 分 30 秒的訓(xùn)練時(shí)間 (26 分鐘除以 4) 。我們非常接近于這個(gè)理想值，這充分展現(xiàn)了該方法很高的擴(kuò)展性。

結(jié)論

如你所見，在一個(gè)英特爾至強(qiáng)集群上訓(xùn)練 Hugging Face transformers 模型是一個(gè)靈活，可擴(kuò)展且性價(jià)比高的解決方案，特別是在你的數(shù)據(jù)集和模型是小尺寸或者中等尺寸情況下。

以下列出了一些其他可幫助你起步的資源：

• Intel IPEX[10]?GitHub 代碼倉(cāng)庫

• Hugging Face 文檔:?Efficient training on CPU[11]?及?Efficient training on many CPUs[12]

如你有任何問題或反饋，請(qǐng)通過留言方式告訴我們。

感謝閱讀！

英文原文:?https://huggingface.co/blog/intel-sapphire-rapids

譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向?yàn)?transformer-family?模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。

頭圖: 茶葉蛋蛋

文內(nèi)鏈接

[1]?Hugging Face 博客文章: 使用英特爾技術(shù)加速 PyTorch 分布式微調(diào):?https://huggingface.co/blog/accelerating-pytorch