add shared_codebooks

a-kore · Jun 22, 2023 · bca4dee · bca4dee
1 parent 479ae5a
commit bca4dee
Show file tree

Hide file tree

Showing 7 changed files with 1,517 additions and 399 deletions.
diff --git a/rpq/models/rpqopt.py b/rpq/models/rpqopt.py
@@ -341,8 +341,6 @@ def forward(
             hidden_states = self.self_attn_layer_norm(hidden_states)
 
         # Fully Connected
-        hidden_states_shape = hidden_states.shape
-        hidden_states = hidden_states.reshape(-1, hidden_states.size(-1))
         residual = hidden_states
 
         # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention
@@ -355,7 +353,7 @@ def forward(
         hidden_states = self.fc2(hidden_states)
         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
 
-        hidden_states = (residual + hidden_states).view(hidden_states_shape)
+        hidden_states = (residual + hidden_states)
 
         # 350m applies layer norm AFTER attention
         if not self.do_layer_norm_before: