LHS 费用模型

要点；

本页介绍了 Latency Hiding Scheduler 使用的费用模型的内部结构。如果您有兴趣调整模型，请直接前往“调整”部分。

延迟隐藏调度器 (LHS) 是一种编译器传递，用于以最大限度减少实际运行时间的方式调度 HLO DAG。

其决策由统一费用模型指导，该模型使用性能表和分析模型的混合体。具体而言，XLA 会嵌入 GEMM 和快速互连集合的性能表，并在其他情况下使用分析性网络和融合成本模型。本文档的其余部分将简要介绍这些功能的内部运作方式。

效果表 - ICI 集体

性能表由两个主要组件组成：收集器和插值器。

收集器

收集器是一种 C++ 工具，负责为集体操作生成性能表。它用于衡量各个 HLO 操作的性能（例如，all-gather、all-reduce）在静态定义的形参空间中。

工作方式

该工具会针对给定的集群，对一系列集体操作、传输大小和传输方案执行扫描。它使用现有的多主机 HLO 运行程序基础架构和 ExecutionProfile 数据来运行生成的 HLO 并收集性能指标。

数据收集参数

系统会针对以下参数的交叉组合收集延迟时间表：

  • 集体类型：
    • all-reduce
    • all-gather
    • reduce-scatter
  • 转移大小：
    • 从 1024B 到 2GiB 的对数刻度（例如，1024B、2048B、4096B 等）
  • 转移方案：
    • rail-aligned
    • non-rail-aligned

此扫描针对具有 2、4 和 8 个设备的节点内集群运行。

输出

收集运行的结果是一个 .pbtxt 格式的延迟时间表（每个平台大约 116 KB）。

Interpolator

插值器是一种编译器组件，它会使用生成的性能表在编译期间提供运行时估计值。

内部数据结构

在初始化时，插值器会将性能表处理为映射。 此映射使用 (collective_type, transfer_scheme) 的元组作为其键。

与每个键关联的值是一个二维欧几里得平面。此平面根据两个轴对网络吞吐量（由收集器测量）进行索引：

  1. 转移大小。
  2. 涉及的设备数量。

查找和插值

当编译器遇到集体操作时，插值器会执行以下步骤：

  1. 它使用操作的 (collective_type, transfer_scheme) 作为映射键来确定正确的 2D 吞吐量平面。
  2. 然后，它会使用操作的 (transfer_size, num_devices) 作为查询点，在该二维平面内使用加权平均检索（基于欧几里得距离）。
  3. 此查找的结果是单个唯一的网络吞吐量值。

理由：吞吐量和外推

该系统旨在存储网络吞吐量，而不是原始延迟时间。这种设计选择可显著简化对表中未明确显示的传输大小的性能外推。

如果延迟时间表捕获到总大小为 S 时的网络带宽饱和度，则该点的吞吐量 T 被视为最大吞吐量。对于任何规模为 S' > S 的新集体，运行时可以估计为：

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

这样，模型就可以估计任意大小的集合的性能，即使这些集合的大小超过了收集器测量的 2GiB 上限。

  • 低估最大吞吐量。
  • 因此，对于大型传输，系统会高估运行时。

一般来说，XLA:GPU 团队会维护性能表，但如果用户决定提供自己的性能表，则生成这些表的用户有责任确保这些表具有代表性，并且包含目标硬件在带宽饱和区域的测量结果。

效果表 - GEMM

与集合系统类似，GEMM 延迟时间表由两个组件支持：收集器和插值器。

收集器

收集器是一种 C++ 工具，用于计算通用矩阵乘法 (GEMM) 的性能表。它用于衡量 HLO dot op 级别的矩阵乘法性能。

工作方式

该工具会对 GEMM 维度（批次、两个非收缩维度和一个收缩维度）和数据类型的静态空间进行扫描。

  • 默认数据类型： LHS = bf16,f32、RHS = bf16,f32、OUT = bf16,f32。
  • 基础设施：重复使用 HLO 操作分析器。

集合参数

系统会针对以下维度的交叉组合收集延迟时间表：

  • 批次： {1, 2, 4}
  • m（非收缩）： {256, 512, ..., 4096}
  • n（非合约）： {256, 512, ..., 4096}
  • k（合约）： {256, 512, ..., 4096}

输出和存储

完整扫描会生成一个 .pbtxt 延迟时间表，可供插值器使用。

Interpolator

插值器是使用生成的表格来估计 GEMM 性能的编译器组件。

理由：FLOPS 饱和度

收集的延迟时间表使插值器能够为每个条目重建 FLOPS：

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

一个关键的观察结果是，FLOPS 在某个点达到饱和；也就是说，硬件在超过某个矩阵形状后达到峰值 FLOPS。这种饱和度允许使用与集合相同的外推方法。

查找和插值

插值器会根据表格数据构建一个 4D 欧氏空间。为了提供效果估算值，它会在这个 4D 空间内执行加权平均插值。如果某种数据类型没有对应的表格，则每个维度都会根据字节数进行归一化处理。

分析成本模型 - DCN

S 形集体成本模型

S 形曲线模型是一种完全分析型的网络屋顶线模型。

概览

该模型旨在根据一组固定的网络属性来估计集体操作的性能。

模型输入

该模型需要两类输入：

  1. 固定网络属性（用户定义）：

    • 集体启动开销
    • 网卡速度
    • RTT（往返时间）

    默认情况下，XLA 会自动检测平台并使用最常见架构的值。这些属性可由用户配置。如需了解详情，请参阅调优部分。

  2. 每个集合的输入：

    • 集体类型（例如，AllGather、ReduceScatter）
    • 转移大小
    • 参与通信的节点数量

集成

S 形曲线模型已集成到 XLA:GPU 中，并已在 Hopper 和 Blackwell 上使用。

分析性成本模型 - Fusions

对于其他内核，我们依靠 GPU 性能成本模型来估计合适的运行时。如需了解详情，请点击此处。

调优

可以通过发出正确的 XLA 标志来调整 S 形曲线模型。在大多数情况下，默认配置应该足够好，但在其他情况下，模型控制会公开。

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"