[GPT-3] Add support for pure fp16 in dygraph mode (#1176)

Author: haohongxiang

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

@@ -22,10 +22,14 @@ def process_batch_size(args):
     if args.global_batch_size is None and args.local_batch_size is None:
         raise ValueError("global_batch_size or local_batch_size should be set.")
     elif args.global_batch_size is not None and args.local_batch_size is not None:
-        assert args.global_batch_size // args.local_batch_size == args.dp_degree, \
-            "global_batch_size[{}] should be divided by local_batch_size[{}] when dp_degree is [{}]"\
-                .format(args.global_batch_size, args.local_batch_size, args.dp_degree)
+        assert args.global_batch_size // args.local_batch_size == (args.dp_degree *
+            args.sharding_degree), "global_batch_size[{}] should be divided by local_batch_size[{}] "\
+            "when dp_degree is [{}] and sharding_degree is [{}]".format(args.global_batch_size,
+            args.local_batch_size, args.dp_degree, args.sharding_degree)
     elif args.global_batch_size is not None and args.local_batch_size is None:
+        assert args.global_batch_size % (args.dp_degree * args.sharding_degree) == 0, \
+            "global_batch_size[{}] should be divided by dp_degree[{}] times sharding_degree[{}]"\
+            .format(args.global_batch_size, args.dp_degree, args.sharding_degree)
         args.local_batch_size = args.global_batch_size // (args.dp_degree *
                                                            args.sharding_degree)
     else:
@@ -220,13 +224,13 @@ def parse_args(MODEL_CLASSES):
         const=False,
         help="Using the recompute to save the memory.")
 
-    # AMP config
+    # Pure FP16 config
     parser.add_argument(
-        "--use_amp",
+        "--use_pure_fp16",
         type=str2bool,
         nargs='?',
         const=False,
-        help="Enable mixed precision training.")
+        help="Enable pure fp16 precision training.")
 
     parser.add_argument(
         "--scale_loss",

examples/language_model/gpt-3/dygraph/run.sh

@@ -21,6 +21,6 @@ python -m paddle.distributed.launch --log_dir $log_dir --gpus "0,1,2,3,4,5,6,7"
     --mp_degree 2\
     --pp_degree 2\
     --sharding_degree 1\
-    --use_amp True\
+    --use_pure_fp16 True\
     --use_recompute False
 

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

@@ -126,8 +126,9 @@ def do_train(args):
         "sharding_degree": args.sharding_degree
     }
 
+    accumulate_steps = args.local_batch_size // args.micro_batch_size
     strategy.pipeline_configs = {
-        "accumulate_steps": args.local_batch_size // args.micro_batch_size,
+        "accumulate_steps": accumulate_steps,
         "micro_batch_size": args.micro_batch_size
     }
 
@@ -160,8 +161,8 @@ def do_train(args):
     # Define log writer
     log_writer_path = os.path.join(
         args.output_dir, "train_log",
-        "{}_globalbsz_{}_amp_{}_recompute_{}_card_{}".format(
-            args.model_name_or_path, args.global_batch_size, args.use_amp,
+        "{}_globalbsz_{}_pure_fp16_{}_recompute_{}_card_{}".format(
+            args.model_name_or_path, args.global_batch_size, args.use_pure_fp16,
             False, global_rank).lower())
 
     if os.path.exists(log_writer_path):
@@ -246,16 +247,25 @@ def do_train(args):
             parameters=model.parameters(),
             weight_decay=args.weight_decay,
             grad_clip=clip,
-            apply_decay_param_fun=lambda x: x in decay_params)
+            apply_decay_param_fun=lambda x: x in decay_params,
+            # TODO: remove 'multi_precision' in definition of optimizer
+            # and add it to 'paddle.amp.decorate'
+            multi_precision=args.use_pure_fp16)
+
+    if args.use_pure_fp16:
+        scaler = paddle.amp.GradScaler(init_loss_scaling=args.scale_loss)
+        scaler = fleet.distributed_scaler(scaler)
+        # level O2 means converting the network to FP16
+        model, optimizer = paddle.amp.decorate(
+            models=model,
+            optimizers=optimizer,
+            level='O2',
+            save_dtype='float32')
 
     if paddle.distributed.get_world_size() > 1:
         model = fleet.distributed_model(model)
         optimizer = fleet.distributed_optimizer(optimizer)
 
-    if args.use_amp:
-        scaler = paddle.amp.GradScaler(init_loss_scaling=args.scale_loss)
-        scaler = fleet.distributed_scaler(scaler)
-
     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")
@@ -294,23 +304,36 @@ def do_train(args):
                 position_ids.stop_gradient = True
 
                 if args.pp_degree == 1:
-                    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)
-                        loss = criterion(preds, labels, loss_mask)
-
-                    if args.use_amp:
-                        scaler.scale(loss).backward()
+                    # In ParallelMode of DataParallel, 'no_sync' can be used for improving
+                    # performance of model by gradient accumulation.
+                    loss = 0.0
+                    for i in range(accumulate_steps):
+                        start_index = i * args.micro_batch_size
+                        end_index = start_index + args.micro_batch_size
+                        with paddle.amp.auto_cast(
+                                args.use_pure_fp16,
+                                custom_black_list=[
+                                    "reduce_sum",
+                                    "c_softmax_with_cross_entropy",
+                                    "elementwise_div"
+                                ],
+                                level='O2'):
+                            preds = model(
+                                tokens[start_index:end_index, :],
+                                position_ids[start_index:end_index, :])
+                            loss_mbs = criterion(
+                                preds, labels[start_index:end_index, :],
+                                loss_mask[start_index:end_index, :])
+                        loss_mbs = loss_mbs / accumulate_steps
+                        if args.use_pure_fp16:
+                            scaler.scale(loss_mbs).backward()
+                        else:
+                            loss_mbs.backward()
+                        loss = loss + loss_mbs
+
+                    if args.use_pure_fp16:
                         scaler.minimize(optimizer, loss)
                     else:
-                        loss.backward()
                         optimizer.step()
 
                     if lr_scheduler is not None:
@@ -320,19 +343,17 @@ def do_train(args):
                 else:
                     data = [(tokens, position_ids), (labels, loss_mask)]
                     with paddle.amp.auto_cast(
-                            args.use_amp,
-                            custom_white_list=[
-                                "layer_norm", "softmax", "gelu"
-                            ],
+                            args.use_pure_fp16,
                             custom_black_list=[
                                 "reduce_sum", "c_softmax_with_cross_entropy",
-                                "c_embedding"
-                            ]):
+                                "elementwise_div"
+                            ],
+                            level='O2'):
                         loss = model.train_batch(
                             data,
                             optimizer=optimizer,
                             lr_scheduler=lr_scheduler,
-                            scaler=scaler if args.use_amp else None)
+                            scaler=scaler if args.use_pure_fp16 else None)
 
                 if global_step % args.logging_freq == 0:
                     avg_loss = loss.numpy()