version 1

run_pretrain.py

+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from init_env import get_context
+
+context = get_context()
+num_trainers = None
+trainer_id = None
+
+import argparse
+import collections
+import itertools
+import math
+import os
+import random
+import time
+import h5py
+import json
+import distutils.util
+from functools import partial
+from concurrent.futures import ThreadPoolExecutor
+from mlperf_logging.mllog import constants
+import paddle.fluid.profiler as profiler
+#from mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+from models.mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+
+import numpy as np
+
+if context.is_trainer:
+    import utility
+    import dataset
+    import models
+    import paddle
+    import paddle.fluid.core as core
+    import paddle.distributed.fleet as fleet
+    from paddle.fluid.clip import _allow_pure_fp16_global_norm_clip
+    from paddle.fluid.contrib.mixed_precision.fp16_utils import _keep_layer_norm_scale_bias_to_fp32
+    from paddle.io import DataLoader, Dataset
+    from paddle.fluid.layer_helper import LayerHelper
+
+    from models.modeling import BertForPretraining, BertModel, BertPretrainingCriterion
+    from models.modeling import BertConfig
+    from models.optimization import LinearWarmupPolyDecayScheduler
+    from dataset import create_data_holder, create_pretraining_dataset, create_cpu_exchange_padding_pretraining_dataset, create_new_eval_dataset
+
+    try:
+        from paddle.incubate.optimizer import DistributedFusedLamb
+    except ImportError:
+        print('DistributedFusedLamb import error')
+
+    try:
+        from custom_setup_ops import custom_lr
+    except ImportError as e:
+        print('custom_setup_ops import error: {}'.format(e))
+import pdb
+
+def append_lr_op(base_lr,
+                 max_step,
+                 startup_warmup_step=0,
+                 warmup_step=0,
+                 end_lr=0.0,
+                 degree=1.0):
+    helper = LayerHelper('lr_op')
+    step_var = paddle.fluid.layers.create_global_var(
+        shape=[1], value=0, dtype='int64', persistable=True)
+    lr_var = helper.create_variable_for_type_inference(dtype=np.float32)
+    helper.append_op(
+        type='custom_lr',
+        inputs={'X': [step_var]},
+        outputs={'Out': [lr_var]},
+        attrs={
+            'base_lr': base_lr,
+            'end_lr': end_lr,
+            'degree': degree,
+            'start_warmup_step': startup_warmup_step,
+            'warmup_step': warmup_step,
+            'max_step': max_step
+        })
+    return lr_var, step_var
+
+
+def append_acc_merge_op(eval_program, acc, total):
+    block = eval_program.global_block()
+    acc = block.vars.get(acc.name)
+    total = block.vars.get(total.name)
+    startup_program = paddle.static.Program()
+    with paddle.static.program_guard(eval_program, startup_program):
+        acc_step = paddle.static.create_global_var(
+            name="eval_step",
+            shape=[2],
+            value=0,
+            dtype=paddle.int64,
+            persistable=True,
+            force_cpu=True)
+        acc_out = paddle.static.create_global_var(
+            name="eval_acc",
+            shape=[2],
+            value=0,
+            dtype=paddle.float64,
+            persistable=True)
+        helper = LayerHelper('acc_merge_op')
+        helper.append_op(
+            type='acc_merge',
+            inputs={'Acc': [acc],
+                    'Total': [total]},
+            outputs={'Out': [acc_out],
+                     'Step': [acc_step]})
+    scope = utility.get_scope()
+    t = scope.var(acc_step.name).get_tensor()
+    set_eval_step_func = lambda v: t.set(np.array([0, v], dtype=np.int64), paddle.CPUPlace())
+    set_eval_step_func(0)
+    return [acc_out], set_eval_step_func
+
+
+def save_env(local_file):
+    with open(local_file, "w") as f:
+        f.write(json.dumps(dict(os.environ), sort_keys=True, indent=2))
+
+
+def print_flags():
+    print(json.dumps(dict(core.globals()), sort_keys=True, indent=2))
+
+
+def initial_loss_scale():
+    return float(os.getenv("INIT_LOSS_SCALE", 2**20))
+
+
+def str2bool(s):
+    return True if distutils.util.strtobool(s) else False
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_dir",
+        default=None,
+        type=str,
+        required=True,
+        help="The input directory where the data will be read from.", )
+    parser.add_argument(
+        "--eval_dir",
+        default=None,
+        type=str,
+        help="The eval data dir. Should contain .hdf5 files  for the task.")
+    parser.add_argument(
+        "--num_eval_examples",
+        default=10000,
+        type=int,
+        help="number of eval examples to run eval on")
+    parser.add_argument(
+        "--max_predictions_per_seq",
+        default=80,
+        type=int,
+        help="The maximum total of masked tokens in input sequence")
+
+    parser.add_argument(
+        "--train_batch_size",
+        default=8,
+        type=int,
+        help="Batch size per GPU/CPU for training.", )
+    parser.add_argument(
+        "--eval_batch_size",
+        default=8,
+        type=int,
+        help="Batch size per GPU/CPU for evaluation.")
+    parser.add_argument(
+        "--learning_rate",
+        default=5e-5,
+        type=float,
+        help="The initial learning rate for Adam.")
+    parser.add_argument(
+        "--weight_decay_rate",
+        default=0.01,
+        type=float,
+        help="Weight decay if we apply some.")
+    parser.add_argument(
+        "--lamb_epsilon",
+        default=1e-6,
+        type=float,
+        help="Epsilon for Adam optimizer.")
+    parser.add_argument(
+        "--opt_lamb_beta_1", default=0.9, type=float, help="LAMB beta1.")
+    parser.add_argument(
+        "--opt_lamb_beta_2", default=0.999, type=float, help="LAMB beta2.")
+    parser.add_argument(
+        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument(
+        "--max_steps",
+        default=-1,
+        type=int,
+        help="If > 0: set total number of training steps to perform. Override num_train_epochs.",
+    )
+    parser.add_argument(
+        "--warmup_proportion",
+        default=0.0,
+        type=float,
+        help="Linear warmup step proportion.")
+    parser.add_argument(
+        "--start_warmup_step",
+        default=0,
+        type=int,
+        help="The default start warmup steps.")
+    parser.add_argument(
+        "--warmup_steps",
+        default=0,
+        type=int,
+        help="Linear warmup over warmup_steps.")
+    parser.add_argument(
+        "--log_freq", type=int, default=500, help="Log every X updates steps.")
+    parser.add_argument(
+        "--seed", type=int, default=42, help="Random seed for initialization")
+    parser.add_argument(
+        "--use_amp",
+        type=str2bool,
+        default=False,
+        help="Enable mixed precision training.")
+    parser.add_argument(
+        "--enable_addto",
+        type=str2bool,
+        default=False,
+        help="Whether to enable the addto strategy.")
+    parser.add_argument(
+        "--use_pure_fp16",
+        type=str2bool,
+        default=False,
+        help="Whether to use pure fp16 training.")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="gpu",
+        help="Device for selecting for the training.")
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of merge steps before gradient update."
+        "global_batch_size = gradient_accumulation_steps * batch_size.")
+    parser.add_argument(
+        "--tf_ckpt_path",
+        type=str,
+        default=None,
+        help="The pickled checkpoint path of TensorFlow.")
+    parser.add_argument(
+        "--bert_config_path",
+        type=str,
+        default=None,
+        help="The bert config path.")
+    parser.add_argument(
+        "--unpad",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad optimization.")
+    parser.add_argument(
+        "--pad",
+        type=str2bool,
+        default=False,
+        help="Whether to use pad optimization.")
+    parser.add_argument(
+        "--unpad_fmha",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad_fmha optimization.")
+    parser.add_argument(
+        "--unpad_fmha_mke_opt",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad_fmha_mke_opt optimization.")
+    parser.add_argument(
+        "--pad_fmha",
+        type=str2bool,
+        default=False,
+        help="Whether to use pad_fmha optimization.")
+    parser.add_argument(
+        "--fused_bias_mha",
+        type=str2bool,
+        default=False,
+        help="Whether to use fused_bias_mha optimization.")
+    parser.add_argument(
+        "--fused_bias_fc",
+        type=str2bool,
+        default=False,
+        help="Whether to use fused_bias_fc optimization.")
+    parser.add_argument(
+        "--fused_dropout_add_ln",
+        type=str2bool,
+        default=False,
+        help="Whether to use fused_dropout_add_ln optimization.")
+    parser.add_argument(
+        "--weight_transpose",
+        type=str2bool,
+        default=False,
+        help="Whether the weight of linear is stored in tranpose format.")
+    parser.add_argument(
+        "--max_seq_length", type=int, default=512, help="The max seq length.")
+    parser.add_argument(
+        "--batch_size", type=int, default=56, help="The batch size.")
+    parser.add_argument(
+        "--num_epochs_to_generate_seeds_for",
+        type=int,
+        default=2,
+        help="Number of epochs to plan seeds for. Same set across all workers.")
+    parser.add_argument(
+        "--distributed_lamb",
+        type=str2bool,
+        default=False,
+        help="Whether to use distributed LAMB optimizer.")
+    parser.add_argument(
+        "--exchange_padding",
+        type=str2bool,
+        default=False,
+        help="Whether to exchange padding across devices.")
+    parser.add_argument(
+        "--cpu_exchange_padding",
+        type=str2bool,
+        default=False,
+        help="Whether to use CPU to do exchange padding.")
+    parser.add_argument(
+        "--target_mlm_accuracy",
+        type=float,
+        default=0.72,
+        help="The target mlm accuracy to be coveraged.")
+    parser.add_argument(
+        "--max_samples_termination",
+        type=int,
+        default=4500000,
+        help="The max samples threshold to terminate the training process.")
+    parser.add_argument(
+        "--use_uncompressed_dataset",
+        type=str2bool,
+        default=False,
+        help="Whether to use the uncompressed dataset.")
+    parser.add_argument(
+        "--dense_seq_output",
+        type=str2bool,
+        default=False,
+        help="Whether to use the dense_seq_output optimization.")
+    parser.add_argument(
+        "--unpad_embed",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad_embed optimization.")
+    parser.add_argument(
+        "--sort_eval_data",
+        type=str2bool,
+        default=False,
+        help="Whether to sort the eval data.")
+    args = parser.parse_args()
+    return args
+
+
+'''
+def select_dataset_file_for_each_worker(files, f_start_id, worker_num,
+                                        worker_index):
+    """  
+    Spliting the train file according to the worker index.
+    """
+    num_files = len(files)
+    if worker_num > num_files:
+        remainder = worker_num % num_files
+        data_file = files[(
+            f_start_id * worker_num + worker_index + remainder * f_start_id) %
+                          num_files]
+    else:
+        data_file = files[(f_start_id * worker_num + worker_index) % num_files]
+    return data_file
+'''
+
+
+def create_strategy(args):
+    """
+    Create build strategy and exec strategy.
+    """
+    build_strategy = paddle.static.BuildStrategy()
+    exec_strategy = paddle.static.ExecutionStrategy()
+
+    build_strategy.enable_addto = args.enable_addto
+    build_strategy.allow_cuda_graph_capture = False
+    if args.distributed_lamb:
+        build_strategy.fuse_all_reduce_ops = False
+        build_strategy.reduce_strategy = paddle.static.BuildStrategy.ReduceStrategy._NoReduce
+    else:
+        build_strategy.fuse_all_reduce_ops = True
+        build_strategy.reduce_strategy = paddle.static.BuildStrategy.ReduceStrategy.AllReduce
+
+    #build_strategy.fuse_gemm_epilogue = True
+    build_strategy.fix_op_run_order = True
+    exec_strategy.num_threads = 1
+    exec_strategy.num_iteration_per_drop_scope = 10000
+    return build_strategy, exec_strategy
+
+
+def dist_optimizer(args, optimizer):
+    """
+    Create a distributed optimizer based on a normal optimizer
+    """
+    build_strategy, exec_strategy = create_strategy(args)
+
+    dist_strategy = fleet.DistributedStrategy()
+    if args.distributed_lamb:
+        dist_strategy.gradient_scale_configs = {'scale_strategy': 'sum'}
+    dist_strategy.execution_strategy = exec_strategy
+    dist_strategy.build_strategy = build_strategy
+    if args.distributed_lamb:
+        dist_strategy.fuse_all_reduce_ops = False
+    else:
+        dist_strategy.fuse_all_reduce_ops = True
+
+    dist_strategy.fuse_grad_size_in_MB = 0
+    if args.use_amp:
+        dist_strategy.amp = True
+        custom_white_list = [
+          # 'softmax',
+          # 'softmax_grad',
+          # 'softmax_with_cross_entropy',
+          # 'softmax_with_cross_entropy_grad',
+          #  'reduce_sum',
+          #  'reduce_sum_grad',
+        ]
+
+        custom_black_list = [
+            'arg_max',
+            'concat',
+            'cumsum',
+            'dropout',
+            'dropout_grad',
+            'elementwise_add',
+            'elementwise_add_grad',
+            'elementwise_div',
+            'elementwise_div_grad',
+            'elementwise_max',
+            'elementwise_mul',
+            'elementwise_mul_grad',
+            'elementwise_sub',
+            'fill_any_like',
+            'fill_constant',
+            'gather',
+            'gather_grad',
+            'gelu',
+            'gelu_grad',
+            'lamb',
+            'layer_norm',
+            'layer_norm_grad',
+            'lookup_table_v2',
+            'lookup_table_v2_grad',
+            'matmul',
+            'matmul_grad',
+            'matmul_v2',
+            'matmul_v2_grad',
+            'reduce_any',
+            'reduce_sum',
+             'reduce_sum_grad',
+            'reshape2',
+            'reshape2_grad',
+            'scatter',
+            'scatter_grad',
+            'slice',
+            'slice_grad',
+            'softmax',
+            'softmax_grad',
+            'softmax_with_cross_entropy',
+            'softmax_with_cross_entropy_grad',
+            'sqrt',
+            'square',
+            'sum',
+            'tanh',
+            'tanh_grad',
+            'transpose2',
+            'custom_fmha',
+            'custom_fmha_grad',
+            'custom_fused_dropout_residual_ln',
+            'custom_fused_dropout_residual_ln_grad',
+            'custom_fused_dense',
+            'custom_fused_dense_grad',
+            'reduce_mean_grad',
+            'reduce_mean',
+          #  'reduce_sum',
+          #  'reduce_sum_grad',
+            'cast',
+            'uniform_random',
+            'dropout',
+            'scale',
+            'unsqueeze2',
+            'index_select',
+        ]
+
+        dist_strategy.amp_configs = {
+            'custom_white_list': custom_white_list,
+            'custom_black_list': custom_black_list,
+            'init_loss_scaling': initial_loss_scale(),
+            'incr_every_n_steps': 2000,
+            'decr_every_n_nan_or_inf': 1,
+            'incr_ratio': 2.0,
+            'decr_ratio': 0.5,
+            'use_dynamic_loss_scaling': True,
+            'use_pure_fp16': args.use_pure_fp16,
+            'use_fp16_guard': args.use_pure_fp16,
+        }
+    optimizer = fleet.distributed_optimizer(optimizer, strategy=dist_strategy)
+    return optimizer
+
+
+def prune_exchange_padding_op(eval_program):
+    ops = eval_program.global_block().ops
+    assert ops[0].type == "sort_bert_inputs_across_devices"
+    eval_program.global_block()._remove_op(0)
+    eval_program._sync_with_cpp()
+
+
+def do_eval(exe, eval_program, eval_dataloader, eval_dataloader_cpu_tensor,
+            fetch_list):
+    acc = None 
+    n = len(eval_dataloader)
+    for i, batch in enumerate(eval_dataloader):
+        batch[0]['host_prefix_sum_seq_len'] = eval_dataloader_cpu_tensor[i][0]
+        if i + 1 == n:
+            acc = exe.run(eval_program, feed=batch, fetch_list=fetch_list)[0]
+        else:
+            exe.run(eval_program, feed=batch)
+
+    if num_trainers > 1:
+        accl = np.empty(
+            shape=[2],dtype=np.float64)
+        context.trainer_comm.Allreduce(acc, accl)
+    return accl[0] / accl[1]
+
+
+def get_found_nan_inf_flag_var(prog):
+    flag_var = None
+    for op in prog.global_block().ops:
+        if op.type == "update_loss_scaling":
+            cur_flag_var = op.input("FoundInfinite")
+            assert len(cur_flag_var) == 1
+            cur_flag_var = cur_flag_var[0]
+            if flag_var is None:
+                flag_var = cur_flag_var
+                flag_var = prog.global_block().vars.get(flag_var)
+                assert flag_var is not None
+            else:
+                assert flag_var.name == cur_flag_var
+    return flag_var
+
+
+def get_global_step_var(prog):
+    flag_var = None
+    for op in prog.global_block().ops:
+        if op.type == "distributed_fused_lamb":
+            cur_flag_var = op.output("Step")
+            assert len(cur_flag_var) == 1
+            cur_flag_var = cur_flag_var[0]
+            if flag_var is None:
+                flag_var = cur_flag_var
+                flag_var = prog.global_block().vars.get(flag_var)
+                assert flag_var is not None
+            else:
+                assert flag_var.name == cur_flag_var
+    return flag_var
+
+
+def get_loss_scaling_var_names_to_restored(prog):
+    def retrieve_var_names(loss_scale_op):
+        loss_scale_var_in = loss_scale_op.input("PrevLossScaling")
+        loss_scale_var_out = loss_scale_op.output("LossScaling")
+        assert len(loss_scale_var_in) == 1 and len(
+            loss_scale_var_out) and loss_scale_var_in[0] == loss_scale_var_out[
+                0]
+
+        good_steps_var_in = loss_scale_op.input("InGoodSteps")
+        good_steps_var_out = loss_scale_op.output("OutGoodSteps")
+        assert len(good_steps_var_in) == 1 and len(
+            good_steps_var_out) == 1 and good_steps_var_in[
+                0] == good_steps_var_out[0]
+
+        bad_steps_var_in = loss_scale_op.input("InBadSteps")
+        bad_steps_var_out = loss_scale_op.output("OutBadSteps")
+        assert len(bad_steps_var_in) == 1 and len(
+            bad_steps_var_out) == 1 and bad_steps_var_in[
+                0] == bad_steps_var_out[0]
+        return loss_scale_var_in[0], good_steps_var_in[0], bad_steps_var_in[0]
+
+    loss_scale_var, good_steps_var, bad_steps_var = None, None, None
+    for op in prog.global_block().ops:
+        if op.type == 'update_loss_scaling':
+            cur_loss_scale_var, cur_good_steps_var, cur_bad_steps_var = retrieve_var_names(
+                op)
+            if loss_scale_var is None:
+                loss_scale_var = cur_loss_scale_var
+                good_steps_var = cur_good_steps_var
+                bad_steps_var = cur_bad_steps_var
+
+            assert loss_scale_var == cur_loss_scale_var
+            assert good_steps_var == cur_good_steps_var
+            assert bad_steps_var == cur_bad_steps_var
+
+    assert loss_scale_var is not None
+    return loss_scale_var, good_steps_var, bad_steps_var
+
+
+def get_loss_scaling_var_values(prog):
+    scope = utility.get_scope()
+    var_names = get_loss_scaling_var_names_to_restored(prog)
+    data = []
+    for name in var_names:
+        value = np.array(scope.find_var(name).get_tensor())
+        data.append(value)
+        if trainer_id == 0:
+            print('Var {} value is: {}'.format(name, value))
+    return data
+
+
+def restore_loss_scaling_var_values(prog, data, place):
+    scope = utility.get_scope()
+    var_names = get_loss_scaling_var_names_to_restored(prog)
+    assert len(var_names) == len(data)
+    for name, np_arr in zip(var_names, data):
+        var = scope.find_var(name).get_tensor()
+        var.set(np_arr, place)
+
+
+def paddle_dtype_to_np_dtype(pd_dtype):
+    if pd_dtype == paddle.float32:
+        return np.float32
+    elif pd_dtype == paddle.int64:
+        return np.int64
+    else:
+        raise ValueError('Unsupported tensor dtype {}'.format(pd_dtype))
+
+
+def inplace_fill_constant(tensor, value):
+    dtype = paddle_dtype_to_np_dtype(tensor._dtype())
+    shape = tensor.shape()
+    np_value = (np.ones(shape) * value).astype(dtype)
+    place = tensor._place()
+    if place._equals(paddle.CPUPlace()):
+        place = paddle.CPUPlace()
+    elif place._equals(utility.get_place()):
+        place = utility.get_place()
+    else:
+        raise ValueError("Unsupported tensor place {}".format(place))
+
+    old_ptr = tensor._ptr()
+    tensor.set(np_value, place)
+    new_ptr = tensor._ptr()
+    assert old_ptr == new_ptr
+
+
+def recover_lamb_status(train_prog, args):
+    scope = utility.get_scope()
+
+    def fill_input_value(op, input_slot, value):
+        var_names = op.input(input_slot)
+        assert len(var_names) == 1
+        t = scope.find_var(var_names[0]).get_tensor()
+        inplace_fill_constant(t, value)
+
+    def fill_output_value(op, output_slot, value):
+        var_names = op.output(output_slot)
+        assert len(var_names) == 1
+        t = scope.find_var(var_names[0]).get_tensor()
+        inplace_fill_constant(t, value)
+
+    found_cnt = 0
+    for i in range(train_prog.num_blocks):
+        block = train_prog.block(i)
+        for op in block.ops:
+            if args.distributed_lamb:
+                assert op.type != 'lamb'
+            else:
+                assert op.type != 'distributed_fused_lamb'
+
+            if op.type != 'lamb' and op.type != 'distributed_fused_lamb':
+                continue
+
+            fill_input_value(op, "Moment1", 0)
+            fill_input_value(op, "Moment2", 0)
+            fill_input_value(op, "Beta1Pow", args.opt_lamb_beta_1)
+            fill_input_value(op, "Beta2Pow", args.opt_lamb_beta_2)
+            if op.type == 'distributed_fused_lamb':
+                fill_output_value(op, "Step", 0)
+                if args.gradient_accumulation_steps > 1:
+                    fill_output_value(op, "AccStep", 0)
+            found_cnt += 1
+
+    if args.distributed_lamb:
+        assert found_cnt == 1
+    else:
+        assert found_cnt > 1
+
+
+def run_warmup(args,
+               exe,
+               train_prog,
+               lr_scheduler,
+               eval_prog=None,
+               warmup_iter=8):
+    train_warmup_data = dataset.prepare_warmup_data(args, args.train_batch_size,
+                                                    exe.place)
+    list_train_warmup_data = [train_warmup_data]
+
+    paddle.device.cuda.synchronize()
+    restored_values = get_loss_scaling_var_values(train_prog)
+    old_lr_func = lr_scheduler.__call__
+    lr_scheduler.__call__ = lambda self: 0.0
+    for _ in range(warmup_iter):
+        exe.run(train_prog, feed=list_train_warmup_data)
+    lr_scheduler.__call__ = old_lr_func
+    restore_loss_scaling_var_values(train_prog, restored_values, exe.place)
+    paddle.device.cuda.synchronize()
+
+    # extra run as NV code
+    exe.run(train_prog, feed=list_train_warmup_data)
+    paddle.device.cuda.synchronize()
+
+    if eval_prog is not None:
+        eval_warmup_data = dataset.prepare_warmup_data(
+            args, args.eval_batch_size, exe.place)
+        for _ in range(warmup_iter):
+            exe.run(eval_prog, feed=[eval_warmup_data])
+        paddle.device.cuda.synchronize()
+
+    #recover_lamb_moment_pows(train_prog, args)
+    recover_lamb_status(train_prog, args)
+    paddle.device.cuda.synchronize()
+    if utility.get_trainer_id() == 0:
+        print('warmup ends')
+
+
+def do_train(args):
+    # Initialize the paddle and paddle fleet execute enviroment
+    if trainer_id == 0:
+        paddle_bert_print_start(key=constants.INIT_START)
+        paddle_bert_print_event(key=constants.SEED, val=args.seed)
+        paddle_bert_print_event(
+            key=constants.GLOBAL_BATCH_SIZE,
+            val=args.train_batch_size * num_trainers *
+            args.gradient_accumulation_steps)
+        paddle_bert_print_event(key='d_batch_size', val=args.train_batch_size)
+        paddle_bert_print_event(
+            key=constants.GRADIENT_ACCUMULATION_STEPS,
+            val=args.gradient_accumulation_steps)
+        paddle_bert_print_event(
+            key="max_predictions_per_seq", val=args.max_predictions_per_seq)
+        paddle_bert_print_event(
+            key='opt_learning_rate_training_steps', val=args.max_steps)
+        paddle_bert_print_event(key='num_warmup_steps', val=args.warmup_steps)
+
+    place = utility.get_place()
+    fleet.init(is_collective=True)
+
+    worker_num = num_trainers
+    worker_index = trainer_id
+
+    # Define the input data in the static mode
+    main_program = paddle.static.default_main_program()
+    startup_program = paddle.static.default_startup_program()
+
+    data_holders, data_inputs, data_labels, varlen_info, mlm_label_info = create_data_holder(
+        args)
+
+    config = BertConfig.from_json_file(args.bert_config_path)
+    #config.fused_mha = args.fused_mha
+    #config.fused_gelu_bias = args.fused_gelu_bias
+    config.fused_bias_mha = args.fused_bias_mha
+    config.fused_bias_fc = args.fused_bias_fc
+    #config.fused_bias_fc_loss_head = args.fused_bias_fc_loss_head
+    config.dense_seq_output = args.dense_seq_output
+    config.unpad = args.unpad
+    config.unpad_fmha = args.unpad_fmha
+    config.unpad_fmha_mke_opt = args.unpad_fmha_mke_opt
+    config.pad_fmha = args.pad_fmha
+    config.max_seq_length = args.max_seq_length
+    #config.pad = args.pad
+    #config.fuse_qkv = not args.disable_fuse_qkv
+    #config.fuse_scale = not args.disable_fuse_scale
+    #config.fuse_mask = not args.disable_fuse_mask 
+    #config.fuse_dropout = args.enable_import paddle.fluid.profiler as profilerfuse_dropout
+    config.fused_dropout_add_ln = args.fused_dropout_add_ln
+    #config.apex_softmax = not args.disable_apex_softmax
+    #config.enable_stream = args.enable_stream
+    #if config.fuse_mask: config.apex_softmax = True
+    #if not config.pad: config.enable_stream = True
+    #if config.unpad: config.fused_mha = False
+
+    config.unpad_embed = args.unpad_embed
+
+    config.weight_transpose = args.weight_transpose
+    config.batch_size = args.batch_size
+
+    # Padding for divisibility by 8
+    if config.vocab_size % 8 != 0:
+        config.vocab_size += 8 - (config.vocab_size % 8)
+
+    int16_max = np.iinfo(np.int16).max
+    assert config.vocab_size <= int16_max
+    assert args.train_batch_size * args.max_seq_length <= int16_max
+    assert args.eval_batch_size * args.max_seq_length <= int16_max
+
+    def build_model():
+        model = BertForPretraining(BertModel(config), config)
+        criterion = BertPretrainingCriterion(config)
+
+        prediction_scores, seq_relationship_score = model(*data_inputs)
+        func = criterion(prediction_scores, seq_relationship_score,
+                         *data_labels)
+        return model, func
+
+    if args.use_pure_fp16:
+        with paddle.static.amp.fp16_guard():
+            model, func = build_model()
+    else:
+        model, func = build_model()
+    loss, mlm_acc, num_masked = func()
+
+    eval_program = main_program.clone(for_test=True)
+    if args.gpu_exchange_padding and utility.use_nv_input(
+    ) and num_trainers > 1:
+        prune_exchange_padding_op(eval_program)
+
+    # Define the dynamic learing_reate scheduler and optimizer
+    num_training_steps = args.max_steps if args.max_steps > 0 else len(
+        train_dataloader) * args.num_train_epochs
+    if args.warmup_steps == 0:
+        warmup_steps = int(args.max_steps * args.warmup_proportion)
+        warmup_start = 0
+    else:
+        warmup_steps = args.warmup_steps
+        warmup_start = args.start_warmup_step
+
+    lr_scheduler = LinearWarmupPolyDecayScheduler(
+        startup_warmup_steps=warmup_start,
+        warmup_steps=warmup_steps,
+        total_steps=args.max_steps,
+        base_lr=args.learning_rate,
+        end_lr=0.0,
+        degree=1.0)
+
+    # Generate parameter names needed to perform weight decay.
+    # All bias and LayerNorm parameters are excluded.
+    decay_params = [
+        p.name for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    if args.max_grad_norm > 0:
+        grad_clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=args.max_grad_norm)
+    else:
+        grad_clip = None
+    exclude_from_weight_decay_fn = lambda var: var.name not in decay_params
+
+    lr_var, step_var = append_lr_op(
+        base_lr=args.learning_rate,
+        max_step=args.max_steps,
+        startup_warmup_step=warmup_start,
+        warmup_step=warmup_steps,
+        end_lr=0.0,
+        degree=1.0)
+    optimizer_kwargs = {
+        'learning_rate': lr_var,
+        'lamb_weight_decay': args.weight_decay_rate,
+        'epsilon': args.lamb_epsilon,
+        'exclude_from_weight_decay_fn': exclude_from_weight_decay_fn,
+        'grad_clip': grad_clip,
+        'beta1': args.opt_lamb_beta_1,
+        'beta2': args.opt_lamb_beta_2,
+    }
+
+    if trainer_id == 0:
+        paddle_bert_print_event(
+            key=constants.OPT_BASE_LR, val=args.learning_rate)
+        paddle_bert_print_event(
+            key="opt_learning_rate_warmup_steps", val=args.warmup_steps)
+        paddle_bert_print_event(
+            key="start_warmup_step", val=args.start_warmup_step)
+
+    #_keep_layer_norm_scale_bias_to_fp32(False)
+    _allow_pure_fp16_global_norm_clip(True)
+    if args.distributed_lamb:
+        # Note: 
+        # (1) broadcast_master_param is true means use fp32 to update parameter.
+        # (2) if 'is_grad_scaled_by_nranks' is false, must set 'dist_strategy.gradient_scale_configs = {'scale_strategy': 'sum'}'.
+        #optimizer_kwargs.update({"broadcast_master_param": True, "clip_after_allreduce": False, "is_grad_scaled_by_nranks": False})
+        optimizer_kwargs.update({
+            "use_master_param_norm": True,
+            "clip_after_allreduce": False,
+            "is_grad_scaled_by_nranks": False,
+        })
+        if args.gradient_accumulation_steps > 1:
+            optimizer_kwargs.update({
+                "gradient_accumulation_steps": args.gradient_accumulation_steps
+            })
+        optimizer = DistributedFusedLamb(**optimizer_kwargs)
+        optimizer._set_step(step_var)
+    else:
+        optimizer = paddle.optimizer.Lamb(**optimizer_kwargs)
+        if args.use_pure_fp16:
+            optimizer._multi_precision = True
+
+    if hasattr(optimizer, "_get_parameter"):
+        get_param_func = optimizer._get_parameter
+    else:
+        get_param_func = lambda name: utility.get_scope().find_var(name).get_tensor()
+
+
+    if trainer_id == 0:
+        paddle_bert_print_event(key='opt_epsilon', val=args.lamb_epsilon)
+        paddle_bert_print_event(key='opt_lamb_beta_1', val=args.opt_lamb_beta_1)
+        paddle_bert_print_event(key='opt_lamb_beta_2', val=args.opt_lamb_beta_2)
+        paddle_bert_print_event(
+            key='opt_lamb_weight_decay_rate', val=args.weight_decay_rate)
+
+    # Use the fleet api to compile the distributed optimizer
+    print("optimizer",optimizer._get_parameter)
+
+    optimizer = dist_optimizer(args, optimizer)
+    print("loss",loss)
+
+    #pdb.set_trace()
+    optimizer.minimize(loss)
+
+    # Define the Executor for running the static model
+    exe = paddle.static.Executor(place)
+    exe.run(startup_program)
+
+    if args.use_amp:
+        optimizer.amp_init(place, test_program=eval_program, use_fp16_test=True)
+
+    eval_fetch_list, set_eval_step = append_acc_merge_op(eval_program, mlm_acc,
+                                                         num_masked)
+
+    with open('startup_program_{}.txt'.format(trainer_id), 'w') as f:
+        f.write(str(startup_program))
+
+    with open('main_program_{}.txt'.format(trainer_id), 'w') as f:
+        f.write(str(main_program))
+
+    with open('eval_program_{}.txt'.format(trainer_id), 'w') as f:
+        f.write(str(eval_program))
+
+    found_nan_inf = get_found_nan_inf_flag_var(main_program)
+    found_global_step_val = get_global_step_var(main_program)
+    if args.use_amp:
+        assert found_nan_inf is not None
+    else:
+        assert found_nan_inf is None
+    assert found_global_step_val is not None
+
+    build_strategy, exec_strategy = create_strategy(args)
+    eval_program = paddle.static.CompiledProgram(
+        eval_program).with_data_parallel(
+            build_strategy=build_strategy, exec_strategy=exec_strategy)
+    run_warmup(
+        args,
+        exe,
+        main_program,
+        lr_scheduler,
+        eval_prog=eval_program,
+        warmup_iter=8)
+
+    ckpt = None
+    if args.tf_ckpt_path:
+        if trainer_id == 0:
+            print('starts to load tf checkpoint...')
+            start_t = time.time()
+        ckpt = model.load_tf_ckpt(args, get_param_func)
+        paddle.device.cuda.synchronize()
+        if trainer_id == 0:
+            end_t = time.time()
+            print('checkpoint loading ends, time cost: {}s'.format(end_t -
+                                                                   start_t))
+
+    def save_ckpt(ckpt_saved_path):
+        if trainer_id == 0 and ckpt is not None:
+            dirname = 'log_{}'.format(num_trainers)
+            ckpt_saved_path = os.path.join(dirname, ckpt_saved_path)
+            ckpt.save(ckpt_saved_path, get_param_func)
+            print('PaddlePaddle ckpt saved into: {}'.format(ckpt_saved_path))
+
+    # save_ckpt('loaded_ckpt_0.pickled')
+
+    pool = ThreadPoolExecutor(1)
+    # pre-compute eval boundaries
+    samples_trained_per_step = args.train_batch_size * num_trainers * args.gradient_accumulation_steps
+    eval_interval_samples = 0.05 * (230.23 * samples_trained_per_step + 3000000)
+    eval_interval_samples = int(eval_interval_samples / 25000) * 25000
+    start, stop, step = eval_interval_samples, args.max_samples_termination, eval_interval_samples
+    eval_steps = [
+        math.ceil(i / samples_trained_per_step)
+        for i in np.arange(start, stop, step)
+    ]
+    eval_count = 0
+    next_eval_step = eval_steps[eval_count]
+    next_eval_training_step = next_eval_step * args.gradient_accumulation_steps
+    create_train_dataset = create_cpu_exchange_padding_pretraining_dataset if args.cpu_exchange_padding else create_pretraining_dataset
+
+    context.trainer_comm.Barrier()
+    if trainer_id == 0:
+        paddle_bert_print_end(key=constants.INIT_STOP)
+        paddle_bert_print_start(key=constants.RUN_START)
+
+    f_start_id = 0
+    train_dataloader, train_dataloader_cpu_tensor = create_train_dataset(
+        data_holders, f_start_id, tolist=False)
+
+    if args.eval_dir:
+        eval_dataset_future = pool.submit(create_new_eval_dataset, data_holders)
+    need_next_training_shard = args.train_batch_size * args.gradient_accumulation_steps * args.max_steps > 10000
+
+    global_step = 0
+    skipped_steps = 0
+    training_steps = 0
+    now_skipped = 0
+    now_step = 0
+    skip_interval = 0
+    average_loss = 0.0
+    end_training, converged = False, False
+    coveraged = False
+    epoch = 1  # to be consistent with NV
+
+    file_num = context.file_num()
+
+    while global_step < args.max_steps and not end_training:
+        if trainer_id == 0:
+            paddle_bert_print_start(
+                key=constants.EPOCH_START, metadata={'epoch_num': epoch})
+            paddle_bert_print_start(
+                key=constants.BLOCK_START,
+                metadata={'first_epoch_num': epoch,
+                          'epoch_count': 1})
+            now_time = time.time()
+
+        empty_fetch_list = []
+        full_fetch_list = [loss, mlm_acc, found_global_step_val]
+        for f_id in range(f_start_id + 1, file_num):
+            if need_next_training_shard:
+                dataset_future = pool.submit(create_train_dataset, data_holders,
+                                             f_id)
+            # limin-todo:
+            # limin_file_start = time.time()
+            for step, batch in enumerate(train_dataloader):
+                training_steps += 1
+                #if training_steps == 2:
+                #   exit()
+                #print("training_steps",training_steps)
+                update_step = training_steps % args.gradient_accumulation_steps == 0
+                batch[0][
+                    'host_prefix_sum_seq_len'] = train_dataloader_cpu_tensor[
+                        step][0]
+
+                if training_steps < next_eval_training_step and False:
+                    cur_fetch_list = empty_fetch_list
+                else:
+                    cur_fetch_list = full_fetch_list
+                #with profiler.profiler('All', 'total') as prof:
+                return_val = exe.run(main_program,
+                                     feed=batch,
+                                     fetch_list=cur_fetch_list)
+                if training_steps >= next_eval_training_step or True:
+                    loss_return, mlm_acc_return, found_global_step_return = return_val
+                    average_loss += loss_return[0]
+  
+                if update_step:
+                    if training_steps >= next_eval_training_step:
+#                        global_step = found_global_step_return[0]
+                        global_step = int(training_steps /
+                                          args.gradient_accumulation_steps)           
+                        has_nan_inf_step = (
+                            training_steps / args.gradient_accumulation_steps -
+                            global_step)
+                        if has_nan_inf_step:
+                            skipped_steps = has_nan_inf_step
+                    else:
+                        global_step = int(training_steps /
+                                          args.gradient_accumulation_steps)
+
+                    if args.eval_dir and global_step == next_eval_step:
+                        skip_interval = skipped_steps
+                        now_skipped = skip_interval
+                        if True:
+                            if eval_count == 0:
+                                eval_dataloader, eval_dataloader_cpu_tensor = eval_dataset_future.result(
+                                    timeout=None)
+                                eval_dataloader = list(eval_dataloader)
+                                set_eval_step(len(eval_dataloader))
+                            samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * num_trainers
+                            samples_trained_prev = samples_trained
+
+                            eval_start_t = time.time()
+                            eval_avg_mlm_accuracy = do_eval(
+                                exe, eval_program, eval_dataloader,
+                                eval_dataloader_cpu_tensor, eval_fetch_list)
+                            eval_end_t = time.time()
+                            # if trainer_id == 0:
+                            #    print('Eval time: {} s'.format(eval_end_t -
+                            #                                   eval_start_t))
+
+                            if args.target_mlm_accuracy > 0 and eval_avg_mlm_accuracy >= args.target_mlm_accuracy:
+                                end_training, converged = True, True
+                                coveraged = True
+                                if trainer_id == 0:
+                                    print(
+                                        "%f > %f, Target MLM Accuracy reached at global_step %d , training_step %d"
+                                        % (eval_avg_mlm_accuracy,
+                                           args.target_mlm_accuracy,
+                                           global_step, training_steps))
+                            eval_count += 1
+                            next_eval_step = eval_steps[eval_count]
+                            next_eval_training_step = next_eval_step * args.gradient_accumulation_steps
+                            if trainer_id == 0:
+                                paddle_bert_print_event(
+                                    key=constants.EVAL_ACCURACY,
+                                    val=eval_avg_mlm_accuracy,
+                                    metadata={'epoch_num': epoch})
+                                print({
+                                    "global_steps": global_step,
+                                    "eval_mlm_accuracy": eval_avg_mlm_accuracy
+                                })
+
+                samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * num_trainers
+                if args.log_freq > 0 and training_steps % (
+                        args.log_freq * args.gradient_accumulation_steps) == 0:
+                    samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * num_trainers
+                    #just add the following to print loss
+                    cur_fetch_list = full_fetch_list
+                    return_val = exe.run(main_program,
+                                         feed=batch,
+                                         fetch_list=cur_fetch_list)
+                    loss_return, mlm_acc_return, found_global_step_return = return_val
+                    print("loss_return, mlm_acc_return, found_global_step_return:",loss_return, mlm_acc_return, found_global_step_return[0])
+                    if trainer_id == 0:
+                        time_interval = time.time() - now_time
+                        step_interval = global_step - now_step
+                        now_time = time.time()
+                        now_step = global_step
+                        training_perf = args.train_batch_size * args.gradient_accumulation_steps * num_trainers * (
+                            step_interval + skip_interval) / time_interval
+                        skip_interval = 0
+                        divisor = args.gradient_accumulation_steps
+                        now_lr = lr_scheduler()
+                        step_loss_to_print = loss_return[0] * args.gradient_accumulation_steps / divisor
+                        average_loss_to_print = average_loss / (args.log_freq * divisor)
+                        print({
+                            "training_steps": training_steps,
+                            "average_loss": average_loss_to_print,
+                            "step_loss": step_loss_to_print,
+                            "learning_rate": now_lr,
+                            "seq/s": training_perf,
+                            "global_steps": now_step,
+                            "samples_trained": samples_trained,
+                            "skipped_steps": now_skipped,
+                            "timestamp": now_time,
+                        })
+
+                        paddle_bert_print_event(
+                            key='tracked_stats',
+                            val={
+                                'seq/s': training_perf,
+                                'step_loss': step_loss_to_print,
+                                #'avg_loss': average_loss_to_print,
+                                #'lr': now_lr,
+                            },
+                            metadata={"step": (epoch, training_steps)})
+                        paddle_bert_print_event(
+                            key='throughput', val=training_perf)
+
+                    average_loss = 0.0
+
+                if global_step >= args.max_steps or end_training:
+                    status = 'success' if coveraged else 'abort'
+                    end_training = True
+                    break
+
+            # limin-todo:
+            #limin_file_end = time.time()
+            #print("limin: one file time: ", limin_file_end - limin_file_start)
+            del train_dataloader
+            if samples_trained >= args.max_samples_termination or end_training:
+                status = 'success' if converged else 'aborted'
+                end_training = True
+                break
+
+            if not need_next_training_shard:
+                dataset_future = pool.submit(create_train_dataset, f_id)
+            train_dataloader, train_dataloader_cpu_tensor = dataset_future.result(
+                timeout=None)
+
+        if trainer_id == 0:
+            paddle_bert_print_end(
+                key=constants.BLOCK_STOP, metadata={'first_epoch_num': epoch})
+            paddle_bert_print_end(
+                key=constants.EPOCH_STOP, metadata={'epoch_num': epoch})
+
+        epoch += 1
+
+    if trainer_id == 0:
+        paddle_bert_print_event(constants.TRAIN_SAMPLES, val=samples_trained)
+        paddle_bert_print_event(
+            constants.EVAL_SAMPLES, val=args.num_eval_examples)
+        paddle_bert_print_end(
+            key=constants.RUN_STOP, metadata={'status': status})
+
+        #save_ckpt('saved_ckpt_final.pickled')
+
+
+def do_read(args):
+    while context.read_file():
+        pass
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    if not args.exchange_padding:
+        args.cpu_exchange_padding = False
+        args.gpu_exchange_padding = False
+    else:
+        args.gpu_exchange_padding = not args.cpu_exchange_padding
+
+    context.init_args(args)
+    trainer_id = context.trainer_id
+    num_trainers = context.trainer_num
+
+    if context.is_trainer:
+        if trainer_id == 0:
+            paddle_bert_print_event(key=constants.SUBMISSION_ORG, val="Baidu")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_PLATFORM, val="1 x NVIDIA A100 GPU")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_DIVISION, val="closed")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_STATUS, val="onprem")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_BENCHMARK, val="bert")
+            paddle_bert_print_event(key=constants.CACHE_CLEAR, val=True)
+        save_env('./paddle_env_{}.json'.format(trainer_id))
+        if trainer_id == 0:
+            print(args)
+            print_flags()
+        if args.use_pure_fp16:
+            assert args.use_amp, "--use_amp must be True if --use_pure_fp16 is True"
+        paddle.enable_static()
+        with paddle.static.scope_guard(utility.get_scope()):
+            do_train(args)
+    else:
+        if args.cpu_exchange_padding:
+            do_read(args)
+    if args.cpu_exchange_padding:
+        context.stop_reader()
+    print('Process ends')

