test_Embedding_sparse_xpu (and test_EmbeddingBag_sparse_xpu) fail under PYTORCH_TEST_WITH_DYNAMO=1: IndexError: tuple index out of range in dynamo pgo.is_stride_dynamic
- Dominant language
- Python
- Stars
- 113
- Forks
- 129
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
### 🐛 Describe the bug
`TestNN.test_Embedding_sparse_xpu` fails under `PYTORCH_TEST_WITH_DYNAMO=1` with an `IndexError: tuple index out of range` raised from Dynamo's PGO (profile-guided-optimization) automatic-dynamic-shapes stride tracking, while tracing the backward pass of a sparse-gradient `nn.Embedding(sparse=True)`. Like the companion issue for `test_Conv1d_zero_batch_xpu_fp32`, this is triggered by test names introduced in https://github.com/pytorch/pytorch/pull/189653, which generalized `test/test_nn.py`'s `add_test()` to generate device-suffixed test names (e.g. `_xpu`) for any accelerator, not just `_cuda`. The equivalent CPU-named test (`TestNN.test_Embedding_sparse`) already has a skip entry in `test/compiled_autograd_skips/`, but the new `_xpu`-suffixed name does not.
### Reproduction steps
```bash
# Apply/checkout https://github.com/pytorch/pytorch/pull/189653 (or any revision after it lands)
# on an XPU-enabled PyTorch build.
PYTORCH_TEST_WITH_DYNAMO=1 python test/test_nn.py TestNN.test_Embedding_sparse_xpu
```
### Mini-reproducer (no test harness needed)
This standalone script reproduces the crash purely on CPU, with no dependency on `test/test_nn.py` or `common_nn.py`. Save it as `repro_embedding_sparse.py` and run `python repro_embedding_sparse.py`:
```python
import torch
import torch._dynamo
torch._dynamo.config.compiled_autograd = True
def test_fn():
weight = torch.zeros(4, 3, dtype=torch.double, requires_grad=True)
indices = torch.tensor([[0, 1, 0, 1]])
# Repeated forward/backward on a sparse=True embedding, mirroring the
# 5-iteration backward loop in
# torch.testing._internal.common_nn.NewModuleTest.test_device()
for i in range(5):
out = torch.nn.functional.embedding(indices, weight, sparse=True)
grad_out = out.clone().detach().normal_()
out.backward(grad_out, retain_graph=True) # <-- crashes on a later iteration
print(f"iter {i}: weight.grad.is_sparse = {weight.grad.is_sparse}")
# The torch._dynamo.optimize(...) wrapper mirrors how
# torch.testing._internal.common_utils.TestCase._run_custom() wraps the
# entire test method under PYTORCH_TEST_WITH_DYNAMO=1, which is what
# actually triggers the crash (a bare torch._dynamo.config.compiled_autograd=True
# without wrapping the whole call stack in dynamo does not reproduce it).
wrapped = torch._dynamo.optimize("eager_noexcept", nopython=False)(test_fn)
wrapped()
```
### Error log
```text
Traceback (most recent call last):
File ".../torch/_dynamo/variables/builder.py", line 4846, in _wrap_to_fake_tensor_and_record_impl
symbolic_context = _automatic_dynamic(
File ".../torch/_dynamo/variables/builder.py", line 4677, in _automatic_dynamic
config.automatic_dynamic_shapes and frame_state_entry.is_stride_dynamic(i)
File ".../torch/_dynamo/pgo.py", line 340, in is_stride_dynamic
return self.stride[dim] is auto_dynamic
IndexError: tuple index out of range
from user code:
File ".6", line 5, in forward
getitem = inputs[0]
Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"
```
Full pytest traceback:
```text
test/test_nn.py:7310: in with_tf32_off
test.test_device(self, **kwargs)
torch/testing/_internal/common_nn.py:3590: in test_device
cpu_gradInput = test_case._backward(cpu_module, cpu_input_tuple, cpu_output, cpu_gradOutput)
test/test_nn.py:96: in _backward
output.backward(grad_output, retain_graph=True, create_graph=create_graph)
torch/_tensor.py:623: in backward
torch.autograd.backward(...)
torch/autograd/__init__.py:395: in backward
_engine_run_backward(...)
torch/_dynamo/compiled_autograd.py:1273: in runtime_wrapper
out = compiled_fn(...)
...
torch/_dynamo/variables/builder.py:4846: in _wrap_to_fake_tensor_and_record_impl
symbolic_context = _automatic_dynamic(...)
torch/_dynamo/variables/builder.py:4677: in _automatic_dynamic
config.automatic_dynamic_shapes and frame_state_entry.is_stride_dynamic(i)
torch/_dynamo/pgo.py:340: in is_stride_dynamic
return self.stride[dim] is auto_dynamic
IndexError: tuple index out of range
```
### Root cause analysis
1. `PYTORCH_TEST_WITH_DYNAMO=1` sets `torch._dynamo.config.compiled_autograd = True` globally for the test suite (`torch/testing/_internal/common_utils.py`, around line 1990), so `output.backward(...)` in `test/test_nn.py`'s `_backward()` helper traces the backward pass through Dynamo instead of running purely in eager mode.
2. The `Embedding_sparse` test entry (`torch/testing/_internal/common_nn.py:1808-1815`) constructs `nn.Embedding(4, 3, dtype=torch.double, sparse=True)`, whose backward produces a **sparse COO-layout gradient** (`torch.sparse_coo`) for `weight.grad`.
3. `NewModuleTest.test_device()` calls `_backward()` in a loop 5 times (`common_nn.py:3587-3594`), each time re-invoking `output.backward(...)`.
4. Compiled autograd's Dynamo tracer wraps intermediate tensors seen during backward-graph tracing via `wrap_to_fake_tensor_and_record()` → `_automatic_dynamic()` (`torch/_dynamo/variables/builder.py:4677`), which calls `frame_state_entry.is_stride_dynamic(i)` for every `i in range(e.dim())`.
5. `FrameStateSizeEntry.is_stride_dynamic()` (`torch/_dynamo/pgo.py:316-340`) does `return self.stride[dim] is auto_dynamic`. The `self.stride` tuple is built/merged across repeated invocations of the same traced code object via `update_automatic_dynamic()` / `FrameStateSizeEntry.__ior__` (`pgo.py:394-422`), based on **dense-tensor `.stride()` semantics**.
6. For a **sparse-layout tensor**, `.stride()` does not carry the usual meaning — PyTorch returns a stride tuple whose length/values don't line up the same way as a dense tensor's `.dim()` across the repeated backward-loop iterations. The frame-state stride tuple ends up **shorter than `e.dim()`** (or otherwise mismatched) for the sparse gradient tensor's dimensionality on a later loop iteration, so `self.stride[dim]` indexes out of bounds.
This is a Dynamo/PGO bug: `is_stride_dynamic()` has no bounds-check (`if dim >= len(self.stride): return False` or similar) and, more fundamentally, PGO's automatic-dynamic-shapes stride tracking does not account for sparse tensor layouts at all — it assumes `.stride()` is always meaningful and consistently-shaped across all recorded invocations of a traced frame, which is not the case for `torch.sparse_coo` tensors.
Confirmed to be **CPU-only** and **not XPU-specific**:
- The crash happens while computing `cpu_gradInput` in `test_device()`, entirely on CPU tensors, before any GPU/XPU code path is reached.
- `PYTORCH_TEST_WITH_DYNAMO=1 PYTORCH_TEST_WITHOUT_COMPILED_AUTOGRAD=1 python test/test_nn.py TestNN.test_Embedding_sparse_xpu` **passes** — confirming compiled autograd's backward-graph tracing (not plain forward-pass Dynamo tracing) is the trigger.
- Without `PYTORCH_TEST_WITH_DYNAMO` at all, the test **passes**.
### Why this wasn't previously caught
`test/compiled_autograd_skips/` already contains a skip entry for `TestNN.test_Embedding_sparse` (the CPU-only test name, no device suffix), which correctly suppresses this exact compiled-autograd bug on CPU-only runs. However, `add_test()` in `test/test_nn.py` (post-#189653) generates a **second, device-suffixed test** (`test_Embedding_sparse_xpu`) whenever an accelerator (XPU here) is present, and that device-suffixed variant was never added to the skip list — the same class of gap as the companion `test_Conv1d_zero_batch_xpu` issue.
### Suggested fix
1. **Immediate/minimal**: add `TestNN.test_Embedding_sparse_xpu` and `TestNN.test_EmbeddingBag_sparse_xpu` to `test/compiled_autograd_skips/` in pytorch/pytorch, following the existing convention (empty marker file), mirroring the existing `TestNN.test_Embedding_sparse` and `TestNN.test_EmbeddingBag_sparse` entries. (Confirmed: `test_EmbeddingBag_sparse_xpu` fails with the identical `IndexError` and has the identical CPU-name-only skip gap.)
2. **Root fix**: harden `FrameStateSizeEntry.is_stride_dynamic()` in `torch/_dynamo/pgo.py` against out-of-range `dim` (bounds-check before indexing `self.stride[dim]`), and audit whether PGO's automatic-dynamic-shapes stride tracking should skip/special-case sparse-layout tensors entirely, since `.stride()` semantics don't apply to `torch.sparse_coo` (and other sparse layouts) the way they do to dense/strided tensors.
3. Longer term: since #189653 will surface many more `_xpu`-suffixed test names that were never exercised by CPU-only skip-list generation, re-running `scripts/compile_tests/update_failures.py` on an XPU-enabled worker (or auditing `test/compiled_autograd_skips/` for missing device-suffixed siblings of already-skipped tests, e.g. by diffing test names with/without the `_xpu` suffix) would help catch the rest of this class of gap before it lands.
### Versions
```text
PyTorch version: 2.14.0.dev20260715+xpu
Is XPU available: True
XPU used to build PyTorch: 20260000
Intel GPU models onboard: Intel(R) Data Center GPU Max 1550 (x2)
OS: Ubuntu 24.04.4 LTS (x86_64)
Python version: 3.10.20
Reproduced with: PYTORCH_TEST_WITH_DYNAMO=1, applying pytorch/pytorch#189653
```
Contributor guide
Research direction
Start with test/compiled_autograd_skips/ and compare the existing TestNN.test_Embedding_sparse and TestNN.test_EmbeddingBag_sparse entries with the generated _xpu names. Reproduce using the PYTORCH_TEST_WITH_DYNAMO=1 command in the issue, then verify both device-suffixed tests are handled consistently; torch/_dynamo/pgo.py is the entry point for investigating the longer-term crash fix.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- testing-qa, tooling
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100