COMET框架：突破AI加速器性能瓶颈，显式建模集体通信与复合操作数据流

关键词：复合操作数据流建模、集体通信操作、内存层级优化、机器学习加速器、性能建模与优化

在人工智能技术日新月异的今天，大语言模型、状态空间模型等复杂神经网络已成为推动技术发展的核心引擎。然而，这些模型所依赖的复合操作——即由多个基础操作（如矩阵乘法、归一化、逐元素变换）组合而成的结构化模块——正在对现有的硬件加速器数据流设计与性能优化提出严峻挑战。

传统的数据流优化框架往往仅针对单一操作进行建模，或在分布式计算场景中忽视集体通信的成本，导致其难以应对现代大规模、分布式训练的复合操作需求。

更关键的是，随着模型规模持续扩大，计算往往需要分布在多个计算核心甚至多个加速器之间，频繁的集体通信（如 All-Reduce、All-Gather）成为不可忽视的性能瓶颈。

为此，来自普渡大学、d-Matrix 与 Meta 的研究团队提出了一种全新的框架——COMET（Compound Operation Modeling with Explicit Collectives），旨在系统性建模与优化面向机器学习加速器的复合操作数据流。

该框架不仅首次显式建模集体通信成本，还提出了一种细粒度的性能分析模型，能够准确捕捉复合操作内部的数据依赖与通信开销。

图 3 | COMET 工作流程总览。COMET 的工作流程以复合操作描述、硬件架构规格和映射方案为输入。经映射实例生成与验证后，构建包含数据流转和集合操作的中间表示，再通过代价模型精准估算时延与能耗，为分布式深度学习推理的数据流优化提供决策依据。

本文将深入解读 COMET 的核心创新、工作方法、实验评估及其对 AI 硬件设计的影响。

零、关键问题

问题一：COMET 对集体通信的显式建模是否真正解决了分布式复合操作中的动态通信瓶颈？

文中强调 COMET 能够显式建模集体通信（如 All-Reduce、All-Gather），并在 GEMM-Softmax 等复合操作中展示了性能提升。然而，现代大规模模型（如 LLM、SSM）在训练和推理中常涉及动态、非规则通信模式，例如依赖输入序列长度的条件通信、稀疏注意力机制中的非全连接通信等。
COMET 的集体通信节点（CO 节点）是否支持动态通信模式的建模？ 例如，能否处理条件集体操作（如仅在某些条件下触发 Reduce）或非对称通信（如部分节点参与）？如果仅支持静态、规则的集体操作，这是否限制了其在真实动态负载中的适用性？

COMET 主要针对静态、规则的集体通信模式进行显式建模，尚未充分支持动态、非规则的通信模式。其集体通信节点（CO 节点）的属性（如操作类型、张量、规约操作、源/目标）在映射阶段是固定的，依赖预设的循环展开和分块策略。文中案例（如 GEMM-Softmax、GEMM-LayerNorm）使用的 All-Reduce 和 Gather 通信模式均为规则、数据并行的集体操作。

作者提到 CO 节点在映射实例生成时即被确定，依赖“用户提供的映射”和“硬件架构描述”。
distSM 和 distLN 映射使用的 All-Reduce 是在循环分块后触发的固定模式通信。
但是作者未提及条件通信、稀疏通信或动态数据依赖的集体操作支持。

综合来说，COMET 在静态数据流和规则通信模式中表现突出，但对于动态通信瓶颈（如条件集体、稀疏注意力通信）的支持尚未在本文中体现，这可能是其在实际动态负载中的限制。

问题二：COMET 的成本模型是否高估了“数据暂存低效性”对整体延迟的影响？

文中指出 COMET 相比 TileFlow 和 Timeloop 更能捕捉“数据暂存低效性”（如 ramp-up/ramp-down 阶段），并因此给出更高的延迟估计。然而，这类低效性在实际硬件中可能被部分隐藏或优化（如通过预取、投机执行、硬件流水线等）。
COMET 是否过度强调了这些“理论上的低效性”，而在实际硬件平台上（尤其是支持复杂流水线与预取机制的现代加速器如 TPU、NVIDIA H100）中，其预测是否仍具代表性？如果是，这是否会导致 COMET 在搜索最优数据流时偏向于过于保守的映射策略，从而错失更激进的优化机会？

