ASAP: An Agentic Solution to Auto-Optimize Performance of Large-Scale LLM Training¶

会议: NeurIPS 2025
arXiv: 2511.03844
代码: 未公开（Google 内部 ADK）
领域: 推荐系统
关键词: 多Agent系统, 分布式训练优化, sharding配置, 瓶颈分析, Roofline模型

一句话总结¶

ASAP 是一个多 Agent 系统（Coordinator + Analyzer + Proposal），自动化诊断大规模 LLM 分布式训练的瓶颈类型（计算/内存/通信）并提出 sharding 配置方案，在 3 个实验场景中均匹配人类专家方案，实现最高 2.58× 吞吐量提升。

研究背景与动机¶

领域现状：大规模 LLM 训练需要在数百/千个 TPU/GPU 上做并行化（数据并行/模型并行/序列并行/流水线并行等），sharding 配置对性能影响巨大（差配置可导致 50%+ 吞吐量损失）。

现有痛点：(a) 手动试错需要深厚的分布式系统专家知识；(b) 黑盒优化（如贝叶斯优化）在巨大配置空间中收敛慢，且不提供可解释的理由；(c) 不同瓶颈类型（计算密集/内存密集/通信密集）需要不同的优化策略，通用方法效率低。

核心矛盾：诊断需要理解性能分析报告（Xprof profiles），决策需要了解并行化原理和硬件拓扑——这些都是人类专家才有的知识。

本文目标：用 LLM Agent 替代人类专家做性能诊断→方案提出→配置调优的全流程。

切入角度：将 LLM 推理能力与专门工具（Xprof 性能分析、Roofline 分析）和知识库（历史优化记录、JAX 并行化文档）结合为多 Agent 系统。

核心 idea：Coordinator 编排 → Analyzer 诊断瓶颈类型 → Proposal（RAG+知识库）提出多方案 = LLM Agent 做分布式训练专家。

方法详解¶

整体框架¶

三个 Agent 协作：Coordinator（接受实验 ID，编排工作流）→ Analyzer（调用 Xprof 分析 KPI + 操作级性能 + Roofline，诊断计算/内存/通信瓶颈）→ Proposal（用 RAG 检索历史配置+知识库文档，生成 3 个候选方案+理由+预期影响）。

关键设计¶

Analyzer Agent（性能诊断）:
- 读取 KPI：加速器忙碌率、训练步时间、goodput
- 操作级分析：Top-5 耗时操作的 FLOPs 利用率、HBM 带宽利用率
- Roofline 分析：判断每个操作是计算密集、HBM 密集还是通信密集
- 关联 KPI 异常与操作级瓶颈→生成诊断报告
Proposal Agent（RAG 方案生成）:
- 历史数据库：过去的配置、Xprof profile、影响摘要
- 知识库：JAX 并行化文档
- 生成 3 个候选 sharding 配置 + 理由 + 预期 trade-off
- 置信度评分（85-90%）
Sharding 配置空间:
- ici_mesh：{model, data, replica, sequence} 四个并行维度
- 硬件：TPU v5p-512（8×8×8）、TPU v6e-16（4×4）、TPU v5e-256（16×16）

损失函数 / 训练策略¶

无训练——纯 LLM 推理+工具调用。

实验关键数据¶

主实验（3 个场景）¶

场景	硬件	瓶颈类型	Agent 方案	效果	匹配专家？
1	TPU v5p-512	计算密集	model=8,data=16,seq=4	-28% 步时间	✓ 完全匹配
2	TPU v6e-16	HBM 密集	data=4,model=4	1.43× 吞吐量	✓ 完全匹配
3	TPU v5e-256	通信密集	data=8,model=2,seq=16	2.58× 吞吐量	✓ 完全匹配

消融实验¶

配置	关键发现	说明
场景1在历史库中	Agent 检索并改进	检索增强
场景2/3不在历史库	Agent 仍正确诊断+提案	原则推理&泛化
Roofline 分析贡献	区分计算/内存/通信	诊断精度关键
3 方案 vs 1 方案	最优常在第1方案	多样性提供备选
Agent 置信度	85-90%	正确反映不确定性

关键发现¶

