Rudder:利用 LLM 智能体驾驶分布式 GNN 训练中的预取
机器学习
2026-03-02 v1 人工智能
分布式、并行与集群计算
多智能体系统
性能
摘要
大规模图神经网络(GNN)通常通过采样顶点邻居来固定距离进行训练。由于大型输入图分布在不同节点上,训练需要频繁的非规则通信导致前进受阻。此外,获取的数据随图、图分布、采样和批量参数以及缓存策略的不同而变化。因此,任何静态预取方法都会错失适应不同动态条件的关键机会。本文引入 Rudder,作为嵌入在最先进 AWS DistDGL 框架中的软件模块,用于自主预取远程节点并最小化通信。Rudder 的适应性不同于标准启发式方法和传统 ML 分类器。我们观察到,当代大型语言模型(LLM)的生成式 AI 具有如原语学习(ICL)用于零样本任务、逻辑多步骤推理等涌现属性。我们发现这一行为非常适合即使在大量欠训练情况下也适用于自适应控制。使用标准数据集和在 NERSC Perlmutter 超级计算机上未见配置的评估显示,Rudder 相对于基线 DistDGL(无预取)实现了最高 91% 的端到端训练性能提升,相对于静态预取实现了 82% 的提升,将通信减少超过 50%。我们的代码可在 https://github.com/aishwaryyasarkar/rudder-llm-agent 获取。
引用
@article{arxiv.2602.23556,
title = {Rudder: Steering Prefetching in Distributed GNN Training using LLM Agents},
author = {Aishwarya Sarkar and Sayan Ghosh and Nathan Tallent and Aman Chadha and Tanya Roosta and Ali Jannesari},
journal= {arXiv preprint arXiv:2602.23556},
year = {2026}
}
备注
Accepted to the 40th ACM International Conference on Supercomputing (ICS 2026)