如何在共享vGPU环境下限制Spacy、Torch等框架的单进程显存占用

1. PyTorch等效显存占比限制配置 #

PyTorch没有和TensorFlow per_process_gpu_memory_fraction 完全一致的原生参数，但可以通过以下两种方式实现完全等效的显存占比限制：

• 直接设置进程显存占比上限：在进程第一次调用CUDA操作前，执行 torch.cuda.set_per_process_memory_fraction(0.3, device=0)，其中0.3就是你要限制的显存占比，device参数指定对应的GPU设备编号即可。
• 固定显存分配上限：设置环境变量 PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:<你的最大显存MB数>，可以强制PyTorch显存分配器不申请超过指定值的显存，适合需要固定硬限制的场景。

2. spaCy等效显存占比限制配置 #

spaCy的显存限制完全依赖其使用的深度学习后端：

• 若使用spaCy 3.x及以上版本，默认采用基于PyTorch的Thinc后端，直接按上述PyTorch的方法配置即可。
• 若使用旧版本基于TensorFlow的后端，直接使用你已知的 per_process_gpu_memory_fraction 参数配置即可。
  配置操作建议放在加载spaCy模型之前执行，保证限制生效。

3. 共享vGPU多模型混合部署方案 #

你当前使用的vGPU插件本身支持集群层面的显存配额分配，结合框架层面的显存限制可以实现稳定的多模型混合部署：

1. 先在Kubeflow Pod调度层，给每个模型的服务Pod分配对应比例的vGPU显存：比如给TensorFlow模型Pod分配40%显存配额，spaCy模型Pod分配20%，PyTorch模型Pod分配30%，从集群层面先做资源隔离，避免进程间显存抢占。
2. 再在每个模型的业务代码里，配置对应框架的显存上限，和集群分配的配额对齐，避免框架本身的显存预分配逻辑超出集群配额触发OOM错误。
   两层配置配合使用，就能实现多模型同时在同一共享vGPU上稳定运行，不会出现内存不足的问题。

内容的提问来源于stack exchange，提问作者Stephan