cjones/WhisperLive
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #301 from makaveli10/upgrade_tensorrt

Browse Source 
Upgrade tensorrt_llm==0.15.0.
This commit is contained in:
Marcus Edel
2024-12-03 10:04:14 -05:00
committed by GitHub
parent a55b99c11e 49e232bc4d
commit 19c05c8231
6 changed files with 220 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+6 -2
View File
@@ -133,12 +133,16 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt
# Build small.en engine
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en # float16
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8 # int8 weight only quantization
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4 # int4 weight only quantization
# Run server with small.en
python3 run_server.py --port 9090 \
--backend tensorrt \
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int8"
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
```
- CPU
TensorRT_whisper.md
+6 -10
View File
@@ -1,17 +1,11 @@
# WhisperLive-TensorRT
We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
**Note**: We use `tensorrt_llm==0.9.0`
**Note**: We use `tensorrt_llm==0.15.0.dev2024111200`
## Installation
- Install [docker](https://docs.docker.com/engine/install/)
- Install [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
- Clone this repo.
```bash
git clone https://github.com/collabora/WhisperLive.git
cd WhisperLive
```
- Run WhisperLive TensorRT in docker
```bash
docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt:latest
@@ -21,7 +15,9 @@ docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it g
- We build `small.en` and `small` multilingual TensorRT engine as examples below. The script logs the path of the directory with Whisper TensorRT engine. We need that model_path to run the server.
```bash
# convert small.en
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en # float16
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8 # int8 weight only quantization
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4 # int4 weight only quantization
# convert small multilingual model
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
@@ -32,11 +28,11 @@ bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
# Run English only model
python3 run_server.py --port 9090 \
--backend tensorrt \
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
# Run Multilingual model
python3 run_server.py --port 9090 \
--backend tensorrt \
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \
--trt_multilingual
```
docker/Dockerfile.tensorrt
+5 -5
View File
@@ -1,15 +1,16 @@
FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04 AS base
FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04 AS base
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y \
python3.10 python3-pip openmpi-bin libopenmpi-dev git wget \
python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget \
&& rm -rf /var/lib/apt/lists/*
FROM base AS devel
RUN pip3 install --no-cache-dir -U tensorrt_llm==0.10.0 --extra-index-url https://pypi.nvidia.com
RUN pip3 install --no-cache-dir -U tensorrt_llm==0.15.0.dev2024111200 --extra-index-url https://pypi.nvidia.com
WORKDIR /app
RUN git clone -b v0.10.0 --depth 1 https://github.com/NVIDIA/TensorRT-LLM.git && \
RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git && cd TensorRT-LLM && \
git checkout c629546ce429623c8a163633095230154a6f0574 && cd ../ && \
mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
rm -rf TensorRT-LLM
@@ -24,7 +25,6 @@ RUN apt update && bash setup.sh && rm setup.sh
COPY requirements/server.txt .
RUN pip install --no-cache-dir -r server.txt && rm server.txt
RUN pip install -U huggingface_hub tokenizers==0.19.0
COPY whisper_live ./whisper_live
COPY scripts/build_whisper_tensorrt.sh .
COPY run_server.py .
scripts/build_whisper_tensorrt.sh
+51 -6
View File
@@ -38,12 +38,24 @@ download_and_build_model() {
"large-v3" | "large")
model_url="https://openaipublic.azureedge.net/main/whisper/models/e5b1a55b89c1367dacf97e3e19bfd829a01529dbfdeefa8caeb59b3f1b81dadb/large-v3.pt"
;;
"large-v3-turbo" | "turbo")
model_url="https://openaipublic.azureedge.net/main/whisper/models/aff26ae408abcba5fbf8813c21e62b0941638c5f6eebfb145be0c9839262a19a/large-v3-turbo.pt"
;;
*)
echo "Invalid model name: $model_name"
exit 1
;;
esac
if [ "$model_name" == "turbo" ]; then
model_name="large-v3-turbo"
fi
local inference_precision="float16"
local weight_only_precision="${2:-float16}"
local max_beam_width=4
local max_batch_size=1
echo "Downloading $model_name..."
# wget --directory-prefix=assets "$model_url"
# echo "Download completed: ${model_name}.pt"
@@ -54,11 +66,43 @@ download_and_build_model() {
echo "${model_name}.pt already exists in assets directory."
fi
local output_dir="whisper_${model_name//./_}"
local sanitized_model_name="${model_name//./_}"
local checkpoint_dir="whisper_${sanitized_model_name}_weights_${weight_only_precision}"
local output_dir="whisper_${sanitized_model_name}_${weight_only_precision}"
echo "$output_dir"
echo "Running build script for $model_name with output directory $output_dir"
python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin --use_bert_attention_plugin --enable_context_fmha --model_name "$model_name"
echo "Whisper $model_name TensorRT engine built."
echo "Converting model weights for $model_name..."
python3 convert_checkpoint.py \
$( [[ "$weight_only_precision" == "int8" || "$weight_only_precision" == "int4" ]] && echo "--use_weight_only --weight_only_precision $weight_only_precision" ) \
--output_dir "$checkpoint_dir" --model_name "$model_name"
echo "Building encoder for $model_name..."
trtllm-build \
--checkpoint_dir "${checkpoint_dir}/encoder" \
--output_dir "${output_dir}/encoder" \
--moe_plugin disable \
--enable_xqa disable \
--max_batch_size "$max_batch_size" \
--gemm_plugin disable \
--bert_attention_plugin "$inference_precision" \
--max_input_len 3000 \
--max_seq_len 3000
echo "Building decoder for $model_name..."
trtllm-build \
--checkpoint_dir "${checkpoint_dir}/decoder" \
--output_dir "${output_dir}/decoder" \
--moe_plugin disable \
--enable_xqa disable \
--max_beam_width "$max_beam_width" \
--max_batch_size "$max_batch_size" \
--max_seq_len 200 \
--max_input_len 14 \
--max_encoder_input_len 3000 \
--gemm_plugin "$inference_precision" \
--bert_attention_plugin "$inference_precision" \
--gpt_attention_plugin "$inference_precision"
echo "TensorRT LLM engine built for $model_name."
echo "========================================="
echo "Model is located at: $(pwd)/$output_dir"
}
@@ -70,8 +114,9 @@ fi
tensorrt_examples_dir="$1"
model_name="${2:-small.en}"
weight_only_precision="${3:-float16}" # Default to float16 if not provided
cd $1/whisper
cd $tensorrt_examples_dir/whisper
pip install --no-deps -r requirements.txt
download_and_build_model "$model_name"
download_and_build_model "$model_name" "$weight_only_precision"
whisper_live/client.py
+2 -2
View File
@@ -112,9 +112,9 @@ class Client:
for i, seg in enumerate(segments):
if not text or text[-1] != seg["text"]:
text.append(seg["text"])
if i == len(segments) - 1 and not seg["completed"]:
if i == len(segments) - 1 and not seg.get("completed", False):
self.last_segment = seg
elif (self.server_backend == "faster_whisper" and seg["completed"] and
elif (self.server_backend == "faster_whisper" and seg.get("completed", False) and
(not self.transcript or
float(seg['start']) >= float(self.transcript[-1]['end']))):
self.transcript.append(seg)
whisper_live/transcriber_tensorrt.py
+150 -67
View File
@@ -1,5 +1,6 @@
import json
import re
import math
from collections import OrderedDict
from pathlib import Path
from typing import Union
@@ -14,7 +15,8 @@ import tensorrt_llm
import tensorrt_llm.logger as logger
from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt,
trt_dtype_to_torch)
from tensorrt_llm.runtime import ModelConfig, SamplingConfig
from tensorrt_llm.bindings import GptJsonConfig, KVCacheType
from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelConfig, SamplingConfig
from tensorrt_llm.runtime.session import Session, TensorInfo
@@ -24,49 +26,101 @@ HOP_LENGTH = 160
CHUNK_LENGTH = 30
N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000 samples in a 30-second chunk
def read_config(component, engine_dir):
config_path = engine_dir / component / 'config.json'
with open(config_path, 'r') as f:
config = json.load(f)
model_config = OrderedDict()
model_config.update(config['pretrained_config'])
model_config.update(config['build_config'])
return model_config
def remove_tensor_padding(input_tensor,
input_tensor_lengths=None,
pad_value=None):
if pad_value:
assert input_tensor_lengths is None, "input_tensor_lengths should be None when pad_value is provided"
# Text tensor case: batch, seq_len
assert torch.all(
input_tensor[:, 0] != pad_value
), "First token in each sequence should not be pad_value"
assert input_tensor_lengths is None
# Create a mask for all non-pad tokens
mask = input_tensor != pad_value
# Apply the mask to input_tensor to remove pad tokens
output_tensor = input_tensor[mask].view(1, -1)
else:
# Audio tensor case: batch, seq_len, feature_len
# position_ids case: batch, seq_len
assert input_tensor_lengths is not None, "input_tensor_lengths must be provided for 3D input_tensor"
# Initialize a list to collect valid sequences
valid_sequences = []
for i in range(input_tensor.shape[0]):
valid_length = input_tensor_lengths[i]
valid_sequences.append(input_tensor[i, :valid_length])
# Concatenate all valid sequences along the batch dimension
output_tensor = torch.cat(valid_sequences, dim=0)
return output_tensor
class WhisperEncoding:
def __init__(self, engine_dir):
self.session = self.get_session(engine_dir)
config = read_config('encoder', engine_dir)
self.n_mels = config['n_mels']
self.dtype = config['dtype']
self.num_languages = config['num_languages']
self.encoder_config = config
def get_session(self, engine_dir):
config_path = engine_dir / 'encoder_config.json'
with open(config_path, 'r') as f:
config = json.load(f)
use_gpt_attention_plugin = config['plugin_config'][
'gpt_attention_plugin']
dtype = config['builder_config']['precision']
n_mels = config['builder_config']['n_mels']
num_languages = config['builder_config']['num_languages']
self.dtype = dtype
self.n_mels = n_mels
self.num_languages = num_languages
serialize_path = engine_dir / f'whisper_encoder_{self.dtype}_tp1_rank0.engine'
serialize_path = engine_dir / 'encoder' / 'rank0.engine'
with open(serialize_path, 'rb') as f:
session = Session.from_serialized_engine(f.read())
return session
def get_audio_features(self, mel):
input_lengths = torch.tensor(
[mel.shape[2] // 2 for _ in range(mel.shape[0])],
def get_audio_features(self,
mel,
mel_input_lengths,
encoder_downsampling_factor=2):
if isinstance(mel, list):
longest_mel = max([f.shape[-1] for f in mel])
mel = [
torch.nn.functional.pad(f, (0, longest_mel - f.shape[-1]),
mode='constant') for f in mel
]
mel = torch.cat(mel, dim=0).type(
str_dtype_to_torch("float16")).contiguous()
bsz, seq_len = mel.shape[0], mel.shape[2]
position_ids = torch.arange(
math.ceil(seq_len / encoder_downsampling_factor),
dtype=torch.int32,
device=mel.device)
device=mel.device).expand(bsz, -1).contiguous()
if self.encoder_config['plugin_config']['remove_input_padding']:
# mel B,D,T -> B,T,D -> BxT, D
mel = mel.transpose(1, 2)
mel = remove_tensor_padding(mel, mel_input_lengths)
position_ids = remove_tensor_padding(
position_ids, mel_input_lengths // encoder_downsampling_factor)
inputs = OrderedDict()
inputs['x'] = mel
inputs['input_lengths'] = input_lengths
inputs['input_features'] = mel
inputs['input_lengths'] = mel_input_lengths
inputs['position_ids'] = position_ids
output_list = [
TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
TensorInfo('input_features', str_dtype_to_trt(self.dtype),
mel.shape),
TensorInfo('input_lengths', str_dtype_to_trt('int32'),
input_lengths.shape)
mel_input_lengths.shape),
TensorInfo('position_ids', str_dtype_to_trt('int32'),
inputs['position_ids'].shape)
]
output_info = (self.session).infer_shapes(output_list)
@@ -84,48 +138,44 @@ class WhisperEncoding:
stream=stream.cuda_stream)
assert ok, 'Engine execution failed'
stream.synchronize()
audio_features = outputs['output']
return audio_features
encoder_output = outputs['encoder_output']
encoder_output_lengths = mel_input_lengths // encoder_downsampling_factor
return encoder_output, encoder_output_lengths
class WhisperDecoding:
def __init__(self, engine_dir, runtime_mapping, debug_mode=False):
self.decoder_config = self.get_config(engine_dir)
self.decoder_config = read_config('decoder', engine_dir)
self.decoder_generation_session = self.get_session(
engine_dir, runtime_mapping, debug_mode)
def get_config(self, engine_dir):
config_path = engine_dir / 'decoder_config.json'
with open(config_path, 'r') as f:
config = json.load(f)
decoder_config = OrderedDict()
decoder_config.update(config['plugin_config'])
decoder_config.update(config['builder_config'])
return decoder_config
def get_session(self, engine_dir, runtime_mapping, debug_mode=False):
dtype = self.decoder_config['precision']
serialize_path = engine_dir / f'whisper_decoder_{dtype}_tp1_rank0.engine'
serialize_path = engine_dir / 'decoder' / 'rank0.engine'
with open(serialize_path, "rb") as f:
decoder_engine_buffer = f.read()
decoder_model_config = ModelConfig(
max_batch_size=self.decoder_config['max_batch_size'],
max_beam_width=self.decoder_config['max_beam_width'],
num_heads=self.decoder_config['num_heads'],
num_kv_heads=self.decoder_config['num_heads'],
num_heads=self.decoder_config['num_attention_heads'],
num_kv_heads=self.decoder_config['num_attention_heads'],
hidden_size=self.decoder_config['hidden_size'],
vocab_size=self.decoder_config['vocab_size'],
num_layers=self.decoder_config['num_layers'],
gpt_attention_plugin=self.decoder_config['gpt_attention_plugin'],
remove_input_padding=self.decoder_config['remove_input_padding'],
cross_attention=self.decoder_config['cross_attention'],
cross_attention=True,
num_layers=self.decoder_config['num_hidden_layers'],
gpt_attention_plugin=self.decoder_config['plugin_config']
['gpt_attention_plugin'],
remove_input_padding=self.decoder_config['plugin_config']
['remove_input_padding'],
kv_cache_type=KVCacheType.PAGED
if self.decoder_config['plugin_config']['paged_kv_cache'] == True
else KVCacheType.CONTINUOUS,
has_position_embedding=self.
decoder_config['has_position_embedding'],
has_token_type_embedding=self.
decoder_config['has_token_type_embedding'],
dtype=self.decoder_config['dtype'],
has_token_type_embedding=False,
)
decoder_generation_session = tensorrt_llm.runtime.GenerationSession(
decoder_model_config,
@@ -138,14 +188,12 @@ class WhisperDecoding:
def generate(self,
decoder_input_ids,
encoder_outputs,
encoder_max_input_length,
encoder_input_lengths,
eot_id,
max_new_tokens=40,
num_beams=1):
encoder_input_lengths = torch.tensor(
[encoder_outputs.shape[1] for x in range(encoder_outputs.shape[0])],
dtype=torch.int32,
device='cuda')
batch_size = decoder_input_ids.shape[0]
decoder_input_lengths = torch.tensor([
decoder_input_ids.shape[-1]
for _ in range(decoder_input_ids.shape[0])
@@ -154,10 +202,10 @@ class WhisperDecoding:
device='cuda')
decoder_max_input_length = torch.max(decoder_input_lengths).item()
cross_attention_mask = torch.ones(
[encoder_outputs.shape[0], 1,
encoder_outputs.shape[1]]).int().cuda()
cross_attention_mask = torch.ones([
batch_size, decoder_max_input_length + max_new_tokens,
encoder_max_input_length
]).int().cuda()
# generation config
sampling_config = SamplingConfig(end_id=eot_id,
pad_id=eot_id,
@@ -167,11 +215,24 @@ class WhisperDecoding:
decoder_max_input_length,
max_new_tokens,
beam_width=num_beams,
encoder_max_input_length=encoder_outputs.shape[1])
encoder_max_input_length=encoder_max_input_length)
torch.cuda.synchronize()
decoder_input_ids = decoder_input_ids.type(torch.int32).cuda()
if self.decoder_config['plugin_config']['remove_input_padding']:
# 50256 is the index of <pad> for all whisper models' decoder
WHISPER_PAD_TOKEN_ID = 50256
decoder_input_ids = remove_tensor_padding(
decoder_input_ids, pad_value=WHISPER_PAD_TOKEN_ID)
if encoder_outputs.dim() == 3:
encoder_output_lens = torch.full((encoder_outputs.shape[0], ),
encoder_outputs.shape[1],
dtype=torch.int32,
device='cuda')
encoder_outputs = remove_tensor_padding(encoder_outputs,
encoder_output_lens)
output_ids = self.decoder_generation_session.decode(
decoder_input_ids,
decoder_input_lengths,
@@ -196,18 +257,23 @@ class WhisperTRTLLM(object):
runtime_mapping = tensorrt_llm.Mapping(world_size, runtime_rank)
torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node)
engine_dir = Path(engine_dir)
encoder_config = read_config('encoder', engine_dir)
decoder_config = read_config('decoder', engine_dir)
self.n_mels = encoder_config['n_mels']
self.num_languages = encoder_config['num_languages']
is_multilingual = (decoder_config['vocab_size'] >= 51865)
self.encoder = WhisperEncoding(engine_dir)
self.decoder = WhisperDecoding(engine_dir,
runtime_mapping,
debug_mode=False)
runtime_mapping,
debug_mode=False)
self.n_mels = self.encoder.n_mels
# self.tokenizer = get_tokenizer(num_languages=self.encoder.num_languages,
# tokenizer_dir=assets_dir)
self.device = device
self.tokenizer = get_tokenizer(
is_multilingual,
num_languages=self.encoder.num_languages,
num_languages=self.num_languages,
language=language,
task=task,
)
@@ -274,8 +340,10 @@ class WhisperTRTLLM(object):
def process_batch(
self,
mel,
mel_input_lengths,
text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
num_beams=1):
num_beams=1,
max_new_tokens=96):
prompt_id = self.tokenizer.encode(
text_prefix, allowed_special=set(self.tokenizer.special_tokens.keys()))
@@ -283,11 +351,14 @@ class WhisperTRTLLM(object):
batch_size = mel.shape[0]
decoder_input_ids = prompt_id.repeat(batch_size, 1)
encoder_output = self.encoder.get_audio_features(mel)
encoder_output, encoder_output_lengths = self.encoder.get_audio_features(mel, mel_input_lengths)
encoder_max_input_length = torch.max(encoder_output_lengths).item()
output_ids = self.decoder.generate(decoder_input_ids,
encoder_output,
encoder_max_input_length,
encoder_output_lengths,
self.tokenizer.eot,
max_new_tokens=96,
max_new_tokens=max_new_tokens,
num_beams=num_beams)
texts = []
for i in range(len(output_ids)):
@@ -302,10 +373,22 @@ class WhisperTRTLLM(object):
dtype='float16',
batch_size=1,
num_beams=1,
padding_strategy="max",
):
mel = mel.type(str_dtype_to_torch(dtype))
mel = mel.unsqueeze(0)
predictions = self.process_batch(mel, text_prefix, num_beams)
# repeat the mel spectrogram to match the batch size
mel = mel.repeat(batch_size, 1, 1)
if padding_strategy == "longest":
pass
else:
mel = torch.nn.functional.pad(mel, (0, 3000 - mel.shape[2]))
features_input_lengths = torch.full((mel.shape[0], ),
mel.shape[2],
dtype=torch.int32,
device=mel.device)
predictions = self.process_batch(mel, features_input_lengths, text_prefix, num_beams)
prediction = predictions[0]
# remove all special tokens in the prediction