Add training benchmark for gpt / gpt-3. (#1091)

Author: ZHUI

examples/language_model/gpt-3/dygraph/dataset.py

@@ -74,15 +74,30 @@ def construct_samples_and_shuffle_data(name, data_prefix, documents, sizes,
                 separate_last_epoch = (
                     last_epoch_num_samples < int(0.80 * num_samples_per_epoch))
             # Note. len(doc_idx) = num_epochs * len(doc)
+            start_time = time.time()
             doc_idx = _build_doc_idx(documents, num_epochs, np_rng,
                                      separate_last_epoch)
             np.save(doc_idx_filename, doc_idx, allow_pickle=True)
-
+            print(' > elasped time to build and save doc-idx mapping '
+                  '(seconds): {:4f}'.format(time.time() - start_time))
             # sample-idx. pos of each seq_len of data.
+            start_time = time.time()
             assert doc_idx.dtype == np.int32
-            sample_idx = _build_sample_idx(sizes, doc_idx, seq_length,
-                                           num_epochs, tokens_per_epoch)
+            assert sizes.dtype == np.int32
+
+            import data_tools.helpers as helpers
+
+            sample_idx = helpers.build_sample_idx(sizes, doc_idx, seq_length,
+                                                  num_epochs, tokens_per_epoch)
+            # sample_idx = _build_sample_idx(sizes, doc_idx, seq_length,
+            #                                num_epochs, tokens_per_epoch)
+
             np.save(sample_idx_filename, sample_idx, allow_pickle=True)
+            print(' > elasped time to build and save sample-idx mapping '
+                  '(seconds): {:4f}'.format(time.time() - start_time))
+
+            # shuffle-idx.
+            start_time = time.time()
 
             if separate_last_epoch:
                 num_samples_ = num_samples_from_epochs_minus_one
@@ -93,14 +108,25 @@ def construct_samples_and_shuffle_data(name, data_prefix, documents, sizes,
             shuffle_idx = _build_shuffle_idx(num_samples_,
                                              sample_idx.shape[0] - 1, np_rng)
             np.save(shuffle_idx_filename, shuffle_idx, allow_pickle=True)
+            print(' > elasped time to build and save shuffle-idx mapping'
+                  ' (seconds): {:4f}'.format(time.time() - start_time))
+
     else:
         while True:
             if (not os.path.isfile(doc_idx_filename)) or \
                (not os.path.isfile(sample_idx_filename)) or \
                (not os.path.isfile(shuffle_idx_filename)):
                 time.sleep(3)
             else:
-                break
+                try:
+                    np.load(
+                        shuffle_idx_filename, allow_pickle=True, mmap_mode='r')
+                    break
+                except Exception as e:
+                    print(
+                        "%s file is still writing or damaged, please wait a moment."
+                        % shuffle_idx_filename)
+                    time.sleep(3)
 
     # Restore random state
     np_rng.set_state(savedState)
@@ -241,19 +267,48 @@ def create_pretrained_dataset(args,
                               max_seq_len=1024,
                               places=None,
                               data_holders=None):
+    if local_rank == 0:
+        start_time = time.time()
+        print('> compiling dataset index builder ...')
+        from data_tools.dataset_utils import compile_helper
+        compile_helper()
+        print(
+            '>>> done with dataset index builder. Compilation time: {:.3f} '
+            'seconds'.format(time.time() - start_time),
+            flush=True)
+
     device_world_size = paddle.distributed.get_world_size()
     device_world_rank = paddle.distributed.get_rank()
 
     logger.info(
         "The distributed run, total device num:{}, distinct dataflow num:{}.".
         format(device_world_size, data_world_size))
 
-    process_datas = np.load(input_path, mmap_mode="r+", allow_pickle=True)
-    # All documment ids, extend as 1-D array.
-    sample_ids = process_datas["ids"]
-    # The len(sample_lens) num of docs
-    # The sum(sample_lens) should equal len(sample_ids)
-    sample_lens = process_datas["lens"]
+    assert len(input_path) == 1, "GPT only support one dataset for now."
+
+    input_prefix = input_path[0]
+
+    if os.path.isfile(input_prefix + "_ids.npz"):
+        logger.warning(
+            "You are using compatible dataset, please make new dataset as the readme!"
+        )
+        process_datas = np.load(
+            input_prefix + "_ids.npz", mmap_mode="r+", allow_pickle=True)
+        sample_ids = process_datas["ids"]
+        sample_lens = process_datas["lens"].astype("int32")
+    else:
+        for suffix in ["_ids.npy", "_idx.npz"]:
+            if not os.path.isfile(input_prefix + suffix):
+                raise ValueError("File Not found, %s" % (path + suffix))
+
+        sample_ids = np.load(
+            input_prefix + "_ids.npy", mmap_mode="r", allow_pickle=True)
+        # All documment ids, extend as 1-D array.
+
+        process_datas = np.load(input_prefix + "_idx.npz")
+        # The len(sample_lens) num of docs
+        # The sum(sample_lens) should equal len(sample_ids)
+        sample_lens = process_datas["lens"]
 
     splits = get_train_valid_test_split_(args.split, len(sample_lens))
     assert len(sample_lens) >= splits[
@@ -262,7 +317,7 @@ def create_pretrained_dataset(args,
 
     def build_dataset(index, name, num_samples):
         dataset = GPTDataset(
-            file_path=input_path,
+            file_path=input_prefix,
             build_data_file=local_rank == 0,
             name="gpt_" + name,
             max_seq_len=max_seq_len,

examples/language_model/gpt-3/dygraph/run_pretrain.py

@@ -17,6 +17,7 @@
 import os
 import random
 import time
+import sys
 
 import numpy as np
 import paddle
@@ -25,6 +26,10 @@
 from paddlenlp.transformers import GPTTokenizer, GPTChineseTokenizer
 from paddlenlp.utils.log import logger
 
+# to import data_tools
+filepath = os.path.abspath(os.path.dirname(__file__))
+sys.path.insert(0, os.path.join(filepath, "../../"))
+
 from dataset import create_pretrained_dataset
 from args import parse_args
 import lr
@@ -87,6 +92,30 @@ def run_evaluate(args,
     model.train()
 
 
+def get_train_data_file(args):
+    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 str(f).endswith(
+            "_idx.npz"))
+    ]
+    files = [x.replace("_idx.npz", "") for x in files]
+    if len(files) == 0:
+        logger.warning(
+            "Not found dataset with name of xxx_ids.npy and xxx_idx.npz! Try to found old compatible xxx_ids.npz file."
+        )
+    else:
+        return files
+
+    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 str(f).endswith(
+            "_ids.npz"))
+    ]
+
+    files = [x.replace("_ids.npz", "") for x in files]
+    return files
+
+
 def do_train(args):
     paddle.set_device(args.device)
     strategy = fleet.DistributedStrategy()
@@ -240,18 +269,13 @@ def do_train(args):
     global_step = 0
     tic_train = time.time()
     for epoch in range(args.num_train_epochs):
-        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 "npz_"
-                not in str(f))
-        ]
+        files = get_train_data_file(args)
         files.sort()
         num_files = len(files)
         for f_id in range(num_files):
             data_file = files[f_id]
             train_data_loader, valid_data_loader, test_data_loader = create_pretrained_dataset(
-                args,
-                data_file,
+                args, [data_file],
                 local_rank=local_rank,
                 data_world_size=data_world_size,
                 data_world_rank=data_world_rank,

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

@@ -36,7 +36,8 @@
 from visualdl import LogWriter
 
 # Used to load the data_tools path, should import before dataset
-sys.path.insert(0, "../../")
+filepath = os.path.abspath(os.path.dirname(__file__))
+sys.path.insert(0, os.path.join(filepath, "../../"))
 from dataset import create_pretrained_dataset
 from args import parse_args
 import lr
@@ -448,7 +449,7 @@ def do_train(args):
                 save_persistables(exe,
                                   os.path.join(output_dir, "static_vars"),
                                   main_program)
-                if global_step == args.save_steps:
+                if global_step <= args.save_steps:
                     model.init_config["init_args"][0].init_config.pop("topo",
                                                                       None)
                 model.save_pretrained(output_dir)

examples/language_model/gpt/args.py

@@ -211,7 +211,7 @@ def parse_args(MODEL_CLASSES):
     parser.add_argument(
         "--scale_loss",
         type=float,
-        default=128,
+        default=32768,
         help="The value of scale_loss for fp16. This is only used for AMP training."
     )
     parser.add_argument(
@@ -245,6 +245,13 @@ def parse_args(MODEL_CLASSES):
         default="cosine",
         choices=["cosine", "none"],
         help="Learning rate decay style.")
+    parser.add_argument(
+        '-p',
+        '--profiler_options',
+        type=str,
+        default=None,
+        help='The option of profiler, which should be in format \"key1=value1;key2=value2;key3=value3\".'
+    )
     args = parser.parse_args()
     args.test_iters = args.eval_iters * 10
 

examples/language_model/gpt/run_pretrain.py

@@ -29,6 +29,7 @@
 from dataset import create_pretrained_dataset
 from args import parse_args
 import lr
+from paddle.distributed import fleet
 
 MODEL_CLASSES = {
     "gpt": (GPTForPretraining, GPTTokenizer),
@@ -183,6 +184,9 @@ def do_train(args):
         grad_clip=clip,
         apply_decay_param_fun=lambda x: x in decay_params)
 
+    if args.use_amp:
+        scaler = paddle.amp.GradScaler(init_loss_scaling=args.scale_loss)
+
     if args.model_name_or_path not in pretrained_models_list:
         logger.info("Try to load checkpoint from %s " % args.model_name_or_path)
         opt_path = os.path.join(args.model_name_or_path, "model_state.pdopt")
@@ -218,9 +222,27 @@ def do_train(args):
                 tokens, loss_mask, attention_mask, position_ids, labels = batch
                 loss_mask.stop_gradient = True
                 attention_mask.stop_gradient = True
+                with paddle.amp.auto_cast(
+                        args.use_amp,
+                        custom_white_list=["layer_norm", "softmax", "gelu"],
+                        custom_black_list=[
+                            "reduce_sum", "c_softmax_with_cross_entropy",
+                            "c_embedding"
+                        ]):
+
+                    preds = model(tokens, position_ids, attention_mask)
+                    loss = criterion(preds, labels, loss_mask)
+
+                if args.use_amp:
+                    scaler.scale(loss).backward()
+                    scaler.minimize(optimizer, loss)
+                else:
+                    loss.backward()
+                    optimizer.step()
 
-                preds = model(tokens, position_ids, attention_mask)
-                loss = criterion(preds, labels, loss_mask)
+                if lr_scheduler is not None:
+                    lr_scheduler.step()
+                optimizer.clear_grad()
 
                 if global_step % args.logging_freq == 0:
                     speed = args.logging_freq / (time.time() - tic_train)
@@ -233,11 +255,6 @@ def do_train(args):
                                           optimizer.get_lr(), global_step)
 
                     tic_train = time.time()
-                loss.backward()
-                optimizer.step()
-                if lr_scheduler is not None:
-                    lr_scheduler.step()
-                optimizer.clear_grad()
 
                 if args.check_accuracy:
                     if global_step >= args.max_steps:

examples/language_model/gpt/run_pretrain_static.py

@@ -31,6 +31,7 @@
 from paddlenlp.transformers import GPTTokenizer, GPTChineseTokenizer
 from paddlenlp.ops import Topology, get_rng_state_tracker
 from paddlenlp.utils.log import logger
+from paddlenlp.utils import profiler
 import paddlenlp.ops as ops
 from visualdl import LogWriter
 
@@ -92,7 +93,7 @@ def dist_optimizer(args, topo):
                 'gelu',
             ],
             "custom_black_list": ['c_softmax_with_cross_entropy'],
