Scattered Mixture-of-Experts Implementation

Summary

This paper presents ScatterMoE, an SMoE implementation that minimizes memory overhead. Since ParallelLinear allows for performing different combinations of SMoE transformations allowing for the input and output to be either grouped or scattered, ScatterMoE relies on ParallelLinear to achieve both forward and backward passes. Generally, these operations are performed without additional copying (or padding) of the input and output tensors involved in the existing implementations. Experiments validate the effectiveness of ScatterMoE on throughput and memory use.

Rating

8

Confidence

3

Ethics flag

1

Reasons to accept

1) The implementation proposed in this study is beneficial for various real-world applications. 2) The paper is written-well with a good presentation style and experimental results.

Reasons to reject

I have several questions to authors, as listed in the following part.

Questions to authors

1) 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? 2) In the comparative experiments, what is the performance of PIT?

Summary

This paper introduces a new method called ScatterMoE for the efficient implementation of MoE. It proposes a new kernel called scatter2scatter which supports group-to-group, group-to-scatter, scatter-to-scatter, and scatter-to-group operation. The benchmark against the recent SOTA work MegaBlocks demonstrates its higher throughput and lower memory footprint.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

1. From the interpretation in the paper, this method is simple by fusing fuses grouped GeMMs and scattered read and write operations and effective by higher throughput and lower memory footprint. 2. The paper is clearly written and easy to read.

Reasons to reject

1. 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. 2. 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.

Summary

This paper proposes ScatterMoE which optimizes the memory cost and the throughput of tokens for SMoE models. Specially, ScatterMoE focuses on the issue on extra memory copies of token representations during the operation of MoE scattering and gathering. By implementing custom kernels using Triton or CUTLASS, ScatterMoE eliminates the extra memory copies and improves the overall training & inference performance, compared with existing SMoE frameworks like Megablocks. The work provides the detailed implementation of ScatterMoE and extensively compares the performance with baselines (Megablocks and naive PyTorch implementation) considering a variety of experimental settings (granularity, sparsity and MoA).

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

1. The motivation is straightforward and reasonable, which directly eliminates the excessive memory consumption of representation tensor copies and token paddings. 2. The paper describes the implementation clearly including the forward & backward passes of ParallelLinear and its extension to MoA modules. 3. The comparison in performance provided is extensive and includes sufficient conditions (like granularity and sparsity) which affects the experimental results.

Reasons to reject

1. As a competitor of Megablocks, this work should better be accomplished into a usable package (instead of just raw fused kernels) and provide a more end-to-end time-basis training curve for current prevalent LLMs. If achieved, I think the work will be much more influential and have more significance to be accepted. 2. It will be much better if larger scales of MoE models are benchmarked, in fact a 1.5B parameter configuration is more like a toy setting.

Questions to authors

For some recent variants of MoE like shared experts (DeepSeekMoE for e.g.), will ScatterMoE still be supported and in advantage in performance?