Query about Sparsity Implementation and Acceleration Mechanism

[huchenz1]

Hi, thank you for your wonderful work!

I’ve been exploring the implementation of sparsity in the codebase and noticed that sparsity is achieved through masks such as attn_weight_mask, mlp_weight_mask, and token_select. However, during the forward pass, these masks are represented as binary values (0s and 1s), and the tensor dimensions remain unchanged.

Could you kindly clarify how the actual acceleration effect is achieved during runtime?

[wangbo-zhao]

Hi, we can skip those masked heads/channels during inference to achieve actual acceleration.

[wangbo-zhao]

Here is an example to implement this:

[EricLina]

How many channels will be skipped, it seems like the probability of channel_mask.sum() == 0 is very low.

[EricLina]

I think MoE is closer to truly sparse skip.

[wangbo-zhao]

It depends on the specific block and timestep. At certain timesteps or certain blocks, the probability of channel_mask.sum() == 0 can be very high, or only a few heads/channels may be active.

[MiaApr]

Hi,

First, thank you for your excellent work on this project; it has been very insightful.

I have a question regarding the token selection mechanism. When looking at the code, specifically in models.py at line 311, it appears that the MLP layer still processes all tokens during inference.

The token mask is multiplied with the output of the MLP, not the input. My understanding is that this approach might not achieve the goal of accelerating inference, as the computation for all tokens in the MLP still needs to be performed.

I am a bit confused by this and would greatly appreciate it if you could clarify this implementation detail.

Thank you for your time!