mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
a5eff139f1
Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com> Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com> Co-authored-by: Erin Ho <14718778+hchings@users.noreply.github.com>
58 lines
2.1 KiB
Bash
58 lines
2.1 KiB
Bash
#!/bin/bash
|
|
#SBATCH -A <account>
|
|
#SBATCH -p <partition>
|
|
#SBATCH -t 01:00:00
|
|
#SBATCH -N 2
|
|
#SBATCH --ntasks-per-node=8
|
|
#SBATCH -o logs/trtllm-serve.out
|
|
#SBATCH -e logs/trtllm-serve.err
|
|
#SBATCH -J trtllm-serve
|
|
|
|
### :title Run trtllm-serve with pytorch backend on Slurm
|
|
### :order 2
|
|
### :section Slurm
|
|
|
|
# NOTE, this feature is experimental and may not work on all systems.
|
|
# The trtllm-llmapi-launch is a script that launches the LLM-API code on
|
|
# Slurm-like systems, and can support multi-node and multi-GPU setups.
|
|
|
|
# Note that, the number of MPI processes should be the same as the model world
|
|
# size. e.g. For tensor_parallel_size=16, you may use 2 nodes with 8 gpus for
|
|
# each, or 4 nodes with 4 gpus for each or other combinations.
|
|
|
|
# This docker image should have tensorrt_llm installed, or you need to install
|
|
# it in the task.
|
|
|
|
# The following variables are expected to be set in the environment:
|
|
# You can set them via --export in the srun/sbatch command.
|
|
# CONTAINER_IMAGE: the docker image to use, you'd better install tensorrt_llm in it, or install it in the task.
|
|
# MOUNT_DIR: the directory to mount in the container
|
|
# MOUNT_DEST: the destination directory in the container
|
|
# WORKDIR: the working directory in the container
|
|
# SOURCE_ROOT: the path to the TensorRT-LLM source
|
|
# PROLOGUE: the prologue to run before the script
|
|
# LOCAL_MODEL: the local model directory to use, NOTE: downloading from HF is
|
|
# not supported in Slurm mode, you need to download the model and put it in
|
|
# the LOCAL_MODEL directory.
|
|
|
|
echo "Starting trtllm-serve..."
|
|
# Just launch trtllm-serve job with trtllm-llmapi-launch command.
|
|
srun -l \
|
|
--container-image=${CONTAINER_IMAGE} \
|
|
--container-mounts=${MOUNT_DIR}:${MOUNT_DEST} \
|
|
--container-workdir=${WORKDIR} \
|
|
--export=ALL,PYTHONPATH=${SOURCE_ROOT} \
|
|
--mpi=pmix \
|
|
bash -c "
|
|
set -ex
|
|
$PROLOGUE
|
|
export PATH=$PATH:~/.local/bin
|
|
|
|
trtllm-llmapi-launch \
|
|
trtllm-serve $LOCAL_MODEL \
|
|
--tp_size 16 \
|
|
--backend pytorch \
|
|
--host 0.0.0.0 \
|
|
${ADDITIONAL_OPTIONS}
|
|
"
|