reinstall fa, te, apex, gg

Author: wenxie-amd

.github/workflows/docker/Dockerfile

@@ -25,6 +25,64 @@ RUN pip3 uninstall aiter primus_turbo -y && \
     pip3 install -r requirements.txt && \
     GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation .
 
+RUN pip uninstall flash_attn transformer_engine grouped_gemm -y
+
+# FA v3.0.0.r1
+ARG FA_REPO=https://github.com/ROCm/flash-attention.git
+ARG FA_BRANCH=v3.0.0.r1-cktile
+ARG GPU_ARCHS=gfx942
+
+RUN git clone --recursive ${FA_REPO} -b ${FA_BRANCH} \
+    && cd flash-attention \
+    && python setup.py install \
+    && cd .. \
+    && rm -rf flash-attention
+
+# Latest TransformerEngine
+ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git
+ARG TE_BRANCH=ec21e1f35f5a36c1def9c89c4865eba3947f0665
+
+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_USES_BWD_V3=1
+ENV NVTE_CK_HOW_V3_BF16_CVT=2
+ENV GPU_ARCHS=gfx942
+ENV CU_NUM=304
+
+RUN git clone --recursive ${TE_REPO} \
+    && cd TransformerEngine \
+    && git checkout ${TE_BRANCH} \
+    && git submodule update --init --recursive \
+    && 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
+
+# APEX 1.7.0
+ARG APEX_REPO=https://github.com/rocm/apex
+ARG APEX_BRANCH=release/1.7.0
+
+RUN pip uninstall -y apex \
+    && git clone ${APEX_REPO} -b ${APEX_BRANCH} \
+    && cd 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