ASAP: Auto-optimized LLM Training
- ASAP is a multi-agent system that automates performance engineering of large-scale LLM training by diagnosing bottlenecks and recommending optimized sharding configurations.
- It leverages three specialized agents—Coordinator, Analyzer, and Proposal—coupled with persistent logging in Sharding Memory for traceability and explainability.
- Experimental results demonstrate up to a 28% reduction in training step time and a 1.43× throughput improvement, with further gains from human expert intervention.
ASAP, short for Agentic Solution to Auto-optimize Performance of Large-Scale LLM Training, is a multi-agent system for automated performance engineering of distributed large-language-model training on domain-specific accelerator systems. It addresses an optimization setting characterized by a large and rapidly changing search space in which manual tuning is slow and black-box search is resource-intensive. The system combines LLM-based reasoning with performance profiling tools, roofline analysis, and a knowledge base of best practices and past successful optimizations, with the specific aim of diagnosing bottlenecks and recommending optimized sharding configurations for distributed training (Ding et al., 5 Nov 2025).
1. Definition and problem setting
ASAP is designed for the optimization of large-scale LLM training on distributed accelerator platforms, where performance depends on interacting choices over sharding, communication, memory use, and hardware topology. The motivating problem is not merely parameter tuning in isolation, but the diagnosis of system-level bottlenecks from heterogeneous evidence such as step-time metrics, device utilization indicators, HLO operation profiles, and roofline-derived limits (Ding et al., 5 Nov 2025).
The system is framed as an agentic approach because optimization is decomposed across cooperating LLM agents rather than handled by a single monolithic optimizer. In the reported design, this decomposition serves two purposes. First, it separates orchestration, diagnosis, and proposal generation into distinct roles. Second, it preserves traceability by recording requests, tool invocations, model responses, and user dialogs in persistent memory. This suggests a design emphasis not only on automation but also on explainability and post hoc inspection (Ding et al., 5 Nov 2025).
The immediate optimization target in ASAP is the sharding configuration for distributed training. In the reported implementation, the Proposal Agent synthesizes candidate values for ici_mesh = {model, data, replica, seq}, together with expected impact, trade-off analysis, confidence, and supporting citations from either historical records or a document knowledge base. The central claim is therefore narrower than full-stack compiler or kernel autotuning: ASAP primarily automates diagnosis and sharding recommendation, while still allowing later intervention by human experts (Ding et al., 5 Nov 2025).
2. Multi-agent system architecture
ASAP consists of three collaborating LLM agents plus a persistent memory module, all running on Google’s Agent Development Kit (ADK) and communicating through simple JSON messages stored in Sharding Memory (Ding et al., 5 Nov 2025).
| Component | Role | Primary inputs or outputs |
|---|---|---|
| Coordinator Agent | System entry point and orchestration | Experiment record, Analyzer report, final recommendations |
| Analyzer Agent | Bottleneck diagnosis | Xprof KPIs, HLO-op profiles, roofline data, structured bottleneck report |
| Proposal Agent | Sharding proposal synthesis | Bottleneck report, current sharding config, PastOpsDB, DocKB |
| Sharding Memory | Traceability and logging | Agent requests, tool invocations, LLM responses, user dialogs |
The Coordinator Agent is invoked with an experiment ID. It fetches the experiment record, including sharding configuration, code version, hyperparameters, and hardware topology. It forwards raw profiling data to the Analyzer, receives the bottleneck report, sends that report to the Proposal Agent, and returns final recommendations. In architectural terms, it provides workflow control but not substantive diagnosis or optimization logic (Ding et al., 5 Nov 2025).
The Analyzer Agent ingests raw Xprof KPI metrics, HLO-op profiles, and on-device roofline data. Its function is to correlate high-level KPIs such as step_time_ms, device_duty_cycle%, and communication_overhead% with operation-level statistics including FLOPs, bytes accessed, and operation execution time. It then emits a structured report with bottleneck labels such as “HBM bottleneck at all-reduce” or “compute-bound at matmul” (Ding et al., 5 Nov 2025).
The Proposal Agent receives the bottleneck report and the current sharding configuration. It uses a retrieval-augmented generation workflow that first queries a historical database of prior optimizations and then falls back, when needed, to a document knowledge base based on “How To Scale Large Models.” It synthesizes three distinct sharding proposals, each containing a new ici_mesh, expected impact, trade-off analysis, confidence score, and citations to the supporting material (Ding et al., 5 Nov 2025).
The Sharding Memory is a file-based log of all agent requests, tool invocations, LLM responses, and user dialogs. Its stated purpose is traceability, debugging, and re-analysis. A plausible implication is that ASAP is intended to be auditable in a way that pure black-box search systems often are not, although the paper’s concrete description remains at the level of logging and replay support (Ding et al., 5 Nov 2025).
3. Analytical core: profiling, roofline analysis, and retrieval
ASAP’s analytical mechanism is built on a tight coupling between quantitative system telemetry and natural-language LLM reasoning. The profiling layer is provided through the Xprof API, which exposes KPIs such as step_time_ms, device_duty_cycle%, and communication_overhead%, alongside per-operation HLO records of FLOPs_i, Bytes_i, and Time_i (Ding et al., 5 Nov 2025).
The Analyzer performs a roofline analysis per operation. In the exposition of the system, the relevant definitions are
and
An equivalent notation later defines
Using these quantities, the Analyzer labels operations as compute-bound, memory-bound, or communication-bound. The paper also states decision heuristics in time-domain terms: if Time_i ≈ FLOPs_i / PeakFLOPS, the operation is compute-bound, whereas if Time_i ≈ Bytes_i / PeakBandwidth, it is memory-bound (Ding et al., 5 Nov 2025).
This numerical analysis is then translated into text representations consumable by the LLM, for example: “Op all-reduce: 11.02% roofline efficiency, bottleneck=HBM.” The Proposal Agent uses chain-of-thought prompts to map such diagnoses to sharding adjustments, including examples such as increasing model parallelism to reduce per-chip HBM footprint. The paper’s core methodological point is therefore not LLM reasoning alone, but LLM reasoning constrained by profiler outputs, roofline structure, and retrieved optimization precedents (Ding et al., 5 Nov 2025).
The retrieval layer is split into two stores. The Historical Database, named “PastOpsDB,” contains records with schema {experiment_hash, hardware_topology, model_size, baseline_config, optimized_config, Δstep_time%, Δthroughput×} and is indexed by model-size range, device type, and bottleneck type. The Document Knowledge Base, “DocKB,” is based on “How To Scale Large Models” chapters on model, data, and sequence parallelism and is stored as vector embeddings for similarity search. PastOpsDB is updated automatically each time an engineer approves a proposal and logs the outcome (Ding et al., 5 Nov 2025).
4. End-to-end workflow and proposal generation
The end-to-end logic of ASAP is given in high-level pseudocode. The Coordinator first retrieves the experiment record and Xprof profile, then invokes the Analyzer, then the Proposal Agent, then appends record, report, and proposals into Sharding Memory, and finally returns the highest-confidence proposal (Ding et al., 5 Nov 2025).
The Analyzer’s internal loop iterates over HLO profiles, computes intensity, derives roofline efficiency, labels each operation, and synthesizes a bottleneck report. The Proposal Agent then executes a retrieval path over PastOpsDB, optionally retrieves documents from DocKB when similarity is below threshold, and generates three proposals through the LLM using a prompt template of the form “Given a {bottleneck}, suggest ici_mesh changes...” (Ding et al., 5 Nov 2025).
Several features of this workflow are structurally important. First, proposal generation is conditioned on the current sharding configuration rather than performed from scratch. Second, historical reuse is preferred when there is a sufficiently similar prior case. Third, document retrieval functions as a backstop when historical precedent is weak. This suggests a hybrid between case-based reasoning and principle-based synthesis, though that characterization is interpretive rather than explicitly named in the system description (Ding et al., 5 Nov 2025).
The output of the Proposal Agent is not a single recommendation but three distinct sharding proposals. Each includes a new ici_mesh, an expected impact statement, trade-off analysis, a confidence score, and citations to the knowledge source. The system then selects the best candidate through choose_highest_confidence(proposals). The explicit presence of confidence and citation fields is central to ASAP’s explainability claim: recommendations are presented with rationale rather than as opaque search outputs (Ding et al., 5 Nov 2025).
5. Experimental setting and quantitative performance
The reported evaluation uses three TPU platforms: TPU v5p-512 (8×8×8 mesh), TPU v6e-16 (4×4 mesh), and TPU v5e-256 (16×16 mesh). The software stack consists of Google ADK agents built on Gemini Pro LLM, the Xprof profiling API, and a retrieval stack using a Faiss index with Gemini for document retrieval. The evaluated models are LLMs in the 10B–100B parameter range, including proprietary Gemini-family models. The benchmarks track training step time, device duty cycle, throughput measured as tokens/sec or goodput, and communication overhead (Ding et al., 5 Nov 2025).
| Experiment | Baseline Step Time (ms) | ASAP Step Time (ms) |
|---|---|---|
| 1: Compute-Bound (v5p) | 35656.45 | 25632.94 |
| 2: HBM-Bound (v6e) | 2516.50 | 1758.32 |
| 3: Comm-Bound (v5e) | 7354.69 | 4355.15 |
The aggregate results are reported as up to 28% training step time reduction and 1.43 times throughput improvement from ASAP-generated sharding configurations. When combined with additional optimization from human experts, throughput is reported to increase further to 2.58 times (Ding et al., 5 Nov 2025).
The three named experiments clarify how these gains arise. In Experiment 1, a compute-bound workload on v5p improved from 35656.45 ms to 25632.94 ms, corresponding to –28%, and the paper states that ASAP’s top proposal matched the human-engineered solution exactly. In Experiment 2, an HBM-bound workload on v6e improved from 2516.50 ms to 1758.32 ms, with a 1.43× throughput increase, and the text notes that the model-parallel sharding solution was novel relative to the history database. In Experiment 3, a communication-bound workload on v5e improved from 7354.69 ms to 4355.15 ms, and total throughput reached 2.58× after additional batch-size tuning by human experts on top of ASAP’s proposal (Ding et al., 5 Nov 2025).
The article’s summary numbers should be read with care. The table includes per-experiment values of 1.28×, 1.43×, and 2.58×, and reports an average or total line with ≈28%, 1.43×, and 2.58×. The text specifically emphasizes the 28% step-time reduction and 1.43× throughput improvement attributable to ASAP-generated sharding configurations, with 2.58× attributed to the combined human-plus-ASAP setting (Ding et al., 5 Nov 2025).
6. Limitations, failure modes, and prospective extensions
The stated limitations of ASAP are concrete and system-specific. It depends on accurate and complete Xprof data, so sparse or skewed profiling can lead to misdiagnosis. It currently considers only a single global sharding configuration and cannot yet produce per-layer heterogeneous sharding. It also lacks automated closed-loop feedback: proposals must be validated by engineers or by external scripts rather than by an integrated evaluator (Ding et al., 5 Nov 2025).
The described failure modes follow directly from these design constraints. Over-reliance on the history database may cause repetition of suboptimal prior solutions when context has drifted. Low-quality retrieval can induce LLM hallucinations in reasoning steps. Proposed configurations may also trigger out-of-memory errors when they increase on-device state without explicit checks. These are not framed as theoretical edge cases but as practical operating risks within the present system design (Ding et al., 5 Nov 2025).
The future-work agenda extends the architecture in four directions. The paper proposes an Evaluator Agent to run quick synthetic tests and feed results back, thereby closing the loop through reinforcement learning. It also proposes extension to per-layer or operator-level sharding strategies, incorporation of compiler-level optimizations such as operator fusion and memory layout, and the addition of cross-layer roofline models to predict global end-to-end throughput trade-offs (Ding et al., 5 Nov 2025).
Taken together, these limitations and extensions delineate ASAP’s current scope. It is already positioned as a scalable and explainable methodology for AI-assisted performance engineering in large-scale LLM training, but its present implementation remains centered on profiler-informed diagnosis and sharding recommendation rather than full autonomous optimization. This suggests that ASAP is best understood as a structured co-optimization system in which LLM agents, system telemetry, and human expertise remain mutually coupled rather than fully substitutive (Ding et al., 5 Nov 2025).