From 04fc52650de3b5b95b0f1f217151ab21ca1e0410 Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Tue, 28 Jul 2026 21:14:02 -0700 Subject: [PATCH] Fix pyre type errors in parallelism.py (reap OMH type-checking FAILURE) Summary: moe:parallelism-type-checking has 42 consecutive FAILUREs (DISABLED_FAILING) on the fbgemm_dev OMH dashboard with detail 'type errors!'. torch.distributed.new_group is typed Union[int, ProcessGroup] in the stubs, but _ROUTED_EXPERTS_MP_GROUP and _EP_GROUP are Optional[ProcessGroup], so the assignments failed pyre [9] Incompatible variable type. Narrow the returned group with assert isinstance(group, ProcessGroup) at both call sites (matches the file's existing assert style; new_group returns a ProcessGroup here). Differential Revision: D113992766 --- fbgemm_gpu/experimental/gen_ai/test/moe/parallelism.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/fbgemm_gpu/experimental/gen_ai/test/moe/parallelism.py b/fbgemm_gpu/experimental/gen_ai/test/moe/parallelism.py index b4a51576b0..33fae8dfc1 100644 --- a/fbgemm_gpu/experimental/gen_ai/test/moe/parallelism.py +++ b/fbgemm_gpu/experimental/gen_ai/test/moe/parallelism.py @@ -115,6 +115,9 @@ def init_parallel( ): ranks = list(range(base_rank, base_rank + mp_size_for_routed_experts)) group = torch.distributed.new_group(ranks, timeout=timeout) + # new_group is typed Union[int, ProcessGroup] in the stubs; it returns a + # ProcessGroup here. Narrow it so the Optional[ProcessGroup] global assigns. + assert isinstance(group, ProcessGroup) if global_rank in ranks: _ROUTED_EXPERTS_MP_GROUP = group @@ -126,5 +129,8 @@ def init_parallel( for i in range(num_ep_groups): ranks = list(range(i, get_world_size(), num_ep_groups)) group = torch.distributed.new_group(ranks, timeout=timeout) + # new_group is typed Union[int, ProcessGroup] in the stubs; it returns a + # ProcessGroup here. Narrow it so the Optional[ProcessGroup] global assigns. + assert isinstance(group, ProcessGroup) if global_rank in ranks: _EP_GROUP = group