Updated default config.

configuration_grok.py

@@ -96,8 +96,8 @@ class GrokConfig(PretrainedConfig):
         num_hidden_layers=64,
         num_attention_heads=48,
         num_key_value_heads=8,
-        hidden_act="silu",
-        max_position_embeddings=4096,
+        hidden_act="gelu_new",
+        max_position_embeddings=8192,
         initializer_range=0.02,
         rms_norm_eps=1e-5,
         use_cache=True,
@@ -105,7 +105,7 @@ class GrokConfig(PretrainedConfig):
         bos_token_id=1,
         eos_token_id=2,
         tie_word_embeddings=True,
-        rope_theta=1e5,
+        rope_theta=1e4,
         attention_dropout=0.0,
         num_experts_per_tok=2,
         num_local_experts=8,

modeling_grok.py

@@ -338,8 +338,11 @@ class GrokDecoderLayer(nn.Module):
        	self.top_k = config.num_experts_per_tok
 
         self.multi_head_attention = GrokAttention(config, layer_idx)
-       	self.router = nn.Linear(self.hidden_size, self.num_experts, dtype=torch.float32, bias=False)
-        self.moe = nn.ModuleList([GrokBlockSparseTop2MLP(config) for _ in range(self.num_experts)])
+        if self.num_experts > 1:
+            self.router = nn.Linear(self.hidden_size, self.num_experts, dtype=torch.float32, bias=False)
+            self.moe = nn.ModuleList([GrokBlockSparseTop2MLP(config) for _ in range(self.num_experts)])
+        else:
+            self.mlp = GrokBlockSparseTop2MLP(config)
 
         self.rms_norm = GrokRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rms_norm_1 = GrokRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -397,6 +400,18 @@ class GrokDecoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.rms_norm_2(hidden_states)
 
+        if self.num_experts <= 1:
+            hidden_states = self.mlp(hidden_states)
+            hidden_states = residual + self.rms_norm_3(hidden_states)
+            outputs = (hidden_states,)
+            if output_attentions:
+                outputs += (self_attn_weights,)
+            if use_cache:
+                outputs += (present_key_value,)
+            if output_router_logits:
+                outputs += (None,)
+            return outputs
+
         batch_size, sequence_length, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
         # router_logits: (batch * sequence_length, n_experts)