vscode 调试 openmpi 后端的 pytorch 分布式程序

vscode 调试 openmpi 后端的 pytorch 分布式程序

1 验证 openmpi 是否支持 CUDA

ompi_info --parsable --all | grep mpi_built_with_cuda_support:value

输出：mca:mpi:base:param:mpi_built_with_cuda_support:value:true 表示 openmpi 支持 CUDA，否则需要重新编译 openmpi。

2 终端运行 openmpi 后端的 pytorch 分布式程序

pytorch 测试代码 run.py :

import torch
import torch.distributed as dist
import platform

def run(rank, size):
    tensor = torch.zeros(1).cuda(0)
    if rank == 0:
        # Send the tensor to process 1
        tensor += 1
        dist.send(tensor=tensor, dst=1)
    else:
        # Receive tensor from process 0
        dist.recv(tensor=tensor, src=0)
    print('Rank ', rank, ' has data ', tensor[0])

def init_process(fn):
    dist.init_process_group(backend='mpi')
    rank = dist.get_rank()
    size = dist.get_world_size()
    print('I am rank ', rank, ' on ', platform.node())
    fn(rank, size)

if __name__ == '__main__':
    init_process(run)

终端运行 mpirun -np 2 python run.py 输出：

$ mpirun -np 2 python run.py
I am rank  1  on  wh9
I am rank  0  on  wh9
Rank  0  has data  tensor(1., device='cuda:0')
Rank  1  has data  tensor(1., device='cuda:0')

注意：pytorch 默认不支持 openmpi 后端，需要从源码编译 pytorch。

3 vscode 调试 openmpi 后端的 pytorch 分布式程序

• 安装最新版本的 ptvsd, 同时将 vscode 更新到最新版本

pip install ptvsd

• 修改 run.py 代码，使支持 ptvsd 调试，每一个进程都添加以下代码：

ptvsd.enable_attach(address =('127.0.0.1', 3000))
ptvsd.wait_for_attach()

127.0.0.1 为本机IP地址，3000为调试端口号，不同的进程设置不同的端口号，用以调试不同的进程。

修改后 run.py 代码为：

import torch
import ptvsd
import torch.distributed as dist
import platform

def run(rank, size):
	ptvsd.enable_attach(address =('127.0.0.1', 3000+rank))
    ptvsd.wait_for_attach()
    tensor = torch.zeros(1).cuda(0)
    if rank == 0:
        # Send the tensor to process 1
        tensor += 1
        dist.send(tensor=tensor, dst=1)
    else:
        # Receive tensor from process 0
        dist.recv(tensor=tensor, src=0)
    print('Rank ', rank, ' has data ', tensor[0])

def init_process(fn):
    dist.init_process_group(backend='mpi')
    rank = dist.get_rank()
    size = dist.get_world_size()
    print('I am rank ', rank, ' on ', platform.node())
    fn(rank, size)
    
if __name__ == '__main__':
    init_process(run)

上述程序中: rank 0 的调试端口为 3000，rank 1 的调试端口为 3001。

• 运行 mpirun -np 2 python run.py，程序将阻塞到 wait_for_attach() 处，等待 vscode 连接。

• 打开 vscode， 在 launch.json 文件中新建 python 远程连接任务

这里调试端口 3000 表示调试 rank 0，同理，设置为 3001 就可以调试 rank 1。

• 添加断点

• 按 F5 远程调试程序，程序运行到断点处

注意：vscode 远程调试前，一定要先在终端中运行：mpirun -np 2 python run.py

参考资料

asroy/Debugging Mixed Python C++ code in Visual Studio Code