Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
185 changes: 80 additions & 105 deletions docker/Dockerfile
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 \

@aoshen02 aoshen02 May 18, 2026

Copy link
Copy Markdown
Collaborator

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

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This check for conflict markers is redundant and unreachable. The SHELL instruction on line 12 enables pipefail and errexit (-e), so the build will fail immediately if git apply returns a non-zero exit code. Additionally, git apply --3way does not leave conflict markers in files upon failure; it simply fails the operation.

    git apply megatron.patch --3way && \
    rm megatron.patch && \

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"]
12 changes: 3 additions & 9 deletions docker/README.md
Original file line number Diff line number Diff line change
@@ -1,17 +1,11 @@
# Docker release rule

We will publish 2 kinds of docker images:
1. stable version, which based on official sglang release. We will store the patch on those versions.
2. latest version, which aligns to `lmsysorg/sglang:latest`.
1. stable version, which is based on an official vLLM release.
2. latest version, which aligns to the current vLLM-based Dockerfile.

current stable version is:
- sglang v0.5.9 (bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2), megatron dev 3714d81d418c9f1bca4594fc35f9e8289f652862

history versions:
- sglang v0.5.7 nightly-dev-20260107-dce8b060 (dce8b0606c06d3a191a24c7b8cbe8e238ab316c9), megatron dev 3714d81d418c9f1bca4594fc35f9e8289f652862
- sglang v0.5.6 nightly-dev-20251208-5e2cda61 (5e2cda6158e670e64b926a9985d65826c537ac82), megatron v0.14.0 (23e00ed0963c35382dfe8a5a94fb3cda4d21e133)
- sglang v0.5.5.post1 (303cc957e62384044dfa8e52d7d8af8abe12f0ac), megatron v0.14.0 (23e00ed0963c35382dfe8a5a94fb3cda4d21e133)
- sglang v0.5.0rc0-cu126 (8ecf6b9d2480c3f600826c7d8fef6a16ed603c3f), megatron 48406695c4efcf1026a7ed70bb390793918dd97b
- vLLM v0.21.0, megatron dev 3714d81d418c9f1bca4594fc35f9e8289f652862

The command to build:

Expand Down
Loading