yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-12-12 15:51:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Flash attention as part of the image

Browse Source
This commit is contained in:
Jake Poznanski 2024-09-24 11:57:56 -07:00
parent 3c8e05362f
commit 0d9917367b
1 changed files with 18 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
scripts/beaker/Dockerfile
View File

@ -1,3 +1,20 @@
FROM gcr.io/ai2-beaker-core/public/cqgl31u2ba5vrtuc91og:latest FROM gcr.io/ai2-beaker-core/public/cqgl31u2ba5vrtuc91og:latest
RUN apt update && apt-get install -y libaio-dev # Update the package list and install libaio-dev and gnupg2
RUN apt update && apt-get install -y libaio-dev gnupg2
# Add NVIDIA package repository keys
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub \
&& apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub \
&& apt-get -y update
# Set up the NVIDIA CUDA repository
RUN apt-get install -y software-properties-common \
&& add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /" \
&& apt-get update
# Install CUDA toolkit and nvcc 12.1
RUN apt-get install -y cuda-nvcc-12-1
# Get flash attention setup
RUN pip install flash-attn --no-build-isolation