layer id init

myolmoe/modeling_myolmoe.py

@@ -530,7 +530,7 @@ class OlmoeDecoderLayer(nn.Module):
         self.self_attn = OLMOE_ATTENTION_CLASSES[config._attn_implementation](
             config=config, layer_idx=layer_idx
         )
-        self.mlp = OlmoeSparseMoeBlock(config)
+        self.mlp = OlmoeSparseMoeBlock(config, layer_idx)
         self.input_layernorm = OlmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = OlmoeRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps