refactor

Author: dreaming-panda

examples/batch_generation.py

@@ -0,0 +1,64 @@
+import sys
+sys.path.append("..")
+from models.llama import LLM
+import argparse
+import torch
+from transformers import AutoTokenizer
+import jsonlines
+parser = argparse.ArgumentParser()
+parser.add_argument('--model', type=str, default="meta-llama/Meta-Llama-3.1-8B-Instruct",help='model')
+parser.add_argument('--T', type=int, default=2000, help='repeat times')
+parser.add_argument('--B', type=int, default=2, help='batch size')
+parser.add_argument('--M', type=int, default=4096, help='max length')
+parser.add_argument('--D', type=int, default=1, help='dec length')
+parser.add_argument('--G', type=int, default=32, help='generation length')
+parser.add_argument('--K', type=int, default=10, help='K')
+parser.add_argument('--L', type=int, default=150, help='K')
+args = parser.parse_args()
+print(args)
+MAX_LEN = args.M
+DEC_LEN = args.D
+GEN_LEN = args.G
+BATCH_SIZE = args.B
+MODEL_NAME = args.model
+DTYPE = torch.bfloat16
+DEVICE = "cuda:0"
+T = args.T
+WARM_UP = 10
+
+with open("../data/data4k.jsonl") as f:
+    d = jsonlines.Reader(f)
+    for idx, item in enumerate(d):
+        data = item
+        break
+
+llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=BATCH_SIZE, device=DEVICE, dtype=DTYPE)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+text = data["input"]
+input_ids = tokenizer.encode(text=text, return_tensors="pt").to(device=DEVICE)
+PREFIX_LEN = input_ids.shape[1]
+
+position_ids = torch.arange(MAX_LEN, device=DEVICE).unsqueeze(0).repeat(BATCH_SIZE, 1)
+
+batch_logits = []
+for i  in range(BATCH_SIZE):
+    logits = llm.prefill(input_ids, i)
+    batch_logits.append(logits)
+
+logits = torch.cat(batch_logits, dim=0)
+generated_tokens = []
+prefix_len = input_ids.shape[1]
+for k in range(GEN_LEN):
+    input_ids = logits.argmax(dim=-1)
+    logits = llm.inference(input_ids=input_ids, position_ids=position_ids[:,prefix_len + k:prefix_len + k + 1])
+    generated_tokens.append(input_ids)
+    if input_ids[0].item() in [128000, 128001, 128008, 128009]:
+                break
+generated_tokens = torch.cat(generated_tokens, dim=1).to(device="cpu")
+decoded_texts = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+print(decoded_texts)
+    
+    
+    
+
+

examples/bench.py

@@ -0,0 +1,59 @@
+import sys
+sys.path.append("..")
+from models.llama import LLM
+import argparse
+import torch
+from transformers import AutoTokenizer
+import jsonlines
+import time
+parser = argparse.ArgumentParser()
+parser.add_argument('--model', type=str, default="meta-llama/Meta-Llama-3.1-8B-Instruct",help='model')
+parser.add_argument('--B', type=int, default=1, help='batch size')
+parser.add_argument('--M', type=int, default=98304, help='max length')
+parser.add_argument('--D', type=int, default=1, help='dec length')
+parser.add_argument('--P', type=int, default=98000, help='prefill length')
+parser.add_argument('--G', type=int, default=128, help='generation length')
+parser.add_argument('--K', type=int, default=10, help='K')
+parser.add_argument('--L', type=int, default=150, help='L')
+args = parser.parse_args()
+print(args)
+MAX_LEN = args.M
+DEC_LEN = args.D
+GEN_LEN = args.G
+B = args.B
+MODEL_NAME = args.model
+DTYPE = torch.bfloat16
+PREFIX_LEN = args.P
+DEVICE = "cuda:0"
+WARM_UP = 32
+
+with open("../data/data.jsonl") as f:
+    d = jsonlines.Reader(f)
+    for idx, item in enumerate(d):
+        data = item
+        break
+
+llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=B, device=DEVICE, dtype=DTYPE)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+text = data["input"]
+input_ids = tokenizer.encode(text=text, return_tensors="pt").to(device=DEVICE)
+input_ids = input_ids[:,:PREFIX_LEN].repeat(B, 1)
+position_ids = torch.arange(MAX_LEN, device=DEVICE).unsqueeze(0).repeat(B, 1)
+
+for i in range(B):
+    logits = llm.prefill(input_ids=input_ids[i:i+1], request_id=i)
+
+generated = input_ids[0].tolist()
+for k in range(WARM_UP):
+    logits = llm.inference(input_ids=input_ids[:, 128+k:128+k+1], position_ids=position_ids[:,PREFIX_LEN + k:PREFIX_LEN + k + 1])
+
+torch.cuda.synchronize()
+t1 = time.time()
+for k in range(GEN_LEN):
+    logits = llm.inference(input_ids=input_ids[:, WARM_UP+k:WARM_UP+k+1], position_ids=position_ids[:,WARM_UP + PREFIX_LEN + k: WARM_UP + PREFIX_LEN + k + 1])
+
+torch.cuda.synchronize()
+t2 = time.time()
+
+print("Decoding Latency {:.2f} ms/token".format(1000 * (t2 - t1)/GEN_LEN))
+print("Decoding Throughput {:.2f} token/s".format(B * GEN_LEN / (t2 - t1)))

