Add position_ids arg to OPTEmbedding forward function (#6939)

This PR updates the DeepSpeed `OPTEmbedding` forward function to include a new `positions_ids` argument. --------- Co-authored-by: Logan Adams <loadams@microsoft.com>
microsoft · Jan 9, 2025 · 0fc3daa · 0fc3daa
1 parent 45fce45
commit 0fc3daa
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/.github/workflows/nv-ds-chat.yml b/.github/workflows/nv-ds-chat.yml
@@ -43,7 +43,7 @@ jobs:
 
       - name: Install deepspeed
         run: |
-          pip install transformers==4.45.2
+          pip install transformers
           pip install .[dev]
           ds_report
 

diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
@@ -191,7 +191,7 @@ def __init__(self, weight_shape=None, weight=None, bias=None):
         self.offset = 2
         super().__init__(weight_shape, weight=weight)
 
-    def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0):
+    def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0, position_ids: int = 0):
         """`input_ids_shape` is expected to be [bsz x seqlen]."""
         attention_mask = attention_mask.long()