(m)t5 lora parallel

adapter-hub · Dec 21, 2024 · 4e40256 · 4e40256
1 parent 18f9053
commit 4e40256
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/src/adapters/models/mt5/modeling_mt5.py b/src/adapters/models/mt5/modeling_mt5.py
@@ -77,7 +77,12 @@ def forward(
         is_cross_attention = key_value_states is not None
 
         query_states = self.q(hidden_states)
-        query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
+        # >>> START AH Changes <<<
+        # adapt bsz for lora parallel
+        query_states = query_states.view(query_states.shape[0], -1, self.n_heads, self.key_value_proj_dim).transpose(
+            1, 2
+        )
+        # >>> END AH Changes <<<
 
         if past_key_value is not None:
             is_updated = past_key_value.is_updated.get(self.layer_idx)
@@ -95,8 +100,12 @@ def forward(
         else:
             key_states = self.k(current_states)
             value_states = self.v(current_states)
-            key_states = key_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
-            value_states = value_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
+            key_states = key_states.view(key_states.shape[0], -1, self.n_heads, self.key_value_proj_dim).transpose(
+                1, 2
+            )
+            value_states = value_states.view(
+                value_states.shape[0], -1, self.n_heads, self.key_value_proj_dim
+            ).transpose(1, 2)
 
             if past_key_value is not None:
                 # save all key/value_states to cache to be re-used for fast auto-regressive generation

diff --git a/src/adapters/models/t5/modeling_t5.py b/src/adapters/models/t5/modeling_t5.py
@@ -77,7 +77,12 @@ def forward(
         is_cross_attention = key_value_states is not None
 
         query_states = self.q(hidden_states)
-        query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
+        # >>> START AH Changes <<<
+        # adapt bsz for lora parallel
+        query_states = query_states.view(query_states.shape[0], -1, self.n_heads, self.key_value_proj_dim).transpose(
+            1, 2
+        )
+        # >>> END AH Changes <<<
 
         if past_key_value is not None:
             is_updated = past_key_value.is_updated.get(self.layer_idx)
@@ -95,8 +100,12 @@ def forward(
         else:
             key_states = self.k(current_states)
             value_states = self.v(current_states)
-            key_states = key_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
-            value_states = value_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
+            key_states = key_states.view(key_states.shape[0], -1, self.n_heads, self.key_value_proj_dim).transpose(
+                1, 2
+            )
+            value_states = value_states.view(
+                value_states.shape[0], -1, self.n_heads, self.key_value_proj_dim
+            ).transpose(1, 2)
 
             if past_key_value is not None:
                 # save all key/value_states to cache to be re-used for fast auto-regressive generation