run_pretrain_fp32.py

+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from init_env import get_context
+
+context = get_context()
+num_trainers = None
+trainer_id = None
+
+import argparse
+import collections
+import itertools
+import math
+import os
+import random
+import time
+import h5py
+import json
+import distutils.util
+from functools import partial
+from concurrent.futures import ThreadPoolExecutor
+from mlperf_logging.mllog import constants
+import paddle.fluid.profiler as profiler
+#from mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+from models.mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+
+import numpy as np
+
+if context.is_trainer:
+    import utility
+    import dataset
+    import models
+    import paddle
+    import paddle.fluid.core as core
+    import paddle.distributed.fleet as fleet
+    from paddle.fluid.clip import _allow_pure_fp16_global_norm_clip
+    from paddle.fluid.contrib.mixed_precision.fp16_utils import _keep_layer_norm_scale_bias_to_fp32
+    from paddle.io import DataLoader, Dataset
+    from paddle.fluid.layer_helper import LayerHelper
+
+    from models.modeling import BertForPretraining, BertModel, BertPretrainingCriterion
+    from models.modeling import BertConfig
+    from models.optimization import LinearWarmupPolyDecayScheduler
+    from dataset import create_data_holder, create_pretraining_dataset, create_cpu_exchange_padding_pretraining_dataset, create_new_eval_dataset
+
+    try:
+        from paddle.incubate.optimizer import DistributedFusedLamb
+    except ImportError:
+        print('DistributedFusedLamb import error')
+
+    try:
+        from custom_setup_ops import custom_lr
+    except ImportError as e:
+        print('custom_setup_ops import error: {}'.format(e))
+import pdb
+
+def append_lr_op(base_lr,
+                 max_step,
+                 startup_warmup_step=0,
+                 warmup_step=0,
+                 end_lr=0.0,
+                 degree=1.0):
+    helper = LayerHelper('lr_op')
+    step_var = paddle.fluid.layers.create_global_var(
+        shape=[1], value=0, dtype='int64', persistable=True)
+    lr_var = helper.create_variable_for_type_inference(dtype=np.float32)
+    helper.append_op(
+        type='custom_lr',
+        inputs={'X': [step_var]},
+        outputs={'Out': [lr_var]},
+        attrs={
+            'base_lr': base_lr,
+            'end_lr': end_lr,
+            'degree': degree,
+            'start_warmup_step': startup_warmup_step,
+            'warmup_step': warmup_step,
+            'max_step': max_step
+        })
+    return lr_var, step_var
+
+
+def append_acc_merge_op(eval_program, acc, total):
+    block = eval_program.global_block()
+    acc = block.vars.get(acc.name)
+    total = block.vars.get(total.name)
+    startup_program = paddle.static.Program()
+    with paddle.static.program_guard(eval_program, startup_program):
+        acc_step = paddle.static.create_global_var(
+            name="eval_step",
+            shape=[2],
+            value=0,
+            dtype=paddle.int64,
+            persistable=True,
+            force_cpu=True)
+        acc_out = paddle.static.create_global_var(
+            name="eval_acc",
+            shape=[2],
+            value=0,
+            dtype=paddle.float64,
+            persistable=True)
+        helper = LayerHelper('acc_merge_op')
+        helper.append_op(
+            type='acc_merge',
+            inputs={'Acc': [acc],
+                    'Total': [total]},
+            outputs={'Out': [acc_out],
+                     'Step': [acc_step]})
+    scope = utility.get_scope()
+    t = scope.var(acc_step.name).get_tensor()
+    set_eval_step_func = lambda v: t.set(np.array([0, v], dtype=np.int64), paddle.CPUPlace())
+    set_eval_step_func(0)
+    return [acc_out], set_eval_step_func
+
+
+def save_env(local_file):
+    with open(local_file, "w") as f:
+        f.write(json.dumps(dict(os.environ), sort_keys=True, indent=2))
+
+
+def print_flags():
+    print(json.dumps(dict(core.globals()), sort_keys=True, indent=2))
+
+
+def initial_loss_scale():
+    return float(os.getenv("INIT_LOSS_SCALE", 2**20))
+
+
+def str2bool(s):
+    return True if distutils.util.strtobool(s) else False
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_dir",
+        default=None,
+        type=str,
+        required=True,
+        help="The input directory where the data will be read from.", )
+    parser.add_argument(
+        "--eval_dir",
+        default=None,
+        type=str,
+        help="The eval data dir. Should contain .hdf5 files  for the task.")
+    parser.add_argument(
+        "--num_eval_examples",
+        default=10000,
+        type=int,
+        help="number of eval examples to run eval on")
+    parser.add_argument(
+        "--max_predictions_per_seq",
+        default=80,
+        type=int,
+        help="The maximum total of masked tokens in input sequence")
+
+    parser.add_argument(
+        "--train_batch_size",
+        default=8,
+        type=int,
+        help="Batch size per GPU/CPU for training.", )
+    parser.add_argument(
+        "--eval_batch_size",
+        default=8,
+        type=int,
+        help="Batch size per GPU/CPU for evaluation.")
+    parser.add_argument(
+        "--learning_rate",
+        default=5e-5,
+        type=float,
+        help="The initial learning rate for Adam.")
+    parser.add_argument(
+        "--weight_decay_rate",
+        default=0.01,
+        type=float,
+        help="Weight decay if we apply some.")
+    parser.add_argument(
+        "--lamb_epsilon",
+        default=1e-6,
+        type=float,
+        help="Epsilon for Adam optimizer.")
+    parser.add_argument(
+        "--opt_lamb_beta_1", default=0.9, type=float, help="LAMB beta1.")
+    parser.add_argument(
+        "--opt_lamb_beta_2", default=0.999, type=float, help="LAMB beta2.")
+    parser.add_argument(
+        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument(
+        "--max_steps",
+        default=-1,
+        type=int,
+        help="If > 0: set total number of training steps to perform. Override num_train_epochs.",
+    )
+    parser.add_argument(
+        "--warmup_proportion",
+        default=0.0,
+        type=float,
+        help="Linear warmup step proportion.")
+    parser.add_argument(
+        "--start_warmup_step",
+        default=0,
+        type=int,
+        help="The default start warmup steps.")
+    parser.add_argument(
+        "--warmup_steps",
+        default=0,
+        type=int,
+        help="Linear warmup over warmup_steps.")
+    parser.add_argument(
+        "--log_freq", type=int, default=500, help="Log every X updates steps.")
+    parser.add_argument(
+        "--seed", type=int, default=42, help="Random seed for initialization")
+    parser.add_argument(
+        "--use_amp",
+        type=str2bool,
+        default=False,
+        help="Enable mixed precision training.")
+    parser.add_argument(
+        "--enable_addto",
+        type=str2bool,
+        default=False,
+        help="Whether to enable the addto strategy.")
+    parser.add_argument(
+        "--use_pure_fp16",
+        type=str2bool,
+        default=False,
+        help="Whether to use pure fp16 training.")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="gpu",
+        help="Device for selecting for the training.")
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of merge steps before gradient update."
+        "global_batch_size = gradient_accumulation_steps * batch_size.")
+    parser.add_argument(
+        "--tf_ckpt_path",
+        type=str,
+        default=None,
+        help="The pickled checkpoint path of TensorFlow.")
+    parser.add_argument(
+        "--bert_config_path",
+        type=str,
+        default=None,
+        help="The bert config path.")
+    parser.add_argument(
+        "--unpad",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad optimization.")
+    parser.add_argument(
+        "--pad",
+        type=str2bool,
+        default=False,
+        help="Whether to use pad optimization.")
+    parser.add_argument(
+        "--unpad_fmha",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad_fmha optimization.")
+    parser.add_argument(
+        "--unpad_fmha_mke_opt",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad_fmha_mke_opt optimization.")
+    parser.add_argument(
+        "--pad_fmha",
+        type=str2bool,
+        default=False,
+        help="Whether to use pad_fmha optimization.")
+    parser.add_argument(
+        "--fused_bias_mha",
+        type=str2bool,
+        default=False,
+        help="Whether to use fused_bias_mha optimization.")
+    parser.add_argument(
+        "--fused_bias_fc",
+        type=str2bool,
+        default=False,
+        help="Whether to use fused_bias_fc optimization.")
+    parser.add_argument(
+        "--fused_dropout_add_ln",
+        type=str2bool,
+        default=False,
+        help="Whether to use fused_dropout_add_ln optimization.")
+    parser.add_argument(
+        "--weight_transpose",
+        type=str2bool,
+        default=False,
+        help="Whether the weight of linear is stored in tranpose format.")
+    parser.add_argument(
+        "--max_seq_length", type=int, default=512, help="The max seq length.")
+    parser.add_argument(
+        "--batch_size", type=int, default=56, help="The batch size.")
+    parser.add_argument(
+        "--num_epochs_to_generate_seeds_for",
+        type=int,
+        default=2,
+        help="Number of epochs to plan seeds for. Same set across all workers.")
+    parser.add_argument(
+        "--distributed_lamb",
+        type=str2bool,
+        default=False,
+        help="Whether to use distributed LAMB optimizer.")
+    parser.add_argument(
+        "--exchange_padding",
+        type=str2bool,
+        default=False,
+        help="Whether to exchange padding across devices.")
+    parser.add_argument(
+        "--cpu_exchange_padding",
+        type=str2bool,
+        default=False,
+        help="Whether to use CPU to do exchange padding.")
+    parser.add_argument(
+        "--target_mlm_accuracy",
+        type=float,
+        default=0.72,
+        help="The target mlm accuracy to be coveraged.")
+    parser.add_argument(
+        "--max_samples_termination",
+        type=int,
+        default=4500000,
+        help="The max samples threshold to terminate the training process.")
+    parser.add_argument(
+        "--use_uncompressed_dataset",
+        type=str2bool,
+        default=False,
+        help="Whether to use the uncompressed dataset.")
+    parser.add_argument(
+        "--dense_seq_output",
+        type=str2bool,
+        default=False,
+        help="Whether to use the dense_seq_output optimization.")
+    parser.add_argument(
+        "--unpad_embed",
+        type=str2bool,
+        default=False,
+        help="Whether to use unpad_embed optimization.")
+    parser.add_argument(
+        "--sort_eval_data",
+        type=str2bool,
+        default=False,
+        help="Whether to sort the eval data.")
+    args = parser.parse_args()
+    return args
+
+
+'''
+def select_dataset_file_for_each_worker(files, f_start_id, worker_num,
+                                        worker_index):
+    """  
+    Spliting the train file according to the worker index.
+    """
+    num_files = len(files)
+    if worker_num > num_files:
+        remainder = worker_num % num_files
+        data_file = files[(
+            f_start_id * worker_num + worker_index + remainder * f_start_id) %
+                          num_files]
+    else:
+        data_file = files[(f_start_id * worker_num + worker_index) % num_files]
+    return data_file
+'''
+
+
+def create_strategy(args):
+    """
+    Create build strategy and exec strategy.
+    """
+    build_strategy = paddle.static.BuildStrategy()
+    exec_strategy = paddle.static.ExecutionStrategy()
+
+    build_strategy.enable_addto = args.enable_addto
+    build_strategy.allow_cuda_graph_capture = False
+    if args.distributed_lamb:
+        build_strategy.fuse_all_reduce_ops = False
+        build_strategy.reduce_strategy = paddle.static.BuildStrategy.ReduceStrategy._NoReduce
+    else:
+        build_strategy.fuse_all_reduce_ops = True
+        build_strategy.reduce_strategy = paddle.static.BuildStrategy.ReduceStrategy.AllReduce
+
+    #build_strategy.fuse_gemm_epilogue = True
+    build_strategy.fix_op_run_order = True
+    exec_strategy.num_threads = 1
+    exec_strategy.num_iteration_per_drop_scope = 10000
+    return build_strategy, exec_strategy
+
+
+def dist_optimizer(args, optimizer):
+    """
+    Create a distributed optimizer based on a normal optimizer
+    """
+    build_strategy, exec_strategy = create_strategy(args)
+
+    dist_strategy = fleet.DistributedStrategy()
+    if args.distributed_lamb:
+        dist_strategy.gradient_scale_configs = {'scale_strategy': 'sum'}
+    dist_strategy.execution_strategy = exec_strategy
+    dist_strategy.build_strategy = build_strategy
+    if args.distributed_lamb:
+        dist_strategy.fuse_all_reduce_ops = False
+    else:
+        dist_strategy.fuse_all_reduce_ops = True
+
+    dist_strategy.fuse_grad_size_in_MB = 0
+    if args.use_amp:
+        dist_strategy.amp = True
+        custom_white_list = [
+          # 'softmax',
+          # 'softmax_grad',
+          # 'softmax_with_cross_entropy',
+          # 'softmax_with_cross_entropy_grad',
+          #  'reduce_sum',
+          #  'reduce_sum_grad',
+        ]
+
+        custom_black_list = [
+            'arg_max',
+            'concat',
+            'cumsum',
+            'dropout',
+            'dropout_grad',
+            'elementwise_add',
+            'elementwise_add_grad',
+            'elementwise_div',
+            'elementwise_div_grad',
+            'elementwise_max',
+            'elementwise_mul',
+            'elementwise_mul_grad',
+            'elementwise_sub',
+            'fill_any_like',
+            'fill_constant',
+            'gather',
+            'gather_grad',
+            'gelu',
+            'gelu_grad',
+            'lamb',
+            'layer_norm',
+            'layer_norm_grad',
+            'lookup_table_v2',
+            'lookup_table_v2_grad',
+            'matmul',
+            'matmul_grad',
+            'matmul_v2',
+            'matmul_v2_grad',
+            'reduce_any',
+            'reduce_sum',
+             'reduce_sum_grad',
+            'reshape2',
+            'reshape2_grad',
+            'scatter',
+            'scatter_grad',
+            'slice',
+            'slice_grad',
+            'softmax',
+            'softmax_grad',
+            'softmax_with_cross_entropy',
+            'softmax_with_cross_entropy_grad',
+            'sqrt',
+            'square',
+            'sum',
+            'tanh',
+            'tanh_grad',
+            'transpose2',
+            'custom_fmha',
+            'custom_fmha_grad',
+            'custom_fused_dropout_residual_ln',
+            'custom_fused_dropout_residual_ln_grad',
+            'custom_fused_dense',
+            'custom_fused_dense_grad',
+            'reduce_mean_grad',
+            'reduce_mean',
+          #  'reduce_sum',
+          #  'reduce_sum_grad',
+            'cast',
+            'uniform_random',
+            'dropout',
+            'scale',
+            'unsqueeze2',
+            'index_select',
+        ]
+
+        dist_strategy.amp_configs = {
+            'custom_white_list': custom_white_list,
+            'custom_black_list': custom_black_list,
+            'init_loss_scaling': initial_loss_scale(),
+            'incr_every_n_steps': 2000,
+            'decr_every_n_nan_or_inf': 1,
+            'incr_ratio': 2.0,
+            'decr_ratio': 0.5,
+            'use_dynamic_loss_scaling': True,
+            'use_pure_fp16': args.use_pure_fp16,
+            'use_fp16_guard': args.use_pure_fp16,
+        }
+    optimizer = fleet.distributed_optimizer(optimizer, strategy=dist_strategy)
+    return optimizer
+
+
+def prune_exchange_padding_op(eval_program):
+    ops = eval_program.global_block().ops
+    assert ops[0].type == "sort_bert_inputs_across_devices"
+    eval_program.global_block()._remove_op(0)
+    eval_program._sync_with_cpp()
+
+
+def do_eval(exe, eval_program, eval_dataloader, eval_dataloader_cpu_tensor,
+            fetch_list):
+    acc = None 
+    n = len(eval_dataloader)
+    for i, batch in enumerate(eval_dataloader):
+        batch[0]['host_prefix_sum_seq_len'] = eval_dataloader_cpu_tensor[i][0]
+        if i + 1 == n:
+            acc = exe.run(eval_program, feed=batch, fetch_list=fetch_list)[0]
+        else:
+            exe.run(eval_program, feed=batch)
+
+    if num_trainers > 1:
+        accl = np.empty(
+            shape=[2],dtype=np.float64)
+        context.trainer_comm.Allreduce(acc, accl)
+    return accl[0] / accl[1]
+
+
+def get_found_nan_inf_flag_var(prog):
+    flag_var = None
+    for op in prog.global_block().ops:
+        if op.type == "update_loss_scaling":
+            cur_flag_var = op.input("FoundInfinite")
+            assert len(cur_flag_var) == 1
+            cur_flag_var = cur_flag_var[0]
+            if flag_var is None:
+                flag_var = cur_flag_var
+                flag_var = prog.global_block().vars.get(flag_var)
+                assert flag_var is not None
+            else:
+                assert flag_var.name == cur_flag_var
+    return flag_var
+
+
+def get_global_step_var(prog):
+    flag_var = None
+    for op in prog.global_block().ops:
+        if op.type == "distributed_fused_lamb":
+            cur_flag_var = op.output("Step")
+            assert len(cur_flag_var) == 1
+            cur_flag_var = cur_flag_var[0]
+            if flag_var is None:
+                flag_var = cur_flag_var
+                flag_var = prog.global_block().vars.get(flag_var)
+                assert flag_var is not None
+            else:
+                assert flag_var.name == cur_flag_var
+    return flag_var
+
+
+def get_loss_scaling_var_names_to_restored(prog):
+    def retrieve_var_names(loss_scale_op):
+        loss_scale_var_in = loss_scale_op.input("PrevLossScaling")
+        loss_scale_var_out = loss_scale_op.output("LossScaling")
+        assert len(loss_scale_var_in) == 1 and len(
+            loss_scale_var_out) and loss_scale_var_in[0] == loss_scale_var_out[
+                0]
+
+        good_steps_var_in = loss_scale_op.input("InGoodSteps")
+        good_steps_var_out = loss_scale_op.output("OutGoodSteps")
+        assert len(good_steps_var_in) == 1 and len(
+            good_steps_var_out) == 1 and good_steps_var_in[
+                0] == good_steps_var_out[0]
+
+        bad_steps_var_in = loss_scale_op.input("InBadSteps")
+        bad_steps_var_out = loss_scale_op.output("OutBadSteps")
+        assert len(bad_steps_var_in) == 1 and len(
+            bad_steps_var_out) == 1 and bad_steps_var_in[
+                0] == bad_steps_var_out[0]
+        return loss_scale_var_in[0], good_steps_var_in[0], bad_steps_var_in[0]
+
+    loss_scale_var, good_steps_var, bad_steps_var = None, None, None
+    for op in prog.global_block().ops:
+        if op.type == 'update_loss_scaling':
+            cur_loss_scale_var, cur_good_steps_var, cur_bad_steps_var = retrieve_var_names(
+                op)
+            if loss_scale_var is None:
+                loss_scale_var = cur_loss_scale_var
+                good_steps_var = cur_good_steps_var
+                bad_steps_var = cur_bad_steps_var
+
+            assert loss_scale_var == cur_loss_scale_var
+            assert good_steps_var == cur_good_steps_var
+            assert bad_steps_var == cur_bad_steps_var
+
+    assert loss_scale_var is not None
+    return loss_scale_var, good_steps_var, bad_steps_var
+
+
+def get_loss_scaling_var_values(prog):
+    scope = utility.get_scope()
+    var_names = get_loss_scaling_var_names_to_restored(prog)
+    data = []
+    for name in var_names:
+        value = np.array(scope.find_var(name).get_tensor())
+        data.append(value)
+        if trainer_id == 0:
+            print('Var {} value is: {}'.format(name, value))
+    return data
+
+
+def restore_loss_scaling_var_values(prog, data, place):
+    scope = utility.get_scope()
+    var_names = get_loss_scaling_var_names_to_restored(prog)
+    assert len(var_names) == len(data)
+    for name, np_arr in zip(var_names, data):
+        var = scope.find_var(name).get_tensor()
+        var.set(np_arr, place)
+
+
+def paddle_dtype_to_np_dtype(pd_dtype):
+    if pd_dtype == paddle.float32:
+        return np.float32
+    elif pd_dtype == paddle.int64:
+        return np.int64
+    else:
+        raise ValueError('Unsupported tensor dtype {}'.format(pd_dtype))
+
+
+def inplace_fill_constant(tensor, value):
+    dtype = paddle_dtype_to_np_dtype(tensor._dtype())
+    shape = tensor.shape()
+    np_value = (np.ones(shape) * value).astype(dtype)
+    place = tensor._place()
+    if place._equals(paddle.CPUPlace()):
+        place = paddle.CPUPlace()
+    elif place._equals(utility.get_place()):
+        place = utility.get_place()
+    else:
+        raise ValueError("Unsupported tensor place {}".format(place))
+
+    old_ptr = tensor._ptr()
+    tensor.set(np_value, place)
+    new_ptr = tensor._ptr()
+    assert old_ptr == new_ptr
+
+
+def recover_lamb_status(train_prog, args):
+    scope = utility.get_scope()
+
+    def fill_input_value(op, input_slot, value):
+        var_names = op.input(input_slot)
+        assert len(var_names) == 1
+        t = scope.find_var(var_names[0]).get_tensor()
+        inplace_fill_constant(t, value)
+
+    def fill_output_value(op, output_slot, value):
+        var_names = op.output(output_slot)
+        assert len(var_names) == 1
+        t = scope.find_var(var_names[0]).get_tensor()
+        inplace_fill_constant(t, value)
+
+    found_cnt = 0
+    for i in range(train_prog.num_blocks):
+        block = train_prog.block(i)
+        for op in block.ops:
+            if args.distributed_lamb:
+                assert op.type != 'lamb'
+            else:
+                assert op.type != 'distributed_fused_lamb'
+
+            if op.type != 'lamb' and op.type != 'distributed_fused_lamb':
+                continue
+
+            fill_input_value(op, "Moment1", 0)
+            fill_input_value(op, "Moment2", 0)
+            fill_input_value(op, "Beta1Pow", args.opt_lamb_beta_1)
+            fill_input_value(op, "Beta2Pow", args.opt_lamb_beta_2)
+            if op.type == 'distributed_fused_lamb':
+                fill_output_value(op, "Step", 0)
+                if args.gradient_accumulation_steps > 1:
+                    fill_output_value(op, "AccStep", 0)
+            found_cnt += 1
+
+    if args.distributed_lamb:
+        assert found_cnt == 1
+    else:
+        assert found_cnt > 1
+
+
+def run_warmup(args,
+               exe,
+               train_prog,
+               lr_scheduler,
+               eval_prog=None,
+               warmup_iter=8):
+    train_warmup_data = dataset.prepare_warmup_data(args, args.train_batch_size,
+                                                    exe.place)
+    list_train_warmup_data = [train_warmup_data]
+
+    paddle.device.cuda.synchronize()
+    restored_values = get_loss_scaling_var_values(train_prog)
+    old_lr_func = lr_scheduler.__call__
+    lr_scheduler.__call__ = lambda self: 0.0
+    for _ in range(warmup_iter):
+        exe.run(train_prog, feed=list_train_warmup_data)
+    lr_scheduler.__call__ = old_lr_func
+    restore_loss_scaling_var_values(train_prog, restored_values, exe.place)
+    paddle.device.cuda.synchronize()
+
+    # extra run as NV code
+    exe.run(train_prog, feed=list_train_warmup_data)
+    paddle.device.cuda.synchronize()
+
+    if eval_prog is not None:
+        eval_warmup_data = dataset.prepare_warmup_data(
+            args, args.eval_batch_size, exe.place)
+        for _ in range(warmup_iter):
+            exe.run(eval_prog, feed=[eval_warmup_data])
+        paddle.device.cuda.synchronize()
+
+    #recover_lamb_moment_pows(train_prog, args)
+    recover_lamb_status(train_prog, args)
+    paddle.device.cuda.synchronize()
+    if utility.get_trainer_id() == 0:
+        print('warmup ends')
+
+
+def do_train(args):
+    # Initialize the paddle and paddle fleet execute enviroment
+    if trainer_id == 0:
+        paddle_bert_print_start(key=constants.INIT_START)
+        paddle_bert_print_event(key=constants.SEED, val=args.seed)
+        paddle_bert_print_event(
+            key=constants.GLOBAL_BATCH_SIZE,
+            val=args.train_batch_size * num_trainers *
+            args.gradient_accumulation_steps)
+        paddle_bert_print_event(key='d_batch_size', val=args.train_batch_size)
+        paddle_bert_print_event(
+            key=constants.GRADIENT_ACCUMULATION_STEPS,
+            val=args.gradient_accumulation_steps)
+        paddle_bert_print_event(
+            key="max_predictions_per_seq", val=args.max_predictions_per_seq)
+        paddle_bert_print_event(
+            key='opt_learning_rate_training_steps', val=args.max_steps)
+        paddle_bert_print_event(key='num_warmup_steps', val=args.warmup_steps)
+
+    place = utility.get_place()
+    fleet.init(is_collective=True)
+
+    worker_num = num_trainers
+    worker_index = trainer_id
+
+    # Define the input data in the static mode
+    main_program = paddle.static.default_main_program()
+    startup_program = paddle.static.default_startup_program()
+
+    data_holders, data_inputs, data_labels, varlen_info, mlm_label_info = create_data_holder(
+        args)
+
+    config = BertConfig.from_json_file(args.bert_config_path)
+    #config.fused_mha = args.fused_mha
+    #config.fused_gelu_bias = args.fused_gelu_bias
+    config.fused_bias_mha = args.fused_bias_mha
+    config.fused_bias_fc = args.fused_bias_fc
+    #config.fused_bias_fc_loss_head = args.fused_bias_fc_loss_head
+    config.dense_seq_output = args.dense_seq_output
+    config.unpad = args.unpad
+    config.unpad_fmha = args.unpad_fmha
+    config.unpad_fmha_mke_opt = args.unpad_fmha_mke_opt
+    config.pad_fmha = args.pad_fmha
+    config.max_seq_length = args.max_seq_length
+    #config.pad = args.pad
+    #config.fuse_qkv = not args.disable_fuse_qkv
+    #config.fuse_scale = not args.disable_fuse_scale
+    #config.fuse_mask = not args.disable_fuse_mask 
+    #config.fuse_dropout = args.enable_import paddle.fluid.profiler as profilerfuse_dropout
+    config.fused_dropout_add_ln = args.fused_dropout_add_ln
+    #config.apex_softmax = not args.disable_apex_softmax
+    #config.enable_stream = args.enable_stream
+    #if config.fuse_mask: config.apex_softmax = True
+    #if not config.pad: config.enable_stream = True
+    #if config.unpad: config.fused_mha = False
+
+    config.unpad_embed = args.unpad_embed
+
+    config.weight_transpose = args.weight_transpose
+    config.batch_size = args.batch_size
+
+    # Padding for divisibility by 8
+    if config.vocab_size % 8 != 0:
+        config.vocab_size += 8 - (config.vocab_size % 8)
+
+    int16_max = np.iinfo(np.int16).max
+    assert config.vocab_size <= int16_max
+    assert args.train_batch_size * args.max_seq_length <= int16_max
+    assert args.eval_batch_size * args.max_seq_length <= int16_max
+
+    def build_model():
+        model = BertForPretraining(BertModel(config), config)
+        criterion = BertPretrainingCriterion(config)
+
+        prediction_scores, seq_relationship_score = model(*data_inputs)
+        func = criterion(prediction_scores, seq_relationship_score,
+                         *data_labels)
+        return model, func
+
+    if args.use_pure_fp16:
+        with paddle.static.amp.fp16_guard():
+            model, func = build_model()
+    else:
+        model, func = build_model()
+    loss, mlm_acc, num_masked = func()
+
+    eval_program = main_program.clone(for_test=True)
+    if args.gpu_exchange_padding and utility.use_nv_input(
+    ) and num_trainers > 1:
+        prune_exchange_padding_op(eval_program)
+
+    # Define the dynamic learing_reate scheduler and optimizer
+    num_training_steps = args.max_steps if args.max_steps > 0 else len(
+        train_dataloader) * args.num_train_epochs
+    if args.warmup_steps == 0:
+        warmup_steps = int(args.max_steps * args.warmup_proportion)
+        warmup_start = 0
+    else:
+        warmup_steps = args.warmup_steps
+        warmup_start = args.start_warmup_step
+
+    lr_scheduler = LinearWarmupPolyDecayScheduler(
+        startup_warmup_steps=warmup_start,
+        warmup_steps=warmup_steps,
+        total_steps=args.max_steps,
+        base_lr=args.learning_rate,
+        end_lr=0.0,
+        degree=1.0)
+
+    # Generate parameter names needed to perform weight decay.
+    # All bias and LayerNorm parameters are excluded.
+    decay_params = [
+        p.name for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    if args.max_grad_norm > 0:
+        grad_clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=args.max_grad_norm)
+    else:
+        grad_clip = None
+    exclude_from_weight_decay_fn = lambda var: var.name not in decay_params
+
+    lr_var, step_var = append_lr_op(
+        base_lr=args.learning_rate,
+        max_step=args.max_steps,
+        startup_warmup_step=warmup_start,
+        warmup_step=warmup_steps,
+        end_lr=0.0,
+        degree=1.0)
+    optimizer_kwargs = {
+        'learning_rate': lr_var,
+        'lamb_weight_decay': args.weight_decay_rate,
+        'epsilon': args.lamb_epsilon,
+        'exclude_from_weight_decay_fn': exclude_from_weight_decay_fn,
+        'grad_clip': grad_clip,
+        'beta1': args.opt_lamb_beta_1,
+        'beta2': args.opt_lamb_beta_2,
+    }
+
+    if trainer_id == 0:
+        paddle_bert_print_event(
+            key=constants.OPT_BASE_LR, val=args.learning_rate)
+        paddle_bert_print_event(
+            key="opt_learning_rate_warmup_steps", val=args.warmup_steps)
+        paddle_bert_print_event(
+            key="start_warmup_step", val=args.start_warmup_step)
+
+    #_keep_layer_norm_scale_bias_to_fp32(False)
+    _allow_pure_fp16_global_norm_clip(True)
+    if args.distributed_lamb:
+        # Note: 
+        # (1) broadcast_master_param is true means use fp32 to update parameter.
+        # (2) if 'is_grad_scaled_by_nranks' is false, must set 'dist_strategy.gradient_scale_configs = {'scale_strategy': 'sum'}'.
+        #optimizer_kwargs.update({"broadcast_master_param": True, "clip_after_allreduce": False, "is_grad_scaled_by_nranks": False})
+        optimizer_kwargs.update({
+            "use_master_param_norm": True,
+            "clip_after_allreduce": False,
+            "is_grad_scaled_by_nranks": False,
+        })
+        if args.gradient_accumulation_steps > 1:
+            optimizer_kwargs.update({
+                "gradient_accumulation_steps": args.gradient_accumulation_steps
+            })
+        optimizer = DistributedFusedLamb(**optimizer_kwargs)
+        optimizer._set_step(step_var)
+    else:
+        optimizer = paddle.optimizer.Lamb(**optimizer_kwargs)
+        if args.use_pure_fp16:
+            optimizer._multi_precision = True
+
+    if hasattr(optimizer, "_get_parameter"):
+        get_param_func = optimizer._get_parameter
+    else:
+        get_param_func = lambda name: utility.get_scope().find_var(name).get_tensor()
+
+
+    if trainer_id == 0:
+        paddle_bert_print_event(key='opt_epsilon', val=args.lamb_epsilon)
+        paddle_bert_print_event(key='opt_lamb_beta_1', val=args.opt_lamb_beta_1)
+        paddle_bert_print_event(key='opt_lamb_beta_2', val=args.opt_lamb_beta_2)
+        paddle_bert_print_event(
+            key='opt_lamb_weight_decay_rate', val=args.weight_decay_rate)
+
+    # Use the fleet api to compile the distributed optimizer
+    print("optimizer",optimizer._get_parameter)
+
+    optimizer = dist_optimizer(args, optimizer)
+    print("loss",loss)
+
+    #pdb.set_trace()
+    optimizer.minimize(loss)
+
+    # Define the Executor for running the static model
+    exe = paddle.static.Executor(place)
+    exe.run(startup_program)
+
+    if args.use_amp:
+        optimizer.amp_init(place, test_program=eval_program, use_fp16_test=True)
+
+    eval_fetch_list, set_eval_step = append_acc_merge_op(eval_program, mlm_acc,
+                                                         num_masked)
+
+    with open('startup_program_{}.txt'.format(trainer_id), 'w') as f:
+        f.write(str(startup_program))
+
+    with open('main_program_{}.txt'.format(trainer_id), 'w') as f:
+        f.write(str(main_program))
+
+    with open('eval_program_{}.txt'.format(trainer_id), 'w') as f:
+        f.write(str(eval_program))
+
+    found_nan_inf = get_found_nan_inf_flag_var(main_program)
+    found_global_step_val = get_global_step_var(main_program)
+    if args.use_amp:
+        assert found_nan_inf is not None
+    else:
+        assert found_nan_inf is None
+    assert found_global_step_val is not None
+
+    build_strategy, exec_strategy = create_strategy(args)
+    eval_program = paddle.static.CompiledProgram(
+        eval_program).with_data_parallel(
+            build_strategy=build_strategy, exec_strategy=exec_strategy)
+    run_warmup(
+        args,
+        exe,
+        main_program,
+        lr_scheduler,
+        eval_prog=eval_program,
+        warmup_iter=8)
+
+    ckpt = None
+    if args.tf_ckpt_path:
+        if trainer_id == 0:
+            print('starts to load tf checkpoint...')
+            start_t = time.time()
+        ckpt = model.load_tf_ckpt(args, get_param_func)
+        paddle.device.cuda.synchronize()
+        if trainer_id == 0:
+            end_t = time.time()
+            print('checkpoint loading ends, time cost: {}s'.format(end_t -
+                                                                   start_t))
+
+    def save_ckpt(ckpt_saved_path):
+        if trainer_id == 0 and ckpt is not None:
+            dirname = 'log_{}'.format(num_trainers)
+            ckpt_saved_path = os.path.join(dirname, ckpt_saved_path)
+            ckpt.save(ckpt_saved_path, get_param_func)
+            print('PaddlePaddle ckpt saved into: {}'.format(ckpt_saved_path))
+
+    # save_ckpt('loaded_ckpt_0.pickled')
+
+    pool = ThreadPoolExecutor(1)
+    # pre-compute eval boundaries
+    samples_trained_per_step = args.train_batch_size * num_trainers * args.gradient_accumulation_steps
+    eval_interval_samples = 0.05 * (230.23 * samples_trained_per_step + 3000000)
+    eval_interval_samples = int(eval_interval_samples / 25000) * 25000
+    start, stop, step = eval_interval_samples, args.max_samples_termination, eval_interval_samples
+    eval_steps = [
+        math.ceil(i / samples_trained_per_step)
+        for i in np.arange(start, stop, step)
+    ]
+    eval_count = 0
+    next_eval_step = eval_steps[eval_count]
+    next_eval_training_step = next_eval_step * args.gradient_accumulation_steps
+    create_train_dataset = create_cpu_exchange_padding_pretraining_dataset if args.cpu_exchange_padding else create_pretraining_dataset
+
+    context.trainer_comm.Barrier()
+    if trainer_id == 0:
+        paddle_bert_print_end(key=constants.INIT_STOP)
+        paddle_bert_print_start(key=constants.RUN_START)
+
+    f_start_id = 0
+    train_dataloader, train_dataloader_cpu_tensor = create_train_dataset(
+        data_holders, f_start_id, tolist=False)
+
+    if args.eval_dir:
+        eval_dataset_future = pool.submit(create_new_eval_dataset, data_holders)
+    need_next_training_shard = args.train_batch_size * args.gradient_accumulation_steps * args.max_steps > 10000
+
+    global_step = 0
+    skipped_steps = 0
+    training_steps = 0
+    now_skipped = 0
+    now_step = 0
+    skip_interval = 0
+    average_loss = 0.0
+    end_training, converged = False, False
+    coveraged = False
+    epoch = 1  # to be consistent with NV
+
+    file_num = context.file_num()
+
+    while global_step < args.max_steps and not end_training:
+        if trainer_id == 0:
+            paddle_bert_print_start(
+                key=constants.EPOCH_START, metadata={'epoch_num': epoch})
+            paddle_bert_print_start(
+                key=constants.BLOCK_START,
+                metadata={'first_epoch_num': epoch,
+                          'epoch_count': 1})
+            now_time = time.time()
+
+        empty_fetch_list = []
+        full_fetch_list = [loss, mlm_acc, found_global_step_val]
+        for f_id in range(f_start_id + 1, file_num):
+            if need_next_training_shard:
+                dataset_future = pool.submit(create_train_dataset, data_holders,
+                                             f_id)
+            # limin-todo:
+            # limin_file_start = time.time()
+            for step, batch in enumerate(train_dataloader):
+                training_steps += 1
+                #if training_steps == 2:
+                #   exit()
+                #print("training_steps",training_steps)
+                update_step = training_steps % args.gradient_accumulation_steps == 0
+                batch[0][
+                    'host_prefix_sum_seq_len'] = train_dataloader_cpu_tensor[
+                        step][0]
+
+                if training_steps < next_eval_training_step and False:
+                    cur_fetch_list = empty_fetch_list
+                else:
+                    cur_fetch_list = full_fetch_list
+                #with profiler.profiler('All', 'total') as prof:
+                return_val = exe.run(main_program,
+                                     feed=batch,
+                                     fetch_list=cur_fetch_list)
+                if training_steps >= next_eval_training_step or True:
+                    loss_return, mlm_acc_return, found_global_step_return = return_val
+                    average_loss += loss_return[0]
+  
+                if update_step:
+                    if training_steps >= next_eval_training_step:
+#                        global_step = found_global_step_return[0]
+                        global_step = int(training_steps /
+                                          args.gradient_accumulation_steps)           
+                        has_nan_inf_step = (
+                            training_steps / args.gradient_accumulation_steps -
+                            global_step)
+                        if has_nan_inf_step:
+                            skipped_steps = has_nan_inf_step
+                    else:
+                        global_step = int(training_steps /
+                                          args.gradient_accumulation_steps)
+
+                    if args.eval_dir and global_step == next_eval_step:
+                        skip_interval = skipped_steps
+                        now_skipped = skip_interval
+                        if True:
+                            if eval_count == 0:
+                                eval_dataloader, eval_dataloader_cpu_tensor = eval_dataset_future.result(
+                                    timeout=None)
+                                eval_dataloader = list(eval_dataloader)
+                                set_eval_step(len(eval_dataloader))
+                            samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * num_trainers
+                            samples_trained_prev = samples_trained
+
+                            eval_start_t = time.time()
+                            eval_avg_mlm_accuracy = do_eval(
+                                exe, eval_program, eval_dataloader,
+                                eval_dataloader_cpu_tensor, eval_fetch_list)
+                            eval_end_t = time.time()
+                            # if trainer_id == 0:
+                            #    print('Eval time: {} s'.format(eval_end_t -
+                            #                                   eval_start_t))
+
+                            if args.target_mlm_accuracy > 0 and eval_avg_mlm_accuracy >= args.target_mlm_accuracy:
+                                end_training, converged = True, True
+                                coveraged = True
+                                if trainer_id == 0:
+                                    print(
+                                        "%f > %f, Target MLM Accuracy reached at global_step %d , training_step %d"
+                                        % (eval_avg_mlm_accuracy,
+                                           args.target_mlm_accuracy,
+                                           global_step, training_steps))
+                            eval_count += 1
+                            next_eval_step = eval_steps[eval_count]
+                            next_eval_training_step = next_eval_step * args.gradient_accumulation_steps
+                            if trainer_id == 0:
+                                paddle_bert_print_event(
+                                    key=constants.EVAL_ACCURACY,
+                                    val=eval_avg_mlm_accuracy,
+                                    metadata={'epoch_num': epoch})
+                                print({
+                                    "global_steps": global_step,
+                                    "eval_mlm_accuracy": eval_avg_mlm_accuracy
+                                })
+
+                samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * num_trainers
+                if args.log_freq > 0 and training_steps % (
+                        args.log_freq * args.gradient_accumulation_steps) == 0:
+                    samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * num_trainers
+                    #just add the following to print loss
+                    cur_fetch_list = full_fetch_list
+                    return_val = exe.run(main_program,
+                                         feed=batch,
+                                         fetch_list=cur_fetch_list)
+                    loss_return, mlm_acc_return, found_global_step_return = return_val
+                    print("loss_return, mlm_acc_return, found_global_step_return:",loss_return, mlm_acc_return, found_global_step_return[0])
+                    if trainer_id == 0:
+                        time_interval = time.time() - now_time
+                        step_interval = global_step - now_step
+                        now_time = time.time()
+                        now_step = global_step
+                        training_perf = args.train_batch_size * args.gradient_accumulation_steps * num_trainers * (
+                            step_interval + skip_interval) / time_interval
+                        skip_interval = 0
+                        divisor = args.gradient_accumulation_steps
+                        now_lr = lr_scheduler()
+                        step_loss_to_print = loss_return[0] * args.gradient_accumulation_steps / divisor
+                        average_loss_to_print = average_loss / (args.log_freq * divisor)
+                        print({
+                            "training_steps": training_steps,
+                            "average_loss": average_loss_to_print,
+                            "step_loss": step_loss_to_print,
+                            "learning_rate": now_lr,
+                            "seq/s": training_perf,
+                            "global_steps": now_step,
+                            "samples_trained": samples_trained,
+                            "skipped_steps": now_skipped,
+                            "timestamp": now_time,
+                        })
+
+                        paddle_bert_print_event(
+                            key='tracked_stats',
+                            val={
+                                'seq/s': training_perf,
+                                'step_loss': step_loss_to_print,
+                                #'avg_loss': average_loss_to_print,
+                                #'lr': now_lr,
+                            },
+                            metadata={"step": (epoch, training_steps)})
+                        paddle_bert_print_event(
+                            key='throughput', val=training_perf)
+
+                    average_loss = 0.0
+
+                if global_step >= args.max_steps or end_training:
+                    status = 'success' if coveraged else 'abort'
+                    end_training = True
+                    break
+
+            # limin-todo:
+            #limin_file_end = time.time()
+            #print("limin: one file time: ", limin_file_end - limin_file_start)
+            del train_dataloader
+            if samples_trained >= args.max_samples_termination or end_training:
+                status = 'success' if converged else 'aborted'
+                end_training = True
+                break
+
+            if not need_next_training_shard:
+                dataset_future = pool.submit(create_train_dataset, f_id)
+            train_dataloader, train_dataloader_cpu_tensor = dataset_future.result(
+                timeout=None)
+
+        if trainer_id == 0:
+            paddle_bert_print_end(
+                key=constants.BLOCK_STOP, metadata={'first_epoch_num': epoch})
+            paddle_bert_print_end(
+                key=constants.EPOCH_STOP, metadata={'epoch_num': epoch})
+
+        epoch += 1
+
+    if trainer_id == 0:
+        paddle_bert_print_event(constants.TRAIN_SAMPLES, val=samples_trained)
+        paddle_bert_print_event(
+            constants.EVAL_SAMPLES, val=args.num_eval_examples)
+        paddle_bert_print_end(
+            key=constants.RUN_STOP, metadata={'status': status})
+
+        #save_ckpt('saved_ckpt_final.pickled')
+
+
+def do_read(args):
+    while context.read_file():
+        pass
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    if not args.exchange_padding:
+        args.cpu_exchange_padding = False
+        args.gpu_exchange_padding = False
+    else:
+        args.gpu_exchange_padding = not args.cpu_exchange_padding
+
+    context.init_args(args)
+    trainer_id = context.trainer_id
+    num_trainers = context.trainer_num
+
+    if context.is_trainer:
+        if trainer_id == 0:
+            paddle_bert_print_event(key=constants.SUBMISSION_ORG, val="Baidu")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_PLATFORM, val="1 x NVIDIA A100 GPU")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_DIVISION, val="closed")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_STATUS, val="onprem")
+            paddle_bert_print_event(
+                key=constants.SUBMISSION_BENCHMARK, val="bert")
+            paddle_bert_print_event(key=constants.CACHE_CLEAR, val=True)
+        save_env('./paddle_env_{}.json'.format(trainer_id))
+        if trainer_id == 0:
+            print(args)
+            print_flags()
+        if args.use_pure_fp16:
+            assert args.use_amp, "--use_amp must be True if --use_pure_fp16 is True"
+        paddle.enable_static()
+        with paddle.static.scope_guard(utility.get_scope()):
+            do_train(args)
+    else:
+        if args.cpu_exchange_padding:
+            do_read(args)
+    if args.cpu_exchange_padding:
+        context.stop_reader()
+    print('Process ends')

