图神经网络中的协作式小批量处理

训练大规模图神经网络（GNNs）需要大量计算资源，且过程高度数据密集。降低资源需求最有效的方式之一是小批量训练结合图采样。GNNs 具有独特性质：小批量中各项存在重叠数据。然而，常用的独立小批量（Independent Minibatching）方法为每个处理单元（PE，即核心和/或 GPU）分配自有小批量处理，导致跨 PE 的重复计算与输入数据访问。这放大了邻域爆炸现象（NEP），即限制扩展的主要瓶颈。为在多 PE 设置中减弱 NEP 影响，我们提出称为协作式小批量（Cooperative Minibatching）的新方法。我们的方法利用采样子图大小是批量大小的凹函数这一事实，从而随批量增大显著减少工作量。因此，即便全局批量大小相同，配备快速互连的处理器作为单一更大型处理器协同处理一个大批量，优于分别处理较小批量。我们还展示在串行执行中通过生成依赖的连续小批量利用同一现象。实验评估显示，仅增加此依赖性便可在不损害模型收敛的前提下为获取顶点嵌入节省至多 4 倍带宽。结合所提方法，我们在单节点多 GPU 系统上使用相同资源相较独立小批量实现至多 64% 加速。