From c1702d851d9086415e6a942224220ddfdca57e60 Mon Sep 17 00:00:00 2001 From: richardoo-707 Date: Tue, 27 Oct 2026 10:15:37 +0800 Subject: [PATCH] [moe] resolve the default EP group before tagging parameters --- .github/scripts/ci/run_colossalai_batch.sh | 1 + .github/scripts/ci/run_cpu_unit_tests.sh | 1 + colossalai/moe/layer.py | 2 ++ tests/test_moe/test_ep_moe_layer.py | 26 ++++++++++++++----- tests/test_moe/test_ep_moe_layer_cpu.py | 29 ++++++++++++++++++++++ 5 files changed, 53 insertions(+), 6 deletions(-) create mode 100644 tests/test_moe/test_ep_moe_layer_cpu.py diff --git a/.github/scripts/ci/run_colossalai_batch.sh b/.github/scripts/ci/run_colossalai_batch.sh index 01db482cca27..fd91efd81b40 100644 --- a/.github/scripts/ci/run_colossalai_batch.sh +++ b/.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 diff --git a/.github/scripts/ci/run_cpu_unit_tests.sh b/.github/scripts/ci/run_cpu_unit_tests.sh index 2677eccf2d84..5d2d08071a88 100644 --- a/.github/scripts/ci/run_cpu_unit_tests.sh +++ b/.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 diff --git a/colossalai/moe/layer.py b/colossalai/moe/layer.py index 7dbffa31d11b..4e5d6745a0f9 100644 --- a/colossalai/moe/layer.py +++ b/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 diff --git a/tests/test_moe/test_ep_moe_layer.py b/tests/test_moe/test_ep_moe_layer.py index c95f55d04f9d..a1c35b1fcf07 100644 --- a/tests/test_moe/test_ep_moe_layer.py +++ b/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( diff --git a/tests/test_moe/test_ep_moe_layer_cpu.py b/tests/test_moe/test_ep_moe_layer_cpu.py new file mode 100644 index 000000000000..3a716d4150a3 --- /dev/null +++ b/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()