ASAP: an Agentic Solution to Auto-optimize Performance of Large-Scale LLM Training

AI summary

Overview

Research area: Multi-agent systems applied to performance engineering for large-scale distributed LLM training (the paper is listed under cs.MA on arXiv, arXiv:2511.03844v1).

Technical level: Advanced. The paper assumes familiarity with distributed training concepts (data, model, and sequence parallelism), accelerator profiling vocabulary (HLO operation profiles, roofline analysis, duty cycle), and multi-agent LLM system design.

Scope: The paper introduces ASAP, a three-agent LLM framework that diagnoses training bottlenecks from profiling data and recommends optimized sharding configurations on Google TPU systems.

What This Paper Is About

Training large language models across many accelerators requires choosing a "sharding" configuration — how the model, data, and sequence are split across chips — and finding a good configuration currently relies on slow manual trial and error or expensive black-box search. The paper builds a multi-agent system (ASAP) that reads profiling data and a curated knowledge base, diagnoses whether a workload is limited by compute, memory bandwidth, or communication, and proposes improved sharding configurations with written justifications. The goal is to make this optimization loop faster, more explainable, and less dependent on scarce human expertise.

Key Contributions

1. A three-agent architecture for sharding optimization. ASAP consists of a Coordinator Agent (orchestrates the workflow and routes data), an Analyzer Agent (diagnoses bottlenecks), and a Proposal Agent (generates optimized sharding configurations), plus a Sharding Memory module that persistently stores tool calls, logs, and agent responses for context and traceability.

2. Integration of LLM reasoning with real profiling infrastructure. The Analyzer Agent consumes raw experiment configurations and data from the Xprof profiling API, including KPIs such as accelerator busy time, training step time, and goodput, HLO operation profiles (execution time, FLOPs utilization, memory bandwidth usage, tensor shapes), and on-device roofline data (total FLOPs, bytes accessed, execution time).

3. A retrieval-augmented proposal workflow grounded in expert knowledge. The Proposal Agent analyzes the Analyzer's report, queries a database of past optimizations (before-and-after configurations, Xprof links, summaries of changes and impact), queries a document knowledge base adapted from "How to Scale Large Models," and synthesizes three distinct sharding proposals, each with justification, expected impact, trade-offs, and citations.

4. A three-scenario empirical evaluation on three TPU generations. The framework was tested on TPU v5p, v6e, and v5e platforms against human-engineered solutions, with reported reductions in step time and gains in throughput.

Main Findings

• Compute-bound workload on TPU v5p-512 (Experiment 1): The baseline ici_mesh was {'model': 8, 'data': 8, 'seq': 8} with a device duty cycle of 99.73%, communication overhead of 2.57%, step time of 35656.45 ms, batch size 256, sequence length 8192, and an 8x8x8 TPU topology. The agent recommended {'model': 8, 'data': 16, 'seq': 4}, which precisely matched the human-engineered solution and decreased TPU compute time by 28% while improving FLOPs and HBM utilization. This case was present in the agent's historical optimization database, and the agent expressed roughly 85–90% confidence.

• HBM-bound workload on TPU v6e-16 (Experiment 2): The baseline was {'data': 4, 'model': 1, 'replica': 1, 'seq': 4} on a 4x4 topology with a duty cycle of 99.95%, communication overhead of 7.52%, and step time of 2516.50 ms. The top time-consuming operations, including all-reduce and memory-bound fusions, were limited by HBM bandwidth rather than inter-chip interconnect. The agent's proposal {'replica': 1, 'data': 4, 'model': 4, 'seq': 1} was identical to the human-engineered solution, which improved throughput by 1.43x. This optimization was not in the historical database, and the agent maintained its stance with roughly 90% confidence when challenged.

• Communication-bound workload on TPU v5e-256 (Experiment 3): The baseline was {'data': 4, 'model': 4, 'seq': 16}. Despite a near-100% duty cycle, the top five most time-consuming operations were all collective communication operations (all-reduce, all-reduce-scatter, collective-permute) bound by memory bandwidth during data exchange. The agent recommended {'replica': 1, 'data': 8, 'seq': 16, 'model': 2} to cut model parallelism from 4 to 2 and raise data parallelism from 4 to 8. The proposal matched the human expert's final configuration; combined with an increased batch size, it led to a 2.58x increase in total throughput. This optimization was also novel, with confidence stated at 85–90%.

• Generalization beyond memorized examples. Experiments 2 and 3 involved configurations absent from the historical optimization database, and the agent still produced correct solutions by reasoning over analysis reports and knowledge-base documents.

• Summary performance claims. The abstract reports that ASAP-generated sharding configurations can contribute up to 28% training step time reduction and 1.43x throughput improvement, and that when combined with additional optimization from human experts, throughput can be further increased to 2.58x.

