PyTorch训练随机卡在DataLoader加载阶段问题排查求助

问题背景 #

原8卡分布式训练代码改为单卡（--nproc_per_node=1）运行后，训练初始正常，但会随机卡在DataLoader数据加载环节，GPU使用率降至0%。即使将num_workers设为0仍出现卡顿，调试发现代码永久阻塞在torch.cuda.synchronize()处。相关代码与报错信息如下：

DataLoader初始化代码 #

data_loader = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train,
        batch_size=config.DATA.BATCH_SIZE,
        num_workers=config.DATA.NUM_WORKERS,
        pin_memory=config.DATA.PIN_MEMORY,
        drop_last=True,
    )

其中dataset_train为torchvision.datasets.ImageFolder实例，batch_size=32。

阻塞位置代码 #

for idx, (samples, targets) in enumerate(data_loader):
    outputs = model(samples)
       
    print(f"before synchronize ...",end="")
    torch.cuda.synchronize()
    print(f" synchronized ...",end="")

    # more code below ...

原因分析 #

1. 分布式残留逻辑冲突：原分布式代码中未清理的进程组初始化、分布式采样器等逻辑，在单卡运行时会引发异步通信阻塞，最终导致CUDA同步卡住。
2. 隐性CUDA错误：数据加载或模型前向过程中可能存在设备不匹配、内存越界等隐性错误，这类错误不会直接抛出异常，但会导致后续synchronize()永久阻塞。
3. 弹性训练框架兼容性问题：使用torch.distributed.run/launch启动单卡训练时，elastic的进程监控逻辑可能与单卡运行环境不兼容，引发进程间死锁。

可行解决方案 #

1. 彻底清理分布式残留代码 #

• 跳过分布式初始化：仅在多卡训练时执行分布式进程组初始化，单卡时完全跳过：

  import torch.distributed as dist
  if config.DISTRIBUTED:
      dist.init_process_group(backend='nccl')
      # 训练结束后主动销毁进程组
      dist.destroy_process_group()
  

• 替换分布式采样器：单卡运行时禁用DistributedSampler，改用默认随机采样或直接开启shuffle：

  data_loader = torch.utils.data.DataLoader(
      dataset_train, 
      sampler=sampler_train if config.DISTRIBUTED else None,
      batch_size=config.DATA.BATCH_SIZE,
      num_workers=config.DATA.NUM_WORKERS,
      pin_memory=config.DATA.PIN_MEMORY,
      drop_last=True,
      shuffle=not config.DISTRIBUTED  # 单卡模式下开启shuffle
  )
  

2. 排查并修复隐性CUDA错误 #

• 添加CUDA错误检查：在关键步骤后检查CUDA状态，定位触发阻塞的操作：

  torch.cuda.empty_cache()
  for idx, (samples, targets) in enumerate(data_loader):
      # 确保数据已正确转移到GPU
      samples = samples.cuda(non_blocking=True)
      targets = targets.cuda(non_blocking=True)
      
      outputs = model(samples)
      # 检查前向传播是否存在CUDA错误
      cuda_err = torch.cuda.get_last_error()
      if cuda_err:
          print(f"Iter {idx} forward pass CUDA error: {cuda_err}")
          exit()
      
      print(f"before synchronize ...", end="")
      torch.cuda.synchronize()
      print(f" synchronized ...", end="")
  

• 禁用pin_memory：尝试将pin_memory设为False，排除内存拷贝环节的阻塞问题。

3. 调整单卡训练启动方式 #

• 直接用Python命令启动：放弃torch.distributed.launch/run，直接执行python train.py，避免elastic框架的进程监控干扰。
• 若必须用分布式启动器：添加--standalone参数，强制单节点独立运行：

  torchrun --nproc_per_node=1 --standalone train.py
  

4. 其他调试与优化手段 #

• 用nvidia-smi实时监控GPU内存，排查是否存在内存泄漏或突然占满的情况。
• 临时降低batch_size至16或8，排除内存不足导致的隐性阻塞。
• 更新PyTorch版本到稳定版，避免旧版本的CUDA同步bug。

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