slime-ascend/docker/Dockerfile.rocm-代码预览-slime-ascend:基于昇腾的LLM后训练框架项目

#### Use the base image

# The Docker image built with this Dockerfile:
# Supports at least up to slime commit ID: 2710445 (Oct 9, 2025) - supported by amd_patch/sglv0.5.0rc0 
# Still need to update amd_patch

# 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


# The Docker image built with this Dockerfile:
# PR: commit ID 36711aa (Aug 22, 2025) dockerfile - Supports at least up to slime commit ID: d4a7741 (Sep 7, 2025)


# 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 SGLang###############
###########################################

# This is necessary for scope purpose
# ARG GPU_ARCH=gfx942
ENV GPU_ARCH=gfx942

# ===============================
# Base image 942 and args
# FROM $BASE_IMAGE_942 AS gfx942
ENV BUILD_VLLM="0"
ENV BUILD_TRITON="1"
ENV BUILD_AITER_ALL="1"
ENV AITER_COMMIT="v0.1.4"

# # ===============================
# # Base image 950 and args
# FROM $BASE_IMAGE_950 AS gfx950
# ENV BUILD_VLLM="0"
# ENV BUILD_TRITON="0"
# ENV BUILD_AITER_ALL="1"
# ENV AITER_COMMIT="v0.1.4"

# ===============================
# Chosen arch and args
# FROM ${GPU_ARCH}

# This is necessary for scope purpose, again
# ARG GPU_ARCH=gfx950
ENV GPU_ARCH_LIST=${GPU_ARCH:-${PYTORCH_ROCM_ARCH}}

ARG SGL_REPO="https://github.com/sgl-project/sglang.git"
ARG SGL_DEFAULT="main"
# ARG SGL_BRANCH=${SGL_DEFAULT}
ARG SGL_BRANCH="8ecf6b9d2480c3f600826c7d8fef6a16ed603c3f"

ARG TRITON_REPO="https://github.com/ROCm/triton.git"
ARG TRITON_COMMIT="improve_fa_decode_3.0.0"

ARG AITER_REPO="https://github.com/ROCm/aiter.git"


WORKDIR /workspace
# -----------------------
# AITER
RUN pip uninstall -y aiter
RUN git clone ${AITER_REPO} \
 && cd aiter \
 && git checkout ${AITER_COMMIT} \
 && git submodule update --init --recursive
RUN cd aiter \
     && if [ "$BUILD_AITER_ALL" = "1" ]; then \
          PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \
        else \
          GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \
        fi

# -----------------------
# Triton
RUN if [ "$BUILD_TRITON" = "1" ]; then \
        pip uninstall -y triton \
     && git clone ${TRITON_REPO} \
     && cd triton \
     && git checkout ${TRITON_COMMIT} \
     && cd python \
     && python setup.py install; \
    fi

# -----------------------
# Build vLLM
ARG VLLM_REPO="https://github.com/ROCm/vllm.git"
ARG VLLM_BRANCH="9f6b92db47c3444b7a7d67451ba0c3a2d6af4c2c"
RUN if [ "$BUILD_VLLM" = "1" ]; then \
        git clone ${VLLM_REPO} \
     && cd vllm \
     && git checkout ${VLLM_BRANCH} \
     && python -m pip install -r requirements/rocm.txt \
     && python setup.py clean --all \
     && python setup.py develop; \
    fi

# -----------------------
# Build SGLang
ARG BUILD_TYPE=all

RUN pip install IPython \
    && pip install orjson \
    && pip install python-multipart \
    && pip install torchao \
    && pip install pybind11

RUN pip install "setuptools<70.0.0" --force-reinstall
RUN pip uninstall -y sgl_kernel sglang
RUN git clone ${SGL_REPO} \
    && cd sglang \
    && if [ "${SGL_BRANCH}" = ${SGL_DEFAULT} ]; then \
         echo "Using ${SGL_DEFAULT}, default branch."; \
         git checkout ${SGL_DEFAULT}; \
       else \
         echo "Using ${SGL_BRANCH} branch."; \
         git checkout ${SGL_BRANCH}; \
       fi \
    && cd sgl-kernel \
    && rm -f pyproject.toml \
    && mv pyproject_rocm.toml pyproject.toml \
    && AMDGPU_TARGET=$GPU_ARCH_LIST 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

RUN python -m pip cache purge

# Copy config files to support MI300X in virtualized environments (MI300X_VF).  Symlinks will not be created in image build.
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"' -- {}

# Performance environment variable.
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 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


# sglang patch
# COPY patch/${SGLANG_VERSION}/sglang.patch /sgl-workspace/sglang/
COPY amd_patch/sglv0.5.0rc0 /workspace/patch
RUN cd /workspace/sglang && \
  git apply /workspace/patch/sglang.patch && \
  if grep -R -n '^<<<<<<< ' .; then \
    echo "Patch failed to apply cleanly. Please resolve conflicts." && \
    exit 1; \
  fi 



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

# Fix importlib-metadata version conflict before vllm installation
RUN pip install "importlib-metadata>=6.0,<=8.0.0" --force-reinstall

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

COPY amd_patch/sglv0.5.0rc0 /workspace/patch

RUN pip uninstall -y megatron-core && \
    git clone https://github.com/NVIDIA/Megatron-LM && \
    cd Megatron-LM && \
    git checkout 48406695c4efcf1026a7ed70bb390793918dd97b && \
    git apply /workspace/patch/amd_megatron_fused_kernels_init.patch && \
    pip install -vvv -e . && \
    cd /workspace/

# sandwitch norm for GLM models
RUN cd Megatron-LM && \
    git apply /workspace/patch/megatron.patch --3way && \
    if grep -R -n '^<<<<<<< ' .; then \
      echo "Patch failed to apply cleanly. Please resolve conflicts." && \
      exit 1; \
    fi

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




#########################################
###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
########################################
########################################
########################################


########################################
########Additional packages#############
########################################
RUN pip install tensorboard
########################################
########################################
########################################


WORKDIR /workspace/

CMD ["/usr/bin/bash"]