Update Dockerfile

Dockerfile

@@ -85,15 +85,15 @@ RUN /opt/conda/bin/conda install -c "nvidia/label/cuda-11.8.0"  cuda==11.8 && \
     /opt/conda/bin/conda clean -ya
 
 
-# Build Flash Attention CUDA kernels
-FROM kernel-builder as flash-att-builder
+# # Build Flash Attention CUDA kernels
+# FROM kernel-builder as flash-att-builder
 
-WORKDIR /usr/src
+# WORKDIR /usr/src
 
-COPY server/Makefile-flash-att Makefile
+# COPY server/Makefile-flash-att Makefile
 
-# Build specific version of flash attention
-RUN make build-flash-attention
+# # Build specific version of flash attention
+# RUN make build-flash-attention
 
 # Build Transformers CUDA kernels
 FROM kernel-builder as transformers-builder
@@ -142,9 +142,9 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
 COPY --from=pytorch-install /opt/conda /opt/conda
 
 # Copy build artifacts from flash attention builder
-COPY --from=flash-att-builder /usr/src/flash-attention/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
-COPY --from=flash-att-builder /usr/src/flash-attention/csrc/layer_norm/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
-COPY --from=flash-att-builder /usr/src/flash-attention/csrc/rotary/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
+# COPY --from=flash-att-builder /usr/src/flash-attention/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
+# COPY --from=flash-att-builder /usr/src/flash-attention/csrc/layer_norm/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
+# COPY --from=flash-att-builder /usr/src/flash-attention/csrc/rotary/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
 
 # Copy build artifacts from transformers builder
 COPY --from=transformers-builder /usr/src/transformers /usr/src/transformers