Rudder: Steering Prefetching in Distributed GNN Training using LLM Agents

This paper introduces Rudder, an LLM-agent-based software module integrated into the AWS DistDGL framework that autonomously adapts prefetching strategies for distributed Graph Neural Network training, achieving up to 91% performance improvement and over 50% communication reduction compared to baseline methods.

The Gist

The Big Picture: The "Data Delivery" Problem

Imagine you are running a massive, high-speed pizza delivery service for a city (this is Distributed Graph Neural Network Training).

• The Pizzas: These are the data points (nodes) from a giant map (the graph).
• The Drivers: These are your computer processors (GPUs) working in parallel.
• The Problem: The city is huge, and the drivers are scattered across different neighborhoods. To make a pizza, a driver needs ingredients (data) that might be stored in a warehouse in a completely different part of town.

In traditional systems, every time a driver needs an ingredient, they have to stop, call the warehouse, wait for the truck to drive over, and then continue cooking. This "waiting time" (communication) is the biggest bottleneck. It slows down the whole operation.

The Old Solutions: The "Static" and "Guessing" Approaches

To fix this, engineers tried two main things:

1. The Static Plan: "Let's just bring everything to the warehouse now."
   • The Flaw: The city is too big. You can't store every single pizza ingredient in one place. You run out of space (memory).
2. The Fixed Schedule: "Let's bring ingredients every 10 minutes, no matter what."
   • The Flaw: Sometimes you need ingredients now, and sometimes you don't need them for an hour. A rigid schedule wastes time and space. It's like a delivery driver bringing a fresh tomato when the recipe calls for cheese.

The New Solution: Rudder (The Smart Captain)

The authors introduce Rudder. Think of Rudder as a super-smart, autonomous ship captain (an AI Agent) steering the delivery fleet.

Instead of following a rigid schedule or trying to carry everything, Rudder uses a Large Language Model (LLM)—the same kind of AI that writes poems or answers questions—to make real-time decisions.

How Rudder Works (The Analogy)

Imagine the ship's captain (Rudder) has a small, magical pantry (the Persistent Buffer) on the deck.

1. Observation: The captain constantly checks the weather, the speed of the ship, and the current menu (the Metrics).
2. Reasoning: Instead of just following a rulebook, the captain thinks.
   • Scenario A: "The ship is moving fast, and we are about to need cheese. But the pantry is full of stale bread we won't use. Let's throw out the bread and fetch fresh cheese before we even ask for it."
   • Scenario B: "The ship is slowing down. We have plenty of time. Let's wait and not waste fuel fetching ingredients we might not need."
3. Action: The captain tells the crew exactly what to swap in the pantry and what to fetch next.

Why Use an AI Captain (LLM) Instead of a Human or a Robot?

The paper compares Rudder to two other types of managers:

1. The Robot (Traditional Machine Learning):

   • How it works: You have to teach the robot for weeks using old delivery logs. It learns a pattern: "If it's Tuesday, bring cheese."
   • The Problem: If the city layout changes or the weather gets weird (a new dataset), the robot gets confused because it was only trained on old data. It needs to be re-taught every time things change.

2. The Human Captain (The LLM Agent):

   • How it works: This captain has read millions of books about logistics, cooking, and traffic. You don't need to train them on your specific city. You just hand them the current map and say, "We need to optimize for speed right now."
   • The Magic: They use In-Context Learning. They look at the current situation, reason through it step-by-step (like a human thinking), and make a smart guess immediately. They adapt to any city, even one they've never seen before.

The Results: Why Rudder Wins

The researchers tested Rudder on a supercomputer (NERSC Perlmutter) with real-world data. Here is what happened:

• Speed: Rudder made the training 91% faster than doing nothing, and 82% faster than the old "fixed schedule" methods.
• Efficiency: It cut the "waiting for trucks" (communication) time by over 50%.
• Adaptability: When the researchers changed the rules of the game (different graph sizes, different computer setups), Rudder didn't break. It just adjusted its strategy on the fly. The old robots (Machine Learning classifiers) struggled because the new data didn't match their old training.

The "Secret Sauce"

The paper highlights a fascinating discovery: You don't need a giant AI to do this.

• They found that a small, lightweight AI (a "Small Language Model") works just as well as a massive, expensive one.
• It's like having a sharp, experienced local captain rather than a slow, over-thinking giant. The small captain is fast, uses less fuel (computer memory), and makes great decisions because it can "reason" about the immediate situation.

Summary

Rudder is a smart tool that helps computers learn from massive, complex maps (graphs) much faster. Instead of using rigid rules or slow, pre-trained robots, it uses a smart AI captain that looks at the current situation, thinks about what to do next, and swaps out old data for new data just in time. This saves massive amounts of time and computing power, making AI training faster and more efficient.

Technical Summary

1. Problem Statement

Large-scale Graph Neural Network (GNN) training on distributed systems faces significant performance bottlenecks due to irregular and unpredictable communication.

