PyTorch DataParallel与DistributedDataParallel运行无限挂起问题

尝试用torch.nn.parallel.DistributedDataParallel运行基础训练脚本，无报错直接无限挂起。改用nn.DataParallel运行简单网络也遇到同样问题。已尝试设置NCCL_LL_THRESHOLD=0和NCCL_P2P_DISABLE=1，问题未解决，需要排查思路。

复现代码 #

# script.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel
import types
import argparse
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 2, bias=False)
        self.drop = nn.Dropout(p=0.5)
    def forward(self, x):
        print('fc1.weight {}'.format(self.fc1.weight))
        x = self.fc1(x)
        x = self.drop(x)
        print('x {}'.format(x))
        return x
def main():
    parser = argparse.ArgumentParser(description='fdsa')
    parser.add_argument("--local_rank", default=0, type=int)
    args = parser.parse_args()
    args.gpu = args.local_rank
    torch.cuda.set_device(args.gpu)
    torch.distributed.init_process_group(backend='nccl',
                                         init_method='env://')
    args.world_size = torch.distributed.get_world_size()
    model = MyModel().to(args.gpu)
    model = DistributedDataParallel(
        model,
        device_ids=[args.gpu],
        output_device=args.local_rank,
    )
    for i in range(2):
        model.zero_grad()
        x = torch.randn(1, 2, device=args.gpu)
        out = model(x)
        print('iter {}, out {}'.format(i, out))
        out.mean().backward()
if __name__ == "__main__":
    main()

执行命令及输出 #

执行命令：

python -m torch.distributed.launch --nproc_per_node=2 script.py

输出：

FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use_env is set by default in torchrun.
If your script expects `--local_rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See 
https://pytorch.org/docs/stable/distributed.html#launch-utility for 
further instructions

  warnings.warn(
WARNING:torch.distributed.run:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
objdet1-0:9671:9671 [0] NCCL INFO Bootstrap : Using eth0:10.233.113.162<0>
objdet1-0:9671:9671 [0] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
objdet1-0:9671:9671 [0] NCCL INFO P2P plugin IBext
objdet1-0:9671:9671 [0] NCCL INFO NET/IB : No device found.
objdet1-0:9671:9671 [0] NCCL INFO NET/IB : No device found.
objdet1-0:9671:9671 [0] NCCL INFO NET/Socket : Using [0]eth0:10.233.113.162<0>
objdet1-0:9671:9671 [0] NCCL INFO Using network Socket
NCCL version 2.10.3+cuda11.6
objdet1-0:9672:9672 [1] NCCL INFO Bootstrap : Using eth0:10.233.113.162<0>
objdet1-0:9672:9672 [1] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
objdet1-0:9672:9672 [1] NCCL INFO P2P plugin IBext
objdet1-0:9672:9672 [1] NCCL INFO NET/IB : No device found.
objdet1-0:9672:9672 [1] NCCL INFO NET/IB : No device found.
objdet1-0:9672:9672 [1] NCCL INFO NET/Socket : Using [0]eth0:10.233.113.162<0>
objdet1-0:9672:9672 [1] NCCL INFO Using network Socket

nvidia-smi输出 #

Every 2.0s: nvidia-smi                                                                                                                                                                                    objdet1-0: Fri Jul 29 02:11:00 2022
Fri Jul 29 02:11:01 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.129.06   Driver Version: 470.129.06   CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA A100-SXM...  On   | 00000000:07:00.0 Off |                    0 |
| N/A   26C    P0    61W / 400W |    858MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA A100-SXM...  On   | 00000000:0F:00.0 Off |                    0 |
| N/A   25C    P0    61W / 400W |    858MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA A100-SXM...  On   | 00000000:47:00.0 Off |                    0 |
| N/A   25C    P0    52W / 400W |      3MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA A100-SXM...  On   | 00000000:4E:00.0 Off |                    0 |
| N/A   26C    P0    52W / 400W |      3MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

排查思路 #

• 检查进程间通信端口是否被防火墙拦截：NCCL默认使用随机端口，可手动设置NCCL_SOCKET_IFNAME=eth0指定网卡，同时确认该网卡的端口未被防火墙封禁。
• 验证NCCL通信可用性：运行NCCL自带的测试工具nccl-tests，比如all_reduce_perf，确认多GPU间能正常完成通信操作。
• 检查PyTorch与CUDA、NCCL版本兼容性：确保PyTorch版本与CUDA 11.6、NCCL 2.10.3匹配，版本不兼容可能导致底层通信死锁。
• 禁用NCCL插件：当前日志显示加载了nccl_rdma_sharp_plugin，尝试设置NCCL_NET_PLUGIN=none禁用插件，避免插件引入的问题。
• 测试单GPU运行：先注释掉分布式相关代码，确认单GPU下脚本能正常运行，排除模型本身或训练逻辑的问题。
• 检查进程状态：用ps -ef | grep script.py查看进程是否存在，用gdb或strace跟踪系统调用，定位挂起时的具体操作。
• 确认GPU P2P功能：运行nvidia-smi topo -m查看GPU间的拓扑连接，若P2P不可用，尝试强制禁用P2P（已试过NCCL_P2P_DISABLE=1，可再确认是否生效）。

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