Fix AG groups for custom parallel order

[MrAnayDongre]

Summary

Fixes #4261.

create_all_gather_groups() was reconstructing RankGenerator instances with hardcoded order='tp-cp-ep-dp-pp' and rank_offset=0. This could produce AG groups whose rank membership did not match the process groups created by initialize_model_parallel() when model parallel was initialized with a non-default order or rank offset.

This change caches the decoder and expert decoder RankGenerator instances created during initialize_model_parallel() and reuses them when creating regular and expert AG groups.

Changes

• Cache init-time decoder and expert decoder RankGenerator instances in parallel_state.py.
• Reuse the cached generators in create_all_gather_groups() instead of reconstructing them with hardcoded order/rank offset.
• Reset the cached generators in destroy_model_parallel().
• Add regression coverage for custom rank-generation state.
• Extend existing AG group tests over non-default parallel order.
• Update the expert AG test to exercise create_all_gather_groups(for_expert_parallelism=True) directly.

Testing

• python -m py_compile megatron/core/parallel_state.py tests/unit_tests/test_parallel_state.py
• BASE_REF=main CHECK_ONLY=true bash tools/autoformat.sh
• tests/unit_tests/test_parallel_state.py::test_create_all_gather_groups_uses_cached_rank_generators in an NVIDIA PyTorch base container using system Python/system torch
• Collect-only for all-gather/cached-rank-generator tests in the NVIDIA PyTorch base container
• Single-rank torch.distributed.run smoke for cached_rank_generators in the NVIDIA PyTorch base container
• Multi-rank NCCL all-gather tests not run locally because local GPU_COUNT=1; requires Megatron CI
• Full test suite / functional tests not run

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[MrAnayDongre]

Ready for review.

Fixes #4261. Local validation passed for the CPU-safe regression and single-rank smoke path; multi-rank NCCL validation needs Megatron CI.

[guihong-nv]

No blocking issue found; I noted that multi-rank NCCL validation remains outstanding.

[MrAnayDongre]

Thanks for the review.

Agreed that multi-rank NCCL validation remains the outstanding piece. I do not have local multi-GPU hardware to run that path, but the targeted CPU-safe regression and single-rank distributed smoke passed locally. I am happy to address any CI failures or requested changes once NVIDIA CI is approved and run.

[guihong-nv]

/ok to test 27891de

[MrAnayDongre]

Rebased on the latest main and resolved the parallel_state.py conflict while preserving the initialized RankGenerator behavior for regular and expert AG groups.

The conflict came from #4967 (GTP), which added the gtp_remat axis and _inject_gtp_remat_axis(...) to the generators rebuilt inside create_all_gather_groups(). Reusing the generators cached by initialize_model_parallel() subsumes that change: they are already built with gtp_remat/expert_gtp_remat, the injected order, and the caller's order and rank_offset, so the GTP behavior is preserved and no longer depends on a hardcoded order string.

New HEAD: 9439662c2

Local validation passed (run in nvcr.io/nvidia/pytorch:26.06-py3):

• python3 -m py_compile megatron/core/parallel_state.py tests/unit_tests/test_parallel_state.py
• CHECK_ONLY=true bash tools/autoformat.sh — black/isort clean, pylint 10.00/10, ruff clean
• targeted CPU-safe cached-rank-generator regression
• relevant test collection
• single-rank distributed smoke

I also re-checked the regression property against the new upstream structure: restoring the hardcoded order='tp-cp-ep-dp-pp' / rank_offset=0 reconstruction makes the test fail on real rank membership (expected (4, 8, 12, 16) vs default-order (0, 2, 4, 6)), so it still guards the original bug.

Multi-rank NCCL validation still requires NVIDIA CI.