ML生产力Goodput：度量与优化大规模机器学习集群效率的三层框架

1. 项目概述：为什么我们需要重新审视ML集群的效率度量？

在管理一个由成千上万块TPU、GPU等加速器芯片组成的大规模机器学习集群时，我们每天都会面临一个灵魂拷问：我们的硬件资源，真的被高效利用了吗？你可能会盯着集群监控面板上那些常年保持在90%以上的“设备占用率”或“容量利用率”图表，感觉一切良好。但当你深入去看，却发现一个需要数千芯片、以特定拓扑（比如4x4x16的3D环面）启动的大模型训练任务，在队列里等待了数小时甚至数天；或者发现一块昂贵的TPU芯片虽然显示“忙碌”，但其核心的矩阵乘法单元（MXU）利用率可能还不到30%。这种“虚假繁荣”在传统指标下被完全掩盖了。

这就是我们引入“ML生产力Goodput”（ML Productivity Goodput， 简称MPG）这个度量框架的根本原因。它不是一个单一的百分比数字，而是一把“手术刀”，旨在精准解剖大规模ML集群的效率瓶颈。传统的“利用率”指标（Utilization）只关心“设备是否在忙”，而MPG关注的是“设备是否在高效地做有用功”。这其中的差别，就好比衡量一个工厂，是只看机器是否在转（利用率），还是看它生产合格产品的速度和数量（生产力Goodput）。

MPG将整体效率分解为三个相互关联又彼此独立的维度：调度Goodput（Scheduling Goodput， SG）、运行时Goodput（Runtime Goodput， RG）和程序Goodput（Program Goodput， PG）。这三者的乘积，最终构成了我们评估集群真实生产力的核心指标。这个框架的价值在于，它迫使我们从系统全栈的视角去思考问题：一个任务的低效，究竟是调度器找不到合适的硬件拓扑（SG低），是数据加载或检查点恢复太慢导致设备空等（RG低），还是编译器生成的代码本身就没能榨干硬件的算力（PG低）？只有精准定位，优化才能有的放矢。

2. MPG度量框架全栈拆解：调度、运行时与程序的效率三角

要理解MPG，我们必须先摆脱“集群是一堆同质化计算单元”的简化思维。一个现代的大规模ML集群，更像一个由不同规格、不同拓扑互联的“计算岛屿”组成的 archipelago（群岛）。任务（Job）对硬件不仅有数量要求，更有严格的拓扑形状要求。MPG的三层分解，正是对应了任务从提交到完成所经历的三个关键阶段。

2.1 调度Goodput（SG）：你的任务何时能“上岛”？

调度Goodput衡量的是任务获取所需计算资源的能力。它的核心挑战是“拓扑匹配”和“碎片化”。假设你的集群有10000块芯片，总体空闲率很低，看起来容量充足。但如果这些空闲芯片分散在几百个不同规格、不同互联状态的小块中，而你的任务需要512块芯片以特定的2D网格拓扑（例如16x32）同时启动，那么你可能依然无法被调度。SG低下，直接表现为任务在队列中长时间等待。

SG的优化，本质上是集群的“城市规划”问题。调度器需要在两个矛盾的目标间权衡：一是尽可能提高全局芯片的占用率（传统指标），二是减少因频繁抢占（Preemption）和资源碎片化导致的大任务无法启动。从我们实际运营的经验来看，调度器会形成一种“欺软怕硬”的偏好：对于需要整个Pod或多个Pod的“超大型任务”，调度器会尽量避免抢占它们，因为其启动开销（加载模型、初始化通信）巨大，抢占带来的整体MPG损失是灾难性的。相反，对于只需要几个芯片的“小型任务”，调度器则有很高的灵活性，可以将其填充到集群的各个碎片化角落，起到“填缝剂”的作用，从而提升整体SG。而对于中等规模的任务，它们往往成为调度策略调整的“牺牲品”，在资源紧张时最先被考虑驱逐。

注意：高SG并不意味着所有任务都能立即被调度。它反映的是一种全局最优的权衡。一个健康的SG指标应该显示，超大型任务和小型任务的SG都接近理想值（如>95%），而中型任务的SG可能会有所波动，这通常是调度策略主动选择的结果，而非系统故障。

2.2 运行时Goodput（RG）：设备拿到手后，真的在全力计算吗？

当任务成功被调度，获取了梦寐以求的硬件后，运行时Goodput就开始发挥作用了。RG衡量的是，在任务的生命周期内，加速器实际用于执行有效计算的时间比例。导致RG降低的“时间杀手”主要来自主机侧（Host）的 overhead 和任务内部的中断。