• Qualitative rather than automated evaluation. For each scenario, the agent received an experiment ID and Xprof profiling data from a baseline run, and its proposed configurations were compared qualitatively against those developed by human performance experts; the authors state the goal was to assess alignment with expert-derived solutions.

Methodology in Plain English

The system is built on Google's Agent Development Kit (ADK), uses Gemini Pro as the underlying model, and is invoked from the command line with an experiment ID — a key that retrieves the experiment's full record, including its sharding configuration, hyperparameters, code version, and system environment.

The Coordinator Agent receives that ID and directs the workflow. The Analyzer Agent pulls raw data from the experiment configuration and the Xprof profiling API, fetching KPIs, HLO operation profiles, and on-device roofline analysis data supplied as JSON strings. By correlating high-level metrics with operation-level detail, it identifies whether the workload is limited by compute, memory, or communication, and names the specific operations responsible.

The Proposal Agent then runs a retrieval-augmented generation workflow: it interprets the Analyzer's report, looks up similar model and hardware setups in a database of past optimizations, consults a document knowledge base of scaling principles, and outputs three candidate sharding configurations. Each ici_mesh proposal specifies data, model, replica, and sequence parallelism dimensions and comes with a justification, an expected performance impact, a discussion of trade-offs, and citations. The Sharding Memory module keeps a persistent, file-based worklog of all tool calls, LLM responses, and user inputs.

In the experiments, the authors fed each baseline's profiling data to the agent, collected its proposals, and compared the recommended sharding configuration against the one human experts had actually adopted. No dataset sizes or standard benchmarking suites are reported.

Why This Matters

Distributed LLM training is expensive, and a suboptimal sharding choice translates directly into wasted accelerator time. The paper argues that reasoning-based agents can navigate this complex search space more cheaply and explainably than black-box optimization, and that the resulting recommendations can be reviewed by a human engineer or evaluated on hardware.

Impact on research: The work positions agentic LLM reasoning as a complement to analytical performance models and learned performance predictors, and offers a template for AI-assisted performance engineering where explanations and citations matter as much as the recommendation.

Real-world applications:

• Automating sharding configuration selection for teams training large models on multi-chip TPU pods.
• Triaging profiling output (KPIs, HLO profiles, roofline data) that normally requires deep expert interpretation.
• Building institutional knowledge bases of successful optimizations that persist and generalize across projects.
• Reducing trial-and-error experiment cycles in accelerator capacity planning, where each wasted run has a direct dollar cost.

Industry relevance: The reported gains (up to 28% step time reduction, 1.43x throughput from the agent alone, 2.58x with additional human optimization) target the exact cost center of frontier model development, and the framework is built on commercially available tooling (ADK, Gemini Pro, Xprof) rather than a research-only stack.

Future Directions

• Remove the dependency on accurate profiling data, which the authors identify as a current limitation, so the agent remains useful when instrumentation is incomplete or noisy.
• Move beyond a single global sharding configuration to support more sophisticated techniques such as per-layer sharding.
• Add an Evaluator Agent to close the feedback loop, letting the system refine strategies through reinforcement learning and discover novel configurations rather than only proposing them.
• Broaden the optimization scope to jointly balance on-device computation against network communication, and to incorporate compiler-level optimizations such as operator fusion and memory layout selection.

Target Audience

Performance engineers and ML systems researchers working on large-scale distributed training, especially those responsible for parallelism and sharding decisions on accelerator pods. It also suits multi-agent systems researchers looking for a concrete domain where agentic reasoning plus retrieval is evaluated against expert ground truth. Readers need grounding in distributed training terminology and accelerator profiling; those without it will still follow the architecture and results but may find the analysis details harder to interpret.

Authors’ abstract

Optimizing large-language model (LLM) training on distributed domain-specific accelerator systems presents significant challenges due to its complex optimization space. Existing optimization methods, however, rely on time-consuming manual tuning or resource-intensive black-box searches, which struggle to keep pace with the rapidly evolving LLM domain, leading to slow development and underutilized resources. To address this, we introduce ASAP, an Agentic Solution to Auto-optimize Performance of Large-Scale LLM Training. It is a multi-agent system, featuring Coordinator, Analyzer, and Proposal agents, which integrates LLM reasoning with insights from performance profiling tools, roofline analysis, and a knowledge base of best practices and successful past optimizations from human experts. Our proposed design can automate the diagnosis of performance bottlenecks and recommend optimized sharding configurations with reasoning, thus effectively improving the efficiency of distributed LLM training. Experiments have shown that the ASAP-generated sharding configurations can contribute up to 28% training step time reduction and 1.43 times throughput improvement. When combined with additional optimization from human experts, throughput can be further increased to 2.58 times. The proposed ASAP promises to provide a scalable and explainable methodology for AI-assisted performance engineering in large-scale LLM training.