reinstall fa, te, apex, gg

Author: wenxie-amd

.github/workflows/docker/Dockerfile

@@ -25,31 +25,36 @@ RUN pip3 uninstall aiter primus_turbo -y && \
     pip3 install -r requirements.txt && \
     GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation .
 
+ENV MAX_JOBS=128
+ENV PYTORCH_ROCM_ARCH="gfx942;gfx950"
 RUN pip uninstall flash_attn transformer_engine grouped_gemm -y
 
-# FA v3.0.0.r1
+# FA
 ARG FA_REPO=https://github.com/ROCm/flash-attention.git
-ARG FA_BRANCH=v3.0.0.r1-cktile
-ARG GPU_ARCHS=gfx942
+ARG FA_BRANCH=6387433156558135a998d5568a9d74c1778666d8
+ENV GPU_ARCHS="${PYTORCH_ROCM_ARCH}"
 
-RUN git clone --recursive ${FA_REPO} -b ${FA_BRANCH} \
+RUN git clone --recursive ${FA_REPO} \
     && cd flash-attention \
+    && git checkout ${FA_BRANCH} \
     && python setup.py install \
     && cd .. \
     && rm -rf flash-attention
 
 # Latest TransformerEngine
 ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git
-ARG TE_BRANCH=ec21e1f35f5a36c1def9c89c4865eba3947f0665
+ARG TE_BRANCH=1834247827f47f5eab76ea790841168ac778592e
 
 ENV NVTE_USE_HIPBLASLT=1
 ENV NVTE_FRAMEWORK=pytorch
 ENV NVTE_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
 ENV NVTE_USE_CAST_TRANSPOSE_TRITON=0
+ENV NVTE_CK_IS_V3_ATOMIC_FP32=0
 ENV NVTE_CK_USES_BWD_V3=1
+ENV NVTE_CK_USES_FWD_V3=1
+ENV CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=2
 ENV NVTE_CK_HOW_V3_BF16_CVT=2
-ENV GPU_ARCHS=gfx942
-ENV CU_NUM=304
+ENV NVTE_USE_ROCM=1
 
 RUN git clone --recursive ${TE_REPO} \
     && cd TransformerEngine \
@@ -58,20 +63,15 @@ RUN git clone --recursive ${TE_REPO} \
     && MAX_JOBS=${MAX_JOBS} pip install . \
     && cd ..
 
-RUN apt --fix-broken install -y
-RUN pip install datasets numpy==1.26.4 transformers
-RUN pip install --upgrade 'optree>=0.13.0'
-
-WORKDIR /workspace/
-
 # Groupped GEMM
 ARG GROUPED_GEMM_REPO=https://github.com/caaatch22/grouped_gemm.git
 ARG GROUPED_GEMM_BRANCH=rocm
+
 RUN git clone ${GROUPED_GEMM_REPO} &&\
     cd grouped_gemm &&\
     git checkout ${GROUPED_GEMM_BRANCH} &&\
     git submodule update --init --recursive &&\
-    pip install . && cd .. && rm -rf grouped_gemm
+    pip install . && cd ../ && rm -rf grouped_gemm
 
 # APEX 1.7.0
 ARG APEX_REPO=https://github.com/rocm/apex
@@ -83,7 +83,6 @@ RUN pip uninstall -y apex \
     && 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" ./ \
     && cd .. && rm -r apex
 
-
 # Set the default working directory
 WORKDIR /opt