最常见的影响RG的因素包括：

1. 数据输入管道（Input Pipeline）瓶颈：特别是在训练任务中，如果数据预处理、加载的速度跟不上GPU/TPU的计算速度，加速器就会陷入“饥饿”等待状态。这就像一条高速装配线，因为零件供应不上而频繁停工。
2. 检查点（Checkpoint）开销：保存和恢复模型检查点是容错和模型选择的必需操作，但同步的检查点会强制暂停所有计算进程，导致设备完全空闲。异步检查点技术可以将数据写入与计算重叠，是提升RG的关键。
3. 编译等待：对于JIT（即时编译）框架，每次任务启动或遇到新计算图时，都可能需要编译，这段时间设备是闲置的。采用AOT（提前编译）策略，在廉价的CPU上预先编译好计算图并缓存，可以显著减少运行时的编译延迟。
4. 通信同步：在分布式训练中，All-Reduce等集合通信操作如果与计算是串行的，也会产生空等。

提升RG的策略，核心思想是“异步化”和“流水线化”。例如，Google的Pathways系统采用单客户端异步数据流，允许计算和通信更好地重叠。Plumber等工具则用于诊断和优化数据管道中的瓶颈。从我们的观测数据看，训练任务通常比在线推理服务有更高的RG，因为训练的计算图相对稳定，计算需求持续；而在线服务流量波动大，更容易出现资源利用不均衡的情况。

2.3 程序Goodput（PG）：你的代码，榨干了硬件的每一分算力吗？

这是最接近硬件微架构的一层。程序Goodput衡量的是，在加速器核心执行计算指令的每一个时钟周期里，有多少比例是在执行对模型训练/推理有实际贡献的浮点运算（FLOPS）。它直接反映了编译器优化和模型实现的质量。

影响PG的因素包括：

1. 计算与通信的重叠：在模型并行训练中，通信（梯度同步）往往是瓶颈。通过将通信操作分解为更细的粒度并与计算操作交错执行，可以隐藏通信延迟，提升整体PG。例如，将一次大的All-Reduce分解为多次小的、与计算层重叠的通信。
2. 编译器优化水平：编译器决策，如算子融合（Operator Fusion）、张量布局（Tensor Layout）、循环分块（Tiling）策略等，对性能有决定性影响。一个差的布局决策可能导致内存访问效率低下，即使计算单元空闲，也在等待数据。
3. 内核（Kernel）效率：即便是最简单的矩阵乘法，针对不同尺寸和硬件，也有数十种实现方案。选择最优的内核至关重要。
4. 硬件适配度：当新一代加速器（如从TPU v4到v5）上市初期，现有的模型代码和编译器可能还未针对其新特性（如新的指令集、内存层次）进行充分优化，导致PG偏低。随着软件栈的成熟，PG会逐步提升。

提升PG需要深入的编译器和体系结构协同优化。自动调优工具如XTAT，可以系统性地搜索庞大的编译参数空间，为特定模型和硬件找到最优配置。跟踪PG随时间的变化��可以清晰量化每一项编译器优化的实际收益。

3. 实操：如何度量与分析三层Goodput

理论框架清晰后，下一步就是落地。度量MPG需要在系统栈的各个层面植入探针（Instrumentation），并建立一个统一的分析视图。这不是一个开箱即用的工具，而是一个需要定制的度量体系。

3.1 数据采集与指标定义

首先，我们需要为每一层定义可度量的指标：

• 调度Goodput (SG)：SG = 任务实际获得的、符合拓扑要求的加速器时间 / 任务从提交到结束的总时间。这里的分母是“墙上时钟时间”，分子是“有效占用时间”。调度器的排队时间、因拓扑不匹配导致的等待时间，都会降低SG。
• 运行时Goodput (RG)：RG = 加速器核心执行计算的时间 / 任务占用加速器的总时间。需要从硬件性能计数器（如TPU的MXU利用率计数器、GPU的SM Active Cycles）和运行时日志（检查点起止时间、数据加载阻塞时间）中综合计算。主机开销、通信等待、检查点时间都需要从分母中扣除。
• 程序Goodput (PG)：PG = 实际完成的、有效的FLOPS / 硬件理论峰值FLOPS。这需要结合性能剖析工具（如TPU的Profile工具、NVIDIA Nsight）来获取实际FLOPS，并与硬件的理论算力对比。

3.2 建立分析仪表盘与细分维度

