gpu - 分析 JAX 代码：什么是 redzone_checker，为什么要花这么多时间？

Question

我找到了这篇文章，但仍不清楚redzone_checker内核在做什么以及为什么。具体来说，它是否应该占用我应用程序运行时间的 90% 以上？TensorBoard 报告说它占用了我的 JAX 代码的绝大部分运行时间，我想知道

1. 实际上是这个内核花费了太多时间，还是这是使用 TensorBoard 分析 JAX 的副作用（即，输出在某种程度上具有误导性）？
2. 有没有办法减少redzone_checker内核花费的时间？这甚至是个好主意吗？

提前感谢您的任何见解。

Answer 1

在分析之前确保预热。

它可能是 jit 编译时间。