[hotfix] update modeling

configuration_grok1.py

@@ -16,6 +16,8 @@ class Grok1Config(PretrainedConfig):
         attn_output_multiplier=1.0,
         max_attn_value=1.0,
         max_position_embeddings=4096,
+        embedding_multiplier_scale: float = 1.0,
+        output_multiplier_scale: float = 1.0,
         rms_norm_eps=1e-5,
         use_cache=True,
         pad_token_id=None,
@@ -32,6 +34,8 @@ class Grok1Config(PretrainedConfig):
         self.attn_output_multiplier = attn_output_multiplier
         self.max_attn_value = max_attn_value
         self.max_position_embeddings = max_position_embeddings
+        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.num_hidden_layers = num_hidden_layers

modeling_grok1.py

@@ -259,8 +259,6 @@ class MultiHeadAttention(nn.Module):
 
         past_key_value = (key_states, value_states) if use_cache else None
 
-        # TODO: repeat kv
-
         attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)).to(
             torch.float
         )
@@ -536,6 +534,7 @@ class Grok1Model(Grok1PretrainedModel):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
+        self.embedding_multiplier_scale = config.embedding_multiplier_scale
 
         self.embed_tokens = nn.Embedding(
             config.vocab_size, config.hidden_size, self.padding_idx
@@ -654,6 +653,7 @@ class Grok1Model(Grok1PretrainedModel):
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
+            inputs_embeds = inputs_embeds * self.embedding_multiplier_scale
 
         if HAS_MASK_UTILS:
             # 4d mask is passed through the layers
@@ -772,6 +772,7 @@ class Grok1ModelForCausalLM(Grok1PretrainedModel):
         super().__init__(config)
         self.model = Grok1Model(config)
         self.vocab_size = config.vocab_size
+        self.output_multiplier_scale = config.output_multiplier_scale
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.router_aux_loss_coef = config.router_aux_loss_coef
         self.num_experts = config.num_experts
@@ -846,6 +847,7 @@ class Grok1ModelForCausalLM(Grok1PretrainedModel):
 
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
+        logits = logits * self.output_multiplier_scale
         logits = logits.float()
 
         loss = None