From 722f880e3083d443ba5c115cf3900c288792a059 Mon Sep 17 00:00:00 2001
From: zhaoyingli <86812880+zhaoyinglia@users.noreply.github.com>
Date: Mon, 3 Apr 2023 15:36:51 +0800
Subject: [PATCH] make micro bsz configurable (#52447)

---
 python/paddle/distributed/auto_parallel/parallelizer_v2.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python/paddle/distributed/auto_parallel/parallelizer_v2.py b/python/paddle/distributed/auto_parallel/parallelizer_v2.py
index 6f77dbd4e07..1e30467c4f7 100644
--- a/python/paddle/distributed/auto_parallel/parallelizer_v2.py
+++ b/python/paddle/distributed/auto_parallel/parallelizer_v2.py
@@ -155,6 +155,11 @@ class Parallelizer:
                     time.time() - time0, self._mode
                 )
             )
+            micro_bsz = (
+                1
+                if not self._strategy.pipeline.enable
+                else self._strategy.pipeline.micro_batch_size
+            )
             time0 = time.time()
             resharder = Resharder(
                 dist_main_prog,
@@ -162,7 +167,7 @@ class Parallelizer:
                 rank,
                 self._dist_context,
                 [],
-                1,
+                micro_bsz,
             )
             resharder.reshard()
             self._logger.debug(
-- 
GitLab