Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/examples/language/grok-1/inference.py b/examples/language/grok-1/inference.py
--- a/examples/language/grok-1/inference.py
+++ b/examples/language/grok-1/inference.py
@@ -1,8 +1,7 @@
 import time
 
 import torch
-from sentencepiece import SentencePieceProcessor
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, LlamaTokenizerFast
 from utils import get_defualt_parser, inference, print_output
 
 if __name__ == "__main__":
@@ -16,11 +15,17 @@
         device_map="auto",
         torch_dtype=torch.bfloat16,
     )
-    sp = SentencePieceProcessor(model_file=args.tokenizer)
+    model.eval()
+    init_time = time.time() - start
+
+    # A transformers-compatible version of the grok-1 tokenizer by Xenova
+    # https://huggingface.co/Xenova/grok-1-tokenizer
+    tokenizer = LlamaTokenizerFast.from_pretrained("Xenova/grok-1-tokenizer")
+
     for text in args.text:
         output = inference(
             model,
-            sp,
+            tokenizer,
             text,
             max_new_tokens=args.max_new_tokens,
             do_sample=args.do_sample,
@@ -28,5 +33,14 @@
             top_k=args.top_k,
             top_p=args.top_p,
         )
-        print_output(text, sp.decode(output))
-    print(f"Overall time: {time.time() - start} seconds.")
+        print_output(text, tokenizer.decode(output))
+
+    overall_time = time.time() - start
+    gen_latency = overall_time - init_time
+    avg_gen_latency = gen_latency / len(args.text)
+    print(
+        f"Initializing time: {init_time:.2f} seconds.\n"
+        f"Overall time: {overall_time:.2f} seconds. \n"
+        f"Generation latency: {gen_latency:.2f} seconds. \n"
+        f"Average generation latency: {avg_gen_latency:.2f} seconds. \n"
+    )
diff --git a/examples/language/grok-1/inference_tp.py b/examples/language/grok-1/inference_tp.py
--- a/examples/language/grok-1/inference_tp.py
+++ b/examples/language/grok-1/inference_tp.py
@@ -2,8 +2,7 @@
 
 import torch
 from grok1_policy import Grok1ForCausalLMPolicy
-from sentencepiece import SentencePieceProcessor
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, LlamaTokenizerFast
 from utils import get_defualt_parser, inference, print_output
 
 import colossalai
@@ -33,11 +32,17 @@
             args.pretrained, trust_remote_code=True, torch_dtype=torch.bfloat16
         )
     model, *_ = booster.boost(model)
-    sp = SentencePieceProcessor(model_file=args.tokenizer)
+    model.eval()
+    init_time = time.time() - start
+
+    # A transformers-compatible version of the grok-1 tokenizer by Xenova
+    # https://huggingface.co/Xenova/grok-1-tokenizer
+    tokenizer = LlamaTokenizerFast.from_pretrained("Xenova/grok-1-tokenizer")
+
     for text in args.text:
         output = inference(
             model.unwrap(),
-            sp,
+            tokenizer,
             text,
             max_new_tokens=args.max_new_tokens,
             do_sample=args.do_sample,
@@ -46,5 +51,14 @@
             top_p=args.top_p,
         )
         if coordinator.is_master():
-            print_output(text, sp.decode(output))
-    coordinator.print_on_master(f"Overall time: {time.time() - start} seconds.")
+            print_output(text, tokenizer.decode(output))
+
+    overall_time = time.time() - start
+    gen_latency = overall_time - init_time
+    avg_gen_latency = gen_latency / len(args.text)
+    coordinator.print_on_master(
+        f"Initializing time: {init_time:.2f} seconds.\n"
+        f"Overall time: {overall_time:.2f} seconds. \n"
+        f"Generation latency: {gen_latency:.2f} seconds. \n"
+        f"Average generation latency: {avg_gen_latency:.2f} seconds. \n"
+    )
diff --git a/examples/language/grok-1/utils.py b/examples/language/grok-1/utils.py
--- a/examples/language/grok-1/utils.py
+++ b/examples/language/grok-1/utils.py
@@ -20,9 +20,9 @@ def print_output(text, output):
 
 
 @torch.no_grad()
-def inference(model, sp, text, **generate_kwargs):
-    input_ids = sp.encode(text)
-    input_ids = torch.tensor([input_ids]).cuda()
+def inference(model, tokenizer, text, **generate_kwargs):
+    input_ids = tokenizer(text, return_tensors="pt").input_ids
+    input_ids = input_ids.cuda()
     attention_mask = torch.ones_like(input_ids)
     inputs = {
         "input_ids": input_ids,


Test patch: