Repository: pytorch
Issue:  UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance. grad.sizes() = [1, 64, 64, 64], strides() = [1, 4096, 64, 1] bucket_view.sizes() = [1, 64, 64, 64], strides() = [262144, 4096, 64, 1]
Labels: oncall: distributed,module: ddp

### 🐛 Describe the bug

class RFNO(nn.Module):
    def __init__(self, out_channels=64, modes1=64, modes2=64):
        super(RFNO, self).__init__()
        self.out_channels = out_channels
        self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1
        self.modes2 = modes2

        self.scale = (1 / out_channels)
        self.weights0 = nn.Parameter(self.scale * torch.rand(1, out_channels, 1, 1, dtype=COMPLEXTYPE))
        self.weights1 = nn.Parameter(self.scale * torch.rand(1, out_channels, modes1, modes2, dtype=COMPLEXTYPE))
        self.weights2 = nn.Parameter(self.scale * torch.rand(1, out_channels, modes1, modes2, dtype=COMPLEXTYPE))

    def compl_mul2d(self, input, weights):
        # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y)
        return torch.einsum("bixy,ioxy->boxy", input, weights)

    def forward(self, x):
        batchsize = x.shape[0]
        print(self.weights1.stride()) 

        x_ft = torch.fft.rfft2(x)
        x_ft = x_ft * self.weights0
        out_ft = torch.zeros(batchsize, self.out_channels,  x.size(-2), x.size(-1)//2 + 1, dtype=COMPLEXTYPE, device=x.device)
        
        out_ft[:, :, :self.modes1, :self.modes2] = \
            self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
        out_ft[:, :, -self.modes1:, :self.modes2] = \
            self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)
       
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x
        
        
        When I use DDP for parallel training it gives warnings, but not for single card training.  “ UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance. grad.sizes() = [1, 64, 64, 64], strides() = [1, 4096, 64, 1] bucket_view.sizes() = [1, 64, 64, 64], strides() = [262144, 4096, 64, 1]” 

        When I set def __init__(self, out_channels=64, modes1=64, modes2=64) to def __init__(self, out_channels=64, modes1=32, modes2=32).  "UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance. grad.sizes() = [1, 64, 32, 32], strides() = [1, 1024, 32, 1]  bucket_view.sizes() = [1, 64, 32, 32], strides() = [65536, 1024, 32, 1] "
        
        

### Versions

class RFNO(nn.Module):
    def __init__(self, out_channels=64, modes1=64, modes2=64):
        super(RFNO, self).__init__()
        self.out_channels = out_channels
        self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1
        self.modes2 = modes2

        self.scale = (1 / out_channels)
        self.weights0 = nn.Parameter(self.scale * torch.rand(1, out_channels, 1, 1, dtype=COMPLEXTYPE))
        self.weights1 = nn.Parameter(self.scale * torch.rand(1, out_channels, modes1, modes2, dtype=COMPLEXTYPE))
        self.weights2 = nn.Parameter(self.scale * torch.rand(1, out_channels, modes1, modes2, dtype=COMPLEXTYPE))

    def compl_mul2d(self, input, weights):
        # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y)
        return torch.einsum("bixy,ioxy->boxy", input, weights)

    def forward(self, x):
        batchsize = x.shape[0]
        print(self.weights1.stride()) 

        x_ft = torch.fft.rfft2(x)
        x_ft = x_ft * self.weights0
        out_ft = torch.zeros(batchsize, self.out_channels,  x.size(-2), x.size(-1)//2 + 1, dtype=COMPLEXTYPE, device=x.device)
        
        out_ft[:, :, :self.modes1, :self.modes2] = \
            self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
        out_ft[:, :, -self.modes1:, :self.modes2] = \
            self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)
       
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x
        
        
        When I use DDP for parallel training it gives warnings, but not for single card training.  “ UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance. grad.sizes() = [1, 64, 64, 64], strides() = [1, 4096, 64, 1] bucket_view.sizes() = [1, 64, 64, 64], strides() = [262144, 4096, 64, 1]” 

        When I set def __init__(self, out_channels=64, modes1=64, modes2=64) to def __init__(self, out_channels=64, modes1=32, modes2=32).  "UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance. grad.sizes() = [1, 64, 32, 32], strides() = [1, 1024, 32, 1]  bucket_view.sizes() = [1, 64, 32, 32], strides() = [65536, 1024, 32, 1] "
        
        

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o