Weight-Tie MoE

[ClashLuke]

• is_stacked: check for _stacked parameter-name suffix instead of dim size
• pattern_match: looped lax.cond to single lax.switch call

This PR reduces the number of parameters from 9.6B to 1.6B and performs as well as untied MoE for the first 2 billion tokens:

With that, both are significantly better than the baseline:

[ClashLuke]

Better than MoE after 10B tokens