Deploy Gemma 7B with TensorRT-LLM

import os
import subprocess
from pathlib import Path

import kagglehub

# Update the following with your Kaggle username and key
os.environ["KAGGLE_USERNAME"] = "YOUR_USERNAME"
os.environ["KAGGLE_KEY"] = "YOUR_KEY"

torch_weights_dir = kagglehub.model_download(f"google/gemma/pyTorch/7b-it")

# Create the directory where we will save the TRT engine. This is the file we will be deploying.
model_dir = Path(os.getcwd() + "/trt-engine").expanduser()
model_dir.mkdir(parents=True, exist_ok=True)

# Build the TRT engine
part1 = f"""
python ./TensorRT-LLM/examples/gemma/convert_checkpoint.py \
    --ckpt-type torch \
    --model-dir {str(torch_weights_dir)} \
    --dtype bfloat16 \
    --world-size 1 \
    --output-model-dir {os.getcwd()}
"""

part2 = f"""
trtllm-build --checkpoint_dir {os.getcwd()} \
             --gemm_plugin bfloat16 \
             --gpt_attention_plugin bfloat16 \
             --max_batch_size 8 \
             --max_input_len 3000 \
             --max_output_len 100 \
             --context_fmha enable \
             --output_dir {str(model_dir)}
"""

print("Building engine...")
subprocess.run(part1, shell=True, check=True)
subprocess.run(part2, shell=True, check=True)
print("\nEngine built successfully! You can find it at: ", model_dir)

import csv
import sys
import time
from typing import Optional

import numpy as np
import torch

sys.path.insert(0, "/app/TensorRT-LLM/examples")
import tensorrt_llm
import tensorrt_llm.profiler
from tensorrt_llm.logger import logger
from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner
from utils import load_tokenizer, read_model_name

from pipeline import Pipeline, Variable, entity, pipe
from pipeline.objects.graph import InputField, InputSchema

if PYTHON_BINDINGS:
    from tensorrt_llm.runtime import ModelRunnerCpp
# logger.set_level('error')

# TensorRT params
engine_dir = "/app/trt-engine"
tokenizer_dir = "/app/gemma_pytorch/tokenizer/tokenizer.model"


class Kwargs(InputSchema):
    temperature: Optional[float] = InputField(
        0.95, title="Temperature", description="The temperature to use for sampling."
    )
    top_k: Optional[int] = InputField(
        100, title="Top k", description="The number of tokens to sample from."
    )
    top_p: Optional[float] = InputField(
        1.0, title="Top p", description="The cumulative probability to sample from."
    )
    repetition_penalty: Optional[float] = InputField(
        1.05, title="Repetition penalty", description="The repetition penalty to use."
    )
    prompt_template: Optional[str] = InputField(
        "<start_of_turn>user\n{input_text}<end_of_turn>\n<start_of_turn>model\n",
        title="Prompt template",
        description="The template to use for generation.",
        optional=True,
    )


