diff --git a/environment_setup.sh b/environment_setup.sh
index c0a79012..ef748e27 100755
--- a/environment_setup.sh
+++ b/environment_setup.sh
@@ -21,7 +21,7 @@ pip install --upgrade pip setuptools
 pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.8/flash_attn-2.5.8+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 
 # Install VILA
-pip install -e ".[train,eval]"
+pip install -e ".[train,eval]" --use-pep517
 
 # Quantization requires the newest triton version, and introduce dependency issue
 pip install triton==3.1.0
diff --git a/pyproject.toml b/pyproject.toml
index 099564c7..d7e8b3b1 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,7 +34,7 @@ vila-infer = "llava.cli.infer:main"
 vila-upload = "llava.cli.upload2hf:main"
 
 [project.optional-dependencies]
-train = ["deepspeed==0.9.5", "ninja", "wandb"]
+train = ["deepspeed>=0.9.5", "ninja", "wandb"]
 eval = ["word2number", "Levenshtein", "nltk", "pywsd"]
 
 [project.urls]