遭遇ProcessExitedException，如何在Databricks Notebook用torch.multiprocessing启动多进程

问题根因 #

Databricks Notebook的代码运行在交互式解释器上下文当中，不存在标准Python脚本的__main__模块入口，代码中编写的if __name__ == '__main__'判断条件在Notebook环境下永远不生效。当调用mp.spawn以spawn方式启动子进程时，子进程会重新导入当前Notebook的所有全局代码，陷入逻辑循环最终异常退出，触发你遇到的ProcessExitedException报错。

解决方案 #

方案1（推荐）：使用Databricks官方适配的TorchDistributor工具 #

Databricks Runtime ML 11.0及以上版本内置了TorchDistributor，专门适配Notebook环境下的PyTorch分布式训练，无需自行处理多进程启动逻辑，示例代码如下：

import torch.distributed as dist
from pyspark.ml.torch.distributor import TorchDistributor

def train(rank, world_size):
    # 初始化分布式进程组，CPU训练用gloo后端，GPU训练用nccl后端
    dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
    print(f"hello from process {rank}")
    # 此处写入你的训练逻辑
    dist.destroy_process_group()

if __name__ == '__main__':
    processes = 4
    # 单节点训练设置local_mode=True，多节点集群训练设置local_mode=False
    TorchDistributor(num_processes=processes, local_mode=True).run(train)
    print("completed")

方案2：将训练逻辑封装为独立Python脚本调用 #

如果你需要使用原生torch.multiprocessing逻辑，可以把训练代码保存为独立的.py脚本，再在Notebook中调用该脚本，独立脚本的__name__ == '__main__'判断可以正常生效：

1. 将你的训练代码写入/dbfs/tmp/train.py
2. 在Notebook单元格中执行以下命令调用脚本：

import subprocess
subprocess.run(["python", "/dbfs/tmp/train.py"], check=True)

注意事项 #

• 多节点分布式训练时，将TorchDistributor的local_mode参数设置为False，工具会自动调度集群所有worker节点的资源
• 分布式训练初始化进程组时，CPU训练选择gloo后端，GPU训练选择nccl后端

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