AI21 Labs 使用 Amazon EC2 P4d 執行個體和 PyTorch 訓練了具有 1,780 億個參數的語言模型

AI21 Labs 使用機器學習技術開發專注於理解語義的語言模型，並在 2021 年設立了訓練當時剛發布不久的 Jurassic-1 Jumbo 的目標。Jurassic-1 Jumbo 是一種具有 1,780 億個參數的自動迴歸語言模型。註冊參加 Beta 測試的開發人員可以存取 Jurassic-1 Jumbo，並可以立即開始為自己的使用案例自訂模型。該軟體新創公司希望有效地訓練該模型，因此尋求與 Amazon Web Services (AWS) 合作，並使用 Amazon Elastic Compute Cloud (Amazon EC2) 建置了一個解決方案。Amazon EC2 是一種提供安全且可調整大小之雲端運算容量的 Web 服務。選擇 Amazon EC2 讓該公司能夠控制訓練過程，包括節點分配。

為了取得強大的運算和聯網功能，該公司選擇使用 Amazon EC2 P4d 執行個體。該執行個體可為雲端中的機器學習訓練和高效能運算應用程式提供高輸送量和低延遲網路。AI21 Labs 使用 Amazon EC2 P4d 執行個體，透過將模型訓練分散到數百個 GPU 上來取得所需的效能和記憶體，並透過其 Jurassic-1 Jumbo 模型將自然語言處理技術作為服務提供給開發人員。由於該公司現在訓練和控制其大型模型，因此他們可以更輕鬆地開發相同規模的新模型並進行創新。

AI21 Labs 成立於 2017 年，他們有雙重使命：進行自然語言處理研究和開發採用人工智慧技術的閱讀和寫作產品。其旗艦產品 Wordtune 是一款智慧型寫作和編輯助理，於 2020 年 10 月推出，現已發展到支援近百萬使用者。其另一個主要產品 AI21 Studio 提供對其 Jurassic-1 語言模型以及自訂模型開發的 API 存取。AI21 Labs 共同創辦人兼聯合執行長 Yooav Shoham 表示：「我們是一小撮提供語言模型服務的公司之一，我們讓包括獨立開發人員和跨國企業在內的任何人都能夠利用先進的自然語言處理技術構建應用程式和服務。此外，我們追求科學創新，並研究解決方案以克服具有如此規模和複雜性的模型所帶來的軟體工程挑戰。」

為了有效地訓練其第一個深度學習巨型模型並支援模型的高擴展性和效能需求，AI21 Labs 需要強大的運算能力、高速網路以及技術支援和指導。出於這些原因，該公司於 2021 年初開始在 AWS 上實作解決方案，並選擇使用 Amazon EC2 P4d 執行個體訓練模型。這些執行個體部署在稱為 Amazon EC2 UltraCluster 的超大規模叢集中，這些叢集提供超過 4,000 個 NVIDIA A100 GPU、PB 級不阻塞式網路基礎設施和高輸送量低延遲的儲存。 

該公司還進行進一步的優化，採用低延遲高頻寬的 GPUDirectRDMA 以及 Elastic Fabric Adapter (EFA)，後者是一種用於 Amazon EC2 執行個體的網路介面，可讓客戶在 AWS 上大規模執行需要進行高度節點間通訊的應用程式。鑑於模型的規模，該團隊需要使用平行處理來縮短訓練時間，因此他們希望 AWS 上的網路功能能夠支援其分散式訓練和模型平行處理。AI21 Labs 的 Jurassic 技術主管 Opher Lieber 表示：「Amazon EC2 P4d 執行個體在 EFA 上提供 400 Gbps 的高效能網路。當擴展到數百個 GPU 時，GPU 與 GPU 之間的聯網速度會直接影響有效擴展和保持成本效益的能力。」

AI21 Labs 首先在啟用 EFA 的 Amazon EC2 P4d 執行個體上建置其程式碼庫。然後他們測試並驗證了其多節點訓練方法的效能和有效擴展。接下來，該團隊啟動了全規模模型的快速訓練 (使用數百個 GPU) 來驗證功能和效能。之後，他們能夠在 AWS 上訓練其 Jurassic-1 Jumbo 模型。在協同運作方面，該公司選擇使用內部解決方案。該解決方案使用 AWS 軟體開發套件適用於 Python 的 AWS SDK (Boto3) 分配執行個體。適用於 Python 的 AWS SDK (Boto3) 讓客戶可以輕鬆地將自己的 Python 應用程式、程式庫或指令碼與各種 AWS 服務相整合。

在儲存方面，AI21 Labs 選擇使用 Amazon Simple Storage Service (Amazon S3)。該服務提供領先業界的可擴展性、資料可用性、安全性和效能。「在 AWS 團隊的協助下，我們能夠在 Amazon S3 上實現非常好的效能。不管從效能還是價格來說，這都是一個顯而易見的選擇，」Lieber 說道。該團隊使用 Amazon S3 儲存貯體以分散式方式高效地儲存和載入檢查點。為了記錄訓練進度和事件，該團隊使用了 Amazon CloudWatch (一項監控與可觀察性服務)。 

在實作其解決方案的過程中，AI21 Labs 獲得了 AWS 的協助。其團隊諮詢了 AWS 專家，專家就服務水準、架構和硬體相關的問題及疑慮提供了指導。此外，該公司使用 AWS 上的 PyTorch (一種開放原始碼深度學習架構，可用於輕鬆開發機器學習模型並將其部署到生產中) 提高了 Jurassic-1 Jumbo 的效能。 

AI21 Labs 在幾個月的時間內便完成了模型的訓練，並於 2021 年 6 月結束訓練工作。該公司新的 megamodel 模型是一種自動迴歸語言模型，具有 1,780 億個參數，可與競爭對手提供的產品相媲美。該模型還提供一個具有 256,000 個項目的差異化詞彙，從而擴展文字表示能力和對命名實體的支援。該公司現在透過其 AI21 Studio 產品在公開測試中提供 Jurassic-1 Jumbo (及相關的 Jurassic-1 Large，後者具有 70 億個參數)。藉助該服務，廣大開發人員可以在 Jurassic-1 Jumbo 模型的基礎上建置產品。AI21 Labs 已經看到其產品在許多產業得到採用，包括行銷、內容創作、遊戲、醫學研究、汽車、電信和金融產業。

由於 AI21 Labs 擁有並可以直接存取自己的模型，因此其能夠在不依賴第三方的情況下進行調整和創新，並且可以持續探索創新目標，這對他們履行使命至關重要。AI21 Labs 目前正在開發其他模型的原型，並計畫對這些模型進行大規模訓練。「訓練和擁有我們自己的巨型模型將繼續成為我們 Wordtune 和 AI21 Studio 產品的關鍵差異化因素，」Shoham 說道。