trending/miles
1
0
Fork 0
mirror of https://github.com/radixark/miles.git synced 2026-10-02 07:14:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Hardware] AMD Dockerfile update and test (#158)

Browse Source 
* update the dockerfile

* update the doc

* revise aiter and sglang version

* update the dockerfile

* update

* update the training script

* update docker/Dockerfile.rocm

* update

* support llama3

* update

* update

* del model_convert.sh

* add the converted model resource

* update dockerfile

* update doc

* update

* formate

* add the docker/dockerfile supportted commit ID - Need to fix commit ID:c22f55b arised issue on AMD GPU in future
This commit is contained in:
Ethan (Yusheng) Su
2025-08-23 09:09:55 +08:00
committed by GitHub
parent 68f040a86e
commit 36711aa668
7 changed files with 483 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/Dockerfile.rocm
+15 -3
View File
@@ -1,6 +1,16 @@
#### Use the base image
# The Docker image built with this 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 "yushengsuthu/verl:verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4"
# 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"]
@@ -23,8 +33,10 @@ WORKDIR /workspace/
RUN rm -rf TransformerEngine
RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git
WORKDIR /workspace/TransformerEngine
# git checkout bb061ade
# git checkout 864405c
RUN git checkout 236178e5
# RUN git checkout bb061ade
# RUN git checkout 864405c
ENV NVTE_FRAMEWORK=pytorch
ENV NVTE_ROCM_ARCH=gfx942
docs/en/amd_tutorial.md
+6 -4
View File
@@ -14,15 +14,15 @@ If you are running Slime on AMD's Instinct, please refer to the following materi
You can download the prebuilt image from DockerHub: [rlsys/slime](https://hub.docker.com/r/rlsys/slime/tags).
```bash
docker pull rlsys/slime:slime_ubuntu22.04_rocm6.3.4-patch-numa-patch_sglang0.4.9_megatron-core-patch_ray2.47.1_apex_torch-memory-saver0.0.8-patch
docker pull 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
```
Or you can use the [Dockerfile.rocm](docker/Dockerfile.rocm) to build it on your side.
```bash
cd docker
docker build -f Dockerfile.rocm -t slime_ubuntu22.04_rocm6.3.4-patch-numa-patch_sglang0.4.9_megatron-core-patch_ray2.47.1_apex_torch-memory-saver0.0.8-patch .
docker build -f Dockerfile.rocm -t 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 .
```
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 to support virtual memory management on MI300X.
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
@@ -46,7 +46,7 @@ docker run --rm -it \
--ulimit memlock=-1 \
--ulimit stack=67108864 \
-w $PWD \
rlsys/slime:slime_ubuntu22.04_rocm6.3.4-patch-numa-patch_sglang0.4.9_megatron-core-patch_ray2.47.1_apex_torch-memory-saver0.0.8-patch \
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 \
/bin/bash
```
@@ -76,6 +76,8 @@ PYTHONPATH=${MEGATRON_LM_PATH} python tools/convert_hf_to_torch_dist.py \
--save model/Qwen3-4B_torch_dist
```
Note: You might encounter some issue in the current model convert script on AMD GPUs. You can go [here](https://huggingface.co/zyzshishui0627/models) to dowload the converted models.
⚠️ If you encounter an issue where slime cannot be found, please run `pip install -e .` in the slime directory.
install_megatron_core.sh
+30
View File
@@ -0,0 +1,30 @@
#!/bin/bash
# Script to install megatron-core from AMD version
# This script converts the Docker RUN command into a standalone shell script
set -e # Exit on any error
echo "Starting megatron-core installation..."
# Uninstall existing megatron-core if present
echo "Uninstalling existing megatron-core..."
pip uninstall -y megatron-core || true
# Clone the AMD version of Megatron-LM
echo "Cloning Megatron-LM AMD version..."
git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git
# Change to the cloned directory
echo "Changing to Megatron-LM-amd_version directory..."
cd Megatron-LM-amd_version
# Install the package in editable mode with verbose output
echo "Installing megatron-core in editable mode..."
pip install -vvv -e .
# Return to the original directory (equivalent to cd /workspace/ in Docker)
echo "Returning to original directory..."
cd ..
echo "megatron-core installation completed successfully!"
scripts/models/llama3.2-3B-Instruct-amd.sh
+19
View File
@@ -0,0 +1,19 @@
MODEL_ARGS=(
--swiglu
--num-layers 28
--hidden-size 3072
--ffn-hidden-size 8192
--num-attention-heads 24
--group-query-attention
--num-query-groups 8
--max-position-embeddings 131072
--use-rotary-position-embeddings
--disable-bias-linear
--normalization "RMSNorm"
--norm-epsilon 1e-5
--rotary-base 500000
--vocab-size 128256
--kv-channels 128
--use-rope-scaling
--rotary-scaling-factor 32.0
)
scripts/run-llama3.2-3B-Instruct-amd.sh
+183
View File
@@ -0,0 +1,183 @@
#!/bin/bash
# hf download meta-llama/Llama-3.2-3B-Instruct --local-dir /root/Llama-3.2-3B-Instruct
# 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="/home/yushensu/projects/slime" # Need to change to your own path
export SLIME_DIR=$SLIME_DIR
MODEL_DIR="/home/yushensu/projects/model" # Need to change to your own path
export MODEL_DIR=$MODEL_DIR
DATA_DIR="/home/yushensu/projects/data" # Need to change to your own path
export DATA_DIR=$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
# NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l)
# if [ "$NVLINK_COUNT" -gt 0 ]; then
# HAS_NVLINK=1
# else
# HAS_NVLINK=0
# fi
# echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
source "${SCRIPT_DIR}/models/llama3.2-3B-Instruct-amd.sh"
CKPT_ARGS=(
--hf-checkpoint ${MODEL_DIR}/Llama-3.2-3B-Instruct
--ref-load ${MODEL_DIR}/Llama-3.2-3B-Instruct_torch_dist
--load ${MODEL_DIR}/Llama-3.2-3B-Instruct_slime/
--save ${MODEL_DIR}/Llama-3.2-3B-Instruct_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 math
--num-epoch 1
--rollout-batch-size 32
--n-samples-per-prompt 8
--rollout-max-response-len 16384
--rollout-temperature 0.8
--global-batch-size 256
--balance-data
)
EVAL_ARGS=(
--eval-interval 10
--eval-prompt-data aime ${DATA_DIR}/aime-2024/aime-2024.jsonl
--n-samples-per-eval-prompt 8
--eval-max-response-len 16384
--eval-top-p 0.7
)
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 llama3.2-3B
# --wandb-key ${WANDB_API_KEY}
)
SGLANG_ARGS=(
--rollout-num-gpus-per-engine 2
--sglang-mem-fraction-static 0.4
)
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
# Build the runtime environment JSON with proper variable substitution
RUNTIME_ENV_JSON="{
\"env_vars\": {
\"PYTHONPATH\": \"/workspace/Megatron-LM-amd_version/\",
\"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\"
}
}"
ray job submit --address="http://127.0.0.1:8265" \
--runtime-env-json="${RUNTIME_ENV_JSON}" \
-- 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
scripts/run-qwen3-4B-amd.sh
+5 -5
View File
@@ -57,6 +57,7 @@ export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-"0,1,2,3,4,5,6,7"} #You can ch
# will prevent ray from buffering stdout/stderr
export PYTHONBUFFERED=16
# Current Model convert script on AMD GPU has some issue, please download the converted model from here: https://huggingface.co/zyzshishui0627/models
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
source "${SCRIPT_DIR}/models/qwen3-4B.sh"
@@ -64,7 +65,8 @@ source "${SCRIPT_DIR}/models/qwen3-4B.sh"
CKPT_ARGS=(
--hf-checkpoint ${MODEL_DIR}/Qwen3-4B
#--hf-checkpoint /root/Qwen3-4B-FP8
--ref-load ${MODEL_DIR}/Qwen3-4B_torch
--ref-load ${MODEL_DIR}/Qwen3-4B_torch_dist
# --ref-load ${MODEL_DIR}/Qwen3-4B_torch_dist_amd_new
--load ${MODEL_DIR}/Qwen3-4B_slime/
--save ${MODEL_DIR}/Qwen3-4B_slime/
--save-interval 20
@@ -76,9 +78,7 @@ ROLLOUT_ARGS=(
--label-key label
--apply-chat-template
--rollout-shuffle
--rm-type deepscaler
--num-rollout 3000
--rollout-batch-size 32
--n-samples-per-prompt 8
@@ -148,7 +148,7 @@ WANDB_ARGS=(
# )
SGLANG_ARGS=(
--rollout-num-gpus-per-engine 2
--sglang-mem-fraction-static 0.4
--sglang-mem-fraction-static 0.7
)
####################
@@ -164,7 +164,7 @@ MISC_ARGS=(
--attention-backend flash
### AMD Support ###
# disable gradient accumulation fusion: Need to add apex to enable this
# --no-gradient-accumulation-fusion
--no-gradient-accumulation-fusion
###################
)
scripts/run-qwen3-8B-amd.sh
+225
View File
@@ -0,0 +1,225 @@
#!/bin/bash
# bash scripts/run-qwen3-4B-amd.sh
####clear before training
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="/home/yushensu/projects/slime" # Need to change to your own path
export SLIME_DIR=$SLIME_DIR
MODEL_DIR="/home/yushensu/projects/model" # Need to change to your own path
export MODEL_DIR=$MODEL_DIR
DATA_DIR="/home/yushensu/projects/data" # Need to change to your own path
export DATA_DIR=$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
####################
# ### AMD Support ### (If you do not istall, please install them)
# # # Clone and install Megatron-LMi-amd_version
# export MAX_JOBS=512
# cd $SLIME_DIR
# pip uninstall megatron-core -y
# if [ ! -d "Megatron-LM-amd_version" ]; then
# git clone git@github.com:yushengsu-thu/Megatron-LM-amd_version.git
# else
# echo "Megatron-LM-amd_version directory already exists, skipping clone"
# fi
# cd Megatron-LM-amd_version
# pip install -vvv -e .
# cd $SLIME_DIR
# # Install slime
# pip install -e .
# ####################
# will prevent ray from buffering stdout/stderr
export PYTHONBUFFERED=16
# Current Model convert script on AMD GPU has some issue, please download the converted model from here: https://huggingface.co/zyzshishui0627/models
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
source "${SCRIPT_DIR}/models/qwen3-8B.sh"
CKPT_ARGS=(
--hf-checkpoint ${MODEL_DIR}/Qwen3-8B
#--hf-checkpoint /root/Qwen3-4B-FP8
--ref-load ${MODEL_DIR}/Qwen3-8B_torch_dist
# --ref-load ${MODEL_DIR}/Qwen3-8B_torch_dist_amd_new
--load ${MODEL_DIR}/Qwen3-8B_slime/
--save ${MODEL_DIR}/Qwen3-8B_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 0.8
--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 0.7
)
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}
)
### AMD Support ###
# Need to fix some issue with torch_memory_saver in rocm to support larger --sglang-mem-fraction-static
# SGLANG_ARGS=(
# --rollout-num-gpus-per-engine 2
# --sglang-mem-fraction-static 0.7
# )
SGLANG_ARGS=(
--rollout-num-gpus-per-engine 2
--sglang-mem-fraction-static 0.4
)
####################
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-amd_version/",
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-amd_version/",
"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[@]} \
${DISTRIBUTED_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