fix docker file bug

This commit is contained in:
mulin.lyh
2023-09-01 18:05:11 +08:00
parent ad1081b08e
commit c11bbb03fb

View File

@@ -50,11 +50,8 @@ RUN CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6" pip
# torchmetrics==0.11.4 for ofa
RUN pip install --no-cache-dir tiktoken torchmetrics==0.11.4 'transformers<4.31.0' transformers_stream_generator 'protobuf<=3.20.0' bitsandbytes basicsr
COPY docker/scripts/install_flash_attension.sh /tmp/install_flash_attension.sh
RUN bash /tmp/install_flash_attension.sh
&& \
git clone -b v1.0.8 https://github.com/Dao-AILab/flash-attention && \
cd flash-attention && pip install . && \
pip install csrc/layer_norm && \
pip install csrc/rotary && \
cd .. && \
rm -rf flash-attention
RUN if [ "$USE_GPU" = "True" ] ; then \
bash /tmp/install_flash_attension.sh; \
else \
echo 'cpu unsupport flash attention'; \
fi