run_with_docker.sh

+#!/bin/bash
+
+# Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euxo pipefail
+
+func_get_cuda_visible_devices() {
+  local NGPUS=`nvidia-smi --list-gpus | wc -l` 
+  local NGPUS_M1=$(($NGPUS-1))
+  seq -s , 0 $NGPUS_M1
+}
+
+# Vars without defaults
+: "${STAGE:?STAGE not set}"
+: "${CONT:?CONT not set}"
+: "${BASE_DATA_DIR:?BASE_DATA_DIR not set}"
+
+# Vars with defaults
+: "${CLEAR_CACHES:=1}"
+: "${DATESTAMP:=$(date +'%y%m%d%H%M%S%N')}"
+: "${LOG_DIR:=$(pwd)/results}"
+: "${MLPERF_MODEL_CONSTANT:=constants.BERT}"
+: "${NEXP:=5}"
+
+: "${DGXNGPU:=`nvidia-smi --list-gpus | wc -l`}"
+: "${CUDA_VISIBLE_DEVICES:=`func_get_cuda_visible_devices`}"
+: "${DGXSYSTEM:="DGXA100_1x8x56x1"}"
+: "${CONFIG_FILE:="./config_${DGXSYSTEM}.sh"}"
+: "${LOG_FILE_BASE:="${LOG_DIR}/${DATESTAMP}"}"
+: "${CONT_NAME:=language_model}"
+: "${NV_GPU:="${CUDA_VISIBLE_DEVICES}"}"
+: "${MASTER_PORT:="29500"}"
+
+export DGXNGPU
+export CUDA_VISIBLE_DEVICES
+
+readonly docker_image=${CONT}
+
+# Setup directories
+mkdir -p "${LOG_DIR}"
+
+# Get list of envvars to pass to docker
+mapfile -t _config_env < <(env -i bash -c ". ${CONFIG_FILE} && compgen -e" | grep -E -v '^(PWD|SHLVL)')
+_config_env+=(SEED)
+_config_env+=(MASTER_PORT)
+mapfile -t _config_env < <(for v in "${_config_env[@]}"; do echo "--env=$v"; done)
+
+if [[ "$STAGE" == "LOGIN" ]]; then
+    docker exec -it "${_config_env[@]}" "${CONT_NAME}" bash
+    exit 0
+fi
+
+# Cleanup container
+cleanup_docker() {
+    docker container rm -f "${CONT_NAME}" || true
+}
+cleanup_docker
+
+if [[ "$STAGE" == "RUN" || "$STAGE" == "run" ]]; then
+    trap 'set -eux; cleanup_docker' EXIT
+fi
+
+NVIDIA_SMI=`which nvidia-smi`
+
+# Setup container
+nvidia-docker run --rm --init --detach \
+    --net=host --uts=host --ipc=host --security-opt=seccomp=unconfined \
+    --ulimit=stack=67108864 --ulimit=memlock=-1 \
+    --ulimit=nofile=1000000 \
+    --name="${CONT_NAME}" \
+    --env BASE_DATA_DIR=$BASE_DATA_DIR \
+    -v $NVIDIA_SMI:$NVIDIA_SMI \
+    -v $PWD:/workspace/bert \
+    -v $BASE_DATA_DIR:$BASE_DATA_DIR \
+    -w /workspace/bert \
+    "${CONT}" sleep infinity
+
+#make sure container has time to finish initialization
+sleep 30
+docker exec "${CONT_NAME}" true
+
+# Run experiments
+for _experiment_index in $(seq 1 "${NEXP}"); do
+    (
+        echo "Beginning trial ${_experiment_index} of ${NEXP}"
+
+        # Clear caches
+        if [ "${CLEAR_CACHES}" -eq 1 ]; then
+            sync && sudo /sbin/sysctl vm.drop_caches=3
+        fi
+
+        # Run experiment
+	export SEED=${SEED:-"$RANDOM"}
+        docker exec "${_config_env[@]}" "${CONT_NAME}" bash "./run_and_time.sh"
+    ) |& tee "${LOG_FILE_BASE}_${_experiment_index}.log"
+done

