如何调试PyTorch分布式ProcessGroup？寻求更优IDE调试方法

问题背景 #

需要分析PyTorch分布式后端接口，但调试时遇到子进程无法调试的问题。试过VSCode Python调试+GDB附加、Python C++调试器，均未解决子进程调试难题，寻求更便捷的IDE调试方法。

测试代码：

import os
import time
import numpy as np
import torch
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.distributed._tensor import DTensor, DeviceMesh, Shard, Replicate, distribute_tensor,zeros

def run(rank, size):
    a = torch.tensor([[0, 2.], [3, 0]])
    a.to_sparse()
    if rank == 0:
        print(a.to_dense())
    dist.barrier()
    dist.all_reduce(a)
    dist.barrier()
    if rank == 0 :
        print(a.to_dense())

def init_process(rank_id, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '12347'
    dist.init_process_group(backend, rank=rank_id, world_size=size)
    fn(rank_id, size)

if __name__ == "__main__":
    big_tensor = torch.arange(0,16).reshape(4,4)
    size = 1
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, run))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

可行调试方案 #

1. VSCode 多进程调试配置 #

修改.vscode/launch.json，开启子进程自动附着并配置必要环境变量：

{
    "version": "0.2.0",
    "configurations": [
        {
            "name": "Python: Distributed Debug",
            "type": "python",
            "request": "launch",
            "program": "${file}",
            "env": {
                "MASTER_ADDR": "127.0.0.1",
                "MASTER_PORT": "12347",
                "PYDEVD_WARN_EVALUATION_TIMEOUT": "1000"
            },
            "subProcess": true,
            "justMyCode": false
        }
    ]
}

若自动附着失效，可在init_process开头添加手动附着逻辑：

import debugpy
debugpy.listen(('127.0.0.1', 5678 + rank_id))
debugpy.wait_for_client()

再在VSCode中添加对应端口的调试配置，逐个附着子进程。

2. PyCharm Professional 原生多进程调试 #

PyCharm支持自动附着子进程，只需简单配置：

• 打开Run -> Edit Configurations
• 找到目标Python配置，勾选Attach to subprocess automatically while debugging
• 直接启动调试，PyCharm会自动捕获所有spawn生成的子进程，无需修改代码。

3. 手动端口分配 + IDE附着 #

为每个子进程分配独立调试端口，再用IDE逐个连接：
在init_process开头添加：

import debugpy
# 每个rank使用不同端口避免冲突
debugpy.listen(('127.0.0.1', 5678 + rank_id))
print(f"Rank {rank_id} 等待调试器连接，端口：{5678 + rank_id}")
debugpy.wait_for_client()

启动主进程后，在VSCode/PyCharm中创建对应端口的调试配置，逐个附着子进程。

4. TorchRun + IDE调试 #

改用torchrun启动分布式进程，配合IDE调试：
将启动命令改为：

torchrun --nproc_per_node=2 your_script.py

在IDE中配置调试命令为上述torchrun指令，开启子进程自动附着，IDE可识别所有由torchrun启动的进程。

注意事项 #

• 开启justMyCode: false，否则无法调试PyTorch后端内部代码
• 调试C++后端（如NCCL、Gloo）时，可在子进程中添加os.system("gdb -p {}".format(os.getpid()))手动触发GDB附着
• 调试初期建议从单进程（size=1）开始，逐步扩展到多进程场景

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