Add the deepspeed injection_policy of mistral (#1309)

Signed-off-by: yuanwu <[email protected]>
huggingface · Sep 6, 2024 · bcd73f8 · bcd73f8
1 parent 5164d51
commit bcd73f8
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/optimum/habana/checkpoint_utils.py b/optimum/habana/checkpoint_utils.py
@@ -149,4 +149,9 @@ def get_ds_injection_policy(config):
 
             policy = {LlamaDecoderLayer: ("self_attn.o_proj", "mlp.down_proj")}
 
+        if model_type == "mistral":
+            from transformers.models.mistral.modeling_mistral import MistralDecoderLayer
+
+            policy = {MistralDecoderLayer: ("self_attn.o_proj", "mlp.down_proj")}
+
     return policy