feat: support shared experts layer and tests

tests/test_mb_moe_shared_expert.py

@@ -0,0 +1,139 @@
+import torch
+import megablocks
+from megablocks.layers import MegaBlocksMoeMLPWithSharedExpert, create_shared_expert_weights
+
+
+def test_megablocks_moe_mlp_with_shared_expert_import():
+    mlp = MegaBlocksMoeMLPWithSharedExpert()
+    assert hasattr(mlp, 'shared_up_proj_weight')
+    assert hasattr(mlp, 'shared_down_proj_weight')
+    assert hasattr(mlp, 'set_shared_expert_weights')
+
+
+def test_set_shared_expert_weights():
+    mlp = MegaBlocksMoeMLPWithSharedExpert()
+    
+    hidden_size = 128
+    shared_expert_hidden_size = 256
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    dtype = torch.float32
+    
+    up_proj_weight = torch.randn(shared_expert_hidden_size, hidden_size, device=device, dtype=dtype)
+    down_proj_weight = torch.randn(hidden_size, shared_expert_hidden_size, device=device, dtype=dtype)
+    up_proj_bias = torch.randn(shared_expert_hidden_size, device=device, dtype=dtype)
+    down_proj_bias = torch.randn(hidden_size, device=device, dtype=dtype)
+    
+    mlp.set_shared_expert_weights(
+        up_proj_weight=up_proj_weight,
+        down_proj_weight=down_proj_weight,
+        up_proj_bias=up_proj_bias,
+        down_proj_bias=down_proj_bias,
+        weighted_sum=True,
+        activation_fn=torch.nn.functional.gelu
+    )
+    
+    assert torch.equal(mlp.shared_up_proj_weight, up_proj_weight)
+    assert torch.equal(mlp.shared_down_proj_weight, down_proj_weight)
+    assert torch.equal(mlp.shared_up_proj_bias, up_proj_bias)
+    assert torch.equal(mlp.shared_down_proj_bias, down_proj_bias)
+    assert mlp.shared_expert_weighted_sum == True
+    assert mlp.shared_activation_fn == torch.nn.functional.gelu
+
+
+def test_create_shared_expert_weights():
+    hidden_size = 128
+    shared_expert_hidden_size = 256
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    dtype = torch.float32
+    
+    def init_method(tensor):
+        torch.nn.init.xavier_uniform_(tensor)
+    
+    up_proj_weight, down_proj_weight, up_proj_bias, down_proj_bias = create_shared_expert_weights(
+        hidden_size=hidden_size,
+        shared_expert_hidden_size=shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+        init_method=init_method
+    )
+    
+    assert up_proj_weight.shape == (shared_expert_hidden_size, hidden_size)
+    assert down_proj_weight.shape == (hidden_size, shared_expert_hidden_size)
+    assert up_proj_weight.device.type == device.type
+    assert down_proj_weight.device.type == device.type
+    assert up_proj_weight.dtype == dtype
+    assert down_proj_weight.dtype == dtype
+    assert up_proj_bias is None
+    assert down_proj_bias is None
+
+
+def test_shared_expert_weights_none_by_default():
+    mlp = MegaBlocksMoeMLPWithSharedExpert()
+    
+    assert mlp.shared_up_proj_weight is None
+    assert mlp.shared_down_proj_weight is None
+    assert mlp.shared_up_proj_bias is None
+    assert mlp.shared_down_proj_bias is None
+    assert mlp.shared_expert_weighted_sum == False
+    assert mlp.shared_activation_fn is None
+
+
+def test_inheritance_from_megablocks_moe_mlp():
+    mlp = MegaBlocksMoeMLPWithSharedExpert()
+    
+    from megablocks.layers import MegaBlocksMoeMLP
+    assert isinstance(mlp, MegaBlocksMoeMLP)
+    assert hasattr(mlp, 'forward')
+
+
+def test_shared_expert_weights_custom_init():
+    hidden_size = 64
+    shared_expert_hidden_size = 128
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    dtype = torch.float16
+    
+    def custom_init(tensor):
+        torch.nn.init.constant_(tensor, 0.5)
+    
+    def custom_output_init(tensor):
+        torch.nn.init.constant_(tensor, 0.1)
+    
+    up_proj_weight, down_proj_weight, up_proj_bias, down_proj_bias = create_shared_expert_weights(
+        hidden_size=hidden_size,
+        shared_expert_hidden_size=shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+        init_method=custom_init,
+        output_layer_init_method=custom_output_init
+    )
+    
+    assert torch.all(up_proj_weight == 0.5)
+    assert torch.all(down_proj_weight == 0.1)
+    assert up_proj_weight.dtype == dtype
+    assert down_proj_weight.dtype == dtype
+
+
+def test_shared_expert_weights_dimensions():
+    mlp = MegaBlocksMoeMLPWithSharedExpert()
+    
+    batch_size = 4
+    seq_len = 16
+    hidden_size = 128
+    shared_expert_hidden_size = 256
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    up_proj_weight = torch.randn(shared_expert_hidden_size, hidden_size, device=device)
+    down_proj_weight = torch.randn(hidden_size, shared_expert_hidden_size, device=device)
+    
+    mlp.set_shared_expert_weights(
+        up_proj_weight=up_proj_weight,
+        down_proj_weight=down_proj_weight
+    )
+    
+    x = torch.randn(seq_len, batch_size, hidden_size, device=device)
+    
+    expected_up_output_shape = (seq_len, batch_size, shared_expert_hidden_size)
+    expected_down_output_shape = (seq_len, batch_size, hidden_size)
+    
+    assert up_proj_weight.shape[1] == x.shape[-1]
+    assert down_proj_weight.shape[0] == x.shape[-1]

