mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2026-07-28 11:46:09 +08:00
86 lines
3.3 KiB
Docker
86 lines
3.3 KiB
Docker
# LLaMA-Factory + Megatron Bridge (CUDA) runtime
|
|
# Mirrors the verified host env: Python 3.12, PyTorch 2.12.1+cu126,
|
|
# TransformerEngine 2.17, megatron-core 0.18, megatron-bridge 0.5.
|
|
#
|
|
# CUDA user-mode libs come from the PyTorch cu126 wheels (same as the host venv).
|
|
# Layers are intentionally few (helps vfs / nested-docker disk usage).
|
|
#
|
|
# Build from repo root:
|
|
# bash docker/docker-cuda/build-megatron.sh
|
|
|
|
ARG BASE_IMAGE=ubuntu:24.04
|
|
FROM ${BASE_IMAGE}
|
|
|
|
ARG PIP_INDEX=https://mirrors.aliyun.com/pypi/simple
|
|
ARG PYPI_TRUSTED_HOST=mirrors.aliyun.com
|
|
ARG TORCH_INDEX=https://download.pytorch.org/whl/cu126
|
|
ARG APT_MIRROR=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/
|
|
|
|
ENV DEBIAN_FRONTEND=noninteractive \
|
|
PIP_ROOT_USER_ACTION=ignore \
|
|
PIP_DISABLE_PIP_VERSION_CHECK=1 \
|
|
PIP_BREAK_SYSTEM_PACKAGES=1 \
|
|
PIP_CONSTRAINT="" \
|
|
MAX_JOBS=8 \
|
|
PYTHONDONTWRITEBYTECODE=1 \
|
|
PYTHONUNBUFFERED=1 \
|
|
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
|
DISABLE_VERSION_CHECK=1 \
|
|
USE_MEGATRON_BRIDGE=1 \
|
|
GRADIO_SERVER_PORT=7860 \
|
|
API_PORT=8000 \
|
|
http_proxy= \
|
|
https_proxy=
|
|
|
|
SHELL ["/bin/bash", "-c"]
|
|
WORKDIR /app
|
|
|
|
# System deps + Ubuntu Python 3.12
|
|
RUN if [ -f /etc/apt/sources.list.d/ubuntu.sources ]; then \
|
|
sed -i "s|http://archive.ubuntu.com/ubuntu|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu|${APT_MIRROR}|g" /etc/apt/sources.list.d/ubuntu.sources; \
|
|
elif [ -f /etc/apt/sources.list ]; then \
|
|
sed -i "s|http://archive.ubuntu.com/ubuntu/|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu/|${APT_MIRROR}|g" /etc/apt/sources.list; \
|
|
fi && \
|
|
apt-get update && apt-get install -y --no-install-recommends \
|
|
ca-certificates curl git vim wget \
|
|
build-essential ninja-build cmake libgomp1 zip unzip \
|
|
python3 python3-pip python3-dev python3-venv && \
|
|
ln -sf /usr/bin/python3 /usr/bin/python && \
|
|
rm -rf /var/lib/apt/lists/* && \
|
|
pip install --no-cache-dir --upgrade pip setuptools wheel packaging ninja pybind11 \
|
|
--trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX}
|
|
|
|
COPY . /app
|
|
|
|
# PyTorch + TE + Megatron Bridge + LLaMA-Factory in one layer
|
|
RUN pip install --no-cache-dir \
|
|
torch==2.12.1 torchvision==0.27.1 torchaudio==2.11.0 \
|
|
--index-url ${TORCH_INDEX} && \
|
|
pip install --no-cache-dir --no-build-isolation \
|
|
"transformer-engine[pytorch]==2.17.0" \
|
|
--trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
|
|
pip install --no-cache-dir --no-build-isolation \
|
|
"megatron-bridge==0.5.0" \
|
|
--trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
|
|
pip install --no-cache-dir --no-build-isolation -e . \
|
|
--trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
|
|
pip install --no-cache-dir --no-build-isolation \
|
|
-r requirements/metrics.txt \
|
|
--trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
|
|
pip install --no-cache-dir --no-build-isolation \
|
|
"megatron-bridge==0.5.0" \
|
|
--trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
|
|
python - <<'PY'
|
|
import torch
|
|
import megatron.core
|
|
import transformer_engine
|
|
from megatron.bridge import AutoBridge # noqa: F401
|
|
import llamafactory
|
|
print("torch", torch.__version__, "cuda", torch.version.cuda)
|
|
print("te", transformer_engine.__version__)
|
|
print("megatron-bridge import ok")
|
|
PY
|
|
|
|
EXPOSE 7860 8000
|
|
CMD ["bash"]
|