最近在做一个多工具调用的 Agent 项目,底层用 PyTorch 2.1,想着用 torch.compile 加速一下推理。但发现 Agent 每步的输入长度和工具调用次数都不固定,动态 shape 特别多,compile 之后要么重新编译好几次,要么直接 fallback 回 eager,速度反而更慢了。我看官方说支持 dynamic shape,但实际效果感觉不太行。有没有大佬在生产环境试过?是不是 Agent 这种控制流复杂的场景就不适合上 compile?还是我用法有问题?
楼主
29天前
PyTorch 2.x 的 torch.compile 在 Agent 推理场景下到底能不能打?
请 登录 后发表回复
全部回复
共 3 条
2楼
11天前
写得挺好,建议补充一些性能数据。
3楼
4天前
刚入门,这个对我帮助很大。
4楼
23小时前
动态shape确实是个坑,我们项目里也踩过。Agent这种每步输入长度飘忽不定的场景,torch.compile的guard机制很容易反复触发重编译,最后性能还不如老老实实跑eager。可以试试把compile范围收窄到单个工具调用的子图,或者用dynamic=True配合mark_dynamic把常见shape分桶,别指望一把梭整个Agent都编译进去。另外2.1对控制流的支持确实一般,升到2.3+再试试可能体感会好点。