-
Notifications
You must be signed in to change notification settings - Fork 86
[dockerfile]: swich to vllm==0.21.0 #9
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,122 +1,97 @@ | ||
| ARG SGLANG_IMAGE_TAG=v0.5.10.post1 | ||
| FROM slimerl/sglang:${SGLANG_IMAGE_TAG} AS sglang | ||
| ARG BASE_IMAGE=vllm/vllm-openai:v0.20.2 | ||
| FROM ${BASE_IMAGE} | ||
|
|
||
| # ======================================== Arguments ============================================= | ||
|
|
||
| ARG PATCH_VERSION=latest | ||
| ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087 | ||
| ARG VLLM_VERSION=0.20.2 | ||
| ARG MEGATRON_COMMIT=3714d81d418c9f1bca4594fc35f9e8289f652862 | ||
| ARG PATCH_VERSION=v0.5.9 | ||
| ARG INSTALL_VLLM=0 | ||
| ARG VLLM_PIP_INSTALL_ARGS= | ||
| ARG ENABLE_CUDA_13=1 | ||
| ARG SGLANG_VERSION=0.5.9 | ||
|
|
||
| ARG ENABLE_CUDA_13=0 | ||
| SHELL ["/bin/bash", "-euxo", "pipefail", "-c"] | ||
|
|
||
| # ======================================== Setup ============================================= | ||
| WORKDIR /root | ||
|
|
||
| WORKDIR /root/ | ||
| RUN apt-get update && \ | ||
| apt-get install -y --no-install-recommends git rsync dnsutils nvtop && \ | ||
| rm -rf /var/lib/apt/lists/* | ||
|
|
||
| # ======================================== Apt dependencies ============================================= | ||
|
|
||
| RUN apt update | ||
| RUN apt install -y nvtop rsync dnsutils | ||
|
|
||
| # ====================================== Python dependencies ============================================ | ||
|
|
||
| # The compilation is slow, thus should be put at top | ||
| # TransformerEngines does not support too high FA2 | ||
| RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation | ||
|
|
||
| # The compilation is slow, thus should be put at top | ||
| RUN git clone https://github.com/Dao-AILab/flash-attention.git && \ | ||
| cd flash-attention/ && git checkout fbf24f67cf7f6442c5cfb2c1057f4bfc57e72d89 && git submodule update --init && cd hopper/ && \ | ||
| MAX_JOBS=96 python setup.py install && \ | ||
| export python_path=`python -c "import site; print(site.getsitepackages()[0])"` && \ | ||
| mkdir -p $python_path/flash_attn_3 && \ | ||
| cp flash_attn_interface.py $python_path/flash_attn_3/flash_attn_interface.py && \ | ||
| rm -rf flash-attention/ | ||
|
|
||
| RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps | ||
|
|
||
| RUN pip install flash-linear-attention==0.4.1 | ||
| RUN pip install tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ | ||
| # If the base image is not already vllm/vllm-openai:<version>, set | ||
| # INSTALL_VLLM=1 to install the requested wheel explicitly. | ||
| RUN if [ "${INSTALL_VLLM}" = "1" ]; then \ | ||
| python3 -m pip install --force-reinstall "vllm==${VLLM_VERSION}" ${VLLM_PIP_INSTALL_ARGS}; \ | ||
| fi | ||
|
|
||
| # TE does not have wheel on cuda 13 yet, thus need to install from source | ||
| RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ | ||
| pip install nvidia-mathdx==26.6.0 && \ | ||
| pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10; \ | ||
| pip install --no-deps nvidia-mathdx==26.6.0 && \ | ||
| pip -v install --no-build-isolation --no-deps git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10; \ | ||
| else \ | ||
| pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ | ||
| fi | ||
|
|
||
| RUN NVCC_APPEND_FLAGS="--threads 4" \ | ||
| pip -v install --disable-pip-version-check --no-cache-dir \ | ||
| --no-build-isolation \ | ||
| --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 8" git+https://github.com/NVIDIA/apex.git@10417aceddd7d5d05d7cbf7b0fc2daad1105f8b4 | ||
|
|
||
| RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ | ||
| cd Megatron-LM && git checkout ${MEGATRON_COMMIT} && \ | ||
| pip install -e . | ||
|
|
||
| RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@d64a639 --no-cache-dir --force-reinstall | ||
| RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@bridge --no-deps --no-build-isolation | ||
| RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation | ||
|
|
||
| # This patch from masahi will be included in later Triton releases | ||
| RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ | ||
| (cd /root && git clone -b feat/v350_plus_8045 https://github.com/fzyzcjy/triton.git && cd triton && pip install -r python/requirements.txt && pip install --verbose -e .); \ | ||
| fi | ||
|
|
||
| COPY requirements.txt /tmp/requirements.txt | ||
| RUN pip install --ignore-installed PyJWT && \ | ||
| pip install -r /tmp/requirements.txt | ||
|
|
||
| # Temporarily install another sgl-kernel version for GB300 without rebuilding the whole image | ||
| RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ | ||
| SGL_KERNEL_VERSION=0.3.17.post2 && \ | ||
| python3 -m pip install https://github.com/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sgl_kernel-${SGL_KERNEL_VERSION}+cu130-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps; \ | ||
| fi | ||
| pip -v install --no-build-isolation --no-deps "transformer_engine[pytorch]==2.10.0"; \ | ||
| fi && \ | ||
| pip install --no-deps onnxscript onnx_ir onnx ml_dtypes | ||
|
|
||
| # Keep sglang installed (with --no-deps) so training-side imports keep working | ||
| RUN pip install --no-deps "sglang==${SGLANG_VERSION}" sglang-router==0.3.2 torch_memory_saver==0.0.9 | ||
|
|
||
| # Install slime runtime requirements while excluding SGLang packages. Future | ||
| # code paths should provide vLLM equivalents instead of importing SGLang. | ||
| COPY requirements.txt /tmp/slime-requirements.txt | ||
| RUN python3 - <<'PY' && \ | ||
| python3 -m pip install -r /tmp/slime-requirements-vllm.txt | ||
| from pathlib import Path | ||
|
|
||
| source = Path("/tmp/slime-requirements.txt") | ||
| target = Path("/tmp/slime-requirements-vllm.txt") | ||
| skip_prefixes = ("sglang", "sglang-router") | ||
|
|
||
| lines = [] | ||
| for raw_line in source.read_text().splitlines(): | ||
| normalized = raw_line.split("#", 1)[0].strip().lower() | ||
| if normalized.startswith(skip_prefixes): | ||
| continue | ||
| lines.append(raw_line) | ||
|
|
||
| target.write_text("\n".join(lines) + "\n") | ||
| PY | ||
|
|
||
| # slime's Megatron integration needs the local Megatron patch. This build does | ||
| # not apply docker/patch/*/sglang.patch. | ||
| RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive /root/Megatron-LM && \ | ||
| cd /root/Megatron-LM && \ | ||
| git checkout "${MEGATRON_COMMIT}" | ||
|
|
||
| COPY docker/patch/${PATCH_VERSION}/megatron.patch /root/Megatron-LM/megatron.patch | ||
| RUN cd /root/Megatron-LM && \ | ||
| git update-index --refresh && \ | ||
| git apply megatron.patch --3way && \ | ||
| if grep -R -n '^<<<<<<< ' .; then \ | ||
| echo "Megatron patch failed to apply cleanly. Please resolve conflicts." && \ | ||
| exit 1; \ | ||
| fi && \ | ||
| rm megatron.patch && \ | ||
|
Comment on lines
+70
to
+75
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This check for conflict markers is redundant and unreachable. The |
||
| python3 -m pip install -e . --no-deps | ||
|
|
||
| # https://github.com/pytorch/pytorch/issues/168167 | ||
| RUN pip install nvidia-cudnn-cu12==9.16.0.29 | ||
| ENV PYTHONPATH="/root/Megatron-LM:${PYTHONPATH}" | ||
|
|
||
| # reinstall numpy 1.x for megatron | ||
| RUN pip install "numpy<2" | ||
| # Install slime itself without dependency resolution so the vLLM base image | ||
| # keeps control of CUDA/PyTorch/Triton/transformers compatibility. | ||
| COPY . /root/slime | ||
| RUN cd /root/slime && \ | ||
| python3 -m pip install -e . --no-deps | ||
|
|
||
| RUN pip install https://github.com/zhuzilin/sgl-router/releases/download/v0.3.2-5f8d397/sglang_router-0.3.2-cp38-abi3-manylinux_2_28_x86_64.whl --force-reinstall | ||
| RUN python -c "import sglang_router; assert 'slime' in sglang_router.__version__" | ||
| WORKDIR /root/slime | ||
|
|
||
| RUN rm -rf /root/.cache/pip /root/flash-attention | ||
| RUN python3 - <<'PY' | ||
| import importlib.metadata as metadata | ||
| import vllm | ||
|
|
||
| # ====================================== Patches ============================================ | ||
| print("vllm import ok", getattr(vllm, "__version__", metadata.version("vllm"))) | ||
| print("sglang version", metadata.version("sglang")) | ||
| print("transformer_engine version", metadata.version("transformer_engine")) | ||
| PY | ||
|
|
||
| COPY docker/patch/${PATCH_VERSION}/megatron.patch /root/Megatron-LM/ | ||
| RUN cd Megatron-LM && \ | ||
| git update-index --refresh && \ | ||
| git apply megatron.patch --3way && \ | ||
| if grep -R -n '^<<<<<<< ' .; then \ | ||
| echo "Patch failed to apply cleanly. Please resolve conflicts." && \ | ||
| exit 1; \ | ||
| fi && \ | ||
| rm megatron.patch | ||
|
|
||
| # TODO temporarily skip patching for GB200/GB300 (and require users to bring their own sglang version). should add back later. | ||
| ARG ENABLE_SGLANG_PATCH=1 | ||
| COPY docker/patch/${PATCH_VERSION}/sglang.patch /sgl-workspace/sglang/ | ||
| RUN if [ "$ENABLE_SGLANG_PATCH" = "1" ]; then \ | ||
| cd /sgl-workspace/sglang && \ | ||
| git update-index --refresh && \ | ||
| git apply sglang.patch --3way && \ | ||
| if grep -R -n '^<<<<<<< ' .; then \ | ||
| echo "Patch failed to apply cleanly. Please resolve conflicts." && \ | ||
| exit 1; \ | ||
| fi && \ | ||
| rm sglang.patch; \ | ||
| fi | ||
|
|
||
| # ====================================== Install main package ============================================ | ||
|
|
||
| ARG SLIME_COMMIT=main | ||
| RUN git clone https://github.com/THUDM/slime.git /root/slime && \ | ||
| cd /root/slime && \ | ||
| git checkout ${SLIME_COMMIT} && \ | ||
| pip install -e . --no-deps | ||
|
|
||
| RUN cd /root/slime/slime/backends/megatron_utils/kernels/int4_qat && \ | ||
| pip install . --no-build-isolation | ||
| CMD ["/bin/bash"] | ||
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We need TE, it's a dependency of mcore