rundir_1gpu/init_env.py

+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from mpi4py import MPI
+import numpy as np
+import time
+import paddle
+from pybind.functions import process_allgathered_inputs as process_bert_inputs
+from pybind.functions import process_eval_inputs as process_bert_eval_inputs
+import h5py
+import random
+
+global_comm = MPI.COMM_WORLD
+global_rank = global_comm.rank
+global_world_size = global_comm.size
+#assert global_world_size % 2 == 0
+
+
+def create_group_comm(ranks):
+    ranks = list(ranks)
+    new_group = global_comm.group.Incl(ranks)
+    new_comm = global_comm.Create_group(new_group)
+    return new_comm
+
+
+def generate_seeds(rng, size):
+    """
+    Generate list of random seeds
+
+    :param rng: random number generator
+    :param size: length of the returned list
+    """
+    seeds = [rng.randint(0, 2**32 - 1) for _ in range(size)]
+    return seeds
+
+
+def broadcast_seeds(comm, seeds, root=0):
+    seeds = np.array(seeds).astype(np.int64)
+    comm.Bcast(seeds, root=root)
+    return seeds.tolist()
+
+
+def select_dataset_file_for_each_worker(files, f_start_id, worker_num,
+                                        worker_index):
+    """
+    Spliting the train file according to the worker index.
+    """
+    num_files = len(files)
+    if worker_num > num_files:
+        remainder = worker_num % num_files
+        data_file = files[(
+            f_start_id * worker_num + worker_index + remainder * f_start_id) %
+                          num_files]
+    else:
+        data_file = files[(f_start_id * worker_num + worker_index) % num_files]
+        tmp = (f_start_id * worker_num + worker_index) % num_files
+        #print("tmp",tmp)
+    # limin-todo: 
+    #data_file = "/data2/zengjinle/dataset/bert_data/hdf5/training-4320/hdf5_4320_shards_uncompressed/part_01799_of_04320.hdf5"
+    #print("data_file: ", data_file)
+    return data_file
+
+
+def read_hdf5_file(input_file, dtype=np.int16):
+    keys = [
+        'input_ids',
+        'input_mask',
+        'segment_ids',
+        'masked_lm_positions',
+        'masked_lm_ids',
+        'next_sentence_labels',
+    ]
+    if not os.path.exists(input_file):
+        return None
+    with h5py.File(input_file, 'r') as f:
+        outputs = [np.array(f[key], dtype=dtype) for key in keys]
+        n = outputs[0].shape[0]
+        masked_lm_labels = np.zeros(outputs[0].shape, dtype=dtype)
+        lengths = np.zeros(n, dtype=dtype)
+        for i in range(n):
+            masked_lm_positions = outputs[3][i]
+            masked_lm_ids = outputs[4][i]
+            length = np.count_nonzero(masked_lm_positions)
+            masked_lm_labels[i][
+                masked_lm_positions[:length]] = masked_lm_ids[:length]
+            lengths[i] = np.count_nonzero(outputs[1][i])
+        outputs = [
+            outputs[0], outputs[2], outputs[1], masked_lm_labels, outputs[-1],
+            lengths
+        ]
+        idx = np.random.choice(np.arange(n), n, replace=False)
+        for i in range(len(outputs)):
+            outputs[i] = outputs[i][idx]
+    return outputs
+
+
+def read_eval_hdf5_file(input_file, dtype=np.int16):
+    keys = [
+        'input_ids',
+        'input_mask',
+        'segment_ids',
+        'masked_lm_positions',
+        'masked_lm_ids',
+        'next_sentence_labels',
+    ]
+    if not os.path.exists(input_file):
+        return None
+    with h5py.File(input_file, 'r') as f:
+        outputs = [np.asarray(f[key][:]) for key in keys]
+        nsamples = outputs[0].shape[0]
+
+        all_data = []
+        for index in range(nsamples):
+            [
+                input_ids, input_mask, segment_ids, masked_lm_positions,
+                masked_lm_ids, next_sentence_labels
+            ] = [
+                input[index].astype(dtype)
+                if indice < 5 else np.asarray(input[index].astype(dtype))
+                for indice, input in enumerate(outputs)
+            ]
+
+            length = np.count_nonzero(masked_lm_positions)
+            masked_lm_positions = masked_lm_positions[:length]
+            masked_lm_ids = masked_lm_ids[:length]
+
+            masked_lm_labels = np.zeros(input_ids.shape, dtype=dtype)
+            masked_lm_labels[masked_lm_positions] = masked_lm_ids
+
+            #if index == 0:
+            #    print("masked_lm_labels = ", masked_lm_labels)
+            #    print("masked_lm_positions = ", masked_lm_positions)
+            #    print("masked_lm_ids = ", masked_lm_ids)
+            seq_len = np.asarray(np.count_nonzero(input_mask))
+
+            data = [
+                input_ids,
+                segment_ids,
+                input_mask,
+                masked_lm_labels,
+                next_sentence_labels,
+                seq_len,
+            ]
+            # (2050, ), i.e., 512 * 4 + 1 + 1
+            one_sample_data = np.concatenate([d.flatten() for d in data])
+
+            all_data.extend(one_sample_data)
+
+        # (2050000, ) -> (10000, 2050)
+    return np.asarray(all_data).reshape((nsamples, -1))
+
+
+class WorkerInitObj(object):
+    "Construct the object with different seed, and the Dataloader will generate the data "
+    "with different seed in each worker."
+
+    def __init__(self, seed):
+        self.seed = seed
+
+    def __call__(self, id):
+        np.random.seed(seed=self.seed + id)
+        random.seed(self.seed + id)
+
+
+class Context:
+    def __init__(self):
+        half_size = int(global_world_size / 1)
+        self.trainer_id = global_rank % half_size
+        self.trainer_num = half_size
+        self.is_trainer = (global_rank < half_size)
+
+        #self.reader_id = self.trainer_id
+        #self.reader_num = self.trainer_num
+        #self.is_reader = not self.is_trainer
+
+        self.trainer_comm = global_comm #create_group_comm(range(0, half_size))
+        #self.reader_comm = create_group_comm(
+        #    range(half_size, global_world_size))
+        #self.trainer_reader_comm = create_group_comm(
+        #    [self.trainer_id, self.trainer_id + half_size])
+        self.global_comm = global_comm
+
+    def init_args(self, args, dtype=np.int16):
+        self.args = args
+        self.files = [
+            os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir)
+            if os.path.isfile(os.path.join(args.input_dir, f)) and "part" in f
+        ]
+        self.files.sort()
+        self.fid_buf = np.array([1], dtype=np.int64)
+        with h5py.File(self.files[0], 'r') as f:
+            self.num_samples = np.array(f["next_sentence_labels"][:]).size
+
+        self.batch_size = args.train_batch_size
+        self.max_seq_length = args.max_seq_length
+        self.worker_seeds, self.shuffling_seeds = self._setup_seeds(
+            args.seed, args.num_epochs_to_generate_seeds_for)
+        self.epoch_idx = 0
+
+        data_buf_size = self.num_samples * 4 * self.max_seq_length + self.num_samples * 2
+        self.data_buf = np.empty(
+            shape=[ data_buf_size], dtype=dtype)
+
+        self.eval_dir = args.eval_dir
+        self.num_eval_examples = args.num_eval_examples
+        self.eval_batch_size = args.eval_batch_size
+
+        cur_seed = self.worker_seeds[self.trainer_id]
+        np.random.seed(cur_seed)
+        random.seed(cur_seed)
+        paddle.seed(cur_seed)
+        self.worker_init = WorkerInitObj(cur_seed)
+        self.barrier()
+
+    def shuffle_files(self):
+        random.Random(self.shuffling_seeds[self.epoch_idx]).shuffle(self.files)
+        self.epoch_idx += 1
+
+    def _setup_seeds(self, master_seed, epochs):
+        if master_seed is None:
+            master_seed = random.SystemRandom().randint(0, 2**32 - 1)
+            if self.trainer_id == 0:
+                print('Using random master seed: {}'.format(master_seed))
+        else:
+            print('Using master seed from command line: {}'.format(master_seed))
+
+        # initialize seeding RNG
+        seeding_rng = random.Random(master_seed)
+
+        # generate worker seeds, one seed for every distributed worker
+        worker_seeds = generate_seeds(seeding_rng, self.trainer_num)
+
+        # generate seeds for data shuffling, one seed for every epoch
+        shuffling_seeds = generate_seeds(seeding_rng, epochs)
+
+        worker_seeds = broadcast_seeds(self.global_comm, worker_seeds)
+        shuffling_seeds = broadcast_seeds(self.global_comm, shuffling_seeds)
+        return worker_seeds, shuffling_seeds
+
+    def worker_seed(self):
+        return self.worker_seeds[self.trainer_id]
+
+    def barrier(self):
+        self.global_comm.barrier()
+
+    def stop_reader(self):
+        if self.is_trainer:
+            self.read_file(-1)
+
+    def file_num(self):
+        return len(self.files)
+
+    def read_file(self, f_id=None):
+        if self.is_trainer:
+            self.fid_buf[0] = f_id
+           # self.trainer_reader_comm.Isend(self.fid_buf, dest=1)
+            if f_id == 0:
+                self.shuffle_files()
+            elif f_id < 0:
+                return
+
+#            self.trainer_reader_comm.Recv(self.data_buf, source=1)
+#            results = process_bert_inputs(self.data_buf, self.num_samples,
+#                                          self.max_seq_length, self.batch_size,
+#                                          self.trainer_id, self.trainer_num)
+
+#            return results
+#        else:
+#            self.trainer_reader_comm.Recv(self.fid_buf, 0)
+#            f_id = self.fid_buf[0]
+#            if f_id == 0:
+#                self.shuffle_files()
+#            elif f_id < 0:
+#                return False
+            #print("self.file",self.files)
+            #print("f_id",f_id)
+            fname = select_dataset_file_for_each_worker(
+                self.files, f_id, self.trainer_num, self.trainer_id)
+            #print("fname",fname)
+            data = read_hdf5_file(fname, dtype=self.data_buf.dtype)
+            #print("data",data)
+            #send_buf = np.concatenate([d.flatten() for d in data])
+            self.data_buf = np.concatenate([d.flatten() for d in data])
+            #print("send_buf",send_buf)
+            #self.trainer_comm.Allgather(send_buf, self.data_buf)
+            #print("self.data_buf",self.data_buf)
+
+            results = process_bert_inputs(self.data_buf, self.num_samples,
+                                          self.max_seq_length, self.batch_size,
+                                          0, 1)
+                                          #self.trainer_id, self.trainer_num)
+#            self.trainer_reader_comm.Send(self.data_buf, dest=0)
+            return results
+#            return True
+
+    def read_eval_file(self):
+        if self.is_trainer:
+
+            eval_data = []
+            for eval_file in sorted(os.listdir(self.eval_dir)):
+                eval_file_path = os.path.join(self.eval_dir, eval_file)
+                if os.path.isfile(eval_file_path) and 'part' in eval_file_path:
+                    data = read_eval_hdf5_file(
+                        eval_file_path, dtype=self.data_buf.dtype)
+                    eval_data.extend(data)
+                    if len(eval_data) > self.num_eval_examples:
+                        break
+
+            chunk_size = self.num_eval_examples // self.trainer_num
+            rank = self.trainer_id
+            remainder = self.num_eval_examples % self.trainer_num
+            if rank < remainder:
+                eval_data = eval_data[(chunk_size + 1) * rank:(chunk_size + 1) *
+                                      (rank + 1)]
+            else:
+                eval_data = eval_data[chunk_size * rank + remainder:chunk_size *
+                                      (rank + 1) + remainder]
+
+            results = process_bert_eval_inputs(eval_data, self.max_seq_length,
+                                               self.eval_batch_size,
+                                               self.args.sort_eval_data)
+            return results
+
+
+_context = Context()
+
+
+def get_context():
+    return _context

