NVIDIA / NVIDIA/Megatron-LM

[QUESTION] Should Latent MoE scale both expert `linear_fc2` and `fc2_latent_proj`?

Open
#5,225 1 comment 0 reactions 1 assignee Claimed by @yashaswikarnati View on GitHub
community-request module: moe question waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.