Allow Flash attention to build

Author: Doug Lehr

docker/Dockerfile.rocm_base

@@ -78,13 +78,13 @@ RUN cd /opt/rocm/share/amd_smi \
     && pip wheel . --wheel-dir=dist
 RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
 
-#FROM base AS build_pytorch
+FROM base AS build_pytorch
 #ARG PYTORCH_BRANCH
 #ARG PYTORCH_VISION_BRANCH
 #ARG PYTORCH_REPO
 #ARG PYTORCH_VISION_REPO
-#ARG FA_BRANCH
-#ARG FA_REPO
+ARG FA_BRANCH
+ARG FA_REPO
 #RUN git clone ${PYTORCH_REPO} pytorch
 #RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
 #    pip install -r requirements.txt && git submodule update --init --recursive \
@@ -127,8 +127,8 @@ RUN mkdir /app/debs
 #    cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
-#RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-#    cp /install/*.whl /app/debs
+RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs
 # RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
 #     cp /install/*.whl /app/debs
 
@@ -142,8 +142,8 @@ FROM base AS final
 #    pip install /install/*.whl
 RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
     pip install /install/*.whl
-#RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-#    pip install /install/*.whl
+RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
+    pip install /install/*.whl
 # RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
 #     pip install /install/*.whl