训练负载不平衡
3DGS是一项创新的3D图形技术，虽然这项技术具有巨大的潜力，但也面临着显著的计算挑战，因此优化显得尤为重要。为了提升3DGS的性能，目前主要探索了
两种优化策略
。
-
第一种策略着眼于算法本身的修改，以实现更高的精度、更少的存储需求或更高的计算效率。
-
第二种策略则不改变算法本身，而是通过各种方法提升原始算法的计算效率，例如智能任务调度或CUDA内核的优化。
尽管3DGS技术取得了很大的进步，训练3DGS模型仍然是一项充满挑战且耗时的任务，
尤其是在负载不平衡的情况下。当像素和高斯球之间的工作负载不均衡时，会影响训练性能
。具体来说，图形处理器（GPU）设计用于处理规则且同质化的任务，并以SIMT（单指令多线程）方式运行。因此，当像素和高斯球之间的工作负载存在差异时，一些线程会处于活动状态，而另一些线程则会处于空闲状态，从而降低整体性能。据我们所知，
目前尚无研究在内核层面解决这些不平衡问题
。
在对3DGS训练过程进行深入分析和实验的基础上，我们确定并归纳了三种显著的负载不平衡问题：
-
CUDA静态分配导致的流处理器（SM）负载不平衡
：如果使用静态分配方法，在不同的SM之间会出现负载不平衡问题。尽管在不同的SM之间线程块任务的数量相同，但每个线程块的工作负载可能会有很大的差异。这是一个被忽视的性能因素，导致了SM之间的负载不均衡。
-
分块中的负载不平衡
：在3DGS中，为避免每个像素派生高斯球的计算成本，图像最初被划分为多个不重叠的分块。每个分块包含16×16个像素。然而，如果基本的CUDA分配和分块策略保持不变，一些分块可能会拥有极大的工作负载，而其他分块的工作负载则很轻，从而导致SM之间出现严重的负载不平衡。
-
训练阶段的负载不平衡
：在3DGS训练过程中，不同阶段的数据特性会发生显著变化。在初始阶段，像高斯分布这样的度量在不同块或线程之间存在显著差异。随着训练的进行，这些不平衡随着数据特性的变化而改善。在经过多次迭代后，极端的数据特性变得更加平衡。
Balanced 3DGS
[1]其亮点包括：
-
首次提出了基于高斯的并行负载均衡方法，以优化Warp内正向计算的CUDA渲染内核。
-
此外，创新性地提出了块间动态工作负载分配技术，通过在计算块之间均匀分配任务，最大限度地减少线程的空闲时间并最大化资源利用率。
-
精细化组合负载均衡方法结合了这两种技术，提供了一个完整的负载不平衡问题解决方案。
-
提出了一种基于实验的自适应内核选择策略。这种策略克服了手动选择度量标准的局限性，能够更准确地反映训练过程的实际情况，并确保更好的性能和效率。
具体方法
为了应对在流处理器 (SM)、图像分块以及训练阶段中提到的负载不平衡问题，我们提出了 Balanced 3DGS，这是一种基于高斯的并行渲染与精细化分块的方法，专注于3DGS训练中的负载均衡。该方法的关键在于如何在不同的SM之间调度工作负载，考虑高斯在不同图像分块中的分布不均，从而在高斯渲染中实现最佳利用率。值得注意的是，这种解决方案不会导致精度损失。
块间动态工作负载分配
如果图像分块被静态方法映射到线程块（SM），不同图像分块中的高斯数量分布不均将导致严重的负载不平衡：某些线程块将较早完成其任务，而必须等待负载最重线程块完成，这会引发明显的整体性能下降和计算资源浪费。
静态分配中，线程块按固定顺序处理分块，高斯数量的差异会导致负载不平衡。
为了解决这个问题，我们使用一种动态任务映射（分配）方法：
-
-
-
一旦线程块完成其当前工作，它会立即从任务池中获取下一个分块并运行；
-
通过这种方式，某些线程块可以集中处理少量高负载分块，而其他线程块可以处理许多低负载分块，从而避免由于工作负载分配不平衡造成的SM空闲。图2展示了静态和动态工作负载分配的SM工作负载差异。
SM执行顺序中，动态分配通过任务池均匀分配工作负载，解决分块负载差异问题。
动态分配的具体实现通过以下步骤完成：
-
-
-
如果线程块完成了任务，它会返回任务池并获取下一个分块；
-
任务分配逻辑使用原子操作以确保线程间的任务调度不会冲突。
算法1展示了具体的实现细节，包括线程块如何动态获取任务以及如何同步任务状态。
这种动态分配策略显著提高了GPU资源的利用率，最大化了计算效率，尤其是在负载差异较大的场景中。
基于高斯的并行渲染
当前，一个线程负责处理一个像素的工作负载。传统的渲染内核采用像素级并行方式，同时串行处理高斯。然而，即使在同一个图像分块内，像素共享同一组高斯，不同像素的实际工作负载可能差异很大。因为，如果某些像素在特定高斯的渲染过程中完成了计算，它们不需要遍历剩余的高斯，这种情况称为“提前停止”（early stop）。提前停止可能导致分块内像素之间的工作负载差异显著。因此，简单地将像素映射到线程会导致 Warp 内线程负载不均，而 Warp 内线程以 SIMT（单指令多线程）方式运行，这种不均衡会导致性能显著下降。
因此，我们提出了一种 Warp 协作（warp-collaboration）方法：
-
Warp 内的所有线程将协作处理单个像素，通过基于高斯的并行方式来避免 Warp 内的工作负载分歧。
-
在这种方法中，传统方法中由一个 Warp 并行处理的 32 个像素将改为串行处理。
-
换句话说，我们用高斯级并行替代了像素级并行，从而减少了 Warp 内的工作负载分歧（见算法2）。
实际上，不同高斯的计算之间存在依赖关系，如公式（4）所示。由于我们以 32 路并行方式处理高斯，因此需要额外执行一次前缀乘积操作，以确保每个高斯能够从其前序高斯中获取所有计算结果。在图3中，展示了 Warp 级像素级并行和 Warp 级高斯级并行的区别。
精细化组合负载均衡
这两种线程块之间和线程之间的负载均衡方法是正交的，我们可以将它们结合起来以实现更好的性能。在此之前，我们发现图像分块的尺寸对负载均衡和整体性能有显著影响。“分块尺寸”指分配给单个线程块的像素数量。
为了分析分块尺寸对性能的影响，我们使用了 A100 GPU 进行实验。假设一个线程块包含 128 个线程以简化分析。A100 具有 108 个流处理器（SM），每个流处理器可以同时调度 64 个 Warp（或 16 个线程块，每个线程块包含 4 个 Warp）。因此，一个波次可以包含多达 16 × 108 = 1728 个线程块。
在原始方法中，一个线程块处理一个 16×8 像素大小的分块。考虑到图像尺寸为 960×540 像素，我们可以发射 ⌈960/16⌉ × ⌈540/8⌉ = 4080 个线程块，这意味着渲染内核中只有 4080/1728 ≈ 2.36 个波次。在这种情况下，分块之间的负载不平衡可能会非常显著。我们需要更多的波次，以确保轻负载分块的工作量能够积累并与极重负载的分块相当。
解决方案：减小分块尺寸
为了解决上述问题，我们使用更小的分块尺寸分配给每个线程块（参见算法 3）。具体来说：
-
渲染内核中的初始分块尺寸和线程块内的线程大小仍保持为 16×8。
-
但是，我们将每个线程块处理的分块从 128 像素的分块减少到仅 4 像素的分块。
这样：
-
可分配的任务组数量增加了 32 倍（参见算法 3 第 23 行）。
-
在原始渲染内核中，128 像素的分块只需要将高斯数据加载到共享内存一次；而现在，由于每个线程块只处理 4 个像素，高斯数据需要多次加载到共享内存（参见算法 3 第 13 行）。显然，分块尺寸是负载均衡和内存访问效率之间的权衡。
性能优势
尽管更小的分块尺寸会显著增加内存访问的开销，但我们的实验结果（见图 5，高斯分布）表明，解决负载不平衡问题能带来更大的性能提升。
通过减小线程块的任务分配尺寸，我们能够直接结合介绍的两种负载均衡方法。图 4 展示了我们的组合负载均衡方法。
自适应内核选择策略
3D高斯分割模型训练的初始数据主要来自于处理过的LiDAR稀疏SfM点云数据。在训练的不同阶段，数据特性变化显著。分块中的高斯数量是负载平衡的关键指标。图5显示了随着训练的进行，每个块/线程中的高斯分布如何变化。
在训练开始时，不同块/线程中的高斯数量差异很大。一些极端负载的块/线程如果未能被适当处理，可能会导致显著的性能下降。在这种情况下，我们应使用提出的组合负载平衡优化方法。
然而，随着训练的进行，高斯数量的不平衡逐渐改善。从图5可以看到，经过若干次训练迭代后，最大高斯数量显著减少，其分布变得更加集中。在这种情况下，继续使用我们的负载平衡优化不会带来显著的性能提升；此外，正如上文所述，我们的优化可能会引入额外的开销。基于这些观察，当数据已经达到良好的平衡状态时，我们应放弃这些优化。
自然地，如何判断数据是否已经达到良好的平衡状态，并据此选择最佳内核，是实现最佳性能的关键。
与传统神经网络训练不同，固定的渲染内核无法在3DGS训练过程中保证负载均衡，因此需要基于实验的自适应内核选择策略。此策略可以避免手动选择度量平衡的指标，并能更好地反映实际情况。
图6展示了我们的训练过程：
-
-
每1000次迭代后，我们分别运行两种内核，并与当前数据单独比较其性能。
-
如果负载平衡优化内核的性能较差，我们假定数据已经达到良好的平衡状态。因此，剩余的训练将使用原始内核。
-
对于原始内核，我们还进行了某些内存访问优化，如第5.2.2节所述。具体来说，将特征和深度存储在共享内存中。尽管这种优化在训练初期因内存访问占比较小而效果不显著，但随着训练的进行，负载逐渐平衡，其效果会变得更加明显。
实验效果