COMET 确实显式建模了数据暂存低效性（如 ramp-up/ramp-down 阶段），并在模型中引入了“强制停顿（CS）”和“可选停顿（OS）”，这可能导致其在某些情况下比实际硬件测量或更乐观的模型（如 Timeloop）给出更高的延迟估计。然而，这种建模是为了更真实地反映在没有完全隐藏内存延迟的硬件上的行为。

COMET 的延迟模型包含强制停顿（CS）和可选停顿（OS），以准确捕捉数据在“父内存到子内存”传输过程中，初始填充和最终排空阶段引入的延迟。
与 Timeloop 相比，COMET 的延迟估计通常更高，这正是因为 Timeloop 假设了“完美流水线”而忽略了这些暂存阶段。
作者采用的硬件模型假设了双缓冲、双端口内存，这在许多边缘或云加速器中是合理的，但可能不适用于具有更复杂预取或流水线机制的高端 GPU 或 TPU。

综合来看，COMET 的成本模型在某些高度优化的硬件平台上可能显得“保守”，但它提供了一个更全面的性能分析视角，尤其适用于内存带宽受限或流水线不完美的场景。其高估的程度取决于目标硬件内存子系统的实际行为。

一、复合操作与集体通信：为何成为性能瓶颈？

1.1 什么是复合操作？

在现代神经网络中，许多层被设计为复合操作，即将多个基础操作组合为一个模块，以提升模型的模块化、效率与表达能力。典型的例子包括：
* 自注意力机制：包含矩阵乘法（GEMM）、Softmax 归一化、缩放操作；
* 归一化层（如 LayerNorm）：包含统计量计算、仿射变换等非 GEMM 操作。

这些复合操作虽提升了模型表达能力，但也引入了数据局部性差、中间张量存储开销大、通信同步复杂等问题。

1.2 集体通信在分布式计算中的关键角色

当复合操作需要在多核或多加速器间分布式执行时，不同计算单元之间需通过集体通信来同步中间结果。常见的集体操作包括：
* All-Reduce：所有节点参与规约，结果广播给所有节点；
* All-Gather：所有节点收集其他节点的数据；
* Reduce-Scatter：规约后分散到不同节点；
* Broadcast：将数据从一节点广播到所有节点。

这些操作在分布式训练与推理中频繁发生，其延迟与能耗直接影响整体系统效率。

图 2 | (a) FlashAttention 中的操作流程示意图；(b) 深度神经网络加速器架构；(c) 含隐式集合操作的矩阵乘法循环嵌套表示。图 2 (a) 拆解了 FlashAttention 的核心步骤，包含矩阵乘法与多个非矩阵乘法操作。图 2 (b) 的加速器架构由分布式计算簇与多级存储构成，支撑数据高效流转。图 2 (c) 的循环嵌套中，空间循环隐含了计算核间的数据交互，为显式集合操作建模提供了基础。

1.3 现有框架的不足

现有的数据流优化框架（如 Timeloop、TileFlow）大多侧重于单一操作或简单融合，缺乏对集体通信的显式建模，也 未能充分考虑复合操作内部的数据依赖与内存访问模式。这导致在分布式、多核场景下的性能预测失真，优化空间受限。

二、COMET 的核心创新：显式集体表示与细粒度成本模型

COMET 的核心贡献可归纳为两点：
1. 提出一种显式集体表示方法，将集体通信作为数据流的一部分进行建模；
2. 设计一种细粒度的成本模型，涵盖计算、内存访问、通信及操作间依赖。

图 1 | (a) 数据流设计方案；(b) 集合操作类型；(c) 含显式集合操作表示的映射变体；(d) 可行的调度策略。该图展示了 COMET 框架的核心设计空间维度。数据流设计需考量分块、循环顺序与空间展开方式；集合操作支撑分布式场景下的数据同步；映射变体明确了操作融合层级与集合操作的关联；调度策略则决定多计算单元的任务执行时序，以此实现复合操作的性能优化。

2.1 显式集体表示

COMET 将复合操作的数据流建模为一个四维设计空间（如图 1 所示），涵盖：
* 循环分块、排序与空间展开策略；
* 集体操作类型与执行位置；
* 操作间执行顺序与融合层级；
* 跨计算单元的调度策略。

在 COMET 的中间表示中，每个张量在每一级内存层次上都拥有独立的循环嵌套，从而支持对不同张量数据移动与复用的精细控制。

集体操作节点被显式插入到数据流图中，并附有以下属性：
* ColOpType：集体操作类型；
* Tensor：操作的张量；
* ReduceOp：规约操作（如 max、add）；
* Src/Dest：源与目标内存层级。

