Add Multi Card Command.

Author: zzjjay

examples/benchmark/mmlu/README.md

@@ -16,14 +16,27 @@ tar xf data.tar
 
 在当前目录下运行以下脚本：
 
+- 单卡运行
 ```
+export CUDA_VISIBLE_DEVICES=0
 python eval.py \
     --model_name_or_path /path/to/your/model \
     --temperature 0.2 \
     --ntrain 5 \
     --output_dir ${output_path} \
     --dtype 'float16'
 ```
+- 多卡运行
+```
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+python -m paddle.distributed.fleet.launch eval.py \
+    --model_name_or_path /path/to/your/model \
+    --temperature 0.2 \
+    --ntrain 5 \
+    --output_dir ${output_path} \
+    --dtype 'float16' \
+    --tensor_parallel_degree 4
+```
 
 参数说明
 

examples/benchmark/mmlu/eval.py

@@ -17,6 +17,7 @@
 import os
 
 import numpy as np
+import paddle
 import pandas as pd
 from categories import categories, subcategories
 from evaluator import ModelEvaluator
@@ -29,9 +30,9 @@ def main(args, evaluator):
         [f.split("_test.csv")[0] for f in os.listdir(os.path.join(args.data_dir, "test")) if "_test.csv" in f]
     )
     if not os.path.exists(args.output_dir):
-        os.makedirs(args.output_dir)
+        os.makedirs(args.output_dir, exist_ok=True)
     if not os.path.exists(os.path.join(args.output_dir, "results_{}".format(args.model_name_or_path))):
-        os.makedirs(os.path.join(args.output_dir, "results_{}".format(args.model_name_or_path)))
+        os.makedirs(os.path.join(args.output_dir, "results_{}".format(args.model_name_or_path)), exist_ok=True)
 
     all_cors = []
     subcat_cors = {subcat: [] for subcat_lists in subcategories.values() for subcat in subcat_lists}
@@ -95,15 +96,25 @@ def main(args, evaluator):
     parser.add_argument("--data_dir", "-d", type=str, default="data")
     parser.add_argument("--output_dir", type=str, default="results")
     parser.add_argument("--dtype", default="float32", type=str)
+    parser.add_argument("--tensor_parallel_degree", default=1, type=int)
 
     args = parser.parse_args()
     print(args)
 
+    if args.tensor_parallel_degree > 1:
+        strategy = paddle.distributed.fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "mp_degree": args.tensor_parallel_degree,
+        }
+        # Set control in tensor parallel
+        strategy.tensor_parallel_configs = {"tensor_init_seed": 1234}
+        paddle.distributed.fleet.init(is_collective=True, strategy=strategy)
     evaluator = ModelEvaluator(
         model_name_or_path=args.model_name_or_path,
         ntrain=args.ntrain,
         temperature=args.temperature,
         dtype=args.dtype,
+        tensor_parallel_degree=args.tensor_parallel_degree,
     )
 
     main(args, evaluator=evaluator)

examples/benchmark/mmlu/evaluator.py

@@ -23,10 +23,21 @@
 
 
 class ModelEvaluator(object):
-    def __init__(self, model_name_or_path, ntrain, temperature=0.2, dtype="float32"):
+    def __init__(self, model_name_or_path, ntrain, temperature=0.2, dtype="float32", tensor_parallel_degree=1):
         self.model_name_or_path = model_name_or_path
         self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path, dtype=dtype, low_cpu_mem_usage=True)
+        self.tensor_parallel_degree = tensor_parallel_degree
+        if self.tensor_parallel_degree > 1:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name_or_path,
+                dtype=dtype,
+                low_cpu_mem_usage=True,
+                tensor_parallel_output=False,
+                tensor_parallel_degree=self.tensor_parallel_degree,
+                tensor_parallel_rank=paddle.distributed.get_rank(),
+            )
+        else:
+            self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path, dtype=dtype, low_cpu_mem_usage=True)
         self.model.eval()
         self.generation_config = dict(
             temperature=temperature,