mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-16 15:55:08 +08:00
44 lines
1.3 KiB
Bash
44 lines
1.3 KiB
Bash
#! /bin/bash
|
|
set -u
|
|
set -e
|
|
set -x
|
|
|
|
role=${1}
|
|
instance_id=${2}
|
|
model_path=${3}
|
|
port=${4}
|
|
numa_bind=${5}
|
|
log_dir=${6}
|
|
enable_nsys=${7}
|
|
config_file=${8}
|
|
|
|
# Clear UCX_TLS for specific clusters
|
|
unset UCX_TLS
|
|
|
|
echo "SLURM_PROCID: ${SLURM_PROCID}, hostname: $(hostname), instance_id: ${instance_id}"
|
|
echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
|
|
|
|
if [ "${numa_bind}" = "true" ]; then
|
|
numa_bind_cmd="numactl -m 0,1"
|
|
echo "numactl -m 0,1 - Only allocate memory from nodes on GB200/GB300 NVL72"
|
|
else
|
|
numa_bind_cmd=""
|
|
echo "Not binding memory. If on GB200/GB300 NVL72, use \"numactl -m 0,1\" to only allocate memory from nodes."
|
|
fi
|
|
|
|
echo "config_file: ${config_file}"
|
|
|
|
nsys_prefix=""
|
|
if [ "${enable_nsys}" != "true" ]; then
|
|
echo "nsys is not enabled, start normal flow"
|
|
else
|
|
nsys_file=${log_dir}/nsys_worker_proc_${role}_${instance_id}_${SLURM_PROCID}
|
|
echo "nsys is enabled on ${role} GPUs, TLLM_PROFILE_START_STOP=${TLLM_PROFILE_START_STOP}"
|
|
nsys_prefix="nsys profile -o ${nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=none"
|
|
fi
|
|
|
|
${nsys_prefix} trtllm-llmapi-launch ${numa_bind_cmd} \
|
|
trtllm-serve ${model_path} \
|
|
--host $(hostname) --port ${port} \
|
|
--config ${config_file}
|