Performance of `torch.compile` is significantly slowed down under `torch.inference_mode`

Crazyai: <p>I have observed that when using <code>torch.compile</code> to optimize a model, the performance significantly degrades during inference under <code>torch.inference_mode</code> . In fact, it is even worse than the performance of the non-optimized model. However, when I place the compilation process within the context of <code>torch.inference_mode</code> , the performance issues are resolved.</p>
<p>My resultI tested in <code>torch==2.1.0+cu118</code> on an A10 GPU:</p>
<pre><code class="lang-auto"># 1. both are disabled, speed up 2.5X, it's great!
enable inference mode when compile: False
enable inference mode when benchmark: False
not compiled latency: 0.002047449684143066
compiled latency: 0.0008150367736816407, speed up: 2.5120948529652662

# 2.  both are enabled, speed up 2.35X, great work!
enable inference mode when compile: True
enable inference mode when benchmark: True
not compiled latency: 0.0019193056106567383
compiled latency: 0.0008156064033508301, speed up: 2.3532252846121366

# 3. enable torch.inference_mode when benchmark, disabled when compiling
# the compiled model is slower than naive torch model ! help me!
# and the naive model's latency is still ~0.002
enable inference mode when compile: False
enable inference mode when benchmark: True
not compiled latency: 0.0019556352615356445
compiled latency: 0.010692658996582031, speed up: 0.18289513040309005

# 4. disable in benchmark and enable in compiling, the result is also disappointed
enable inference mode when compile: True
enable inference mode when benchmark: False
not compiled latency: 0.0020351423263549806
compiled latency: 0.011507730865478516, speed up: 0.17685001067066186
</code></pre>
<p>And here is the code to reproduce my result,</p>
<pre><code class="lang-auto">import torch
from torchvision.models import resnet18
import sys
from contextlib import nullcontext

# enable inference mode when compile or benchmark?
enable_when_compile = sys.argv[1] == "true"
enable_when_benchmark = sys.argv[2] == "true"


def _benchmark(
    iters,
    f,
    context,   # with torch.inference or not
    *args,
    **kwargs,
) -&gt; float:
    """Estimates the average time duration for a single inference call in second
    Returns:
        estimated average time duration in second for a single inference call
    """
    with context():
        f(*args, **kwargs)
    torch.cuda.synchronize()
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    with context():
        start_event.record()
        for _ in range(iters):
            f(*args, **kwargs)
        end_event.record()
    torch.cuda.synchronize()
    elapsed_time_s = start_event.elapsed_time(end_event) * 1.0e-3
    avg_time_s = elapsed_time_s / iters
    print("Estimated average time duration: {:.6f} s".format(avg_time_s))
    return avg_time_s


class BenchmarkRunner(object):
    def __init__(self, use_inference_mode: bool):
        self.context = nullcontext if not use_inference_mode else torch.inference_mode

    def __call__(self, iters, f, *args, **kwargs) -&gt; float:
        return _benchmark(iters, f, self.context, *args, **kwargs)


@torch.no_grad()
def run():
    input = [torch.rand(8, 3, 224, 224).to(torch.device("cuda"), dtype=torch.float16)]
    net = resnet18(pretrained=False).cuda().half()
    net.eval()
    context = nullcontext if not enable_when_compile else torch.inference_mode
    compiled = torch.compile(net, mode="reduce-overhead", backend="inductor")
    with context():
        _ = compiled(*input)

    latency_compiled = BenchmarkRunner(enable_when_benchmark)(10, compiled, *input)
    latency_torch = BenchmarkRunner(enable_when_benchmark)(10, net, *input)

    print(f"enable inference mode when compile: {enable_when_compile}")
    print(f"enable inference mode when benchmark: {enable_when_benchmark}")
    print(f"not compiled latency: {latency_torch}")
    print(f"compiled latency: {latency_compiled}, speed up: {latency_torch / latency_compiled}")

if __name__ == "__main__":
    run()

</code></pre>

