Abuqader/fp8 tp8 (#244)

Co-authored-by: Abu Qader <[email protected]>

Author: bolasim

mistral/mistral-tp8/config.yaml

@@ -0,0 +1,43 @@
+apply_library_patches: true
+base_image:
+  image: baseten/trtllm-server:r23.12_baseten_v0.9.0.dev2024022000
+  python_executable_path: /usr/bin/python3
+build:
+  arguments:
+    engine_repository: baseten/mistral_fp8_tp8_i2048_o2048_bs128-tllm_0.9.0.dev2024022000
+    pipeline_parallel_count: 1
+    tensor_parallel_count: 8
+    tokenizer_repository: mistralai/Mistral-7B-Instruct-v0.2
+bundled_packages_dir: packages
+data_dir: data
+description: Generate text from a prompt with this seven billion parameter language
+  model.
+environment_variables: {}
+examples_filename: examples.yaml
+external_data: null
+external_package_dirs: []
+input_type: Any
+live_reload: false
+model_cache: []
+model_class_filename: model.py
+model_class_name: Model
+model_framework: custom
+model_metadata:
+  tags:
+  - text-generation
+  - openai-compatible
+model_module_dir: model
+model_name: Mistral 8
+model_type: Model
+python_version: py311
+requirements:
+- tritonclient[all]
+- transformers
+- jinja2
+resources:
+  accelerator: H100:8
+  use_gpu: true
+runtime:
+  num_workers: 1
+  predict_concurrency: 512
+secrets: {}

mistral/mistral-tp8/model/__init__.py

@@ -0,0 +1,97 @@
+import os
+from itertools import count
+
+import build_engine_utils
+from constants import (
+    GRPC_SERVICE_PORT,
+    HF_AUTH_KEY_CONSTANT,
+    HTTP_SERVICE_PORT,
+    TOKENIZER_KEY_CONSTANT,
+)
+from schema import ModelInput, TrussBuildConfig
+from transformers import AutoTokenizer
+from triton_client import TritonClient, TritonServer
+
+
+class Model:
+    def __init__(self, data_dir, config, secrets):
+        self._data_dir = data_dir
+        self._config = config
+        self._secrets = secrets
+        self._request_id_counter = count(start=1)
+        self.triton_client = None
+        self.triton_server = None
+        self.tokenizer = None
+        self.uses_openai_api = None
+
+    def load(self):
+        build_config = TrussBuildConfig(**self._config["build"]["arguments"])
+        self.uses_openai_api = "openai-compatible" in self._config.get(
+            "model_metadata", {}
+        ).get("tags", [])
+        hf_access_token = None
+        if "hf_access_token" in self._secrets._base_secrets.keys():
+            hf_access_token = self._secrets["hf_access_token"]
+
+        # TODO(Abu): Move to pre-runtime
+        if build_config.requires_build:
+            build_engine_utils.build_engine_from_config_args(
+                engine_build_args=build_config.engine_build_args,
+                dst=self._data_dir,
+            )
+
+        self.triton_server = TritonServer(
+            grpc_port=GRPC_SERVICE_PORT,
+            http_port=HTTP_SERVICE_PORT,
+        )
+
+        self.triton_server.create_model_repository(
+            truss_data_dir=self._data_dir,
+            engine_repository_path=build_config.engine_repository
+            if not build_config.requires_build
+            else None,
+            huggingface_auth_token=hf_access_token,
+        )
+
+        env = {}
+        if hf_access_token:
+            env[HF_AUTH_KEY_CONSTANT] = hf_access_token
+        env[TOKENIZER_KEY_CONSTANT] = build_config.tokenizer_repository
+
+        self.triton_server.start(
+            world_size=build_config.tensor_parallel_count,
+            env=env,
+        )
+
+        self.triton_client = TritonClient(
+            grpc_service_port=GRPC_SERVICE_PORT,
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            build_config.tokenizer_repository, token=hf_access_token
+        )
+        self.eos_token_id = self.tokenizer.eos_token_id
+
+    async def predict(self, model_input):
+        model_input["request_id"] = str(os.getpid()) + str(
+            next(self._request_id_counter)
+        )
+        model_input["eos_token_id"] = self.eos_token_id
+
+        self.triton_client.start_grpc_stream()
+
+        model_input = ModelInput(**model_input)
+
+        result_iterator = self.triton_client.infer(model_input)
+
+        async def generate():
+            async for result in result_iterator:
+                yield result
+
+        if model_input.stream:
+            return generate()
+        else:
+            if self.uses_openai_api:
+                return "".join(generate())
+            else:
+                return {"text": "".join(generate())}

mistral/mistral-tp8/model/model.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+from schema import EngineBuildArgs
+
+
+def build_engine_from_config_args(
+    engine_build_args: EngineBuildArgs,
+    dst: Path,
+):
+    import os
+    import shutil
+    import sys
+
+    # NOTE: These are provided by the underlying base image
+    # TODO(Abu): Remove this when we have a better way of handling this
+    sys.path.append("/app/baseten")
+    from build_engine import Engine, build_engine
+    from trtllm_utils import docker_tag_aware_file_cache
+
+    engine = Engine(**engine_build_args.model_dump())
+
+    with docker_tag_aware_file_cache("/root/.cache/trtllm"):
+        built_engine = build_engine(engine, download_remote=True)
+
+        if not os.path.exists(dst):
+            os.makedirs(dst)
+
+        for filename in os.listdir(str(built_engine)):
+            source_file = os.path.join(str(built_engine), filename)
+            destination_file = os.path.join(dst, filename)
+            if not os.path.exists(destination_file):
+                shutil.copy(source_file, destination_file)
+
+        return dst

mistral/mistral-tp8/packages/build_engine_utils.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+# If changing model repo path, please updated inside tensorrt_llm config.pbtxt as well
+TENSORRT_LLM_MODEL_REPOSITORY_PATH = Path("/packages/tensorrt_llm_model_repository/")
+GRPC_SERVICE_PORT = 8001
+HTTP_SERVICE_PORT = 8003
+HF_AUTH_KEY_CONSTANT = "HUGGING_FACE_HUB_TOKEN"
+TOKENIZER_KEY_CONSTANT = "TRITON_TOKENIZER_REPOSITORY"
+ENTRYPOINT_MODEL_NAME = "ensemble"