Scattered Mixture-of-Experts Implementation

For Megablocks and PIT, the authors stated that “the sparse matrix format makes the intermediate representation harder to extend upon”. Does this lead to memory overhead? If yes, can you describe the reasons? Otherwise, is there any in-depth analysis or experiments on the aforementioned feature of Megablocks and PIT?

In an MoE MLP, the sparse matrix format does not create the overhead. Megablocks makes a re-ordered copy of the input before performing the GeMM operation, which is the step that creates the additional memory overhead. Figure 1 in the paper illustrates the additional copying that causes this.

However, extending MoE to Attention the memory issue is further compounded:

1. First we need to perform a ParallelLinear transformation of the input, to produce queries, which will require re-ordering to expert order.
2. These queries will need to be re-ordered to chronological order for application of position embeddings (e.g. RoPE) and appropriate masking.
3. Once the standard attention step is done, the output embeddings will have to be reordered to expert order again,
4. before the final output expert transformation, before being reordered to chronological order.

There are 4 re-ordering steps in total, and each will incur an additional duplication of the input. With scatter2scatter, we can bypass the reordering by only performing the expert ordered operation in SRAM. This results in a 42% savings in memory usage.

Both Megablocks (in sparse mode) and PIT rely on a sparse representation of the intermediate hidden state in an MoE MLP. Unless the operation performed on the sparse representation is a point-wise operation (like an activation function, for example) to perform anything that requires order sensitivity (attention) would require converting back to the dense format, then re-ordering, incurring all of the above-mentioned costs AND the sparse-dense-sparse conversion.

In the comparative experiments, what is the performance of PIT?

We were unable to find an open source implementation of PIT to compare against

Thank you for the feedback!

The paper should provide more information about the details of the key innovation, including its detailed implementation, how it avoids padding, instead of using paragraphs to introduce MoE, its integration with MoE and MoA. Also, the pseudo-codes are not informative enough.

We will include more detailed information about the specifics you pointed out. Specifically, we will expand on the following points in the paper:

1. Instead of performing a padded copy in High Bandwidth Memory (HBM), we sort the tokens according to the experts, and pad the indices instead.
2. When loading a tile into Static RAM (SRAM), we load according to the padded indices, resulting in a padded tile. The full padded array of token embeddings is never realised in HBM.
3. We will add the pseudo-code for the scatter2scatter kernel itself, since the current pseudo-code only describes the overarching process using scatter2scatter as a primitive.

It is advised to provide more experimental results about model quality to ensure that using ScatterMoE can achieve similar or even better quality on the LLM evaluation benchmarks.

We converted Mixtral 7B to use ScatterMoE (script included), and ran the lm-evaluation harness on the following benchmarks.

The difference in the results are negligible, which we expect as ScatterMoE is a different implementation and does not make any fundamental changes to MoEs.

Thank you for your feedback!

As a competitor of Megablocks,... , I think the work will be much more influential and have more significance to be accepted.

We agree with this perspective, and ScatterMoE has been integrated into an open source fine-tuning framework and we have plans to include it into a pretraining framework are underway (pull-request submitted).

However, ScatterMoE is a modular, light-weight implementation of MoEs that is meant to be used as a drop-in replacement for MoEs. We’ve added a conversion script for the Mixtral model on huggingface to demonstrate this, and include the lm-evaluation-harness results achieved with both the naive huggingface implementation and using ScatterMoE, and the results show that the differences are negligible.

It will be much better if larger scales of MoE models are benchmarked....

We benchmarked a bigger model of 9B parameters: 32 experts and k=4, d_model = 2048, mlp_size = 1536 per expert, 26 layers. This was the largest we could manage with a single node and 8 X A100 80gb GPUs. At this size, the throughput for a minibatch of 32 * 2048, with 16 acc steps for an actual batch size of 2, was 57957.8 tokens/s for ScatterMoE, and 19924.9 tokens/s for Megablocks scattered. However, this may be due to a memory leak in the recent version of Megablocks.

For some recent variants of MoE like shared experts (DeepSeekMoE for e.g.),...

ScatterMoE can support the shared expert scheme in DeepSeekMoE: one simple way to do this is to either adjust the logit to a large value for that expert, or to append the shared expert index before passing it to the MoE module. However, it is unlikely that there will be an advantage over those already discussed in this report.