謝賽寧團隊突破高斯潑濺內存瓶頸，并行方案實現多顯卡訓練

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

克雷西 發自 凹非寺
量子位 | 公眾號 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在多卡上進行并行訓練不僅帶來量的改變，也帶來了質的突破。

△OOM：Out of memory，內存不足

另外，通過支持更大的批量（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