rundir_1gpu/run_1gpu.sh

+#!/bin/bash
+
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -ex
+export FLAGS_rocm_dir=/opt/dtk-21.04
+export HIP_LAUNCH_BLOCKING=1
+export FLAGS_max_inplace_grad_add=2
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+#export NCCL_DEBUG=INFO 
+export NCCL_P2P_LEVEL=5
+
+export USE_NV_INPUT=1
+
+USE_UNCOMPRESSED_DATASET=1
+
+BASE_DATA_DIR=${BASE_DATA_DIR:-"/public/DL_DATA/mlperf/bert"}
+export USE_NV_INPUT
+UNCOMPRESSED_DATA_DIR=$BASE_DATA_DIR/hdf5/training-4320/hdf5_4320_shards_uncompressed
+VARLENGTH_DATA_DIR=$BASE_DATA_DIR/hdf5/training-4320/hdf5_4320_shards_varlength
+
+export DATA_DIR=$UNCOMPRESSED_DATA_DIR
+export EVAL_DIR=$BASE_DATA_DIR/hdf5/eval
+if [[ "$USE_NV_INPUT" == "1" && "$USE_UNCOMPRESSED_DATASET" == "0" ]]; then
+  export DATA_DIR="$VARLENGTH_DATA_DIR"
+  export EVAL_DIR=$BASE_DATA_DIR/hdf5/eval
+else
+  export USE_UNCOMPRESSED_DATASET=1
+fi
+export USE_UNCOMPRESSED_DATASET
+export TF_CKPT_PATH=$BASE_DATA_DIR/phase1/model.ckpt-28252.tf_pickled
+export BERT_CONFIG_PATH=$BASE_DATA_DIR/phase1/bert_config.json
+
+export PYTHON=python3
+export PADDLE_TRAINER_ID=${OMPI_COMM_WORLD_RANK}
+export PADDLE_TRAINERS_NUM=${PADDLE_TRAINERS_NUM:-"1"}
+export PADDLE_TRAINER_ENDPOINTS=${PADDLE_TRAINER_ENDPOINTS:-"localhost:60045"}
+#export RCCL_TIMELINE_EXPORT=1
+#export HIP_KERNEL_PRINTF=1 
+#export RCCL_TIMELINE_CFG_FILENAME=.timeline.cfg
+#export RCCL_TIMELINE_DUMP_DIR=/public/home/zhangqha/mlperf_last/bert/paddle/fp16/performance.test.fp16.4096gbs/timeline
+#export RCCL_TIMELINE_MAXPROFILING_THREADS=8
+#export RCCL_TIMELINE_EVENT_SKIP=398
+#export RCCL_TIMELINE_MAXPROFILING_EVENTS=2048
+OMPI_COMM_WORLD_RANK=${OMPI_COMM_WORLD_RANK:-"0"}
+lrank=$OMPI_COMM_WORLD_LOCAL_RANK
+
+function get_device_id() {
+$PYTHON <<EOF
+import paddle
+import os
+gpus = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+if gpus is None:
+    print($OMPI_COMM_WORLD_RANK)
+else:
+    gpus = gpus.split(",")
+    print(gpus[$OMPI_COMM_WORLD_RANK])
+EOF
+}
+if [[ $PADDLE_TRAINER_ID -lt $PADDLE_TRAINERS_NUM ]]; then
+  export CUDA_VISIBLE_DEVICES=4 #$(expr $OMPI_COMM_WORLD_RANK % 4) #`get_device_id`
+  export IS_TRAINER=1
+  export IS_READER=0
+else
+  export CUDA_VISIBLE_DEVICES=""
+  export IS_TRAINER=0
+  export IS_READER=1
+fi
+
+echo "Trainer :" $CUDA_VISIBLE_DEVICES $PADDLE_TRAINER_ENDPOINTS $PADDLE_TRAINERS_NUM
+
+export FLAGS_sync_nccl_allreduce=0
+export FLAGS_fraction_of_gpu_memory_to_use=0.99
+#export FLAGS_allocator_strategy=naive_best_fit
+export FLAGS_call_stack_level=2
+export FLAGS_use_fast_math=0
+#export FLAGS_check_nan_inf=1
+#export CUDA_LAUNCH_BLOCKING=1
+#export FLAGS_benchmark=1
+export FLAGS_enable_nvtx=1
+
+#export FLAGS_inplace_addto_external_ops=custom_fused_dense_grad
+
+batch_size=4
+eval_batch_size=9
+#eval_batch_size=16
+use_amp=True
+use_pure_fp16=True
+
+max_steps=820
+log_freq=50
+eval_iter_start_samples=175000
+eval_iter_samples=175000
+max_seq_length=512
+
+dense_seq_output=True
+unpad=False
+unpad_fmha=False
+fused_bias_mha=True
+fused_bias_fc=True
+## can be False or True 
+weight_transpose=True
+
+###fused_dropout_add_ln=True
+fused_dropout_add_ln=False
+exchange_padding=True
+cpu_exchange_padding=True
+
+distributed_lamb=True
+
+unpad_embed=False
+unpad_fmha_mke_opt=True
+
+sort_eval_data=False
+
+LOG_DIR="log_${PADDLE_TRAINERS_NUM}"
+mkdir -p ${LOG_DIR}
+LOG_FILE=${LOG_DIR}/worker.${PADDLE_TRAINER_ID}
+
+#export FLAGS_lamb_allreduce_first=1
+#export FLAGS_use_multi_tensor_apply=1
+export FLAGS_max_inplace_grad_add=2
+
+if [[ "$exchange_padding" == "true" || "$exchange_padding" == "True" ]]; then
+  if [[ "$cpu_exchange_padding" == "true" || "$cpu_exchange_padding" == "True" ]]; then
+    export DATA_DIR="$UNCOMPRESSED_DATA_DIR"
+  fi
+fi
+
+
+#$NSYS_CMD $BIND_CMD $PYTHON -u run_pretrain.py \
+BERT_CMD="run_pretrain.py \
+   --max_predictions_per_seq 76 \
+   --train_batch_size $batch_size   \
+   --eval_batch_size $eval_batch_size \
+   --sort_eval_data $sort_eval_data \
+   --learning_rate 0.00255 \
+   --weight_decay 0.0 \
+   --lamb_epsilon 1e-06 \
+   --start_warmup_step -76 \
+   --warmup_proportion 0.0 \
+   --warmup_steps 256 \
+   --input_dir $DATA_DIR \
+   --log_freq $log_freq \
+   --max_steps $max_steps \
+   --tf_ckpt_path $TF_CKPT_PATH \
+   --bert_config_path $BERT_CONFIG_PATH \
+   --unpad $unpad \
+   --unpad_fmha $unpad_fmha \
+   --unpad_fmha_mke_opt $unpad_fmha_mke_opt \
+   --unpad_embed $unpad_embed \
+   --fused_bias_mha $fused_bias_mha \
+   --fused_bias_fc $fused_bias_fc \
+   --fused_dropout_add_ln $fused_dropout_add_ln \
+   --weight_transpose $weight_transpose \
+   --max_seq_length $max_seq_length \
+   --eval_dir $EVAL_DIR \
+   --distributed_lamb $distributed_lamb \
+   --exchange_padding $exchange_padding \
+   --cpu_exchange_padding $cpu_exchange_padding \
+   --seed $SEED \
+   --use_uncompressed_dataset $USE_UNCOMPRESSED_DATASET \
+   --dense_seq_output $dense_seq_output \
+   --gradient_accumulation_steps 1 \
+   --opt_lamb_beta_1 0.71 \
+   --opt_lamb_beta_2 0.88 \
+   --enable_addto True \
+   --use_pure_fp16 $use_pure_fp16 \
+   --use_amp $use_amp"
+
+## 2>&1 | tee $LOG_FILE"
+
+#Run experiments
+python3 -u $BERT_CMD
+#echo ${APP}
+#export NCCL_IB_HCA=mlx5_0:1
+#case ${lrank} in
+#case $(expr $lrank % 4) in
+#[0])
+#  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+#  export HIP_VISIBLE_DEVICES=0,1,2,3
+#  export FLAGS_selected_gpus=0
+#  export UCX_NET_DEVICES=mlx5_0:1
+#  export UCX_IB_PCI_BW=mlx5_0:50Gbs
+#  numactl --cpunodebind=0 --membind=0 ${APP} >& $LOG_FILE
+#  ;;
+#[1])
+#  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+#  export HIP_VISIBLE_DEVICES=0,1,2,3
+#  export FLAGS_selected_gpus=1
+#  export UCX_NET_DEVICES=mlx5_1:1
+#  export UCX_IB_PCI_BW=mlx5_1:50Gbs
+#  numactl --cpunodebind=1 --membind=1 ${APP} >& $LOG_FILE
+#  ;;
+#[2])
+#  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+#  export HIP_VISIBLE_DEVICES=0,1,2,3
+#  export FLAGS_selected_gpus=2
+#  export UCX_NET_DEVICES=mlx5_2:1
+#  export UCX_IB_PCI_BW=mlx5_2:50Gbs
+#  numactl --cpunodebind=2 --membind=2 ${APP} >& $LOG_FILE
+#  ;;
+#[3])
+#  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+#  export HIP_VISIBLE_DEVICES=0,1,2,3
+#  export FLAGS_selected_gpus=3
+#  export UCX_NET_DEVICES=mlx5_3:1
+#  export UCX_IB_PCI_BW=mlx5_3:50Gbs
+#  numactl --cpunodebind=3 --membind=3 ${APP}  >& $LOG_FILE
+#  ;;
+#esac

rundir_8gpu/bash/run_8gpu.sh

+#!/bin/bash
+export MIOPEN_FIND_MODE=1
+export PADDLE_TRAINERS_NUM=8
+export PADDLE_TRAINER_ENDPOINTS=localhost:60005,localhost:60006,localhost:60007,localhost:60008,localhost:60009,localhost:60010,localhost:60011,localhost:60012
+export PYTHON=python3
+export SEED=${SEED:-"$RANDOM"}
+export LD_LIBRARY_PATH=/opt/dtk-21.04/rccl/lib:$LD_LIBRARY_PATH
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+export NCCL_P2P_LEVEL=5
+export use_hierarchical_allreduce=True
+export num_process=16
+
+if [[ $num_process -gt 1 ]]; then
+  ORTERUN=`which orterun`
+  mpirun="mpirun --allow-run-as-root -np $num_process --bind-to none -x PADDLE_TRAINERS_NUM -x PADDLE_TRAINER_ENDPOINTS -x LD_LIBRARY_PATH -x SEED -x PYTHON -x NPROC_PER_NODE -x use_hierarchical_allreduce ./run_benchmark.sh"
+else
+  mpirun=""
+fi
+
+echo "command is " $mpirun $CMD
+for NPROC_PER_NODE in 8; do
+#  echo "command is " $mpirun $CMD
+  export NPROC_PER_NODE=$NPROC_PER_NODE
+  $mpirun $CMD
+done
+ 

rundir_8gpu/bash/run_benchmark.sh

+#!/bin/bash
+
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -ex
+#export ROCBLAS_LAYER=3
+export FLAGS_rocm_dir=/opt/dtk-21.04/
+export FLAGS_max_inplace_grad_add=2
+export NCCL_P2P_LEVEL=5
+
+export USE_NV_INPUT=1
+
+USE_UNCOMPRESSED_DATASET=1
+
+BASE_DATA_DIR=${BASE_DATA_DIR:-"/root/mlperf-paddle_bert.20220919-training-bert/training/data"}
+export USE_NV_INPUT
+UNCOMPRESSED_DATA_DIR=$BASE_DATA_DIR/hdf5/training-4320/hdf5_4320_shards_uncompressed
+VARLENGTH_DATA_DIR=$BASE_DATA_DIR/hdf5/training-4320/hdf5_4320_shards_varlength
+
+export DATA_DIR=$UNCOMPRESSED_DATA_DIR
+export EVAL_DIR=$BASE_DATA_DIR/hdf5/eval
+if [[ "$USE_NV_INPUT" == "1" && "$USE_UNCOMPRESSED_DATASET" == "0" ]]; then
+  export DATA_DIR="$VARLENGTH_DATA_DIR"
+  export EVAL_DIR=$BASE_DATA_DIR/hdf5/eval
+else
+  export USE_UNCOMPRESSED_DATASET=1
+fi
+export USE_UNCOMPRESSED_DATASET
+export TF_CKPT_PATH=$BASE_DATA_DIR/phase1/model.ckpt-28252.tf_pickled
+export BERT_CONFIG_PATH=$BASE_DATA_DIR/phase1/bert_config.json
+
+export PYTHON=python3
+export PADDLE_TRAINER_ID=${OMPI_COMM_WORLD_RANK}
+export PADDLE_TRAINERS_NUM=${PADDLE_TRAINERS_NUM:-"1"}
+export PADDLE_TRAINER_ENDPOINTS=${PADDLE_TRAINER_ENDPOINTS:-""}
+
+OMPI_COMM_WORLD_RANK=${OMPI_COMM_WORLD_RANK:-"0"}
+lrank=$OMPI_COMM_WORLD_LOCAL_RANK
+
+function get_device_id() {
+$PYTHON <<EOF
+import paddle
+import os
+gpus = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+if gpus is None:
+    print($OMPI_COMM_WORLD_RANK)
+else:
+    gpus = gpus.split(",")
+    print(gpus[$OMPI_COMM_WORLD_RANK])
+EOF
+}
+if [[ $PADDLE_TRAINER_ID -lt $PADDLE_TRAINERS_NUM ]]; then
+  export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 #$(expr $OMPI_COMM_WORLD_RANK % 4) #`get_device_id`
+  export IS_TRAINER=1
+  export IS_READER=0
+else
+  export CUDA_VISIBLE_DEVICES=""
+  export IS_TRAINER=0
+  export IS_READER=1
+fi
+
+echo "Trainer :" $CUDA_VISIBLE_DEVICES $PADDLE_TRAINER_ENDPOINTS $PADDLE_TRAINERS_NUM
+
+export FLAGS_sync_nccl_allreduce=0
+export FLAGS_fraction_of_gpu_memory_to_use=0.99
+#export FLAGS_allocator_strategy=naive_best_fit
+export FLAGS_call_stack_level=2
+export FLAGS_use_fast_math=0
+#export FLAGS_check_nan_inf=1
+#export PADDLE_INF_NAN_SKIP_OP="matmul_v2_grad,reshape2_grad,cast,elementwise_add_grad,reduce_mean_grad,softmax_with_cross_entropy_grad"
+#export CUDA_LAUNCH_BLOCKING=1
+#export FLAGS_benchmark=1
+export FLAGS_enable_nvtx=1
+
+#export FLAGS_inplace_addto_external_ops=custom_fused_dense_grad
+
+batch_size=4
+eval_batch_size=63
+#eval_batch_size=16
+use_amp=True
+use_pure_fp16=True
+
+max_steps=7100
+log_freq=50
+eval_iter_start_samples=150000
+eval_iter_samples=150000
+max_seq_length=512
+
+dense_seq_output=True
+unpad=False
+unpad_fmha=False
+fused_bias_mha=True
+fused_bias_fc=True
+## can be False or True 
+weight_transpose=True
+
+#fused_dropout_add_ln=True
+fused_dropout_add_ln=False
+exchange_padding=True
+cpu_exchange_padding=True
+
+distributed_lamb=True
+
+unpad_embed=False
+unpad_fmha_mke_opt=True
+
+sort_eval_data=False
+
+LOG_DIR="log_${PADDLE_TRAINERS_NUM}"
+mkdir -p ${LOG_DIR}
+LOG_FILE=${LOG_DIR}/worker.${PADDLE_TRAINER_ID}
+
+#export FLAGS_lamb_allreduce_first=1
+#export FLAGS_use_multi_tensor_apply=1
+export FLAGS_max_inplace_grad_add=2
+
+if [[ "$exchange_padding" == "true" || "$exchange_padding" == "True" ]]; then
+  if [[ "$cpu_exchange_padding" == "true" || "$cpu_exchange_padding" == "True" ]]; then
+    export DATA_DIR="$UNCOMPRESSED_DATA_DIR"
+  fi
+fi
+
+
+#$NSYS_CMD $BIND_CMD $PYTHON -u run_pretrain.py \
+BERT_CMD="run_pretrain.py \
+   --max_predictions_per_seq 76 \
+   --train_batch_size $batch_size   \
+   --eval_batch_size $eval_batch_size \
+   --sort_eval_data $sort_eval_data \
+   --learning_rate 0.000425 \
+   --weight_decay 1e-2 \
+   --lamb_epsilon 1e-6 \
+   --start_warmup_step 0 \
+   --warmup_proportion 0.0 \
+   --warmup_steps 0 \
+   --input_dir $DATA_DIR \
+   --log_freq $log_freq \
+   --max_steps $max_steps \
+   --tf_ckpt_path $TF_CKPT_PATH \
+   --bert_config_path $BERT_CONFIG_PATH \
+   --unpad $unpad \
+   --unpad_fmha $unpad_fmha \
+   --unpad_fmha_mke_opt $unpad_fmha_mke_opt \
+   --unpad_embed $unpad_embed \
+   --fused_bias_mha $fused_bias_mha \
+   --fused_bias_fc $fused_bias_fc \
+   --fused_dropout_add_ln $fused_dropout_add_ln \
+   --weight_transpose $weight_transpose \
+   --max_seq_length $max_seq_length \
+   --eval_dir $EVAL_DIR \
+   --distributed_lamb $distributed_lamb \
+   --exchange_padding $exchange_padding \
+   --cpu_exchange_padding $cpu_exchange_padding \
+   --seed $SEED \
+   --use_uncompressed_dataset $USE_UNCOMPRESSED_DATASET \
+   --dense_seq_output $dense_seq_output \
+   --gradient_accumulation_steps 14 \
+   --opt_lamb_beta_1 0.9 \
+   --opt_lamb_beta_2 0.999 \
+   --enable_addto True \
+   --use_pure_fp16 $use_pure_fp16 \
+   --use_amp $use_amp"
+
+## 2>&1 | tee $LOG_FILE"
+
+APP="python3 -u $BERT_CMD"
+case $(expr $lrank % 8) in 
+[0])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=0
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_0:50Gbs
+  numactl --cpunodebind=0 --membind=0 ${APP} >& $LOG_FILE
+  ;;
+[1])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=1
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_1:50Gbs
+  numactl --cpunodebind=1 --membind=1 ${APP} >& $LOG_FILE
+  ;;
+[2])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=2
+ # export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_2:50Gbs
+  numactl --cpunodebind=2 --membind=2 ${APP} >& $LOG_FILE
+  ;;
+[3])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=3
+#  export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_3:50Gbs
+  numactl --cpunodebind=3 --membind=3 ${APP}  >& $LOG_FILE
+  ;;
+[4])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=4
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_0:50Gbs
+  numactl --cpunodebind=4 --membind=4 ${APP} >& $LOG_FILE
+  ;;
+[5])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=5
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_1:50Gbs
+  numactl --cpunodebind=5 --membind=5 ${APP} >& $LOG_FILE
+  ;;
+[6])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=6
+ # export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_2:50Gbs
+  numactl --cpunodebind=6 --membind=6 ${APP} >& $LOG_FILE
+  ;;
+[7])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+  export FLAGS_selected_gpus=7
+#  export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_3:50Gbs
+  numactl --cpunodebind=7 --membind=7 ${APP}  >& $LOG_FILE
+  ;;
+
+esac
+

