Rudder: Steering Prefetching in Distributed GNN Training using LLM Agents

TL;DR

Rudder is an adaptive prefetching module for distributed GNN training that leverages LLMs to dynamically optimize remote node fetching, significantly improving training performance and reducing communication overhead.

Contribution

This paper introduces Rudder, a novel LLM-based adaptive prefetching system integrated into AWS DistDGL for efficient distributed GNN training.

Findings

Achieves up to 91% performance improvement over baseline

Reduces communication by over 50%

Outperforms static prefetching strategies

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.