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億個參數模型的用例，而這將有助於其他組織以更具成本效益的方式，在雲端訓練大型模型。