Local AI & Edge

AMD ROCm team open-sources distributed 3D Gaussian Splatting training for multi-GPU scaling

Key Takeaway

AMD's gsplat library implements an unconventional distributed training strategy for 3D Gaussian Splatting that avoids all-reducing gradients entirely, enabling scene reconstruction beyond single-GPU memory limits on AMD Instinct hardware.

Summary

AMD's ROCm blog has published a technical deep-dive on scaling gsplat, an open-source GPU-optimized Python library for 3D Gaussian Splatting (3DGS), across multiple AMD Instinct GPUs. The post details why production-scale 3DGS scenes—such as city blocks captured at 4K resolution—quickly exhaust single-GPU memory, with a single Gaussian requiring up to 708 bytes including Adam optimizer state at default spherical harmonics degree 3. Rather than applying standard PyTorch DistributedDataParallel, gsplat uses a genuinely unusual parallelization approach that is data-parallel in spirit but never all-reduces scene gradients, instead distributing the parameter-heavy scene representation itself.

The article walks through the collective communication pattern call-by-call and validates the design against a rocprofv3 trace from a real four-GPU run, revealing that the expected communication bottleneck is not the actual performance limiter. The work is published as part of AMD's open-source ROCm ecosystem, with gsplat available for independent deployment on AMD GPU hardware without proprietary restrictions.

Distributed training architecture and memory scaling

  • Memory footprint per Gaussian: At gsplat's default SH degree 3, each Gaussian requires approximately 708 bytes including Adam optimizer state (59 learnable parameters plus two moment vectors), limiting a single MI300X with 192 GB HBM3 to roughly 271 million Gaussians before accounting for transient gradient and intersection buffers.
  • Non-standard parallelization strategy: Unlike typical PyTorch DDP which all-reduces gradients across ranks while keeping parameter replicas identical, gsplat's approach distributes the scene parameters themselves and avoids all-reducing gradients of the training scene entirely.
  • Per-tile intersection bottleneck: The dominant transient memory cost during training comes from per-tile sorted lists of Gaussian-to-pixel intersections, which scale with image resolution and Gaussian density and frequently trigger out-of-memory errors at 4K resolution before parameter storage limits are reached.
  • rocprofv3 validation: The implementation was verified against a real four-GPU rocprofv3 trace, which confirmed the design's correctness and identified that the obvious suspected bottleneck was not the actual performance limiter.
  • Three-limit attack: Distributed execution addresses capacity limits (scene fits), quality limits (less aggressive pruning/densification), and throughput limits (more compute units and memory bandwidth working in parallel), whereas larger single GPUs only address capacity.

Why it matters

This release strengthens AMD's open-source AI tooling story for workloads beyond large language models, giving researchers and autonomous-driving developers a vendor-agnostic Python path to train production-scale 3DGS scenes on AMD Instinct hardware without CUDA lock-in. For the Linux AI ecosystem, gsplat's unconventional distributed strategy offers a reference implementation for memory-dominated computer-vision workloads where standard DDP assumptions break down, potentially informing similar adaptations in other PyTorch-based 3D reconstruction tools.

Original reference: rocm.blogs.amd.com →