Merge branch 'moe'

pipegoose/nn/expert_parallel/layers.py

@@ -0,0 +1,5 @@
+from torch import nn
+
+
+class MoELayer(nn.Module):
+    pass

pipegoose/nn/expert_parallel/routers.py

@@ -1,6 +1,28 @@
 """DON'T USE THIS MODULE: Under development."""
 from abc import ABC
+from enum import Enum, auto
 
+from torch import nn
 
-class MoERouter(ABC):
+
+class RouterType(Enum):
+    """An enum for router types."""
+
+    TOP_1 = auto()
+    TOP_2 = auto()
+
+
+class Router(ABC):
+    pass
+
+
+class Top1Router(Router, nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, inputs):
+        pass
+
+
+def get_router(router_type: RouterType) -> Router:
     pass

tests/nn/expert_parallel/test_router.py

@@ -0,0 +1,20 @@
+import pytest
+import torch
+
+from pipegoose.nn.expert_parallel.routers import RouterType, get_router
+
+
+@pytest.mark.parametrize("router_type", [RouterType.TOP_1])
+def test_topk_router(router_type):
+    SEQ_LEN = 10
+    HIDDEN_DIM = 64
+    N_EXPERTS = 5
+
+    inputs = torch.randn(SEQ_LEN, HIDDEN_DIM)
+    router = get_router(router_type)(
+        n_experts=N_EXPERTS,
+    )
+
+    outputs = router(inputs)
+
+    assert isinstance(outputs, torch.Tensor)