[GPT-3] Support tensor model parallel in static graph. (#2245)

Co-authored-by: Zhong Hui <[email protected]>

Author: wangxicoding

examples/language_model/gpt-3/static/modeling.py

@@ -910,17 +910,20 @@ def _init_generation_caches(self, src_ids):
 
     def parallel_matmul(self, lm_output, logit_weights, parallel_output, topo):
         if topo is not None and topo.mp_info.size > 1:
+            hybrid_groups = fleet.get_hybrid_communicate_group()
+            model_parallel_group = hybrid_groups.get_model_parallel_group()
+
             input_parallel = paddle.distributed.collective._c_identity(
-                lm_output, group=None)
+                lm_output, group=model_parallel_group)
 
             logits = paddle.matmul(
                 input_parallel, logit_weights, transpose_y=True)
 
             if parallel_output:
                 return logits
 
-            # TODO(qinqing): collective._c_concat is not support in static graph now
-            return paddle.distributed.collective._c_concat(logits, group=None)
+            return paddle.distributed.collective._c_concat(
+                logits, group=model_parallel_group)
         else:
             logits = paddle.matmul(lm_output, logit_weights, transpose_y=True)
             return logits

examples/language_model/gpt-3/static/run_gen.sh

@@ -19,10 +19,6 @@ python -u  -m paddle.distributed.fleet.launch \
     --max_seq_len 1024 \
     --micro_batch_size 2 \
     --global_batch_size 2 \
-    --sharding_degree 1 \
-    --mp_degree 1 \
-    --dp_degree 1 \
-    --pp_degree 1 \
     --max_dec_len 20 \
     --decoding_strategy 'topk_sampling' \
     --topp 0.9 \

examples/language_model/gpt-3/static/run_generation.py

@@ -114,11 +114,26 @@ def do_generation(args):
     # Initialize the paddle and paddle fleet execute environment
     paddle.enable_static()
 
+    assert args.dp_degree == 1, "Data parallel is not supported in inference"
+    assert args.sharding_degree == 1, "Sharding parallel is temporarily not supported in inference"
+    assert args.pp_degree == 1, "Pipeline parallel will be supported later"
+
+    if args.mp_degree == 1:
+        args.mp_degree = paddle.distributed.get_world_size()
+    else:
+        assert args.mp_degree == paddle.distributed.get_world_size(), \
+            "If mp_degree is specified, the size must be the same as world_size"
+
     strategy = fleet.DistributedStrategy()
-    strategy.hybrid_configs = {"dp_degree": 1, "mp_degree": 2, "pp_degree": 1}
+    strategy.tensor_parallel = True
+    strategy.tensor_parallel_configs = {
+        "tensor_parallel_degree": args.mp_degree
+    }
+
     fleet.init(is_collective=True, strategy=strategy)
 
-    group = paddle.distributed.init_parallel_env()
+    # temp use dynamic init, use HybridParallelInferenceHelper in future?
+    paddle.distributed.init_parallel_env()
 
     # Create the random seed for the worker
     random.seed(args.seed)