-

mila-iqia · Oct 10, 2024 · baf5304 · baf5304
1 parent db19114
commit baf5304
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 6 deletions.
diff --git a/benchmarks/rlhf/requirements.in b/benchmarks/rlhf/requirements.in
@@ -4,3 +4,4 @@ trl
 accelerate
 transformers
 datasets
+einops
diff --git a/benchmarks/rlhf/requirements.rocm.txt b/benchmarks/rlhf/requirements.rocm.txt
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
@@ -18,6 +18,13 @@ else
     export MILABENCH_CONFIG="$MILABENCH_SOURCE/config/standard.yaml"
 fi
 
+
+export GPU="$(/opt/rocm/lib/llvm/bin/amdgpu-arch | head -n 1)"
+export TORCH_ROCM_ARCH_LIST="$GPU"
+export ROCM_TARGETS="$GPU"
+export PYTORCH_ROCM_ARCH="$GPU"
+
+
 ARGS="$@"
 
 install_prepare() {
@@ -68,21 +75,26 @@ install_prepare() {
 
         # https://github.com/ROCm/jax/releases/tag/rocm-jaxlib-v0.4.30
         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
+        pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.30.tar.g
 
-        # 
-        FORCE_ONLY_CUDA=1 pip install -U -v --no-build-isolation git+https://github.com/rusty1s/pytorch_cluster.git
-        FORCE_ONLY_CUDA=1 pip install -U -v --no-build-isolation git+https://github.com/rusty1s/pytorch_scatter.git
-        FORCE_ONLY_CUDA=1 pip install -U -v --no-build-isolation git+https://github.com/rusty1s/pytorch_sparse.git
+        pip uninstall torch_cluster torch_scatter torch_sparse -y
+        FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_cluster.git
+        FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_scatter.git
+        FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_sparse.git
 
         # takes forever to compile
         # https://github.com/ROCm/xformers
+        pip uninstall xformers
         pip install xformers --index-url https://download.pytorch.org/whl/rocm6.1
         # 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/facebookresearch/xformers.git
         # pip install xformers -U --index-url https://download.pytorch.org/whl/rocm6.1
+
+        pip uninstall flash-attention
         pip install -v -U --no-build-isolation --use-pep517 --no-deps git+https://github.com/ROCm/flash-attention.git 
         pip uninstall pynvml nvidia-ml-py -y
+
+        pip install einops
     )
 
     pip uninstall pynvml nvidia-ml-py -y
@@ -99,7 +111,9 @@ else
     . $MILABENCH_WORDIR/env/bin/activate
 fi
 
-cd $MILABENCH_WORDIR
+
+milabench prepare $ARGS 
+
 #
 #   Run the benchmakrs
 milabench run $ARGS