Skip to content

[BUG] Default EPMoELayer group is rejected by get_ep_size #6485

Open

Description

@richardoo-707

The new EPMoELayer in #6477 accepts ep_group=None as the default WORLD group, but writes None into its expert parameters' ep_group attribute. get_ep_size(param) rejects that attribute with an assertion, even after the distributed process group has been initialized.

This affects feat/moe-ep-layer at 055853017e23c778402b2c8d965838cb09fb43e4; the new layer is not yet on main.

Reproduction in a CPU/Gloo environment:

import tempfile
from pathlib import Path

import torch.distributed as dist
from colossalai.moe import EPMoELayer, ReferenceExperts
from colossalai.tensor.moe_tensor.api import get_ep_size

with tempfile.TemporaryDirectory() as directory:
    dist.init_process_group(
        "gloo", init_method=(Path(directory) / "rendezvous").as_uri(), rank=0, world_size=1
    )
    try:
        layer = EPMoELayer(4, 8, 4, experts=ReferenceExperts(), device="cpu")
        assert get_ep_size(layer.gate_up_proj) == 1
    finally:
        dist.destroy_process_group()

Expected: the layer, its default dispatcher and its expert parameter tags consistently identify WORLD, and the getter returns 1. Actual: AssertionError: The tensor does not have expert parallel group.

The correction should preserve explicit subgroups and construction before distributed initialization. Verified environment: Ubuntu 22.04.5, Python 3.10.12, Torch 2.13.0+cu130, zero visible GPUs.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions