[AutoParallel]fix dist_loader when batch size is None (#60234)

zhaoyinglia · web-flow · commit 81d800bb9ab1 · 2023-12-22T11:02:11.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/dist_loader.py b/python/paddle/distributed/auto_parallel/static/dist_loader.py
@@ -254,21 +254,26 @@ def __init__(
         self.dp_world_sizes = data_parallel_world_size
         self.dp_ranks = data_parallel_rank
         self.split_data = split_data
-        # TODO: rank info
-        self.batch_sampler = DistributedBatchSampler(
-            dataset=self.dataset,
-            batch_size=self.batch_size,
-            num_replicas=self.dp_world_sizes[0],
-            rank=self.dp_ranks[0],
-            shuffle=self.shuffle,
-            drop_last=self.drop_last,
-        )
+
+        if self.batch_size is None:
+            self.batch_sampler = None
+        else:
+            self.batch_sampler = DistributedBatchSampler(
+                dataset=self.dataset,
+                batch_size=self.batch_size,
+                num_replicas=self.dp_world_sizes[0],
+                rank=self.dp_ranks[0],
+                shuffle=self.shuffle,
+                drop_last=self.drop_last,
+            )
+
         self._dataloader = paddle.io.DataLoader(
             self.dataset,
             feed_list=self.feed_list,
             places=self.places,
             return_list=self.return_list,
             batch_sampler=self.batch_sampler,
+            batch_size=1 if self.batch_sampler else self.batch_size,
             collate_fn=self.collate_fn,
             num_workers=self.num_workers,
             use_buffer_reader=self.use_buffer_reader,
