[Bug Skip]: DumpError can't serialize torch.nn.attention._SDPBackend in test_ac_rematerialize_with_sdpa_dropout_zero
- 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
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