ASAP: an Agentic Solution to Auto-optimize Performance of Large-Scale LLM Training
Yuran Ding, Xinwei Chen, Xiaofan Zhang, Zongwei Zhou
TL;DR
ASAP tackles the challenge of efficiently optimizing large-scale LLM training on distributed accelerators by introducing a multi-agent framework that leverages LLM reasoning, profiling data, roofline analysis, and a knowledge base to automatically diagnose bottlenecks and propose explainable sharding configurations. The Coordinator, Analyzer, and Proposal agents, together with a persistent Sharding Memory, implement a Retrieval-Augmented Generation workflow to produce three candidate sharding plans per experiment. In three TPU-based experiments, ASAP either matched or generalized beyond human-optimized baselines, achieving up to 28% step-time reduction and up to 2.58× throughput improvements, highlighting its potential to accelerate performance engineering. The work demonstrates a scalable, explainable approach to AI-assisted performance tuning for large-scale LLM training, with planned enhancements to broaden optimization scope and integrate closed-loop learning.
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.