fix(granitemoe): Use an inline conditional for creation of self.block_sparse_moe

This will prevent additional SparseMoeBlock code from being injected during
modelin_* generation.

Branch: GraniteMoeAsDenseFix

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>

Author: gabe-l-hart

src/transformers/models/granitemoe/modular_granitemoe.py

@@ -102,10 +102,9 @@ class GraniteMoeDecoderLayer(MixtralDecoderLayer):
     def __init__(self, config: GraniteMoeConfig, layer_idx: int):
         super().__init__(config, layer_idx)
         self.self_attn = GraniteMoeAttention(config=config, layer_idx=layer_idx)
-        if config.num_local_experts > 0:  # Diff with mixtral!
-            self.block_sparse_moe = GraniteMoeMoE(config)
-        elif self.block_sparse_moe:
-            delattr(self, "block_sparse_moe")  # Added by mixtral parent class
+        self.block_sparse_moe = (
+            GraniteMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) if config.num_local_experts > 0 else None
+        )  # Diff with mixtral!
         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)