用 Profiler 追踪 ops-transformer 算子:GE 融合与 Runtime 调度的实战调试

用 Profiler 追踪 ops-transformer 算子:GE 融合与 Runtime 调度的实战调试

大模型训练跑不动,大多数人第一反应是"算力不够"。但我见过的实际情况里,80% 以上的性能问题出在算子调度和数据搬运上,不是算力本身。解决这个问题最有用的工具,就是 CANN 内置的 Profiler。

这篇文章不聊架构理论,手把手带你在真实环境里用 Profiler 追踪 ops-transformer 的算子,理解 GE 的融合决策和 Runtime 的调度行为。

环境确认:先检查工具是否就绪

bash 复制代码
# 确认 CANN 已安装
source /usr/local/Ascend/ascend-toolkit/set_env.sh
which ascend snorkel-validate

# 确认昇腾 NPU 可用
npu-smi info

npu-smi info 输出 NPU 型号和状态后,继续下一步。如果这一步就报错,先把 CANN 环境装好再往下走。

第一步:写一个带 ops-transformer 算子的测试脚本

找个能复现问题的地方,先写一个简单的测试脚本:

python 复制代码
# test_attention_profile.py
import torch
import torch_npu
from flash_attention_ops import flash_attention_npu

# 构造一个 Attention 计算
batch, heads, seq_len, dim = 4, 32, 2048, 64
q = torch.randn(batch, heads, seq_len, dim, dtype=torch.float16).npu()
k = torch.randn(batch, heads, seq_len, dim, dtype=torch.float16).npu()
v = torch.randn(batch, heads, seq_len, dim, dtype=torch.float16).npu()

# 跑 10 次取平均时间
import time
for _ in range(3):
    _ = flash_attention_npu(q, k, v, causal=True)  # 先预热,有 JIT 编译

torch.npu.synchronize()
start = time.perf_counter()
for _ in range(10):
    out = flash_attention_npu(q, k, v, causal=True)
torch.npu.synchronize()
elapsed = (time.perf_counter() - start) / 10
print(f"单次耗时: {elapsed*1000:.2f}ms")

跑通之后记录下单次耗时。后面开 Profiler 再跑一次,对比有没有性能下降。

第二步:开启 Profiler 数据采集

CANN 的 Profiler 会抓 GPU Trace 和 AI Trace 两类数据。GPU Trace 记录每个算子在 NPU 上的执行时间,AI Trace 记录 Python 层的调用链路。

python 复制代码
# test_attention_profile.py 中加入 Profiler
from torch_npu.profiler import profile

with profile(
    activities=[
        torch_npu.profiler.ProfilerActivity.NPU,  # 记录 NPU 层
        torch_npu.profiler.ProfilerActivity.CPU,  # 记录 CPU 层调用
    ],
    record_shapes=True,        # 记录输入 shape,用于分析算子融合
    profile_memory=True,      # 记录显存使用
    with_stack=True,          # 记录 Python 调用栈
    export_name="attention_trace.json"  # 输出文件名
):
    for _ in range(10):
        out = flash_attention_npu(q, k, v, causal=True)

执行完成后,当前目录下会生成 attention_trace.json

第三步:用 CANN 的 Profiler UI 分析数据

attention_trace.json 拉到有 GUI 的机器上分析,或者在服务器上用内置的查看工具:

bash 复制代码
# 如果服务器有图形界面,直接用 HUD 打开
# CANN Profiler GUI 的启动命令
ascend-profile -i attention_trace.json

# 如果是纯命令行机器,导出关键数据
ascend-profile -i attention_trace.json -o summary.csv --format csv

在 Profiler GUI 里,重点关注这几个视图:

Timeline 视图(GPU Trace)

找到 ops-transformer 的 FlashAttention 算子,看它在整个 Timeline 里占的宽度。宽说明执行时间长,窄说明快。同时注意它旁边有没有其他算子紧挨着------如果 MatMul → Softmax → MatMul 三个算子紧紧挨着,说明没有被融合;如果中间有大段空白,说明有数据搬运的等待时间。

Operator 视图(AI Trace)

这个视图列出每个算子的总耗时和调用次数。重点找 ops-transformer 相关的算子,看单次平均耗时是多少,有没有异常值。如果某个算子的耗时波动很大(最大值是最小值的 5 倍以上),大概率是数据排布不对触发了额外的格式转换。

Memory 视图

看 HBM 的读写量。ops-transformer 的融合算子应该显著减少中间结果的写入。如果发现大量小块的 HBM 读写,说明算子没走融合路径,GE 的优化没有生效。

第四步:分析 GE 的融合决策

GE 的融合决策记录在编译日志里。要看这个日志,先在训练脚本里加上环境变量让 GE 输出详细的融合信息:

bash 复制代码
# 设置 GE 融合日志级别
export ASCEND_GLOBAL_LOG_LEVEL=3
export ENABLE_OOL_OP痕过融合LOG=1

# 重新跑你的训练脚本,日志会输出 GE 的融合决策过程
python train_llama.py 2>&1 | tee ge_fusion_log.txt

