mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
23 lines
496 B
Bash
Executable File
23 lines
496 B
Bash
Executable File
#!/bin/bash
|
|
|
|
set -euo pipefail
|
|
|
|
CONTAINER_NAME=${CONTAINER_NAME:-layer_wise_benchmarks}
|
|
CONTAINER_MOUNTS=$(realpath "$(pwd)/../.."):$(realpath "$(pwd)/../..")
|
|
|
|
if [ "${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" \
|
|
"$@"
|