3/3 场景完全匹配人类专家方案——说明 LLM 推理+工具可以替代分布式系统专家
场景 2/3 不在历史库中——Agent 不是记忆而是推理，能泛化到新场景
2.58× 吞吐量提升是最大收益——来自正确识别通信瓶颈并减少模型并行度

亮点与洞察¶

LLM 推理 + 专门工具 = 领域专家：ASAP 证明了 LLM 不是简单地检索历史配置，而是理解并行化原理后做推理——能处理训练库中没有的新情况。
可解释的决策：每个方案都附带详细的理由和性能引用，比黑盒优化更有工程价值。
三种瓶颈类型的全覆盖：不同场景需要截然不同的策略（计算密集→增大并行度 vs 通信密集→减少并行度），Agent 正确区分了三者。

局限与展望¶

仅支持全局 sharding（不支持逐层 sharding）
无闭环反馈——提出方案但不自动验证和迭代
缺少编译器级优化（算子融合、内存布局选择）
依赖准确的 Xprof 数据，噪声数据的鲁棒性未测试
计算与通信的整体 trade-off 优化不够

评分¶

新颖性: ⭐⭐⭐⭐ 多 Agent + 性能分析工具的组合是合理且有效的
实验充分度: ⭐⭐⭐ 仅 3 个场景，虽然都匹配专家但统计量少
写作质量: ⭐⭐⭐⭐ 诊断过程描述详细，推理链路清晰
价值: ⭐⭐⭐⭐ 对大规模 LLM 训练的自动化调优有实际工程价值\n
对MLOps的启示: ASAP的多Agent诊断-提案模式可推广到超参搜索、数据管道优化等
与MLSys社区的交叉: 将LLM Agent引入系统优化是有前景的新方向
商业价值: 大规模训练集群的配置优化直接关系到数百万美元级的计算成本
跨硬件泛化: 3种TPU拓扑上全部成功，说明推理能力不限于特定硬件
对开源社区: 诊断方法论（Roofline+操作级分析）可迁移到PyTorch/DeepSpeed
与auto-tuning的对比: 传统auto-tuning不理解为什么，ASAP给出可解释的决策链路
可扩展架构: 三Agent模式可增加Evaluator Agent实现闭环迭代优化
Agent置信度的价值: 85-90%的置信度评估帮助用户决定是否需要额外验证
知识库的可维护性: RAG方案使知识库更新只需添加新记录，无需重训Agent
对MLOps的启示: ASAP的多Agent诊断-提案模式可推广到超参搜索、数据管道优化等
与MLSys社区的交叉: 将LLM Agent引入系统优化是有前景的新方向
商业价值: 大规模训练集群的配置优化直接关系到数百万美元级的计算成本
跨硬件泛化: 3种TPU拓扑上全部成功，说明推理能力不限于特定硬件
对开源社区: 诊断方法论（Roofline+操作级分析）可迁移到PyTorch/DeepSpeed
与auto-tuning的对比: 传统auto-tuning不理解为什么，ASAP给出可解释的决策链路
可扩展架构: 三Agent模式可增加Evaluator Agent实现闭环迭代优化
Agent置信度的价值: 85-90%的置信度评估帮助用户决定是否需要额外验证
知识库的可维护性: RAG方案使知识库更新只需添加新记录，无需重训Agent
对MLOps的启示: ASAP的多Agent诊断-提案模式可推广到超参搜索、数据管道优化等
与MLSys社区的交叉: 将LLM Agent引入系统优化是有前景的新方向
商业价值: 大规模训练集群的配置优化直接关系到数百万美元级的计算成本
跨硬件泛化: 3种TPU拓扑上全部成功，说明推理能力不限于特定硬件
对开源社区: 诊断方法论（Roofline+操作级分析）可迁移到PyTorch/DeepSpeed
与auto-tuning的对比: 传统auto-tuning不理解为什么，ASAP给出可解释的决策链路
可扩展架构: 三Agent模式可增加Evaluator Agent实现闭环迭代优化
Agent置信度的价值: 85-90%的置信度评估帮助用户决定是否需要额外验证
知识库的可维护性: RAG方案使知识库更新只需添加新记录，无需重训Agent