diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c42b663..311d8bb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -99,6 +99,35 @@ jobs: push: true tags: ghcr.io/collabora/whisperlive-cpu:latest + build-and-push-docker-tensorrt: + needs: [run-tests, check-code-format] + timeout-minutes: 20 + runs-on: ubuntu-22.04 + if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/')) + steps: + - uses: actions/checkout@v2 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@v1 + with: + registry: ghcr.io + username: ${{ github.repository_owner }} + password: ${{ secrets.GHCR_TOKEN }} + + - name: Docker Prune + run: docker system prune -af + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v1 + + - name: Build and push Docker GPU image + uses: docker/build-push-action@v2 + with: + context: . + file: docker/Dockerfile.tensorrt + push: true + tags: ghcr.io/collabora/whisperlive-tensorrt:latest + build-and-push-docker-gpu: needs: [run-tests, check-code-format, build-and-push-docker-cpu] timeout-minutes: 20 diff --git a/README.md b/README.md index 53de8d6..6637309 100644 --- a/README.md +++ b/README.md @@ -115,7 +115,18 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu:latest ``` - - TensorRT. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) in order to setup docker and use TensorRT backend. We provide a pre-built docker image which has TensorRT-LLM built and ready to use. + - TensorRT. + ```bash + docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt + + # Build tiny.en engine + bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en + + # Run server with tiny.en + python3 run_server.py --port 9090 \ + --backend tensorrt \ + --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en" + ``` - CPU ```bash diff --git a/TensorRT_whisper.md b/TensorRT_whisper.md index 1bc303f..8baf0bb 100644 --- a/TensorRT_whisper.md +++ b/TensorRT_whisper.md @@ -1,6 +1,6 @@ -# Whisper-TensorRT +# WhisperLive-TensorRT We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup. -**Note**: We use [our fork to setup TensorRT](https://github.com/makaveli10/TensorRT-LLM) +**Note**: We use `tensorrt_llm==0.9.0` ## Installation - Install [docker](https://docs.docker.com/engine/install/) @@ -12,56 +12,31 @@ git clone https://github.com/collabora/WhisperLive.git cd WhisperLive ``` -- Pull the TensorRT-LLM docker image which we prebuilt for WhisperLive TensorRT backend. +- Run WhisperLive TensorRT in docker ```bash -docker pull ghcr.io/collabora/whisperbot-base:latest +docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt:latest ``` -- Next, we run the docker image and mount WhisperLive repo to the containers `/home` directory. -```bash -docker run -it --gpus all --shm-size=8g \ - --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \ - -p 9090:9090 -v /path/to/WhisperLive:/home/WhisperLive \ - ghcr.io/collabora/whisperbot-base:latest -``` - -- Make sure to test the installation. -```bash -# export ENV=${ENV:-/etc/shinit_v2} -# source $ENV -python -c "import torch; import tensorrt; import tensorrt_llm" -``` -**NOTE**: Uncomment and update library paths if imports fail. - ## Whisper TensorRT Engine -- We build `small.en` and `small` multilingual TensorRT engine. The script logs the path of the directory with Whisper TensorRT engine. We need the model_path to run the server. +- 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 scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small.en +bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en # convert small multilingual model -bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small +bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small ``` ## Run WhisperLive Server with TensorRT Backend ```bash -cd /home/WhisperLive - -# Install requirements -apt update && bash scripts/setup.sh -pip install -r requirements/server.txt - -# Required to create mel spectogram -wget --directory-prefix=assets assets/mel_filters.npz https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz - # Run English only model python3 run_server.py --port 9090 \ --backend tensorrt \ - --trt_model_path "path/to/whisper_trt/from/build/step" + --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en" # Run Multilingual model python3 run_server.py --port 9090 \ --backend tensorrt \ - --trt_model_path "path/to/whisper_trt/from/build/step" \ + --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \ --trt_multilingual ``` diff --git a/docker/Dockerfile.tensorrt b/docker/Dockerfile.tensorrt new file mode 100644 index 0000000..de6027e --- /dev/null +++ b/docker/Dockerfile.tensorrt @@ -0,0 +1,28 @@ +FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04 + +ARG DEBIAN_FRONTEND=noninteractive + +RUN apt-get update && apt-get install -y \ + python3.10 python3-pip openmpi-bin libopenmpi-dev git wget \ + && rm -rf /var/lib/apt/lists/* + +RUN pip3 install --no-cache-dir -U tensorrt_llm==0.9.0 --extra-index-url https://pypi.nvidia.com + +WORKDIR /app + +RUN git clone -b v0.9.0 --depth 1 https://github.com/NVIDIA/TensorRT-LLM.git && \ + mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \ + rm -rf TensorRT-LLM + +COPY assets/ ./assets +RUN wget -nc -P assets/ https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz + +COPY scripts/setup.sh ./ +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 + +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 9824803..c4d4856 100644 --- a/scripts/build_whisper_tensorrt.sh +++ b/scripts/build_whisper_tensorrt.sh @@ -57,7 +57,7 @@ download_and_build_model() { local output_dir="whisper_${model_name//./_}" 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 --model_name "$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" echo "Whisper $model_name TensorRT engine built." echo "=========================================" echo "Model is located at: $(pwd)/$output_dir" diff --git a/whisper_live/transcriber_tensorrt.py b/whisper_live/transcriber_tensorrt.py index aaa8cc1..8f427fc 100644 --- a/whisper_live/transcriber_tensorrt.py +++ b/whisper_live/transcriber_tensorrt.py @@ -35,6 +35,8 @@ class WhisperEncoding: 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'] @@ -51,12 +53,21 @@ class WhisperEncoding: return session def get_audio_features(self, mel): - inputs = OrderedDict() - output_list = [] - inputs.update({'x': mel}) - output_list.append( - TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape)) + input_lengths = torch.tensor( + [mel.shape[2] // 2 for _ in range(mel.shape[0])], + dtype=torch.int32, + device=mel.device) + + inputs = OrderedDict() + inputs['x'] = mel + inputs['input_lengths'] = input_lengths + + output_list = [ + TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape), + TensorInfo('input_lengths', str_dtype_to_trt('int32'), + input_lengths.shape) + ] output_info = (self.session).infer_shapes(output_list) @@ -101,6 +112,8 @@ class WhisperDecoding: 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'], hidden_size=self.decoder_config['hidden_size'], @@ -141,6 +154,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() + # generation config sampling_config = SamplingConfig(end_id=eot_id, pad_id=eot_id, @@ -161,6 +178,7 @@ class WhisperDecoding: sampling_config, encoder_output=encoder_outputs, encoder_input_lengths=encoder_input_lengths, + cross_attention_mask=cross_attention_mask, ) torch.cuda.synchronize()