rlsys/slime

By rlsys

•Updated 12 months ago

Support slime on ROCm for AMD GPUs

Image

4.1K

Overview Tags

rlsys/slime repository overview

⁠Support slime on ROCm for AMD GPUs

⁠Stable version:

==================

rlsys/slime:slime_ubuntu22.04_rocm6.3.4-patch-numa-patch_sglang0.4.9.post6_aiter0.1.3_megatron-patch_apex_torch-memory-saver0.0.8-patch-vim

Issue: [Current has a convergence issue and will hang sometimes]
Update the SGLang to v0.4.9.post6; 
Dockerfile (Have not pushed yet); supports up to (Aug 22, 2025, commit ID: c7951ed)

rlsys/slime:slime_ubuntu22.04_rocm6.3.4-patch-numa-patch_sglang0.4.9_megatron-patch_ray2.47.1_apex_torch-memory-saver0.0.8-patch-vim

Issue: [Current has a convergence issue and will hang sometimes]
Update the SGLang to v0.4.9; AITER: v0.1.3, fixed torch_dist ckp loading
Dockerfile (Aug 22, 2025, commit ID: c7951ed); supports up to (Aug 10, 2025, commit ID: 36711aa)

rlsys/slime:slime_ubuntu22.04_rocm6.3.4-patch-numa_vllm0.8.5-patch_sglang0.4.7_megatron-core-patch_ray0.47-patch_apex

Issue: "C"PU memory issue (will take 1.x T)
Update the SGLang: v0.4.7; AITER: v0.1.2
Dockerfile (Jul 30, 2025, commit ID: a5b1bd3); supports up to (Aug 10, 2025, commit ID: 36711aa)

==================

⁠The current latest Docker image was built by:

#### Use the base image

# The Docker image built with this Dockerfile:
# PR: commit ID 36711aa (Aug 22, 2025) dockerfile - Supports up to slime commit ID: 9a48ba0 (Aug 10, 2025)

# Start to failfrom c22f55b (Aug 10, 2025) - Need to fix the bug from here

# You can find the latest pre-built Docker image from here: https://hub.docker.com/r/rlsys/slime/tags
# Current latest docker img: `rlsys/slime:slime_ubuntu22.04_rocm6.3.4-patch-numa-patch_sglang0.4.9_megatron-patch_ray2.47.1_apex_torch-memory-saver0.0.8-patch-vim` manually add the patch to mitigate checkpoint loading issue. (vim /workspace/Megatron-LM-amd_version/megatron/training/checkpointing.py. Line: 1449 ~ 1457 - comment out if becasue of dismatch number of dist checkpoints

# Thanks to Yang Wang (https://www.microsoft.com/en-us/research/people/yangwang5/) for working on the patch for this ROCm base Docker image to support virtual memory management on MI300X.

# FROM "rlfoundation.azurecr.io/rocm6.3.4:vllm-0.8.5-numa-patch-ubuntu-22.04"
FROM "rlsys/rocm-6.3.4-patch:rocm6.3.4-numa-patch_ubuntu-22.04"

SHELL ["/bin/bash", "-ceuxo", "pipefail"]

ARG MAX_JOBS=512
ENV MAX_JOBS=${MAX_JOBS}

ENV PATH="/usr/local/python3.12/bin:$PATH"
RUN ln -sf /usr/bin/python3.12 /usr/bin/python && \
    ln -sf /usr/bin/pip3.12 /usr/bin/pip

RUN apt-get update
RUN apt-get install -y pkg-config liblzma-dev


###########################################
##########Install TransformerEngine########
###########################################
WORKDIR /workspace/

RUN rm -rf TransformerEngine 
# RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git
RUN git clone https://github.com/ROCm/TransformerEngine.git
WORKDIR /workspace/TransformerEngine

RUN git checkout 236178e
# RUN git checkout bb061ad
# RUN git checkout 864405c

RUN git submodule update --init --recursive

ENV NVTE_FRAMEWORK=pytorch 
ENV NVTE_ROCM_ARCH=gfx942 
ENV NVTE_USE_HIPBLASLT=1
ENV NVTE_USE_ROCM=1  

# export CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr:${CMAKE_PREFIX_PATH:-}"
ENV CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr"
RUN MAX_JOBS=${MAX_JOBS} pip install . -vvv 
WORKDIR /workspace/
###########################################
###########################################
###########################################



###########################################################
####Install vllm - sglang require vllm 0.6.7 dependency####
# #########################################################
#### Require vllm 0.6.7 - checkout 113274a0
WORKDIR /workspace/
RUN rm -rf vllm
RUN pip uninstall -y vllm
# Refer to here (down-grade vllm to 0.6.3): https://docs.vllm.ai/en/v0.6.3/getting_started/amd-installation.html
RUN git clone https://github.com/ROCm/vllm.git
# git clone https://github.com/vllm-project/vllm.git
WORKDIR /workspace/vllm
RUN git checkout 113274a0
ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942"
ENV MAX_JOBS=${MAX_JOBS}
RUN pip install "boto3>=1.26.0"
RUN pip install setuptools_scm
# will add src into py. You can delete the repo
RUN python3 setup.py install
WORKDIR /workspace/
###########################################
###########################################


