Support of Dynamic and Static Shape in One Module

username: chenghua_wang
cooked: <p>I’m trying to use ExecuTorch to deploy a LLM with a QNN backend. Unfortunately, QNN doesn’t support dynamic shapes, so some trick is needed to utilise up the NPU. A simple idea is shown below:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/7/279bb8538481273f926bbe97d9f7b5e05918a696.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/279bb8538481273f926bbe97d9f7b5e05918a696" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/279bb8538481273f926bbe97d9f7b5e05918a696_2_651x500.jpeg" alt="image" data-base62-sha1="5EofuQIEtEd60WVPXG0evSGLKfk" width="651" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/279bb8538481273f926bbe97d9f7b5e05918a696_2_651x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/279bb8538481273f926bbe97d9f7b5e05918a696_2_976x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/279bb8538481273f926bbe97d9f7b5e05918a696_2_1302x1000.jpeg 2x" data-dominant-color="F0F1F1"><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">1920×1473 167 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>
<p>Divide the Prompt into fixed-size Chunks, and these fixed Chunks first enter the Static Shape QKV Proj (run using QNN-NPU), then proceed to the Dynamic Shape Attention (run using CPU). Because of the chunking, the computation of adjacent Tokens can be pipelined.</p>
<p>I am not familiar with ExecuTorch. In the examples of Dynamic Shape provided by ExecuTorch, it seems that dynamic shapes are specified within the <code>export</code> function. such as:</p>
<pre data-code-wrap="python"><code class="lang-python">aten_dialect: ExportedProgram = export(
    Basic(), example_args, dynamic_shapes=dynamic_shapes
)
</code></pre>
<p>However, it seems that I can only specify dynamic shapes for the Top-Level Module, but not for the nested Modules within it. such as:</p>
<pre data-code-wrap="python"><code class="lang-python">class MatMul(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, a: torch.Tensor, b: torch.Tensor) -&gt; torch.Tensor:
        return a @ b

class Basic(nn.Module):
    def __init__(self):
        super().__init__()
        self.MatMul = MatMul()

    def forward(self, x: torch.Tensor, y: torch.Tensor) -&gt; torch.Tensor:
        out1 = self.MatMul(x, y) # 3x6x3
        out2 = self.MatMul(out1, x) # 3x3x6
        return out2

example_args = (torch.randn(3, 6), torch.randn(6, 3))
dim1_x = Dim("dim1_x", min=1, max=10)
dynamic_shapes = {"a": {0: dim1_x, 1: dim1_x}, "b": {0: dim1_x, 1: dim1_x}}
aten_dialect: ExportedProgram = export(
    Basic(), example_args, dynamic_shapes=dynamic_shapes
)
</code></pre>
<p>A feasible approach is to compile and call each Module separately; for instance, compile the QKV-Proj Module/O-Proj Module with a Static Shape and the Attention Module with a Dynamic Shape. However, this method is quite cumbersome because it requires disassembling and compiling each Transformer Block within the Large Language Model (LLM) individually.</p>
<ul>
<li>So my first question is, how can I make the nested Modules support dynamic shapes?</li>
</ul>
<p>After the model is compiled, ExecuTorch can dispatch different subgraphs to run on various hardware (as I mentioned, dispatch QKV/O Proj to the NPU and Attention to the CPU/GPU), which is very convenient. But how can I achieve the pipeline form I mentioned above? For example, once the QKV Proj computation for the first token is completed, the computation kernel for the QKV Proj of the second token can be launched.</p>
<ul>
<li>My second question is, can I call a compiled Module individually, rather than starting from the Top-Level Module? In this way, I can lock different Modules to achieve a pipeline.</li>
</ul>
<p>The Dynamic KV Caches seems also unsolved? <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"> see <a href="https://github.com/pytorch/executorch/issues/4740" class="inline-onebox" rel="noopener nofollow ugc">Support for dynamic caches · Issue #4740 · pytorch/executorch · GitHub</a></p>

username: jackzhxng
cooked: <p>Hi <a class="mention" href="/u/chenghua_wang">@chenghua_wang</a>,</p>
<p>Thanks for providing this diagram! Firstly, I would like to offer an update for our progress on dynamic kv caches <a href="https://github.com/pytorch/executorch/issues/4740" rel="noopener nofollow ugc">here</a>.</p>
<p>Next - do you need to have a dynamic cache? If not then we can keep the entire model static, in which case you can go back to using a QNN backend. <a href="https://pytorch.org/executorch/main/build-run-qualcomm-ai-engine-direct-backend.html" rel="noopener nofollow ugc">Here</a> is an example.</p>
<p>If you do need a dynamic cache and thus dynamic shapes, then the entire model would need to support dynamic shape, including your Q/K/V and O proj layers.</p>

