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

Upgrade tensorrt_llm to 0.15.0

Browse Source 
Signed-off-by: makaveli10 <vineet.suryan@collabora.com>
This commit is contained in:
makaveli10
2024-11-21 06:13:00 +00:00
parent 446fc6e835
commit c25a036c02
5 changed files with 218 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
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