[moe] resolve the default EP group before tagging parameters

.github/scripts/ci/run_colossalai_batch.sh

@@ -75,6 +75,7 @@ case "${BATCH}" in
             tests/test_config/test_load_config.py
             tests/test_infer/test_async_engine/test_request_tracer.py
             tests/test_lazy/test_chatglm_cpu.py
+            tests/test_moe/test_ep_moe_layer_cpu.py
             tests/test_optimizer/test_lr_scheduler.py
             tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py
             tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py

.github/scripts/ci/run_cpu_unit_tests.sh

@@ -13,6 +13,7 @@ TESTS=(
     tests/test_config/test_load_config.py
     tests/test_infer/test_async_engine/test_request_tracer.py
     tests/test_lazy/test_chatglm_cpu.py
+    tests/test_moe/test_ep_moe_layer_cpu.py
     tests/test_optimizer/test_lr_scheduler.py
     tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py
     tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py

colossalai/moe/layer.py

@@ -67,6 +67,8 @@ def __init__(
         dtype: Optional[torch.dtype] = None,
     ):
         super().__init__()
+        if ep_group is None and dist.is_initialized():
+            ep_group = dist.group.WORLD
         self.ep_group = ep_group
         self.ep_size = dist.get_world_size(ep_group) if dist.is_initialized() else 1
         self.ep_rank = dist.get_rank(ep_group) if dist.is_initialized() else 0

tests/test_moe/test_ep_moe_layer.py

@@ -5,7 +5,13 @@
 
 import colossalai
 from colossalai.moe import EPMoELayer, GroupedMMExperts, NCCLAllToAllDispatcher, ReferenceExperts
-from colossalai.tensor.moe_tensor.api import ExpertGradScale, get_ep_group, get_expert_grad_scale, is_moe_tensor
+from colossalai.tensor.moe_tensor.api import (
+    ExpertGradScale,
+    get_ep_group,
+    get_ep_size,
+    get_expert_grad_scale,
+    is_moe_tensor,
+)
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
 
 # Relative L2 error allowed between the bf16 layer and an fp32 single-GPU reference. The largest error over all checks
@@ -152,15 +158,23 @@ def run_layer_vs_reference(ep_size: int, backend: str, grad_scale: ExpertGradSca
                 assert err < RTOL, f"{config}: {name}[{expert}] grad {err:.3e}"
 
 
-def run_param_contract():
+@parameterize("use_default_group", [False, True])
+def run_param_contract(use_default_group: bool):
     """Expert params carry their EP tags, and loading weights keeps the parameter objects."""
-    group = ep_groups(2)
+    group = None if use_default_group else ep_groups(2)
+    expected_group = dist.group.WORLD if use_default_group else group
+    ep_size = dist.get_world_size(expected_group)
+    ep_rank = dist.get_rank(expected_group)
     layer = EPMoELayer(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, ep_group=group, device="cuda")
-    assert layer.num_local_experts == NUM_EXPERTS // 2
-    assert layer.local_expert_range == range(dist.get_rank(group) * 4, dist.get_rank(group) * 4 + 4)
+    assert layer.ep_group is expected_group
+    assert layer.dispatcher.ep_group is expected_group
+    assert layer.num_local_experts == NUM_EXPERTS // ep_size
+    start = ep_rank * layer.num_local_experts
+    assert layer.local_expert_range == range(start, start + layer.num_local_experts)
     params = (layer.gate_up_proj, layer.down_proj)
     for param in params:
-        assert is_moe_tensor(param) and get_ep_group(param) is group
+        assert is_moe_tensor(param) and get_ep_group(param) is expected_group
+        assert get_ep_size(param) == ep_size
         assert get_expert_grad_scale(param) == ExpertGradScale.IN_FORWARD
     ids = [id(p) for p in params]
     layer.load_full_weights(

tests/test_moe/test_ep_moe_layer_cpu.py

@@ -0,0 +1,29 @@
+import pytest
+import torch.distributed as dist
+
+from colossalai.moe import EPMoELayer, ReferenceExperts
+from colossalai.tensor.moe_tensor.api import ExpertGradScale, get_ep_group, get_ep_size, get_expert_grad_scale
+
+
+@pytest.mark.parametrize("initialized", [False, True])
+def test_ep_moe_layer_default_process_group(initialized: bool, tmp_path):
+    """Default-group expert tags work with the parameter API without changing local construction."""
+    assert not dist.is_initialized()
+    if initialized:
+        dist.init_process_group("gloo", init_method=(tmp_path / "rendezvous").as_uri(), rank=0, world_size=1)
+    try:
+        layer = EPMoELayer(4, 8, 4, experts=ReferenceExperts(), device="cpu")
+        expected_group = dist.group.WORLD if initialized else None
+        assert layer.ep_group is expected_group
+        assert layer.dispatcher.ep_group is expected_group
+        assert layer.ep_size == 1
+        assert layer.ep_rank == 0
+        assert layer.num_local_experts == 4
+        for param in layer.parameters():
+            assert get_ep_group(param) is expected_group
+            assert get_expert_grad_scale(param) == ExpertGradScale.IN_FORWARD
+            if initialized:
+                assert get_ep_size(param) == 1
+    finally:
+        if initialized:
+            dist.destroy_process_group()