[optimizer] Fix LayerWise expert ownership

[yushengsu-thu]

Summary

Fix LayerWiseDistributedOptimizer when dense and expert parameters have different data-parallel ownership groups.

In the reproduced TP2 / ETP1 / EP1 / DP1 setup, dp_cp has one rank while expt_dp has two. The old singleton fast path skipped sharding entirely, leaving a full expert optimizer copy on both ranks. World-reduced gradient statistics then counted the same expert replicas twice, making the norm sqrt(2) too large and clipping too strong.

The fix shards dense and expert parameters independently, handles pure-expert and empty ownership shards in all-gather/broadcast, filters expert replicas with the ETP group, and preserves the required topology metadata on child optimizers and BF16 master parameters. It assumes source expert parameters are already tagged correctly and is not a full upstream backport.

Validation

Manually tested on 4×H200 at 3d107754:

• CPU ownership and empty-shard regressions;
• NCCL topologies: TP2 / ETP1, TP4 / ETP2 / EP1, and TP2 / ETP1 / EP2;
• full 4-GPU LayerWise suite: 21 passed per rank;
• single-GPU optimizer suite: 16 passed, 40 skipped.

The distributed tests cover unique ownership, exact norm and clipping, zero counts, parameter updates, and expert-DP replica equality.

Related work

• Fix gradient-norm undercounting when using EP and TP NVIDIA/Megatron-LM#5916 ：already merge to upstream main
  expert-topology correctness fix。 (on main)
• radixark/Megatron-LM#80 is the corresponding standard split-optimizer fix on main; this PR is the LayerWise counterpart on miles-main.
• NVIDIA/Megatron-LM#1965 introduced the LayerWise sharding/all-gather path later promoted to main in #2241. This PR fixes the singleton and empty-shard cases in that path.
• NVIDIA/Megatron-LM#5916 is the complete upstream expert-topology fix. This PR ports only the LayerWise-relevant pieces.

[yushengsu-thu]

@codex review