[Bug]: Dynamo is activating profiler multiple times while native does it only once

[anijain2305]

🐛 Describe the bug

PR - pytorch/pytorch#87311

Command - python benchmarks/dynamo/torchbench.py --float32 --inductor --training --performance --only=resnet18

There are lots of profiler activity for Dynamo, but not native pytorch. This could impact speedpup

Error logs

STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:294] Completed Stage: Warm Up
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:300] Completed Stage: Collection
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:294] Completed Stage: Warm Up
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:300] Completed Stage: Collection
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:294] Completed Stage: Warm Up
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:300] Completed Stage: Collection
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:294] Completed Stage: Warm Up
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:300] Completed Stage: Collection
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:294] Completed Stage: Warm Up
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:300] Completed Stage: Collection
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:294] Completed Stage: Warm Up
STAGE:2022-10-19 18:38:51 82265:82265 ActivityProfilerController.cpp:300] Completed Stage: Collection

Did Dynamo succeed?

• Does dynamo.optimize("eager") succeed?

Did AOT succeed?

• Did dynamo.optimize("aot_eager") succeed?

Did Inductor succeed?

• Does dynamo.optimize("inductor") succeed?

Minified repro

N/A

[yanboliang]

This root cause is torch.profiler.profile.__init__ starts ActivityProfilerController at its construction. TorchDynamo generates the torch.profiler.profile object during tracing, which leads to the unintended call. But actually we don't use the profile, so the ActivityProfilerController exits just after the generation. I think we can fix it by not generating the torch.profiler.profile object, as dynamo ignore all functions under torch.profiler. I'll send a fix soon.

[anijain2305]

Repro - https://www.internalfb.com/phabricator/paste/view/P550220100

cc @yanboliang

[yanboliang]

This works well in most cases, except for some fallback scenarios. Let me explain use the following example:

import torch
import torch._dynamo
import logging
from torch.autograd.profiler import record_function

# torch._dynamo.config.log_level = logging.DEBUG
# torch._dynamo.config.verbose = True
# torch._dynamo.config.cache_size_limit = 1

def fn(x, y):
    a = x**2
    with torch.autograd.profiler.record_function("my_function1"):
        a = a + 2
        torch._dynamo.graph_break()
        b = a + y
    return b

x = torch.ones((2, 2), requires_grad=True)

opt_fn = torch._dynamo.optimize("eager")(fn)
for i in range(123, 125):
    res = opt_fn(x, i)
    print(res)

There is a graph break inside of the context manager, so there is a WithExitFunctionVariable in the stack and we need to reconstruct it after the graph break. During the reconstruction, actually we call autograd.profiler.profile() which is fine, as we trace it and wrap it into the ProfilerContextWrapperVariable at here. We don't execute this function call, it works well so far.
However, in some scenarios(e.g, hitting cache limit), we fallback to original python code, but the python code is a resume execution not the original user's code. And at the resume execution, we set up a context manager with autograd.profiler.profile().
This is why we step into the autograd.profiler.profile() unexpected. And this behavior is flaky, it's only triggered in some scenarios, not always. In my example, if you change the cache_size_limit to 1 and re-run, this bug will be triggered.
The fix is to have a dummy context manager and do nothing, so even dynamo fallback to python code, it still do nothing as expected.

[yanboliang]

Fixed by pytorch/pytorch#88930