diff --git a/README.md b/README.md index 19bb230..7475df6 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/TensorRT_whisper.md b/TensorRT_whisper.md index 8baf0bb..56a0318 100644 --- a/TensorRT_whisper.md +++ b/TensorRT_whisper.md @@ -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 ``` diff --git a/docker/Dockerfile.tensorrt b/docker/Dockerfile.tensorrt index 3607ab3..306b015 100644 --- a/docker/Dockerfile.tensorrt +++ b/docker/Dockerfile.tensorrt @@ -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 --depth 1 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 . \ No newline at end of file diff --git a/scripts/build_whisper_tensorrt.sh b/scripts/build_whisper_tensorrt.sh index c4d4856..0b5c22c 100644 --- a/scripts/build_whisper_tensorrt.sh +++ b/scripts/build_whisper_tensorrt.sh @@ -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" diff --git a/whisper_live/transcriber_tensorrt.py b/whisper_live/transcriber_tensorrt.py index 8f427fc..8742e69 100644 --- a/whisper_live/transcriber_tensorrt.py +++ b/whisper_live/transcriber_tensorrt.py @@ -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 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