[bert][gpu][graph_kernel]set AllReduce split

model_zoo/official/nlp/bert/run_pretrain.py

@@ -40,13 +40,15 @@ from src.utils import LossCallBack, BertLearningRate
 _current_dir = os.path.dirname(os.path.realpath(__file__))
 
 
-def _set_bert_all_reduce_split():
+def _set_bert_all_reduce_split(device_target='Ascend', enable_graph_kernel=False):
     """set bert all_reduce fusion split, support num_hidden_layers is 12 and 24."""
     if bert_net_cfg.num_hidden_layers == 12:
         if bert_net_cfg.use_relative_positions:
             context.set_auto_parallel_context(all_reduce_fusion_config=[29, 58, 87, 116, 145, 174, 203, 217])
         else:
             context.set_auto_parallel_context(all_reduce_fusion_config=[28, 55, 82, 109, 136, 163, 190, 205])
+            if device_target == 'GPU' and enable_graph_kernel:
+                context.set_auto_parallel_context(all_reduce_fusion_config=[180, 205])
     elif bert_net_cfg.num_hidden_layers == 24:
         if bert_net_cfg.use_relative_positions:
             context.set_auto_parallel_context(all_reduce_fusion_config=[30, 90, 150, 210, 270, 330, 390, 421])
@@ -137,6 +139,9 @@ def run_pretrain():
     args_opt = parser.parse_args()
     context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
     context.set_context(reserve_class_name_in_scope=False)
+    is_auto_enable_graph_kernel = _auto_enable_graph_kernel(args_opt.device_target, args_opt.enable_graph_kernel)
+    if args_opt.enable_graph_kernel == "true" or is_auto_enable_graph_kernel:
+        context.set_context(enable_graph_kernel=True)
     ckpt_save_dir = args_opt.save_checkpoint_path
     if args_opt.distribute == "true":
         if args_opt.device_target == 'Ascend':
@@ -152,17 +157,11 @@ def run_pretrain():
         context.reset_auto_parallel_context()
         context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True,
                                           device_num=device_num)
-        if args_opt.device_target == 'Ascend':
-            _set_bert_all_reduce_split()
+        _set_bert_all_reduce_split(args_opt.device_target, context.get_context('enable_graph_kernel'))
     else:
         rank = 0
         device_num = 1
 
-    is_auto_enable_graph_kernel = _auto_enable_graph_kernel(args_opt.device_target, args_opt.enable_graph_kernel)
-
-    if args_opt.enable_graph_kernel == "true" or is_auto_enable_graph_kernel:
-        context.set_context(enable_graph_kernel=True)
-
     if args_opt.device_target == 'GPU' and bert_net_cfg.compute_type != mstype.float32 and \
         not is_auto_enable_graph_kernel:
         logger.warning('Gpu only support fp32 temporarily, run with fp32.')