improve MoE bias update logic in optimizer

[rakkit]

We put all experts' usage into a buffer such that we only need one reduce rather than #number-of-layers times

Additionally, handle cases where tokens per expert are counted twice during full recompute.

[tianyu-l]

Thank you for the PR! I left some comments.

[tianyu-l]

Thanks, had some more comments.

[rakkit]

for moe ep usage and/or bias . Here we need to do smth like

 expert_usage_metrics = {
     f"moe_ep_usage/L-{layer_id}_EP-{ep_idx}": usage / sum_tokens
     for ep_idx, usage in enumerate(tokens_per_expert)
 }
 
 model_part._metrics_to_log.update{expert_usage_metrics}

and once we finalize RP#1578

for Moe model we can have

    def get_extra_metrics(self, model_parts: list[nn.Module], *args, **kwargs) -> None  | dict[str, Any]:
        return model_parts._metrics_to_log

[tianyu-l]

Thanks! Had some final comments.

[rakkit]

Removed the comment (for ep-usage) and added the early exit in the first loop

[tianyu-l]

LGTM, thank you! Please fix linting so we can merge.

[rakkit]

format fixed, thanks a lot for the discussion.

It's important to know that this PR only improves the code on COMM parts -> reduced to only once.
In practice, from the profiler, the second loop will launch lots of kernels, [num moe layer] * [slice, mean, sign, multi, add, zeros), unless one makes everything there vectorized.