tests/test_mb_moe_shared_expert_multi.py

@@ -0,0 +1,200 @@
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import os
+import pytest
+from megablocks.layers import MegaBlocksMoeMLPWithSharedExpert, create_shared_expert_weights
+
+
+def run_distributed_shared_expert_test(rank, world_size):
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12356"
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+
+    dist.init_process_group(
+        backend="gloo",
+        rank=rank,
+        world_size=world_size,
+    )
+
+    model = MegaBlocksMoeMLPWithSharedExpert()
+
+    hidden_size = 128
+    shared_expert_hidden_size = 192
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def simple_init(tensor):
+        torch.nn.init.xavier_uniform_(tensor)
+
+    shared_up_proj_weight, shared_down_proj_weight, shared_up_proj_bias, shared_down_proj_bias = create_shared_expert_weights(
+        hidden_size=hidden_size,
+        shared_expert_hidden_size=shared_expert_hidden_size,
+        device=torch.device(device),
+        dtype=torch.float32,
+        init_method=simple_init
+    )
+
+    model.set_shared_expert_weights(
+        up_proj_weight=shared_up_proj_weight,
+        down_proj_weight=shared_down_proj_weight,
+        up_proj_bias=shared_up_proj_bias,
+        down_proj_bias=shared_down_proj_bias,
+        weighted_sum=True,
+        activation_fn=torch.nn.functional.gelu
+    )
+
+    assert model.shared_up_proj_weight is not None, f"Shared up proj weight not set on rank {rank}"
+    assert model.shared_down_proj_weight is not None, f"Shared down proj weight not set on rank {rank}"
+    assert model.shared_expert_weighted_sum == True, f"Weighted sum not set correctly on rank {rank}"
+    
+    print(f"Rank {rank}: Shared expert setup test passed!")
+
+    dist.destroy_process_group()
+
+
+def run_distributed_shared_expert_weighted_sum_test(rank, world_size):
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12357"
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+
+    dist.init_process_group(
+        backend="gloo",
+        rank=rank,
+        world_size=world_size,
+    )
+
+    model = MegaBlocksMoeMLPWithSharedExpert()
+
+    hidden_size = 64
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def simple_init(tensor):
+        torch.nn.init.xavier_uniform_(tensor)
+
+    shared_up_proj_weight, shared_down_proj_weight, _, _ = create_shared_expert_weights(
+        hidden_size=hidden_size,
+        shared_expert_hidden_size=96,
+        device=torch.device(device),
+        dtype=torch.float32,
+        init_method=simple_init
+    )
+
+    model.set_shared_expert_weights(
+        up_proj_weight=shared_up_proj_weight,
+        down_proj_weight=shared_down_proj_weight,
+        weighted_sum=False,
+        activation_fn=torch.nn.functional.relu
+    )
+
+    assert model.shared_up_proj_weight is not None, f"Shared up proj weight not set on rank {rank}"
+    assert model.shared_down_proj_weight is not None, f"Shared down proj weight not set on rank {rank}"
+    assert model.shared_expert_weighted_sum == False, f"Weighted sum not set correctly on rank {rank}"
+    assert model.shared_activation_fn == torch.nn.functional.relu, f"Activation function not set correctly on rank {rank}"
+    
+    print(f"Rank {rank}: Weighted sum setup test passed!")
+
+    dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("world_size", [1, 2, 4, 8])
+def test_shared_expert_distributed_functionality(world_size):
+    if world_size == 1:
+        # Single process test
+        model = MegaBlocksMoeMLPWithSharedExpert()
+        
+        hidden_size = 128
+        shared_expert_hidden_size = 192
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        def simple_init(tensor):
+            torch.nn.init.xavier_uniform_(tensor)
+
+        shared_up_proj_weight, shared_down_proj_weight, shared_up_proj_bias, shared_down_proj_bias = create_shared_expert_weights(
+            hidden_size=hidden_size,
+            shared_expert_hidden_size=shared_expert_hidden_size,
+            device=torch.device(device),
+            dtype=torch.float32,
+            init_method=simple_init
+        )
+
+        model.set_shared_expert_weights(
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            weighted_sum=True,
+            activation_fn=torch.nn.functional.gelu
+        )
+
+        assert model.shared_up_proj_weight is not None, "Shared up proj weight not set"
+        assert model.shared_down_proj_weight is not None, "Shared down proj weight not set"
+        assert model.shared_expert_weighted_sum == True, "Weighted sum not set correctly"
+        
+        print("Single process shared expert setup test passed!")
+    else:
+        # Multi-process test
+        mp.spawn(run_distributed_shared_expert_test, args=(world_size,), nprocs=world_size, join=True)
+        print("Multi-process shared expert test completed successfully!")
+
+
+@pytest.mark.parametrize("world_size", [1, 2, 4, 8])
+def test_shared_expert_distributed_weighted_sum(world_size):
+    if world_size == 1:
+        # Single process test
+        model = MegaBlocksMoeMLPWithSharedExpert()
+
+        hidden_size = 64
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        def simple_init(tensor):
+            torch.nn.init.xavier_uniform_(tensor)
+
+        shared_up_proj_weight, shared_down_proj_weight, _, _ = create_shared_expert_weights(
+            hidden_size=hidden_size,
+            shared_expert_hidden_size=96,
+            device=torch.device(device),
+            dtype=torch.float32,
+            init_method=simple_init
+        )
+
+        model.set_shared_expert_weights(
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            weighted_sum=False,
+            activation_fn=torch.nn.functional.relu
+        )
+
+        assert model.shared_up_proj_weight is not None, "Shared up proj weight not set"
+        assert model.shared_down_proj_weight is not None, "Shared down proj weight not set"
+        assert model.shared_expert_weighted_sum == False, "Weighted sum not set correctly"
+        assert model.shared_activation_fn == torch.nn.functional.relu, "Activation function not set correctly"
+        
+        print("Single process weighted sum setup test passed!")
+    else:
+        # Multi-process test
+        mp.spawn(run_distributed_shared_expert_weighted_sum_test, args=(world_size,), nprocs=world_size, join=True)
+        print("Multi-process shared expert weighted sum test completed successfully!")
+
+
+def test_shared_expert_single_process():
+    model = MegaBlocksMoeMLPWithSharedExpert()
+    
+    assert model.shared_up_proj_weight is None
+    assert model.shared_down_proj_weight is None
+    assert hasattr(model, 'set_shared_expert_weights')
+    
+    print("Single process shared expert basic test passed!")
+
+
+if __name__ == "__main__":
+    test_shared_expert_single_process()
+    print("Single process test passed!")
+    
+    os.environ['WORLD_SIZE'] = '2'
+    test_shared_expert_distributed_functionality()
+    print("Distributed functionality test passed!")
+    
+    test_shared_expert_distributed_weighted_sum()
+    print("Distributed weighted sum test passed!")

torch-ext/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,125 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
 def get_device_mesh(model):
     # Extract device_mesh from child's unused pre_hook closure
     try:
@@ -687,7 +866,7 @@ def get_device_mesh(model):
         hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
         # Extract the device_mesh from the closure
         return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
-    except:
+    except Exception:
         return None
 
 
@@ -703,8 +882,11 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
 
-        device_mesh = get_device_mesh(self)
-        expert_parallel_group = device_mesh.get_group() if device_mesh else None
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -734,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out