图神经网络（GNNs）通过有效捕获图结构数据中的关系模式和层次依赖性，已迅速从理论构建发展成为多个领域的基础工具。例如，谷歌地图利用GNNs对复杂的交通网络进行建模，从而显著提高了到达时间（ETA）预测的准确性[1]。在社交网络分析中，GNNs通过建模复杂的用户互动来实现高级社区检测和影响力预测[2]。同时，推荐系统使用GNNs捕获高阶用户-项目关系，相比传统的矩阵分解方法提高了15-30%的准确性[3]。此外，GNNs越来越多地应用于其他具有挑战性的领域，如知识图谱[4]，进一步凸显了它们在处理复杂网络数据方面的变革性影响。
随着GNNs越来越多地应用于具有数十亿节点和边的大规模图，高效训练它们已成为一个关键的系统挑战。为了应对这种规模，现代框架通常采用分布式数据并行训练结合小批量采样，以提高内存效率和可扩展性[5]、[6]、[7]。在这种范式中，输入图被划分为逻辑子图，每个子图分配给一个专用的工作进程。这些分区并行处理小批量数据，同时定期在设备之间同步模型参数。然而，这种设计引入了结构和随机的工作负载不平衡来源。现实世界中的图通常表现出高度偏斜的度分布，小批量采样产生的子图具有不可预测的大小和拓扑结构。随着时间的推移，频繁采样或高度节点往往会主导某些分区，导致持续的工作负载不平衡。这种不平衡会导致同步延迟、设备利用率降低和训练成本增加[8]。
在小批量采样的分布式GNN训练中，工作负载不平衡尤为明显，因为计算图在每个训练步骤中都会动态变化。尽管底层图是静态的，但每步的工作负载分布却高度可变。这种动态特性导致每个GPU的计算和通信需求波动，随着时间的推移积累成显著的不平衡。尽管METIS[9]等静态图划分方法被广泛使用，但它们假设工作负载模式是稳定的，因此在实践中难以保持平衡的执行。即使定期重新应用，这些方法在适应不可预测的采样动态方面也存在固有的局限性，导致训练期间持续的低效率。
为了解决这些限制，最近的系统提出了自适应策略来减轻工作负载不平衡。这些策略包括采样感知的批量构建、执行配置调整和动态任务迁移[8]、[10]、[11]、[12]、[13]。虽然这些方法显示出有希望的改进，但它们主要关注系统级策略，如参数调整、运行时调度和延迟更新。然而，它们往往忽略了由采样计算图的动态拓扑结构引起的细粒度、结构诱导的不平衡。为了解决这个问题，我们采用了一种结构感知的节点迁移策略，该策略明确考虑了拓扑不平衡和批次间的持续不平衡。尽管一些现有框架也执行节点迁移，但它们通常依赖于即时采样或训练成本来触发决策，这使得它们对短期波动非常敏感，容易发生过早或无效的迁移。例如，FlexGraph[14]直接从批量级成本信号估计不平衡，而不建模长期模式，通常导致不必要的或时机不当的迁移。此外，其迁移策略基于贪婪选择，不考虑图结构，导致次优的放置和过多的跨GPU通信。这些局限性突显了在分布式GNN训练中需要更稳定的不平衡检测和结构感知调度。
为此，我们提出了AFS-GNN，这是一个具有调度意识的框架，旨在在动态和不平衡的工作负载下提高分布式GNN训练的效率和稳定性。AFS-GNN将轻量级运行时监控、自适应不平衡检测和结构保留的任务迁移整合到一个统一的流程中。该系统的关键特性是其能够基于工作负载变化点检测触发迁移决策，从而及时响应运行时的变化。为了指导成本效益高的迁移，AFS-GNN量化了节点级别的工作负载，并选择性地识别出需要重新分配的高影响计算块。这些候选子图使用谱方法进行划分，以最小化GPU间的依赖性，同时匹配目标迁移量。通过将迁移规划与GPU执行解耦，并通过节点内通信高效执行数据传输，AFS-GNN实现了低开销，而不干扰正在进行的训练。实验结果表明，AFS-GNN相比Euler加速了21.7%，相比FlexGraph加速了13.7%，并在不同的批量大小和分区数量下保持了稳定的性能。消融研究进一步验证了其工作负载跟踪器和谱划分组件的有效性。
总体而言，AFS-GNN提供了四个关键贡献：（1）基于轻量级变化点的工作负载监控机制，用于实时不平衡检测；（2）在层次依赖图（HDGs）上基于成本的谱划分方法，以生成最小化GPU间通信的迁移计划；（3）结构感知的任务迁移设计，将调度与执行解耦，提高了响应性和稳定性；（4）在大型基准测试上的广泛实验验证，证明了其在不同规模和批量大小下的训练延迟减少、稳定性，以及通过消融研究验证了每个组件的有效性。
