Add DeepEP fallback logic and tests

Signed-off-by: ooooo <[email protected]>

Author: ooooo-create

nemo_automodel/components/moe/layers.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import warnings
 from dataclasses import dataclass
 from functools import partial
 from typing import Literal, Optional
@@ -20,6 +20,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from nemo_automodel.components.distributed.init_utils import get_world_size_safe
 from nemo_automodel.components.moe.utils import BackendConfig, initialize_linear_module
 from nemo_automodel.shared.utils import dtype_from_str
 
@@ -914,7 +915,15 @@ def __init__(self, config: MoEConfig, backend: BackendConfig):
             self.gate = FakeBalancedGate(config)
         else:
             self.gate = Gate(config, gate_precision=backend.gate_precision)
-        if backend.enable_deepep:
+        if backend.enable_deepep and get_world_size_safe() == 1:
+            warnings.warn(
+                "DeepEP is enabled in config, but world size is 1. "
+                "DeepEP requires multiple GPUs. Falling back to standard GroupedExperts.",
+                category=UserWarning,
+                stacklevel=2,
+            )
+            self.experts = GroupedExperts(config)
+        elif backend.enable_deepep:
             self.experts = GroupedExpertsDeepEP(config)
         else:
             self.experts = GroupedExperts(config)

tests/unit_tests/moe/test_layers.py

@@ -1258,10 +1258,20 @@ def test_moe_init_with_fake_balanced_gate(self, moe_config, backend_config):
         assert isinstance(moe.gate, FakeBalancedGate)
         assert isinstance(moe.experts, GroupedExperts)
 
-    def test_moe_init_with_deepep(self, moe_config, backend_config):
-        """Test MoE initialization with DeepEP."""
+    def test_moe_init_with_deepep_single_device(self, moe_config, backend_config):
+        """DeepEP enabled but world size == 1 should fall back to GroupedExperts."""
         backend_config.enable_deepep = True
-        moe = MoE(moe_config, backend_config)
+        with patch("nemo_automodel.components.moe.layers.get_world_size_safe", return_value=1):
+            moe = MoE(moe_config, backend_config)
+
+        assert isinstance(moe.gate, Gate)
+        assert isinstance(moe.experts, GroupedExperts)
+
+    def test_moe_init_with_deepep_multi_device(self, moe_config, backend_config):
+        """DeepEP enabled and world size > 1 should use GroupedExpertsDeepEP."""
+        backend_config.enable_deepep = True
+        with patch("nemo_automodel.components.moe.layers.get_world_size_safe", return_value=2):
+            moe = MoE(moe_config, backend_config)
 
         assert isinstance(moe.gate, Gate)
         assert isinstance(moe.experts, GroupedExpertsDeepEP)