这种表示方法使得 COMET 能够灵活探索不同的集体策略，例如在全局缓冲区（GB）级别执行 All-Reduce，或在输出缓冲区（OB）级别执行更轻量的 Gather 操作。

2.2 复合操作示例：GEMM-Softmax 的显式集体表示

以典型的复合操作GEMM 后接 Softmax 为例（如图 4 所示），其 Softmax 被分解为多个子操作，并在多集群间分布式执行。

图 4 | (a) 典型的复合操作，包含矩阵乘法及分解为多个基础操作的 Softmax（由单指令多数据单元执行）；(b) 复合操作在图 2 (b) 所示架构上的分块示例；(c) 矩阵乘法 – Softmax 复合操作的显式集合操作表示。图 4 (a) 将 Softmax 拆解为求行最大值、指数运算等基础操作，体现了复合操作的内部结构。图 4 (b) 展示了张量在多级存储间的分块传输方式。图 4 (c) 通过时间循环与空间循环的组合，结合显式集合操作节点，清晰刻画了分布式执行中数据同步的关键环节。

上图展示了 GEMM-Softmax 复合操作的树状表示，其中显式插入了 All-Reduce 集体节点。

在分布式执行中，Softmax 需要获取整行张量的最大值与求和值，这要求跨集群执行两次 All-Reduce。COMET 通过显式插入集体节点（如CO₁⁰和CO₁¹）来建模这一通信过程，从而能够在后续优化中调整其执行位置或替换为其他集体类型（如 Gather），以权衡通信开销与内存占用。

2.3 细粒度成本模型

COMET 的成本模型涵盖延迟与能耗两方面，并特别强化了对以下因素的建模：

（1）内存传输延迟

COMET 采用双缓冲假设，支持内存传输与计算重叠。其内存传输延迟模型如下：

L_mem = DV / BW + CS + OS

其中DV为每迭代传输数据量，BW为带宽。总延迟还包括强制停顿（Compulsory Stall, CS）与可选停顿（Optional Stall, OS），以捕捉数据填充与排空阶段的空闲时间。

（2）集体操作延迟

集体操作的延迟 L_col，由内存访问延迟 L_mem 与片上网络（NoC）延迟 L_noc 共同构成：

L_col = L_mem + L_noc
L_noc = hops * (t_route + t_queue)

其中hops为通信跳数，W为 NoC 通道宽度，t_route与t_queue分别为路由与排队延迟。

执行一次集合操作的总耗时，等于数据在内存中读写的时间，加上数据在片上网络传输的时间。这个公式是 COMET 框架中性能建模的关键部分，用来量化分布式计算里集合通信的开销。

（3）操作间依赖与调度

COMET 支持多种调度策略（顺序、流水线、并行），并通过冲突检测模型捕捉资源共享竞争带来的额外延迟。例如，在顺序调度中，节点总延迟为子节点延迟之和；在流水线调度中，则为最慢子节点延迟加上冲突停顿时间。

三、COMET 的工作流程与实现

COMET 的整体工作流程如图 3 所示，用户需提供三方面输入：
1. 工作负载描述：复合操作的结构与张量形状；
2. 硬件架构描述：内存层次、计算单元、NoC 拓扑；
3. 映射约束：循环分块、展开、融合等限制。

系统首先生成满足约束的映射实例，验证其内存容量是否合规，随后构建包含集体操作的中间表示，最终通过成本模型评估延迟与能耗。

四、相关工作比较

COMET 并非首个关注复合操作优化的框架，但其在集体通信建模与细粒度依赖分析方面实现了显著突破。

| 框架 | 目标 | 是否支持集体建模 | 是否支持非 GEMM 单元 | 主要局限 |
| :— | :— | :— | :— | :— |
| Timeloop | 单操作数据流优化 | 否 | 有限 | 无法处理复合操作 |
| TileFlow | 复合操作融合优化 | 隐式 | 否 | 忽略数据暂存低效 |
| LoopTree | 融合中间张量权衡 | 否 | 否 | 缺乏通信成本建模 |
| COMET | 复合操作+集体通信 | 显式 | 是 | 全面覆盖计算、内存、通信 |

表：COMET 与现有框架对比

Timeloop：专注于单操作映射优化，假设理想流水线，忽略启动与排空阶段延迟；
TileFlow：通过树状分析建模融合数据流，但未显式建模集体操作，且假设单一计算单元；
LoopTree：探索中间张量保留与重计算的权衡，适用于内存受限场景，但未扩展至分布式执行。

