PAYATHON 2026

PyTorch 子进程发生 CUDA OOM 后如何安全重试？

架构老王

结论

通常不必为了处理 CUDA OOM 把 subprocess 换成 multiprocessing。应先等训练进程退出，再重试，并确认没有仍在使用 GPU 的后代进程。进程退出后，驱动会回收它持有的 CUDA context 和显存分配。torch.cuda.empty_cache() 只能释放当前存活进程中 PyTorch 缓存但未使用的显存，无法替其他进程释放显存。

训练脚本可以捕获 OOM 并把结构化信息传给父进程，但这些信息通常不足以准确推算下一次需要多少显存或合适的 batch_size。异常中的申请大小和显存统计可用于诊断，实际需求还受模型、输入、优化器状态和运行时分配影响。

等子进程退出后再重试

Popen.communicate() 会等待直接子进程结束。训练脚本因未捕获的异常退出时，父进程可以等 communicate() 返回后再启动重试。也可以用 subprocess.run() 简化等待和输出收集：

import subprocess

result = subprocess.run(
    ["python", "train.py", config_path],
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE,
    text=True,
    check=False,
)

if result.returncode != 0 and "CUDA out of memory" in result.stderr:
    reduce_batch_size(config_path)
    # 到这里直接子进程已经退出，再决定是否重试

这段示例只检查直接子进程。如果训练脚本还会创建训练 worker 或其他后代进程，直接子进程退出并不代表这些后代也已退出。需要按操作系统管理进程组或作业，并在重试前确认相关进程都已结束。僵尸进程已经退出，不会继续执行 CUDA 工作；仍存活的子进程才可能继续占用显存。

在训练脚本中传回 OOM 信息

可以在训练入口捕获异常，把诊断信息写入 stderr 或专用结果文件，再以非零状态退出。父进程仍应等进程结束后再重试：

# train.py
import json
import sys
import torch

try:
    train()
except RuntimeError as exc:
    if "out of memory" in str(exc).lower():
        info = {
            "type": "cuda_oom",
            "message": str(exc),
        }
        print("OOM_INFO=" + json.dumps(info), file=sys.stderr, flush=True)
    raise

父进程可以在 stderr 中查找 OOM_INFO=，再解析后面的 JSON。若要传递更稳定的结构化数据，也可以让子进程写入约定的结果文件。无论采用哪种方式，都要检查退出码，避免把日志文本误判为 OOM。

PyTorch 的 OOM 通常以 RuntimeError 抛出，但异常类型和错误文本可能随 PyTorch 版本变化。上例根据异常文本识别 OOM，适合用作诊断通道。应结合退出状态判断，不要把它当作长期稳定的协议。

如何设置重试策略

父进程可以对 OOM 设置有限次数的退避，例如逐步降低 batch_size。每次重试前保存配置并记录实际值。如果最小批量仍然 OOM、配置解析失败，或出现其他异常，就应停止重试并保留完整日志，避免陷入无限循环或掩盖非 OOM 错误。

OOM 消息中的显存数字可以帮助排查问题，但不足以直接推算合适的 batch_size。模型结构、序列长度或图像尺寸、梯度累积和优化器状态等因素都可能改变峰值显存。自动调整时，应设置可配置的最小值和最大重试次数。

multiprocessing 是否更合适

multiprocessing 不会自动解决显存回收问题。只有训练工作进程退出后，CUDA context 才会随进程生命周期释放；父进程仍需等待子进程退出。使用 multiprocessing 时，还要谨慎处理 CUDA 初始化和进程启动方式，并通过队列或管道传递结果。对于已经通过命令行启动训练脚本的工作流，继续使用 subprocess 通常更直接。

备注：内容仅供参考。