Rudder: Steering Prefetching in Distributed GNN Training using LLM Agents

📄 arXiv: 2602.23556v1 📥 PDF

作者: Aishwarya Sarkar, Sayan Ghosh, Nathan Tallent, Aman Chadha, Tanya Roosta, Ali Jannesari

分类: cs.LG, cs.AI, cs.DC, cs.MA, cs.PF

发布日期: 2026-02-26

备注: Accepted to the 40th ACM International Conference on Supercomputing (ICS 2026)

🔗 代码/项目: GITHUB

💡 一句话要点

Rudder：利用LLM Agent在分布式GNN训练中实现自适应预取

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 分布式GNN训练 预取策略 大型语言模型 上下文学习 图神经网络 通信优化 DistDGL

📋 核心要点

  1. 分布式GNN训练中，频繁的通信成为性能瓶颈，静态预取策略难以适应动态变化的数据访问模式。
  2. Rudder利用LLM的上下文学习和推理能力，作为Agent自主学习并预测需要预取的远程节点，从而减少通信。
  3. 实验表明，Rudder在端到端训练性能上显著优于DistDGL和静态预取，通信量也大幅降低。

📝 摘要（中文）

大规模图神经网络（GNN）的训练通常通过采样顶点固定距离的邻居来实现。由于大型输入图是分布式的，训练需要频繁的不规则通信，这会阻碍训练进度。此外，获取的数据随图、图分布、采样和批处理参数以及缓存策略而变化。因此，任何静态预取方法都会错过适应不同动态条件的关键机会。本文介绍Rudder，一个嵌入在最先进的AWS DistDGL框架中的软件模块，用于自主预取远程节点并最小化通信。Rudder的自适应性与标准启发式方法和传统ML分类器形成对比。我们观察到，当代大型语言模型（LLM）中发现的生成式AI表现出涌现属性，如用于零样本任务的上下文学习（ICL），以及逻辑多步推理。我们发现这种行为非常适合自适应控制，即使在大量欠训练的情况下也是如此。在NERSC Perlmutter超级计算机上使用标准数据集和未见配置进行的评估表明，端到端训练性能比基线DistDGL（无预取）提高了高达91%，比静态预取提高了82%，通信减少了50%以上。我们的代码可在https://github.com/aishwaryyasarkar/rudder-llm-agent 获得。

🔬 方法详解

问题定义：在分布式GNN训练中，由于图数据分布在多个计算节点上，训练过程中需要频繁地从远程节点获取数据。这种通信开销成为训练速度的瓶颈。现有的静态预取策略无法根据图结构、数据分布、采样参数等动态因素进行调整，导致预取效率低下，甚至引入额外的开销。

核心思路：Rudder的核心思路是利用大型语言模型（LLM）的上下文学习和推理能力，构建一个智能Agent，使其能够自主学习并预测哪些远程节点需要在训练过程中被预取。通过将预取问题建模为一个序列决策问题，LLM Agent可以根据历史的训练数据和当前的状态，动态地调整预取策略，从而减少通信开销。

技术框架：Rudder作为一个软件模块嵌入到AWS DistDGL框架中。其主要流程包括：1) 收集训练过程中的状态信息，例如图结构、数据分布、采样参数等；2) 将这些状态信息作为LLM Agent的输入，Agent根据上下文学习和推理能力，生成预取决策；3) 根据Agent的决策，Rudder执行预取操作，将远程节点的数据加载到本地；4) 训练完成后，Rudder收集训练结果，并反馈给LLM Agent，用于优化后续的预取策略。

关键创新：Rudder的关键创新在于将LLM引入到分布式GNN训练的预取策略中。与传统的启发式方法和机器学习分类器相比，LLM Agent具有更强的上下文学习和推理能力，能够更好地适应动态变化的数据访问模式。此外，Rudder采用了一种轻量级的训练方法，即使在数据量较少的情况下，也能获得较好的性能。

关键设计：Rudder的关键设计包括：1) LLM Agent的prompt设计，需要充分利用LLM的上下文学习能力，将状态信息有效地编码到prompt中；2) 预取决策的生成方式，Rudder采用了一种基于概率的采样方法，允许Agent探索不同的预取策略；3) 奖励函数的设计，需要综合考虑通信开销、计算开销和预取准确率等因素，以引导Agent学习最优的预取策略。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Rudder在NERSC Perlmutter超级计算机上的实验结果表明，与基线DistDGL（无预取）相比，端到端训练性能提升高达91%，与静态预取相比，提升高达82%，同时通信量减少超过50%。这些结果表明Rudder在分布式GNN训练中具有显著的性能优势。

🎯 应用场景

Rudder可应用于各种大规模图神经网络的分布式训练场景，尤其是在图数据规模庞大、节点间连接复杂、通信开销显著的情况下。该技术能够有效提升训练效率，降低训练成本，加速GNN在推荐系统、社交网络分析、生物信息学等领域的应用。

📄 摘要（原文）

Large-scale Graph Neural Networks (GNNs) are typically trained by sampling a vertex's neighbors to a fixed distance. Because large input graphs are distributed, training requires frequent irregular communication that stalls forward progress. Moreover, fetched data changes with graph, graph distribution, sample and batch parameters, and caching polices. Consequently, any static prefetching method will miss crucial opportunities to adapt to different dynamic conditions. In this paper, we introduce Rudder, a software module embedded in the state-of-the-art AWS DistDGL framework, to autonomously prefetch remote nodes and minimize communication. Rudder's adaptation contrasts with both standard heuristics and traditional ML classifiers. We observe that the generative AI found in contemporary Large Language Models (LLMs) exhibits emergent properties like In-Context Learning (ICL) for zero-shot tasks, with logical multi-step reasoning. We find this behavior well-suited for adaptive control even with substantial undertraining. Evaluations using standard datasets and unseen configurations on the NERSC Perlmutter supercomputer show up to 91% improvement in end-to-end training performance over baseline DistDGL (no prefetching), and an 82% improvement over static prefetching, reducing communication by over 50%. Our code is available at https://github.com/aishwaryyasarkar/rudder-llm-agent.