chore: requirement

gpt4all-api/triton/README.md

@@ -0,0 +1,13 @@
+# To Run Inference Server
+
+docker run --gpus=1 --rm --net=host -v ${PWD}/model_store:/model_store nvcr.io/nvidia/tritonserver:23.01-py3 tritonserver --model-repository=/model_store
+
+python client.py --model=<model_name>
+
+
+## Dynamic Batching
+Need to figure out how to do batching such that we can have dynamic batching
+We're getting 1.3 infer/sec which seems slow....
+
+To test,
+perf_analyzer -m nomic-ai--gpt4all-j --input-data test_data.json --measurement-interval 25000 --request-rate-range=10 -b 8

gpt4all-api/triton/client.py

@@ -0,0 +1,75 @@
+import torch
+import tritonclient.grpc.aio as grpcclient
+
+
+def prepare_inference_inputs(
+    inputs_ids: torch.IntTensor, new_tokens: int = 1, temperature: float = 1.0
+):
+    batch_size = inputs_ids.shape[0]
+
+    input_ids_input = grpcclient.InferInput("input_ids", inputs_ids.shape, "INT32")
+    input_ids_input.set_data_from_numpy(inputs_ids.int().cpu().numpy())
+
+    new_tokens_input = grpcclient.InferInput(
+        "tensor_of_seq_len", [batch_size, new_tokens], "INT32"
+    )
+    new_tokens_input.set_data_from_numpy(
+        torch.zeros(batch_size, new_tokens, dtype=torch.int32).cpu().numpy()
+    )
+
+    temperature_input = grpcclient.InferInput("temperature", [batch_size, 1], "FP32")
+    temperature_input.set_data_from_numpy(
+        torch.full([batch_size, 1], temperature, dtype=torch.float32).cpu().numpy()
+    )
+
+    inputs = [input_ids_input, new_tokens_input, temperature_input]
+    outputs = [
+        grpcclient.InferRequestedOutput("logits"),
+        grpcclient.InferRequestedOutput("output_ids"),
+    ]
+    return inputs, outputs
+
+
+async def infer(
+    triton_client, model_name, input_ids, new_tokens: int = 1, temperature: float = 1.0
+):
+    inputs, outputs = prepare_inference_inputs(input_ids, new_tokens, temperature)
+
+    triton_model_name = model_name.replace("/", "--")
+
+    result = await triton_client.infer(
+        model_name=triton_model_name, inputs=inputs, outputs=outputs
+    )
+
+    logits = torch.tensor(result.as_numpy("logits").copy(), requires_grad=False)
+    output_ids = torch.tensor(result.as_numpy("output_ids").copy(), requires_grad=False)
+
+    return logits, output_ids
+
+def Client(url: str):
+  return grpcclient.InferenceServerClient(url=url)
+
+if __name__ == "__main__":
+    import argparse
+    from transformers import AutoTokenizer
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", type=str, default="localhost:8001")
+    parser.add_argument("--model", type=str, default="gpt2")
+
+    args = parser.parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=False)
+
+    async def main():
+        async with Client(args.url) as triton_client:
+            while True:
+                prompt = input("Prompt: ")
+                input_ids = tokenizer.encode(prompt, return_tensors="pt")
+                last_logits, output_ids = await infer(
+                    triton_client, args.model, input_ids, new_tokens=256, temperature=1.0,
+                )
+                print(tokenizer.decode(output_ids[0]))
+
+    import asyncio
+    asyncio.run(main())

gpt4all-api/triton/convert_to_triton.py

