🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/132926

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (3 Unrelated Failures)

As of commit 440d1bc with merge base b7a5c7d ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Is this backwards compatible?

No, this will be BC-breaking - see discussion in #123771.

i.e. there are users that depend on the old compile/disable interaction.

Hi, sorry, I'm going to thrash you again.

I was reviewing the docs you added and proposing a rewrite of the docs, and while doing so, I realized that I actually think disable as it is used today is different from the "disable" that Horace was proposing in the issue, and we should actually give it a different name. Horace and I came up with "torch.compiler.bypass", but Claude also suggested "force_uncompiled" and "prevent_compilation" which also seemed good. Here are my proposed docs for disable vs these other options.

What is the difference between torch.compiler.disable and torch.compiler.force_uncompiled? It's easiest to consider the use cases for these two functions.

torch.compiler.disable is for when you want to prevent an outer compilation from continuing compiling into an inner function. You typically want to do this if you have some code that is reachable from a torch.compile that triggers a problem in the compiler (e.g., it causes infinite recompilations) and you just want to disable compilation for it. It comes in two flavors: the default is to disable compilation for that function and all recursive calls (unless you explicitly call torch.compile again internally), and recursive=False, where it will only disable compilation for that particular function call, but will attempt to continue compilation for inner function calls.

torch.compiler.force_uncompiled is for when you are in an outer function, and you want to prevent inner compilations from occurring. It's useful if you have some code that has torch.compile directly annotated on functions, but you want to temporarily disable the compiler / any preexisting compiled code so that you can run them under eager semantics, e.g., for benchmarking or if you want to run a TorchDispatchMode (like FlopCounter) which only works in eager semantics. If necessary, you can also undo the effect of this context manager by saying force_uncompiled(False). force_uncompiled is forbidden in graph; you're not allowed to compile code that has a force_uncompiled invocation in it.

I don't think you need to change much in your PR. Here's what changes:

1. The PR is no longer BC breaking. We don't need a JK, we don't need to annotate every HOP with enable. There is a funny interaction with HOPs and FlopCounterMode but we can deal with it when we get there. force_uncompiled won't work too well with things that require mandatory compilation, so we could still have a force_uncompiled(False) decorator on the content store helper.
2. force_uncompiled can be a context manager, because it doesn't rely on setting up an eval frame handler to work, unlike torch.compile. Consequently, it's not necessary to have OptimizedModule logic for it; just do the context manager syntax instead if you need to work with a module.
3. You need to rename your disable implementation, but all the other pieces stay the same. In particular, this controlling TLS that is accessed by eval frame handler still seems like the right implementation strategy to me.

WDYT?

We talked about this today in more general at https://docs.google.com/document/d/1qlQyjbENz45L7uGglmiF6CkRgr_VOdxFAcZGCqdtLTY/edit

SGTM from an implementation perspective, although I'm concerned about sufficiently differentiating between disable and force_uncompiled. It's still not clear to me from a user perspective as to when I should use one over the other.

A few other clarification questions:

• are you suggesting that force_uncompiled should only be a context manager? The function/decorator version is very close in effect to disable, but has small but important differences.
• "force_uncompiled is forbidden in graph" - do you mean we should skip the frame if we encounter the context manager while tracing? Or graph break if we encounter?

SGTM from an implementation perspective, although I'm concerned about sufficiently differentiating between disable and force_uncompiled. It's still not clear to me from a user perspective as to when I should use one over the other.

It was pretty clear to me when I was writing the docs, so I guess the new docs are not good enough 🤣

Here's a shorter pitch: if you are currently in a torch.compile, but you want to get out, use torch.compiler.disable. If you are not in a torch.compile, but you don't want to compile for any of your nested function calls, use torch.compiler.set_phase("force_eager") (my new name for the API LOL).

are you suggesting that force_uncompiled should only be a context manager? The function/decorator version is very close in effect to disable, but has small but important differences.

I don't mind a decorator but it should be a contextlib.contextmanager style decorator. I don't think set_phase needs to support OptimizedModule.

"force_uncompiled is forbidden in graph" - do you mean we should skip the frame if we encounter the context manager while tracing? Or graph break if we encounter?

"forbidden in graph" is a preexisting concept which makes the compiler hard error if it traverses it while tracing. It is saying "this code should never be compiled and if you tried to compile it some big accident happened so raise a major error."

Difference between disable and force_eager makes more sense now.

What other options are you thinking for set_phase? As of now, it looks like there should be at least a force_eager and a default option?

The ones we have right now are: "force_eager", "default", "fail_on_recompile", "eager_on_recompile"

And I will happily bikeshed this forever if needed

To add more color from previous discussion with Ed and others, summary of motivations: https://docs.google.com/document/d/1qlQyjbENz45L7uGglmiF6CkRgr_VOdxFAcZGCqdtLTY/edit#bookmark=id.nl59a72dv4o. Particularly torch.compiler.set_phase("force_eager") will be super useful for Traceable FSDP / Compiled DDP / float8 / optimizer warmup phase, and the recompile-related options will also be helpful.

Moved to a new PR: #137504