examples/bench.sh

@@ -0,0 +1,10 @@
+numactl -C 0-31,52-83 -m 0,1 python bench.py --B 1 --K 0 --L 150 --model codellama/CodeLlama-7b-Instruct-hf --M 16384 --P 16000
+numactl -C 0-31,52-83 -m 0,1 python bench.py --B 4 --K 0 --L 150 --model codellama/CodeLlama-7b-Instruct-hf --M 16384 --P 16000
+numactl -C 0-31,52-83 -m 0,1 python bench.py --B 8 --K 0 --L 150 --model codellama/CodeLlama-7b-Instruct-hf --M 16384 --P 16000
+# numactl -C 0-31,52-83 -m 0,1 python bench.py --B 1 --K 10 --L 170 --model codellama/CodeLlama-7b-Instruct-hf --M 131072 --P 128000
+
+
+# numactl -C 0-31,52-83 -m 0,1 python bench.py --B 12 --K 9 --L 120 --model codellama/CodeLlama-7b-Instruct-hf --M 131072 --P 128000
+
+
+# numactl -C 0-31,52-83 -m 0,1 python bench.py --B 12 --K 8 --L 75 --model codellama/CodeLlama-7b-Instruct-hf --M 131072 --P 128000

examples/generation.py

@@ -0,0 +1,44 @@
+import sys
+sys.path.append("..")
+from models.llama import LLM
+import argparse
+import torch
+from transformers import AutoTokenizer
+import jsonlines
+from models.template import Templates
+parser = argparse.ArgumentParser()
+parser.add_argument('--model', type=str, default="meta-llama/Meta-Llama-3.1-8B-Instruct",help='model')
+parser.add_argument('--M', type=int, default=8192, help='max length')
+parser.add_argument('--D', type=int, default=1, help='dec length')
+parser.add_argument('--G', type=int, default=256, help='generation length')
+parser.add_argument('--K', type=int, default=10, help='K')
+parser.add_argument('--L', type=int, default=150, help='K')
+parser.add_argument('--data', type=str, default="../data/story.txt", help='source data file')
+parser.add_argument('--template', type=str, default="meta-llama3", help='chat template')
+args = parser.parse_args()
+print(args)
+MAX_LEN = args.M
+DEC_LEN = args.D
+GEN_LEN = args.G
+MODEL_NAME = args.model
+DTYPE = torch.bfloat16
+DEVICE = "cuda:0"
+chat_template = Templates[args.template]
+llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=1, device=DEVICE, dtype=DTYPE, generation_buffer=args.G + 32)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+with open(args.data, "r", encoding="utf-8") as file:
+    content = file.read()
+    content = chat_template.format(content)
+    input_ids = tokenizer.encode(text=content, return_tensors="pt")
+    context = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+    print(context)
+    input_ids = input_ids.to(DEVICE)
+    PREFIX_LEN = input_ids.shape[1]
+    position_ids = torch.arange(MAX_LEN, device=DEVICE).unsqueeze(0)
+    generated = llm.generate(input_ids, max_tokens=args.G)
+    text = tokenizer.decode(generated, skip_special_tokens=True)
+    print("\033[32m" + text + "\033[0m")
+    
+
+
+