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
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 12. DistDGL+Rudder’s performance with agents and classifiers across datasets, #trainers, and persistent buffer sizes (5%/25%). The bars show mean epoch time (left Y; lower is better; 5%/25% bars are overlapped), with %-Hits ( × =5%/ × =25%) (right Y; higher is better).
(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
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.
Figure 21
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%.
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× over OpenRLHF and by 1.10--1.47× over Verl across diverse clusters.
Songlin Jiang, Tuo Shi, Sitong Zhang +3
Aalto University · Shenzhen University of Advanced Technology · Zhejiang University
Graph Neural Networks (GNNs) are widely used across domains such as natural sciences, social network analysis, chip design, and recommendation systems. However, as graph sizes grow, storing and processing them entirely on a single-node CPU-GPU system becomes increasingly impractical. A promising approach is to distribute the graph across multiple remote memory nodes, though this introduces a major bottleneck: inter-node network congestion during training. To address this, we propose LGNNIC, a novel inter-node system architecture that leverages SmartNICs co-located with remote memory nodes-a configuration already available in modern systems-to reduce communication overhead in distributed GNN training. LGNNIC offloads key preprocessing tasks to SmartNICs, reducing the volume of data transferred to computational (training) nodes and alleviating network congestion. We introduce two complementary techniques executed on the SmartNICs during the preprocessing phase: Neighbor Sampling, which performs mini-batch sampling, and Quantization of the sampled batches. To evaluate LGNNIC under different communication infrastructures, we designed both an optimized low-overhead DMA-based synchronization mechanism and a high-overhead socket-based alternative used as a benchmark. We evaluate the core SmartNIC offloading mechanisms across standard GNN workloads and sampling hyperparameters using a proof-of-concept (PoC) system comprising one remote-memory node with an NVIDIA BlueField-2 SmartNIC and one compute node with an A100 GPU. Both Neighbor Sampling and Quantization on the remote node demonstrated substantial training speedups in most configurations. Neighbor Sampling achieved up to 62.4x and 17.5x speedups with Sockets and DOCA-DMA, respectively, primarily due to reduced data transaction time. Quantization provided additional speedups of up to 3.6x and 1.3x, respectively, by reducing data transfer.
Liad Gerstman, Aditya Dhakal, Dejan Milojicic +1
Technion Israel Institute of Technology Haifa, Israel · Hewlett Packard Enterprise Labs Fort Collins, CO, USA · Hewlett Packard Enterprise Labs Milpitas, CA, USA
Multi-agent LLM systems coordinate through natural-language messages that consume 40--60% of their token budget. Replacing these with structured graphs reduces cost but fails on tasks requiring adaptive reasoning. We propose Routed Graph Handoff, where a lightweight LLM router (155 tokens, 0.15% overhead) selects between a typed dependency graph and natural language for each delegation. On four benchmarks (1,050+ trajectories), the routed system matches or exceeds NL-only on every task: +12.7,pp on τ-retail at 3.2× compression (p<0.01), +8.7,pp on BrowseComp at 2.2× compression (p<0.05), and parity on BFCL and AppWorld. Without the router, graph-only delegation regresses 14.6,pp on AppWorld; the router eliminates this at near-zero cost. A graph-aware executor prompt is required: the same schema without interpretation guidance yields no gain. An oracle analysis reveals 8.6,pp of additional headroom, motivating execution-time adaptive routing as future work.