intel / intel/torch-xpu-ops

[Bug Skip]: DumpError can't serialize torch.nn.attention._SDPBackend in test_ac_rematerialize_with_sdpa_dropout_zero

Open
#4,945 0 comments 0 reactions 0 assignees View on GitHub
skipped
Dominant language
Python
Stars
113
Forks
128
Avg merge
5d 13h
Merged PRs (30d)
107

Description

### 🐛 Describe the bug with skip template

DumpError occurs while serializing test report data for `third_party.torch-xpu-ops.test.xpu.dynamo.test_activation_checkpointing_xpu.RematerializeACNodesPassTests.test_ac_rematerialize_with_sdpa_dropout_zero`. The failure is:

```xml

...
DumpError: can't serialize

```

-------
### Skip
Cases:
op_ut,third_party.torch-xpu-ops.test.xpu.dynamo.test_activation_checkpointing_xpu.RematerializeACNodesPassTests,test_ac_rematerialize_with_sdpa_dropout_zero

-------

### Error Information
test
```python
@unittest.skipIf(
not torch.get_device_module(GPU_TYPE).is_available(), "CUDA not available"
)
def test_ac_rematerialize_with_sdpa_dropout_zero(self):
from torch.nn.attention import sdpa_kernel, SDPBackend

cases = []
if PLATFORM_SUPPORTS_MEM_EFF_ATTENTION:
cases.append((SDPBackend.EFFICIENT_ATTENTION, torch.float32))
if PLATFORM_SUPPORTS_FLASH_ATTENTION:
cases.append((SDPBackend.FLASH_ATTENTION, torch.float16))
if PLATFORM_SUPPORTS_CUDNN_ATTENTION:
cases.append((SDPBackend.CUDNN_ATTENTION, torch.float16))
if not cases:
self.skipTest("No fused SDPA backends available")
sdpa_ops = {
torch.ops.aten.scaled_dot_product_attention.default,
torch.ops.aten._scaled_dot_product_cudnn_attention.default,
torch.ops.aten._scaled_dot_product_flash_attention.default,
torch.ops.aten._scaled_dot_product_efficient_attention.default,
torch.ops.aten._scaled_dot_product_fused_attention_overrideable.default,
}

def policy_fn(ctx, op, *args, **kwargs):
if op in sdpa_ops:
return torch.utils.checkpoint.CheckpointPolicy.PREFER_RECOMPUTE
return torch.utils.checkpoint.CheckpointPolicy.PREFER_SAVE

context_fn = functools.partial(
torch.utils.checkpoint.create_selective_checkpoint_contexts, policy_fn
)

for backend, dtype in cases:
with self.subTest(backend=backend, dtype=dtype):
torch._dynamo.reset()
q = torch.randn(
2, 4, 128, 64, device=GPU_TYPE, dtype=dtype, requires_grad=True
)
k = torch.randn(
2, 4, 128, 64, device=GPU_TYPE, dtype=dtype, requires_grad=True
)
v = torch.randn(
2, 4, 128, 64, device=GPU_TYPE, dtype=dtype, requires_grad=True
)

def fwd_bwd_with_sdpa(q, k, v):
with sdpa_kernel(backend):
z = torch.utils.checkpoint.checkpoint(
lambda q, k, v: F.scaled_dot_product_attention(
q, k, v, dropout_p=0.0
),
q,
k,
v,
use_reentrant=False,
context_fn=context_fn,
)
loss = z.sum()
dq, dk, dv = _grad(loss, (q, k, v))

return z.detach(), dq, dk, dv

result_with, gm_with = self._compile_and_capture(
fwd_bwd_with_sdpa, True, (q, k, v)
)
torch._dynamo.reset()
result_without, _ = self._compile_and_capture(
fwd_bwd_with_sdpa, False, (q, k, v)
)
eager_inputs = tuple(
t.detach().clone().requires_grad_(True) for t in (q, k, v)
)
result_eager = fwd_bwd_with_sdpa(*eager_inputs)

for actual, expected in zip(result_with, result_without):
self.assertEqual(actual, expected)
for actual, expected in zip(result_with, result_eager):
self.assertEqual(actual, expected)
self.assertEqual(sum(self.count_op(gm_with, op) for op in sdpa_ops), 2)

```

Sample reproduction context from the failure report:

```text
third_party.torch-xpu-ops.test.xpu.dynamo.test_activation_checkpointing_xpu.RematerializeACNodesPassTests
test_ac_rematerialize_with_sdpa_dropout_zero
Traceback ends in xdist serialization during testreport sending:
DumpError: can't serialize
```

### XPU PR
https://github.com/intel/torch-xpu-ops/pull/4894

### Versions
Version context from report: pytorch main 2.15.0a0
pytorch main 2.15.0a0

Contributor guide

Open the contributing guide

Research direction

Start with third_party.torch-xpu-ops.test.xpu.dynamo.test_activation_checkpointing_xpu.RematerializeACNodesPassTests.test_ac_rematerialize_with_sdpa_dropout_zero and reproduce the failure in the test-report or xdist serialization path. Inspect how the test's subTest parameters are included in the report. Done means the test report serializes successfully without DumpError for _SDPBackend.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
testing-qa
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
50/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.