Upgrade tensorrt_llm to 0.15.0

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>

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

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
 ```

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 .

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"
+    echo "Converting model weights for $model_name..."
-    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"
+    python3 convert_checkpoint.py \
-    echo "Whisper $model_name TensorRT engine built."
+        $( [[ "$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/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'
+        serialize_path = engine_dir / 'encoder' / 'rank0.engine'
-        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'
-
         with open(serialize_path, 'rb') as f:
             session = Session.from_serialized_engine(f.read())
-
         return session
 
-    def get_audio_features(self, mel):
+    def get_audio_features(self,
-
+                           mel,
-        input_lengths = torch.tensor(
+                           mel_input_lengths,
-            [mel.shape[2] // 2 for _ in range(mel.shape[0])],
+                           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_features'] = mel
-        inputs['input_lengths'] = input_lengths
+        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']
+        encoder_output = outputs['encoder_output']
-        return audio_features
+        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 / 'decoder' / 'rank0.engine'
-        serialize_path = engine_dir / f'whisper_decoder_{dtype}_tp1_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_heads=self.decoder_config['num_attention_heads'],
-            num_kv_heads=self.decoder_config['num_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'],
+            cross_attention=True,
-            gpt_attention_plugin=self.decoder_config['gpt_attention_plugin'],
+            num_layers=self.decoder_config['num_hidden_layers'],
-            remove_input_padding=self.decoder_config['remove_input_padding'],
+            gpt_attention_plugin=self.decoder_config['plugin_config']
-            cross_attention=self.decoder_config['cross_attention'],
+            ['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.
+            dtype=self.decoder_config['dtype'],
-            decoder_config['has_token_type_embedding'],
+            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(
+        batch_size = decoder_input_ids.shape[0]
-            [encoder_outputs.shape[1] for x in range(encoder_outputs.shape[0])],
-            dtype=torch.int32,
-            device='cuda')
-
         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(
+        cross_attention_mask = torch.ones([
-            [encoder_outputs.shape[0], 1,
+            batch_size, decoder_max_input_length + max_new_tokens,
-             encoder_outputs.shape[1]]).int().cuda()
+            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,
+                                        runtime_mapping,
-                                       debug_mode=False)
+                                        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