一句话总结

来自内布拉斯加大学林肯分校、罗切斯特大学医学中心以及Holland Computing Center的研究人员提出了一种可扩展的多GPU框架，用于三维多细胞生长模拟，该框架结合了GPU加速、空间分箱、域分解和负载感知分区，并采用基于RNN的负载均衡控制器，该控制器观察最近每个rank的执行时间和分区状态，学习对反应式边界调整规则的残差修正，并在可微分代理中离线训练，使用随机化的工作负载动态且无需执行轨迹；该控制器将平均全局不平衡从静态分区下的$11.3\%$11.3%降低到$3.5\%$3.5%，相对于静态分区将端到端运行时间降低了$9.0\%$9.0%，并且与反应式基线相比，切片迁移减少了$7.7\times$7.7×。

核心贡献

• 一种可扩展的多GPU框架，用于三维多细胞生长模拟，结合了GPU加速、空间分箱、域分解和负载感知分区；GPU加速与空间分箱相比串行CPU基线将元素交互时间降低了约三个数量级，并且在630万个元素时，分箱相对于GPU全对全实现进一步提供了9.2倍的加速。
• 一个由RNN引导的负载均衡控制器从最近的每个rank执行时间和分区状态中学习对反应式边界调整规则的残差修正，并在负载均衡循环的可微分代理中离线训练，无需测量执行轨迹。在8个GPU上的端到端43.5万细胞模拟中，它将平均全局不平衡从静态分区下的11.3%降低到3.5%，而反应式规则为3.8%，相对于静态分区将总运行时间降低了9.0%，并且在运行时间相当的情况下，切片迁移比反应式基线少7.7倍。
• 一个涉及增殖、分化和组织重组的代表性胚胎表皮发育用例展示了对空间和时间上不断变化的工作负载的支持，以及显著改善的均衡-迁移权衡。

引言

多细胞组织发育源于局部细胞行为，如生长、分裂、迁移和粘附，而三维细胞模型使这些机制变得明确且可控。亚细胞元素模型通过用多个相互作用的元素表示每个细胞，提供了更精细的力学细节，但这显著增加了力评估成本，并且随着细胞生长、分裂和移动，工作负载在空间上变得不均匀且随时间变化。单GPU实现受到计算能力和内存的限制，而先前的反应式负载均衡方法仅在观察到不平衡后才调整分区边界，并未利用工作负载的时间趋势。作者开发了一种用于表皮亚细胞元素组织生长的可扩展多GPU框架，该框架结合了GPU并行力评估、空间分箱、域分解、负载感知初始分区以及一个由RNN引导的负载均衡控制器，该控制器学习对反应式边界调整规则的残差修正。

方法

作者提出了一种基于亚细胞元素模型（SEM）的多GPU框架，用于加速表皮组织形成模拟。整体工作流程包括负载感知初始化、GPU并行元素更新、幽灵区域交换以及RNN引导的运行时重分区。

在SEM中，每个生物细胞由多个相互作用的亚细胞元素表示。细胞$\alpha$α中元素$\alpha_i$αi​的位置$Y_{\alpha_i}$Yαi​​根据来自附近元素和外部结构的力进行更新：

$$\frac{d Y_{\alpha_i}}{dt} = - \nabla_{\alpha_i} \sum_{\beta_j \in \mathcal{I}} V_{\mathrm{ele}} \left(| Y_{\alpha_i} - Y_{\beta_j} |\right) - \nabla_{\alpha_i} V_{\mathrm{ext}} (Y_{\alpha_i})$$dtdYαi​​​=−∇αi​​βj​∈I∑​Vele​(∣Yαi​​−Yβj​​∣)−∇αi​​Vext​(Yαi​​)

其中$\mathcal{I}$I是邻近元素集合，$V_{\mathrm{ele}}$Vele​包括细胞内和细胞间相互作用，$V_{\mathrm{ext}}$Vext​模拟外部粘附。所有相互作用都在规定的截断半径内评估。为了避免全对全比较的$O(n^2)$O(n2)成本，三维域被组织成由截断距离大小的立方体分箱组成的规则网格。每个元素仅检查自己的分箱和26个相邻分箱，使相互作用计算线性缩放。对于多GPU执行，空间域沿一个维度划分为连续的切片，并分配给各个GPU rank。每个rank维护一个包含相邻分区截断距离内元素的幽灵区域，以处理边界相互作用。

