现在图的op数量很多时,比如大模型达到几千、上万的 op 时,性能定位变成一个难题。
现在最可靠的工具是 nvtx,但是 nvtx 的使用成本高:
- 不带 op name 的nvtx,需要特殊的启动开关来生成,之后用户要看懂也有难度;
- 带 op name 的 nvtx,需要编译特殊的 oneflow 版本
有个潜力是把 opcounter 项目扩展下,支持:
- 计算量;(已经做了,让其可以支持分nn.Module的统计)
- 内存;(导出内存复用的 chunk 和 block 数据给 python 端,op name 对应的内存块大小是知道的)
- 通信;(利用自动并行的 cost mode?l)