intel / intel/torch-xpu-ops

XPU eagerly initializes the device under `FakeTensorMode` tracing, aborting when no XPU device is present

Open
#5,157 0 comments 0 reactions 1 assignee Claimed by @BBBela View on GitHub
Dominant language
Python
Stars
113
Forks
129
Avg merge
5d 9h
Merged PRs (30d)
112

Description

### 🐛 Describe the bug

## Summary

Under `FakeTensorMode` (e.g. during `torch.export.export`), tracing certain ops on a
fake `xpu` tensor eagerly constructs a **real** XPU device guard and calls
`c10::xpu::set_device(0)`. When no XPU device is visible (0 devices), this fails
`check_device_index` and throws a `c10::Error` that escapes through a
`noexcept`/destructor boundary, so the process **aborts via `std::terminate`
(SIGABRT)** instead of raising a catchable Python exception.

This breaks the "fake export onto an accelerator you don't physically have"
workflow that works on CUDA. The whole point of fake export is to be
device-agnostic and not touch a physical device, so tracing should never call
`set_device` on a real XPU device.

## Reproducer

Run on a PyTorch build with XPU support, with the XPU device hidden to simulate a
machine that has the build but no usable device:

```bash
ONEAPI_DEVICE_SELECTOR="*:cpu" python repro.py
```

```python
# repro.py
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.utils import _pytree as pytree

cpu_x = torch.randn(5)
mode = FakeTensorMode(allow_non_fake_inputs=True)
with mode:
(args,) = pytree.tree_map_only(torch.Tensor, lambda x: x.to("xpu:0"), (cpu_x,))

class M(torch.nn.Module):
def forward(self, x):
return torch.nonzero(x)

ep = torch.export.export(M(), (args,)) # <-- aborts here
print("exported ok")
```

## Actual behavior

```
terminate called after throwing an instance of 'c10::Error'
what(): The device index is out of range. It must be in [0, 0), but got 0.
Exception raised from check_device_index at c10/xpu/XPUFunctions.h:38
...
c10::xpu::set_device(signed char)
...
```

Process exits with 134 (SIGABRT). Because it is a C++ `terminate`, a surrounding
Python `try/except` cannot catch it.

## Expected behavior

Tracing a fake `xpu` tensor under `FakeTensorMode`/export should not touch a
physical device (no `set_device`), exactly like CUDA — the same reproducer with
`cuda:0` + `CUDA_VISIBLE_DEVICES=""` succeeds. At minimum, an out-of-range device
index must raise a catchable `RuntimeError`, never `std::terminate`.

## Notes / scope

- It is **not** the fake `.to("xpu:0")` conversion — that stays symbolic and
succeeds. The abort happens later, during export **tracing** of the op.
- Op-specific: a trivial pointwise op (`x + x`) traces fine. Observed aborts with:
`nonzero`, `__getitem__`, `nn.functional.batch_norm`,
`nn.functional.instance_norm`, `nn.functional.multi_margin_loss`,
`nn.functional.scaled_dot_product_attention`.
- CUDA does not hit this (its device guard is NoOp'd / tolerant under fake mode),
so this is an XPU-vs-CUDA gap. Compare `FakeTensorMode.__enter__`, which calls
`torch._C._ensureCUDADeviceGuardSet()` only when `avoid_device_init` is set;
there is no XPU equivalent, and on an XPU build with a device present
`avoid_device_init` is `False`.

## Possible fix directions

- Provide an XPU analog of the NoOp device-guard patching used for CUDA under
`FakeTensorMode` (so fake/meta tracing does not call `c10::xpu::set_device`).
- Ensure `c10::xpu::set_device` / device-guard errors surface as catchable
`RuntimeError` rather than escaping to `std::terminate`.

### Versions

- torch: `2.15.0a0+git4c66a66` (built from source)
- XPU compiled: `True`, CUDA built: `False`
- Repro requires the XPU device hidden via `ONEAPI_DEVICE_SELECTOR="*:cpu"`
(or any config with 0 visible XPU devices).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.