From 50ca04b052f5abf647bbc5adbacad662db457082 Mon Sep 17 00:00:00 2001 From: drbh Date: Tue, 6 Feb 2024 21:15:16 +0000 Subject: [PATCH] feat: update docker for mamba --- Dockerfile | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/Dockerfile b/Dockerfile index b6c5b2ed..cc1a501a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -159,6 +159,12 @@ FROM kernel-builder as megablocks-builder RUN pip install git+https://github.com/OlivierDehaene/megablocks@181709df192de9a941fdf3a641cdc65a0462996e +# Build mamba kernels +FROM kernel-builder as mamba-builder +WORKDIR /usr/src +COPY server/Makefile-selective-scan Makefile +RUN make build-all + # Text Generation Inference base image FROM nvidia/cuda:12.1.0-base-ubuntu20.04 as base @@ -205,6 +211,10 @@ COPY --from=eetq-kernels-builder /usr/src/eetq/build/lib.linux-x86_64-cpython-31 # Copy builds artifacts from vllm builder COPY --from=vllm-builder /usr/src/vllm/build/lib.linux-x86_64-cpython-310 /opt/conda/lib/python3.10/site-packages +# Copy build artifacts from mamba builder +COPY --from=mamba-builder /usr/src/mamba/build/lib.linux-x86_64-cpython-310/ /opt/conda/lib/python3.10/site-packages +COPY --from=mamba-builder /usr/src/causal-conv1d/build/lib.linux-x86_64-cpython-310/ /opt/conda/lib/python3.10/site-packages + # Install flash-attention dependencies RUN pip install einops --no-cache-dir