fix(training): correct parameter norm reporting with GTP

[yashaswikarnati]

What does this PR do ?

Correct parameter-norm reporting when generalized tensor parallelism (GTP) is enabled.

Changelog

• Count replicated dense and expert parameters once across their respective GTP axes.
• Keep every GTP weight shard while bypassing the tensor-parallel duplicate filter.
• Reduce the complete expert contribution across expt_gtp_remat, including replicated expert biases.
• Add focused dense-GTP and expert-GTP parameter-norm regressions, including replica filtering on a nonzero expert-GTP rank.

GitHub Actions CI

Local validation:

• TestCalcParamsL2Norm: 28 passed on each of two local processes.
• Expert-GTP rank-0/rank-1 replica-filter regression: 2 passed on each of two local processes.
• uv run pre-commit run --all-files: passed.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added the necessary tests.
• Documentation is not required because this does not change user-facing configuration or APIs.
• This change does not affect optional components.

Additional Information

The implementation follows the pinned Megatron-LM GTP duplicate-filtering semantics while reducing the full expert norm across EGTP so replicated expert parameters are reported consistently on every rank.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[claude]

Light review: GTP-aware param L2 norm

The change to calc_params_l2_norm looks correct and is well covered by the two new unit tests. The dense path relies on pg_collection.mp already including the GTP axis; the expert path adds an explicit all_reduce over expt_gtp_remat because tp_ep_pp excludes expert GTP. The GTP-vs-replica filtering (skip TP dedup for GTP params; drop non-GTP replicas on GTP rank != 0) is symmetric and reads cleanly.

One thing to confirm: the pre-existing tests now hit the new get_pg_rank calls unpatched. The autouse _PG fixture sets gtp_remat to object() and expt_gtp_remat to None, but test_single_model_fp32, test_list_of_models, and the bf16 tests do NOT monkeypatch get_pg_rank. They only pass if real get_pg_rank(object()) and get_pg_rank(None) both return 0. Please confirm these CUDA-gated tests still pass end-to-end; if get_pg_rank calls .rank() on the fake object() they would raise or silently zero the norm. Patching get_pg_rank in those tests too (as the new tests do) would make this robust.

Suggested test cases:

• test_gtp_param_norm_filters_replicas_and_keeps_shards
• test_expert_gtp_param_norm_reduces_over_expert_gtp rank-zero-keeps-replica
• test_expert_gtp_param_norm_reduces_over_expert_gtp nonzero-rank-drops-replica
• test_single_model_fp32 (regression path, now exercises unpatched get_pg_rank)
• test_list_of_models (regression path)
• test_bf16_mode_without_main_param (regression path)
• test_bf16_mode_with_main_param (regression path)

No perf tests impacted.