username: jackzhxng
cooked: <p>In case you were interested though in calling different modules within your program, <a href="https://github.com/pytorch/executorch/blob/main/examples/models/llava/export_llava.py#L192-L208" rel="noopener nofollow ugc">this</a> is how you would do it. Note that you can only export the <code>forward()</code> method of a module, so you will need to wrap your individual submodules in a <code>nn.module</code> with a <code>forward()</code> method.</p>

username: chenghua_wang
cooked: <p>Hi <a class="mention" href="/u/jackzhxng">@jackzhxng</a>,</p>
<p>Thank you for your helpful reply.</p>
<aside class="quote no-group" data-username="jackzhxng" data-post="3" data-topic="208561">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jackzhxng/48/72367_2.png" class="avatar"> jackzhxng:</div>
<blockquote>
<p>individual submodules in a <code>nn.module</code> with a <code>forward()</code> method.</p>
</blockquote>
</aside>
<p>This example has solved my problems, I will give it a try.</p>
<aside class="quote no-group" data-username="jackzhxng" data-post="2" data-topic="208561">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jackzhxng/48/72367_2.png" class="avatar"> jackzhxng:</div>
<blockquote>
<p>If you do need a dynamic cache and thus dynamic shapes, then the entire model would need to support dynamic shape, including your Q/K/V and O proj layers.</p>
</blockquote>
</aside>
<p>Yes, I need dynamic cache, but during the pre-fill stage, the QKV/O Proj Layers is static due to fixed chunk size.</p>
<hr>
<p>Now I have another question. Because I need to call multiple compiled modules, for example:</p>
<pre data-code-wrap="python"><code class="lang-python">out_1 = module_1.forward(...)
out_2 = module_2.forward(out_1, ...)
</code></pre>
<ul>
<li>The memory of the tensor <code>out_1</code> is managed by <code>module_2</code> , or is it managed by the user?</li>
</ul>
<p>It should be managed by the user? Because during execution, <code>out_1</code> does not know how many successor nodes it has.</p>

username: cccclai
cooked: <p>Calling part of the graph in cpu and part of the graph in npu might cause slower latency. There are two ways to use QNN. One way is to prefill token by token, and you can checkout the instructions here <a href="https://github.com/pytorch/executorch/tree/main/examples/models/llama2#optional-smaller-models-delegated-to-other-backends" class="inline-onebox" rel="noopener nofollow ugc">executorch/examples/models/llama2 at main · pytorch/executorch · GitHub</a></p>
<p>For QNN, the command line for running floating point is</p>
<pre><code class="lang-auto">python -m examples.models.llama2.export_llama --disable_dynamic_shape -kv --qnn -c stories110M.pt -p params.json
</code></pre>
<p>and for 4bit quantized, it’s still wip. The corresponding command line is</p>
<pre><code class="lang-auto">python -m examples.models.llama2.export_llama --disable_dynamic_shape -kv --qnn --pt2e_quantize qnn_16a4w  
</code></pre>

username: chenghua_wang
cooked: <p>Thank you for your answer, this is a very good example, I will try it later:</p>
<aside class="quote no-group" data-username="cccclai" data-post="5" data-topic="208561">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cccclai/48/35398_2.png" class="avatar"> cccclai:</div>
<blockquote>
<p>checkout the instructions here <a href="https://github.com/pytorch/executorch/tree/main/examples/models/llama2#optional-smaller-models-delegated-to-other-backends" rel="noopener nofollow ugc">executorch/examples/models/llama2 at main · pytorch/executorch · GitHub</a></p>
</blockquote>
</aside>
<p>The original intention of dividing the prompt into chunks and executing them separately on the CPU and NPU is to utilize dynamic cache. Indeed, the NPU is much faster, after using the Pipeline, it may not be possible to completely overlap the CPU’s computation time.</p>
<p>It seems that QNN2.21.0+ support dynamic shape? see document: <a href="https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-50/api_usage_guidelines.html#dynamic-dimensions" class="inline-onebox" rel="noopener nofollow ugc">Qualcomm Documentation</a></p>

username: gianfranco
cooked: <p>Hey Guys, I am still having problem with dynamic shape on mps backend.  Can someone say how to export llama3_2 .pte file  on mps and run it with a generic prompt?</p>