mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-13 22:18:36 +08:00
24 lines
516 B
Bash
Executable File
24 lines
516 B
Bash
Executable File
#!/bin/bash
|
|
|
|
set -euo pipefail
|
|
|
|
CONTAINER_NAME=${CONTAINER_NAME:-layer_wise_benchmarks}
|
|
TRTLLM_ROOT=$(realpath "$(dirname -- "$0")"/../..)
|
|
CONTAINER_MOUNTS=$TRTLLM_ROOT:$TRTLLM_ROOT
|
|
|
|
if [ -z "${SLURM_JOB_ID:-}" ]; then
|
|
echo "Please set SLURM_JOB_ID"
|
|
exit 1
|
|
fi
|
|
|
|
WORKDIR=$(realpath "$(pwd)")
|
|
|
|
set -x
|
|
srun --mpi=pmix \
|
|
-N "$NODES" \
|
|
--ntasks-per-node $(($NP / $NODES)) \
|
|
--container-name "$CONTAINER_NAME" \
|
|
--container-mounts "$CONTAINER_MOUNTS" \
|
|
--container-workdir "$WORKDIR" \
|
|
"$@"
|