Merge pull request #158 from InfiniTensor/issue/157

issue/157 - add/adjust cambricon support in scripts

examples/bench.py

@@ -135,6 +135,11 @@ def get_args():
         action="store_true",
         help="Run nvidia test",
     )
+    parser.add_argument(
+        "--cambricon",
+        action="store_true",
+        help="Run cambricon test",
+    )
     parser.add_argument(
         "--model",
         type=str,
@@ -268,6 +273,8 @@ if __name__ == "__main__":
         device_str = "cpu"
     elif args.nvidia:
         device_str = "cuda"
+    elif args.cambricon:
+        device_str = "mlu"
     else:
         print(
             "python examples/bench.py --nvidia --model=~/TinyLlama-1.1B-Chat-v1.0/ --batch-size=2 --tp=1 --input-len=50 --output-len=50"

examples/jiuge.py

@@ -40,6 +40,11 @@ def get_args():
         action="store_true",
         help="Run iluvatar test",
     )
+    parser.add_argument(
+        "--cambricon",
+        action="store_true",
+        help="Run cambricon test",
+    )
     parser.add_argument(
         "--model_path",
         type=str,
@@ -188,6 +193,8 @@ if __name__ == "__main__":
         device_str = "musa"
     elif args.iluvatar:
         device_str = "cuda"
+    elif args.cambricon:
+        device_str = "mlu"
     else:
         print(
             "Usage:  python examples/jiuge.py [--cpu | --nvidia | --metax | --moore | --iluvatar] --model_path=<path/to/model_dir>\n"

test/bench/test_benchmark.py

@@ -63,7 +63,7 @@ class InfiniLMBenchmark(BaseBenchmark):
         device_map = {
             "cpu": "cpu",
             "nvidia": "cuda",
-            "cambricon": "cambricon",
+            "cambricon": "mlu",
             "ascend": "ascend",
             "metax": "metax",
             "moore": "moore",
@@ -210,6 +210,8 @@ class TorchBenchmark(BaseBenchmark):
             self.device = torch.device("cuda")
         elif device_type_str == "cpu":
             self.device = torch.device("cpu")
+        elif device_type_str == "cambricon":
+            self.device = torch.device("mlu")
         else:
             raise ValueError(
                 f"Torch backend unsupported device type: {device_type_str}"