微軟研究人員日前推出了名為 ZeRO++ 的新技術(shù)，用于優(yōu)化在訓(xùn)練大型 AI 模型時(shí)，容易遇到的數(shù)據(jù)傳輸成本和帶寬限制的難題，可顯著減少大模型訓(xùn)練時(shí)間和成本。
據(jù)悉，ZeRO++ 建立在現(xiàn)有的 ZeRO 傳輸技術(shù)基礎(chǔ)上，并提供增強(qiáng)的通信策略，可提高訓(xùn)練效率，同時(shí)減少訓(xùn)練時(shí)間和成本。
▲ 圖源 微軟
為了減少參數(shù)通信量，ZeRO++ 可對(duì)權(quán)重進(jìn)行量化，其利用基于塊的量化方法來(lái)保持訓(xùn)練精度，這種優(yōu)化的量化過(guò)程相對(duì)原始 Zero 傳輸技術(shù)更快更準(zhǔn)確。為了能夠盡量減少通信開(kāi)銷(xiāo)，ZeRO++ 通過(guò)在每臺(tái)機(jī)器上保持完整的模型副本，以向 GPU 顯存換取通信帶寬。而在梯度通信方面，ZeRO++ 引入了一種名為 qgZ 的新的量化梯度通信方式，可以減少跨節(jié)點(diǎn)的流量和延遲。
▲ 圖源 微軟
這些改進(jìn)的通信技術(shù)大大減少了通信量，微軟研究人員表示，與 ZeRO 相比，ZeRO++ 減少了高達(dá) 4 倍的通信量，提高了訓(xùn)練吞吐量和效率。當(dāng)在每個(gè) GPU 上使用小批量大小時(shí)，在高帶寬集群中，ZeRO++ 相比 ZeRO-3 的吞吐量提高了 28% 至 36%。在低帶寬集群中，與 ZeRO-3 相比，ZeRO++ 實(shí)現(xiàn)了平均 2 倍的加速，使得大模型訓(xùn)練在更多種類(lèi)的集群上更為可行。
例如 Turing-NLG、ChatGPT 和 GPT-4 這樣的大型模型，其訓(xùn)練需要跨多個(gè) GPU 設(shè)備占用大量顯存和計(jì)算資源，而 ZeRO++ 引入了通信優(yōu)化策略，以克服在低帶寬集群上進(jìn)行訓(xùn)練時(shí)原有 ZeRO 傳輸技術(shù)的帶寬限制。目前微軟已經(jīng)放出了相關(guān)技術(shù)文檔，研究人員可以利用 ZeRO++ 更有效地訓(xùn)練模型，在 AI 領(lǐng)域探索新的可能性。
