Parallelising expert execution on single GPU using CUDA streams

Shamauk: <p><strong>Keywords</strong>: Mixture-of-Experts, CUDA Streams, Parallelism<br>
<strong>One-liner</strong>: Attempting to execute distinct experts each with their own disjoint subset of tokens concurrently with CUDA Streams.<br>
<strong>Expectation</strong>: Reduction in latency times for inference.<br>
<strong>Hardware/Software</strong>. <a href="http://nvcr.io/nvidia/pytorch:24.04-py3" rel="noopener nofollow ugc">nvcr.io/nvidia/pytorch:24.04-py3</a>, <strong>single</strong> V100 32GB, DGX1 (again only using 1 of the 8)</p>
<p>Note: Realised after composing that I am only able to post a single screenshot, so you will have to scroll to the replies for some.</p>
<p><strong>Full Post</strong><br>
I am currently working with the SwitchTransformer model and began tracing its execution using Nsys. What I have noticed was that the different experts execute <em>sequentially</em>. This is surprising to me since 1. these experts do not depend on one another and 2. GPUs are massively parallel and I would expect it have no problem running an extra matrix multiplication at the same time.</p>
<ol>
<li></li>
</ol>
<p>Here is a showcase of the sequential execution. (Reply)<br>
Here we see the execution of a batch through the first two blocks of the encoder. One can notice that the block without MoE is significantly quicker even though it does the same computation (the gating function is very small). The second block, you can notice the execution of all 64 experts, sequentially causing a major slowdown.</p>
<p>I want to <em>reduce</em> the time for a forward pass through a <em>block with MoE</em>. I figured the best approach is to find a way to run the experts concurrently since they do not depend on each other.</p>
<ol start="2">
<li></li>
</ol>
<p>I firstly did a simple change to the HuggingFace SwitchTransformer model.</p>
<pre data-code-wrap="python"><code class="lang-python">for idx, expert in enumerate(self.experts.values()):
            token_indices = router_mask[:, :, idx].bool()
            with torch.cuda.stream(self.cuda_streams[idx]):
                next_states[token_indices] = expert(hidden_states[token_indices]).to(next_states.dtype)
</code></pre>
<p>At model creation I create a stream for each expert and then simply index in when deciding which stream to use.</p>
<p>(Reply). The result is that they run on separate streams although they appear to still be running sequentially – just on different streams. Furthermore inference takes <strong>68% longer</strong>.</p>
<ol start="3">
<li></li>
</ol>
<p>Furthermore was curious if the assignment was causing the serialised order so I just throw away the values.</p>
<pre data-code-wrap="python"><code class="lang-python">with torch.cuda.stream(self.cuda_streams[idx]):
                # Don't care abt model acc, so throw away the computation result
                expert(hidden_states[token_indices.clone()])
</code></pre>
<p>This does start to get some kernels executing concurrently. (Reply)<br>
However still inference takes 42% longer than serial execution.</p>
<ol start="4">
<li></li>
</ol>
<p>I was thinking maybe it is an issue of having too many streams so I reduced it to two.</p>
<pre data-code-wrap="python"><code class="lang-python">for idx, expert in enumerate(self.experts.values()):
            token_indices = router_mask[:, :, idx].bool()
            if idx % 2 == 0:
                with torch.cuda.stream(self.cuda_streams[0]):
                    expert(hidden_states[token_indices.clone()])
            else:
                with torch.cuda.stream(self.cuda_streams[1]):
                    expert(hidden_states[token_indices.clone()])
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/b/cbef2c1f05f6ba274b355957dab5812b8f57daa8.png" data-download-href="https://discuss.pytorch.org/uploads/default/cbef2c1f05f6ba274b355957dab5812b8f57daa8" title="Screenshot 2024-05-27 at 2.51.34 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cbef2c1f05f6ba274b355957dab5812b8f57daa8_2_690x124.png" alt="Screenshot 2024-05-27 at 2.51.34 PM" data-base62-sha1="t65lKPVNQDjaKjf5eas6Z5nq41q" width="690" height="124" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cbef2c1f05f6ba274b355957dab5812b8f57daa8_2_690x124.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cbef2c1f05f6ba274b355957dab5812b8f57daa8_2_1035x186.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cbef2c1f05f6ba274b355957dab5812b8f57daa8_2_1380x248.png 2x" data-dominant-color="DFE3EB"><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">Screenshot 2024-05-27 at 2.51.34 PM</span><span class="informations">1566×282 28.4 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><br>
The overlap is still not great and the overhead of using streams comes at increasing inference by 20% over the base serial execution.</p>
<p><strong>Action</strong><br>
If anyone has more experience or knowledge, do please let me know on how to tackle this problem.</p>

