Chakra项目使用指南：从安装到可视化工具全解析-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00126/article/details/148682543

Chakra项目使用指南：从安装到可视化工具全解析

chakra Repository for MLCommons Chakra schema and tools 项目地址: https://gitcode.com/gh_mirrors/cha/chakra

项目概述

Chakra是一个用于机器学习工作负载分析和模拟的工具集，由MLCommons社区开发维护。它提供了一套完整的工具链，用于捕获、转换、分析和可视化机器学习模型的执行轨迹。通过Chakra，研究人员和开发者可以深入理解模型在分布式系统中的执行行为，优化性能瓶颈，并进行各种模拟实验。

环境安装与配置

虚拟环境搭建

在开始使用Chakra前，强烈建议创建一个独立的Python虚拟环境。这可以避免与系统Python环境产生冲突，也便于后续的依赖管理。

# 创建名为chakra_env的虚拟环境
python3 -m venv chakra_env

# 激活虚拟环境
source chakra_env/bin/activate

Chakra核心安装

在激活虚拟环境后，可以通过以下方式安装Chakra：

# 从源码安装
pip install .

依赖组件安装

Chakra的正常运行需要两个关键依赖组件：

PARAM：提供基础参数管理功能

git clone git@github.com:facebookresearch/param.git
cd param/et_replay
git checkout 7b19f586dd8b267333114992833a0d7e0d601630
pip install .

Holistic Trace Analysis：用于轨迹分析

git clone https://github.com/facebookresearch/HolisticTraceAnalysis.git
cd HolisticTraceAnalysis
git checkout d731cc2e2249976c97129d409a83bd53d93051f6
git submodule update --init
pip install -r requirements.txt
pip install -e .

核心工具详解

1. 执行轨迹链接器(chakra_trace_link)

这个工具的主要功能是将主机端和设备端的执行轨迹合并为一个完整的执行轨迹文件。

典型使用场景：当你的机器学习模型在GPU上运行时，主机CPU和GPU会分别生成执行轨迹。使用此工具可以将两者关联起来，形成完整的执行视图。

chakra_trace_link \
    --chakra-host-trace /path/to/chakra_host_trace \
    --chakra-device-trace /path/to/chakra_device_trace \
    --rank [RANK] \
    --output-file /path/to/chakra_host_device_trace.json

2. 执行轨迹转换器(chakra_converter)

将JSON格式的执行轨迹转换为protobuf格式，便于后续的模拟和分析。

chakra_converter PyTorch \
    --input /path/to/chakra_host_device_trace.json \
    --output /path/to/chakra_trace \
    [--simulate]

参数说明：

--simulate：可选参数，启用后会进行模拟验证，可以检查转换后的轨迹是否合理

3. 执行轨迹馈送器(et_feeder)

这是一个C++库，用于将Chakra轨迹馈送到兼容的C++模拟器中。目前主要支持ASTRA-sim模拟器。

使用流程：

克隆并构建ASTRA-sim
配置Chakra前端
生成执行轨迹
运行模拟

4. 可视化工具集

Chakra提供了多种可视化工具，帮助用户直观理解执行轨迹：

执行轨迹可视化器(chakra_visualizer)

chakra_visualizer \
    --input_filename /path/to/chakra_et
    --output_filename /path/to/output.[graphml|pdf|dot]

支持多种输出格式，包括GraphML、PDF和DOT格式。

JSON转换器(chakra_jsonizer)

将二进制格式的执行轨迹转换为易读的JSON格式：

chakra_jsonizer \
    --input_filename /path/to/chakra_et \
    --output_filename /path/to/output_json

时间线可视化器(chakra_timeline_visualizer)

生成执行时间线图，直观展示各操作在不同NPU上的执行情况：

chakra_timeline_visualizer \
    --input_filename /path/to/input.csv \
    --output_filename /path/to/output.json \
    --num_npus 4 \
    --npu_frequency 1.5GHz

可视化结果可以通过Chrome浏览器的chrome://tracing工具查看。