XLA 一线实战:大规模 TPU 工作负载的正确性
XLA from the Trenches: Correctness on TPU Workloads at Scale
精选理由
介绍一套在大规模 TPU 工作负载中发现并定位编译器缺陷的方法,帮助保障训练与推理的正确性和稳定性。
AI 摘要
加速现代基础模型的需求,推动领域专用编译器采用日益复杂的多层优化流水线。然而,对优化的高度重视也使峰值执行性能与正确性之间的权衡变得微妙。隐蔽的代码生成缺陷可能引发数值精度下降以及内存或缓冲区损坏,导致耗资数百万美元的训练运行失效,并降低线上服务质量,却不报出明显的编译器错误。本文提出了一套经过实战检验的方法,用于在超大规模环境中验证 Google 的 XLA:TPU 编译器并对其问题进行分诊。我们利用模糊测试技术,将合成程序生成与真实生产环境执行轨迹相结合,同时协调 AI 智能体集群自主生成压力测试、复现故障并拟定修复方案。为定位完整工作负载中的根本原因,我们还配套使用了分层差分流水线,通过二分法将大型计算图中的问题逐步缩小到融合内核中出现故障的单个操作。该框架已部署于生产环境,推动解决了超过 142 个编译器缺陷,另有 103 项发现正在审查中。我们从中提炼了切实可行的经验教训。