Update modeling_motif.py

modeling_motif.py

@@ -328,23 +328,10 @@ class MotifMLP(nn.Module):
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
         self.act_fn = ACT2FN[config.hidden_act]
 
-        if config.wesar_weights:
-            self.gate_up_proj_alpha = nn.Parameter(torch.tensor(1) *config.gate_up_proj_alpha)
-            self.down_proj_alpha = nn.Parameter(torch.tensor(1) * config.down_proj_alpha)
-        else:
-            self.gate_up_proj_alpha=1
-            self.down_proj_alpha=1
-        if config.muP:
-            self.down_proj.__do_scale_tager__ = True
-            self.gate_proj.__do_scale_tager_mu_dim_model__  = True
-            self.up_proj.__do_scale_tager_mu_dim_model__ = True
-            self.down_proj.__do_scale_tager_mu_ffn__ = True
-
-
     def forward(self, hidden_state):
-        hidden_state = hidden_state*self.gate_up_proj_alpha
+        hidden_state = hidden_state
         #hidden_state = self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))*
-        return self.down_proj_alpha*self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
 
 
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
@@ -470,13 +457,6 @@ class MotifAttention(nn.Module):
                                                 max_position_embeddings=self.max_position_embeddings,
                                                 base=self.rope_theta)
 
-        for param in ["q_proj_alpha", "k_proj_alpha", "v_proj_alpha", "o_proj_alpha"]:
-            setattr(
-                self, param,
-                nn.Parameter(torch.tensor(getattr(config, param, 1.0), dtype=torch.float))
-                if config.wesar_weights else 1.0)
-
-
     def forward(
             self,
             hidden_states: torch.Tensor,
@@ -490,9 +470,9 @@ class MotifAttention(nn.Module):
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
 
-        query_states = self.q_proj(hidden_states) * self.q_proj_alpha
-        key_states = self.k_proj(hidden_states) * self.k_proj_alpha
-        value_states = self.v_proj(hidden_states) * self.v_proj_alpha
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
 
         ## bsz, seq, n_heads, head_dim
 
@@ -685,9 +665,9 @@ class MotifFlashAttention2(MotifAttention):
     ):
         bsz, q_len, _ = hidden_states.size()
 
-        query_states = self.q_proj(hidden_states) * self.q_proj_alpha
-        key_states = self.k_proj(hidden_states) * self.k_proj_alpha
-        value_states = self.v_proj(hidden_states) * self.v_proj_alpha
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
 
         query_states = query_states.view(bsz, q_len, 2 * self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, 2 * self.num_key_value_heads, self.head_dim).transpose(1, 2)
@@ -798,7 +778,7 @@ class MotifFlashAttention2(MotifAttention):
                              f" {attn_output.size()}")
 
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-        attn_output = self.o_proj(attn_output) * self.o_proj_alpha
+        attn_output = self.o_proj(attn_output)
 
         return attn_output, None, past_key_value
 
@@ -919,15 +899,6 @@ class MotifDecoderLayer(nn.Module):
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        if config.wesar_weights and config.use_norm_alpha:
-            self.input_layernorm_alpha = nn.Parameter(torch.tensor(1).float())
-        else:
-            self.input_layernorm_alpha = 1
-
-        if config.wesar_weights and config.use_norm_alpha :
-            self.post_attention_layernorm_alpha = nn.Parameter(torch.tensor(1).float())
-        else:
-            self.post_attention_layernorm_alpha = 1
 
     def forward(
         self,
@@ -965,7 +936,7 @@ class MotifDecoderLayer(nn.Module):
 
         residual = hidden_states
 
-        hidden_states = self.input_layernorm(hidden_states) * self.input_layernorm_alpha
+        hidden_states = self.input_layernorm(hidden_states)
 
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
@@ -982,7 +953,7 @@ class MotifDecoderLayer(nn.Module):
 
         # Fully Connected
         residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states) * self.post_attention_layernorm_alpha
+        hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
 
@@ -1199,14 +1170,7 @@ class MotifModel(MotifPreTrainedModel):
         self.post_init()
 
         self.scale_emb = 1
-
-        # Reparameterization <|_1_|>
-        if config.wesar_weights :
-            logger.info(f'config.wesar_weights {config.wesar_weights}')
-            self.norm_alpha = nn.Parameter(torch.tensor(1).float())
-            self.scale_emb = 10
-        else:
-            self.norm_alpha = 1
+        self.norm_alpha = 1
 
     def get_input_embeddings(self):
         return self.embed_tokens