IBM聯合PyTorch大幅改進在云計算訓練大型AI模型的效率

IBM研究院與PyTorch合作，發展一個稱作rate_limiter的控件，能夠通過配置訓練時的內存，使得具有數十億參數的模型，也能夠在標準云計算網絡基礎設施，像是以太網絡上運行。

當前的人工智能模型，只要使用簡單的文本題事，就能創造出圖像、歌曲甚至是網站，而這種具有數十億參數的模型，被稱為基礎模型，不需要大量時長的訓練與標記，基礎模型就可以被簡單地重新用于另一項任務。

基礎模型主要在高端的高性能計算（HPC）基礎設施上運行，雖然這些系統很可靠，但IBM提到，對于許多想要自己訓練基礎模型，并且僅供自己使用的人來說，HPC設備是一個極高的門檻，包括高端GPU搭配低延遲InfiniBand網絡系統，以及各種專門的操作流程等。

而IBM研究人員與PyTorch分布式團隊合作，尋找出可在網絡硬件上，訓練大型人工智能模型的方法，而這個聯合小組也已經驗證，可以在Red Hat的OpenShift平臺上，使用基于以太網絡的常規網絡，來擴展和訓練大型模型。

使用PyTorch的完全分片資料平行技術（FSDP），該團隊可以在IBM Cloud上的標準以太網絡，訓練具有110億參數的模型，實現與HPC網絡系統同等規模的高效率訓練。研究人員提到，過去嘗試使用以太網絡，在PyTorch上訓練數十億參數的模型，但是效果不佳，遠低于訓練基礎模型所需要的性能，隨著模型規模增長，只有當GPU能夠完整包存模型副本和訓練狀態時，資料平行訓練的標準方法才能發揮作用。

雖然FSDP或DeepSpeed等新方法可以在訓練期間，有效地將模型和資料分散到多個GPU上，但也只限于HPC系統，而無法在由以太網絡連接的系統上運行。為此，聯合研究團隊開發了一個FSDP API，并構建rate_limiter控件，該組件可以控制用于發送和接收張量所使用的內存，進而緩解系統的內存壓力，使效率較之前提高達4.5倍。

這項研究使用的基礎設施為現成的硬件，該系統在IBM Cloud上運行，由200個節點組成，每個節點有8張Nvidia A100 80GB顯卡，與96 vCPU、1.2TB CPU內存，節點內的顯卡以NVLink連接，顯卡之間帶寬為600 GBps，而節點之間通過兩條100 Gbps以太網絡連接，提供120 GBps可用帶寬。

該GPU系統自5月來持續運行，研究人員配置Red Hat OpenShift容器平臺執行人工智能負載，用于大型人工智能模型的端到端訓練、微調和推理。研究人員提到，這是目前該產業中，第一個使用Kubernetes、PyTorch FSDP API和標準以太網絡，高效率實現高達110億個參數模型的用例，而這將有助于其他組織以更具成本效益的方式，在云計算訓練大型模型。