`AttentionPairBiasMask.backward()` fails with non-contiguous gradient tensors despite contiguous forward inputs

[gcroci2]

We're encountering a runtime error in AttentionPairBiasMask.backward() during training:

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

The error occurs at line 358 in attention_pair_bias_torch.py:

grad_out_mask = grad_out_mask.view(B, ctx.num_heads, -1)

All input tensors to cuet.attention_pair_bias() in the forward pass are verified contiguous, and the forward pass completes successfully every time. The error occurs during backpropagation when gradient tensors become non-contiguous.

Environment

• Hardware: H100 GPUs
• cuEquivariance version: 0.6.1
• PyTorch version: 2.7
• CUDA version: 12.x

We've implemented a monkey patch that makes gradient tensors contiguous before calling the original backward function, which resolves the issue. However:

• Should AttentionPairBiasMask.backward() handle non-contiguous gradient tensors? (A simple fix would be changing view() to reshape())
• Or is it expected that all gradient inputs will be contiguous? If so, what's the recommended approach for ensuring this?

[hsadasivan]

Hello @gcroci2 ,
Thanks for reporting this.
Could you post a simple reproducer for us to test this?

[gcroci2]

Could you post a simple reproducer for us to test this?

import torch
import cuequivariance_torch as cuet

if __name__ == "__main__":
    device = torch.device("cuda")
    s = torch.rand(*(1, 248, 384), dtype=torch.float32).to(device).requires_grad_(True)
    q = torch.rand(*(1, 16, 248, 32), dtype=torch.float32).to(device).requires_grad_(True)
    k = torch.rand(*(1, 16, 248, 32), dtype=torch.float32).to(device).requires_grad_(True)
    v = torch.rand(*(1, 16, 248, 32), dtype=torch.float32).to(device).requires_grad_(True)
    z = torch.rand(*(1, 248, 248, 128), dtype=torch.float32).to(device).requires_grad_(True)
    mask = (torch.rand(*(1, 248)) > 0.5).to(torch.bool).to(device).requires_grad_(False)
    w_proj_z = torch.rand(*(16, 128), dtype=torch.float32).to(device).requires_grad_(True)
    w_proj_g = torch.rand(*(512, 384), dtype=torch.float32).to(device).requires_grad_(True)
    w_proj_o = torch.rand(*(384, 512), dtype=torch.float32).to(device).requires_grad_(True)
    w_ln_z = torch.rand(*(128,), dtype=torch.float32).to(device).requires_grad_(True)
    b_ln_z = torch.rand(*(128,), dtype=torch.float32).to(device).requires_grad_(True)
    b_proj_z = torch.rand(*(16,), dtype=torch.float32).to(device).requires_grad_(True)
    b_proj_g = torch.rand(*(512,), dtype=torch.float32).to(device).requires_grad_(True)
    b_proj_o = torch.rand(*(384,), dtype=torch.float32).to(device).requires_grad_(True)
    
    output, proj_z = cuet.attention_pair_bias(
        s=s.contiguous(), q=q.contiguous(), k=k.contiguous(), v=v.contiguous(), z=z.contiguous(), 
        mask=mask.contiguous(), num_heads=16, attn_scale=0.1767766952966369, multiplicity=1,
        w_proj_z=w_proj_z.contiguous(), w_proj_g=w_proj_g.contiguous(), w_proj_o=w_proj_o.contiguous(),
        w_ln_z=w_ln_z.contiguous(), b_ln_z=b_ln_z.contiguous(), b_proj_z=b_proj_z.contiguous(),
        b_proj_g=b_proj_g.contiguous(), b_proj_o=b_proj_o.contiguous()
    )
    
    print("Forward pass succeeded")
    loss = output.sum()
    print("Starting backward pass...")
    loss.backward()
    print("Backward pass succeeded")

Output:

Forward pass succeeded
Starting backward pass...
Traceback (most recent call last):
  File "/app/test_cueq.py", line 45, in <module>
    loss.backward()
  File "/opt/conda/lib/python3.11/site-packages/torch/_tensor.py", line 648, in backward
    torch.autograd.backward(
  File "/opt/conda/lib/python3.11/site-packages/torch/autograd/__init__.py", line 353, in backward
    _engine_run_backward(
  File "/opt/conda/lib/python3.11/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/autograd/function.py", line 307, in apply
    return user_fn(self, *args)
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/cuequivariance_ops_torch/attention_pair_bias_torch.py", line 358, in backward
    grad_out_mask = grad_out_mask.view(B, ctx.num_heads, -1)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.