hipblaslt check torch version

Author: lshqqytiger

installer.py

@@ -531,7 +531,7 @@ def install_rocm_zluda(torch_command):
         ort_package = os.environ.get('ONNXRUNTIME_PACKAGE', f"--pre onnxruntime-training{'' if ort_version is None else ('==' + ort_version)} --index-url https://pypi.lsh.sh/{rocm.version[0]}{rocm.version[2]} --extra-index-url https://pypi.org/simple")
         install(ort_package, 'onnxruntime-training')
 
-        if bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1"))) and rocm.version != "6.2":
+        if rocm.version == rocm.version_torch and rocm.get_blaslt_enabled():
             log.debug(f'hipBLASLt arch={hip_default_device.name} available={hip_default_device.blaslt_supported}')
             rocm.set_blaslt_enabled(hip_default_device.blaslt_supported)
     return torch_command

modules/rocm.py

@@ -3,6 +3,7 @@
 import ctypes
 import shutil
 import subprocess
+import importlib.metadata
 from typing import Union, List
 
 
@@ -81,6 +82,13 @@ def get_gfx_version(self) -> Union[str, None]:
         return None
 
 
+def get_version_torch() -> Union[str, None]:
+    version_ = importlib.metadata.version("torch")
+    if "+rocm" not in version_: # unofficial build, non-rocm torch.
+        return None
+    return version_.split("+rocm")[1]
+
+
 if sys.platform == "win32":
     def find() -> Union[str, None]:
         hip_path = shutil.which("hipconfig")
@@ -173,10 +181,14 @@ def set_blaslt_enabled(enabled: bool) -> None:
         else:
             os.environ["TORCH_BLAS_PREFER_HIPBLASLT"] = "0"
 
+    def get_blaslt_enabled() -> bool:
+        return bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1")))
+
     is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', None) is not None
 path = find()
 is_installed = False
 version = None
+version_torch = get_version_torch()
 if path is not None:
     is_installed = True
     version = get_version()