rundir_8gpu/init_env.py

+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from mpi4py import MPI
+import numpy as np
+import time
+import paddle
+from pybind.functions import process_allgathered_inputs as process_bert_inputs
+from pybind.functions import process_eval_inputs as process_bert_eval_inputs
+import h5py
+import random
+
+global_comm = MPI.COMM_WORLD
+global_rank = global_comm.rank
+global_world_size = global_comm.size
+assert global_world_size % 2 == 0
+
+
+def create_group_comm(ranks):
+    ranks = list(ranks)
+    new_group = global_comm.group.Incl(ranks)
+    new_comm = global_comm.Create_group(new_group)
+    return new_comm
+
+
+def generate_seeds(rng, size):
+    """
+    Generate list of random seeds
+
+    :param rng: random number generator
+    :param size: length of the returned list
+    """
+    seeds = [rng.randint(0, 2**32 - 1) for _ in range(size)]
+    return seeds
+
+
+def broadcast_seeds(comm, seeds, root=0):
+    seeds = np.array(seeds).astype(np.int64)
+    comm.Bcast(seeds, root=root)
+    return seeds.tolist()
+
+
+def select_dataset_file_for_each_worker(files, f_start_id, worker_num,
+                                        worker_index):
+    """
+    Spliting the train file according to the worker index.
+    """
+    num_files = len(files)
+    if worker_num > num_files:
+        remainder = worker_num % num_files
+        data_file = files[(
+            f_start_id * worker_num + worker_index + remainder * f_start_id) %
+                          num_files]
+    else:
+        data_file = files[(f_start_id * worker_num + worker_index) % num_files]
+    # limin-todo: 
+    #data_file = "/data2/zengjinle/dataset/bert_data/hdf5/training-4320/hdf5_4320_shards_uncompressed/part_01799_of_04320.hdf5"
+    #print("data_file: ", data_file)
+    return data_file
+
+
+def read_hdf5_file(input_file, dtype=np.int16):
+    keys = [
+        'input_ids',
+        'input_mask',
+        'segment_ids',
+        'masked_lm_positions',
+        'masked_lm_ids',
+        'next_sentence_labels',
+    ]
+    if not os.path.exists(input_file):
+        return None
+    with h5py.File(input_file, 'r') as f:
+        outputs = [np.array(f[key], dtype=dtype) for key in keys]
+        n = outputs[0].shape[0]
+        masked_lm_labels = np.zeros(outputs[0].shape, dtype=dtype)
+        lengths = np.zeros(n, dtype=dtype)
+        for i in range(n):
+            masked_lm_positions = outputs[3][i]
+            masked_lm_ids = outputs[4][i]
+            length = np.count_nonzero(masked_lm_positions)
+            masked_lm_labels[i][
+                masked_lm_positions[:length]] = masked_lm_ids[:length]
+            lengths[i] = np.count_nonzero(outputs[1][i])
+        outputs = [
+            outputs[0], outputs[2], outputs[1], masked_lm_labels, outputs[-1],
+            lengths
+        ]
+        idx = np.random.choice(np.arange(n), n, replace=False)
+        for i in range(len(outputs)):
+            outputs[i] = outputs[i][idx]
+    return outputs
+
+
+def read_eval_hdf5_file(input_file, dtype=np.int16):
+    keys = [
+        'input_ids',
+        'input_mask',
+        'segment_ids',
+        'masked_lm_positions',
+        'masked_lm_ids',
+        'next_sentence_labels',
+    ]
+    if not os.path.exists(input_file):
+        return None
+    with h5py.File(input_file, 'r') as f:
+        outputs = [np.asarray(f[key][:]) for key in keys]
+        nsamples = outputs[0].shape[0]
+
+        all_data = []
+        for index in range(nsamples):
+            [
+                input_ids, input_mask, segment_ids, masked_lm_positions,
+                masked_lm_ids, next_sentence_labels
+            ] = [
+                input[index].astype(dtype)
+                if indice < 5 else np.asarray(input[index].astype(dtype))
+                for indice, input in enumerate(outputs)
+            ]
+
+            length = np.count_nonzero(masked_lm_positions)
+            masked_lm_positions = masked_lm_positions[:length]
+            masked_lm_ids = masked_lm_ids[:length]
+
+            masked_lm_labels = np.zeros(input_ids.shape, dtype=dtype)
+            masked_lm_labels[masked_lm_positions] = masked_lm_ids
+
+            #if index == 0:
+            #    print("masked_lm_labels = ", masked_lm_labels)
+            #    print("masked_lm_positions = ", masked_lm_positions)
+            #    print("masked_lm_ids = ", masked_lm_ids)
+            seq_len = np.asarray(np.count_nonzero(input_mask))
+
+            data = [
+                input_ids,
+                segment_ids,
+                input_mask,
+                masked_lm_labels,
+                next_sentence_labels,
+                seq_len,
+            ]
+            # (2050, ), i.e., 512 * 4 + 1 + 1
+            one_sample_data = np.concatenate([d.flatten() for d in data])
+
+            all_data.extend(one_sample_data)
+
+        # (2050000, ) -> (10000, 2050)
+    return np.asarray(all_data).reshape((nsamples, -1))
+
+
+class WorkerInitObj(object):
+    "Construct the object with different seed, and the Dataloader will generate the data "
+    "with different seed in each worker."
+
+    def __init__(self, seed):
+        self.seed = seed
+
+    def __call__(self, id):
+        np.random.seed(seed=self.seed + id)
+        random.seed(self.seed + id)
+
+
+class Context:
+    def __init__(self):
+        half_size = int(global_world_size / 2)
+        self.trainer_id = global_rank % half_size
+        self.trainer_num = half_size
+        self.is_trainer = (global_rank < half_size)
+
+        self.reader_id = self.trainer_id
+        self.reader_num = self.trainer_num
+        self.is_reader = not self.is_trainer
+
+        self.trainer_comm = create_group_comm(range(0, half_size))
+        self.reader_comm = create_group_comm(
+            range(half_size, global_world_size))
+        self.trainer_reader_comm = create_group_comm(
+            [self.trainer_id, self.trainer_id + half_size])
+        self.global_comm = global_comm
+
+    def init_args(self, args, dtype=np.int16):
+        self.args = args
+        self.files = [
+            os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir)
+            if os.path.isfile(os.path.join(args.input_dir, f)) and "part" in f
+        ]
+        self.files.sort()
+        self.fid_buf = np.array([1], dtype=np.int64)
+        with h5py.File(self.files[0], 'r') as f:
+            self.num_samples = np.array(f["next_sentence_labels"][:]).size
+
+        self.batch_size = args.train_batch_size
+        self.max_seq_length = args.max_seq_length
+        self.worker_seeds, self.shuffling_seeds = self._setup_seeds(
+            args.seed, args.num_epochs_to_generate_seeds_for)
+        self.epoch_idx = 0
+
+        data_buf_size = self.num_samples * 4 * self.max_seq_length + self.num_samples * 2
+        self.data_buf = np.empty(
+            shape=[self.trainer_num * data_buf_size], dtype=dtype)
+
+        self.eval_dir = args.eval_dir
+        self.num_eval_examples = args.num_eval_examples
+        self.eval_batch_size = args.eval_batch_size
+
+        cur_seed = self.worker_seeds[self.trainer_id]
+        np.random.seed(cur_seed)
+        random.seed(cur_seed)
+        paddle.seed(cur_seed)
+        self.worker_init = WorkerInitObj(cur_seed)
+        self.barrier()
+
+    def shuffle_files(self):
+        random.Random(self.shuffling_seeds[self.epoch_idx]).shuffle(self.files)
+        self.epoch_idx += 1
+
+    def _setup_seeds(self, master_seed, epochs):
+        if master_seed is None:
+            master_seed = random.SystemRandom().randint(0, 2**32 - 1)
+            if self.trainer_id == 0:
+                print('Using random master seed: {}'.format(master_seed))
+        else:
+            print('Using master seed from command line: {}'.format(master_seed))
+
+        # initialize seeding RNG
+        seeding_rng = random.Random(master_seed)
+
+        # generate worker seeds, one seed for every distributed worker
+        worker_seeds = generate_seeds(seeding_rng, self.trainer_num)
+
+        # generate seeds for data shuffling, one seed for every epoch
+        shuffling_seeds = generate_seeds(seeding_rng, epochs)
+
+        worker_seeds = broadcast_seeds(self.global_comm, worker_seeds)
+        shuffling_seeds = broadcast_seeds(self.global_comm, shuffling_seeds)
+        return worker_seeds, shuffling_seeds
+
+    def worker_seed(self):
+        return self.worker_seeds[self.trainer_id]
+
+    def barrier(self):
+        self.global_comm.barrier()
+
+    def stop_reader(self):
+        if self.is_trainer:
+            self.read_file(-1)
+
+    def file_num(self):
+        return len(self.files)
+
+    def read_file(self, f_id=None):
+        if self.is_trainer:
+            self.fid_buf[0] = f_id
+            self.trainer_reader_comm.Isend(self.fid_buf, dest=1)
+            if f_id == 0:
+                self.shuffle_files()
+            elif f_id < 0:
+                return
+
+            self.trainer_reader_comm.Recv(self.data_buf, source=1)
+            results = process_bert_inputs(self.data_buf, self.num_samples,
+                                          self.max_seq_length, self.batch_size,
+                                          self.trainer_id, self.trainer_num)
+
+            return results
+        else:
+            self.trainer_reader_comm.Recv(self.fid_buf, 0)
+            f_id = self.fid_buf[0]
+            if f_id == 0:
+                self.shuffle_files()
+            elif f_id < 0:
+                return False
+
+            fname = select_dataset_file_for_each_worker(
+                self.files, f_id, self.trainer_num, self.trainer_id)
+            data = read_hdf5_file(fname, dtype=self.data_buf.dtype)
+            send_buf = np.concatenate([d.flatten() for d in data])
+            self.reader_comm.Allgather(send_buf, self.data_buf)
+            self.trainer_reader_comm.Send(self.data_buf, dest=0)
+            return True
+
+    def read_eval_file(self):
+        if self.is_trainer:
+
+            eval_data = []
+            for eval_file in sorted(os.listdir(self.eval_dir)):
+                eval_file_path = os.path.join(self.eval_dir, eval_file)
+                if os.path.isfile(eval_file_path) and 'part' in eval_file_path:
+                    data = read_eval_hdf5_file(
+                        eval_file_path, dtype=self.data_buf.dtype)
+                    eval_data.extend(data)
+                    if len(eval_data) > self.num_eval_examples:
+                        break
+
+            chunk_size = self.num_eval_examples // self.trainer_num
+            rank = self.trainer_id
+            remainder = self.num_eval_examples % self.trainer_num
+            if rank < remainder:
+                eval_data = eval_data[(chunk_size + 1) * rank:(chunk_size + 1) *
+                                      (rank + 1)]
+            else:
+                eval_data = eval_data[chunk_size * rank + remainder:chunk_size *
+                                      (rank + 1) + remainder]
+
+            results = process_bert_eval_inputs(eval_data, self.max_seq_length,
+                                               self.eval_batch_size,
+                                               self.args.sort_eval_data)
+            return results
+
+
+_context = Context()
+
+
+def get_context():
+    return _context

rundir_8gpu/slurm/run_benchmark.sh

+#!/bin/bash
+
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -ex
+#export ROCBLAS_LAYER=3
+export FLAGS_rocm_dir=/public/home/zhangqha/dtk-21.04/
+export HIP_LAUNCH_BLOCKING=1
+export FLAGS_max_inplace_grad_add=2
+export NCCL_P2P_LEVEL=5
+
+export USE_NV_INPUT=1
+
+USE_UNCOMPRESSED_DATASET=1
+
+BASE_DATA_DIR=${BASE_DATA_DIR:-"/public/share/huchen_ty/data"}
+export USE_NV_INPUT
+UNCOMPRESSED_DATA_DIR=$BASE_DATA_DIR/hdf5/training-4320/hdf5_4320_shards_uncompressed
+VARLENGTH_DATA_DIR=$BASE_DATA_DIR/hdf5/training-4320/hdf5_4320_shards_varlength
+
+export DATA_DIR=$UNCOMPRESSED_DATA_DIR
+export EVAL_DIR=$BASE_DATA_DIR/hdf5/eval
+if [[ "$USE_NV_INPUT" == "1" && "$USE_UNCOMPRESSED_DATASET" == "0" ]]; then
+  export DATA_DIR="$VARLENGTH_DATA_DIR"
+  export EVAL_DIR=$BASE_DATA_DIR/hdf5/eval
+else
+  export USE_UNCOMPRESSED_DATASET=1
+fi
+export USE_UNCOMPRESSED_DATASET
+export TF_CKPT_PATH=$BASE_DATA_DIR/phase1/model.ckpt-28252.tf_pickled
+export BERT_CONFIG_PATH=$BASE_DATA_DIR/phase1/bert_config.json
+
+export PYTHON=python3
+export PADDLE_TRAINER_ID=${OMPI_COMM_WORLD_RANK}
+export PADDLE_TRAINERS_NUM=${PADDLE_TRAINERS_NUM:-"1"}
+export PADDLE_TRAINER_ENDPOINTS=${PADDLE_TRAINER_ENDPOINTS:-""}
+
+OMPI_COMM_WORLD_RANK=${OMPI_COMM_WORLD_RANK:-"0"}
+lrank=$OMPI_COMM_WORLD_LOCAL_RANK
+
+function get_device_id() {
+$PYTHON <<EOF
+import paddle
+import os
+gpus = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+if gpus is None:
+    print($OMPI_COMM_WORLD_RANK)
+else:
+    gpus = gpus.split(",")
+    print(gpus[$OMPI_COMM_WORLD_RANK])
+EOF
+}
+if [[ $PADDLE_TRAINER_ID -lt $PADDLE_TRAINERS_NUM ]]; then
+  export CUDA_VISIBLE_DEVICES=0,1,2,3 #$(expr $OMPI_COMM_WORLD_RANK % 4) #`get_device_id`
+  export IS_TRAINER=1
+  export IS_READER=0
+else
+  export CUDA_VISIBLE_DEVICES=""
+  export IS_TRAINER=0
+  export IS_READER=1
+fi
+
+echo "Trainer :" $CUDA_VISIBLE_DEVICES $PADDLE_TRAINER_ENDPOINTS $PADDLE_TRAINERS_NUM
+
+export FLAGS_sync_nccl_allreduce=0
+export FLAGS_fraction_of_gpu_memory_to_use=0.99
+#export FLAGS_allocator_strategy=naive_best_fit
+export FLAGS_call_stack_level=2
+export FLAGS_use_fast_math=0
+#export FLAGS_check_nan_inf=1
+#export PADDLE_INF_NAN_SKIP_OP="matmul_v2_grad,reshape2_grad,cast,elementwise_add_grad,reduce_mean_grad,softmax_with_cross_entropy_grad"
+#export CUDA_LAUNCH_BLOCKING=1
+#export FLAGS_benchmark=1
+export FLAGS_enable_nvtx=1
+
+#export FLAGS_inplace_addto_external_ops=custom_fused_dense_grad
+
+batch_size=4
+eval_batch_size=63
+#eval_batch_size=16
+use_amp=True
+use_pure_fp16=True
+
+max_steps=7100
+log_freq=50
+eval_iter_start_samples=150000
+eval_iter_samples=150000
+max_seq_length=512
+
+dense_seq_output=True
+unpad=False
+unpad_fmha=False
+fused_bias_mha=True
+fused_bias_fc=True
+## can be False or True 
+weight_transpose=True
+
+#fused_dropout_add_ln=True
+fused_dropout_add_ln=False
+exchange_padding=True
+cpu_exchange_padding=True
+
+distributed_lamb=True
+
+unpad_embed=False
+unpad_fmha_mke_opt=True
+
+sort_eval_data=False
+
+LOG_DIR="log_${PADDLE_TRAINERS_NUM}"
+mkdir -p ${LOG_DIR}
+LOG_FILE=${LOG_DIR}/worker.${PADDLE_TRAINER_ID}
+
+#export FLAGS_lamb_allreduce_first=1
+#export FLAGS_use_multi_tensor_apply=1
+export FLAGS_max_inplace_grad_add=2
+
+if [[ "$exchange_padding" == "true" || "$exchange_padding" == "True" ]]; then
+  if [[ "$cpu_exchange_padding" == "true" || "$cpu_exchange_padding" == "True" ]]; then
+    export DATA_DIR="$UNCOMPRESSED_DATA_DIR"
+  fi
+fi
+
+
+#$NSYS_CMD $BIND_CMD $PYTHON -u run_pretrain.py \
+BERT_CMD="run_pretrain.py \
+   --max_predictions_per_seq 76 \
+   --train_batch_size $batch_size   \
+   --eval_batch_size $eval_batch_size \
+   --sort_eval_data $sort_eval_data \
+   --learning_rate 0.000425 \
+   --weight_decay 1e-2 \
+   --lamb_epsilon 1e-6 \
+   --start_warmup_step 0 \
+   --warmup_proportion 0.0 \
+   --warmup_steps 0 \
+   --input_dir $DATA_DIR \
+   --log_freq $log_freq \
+   --max_steps $max_steps \
+   --tf_ckpt_path $TF_CKPT_PATH \
+   --bert_config_path $BERT_CONFIG_PATH \
+   --unpad $unpad \
+   --unpad_fmha $unpad_fmha \
+   --unpad_fmha_mke_opt $unpad_fmha_mke_opt \
+   --unpad_embed $unpad_embed \
+   --fused_bias_mha $fused_bias_mha \
+   --fused_bias_fc $fused_bias_fc \
+   --fused_dropout_add_ln $fused_dropout_add_ln \
+   --weight_transpose $weight_transpose \
+   --max_seq_length $max_seq_length \
+   --eval_dir $EVAL_DIR \
+   --distributed_lamb $distributed_lamb \
+   --exchange_padding $exchange_padding \
+   --cpu_exchange_padding $cpu_exchange_padding \
+   --seed $SEED \
+   --use_uncompressed_dataset $USE_UNCOMPRESSED_DATASET \
+   --dense_seq_output $dense_seq_output \
+   --gradient_accumulation_steps 14 \
+   --opt_lamb_beta_1 0.9 \
+   --opt_lamb_beta_2 0.999 \
+   --enable_addto True \
+   --use_pure_fp16 $use_pure_fp16 \
+   --use_amp $use_amp"
+
+## 2>&1 | tee $LOG_FILE"
+
+APP="python3 -u $BERT_CMD"
+case $(expr $lrank % 4) in
+[0])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3
+  export FLAGS_selected_gpus=0
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_0:50Gbs
+  numactl --cpunodebind=0 --membind=0 ${APP} >& $LOG_FILE
+  ;;
+[1])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3
+  export FLAGS_selected_gpus=1
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_1:50Gbs
+  numactl --cpunodebind=1 --membind=1 ${APP} >& $LOG_FILE
+  ;;
+[2])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3
+  export FLAGS_selected_gpus=2
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_2:50Gbs
+  numactl --cpunodebind=2 --membind=2 ${APP} >& $LOG_FILE
+  ;;
+[3])
+  echo "work ${lrank} less than ${PADDLE_TRAINERS_NUM} on DCU $(expr $lrank % 4)"
+  export HIP_VISIBLE_DEVICES=0,1,2,3
+  export FLAGS_selected_gpus=3
+  #export UCX_NET_DEVICES=mlx5_0:1
+  #export UCX_IB_PCI_BW=mlx5_3:50Gbs
+  numactl --cpunodebind=3 --membind=3 ${APP}  >& $LOG_FILE
+  ;;
+esac
+

rundir_8gpu/slurm/run_sbatch.sh

+#!/bin/bash
+#SBATCH -p ty_huchen 
+#SBATCH -N 2
+#SBATCH --ntasks-per-node=8
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=dcu:4
+#SBATCH -J mlperf
+source ~/paddle_dtk21.04.sh
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+export NCCL_GRAPH_FILE=./text.xml
+export PYTHON=python3
+export PADDLE_TRAINER_ENDPOINTS=`$PYTHON -c "import list;print(list.get_list())"`
+echo $PADDLE_TRAINER_ENDPOINTS
+#set -e
+hostfile=./hostfile
+scontrol show hostnames $SLURM_JOB_NODELIST > ${hostfile}
+num_node=$(cat $hostfile|sort|uniq |wc -l)
+num_DCU=$(($num_node*4))
+export LD_LIBRARY_PATH=/public/home/zhangqha/dtk-21.04/rccl/lib:$LD_LIBRARY_PATH
+export LD_LIBRARY_PATH=/public/home/zhangqha/ucx/lib:$LD_LIBRARY_PATH
+export NCCL_PLUGIN_P2P=ucx
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+export NCCL_P2P_LEVEL=5
+rm hosts
+p=0
+for i in `cat hostfile`;do
+    for j in `seq 1 4`;do
+        num=$[ $p * 4 + $j - 1]
+        echo "rank ${num}=${i} slot=$[j-1]" >> hosts
+    done
+    p=$(expr $p + 1)
+done
+for i in `cat hostfile`;do
+    for j in `seq 1 4`;do
+        num=$[ $p * 4 + $j - 1]
+        echo "rank ${num}=${i} slot=$[j+3]" >> hosts
+    done
+    p=$(expr $p + 1)
+done
+nodename=$(cat $hostfile |sed -n "1p")
+rm $SLURM_JOB_ID
+
+#ldconfig
+export HIP_LAUNCH_BLOCKING=1
+
+export PADDLE_TRAINERS_NUM=$num_DCU
+
+
+export SEED=${SEED:-"$RANDOM"}
+echo "PADDLE_TRAINER_ENDPOINTS " $PADDLE_TRAINER_ENDPOINTS
+
+CMD="bash run_benchmark.sh"
+
+
+#bash kill_grep.sh $PYTHON || true
+
+num_process=$(($PADDLE_TRAINERS_NUM*1))
+
+if [[ $num_process -gt 1 ]]; then
+  ORTERUN=`which orterun`
+  mpirun="mpirun --allow-run-as-root -np $num_process -machinefile ./hosts -mca pml ucx -x UCX_IB_ADDR_TYPE=ib_global -x UCX_TLS=rc_x,rocm_copy -mca btl_tcp_if_exclude ib0 --bind-to none -x PADDLE_TRAINERS_NUM -x PADDLE_TRAINER_ENDPOINTS -x LD_LIBRARY_PATH -x SEED -x PYTHON -x NPROC_PER_NODE"
+else
+  mpirun=""
+fi
+
+echo "command is " $mpirun $CMD
+for NPROC_PER_NODE in 4; do
+  echo "command is " $mpirun $CMD
+  export NPROC_PER_NODE=$NPROC_PER_NODE
+  $mpirun $CMD
+done
+# bash kill_grep.sh run_and_time || true
+# bash kill_grep.sh python || true

