KEMBAR78
AOTAutograd fails to trace with contiguous tangents when subclass has inner noncontiguous tensors · Issue #124397 · pytorch/pytorch · GitHub
Skip to content

AOTAutograd fails to trace with contiguous tangents when subclass has inner noncontiguous tensors #124397

@bdhirsh

Description

@bdhirsh

As mentioned in #124090, it is possible to end up with a tensor subclass that advertises as contiguous, but secretly has a noncontiguous tensor.

After fixing the FakeTensor problem linked there, we hit another correctness problem: today, AOTAutograd is supposed to be able to always assume that tangents feeding into the backward graph are contiguous. It does so in two steps:
(1) At trace time, when we generate FakeTensor tangents to trace the joint graph, force them to be contiguous
(2) at runtime in the backward, if any of our tangents are not contiguous, force them to be contiguous.

We do this by calling .contiguous(). However, this will incorrectly no-op if our tangent is a contiguous tensor subclass, with a secretly-noncontiguous inner tensor. We should explicitly force the inner tensors to be contiguous as well.

(Technically this will all become unnecessary once we do the work to "trace a pre-dispatch joint graph", although this is still rpetty far away)

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @anijain2305 @chauhang

Metadata

Metadata

Assignees

Labels

high prioritymodule: aotdispatchumbrella label for AOTAutograd issuesmodule: pt2-dispatcherPT2 dispatcher-related issues (e.g., aotdispatch, functionalization, faketensor, custom-op,oncall: pt2triagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions