Linux下PyTorch CPU OOM导致SSH服务器中断的问题求助

问题背景 #

使用PyTorch 2.0.1+cu117版本，在Ubuntu 22.04系统（覆盖一台虚拟机、一台裸机）中，当发生CPU内存不足（OOM）时，系统会完全失去SSH连接及Xserver访问权限，只能通过vSphere或物理电源按钮强制重启：

• GPU OOM时进程会正常终止并返回标准错误，无系统异常
• 确认两台设备均已启用系统级OOM killer

已尝试/考虑过的方案（均未解决问题） #

• 配置RLIMIT内存限制，无效
• 计划在训练代码中添加可用内存检查逻辑，触发后终止训练，但认为该方案不够严谨
• 考虑过通过为Python设置别名（适配conda环境切换），将进程加入cgroup限制内存，但因对操作系统知识有限，且有人提醒修改cgroup存在风险，暂未尝试

需求 #

由于无法精准预估模型所需内存，需要一种能够在CPU OOM时优雅终止进程且不影响SSH服务正常运行的解决方案。

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