Allow passing in lbl_process_group directly

llmfoundry/models/utils/config_moe_args.py

@@ -177,9 +177,9 @@ def config_megablocks_moe_args(
                 lbl_process_group = create_set_process_group(lbl_process_group)
             else:
                 lbl_process_group = None
-        elif lbl_process_group is not None:
+        elif not isinstance(lbl_process_group, distributed.ProcessGroup):
             raise ValueError(
-                f'Unknown {lbl_process_group=}. Options are: none | expert_group | global_group | <GROUP_SIZE>.',
+                f'Unknown {lbl_process_group=}. Options are: none | a process group | ``expert_group`` | ``global_group`` | <GROUP_SIZE>.',
             )
         ffn_config['lbl_process_group'] = lbl_process_group
 

tests/models/utils/test_config_moe_args.py

@@ -0,0 +1,45 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+import pytest
+from torch import distributed
+
+from llmfoundry.models.utils.config_moe_args import (
+    config_megablocks_moe_args,
+    get_megablocks_device_mesh,
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.world_size(2)
+def test_config_megablocks_moe_args_pg():
+    ffn_config_base: dict[str, Any] = {
+        'moe_world_size': 1,
+        'ffn_type': 'mb_moe',
+        'fc_type': 'torch',
+    }
+
+    ffn_config_str = ffn_config_base.copy()
+    ffn_config_str['lbl_process_group'] = 'global_group'
+
+    ffn_config_pg = ffn_config_base.copy()
+    ffn_config_pg['lbl_process_group'] = distributed.group.WORLD
+
+    output_str = config_megablocks_moe_args(
+        ffn_config=ffn_config_str,
+        d_model=128,
+        expansion_ratio=4,
+        n_layers=2,
+        get_device_mesh=get_megablocks_device_mesh,
+    )
+    output_pg = config_megablocks_moe_args(
+        ffn_config=ffn_config_pg,
+        d_model=128,
+        expansion_ratio=4,
+        n_layers=2,
+        get_device_mesh=get_megablocks_device_mesh,
+    )
+
+    assert output_str['lbl_process_group'] == output_pg['lbl_process_group']