Thread custom process groups through MoE grad finalization

[yashaswikarnati]

Summary

• use explicit MoE tensor-parallel groups for TE checkpointing and shared-expert TP collectives
• allow router expert-bias finalization to reduce on an explicit TPxDPxCP process group
• require explicit pg_collection.tp_dp_cp when expert-bias updates run with custom process groups
• add focused coverage for group forwarding and early validation

Testing

• cog batch job 11732315: python -m pytest tests/unit_tests/transformer/moe/test_routers.py::test_get_updated_expert_bias_uses_explicit_group tests/unit_tests/distributed/test_finalize_model_grads.py::test_update_router_expert_bias_uses_explicit_group tests/unit_tests/distributed/test_finalize_model_grads.py::test_finalize_model_grads_uses_pg_collection_tp_dp_cp tests/unit_tests/distributed/test_finalize_model_grads.py::test_finalize_model_grads_requires_tp_dp_cp_for_explicit_groups -q
• pre-push hooks: black, pylint, isort

[yashaswikarnati]

we can remove TODO on 1183 ?

[yashaswikarnati]

i dont like fake tests with monkey patches.. ci actually runs the test on 8 gpus.. is there a way of less verbose and actua testing

[yashaswikarnati]

Addressed in 53f2fa5. Removed the monkeypatch-based finalize tests and replaced them with real distributed coverage: the tests initialize model parallel, build actual ProcessGroupCollection values, run finalize_model_grads, and verify router expert-bias update / early validation behavior. Verified on an 8-GPU cog job: 11732530.

[yashaswikarnati]

here