add the codes to check AMD Instinct GPU number (#22367)

Signed-off-by: Zhang Jason <ning.zhang2@amd.com>
This commit is contained in:
Zhang Jason 2025-08-06 23:58:38 +08:00 committed by GitHub
parent 2cb6ef8996
commit b4b9813b5e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -21,8 +21,14 @@ check_hf_token() {
}
check_num_gpus() {
# can you check if the number of GPUs are >=2 via nvidia-smi?
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
# can you check if the number of GPUs are >=2 via nvidia-smi/rocm-smi?
which rocm-smi > /dev/null 2>&1
if [ $? -ne 0 ]; then
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
else
num_gpus=$(rocm-smi --showid | grep Instinct | wc -l)
fi
if [ "$num_gpus" -lt 2 ]; then
echo "You need at least 2 GPUs to run disaggregated prefill."
exit 1