[QUESTION] The scope of reducing gradient norms

**Your question**
I have a question about the design of gradient-norm computation in the latest Megatron-LM codebase, especially the interplay between the distributed optimizer and all kinds of parallel mechanisms.

Currently, Megatron-LM gathers the gradient norm information by the function `get_grad_stats_parallel_group`, which distinguishes the distributed optimizer case and non-dist case. See: https://github.com/NVIDIA/Megatron-LM/blob/41eecc45ec7b38f7dea237cce6c0df9b8eaaa710/megatron/core/optimizer/optimizer.py#L161-L179

For the non-dist case, the group scope is defined as model parallel group, i.e. the TPxPP rectangle: https://github.com/NVIDIA/Megatron-LM/blob/41eecc45ec7b38f7dea237cce6c0df9b8eaaa710/megatron/core/parallel_state.py#L931-L944 For the dist case, it just returns `None` and reduce over the whole world. https://github.com/NVIDIA/Megatron-LM/blob/41eecc45ec7b38f7dea237cce6c0df9b8eaaa710/megatron/core/optimizer/distrib_optimizer.py#L617-L623 Since the ZeRO-1 optimizer shards the main_grad over DP groups, it should well handle the TPxPPxDP case.

**My question is, if other types of parallelisms are enabled, such as EP, which shards the experts but not attention modules, will the result of the whole-world reduction still be correct?** As I understand it, the EP-only setting may duplicate the gradient information of the attention modules.

Could anyone help me find the answer? Thanks in advance.