Fix 1383 Llama model on transformers=4.41[WIP] (intel-analytics#11280)

RyuKosei · Jul 19, 2024 · b370cef · b370cef
1 parent 4097b46
commit b370cef
Show file tree

Hide file tree

Showing 2 changed files with 790 additions and 84 deletions.
diff --git a/python/llm/src/ipex_llm/transformers/convert.py b/python/llm/src/ipex_llm/transformers/convert.py
@@ -980,19 +980,32 @@ def _optimize_post(model, lightweight_bmm=False):
         convert_forward(model,
                         transformers.models.llama.modeling_llama.LlamaDecoderLayer,
                         llama_decoder_forward)
+
         if version.parse(trans_version) >= version.parse("4.36.0"):
             # transformers version >= 4.36.0
             from ipex_llm.transformers.models.llama import llama_attention_forward_4_38
             if version.parse(trans_version) >= version.parse("4.38.0"):
-                from ipex_llm.transformers.models.llama import llama_model_forward_4_38
-                convert_forward(
-                    model,
-                    transformers.models.llama.modeling_llama.LlamaModel,
-                    llama_model_forward_4_38)
-                convert_forward(
-                    model,
-                    transformers.models.llama.modeling_llama.LlamaAttention,
-                    llama_attention_forward_4_38)
+                if version.parse(trans_version) >= version.parse("4.41.0"):
+                    from ipex_llm.transformers.models.llama import llama_model_forward_4_41
+                    from ipex_llm.transformers.models.llama import llama_attention_forward_4_41
+                    convert_forward(
+                        model,
+                        transformers.models.llama.modeling_llama.LlamaModel,
+                        llama_model_forward_4_41)
+                    convert_forward(
+                        model,
+                        transformers.models.llama.modeling_llama.LlamaAttention,
+                        llama_attention_forward_4_41)
+                else:
+                    from ipex_llm.transformers.models.llama import llama_model_forward_4_38
+                    convert_forward(
+                        model,
+                        transformers.models.llama.modeling_llama.LlamaModel,
+                        llama_model_forward_4_38)
+                    convert_forward(
+                        model,
+                        transformers.models.llama.modeling_llama.LlamaAttention,
+                        llama_attention_forward_4_38)
             else:
                 from ipex_llm.transformers.models.llama import llama_model_forward_4_36
                 convert_forward(