PyTorch 2.0中3D卷积层组合时的AssertionError问题求助

在GPT神经网络中添加两个3D卷积层（分别用于通道扩展和信息聚合），输入序列形状为(128,96,3,3,3,100)（128为batch size，96为序列时间步，100为特征维度，(3,3,3)表示27邻域信息）。为适配nn.Conv3d的输入要求，在forward函数中将数据重塑为(12288,100,3,3,3)。单独使用任一卷积层时代码运行正常，但同时使用两个层时，反向传播阶段出现AssertionError：

assert_size_stride(buf303, (12288, 300, 3, 3, 3), (8100, 1, 2700, 900, 300))
AssertionError: expected size 300==300, stride 27==1 at dim=1

相关代码：

class GPT(nn.Module):

    def __init__(self, config,
                 dim=100,
                 bias=False,):
        super().__init__()
        self.fe_1=nn.Conv3d(dim, dim*3, kernel_size=1,stride=1, bias=bias)
        self.fe_2=nn.Conv3d(dim*3, dim, kernel_size=3,stride=1, bias=bias)
        
        self.transformer = nn.ModuleDict(dict(...))
    def forward(self, idx, targets=None):
        device = idx.device
        b, t, h, w, l, c = idx.shape
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)

        
        idx = idx.permute(0,1,5,2,3,4)
        idx = idx.reshape((-1,c,h,w,l))#(12288,100,3,3,3)
        out = self.fe_1(idx)
        out = self.fe_2(out)
        idx = out.reshape((b,t,c))


        # forward the transformer itself
        tok_emb = self.transformer.wte(idx) 
        pos_emb = self.transformer.wpe(pos)
        x = self.transformer.drop(tok_emb + pos_emb)

错误栈信息：

Traceback (most recent call last):
  File "/home/xxx/nanoGPT-master/nanoGPT-master/traincopy_0726_725_cnn.py", line 334, in <module>
    scaler.scale(loss).backward()
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/_tensor.py", line 487, in backward
    torch.autograd.backward(
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/autograd/__init__.py", line 200, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/autograd/function.py", line 274, in apply
    return user_fn(self, *args)
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 2348, in backward
    out = call_compiled_backward()
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 2324, in call_compiled_backward
    out = call_func_with_args(
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1247, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/home/xxx/.conda/envs/torchreborn/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 248, in run
    return model(new_inputs)
  File "/tmp/torchinductor_xxx/qx/cqxoopct2ddmpulx7yvt4x3iokeyymbcilpaq5z6gnphb3z6p6j5.py", line 2252, in call
    assert_size_stride(buf303, (12288, 300, 3, 3, 3), (8100, 1, 2700, 900, 300))
AssertionError: expected size 300==300, stride 27==1 at dim=1

错误含义解析 #

这里的stride不是卷积操作的步长，而是Tensor内存布局中的步幅——即访问当前维度下一个元素时需要跳过的元素数量。错误提示第1维度（通道维度）的实际步幅为27，但预期是1，说明该维度的内存不连续，导致PyTorch的TorchInductor编译后端在反向传播时无法匹配预期的内存布局。

核心原因 #

两次卷积叠加后，输出Tensor的内存布局被打乱，而直接调用reshape不会自动整理内存连续性。单独使用一层时，内存连续性仍能满足后端要求，但两层叠加后，内存碎片化触发了断言错误。此外，第二个卷积层使用kernel_size=3但未设置padding，输入的(3,3,3)邻域会被卷积为(1,1,1)，后续reshape虽形状匹配，但存在维度逻辑不匹配的潜在风险。

具体解决步骤 #

• 显式保持内存连续：在每次卷积输出后调用.contiguous()，强制Tensor整理内存布局：

  out = self.fe_1(idx).contiguous()
  out = self.fe_2(out).contiguous()
  

• 修正卷积padding设置：如果需要保持空间维度(3,3,3)不变，给第二个卷积层添加padding=1：

  self.fe_2=nn.Conv3d(dim*3, dim, kernel_size=3, stride=1, padding=1, bias=bias)
  

• 临时禁用编译后端排查：若上述方法无效，可临时禁用TorchInductor验证是否为后端兼容性问题，在训练代码开头添加：

  import torch
  torch._dynamo.config.suppress_errors = True
  torch.backends.cudnn.benchmark = False
  

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