fix(optimizer): use expert TP group for gradient stats

[guapisolo]

Summary

Count expert gradients with expert TP ownership in standard split optimizers.

Symptom & Reproduction

• Symptom: With TP4/EP4/ETP1, dense TP ranks 1–3 drop their distinct expert gradients from global statistics.
• Reproduction: The distributed test creates one dense parameter plus one allreduce=False expert parameter per rank under TP=EP=world size and ETP1. Before the fix, the expert child uses dense TP ownership, so only dense TP rank 0 retains its expert gradient.

In a DSv4 training step, the reported norm was 11.5135000849 versus the logical FP64 norm 12.2754343246. With clip_grad=1, this made the clipping coefficient 6.62% too large.

CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.run --standalone --nproc_per_node=4 \
  -m pytest -xvs tests/unit_tests/test_optimizer.py \
  -k expert_optimizer_uses_expert_tp_group_for_grad_norm

The trigger is the standard Adam/SGD non-FSDP path with separate dense and expert child optimizers when dense TP ownership differs from ETP ownership. Backward still computes the omitted expert gradients, and the optimizer still updates those parameters; only gradient statistics and clipping use the wrong ownership set.

Root Cause

1. get_megatron_optimizer creates separate dense and expert child optimizers.
2. _get_megatron_optimizer_based_on_param_groups assigns dense TP ownership to both children.
3. param_is_not_tensor_parallel_duplicate keeps replicated parameters only on that group's rank 0.
4. TP4/EP4/ETP1 places distinct expert shards on ranks filtered as dense replicas.

get_main_grads_for_grad_norm and count_zeros both consume the child optimizer's tp_group, so the same ownership error affects norm computation, clipping, and zero-count logging.

Fix

Add a child-specific tp_group override to _get_megatron_optimizer_based_on_param_groups and pass the expert tensor-parallel group when constructing the split expert optimizer. Dense children retain the dense TP group. The legacy path uses the global ETP group, an explicit ProcessGroupCollection uses expt_tp, and a custom collection without expt_tp retains its existing dense-TP fallback.

NVIDIA/Megatron-LM#5916 is the comprehensive upstream root fix, merged into NVIDIA main as cd4afffa. PR #80's base d075c1e does not contain that commit. This PR adapts only the split-optimizer ownership correction required by this fork; it does not backport NVIDIA#5916's gradient-synchronization, TE/native parameter-tagging, or parameter-norm changes. FSDP and Muon also remain outside this PR.

Verification

• test_expert_optimizer_uses_expert_tp_group_for_grad_norm on eight ranks: legacy and explicit modes bind expert gradients to ETP1 on every rank; custom_without_expt_tp preserves the prior dense-TP fallback.
• DSv4 TP4/EP4/ETP1 source-fix replay: the corrected reported norm was 12.2851015996 versus FP64 logical 12.2851021737 (4.67e-8 relative error), with zero ownership mismatches instead of 512 missing expert tensors on each of ranks 1–3.
• Qwen3-30B-A3B TP2/CP2/EP4/ETP1 integration on a Miles-compatible Megatron base whose optimizer source was byte-identical to this PR: two four-step runs completed; fixed reported norms matched independent FP64 logical reductions within 1.28e-7 relative error, while the old selector was 3.97% to 5.49% low across run 2.

Review Focus

• _get_megatron_optimizer_based_on_param_groups: the optional override changes only the selected child's duplicate-filter ownership.
• get_megatron_optimizer: legacy and explicit process-group paths select ETP for the expert child without changing dense ownership.
• custom_without_expt_tp: the compatibility fallback is intentionally unchanged; broader expert-topology fixes from Fix gradient-norm undercounting when using EP and TP NVIDIA/Megatron-LM#5916 remain outside this PR.