[NVIDIA] Use separate FP8 einsum instances in MoE

[kaixih]

We found in the original version, the same fp8 einsum op is reused in the MoE by dispatch, FFN, combine gemm ops, meaning they will share the same fp8 meta params, which is wrong.

This PR fixed this issue by cloning the fp8 einsum to each of the gemm in the MoE.

cc. @zhangqiaorjc @abhinavgoel95 @hx89

[kaixih]

Sure. Let's move the conversation here.

@kaixih could you add a config option in MoE transformer that preserves the old behavior by default

i don't want to break existing checkpoints

I think by default the einsum_tpl is configed to be jnp.einsum and when users choose to set USE_FP8 in paxml, the fp8 einsum will be configed into the MoE layer.

[kaixih]

@zhangqiaorjc Please let me know if I miss your point of the above comment.

[kaixih]

We noticed that there might be an issue of duplicated quantization over the activation for the first gemm in the gated FFN.

I fixed it by proposing a new gated_einsum where we can control the quantization only applied once.