TensorRT-LLMs/examples/layer_wise_benchmarks/mpi_launch.sh
Tailing Yuan a7fe043b13
[None][feat] Layer-wise benchmarks: support TEP balance, polish slurm scripts (#10237)
Signed-off-by: Tailing Yuan <yuantailing@gmail.com>
2026-01-05 11:23:04 +08:00

15 lines
312 B
Bash
Executable File

#!/bin/bash
set -euo pipefail
# Clear slurm envs
unset $(env | awk -F'=' '{print $1}' | (grep -E "SLURM_|SLURMD_|slurm_|MPI_|PMIX_" || true))
extra_args=
if [ -v TLLM_AUTOTUNER_CACHE_PATH ]; then
extra_args+="-x TLLM_AUTOTUNER_CACHE_PATH"
fi
set -x
mpirun --allow-run-as-root --np ${NP} $extra_args "$@"