LlamaFactory/docker/docker-cuda/Dockerfile.megatron

# NVIDIA official image (ubuntu-24.04 + cuda-12.9.1 + python-3.12)
# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-25-06.html
FROM nvcr.io/nvidia/pytorch:25.06-py3

ENV DEBIAN_FRONTEND=noninteractive
ENV PIP_ROOT_USER_ACTION=ignore
ENV PYPI_MIRROR=https://mirrors.aliyun.com/pypi/simple/
ENV PYPI_TRUSTED_HOST=mirrors.aliyun.com
ENV APT_MIRROR=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/
ENV PIP_CONSTRAINT=""

RUN pip install --upgrade pip setuptools wheel "hatchling>=1.18.0" editables --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PYPI_MIRROR}

RUN pip uninstall -y torch torchvision torch-tensorrt \
    flash_attn transformer-engine \
    cudf dask-cuda cugraph cugraph-service-server cuml raft-dask cugraph-dgl cugraph-pyg dask-cudf

RUN pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu129

RUN pip uninstall -y opencv opencv-python opencv-python-headless && \
    rm -rf /usr/local/lib/python3.12/dist-packages/cv2/ && \
    pip install opencv-python-headless==4.11.0.86 --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PYPI_MIRROR}

RUN pip install --trusted-host mirrors.aliyun.com --index-url ${PYPI_MIRROR} \
    "megatron-core>=0.13.0,<0.14.0" "deepspeed==0.16.4"

WORKDIR /build

ARG apex_url=git+https://github.com/NVIDIA/apex.git@25.04
RUN pip uninstall -y apex && \
    MAX_JOBS=32 NINJA_FLAGS="-j32" NVCC_APPEND_FLAGS="--threads 32" \
    pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
    --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 32" ${apex_url}

RUN pip install --no-build-isolation transformer_engine[pytorch]

RUN rm -rf /build
WORKDIR /workspace

RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \
    { \
    echo "deb ${APT_MIRROR} jammy main restricted universe multiverse"; \
    echo "deb ${APT_MIRROR} jammy-security main restricted universe multiverse"; \
    echo "deb ${APT_MIRROR} jammy-updates main restricted universe multiverse"; \
    echo "deb ${APT_MIRROR} jammy-backports main restricted universe multiverse"; \
    } > /etc/apt/sources.list

RUN apt-get update && apt-get install -y zip

RUN apt-get install -y openjdk-21-jdk
ENV JAVA_HOME /usr/lib/jvm/java-21-openjdk-amd64

ARG REPO_URL=https://github.com/hiyouga/LlamaFactory.git
ARG BRANCH=main
WORKDIR /app

# Clone the repository
RUN git clone --depth 1 --branch ${BRANCH} ${REPO_URL} /app || \
    git clone --depth 1 ${REPO_URL} /app

# Install LLaMA Factory
RUN pip install --no-cache-dir -e . --no-build-isolation && \
    pip install --no-cache-dir -r requirements/metrics.txt --no-build-isolation

RUN pip install "git+https://github.com/alibaba/roll.git#subdirectory=mcore_adapter"

# Expose port 7860 for LLaMA Board
ENV GRADIO_SERVER_PORT=7860
EXPOSE 7860

# Expose port 8000 for API service
ENV API_PORT=8000
EXPOSE 8000

# unset proxy
ENV http_proxy=
ENV https_proxy=