# Put your model inside of the below entity class
@entity
class MyModelClass:
    @pipe(run_once=True, on_startup=True)
    def load(self) -> None:
        # Perform any operations needed to load your model here
        print("Loading model...")
        runtime_rank = tensorrt_llm.mpi_rank()

        use_py_session = False

        self.model_name, self.model_version = read_model_name(engine_dir)

        self.tokenizer, self.pad_id, self.end_id = load_tokenizer(
            vocab_file=tokenizer_dir,
            model_name=self.model_name,
            model_version=self.model_version,
        )

        if not PYTHON_BINDINGS:
            logger.warning(
                "Python bindings of C++ session is unavailable, fallback to Python session."
            )
            use_py_session = True

        runner_cls = ModelRunner if use_py_session else ModelRunnerCpp
        self.runner = runner_cls.from_dir(engine_dir=engine_dir, rank=runtime_rank)

        print("Model loaded!")

    def parse_input(
        self,
        tokenizer,
        input_text=None,
        prompt_template=None,
        input_file=None,
        add_special_tokens=True,
        max_input_length=923,
        pad_id=None,
        num_prepend_vtokens=[],
        model_name=None,
        model_version=None,
    ):
        if pad_id is None:
            pad_id = tokenizer.pad_token_id

        batch_input_ids = []
        if input_file is None:
            for curr_text in input_text:
                if prompt_template is not None:
                    curr_text = prompt_template.format(input_text=curr_text)
                input_ids = tokenizer.encode(
                    curr_text,
                    add_special_tokens=add_special_tokens,
                    truncation=True,
                    max_length=max_input_length,
                )
                batch_input_ids.append(input_ids)
        else:
            if input_file.endswith(".csv"):
                with open(input_file, "r") as csv_file:
                    csv_reader = csv.reader(csv_file, delimiter=",")
                    for line in csv_reader:
                        input_ids = np.array(line, dtype="int32")
                        batch_input_ids.append(input_ids[-max_input_length:])
            elif input_file.endswith(".npy"):
                inputs = np.load(input_file)
                for row in inputs:
                    input_ids = row[row != pad_id]
                    batch_input_ids.append(input_ids[-max_input_length:])
            elif input_file.endswith(".txt"):
                with open(
                    input_file, "r", encoding="utf-8", errors="replace"
                ) as txt_file:
                    input_text = txt_file.read()
                    input_ids = tokenizer.encode(
                        input_text,
                        add_special_tokens=add_special_tokens,
                        truncation=True,
                        max_length=max_input_length,
                    )
                    batch_input_ids.append(input_ids)
            else:
                print("Input file format not supported.")
                raise SystemExit
        if model_name == "GemmaForCausalLM":
            batch_input_ids[0] = [tokenizer.bos_token_id] + batch_input_ids[0]

        if num_prepend_vtokens:
            assert len(num_prepend_vtokens) == len(batch_input_ids)
            base_vocab_size = tokenizer.vocab_size - len(
                tokenizer.special_tokens_map.get("additional_special_tokens", [])
            )
            for i, length in enumerate(num_prepend_vtokens):
                batch_input_ids[i] = (
                    list(range(base_vocab_size, base_vocab_size + length))
                    + batch_input_ids[i]
                )

        if model_name == "ChatGLMForCausalLM" and model_version == "glm":
            for ids in batch_input_ids:
                ids.append(tokenizer.sop_token_id)

        batch_input_ids = [torch.tensor(x, dtype=torch.int32) for x in batch_input_ids]
        return batch_input_ids

    @pipe
    def predict(self, prompt: list, num_tokens: int, kwargs: Kwargs) -> dict:
        # Perform any operations needed to predict with your model here
        print("Predicting...")

        if len(prompt) > 8:
            raise ValueError(
                "The number of batched prompts must be less than or equal to 8."
            )

        # # An example to stop generation when the model generate " London" on first sentence, " eventually became" on second sentence
        # stop_words_list = [[" London"], ["eventually became"]]
        # stop_words_list = tensorrt_llm.runtime.to_word_list_format(stop_words_list, tokenizer)
        # stop_words_list = torch.Tensor(stop_words_list).to(torch.int32).to("cuda").contiguous()
        stop_words_list = None

        # # An example to prevent generating " chef" on first sentence, " eventually" and " chef before" on second sentence
        # bad_words_list = [[" chef"], [" eventually, chef before"]]
        # bad_words_list = tensorrt_llm.runtime.to_word_list_format(bad_words_list, tokenizer)
        # bad_words_list = torch.Tensor(bad_words_list).to(torch.int32).to("cuda").contiguous()
        bad_words_list = None

        batch_input_ids = self.parse_input(
            tokenizer=self.tokenizer,
            input_text=prompt,
            prompt_template=kwargs.prompt_template,
            pad_id=self.pad_id,
            model_name=self.model_name,
            model_version=self.model_version,
        )
        input_lengths = [x.size(0) for x in batch_input_ids]

        time_begin = time.time()
        with torch.no_grad():
            outputs = self.runner.generate(
                batch_input_ids,
                max_new_tokens=num_tokens,
                max_attention_window_size=None,
                sink_token_length=None,
                end_id=self.end_id,
                pad_id=self.pad_id,
                temperature=kwargs.temperature,
                top_k=kwargs.top_k,
                top_p=kwargs.top_p,
                num_beams=1,
                repetition_penalty=kwargs.repetition_penalty,
                stop_words_list=stop_words_list,
                bad_words_list=bad_words_list,
                output_sequence_lengths=True,
                return_dict=True,
            )
            torch.cuda.synchronize()

        time_total = time.time() - time_begin

        output_ids = outputs["output_ids"]
        sequence_lengths = outputs["sequence_lengths"]

        batch_size, num_beams, _ = output_ids.size()
        generated_texts = []
        generated_ids = []
        for batch_idx in range(batch_size):
            inputs = output_ids[batch_idx][0][: input_lengths[batch_idx]].tolist()
            input_text = self.tokenizer.decode(inputs)
            print(f'Input [Text {batch_idx}]: "{input_text}"')
            for beam in range(num_beams):
                output_begin = input_lengths[batch_idx]
                output_end = sequence_lengths[batch_idx][beam]
                outputs = output_ids[batch_idx][beam][output_begin:output_end].tolist()
                output_text = self.tokenizer.decode(outputs)
                generated_ids.append(outputs)
                generated_texts.append(output_text)

        generated_tokens = sum([len(x) for x in generated_ids])
        print(
            f"Response generated in {time_total:.2f} seconds, {generated_tokens} tokens, {generated_tokens / time_total:.2f} tokens/second"
        )
        print(f"Batch size: {batch_size}")

        print("Prediction complete!")

        return {
            "response_txt": generated_texts,
            "throughput_tok_s": generated_tokens / time_total,
            "latency_s": time_total,
        }


with Pipeline() as builder:
    prompt = Variable(
        list,
        default=["What is an LLM?"],
        title="Prompt",
        description="The input prompt to generate a response to. Add multiple \
          prompts to generate multiple responses. Single API calls with multiple\
            prompts is more efficient than multiple API calls with single prompts.",
    )
    num_tokens = Variable(
        int,
        default=100,
        title="Max tokens",
        description="Maximum number of tokens to generate per output sequence.",
    )
    kwargs = Variable(Kwargs, title="Other parameters")

    my_model = MyModelClass()
    my_model.load()

    output_var = my_model.predict(prompt, num_tokens, kwargs)

    builder.output(output_var)

my_new_pipeline = builder.get_pipeline()

runtime:
  container_commands:
  - apt update -y && apt install -y software-properties-common python3-launchpadlib && apt update -y && add-apt-repository -y ppa:git-core/ppa && apt update -y
  - apt-get update
  - apt install -y gcc g++ aria2 git git-lfs wget openmpi-bin libopenmpi-dev
  - wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run && sh cuda_12.1.0_530.30.02_linux.run --silent --toolkit
  - git clone https://github.com/NVIDIA/TensorRT-LLM.git
  - git clone https://github.com/google/gemma_pytorch.git
  python:
    version: '3.10'
    requirements: 
    - pipeline-ai
    - tensorrt_llm -U --pre --extra-index-url https://pypi.nvidia.com
    - flax h5py sentencepiece easydict mpmath==1.3.0
accelerators: [nvidia_a100]
accelerator_memory: 39543
pipeline_graph: new_pipeline:my_new_pipeline
pipeline_name: gemma-7b-it-trt
description: null
readme: README.md
extras: {}
cluster: null