Fix token speed slow in machine has multi gpus (#1157)
* Update bat script windows choose GPU has highest ram to start nitro * Update bash script for linux to choose gpu has highest vram --------- Co-authored-by: Hien To <tominhhien97@gmail.com>
This commit is contained in:
parent
3c90bf8469
commit
4f93e14d16
@ -1,5 +1,38 @@
|
|||||||
#!/bin/bash
|
#!/bin/bash
|
||||||
|
|
||||||
# Attempt to run the nitro_linux_amd64_cuda file and if it fails, run nitro_linux_amd64
|
# Check if nvidia-smi exists and is executable
|
||||||
|
if ! command -v nvidia-smi &> /dev/null; then
|
||||||
|
echo "nvidia-smi not found, proceeding with CPU version..."
|
||||||
|
cd linux-cpu
|
||||||
|
./nitro "$@"
|
||||||
|
exit $?
|
||||||
|
fi
|
||||||
|
|
||||||
|
# Find the GPU with the highest VRAM
|
||||||
|
readarray -t gpus < <(nvidia-smi --query-gpu=index,memory.total --format=csv,noheader,nounits)
|
||||||
|
maxMemory=0
|
||||||
|
selectedGpuId=0
|
||||||
|
|
||||||
|
for gpu in "${gpus[@]}"; do
|
||||||
|
IFS=, read -ra gpuInfo <<< "$gpu"
|
||||||
|
gpuId=${gpuInfo[0]}
|
||||||
|
gpuMemory=${gpuInfo[1]}
|
||||||
|
if (( gpuMemory > maxMemory )); then
|
||||||
|
maxMemory=$gpuMemory
|
||||||
|
selectedGpuId=$gpuId
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
|
||||||
|
echo "Selected GPU: $selectedGpuId"
|
||||||
|
export CUDA_VISIBLE_DEVICES=$selectedGpuId
|
||||||
|
|
||||||
|
# Attempt to run nitro_linux_amd64_cuda
|
||||||
cd linux-cuda
|
cd linux-cuda
|
||||||
./nitro "$@" || (echo "nitro_linux_amd64_cuda encountered an error, attempting to run nitro_linux_amd64..." && cd ../linux-cpu && ./nitro "$@")
|
if ./nitro "$@"; then
|
||||||
|
exit $?
|
||||||
|
else
|
||||||
|
echo "nitro_linux_amd64_cuda encountered an error, attempting to run nitro_linux_amd64..."
|
||||||
|
cd ../linux-cpu
|
||||||
|
./nitro "$@"
|
||||||
|
exit $?
|
||||||
|
fi
|
||||||
@ -1,12 +1,44 @@
|
|||||||
@echo off
|
@echo off
|
||||||
|
setlocal enabledelayedexpansion
|
||||||
|
|
||||||
|
set "maxMemory=0"
|
||||||
|
set "gpuId="
|
||||||
|
|
||||||
|
rem check if nvidia-smi command exist or not
|
||||||
|
where nvidia-smi >nul 2>&1
|
||||||
|
if %errorlevel% neq 0 (
|
||||||
|
echo nvidia-smi not found, proceeding with CPU version...
|
||||||
|
cd win-cuda
|
||||||
|
goto RunCpuVersion
|
||||||
|
)
|
||||||
|
|
||||||
|
set "tempFile=%temp%\nvidia_smi_output.txt"
|
||||||
|
nvidia-smi --query-gpu=index,memory.total --format=csv,noheader,nounits > "%tempFile%"
|
||||||
|
|
||||||
|
for /f "usebackq tokens=1-2 delims=, " %%a in ("%tempFile%") do (
|
||||||
|
set /a memorySize=%%b
|
||||||
|
if !memorySize! gtr !maxMemory! (
|
||||||
|
set "maxMemory=!memorySize!"
|
||||||
|
set "gpuId=%%a"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
rem Echo the selected GPU
|
||||||
|
echo Selected GPU: !gpuId!
|
||||||
|
|
||||||
|
rem Set the GPU with the highest VRAM as the visible CUDA device
|
||||||
|
set CUDA_VISIBLE_DEVICES=!gpuId!
|
||||||
|
|
||||||
rem Attempt to run nitro_windows_amd64_cuda.exe
|
rem Attempt to run nitro_windows_amd64_cuda.exe
|
||||||
cd win-cuda
|
cd win-cuda
|
||||||
nitro.exe %*
|
nitro.exe %*
|
||||||
|
if %errorlevel% neq 0 goto RunCpuVersion
|
||||||
|
goto End
|
||||||
|
|
||||||
rem Check the exit code of the previous command
|
:RunCpuVersion
|
||||||
if %errorlevel% neq 0 (
|
rem Run nitro_windows_amd64.exe...
|
||||||
echo nitro_windows_amd64_cuda.exe encountered an error, attempting to run nitro_windows_amd64.exe...
|
cd ..\win-cpu
|
||||||
cd ..\win-cpu
|
nitro.exe %*
|
||||||
nitro.exe %*
|
|
||||||
)
|
:End
|
||||||
|
endlocal
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user