bdhirsh: <p>This seems like a bug - <a class="mention" href="/u/crazyai">@Crazyai</a> can you file a github issue with your repro instructions? <a href="https://github.com/pytorch/pytorch/issues/new/choose" class="inline-onebox" rel="noopener nofollow ugc">Sign in to GitHub · GitHub</a></p>
<p>As a workaround - you can replace <code>torch.inference_mode</code> with <code>torch.no_grad</code>. For some context:</p>
<p>(1) I tried your repro with <code>torch.no_grad()</code> and I see the expected speedups with <code>torch.compile</code>:</p>
<pre><code class="lang-auto">enable inference mode when compile: True
enable inference mode when benchmark: False
not compiled latency: 0.013454794311523438
compiled latency: 0.0005569119930267334, speed up: 24.159641882371112
...
enable inference mode when compile: False
enable inference mode when benchmark: True
not compiled latency: 0.012717078399658202
compiled latency: 0.0005573919773101807, speed up: 22.81532371712148
</code></pre>
<p>(2) It’s worth calling out what the difference between <code>inference_mode</code> and <code>no_grad</code>, which is a bit subtle.</p>
<p>tldr: <code>inference_mode</code> has the additional benefit of removing some cpu overhead in eager mode. When using <code>torch.compile</code>, the difference between <code>inference_mode</code> and <code>no_grad</code> is <strong>more negligible</strong> - <code>torch.compile</code> itself is already doing compilation to avoid cpu overhead as much as possible.</p>
<p>In the end state, with <code>torch.compile</code>, <code>no_grad</code> and <code>inference_mode</code> should probably be equivalent (both in terms of what they do and what their perf implications are). But if there are any bugs with <code>torch.compile</code> and <code>inference_mode</code>, we should fix them.</p>
<p>In eager mode, <code>no_grad</code> can be used to disable gradient calculation (and for example, avoid saving any activations for the backward if you only care about inference). However, even when you run no_grad, there is still some autograd-related overhead that occurs in eager-mode: autograd will perform view_tracking and version_counter tracking, which has some cpu-side overhead associated with it.</p>
<p>When you torch.compile a model, the resulting generated code that we execute at runtime is (90% of the time) generated cpp code, or generated triton code. None of this code is seen by autograd, so there is no view_tracking and version_counter tracking that we need to be disabled. (The small caveat is that inductor can sometimes generate calls to pytorch operations, like <code>aten.as_strided</code> and <code>aten.addmm</code>, which will have a <em>tiny</em> bit of autograd overhead, that can go away if you’re using inference mode. This is much smaller though, than the amount of overhead you’ll see in eager mode, and is probably not noticeable).</p>

Crazyai: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a> Thanks for the reply. I just posted an issue on Github <a href="https://github.com/pytorch/pytorch/issues/114119" rel="noopener nofollow ugc">see it here</a>. Glad to contribute to PyTorch.</p>

Crazyai: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a> I am curious why the speed up ratio of your code is more than 20X… What is the hardware you tested on?</p>

Crazyai: <p>And when I use <code>torch.no_grad</code> instead of <code>torch.inference_mode</code>, the results are still confusing:</p>
<pre><code class="lang-auto">enable inference mode when compile: True
enable inference mode when benchmark: False
not compiled latency: 0.0024103456497192384
compiled latency: 0.011797164916992188, speed up: 0.20431566962732445
</code></pre>
<p>The diff:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/6/d65819241a215e5606721d6179a38d960e0ef159.png" data-download-href="https://discuss.pytorch.org/uploads/default/d65819241a215e5606721d6179a38d960e0ef159" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d65819241a215e5606721d6179a38d960e0ef159_2_690x368.png" alt="image" data-base62-sha1="uAaVy3tdkMWrfjbjDPh7ISSnOd3" width="690" height="368" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d65819241a215e5606721d6179a38d960e0ef159_2_690x368.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d65819241a215e5606721d6179a38d960e0ef159_2_1035x552.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d65819241a215e5606721d6179a38d960e0ef159_2_1380x736.png 2x" data-dominant-color="150F10"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1912×1022 106 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

bdhirsh: <p>Strange, I’m benchmarking on an A100 using a recent nightly.</p>
<p>(1) Can you try running your benchmark (with <code>torch.no_grad</code>) on either a recent nightly to see if you still don’t see speedups?</p>
<p>(2) Most of the speedup here is probably coming from torch.compile using cudagraphs to remove most of the python overhead. If you use a nightly and run with <code>mode="reduce-overhead"</code>, torch.compile should give a warning if it tried and failed to use cudagraphs, and also explain why.</p>
<p>(3) You can also try running with the pytorch profiler to see where the slowdown is coming from (if you’re seeing a 5x slowdown then the profiler will probably point to something pretty obvious if you can get a profile): <a href="https://pytorch.org/docs/stable/profiler.html" class="inline-onebox" rel="noopener nofollow ugc">torch.profiler — PyTorch 2.1 documentation</a></p>

Crazyai: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a> A quick test after updated to torch2.1.1-cuda121, the speed up with <code>torch.no_grad</code> seems ok. But the speed up with <code>torch.inference_mode</code> is not fixed.</p>