采集到原始数据后，不能只看全集群的聚合平均值。必须进行多维度细分（Segmentation）分析，才能发现隐藏的问题：

1. 按任务规模细分：如前所述，分析超大型、中型、小型任务的SG差异，可以验证调度策略的有效性。
2. 按工作负载类型细分：分别查看训练（Training）、批量推理（Bulk Inference）、在线服务（Real-time Serving）的RG和PG。训练任务通常PG较高但RG可能受数据管道影响；在线服务RG波动大；批量推理的PG可能因模型结构变化（如引入MoE专家网络）而产生剧烈波动。
3. 按硬件类型/代际细分：跟踪不同型号加速器（如TPU v4, v5p, A100, H100）上的PG，可以评估软件栈对不同硬件的优化成熟度，指导资源采购和软件投资方向。
4. 按模型架构细分：对比Transformer、CNN、MoE等不同架构模型的PG，可以发现特定优化（如Attention算子优化、专家路由优化）的潜在空间。
5. 按编译器/框架细分：对比使用XLA、TensorFlow、JAX、PyTorch等不同软件栈的任务PG，可以评估不同技术栈的性能表现。

3.3 一个具体的诊断案例：RG异常波动排查

假设在仪表盘上，你发现“批量推理”工作负载分组的RG在过去一个季度出现了如图14所示的先下降后回升的“V型”波动。全集群的RG均值却变化不大。这就是细分分析的价值。

诊断过程如下：

1. 定位：波动仅出现在“批量推理”组，训练和服务组稳定。说明问题具有特异性。
2. 假设：批量推理的模型或数据特征发生了改变。经查，在波动开始时间点附近，团队开始将一批较大的推荐模型从“单芯片副本”架构迁移到“多芯片权重分片”架构，以支持更大的模型。
3. 根因分析：
   • 单芯片副本时代：每个芯片独立处理请求，模型权重完全驻留于芯片内存，数据读取快，通信开销极低，因此RG很高。
   • 权重分片时代：模型参数被切分到多个芯片上。处理一个请求时，芯片间需要进行频繁的权重聚合或通信（例如在MoE模型中，需要从多个专家中动态选择），引入了额外的通信同步开销。这导致了设备空闲等待，RG下降。
4. 验证与解决：通过剖析工具确认，RG下降的时间段内，设备空闲等待主要发生在集合通信操作上。随后，团队应用了“计算-通信重叠”的优化（如3.3节所述），将通信分解、流水线化，RG随之回升。

这个案例表明，MPG的细分分析能够将宏观的性能波动，精准定位到具体的系统变更和优化机会上。

4. 优化策略实战：从度量到行动

度量本身不是目的，基于度量的精准优化才是。MPG的三层分解，直接对应了三类不同的优化战场和责任团队。

4.1 提升调度Goodput（SG）：智能编排与碎片整理

优化SG是集群调度器团队的核心任务。目标是最大化满足拓扑要求的任务分配成功率。

• 策略一：差异化抢占策略。正如前文所述，实施非对称的抢占策略。为超大型任务设置更高的“免驱逐”优先级，或要求驱逐它们时必须达到极高的SG收益阈值。这保护了高价值、高启动成本的任务流。
• 策略二：主动碎片整理。调度器不应只是被动地响应任务请求，而应主动进行“碎片整理”。例如，可以有意地将一些小型、灵活的任务调度到未来可能被大任务需要的资源块周围，以便在时机成熟时，通过快速完成或迁移这些小任务，释放出连续的大块资源。这类似于磁盘碎片整理中的“空间预留”思想。
• 策略三：拓扑感知的放置算法。调度算法必须深度理解硬件互连拓扑（如NVLink、ICI的带宽和延迟）。将通信密集的任务子单元放置在有高速直连的芯片上，可以极大提升后续的PG。这需要调度器与资源管理元数据深度集成。

4.2 提升运行时Goodput（RG）：消除一切不必要的等待

优化RG是框架和运行时团队，以及应用开发者的共同责任。

• 框架层最佳实践：
  • 使用高效数据管道：利用TensorFlow的tf.data、PyTorch的DataLoader并进行充分调优（预取、并行化、缓存）。使用类似Plumber的工具定期剖析管道，找出瓶颈算子。
  • 启用异步检查点：确保训练框架配置了异步检查点保存。例如，使用支持异步快照的存储后端，让计算在检查点写入时继续进行。
  • 预编译与缓存：对于生产环境模型，务必使用AOT编译（如XLA AOT、TorchScript），并将编译结果缓存。避免在任务运行时产生JIT编译开销。