setup.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import paddle.fluid.core as core
+from paddle.utils.cpp_extension import CUDAExtension, setup
+
+compile_dir = os.environ.get('/public/home/zhangqha/padd_install/padd_build.bert/Paddle', '/public/home/zhangqha/padd_install/padd_build.bert/Paddle/build')
+
+define_macros = []
+if core.is_compiled_with_mkldnn():
+    define_macros.append(('PADDLE_WITH_MKLDNN', None))
+if core.is_compiled_with_nccl():
+    define_macros.append(('PADDLE_WITH_NCCL', None))
+define_macros.append(('PADDLE_WITH_HIP', None))
+define_macros.append(('PADDLE_WITH_RCCL', None))
+
+setup(
+    name='custom_setup_ops',
+    ext_modules=CUDAExtension(
+        sources=[
+            #'./fused_dense_op/fused_dense_cuda.cc',
+            #'./fused_dense_op/fused_dense_cuda.cu',
+            ###'./fused_dropout_residual_ln/fused_dropout_residual_ln_cuda.cc',
+            ###'./fused_dropout_residual_ln/fused_dropout_residual_ln_cuda.cu',
+            #'./fmhalib/fmha_cuda.cc',
+            #'./fmhalib/fmha_cuda.cu',
+            './sort_bert_inputs_across_devices/sort_bert_inputs_across_devices.cc',
+            './sort_bert_inputs_across_devices/sort_bert_inputs_across_devices.cu',
+            './lr_op/lr_op_cuda.cc',
+            './lr_op/lr_op_cuda.cu',
+            './acc_merge/acc_merge.cc',
+            './acc_merge/acc_merge.cu',
+        ],
+        # extra_objects=[os.path.join(apex_lib_dir, 'libfmha.so')],
+        # include_dirs=[apex_dir],
+        # library_dirs=[apex_lib_dir],
+        # extra_link_args=['-lfmha', '-ldl', '-lcublas'],
+        _compile_dir=compile_dir,
+        define_macros=define_macros))

stack.py

+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+__all__ = ['Stack', 'Pad', 'Tuple', 'Dict']
+
+
+class Stack(object):
+    """
+    Stacks the input data samples to construct the batch. The N input samples
+    must have the same shape/length and will be stacked to construct a batch.
+
+    Args:
+        axis (int, optional): The axis in the result data along which the input
+            data are stacked. Default: 0.
+        dtype (str|numpy.dtype, optional): The value type of the output. If it
+            is set to None, the type of input data is used. Default: None.
+    """
+
+    def __init__(self, axis=0, dtype=None):
+        self._axis = axis
+        self._dtype = dtype
+
+    def __call__(self, data):
+        """
+        Batchifies the input data by stacking.
+
+        Args:
+            data (list[numpy.ndarray]): The input data samples. It is a list. 
+                Each element is a numpy.ndarray or list.
+
+        Returns:
+            numpy.ndarray: Stacked batch data.
+
+
+        Example:
+            .. code-block:: python
+
+                from paddlenlp.data import Stack
+                a = [1, 2, 3, 4]
+                b = [3, 4, 5, 6]
+                c = [5, 6, 7, 8]
+                result = Stack()([a, b, c])
+                '''
+                [[1, 2, 3, 4],
+                 [3, 4, 5, 6],
+                 [5, 6, 7, 8]]
+                '''
+        """
+        data = np.stack(
+            data,
+            axis=self._axis).astype(self._dtype) if self._dtype else np.stack(
+                data, axis=self._axis)
+        return data
+
+
+class Pad(object):
+    """
+    Pads the input data samples to the largest length at `axis`.
+
+    Args:
+        pad_val (float|int, optional): The padding value. Default: 0.
+        axis (int, optional): The axis to pad the arrays. The arrays will be
+            padded to the largest length at `axis`. For example, assume the 
+            input arrays have shape (10, 8, 5), (6, 8, 5), (3, 8, 5) and the 
+            axis is 0. Each input will be padded into (10, 8, 5) and then 
+            stacked to form the final output, which has shape (3, 10, 8, 5). 
+            Default: 0.
+        ret_length (bool|numpy.dtype, optional): If it is bool, indicate whether
+            to return the valid length in the output, and the data type of
+            returned length is int32 if True. If it is numpy.dtype, indicate the
+            data type of returned length. Default: None.
+        dtype (numpy.dtype, optional): The value type of the output. If it is
+            set to None, the input data type is used. Default: None.
+        pad_right (bool, optional): Whether the padding direction is right-side. 
+            If True, it indicates we pad to the right side, while False indicates 
+            we pad to the left side. Default: True.
+     """
+
+    def __init__(self,
+                 pad_val=0,
+                 axis=0,
+                 ret_length=None,
+                 dtype=None,
+                 pad_right=True):
+        self._pad_val = pad_val
+        self._axis = axis
+        self._ret_length = ret_length
+        self._dtype = dtype
+        self._pad_right = pad_right
+
+    def __call__(self, data):
+        """
+        Batchifies the input data by padding. The input will be padded to the 
+        largest dimension at `axis` and then stacked to form the final output. 
+        In addition, the function will output the original dimensions at the 
+        `axis` if `ret_length` is not None or False.
+
+        Args:
+            data (list[numpy.ndarray|list]): The input data samples. It is a 
+                list. Each element is a numpy.ndarray or list.
+
+        Returns:
+            numpy.ndarray|tuple[numpy.ndarray]: If `ret_length` is False, it 
+            is a numpy.ndarray representing the padded batch data and the 
+            shape is (N, …). Otherwise, it is a tuple, besides the padded batch 
+            data, the tuple also includes a numpy.ndarray representing original 
+            length at `axis` of all input samples, which shaped `(N,)`. 
+
+        Example:
+            .. code-block:: python
+
+                from paddlenlp.data import Pad
+                a = [1, 2, 3, 4]
+                b = [5, 6, 7]
+                c = [8, 9]
+                result = Pad(pad_val=0)([a, b, c])
+                '''
+                [[1, 2, 3, 4],
+                 [5, 6, 7, 0],
+                 [8, 9, 0, 0]]
+                '''
+        """
+        arrs = [np.asarray(ele) for ele in data]
+        original_length = [ele.shape[self._axis] for ele in arrs]
+        max_size = max(original_length)
+        ret_shape = list(arrs[0].shape)
+        ret_shape[self._axis] = max_size
+        ret_shape = (len(arrs), ) + tuple(ret_shape)
+        ret = np.full(
+            shape=ret_shape,
+            fill_value=self._pad_val,
+            dtype=arrs[0].dtype if self._dtype is None else self._dtype)
+        for i, arr in enumerate(arrs):
+            if arr.shape[self._axis] == max_size:
+                ret[i] = arr
+            else:
+                slices = [slice(None) for _ in range(arr.ndim)]
+                if self._pad_right:
+                    slices[self._axis] = slice(0, arr.shape[self._axis])
+                else:
+                    slices[self._axis] = slice(max_size - arr.shape[self._axis],
+                                               max_size)
+
+                if slices[self._axis].start != slices[self._axis].stop:
+                    slices = [slice(i, i + 1)] + slices
+                    ret[tuple(slices)] = arr
+        if self._ret_length:
+            return ret, np.asarray(
+                original_length,
+                dtype="int32") if self._ret_length == True else np.asarray(
+                    original_length, self._ret_length)
+        else:
+            return ret
+
+
+class Tuple(object):
+    """
+    Wraps multiple batchify functions together. The input functions will be applied
+    to the corresponding input fields.
+    
+    Each sample should be a list or tuple containing multiple fields. The i'th
+    batchify function stored in Tuple will be applied on the i'th field. 
+    
+    For example, when data sample is (nd_data, label), you can wrap two batchify
+    functions using `Tuple(DataBatchify, LabelBatchify)` to batchify nd_data and
+    label correspondingly.
+
+    Args:
+        fn (callable|list[callable]|tuple[callable]): The batchify functions to 
+            wrap. It is a callable function or a list/tuple of callable functions.
+        args (tuple[callable]): The additional batchify functions to wrap.
+    """
+
+    def __init__(self, fn, *args):
+        if isinstance(fn, (list, tuple)):
+            assert len(args) == 0, 'Input pattern not understood. The input of Tuple can be ' \
+                                   'Tuple(A, B, C) or Tuple([A, B, C]) or Tuple((A, B, C)). ' \
+                                   'Received fn=%s, args=%s' % (str(fn), str(args))
+            self._fn = fn
+        else:
+            self._fn = (fn, ) + args
+        for i, ele_fn in enumerate(self._fn):
+            assert callable(
+                ele_fn
+            ), 'Batchify functions must be callable! type(fn[%d]) = %s' % (
+                i, str(type(ele_fn)))
+
+    def __call__(self, data):
+        """
+        Batchifies data samples by applying each function on the corresponding 
+        data field, and each data field is produced by stacking the field data 
+        of samples.
+
+        Args:
+            data (list|tuple): The samples to batchfy. Each sample in list/tuple
+                should contain `N` fields.
+
+        Returns:
+            tuple: A tuple composed of results from all including batchifying 
+            functions.
+
+        Example:
+            .. code-block:: python
+                
+                from paddlenlp.data import Stack, Pad, Tuple
+                data = [
+                        [[1, 2, 3, 4], [1]],
+                        [[5, 6, 7], [0]],
+                        [[8, 9], [1]],
+                       ]
+                batchify_fn = Tuple(Pad(pad_val=0), Stack())
+                ids, label = batchify_fn(data)
+                '''
+                ids:
+                [[1, 2, 3, 4],
+                [5, 6, 7, 0],
+                [8, 9, 0, 0]]
+                label: [[1], [0], [1]]
+                '''
+        """
+
+        assert len(data[0]) == len(self._fn),\
+            'The number of attributes in each data sample should contain' \
+            ' {} elements'.format(len(self._fn))
+        ret = []
+        for i, ele_fn in enumerate(self._fn):
+            result = ele_fn([ele[i] for ele in data])
+            if isinstance(result, (tuple, list)):
+                ret.extend(result)
+            else:
+                ret.append(result)
+        return tuple(ret)
+
+
+class Dict(object):
+    """
+    Wraps multiple batchify functions together. The input functions will be 
+    applied to the corresponding input fields.
+    
+    Each sample should be a dict containing multiple fields. Each batchify 
+    function with key stored in `Dict` will be applied on the field which has 
+    the same key. 
+    
+    For example, when data sample is {'tokens': tokens, 'labels': labels}, you 
+    can wrap two batchify functions using 
+    `Dict({'tokens': DataBatchify, 'labels': LabelBatchify})` to batchify tokens 
+    and labels correspondingly.
+
+    Args:
+        fn (dict): The batchify functions to wrap. It is a dict, which values is 
+            callable functions.
+    """
+
+    def __init__(self, fn):
+        assert isinstance(fn, (dict)), 'Input pattern not understood. The input of Dict must be a dict with key of input column name and value of collate_fn ' \
+                                   'Received fn=%s' % (str(fn))
+
+        self._fn = fn
+
+        for col_name, ele_fn in self._fn.items():
+            assert callable(
+                ele_fn
+            ), 'Batchify functions must be callable! type(fn[%d]) = %s' % (
+                col_name, str(type(ele_fn)))
+
+    def __call__(self, data):
+        """
+        Batchifies data samples by applying each function on the corresponding 
+        data field, and each data field is produced by stacking the field data 
+        with the same key as batchify functions of all samples.
+
+        Args:
+            data (list[dict]|tuple[dict]): The samples to batchfy. Each sample 
+                in list/tuple is a dict with `N` key-values.
+                
+        Returns:
+            tuple: A tuple composed of results from all including batchifying 
+            functions.
+            
+        Example:
+            .. code-block:: python
+
+                from paddlenlp.data import Stack, Pad, Dict
+                data = [
+                        {'labels':[1], 'token_ids':[1, 2, 3, 4]},
+                        {'labels':[0], 'token_ids':[5, 6, 7]},
+                        {'labels':[1], 'token_ids':[8, 9]},
+                       ]
+                batchify_fn = Dict({'token_ids':Pad(pad_val=0), 'labels':Stack()})
+                ids, label = batchify_fn(data)
+                '''
+                ids:
+                [[1, 2, 3, 4],
+                [5, 6, 7, 0],
+                [8, 9, 0, 0]]
+                label: [[1], [0], [1]]
+                '''
+        """
+
+        ret = []
+        for col_name, ele_fn in self._fn.items():
+            result = ele_fn([ele[col_name] for ele in data])
+            if isinstance(result, (tuple, list)):
+                ret.extend(result)
+            else:
+                ret.append(result)
+        return tuple(ret)

test.py

+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+import numpy as np
+
+#a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype = int) 
+y_true = paddle.to_tensor([0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19], "int64")
+#pred = np.array([0.21167, -0.593262, -0.314209, -0.436035, 0.0211029, 0.0109406, -0.0210571, 0.588379, 0.215088, -0.277588, -0.0982666, -0.168579, -0.00405884, -0.290283, 0.105896, 0.0933228, 0.0889893, 0.0363159, 0.53418, 0.552734], dtype = float)
+#y_pred = paddle.to_tensor(pred).astype(np.float32)
+y_pred = paddle.to_tensor([0.21167 -0.593262 -0.314209 -0.436035 0.0211029 0.0109406 -0.0210571 0.588379 0.215088 -0.277588 -0.0982666 -0.168579 -0.00405884 -0.290283 0.105896 0.0933228 0.0889893 0.0363159 0.53418 0.552734], "float32") 
+
+#one_hot = F.one_hot(y_true,num_classes=3)
+#res=paddle.sum(paddle.exp(y_pred),axis=1)
+#res = paddle.reshape(res, [-1, 1])
+#softmax = paddle.exp(y_pred)/res
+#logsoftmax = paddle.log(softmax)
+
+#nllloss = -paddle.sum(one_hot.squeeze(1)*logsoftmax)/y_true.shape[0]
+#print(nllloss)
+
+#logsoftmax = F.log_softmax(y_pred)
+
+#nllloss = F.nll_loss(logsoftmax, y_true, ignore_index=-1)
+#print(nllloss)
+loss_fct = paddle.nn.loss.CrossEntropyLoss(ignore_index=-1) 
+nllloss = loss_fct(input=y_pred,label=y_true)
+print(nllloss)

test_fmha_accuracy.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import collections
+import itertools
+import os
+import sys
+import random
+import time
+import h5py
+import json
+from functools import partial
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import distutils.util
+
+import torch
+sys.path.append(
+    "/limin29/mlperf-1.1-repo/submission_training_1.1/NVIDIA/benchmarks/bert/implementations/pytorch"
+)
+from modeling import BertEncoder
+
+import paddle
+import paddle.nn as nn
+import paddle.distributed.fleet as fleet
+from paddle.io import DataLoader, Dataset
+from paddle.static import Program, program_guard
+
+from models.utils.tools import TimeCostAverage
+from mlperf_logging.mllog import constants
+from bert_utils.mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+import bert_utils.utility as utility
+
+from models.modeling import NewBertEncoder
+from models.modeling import BertConfig
+
+iters = 1
+x_type = np.float16
+attention_mask_type = np.int32
+
+batch_size = 56
+max_seq_len = 512
+embed_dim = 1024
+
+attention_mask = np.random.randint(
+    0, 5, size=[batch_size, max_seq_len]).astype(attention_mask_type)
+pad_input = np.random.rand(batch_size, max_seq_len,
+                           embed_dim).astype(x_type) * 0.25
+dout_np = np.random.random(
+    (batch_size, max_seq_len, embed_dim)).astype(np.float16) * 0.25
+
+# adjust nonzero value to 1.
+# todo: set maxseqlen_in_batch to 512, otherwise will report error.
+for i, val in enumerate(attention_mask):
+    if i == 2:
+        for j, ele in enumerate(val):
+            val[j] = 1
+    for j, ele in enumerate(val):
+        if ele != 0:
+            val[j] = 1
+attention_mask = np.ones((batch_size, max_seq_len)).astype(attention_mask_type)
+
+bert_config_path = '/zengjinle/dataset/bert_data/phase1/bert_config.json'
+config = BertConfig.from_json_file(bert_config_path)
+
+config.num_hidden_layers = 1
+config.hidden_size = embed_dim
+config.max_seq_length = 512
+
+config.num_attention_heads = 16
+config.intermediate_size = 4096
+config.hidden_act = 'gelu'
+config.attention_probs_dropout_prob = 0
+config.hidden_dropout_prob = 0
+
+config.unpad_fmha = True
+config.unpad = True
+config.pad_fmha = False
+config.pad = False
+
+config.fuse_mask = False
+config.enable_stream = False
+config.fused_gelu_bias = False
+config.fused_mha = False
+
+# if true, the linear weight parameter is [out_feature, in_feature]
+# fused qkv + bias
+config.fused_bias_mha = True
+
+# fused linear + bias
+config.fused_bias_fc = True
+if config.fused_bias_fc:
+    config.weight_transpose = True
+else:
+    config.weight_transpose = False
+
+# fused_dropout_reisudal_ln
+config.fused_dropout_add = False
+config.fused_dropout_add_ln = True
+'''
+Wq = (np.ones((config.hidden_size, config.hidden_size)).astype(x_type))*0.001
+Wk = (np.ones((config.hidden_size, config.hidden_size)).astype(x_type))*0.001
+Wv = (np.ones((config.hidden_size, config.hidden_size)).astype(x_type))*0.001
+Linear_weight = (np.ones((config.hidden_size, config.hidden_size)).astype(x_type))*0.001
+Bq = np.zeros((config.hidden_size)).astype(x_type)
+Bk = np.zeros((config.hidden_size)).astype(x_type)
+Bv = np.zeros((config.hidden_size)).astype(x_type)
+Linear_bias = np.zeros((config.hidden_size)).astype(x_type)
+'''
+# fmha params, q, k, v and out_liear
+Wq = np.random.rand(config.hidden_size,
+                    config.hidden_size).astype(x_type) * 0.25
+Wk = np.random.rand(config.hidden_size,
+                    config.hidden_size).astype(x_type) * 0.25
+Wv = np.random.rand(config.hidden_size,
+                    config.hidden_size).astype(x_type) * 0.25
+Linear_weight = np.random.rand(config.hidden_size,
+                               config.hidden_size).astype(x_type) * 0.25
+Bq = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+Bk = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+Bv = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+#Bq = np.zeros((config.hidden_size)).astype(x_type)
+#Bk = np.zeros((config.hidden_size)).astype(x_type)
+#Bv = np.zeros((config.hidden_size)).astype(x_type)
+Linear_bias = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+#Linear_bias = np.zeros((config.hidden_size)).astype(x_type)
+
+# ffn: linear_2 params
+linear_2_weight = np.random.rand(config.hidden_size,
+                                 config.intermediate_size).astype(x_type) * 0.25
+linear_2_bias = np.random.rand(config.intermediate_size).astype(x_type) * 0.25
+#linear_2_bias = np.zeros((config.intermediate_size)).astype(x_type)
+
+# ffn: linear_3 params
+linear_3_weight = np.random.rand(config.intermediate_size,
+                                 config.hidden_size).astype(x_type) * 0.25
+linear_3_bias = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+
+# layer_norm params
+norm_1_weight = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+norm_1_bias = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+norm_2_weight = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+norm_2_bias = np.random.rand(config.hidden_size).astype(x_type) * 0.25
+
+print("weight_transpose:", config.weight_transpose)
+
+
+def run_dynamic():
+    paddle.disable_static(place=paddle.CUDAPlace(0))
+    paddle.set_default_dtype(x_type)
+
+    encoder = NewBertEncoder(config)
+
+    attention = encoder.layers[0].attention
+    intermediate = encoder.layers[0].intermediate
+    last_output = encoder.layers[0].output
+
+    self_output = attention.output
+    fmha = attention.fmha
+    Wqkv = np.concatenate((Wq, Wk, Wv))
+    Bqkv = np.concatenate((Bq, Bk, Bv))
+    fmha.Wqkv.set_value(Wqkv)
+    fmha.Bqkv.set_value(Bqkv)
+
+    self_output.dense.bias.set_value(Linear_bias)
+    if config.fused_dropout_add_ln:
+        self_output.fused_dropout_add_ln.bias.set_value(norm_1_bias)
+    self_output.layer_norm.bias.set_value(norm_1_bias)
+
+    intermediate.dense.bias.set_value(linear_2_bias)
+
+    last_output.dense.bias.set_value(linear_3_bias)
+    last_output.layer_norm.bias.set_value(norm_2_bias)
+
+    if config.fused_dropout_add_ln:
+        self_output.fused_dropout_add_ln.weight.set_value(norm_1_weight)
+    self_output.layer_norm.weight.set_value(norm_1_weight)
+
+    last_output.layer_norm.weight.set_value(norm_2_weight)
+
+    if not config.weight_transpose:
+        self_output.dense.weight.set_value(Linear_weight)
+        intermediate.dense.weight.set_value(linear_2_weight)
+        last_output.dense.weight.set_value(linear_3_weight)
+    else:
+        self_output.dense.weight.set_value(Linear_weight.T)
+        intermediate.dense.weight.set_value(linear_2_weight.T)
+        last_output.dense.weight.set_value(linear_3_weight.T)
+
+    print("new encoder: ", encoder)
+    pad_input_tensor = paddle.to_tensor(pad_input, stop_gradient=False)
+    attention_mask_tensor = paddle.to_tensor(attention_mask)
+    encoder_out = encoder(pad_input_tensor, attention_mask_tensor)
+    '''
+    dout = paddle.to_tensor(dout_np)
+    paddle.autograd.backward([encoder_out[-1]], [dout], retain_graph=True)
+    '''
+    return encoder_out[-1], pad_input_tensor.grad, encoder
+
+
+def run_dynamic_pytorch_ref():
+
+    encoder = BertEncoder(config)
+    print("pytorch encoder:", encoder)
+
+    attention = encoder.layer[0].attention
+    intermediate = encoder.layer[0].intermediate
+    last_output = encoder.layer[0].output
+
+    self_output = attention.output
+    fmha = attention.self
+
+    Wqkv = np.concatenate((Wq, Wk, Wv))
+    Bqkv = np.concatenate((Bq, Bk, Bv))
+    fmha.Wqkv = torch.from_numpy(Wqkv).cuda()
+    fmha.Bqkv = torch.from_numpy(Bqkv).cuda()
+
+    self_output.dense.weight.data = torch.from_numpy(Linear_weight.T).cuda()
+    self_output.dense.bias.data = torch.from_numpy(Linear_bias).cuda()
+    self_output.LayerNorm.weight.data = torch.from_numpy(norm_1_weight).cuda()
+    self_output.LayerNorm.bias.data = torch.from_numpy(norm_1_bias).cuda()
+
+    intermediate.dense.weight.data = torch.from_numpy(linear_2_weight.T).cuda()
+    intermediate.dense.bias.data = torch.from_numpy(linear_2_bias).cuda()
+
+    last_output.dense.weight.data = torch.from_numpy(linear_3_weight.T).cuda()
+    last_output.dense.bias.data = torch.from_numpy(linear_3_bias).cuda()
+    last_output.LayerNorm.weight.data = torch.from_numpy(norm_2_weight).cuda()
+    last_output.LayerNorm.bias.data = torch.from_numpy(norm_2_bias).cuda()
+
+    print("pytorch encoder: ", encoder)
+    pad_input_tensor = torch.from_numpy(pad_input).cuda()
+    attention_mask_tensor = torch.from_numpy(attention_mask).cuda()
+    encoder_out = encoder(
+        pad_input_tensor,
+        attention_mask_tensor,
+        output_all_encoded_layers=False)
+    '''
+    dout = torch.from_numpy(dout_np).cuda()
+    torch.autograd.backward(
+            [encoder_out[-1]], [dout], retain_graph=True)
+    '''
+
+    # return encoder_out[-1], pad_input_tensor.grad
+    return encoder_out[-1]
+
+
+def run_dynamic_paddle_ref(fused_encoder):
+    paddle.disable_static(place=paddle.CUDAPlace(0))
+    paddle.set_default_dtype(x_type)
+    ref_encoder_layer = nn.TransformerEncoderLayer(
+        config.hidden_size,
+        config.num_attention_heads,
+        config.intermediate_size,
+        dropout=config.hidden_dropout_prob,
+        activation=config.hidden_act,
+        attn_dropout=config.attention_probs_dropout_prob,
+        act_dropout=0)
+    ref_encoder = nn.TransformerEncoder(ref_encoder_layer,
+                                        config.num_hidden_layers)
+    print("ref_encoder: ", ref_encoder)
+
+    # generate src_mask, is it right? 
+    input_mask = (1 - np.reshape(
+        attention_mask.astype(np.float32),
+        [attention_mask.shape[0], 1, 1, attention_mask.shape[1]])) * -1e9
+
+    # set params
+    attention = fused_encoder.layer[0].attention
+    fmha = attention.fmha
+    self_output = attention.output
+    intermediate = fused_encoder.layer[0].intermediate
+    last_output = fused_encoder.layer[0].output
+    layer_norm_2 = last_output.layer_norm
+
+    ref_self_attn = ref_encoder.layers[0].self_attn
+    ref_norm_1 = ref_encoder.layers[0].norm1
+    ref_norm_2 = ref_encoder.layers[0].norm2
+    ref_linear_1 = ref_encoder.layers[0].linear1
+    ref_linear_2 = ref_encoder.layers[0].linear2
+
+    ref_self_attn.q_proj.weight.set_value(Wq)
+    ref_self_attn.k_proj.weight.set_value(Wk)
+    ref_self_attn.v_proj.weight.set_value(Wv)
+    ref_self_attn.q_proj.bias.set_value(Bq)
+    ref_self_attn.k_proj.bias.set_value(Bk)
+    ref_self_attn.v_proj.bias.set_value(Bv)
+
+    ref_self_attn.out_proj.weight.set_value(Linear_weight)
+    ref_self_attn.out_proj.bias.set_value(Linear_bias)
+    '''
+    ref_self_attn.q_proj.weight.set_value(fmha.Wqkv[0:config.hidden_size, :])
+    ref_self_attn.k_proj.weight.set_value(fmha.Wqkv[config.hidden_size:2*config.hidden_size, :])
+    ref_self_attn.v_proj.weight.set_value(fmha.Wqkv[2*config.hidden_size:3*config.hidden_size, :])
+    ref_self_attn.q_proj.bias.set_value(fmha.Bqkv[0:config.hidden_size])
+    ref_self_attn.k_proj.bias.set_value(fmha.Bqkv[config.hidden_size:2*config.hidden_size])
+    ref_self_attn.v_proj.bias.set_value(fmha.Bqkv[2*config.hidden_size:3*config.hidden_size])
+    
+    ref_self_attn.out_proj.weight.set_value(self_output.dense.weight)
+    ref_self_attn.out_proj.bias.set_value(self_output.dense.bias)
+    '''
+    '''
+    weight_np = self_output.layer_norm.weight.numpy().astype(np.float16)
+    ref_norm_1.weight.set_value(weight_np)
+    weight_np = self_output.layer_norm.bias.numpy().astype(np.float16)
+    ref_norm_1.bias.set_value(weight_np)
+    #ref_norm_1.weight.set_value(self_output.layer_norm.weight)
+    #ref_norm_1.bias.set_value(self_output.layer_norm.bias)
+    
+    weight_np = layer_norm_2.weight.numpy().astype(np.float16)
+    ref_norm_2.weight.set_value(weight_np)
+    weight_np = layer_norm_2.bias.numpy().astype(np.float16)
+    ref_norm_2.bias.set_value(weight_np)
+    #ref_norm_2.weight.set_value(layer_norm_2.weight)
+    #ref_norm_2.bias.set_value(layer_norm_2.bias)
+    
+    ref_linear_1.weight.set_value(intermediate.dense.weight)
+    ref_linear_1.bias.set_value(intermediate.dense.bias)
+    
+    ref_linear_2.weight.set_value(last_output.dense.weight)
+    ref_linear_2.bias.set_value(last_output.dense.bias)
+    '''
+    ref_norm_1.weight.set_value(norm_1_weight)
+    ref_norm_1.bias.set_value(norm_1_weight)
+    ref_norm_2.weight.set_value(norm_2_weight)
+    ref_norm_2.bias.set_value(norm_2_bias)
+
+    ref_linear_1.weight.set_value(linear_2_weight)
+    ref_linear_1.bias.set_value(linear_2_bias)
+
+    ref_linear_2.weight.set_value(linear_3_weight)
+    ref_linear_2.bias.set_value(linear_3_bias)
+
+    pad_input_tensor = paddle.to_tensor(pad_input, stop_gradient=False)
+    input_mask_tensor = paddle.to_tensor(input_mask)
+    output = pad_input_tensor
+    encoder_outputs = []
+    for mod in ref_encoder.layers:
+        output = mod(output, src_mask=input_mask_tensor)
+        encoder_outputs.append(output)
+
+    return encoder_outputs[-1]
+
+
+def run_static():
+
+    paddle.set_default_dtype(x_type)
+    encoder = NewBertEncoder(config)
+
+    x = paddle.static.data(
+        name='X', shape=[batch_size, max_seq_len, embed_dim], dtype=x_type)
+    attn_mask = paddle.static.data(
+        name='SrcMask',
+        shape=[
+            batch_size,
+            max_seq_len,
+        ],
+        dtype=attention_mask_type)
+
+    encoder_out = encoder(x, attn_mask)
+
+    place = paddle.CUDAPlace(0)
+    exe = paddle.static.Executor(place)
+    exe.run(paddle.static.default_startup_program())
+    out = exe.run(paddle.static.default_main_program(),
+                  feed={"X": pad_input,
+                        "SrcMask": attention_mask},
+                  fetch_list=[encoder_out])
+    return out
+
+
+def test_dynamic(atol=1e-2):
+    for i in range(iters):
+        out, input_grad, fused_encoder = run_dynamic()
+        ref_out = run_dynamic_pytorch_ref()
+        # compared with pytorch 
+        np.testing.assert_allclose(
+            ref_out.cpu().detach().numpy(), out.numpy(), rtol=1e-5, atol=atol)
+        res = np.array_equal(ref_out.cpu().detach().numpy(), out.numpy())
+        if (res):
+            print("array equal")
+        else:
+            print("array not equal!")
+        #np.testing.assert_allclose(
+        #    ref_input_grad.cpu().detach().numpy(), input_grad.numpy(), rtol=1e-5, atol=atol)
+
+        # compared with paddle
+        #ref_paddle_out = run_dynamic_paddle_ref(fused_encoder)
+        #print("ref_paddle_out: ", ref_paddle_out)
+        #np.testing.assert_allclose(
+        #    ref_paddle_out.numpy(), out.numpy(), rtol=1e-5, atol=atol)
+        print("out = ", out)
+
+
+def test_static():
+    paddle.enable_static()
+    with paddle.static.program_guard(Program()):
+        out = run_static()
+    # print("encoder out: ", out[-1])
+
+
+test_dynamic(atol=1e-2)
+# test_static()

