[DO NOT MERGE] Ranran hide a2a

[RissyRan]

No description provided.

[wang2yn84]

What does X, E and M stands for here?

[RissyRan]

What does X, E and M stands for here?

Could you point me to the specific line? X/E should be expert, and M should be model dimension.

[wang2yn84]

What does X, E and M stands for here?

Could you point me to the specific line? X/E should be expert, and M should be model dimension.

Sure, I'm looking at jnp.einsum("BXM,XEM -> BXE") from hide_ff2_a2a.py

[RissyRan]

What does X, E and M stands for here?

Could you point me to the specific line? X/E should be expert, and M should be model dimension.

Sure, I'm looking at jnp.einsum("BXM,XEM -> BXE") from hide_ff2_a2a.py

I think this is one example for @gobbleturk's test? X seems sequence here, M is model dimension, E is MLP dimension

[gobbleturk]

What does X, E and M stands for here?

Could you point me to the specific line? X/E should be expert, and M should be model dimension.

Sure, I'm looking at jnp.einsum("BXM,XEM -> BXE") from hide_ff2_a2a.py

I think this is one example for @gobbleturk's test? X seems sequence here, M is model dimension, E is MLP dimension

This script is testing the second a2a at the end of the feed forward layer where we move from activations sharded on the expert dimension to activations sharded on the batch dimension
This is a toy script so I combine sequence and batch into one dimension "B" representing token batch

• X is for experts
• E is for embed (the smaller dimension, often called "hidden dim" or "model dim")
• M is for MLP (the larger dimension, often called "intermediate dim" , or "FF dim")