謝賽寧團隊突破高斯?jié)姙R內(nèi)存瓶頸，并行方案實現(xiàn)多顯卡訓練

克雷西 2024-07-10 14:37:03 來源：量子位

克雷西 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

高斯?jié)姙R模型訓練的內(nèi)存瓶頸，終于被謝賽寧團隊和NYU系統(tǒng)實驗室打破！

通過設計并行策略，團隊推出了高斯?jié)姙R模型的多卡訓練方案，不必再受限于單張卡的內(nèi)存了。

用這種方法在4張卡上訓練，可以加速3.5倍以上；如果增加到32卡，又能有額外6.8倍的加速。

該團隊提出的是一種名為Grendel的分布式訓練系統(tǒng)，第一作者是清華姚班校友趙和旭。

通過多卡訓練，不僅速度更快了，研究團隊還突破了大場景、高分辨率環(huán)境下的內(nèi)存局限，生成了更多高斯，3D結果質(zhì)量也更高了。

為了體現(xiàn)這個成果是多么的鵝妹子嚶，謝賽寧本人發(fā)了這樣一個表情包：

（大哭）：不！你不能擴大3D高斯?jié)姙R的規(guī)模，不管是場景、分辨率還是批大小，這消耗的算力和內(nèi)存實在太高了
GPU：我就笑笑不說話

還有網(wǎng)友調(diào)侃說，看來老黃的股票又要漲了。

又快又好的3D生成

多卡并行的方式，突破了單卡的算力和內(nèi)存的限制，讓Grendel可以處理極具挑戰(zhàn)性的大場景（更多高斯粒子數(shù)量）渲染任務。

如在Rubble（4K分辨率）和MatrixCity（1080p分辨率）這兩個大型復雜場景中，Grendel使用多達4000萬和2400萬個高斯粒子，生成了高保真的渲染結果。

在鏡頭不斷拉近的動態(tài)過程當中，也能看出Grendel生成結果的細致性和連貫性。

從數(shù)據(jù)上看，在Mip360和TT&DB數(shù)據(jù)集上，4卡批量訓練后的渲染質(zhì)量（PSNR）與單卡相比也幾乎沒有損失，進一步驗證了Grendel的多卡并行在不同場景上的有效性。

在保證質(zhì)量的基礎上，Grendel還在這兩個數(shù)據(jù)集上實現(xiàn)了3-4倍的速度提升。

特別是在4K場景中，單卡訓練不僅速度慢，還容易出現(xiàn)內(nèi)存不足，所以使用Grendel在多卡上進行并行訓練不僅帶來量的改變，也帶來了質(zhì)的突破。

△OOM：Out of memory，內(nèi)存不足

另外，通過支持更大的批量（batch size）和動態(tài)負載均衡，Grendel可以更充分地利用多GPU資源，避免計算力的浪費。

例如在Mip-NeRF360數(shù)據(jù)集上，Grendel通過增加批量和動態(tài)均衡負載，可以將4卡并行的加速比從2倍提高到近4倍。

那么， Grendel究竟是如何實現(xiàn)的呢？

將高斯?jié)姙R過程拆解

在開始介紹Grendel的原理之前，先來解答這樣一個問題：

單張卡不夠用，用多卡似乎是很容易想到的思路，為什么以前沒見到有多卡方案呢？

這就涉及到了高斯?jié)姙R模型獨特的計算方式——高斯?jié)姙R分為多個不同階段，每個階段的并行粒度不同，需要進行切換。

這與大多數(shù)神經(jīng)網(wǎng)絡模型的單一粒度并行方式大相徑庭，甚至高斯?jié)姙R根本沒用到任何神經(jīng)網(wǎng)絡。

這就導致了現(xiàn)有的針對神經(jīng)網(wǎng)絡訓練的多卡并行方案（如數(shù)據(jù)并行、模型并行等），難以直接應用于3D高斯?jié)姙R。

另外，在高斯?jié)姙R模型的訓練過程中，不同粒度的過程之間需要進行大量的數(shù)據(jù)通信，加大了并行方案的難度。

這也正是Grendel的設計當中需要解決的問題。

首先，Grendel將3D高斯?jié)姙R的訓練過程劃分為三個主要階段——高斯變換（Gaussian transformation）、渲染（rendering）和損失計算（loss computation）。

針對這三個階段Grendel采用混合粒度的并行策略，在不同的訓練階段使用不同的并行粒度。

• 高斯變換階段采用高斯粒子級（Gaussian-wise）并行，將高斯粒子均勻分布到各個GPU節(jié)點；
• 渲染和損失計算階段采用像素級（pixel-wise）并行，將圖像分割成連續(xù)的像素塊，分配到各個GPU節(jié)點。

在高斯變換和渲染階段中間，Grendel通過稀疏的全對全通信，將每個GPU節(jié)點上的高斯粒子傳輸?shù)叫枰鼈冞M行渲染的GPU節(jié)點。

由于每個像素塊只依賴于其覆蓋范圍內(nèi)的高斯粒子子集，Grendel利用空間局部性，只傳輸相關的粒子，從而減少了通信量。

完成損失計算后，在每個GPU節(jié)點上，系統(tǒng)會根據(jù)損失函數(shù)計算渲染管線各個參數(shù)的梯度，并通過反向傳播回傳給高斯粒子的各個屬性參數(shù)。

之后，系統(tǒng)將各GPU計算出的梯度進行聚合，得到批量數(shù)據(jù)的總梯度，并據(jù)此更新高斯粒子的屬性參數(shù)。

接著就是重復從高斯變換到參數(shù)更新的步驟，直到模型收斂或達到預設的訓練輪數(shù)。

另外，為了處理渲染階段的負載不均衡問題，Grendel引入了動態(tài)負載均衡機制：

在訓練過程中，Grendel會記錄每個像素塊的渲染時間，用于預測當前迭代的負載分布，然后動態(tài)調(diào)整像素塊到GPU節(jié)點的分配，盡量使各個節(jié)點的渲染時間接近。

為了進一步提高GPU利用率和訓練吞吐量，Grendel支持批量訓練,即在每個訓練迭代中并行處理多個輸入圖像，并根據(jù)批量大小動態(tài)調(diào)整學習率，以保證訓練的穩(wěn)定性和收斂性。

作者簡介

Grendel的第一作者，是紐約大學計算機博士生、清華姚班19級校友趙和旭，主要研究方向是分布式機器學習。

在清華期間，趙和旭曾在清華NLP實驗室孫茂松團隊參與研究，接受劉知遠副教授的指導。

他還曾經(jīng)在Eric Xing組訪問，優(yōu)化了一個分布式機器學習中的通訊問題，論文被MLsys2023接收。

另外三名華人作者，Weng Haoyang（翁顥洋）也來自姚班；Daohan Lu來自紐約大學，是謝賽寧的博士生；還有Ang Li博士，是一名浙大校友，現(xiàn)在美國PNNL實驗室從事研究。

趙和旭在紐約大學的兩位導師Jinyang Li教授和Aurojit Panda助理教授，以及紐大知名學者、ResNeXt一作、DiT（Sora核心架構）共同作者謝賽寧助理教授，都參與指導了這一項目。

論文地址：
https://arxiv.org/abs/2406.18533
項目主頁：
https://daohanlu.github.io/scaling-up-3dgs/
GitHub：
https://github.com/nyu-systems/Grendel-GS