test_fused_dense_layer.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import collections
+import itertools
+import os
+import random
+import time
+import h5py
+import json
+from functools import partial
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import distutils.util
+
+import paddle
+import paddle.nn as nn
+import paddle.distributed.fleet as fleet
+from paddle.io import DataLoader, Dataset
+from paddle.static import Program, program_guard
+
+from models.utils.tools import TimeCostAverage
+from mlperf_logging.mllog import constants
+from bert_utils.mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+import bert_utils.utility as utility
+
+from models.modeling import FusedDense
+from models.modeling import BertConfig
+
+x_type = np.float16
+
+batch_size = 8
+max_seq_len = 512
+embed_dim = 1024
+
+pad_input = np.random.rand(batch_size, max_seq_len, embed_dim).astype(x_type)
+dout = np.random.random((batch_size, max_seq_len, embed_dim)).astype(x_type)
+
+bert_config_path = '/zengjinle/dataset/bert_data/phase1/bert_config.json'
+config = BertConfig.from_json_file(bert_config_path)
+
+config.num_hidden_layer = 2
+config.hidden_size = embed_dim
+
+config.attention_probs_dropout_prob = 0
+config.hidden_dropout_prob = 0
+
+config.fused_dense_weight_transpose = True
+
+
+def run_dynamic():
+    paddle.set_default_dtype(x_type)
+    fused_dense = FusedDense(config.hidden_size, config.hidden_size,
+                             config.fused_dense_weight_transpose)
+    pad_input_tensor = paddle.to_tensor(pad_input, stop_gradient=False)
+    out = fused_dense(pad_input_tensor)
+
+    dout_tensor = paddle.to_tensor(dout)
+    paddle.autograd.backward([out], [dout_tensor], retain_graph=True)
+    return out, pad_input_tensor.grad, fused_dense
+
+
+def run_dynamic_ref(fused_dense):
+    paddle.set_default_dtype(x_type)
+    dense = nn.Linear(config.hidden_size, config.hidden_size)
+
+    # set values for ln.
+    if config.fused_dense_weight_transpose:
+        weight_np = fused_dense.weight.numpy()
+        weight_t_np = weight_np.transpose(1, 0)
+        dense.weight.set_value(weight_t_np)
+    else:
+        dense.weight.set_value(fused_dense.weight)
+    dense.bias.set_value(fused_dense.bias)
+
+    pad_input_tensor = paddle.to_tensor(pad_input, stop_gradient=False)
+    out = dense(pad_input_tensor)
+
+    dout_tensor = paddle.to_tensor(dout)
+    paddle.autograd.backward([out], [dout_tensor], retain_graph=True)
+    return out, pad_input_tensor.grad
+
+
+def run_static():
+    paddle.set_default_dtype(x_type)
+    fused_dense = FusedDense(config.hidden_size, config.hidden_size)
+
+    x = paddle.static.data(
+        name='X', shape=[batch_size, max_seq_len, embed_dim], dtype=x_type)
+
+    fused_out = fused_dense(x)
+
+    place = paddle.CUDAPlace(0)
+    exe = paddle.static.Executor(place)
+    exe.run(paddle.static.default_startup_program())
+    out = exe.run(paddle.static.default_main_program(),
+                  feed={"X": pad_input},
+                  fetch_list=[fused_out])
+    return out
+
+
+def test_dynamic(atol=1e-2):
+    paddle.disable_static(place=paddle.CUDAPlace(0))
+    out, input_grad, fused_dense = run_dynamic()
+    ref_out, ref_input_grad = run_dynamic_ref(fused_dense)
+
+    np.testing.assert_allclose(
+        ref_out.numpy(), out.numpy(), rtol=1e-5, atol=atol)
+    np.testing.assert_allclose(
+        ref_input_grad.numpy(), input_grad.numpy(), rtol=1e-5, atol=atol)
+    print("Dynamic Test success!")
+
+
+def test_static():
+    paddle.enable_static()
+    with paddle.static.program_guard(Program()):
+        out = run_static()
+    # print("static out: ", out)
+
+
+test_dynamic()
+test_static()

test_fused_dropout_layer.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import collections
+import itertools
+import os
+import random
+import time
+import h5py
+import json
+from functools import partial
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import distutils.util
+
+import paddle
+import paddle.nn as nn
+import paddle.distributed.fleet as fleet
+from paddle.io import DataLoader, Dataset
+from paddle.static import Program, program_guard
+
+from models.utils.tools import TimeCostAverage
+from mlperf_logging.mllog import constants
+from bert_utils.mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+import bert_utils.utility as utility
+
+from models.modeling import FusedDropoutResidualLn
+from models.modeling import BertConfig
+
+x_type = np.float16
+
+batch_size = 8
+max_seq_len = 512
+embed_dim = 1024
+
+pad_input = np.random.rand(batch_size, max_seq_len, embed_dim).astype(x_type)
+residual_input = np.random.rand(batch_size, max_seq_len,
+                                embed_dim).astype(x_type)
+dout = np.random.random((batch_size, max_seq_len, embed_dim)).astype(x_type)
+
+bert_config_path = '/zengjinle/dataset/bert_data/phase1/bert_config.json'
+config = BertConfig.from_json_file(bert_config_path)
+
+config.num_hidden_layer = 2
+config.hidden_size = embed_dim
+
+config.attention_probs_dropout_prob = 0
+config.hidden_dropout_prob = 0
+
+config.unpad_fmha = True
+
+
+def run_dynamic():
+    paddle.set_default_dtype(x_type)
+    fused_dropout = FusedDropoutResidualLn(
+        config, config.hidden_size, epsilon=1e-12)
+    pad_input_tensor = paddle.to_tensor(pad_input, stop_gradient=False)
+    residual_tensor = paddle.to_tensor(residual_input, stop_gradient=False)
+    out = fused_dropout(pad_input_tensor, residual_tensor)
+
+    dout_tensor = paddle.to_tensor(dout)
+    paddle.autograd.backward([out], [dout_tensor], retain_graph=True)
+    return out, pad_input_tensor.grad, residual_tensor.grad, fused_dropout
+
+
+def run_dynamic_ref(fused_dropout):
+    paddle.set_default_dtype(x_type)
+    layer_norm = nn.LayerNorm(config.hidden_size, epsilon=1e-12)
+    dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    # set values for ln.
+    layer_norm.weight.set_value(fused_dropout.weight)
+    layer_norm.bias.set_value(fused_dropout.bias)
+
+    pad_input_tensor = paddle.to_tensor(pad_input, stop_gradient=False)
+    residual_tensor = paddle.to_tensor(residual_input, stop_gradient=False)
+    hidden_states = dropout(pad_input_tensor)
+    hidden_states = hidden_states + residual_tensor
+    out = layer_norm(hidden_states)
+
+    dout_tensor = paddle.to_tensor(dout)
+    paddle.autograd.backward([out], [dout_tensor], retain_graph=True)
+    return out, pad_input_tensor.grad, residual_tensor.grad
+
+
+def run_static():
+    paddle.set_default_dtype(x_type)
+    fused_dropout = FusedDropoutResidualLn(
+        config, config.hidden_size, epsilon=1e-12)
+
+    x = paddle.static.data(
+        name='X', shape=[batch_size, max_seq_len, embed_dim], dtype=x_type)
+
+    residual = paddle.static.data(
+        name='Residual',
+        shape=[batch_size, max_seq_len, embed_dim],
+        dtype=x_type)
+
+    fused_out = fused_dropout(x, residual)
+
+    place = paddle.CUDAPlace(0)
+    exe = paddle.static.Executor(place)
+    exe.run(paddle.static.default_startup_program())
+    out = exe.run(paddle.static.default_main_program(),
+                  feed={"X": pad_input,
+                        "Residual": residual_input},
+                  fetch_list=[fused_out])
+    return out
+
+
+def test_dynamic(atol=1e-2):
+    paddle.disable_static(place=paddle.CUDAPlace(0))
+    out, input_grad, residual_grad, fused_dropout = run_dynamic()
+    ref_out, ref_input_grad, ref_residual_grad = run_dynamic_ref(fused_dropout)
+
+    np.testing.assert_allclose(
+        ref_out.numpy(), out.numpy(), rtol=1e-5, atol=atol)
+    np.testing.assert_allclose(
+        ref_input_grad.numpy(), input_grad.numpy(), rtol=1e-5, atol=atol)
+    np.testing.assert_allclose(
+        ref_residual_grad.numpy(), residual_grad.numpy(), rtol=1e-5, atol=atol)
+    print("Dynamic Test success!")
+
+
+def test_static():
+    paddle.enable_static()
+    with paddle.static.program_guard(Program()):
+        out = run_static()
+    # print("static out: ", out)
+
+
+test_dynamic()
+test_static()

test_new_bert_encoder_layer.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import collections
+import itertools
+import os
+import random
+import time
+import h5py
+import json
+from functools import partial
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import distutils.util
+
+import paddle
+import paddle.distributed.fleet as fleet
+from paddle.io import DataLoader, Dataset
+from paddle.static import Program, program_guard
+
+from models.utils.tools import TimeCostAverage
+from mlperf_logging.mllog import constants
+from bert_utils.mlperf_logging_helper import paddle_bert_print_start, paddle_bert_print_end, paddle_bert_print_event
+import bert_utils.utility as utility
+
+from models.modeling import NewBertEncoder
+from models.modeling import BertConfig
+
+x_type = np.float16
+attention_mask_type = np.int32
+
+batch_size = 8
+max_seq_len = 512
+embed_dim = 1024
+
+attention_mask = np.random.randint(
+    0, 5, size=[batch_size, max_seq_len]).astype(attention_mask_type)
+pad_input = np.random.rand(batch_size, max_seq_len, embed_dim).astype(x_type)
+
+# adjust nonzero value to 1.
+# todo: set maxseqlen_in_batch to 512, otherwise will report error.
+for i, val in enumerate(attention_mask):
+    if i == 2:
+        for j, ele in enumerate(val):
+            val[j] = 1
+    for j, ele in enumerate(val):
+        if ele != 0:
+            val[j] = 1
+
+# todo
+bert_config_path = '/zengjinle/dataset/bert_data/phase1/bert_config.json'
+config = BertConfig.from_json_file(bert_config_path)
+
+config.num_hidden_layer = 2
+config.hidden_size = embed_dim
+
+#config.num_attention_heads=16 
+#config.intermediate_size=4096
+#config.hidden_act='gelu'
+config.attention_probs_dropout_prob = 0
+config.hidden_dropout_prob = 0
+
+config.unpad_fmha = True
+
+
+def run_dynamic():
+    paddle.set_default_dtype(x_type)
+
+    encoder = NewBertEncoder(config)
+    pad_input_tensor = paddle.to_tensor(pad_input, stop_gradient=False)
+    attention_mask_tensor = paddle.to_tensor(attention_mask)
+    print("pad_input:", pad_input_tensor)
+    print("attention_mask:", attention_mask_tensor)
+    encoder_out = encoder(pad_input_tensor, attention_mask_tensor)
+    print("encoder_out: ", encoder_out[-1])
+    print("encoder_out.shape: ", encoder_out[-1].shape)
+
+    #
+    dout = np.random.random(
+        (batch_size, max_seq_len, embed_dim)).astype(np.float16)
+    dout = paddle.to_tensor(dout)
+    dout = paddle.cast(dout, 'float16')
+    #dout = paddle.ones(shape=[5, 3], dtype='float16')
+    paddle.autograd.backward([encoder_out[-1]], [dout], retain_graph=True)
+    print("pad_input.grad: ", pad_input_tensor.grad)
+
+
+def run_static():
+
+    paddle.set_default_dtype(x_type)
+    encoder = NewBertEncoder(config)
+
+    x = paddle.static.data(
+        name='X', shape=[batch_size, max_seq_len, embed_dim], dtype=x_type)
+    attn_mask = paddle.static.data(
+        name='SrcMask',
+        shape=[
+            batch_size,
+            max_seq_len,
+        ],
+        dtype=attention_mask_type)
+
+    encoder_out = encoder(x, attn_mask)
+
+    place = paddle.CUDAPlace(0)
+    exe = paddle.static.Executor(place)
+    exe.run(paddle.static.default_startup_program())
+    out = exe.run(paddle.static.default_main_program(),
+                  feed={"X": pad_input,
+                        "SrcMask": attention_mask},
+                  fetch_list=[encoder_out])
+    return out
+
+
+def test_dynamic():
+    paddle.disable_static(place=paddle.CUDAPlace(0))
+    run_dynamic()
+
+
+def test_static():
+    paddle.enable_static()
+    with paddle.static.program_guard(Program()):
+        out = run_static()
+    print("encoder out: ", out[-1])
+
+
+#test_dynamic()
+test_static()

test_padding.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+from bert_padding import generate_mask
+
+# test for generate_mask
+
+# [3, 4]
+attention_mask = [[0, 1, 1, 0], [0, 1, 1, 1], [1, 1, 1, 1]]
+attention_mask_tensor = paddle.to_tensor(attention_mask)
+indices, attention_mask, seqlen, ntokens, cu_seqlens, seqlen, maxseqlen = generate_mask(
+    attention_mask_tensor, True)
+
+print("indices", indices)
+print("attention_mask", attention_mask)
+print("seqlen", seqlen)
+print("ntokens", ntokens)
+print("cu_seqlens", cu_seqlens)
+print("maxseqlen", maxseqlen)
+
+# test for unpad(gather) and pad(scatter) 
+
+embed_dim = 3
+input = [[0, 0, 0], [1, 1, 1], [0, 0, 0], [3, 3, 3], [4, 4, 4]]
+index = [1, 3, 4]
+input = paddle.to_tensor(input)
+index = paddle.to_tensor(index)
+out = paddle.gather(input, index)
+
+print(input)
+print(index)
+print("gather out", out)
+
+batch_size = 1
+max_seq_len = 5
+
+scatter_shape = [batch_size * max_seq_len, embed_dim]
+zero_tensor = paddle.zeros(scatter_shape, dtype='int64')
+scatter_out = paddle.scatter(zero_tensor, index, out)
+print("scatter_out: ", scatter_out)
+
+# error: 
+#out_1 = paddle.scatter_nd(index, out, scatter_shape)
+#print("scatter_nd out: ", out_1)

text.xml

+<graphs version="1">
+  <graph id="0" pattern="4" crossnic="0" nchannels="2" speedintra="12" speedinter="12" typeintra="SYS" typeinter="SYS" samechannels="1">
+    <channel>
+      <net dev="0"/>
+      <gpu dev="0"/>
+      <gpu dev="1"/>
+      <gpu dev="2"/>
+      <gpu dev="3"/>
+      <net dev="3"/>
+    </channel>
+    <channel>
+      <net dev="1"/>
+      <gpu dev="0"/>
+      <gpu dev="1"/>
+      <gpu dev="2"/>
+      <gpu dev="3"/>
+      <net dev="2"/>
+    </channel>
+  </graph>
+  <graph id="1" pattern="3" crossnic="0" nchannels="2" speedintra="6" speedinter="6" typeintra="SYS" typeinter="SYS" samechannels="1">
+    <channel>
+      <net dev="0"/>
+      <gpu dev="0"/>
+      <gpu dev="1"/>
+      <gpu dev="2"/>
+      <gpu dev="3"/>
+      <net dev="0"/>
+    </channel>
+    <channel>
+      <net dev="0"/>
+      <gpu dev="0"/>
+      <gpu dev="1"/>
+      <gpu dev="2"/>
+      <gpu dev="3"/>
+      <net dev="0"/>
+    </channel>
+  </graph>
+  <graph id="2" pattern="3" crossnic="0" nchannels="0" speedintra="0" speedinter="0" typeintra="NVL" typeinter="PIX" samechannels="1"/>
+</graphs>
+

unit_test/test_linear_op.py

+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+import paddle
+
+import torch.nn.functional as F
+
+bsz = 32
+hidden_size = 256
+x_type = np.float16
+
+hidden_states = np.random.rand(bsz, hidden_size).astype(x_type) * 0.25
+Wq = np.random.rand(hidden_size, hidden_size).astype(x_type) * 0.25
+Wk = np.random.rand(hidden_size, hidden_size).astype(x_type) * 0.25
+Wv = np.random.rand(hidden_size, hidden_size).astype(x_type) * 0.25
+Bq = np.random.rand(hidden_size).astype(x_type) * 0.25
+Bk = np.random.rand(hidden_size).astype(x_type) * 0.25
+Bv = np.random.rand(hidden_size).astype(x_type) * 0.25
+
+Wqkv = np.concatenate((Wq, Wk, Wv))
+Bqkv = np.concatenate((Bq, Bk, Bv))
+
+
+def run_paddle_linear():
+    paddle.disable_static(place=paddle.CUDAPlace(0))
+    paddle.set_default_dtype(x_type)
+    '''
+    out = paddle.matmul(
+            paddle.cast(paddle.to_tensor(hidden_states), 'float32'), 
+            paddle.cast(paddle.to_tensor(Wqkv), 'float32'), 
+            transpose_x=False, transpose_y=True)
+    out = out + paddle.cast(paddle.to_tensor(Bqkv), 'float32')
+    '''
+    out = paddle.matmul(
+        paddle.to_tensor(hidden_states),
+        paddle.to_tensor(Wqkv),
+        transpose_x=False,
+        transpose_y=True)
+    out = out + paddle.to_tensor(Bqkv)
+    return out
+
+
+def run_pytorch_linear():
+    out = F.linear(
+        torch.from_numpy(hidden_states).cuda(),
+        torch.from_numpy(Wqkv).cuda(), torch.from_numpy(Bqkv).cuda())
+    #out = F.linear(torch.from_numpy(hidden_states).cuda(), torch.from_numpy(Wqkv).cuda(), None)
+    #out = out + torch.from_numpy(Bqkv).cuda()
+    return out
+
+
+def run_numpy_linear():
+    out = np.matmul(hidden_states, Wqkv.transpose(1, 0))
+    out = out + Bqkv
+    return out
+
+
+paddle_out = run_paddle_linear()
+pytorch_out = run_pytorch_linear()
+np_out = run_numpy_linear()
+
+print("compare with pytorch:")
+np.testing.assert_allclose(
+    pytorch_out.cpu().detach().numpy(),
+    paddle_out.numpy(),
+    rtol=1e-5,
+    atol=1e-2)
+print("Success!")
+print("paddle compare with numpy:")
+np.testing.assert_allclose(np_out, paddle_out.numpy(), rtol=1e-5, atol=1e-2)
+print("Success!")
+print("pytorch compare with numpy:")
+np.testing.assert_allclose(
+    np_out, pytorch_out.cpu().detach().numpy(), rtol=1e-5, atol=1e-2)
+print("Success!")