Rudder: Steering Prefetching in Distributed GNN Training using LLM Agents

Aishwarya Sarkar; Sayan Ghosh; Nathan Tallent; Aman Chadha; Tanya Roosta; Ali Jannesari

cs.LGFeb 26, 2026

Rudder: Steering Prefetching in Distributed GNN Training using LLM Agents

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

Organizations: Iowa State University, Ames, IA, USA · Pacific Northwest National Laboratory, Richland, WA, USA · Amazon GenAI, Cupertino, CA, USA · University of California, Berkeley, Berkeley, CA, USA

Abstract

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.

Figures & tables

Figure 1. Declining unique remote nodes in GNN training.

Figure 2. Prefetching interfaces can range from simple to burdensome . The simplest designs sacrifice performance; while the most burdensome require enormous tuning. Rudder achieves high performance while requiring little tuning.

Figure 3. Adaptive replacement consistently yields best %-Hits ( higher is better), relative to other replacement strategies.

Figure 4. High-level aspects of our replacement strategy based on a scoring policy which tracks recent usage.

Figure 5. LLM agent learns by interacting with environment through auxiliary tools, whereas ML classifiers are trained offline.

Figure 6. Quantifying LLMs on key characteristics.

Figure 7. Agents/classifiers deployed alongside GNN training tasks, trainer processes access persistent buffers.

Figure 8. Tasks in the trainer processes of Rudder .

Figure 9. LLM agent request/response flow and context updates.

Figure 10. LLM agents making decisions.

Figure 11. Asynchronous (default) and Synchronous (trainer waits for agent’s response) modes of Rudder.

Panel of Figure 12 — Algorithm 1 Distributed GNN Training Inputs: $S$ : Sampled Nodes, $BUF$ : Persistent Buffer, $\Theta$ : Pretrained LLM / ML Classifier, $G_{P}$ : Graph Partition, $Q_{M}$ : Minibatch Queue, $Q_{R},Q_{\tilde{R}}$ : Agent’s Request/Response Queue, $\Pi$ : Scoring Policy Output: $\Psi$ : Trained GNN Model

(a) Datasets
Dataset	Nodes $\|V\|$	Edges $\|E\|$	Feat. Dim.
products ( ogb2020 )	2.4M	61.85M	100
reddit ( ogb2020 )	0.23M	114.61M	602
papers100M ( ogb2020 )	111M	1.6B	128
orkut ( snapnets )	3.07M	117.18M	8
friendster ( snapnets )	65.6M	1.8B	128

Table 1. (a) Datasets and (b) Models used.

Figure 13. %-Improvement of Rudder with LLM Agents and ML Classifiers over DistDGL+fixed across all configurations.

Figure 14. (Left) #Remote nodes in buffer (5%/25%). (Right) 99% communication volume ( lower is better).

Figure 15. Communication volume ( lower is better) and %-Hits for MassiveGNN (fixed replacement interval at 32) and DistDGL+Rudder across buffer sizes (5%/25%) on products (64 trainers).

Figure 16. %-Training time (top) and %-communication (bottom). Y1: lower is better and Y2 (%-Hits): higher is better, across buffer capacities (5–25% of remote nodes) in products on 16 trainers. Annotations are %-improvements relative to DistDGL+fixed.

Figure 17. %-Hits comparison between synchronous (left) and asynchronous (right) mode of Rudder ( higher is better).

Models	Asynchronous
Models	Pass@1 %-Hits/Acc	Replacement Interval ( $r$ )	Valid/Invalid Response (%)	+ve/-ve Decisions (%)
Gemma3-4B	79	10	100/0	30/70
Gemma3-1B	16	8	100/0	100/0
Llama3.2-3B	63	6	99/1	29/71
SmolLM2-360M	13	4	87/13	35/65
SmolLM2-1.7B	24	5	92/8	70/30

Table 2. Asynchronous vs. Synchronous Evaluations.

