update rocm dependencies

mila-iqia · Aug 21, 2024 · 321acb6 · 321acb6
1 parent 4873a4a
commit 321acb6
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/benchmarks/huggingface/tunableop_results0.csv b/benchmarks/huggingface/tunableop_results0.csv
@@ -0,0 +1,17 @@
+Validator,PT_VERSION,2.4.0
+Validator,ROCBLAS_VERSION,4.0.0-88df9726-dirty
+Validator,HIPBLASLT_VERSION,0.6.0-592518e7
+Validator,ROCM_VERSION,6.0.0.0-91-08e5094
+Validator,GCN_ARCH_NAME,gfx942:sramecc+:xnack-
+GemmTunableOp_float_NT,nt_768_3072_16384,Gemm_Rocblas_69720,0.751226
+GemmTunableOp_float_NT,nt_3072_768_16384,Gemm_Rocblas_69733,0.684042
+GemmTunableOp_float_NT,nt_768_768_16384,Gemm_Hipblaslt_NT_28806,0.264226
+GemmTunableOp_float_NT,nt_768_30522_16384,Gemm_Hipblaslt_NT_27808,5.73919
+GemmTunableOp_float_NN,nn_768_16384_3072,Gemm_Hipblaslt_NN_33293,0.701076
+GemmTunableOp_float_NN,nn_768_16384_768,Gemm_Hipblaslt_NN_33685,0.209309
+GemmTunableOp_float_NN,nn_3072_16384_768,Gemm_Hipblaslt_NN_33225,0.69655
+GemmTunableOp_float_NN,nn_768_16384_30522,Gemm_Hipblaslt_NN_33924,5.81957
+GemmTunableOp_float_TN,tn_30522_16384_768,Default,6.06459
+GemmTunableOp_float_TN,tn_768_16384_3072,Gemm_Hipblaslt_TN_34830,0.584625
+GemmTunableOp_float_TN,tn_3072_16384_768,Gemm_Rocblas_69037,0.742789
+GemmTunableOp_float_TN,tn_768_16384_768,Gemm_Rocblas_69047,0.211827
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
@@ -43,6 +43,22 @@ install_prepare() {
         pip uninstall torch torchvision torchaudio
         pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1
         pip uninstall pynvml
+
+        # sudo apt-get install lld
+        # https://github.com/ROCm/jax/releases/tag/rocm-jaxlib-v0.4.30
+        # does not really work
+        pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm611-cp310-cp310-manylinux2014_x86_64.whl
+        pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.30.tar.gz
+
+        # 
+        FORCE_CUDA=1 pip install -U -v --no-build-isolation git+https://github.com/rusty1s/pytorch_cluster.git
+        FORCE_CUDA=1 pip install -U -v --no-build-isolation git+https://github.com/rusty1s/pytorch_scatter.git
+        FORCE_CUDA=1 pip install -U -v --no-build-isolation git+https://github.com/rusty1s/pytorch_sparse.git
+
+        # takes forever to compile
+        # https://github.com/ROCm/xformers
+        pip install -v -U --no-build-isolation --no-deps git+https://github.com/ROCm/xformers.git@develop#egg=xformers
+        pip install -v -U --no-build-isolation --no-deps git+https://github.com/ROCm/flash-attention.git
     )
 
     #