def get_megablocks_device_mesh(
    device_mesh_cfg: Optional[tuple[int, ...]],
    moe_world_size: int,
    world_size: int,
) -> DeviceMesh:
    """Helper function to get the device mesh for MegaBlocks MoE.

    Args:
        device_mesh_cfg (Optional[tuple[int, ...]]): The device mesh configuration specification.
        moe_world_size (int): The MoE world size.
        world_size (int): The world size.

    Raises:
        ValueError: If the device mesh configuration is not valid.

    Returns:
        The device mesh for MegaBlocks MoE.
    """
    from torch.distributed._tensor.device_mesh import init_device_mesh

    if device_mesh_cfg is None or len(device_mesh_cfg) == 1:
        if device_mesh_cfg is not None:
            world_size = device_mesh_cfg[0]
        sharding_group_dim = world_size // moe_world_size
        device_mesh = init_device_mesh(
            'cuda',
            (sharding_group_dim, moe_world_size),
            mesh_dim_names=('weight_parallel', 'expert_parallel'),
        )
    else:
        raise ValueError(f'{device_mesh_cfg=} must be length 1')

    return device_mesh
Helper function to get the device mesh for MegaBlocks MoE.

    Args:
        device_mesh_cfg (Optional[tuple[int, ...]]): The device mesh configuration specification.
        moe_world_size (int): The MoE world size.
        world_size (int): The world size.

    Raises:
        ValueError: If the device mesh configuration is not valid.

    Returns:
        The device mesh for MegaBlocks MoE.
    
llmfoundry/models/utils/config_moe_args.py