Update the Gaudi trainer with transformers 4.45.2 (huggingface#1398)

examples/text-generation/README.md

@@ -198,9 +198,9 @@ python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \
 > --bf16
 > ```
 
-To run Mamba-130m inference on 1 Gaudi2 card, use the following command:
+To run Mamba-130m inference on 1 Gaudi2 card, use the following command, for example if default custom kernel path is in /root/.cache/huggingface/hub/models--Habana--mamba/blobs/libcustom_tpc_perf_lib.so, if libcustom_tpc_perf_lib.so is in different folder, set accordingly,
 ```bash
-python run_generation.py \
+GC_KERNEL_PATH=/root/.cache/huggingface/hub/models--Habana--mamba/blobs/libcustom_tpc_perf_lib.so:$GC_KERNEL_PATH python run_generation.py \
 --model_name_or_path state-spaces/mamba-130m-hf \
 --max_input_tokens 128 \
 --max_new_tokens 128 \
@@ -209,13 +209,7 @@ python run_generation.py \
 --use_kv_cache \
 --batch_size 1024 \
 ```
-We need to download two shared libraries and set two env variables for mamba, for 1.19 version use *_119 so files. For different release, use different set of these so files.
-```bash
-wget https://huggingface.co/Habana/mamba/resolve/main/hpu_custom_pscan_all.cpython-310-x86_64-linux-gnu_119.so
-wget https://huggingface.co/Habana/mamba/resolve/main/libcustom_tpc_perf_lib_119.so
-export HABANA_CUSTOM_OP_DIR=/path/of/hpu_custom_pscan_all.cpython-310-x86_64-linux-gnu_119.so located
-export GC_KERNEL_PATH=/path/to/libcustom_tpc_perf_lib_119.so:$GC_KERNEL_PATH
-```
+
 ### Use any dataset from the Hugging Face Hub
 
 You can also provide the name of a dataset from the Hugging Face Hub to perform generation on it with the argument `--dataset_name`.

optimum/habana/transformers/models/mamba/modeling_mamba.py

@@ -12,10 +12,21 @@
     logging,
 )
 
+from .util_mamba import set_mamba_lib
+env_variables = os.environ.copy()
 
-my_dir = os.path.realpath(__file__)
-my_len = my_dir.rfind("/")
-base_dir = os.environ.get("HABANA_CUSTOM_OP_DIR", my_dir[:my_len])
+new_file_op, new_file_kernel = set_mamba_lib()
+realpath_kfn = os.path.realpath(new_file_kernel)
+kfn = os.path.basename(realpath_kfn)
+new_kfn = os.path.join(os.path.dirname(realpath_kfn), "libcustom_tpc_perf_lib.so")
+os.rename(realpath_kfn, new_kfn)
+
+
+env_variables["HABANA_CUSTOM_OP_DIR"] = os.path.dirname(new_file_op)
+default_path = env_variables["GC_KERNEL_PATH"]
+env_variables["GC_KERNEL_PATH"] = new_kfn + os.pathsep + default_path
+
+base_dir = env_variables["HABANA_CUSTOM_OP_DIR"] 
 
 custom_op_lib_path = str(next(Path(base_dir).glob("hpu_custom_pscan_all.cpython-*-x86_64-linux-gnu.so")))
 torch.ops.load_library(custom_op_lib_path)

optimum/habana/transformers/models/mamba/util_mamba.py

@@ -0,0 +1,30 @@
+import os
+from huggingface_hub import hf_hub_download
+
+from optimum.habana.utils import  get_habana_frameworks_version
+
+def set_mamba_lib():
+    version_no = get_habana_frameworks_version()
+    env_variables = os.environ.copy()
+
+    name_op = "hpu_custom_pscan_all.cpython-310-x86_64-linux-gnu.so"
+    name_kernel = "libcustom_tpc_perf_lib.so"
+    if version_no.minor == 19:
+        name_op = "hpu_custom_pscan_all.cpython-310-x86_64-linux-gnu_119.so"
+        name_kernel = "libcustom_tpc_perf_lib_119.so"
+
+    file_op = hf_hub_download(repo_id="Habana/mamba", filename=name_op)
+    file_kernel = hf_hub_download(repo_id="Habana/mamba", filename=name_kernel)
+
+    new_file_op = file_op
+    new_file_kernel = file_kernel
+
+    if version_no.minor == 19:
+        new_file_op = file_op[:-7] + ".so"
+        new_file_kernel = file_kernel[:-7] + ".so"
+        os.rename(file_op, new_file_op)
+        os.rename(file_kernel, new_file_kernel)
+
+    return new_file_op, new_file_kernel
+
+

tests/test_text_generation_example.py

@@ -217,38 +217,6 @@ def _test_text_generation(
     if "gemma" in model_name.lower():
         command += ["--use_flash_attention"]
 
-    if "mamba" in model_name.lower():
-        import subprocess
-
-        from huggingface_hub import hf_hub_download
-
-        cmd1 = subprocess.Popen(["pip", "list"], stdout=subprocess.PIPE)
-        cmd2 = subprocess.Popen(["grep", "habana-torch-plugin"], stdin=cmd1.stdout, stdout=subprocess.PIPE)
-        cmd1.stdout.close()
-        version_no, _ = cmd2.communicate()
-
-        name_op = "hpu_custom_pscan_all.cpython-310-x86_64-linux-gnu.so"
-        name_kernel = "libcustom_tpc_perf_lib.so"
-        if "1.19.0" in version_no.decode():
-            name_op = "hpu_custom_pscan_all.cpython-310-x86_64-linux-gnu_119.so"
-            name_kernel = "libcustom_tpc_perf_lib_119.so"
-
-        file_op = hf_hub_download(repo_id="Habana/mamba", filename=name_op)
-        file_kernel = hf_hub_download(repo_id="Habana/mamba", filename=name_kernel)
-
-        new_file_op = file_op
-        new_file_kernel = file_kernel
-
-        if "1.19.0" in version_no.decode():
-            new_file_op = file_op[:-7] + ".so"
-            new_file_kernel = file_kernel[:-7] + ".so"
-            os.rename(file_op, new_file_op)
-            os.rename(file_kernel, new_file_kernel)
-
-        env_variables["HABANA_CUSTOM_OP_DIR"] = os.path.dirname(new_file_op)
-        default_path = env_variables["GC_KERNEL_PATH"]
-        env_variables["GC_KERNEL_PATH"] = new_file_kernel + os.pathsep + default_path
-
     if (reuse_cache or torch_compile) and not parallel_strategy == "tp" and not is_starcoder_first_gen_model:
         command += ["--reuse_cache"]