Repository: huggingface/transformers
Pull request: GraniteMoeHybrid: Allow for only shared expert case.

# What does this PR do?
Allows for the case where `GraniteMoeHybrid` models have `num_local_experts = 0`


Changed files:
filename: src/transformers/models/granitemoe/modeling_granitemoe.py
patch: @@ -504,7 +504,8 @@ def __init__(self, config: GraniteMoeConfig, layer_idx: int):
         self.hidden_size = config.hidden_size
 
         self.self_attn = GraniteMoeAttention(config=config, layer_idx=layer_idx)
-        self.block_sparse_moe = GraniteMoeMoE(config)
+        if config.num_local_experts > 0:
+            self.block_sparse_moe = GraniteMoeMoE(config)
         self.input_layernorm = GraniteMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = GraniteMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 

filename: src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
patch: @@ -1051,7 +1051,8 @@ def __init__(self, config: GraniteMoeHybridConfig, layer_idx: int):
         self.hidden_size = config.hidden_size
         # Either attention or mamba will be initialized, depending on the layer type.
         self.self_attn = None
-        self.block_sparse_moe = GraniteMoeHybridMoE(config)
+        if config.num_local_experts > 0:
+            self.block_sparse_moe = GraniteMoeHybridMoE(config)
         self.input_layernorm = GraniteMoeHybridRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = GraniteMoeHybridRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
@@ -1065,6 +1066,9 @@ def __init__(self, config: GraniteMoeHybridConfig, layer_idx: int):
             self.self_attn = GraniteMoeHybridAttention(config, layer_idx)
         self.layer_type = config.layers_block_type[layer_idx]
 
+        # Accept 0 experts: skip MoE if num_local_experts == 0
+        self.has_experts = getattr(config, "num_local_experts", 0) > 0
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -1131,9 +1135,14 @@ def forward(
         # Fully Connected
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
-        moe_hidden_states, router_logits = self.block_sparse_moe(hidden_states)
 
-        hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+        if self.has_experts:
+            moe_hidden_states, router_logits = self.block_sparse_moe(hidden_states)
+            hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+        else:
+            hidden_states = self.shared_mlp(hidden_states)
+            router_logits = None
+
         hidden_states = residual + hidden_states * self.residual_multiplier
 
         outputs = (hidden_states,)

filename: src/transformers/models/granitemoehybrid/modular_granitemoehybrid.py
patch: @@ -71,6 +71,9 @@ def __init__(self, config: GraniteMoeHybridConfig, layer_idx: int):
             self.self_attn = GraniteMoeHybridAttention(config, layer_idx)
         self.layer_type = config.layers_block_type[layer_idx]
 
+        # Accept 0 experts: skip MoE if num_local_experts == 0
+        self.has_experts = getattr(config, "num_local_experts", 0) > 0
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -137,9 +140,14 @@ def forward(
         # Fully Connected
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
-        moe_hidden_states, router_logits = self.block_sparse_moe(hidden_states)
 
-        hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+        if self.has_experts:
+            moe_hidden_states, router_logits = self.block_sparse_moe(hidden_states)
+            hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+        else:
+            hidden_states = self.shared_mlp(hidden_states)
+            router_logits = None
+
         hidden_states = residual + hidden_states * self.residual_multiplier
 
         outputs = (hidden_states,)

filename: src/transformers/models/granitemoeshared/modeling_granitemoeshared.py
patch: @@ -412,7 +412,8 @@ def __init__(self, config: GraniteMoeSharedConfig, layer_idx: int):
         self.hidden_size = config.hidden_size
 
         self.self_attn = GraniteMoeSharedAttention(config=config, layer_idx=layer_idx)
-        self.block_sparse_moe = GraniteMoeSharedMoE(config)
+        if config.num_local_experts > 0:
+            self.block_sparse_moe = GraniteMoeSharedMoE(config)
         self.input_layernorm = GraniteMoeSharedRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = GraniteMoeSharedRMSNorm(config.hidden_size, eps=config.rms_norm_eps)