• 系统层支持：
  • 采用异步执行模型：研究和采纳如Pathways之类的单客户端、异步数据流执行模型。它允许更灵活的计算-通信-控制流重叠，尤其适合异构和动态负载。
  • 优化主机-设备通信：确保PCIe通道不被其他无关进程占用，优化主机内存与设备内存间的数据传输（DMA），考虑使用GPUDirect RDMA等技术绕过主机内存。

4.3 提升程序Goodput（PG）：榨干硬件的最后一滴算力

优化PG是编译器团队和性能工程师的主场，需要深入到指令和微架构层面。

• 编译器自动调优：这是提升PG最系统性的方法。部署像XTAT这样的生产级自动调优器。它可以针对海量的模型-硬件组合，自动搜索最优的编译选项组合（如融合策略、分块大小、内存布局）。我们内部的数据显示，即使相对于已经高度优化的XLA基线，XTAT仍能为大量模型带来可观的性能提升（平均百分之几到百分之几十不等）。
• 计算通信重叠优化：对于分布式训练，这是必须攻克的难关。技术思路是将大的通信操作（如All-Reduce）分解为多个小步骤，并将其与计算图中的独立部分交错执行。例如，在Transformer层��，可以在计算当前层的Self-Attention时，同时开始传输下一层所需的梯度。这需要编译器深度参与计算图的重写。
• 面向硬件的内核优化：与硬件团队紧密合作，为新一代加速器的新指令集（如TPU的矩阵乘加单元特定指令）编写或生成高度优化的内核。利用Roofline模型分析应用是计算受限还是内存受限，从而指导优化方向。
• 定期性能回归测试：建立一个涵盖关键模型和硬件配置的基准测试套件，持续跟踪其PG。任何编译器或框架的更新，都需要通过这个测试套件来防止性能回退，并量化新优化的收益。

5. 避坑指南与常见问题

在实际推行MPG度量和优化过程中，我们踩过不少坑，也积累了一些关键经验。

• 误区一：“我们的占用率很高，所以效率没问题”。这是最危险的错觉。高占用率可能意味着任务在低效运行（低PG），或者在排队等待（低SG），或者设备在空转（低RG）。MPG要求我们必须打破这种“自满”，深入三层指标去看问题。
• 误区二：只关注全集群平均值。MPG的力量在于细分。一个稳定的全集群PG均值，可能掩盖了某些硬件型号上PG急剧下降、某些模型类型上RG严重偏低的问题。必须建立多维度的细分视图。
• 挑战一：度量开销。植入详细的性能计数器采集和日志记录，本身会带来一定的开销。需要在度量粒度和开销之间取得平衡。通常采用采样方式，而非全量持续记录。对于关键的生产任务，可以开启详细剖析；对于海量小任务，则采用聚合统计。
• 挑战二：指标归因的复杂性。有时一个性能问题可能同时影响多层。例如，一个通信密集的模型，较差的PG（计算通信无法重叠）可能导致更长的单次迭代时间，这可能会增加被调度器抢占的概率（影响SG），也可能会让检查点间隔内的有效进度变少（间接影响RG）。需要结合时间序列和相关性分析来理清因果关系。
• 实操建议：建立跨团队效率小组。MPG优化涉及调度、运行时、编译器、硬件、算法等多个团队。建立一个常设的“效率工程”小组，定期review各层的MPG指标，针对发现的瓶颈发起跨团队联合攻关，是确保优化持续落地的有效组织形式。这个小组负责维护统一的度量仪表盘和基准测试。
• 关于工具选型：目前并没有一个名为“MPG Meter”的开源工具。实现MPG度量需要结合现有监控栈（如Prometheus用于资源采集）、分布式追踪系统（如OpenTelemetry用于跟踪任务生命周期）、硬件性能剖析工具（如TPU的Cloud TPU Profiler、NVIDIA Nsight）以及自定义的指标聚合逻辑。核心是建立一个统一的数据管道和分析平台。

最终，ML生产力Goodput不仅仅是一个度量标准，它更是一种系统优化的哲学。它要求我们放弃对单一、笼统的“利用率”数字的迷信，转而采用一种分层、细分、归因的严谨方法，像对待一个精密仪器一样，去调试和优化我们日益复杂的大规模机器学习基础设施。这个过程没有银弹，但有了MPG这幅“地图”，我们至少能清楚地知道，效率的宝藏埋在哪一层，以及该从哪里开始挖掘。