Speed up for hybrid parallel (PaddlePaddle#1056)

Author: ForFishes

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

@@ -286,10 +286,10 @@ def build_dataset(index, name, num_samples):
             places=places,
             feed_list=data_holders,
             batch_sampler=batch_sampler,
-            num_workers=0,
+            num_workers=1,
             worker_init_fn=worker_init,
             # collate_fn=Tuple(Stack(), Stack(), Stack(), Stack(), Stack()),
-            collate_fn=Tuple(Stack(), Stack(), Stack()),
+            collate_fn=Tuple(Stack(), Stack(), Stack(), Stack()),
             return_list=False)
         return data_loader
 
@@ -349,12 +349,12 @@ def _construct_sample(self, tokens):
         # The pad and eos tokens do not contribute the loss
         loss_mask = np.ones(seq_length, dtype="float32")
         loss_mask[np.where(np.array(tokens) == self.eos_id)] = 0.0
-        # position_ids = np.arange(0, seq_length, dtype="int64")
+        position_ids = np.arange(0, seq_length, dtype="int64")
 
         # attention_mask = (attention_mask - 1.0) * 1e9
         # attention_mask = attention_mask.astype("float32")
         # return [tokens, loss_mask, attention_mask, position_ids, labels]
-        return [tokens, loss_mask, labels]
+        return [tokens, loss_mask, position_ids, labels]
 
     def _get_single_sample_from_idx(self, doc_index_f, doc_index_l, offset_f,
                                     offset_l):

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

@@ -497,8 +497,8 @@ def forward(self, input_ids, position_ids=None):
         position_embeddings = self.position_embeddings(position_ids)
         embeddings = input_embedings + position_embeddings
 
-        with get_rng_state_tracker().rng_state('global_seed'):
-            embeddings = self.dropout(embeddings)
+        #with get_rng_state_tracker().rng_state('global_seed'):
+        embeddings = self.dropout(embeddings)
 
         return embeddings
 
@@ -754,23 +754,24 @@ def forward(self,
             input_ids=input_ids, position_ids=position_ids)
 
         # TODO, use registered buffer
-        causal_mask = paddle.tensor.triu(
-            paddle.ones((paddle.shape(input_ids)[-1],
-                         paddle.shape(input_ids)[-1])) * -1e9,
-            diagonal=1)
+        # causal_mask = paddle.tensor.triu(
+        #     paddle.ones((paddle.shape(input_ids)[-1],
+        #                  paddle.shape(input_ids)[-1])) * -1e9,
+        #     diagonal=1)
 
-        if attention_mask is not None:
-            attention_mask = attention_mask + causal_mask
-        else:
-            attention_mask = causal_mask
+        # if attention_mask is not None:
+        #     attention_mask = attention_mask + causal_mask
+        # else:
+        #     attention_mask = causal_mask
 
         # The tensor returned by triu not in static graph.
-        attention_mask.stop_gradient = True
+        # attention_mask.stop_gradient = True
 
         encoder_outputs = self.decoder(
             embedding_output,
             memory=None,
-            tgt_mask=attention_mask,
+            # tgt_mask=attention_mask,
+            tgt_mask=None,
             use_cache=use_cache,
             cache=cache)
         self.checkpoints.extend(self.decoder.checkpoints)

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

@@ -66,12 +66,12 @@ def run_evaluate(args,
     all_loss = []
     local_time = time.time()
     for eval_step, batch in enumerate(data_loader):
-        tokens, loss_mask, labels = batch
+        tokens, loss_mask, position_ids, labels = batch
         if args.pp_degree < 2:
-            preds = model(tokens)
+            preds = model(tokens, position_ids)
             loss = criterion(preds, labels, loss_mask)
         else:
-            data = [tokens, (labels, loss_mask)]
+            data = [(tokens, position_ids), (labels, loss_mask)]
             loss = model.eval_batch(data, compute_loss=True)
 
         all_loss.append(float(loss))
@@ -237,10 +237,11 @@ def do_train(args):
 
             for step, batch in enumerate(train_data_loader()):
                 global_step += 1
-                tokens, loss_mask, labels = batch
+                tokens, loss_mask, position_ids, labels = batch
 
                 loss_mask.stop_gradient = True
                 labels.stop_gradient = True
+                position_ids.stop_gradient = True
 
                 if args.pp_degree == 1:
                     with paddle.amp.auto_cast(
@@ -252,7 +253,7 @@ def do_train(args):
                                 "reduce_sum", "c_softmax_with_cross_entropy",
                                 "c_embedding"
                             ]):
-                        preds = model(tokens)
+                        preds = model(tokens, position_ids)
                         loss = criterion(preds, labels, loss_mask)
 
                     if args.use_amp:
@@ -267,7 +268,7 @@ def do_train(args):
                     optimizer.clear_grad()
 
                 else:
-                    data = [tokens, (labels, loss_mask)]
+                    data = [(tokens, position_ids), (labels, loss_mask)]
                     with paddle.amp.auto_cast(
                             args.use_amp,
                             custom_white_list=[