DS-MoE: Making MoE Models More Efficient and Less Memory-Intensive

<p>
<em>Estimated reading time: 4 minutes</em> </p>
<p>Mixture-of-Experts (MoE) language models are known for their ability to reduce computing needs by 2 to 4 times compared to traditional dense models, without sacrificing performance. This makes them especially useful in situations where computing resources are limited. However, MoE models typically need 2 to 4 times more parameters to perform as well as a dense model. For example, models like DeepSeekMoE-16B and Qwen1.5-MoE-A2.7B which has 16B parameters were designed to match the performance of a 7B model. The large number of parameters in MoE models incurs larger GPU memory requirements which makes them less efficient in I/O-bounded scenarios like autoregressive generation.</p>
<figure style="text-align: center;">
<img alt="Alternative Text" height="450" src="https://cdn-uploads.huggingface.co/production/uploads/640e3a753830fd441c2c768d/POhcBNq6Gj_rHzp6a8H_h.png" style="margin: 0 auto; display: block;" width="400"/>
<figcaption style="text-align: center;">Figure 1: Decoding Throughput of Dense Models and the SMoE Models with Similar Performance. We test the performance under the setup where only the input length is 1 and output length is 512. This study reveals that traditional Sparse Mixture of Experts (SMoE) models exhibit reduced output throughput in I/O-bounded situations despite their lower computational demands. The models are tested with HuggingFace Transformers.</figcaption>
</figure>
<p>Is it necessary for MoE models to be so large to achieve high performance? Can we create an MoE model that maintains performance but uses fewer parameters and less computational power? Enter DS-MoE. This model achieves similar performance to dense models but uses about one-third of the computational resources and only half as many parameters as other MoE models.</p>
<figure style="text-align: center;">
<img alt="Alternative Text" height="450" src="https://cdn-uploads.huggingface.co/production/uploads/640e3a753830fd441c2c768d/-cktt05v8p7-U-L33rIgs.png" style="margin: 0 auto; display: block;" width="400"/>
<figcaption style="text-align: center;">Figure 2: Number of Parameters for Performance-Matched Models. We plot the size and computational profiles of the Dense-3B, SMoE-5B, and DS-MoE-3B models trained with 100B tokens, each achieving a comparable averaged task performance. DS-MoE demonstrates both computational efficiency and parameter efficiency, where the computational cost is quantified by counting the number of active parameters engaged during inference.</figcaption>
</figure>
<p>The concept of DS-MoE involves densely training the experts and forcing the model's routers to gradually ignore unnecessary experts for a given token. We employ the Mutual Information (MI) loss to the training process, which balances the load of each expert across the entire batch, but also encourages each input token to concentrate their gating probability to fewer experts. </p>
<figure style="text-align: center;">
<img alt="Alternative Text" src="https://cdn-uploads.huggingface.co/production/uploads/640e3a753830fd441c2c768d/dCzK6pAcxQDsb-70uPDIe.png" style="margin: 0 auto; display: block;"/>
<figcaption style="text-align: center;">Figure 3: Subfigure (a) illustrates the conventional sparse training method in MoE models, characterized by sparse gradient propagation in both the router and the experts. Subfigure (b) details the dense training strategy in DS-MoE, which involves dense propagation of gradients for both routers and experts.</figcaption>
</figure>
<p>The MI loss is defined as 
<span class="katex-display"><span class="katex"><span class="katex-mathml"><math display="block" xmlns="http://www.w3.org/1998/Math/MathML"><mrow><msub><mi>L</mi><mrow><mi>M</mi><mi>I</mi></mrow></msub><mo>=</mo><mo>−</mo><mi>H</mi><mo stretchy="false">(</mo><mi>e</mi><mo stretchy="false">)</mo><mo>+</mo><mfrac><mn>1</mn><mrow><mi mathvariant="normal">∣</mi><mi>X</mi><mi mathvariant="normal">∣</mi></mrow></mfrac><munder><mo>∑</mo><mrow><mi>x</mi><mo>∈</mo><mi>X</mi></mrow></munder><mi>H</mi><mo stretchy="false">(</mo><mi>e</mi><mi mathvariant="normal">∣</mi><mi>x</mi><mo stretchy="false">)</mo><mo separator="true">,</mo><mspace width="1em"></mspace><mi>H</mi><mo stretchy="false">(</mo><mi>e</mi><mo stretchy="false">)</mo><mo>=</mo><mo>−</mo><munderover><mo>∑</mo><mrow><mi>i</mi><mo>=</mo><mn>1</mn></mrow><mi>N</mi></munderover><mi>p</mi><mo stretchy="false">(</mo><mi>e</mi><mo stretchy="false">)</mo><mi>log</mi><mo>⁡</mo><mo stretchy="false">(</mo><mi>p</mi><mo stretchy="false">(</mo><mi>e</mi><mo stretchy="false">)</mo><mo stretchy="false">)</mo><mo separator="true">,</mo></mrow> L_{MI} = - H(e) + {1 \over |X|}\sum_{x\in X} H(e|x), \quad H(e) = - \sum_{i=1}^{N}p(e)\log(p(e)),</math></span><span aria-hidden="true" class="katex-html"><span class="base"><span class="strut" style="height:0.8333em;vertical-align:-0.15em;"></span><span class="mord"><span class="mord mathnormal">L</span><span class="msupsub"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height:0.3283em;"><span style="top:-2.55em;margin-left:0em;margin-right:0.05em;"><span class="pstrut" style="height:2.7em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mathnormal mtight" style="margin-right:0.10903em;">M</span><span class="mord mathnormal mtight" style="margin-right:0.07847em;">I</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height:0.15em;"><span></span></span></span></span></span></span><span class="mspace" style="margin-right:0.2778em;"></span><span class="mrel">=</span><span class="mspace" style="margin-right:0.2778em;"></span></span><span class="base"><span class="strut" style="height:1em;vertical-align:-0.25em;"></span><span class="mord">−</span><span class="mord mathnormal" style="margin-right:0.08125em;">H</span><span class="mopen">(</span><span class="mord mathnormal">e</span><span class="mclose">)</span><span class="mspace" style="margin-right:0.2222em;"></span><span class="mbin">+</span><span class="mspace" style="margin-right:0.2222em;"></span></span><span class="base"><span class="strut" style="height:2.6431em;vertical-align:-1.3217em;"></span><span class="mord"><span class="mord"><span class="mopen nulldelimiter"></span><span class="mfrac"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height:1.3214em;"><span style="top:-2.314em;"><span class="pstrut" style="height:3em;"></span><span class="mord"><span class="mord">∣</span><span class="mord mathnormal" style="margin-right:0.07847em;">X</span><span class="mord">∣</span></span></span><span style="top:-3.23em;"><span class="pstrut" style="height:3em;"></span><span class="frac-line" style="border-bottom-width:0.04em;"></span></span><span style="top:-3.677em;"><span class="pstrut" style="height:3em;"></span><span class="mord"><span class="mord">1</span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height:0.936em;"><span></span></span></span></span></span><span class="mclose nulldelimiter"></span></span></span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mop op-limits"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height:1.05em;"><span style="top:-1.8557em;margin-left:0em;"><span class="pstrut" style="height:3.05em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mathnormal mtight">x</span><span class="mrel mtight">∈</span><span class="mord mathnormal mtight" style="margin-right:0.07847em;">X</span></span></span></span><span style="top:-3.05em;"><span class="pstrut" style="height:3.05em;"></span><span><span class="mop op-symbol large-op">∑</span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height:1.3217em;"><span></span></span></span></span></span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mord mathnormal" style="margin-right:0.08125em;">H</span><span class="mopen">(</span><span class="mord mathnormal">e</span><span class="mord">∣</span><span class="mord mathnormal">x</span><span class="mclose">)</span><span class="mpunct">,</span><span class="mspace" style="margin-right:1em;"></span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mord mathnormal" style="margin-right:0.08125em;">H</span><span class="mopen">(</span><span class="mord mathnormal">e</span><span class="mclose">)</span><span class="mspace" style="margin-right:0.2778em;"></span><span class="mrel">=</span><span class="mspace" style="margin-right:0.2778em;"></span></span><span class="base"><span class="strut" style="height:3.106em;vertical-align:-1.2777em;"></span><span class="mord">−</span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mop op-limits"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height:1.8283em;"><span style="top:-1.8723em;margin-left:0em;"><span class="pstrut" style="height:3.05em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mathnormal mtight">i</span><span class="mrel mtight">=</span><span class="mord mtight">1</span></span></span></span><span style="top:-3.05em;"><span class="pstrut" style="height:3.05em;"></span><span><span class="mop op-symbol large-op">∑</span></span></span><span style="top:-4.3em;margin-left:0em;"><span class="pstrut" style="height:3.05em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mathnormal mtight" style="margin-right:0.10903em;">N</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height:1.2777em;"><span></span></span></span></span></span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mord mathnormal">p</span><span class="mopen">(</span><span class="mord mathnormal">e</span><span class="mclose">)</span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mop">lo<span style="margin-right:0.01389em;">g</span></span><span class="mopen">(</span><span class="mord mathnormal">p</span><span class="mopen">(</span><span class="mord mathnormal">e</span><span class="mclose">))</span><span class="mpunct">,</span></span></span></span></span>
where <strong>X</strong> denotes the tokens in a minibatch, and <strong>e</strong> denotes the experts. Intuitively, maximizing H(e) balances the load of each expert across the entire batch, and minimizing H(e|x) encourages each input <strong>x</strong> to concentrate their gating probability to fewer experts. </p>
<p>During inference, DS-MoE chooses only the top K experts based on their scores. The determination of the number of K is based on either a predefined value or an adaptive method, contingent upon the count of experts with scores surpassing a certain threshold. As a result, DS-MoE can perform as well as similarly sized dense models while using far fewer active parameters, as demonstrated in the table.</p>
<div class="max-w-full overflow-auto">
<table>
<thead><tr>
<th>Model</th>
<th>HellaSwag</th>
<th>PIQA</th>
<th>WinoGrande</th>
<th>SciQ</th>
<th>Arc-e</th>
<th>Arc-c</th>
<th>Avg. Perf.</th>
<th>Active Params</th>
</tr>
</thead><tbody><tr>
<td>Dense-3B</td>
<td>40.4</td>
<td>71.4</td>
<td>58.7</td>
<td>86.0</td>
<td>59.6</td>
<td>26.1</td>
<td>57.0</td>
<td>2705M</td>
</tr>
<tr>
<td>SMoE-5B</td>
<td>40.1</td>
<td>70.7</td>
<td>56.5</td>
<td>85.6</td>
<td>58.4</td>
<td>24.8</td>
<td>56.0</td>
<td>1212M</td>
</tr>
<tr>
<td>DS-MoE-3B</td>
<td>39.3</td>
<td>71.6</td>
<td>57.9</td>
<td>85.6</td>
<td>57.7</td>
<td>24.9</td>
<td>56.2</td>
<td>934M</td>
</tr>
<tr>
<td>Dense-6B</td>
<td>44.3</td>
<td>72.2</td>
<td>59.9</td>
<td>88.0</td>
<td>62.9</td>
<td>27.9</td>
<td>59.2</td>
<td>6186M</td>
</tr>
<tr>
<td>DS-MoE-6B</td>
<td>43.5</td>
<td>73.0</td>
<td>57.9</td>
<td>86.9</td>
<td>61.9</td>
<td>27.9</td>
<td>58.5</td>
<td>1813M</td>
</tr>
</tbody>
</table>
</div>
<p>We also tested DS-MoE with vLLM to see how it compares to other models in terms of processing speed and memory usage at the 7B performance tier. We looked at how many requests and tokens it could handle per second, using a setup where each input and output consisted of 1,000 tokens and the GPU memory usage was capped at 90%.</p>
<div class="max-w-full overflow-auto">
<table>
<thead><tr>
<th>Model</th>
<th>Total Params</th>
<th>Active Params</th>
<th>Model Memory</th>
<th>A100 Throughput</th>
<th>A100 TPS</th>
<th>H100 Throughput</th>
<th>H100 TPS</th>
</tr>
</thead><tbody><tr>
<td>Dense-6B</td>
<td>6.4B</td>
<td>6.4B</td>
<td>12.3 GiB</td>
<td>1.04</td>
<td>2079.8</td>
<td>1.40</td>
<td>2808.7</td>
</tr>
<tr>
<td>Mistral-7B</td>
<td>7.2B</td>
<td>7.2B</td>
<td>13.5 GiB</td>
<td>1.07</td>
<td>2140.8</td>
<td>1.52</td>
<td>3047.4</td>
</tr>
<tr>
<td>DeepSeekMoE</td>
<td>17.3B</td>
<td>2.8B</td>
<td>30.5 GiB</td>
<td>1.17</td>
<td>2330.1</td>
<td>1.57</td>
<td>3144.1</td>
</tr>
<tr>
<td>Qwen1.5-MoE</td>
<td>16.4B</td>
<td>2.7B</td>
<td>26.7 GiB</td>
<td>1.33</td>
<td>2665.7</td>
<td>1.81</td>
<td>3616.9</td>
</tr>
<tr>
<td>DS-MoE-6B</td>
<td>6.5B</td>
<td>2.2B</td>
<td>12.6 GiB</td>
<td>2.00</td>
<td>3992.8</td>
<td>2.30</td>
<td>4603.9</td>
</tr>
</tbody>
</table>
</div>
<p>The test shows that DS-MoE outperforms both dense models in terms of computational cost and sparsely trained MoEs in model memory, leading to faster processing in the computation-bounded scenarios as well the I/O bounded scenarios. Note that DS-MoE-6B is not yet comparable with other models regarding downstream performance, due to its training on merely 100 billion tokens (versus the trillions for other models). Nevertheless, DS-MoE has demonstrated significant promise in achieving the performance levels of dense models with a comparable volume of training data.</p>
<p><a href="https://arxiv.org/pdf/2404.05567.pdf" rel="nofollow"><em>Read More in the Paper</em></a></p>
