File size: 2,650 Bytes
b0734ca | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | # ROMA inference image for Quadro RTX 6000 (Turing sm_75, x86_64, 24GB ×4 = 96GB).
#
# Base: CUDA 12.4 runtime on Ubuntu 22.04. CUDA 12.4 matches the repo's torch 2.6.0+cu124, and runs
# fine on the host's 570.x / CUDA-12.8 driver (drivers are backward compatible). We install the
# OFFICIAL PyTorch cu124 wheels (via download.pytorch.org), which include Turing (sm_75) kernels --
# NGC datacenter images may omit sm_75 and produce "no kernel image is available" at runtime, so we
# deliberately do NOT use an NGC base here.
#
# Turing notes (handled in the demo scripts, not here): FlashAttention-2 is unsupported on sm_75 so
# the demos default to attn_implementation=sdpa; bf16 is not accelerated so they default to fp16.
ARG BASE_IMAGE=nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
FROM ${BASE_IMAGE}
ENV DEBIAN_FRONTEND=noninteractive
ENV PIP_INDEX=https://pypi.org/simple
ENV TORCH_INDEX=https://download.pytorch.org/whl/cu124
ENV HF_HUB_ENABLE_HF_TRANSFER=1
# The gradio demos call demo.launch() without a server name; bind to all interfaces so the UI is
# reachable from the host (the demos read this env, added in the gradio edits).
ENV GRADIO_SERVER_NAME=0.0.0.0
ENV PYTHONUNBUFFERED=1
ARG TRANSFORMERS_REF=roma_patch
# System deps: Python 3.10, ffmpeg (decord/av/moviepy), git (to pip-install the transformers fork).
RUN apt-get update && apt-get install -y --no-install-recommends \
python3.10 python3.10-dev python3-pip \
ffmpeg git ca-certificates && \
ln -sf /usr/bin/python3.10 /usr/bin/python && \
rm -rf /var/lib/apt/lists/*
WORKDIR /app
# Install the runtime deps. The +cu124 torch/vision/audio pins resolve from the PyTorch index; the
# transformers@roma_patch fork is fetched from git (already pinned in the requirements file).
COPY requirements-rtx6000.txt /app/requirements-rtx6000.txt
RUN python -m pip install --upgrade pip && \
python -m pip install \
--index-url "$PIP_INDEX" \
--extra-index-url "$TORCH_INDEX" \
-r /app/requirements-rtx6000.txt
# Register the llamafactory package editable from the local checkout (deps already satisfied above).
COPY . /app
RUN python -m pip install --no-deps -e .
# Build-time sanity import (GPU is not visible during build, so we don't assert cuda here).
RUN python -c "import torch, transformers, bitsandbytes; print('torch', torch.__version__, '| transformers', transformers.__version__)"
ENV GRADIO_SERVER_PORT=7860
EXPOSE 7860
ENV API_PORT=8000
EXPOSE 8000
# Persist the checkpoint, HF cache and demo media across container runs.
VOLUME [ "/app/whole_model", "/root/.cache/huggingface", "/app/demo_media" ]
CMD ["bash"]
|