OptiCo: Adaptive Distributed Training Optimization via Collaborative Agent Reasoning

Idea

OptiCo 把 distributed training configuration tuning 组织成四个专职 agent 的闭环：定位瓶颈、提出参数修改、真实启动验证、分析失败，并把所有结果保存在 Global Message Pool。它用 LLM 的结构化推理结合运行反馈，搜索更高 MFU 的配置，而不是预先穷举大量 profiling 组合。

Research Problem

TP、PP、DP、EP、micro batch size、recompute 和 overlap 等参数相互影响，性能还取决于模型与硬件。传统自动并行方案依赖预 profiling、固定规则或大规模组合搜索，配置迁移困难；纯语言建议又无法确认 OOM、head divisibility 等约束及真实吞吐，因此需要把专家式迭代与执行测量连接。

Insight

性能改善和失败日志都是可复用的搜索知识。把运行可行性与失败解释分离，并让所有角色读取同一个配置—性能/错误记录，可以减少重复无效尝试；按瓶颈排序只修改有限参数，也让每轮探索更容易解释。其自适应来自外部反馈和历史上下文，并非训练一个新的 LLM policy。

Method

1. Global Message Pool 初始化为空时先运行随机配置，记录训练参数和 performance。Bottleneck Analysis Agent 读取历史配置/MFU 等，用 CoT 分析最可能限制效率的因素，输出按预期影响排序的参数列表。
2. Strategy Generation Agent 根据排序和明确修改规则调整有限的 MBS、TP、PP、recompute 等字段，保持完整 configuration schema，避免重试历史无效或重复组合。
3. Strategy Validation Agent 在目标训练环境实际启动候选并采集指标；成功结果写回 pool，失败则生成包含修改参数和失败阶段的结构化 log，交给 Failure Analysis Agent 归类 OOM、无效并行设置或框架约束，并记录诊断 confidence。
4. 四个角色通过 pool 迭代到停止条件，返回性能最好的配置；以 inception prompting 固定职责、上下文、输出格式与行为规则。作者在 25 多种模型/GPU/规模配置中，与专家策略比较有限迭代内的 MFU 和吞吐改善。