How to Persist JIT Fused CUDA Kernels State for Efficient Inference on Multi-GPU Setup?

jonf: <p>Hi, PyTorch Community!</p>
<p>I’m currently working on a deep learning project focused on computer vision, utilizing CNNs for inference. Our setup involves a multi-GPU environment, specifically using 8 GPUs, to handle our workload. While we’ve been able to achieve the performance we aimed for, there’s a notable challenge we’re facing: the initial “warm-up” time.</p>
<p><strong>Problem Description:</strong> Our program takes approximately 50 seconds to “warm-up” every time we start it. This warm-up phase significantly impacts our overall efficiency, especially in scenarios where quick restarts are necessary. Upon investigation, we’ve identified that a substantial portion of this time is likely due to the JIT (Just-In-Time) compilation process, particularly the fusion and optimization of CUDA kernels.</p>
<p><strong>Environment Details:</strong></p>
<ul>
<li>PyTorch version: 2.2.2</li>
<li>CUDA version: 12.1</li>
<li>Number of GPUs: 8</li>
<li>GPU model: NVIDIA A10G</li>
<li>Nature of the workload: Inference using deep CNNs for computer vision tasks</li>
</ul>
<p><strong>Main Question:</strong> Is there a way to serialize and persist the state of the JIT fused CUDA kernels to disk after the first compilation? Our goal is to do this “warm-up” process once and then reuse the optimized state in subsequent runs to eliminate or significantly reduce the warm-up time.</p>
<p><strong>Sub-questions:</strong></p>
<ol>
<li>If this serialization and persistence is possible, what are the steps to achieve it?</li>
<li>How can we then reload this saved state from the disk before running our inference workload to bypass the JIT compilation process?</li>
</ol>
<p><strong>Why This Matters:</strong> Reducing or eliminating the warm-up time in our environment could improve our system’s responsiveness and overall throughput. Given the scale at which we’re operating, even small efficiency improvements can lead to significant benefits.</p>
<p>I’d appreciate any guidance or suggestions. If anyone has tackled similar challenges or knows of potential solutions, your input would be appreciated.</p>
<p>Thank you in advance!</p>

ptrblck: <p>Are you using the deprecated <code>torch.jit.script</code> or the newer <code>torch.compile</code> approach?<br>
The latter should have use a cache, if I’m not mistaken, but adding <a class="mention" href="/u/marksaroufim">@marksaroufim</a> just in case I’m missing something.</p>

marksaroufim: <p>In the <code>torch.compile()</code> world you can make warm startup times much faster  with <code>torch._inductor.config.fx_graph_cache = True</code>. That said warm startup times in torch.compile will be significantly improved soon <a href="https://dev-discuss.pytorch.org/t/pt2-core-h1-2024-public-roadmap/1943" class="inline-onebox">PT2 Core - H1 2024 Public Roadmap - PyTorch Dev Discussions</a></p>
<p>And if waiting is not an option you can instead try out <a href="https://pytorch.org/docs/main/torch.compiler_aot_inductor.html" class="inline-onebox">AOTInductor: Ahead-Of-Time Compilation for Torch.Export-ed Models — PyTorch main documentation</a></p>

jonf: <p>Awesome, <a class="mention" href="/u/ptrblck">@ptrblck</a> and <a class="mention" href="/u/marksaroufim">@marksaroufim</a>!</p>
<p>So, we are currently using <code>torch.jit.script</code> to generate TorchScript and running inference via LibTorch/C++, but thanks to you I’ve just realized that we should switch over to <code>torch.compile()</code> to leverage the AOTInductor introduced in PyTorch 2.2. It seems to be precisely what we need. We’ll give it a try.</p>

jonf: <p>So, I’ve just encountered the same issue as described in <a href="https://github.com/pytorch/pytorch/issues/122313" rel="noopener nofollow ugc">Missing Symbols When running AOTInductor example with Libtorch c++11 ABI</a>. I suspect <code>-D_GLIBCXX_USE_CXX11_ABI=1</code> needs to be passed to the AOT compiler. Is there any way to supply additional compiler options manually?</p>

jonf: <p>Setting <code>aot_inductor.abi_compatible</code> to <code>True</code> resolves the problem in our environment. For example:</p>
<pre><code class="lang-auto">    so_path = torch._export.aot_compile(
        model,
        example_inputs,
        # Specify the first dimension of the input x as dynamic
        dynamic_shapes={"x": {0: batch_dim}},
        # Specify the generated shared library path
        options={"aot_inductor.output_path": os.path.join(os.getcwd(), "model.so"), "aot_inductor.abi_compatible": True},
    )
</code></pre>
<p>I have also added a comment to the GitHub issue.</p>