RUN pip install setuptools==75.8.0


###########################################
############build sgalng###################
###########################################
# Set environment variables
ENV BASE_DIR=/workspace
# ENV BASE_DIR=/sgl-workspace
ENV BUILD_TYPE=all
ENV SGL_REPO=https://github.com/sgl-project/sglang
# ENV SGL_BRANCH=v0.4.9
ENV SGL_BRANCH=0.4.9.post6
# ENV SGL_BRANCH=v0.4.10
ENV TRITON_REPO=https://github.com/ROCm/triton.git
ENV TRITON_COMMIT=improve_fa_decode_3.0.0
ENV AITER_REPO=https://github.com/ROCm/aiter.git
ENV AITER_COMMIT=v0.1.3
# ENV AITER_COMMIT=v0.1.4
# v0.1.2 version - commit id: 9d11f47
# ENV AITER_COMMIT=9d11f47

ENV HIP_FORCE_DEV_KERNARG=1
ENV HSA_NO_SCRATCH_RECLAIM=1
ENV SGLANG_SET_CPU_AFFINITY=1
ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1
ENV NCCL_MIN_NCHANNELS=112

ENV SGLANG_USE_AITER=1
ENV SGLANG_MOE_PADDING=1
# ENV MOE_PADDING=1
ENV VLLM_FP8_PADDING=1
ENV VLLM_FP8_ACT_PADDING=1
ENV VLLM_FP8_WEIGHT_PADDING=1
ENV VLLM_FP8_REDUCE_CONV=1
ENV TORCHINDUCTOR_MAX_AUTOTUNE=1
ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1
ENV HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=gfx942"
ENV AMDGPU_TARGETS=gfx942
ENV ROCM_ARCH=gfx942
ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942"

# Switch to working directory
# WORKDIR /sgl-workspace
WORKDIR /workspace

# Clean and create directory
# RUN rm -rf /sgl-workspace && mkdir -p /sgl-workspace
RUN rm -rf /workspace && mkdir -p /workspace

# # Clone and build sglang
# RUN git clone ${SGL_REPO} \
#     && cd sglang \
#     && git checkout ${SGL_BRANCH} || echo "Using default branch" \
#     && cd sgl-kernel \
#     && rm -f pyproject.toml \
#     && mv pyproject_rocm.toml pyproject.toml \
#     && python setup_rocm.py install \
#     && cd .. \
#     && if [ "$BUILD_TYPE" = "srt" ]; then \
#          python -m pip --no-cache-dir install -e "python[srt_hip]"; \
#        else \
#          python -m pip --no-cache-dir install -e "python[all_hip]"; \
#        fi \
#     && cd /sgl-workspace \
#     && cp -r /sgl-workspace/sglang /sglang \
#     && python -m pip cache purge

# Install common Python packages
RUN pip install IPython orjson python-multipart torchao pybind11

# Rebuild Triton
RUN pip uninstall -y triton || true \
    && git clone ${TRITON_REPO} \
    && cd triton \
    && git checkout ${TRITON_COMMIT} \
    && cd python \
    && python3 setup.py install \
    && cd /workspace
    # && cd /sgl-workspace

# Build aiter
#version: Commit 9d11f47
    # && git checkout ${AITER_COMMIT} \
RUN pip uninstall -y aiter || true
RUN git clone ${AITER_REPO} \
    && cd aiter \
    && git checkout ${AITER_COMMIT} \
    && git submodule sync \
    && git submodule update --init --recursive \
    && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \
    && cd /workspace
    # && cd /sgl-workspace
    # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \


###########################################
# Clone and build sglang
RUN git clone ${SGL_REPO} \
    && cd sglang \
    && git checkout ${SGL_BRANCH} || echo "Using default branch" \
    && cd sgl-kernel \
    && rm -f pyproject.toml \
    && mv pyproject_rocm.toml pyproject.toml \
    && python setup_rocm.py install \
    && cd .. \
    && if [ "$BUILD_TYPE" = "srt" ]; then \
         python -m pip --no-cache-dir install -e "python[srt_hip]"; \
       else \
         python -m pip --no-cache-dir install -e "python[all_hip]"; \
       fi \
    && cd /workspace \
    && cp -r /workspace/sglang /sglang \
    && python -m pip cache purge
    # && cd /sgl-workspace \
    # && cp -r /sgl-workspace/sglang /sglang \
    # && python -m pip cache purge
###########################################


# Copy MI300X config 
# RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \
#          /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \
RUN find /workspace/sglang/python/sglang/srt/layers/quantization/configs/ \
         /workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \
         -type f -name '*MI300X*' | \
         xargs -I {} sh -c 'vf_config=$(echo "$1" | sed "s/MI300X/MI300X_VF/"); cp "$1" "$vf_config"' -- {}