由于细胞密度和相互作用数量在空间上变化，均匀几何分解往往导致负载不平衡。作者为每个细胞$i$i计算一个工作负载分数：

$$w_i = n_i n_{\mathcal{N}(i)}$$wi​=ni​nN(i)​

其中$n_i$ni​是细胞$i$i中亚细胞元素的数量，$n_{\mathcal{N}(i)}$nN(i)​是周围$3 \times 3 \times 3$3×3×3分箱邻域中的元素数量。域被划分为薄切片，切片$s$s的工作负载为$W_s = \sum_{i \in s} w_i$Ws​=∑i∈s​wi​。然后将连续切片分组，使得每个GPU rank的累积工作负载大致相等。

随着细胞生长和分裂，工作负载分布动态变化。为了解决这个问题，作者引入了一种RNN引导的动态负载均衡策略，该策略利用最近的工作负载历史和分区状态信息来调整空间分解。

该方法从一个反应式基线开始，该基线监控与分区边界相邻的rank的计算时间$T_t^-$Tt−​和$T_t^+$Tt+​。局部不平衡$I_t = (T_t^+ - T_t^-) / ((T_t^+ + T_t^-) / 2)$It​=(Tt+​−Tt−​)/((Tt+​+Tt−​)/2)在$|I_t| > \tau$∣It​∣>τ时触发边界移动。反应式调整幅度为：

$$|\Delta s_t^{\mathrm{r}}| = \alpha N_{\mathrm{slow}} \frac{|T_t^+ - T_t^-|}{2 T_t^{\mathrm{slow}}}$$∣Δstr​∣=αNslow​2Ttslow​∣Tt+​−Tt−​∣​

其中$N_{\mathrm{slow}}$Nslow​是较慢rank上的切片数量，$\alpha$α限制调整幅度。

为了改进这一反应式规则，一个共享的循环神经网络（RNN）作为分区边界上的控制器。在每个重平衡步骤$t$t，特征向量$\mathbf{x}_{b,t}$xb,t​（包括时间不平衡、分区宽度、生长周期阶段和反应式建议）被输入到门控循环单元（GRU）中以更新隐藏状态$\mathbf{h}_{b,t}$hb,t​。循环状态被映射为标量残差$r_{b,t}$rb,t​，该残差在有界步长空间中修正反应式建议：

$$\Delta s_t = \Delta s_{\max} \tanh \left(\tanh^{-1} \left(\frac{\Delta s_t^{\mathrm{r}}}{\Delta s_{\max}}\right) + r_{b,t}\right)$$Δst​=Δsmax​tanh(tanh−1(Δsmax​Δstr​​)+rb,t​)

这种残差公式允许学习到的控制器在保持有界调整幅度的同时细化反应式规则。修正后的调整被联合应用于所有内部边界，将其向较慢的rank移动，同时遵守最小分区宽度约束。

RNN控制器在负载均衡循环的可微分模拟器中使用合成工作负载轨迹进行训练。这使得控制器能够学习边界调整策略的结构，而不是特定的生物动力学。输入特征是无量纲的局部量，能够泛化到不同的问题几何形状和rank数量。网络通过时间反向传播在50次重平衡事件的展开序列上进行训练。训练损失平均了每个rank时间的相对不平衡及其归一化散布，并添加了$L_1$L1​迁移惩罚以阻止不必要的边界移动。优化器使用Adam与余弦退火，并且在训练期间应用软松弛以保持整数边界分配的可微分性。

实验

评估在GPU集群上使用表皮亚细胞元素模型，以评估单GPU加速、多GPU可扩展性、动态负载均衡以及基于RNN的控制器。GPU加速与空间分箱在大元素数量下实现了接近线性的扩展，多GPU执行有利于更大的工作负载。RNN引导的负载均衡减少了全局时间不平衡，且切片迁移远少于反应式或基于预测的方法，改善了均衡-迁移权衡。一个发育组织生长模拟展示了该框架处理空间演化多细胞工作负载的能力。