最近在折腾MCP Server,想把一个用PyTorch训练的文本分类模型封装成工具,让LLM能直接调用。本地跑推理没问题,但一接到MCP请求,推理几次后显存就爆了。我用了torch.no_grad(),也试了del model和torch.cuda.empty_cache(),但好像没完全释放。是不是MCP每次调用都会重新加载模型?还是需要自己搞个模型池管理?有没有踩过坑的大佬指点一下,或者推荐个轻量级的推理框架结合MCP用?谢了!
楼主
9小时前
MCP Server接入PyTorch模型,推理时总报显存泄漏怎么排查?
请 登录 后发表回复
全部回复
共 2 条
2楼
9小时前
跟你遇到一模一样的问题,核心就是MCP每次请求都会重新加载模型,而且Python的GC不一定及时回收,显存就慢慢垒上去了。我后来是用进程池+预加载模型对象,启动时一次性把模型加载到子进程,MCP请求走RPC调用,推理完只清中间变量不销毁进程,基本没再炸过。你也可以试试用FastAPI搭个最小推理服务,MCP只做HTTP转发,这样模型生命周期完全可控,排查也简单。
3楼
9小时前
这问题我碰过,典型的是MCP Server每次请求都fork子进程加载模型,但父进程的CUDA上下文没清理干净,导致显存越积越多。建议把模型做成常驻的独立进程,用IPC或共享内存跟MCP通信,或者直接用Ray Serve这种带自动GC的推理框架接管模型生命周期,能省很多排查心智负担。