@@ -0,0 +1,149 @@
+import argparse
+import os
+from string import Template
+
+import torch
+from torch import nn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument(
+    "--model", type=str, required=True, help="Path to HF checkpoint with the base model"
+)
+
+parser.add_argument(
+    "--max-batch-size", type=int, default=64, help="Maximum batch size for inference"
+)
+
+parser.add_argument(
+    "--revision",
+    type=str,
+    required=False,
+    help="Optional branch/commit of the HF checkpoint",
+)
+
+parser.add_argument("--device", type=int, default=0)
+args = parser.parse_args()
+
+device = torch.device(args.device)
+
+
+class ModelLogits(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    @torch.inference_mode()
+    def forward(self, input_ids: torch.Tensor):
+        return self.model(input_ids).logits
+
+
+class InferModel(nn.Module):
+    def __init__(self, traced_model, eos_token_id):
+        super().__init__()
+        self.traced_model = traced_model
+        self.eos_token_id = eos_token_id
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        tensor_of_seq_len: torch.Tensor,
+        temperature: torch.Tensor,
+    ):
+        # this has mostly been adapted from huggingface generate
+        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
+        eos_token_id_tensor = torch.tensor([self.eos_token_id]).to(input_ids.device)
+
+        with torch.no_grad():
+            for _ in range(tensor_of_seq_len.shape[1] - 1):
+                logits = self.traced_model(input_ids).float()
+                next_token_logits = logits[:, -1, :]
+                next_token_logits = next_token_logits / temperature
+                
+                next_tokens = torch.multinomial(
+                    torch.softmax(next_token_logits, dim=-1), input_ids.shape[0] 
+                )
+
+                next_tokens = next_tokens * unfinished_sequences + self.eos_token_id * (1 - unfinished_sequences)
+
+                unfinished_sequences = unfinished_sequences.mul(
+                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
+                )
+
+                # stop when each sentence is finished
+                if unfinished_sequences.max() == 0:
+                    return input_ids.int(), logits
+
+                input_ids = torch.cat([input_ids, next_tokens], dim=-1)
+
+                unfinished_sequences = unfinished_sequences.mul(
+                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
+                )
+
+            # in TorchScript, the above logits var lifetime doesn't escape the loop's scope
+            logits = self.traced_model(input_ids).float()
+            next_token_logits = logits[:, -1, :]
+            next_token_logits = next_token_logits / temperature
+
+            next_tokens = torch.multinomial(
+                    torch.softmax(next_token_logits, dim=-1), input_ids.shape[0] 
+                )
+
+            next_tokens = next_tokens * unfinished_sequences + self.eos_token_id * (1 - unfinished_sequences)
+
+            input_ids = torch.cat([input_ids, next_tokens], dim=-1)
+
+            return input_ids.int(), logits
+            
+
+print(f"Converting {args.model} to TorchScript...")
+tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=False)
+model = ModelLogits(AutoModelForCausalLM.from_pretrained(args.model,
+                                                         trust_remote_code=True,
+                                                         revision=args.revision,
+                                                         torch_dtype=torch.float16,
+                                                         use_cache=False))
+
+model.eval()
+model.requires_grad_(False)
+model = model.to(device)
+
+
+input = tokenizer("annotator model's hash is 0x", return_tensors="pt").to(device)
+print(f"{model(input.input_ids)=}")
+
+traced_script_module = torch.jit.trace(model, input.input_ids)
+print("Tracing...")
+print(f"{traced_script_module(input.input_ids)=}")
+
+print("Scripting generation wrapper...")
+# need to script this as we have data conditional flow
+scripted_generator_model = torch.jit.script(InferModel(traced_script_module, tokenizer.eos_token_id))
+print(scripted_generator_model.code)
+
+print(f"{input.input_ids=}")
+# x = input.input_ids, torch.empty(1, 5), torch.full([1, 1], 1.0).cuda(), torch.full([1, 1], len(tokenizer) // 2).cuda(), torch.full([1, 1], 0.9).cuda()
+x = input.input_ids, torch.empty(1, 5),  torch.full([1, 1], 0.9).cuda()
+print(x[0].shape)
+
+print(f"{tokenizer.decode(scripted_generator_model(*x)[0][0])=}")
+
+sanitized_name = args.model.replace("/", "--")
+print("Model renamed to ", sanitized_name)
+
+print("Saving TorchScript model...")
+
+os.makedirs(f"model_store/{sanitized_name}/1", exist_ok=True)
+scripted_generator_model.save(f"model_store/{sanitized_name}/1/traced-model.pt")
+
+config_path = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "triton_config.pbtxt"
+)
+with open(config_path) as f:
+    template = Template(f.read())
+config = template.substitute(
+    {"model_name": sanitized_name, "max_batch_size": args.max_batch_size}
+)
+with open(f"model_store/{sanitized_name}/config.pbtxt", "w") as f:
+    f.write(config)

gpt4all-api/triton/requirements.txt

@@ -0,0 +1,5 @@
+transformers
+triton
+einops
+pandas
+sentencepiece

gpt4all-api/triton/test_data.json

@@ -0,0 +1,34 @@
+{
+    "data":
+      [
+        {
+          "input_ids": {
+            "content": [17250, 11, 703, 389, 345, 30],
+            "shape": [6]
+          },
+          "tensor_of_seq_len": {
+            "content": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+            "shape": [17]
+          },
+          "temperature": {
+            "content": [1.0],
+            "shape": [1]
+          }
+        },
+        {
+            "input_ids": {
+              "content": [17250, 11, 703, 389, 345, 30],
+              "shape": [6]
+            },
+            "tensor_of_seq_len": {
+              "content": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+              "shape": [17]
+            },
+            "temperature": {
+              "content": [1.0],
+              "shape": [1]
+            }
+          }
+
+      ]
+  }

gpt4all-api/triton/triton_config.pbtxt

@@ -0,0 +1,69 @@
+name: "${model_name}"
+backend: "pytorch"
+default_model_filename: "traced-model.pt"
+max_batch_size: ${max_batch_size}
+
+dynamic_batching {
+}
+
+parameters {
+  key: "model_name"
+  value: {
+    string_value: "${model_name}"
+  }
+}
+
+instance_group [
+  {
+    count: 1
+    kind: KIND_GPU
+    gpus: [0]
+  }
+]
+
+input [
+  {
+    name: "input_ids"
+    data_type: TYPE_INT32
+    dims: [-1]
+  },
+  {
+    name: "tensor_of_seq_len"
+    data_type: TYPE_INT32
+    dims: [-1]
+  },
+  {
+    name: "temperature"
+    data_type: TYPE_FP32
+    dims: [-1]
+  }
+]
+
+output [
+  {
+    name: "output_ids"
+    data_type: TYPE_INT32
+    dims: [-1]
+  },
+  {
+    name: "logits"
+    data_type: TYPE_FP32
+    dims: [-1]
+  }
+]
+
+parameters {
+  key: "data_type"
+  value: {
+    string_value: "fp16"
+  }
+}
+
+parameters: {
+  key: "INFERENCE_MODE"
+  value: {
+    string_value: "true"
+  }
+}
+
+version_policy: {specific: {versions: [1]}}