文献阅读《Chakra: Advancing performance benchmarking and co-design using standardized execution traces》

Sat, 15 Jun 2024 00:00:00 +0000

1. 文章简介

1.1 摘要

基准测试和协同设计对于推动 ML 模型、ML 软件和下一代硬件的优化和创新至关重要。全工作量基准（如 MLPerf）在实现不同软件和硬件堆栈之间的公平比较方面发挥着至关重要的作用，尤其是在系统完全设计和部署之后。然而，人工智能创新的步伐要求模拟器和仿真器采用更加敏捷的方法来创建和使用基准，以实现未来的系统协同设计。我们提出了 Chakra，这是一种开放式图模式，用于标准化工作负载规范，捕捉关键操作和依赖关系，也称为执行跟踪（ET）。此外，我们还提出了一套互补的工具/功能，使各种模拟器、仿真器和基准能够收集、生成和采用 Chakra ET。例如，我们使用生成式人工智能模型学习数千种 Chakra ET 的潜在统计属性，并使用这些模型合成 Chakra ET。这些合成的 ET 可以混淆关键的专有信息，还可以针对未来的假设场景进行设计。举例来说，我们展示了一个端到端的概念验证，它能将 PyTorch ET 转换为 Chakra ET，并以此驱动一个开源训练系统模拟器（ASTRA-sim）。我们的最终目标是建立一个充满活力的全行业敏捷基准和工具生态系统，以推动未来的人工智能系统协同设计。

1.2 研究动机

Full Workload Benchmarks对于评价分布式训练的性能虽然非常重要，但是需要更加敏捷的方式
评价分布式训练的软硬件联合设计的效果
不同的AI/ML模型通常由不同的公司在不同的平台上构建，通常不能提供公开模型的细节
缺少在不同的组织之间交换ML模型ET的统一架构
缺少完善的性能评估工具链
目前的方法缺乏合成ET的能力
以Pytorch为例，可以使用Execution Graph Observer获取直接的ET，但是这种方法得到的结果受真实的系统的限制（例如NPU数量，计算时间，网络延迟等），使用合成的ET可以避免这些限制

1.3 主要贡献

捕捉相关执行和依赖信息的 Chakra ET
可视化器、转换器和模拟器等开源工具链
人工智能生成模型的ET合成

2. 实现方法

2.1 设计需求

应该包含内存访问、计算负载、网络通信、并行策略
应该最小化且可扩展
应该从性能建模的角度对 ML 执行的各个方面进行建模
应该能显示不同执行阶段的ET，且与设备解耦

2.2 实现方案

2.2.1 ET的表示

A. Chakra Node

type的值范围

INVALID 字段的初始值，用于忽略非关键节点
MEM_LOAD、MEM_STORE
COMP
COMM_SEND、COMM_RECV、COMM_COLL

B. Chakra Attribute