克雷西 發自 凹非寺
量子位 | 公衆号 QbitAI
高斯潑濺模型訓練的記憶體瓶頸，終于被謝賽甯團隊和NYU系統實驗室打破！
通過設計并行政策，團隊推出了高斯潑濺模型的多卡訓練方案，不必再受限于單張卡的記憶體了。
用這種方法在4張卡上訓練，可以加速3.5倍以上；如果增加到32卡，又能有額外6.8倍的加速。
該團隊提出的是一種名為Grendel的分布式訓練系統，第一作者是清華姚班校友趙和旭。
通過多卡訓練，不僅速度更快了，研究團隊還突破了大場景、高分辨率環境下的記憶體局限，生成了更多高斯，3D結果品質也更高了。
為了展現這個成果是多麼的鵝妹子嘤，謝賽甯本人發了這樣一個表情包：
（大哭）：不！你不能擴大3D高斯潑濺的規模，不管是場景、分辨率還是批大小，這消耗的算力和記憶體實在太高了
GPU：我就笑笑不說話
還有網友調侃說，看來老黃的股票又要漲了。
又快又好的3D生成
多卡并行的方式，突破了單卡的算力和記憶體的限制，讓Grendel可以處理極具挑戰性的大場景（更多高斯粒子數量）渲染任務。
如在Rubble（4K分辨率）和MatrixCity（1080p分辨率）這兩個大型複雜場景中，Grendel使用多達4000萬和2400萬個高斯粒子，生成了高保真的渲染結果。
在鏡頭不斷拉近的動态過程當中，也能看出Grendel生成結果的細緻性和連貫性。
從資料上看，在Mip360和TT&DB資料集上，4卡批量訓練後的渲染品質（PSNR）與單卡相比也幾乎沒有損失，進一步驗證了Grendel的多卡并行在不同場景上的有效性。
在保證品質的基礎上，Grendel還在這兩個資料集上實作了3-4倍的速度提升。
特别是在4K場景中，單卡訓練不僅速度慢，還容易出現記憶體不足，是以使用Grendel在多卡上進行并行訓練不僅帶來量的改變，也帶來了質的突破。
另外，通過支援更大的批量（batch size）和動态負載均衡，Grendel可以更充分地利用多GPU資源，避免計算力的浪費。
例如在Mip-NeRF360資料集上，Grendel通過增加批量和動态均衡負載，可以将4卡并行的加速比從2倍提高到近4倍。
那麼， Grendel究竟是如何實作的呢？
将高斯潑濺過程拆解
在開始介紹Grendel的原理之前，先來解答這樣一個問題：
單張卡不夠用，用多卡似乎是很容易想到的思路，為什麼以前沒見到有多卡方案呢？
這就涉及到了高斯潑濺模型獨特的計算方式——高斯潑濺分為多個不同階段，每個階段的并行粒度不同，需要進行切換。
這與大多數神經網絡模型的單一粒度并行方式大相徑庭，甚至高斯潑濺根本沒用到任何神經網絡。
這就導緻了現有的針對神經網絡訓練的多卡并行方案（如資料并行、模型并行等），難以直接應用于3D高斯潑濺。
另外，在高斯潑濺模型的訓練過程中，不同粒度的過程之間需要進行大量的資料通信，加大了并行方案的難度。
這也正是Grendel的設計當中需要解決的問題。
首先，Grendel将3D高斯潑濺的訓練過程劃分為三個主要階段——高斯變換（Gaussian transformation）、渲染（rendering）和損失計算（loss computation）。
針對這三個階段Grendel采用混合粒度的并行政策，在不同的訓練階段使用不同的并行粒度。
- 高斯變換階段采用高斯粒子級（Gaussian-wise）并行，将高斯粒子均勻分布到各個GPU節點；
- 渲染和損失計算階段采用像素級（pixel-wise）并行，将圖像分割成連續的像素塊，配置設定到各個GPU節點。
在高斯變換和渲染階段中間，Grendel通過稀疏的全對全通信，将每個GPU節點上的高斯粒子傳輸到需要它們進行渲染的GPU節點。
由于每個像素塊隻依賴于其覆寫範圍内的高斯粒子子集，Grendel利用空間局部性，隻傳輸相關的粒子，進而減少了通信量。
完成損失計算後，在每個GPU節點上，系統會根據損失函數計算渲染管線各個參數的梯度，并通過反向傳播回傳給高斯粒子的各個屬性參數。
之後，系統将各GPU計算出的梯度進行聚合，得到批量資料的總梯度，并據此更新高斯粒子的屬性參數。
接着就是重複從高斯變換到參數更新的步驟，直到模型收斂或達到預設的訓練輪數。
另外，為了處理渲染階段的負載不均衡問題，Grendel引入了動态負載均衡機制：
在訓練過程中，Grendel會記錄每個像素塊的渲染時間，用于預測目前疊代的負載分布，然後動态調整像素塊到GPU節點的配置設定，盡量使各個節點的渲染時間接近。
為了進一步提高GPU使用率和訓練吞吐量，Grendel支援批量訓練,即在每個訓練疊代中并行處理多個輸入圖像，并根據批量大小動态調整學習率，以保證訓練的穩定性和收斂性。
作者簡介
Grendel的第一作者，是紐約大學計算機博士生、清華姚班19級校友趙和旭，主要研究方向是分布式機器學習。
在清華期間，趙和旭曾在清華NLP實驗室孫茂松團隊參與研究，接受劉知遠副教授的指導。
他還曾經在Eric Xing組通路，優化了一個分布式機器學習中的通訊問題，論文被MLsys2023接收。
另外三名華人作者，Weng Haoyang（翁颢洋）也來自姚班；Daohan Lu來自紐約大學，是謝賽甯的博士生；還有Ang Li博士，是一名浙大校友，現在美國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
— 完 —
量子位 QbitAI · 頭條号簽約
關注我們，第一時間獲知前沿科技動态