fix(granitemoehybid): Only set self.block_sparse_moe if num_local_experts > 0

Branch: GraniteMoeAsDenseFix

Signed-off-by: Gabe Goodhart <[email protected]>

Author: gabe-l-hart

src/transformers/models/granitemoehybrid/modular_granitemoehybrid.py

@@ -35,6 +35,7 @@
     GraniteMoeSharedDecoderLayer,
     GraniteMoeSharedForCausalLM,
     GraniteMoeSharedMLP,
+    GraniteMoeSharedMoE,
     GraniteMoeSharedModel,
     GraniteMoeSharedPreTrainedModel,
     eager_attention_forward,
@@ -107,6 +108,10 @@ class GraniteMoeHybridRotaryEmbedding(Gemma2RotaryEmbedding):
     pass
 
 
+class GraniteMoeHybridMoE(GraniteMoeSharedMoE):
+    pass
+
+
 class GraniteMoeHybridDecoderLayer(GraniteMoeSharedDecoderLayer):
     def __init__(self, config: GraniteMoeHybridConfig, layer_idx: int):
         super().__init__(config, layer_idx)
@@ -121,6 +126,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]
 
+        # Allow non-MoE (dense)
+        self.block_sparse_moe = GraniteMoeHybridMoE(config) if config.num_local_experts > 0 else None
+
         # Accept 0 experts: skip MoE if num_local_experts == 0
         self.has_experts = getattr(config, "num_local_experts", 0) > 0