Mistral with TensorRT-LLM
Build an optimized inference engine for Mistral
This configuration builds an inference engine to serve Mistral 7B on an H100 GPU. It is very similar to the configuration for any other Mistral model, including fine-tuned variants.
Setup
See the end-to-end engine builder tutorial prerequisites for full setup instructions.
Make sure you have accessed the gated model on Hugging Face and set your hf_access_token in your Baseten workspace secrets.
pip install --upgrade truss
truss init mistral-7b-trt-llm
cd mistral-7b-trt-llm
rm model/model.py
Configuration
This configuration is optimized for low latency, with a batch size of 8. It applies post-training quantization to fp8 for further speed gains.
model_name: Mistral Engine 7B v0.2 6
python_version: py39
resources:
accelerator: H100:1
use_gpu: true
secrets:
hf_access_token: "set token in baseten workspace"
trt_llm:
build:
base_model: mistral
checkpoint_repository:
repo: mistralai/Mistral-7B-Instruct-v0.2
source: HF
max_batch_size: 8
max_beam_width: 1
max_input_len: 7168
max_output_len: 1024
num_builder_gpus: 1
quantization_type: fp8
tensor_parallel_count: 1
Deployment
truss push --publish --trusted
Usage
import requests
import os
# Model ID for production deployment
model_id = ""
# Read secrets from environment variables
baseten_api_key = os.environ["BASETEN_API_KEY"]
# Call model endpoint
resp = requests.post(
f"https://model-{model_id}.api.baseten.co/production/predict",
headers={"Authorization": f"Api-Key {baseten_api_key}"},
json={
"messages": [{"role": "user", "content": "How awesome is TensorRT-LLM?"}],
"max_tokens": 1024
},
stream=True
)
# Print the generated tokens as they get streamed
for content in resp.iter_content():
print(content.decode("utf-8"), end="", flush=True)
The input text prompt to guide the language model’s generation.
One of prompt XOR messages is required.
A list of dictionaries representing the message history, typically used in conversational contexts.
One of prompt XOR messages is required.
The maximum number of tokens to generate in the output. Controls the length of the generated text.
The number of beams used in beam search. Maximum of 1.
A penalty applied to repeated tokens to discourage the model from repeating the same words or phrases.
A penalty applied to tokens already present in the prompt to encourage the generation of new topics.
Controls the randomness of the output. Lower values make the output more deterministic, while higher values increase randomness.
A penalty applied to the length of the generated sequence to control verbosity. Higher values make the model favor shorter outputs.
The token ID that indicates the end of the generated sequence.
The token ID used for padding sequences to a uniform length.
Limits the sampling pool to the top k tokens, ensuring the model only considers the most likely tokens at each step.
Applies nucleus sampling to limit the sampling pool to a cumulative probability p, ensuring only the most likely tokens are considered.