Fix pipeline

Author: DesmonDay

paddlenlp/data/dist_dataloader.py

@@ -23,8 +23,6 @@
     nested_reduce_tensor,
 )
 
-_MAX_DATA_DIM = 64
-
 
 class DummyDataset(paddle.io.Dataset):
     """
@@ -58,6 +56,7 @@ def __init__(
         timeout=0,
         worker_init_fn=None,
         persistent_workers=False,
+        eval=False,
     ):
 
         if dataset is None:
@@ -67,6 +66,7 @@ def __init__(
         super().__init__(dataset=dataset, batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=num_workers)
 
         self._hcg = fleet.get_hybrid_communicate_group()
+        self.eval = eval
 
         # Init pp data comm group.
         if self._hcg.get_pipe_parallel_world_size() > 1:
@@ -128,8 +128,11 @@ def _init_dataloader_comm_group(self):
         parallel_groups = topo.get_comm_list("pipe")
 
         for group in parallel_groups:
-            # only first rank and last rank
-            ranks = [group[0], group[-1]]
+            if not self.eval:
+                # only first rank and last rank
+                ranks = [group[0], group[-1]]
+            else:
+                ranks = group
             comm_group = paddle.distributed.new_group(ranks=ranks)
             if paddle.distributed.get_rank() in ranks:
                 parallel_comm_group = comm_group
@@ -170,6 +173,9 @@ def _broadcast_data(self, data):
                 src=self._pp_data_group.ranks[0],
                 group=self._pp_data_group,
             )
+        else:
+            fake_data = [None]
+
         fake_data = fake_data[0]
 
         if self.mp_group.nranks > 1:
@@ -178,7 +184,6 @@ def _broadcast_data(self, data):
         if self._pp_data_group is not None:
             if process_rank != self._pp_data_group.ranks[0]:
                 data = nested_empty_tensor(fake_data)
-        data = nested_copy_place(data, place=paddle.framework._current_expected_place())
 
         if self.mp_group.nranks > 1 and self.pp_rank == 0:
             data = nested_broadcast_tensor(data, src=self.mp_src_rank, group=self.mp_group)
@@ -195,6 +200,7 @@ def __next__(self):
         if self._need_data:
             try:
                 data = next(self._dataloader_iter)
+                data = nested_copy_place(data, place=paddle.framework._current_expected_place())
             except:
                 pass
         data = self._broadcast_data(data)

paddlenlp/trainer/trainer.py

@@ -1447,24 +1447,41 @@ def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoa
                     process_index=self.args.dataset_rank,
                 )
 
-            return _DataLoader(
-                eval_dataset,
-                batch_size=self.args.per_device_eval_batch_size,
-                collate_fn=self.data_collator,
-                num_workers=self.args.dataloader_num_workers,
-            )
+            if self.args.distributed_dataloader:
+                return _DataLoader(
+                    eval_dataset,
+                    batch_size=self.args.per_device_eval_batch_size,
+                    collate_fn=self.data_collator,
+                    num_workers=self.args.dataloader_num_workers,
+                    eval=True,
+                )
+            else:
+                return _DataLoader(
+                    eval_dataset,
+                    batch_size=self.args.per_device_eval_batch_size,
+                    collate_fn=self.data_collator,
+                    num_workers=self.args.dataloader_num_workers,
+                )
 
         eval_sampler = self._get_eval_sampler(eval_dataset)
 
         if self.args.distributed_dataloader:
             logger.info("Eval using DistDataLoader.")
 
-        return _DataLoader(
-            eval_dataset,
-            batch_sampler=eval_sampler,
-            collate_fn=self.data_collator,
-            num_workers=self.args.dataloader_num_workers,
-        )
+            return _DataLoader(
+                eval_dataset,
+                batch_sampler=eval_sampler,
+                collate_fn=self.data_collator,
+                num_workers=self.args.dataloader_num_workers,
+                eval=True,
+            )
+        else:
+            return _DataLoader(
+                eval_dataset,
+                batch_sampler=eval_sampler,
+                collate_fn=self.data_collator,
+                num_workers=self.args.dataloader_num_workers,
+            )
 
     def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
         """
@@ -1497,25 +1514,42 @@ def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
                     process_index=self.args.dataset_rank,
                 )
 
-            return _DataLoader(
-                test_dataset,
-                batch_size=self.args.per_device_eval_batch_size * self.world_size,
-                collate_fn=self.data_collator,  # _get_collator_with_removed_columns
-                num_workers=self.args.dataloader_num_workers,
-            )
+            if self.args.distributed_dataloader:
+                return _DataLoader(
+                    test_dataset,
+                    batch_size=self.args.per_device_eval_batch_size * self.world_size,
+                    collate_fn=self.data_collator,  # _get_collator_with_removed_columns
+                    num_workers=self.args.dataloader_num_workers,
+                    eval=True,
+                )
+            else:
+                return _DataLoader(
+                    test_dataset,
+                    batch_size=self.args.per_device_eval_batch_size * self.world_size,
+                    collate_fn=self.data_collator,  # _get_collator_with_removed_columns
+                    num_workers=self.args.dataloader_num_workers,
+                )
 
         test_sampler = self._get_eval_sampler(test_dataset)
 
         if self.args.distributed_dataloader:
             logger.info("Test using DistDataLoader.")
 
-        # We use the same batch_size as for eval.
-        return _DataLoader(
-            test_dataset,
-            batch_sampler=test_sampler,
-            collate_fn=self.data_collator,
-            drop_last=self.args.dataloader_drop_last,
-        )
+            # We use the same batch_size as for eval.
+            return _DataLoader(
+                test_dataset,
+                batch_sampler=test_sampler,
+                collate_fn=self.data_collator,
+                drop_last=self.args.dataloader_drop_last,
+                eval=True,
+            )
+        else:
+            return _DataLoader(
+                test_dataset,
+                batch_sampler=test_sampler,
+                collate_fn=self.data_collator,
+                drop_last=self.args.dataloader_drop_last,
+            )
 
     def create_optimizer_and_scheduler(self, num_training_steps: int):
         """