COMET 在继承这些框架优点的同时，通过显式集体表示与增强的成本模型，实现了对现代分布式 AI 工作负载更准确的性能预测与优化。

五、实验评估：COMET 如何提升性能？

研究团队在边缘与云两种加速器配置上评估了 COMET，覆盖三种典型复合操作：GEMM-Softmax、GEMM-LayerNorm与自注意力。

表 1 | 边缘架构的矩阵乘法形状。该表列出适配边缘端加速器的 6 种矩阵乘法维度，涵盖从 1×1024×64 到 512×1024×64 的不同规模。这些形状源于 Llama-3 和 GPT-2 模型的推理负载，覆盖预填充与解码阶段，用于验证 COMET 框架在边缘端资源受限场景下的性能优化能力。

表 2 | 云端架构的矩阵乘法形状。此表包含 6 种面向云端大规模加速器的矩阵乘法维度，尺寸显著大于边缘端，如 1×16384×128、512×4096×128 等。这些形状对应大语言模型的大张量计算场景，用于测试 COMET 在分布式集群下，对集体通信开销和数据流转的建模精度。

表 3 | 边缘架构的注意力操作。该表定义了边缘端加速器的 6 种注意力操作张量维度，查询、键、值的尺寸适配边缘硬件的内存带宽。实验设定批量大小与注意力头数均为 1，模拟边缘端轻量级推理场景，以此评估 COMET 对注意力机制中复合操作的优化效果。

表 4 | 云端架构的注意力操作。此表呈现 6 种云端加速器的注意力操作张量配置，张量维度远大于边缘端，如查询维度达 2048×256。实验同样设置批量大小和注意力头数为 1，聚焦云端大尺寸张量的分布式计算，验证 COMET 在显式集体通信建模下的性能优势。

5.1 硬件配置

表 5 | 边缘与云加速器的硬件配置参数。该表对比边缘与云端加速器的核心硬件参数，包括内存容量、集群规模、缓存带宽等。边缘端采用 2×2 集群与 1GB DRAM，云端为 4×4 集群与 4GB DRAM，信道带宽差距达 8 倍。这些配置为 COMET 的成本模型提供硬件基准，用于精准测算不同场景下的时延与能耗。

5.2 成本模型验证

图 6 | COMET 与主流框架的能耗及时延对比。(a)(b) 分别为与 Timeloop 的能耗、时延对比；(c)(d) 分别为与 TileFlow 的能耗、时延对比。对比结果显示，COMET 与 Timeloop、TileFlow 的能耗和时延高度相关。COMET 的时延估算略高，原因是其纳入了数据分段传输的启动与收尾开销，以及操作间的依赖停顿，能更精准反映实际硬件执行状态。

与 Timeloop 和 TileFlow 的对比如图 6 所示：

与 Timeloop 对比：COMET 在单操作能量预测上高度一致，但延迟预测略高，因其包含了启动与排空阶段的停顿；
与 TileFlow 对比：COMET 在复合操作中预测能量略低，因其更准确地捕捉了中间张量复用；延迟预测更高，因其建模了操作间依赖导致的强制停顿。

5.3 集体策略的影响

COMET 支持两种集体策略：

分布式映射：在多个集群间执行集体操作；
标准映射：将非 GEMM 操作限制在单个集群内，避免跨集群集体通信。

图 7 | 不同数据流的时延与能耗变化：(a) 矩阵乘法 – Softmax 的时延；(b) 矩阵乘法 – Softmax 的能耗；(c) 矩阵乘法 – 层归一化的时延；(d) 矩阵乘法 – 层归一化的能耗。该图直观呈现了分布式与非分布式映射对性能的影响。分布式映射虽能提升计算并行度，但频繁的集合操作会增加开销；非分布式映射可避免通信成本，却可能因单节点计算负载过高导致时延上升，需结合硬件配置选择最优方案。

图 7 显示，分布式映射在大型 GEMM 中因频繁集体操作导致延迟上升；而标准映射虽减少通信，但可能因单核执行成为 SIMD 瓶颈。

图 8 | (a) 分布式 Softmax 与标准 Softmax 映射下，矩阵乘法 – Softmax 复合操作的归一化时延拆解；(b) 分布式层归一化与标准层归一化映射下，矩阵乘法 – 层归一化复合操作的归一化时延拆解。横轴编号对应矩阵乘法编号。时延拆解结果表明，大尺寸矩阵乘法在标准映射下，时延由单指令多数据单元的计算主导；分布式映射下则受集合操作开销制约。小尺寸矩阵乘法的时延瓶颈为强制性停顿，源于数据复用率低，凸显了数据流优化需按需调整的特点。

