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

79 lines
3.1 KiB
Docker

ARG BUILD_BASE_IMAGE=mx-devops-acr-cn-shanghai.cr.volces.com/opensource/public-ai-release/maca/sglang:0.5.10-maca.ai3.7.1.12-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 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 && \
git clone --depth 1 --branch ${SWIFT_VERSION} https://github.com/modelscope/ms-swift.git && \
git clone --depth 1 --branch ${VLLM_METAX_VERSION} https://github.com/MetaX-MACA/vLLM-metax.git && \
git clone --depth 1 --branch ${VLLM_VERSION} https://github.com/vllm-project/vllm.git && \
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/cuda.txt && \
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 cd /workspace/Megatron-LM && \
if [ -f megatron/legacy/fused_kernels/__init__.py ]; then \
sed -i 's/nvcc/cucc/g' megatron/legacy/fused_kernels/__init__.py; \
fi && \
pip install .
# Step 5: install transformer-engine
ARG TE_VERSION=2.8.0+4a002bf5.maca3.5.3.105.torch2.8
ARG UV_TRUSTED_HOST=repos.metax-tech.com
RUN pip install transformer_engine==${TE_VERSION} -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host ${UV_TRUSTED_HOST}
# Step 6: patch and install ms-swift v4.2.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
RUN pip uninstall torchao -y
CMD ["bash"]