打开 ge_fusion_log.txt,搜索 ops-transformerflash_attention,会看到类似这样的输出:

复制代码
[GE Fusion] 子图 #15 检测到算子序列:
  MatMul[qkt] + Softmax + MatMul[pv]
  匹配融合规则: flash_attention_fusion_pass
  融合为单一算子: FlashAttentionKernel
  输出 shape: (batch, heads, seq_len, dim)

[GE Fusion] 子图 #23 检测到算子序列:
  GeLU + Add
  匹配融合规则: activation_fusion_pass
  融合为: FastGeLU

如果你的 FlashAttention 没有出现在融合日志里,说明 GE 没有识别到这个算子序列。可能的原因:输入的 dtype 是 float32 而不是 float16(某些融合规则需要 dtype 匹配)、Tensor 的 shape 没有对齐到融合规则要求的倍数、算子没有被正确注册到 Framework Adaptor。

第五步:看 Runtime 的调度行为

Runtime 的调度信息也在 Profiler 里,但要看得更细可以用 npu-smi 实时监控:

bash 复制代码
# 另开一个终端,持续采样 NPU 利用率和 HBM 带宽
watch -n 0.5 npu-smi dmon -c 1 -s puc -d 10

这个命令每 0.5 秒输出一次 NPU 利用率和显存带宽。

如果 NPU 利用率长期低于 60%,但模型在正常跑,问题大概率出在数据搬运而不是计算------可能是输入数据没有提前搬到 NPU、算子之间有依赖等待、或者 batch size 太小导致计算密度不够。

如果 HBM 带宽利用率很高(>80%),但 NPU 利用率不高,说明带宽成了瓶颈------这是经典的分块策略问题,ops-transformer 的 FlashAttention 在这种场景下尤其容易出问题,因为它本来就是为了解决带宽瓶颈而设计的,如果带宽还是不够,可能是 tile 大小没有针对当前 shape 做优化。

第六步:调优之后的验证

根据 Profiler 的分析结果做调优,改完之后再跑一次带 Profiler 的测试,对比两张 Timeline 图:

python 复制代码
# 对比脚本
import json

def load_profile(path):
    with open(path) as f:
        return json.load(f)

before = load_profile("attention_trace_before.json")
after = load_profile("attention_trace_after.json")

# 提取 ops-transformer 相关算子的耗时
def get_op_time(profile, op_name_prefix):
    return sum(
        evt["duration"] for evt in profile["traceEvents"]
        if evt["name"].startswith(op_name_prefix)
    )

op_name = "FlashAttention"  # 根据实际日志里的名字调整
t_before = get_op_time(before, op_name)
t_after = get_op_time(after, op_name)
speedup = t_before / t_after
print(f"加速比: {speedup:.2f}x")

如果加速比 > 1.5x,说明调优方向对了。如果没变化,说明改动没有触及真正的瓶颈。

系统学习 GE 和 Runtime 的调试方法

Profiler 只是一个工具,真正用好它需要理解 GE 和 Runtime 的设计逻辑。cann-learning-hub 里有关于 Profiler 使用和结果解读的专项教程,比官方文档更贴近实战:

建议按这个顺序学:

  1. cann-learning-hub → CANN Profiler 使用指南 → 学会解读 Timeline 和 Operator 视图
  2. cann-learning-hub → GE 图引擎调优 → 理解融合规则怎么写、怎么调试
  3. cann-learning-hub → Runtime 性能调优 → 理解调度策略和显存管理
  4. 回到你的训练脚本,用 Profiler 验证学到的调优方法

相关仓库:

https://atomgit.com/cann/ops-transformer

https://atomgit.com/cann/cann-learning-hub

https://atomgit.com/cann/ge

相关推荐
后端小肥肠2 小时前
一人公司如何用 WorkBuddy + Obsidian 搭一套长期记忆系统?
人工智能·aigc·agent
RFID舜识物联网2 小时前
破局“信息孤岛”:RFID耐高温标签重塑汽车喷漆车间可视化
大数据·人工智能·科技·物联网·安全·汽车
05大叔2 小时前
预训练模型演化,提示词工程
人工智能·深度学习·自然语言处理
BU摆烂会噶2 小时前
【LangGraph】House_Agent 实战(一):架构与环境配置
人工智能·vscode·python·架构·langchain·人机交互
小小测试开发2 小时前
OpenAI 模型攻克离散几何 80 年难题:Erdős 单位距离猜想被 AI 证明
人工智能·算法·机器学习
moonsims2 小时前
从“传感器融合”升级为“多机器人约束融合系统”-Factor Graph 多约束融合
人工智能·算法
tedcloud1232 小时前
agent-skills部署教程:打造工程化AI Agent系统
服务器·人工智能·系统架构·powerpoint·dreamweaver
测试员周周2 小时前
【Appium 系列】第15节-视觉测试 — 截图、对比、视觉回归
人工智能·python·数据挖掘·回归·appium·测试用例·测试覆盖率
Dfreedom.2 小时前
模型剪枝完全指南:从理论到实践,打造高效深度学习模型
人工智能·算法·机器学习·剪枝·模型加速