[hotfix] update ffn dim

configuration_grok1.py

@@ -9,7 +9,7 @@ class Grok1Config(PretrainedConfig):
         self,
         vocab_size=32000,
         hidden_size=4096,
-        widening_factor=4.0,
+        intermediate_size=32768,
         num_hidden_layers=32,
         num_attention_heads=32,
         num_key_value_heads=32,
@@ -37,7 +37,7 @@ class Grok1Config(PretrainedConfig):
         self.embedding_multiplier_scale = embedding_multiplier_scale
         self.output_multiplier_scale = output_multiplier_scale
         self.hidden_size = hidden_size
-        self.widening_factor = widening_factor
+        self.intermediate_size = intermediate_size
         self.num_hidden_layers = num_hidden_layers
         self.num_attention_heads = num_attention_heads
 

modeling_grok1.py

@@ -395,11 +395,11 @@ class DecoderLayer(nn.Module):
     def __init__(
         self,
         hidden_size: int,
+        intermediate_size: int,
         num_heads: int,
         num_key_value_heads: int,
         num_experts: int,
         top_k: int,
-        widening_factor: float = 4.0,
         max_position_embeddings: int = 2048,
         attn_output_multiplier: float = 1.0,
         max_attn_val: float = 30.0,
@@ -414,8 +414,7 @@ class DecoderLayer(nn.Module):
             attn_output_multiplier=attn_output_multiplier,
             max_attn_val=max_attn_val,
         )
-        ffn_dim = int(hidden_size * widening_factor)
-        self.moe_block = MoeBlock(hidden_size, ffn_dim, num_experts, top_k)
+        self.moe_block = MoeBlock(hidden_size, intermediate_size, num_experts, top_k)
         self.pre_attn_norm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.post_attn_norm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.pre_moe_norm = RMSNorm(hidden_size, eps=rms_norm_eps)
@@ -543,11 +542,11 @@ class Grok1Model(Grok1PretrainedModel):
             [
                 DecoderLayer(
                     hidden_size=config.hidden_size,
+                    intermediate_size=config.intermediate_size,
                     num_heads=config.num_attention_heads,
                     num_key_value_heads=config.num_key_value_heads,
                     num_experts=config.num_experts,
                     top_k=config.num_experts_per_tok,
-                    widening_factor=config.widening_factor,
                     max_position_embeddings=config.max_position_embeddings,
                     attn_output_multiplier=config.attn_output_multiplier,
                     max_attn_val=config.max_attn_value,