fix HF infer & update mpsize arg (deepspeedai#262)

satpalsr · mrwyattii · web-flow · commit 85a69e8f8130 · 2023-04-11T13:03:48.000-07:00
Co-authored-by: Michael Wyatt &lt;mrwyattii@gmail.com&gt;
diff --git a/benchmarks/inference/bert-bench.py b/benchmarks/inference/bert-bench.py
@@ -70,7 +70,7 @@ def print_latency(latency_set, title, warmup=3):
 if args.deepspeed:
     pipe.model = deepspeed.init_inference(pipe.model,
                                           dtype=dtype,
-                                          mp_size=1,
+                                          tensor_parallel={"tp_size": 1},
                                           replace_with_kernel_inject=args.kernel_inject,
                                           enable_cuda_graph=args.graphs)
     pipe.model.profile_model_time()
@@ -90,6 +90,7 @@ def print_latency(latency_set, title, warmup=3):
         mtimes += pipe.model.model_times()
 
 print_latency(times, "e2e latency")
-print_latency(mtimes, "model latency")
+if args.deepspeed:
+    print_latency(mtimes, "model latency")
 
 print(responses[0:3])
diff --git a/benchmarks/inference/gpt-bench.py b/benchmarks/inference/gpt-bench.py
@@ -79,7 +79,7 @@ def print_latency(latency_set, title, warmup=3):
     pipe.model = deepspeed.init_inference(
         pipe.model,
         dtype=dtype,
-        mp_size=args.world_size,
+        tensor_parallel={"tp_size": args.world_size},
         replace_with_kernel_inject=args.kernel_inject,
         enable_cuda_graph=args.graphs,
     )
@@ -101,7 +101,8 @@ def print_latency(latency_set, title, warmup=3):
 
 if args.local_rank == 0:
     print_latency(times, "(e2e) latency")
-    print_latency(mtimes, "(model-only) latency")
+    if args.deepspeed:
+        print_latency(mtimes, "(model-only) latency")
     print_latency(map(lambda t: t / (args.max_tokens - 3), times), "(e2e) per token latency")
     print(f"RESPONSE 0:")
     print("-" * 30)
