mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-05 02:31:33 +08:00
24 lines
522 B
Bash
Executable File
24 lines
522 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" \
|
|
"$@"
|