[QUESTION] Should Latent MoE scale both expert `linear_fc2` and `fc2_latent_proj`?
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
Hi Megatron team @NVIDIA/mcore-oncall,
I have a question about the initialization strategy used by Latent MoE.
From the current implementation, it looks like the routed expert branch in Latent MoE applies `output_layer_init_method` twice:
1. The expert MLP down projection / `linear_fc2`
2. The latent-to-hidden projection / `fc2_latent_proj`
In the standard non-latent MLP / MoE case, using output_layer_init_method on linear_fc2 is intuitive because linear_fc2 directly produces the output that goes into the residual branch.
However, in Latent MoE, expert linear_fc2 only projects back to the latent dimension. The actual projection back to the residual hidden dimension is fc2_latent_proj.
From a residual-variance perspective, one might expect only the final projection into the residual branch to use the depth-scaled output initialization, for example:
fc1_latent_proj: init_method
expert linear_fc1: init_method
expert linear_fc2: init_method
fc2_latent_proj: output_layer_init_method
But the current implementation seems to use:
fc1_latent_proj: init_method
expert linear_fc1: init_method
expert linear_fc2: output_layer_init_method
fc2_latent_proj: output_layer_init_method
This makes the routed latent MoE branch more conservatively initialized than a single scaled output projection.
Contributor guide
Assessment
This issue has not been assessed yet.