# Environment setup complete.
RUN echo "Environment setup complete."
WORKDIR /workspace/


# # patch
# # Date: Jul 5, 2025
# ENV SLIME_COMMIT=9ddbdbd 
# RUN git clone https://github.com/THUDM/slime.git \
#     && cd slime \
#     && git checkout ${SLIME_COMMIT} \
#     && cp docker/patch/sglang.patch  /workspace/sglang/
#     # && cp docker/patch/sglang.patch  /sgl-workspace/sglang/
#     # && cp docker/patch/v0.4.10-cu126/sglang.patch  /sgl-workspace/sglang/
#     # && cp docker/patch/latest/sglang.patch  /sgl-workspace/sglang/
#     # && cp docker/patch/sglang.patch  /sgl-workspace/sglang/
# # COPY /home/yushensu/projects/slime/docker/patch/sglang.patch /sgl-workspace/sglang/
# WORKDIR /workspace/sglang/ 
# # WORKDIR /sgl-workspace/sglang/ 
# RUN git apply sglang.patch && rm sglang.patch
# # WORKDIR /workspace/

# # sgl-router
# # WORKDIR /sgl-workspace/sglang/ 
# RUN apt-get update && apt-get install -y pkg-config libssl-dev
# RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
# RUN pip install setuptools-rust wheel build
# RUN source $HOME/.cargo/env && \
#     mkdir -p sgl-router && \
#     cd sgl-router && \
#     cargo build -j 64 && \
#     python3 -m build && \
#     pip install dist/*.whl --force-reinstall

RUN pip install sglang-router --force-reinstall

###########################################
###########################################
###########################################

RUN pip install transformers==4.51.1


#########################################
#####Install vllm v0.8.5#################
#########################################

WORKDIR /workspace/

ENV VLLM_TARGET_DEVICE=rocm 
ENV ROCM_PATH=/opt/rocm 
ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev

RUN pip uninstall -y vllm || true
RUN rm -rf vllm-patch

RUN git clone https://github.com/RLFoundation/vllm-patch.git \
    && cd vllm-patch \
    && git checkout v0.8.5-sleep-numa \
    && rm -rf build/ dist/ *.egg-info \
    && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \
    && SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH="gfx90a;gfx942" MAX_JOBS=${MAX_JOBS} python3 setup.py install

WORKDIR /workspace/
###########################################
###########################################


#########################################
#### Install megatron-core###############
#########################################
# Can be removed just the current megatron-lm dependency
RUN pip install "numpy>=1.21.0,<2.0" --force-reinstall

RUN pip uninstall -y megatron-core && \
    git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \
    cd Megatron-LM-amd_version && \
    pip install -vvv -e . && \
    cd /workspace/
#########################################
#########################################
#########################################




#########################################
###Add torch_memory_saver################
#########################################
# Set environment variables
ENV HIPCC_COMPILE_FLAGS_APPEND="--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__"
ENV CFLAGS="-D__HIP_PLATFORM_AMD__"
ENV CXXFLAGS="-D__HIP_PLATFORM_AMD__"
# Install torch_memory_saver
# RUN pip install git+https://github.com/YangWang92/torch_memory_saver_numa.git --no-deps
# RUN pip install "git+https://github.com/YangWang92/torch_memory_saver_numa.git@numa"
RUN pip install "git+https://github.com/yushengsu-thu/torch_memory_saver.git"
# pip install git+https://github.com/fzyzcjy/torch_memory_saver.git --no-deps
#########################################
#########################################




########################################
######Install ray#######################
########################################
# need to add this patch manually: https://github.com/ray-project/ray/pull/53531/files
RUN pip uninstall ray -y
# RUN pip install "ray[data,train,tune,serve]>=2.47.0" 
RUN pip install "ray[data,train,tune,serve]==2.47.1" 
########################################
########################################
########################################


### Need to verify whether numerical/convergence issue
#######################################
################apex###################
#######################################
WORKDIR /workspace/
RUN pip uninstall -y apex && \
    git clone https://github.com/ROCm/apex.git && \
    cd apex && \
    python setup.py install && \
    cd /workspace/ 
#######################################
#######################################
#######################################


########################################
############ mbridge####################
########################################
RUN pip install git+https://github.com/ISEEKYAN/mbridge.git --no-deps
########################################
########################################
########################################



########################################
########slime agent framewrok need######
########################################
RUN pip install pydra_config==0.0.15
RUN pip install together
RUN pip install google-generativeai
########################################
########################################
########################################


WORKDIR /workspace/

CMD ["/usr/bin/bash"]

==================

Tag summary

Recent tags

Content type

Image

Digest

sha256:70a54aecd…

Size

30.8 GB

Last updated

12 months ago

docker pull rlsys/slime