图 9 | (a) 分布式 Softmax 与标准 Softmax 映射下，矩阵乘法 – Softmax 复合操作的归一化能耗拆解；(b) 分布式层归一化与标准层归一化映射下，矩阵乘法 – 层归一化复合操作的归一化能耗拆解。横轴编号对应矩阵乘法编号。能耗拆解显示，两种映射的硬件访问次数差异不大，能耗主要由片外存储器访问主导。大尺寸矩阵乘法的集合操作能耗占比明显上升，这是因为跨计算簇的数据传输需消耗大量片上网络资源，进一步验证了集合操作优化的必要性。

5.4 融合映射的收益

图 10 | 不同融合映射相对非融合基线的归一化时延（数值越低性能越佳）：(a) 矩阵乘法 – Softmax 复合操作；(b) 矩阵乘法 – 层归一化复合操作。所有融合映射的时延均低于非融合基线，原因是融合减少了中间张量的片外存取次数。全操作融合的映射方案时延最优，边缘平台的标准融合映射因单节点计算压力大，时延高于分布式融合，而云端平台则需权衡通信与计算开销。

图 11 | 不同融合映射相对非融合基线的归一化能耗（数值越低性能越佳）：(a) 矩阵乘法 – Softmax 复合操作；(b) 矩阵乘法 – 层归一化复合操作。融合映射通过降低片外存储器访问频次，有效降低了能耗。分布式与非分布式融合映射的能耗差异较小，因为两者的存储器访问总量相近。该结果表明，在能耗优化目标下，可优先选择实现简单的映射方案，无需过度追求分布式部署。

图 10 与图 11 对比了不同融合策略与未融合基线的性能：

全融合在大多数情况下实现最低延迟；
部分融合因仍需跨操作数据传输，性能次之；
未融合基线因频繁读写 DRAM，延迟与能耗最高。

COMET 优化后的数据流在几何平均上实现了显著的性能提升：

GEMM-Softmax：1.42 倍 加速；
GEMM-LayerNorm：3.46 倍 加速；
自注意力：1.82 倍 加速。

5.5 自注意力案例研究

图 12 | 不同注意力变体的归一化 (a) 时延与 (b) 能耗对比。对比了非融合注意力、部分融合注意力及 Flash 注意力。数值越低表示性能越优。Flash 注意力凭借全操作融合策略，实现了 1.82 倍的时延优化与 1.54 倍的能耗优化。在边缘平台上，低复用性注意力操作的融合收益有限，性能瓶颈主要在于片外存储器访问；而在云端平台上，大尺寸张量的融合能显著降低数据传输开销，收益更为明显。

图 12 显示，完全融合的 FlashAttention（FA）在延迟与能耗上均优于未融合（UA）与部分融合（PFA）版本，尤其在云平台上收益显著。

图 13 | 不同注意力变体的归一化时延拆解。注意力操作 1-6 部署于边缘平台，操作 7-12 部署于云端平台。时延拆解显示，边缘平台的时延主要由片外存储器访问主导，融合对性能提升有限。云端平台的融合操作虽增加了单指令多数据单元的计算时延，但大幅降低了集合操作开销，整体时延更优，体现了软硬件协同优化的重要性。

图 14 | 不同注意力变体的归一化能耗拆解。注意力操作 1-6 部署于边缘平台，操作 7-12 部署于云端平台。能耗拆解表明，Flash 注意力的片外存储器访问能耗显著低于其他变体，这是其能耗优势的核心来源。融合操作引入的额外计算能耗，远低于数据传输能耗的节省量。该结果为深度学习加速器的注意力操作优化提供了明确方向。

六、总结与展望

COMET 通过显式集体表示与细粒度成本建模，为复合操作在分布式机器学习加速器上的数据流优化提供了系统性框架。它不仅填补了现有工具在集体通信建模方面的空白，还为硬件-软件协同设计提供了重要指导。

未来工作可沿以下方向拓展：

支持更多集体算法：如拓扑感知集体、分层集体等；
扩展至训练场景：涵盖反向传播与梯度同步；
集成自动搜索算法：结合强化学习或进化算法进行映射空间探索；
支持新型硬件：如存内计算、光互连等新兴架构。

COMET 的出现标志着 AI 硬件优化进入集体感知时代，为下一代大模型的高效部署奠定了重要基础。

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：http://www.itsolotime.com/archives/18664