Torch编译缓存,提升推理速度
很多模型,特别是FLUX系列的模型,会应用各种torch.compile技术/技巧来提高推理速度。
对编译函数的第一次调用会跟踪并编译代码,这会增加开销。随后的调用则运行优化后的代码,速度会显著加快。
提示在我们对black-forest-labs/flux-kontext-dev模型进行推理速度测试时,编译版本比未编译版本的运行速度快了30%以上。
性能提升
通过在模型容器生命周期之间缓存编译产物,我们看到了冷启动时间的显著改善:
- black-forest-labs/flux-kontext-dev: ~120秒 → ~60秒(提速50%)
- prunaai/flux-schnell: ~150秒 → ~70秒(提速53%)
- prunaai/flux.1-dev-lora: ~400秒 → ~150秒(提速62%)
该缓存还改善了所有使用torch.compile的模型从容器的启动到首次成功预测所需的时间。
工作原理
该缓存系统的工作原理与许多CI/CD缓存系统类似:
- 当模型容器启动时,它会寻找缓存的编译产物。
- 如果找到,Torch会复用它们,而不是从头开始重新编译。
- 当容器正常关闭时,如果需要,它们会更新缓存。
- 缓存文件以模型版本为键进行存储,并放置在靠近GPU节点的位置。
了解更多
要了解更多关于如何使用torch.compile的信息,请查阅相关的技术文档和官方PyTorchtorch.compile教程。FINISHED
更多精彩内容 请关注我的个人公众号 公众号(办公AI智能小助手)或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号(网络安全技术点滴分享)