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

TL;DR

ASAP is a multi-agent system that automates performance optimization in large-scale LLM training, reducing training time and increasing throughput by integrating reasoning, profiling, and expert knowledge.

Contribution

It introduces a novel multi-agent framework that automates diagnosis and optimization of LLM training performance, combining LLM reasoning with profiling insights and expert knowledge.

Findings

Up to 28% reduction in training step time.

1.43 times increase in throughput with ASAP.

Further throughput increase to 2.58 times with human expert input.

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.