File size: 1,249 Bytes
fb11af9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | # Start from the NVIDIA official image (ubuntu-22.04 + python-3.10)
# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html
FROM nvcr.io/nvidia/pytorch:24.08-py3
# Define environments
ENV MAX_JOBS=32
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
ENV DEBIAN_FRONTEND=noninteractive
ENV NODE_OPTIONS=""
# Install systemctl and tini
RUN apt-get update && \
apt-get install -y -o Dpkg::Options::="--force-confdef" systemd tini && \
apt-get clean || { echo "Installation failed"; exit 1; }
RUN apt-get install -y tzdata \
&& ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime \
&& dpkg-reconfigure -f noninteractive tzdata
# Change pip source
RUN python -m pip install --upgrade pip
# Install torch-2.5.1 + vllm-0.7.3
RUN pip install --no-cache-dir vllm==0.7.3 torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 tensordict torchdata \
transformers>=4.49.0 accelerate datasets peft hf-transfer diffusers \
codetiming hydra-core pandas pyarrow>=15.0.0 pylatexenc qwen-vl-utils wandb ninja liger-kernel \
pytest yapf py-spy pyext pre-commit ruff packaging
# Install flux
RUN pip install --no-cache-dir byte-flux
# Install flash-attn and triton
RUN pip install --no-cache-dir flash-attn triton>=3.1.0
|