A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 12:45 UTC
[ROCm][CI] Add rocm support for run-multi-node-test.sh (#31922)
Signed-off-by: charlifu <charlifu@amd.com> Signed-off-by: Charlie Fu <Charlie.Fu@amd.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Charlie Fu committed
Jan 7, 2026 at 22:36 UTC
cddbc2b4b2547c681d1bdb876fdd6a7b8e0ec58d
1 file changed
+21
-3
.buildkite/scripts/run-multi-node-test.sh
+21
-3
@@ -2,6 +2,17 @@
2
3
set -euox pipefail
4
5
+# To detect ROCm
6
+# Check multiple indicators:
7
+if [ -e /dev/kfd ] || \
8
+ [ -d /opt/rocm ] || \
9
+ command -v rocm-smi &> /dev/null || \
10
+ [ -n "$ROCM_HOME" ]; then
11
+ IS_ROCM=1
12
+else
13
+ IS_ROCM=0
14
+fi
15
+
16
if [[ $# -lt 4 ]]; then
17
echo "Usage: .buildkite/scripts/run-multi-node-test.sh WORKING_DIR NUM_NODES NUM_GPUS DOCKER_IMAGE COMMAND1 COMMAND2 ... COMMANDN"
18
exit 1
@@ -26,13 +37,18 @@ for command in "${COMMANDS[@]}"; do
37
echo "$command"
38
done
39
40
+
41
start_network() {
42
docker network create --subnet=192.168.10.0/24 docker-net
43
}
44
45
start_nodes() {
46
for node in $(seq 0 $(($NUM_NODES-1))); do
35
- GPU_DEVICES='"device='
47
+ if [ "$IS_ROCM" -eq 1 ]; then
48
+ GPU_DEVICES='--device /dev/kfd --device /dev/dri -e HIP_VISIBLE_DEVICES='
49
+ else
50
+ GPU_DEVICES='--gpus "device='
51
+ fi
52
for node_gpu in $(seq 0 $(($NUM_GPUS - 1))); do
53
DEVICE_NUM=$(($node * $NUM_GPUS + $node_gpu))
54
GPU_DEVICES+=$(($DEVICE_NUM))
@@ -40,7 +56,9 @@ start_nodes() {
56
GPU_DEVICES+=','
57
fi
58
done
43
- GPU_DEVICES+='"'
59
+ if [ "$IS_ROCM" -eq 0 ]; then
60
+ GPU_DEVICES+='"'
61
+ fi
62
63
# start the container in detached mode
64
# things to note:
@@ -49,7 +67,7 @@ start_nodes() {
67
# 3. map the huggingface cache directory to the container
68
# 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
69
# starting from 192.168.10.11)
52
- docker run -d --gpus "$GPU_DEVICES" --shm-size=10.24gb -e HF_TOKEN \
70
+ docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \
71
-v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \
72
--network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
73
/bin/bash -c "tail -f /dev/null"