Add workaround for FusedRoPE (#473)

optimum/habana/transformers/models/falcon/modeling_falcon.py

@@ -3,17 +3,23 @@
 
 import torch
 
+from ....utils import get_device_name
+
 
 try:
     from habana_frameworks.torch.hpex.kernels import FusedSDPA
 except ImportError:
     print("Not using HPU fused kernel for scaled_dot_product_attention")
     FusedSDPA = None
 
-try:
-    from habana_frameworks.torch.hpex.kernels import RotaryPosEmbeddingHelperV1 as FusedRoPE
-except ImportError:
-    print("Not using HPU fused kernel for apply_rotary_pos_emb")
+# TODO: remove this workaround when FusedRoPE properly works on Gaudi
+if get_device_name() == "gaudi2":
+    try:
+        from habana_frameworks.torch.hpex.kernels import RotaryPosEmbeddingHelperV1 as FusedRoPE
+    except ImportError:
+        print("Not using HPU fused kernel for apply_rotary_pos_emb")
+        FusedRoPE = None
+else:
     FusedRoPE = None
 
 import habana_frameworks.torch.core as htcore

optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py

@@ -5,11 +5,17 @@
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM, apply_rotary_pos_emb, logger
 
+from ....utils import get_device_name
 
-try:
-    from habana_frameworks.torch.hpex.kernels import RotaryPosEmbeddingHelperV2 as FusedRoPE
-except ImportError:
-    print("Not using HPU fused kernel for apply_rotary_pos_emb")
+
+# TODO: remove this workaround when FusedRoPE properly works on Gaudi
+if get_device_name() == "gaudi2":
+    try:
+        from habana_frameworks.torch.hpex.kernels import RotaryPosEmbeddingHelperV2 as FusedRoPE
+    except ImportError:
+        print("Not using HPU fused kernel for apply_rotary_pos_emb")
+        FusedRoPE = None
+else:
     FusedRoPE = None
 
 

optimum/habana/transformers/models/llama/modeling_llama.py

@@ -14,11 +14,17 @@
     logger,
 )
 
+from ....utils import get_device_name
 
-try:
-    from habana_frameworks.torch.hpex.kernels import RotaryPosEmbeddingHelperV2 as FusedRoPE
-except ImportError:
-    print("Not using HPU fused kernel for apply_rotary_pos_emb")
+
+# TODO: remove this workaround when FusedRoPE properly works on Gaudi
+if get_device_name() == "gaudi2":
+    try:
+        from habana_frameworks.torch.hpex.kernels import RotaryPosEmbeddingHelperV2 as FusedRoPE
+    except ImportError:
+        print("Not using HPU fused kernel for apply_rotary_pos_emb")
+        FusedRoPE = None
+else:
     FusedRoPE = None
 
 try:

optimum/habana/utils.py

@@ -318,3 +318,21 @@ def check_optimum_habana_min_version(min_version):
                 "`pip install git+https://github.com/huggingface/optimum-habana.git`."
             )
         raise ImportError(error_message)
+
+
+def get_device_name():
+    """
+    Returns the name of the current device: Gaudi or Gaudi2.
+
+    Inspired from: https://github.com/HabanaAI/Model-References/blob/a87c21f14f13b70ffc77617b9e80d1ec989a3442/PyTorch/computer_vision/classification/torchvision/utils.py#L274
+    """
+    import habana_frameworks.torch.utils.experimental as htexp
+
+    device_type = htexp._get_device_type()
+
+    if device_type == htexp.synDeviceType.synDeviceGaudi:
+        return "gaudi"
+    elif device_type == htexp.synDeviceType.synDeviceGaudi2:
+        return "gaudi2"
+    else:
+        raise ValueError(f"Unsupported device: the device type is {device_type}.")