diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 36ff4770c..0d98dd97b 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -65,7 +65,7 @@ body: - slime version: - Python version: - PyTorch version: - - CUDA/ROCm version: + - CUDA version: - GPU type and count: - OS: - SGLang version (if relevant): diff --git a/.github/ISSUE_TEMPLATE/question.yml b/.github/ISSUE_TEMPLATE/question.yml index 7c9c4dac2..a6bd52f97 100644 --- a/.github/ISSUE_TEMPLATE/question.yml +++ b/.github/ISSUE_TEMPLATE/question.yml @@ -49,7 +49,7 @@ body: - slime version: - Python version: - PyTorch version: - - CUDA/ROCm version: + - CUDA version: - GPU type and count: - OS: diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm deleted file mode 100644 index adf84c240..000000000 --- a/docker/Dockerfile.rocm +++ /dev/null @@ -1,363 +0,0 @@ -#### 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"] diff --git a/docker/Dockerfile.rocm_MI350-5 b/docker/Dockerfile.rocm_MI350-5 deleted file mode 100644 index caf0faf31..000000000 --- a/docker/Dockerfile.rocm_MI350-5 +++ /dev/null @@ -1,252 +0,0 @@ -#### Use the base image for ROCm 7 / gfx950 (MI355) - -# The Docker image built with this Dockerfile: -# Base image: ROCm 7 with vllm pre-built for gfx950 -# Target GPU: MI355 (gfx950) - - -FROM rocm/sgl-dev:rocm7-vllm-20250904 - -SHELL ["/bin/bash", "-ceuxo", "pipefail"] - -ARG MAX_JOBS=128 -ENV MAX_JOBS=${MAX_JOBS} - -# Set environment variables for gfx950 -ENV GPU_ARCH=gfx950 -ENV PYTORCH_ROCM_ARCH=gfx950 -ENV GPU_ARCH_LIST=gfx950 -ENV AMDGPU_TARGET=gfx950 - - -########################################### -##############1. Install AITER############# -########################################### -WORKDIR /app - -RUN pip uninstall -y aiter || true -RUN rm -rf aiter -RUN git clone https://github.com/ROCm/aiter.git \ - && cd aiter \ - && git checkout v0.1.7.post2 \ - && git submodule update --init --recursive \ - && GPU_ARCHS=gfx950 python setup.py develop -########################################### -########################################### -########################################### - - -########################################### -####2. Install TransformerEngine for gfx950 -########################################### -WORKDIR /app - -RUN rm -rf TransformerEngine -RUN git clone https://github.com/ROCm/TransformerEngine.git \ - && cd TransformerEngine \ - && git checkout 90c04bcdc3c109505b318f40a39680263af55edf \ - && git submodule update --init --recursive - -ENV NVTE_FRAMEWORK=pytorch -ENV NVTE_ROCM_ARCH=gfx950 -ENV NVTE_USE_HIPBLASLT=1 -ENV NVTE_USE_ROCM=1 -ENV CMAKE_PREFIX_PATH="/opt/rocm:/opt/rocm/hip:/usr/local:/usr" - -RUN cd TransformerEngine && pip install . -v -########################################### -########################################### -########################################### - - -######################################### -####3. Install Megatron-LM (NVIDIA version) -######################################### -WORKDIR /app - -RUN pip install "numpy>=1.21.0,<2.0" --force-reinstall - -RUN pip uninstall -y megatron-core || true -RUN rm -rf Megatron-LM -RUN git clone https://github.com/NVIDIA/Megatron-LM \ - && cd Megatron-LM \ - && git checkout 48406695c4efcf1026a7ed70bb390793918dd97b \ - && pip install -e . -######################################### -######################################### -######################################### - - -######################################## -############ 4. Install mbridge######### -######################################## -RUN pip install git+https://github.com/ISEEKYAN/mbridge.git --no-deps -######################################## -######################################## -######################################## - - -######################################## -######5. Install Ray#################### -######################################## -RUN pip uninstall ray -y || true -RUN pip install "ray[data,train,tune,serve]==2.47.1" -######################################## -######################################## -######################################## - - -######################################### -###6. Install torch_memory_saver######### -######################################### -RUN pip install torch_memory_saver -######################################### -######################################### - - -####################################### -####7. Install Apex for ROCm########### -####################################### -WORKDIR /app - -RUN pip uninstall -y apex || true -RUN rm -rf apex -RUN git clone https://github.com/ROCm/apex.git \ - && cd apex \ - && python setup.py install -####################################### -####################################### -####################################### - - -######################################## -###8. Install slime agent framework deps -######################################## -RUN pip install pydra_config==0.0.15 -RUN pip install together -RUN pip install google-generativeai -RUN pip install tensorboard -######################################## -######################################## -######################################## - - -######################################## -###9. Set performance environment vars## -######################################## -ENV HIP_FORCE_DEV_KERNARG=1 -ENV HSA_NO_SCRATCH_RECLAIM=1 -ENV SGLANG_USE_AITER=1 -ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 -ENV SGLANG_MOE_PADDING=1 -ENV SGLANG_SET_CPU_AFFINITY=1 -ENV SGLANG_ROCM_FUSED_DECODE_MLA=1 -ENV SGLANG_USE_ROCM700A=1 -ENV NCCL_MIN_NCHANNELS=112 -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 -######################################## -######################################## -######################################## - - -########################################### -##############Install SGLang############### -########################################### -WORKDIR /app - -# Install prerequisites -RUN pip install IPython orjson python-multipart torchao==0.9.0 pybind11 - -# Clone SGLang -RUN pip uninstall -y sgl_kernel sglang || true -RUN rm -rf sglang -RUN git clone https://github.com/sgl-project/sglang.git \ - && cd sglang \ - && git checkout v0.5.6 - -# Build sgl-kernel for gfx950 -RUN cd sglang/sgl-kernel \ - && rm -f pyproject.toml \ - && mv pyproject_rocm.toml pyproject.toml \ - && AMDGPU_TARGET=gfx950 python setup_rocm.py install - -# Install SGLang -RUN cd sglang \ - && rm -rf python/pyproject.toml \ - && mv python/pyproject_other.toml python/pyproject.toml \ - && pip install -e "python[all_hip]" - -# Test SGLang installation -RUN python -c "import sglang; import sgl_kernel; print('SGLang + sgl_kernel: OK')" - -RUN python -m pip cache purge -########################################### -########################################### -########################################### - - -########################################### -#### APPLY PATCHES (gfx950/MI355) ######### -########################################### - -# Copy patches from slime repo -COPY amd_patch/latest /app/patch - -# Apply Megatron patches -RUN cd /app/Megatron-LM \ - && git apply /app/patch/amd_megatron_fused_kernels_init.patch \ - && git apply /app/patch/megatron.patch --3way \ - && if grep -R -n '^<<<<<<< ' .; then \ - echo "Patch failed to apply cleanly. Please resolve conflicts." && \ - exit 1; \ - fi \ - && pip install -e . -v - -# Apply SGLang patch -RUN cd /app/sglang \ - && git apply /app/patch/sglang.patch || echo "Check patch compatibility with v0.5.6" \ - && if grep -R -n '^<<<<<<< ' .; then \ - echo "Patch failed to apply cleanly. Please resolve conflicts." && \ - exit 1; \ - fi - -# Copy MOE configs for gfx950/MI355 -RUN find /app/sglang/python/sglang/srt/layers/quantization/configs/ \ - /app/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \ - -type f -name '*MI300X*' 2>/dev/null | while read f; do \ - cp "$f" "$(echo $f | sed 's/MI300X/MI300X_VF/')" 2>/dev/null || true; \ - cp "$f" "$(echo $f | sed 's/MI300X/MI355/')" 2>/dev/null || true; \ -done - -########################################### -########################################### -########################################### - - -######################################## -#### Install additional packages######## -######################################## -RUN pip install sglang-router --force-reinstall -######################################## -######################################## -######################################## - - -######################################## -# Fix click/ray incompatibility with Python 3.10 -######################################## -RUN pip install click==8.2.1 -######################################## -######################################## -######################################## - - -WORKDIR /app - -CMD ["/usr/bin/bash"] - diff --git a/docker/Dockerfile_20250810_9a48ba0.rocm b/docker/Dockerfile_20250810_9a48ba0.rocm deleted file mode 100644 index a42f505a5..000000000 --- a/docker/Dockerfile_20250810_9a48ba0.rocm +++ /dev/null @@ -1,361 +0,0 @@ -#### 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) - -# 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 sglang################### -########################################### -# 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.7 -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.2 -# 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"] diff --git a/docker/Dockerfile_20250810_c22f55b.rocm b/docker/Dockerfile_20250810_c22f55b.rocm deleted file mode 100644 index af651a028..000000000 --- a/docker/Dockerfile_20250810_c22f55b.rocm +++ /dev/null @@ -1,374 +0,0 @@ -#### Use the base image - -# The Docker image built with this Dockerfile: -# Supports up to slime commit ID: 9a48ba0 (Aug 10, 2025) -# Start to fail from 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 - - -# 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 sglang################### -########################################### -# 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"] diff --git a/docker/amd_patch/latest/amd_megatron_fused_kernels_init.patch b/docker/amd_patch/latest/amd_megatron_fused_kernels_init.patch deleted file mode 100644 index f6efca346..000000000 --- a/docker/amd_patch/latest/amd_megatron_fused_kernels_init.patch +++ /dev/null @@ -1,51 +0,0 @@ -diff --git a/megatron/legacy/fused_kernels/__init__.py b/megatron/legacy/fused_kernels/__init__.py -index 87cceac3..ac686d74 100644 ---- a/megatron/legacy/fused_kernels/__init__.py -+++ b/megatron/legacy/fused_kernels/__init__.py -@@ -3,6 +3,7 @@ - import os - import pathlib - import subprocess -+import torch - - from torch.utils import cpp_extension - -@@ -15,23 +16,23 @@ os.environ["TORCH_CUDA_ARCH_LIST"] = "" - - - def load(args): -- -- # Check if cuda 11 is installed for compute capability 8.0 -- cc_flag = [] -- _, bare_metal_major, bare_metal_minor = _get_cuda_bare_metal_version( -- cpp_extension.CUDA_HOME -- ) -- if int(bare_metal_major) >= 11: -- cc_flag.append('-gencode') -- cc_flag.append('arch=compute_80,code=sm_80') -- if int(bare_metal_minor) >= 8: -+ if torch.cuda.is_available() and torch.version.cuda: -+ # Check if cuda 11 is installed for compute capability 8.0 -+ cc_flag = [] -+ _, bare_metal_major, bare_metal_minor = _get_cuda_bare_metal_version( -+ cpp_extension.CUDA_HOME -+ ) -+ if int(bare_metal_major) >= 11: - cc_flag.append('-gencode') -- cc_flag.append('arch=compute_90,code=sm_90') -+ cc_flag.append('arch=compute_80,code=sm_80') -+ if int(bare_metal_minor) >= 8: -+ cc_flag.append('-gencode') -+ cc_flag.append('arch=compute_90,code=sm_90') - -- # Build path -- srcpath = pathlib.Path(__file__).parent.absolute() -- buildpath = srcpath / "build" -- _create_build_dir(buildpath) -+ # Build path -+ srcpath = pathlib.Path(__file__).parent.absolute() -+ buildpath = srcpath / "build" -+ _create_build_dir(buildpath) - - # Helper function to build the kernels. - def _cpp_extention_load_helper(name, sources, extra_cuda_flags): diff --git a/docker/amd_patch/latest/megatron.patch b/docker/amd_patch/latest/megatron.patch deleted file mode 100644 index c840133ce..000000000 --- a/docker/amd_patch/latest/megatron.patch +++ /dev/null @@ -1,438 +0,0 @@ -diff --git a/megatron/core/distributed/__init__.py b/megatron/core/distributed/__init__.py -index fe26e8b4..4451f277 100644 ---- a/megatron/core/distributed/__init__.py -+++ b/megatron/core/distributed/__init__.py -@@ -11,3 +11,15 @@ from .finalize_model_grads import finalize_model_grads - from .fsdp.mcore_fsdp_adapter import FullyShardedDataParallel - from .torch_fully_sharded_data_parallel import TorchFullyShardedDataParallel - from .torch_fully_sharded_data_parallel_config import TorchFullyShardedDataParallelConfig -+ -+# Backward compatibility patch for FSDP module reorganization -+import sys -+import importlib.util -+ -+spec = importlib.util.find_spec('megatron.core.distributed.fsdp.src.megatron_fsdp') -+if spec: -+ custom_fsdp = importlib.util.module_from_spec(spec) -+ spec.loader.exec_module(custom_fsdp) -+ sys.modules['megatron.core.distributed.custom_fsdp'] = custom_fsdp -+ if hasattr(custom_fsdp, 'MegatronFSDP'): -+ custom_fsdp.FullyShardedDataParallel = custom_fsdp.MegatronFSDP -diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py -index 99c3edc0..26ea5cb4 100644 ---- a/megatron/core/extensions/transformer_engine.py -+++ b/megatron/core/extensions/transformer_engine.py -@@ -404,6 +404,7 @@ class TELinear(te.pytorch.Linear): - ) - - for param in self.parameters(): -+ setattr(param, "parallel_mode", parallel_mode) - if is_expert: - # Reduce the gradient on the expert_data_parallel group for expert linear layers - setattr(param, "allreduce", not self.expert_parallel) -diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py -index 002edb92..f7273488 100755 ---- a/megatron/core/models/gpt/gpt_layer_specs.py -+++ b/megatron/core/models/gpt/gpt_layer_specs.py -@@ -80,6 +80,8 @@ def get_gpt_layer_with_transformer_engine_spec( - use_te_op_fuser: Optional[bool] = False, - use_kitchen: bool = False, - use_te_activation_func: bool = False, -+ post_self_attn_layernorm: bool = False, -+ post_mlp_layernorm: bool = False, - ) -> ModuleSpec: - """Use this spec to use lower-level Transformer Engine modules (required for fp8 training). - -@@ -182,9 +184,11 @@ def get_gpt_layer_with_transformer_engine_spec( - ), - ), - self_attn_bda=get_bias_dropout_add, -+ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, - pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp, - mlp=mlp, - mlp_bda=get_bias_dropout_add, -+ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, - sharded_state_dict_keys_map={ - "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight", - "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias", -diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py -index df9adc3e..2f4f544a 100644 ---- a/megatron/core/models/gpt/gpt_model.py -+++ b/megatron/core/models/gpt/gpt_model.py -@@ -443,7 +443,7 @@ class GPTModel(LanguageModule): - if self.share_embeddings_and_output_weights: - output_weight = self.shared_embedding_or_output_weight() - -- if mtp_in_postprocess: -+ if mtp_in_postprocess and labels is not None: - hidden_states = self.mtp( - input_ids=input_ids, - position_ids=position_ids, -diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py -index 57332ac3..f3abd642 100644 ---- a/megatron/core/parallel_state.py -+++ b/megatron/core/parallel_state.py -@@ -9,6 +9,7 @@ from typing import Callable, List, Optional - - import numpy as np - import torch -+import torch.distributed as dist - - from .utils import GlobalMemoryBuffer, is_torch_min_version - -@@ -163,6 +164,213 @@ def get_nccl_options(pg_name, nccl_comm_cfgs): - return None - - -+old_new_group = None -+ -+ -+def monkey_patch_torch_dist(): -+ print("Applying monkey patch to torch.distributed", flush=True) -+ global old_new_group -+ if old_new_group is not None: -+ return -+ -+ old_new_group = dist.new_group -+ -+ def new_group(*args, **kwargs): -+ group = old_new_group(*args, **kwargs) -+ # skip none nccl group. -+ if ( -+ len(args) >= 3 and args[2] == "gloo" or -+ "backend" in kwargs and kwargs["backend"] == "gloo" -+ ): -+ return group -+ -+ # Get ranks from arguments -+ if len(args) >= 1 and args[0] is not None: -+ ranks = args[0] -+ elif "ranks" in kwargs and kwargs["ranks"] is not None: -+ ranks = kwargs["ranks"] -+ else: -+ # If no ranks specified, use all ranks in world -+ ranks = list(range(dist.get_world_size())) -+ -+ if len(ranks) == 1: -+ return group -+ -+ group = ReloadableProcessGroup(group, ranks) -+ return group -+ -+ dist.new_group = new_group -+ -+ def get_new_function(func): -+ def new_function(*args, **kwargs): -+ args = ( -+ arg.group if isinstance(arg, ReloadableProcessGroup) else arg -+ for arg in args -+ ) -+ kwargs = { -+ k: (v.group if isinstance(v, ReloadableProcessGroup) else v) -+ for k, v in kwargs.items() -+ } -+ return func(*args, **kwargs) -+ return new_function -+ -+ dist.get_rank = get_new_function(dist.get_rank) -+ dist.get_world_size = get_new_function(dist.get_world_size) -+ dist.get_backend = get_new_function(dist.get_backend) -+ dist.get_global_rank = get_new_function(dist.get_global_rank) -+ dist.get_group_rank = get_new_function(dist.get_group_rank) -+ dist.get_process_group_ranks = get_new_function(dist.get_process_group_ranks) -+ -+ dist.all_reduce = get_new_function(dist.all_reduce) -+ dist.all_gather = get_new_function(dist.all_gather) -+ dist.all_gather_into_tensor = get_new_function(dist.all_gather_into_tensor) -+ dist.all_gather_object = get_new_function(dist.all_gather_object) -+ dist.all_to_all = get_new_function(dist.all_to_all) -+ dist.all_to_all_single = get_new_function(dist.all_to_all_single) -+ dist.broadcast = get_new_function(dist.broadcast) -+ dist.reduce = get_new_function(dist.reduce) -+ dist.reduce_scatter = get_new_function(dist.reduce_scatter) -+ dist.reduce_scatter_tensor = get_new_function(dist.reduce_scatter_tensor) -+ dist.scatter = get_new_function(dist.scatter) -+ dist.gather = get_new_function(dist.gather) -+ dist.barrier = get_new_function(dist.barrier) -+ dist.send = get_new_function(dist.send) -+ dist.recv = get_new_function(dist.recv) -+ dist._coalescing_manager = get_new_function(dist._coalescing_manager) -+ -+ # p2p -+ old_isend = dist.isend -+ old_irecv = dist.irecv -+ -+ dist.isend = get_new_function(dist.isend) -+ dist.irecv = get_new_function(dist.irecv) -+ -+ def get_new_p2pop_function(func): -+ def new_function(*args, **kwargs): -+ def convert(arg): -+ if isinstance(arg, ReloadableProcessGroup): -+ return arg.group -+ elif arg == dist.isend: -+ arg = old_isend -+ elif arg == dist.irecv: -+ arg = old_irecv -+ return arg -+ -+ args = (convert(arg) for arg in args) -+ kwargs = { -+ k: convert(v) -+ for k, v in kwargs.items() -+ } -+ return func(*args, **kwargs) -+ return new_function -+ -+ dist.P2POp.__new__ = get_new_p2pop_function(dist.P2POp.__new__) -+ dist.P2POp.__init__ = get_new_p2pop_function(dist.P2POp.__init__) -+ -+ -+ -+class ReloadableProcessGroup(torch.distributed.ProcessGroup): -+ GROUPS = [] -+ -+ def __init__(self, group, ranks): -+ super().__init__( -+ rank=dist.get_rank(group), -+ size=dist.get_world_size(group), -+ ) -+ #print(f"Creating ReloadableProcessGroup with ranks: {ranks}", flush=True) -+ self.group = group -+ self.group_info = { -+ "ranks": ranks, -+ } -+ ReloadableProcessGroup.GROUPS.append(self) -+ -+ def __getattr__(self, name): -+ return getattr(self.group, name) -+ -+ @staticmethod -+ def destroy_process_groups(): -+ for reloadable_group in ReloadableProcessGroup.GROUPS: -+ if reloadable_group.group is None: -+ continue -+ #print(f"Destroying process group: {reloadable_group.group_info['ranks']}") -+ dist.destroy_process_group(reloadable_group.group) -+ del reloadable_group.group -+ reloadable_group.group = None -+ -+ @staticmethod -+ def reload_process_groups(): -+ for reloadable_group in ReloadableProcessGroup.GROUPS: -+ if reloadable_group.group is not None: -+ continue -+ #print(f"Reloading process group: {reloadable_group.group_info['ranks']}") -+ group = old_new_group( -+ ranks=reloadable_group.group_info["ranks"], -+ backend="nccl" -+ ) -+ reloadable_group.group = group -+ -+ def rank(self) -> int: return self.group.rank() -+ def size(self) -> int: return self.group.size() -+ def name(self) -> str: return self.group.name() -+ -+ def shutdown(self) -> None: -+ if self.group is not None: -+ self.group.shutdown() -+ -+ def abort(self) -> None: -+ if self.group is not None: -+ self.group.abort() -+ -+ def _fwd(self, method, *args, **kwargs): -+ inner = self.group -+ if inner is None: -+ raise RuntimeError("ReloadableProcessGroup: inner PG is None, call reload() first.") -+ return getattr(inner, method)(*args, **kwargs) -+ -+ def barrier(self, *a, **kw): return self._fwd("barrier", *a, **kw) -+ def broadcast(self, *a, **kw): return self._fwd("broadcast", *a, **kw) -+ def allreduce(self, *a, **kw): return self._fwd("allreduce", *a, **kw) -+ def allreduce_coalesced(self, *a, **kw): return self._fwd("allreduce_coalesced", *a, **kw) -+ def reduce(self, *a, **kw): return self._fwd("reduce", *a, **kw) -+ def allgather(self, *a, **kw): return self._fwd("allgather", *a, **kw) -+ def _allgather_base(self, *a, **kw): return self._fwd("_allgather_base", *a, **kw) -+ def allgather_coalesced(self, *a, **kw): return self._fwd("allgather_coalesced", *a, **kw) -+ def allgather_into_tensor_coalesced(self, *a, **kw): return self._fwd("allgather_into_tensor_coalesced", *a, **kw) -+ def gather(self, *a, **kw): return self._fwd("gather", *a, **kw) -+ def scatter(self, *a, **kw): return self._fwd("scatter", *a, **kw) -+ def reduce_scatter(self, *a, **kw): return self._fwd("reduce_scatter", *a, **kw) -+ def _reduce_scatter_base(self, *a, **kw): return self._fwd("_reduce_scatter_base", *a, **kw) -+ def reduce_scatter_tensor_coalesced(self, *a, **kw): return self._fwd("reduce_scatter_tensor_coalesced", *a, **kw) -+ def alltoall_base(self, *a, **kw): return self._fwd("alltoall_base", *a, **kw) -+ def alltoall(self, *a, **kw): return self._fwd("alltoall", *a, **kw) -+ def send(self, *a, **kw): return self._fwd("send", *a, **kw) -+ def recv(self, *a, **kw): return self._fwd("recv", *a, **kw) -+ def recv_anysource(self, *a, **kw): return self._fwd("recv_anysource", *a, **kw) -+ -+ def _start_coalescing(self, *a, **kw): return self._fwd("_start_coalescing", *a, **kw) -+ def _end_coalescing(self, *a, **kw): return self._fwd("_end_coalescing", *a, **kw) -+ def _get_backend_name(self): return self._fwd("_get_backend_name") -+ def _get_backend(self, *a, **kw): return self._fwd("_get_backend", *a, **kw) -+ def _set_default_backend(self, *a, **kw): return self._fwd("_set_default_backend", *a, **kw) -+ @property -+ def bound_device_id(self): return self.group.bound_device_id -+ @bound_device_id.setter -+ def bound_device_id(self, dev): self.group.bound_device_id = dev -+ -+ -+def destroy_process_groups(): -+ """Destroy all reloadable process groups.""" -+ ReloadableProcessGroup.destroy_process_groups() -+ -+ -+def reload_process_groups(): -+ """Reload all reloadable process groups.""" -+ ReloadableProcessGroup.reload_process_groups() -+ -+ -+monkey_patch_torch_dist() -+ -+ - def create_group( - ranks=None, - timeout=None, -diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py -index 63ee9d1f..b90b744c 100644 ---- a/megatron/core/pipeline_parallel/p2p_communication.py -+++ b/megatron/core/pipeline_parallel/p2p_communication.py -@@ -26,22 +26,22 @@ def _batched_p2p_ops( - ops = [] - if tensor_send_prev is not None: - send_prev_op = torch.distributed.P2POp( -- torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, group -+ torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, - ) - ops.append(send_prev_op) - if tensor_recv_prev is not None: - recv_prev_op = torch.distributed.P2POp( -- torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, group -+ torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, - ) - ops.append(recv_prev_op) - if tensor_send_next is not None: - send_next_op = torch.distributed.P2POp( -- torch.distributed.isend, tensor_send_next, next_pipeline_rank, group -+ torch.distributed.isend, tensor_send_next, next_pipeline_rank, - ) - ops.append(send_next_op) - if tensor_recv_next is not None: - recv_next_op = torch.distributed.P2POp( -- torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, group -+ torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, - ) - ops.append(recv_next_op) - if len(ops) > 0: -diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py -index 6f557e1f..b295fd35 100644 ---- a/megatron/core/transformer/transformer_config.py -+++ b/megatron/core/transformer/transformer_config.py -@@ -173,6 +173,9 @@ class TransformerConfig(ModelParallelConfig): - qk_layernorm: bool = False - """Whether to apply `normalization` type of normalization to the query and key embeddings.""" - -+ post_self_attn_layernorm: bool = False -+ post_mlp_layernorm: bool = False -+ - test_mode: bool = False - """Whether to run real-time tests.""" - -diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py -index 84f22bde..b4807d26 100644 ---- a/megatron/core/transformer/transformer_layer.py -+++ b/megatron/core/transformer/transformer_layer.py -@@ -224,6 +224,7 @@ class TransformerLayerSubmodules: - input_layernorm: Union[ModuleSpec, type] = IdentityOp - self_attention: Union[ModuleSpec, type] = IdentityOp - self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp -+ post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp - - pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp - cross_attention: Union[ModuleSpec, type] = IdentityOp -@@ -232,6 +233,7 @@ class TransformerLayerSubmodules: - pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp - mlp: Union[ModuleSpec, type] = IdentityOp - mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp -+ post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp - - # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method - sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict) -@@ -336,6 +338,14 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - # [Module 3: BiasDropoutFusion] - self.self_attn_bda = build_module(submodules.self_attn_bda) - -+ self.post_self_attn_layernorm = build_module( -+ submodules.post_self_attn_layernorm, -+ config=self.config, -+ hidden_size=self.config.hidden_size, -+ eps=self.config.layernorm_epsilon, -+ ) -+ -+ - # [Module 4: Post SelfAttention] Optional Layernorm after self-attn - self.pre_cross_attn_layernorm = build_module( - submodules.pre_cross_attn_layernorm, -@@ -399,6 +409,13 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - # [Module 9: BiasDropoutFusion] - self.mlp_bda = build_module(submodules.mlp_bda) - -+ self.post_mlp_layernorm = build_module( -+ submodules.post_mlp_layernorm, -+ config=self.config, -+ hidden_size=self.config.hidden_size, -+ eps=self.config.layernorm_epsilon -+ ) -+ - self.recompute_input_layernorm = False - self.recompute_pre_mlp_layernorm = False - self.recompute_mlp = False -@@ -535,6 +552,11 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - attention_output_with_bias[0] - ) - -+ attention_output, attention_output_bias = attention_output_with_bias -+ attention_output = self.post_self_attn_layernorm(attention_output) -+ attention_output_with_bias = (attention_output, attention_output_bias) -+ -+ - # TODO: could we move `bias_dropout_add_exec_handler` itself - # inside the module provided in the `bias_dropout_add_spec` module? - nvtx_range_push(suffix="self_attn_bda") -@@ -635,6 +657,10 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - else: - mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output) - -+ mlp_output, mlp_output_bias = mlp_output_with_bias -+ mlp_output = self.post_mlp_layernorm(mlp_output) -+ mlp_output_with_bias = (mlp_output, mlp_output_bias) -+ - if self.recompute_pre_mlp_layernorm: - # discard the output of the pre-mlp layernorm and register the recompute - # as a gradient hook of mlp_output_with_bias[0] -diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py -index 24ba8926..4f039fd4 100644 ---- a/megatron/training/arguments.py -+++ b/megatron/training/arguments.py -@@ -1191,6 +1191,9 @@ def core_transformer_config_from_args(args, config_class=None): - if args.is_hybrid_model: - kw_args['is_hybrid_model'] = args.is_hybrid_model - -+ kw_args['post_self_attn_layernorm'] = args.post_self_attn_layernorm -+ kw_args['post_mlp_layernorm'] = args.post_mlp_layernorm -+ - # handle quantization config - # NOTE: Kitchen arguments are only added to the namespace when - # Kitchen library is available. -@@ -1481,6 +1484,10 @@ def _add_network_size_args(parser): - action='store_true', - help='If set, use original BERT residula connection ' - 'ordering.') -+ group.add_argument('--post-self-attn-layernorm', action='store_true', -+ help='If set, use post self attention layernorm.') -+ group.add_argument('--post-mlp-layernorm', action='store_true', -+ help='If set, use post MLP layernorm.') - group.add_argument('--openai-gelu', action='store_true', - help='Use OpenAIs GeLU implementation. This option' - 'should not be used unless for backward compatibility' diff --git a/docker/amd_patch/latest/sglang.patch b/docker/amd_patch/latest/sglang.patch deleted file mode 100644 index 990c2e628..000000000 --- a/docker/amd_patch/latest/sglang.patch +++ /dev/null @@ -1,203 +0,0 @@ -diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py -index bdb124e51..3edf30ab1 100644 ---- a/python/sglang/srt/configs/model_config.py -+++ b/python/sglang/srt/configs/model_config.py -@@ -454,14 +454,14 @@ class ModelConfig: - ).lower() - - # Detect which checkpoint is it -- for _, method in QUANTIZATION_METHODS.items(): -- quantization_override = method.override_quantization_method( -- quant_cfg, self.quantization -- ) -- if quantization_override: -- quant_method = quantization_override -- self.quantization = quantization_override -- break -+ # for _, method in QUANTIZATION_METHODS.items(): -+ # quantization_override = method.override_quantization_method( -+ # quant_cfg, self.quantization -+ # ) -+ # if quantization_override: -+ # quant_method = quantization_override -+ # self.quantization = quantization_override -+ # break - - # Verify quantization configurations. - if self.quantization is None: -diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py -index 2dd2c75f1..f2adb18f8 100644 ---- a/python/sglang/srt/entrypoints/http_server.py -+++ b/python/sglang/srt/entrypoints/http_server.py -@@ -264,6 +264,10 @@ async def validate_json_request(raw_request: Request): - - - @app.get("/health") -+async def health(request: Request) -> Response: -+ return Response(status_code=200) -+ -+ - @app.get("/health_generate") - async def health_generate(request: Request) -> Response: - """ -diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -index 372717bf9..40665cc90 100644 ---- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -+++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -@@ -190,6 +190,7 @@ class DeepEPBuffer: - f"Consider using --deepep-config to change the behavior." - ) - -+ num_qps_per_rank = 20 - cls._buffer = Buffer( - group, - num_nvl_bytes, -diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py -index 956264fc9..69f729336 100644 ---- a/python/sglang/srt/layers/quantization/fp8.py -+++ b/python/sglang/srt/layers/quantization/fp8.py -@@ -351,10 +351,10 @@ class Fp8LinearMethod(LinearMethodBase): - return - else: - weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data -- layer.weight = torch.nn.Parameter(weight, requires_grad=False) -- layer.weight_scale_inv = torch.nn.Parameter( -- weight_scale, requires_grad=False -- ) -+ # layer.weight = torch.nn.Parameter(weight, requires_grad=False) -+ # layer.weight_scale_inv = torch.nn.Parameter( -+ # weight_scale, requires_grad=False -+ # ) - return - - layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) -diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 95a529c89..758fbfd5f 100644 ---- a/python/sglang/srt/managers/scheduler.py -+++ b/python/sglang/srt/managers/scheduler.py -@@ -1359,7 +1359,7 @@ class Scheduler( - - if memory_leak: - msg = "token_to_kv_pool_allocator memory leak detected! " f"{token_msg}" -- raise ValueError(msg) -+ # raise ValueError(msg) - - if self.disaggregation_mode == DisaggregationMode.DECODE: - req_total_size = ( -@@ -1374,7 +1374,7 @@ class Scheduler( - f"available_size={len(self.req_to_token_pool.free_slots)}, " - f"total_size={self.req_to_token_pool.size}\n" - ) -- raise ValueError(msg) -+ # raise ValueError(msg) - - if ( - self.enable_metrics -@@ -1830,6 +1830,7 @@ class Scheduler( - deepep_mode=DeepEPMode(self.server_args.deepep_mode), - require_mlp_tp_gather=require_mlp_tp_gather(self.server_args), - disable_overlap_schedule=self.server_args.disable_overlap_schedule, -+ offload_tags=self.offload_tags, - ) - - def handle_dp_balance_data(self, local_batch: ScheduleBatch): -@@ -1927,6 +1928,7 @@ class Scheduler( - deepep_mode: DeepEPMode, - require_mlp_tp_gather: bool, - disable_overlap_schedule: bool, -+ offload_tags: set[str], - ): - # Check if other DP workers have running batches - if local_batch is None: -@@ -1957,7 +1959,7 @@ class Scheduler( - ) - - tbo_preparer = TboDPAttentionPreparer() -- if disable_overlap_schedule: -+ if len(offload_tags) == 0 and disable_overlap_schedule: - group = tp_group.device_group - device = tp_group.device - else: -diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 58220b1d6..3c3d081a8 100644 ---- a/python/sglang/srt/managers/tokenizer_manager.py -+++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -1044,10 +1044,15 @@ class TokenizerManager: - request: Optional[fastapi.Request] = None, - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() -- assert ( -- self.server_args.dp_size == 1 -- ), "dp_size must be 1 for init parameter update group" -- result = (await self.init_weights_update_group_communicator(obj))[0] -+ results = await self.init_weights_update_group_communicator(obj) -+ if self.server_args.dp_size == 1: -+ result = results[0] -+ return result.success, result.message -+ else: -+ all_success = all([r.success for r in results]) -+ all_message = [r.message for r in results] -+ all_message = " | ".join(all_message) -+ return all_success, all_message - return result.success, result.message - - async def update_weights_from_distributed( -@@ -1056,9 +1061,6 @@ class TokenizerManager: - request: Optional[fastapi.Request] = None, - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() -- assert ( -- self.server_args.dp_size == 1 or self.server_args.enable_dp_attention -- ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed" - - if obj.abort_all_requests: - self.abort_request(abort_all=True) -@@ -1066,8 +1068,15 @@ class TokenizerManager: - # This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: -- result = (await self.update_weights_from_distributed_communicator(obj))[0] -- return result.success, result.message -+ results = await self.update_weights_from_distributed_communicator(obj) -+ if self.server_args.dp_size == 1: -+ result = results[0] -+ return result.success, result.message -+ else: -+ all_success = all([r.success for r in results]) -+ all_message = [r.message for r in results] -+ all_message = " | ".join(all_message) -+ return all_success, all_message - - async def update_weights_from_tensor( - self, -diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 5222bff0a..ff0bbc62a 100644 ---- a/python/sglang/srt/model_executor/model_runner.py -+++ b/python/sglang/srt/model_executor/model_runner.py -@@ -22,6 +22,7 @@ import os - import time - from dataclasses import dataclass - from typing import List, Optional, Tuple, Union -+from contextlib import nullcontext - - import torch - import torch.distributed as dist -@@ -675,7 +676,7 @@ class ModelRunner: - monkey_patch_vllm_parallel_state() - monkey_patch_isinstance_for_vllm_base_layer() - -- with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_WEIGHTS): -+ with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_WEIGHTS) if not self.is_draft_worker else nullcontext(): - self.model = get_model( - model_config=self.model_config, - load_config=self.load_config, -diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py -index e0f0b373d..a18ac10f1 100644 ---- a/python/sglang/srt/models/glm4_moe.py -+++ b/python/sglang/srt/models/glm4_moe.py -@@ -1108,5 +1108,4 @@ class Glm4MoeForCausalLM(DeepseekV2ForCausalLM): - ) - weight_loader(param, loaded_weight) - -- - EntryClass = [Glm4MoeForCausalLM] diff --git a/docker/amd_patch/sglv0.5.0rc0/amd_megatron_fused_kernels_init.patch b/docker/amd_patch/sglv0.5.0rc0/amd_megatron_fused_kernels_init.patch deleted file mode 100644 index f6efca346..000000000 --- a/docker/amd_patch/sglv0.5.0rc0/amd_megatron_fused_kernels_init.patch +++ /dev/null @@ -1,51 +0,0 @@ -diff --git a/megatron/legacy/fused_kernels/__init__.py b/megatron/legacy/fused_kernels/__init__.py -index 87cceac3..ac686d74 100644 ---- a/megatron/legacy/fused_kernels/__init__.py -+++ b/megatron/legacy/fused_kernels/__init__.py -@@ -3,6 +3,7 @@ - import os - import pathlib - import subprocess -+import torch - - from torch.utils import cpp_extension - -@@ -15,23 +16,23 @@ os.environ["TORCH_CUDA_ARCH_LIST"] = "" - - - def load(args): -- -- # Check if cuda 11 is installed for compute capability 8.0 -- cc_flag = [] -- _, bare_metal_major, bare_metal_minor = _get_cuda_bare_metal_version( -- cpp_extension.CUDA_HOME -- ) -- if int(bare_metal_major) >= 11: -- cc_flag.append('-gencode') -- cc_flag.append('arch=compute_80,code=sm_80') -- if int(bare_metal_minor) >= 8: -+ if torch.cuda.is_available() and torch.version.cuda: -+ # Check if cuda 11 is installed for compute capability 8.0 -+ cc_flag = [] -+ _, bare_metal_major, bare_metal_minor = _get_cuda_bare_metal_version( -+ cpp_extension.CUDA_HOME -+ ) -+ if int(bare_metal_major) >= 11: - cc_flag.append('-gencode') -- cc_flag.append('arch=compute_90,code=sm_90') -+ cc_flag.append('arch=compute_80,code=sm_80') -+ if int(bare_metal_minor) >= 8: -+ cc_flag.append('-gencode') -+ cc_flag.append('arch=compute_90,code=sm_90') - -- # Build path -- srcpath = pathlib.Path(__file__).parent.absolute() -- buildpath = srcpath / "build" -- _create_build_dir(buildpath) -+ # Build path -+ srcpath = pathlib.Path(__file__).parent.absolute() -+ buildpath = srcpath / "build" -+ _create_build_dir(buildpath) - - # Helper function to build the kernels. - def _cpp_extention_load_helper(name, sources, extra_cuda_flags): diff --git a/docker/amd_patch/sglv0.5.0rc0/megatron.patch b/docker/amd_patch/sglv0.5.0rc0/megatron.patch deleted file mode 100644 index c840133ce..000000000 --- a/docker/amd_patch/sglv0.5.0rc0/megatron.patch +++ /dev/null @@ -1,438 +0,0 @@ -diff --git a/megatron/core/distributed/__init__.py b/megatron/core/distributed/__init__.py -index fe26e8b4..4451f277 100644 ---- a/megatron/core/distributed/__init__.py -+++ b/megatron/core/distributed/__init__.py -@@ -11,3 +11,15 @@ from .finalize_model_grads import finalize_model_grads - from .fsdp.mcore_fsdp_adapter import FullyShardedDataParallel - from .torch_fully_sharded_data_parallel import TorchFullyShardedDataParallel - from .torch_fully_sharded_data_parallel_config import TorchFullyShardedDataParallelConfig -+ -+# Backward compatibility patch for FSDP module reorganization -+import sys -+import importlib.util -+ -+spec = importlib.util.find_spec('megatron.core.distributed.fsdp.src.megatron_fsdp') -+if spec: -+ custom_fsdp = importlib.util.module_from_spec(spec) -+ spec.loader.exec_module(custom_fsdp) -+ sys.modules['megatron.core.distributed.custom_fsdp'] = custom_fsdp -+ if hasattr(custom_fsdp, 'MegatronFSDP'): -+ custom_fsdp.FullyShardedDataParallel = custom_fsdp.MegatronFSDP -diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py -index 99c3edc0..26ea5cb4 100644 ---- a/megatron/core/extensions/transformer_engine.py -+++ b/megatron/core/extensions/transformer_engine.py -@@ -404,6 +404,7 @@ class TELinear(te.pytorch.Linear): - ) - - for param in self.parameters(): -+ setattr(param, "parallel_mode", parallel_mode) - if is_expert: - # Reduce the gradient on the expert_data_parallel group for expert linear layers - setattr(param, "allreduce", not self.expert_parallel) -diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py -index 002edb92..f7273488 100755 ---- a/megatron/core/models/gpt/gpt_layer_specs.py -+++ b/megatron/core/models/gpt/gpt_layer_specs.py -@@ -80,6 +80,8 @@ def get_gpt_layer_with_transformer_engine_spec( - use_te_op_fuser: Optional[bool] = False, - use_kitchen: bool = False, - use_te_activation_func: bool = False, -+ post_self_attn_layernorm: bool = False, -+ post_mlp_layernorm: bool = False, - ) -> ModuleSpec: - """Use this spec to use lower-level Transformer Engine modules (required for fp8 training). - -@@ -182,9 +184,11 @@ def get_gpt_layer_with_transformer_engine_spec( - ), - ), - self_attn_bda=get_bias_dropout_add, -+ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, - pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp, - mlp=mlp, - mlp_bda=get_bias_dropout_add, -+ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, - sharded_state_dict_keys_map={ - "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight", - "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias", -diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py -index df9adc3e..2f4f544a 100644 ---- a/megatron/core/models/gpt/gpt_model.py -+++ b/megatron/core/models/gpt/gpt_model.py -@@ -443,7 +443,7 @@ class GPTModel(LanguageModule): - if self.share_embeddings_and_output_weights: - output_weight = self.shared_embedding_or_output_weight() - -- if mtp_in_postprocess: -+ if mtp_in_postprocess and labels is not None: - hidden_states = self.mtp( - input_ids=input_ids, - position_ids=position_ids, -diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py -index 57332ac3..f3abd642 100644 ---- a/megatron/core/parallel_state.py -+++ b/megatron/core/parallel_state.py -@@ -9,6 +9,7 @@ from typing import Callable, List, Optional - - import numpy as np - import torch -+import torch.distributed as dist - - from .utils import GlobalMemoryBuffer, is_torch_min_version - -@@ -163,6 +164,213 @@ def get_nccl_options(pg_name, nccl_comm_cfgs): - return None - - -+old_new_group = None -+ -+ -+def monkey_patch_torch_dist(): -+ print("Applying monkey patch to torch.distributed", flush=True) -+ global old_new_group -+ if old_new_group is not None: -+ return -+ -+ old_new_group = dist.new_group -+ -+ def new_group(*args, **kwargs): -+ group = old_new_group(*args, **kwargs) -+ # skip none nccl group. -+ if ( -+ len(args) >= 3 and args[2] == "gloo" or -+ "backend" in kwargs and kwargs["backend"] == "gloo" -+ ): -+ return group -+ -+ # Get ranks from arguments -+ if len(args) >= 1 and args[0] is not None: -+ ranks = args[0] -+ elif "ranks" in kwargs and kwargs["ranks"] is not None: -+ ranks = kwargs["ranks"] -+ else: -+ # If no ranks specified, use all ranks in world -+ ranks = list(range(dist.get_world_size())) -+ -+ if len(ranks) == 1: -+ return group -+ -+ group = ReloadableProcessGroup(group, ranks) -+ return group -+ -+ dist.new_group = new_group -+ -+ def get_new_function(func): -+ def new_function(*args, **kwargs): -+ args = ( -+ arg.group if isinstance(arg, ReloadableProcessGroup) else arg -+ for arg in args -+ ) -+ kwargs = { -+ k: (v.group if isinstance(v, ReloadableProcessGroup) else v) -+ for k, v in kwargs.items() -+ } -+ return func(*args, **kwargs) -+ return new_function -+ -+ dist.get_rank = get_new_function(dist.get_rank) -+ dist.get_world_size = get_new_function(dist.get_world_size) -+ dist.get_backend = get_new_function(dist.get_backend) -+ dist.get_global_rank = get_new_function(dist.get_global_rank) -+ dist.get_group_rank = get_new_function(dist.get_group_rank) -+ dist.get_process_group_ranks = get_new_function(dist.get_process_group_ranks) -+ -+ dist.all_reduce = get_new_function(dist.all_reduce) -+ dist.all_gather = get_new_function(dist.all_gather) -+ dist.all_gather_into_tensor = get_new_function(dist.all_gather_into_tensor) -+ dist.all_gather_object = get_new_function(dist.all_gather_object) -+ dist.all_to_all = get_new_function(dist.all_to_all) -+ dist.all_to_all_single = get_new_function(dist.all_to_all_single) -+ dist.broadcast = get_new_function(dist.broadcast) -+ dist.reduce = get_new_function(dist.reduce) -+ dist.reduce_scatter = get_new_function(dist.reduce_scatter) -+ dist.reduce_scatter_tensor = get_new_function(dist.reduce_scatter_tensor) -+ dist.scatter = get_new_function(dist.scatter) -+ dist.gather = get_new_function(dist.gather) -+ dist.barrier = get_new_function(dist.barrier) -+ dist.send = get_new_function(dist.send) -+ dist.recv = get_new_function(dist.recv) -+ dist._coalescing_manager = get_new_function(dist._coalescing_manager) -+ -+ # p2p -+ old_isend = dist.isend -+ old_irecv = dist.irecv -+ -+ dist.isend = get_new_function(dist.isend) -+ dist.irecv = get_new_function(dist.irecv) -+ -+ def get_new_p2pop_function(func): -+ def new_function(*args, **kwargs): -+ def convert(arg): -+ if isinstance(arg, ReloadableProcessGroup): -+ return arg.group -+ elif arg == dist.isend: -+ arg = old_isend -+ elif arg == dist.irecv: -+ arg = old_irecv -+ return arg -+ -+ args = (convert(arg) for arg in args) -+ kwargs = { -+ k: convert(v) -+ for k, v in kwargs.items() -+ } -+ return func(*args, **kwargs) -+ return new_function -+ -+ dist.P2POp.__new__ = get_new_p2pop_function(dist.P2POp.__new__) -+ dist.P2POp.__init__ = get_new_p2pop_function(dist.P2POp.__init__) -+ -+ -+ -+class ReloadableProcessGroup(torch.distributed.ProcessGroup): -+ GROUPS = [] -+ -+ def __init__(self, group, ranks): -+ super().__init__( -+ rank=dist.get_rank(group), -+ size=dist.get_world_size(group), -+ ) -+ #print(f"Creating ReloadableProcessGroup with ranks: {ranks}", flush=True) -+ self.group = group -+ self.group_info = { -+ "ranks": ranks, -+ } -+ ReloadableProcessGroup.GROUPS.append(self) -+ -+ def __getattr__(self, name): -+ return getattr(self.group, name) -+ -+ @staticmethod -+ def destroy_process_groups(): -+ for reloadable_group in ReloadableProcessGroup.GROUPS: -+ if reloadable_group.group is None: -+ continue -+ #print(f"Destroying process group: {reloadable_group.group_info['ranks']}") -+ dist.destroy_process_group(reloadable_group.group) -+ del reloadable_group.group -+ reloadable_group.group = None -+ -+ @staticmethod -+ def reload_process_groups(): -+ for reloadable_group in ReloadableProcessGroup.GROUPS: -+ if reloadable_group.group is not None: -+ continue -+ #print(f"Reloading process group: {reloadable_group.group_info['ranks']}") -+ group = old_new_group( -+ ranks=reloadable_group.group_info["ranks"], -+ backend="nccl" -+ ) -+ reloadable_group.group = group -+ -+ def rank(self) -> int: return self.group.rank() -+ def size(self) -> int: return self.group.size() -+ def name(self) -> str: return self.group.name() -+ -+ def shutdown(self) -> None: -+ if self.group is not None: -+ self.group.shutdown() -+ -+ def abort(self) -> None: -+ if self.group is not None: -+ self.group.abort() -+ -+ def _fwd(self, method, *args, **kwargs): -+ inner = self.group -+ if inner is None: -+ raise RuntimeError("ReloadableProcessGroup: inner PG is None, call reload() first.") -+ return getattr(inner, method)(*args, **kwargs) -+ -+ def barrier(self, *a, **kw): return self._fwd("barrier", *a, **kw) -+ def broadcast(self, *a, **kw): return self._fwd("broadcast", *a, **kw) -+ def allreduce(self, *a, **kw): return self._fwd("allreduce", *a, **kw) -+ def allreduce_coalesced(self, *a, **kw): return self._fwd("allreduce_coalesced", *a, **kw) -+ def reduce(self, *a, **kw): return self._fwd("reduce", *a, **kw) -+ def allgather(self, *a, **kw): return self._fwd("allgather", *a, **kw) -+ def _allgather_base(self, *a, **kw): return self._fwd("_allgather_base", *a, **kw) -+ def allgather_coalesced(self, *a, **kw): return self._fwd("allgather_coalesced", *a, **kw) -+ def allgather_into_tensor_coalesced(self, *a, **kw): return self._fwd("allgather_into_tensor_coalesced", *a, **kw) -+ def gather(self, *a, **kw): return self._fwd("gather", *a, **kw) -+ def scatter(self, *a, **kw): return self._fwd("scatter", *a, **kw) -+ def reduce_scatter(self, *a, **kw): return self._fwd("reduce_scatter", *a, **kw) -+ def _reduce_scatter_base(self, *a, **kw): return self._fwd("_reduce_scatter_base", *a, **kw) -+ def reduce_scatter_tensor_coalesced(self, *a, **kw): return self._fwd("reduce_scatter_tensor_coalesced", *a, **kw) -+ def alltoall_base(self, *a, **kw): return self._fwd("alltoall_base", *a, **kw) -+ def alltoall(self, *a, **kw): return self._fwd("alltoall", *a, **kw) -+ def send(self, *a, **kw): return self._fwd("send", *a, **kw) -+ def recv(self, *a, **kw): return self._fwd("recv", *a, **kw) -+ def recv_anysource(self, *a, **kw): return self._fwd("recv_anysource", *a, **kw) -+ -+ def _start_coalescing(self, *a, **kw): return self._fwd("_start_coalescing", *a, **kw) -+ def _end_coalescing(self, *a, **kw): return self._fwd("_end_coalescing", *a, **kw) -+ def _get_backend_name(self): return self._fwd("_get_backend_name") -+ def _get_backend(self, *a, **kw): return self._fwd("_get_backend", *a, **kw) -+ def _set_default_backend(self, *a, **kw): return self._fwd("_set_default_backend", *a, **kw) -+ @property -+ def bound_device_id(self): return self.group.bound_device_id -+ @bound_device_id.setter -+ def bound_device_id(self, dev): self.group.bound_device_id = dev -+ -+ -+def destroy_process_groups(): -+ """Destroy all reloadable process groups.""" -+ ReloadableProcessGroup.destroy_process_groups() -+ -+ -+def reload_process_groups(): -+ """Reload all reloadable process groups.""" -+ ReloadableProcessGroup.reload_process_groups() -+ -+ -+monkey_patch_torch_dist() -+ -+ - def create_group( - ranks=None, - timeout=None, -diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py -index 63ee9d1f..b90b744c 100644 ---- a/megatron/core/pipeline_parallel/p2p_communication.py -+++ b/megatron/core/pipeline_parallel/p2p_communication.py -@@ -26,22 +26,22 @@ def _batched_p2p_ops( - ops = [] - if tensor_send_prev is not None: - send_prev_op = torch.distributed.P2POp( -- torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, group -+ torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, - ) - ops.append(send_prev_op) - if tensor_recv_prev is not None: - recv_prev_op = torch.distributed.P2POp( -- torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, group -+ torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, - ) - ops.append(recv_prev_op) - if tensor_send_next is not None: - send_next_op = torch.distributed.P2POp( -- torch.distributed.isend, tensor_send_next, next_pipeline_rank, group -+ torch.distributed.isend, tensor_send_next, next_pipeline_rank, - ) - ops.append(send_next_op) - if tensor_recv_next is not None: - recv_next_op = torch.distributed.P2POp( -- torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, group -+ torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, - ) - ops.append(recv_next_op) - if len(ops) > 0: -diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py -index 6f557e1f..b295fd35 100644 ---- a/megatron/core/transformer/transformer_config.py -+++ b/megatron/core/transformer/transformer_config.py -@@ -173,6 +173,9 @@ class TransformerConfig(ModelParallelConfig): - qk_layernorm: bool = False - """Whether to apply `normalization` type of normalization to the query and key embeddings.""" - -+ post_self_attn_layernorm: bool = False -+ post_mlp_layernorm: bool = False -+ - test_mode: bool = False - """Whether to run real-time tests.""" - -diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py -index 84f22bde..b4807d26 100644 ---- a/megatron/core/transformer/transformer_layer.py -+++ b/megatron/core/transformer/transformer_layer.py -@@ -224,6 +224,7 @@ class TransformerLayerSubmodules: - input_layernorm: Union[ModuleSpec, type] = IdentityOp - self_attention: Union[ModuleSpec, type] = IdentityOp - self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp -+ post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp - - pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp - cross_attention: Union[ModuleSpec, type] = IdentityOp -@@ -232,6 +233,7 @@ class TransformerLayerSubmodules: - pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp - mlp: Union[ModuleSpec, type] = IdentityOp - mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp -+ post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp - - # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method - sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict) -@@ -336,6 +338,14 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - # [Module 3: BiasDropoutFusion] - self.self_attn_bda = build_module(submodules.self_attn_bda) - -+ self.post_self_attn_layernorm = build_module( -+ submodules.post_self_attn_layernorm, -+ config=self.config, -+ hidden_size=self.config.hidden_size, -+ eps=self.config.layernorm_epsilon, -+ ) -+ -+ - # [Module 4: Post SelfAttention] Optional Layernorm after self-attn - self.pre_cross_attn_layernorm = build_module( - submodules.pre_cross_attn_layernorm, -@@ -399,6 +409,13 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - # [Module 9: BiasDropoutFusion] - self.mlp_bda = build_module(submodules.mlp_bda) - -+ self.post_mlp_layernorm = build_module( -+ submodules.post_mlp_layernorm, -+ config=self.config, -+ hidden_size=self.config.hidden_size, -+ eps=self.config.layernorm_epsilon -+ ) -+ - self.recompute_input_layernorm = False - self.recompute_pre_mlp_layernorm = False - self.recompute_mlp = False -@@ -535,6 +552,11 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - attention_output_with_bias[0] - ) - -+ attention_output, attention_output_bias = attention_output_with_bias -+ attention_output = self.post_self_attn_layernorm(attention_output) -+ attention_output_with_bias = (attention_output, attention_output_bias) -+ -+ - # TODO: could we move `bias_dropout_add_exec_handler` itself - # inside the module provided in the `bias_dropout_add_spec` module? - nvtx_range_push(suffix="self_attn_bda") -@@ -635,6 +657,10 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer): - else: - mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output) - -+ mlp_output, mlp_output_bias = mlp_output_with_bias -+ mlp_output = self.post_mlp_layernorm(mlp_output) -+ mlp_output_with_bias = (mlp_output, mlp_output_bias) -+ - if self.recompute_pre_mlp_layernorm: - # discard the output of the pre-mlp layernorm and register the recompute - # as a gradient hook of mlp_output_with_bias[0] -diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py -index 24ba8926..4f039fd4 100644 ---- a/megatron/training/arguments.py -+++ b/megatron/training/arguments.py -@@ -1191,6 +1191,9 @@ def core_transformer_config_from_args(args, config_class=None): - if args.is_hybrid_model: - kw_args['is_hybrid_model'] = args.is_hybrid_model - -+ kw_args['post_self_attn_layernorm'] = args.post_self_attn_layernorm -+ kw_args['post_mlp_layernorm'] = args.post_mlp_layernorm -+ - # handle quantization config - # NOTE: Kitchen arguments are only added to the namespace when - # Kitchen library is available. -@@ -1481,6 +1484,10 @@ def _add_network_size_args(parser): - action='store_true', - help='If set, use original BERT residula connection ' - 'ordering.') -+ group.add_argument('--post-self-attn-layernorm', action='store_true', -+ help='If set, use post self attention layernorm.') -+ group.add_argument('--post-mlp-layernorm', action='store_true', -+ help='If set, use post MLP layernorm.') - group.add_argument('--openai-gelu', action='store_true', - help='Use OpenAIs GeLU implementation. This option' - 'should not be used unless for backward compatibility' diff --git a/docker/amd_patch/sglv0.5.0rc0/sglang.patch b/docker/amd_patch/sglv0.5.0rc0/sglang.patch deleted file mode 100644 index 990c2e628..000000000 --- a/docker/amd_patch/sglv0.5.0rc0/sglang.patch +++ /dev/null @@ -1,203 +0,0 @@ -diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py -index bdb124e51..3edf30ab1 100644 ---- a/python/sglang/srt/configs/model_config.py -+++ b/python/sglang/srt/configs/model_config.py -@@ -454,14 +454,14 @@ class ModelConfig: - ).lower() - - # Detect which checkpoint is it -- for _, method in QUANTIZATION_METHODS.items(): -- quantization_override = method.override_quantization_method( -- quant_cfg, self.quantization -- ) -- if quantization_override: -- quant_method = quantization_override -- self.quantization = quantization_override -- break -+ # for _, method in QUANTIZATION_METHODS.items(): -+ # quantization_override = method.override_quantization_method( -+ # quant_cfg, self.quantization -+ # ) -+ # if quantization_override: -+ # quant_method = quantization_override -+ # self.quantization = quantization_override -+ # break - - # Verify quantization configurations. - if self.quantization is None: -diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py -index 2dd2c75f1..f2adb18f8 100644 ---- a/python/sglang/srt/entrypoints/http_server.py -+++ b/python/sglang/srt/entrypoints/http_server.py -@@ -264,6 +264,10 @@ async def validate_json_request(raw_request: Request): - - - @app.get("/health") -+async def health(request: Request) -> Response: -+ return Response(status_code=200) -+ -+ - @app.get("/health_generate") - async def health_generate(request: Request) -> Response: - """ -diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -index 372717bf9..40665cc90 100644 ---- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -+++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -@@ -190,6 +190,7 @@ class DeepEPBuffer: - f"Consider using --deepep-config to change the behavior." - ) - -+ num_qps_per_rank = 20 - cls._buffer = Buffer( - group, - num_nvl_bytes, -diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py -index 956264fc9..69f729336 100644 ---- a/python/sglang/srt/layers/quantization/fp8.py -+++ b/python/sglang/srt/layers/quantization/fp8.py -@@ -351,10 +351,10 @@ class Fp8LinearMethod(LinearMethodBase): - return - else: - weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data -- layer.weight = torch.nn.Parameter(weight, requires_grad=False) -- layer.weight_scale_inv = torch.nn.Parameter( -- weight_scale, requires_grad=False -- ) -+ # layer.weight = torch.nn.Parameter(weight, requires_grad=False) -+ # layer.weight_scale_inv = torch.nn.Parameter( -+ # weight_scale, requires_grad=False -+ # ) - return - - layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) -diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 95a529c89..758fbfd5f 100644 ---- a/python/sglang/srt/managers/scheduler.py -+++ b/python/sglang/srt/managers/scheduler.py -@@ -1359,7 +1359,7 @@ class Scheduler( - - if memory_leak: - msg = "token_to_kv_pool_allocator memory leak detected! " f"{token_msg}" -- raise ValueError(msg) -+ # raise ValueError(msg) - - if self.disaggregation_mode == DisaggregationMode.DECODE: - req_total_size = ( -@@ -1374,7 +1374,7 @@ class Scheduler( - f"available_size={len(self.req_to_token_pool.free_slots)}, " - f"total_size={self.req_to_token_pool.size}\n" - ) -- raise ValueError(msg) -+ # raise ValueError(msg) - - if ( - self.enable_metrics -@@ -1830,6 +1830,7 @@ class Scheduler( - deepep_mode=DeepEPMode(self.server_args.deepep_mode), - require_mlp_tp_gather=require_mlp_tp_gather(self.server_args), - disable_overlap_schedule=self.server_args.disable_overlap_schedule, -+ offload_tags=self.offload_tags, - ) - - def handle_dp_balance_data(self, local_batch: ScheduleBatch): -@@ -1927,6 +1928,7 @@ class Scheduler( - deepep_mode: DeepEPMode, - require_mlp_tp_gather: bool, - disable_overlap_schedule: bool, -+ offload_tags: set[str], - ): - # Check if other DP workers have running batches - if local_batch is None: -@@ -1957,7 +1959,7 @@ class Scheduler( - ) - - tbo_preparer = TboDPAttentionPreparer() -- if disable_overlap_schedule: -+ if len(offload_tags) == 0 and disable_overlap_schedule: - group = tp_group.device_group - device = tp_group.device - else: -diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 58220b1d6..3c3d081a8 100644 ---- a/python/sglang/srt/managers/tokenizer_manager.py -+++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -1044,10 +1044,15 @@ class TokenizerManager: - request: Optional[fastapi.Request] = None, - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() -- assert ( -- self.server_args.dp_size == 1 -- ), "dp_size must be 1 for init parameter update group" -- result = (await self.init_weights_update_group_communicator(obj))[0] -+ results = await self.init_weights_update_group_communicator(obj) -+ if self.server_args.dp_size == 1: -+ result = results[0] -+ return result.success, result.message -+ else: -+ all_success = all([r.success for r in results]) -+ all_message = [r.message for r in results] -+ all_message = " | ".join(all_message) -+ return all_success, all_message - return result.success, result.message - - async def update_weights_from_distributed( -@@ -1056,9 +1061,6 @@ class TokenizerManager: - request: Optional[fastapi.Request] = None, - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() -- assert ( -- self.server_args.dp_size == 1 or self.server_args.enable_dp_attention -- ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed" - - if obj.abort_all_requests: - self.abort_request(abort_all=True) -@@ -1066,8 +1068,15 @@ class TokenizerManager: - # This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: -- result = (await self.update_weights_from_distributed_communicator(obj))[0] -- return result.success, result.message -+ results = await self.update_weights_from_distributed_communicator(obj) -+ if self.server_args.dp_size == 1: -+ result = results[0] -+ return result.success, result.message -+ else: -+ all_success = all([r.success for r in results]) -+ all_message = [r.message for r in results] -+ all_message = " | ".join(all_message) -+ return all_success, all_message - - async def update_weights_from_tensor( - self, -diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 5222bff0a..ff0bbc62a 100644 ---- a/python/sglang/srt/model_executor/model_runner.py -+++ b/python/sglang/srt/model_executor/model_runner.py -@@ -22,6 +22,7 @@ import os - import time - from dataclasses import dataclass - from typing import List, Optional, Tuple, Union -+from contextlib import nullcontext - - import torch - import torch.distributed as dist -@@ -675,7 +676,7 @@ class ModelRunner: - monkey_patch_vllm_parallel_state() - monkey_patch_isinstance_for_vllm_base_layer() - -- with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_WEIGHTS): -+ with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_WEIGHTS) if not self.is_draft_worker else nullcontext(): - self.model = get_model( - model_config=self.model_config, - load_config=self.load_config, -diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py -index e0f0b373d..a18ac10f1 100644 ---- a/python/sglang/srt/models/glm4_moe.py -+++ b/python/sglang/srt/models/glm4_moe.py -@@ -1108,5 +1108,4 @@ class Glm4MoeForCausalLM(DeepseekV2ForCausalLM): - ) - weight_loader(param, loaded_weight) - -- - EntryClass = [Glm4MoeForCausalLM] diff --git a/docs/en/get_started/quick_start.md b/docs/en/get_started/quick_start.md index 5d3856fde..ea792928e 100644 --- a/docs/en/get_started/quick_start.md +++ b/docs/en/get_started/quick_start.md @@ -20,8 +20,7 @@ Since slime may contain temporary patches for sglang/megatron, to avoid potentia - B-series basic functionality is stable and suitable for development/testing, but currently lacks CI protection - Both hardware platforms use identical installation and startup procedures -- For scenarios where Docker is not convenient, please refer to [build_conda.sh](https://github.com/THUDM/slime/blob/main/build_conda.sh); -- For AMD support, please refer to [AMD Usage Tutorial](../platform_support/amd_tutorial.md). +- For scenarios where Docker is not convenient, please refer to [build_conda.sh](https://github.com/THUDM/slime/blob/main/build_conda.sh). ### Pull and Start Docker Container diff --git a/docs/en/platform_support/amd_tutorial.md b/docs/en/platform_support/amd_tutorial.md deleted file mode 100644 index 73d690fb1..000000000 --- a/docs/en/platform_support/amd_tutorial.md +++ /dev/null @@ -1,294 +0,0 @@ -# AMD - -⚠️ If you encounter problems on AMD instinct, feel free to reach out [Yusheng Su](https://yushengsu-thu.github.io/). - - -## Introduction - -If you are running slime on AMD's Instinct, please refer to the following materials. This tutorial will explain how to set up the development environment (Docker), use the modified ROCm dependencies, and provide an example for running the experiments. The current rocm docker only support AMD's MI300 and MI325 GPUs. - - - - -## Docker - -You can download the prebuilt image from DockerHub: [rlsys/slime](https://hub.docker.com/r/rlsys/slime/tags). -```bash -docker pull rlsys/slime:latest -``` -Or you can use the [Dockerfile.rocm](https://github.com/THUDM/slime/blob/main/docker/Dockerfile.rocm) to build it on your side. -```bash -cd docker -docker build -f Dockerfile.rocm -t rlsys/slime:latest . -``` - -Acknowledgement: Thanks to [Yang Wang](https://www.microsoft.com/en-us/research/people/yangwang5/) for working on the patch for this [ROCm base Docker image](https://hub.docker.com/r/rlsys/rocm-6.3.4-patch/tags) to support virtual memory management on MI300X. - - -## Quick Start - -### Environment Setup - -Based on the [rlsys/slime](https://hub.docker.com/r/rlsys/slime/tags) image (pre-installed with SGLang and Megatron-LM): -```bash -docker run --rm -it \ - --device /dev/dri \ - --device /dev/kfd \ - -p 8265:8265 \ - --group-add video \ - --cap-add SYS_PTRACE \ - --security-opt seccomp=unconfined \ - --privileged \ - -v $HOME/.ssh:/root/.ssh \ - -v $HOME:$HOME \ - --shm-size 128G \ - --name slime_dev \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - -w $PWD \ - rlsys/slime:latest \ - /bin/bash -``` - -Then, download and install slime: -```bash -git clone https://github.com/THUDM/slime.git -cd slime -pip install -e . --no-deps -``` - -Download the model and data: - -```bash -# hf checkpoint -hf download Qwen/Qwen3-4B --local-dir /root/Qwen3-4B - -# train data -hf download --repo-type dataset zhuzilin/dapo-math-17k \ - --local-dir /root/dapo-math-17k - -# eval data -hf download --repo-type dataset zhuzilin/aime-2024 \ - --local-dir /root/aime-2024 -``` - -### Checkpoint Format Conversion - -Since slime uses Megatron, and Megatron does not support loading Hugging Face checkpoints directly, we need to convert the model to the `torch_dist` format that Megatron supports. - -#### HF → Megatron torch\_dist ckpt - -Use [mbridge](https://github.com/ISEEKYAN/mbridge.git) or [Megatron-LM-amd_version-amd](https://github.com/yushengsu-thu/Megatron-LM-amd_version.git) for conversion: - -```bash -cd slime/ -source scripts/models/qwen3-4B.sh -MEGATRON_LM_PATH=$(pip list | grep megatron-core | awk '{print $NF}') -PYTHONPATH=${MEGATRON_LM_PATH} python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --no-gradient-accumulation-fusion \ - --hf-checkpoint /root/Qwen3-4B \ - --save /root/Qwen3-4B_torch_dist -``` - -Note: We implemented a dedicated AMD conversion script that forces a CPU-only conversion workflow using the Gloo backend to bypass hardware-specific issues. A GPU-based script for ROCm is currently in development. - -⚠️ If you encounter an issue where slime cannot be found, please run `pip install -e . --no-deps` in the slime directory. - - -### Example: Qwen3-4B - -We provide examples to use [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B), please refer to: -- [Example: Qwen3-4B Model](https://github.com/THUDM/slime/blob/main/scripts/run-qwen3-4B-amd.sh): Just run - -```bash -SLIME_DIR=/root \ -MODEL_DIR=/root \ -DATA_DIR=/root \ -bash scripts/run-qwen3-4B-amd.sh -``` - -⚠️ TODO: ROCM seems to not support `apex` yet. Thus, we need to disable gradient accumulation fusionby adding the `--no-gradient-accumulation-fusion` flag in the training script currently. We will continue investigating how to enable this. - -⚠️ Note: The main difference between ROCm's training script and NVIDIA's script is that you need to set `RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` for ray to function properly on AMD GPUs. - -- We show the training script below: - -```bash -#!/bin/bash - -# for rerun the task -pkill -9 sglang -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - - -set -euxo pipefail - - -### AMD Support ### -SLIME_DIR="${SLIME_DIR:-/home/yushensu/projects/slime}" # Default path if not set in environment -export SLIME_DIR - -MODEL_DIR="${MODEL_DIR:-/home/yushensu/projects/model}" # Default path if not set in environment -export MODEL_DIR - -DATA_DIR="${DATA_DIR:-/home/yushensu/projects/data}" # Default path if not set in environment -export DATA_DIR - -# For AMD GPU -export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=${RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES:-"1"} # Must set to 1 -export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-"0,1,2,3,4,5,6,7"} #You can choose which gpus to use -#################### - - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-4B.sh" - -CKPT_ARGS=( - --hf-checkpoint ${MODEL_DIR}/Qwen3-4B - --ref-load ${MODEL_DIR}/Qwen3-4B_torch_dist - --load ${MODEL_DIR}/Qwen3-4B_slime/ - --save ${MODEL_DIR}/Qwen3-4B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data ${DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type deepscaler - --num-rollout 3000 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - - --global-batch-size 256 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime ${DATA_DIR}/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project slime-dev - # --wandb-group qwen3-4B-test - # --wandb-key ${WANDB_KEY} -) - -SGLANG_ARGS=( - --rollout-num-gpus-per-engine 2 - --sglang-mem-fraction-static 0.7 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - ### AMD Support ### - # disable gradient accumulation fusion: Need to add apex to enable this - --no-gradient-accumulation-fusion - ################### -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} - -NUM_GPUS=$(echo ${HIP_VISIBLE_DEVICES} | tr ',' '\n' | wc -l) -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - - -# "PYTHONPATH": "/workspace/Megatron-LM/", -MEGATRON_LM_PATH=$(pip list | grep megatron-core | awk '{print $NF}') - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json='{ - "env_vars": { - "PYTHONPATH": "/workspace/Megatron-LM/", - "CUDA_DEVICE_MAX_CONNECTIONS": "1" - } - }' \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${SGLANG_ARGS[@]} \ - ${MISC_ARGS[@]} - - -####clear after training -pkill -9 sglang -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python -``` diff --git a/docs/zh/get_started/quick_start.md b/docs/zh/get_started/quick_start.md index 463af729f..a57d93452 100644 --- a/docs/zh/get_started/quick_start.md +++ b/docs/zh/get_started/quick_start.md @@ -19,8 +19,7 @@ - B 卡基本功能稳定，可作为开发和测试参考，但暂无 CI 保护 - 两种硬件平台使用完全相同的安装和启动流程 -- 对于不方便使用 docker 的场景，请参考 [build_conda.sh](https://github.com/THUDM/slime/blob/main/build_conda.sh)； -- 对于 AMD 支持，请参考 [AMD 使用教程](../../en/platform_support/amd_tutorial.md)。 +- 对于不方便使用 docker 的场景，请参考 [build_conda.sh](https://github.com/THUDM/slime/blob/main/build_conda.sh)。 ### 拉取并启动 Docker 容器 diff --git a/scripts/run-qwen3-4B-amd.sh b/scripts/run-qwen3-4B-amd.sh deleted file mode 100755 index 8452d8e5f..000000000 --- a/scripts/run-qwen3-4B-amd.sh +++ /dev/null @@ -1,167 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 sglang -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - - -set -euxo pipefail - - -### AMD Support ### -SLIME_DIR="${SLIME_DIR:-/root}" # Default path if not set in environment -export SLIME_DIR - -MODEL_DIR="${MODEL_DIR:-/root}" # Default path if not set in environment -export MODEL_DIR - -DATA_DIR="${DATA_DIR:-/root}" # Default path if not set in environment -export DATA_DIR - -# For AMD GPU -export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=${RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES:-"1"} # Must set to 1 -export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-"0,1,2,3,4,5,6,7"} #You can choose which gpus to use -#################### - - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-4B.sh" - -CKPT_ARGS=( - --hf-checkpoint ${MODEL_DIR}/Qwen3-4B - --ref-load ${MODEL_DIR}/Qwen3-4B_torch_dist - --load ${MODEL_DIR}/Qwen3-4B_slime/ - --save ${MODEL_DIR}/Qwen3-4B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data ${DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type deepscaler - --num-rollout 3000 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - - --global-batch-size 256 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime ${DATA_DIR}/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project slime-dev - # --wandb-group qwen3-4B-test - # --wandb-key ${WANDB_KEY} -) - -SGLANG_ARGS=( - --rollout-num-gpus-per-engine 2 - --sglang-mem-fraction-static 0.7 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - ### AMD Support ### - # disable gradient accumulation fusion: Need to add apex to enable this - --no-gradient-accumulation-fusion - ################### -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} - -NUM_GPUS=$(echo ${HIP_VISIBLE_DEVICES} | tr ',' '\n' | wc -l) -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - - -# "PYTHONPATH": "/workspace/Megatron-LM/", -MEGATRON_LM_PATH=$(pip list | grep megatron-core | awk '{print $NF}') - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json='{ - "env_vars": { - "PYTHONPATH": "/workspace/Megatron-LM/", - "CUDA_DEVICE_MAX_CONNECTIONS": "1" - } - }' \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - --no-offload-train \ - --no-offload-rollout \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${SGLANG_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index f326b1d0d..966bc3459 100644 --- a/slime/backends/megatron_utils/model.py +++ b/slime/backends/megatron_utils/model.py @@ -884,14 +884,6 @@ def initialize_model_and_optimizer( DDP-wrapped model chunks, optimizer, scheduler, and iteration index. """ - if torch.version.hip: - import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module - - from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync - - filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync - print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility") - model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role) model[0].role = role reinit_critic_output_layer = _critic_output_layer_needs_reinit(args, model, role) diff --git a/slime/ray/train_actor.py b/slime/ray/train_actor.py index a8ba6ddc6..4ad29946d 100644 --- a/slime/ray/train_actor.py +++ b/slime/ray/train_actor.py @@ -70,21 +70,17 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): args.world_size = dist.get_world_size() try: - if torch.version.hip is not None: - logger.info("Detected ROCm/HIP environment, skipping NUMA affinity setup") - # will find the coresponding API to implement ROCm version as below - else: - import pynvml + import pynvml - pynvml.nvmlInit() + pynvml.nvmlInit() - local_rank = int(os.environ["RANK"]) % args.num_gpus_per_node + local_rank = int(os.environ["RANK"]) % args.num_gpus_per_node - handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank) - pynvml.nvmlDeviceSetCpuAffinity(handle) + handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank) + pynvml.nvmlDeviceSetCpuAffinity(handle) - logger.info(f"Set NUMA affinity for GPU {local_rank}") - pynvml.nvmlShutdown() + logger.info(f"Set NUMA affinity for GPU {local_rank}") + pynvml.nvmlShutdown() except ImportError: logger.info("Warning: pynvml not available, skipping NUMA affinity setup") diff --git a/slime/utils/rocm_checkpoint_writer.py b/slime/utils/rocm_checkpoint_writer.py deleted file mode 100644 index 7a8a1be2c..000000000 --- a/slime/utils/rocm_checkpoint_writer.py +++ /dev/null @@ -1,27 +0,0 @@ -import torch -from megatron.core.dist_checkpointing.strategies.filesystem_async import FileSystemWriterAsync - - -class ROCmFileSystemWriterAsync(FileSystemWriterAsync): - """ - FileSystemWriterAsync wrapper for ROCm compatibility. - - On ROCm/HIP, using non_blocking=True causes tensors to be stored in pinned memory, - which triggers segmentation faults when forking subprocesses afterward. - """ - - @staticmethod - def preload_tensors(*args, **kwargs): - # Change argument non_blocking to False on HIP platform - # The tensors will be stored in pinned memory if non_blocking=True - # Currently on the ROCm platform, forking a subprocess afterward - # with pinned_memory=True will trigger segmentation fault - if torch.version.hip: - print("HIP/ROCm detected: setting non_blocking=False in preload_tensors") - if "non_blocking" in kwargs: - kwargs["non_blocking"] = False - elif len(args) > 1 and isinstance(args[-1], bool): - # non_blocking is typically the last argument - args = args[:-1] + (False,) - - return FileSystemWriterAsync.preload_tensors(*args, **kwargs) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 8d2758947..d14c22fc9 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -78,13 +78,6 @@ def ceildiv(a, b): def main(): - if torch.version.hip: - import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module - from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync - - filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync - print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility") - configure_logger() # Initialize distributed environment @@ -107,10 +100,6 @@ def main(): args = get_args() init(args) - # if using AMD gpus, we have to do the conversion in cpu - if hasattr(torch.version, "hip") and torch.version.hip is not None: - assert args.use_cpu_initialization, "AMD GPU requires --use_cpu_initialization=True" - model = get_model(get_model_provider_func(args), ModelType.encoder_or_decoder, wrap_with_ddp=False) # Load model