-            "init_loss_scaling": 32768,
+            "init_loss_scaling": args.scale_loss,
             "use_dynamic_loss_scaling": True,
         }
     if args.use_sharding:
@@ -173,7 +174,7 @@ def run_evaluate(data_loader,
                 break
             average_loss = sum(all_loss) / len(all_loss)
             logger.info(
-                "%s step %d, epoch: %d, batch: %d, loss: %f, speed: %.0f tokens/s"
+                "%s step %d, epoch: %d, batch: %d, loss: %f, eval_ips: %.0f tokens/s"
                 % (task_name, global_step, epoch, eval_step, average_loss,
                    iter_steps * args.micro_batch_size * args.max_seq_len /
                    (time.time() - local_time)))
@@ -407,6 +408,7 @@ def do_train(args):
                           use_program_cache=True)
             # In the new 2.0 api, must call this function to change the learning_rate
             lr_scheduler.step()
+            profiler.add_profiler_step(args.profiler_options)
 
             if global_step % args.logging_freq == 0:
                 if topo.is_last:
@@ -446,7 +448,8 @@ def do_train(args):
                 save_persistables(exe,
                                   os.path.join(output_dir, "static_vars"),
                                   main_program)
-                if global_step == args.save_steps:
+
+                if global_step <= args.save_steps:
                     model.init_config["init_args"][0].init_config.pop("topo",
                                                                       None)
                 model.save_pretrained(output_dir)

examples/language_model/gpt/scripts/run.sh

@@ -18,4 +18,6 @@ python -u run_pretrain.py \
     --grad_clip 1.0\
     --logging_freq 1\
     --eval_freq 1000\
-    --device "gpu"
+    --device "gpu" \
+
+#    --use_amp true