fix circular import issue (#3445)

Author: blzheng

intel_extension_for_pytorch/cpu/tpp/fused_bert.py

@@ -13,8 +13,19 @@
 from ...utils._logger import logger, WarningType
 
 try:
-    from transformers.modeling_utils import apply_chunking_to_forward
-    from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
+    import transformers
+
+    transformers_orig_is_tensor = transformers.file_utils.is_tensor
+
+    def is_tensor(x):
+        """Tests if ``x`` is a :obj:`torch.Tensor`, :obj:`tf.Tensor` or :obj:`np.ndarray`."""
+        if transformers_orig_is_tensor(x):
+            return True
+        if isinstance(x, BlockedTensor):
+            return True
+        return False
+
+    transformers.file_utils.is_tensor = is_tensor
 except ImportError:
     pass
 USE_BF16_PARAMS = True
@@ -976,7 +987,7 @@ def forward(
             cross_attn_present_key_value = cross_attention_outputs[-1]
             present_key_value = present_key_value + cross_attn_present_key_value
 
-        layer_output = apply_chunking_to_forward(
+        layer_output = transformers.modeling_utils.apply_chunking_to_forward(
             self.feed_forward_chunk,
             self.chunk_size_feed_forward,
             self.seq_len_dim,
@@ -1109,7 +1120,7 @@ def custom_forward(*inputs):
                 ]
                 if v is not None
             )
-        return BaseModelOutputWithPastAndCrossAttentions(
+        return transformers.modeling_outputs.BaseModelOutputWithPastAndCrossAttentions(
             last_hidden_state=hidden_states,
             past_key_values=next_decoder_cache,
             hidden_states=all_hidden_states,
@@ -1178,23 +1189,6 @@ def forward(self, hidden_states):
 #     return bm_default_blocking_factors
 # BlockedModule.default_blocking_factors = custom_blocking_factors
 
-try:
-    import transformers
-
-    transformers_orig_is_tensor = transformers.file_utils.is_tensor
-
-    def is_tensor(x):
-        """Tests if ``x`` is a :obj:`torch.Tensor`, :obj:`tf.Tensor` or :obj:`np.ndarray`."""
-        if transformers_orig_is_tensor(x):
-            return True
-        if isinstance(x, BlockedTensor):
-            return True
-        return False
-
-    transformers.file_utils.is_tensor = is_tensor
-except ImportError:
-    pass
-
 
 def block(model):
     for m in model.modules():

intel_extension_for_pytorch/llm/__init__.py

@@ -29,12 +29,15 @@
     transformers.dynamic_module_utils.get_class_from_dynamic_module = (
         _get_class_from_dynamic_module
     )
-    transformers.modeling_utils.PreTrainedModel.gradient_checkpointing_disable = (
-        _gradient_checkpointing_disable
-    )
-    transformers.modeling_utils.PreTrainedModel.gradient_checkpointing_enable = (
-        _gradient_checkpointing_enable
-    )
+    from packaing import version
+
+    if version.parse(transformers.__version__) < version.parse("4.36"):
+        transformers.modeling_utils.PreTrainedModel.gradient_checkpointing_disable = (
+            _gradient_checkpointing_disable
+        )
+        transformers.modeling_utils.PreTrainedModel.gradient_checkpointing_enable = (
+            _gradient_checkpointing_enable
+        )
     transformers.tokenization_utils_base.PreTrainedTokenizerBase.pad = _pad
 except ImportError:
     pass

intel_extension_for_pytorch/transformers/models/reference/models.py

@@ -23,10 +23,6 @@
         _prepare_4d_causal_attention_mask,
     )
 
-    if hasattr(transformers.models, "mixtral"):
-        from transformers.models.mixtral.modeling_mixtral import (
-            load_balancing_loss_func,
-        )
     from transformers.modeling_outputs import (
         MoeCausalLMOutputWithPast,
         MoeModelOutputWithPast,
@@ -3277,10 +3273,12 @@ def MixtralForCausalLM_forward(
 
     aux_loss = None
     if output_router_logits:
-        aux_loss = load_balancing_loss_func(
-            outputs.router_logits if return_dict else outputs[-1],
-            self.num_experts,
-            self.num_experts_per_tok,
+        aux_loss = (
+            transformers.models.mixtral.modeling_mixtral.load_balancing_loss_func(
+                outputs.router_logits if return_dict else outputs[-1],
+                self.num_experts,
+                self.num_experts_per_tok,
+            )
         )
         if labels is not None:
             loss += self.router_aux_loss_coef * aux_loss
@@ -5828,7 +5826,7 @@ def JambaForCausalLM_forward(
 
     aux_loss = None
     if output_router_logits:
-        aux_loss = load_balancing_loss_func(
+        aux_loss = transformers.models.jamba.modeling_jamba.load_balancing_loss_func(
             outputs.router_logits if return_dict else outputs[-1],
             self.num_experts,
             self.num_experts_per_tok,