Models	products	reddit	papers	orkut	friendster
Gemma3-4B	76 (-9/11)	75 (-12/15)	74 (-8/9)	82 (-10/12)	77 (-7/8)
Gemma3-1B	20 (-16/19)	19 (-12/15)	13 (-11/12)	9 (-13/15)	82 (-7/8)
Llama3.2-3B	62 (-11/13)	63 (-10/13)	64 (-10/11)	64 (-11/14)	58 (-11/13)
SmolLM2-360M	14 (-20/24)	13 (-14/17)	13 (-18/20)	9 (-14/17)	52 (-11/13)
SmolLM2-1.7B	25 (-13/16)	24 (-13/16)	25 (-13/14)	20 (-14/17)	68 (-12/13)
Qwen-1.5B	36 (-21/26)	34 (-20/24)	38 (-15/17)	36 (-17/20)	35 (-22/24)

Table 4. Pass@1 %-Hits/Accuracy ( higher is better, +95% CI) reflecting per-run variability for Rudder (async).

Figure 20. Trajectories of %-Hits (left) and communication volume (right) of papers100M on 64 nodes with LLM and MLP.

MoE LLMs	Pass@1 %-Hits	Replacement Interval ( $r$ )	Valid/Invalid Response (%)	(+ve/-ve) Decisions (%)
Granite3.1-3B	44	21	99/1	61/39
Mixtral-8x7B	50	21	94/6	56/44
Mixtral-8x22B	56	42	100/0	86/14

Table 5. Pass@1 %-Hits ( higher is better) of different-sized MoE models using products on 16 trainers.

Figure 21. Training times ( lower is better) for MoEs across 5–25% buffer sizes on 16 trainers. Note: Mixtral-8x22B stalled in 10%.

Explore similar work

Nereus: Adaptive Parallelism for LLM Post-Training
Sep 28, 2026Songlin Jiang, Tuo Shi, Sitong Zhang +3LLM Inference Scheduling RL Post-Training
Aalto University · Shenzhen University of Advanced Technology · Zhejiang University
LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs
Aug 7, 2026Liad Gerstman, Aditya Dhakal, Dejan Milojicic +1Graph Neural Networks Communication-Efficient Distributed Training
Technion Israel Institute of Technology Haifa, Israel · Hewlett Packard Enterprise Labs Fort Collins, CO, USA · Hewlett Packard Enterprise Labs Milpitas, CA, USA
Routed Graph Handoff: Adaptive Format Selection for Multi-Agent LLM Delegation
Aug 26, 2026Pratyay Banerjee, Ankit ChadhaMulti-Agent LLM Systems LLM Agents
Amazon, AGI / Sunnyvale, USA

Sep 28, 2026J/K move · Enter open · S save

Nereus: Adaptive Parallelism for LLM Post-Training

Songlin Jiang, Tuo Shi, Sitong Zhang, Zeke Wang+2

Aalto University · Shenzhen University of Advanced Technology · Zhejiang University

LLM Inference Scheduling RL Post-Training LLM Training

Open paper AskPDF arXiv

Reinforcement learning (RL) post-training for large language models (LLMs) coordinates multiple models across generation, inference, and training on GPU clusters. Several factors may change during a run, including resource availability, sequence length, memory pressure, and stage bottlenecks. As a consequence, an execution plan that was initially suitable can then become slow or even infeasible over time. However, adapting a job whose models share GPUs entails significant challenges: deciding whether a new plan is worth the transition cost, reusing the job's distributed state, and coordinating GPU transfers across models and stages. Nereus targets these challenges as a cost-aware runtime that adapts RL post-training jobs into efficient execution plans. Its low-overhead controller selects a memory-feasible global plan and admits the transition using a cost model calibrated against the running job. To estimate and execute a transition, Nereus represents the distributed state of each replica of a model-stage (one model in one stage) as an Elastic Model Unit. It then employs a global transition graph to order the transformations and GPU transfers of these units. In a trace built from real data, online TP/PP adaptation reduces average step latency by 27.7% relative to the initial fixed TP/PP layout with DP scaling. In a 1,000-step run reaching 1,024 GPUs, six transitions consume 0.079% of total run time. Nereus improves end-to-end 8B PPO throughput by 2.14--7.27

\times

over OpenRLHF and by 1.10--1.47

\times

over Verl across diverse clusters.

Abstract

Explore similar work

Nereus: Adaptive Parallelism for LLM Post-Training

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

Routed Graph Handoff: Adaptive Format Selection for Multi-Agent LLM Delegation