PyTorch Lightning分布式训练NCCL后端超时问题求助

使用PyTorch Lightning运行基于DDP（分布式数据并行）的跨节点训练，配置了自定义DDPStrategy（指定NCCL后端、自定义超时时间、自定义ClusterEnvironment实现类），但出现分布式连接超时错误：

• 主节点报错：torch.distributed.DistStoreError: Timed out after 1801 seconds waiting for clients. 2/4 clients joined.
• 从节点报错：torch.distributed.DistBackendError: [4] is setting up NCCL communicator and retrieving ncclUniqueId from [0] via c10d key-value store by key '0', but store->get('0') got error: Connection reset by peer

训练配置详情 #

核心训练代码 #

strategy = DDPStrategy(
        cluster_environment=CustomEnvironment(),
        process_group_backend="nccl",
        timeout=CUSTOM_TIMEOUT,
        find_unused_parameters=True)

# Initialize a trainer
trainer = pl.Trainer(logger=logger,
                    callbacks=[checkpoint_callback],
                     max_epochs=hparams["epochs"],
                     devices=4,  # 单节点4块GPU
                     accelerator="gpu",
                     strategy=strategy)

自定义CustomEnvironment实现 #

from typing import Union, Any, Dict
from pytorch_lightning.plugins.environments import ClusterEnvironment
from datetime import timedelta
import os
import logging

log = logging.getLogger(__name__)

DEFAULT_TIMEOUT = timedelta(seconds=1800)
CUSTOM_TIMEOUT = timedelta(seconds=3600)

class CustomEnvironment(ClusterEnvironment):
    def __init__(self, num_nodes=2):
        super().__init__()
        self._num_nodes = num_nodes
        self._master_port = None
        self._world_size = None
        self._global_rank = None

    def creates_processes_externally(self):
        # 由PyTorch Lightning内部管理进程
        return False

    def detect(self):
        # 按需实现节点与进程检测逻辑
        log.debug("Detect method is called.")

    def global_rank(self):
        if self._global_rank is None:
            self._global_rank = int(os.getenv("RANK", 0))
        log.debug(f"GLOBAL_RANK: {self._global_rank}")
        return self._global_rank

    @property
    def main_address(self):
        return self.master_address()

    @property
    def main_port(self):
        return self.master_port()

    def set_global_rank(self, rank: int):
        self._global_rank = rank
        log.debug(f"Set GLOBAL_RANK: {self._global_rank}")

    def set_world_size(self, world_size: int):
        self._world_size = world_size
        log.debug(f"Set WORLD_SIZE: {self._world_size}")

    def master_address(self):
        MASTER_ADDR = os.getenv("MASTER_ADDR")
        log.debug(f"MASTER_ADDR: {MASTER_ADDR}")
        return MASTER_ADDR

    def master_port(self):
        if self._master_port is None:
            self._master_port = os.getenv("MASTER_PORT")
        log.debug(f"MASTER_PORT: {self._master_port}")
        return int(self._master_port)

    def world_size(self):
        if self._world_size is None:
            log.debug("WORLD_SIZE is not set.")
        return self._world_size

    def node_rank(self):
        MY_RANK = int(os.getenv("NODE_RANK", "0"))
        log.debug(f"NODE_RANK: {MY_RANK}")
        return int(MY_RANK)

    def local_rank(self) -> int:
        LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
        log.debug(f"LOCAL_RANK: {LOCAL_RANK}")
        return LOCAL_RANK

环境变量配置 #

已在所有节点手动设置：

export NCCL_SOCKET_IFNAME=eth0
export NCCL_P2P_DISABLE=1

且所有节点的MASTER_ADDR、MASTER_PORT、NODE_RANK、WORLD_SIZE=8（2节点×4GPU）环境变量配置一致。

1. 修复CustomEnvironment的world_size逻辑漏洞 #

当前world_size方法在_world_size未设置时返回None，会导致PyTorch Lightning无法正确识别全局进程数，引发节点连接不完整问题。修改该方法：

def world_size(self):
    if self._world_size is None:
        # 优先从环境变量读取，避免返回None
        self._world_size = int(os.getenv("WORLD_SIZE", 1))
        log.debug(f"WORLD_SIZE loaded from env: {self._world_size}")
    return self._world_size

2. 验证超时配置是否生效 #

主节点报错超时时间接近默认的1800秒，说明自定义的3600秒超时可能未生效，需检查：

• 确认CUSTOM_TIMEOUT被正确传入DDPStrategy，无代码覆盖逻辑
• 若使用PyTorch Lightning版本<2.0，可尝试通过环境变量TORCH_DIST_TIMEOUT=3600强制设置超时时间

3. 网络连通性与防火墙检查 #

• 验证所有节点通过eth0网卡互通：在主/从节点分别执行ping <对方节点eth0地址>，确保无丢包
• 开放主节点的MASTER_PORT端口：执行ufw allow <MASTER_PORT>/tcp（依防火墙工具调整），避免端口被拦截
• 确认NCCL_SOCKET_IFNAME=eth0生效：所有节点执行echo $NCCL_SOCKET_IFNAME，确保输出为eth0且是节点间通信的正确网卡

4. 进程启动方式与版本验证 #

• 所有节点使用相同版本的PyTorch、PyTorch Lightning、NCCL，版本不兼容会导致通信失败
• 跨节点启动时，确保主/从节点使用相同启动命令，且从节点指定了正确的NODE_RANK

5. 调整NCCL通信参数（可选） #

若以上步骤无效，开启NCCL调试日志并调整参数：

# 开启NCCL调试日志，定位具体通信问题
export NCCL_DEBUG=INFO
# 禁用NCCL重试机制，避免连接重置循环
export NCCL_MAX_RETRIES=0

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