🔗 Helpful links

• 🧪  See artifacts and rendered test results at hud.pytorch.org/pr/73925
• 📄  Preview docs built from this PR
• 📄  Preview C++ docs built from this PR
• ↩️  [fb-only] Re-run with SSH instructions
• 🔧  Opt-in to CIFlow to control what jobs run on your PRs

💊 CI failures summary and remediations

As of commit 3477503 (more details on the Dr. CI page):

• 1/1 failures introduced in this PR

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

linux-xenial-py3.7-gcc5.4 / test (backwards_compat, 1, 1, linux.2xlarge) (1/1)

Step: "Test" (full log | diagnosis details | 🔁 rerun)

2022-03-11T19:23:37.3226402Z Author: PyTorch Team
2022-03-11T19:23:37.3226648Z Author-email: packages@pytorch.org
2022-03-11T19:23:37.3226910Z License: BSD-3
2022-03-11T19:23:37.3227172Z Location: /opt/conda/lib/python3.7/site-packages
2022-03-11T19:23:37.3227422Z Requires: typing-extensions
2022-03-11T19:23:37.3227629Z Required-by: 
2022-03-11T19:23:37.3540424Z + python check_forward_backward_compatibility.py --existing-schemas nightly_schemas.txt
2022-03-11T19:23:37.9640010Z Traceback (most recent call last):
2022-03-11T19:23:37.9640525Z   File "check_forward_backward_compatibility.py", line 308, in <module>
2022-03-11T19:23:37.9640947Z     s = parse_schema(line.strip())
2022-03-11T19:23:37.9641135Z RuntimeError: 
2022-03-11T19:23:37.9641379Z Unknown custom class type profiler._RecordFunction. Please ensure it is registered.:
2022-03-11T19:23:37.9641956Z profiler::_record_function_exit._RecordFunction(__torch__.torch.classes.profiler._RecordFunction _0) -> ()
2022-03-11T19:23:37.9642344Z                                                                                  ~~~~~~~~~~~~~~~ <--- HERE
2022-03-11T19:23:37.9642476Z 
2022-03-11T19:23:38.0400911Z + cleanup
2022-03-11T19:23:38.0401143Z + retcode=1
2022-03-11T19:23:38.0401290Z + set +x
2022-03-11T19:23:38.0442969Z ##[error]Process completed with exit code 1.
2022-03-11T19:23:38.0472466Z ##[group]Run # Ensure the working directory gets chowned back to the current user
2022-03-11T19:23:38.0472784Z �[36;1m# Ensure the working directory gets chowned back to the current user�[0m

Also TLSState = Thread Local State State? :P

FYI @ezyang as well as you're thinking on how to replace this system altogether.

This is ready for a final review

@albanD has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Hey @albanD.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

@albanD @soulitzer this PR broke functorch: https://app.circleci.com/pipelines/github/pytorch/functorch/2032/workflows/2a0f2124-3870-4c6e-b95b-00951c1e9332/jobs/12098 (I confirmed by bisection to this PR that it is the one that causes the failure)

I don't know the exact reason, but I suspect it has to do with the following:

1. functorch's DynamicLayerFrontMode is active inside of vmap.
2. PythonTLS saves DynamicLayerFrontMode as a part of the "saved TLS".
3. Somehow this goes wrong.

This doesn't seem easy to resolve. Could you help take a look please?

Minimal repro:

import torch
import functorch
from functorch import vmap, make_fx

def f(x):
    return torch.sin(x)
inp = torch.randn(5, 3)
f = vmap(f)
fx_f = make_fx(f)(inp)
new_inp = torch.randn(5, 3)

Hmmm not sure why this PR would change this behavior as this was happening before?
But yes, functorch breaks some of the assumptions we make here about how the dispatcher is used by jumping at specific keys. So I am not very surprised :/

It might not look this way because of our lack of test coverage for this, but specifically this PR seems to break the torch_dispatch interaction with functorch. The main consequence of this right now is that we cannot compose AOTAutograd with functorch eager-mode transforms (e.g. vmap, grad, etc), so this is a really large breakage for us.

It would be a huge favor to us if we could figure out how to resolve this soon.

Here's my current understanding to what is happening.

The purpose of this PR is to save the TLS when going from Python -> into the dispatcher so that when __torch_dispatch__ is invoked, the TLS may be set to the saved one.

In the example,

f = vmap(torch.sin)
fx_f = make_fx(f)(inp)

we are:

• creating a ProxyTensor object, then calling vmap(f)(ProxyTensor)
• vmap adds DynamicLayerFrontMode to the dispatch local include key set.
• DynamicLayerFrontMode gets removed from the local include set after the batching rule for sin gets called.
• ProxyTensor.sin re-enables DynamicLayerFrontMode, which is a problem because we've already handled vmap.

So I think the question now is how should the TLS logic in this PR interact with mode-style dispatch keys?

• Let's say we have a HypotheticalMode dispatch key, and it is sandwiched between [PythonTLS, HypotheticalMode, Python].
• Furthermore, let's say that HypotheticalMode::sin (1) prints "sin" (2) temporarily removes HypotheticalMode from the local include set (in other words, it is "done"), and (3) dispatches down the line (either call or redispatch).

Imagine some Python code that looks like:

subclass = TensorSubclass(x)
add_to_tls_local_include_set(HypotheticalMode)
subclass.sin()

This would end up running through HypotheticalMode twice. Would that be the desired behavior?

Would that be the desired behavior?

Yes. The high level idea of this TLS logic is the following: if you call torch.foo(*args) from top python or from within the torch_dispatch of torch.foo for one of the args. You should get the exact same result.

So in this case, that looks like the expected behavior to me? (not that we want to keep it necessarily but at least it does follow the high level idea)