install flash-attn with cuda

app.py

@@ -11,15 +11,15 @@ CHECKPOINTS_PATH = os.path.join(PWD, "checkpoints")
 # import subprocess
 
 # copy cudnn files
-# subprocess.run("cp /usr/local/lib/python3.10/site-packages/nvidia/cudnn/include/*.h /usr/local/cuda/include", env={}, shell=True)
-# subprocess.run("cp /usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib/*.so* /usr/local/cuda/lib64", env={}, shell=True)
+os.system("cp /root/.pyenv/versions/3.10.18/lib/python3.10/site-packages/nvidia/cudnn/include/*.h /usr/local/cuda/include")
+os.system("cp /root/.pyenv/versions/3.10.18/lib/python3.10/site-packages/nvidia/cudnn/lib/*.so* /usr/local/cuda/lib64")
 
 os.system("apt-get update && apt-get install -qqy libmagickwand-dev")
 
 # install packages
-os.system('export FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE && pip install --no-build-isolation "flash-attn<=2.7.4.post1"')
+# os.system('export FLASH_ATTENTION_SKIP_CUDA_BUILD=FALSE && pip install --timeout=1000000000 --no-build-isolation "flash-attn<=2.7.4.post1"')
 os.system(
-    "pip install https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl"
+    "pip install --timeout=1000000000 https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl"
 )
 os.system('export VLLM_ATTENTION_BACKEND=FLASHINFER && pip install "vllm==0.9.0"')
 os.system('pip install "decord==0.6.0"')
@@ -31,10 +31,12 @@ os.system(
     "export CONDA_PREFIX=/usr/local/cuda && ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10"
 )
 
-os.system('pip install --no-build-isolation "transformer-engine[pytorch]"')
-os.system('pip install "decord==0.6.0"')
+os.system('pip install --timeout=1000000000 --no-build-isolation "transformer-engine[pytorch]"')
+os.system('pip install --timeout=1000000000 "decord==0.6.0"')
 
-os.system('pip install "git+https://github.com/nvidia-cosmos/cosmos-transfer1"')
+os.system(
+    'pip install --timeout=1000000000 "git+https://github.com/nvidia-cosmos/cosmos-transfer1@e4055e39ee9c53165e85275bdab84ed20909714a"'
+)
 
 # setup env
 os.environ["CUDA_HOME"] = "/usr/local/cuda"

requirements.txt

@@ -23,7 +23,7 @@ iopath==0.1.10
 loguru==0.7.3
 matplotlib==3.10.3
 mediapy==1.2.4
-megatron_core
+megatron_core==0.10.0
 natsort==8.4.0
 nltk==3.9.1
 numpy==2.2.6