Issue Station
Sign In
Tools hereToolbox

pytorch - ✅(Solved) Fix FakeTensor.avoid_device_init not respecting privateuse1 backend [1 pull requests, 1 participants]

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…
GitHub stats
pytorch/pytorch#178019•Fetched 2026-04-08 01:07:38
View on GitHub
Comments
0
Participants
1
Timeline
84
Reactions
0
Author
arfaianarfaian
Participants
arfaianarfaian
Timeline (top)
mentioned ×35subscribed ×35labeled ×8referenced ×3

Error Message

Error logs

Fix Action

Fixed

PR fix notes

PR #178020: [torch.compile] Enable initializing embedded constants on privateuse1…

Description (problem / solution / changelog)

… backend if registered

Fixes pytorch/pytorch#178019

Changed files

  • torch/_C/__init__.pyi.in (modified, +1/-0)
  • torch/_subclasses/fake_tensor.py (modified, +1/-0)
  • torch/csrc/Module.cpp (modified, +13/-0)

Code Example

def simple(x):
  c = torch.tensor([1, 2, 3, 4, 5], device=x.device)
  return x + c

x = torch.tensor([1, 2, 3, 4, 5], device=my_custom_device)
model = torch.compile(simple, backend=my_custom_backend)
out = model(x)

---

def forward(self, arg0_1: "i64[5][1]my_custom_device:0"):
    _tensor_constant0: "i64[5][1]cpu" = self._tensor_constant0
    lift_fresh_copy: "i64[5][1]cpu" = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0);  _tensor_constant0 = None
    _to_copy: "i64[5][1]my_custom_device:0" = torch.ops.aten._to_copy.default(lift_fresh_copy, device = device(type='my_custom_device', index=0));  lift_fresh_copy = None
    add: "i64[5][1]my_custom_device:0" = torch.ops.aten.add.Tensor(arg0_1, _to_copy);  arg0_1 = _to_copy = None
    return (add,)

---

def forward(self, arg0_1: "i64[5][1]cuda:0"):
    _tensor_constant0: "i64[5][1]cuda:0" = self._tensor_constant0
    lift_fresh_copy: "i64[5][1]cuda:0" = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0);  _tensor_constant0 = None
    add: "i64[5][1]cuda:0" = torch.ops.aten.add.Tensor(lift_fresh_copy, arg0_1);  lift_fresh_copy = arg0_1 = None
    return (add,)
RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

I’m running into an issue where embedded constants are getting placed on the CPU, even though I’ve specified a custom device. Consider the following example:

def simple(x):
  c = torch.tensor([1, 2, 3, 4, 5], device=x.device)
  return x + c

x = torch.tensor([1, 2, 3, 4, 5], device=my_custom_device)
model = torch.compile(simple, backend=my_custom_backend)
out = model(x)

In this example, my_custom_backend wraps aot_autograd. This is the post aot_autograd graph I get:

def forward(self, arg0_1: "i64[5][1]my_custom_device:0"):
    _tensor_constant0: "i64[5][1]cpu" = self._tensor_constant0
    lift_fresh_copy: "i64[5][1]cpu" = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0);  _tensor_constant0 = None
    _to_copy: "i64[5][1]my_custom_device:0" = torch.ops.aten._to_copy.default(lift_fresh_copy, device = device(type='my_custom_device', index=0));  lift_fresh_copy = None
    add: "i64[5][1]my_custom_device:0" = torch.ops.aten.add.Tensor(arg0_1, _to_copy);  arg0_1 = _to_copy = None
    return (add,)

My expectation is that _tensor_constant0 should be initialized on my_custom_device:0 instead (which will also remove the additional _to_copy). Running this same example with cuda/inductor shows that the constant is correctly placed on the cuda device:

def forward(self, arg0_1: "i64[5][1]cuda:0"):
    _tensor_constant0: "i64[5][1]cuda:0" = self._tensor_constant0
    lift_fresh_copy: "i64[5][1]cuda:0" = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0);  _tensor_constant0 = None
    add: "i64[5][1]cuda:0" = torch.ops.aten.add.Tensor(lift_fresh_copy, arg0_1);  lift_fresh_copy = arg0_1 = None
    return (add,)

I’ve narrowed down the culprit to constant lifting logic, specifically in FakeTensor there’s a method called avoid_device_init that determines whether to init on cpu or on a device but out-of-tree backends/privateuse1 are not respected: https://github.com/pytorch/pytorch/blob/main/torch/_subclasses/fake_tensor.py#L1431-L1440

Error logs

No response

Versions

Excluding environment information for security reasons.

cc @bdhirsh @chauhang @penguinwu @eellison @aorenste @bobrenjc93 @NmomoN @mengpenghui @fwenguang @cdzhan @1274085042 @PHLens @albanD

extent analysis

Fix Plan

To fix the issue of embedded constants being placed on the CPU instead of the custom device, we need to modify the avoid_device_init method in FakeTensor to respect out-of-tree backends and private use cases.

Here are the steps:

  • Modify the avoid_device_init method to check for custom devices and backends.
  • If a custom device is detected, initialize the constant on that device instead of the CPU.

Example code:

def avoid_device_init(self, device):
    # Check if the device is a custom device
    if device.type == 'my_custom_device':
        # Initialize the constant on the custom device
        return False
    # Default behavior for other devices
    return True

Additionally, you may need to modify the torch.compile function to pass the custom device to the avoid_device_init method.

Verification

To verify that the fix worked, you can run the same example code and check the generated graph. The constant should now be initialized on the custom device instead of the CPU.

Example verification code:

x = torch.tensor([1, 2, 3, 4, 5], device=my_custom_device)
model = torch.compile(simple, backend=my_custom_backend)
out = model(x)
print(model.graph)

The output graph should show that the constant is initialized on the custom device.

Extra Tips

  • Make sure to test the fix with different custom devices and backends to ensure that it works correctly in all cases.
  • Consider submitting a pull request to the PyTorch repository to include the fix in the main codebase.

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…
#api#ssr#installation#tensor shape#autograd error#container setup#orchestration issue#cache issue#memory leak#API versioning

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING

pytorch - ✅(Solved) Fix FakeTensor.avoid_device_init not respecting privateuse1 backend [1 pull requests, 1 participants]