Shamauk: <ol>
<li></li>
</ol>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/6/864a1184b197e20cf973e834bf33e34991ed4a30.png" data-download-href="https://discuss.pytorch.org/uploads/default/864a1184b197e20cf973e834bf33e34991ed4a30" title="Screenshot 2024-05-27 at 2.33.34 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/6/864a1184b197e20cf973e834bf33e34991ed4a30_2_690x131.png" alt="Screenshot 2024-05-27 at 2.33.34 PM" data-base62-sha1="j9YFERKTTreQYe7nYHFMmYTVlC0" width="690" height="131" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/6/864a1184b197e20cf973e834bf33e34991ed4a30_2_690x131.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/6/864a1184b197e20cf973e834bf33e34991ed4a30_2_1035x196.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/6/864a1184b197e20cf973e834bf33e34991ed4a30_2_1380x262.png 2x" data-dominant-color="9CCD9C"><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">Screenshot 2024-05-27 at 2.33.34 PM</span><span class="informations">1964×374 61.5 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>

Shamauk: <ol start="2">
<li></li>
</ol>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/2/a2436c70cb6bd3fb8a86e370d6eba6c01d08a562.png" data-download-href="https://discuss.pytorch.org/uploads/default/a2436c70cb6bd3fb8a86e370d6eba6c01d08a562" title="Screenshot 2024-05-27 at 2.41.15 PM"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/2/a2436c70cb6bd3fb8a86e370d6eba6c01d08a562.png" alt="Screenshot 2024-05-27 at 2.41.15 PM" data-base62-sha1="n9rNSSZAzc598DYjJl9lSU4aVzk" width="429" height="500" data-dominant-color="EDEDEE"><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">Screenshot 2024-05-27 at 2.41.15 PM</span><span class="informations">732×852 5.01 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>

Shamauk: <ol start="3">
<li></li>
</ol>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/7/d751ca6d6ffbd1adb9b005a3d02851a2fc569590.png" data-download-href="https://discuss.pytorch.org/uploads/default/d751ca6d6ffbd1adb9b005a3d02851a2fc569590" title="Screenshot 2024-05-27 at 2.48.02 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d751ca6d6ffbd1adb9b005a3d02851a2fc569590_2_321x500.png" alt="Screenshot 2024-05-27 at 2.48.02 PM" data-base62-sha1="uINTiMp5rWnvOoa0aZWtFGgWAlW" width="321" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d751ca6d6ffbd1adb9b005a3d02851a2fc569590_2_321x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d751ca6d6ffbd1adb9b005a3d02851a2fc569590_2_481x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d751ca6d6ffbd1adb9b005a3d02851a2fc569590_2_642x1000.png 2x" data-dominant-color="EBECED"><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">Screenshot 2024-05-27 at 2.48.02 PM</span><span class="informations">848×1318 23.6 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>

Shamauk: <p>4 closeup.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/9/e9e6f7dcec1fff62b3eec12f7b3e802d4d58a7d1.png" data-download-href="https://discuss.pytorch.org/uploads/default/e9e6f7dcec1fff62b3eec12f7b3e802d4d58a7d1" title="Screenshot 2024-05-27 at 2.51.44 PM"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/9/e9e6f7dcec1fff62b3eec12f7b3e802d4d58a7d1.png" alt="Screenshot 2024-05-27 at 2.51.44 PM" data-base62-sha1="xnc63GoOebMHxDFF4SbLO7UW7wB" width="690" height="49" data-dominant-color="EBECED"><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">Screenshot 2024-05-27 at 2.51.44 PM</span><span class="informations">1940×140 4.62 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>

ptrblck: <aside class="quote no-group" data-username="Shamauk" data-post="1" data-topic="203553">
<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/shamauk/48/70329_2.png" class="avatar"> Shamauk:</div>
<blockquote>
<p>GPUs are massively parallel and I would expect it have no problem running an extra matrix multiplication at the same time.</p>
</blockquote>
</aside>
<p><a href="https://discuss.pytorch.org/t/running-independent-nn-module-instances-in-nn-modulelist-truly-in-parallel-in-pytorch/188166/3">This post</a> might be helpful explaining the memory and compute resources in the attached GTC talk.</p>