RL-Insight - GMM 专家负载热力图可视化工具

一、简介

RL-Insight 是一个强化学习性能数据快速分析的可视化工具，基于 VeRL 框架采集的 profiling 数据、GMM 专家负载数据进行解析。其中 GMM 专家负载热力图功能用于可视化 MoE（Mixture of Experts）模型中专家的负载分布情况。

模块划分、流水线与扩展步骤见 架构说明。更完整的数据目录与 JSON 字段约定见 数据规格与格式说明。

主要功能

  • GMM 数据解析：支持解析 VeRL 框架采集的 GMM 专家负载数据

  • 专家负载热力图：生成热力图展示不同专家在不同阶段的负载情况

  • 灵活的过滤功能：支持按 step、role、rank 进行数据过滤，自定义可视化范围

安装

需要 Python >= 3.10。

pip install rl-insight

若需使用最新未发布功能，可从源码安装：

git clone https://github.com/verl-project/rl-insight.git
cd rl-insight
pip install -r requirements.txt
pip install -e .

二、快速使用

2.1 采集 GMM 数据

使用 VeRL 框架，基于 msprobe 采集 GMM 专家负载数据，详细参考：

VeRL 集成 msprobe 数据采集 VeRL 采集专家负载数据实践

输入目录结构说明

gmm_parser 当前按 step -> role -> rank -> dump_tensor_data -> group_list.pt 的目录语义解析数据。推荐输入目录结构如下：

gmm_dump/
├── step_1/                                # 训练步骤
│   ├── actor_compute_log_prob/            # 角色/阶段
│   │   └── rank0/                         # Rank ID
│   │       └── dump_tensor_data/          # 张量数据目录
│   │           ├── NPU.npu_grouped_matmul.0.forward.kwargs.group_list.pt
│   │           ├── NPU.npu_grouped_matmul.1.forward.kwargs.group_list.pt
│   │           └── ...                    # 多个 GMM 算子的 group_list.pt 文件
│   └── actor_update/
│       └── rank0/
│           └── dump_tensor_data/
│               ├── NPU.npu_grouped_matmul.0.forward.kwargs.group_list.pt
│               └── ...
├── step_2/
└── ...

路径字段含义：

  • step_<n>：训练 step（对应 heatmap.parser.step 过滤）

  • <role>：角色名（对应 heatmap.parser.role 过滤）

  • rank<n>：rank id（对应 input.rank_list 过滤）

  • dump_tensor_data/*.group_list.pt：MoE grouped_matmul 的专家负载；典型为一维整型张量，第 i 个元素表示第 i 个 expert 分到的 token 数

2.2 执行分析脚本

热力图使用示例

python -m rl_insight.main \
    input.path=<gmm_data_path> \
    input.rank_list=all \
    heatmap.parser.type=gmm \
    heatmap.visualizer.type=gmm_heatmap \
    heatmap.parser.step=1 \
    heatmap.parser.role=actor_compute_log_prob \
    output.path=<output_path>

或修改并直接使用 examples/gmm_exec.sh 脚本:

bash examples/gmm_exec.sh

三、命令行参数

以下说明与 python -m rl_insight.main -h 保持一致；若有出入以命令行帮助为准。

3.1 公共参数

参数

默认值

说明

input.path

（必填）

GMM 数据的根目录路径

output.path

output

输出路径，若为文件夹则在其中生成 gmm_heatmap.png

input.rank_list

all

Rank ID 列表，默认 all 表示所有 rank，可指定多个 rank 用逗号分隔

preset

自动推断

预设名称：timeline、heatmap（根据 CLI 参数自动推断）

config_path

无

YAML 配置文件路径

3.2 Heatmap 专属参数

参数

默认值

说明

heatmap.parser.type

gmm

解析器类型：gmm

heatmap.visualizer.type

gmm_heatmap

可视化类型

heatmap.parser.step

无

特定的 step 进行可视化（可选，支持 1 或 1,2）

heatmap.parser.role

无

特定的 role 进行可视化（可选）

heatmap.visualizer.dpi

200

热力图输出的 DPI

heatmap.visualizer.cmap

viridis

热力图的颜色映射

heatmap.visualizer.gmm_per_layer

3

每个 MoE layer 前向阶段预期的 grouped_matmul 次数

四、输出说明

工具会在指定的输出路径下生成 PNG 文件（文件名默认为 gmm_heatmap.png），包含：

  • 专家负载热力图：横轴为专家索引，纵轴为合并后的 layer；颜色表示各 expert 的 token 负载

  • Segments 色条：位于热力图左侧，与各行对齐，按 step、role、rank 分段着色

  • 颜色编码：使用 viridis 颜色映射，负载越高颜色越深

  • 分隔线：清晰分隔不同的 step、role、rank 组合

  • 图例：显示各个 segment 对应的 step、role、rank 信息

图表解读

  1. 热力图区域：

    • 横轴（X）：专家索引（expert_index）

    • 纵轴（Y）：layer（图中刻度为 layer0、layer1、…，表示解析阶段合并后的层序号）

  2. 左侧 Segments 色条：

    • 与热力图 逐行对齐（共用同一垂直范围）

    • 按 step、role、rank 分段着色；段内文字为 step · role · rank 摘要

    • 段与段之间有分隔线，便于区分不同上下文

  3. 刻度密度：

    • X 轴：专家数较多时会稀疏采样刻度，避免重叠

    • Y 轴：layer / 行数较多时会自动降低 layerK 标签数量（约上限 40），保持可读

五、注意事项

  1. 热力图功能使用 preset=heatmap 或在 CLI 中指定 heatmap. 开头的参数即可启用

  2. 当 output.path 只指定文件夹路径时，会在该文件夹中生成 gmm_heatmap.png 文件

  3. 当不指定 heatmap.parser.step、heatmap.parser.role 或 input.rank_list 参数时，默认显示所有数据

  4. 对于大量数据，工具会自动调整图表大小和标签显示密度，确保可读性

  5. 数据文件需包含有效的专家负载数据，包括 step、role、rank_id、stage、expert_index 和 load 等字段

  6. 若你的模型实现中每层 grouped_matmul 次数不等于 3，请显式设置 heatmap.visualizer.gmm_per_layer 以获得更准确的 actor_update 前向阶段截断结果

目录与 JSON 字段的集中说明另见 数据规格与格式说明。运行时校验逻辑以 rl_insight.data.DataChecker 及 rl_insight/data/rules.py 中的规则定义为准。