谢赛宁团队与纽约大学系统实验室破解了高斯泼溅模型训练时的内存瓶颈。这次合作设计了一种并行策略,推出了名为 Grendel 的分布式训练系统。Grendel 系统的第一作者是清华大学姚班的校友赵和旭。这个系统在使用 4 张显卡时,训练速度可以提升 3.5 倍,而当显卡数量增加到 32 张时,速度提升了 6.8 倍。这种多卡训练方案不仅让训练速度提升,还解决了大场景和高分辨率环境下的内存限制问题。
高斯泼溅模型在处理大规模 3D 渲染任务时,对内存的需求极高。传统的单卡训练方法不仅速度慢,还容易在高分辨率和大场景下出现内存不足的情况。内存瓶颈限制了模型的扩展性和应用场景。为了突破内存瓶颈,Grendel 系统设计了一套高效的并行策略。训练过程被划分为高斯变换、渲染和损失计算三个主要阶段,每个阶段采用不同的并行策略:
- 高斯变换阶段:采用高斯粒子级并行策略,将高斯粒子均匀分布到各个 GPU 节点上进行处理。
- 渲染和损失计算阶段:采用像素级并行策略,将图像分割成连续的像素块,分配到各个 GPU 节点上进行渲染和计算。
这种混合粒度的并行策略,能够在不同的训练阶段灵活调整并行粒度,利用 GPU 资源,提高训练效率。为了减少数据通信量,Grendel 系统采用了稀疏的全对全通信策略。通过利用空间局部性,系统只传输相关的高斯粒子,避免了不必要的数据传输,从而减少了通信量,提高了整体训练效率。
通过这些技术细节的优化,Grendel 系统在处理大规模 3D 渲染任务时,提高了训练速度和效率,突破了传统单卡训练的内存瓶颈,为高斯泼溅模型的应用提供了支持。