On Scaling Up 3D Gaussian Splatting Training

2024-06-26

International Conference on Learning Representations

AI Summary

3D Gaussian Splatting (3DGS) is constrained by single-GPU memory capacity, limiting scalability to high-resolution and large-scale scene reconstruction. To address this, we present Grendel—the first multi-GPU distributed training system for 3DGS. Our method introduces three core innovations: (1) sparse all-to-all communication coupled with dynamic pixel partitioning to minimize inter-GPU data transfer; (2) multi-view batch rendering with adaptive Gaussian density control; and (3) a √batch-size-based parameter scaling rule. Grendel employs distributed parameter sharding and dynamic load balancing to ensure efficient parallel training. Evaluated on the Rubble dataset, Grendel trains 40.4 million Gaussians across 16 GPUs, achieving a PSNR of 27.28—outperforming the single-GPU baseline (11.2 million Gaussians) by +1.0 dB. The implementation is open-sourced.

Technology Category

• Computer Vision: 3D Computer Vision
• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Distributed Search

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

Abstract

3D Gaussian Splatting (3DGS) is increasingly popular for 3D reconstruction due to its superior visual quality and rendering speed. However, 3DGS training currently occurs on a single GPU, limiting its ability to handle high-resolution and large-scale 3D reconstruction tasks due to memory constraints. We introduce Grendel, a distributed system designed to partition 3DGS parameters and parallelize computation across multiple GPUs. As each Gaussian affects a small, dynamic subset of rendered pixels, Grendel employs sparse all-to-all communication to transfer the necessary Gaussians to pixel partitions and performs dynamic load balancing. Unlike existing 3DGS systems that train using one camera view image at a time, Grendel supports batched training with multiple views. We explore various optimization hyperparameter scaling strategies and find that a simple sqrt(batch size) scaling rule is highly effective. Evaluations using large-scale, high-resolution scenes show that Grendel enhances rendering quality by scaling up 3DGS parameters across multiple GPUs. On the Rubble dataset, we achieve a test PSNR of 27.28 by distributing 40.4 million Gaussians across 16 GPUs, compared to a PSNR of 26.28 using 11.2 million Gaussians on a single GPU. Grendel is an open-source project available at: https://github.com/nyu-systems/Grendel-GS

Problem

• Scaling 3D Gaussian Splatting training across multiple GPUs
• Handling high-resolution large-scale 3D reconstruction tasks
• Overcoming memory limitations in distributed 3DGS parameter optimization

Innovation

• Distributed system partitioning 3DGS parameters across GPUs
• Sparse all-to-all communication for Gaussian transfer
• Batched training with multiple camera views simultaneously