PyTorch DistributedDataParallel初始化广播阶段张量尺寸不匹配报错求助

问题现象 #

• 基于GroupFormer项目采用DistributedDataParallel训练时，初始化广播阶段触发张量尺寸不匹配的RuntimeError
• 单机器2GPU（--nnodes=1 --nproc_per_node=2）和单GPU（--nnodes=1 --nproc_per_node=1）场景均出现相同错误，按逻辑单GPU场景不应触发参数广播同步逻辑，但仍报错

报错堆栈 #

Traceback (most recent call last):
  File "main.py", line 53, in <module>
    main()
  File "main.py", line 43, in main
    group_helper = Group(config, work_dir=config['basedir'])
  File "/home/disk1/wgf/project/GroupFormer/group/group.py", line 54, in __init__
    self._build()
  File "/home/disk1/wgf/project/GroupFormer/group/group.py", line 59, in _build
    self._build_model()
  File "/home/disk1/wgf/project/GroupFormer/group/group.py", line 104, in _build_model
    self.model = DistributedDataParallel(model.cuda(), device_ids=[self.rank % torch.cuda.device_count()],
  File "/home/disk1/wgf/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 648, in __init__
    _sync_module_states(
  File "/home/disk1/wgf/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/distributed/utils.py", line 113, in _sync_module_states
    _sync_params_and_buffers(
  File "/home/disk1/wgf/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/distributed/utils.py", line 131, in _sync_params_and_buffers
    dist._broadcast_coalesced(
RuntimeError: The size of tensor a (64) must match the size of tensor b (0) at non-singleton dimension 3

训练命令 #

python -m torch.distributed.run --nnodes=1 --nproc_per_node=1 --node_rank=0  --master_port=22332 main.py

测试验证 #

使用简单ToyModel测试相同分布式训练流程，未触发错误：

import torch.nn as nn
import torch.nn.functional as F

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

问题分析与修复方案 #

可能原因 #

1. 模型存在动态初始化的张量/缓冲区：GroupFormer的核心分组机制可能存在延迟初始化逻辑（比如在forward首次执行时才创建张量），导致DDP初始化时，模型参数/缓冲区未完全初始化，部分张量尺寸为0。即使单GPU，DDP仍会执行参数同步检查，触发尺寸不匹配错误。
2. 空张量未被正确初始化：模型中某些模块（如分组注意力、特征聚合层）的参数或缓冲区初始化时，因配置错误（如分组数设为0、特征维度计算错误）生成了维度为0的空张量。
3. 分布式环境下配置未同步：即使单GPU启动分布式进程，若模型初始化依赖的关键配置（如输入尺寸、分组数）未通过分布式广播同步，可能导致进程内张量尺寸不一致（逻辑上的主进程与自身同步时出现矛盾）。

修复步骤 #

• 排查动态初始化逻辑：检查GroupFormer模型代码，找到所有在__init__方法外创建参数/缓冲区的代码，确保这些逻辑在DDP初始化之前执行完成，且所有进程中张量尺寸一致。
• 定位空张量：在DDP初始化前添加调试代码，遍历模型所有参数和缓冲区，打印形状找出维度为0的张量：

for name, param in model.named_parameters():
    if 0 in param.shape:
        print(f"空参数：{name}，形状：{param.shape}")
for name, buf in model.named_buffers():
    if 0 in buf.shape:
        print(f"空缓冲区：{name}，形状：{buf.shape}")

针对找到的空张量，检查其初始化逻辑，修正配置错误（如分组数、特征维度的计算）。

• 同步分布式配置：在模型初始化前，使用torch.distributed.broadcast同步所有关键配置参数，确保所有进程（即使单进程）使用相同的配置生成模型张量。
• 检查DDP初始化时机：确保模型完全初始化并移至GPU后，再创建DistributedDataParallel实例，避免DDP捕捉到未完全初始化的模型状态。

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