Files
modelscope/docker/Metax/4.1/Dockerfile.with_metax_image

74 lines
2.8 KiB
Docker

ARG BUILD_BASE_IMAGE=mx-devops-acr-cn-shanghai.cr.volces.com/opensource/public-ai-release/maca/ms-swift:4.0.4-maca.ai3.5.3.5-torch2.8-py312-ubuntu22.04-amd64
ARG PYTHON_VERSION=3.12
FROM ${BUILD_BASE_IMAGE} AS base
# NOTE:
# This fast-build path inherits Python/Torch/TE from a prebuilt Metax release image.
# We keep the verified base image tag here instead of guessing a newer one.
# As a result, this path may lag behind the Megatron-SWIFT Quick Start recommendations.
# may need passing a particular vllm version during build
ARG VLLM_VERSION
ARG VLLM_METAX_VERSION
ARG MEGATRON_VERSION
ARG SWIFT_VERSION
ENV MACA_PATH=/opt/maca
ENV CUCC_CMAKE_ENTRY=2
ENV CUDA_PATH=/root/cu-bridge/CUDA_DIR
ENV CUCC_PATH=${MACA_PATH}/tools/cu-bridge
ENV PATH=/opt/conda/bin:/opt/conda/condabin:${CUDA_PATH}/bin:${CUCC_PATH}/tools:${CUCC_PATH}/bin:${MACA_PATH}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_PATH}/lib64:${MACA_PATH}/lib:${MACA_PATH}/mxgpu_llvm/lib:${LD_LIBRARY_PATH}
WORKDIR /workspace
COPY requirements_extra.txt /workspace/requirements_extra.txt
RUN echo $PATH
RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
# Initialize cu-bridge if it is not already prepared in the base image.
RUN if [ ! -d /root/cu-bridge ]; then \
${MACA_PATH}/tools/cu-bridge/tools/pre_make; \
fi
# Clone all GitHub sources while the external proxy is enabled.
RUN rm -rf /workspace/ms-swift /workspace/vLLM-metax /workspace/vllm /workspace/Megatron-LM
RUN git clone --depth 1 --branch ${SWIFT_VERSION} https://github.com/modelscope/ms-swift.git
RUN git clone --depth 1 --branch ${VLLM_METAX_VERSION} https://github.com/MetaX-MACA/vLLM-metax.git
RUN git clone --depth 1 --branch ${VLLM_VERSION} https://github.com/vllm-project/vllm.git
RUN git clone --depth 1 --branch ${MEGATRON_VERSION} https://github.com/NVIDIA/Megatron-LM.git
# install cmake
RUN pip install cmake ninja
# Step 1: build original vLLM for torch setup
RUN cd vllm && \
python3 use_existing_torch.py && \
pip install -r requirements/build.txt
# Step 2: build vLLM with empty device to avoid CUDA dependency
RUN cd vllm && \
VLLM_TARGET_DEVICE=empty pip install -v . --no-build-isolation
# Step 3: build vLLM-metax
RUN cd vLLM-metax && \
python3 use_existing_metax.py && \
pip install -r requirements/build.txt && \
python3 -m build -w -n && \
pip install dist/*.whl
# Step 4: patch and install Megatron-LM
RUN sed -i 's/nvcc/cucc/g' /workspace/Megatron-LM/megatron/legacy/fused_kernels/__init__.py && \
cd /workspace/Megatron-LM && \
pip install .
# Step 5: patch and install ms-swift v4.1.0 with its Megatron extra
RUN sed -i '0,/^\(from \|import \)/{s//import vllm_metax.patch\n&/}' ms-swift/swift/__init__.py && \
cd ms-swift && \
pip install "transformers<5.4.0" && \
pip install '.[megatron]' && \
pip install -r /workspace/requirements_extra.txt
CMD ["bash"]