Lei Mao's Log Book
·
PyTorch CUDA图捕获
💡
原文英文，约3500词，阅读约需13分钟。
📝
内容提要
CUDA图是提升GPU性能的有效工具，PyTorch提供了torch.cuda.graph和torch.cuda.make_graphed_callables两个API来捕获和重放CUDA图。前者需要手动管理，后者简化了流程，适用于不同场景以提升PyTorch模型性能。
🎯
关键要点
- CUDA图是优化GPU性能的有效工具，减少CPU开销。
- PyTorch提供torch.cuda.graph和torch.cuda.make_graphed_callables两个API来捕获和重放CUDA图。
- torch.cuda.graph需要手动管理，而torch.cuda.make_graphed_callables简化了流程。
- 使用torch.cuda.graph时，用户可以完全控制图中包含的操作。
- torch.cuda.make_graphed_callables自动处理热身、静态缓冲区、图捕获和重放。
- torch.cuda.graph API提供最佳性能，适合捕获整个训练迭代。
- torch.cuda.make_graphed_callables API允许部分图捕获，适应动态工作负载。
- 在复杂模型中，可能无法将整个训练步骤捕获为单个图。
- 不同的CUDA图集成级别对训练模型的CPU时间有显著影响。
➡️
