update dockerfile for vllm 011

Author: pan-x-c

pyproject.toml

@@ -73,7 +73,7 @@ dev = [
 ]
 megatron = [
     "megatron-core[mlm]==0.13.1",
-    "transformer_engine[pytorch]==2.6.0.post1",
+    "transformer_engine[pytorch]==2.8.0",
     "mbridge>=0.13.0",
 ]
 

scripts/docker/Dockerfile

@@ -26,7 +26,7 @@ RUN apt update && apt install -y \
 # copy the Trinity-RFT dir into the workspace
 COPY . .
 
-RUN pip install --upgrade pip && pip install -e .[dev] && pip install -e .[flash_attn]
+RUN pip install --upgrade pip && pip install -e .[dev] && pip install flash_attn==2.8.1 --no-build-isolation
 
 # Set Env variables
 

scripts/docker_for_megatron/Dockerfile

@@ -27,12 +27,13 @@ COPY . .
 
 # Install Trinity-RFT with Megatron
 RUN pip install --upgrade pip \
+    && pip install -e .[dev] \
+    && pip install flash_attn==2.8.1 --no-build-isolation \
     && pip install -e .[megatron] \
-    && pip install -e .[flash_attn] \
     && pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
         --config-settings "--build-option=--cpp_ext" \
         --config-settings "--build-option=--cuda_ext" \
-        --resume-retries 999 git+https://github.com/NVIDIA/apex.git
+        --resume-retries 20 git+https://github.com/NVIDIA/apex.git
 
 # Set Env variables