Merge pull request #144 from dongs0104/patch-1

Add OPT Model mapping
tunib-ai · Aug 25, 2022 · 0538469 · 0538469
2 parents a4cad0e + 3743a89
commit 0538469
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/oslo/transformers/mapping_utils.py b/oslo/transformers/mapping_utils.py
@@ -138,6 +138,12 @@ class _TensorParallelMappingForHuggingFace(_ParallelMappingForHuggingFace):
             Update("embed_dim", "num_attention_heads"),
             Head("lm_head", "score", gather_output=True),
         ],
+        OPT=[
+            Column("q_proj", "k_proj", "v_proj", "fc1"),
+            Row("out_proj", "fc2"),
+            Update("embed_dim", "num_heads"),
+            Head("lm_head", "score"),
+        ],
         "Electra": [
             Column("query", "key", "value", "intermediate.dense"),
             Column(