[dynamo] Turn on guard_nn_modules

[anijain2305]

Stack from ghstack (oldest at bottom):

• [dynamo] Add ID_MATCH guards on inlined functions to force compilation on monkeypatching #124975
• -> [dynamo] Turn on guard_nn_modules #125202

Turning on guard_nn_modules adds large number of guards, so we are bound to take a perf hit. But the perf hit is small. These are the numbers

First we observe that compared to Python guards, C++ guards give around 6x speedup. This reduces the total time spent in guards. This is shown in the last column (cpp_guards/inductor_optimized_latency). The worst model is around 1.61%, with most of the models below 1%. I think this is good enough signal to turn the config on.

One might also wonder how much guard slowdown occurs with guard_nn_modules=True. This is the table

For most models, the guard overhead with nn module guards is under 2x. There are a few outliers, where the slowdown is really high and for those models we spend 1%-2% time in C++ guards as shown in first table.

cc @ezyang @msaroufim @bdhirsh @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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! (2 Unrelated Failures)

As of commit bca250a with merge base ae5e2ab ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  test_modules.py::TestModuleCUDA::test_memory_format_nn_BatchNorm2d_eval_mode_cuda_float64
• inductor / rocm6.1-py3.8-inductor / test (inductor, 1, 1, linux.rocm.gpu.2) (gh) (similar failure)
  test_modules.py::TestModuleCUDA::test_memory_format_nn_BatchNorm2d_eval_mode_cuda_float64

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

[ezyang]

You want a JK for internal rollout to killswitch it

[ezyang]

🤔 which we don't have precedent for in this module

[ezyang]

You want to call torch._utils_internal.justknobs_check to read out the default, but you can't do the obvious thing of doing it in the config module as that will cause the JK check to happen at module import time and @oulgen and I know that you can't do that because it will poison the process for forks. So you want this to happen on the first time the config is accessed. The low tech way is to default this to None, and then at the read site, if it is None, pong the JK for the real value. The high tech way is to add some capability to the config module getter (I think we've got a accessor function where you can customize) so that it is able to lazily query JK.

[ezyang]

Another deployment strategy that doesn't involve JKs is to turn it on in OSS only but not fbcode, then enable it on a per PG basis by twiddling it, and then later switch the fbcode default to true.

[anijain2305]

Going with fbcode off and OSS on for now. Will figure out internal rollout strategy.

[anijain2305]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team

Raised by workflow job

[anijain2305]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here