Dockerfile

FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04 AS base

WORKDIR /workspace

ENV DEBIAN_FRONTEND=noninteractive
ENV LANG=C.UTF-8
ENV LC_ALL=C.UTF-8

# use tsinghua source
RUN sed -i 's|http://archive.ubuntu.com/ubuntu/|https://mirrors.tuna.tsinghua.edu.cn/ubuntu/|g' /etc/apt/sources.list \
    && sed -i 's|http://security.ubuntu.com/ubuntu/|https://mirrors.tuna.tsinghua.edu.cn/ubuntu/|g' /etc/apt/sources.list

RUN apt-get update && apt install -y software-properties-common  \
    && add-apt-repository ppa:deadsnakes/ppa \
    && apt-get update \
    && apt-get install -y vim tmux zip unzip wget git cmake build-essential \
     curl libibverbs-dev ca-certificates iproute2 \
     ffmpeg libsm6 libxext6 \
    && apt-get install -y python3.11 python3.11-venv python3.11-dev python3-pip \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1

RUN pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple \
    && pip install packaging ninja vllm torch torchvision diffusers transformers \
     tokenizers accelerate safetensors opencv-python numpy imageio imageio-ffmpeg \
     einops loguru sgl-kernel qtorch ftfy

# please download flash-attention source code first
# git clone https://github.com/Dao-AILab/flash-attention.git --recursive
# todo: add third party repo feature
COPY flash-attention /workspace/flash-attention

# install flash-attention 2
RUN cd flash-attention && pip install --no-cache-dir -v -e .

# install flash-attention 3, only if hopper
RUN cd flash-attention/hopper && pip install --no-cache-dir -v -e .