• The Challenge: GNNs require knowledge of a vertex's neighborhood. In distributed settings, graphs are partitioned across Processing Elements (PEs). Sampling neighbors often requires fetching data from remote partitions, leading to frequent, irregular communication that stalls forward progress.
• Limitations of Existing Solutions:
  • Static Prefetching: Traditional methods use fixed heuristics or static policies to prefetch remote nodes. These fail to adapt to dynamic conditions such as varying graph structures, sampling parameters, batch sizes, and network contention.
  • Overhead of Tuning: Finding optimal static parameters requires costly trial-and-error for every new dataset or configuration.
  • ML Classifier Limitations: While Machine Learning (ML) classifiers can learn from traces, they require extensive offline training on labeled data. They struggle with distribution shifts (e.g., unseen graph topologies or batch sizes) and lack the ability to reason about long-term trade-offs without retraining.

2. Methodology: Rudder

The authors propose Rudder, a software module embedded in the AWS DistDGL framework that uses Large Language Model (LLM) Agents to autonomously steer the prefetching and replacement of remote nodes.

Core Architecture

• Adaptive Replacement Strategy: Rudder maintains a fixed-size local persistent buffer. Instead of static rules, it uses an agent to decide when and what to replace based on real-time metrics.
• The Agent Loop:
  1. Metrics Collector: Streams runtime metrics (e.g., % Hits, communication volume, minibatch progress) to the agent.
  2. Context Builder: Tracks historical replacement decisions and their outcomes to provide temporal context.
  3. Decision Maker (LLM/ML): Analyzes the current state and history to output a binary decision: Replace (fetch new nodes) or Skip.
• Execution Model:
  • Asynchronous (Default): The prefetching and inference tasks run in background threads, overlapping with the main GNN training loop. This minimizes latency impact.
  • Synchronous: The trainer waits for the agent's decision (higher accuracy but potential stalling).
• Scoring Policy: A frequency-tracking mechanism penalizes "stale" nodes (those not accessed recently) to prioritize refreshing the buffer with relevant data.

Why LLM Agents?

The paper argues that LLMs offer unique advantages over traditional ML classifiers for this control problem:

• In-Context Learning (ICL): LLMs can adapt to new tasks (unseen graphs/configurations) via zero-shot prompting without gradient updates or offline training.
• Reasoning: LLMs exhibit logical multi-step reasoning, allowing them to anticipate the consequences of a replacement decision (e.g., "If I replace now, communication will drop, but hits might stagnate").
• Resilience: They handle distribution shifts better than supervised classifiers trained on specific traces.

3. Key Contributions

1. LLM-Driven Adaptive Prefetching: Introduction of Rudder, the first system to use LLM agents for autonomous, real-time replacement decisions in distributed GNN training.
2. Comparative Design Study: A comprehensive evaluation comparing LLM agents against traditional ML classifiers (MLP, XGBoost, SVM, etc.) and static heuristics.
3. Zero-Shot Adaptation: Demonstration that LLM agents can generalize to unseen datasets and configurations without retraining, whereas ML classifiers suffer from performance degradation on out-of-distribution data.
4. System Integration: Implementation within the state-of-the-art DistDGL framework, utilizing asynchronous threading to hide inference overhead.

4. Experimental Results

The system was evaluated on the NERSC Perlmutter supercomputer using diverse datasets (OGB, SNAP, Yelp) and various LLMs (Gemma, Llama, SmolLM, Mixtral).

• Performance Gains:
  • vs. Baseline (No Prefetch): Up to 91% improvement in end-to-end training time.
  • vs. Static Prefetching: Up to 82% improvement over static policies.
  • Communication Reduction: Reduced remote node communication by over 50%.
• Data Persistence: Rudder consistently achieved 20–50% higher % Hits (percentage of required nodes found in the local buffer) compared to static methods.
• LLM vs. ML Classifiers:
  • LLM Agents (e.g., Gemma3-4B): Achieved the highest Pass@1 scores (~76–82%), indicating high accuracy in predicting the system state after a decision. They required no offline training.
  • ML Classifiers: Required extensive offline training and showed lower accuracy on unseen datasets (distribution shifts). They tended to make replacement decisions too frequently, leading to diminishing returns and higher communication overhead.
• Model Size Analysis:
  • Small/Medium LLMs (<5B params): Performed best, balancing reasoning capability with low inference latency and memory footprint.
  • Large/MoE Models: Larger models (e.g., Mixtral-8x22B) did not necessarily perform better and often suffered from high latency or quantization-induced reasoning degradation.
• Trade-offs: The system allows tuning between data persistence (buffer size) and communication volume. Smaller buffers (5%) with Rudder still outperformed larger buffers with static policies.

5. Significance and Impact

• Paradigm Shift: Rudder moves GNN optimization from static, heuristic-based tuning to dynamic, agent-based control. It treats the prefetching problem as a sequential decision-making task solvable by generative AI.
• Efficiency: It demonstrates that small, quantized LLMs can effectively manage complex system resources in High-Performance Computing (HPC) environments, offering a viable alternative to heavy reinforcement learning or manual tuning.
• Scalability: The asynchronous design ensures that the control overhead does not bottleneck the training process, making it suitable for large-scale distributed training on thousands of GPUs.
• Generalizability: The approach suggests that LLM agents are well-suited for other HPC optimization problems where the search space is massive, ill-defined, and subject to dynamic changes, as they can perform approximate search without formal goodness functions.

In summary, Rudder leverages the emergent reasoning and in-context learning capabilities of LLMs to solve the dynamic optimization problem of distributed GNN prefetching, achieving significant reductions in communication overhead and training time without the need for expensive offline training or manual parameter tuning.