[AutoParallel] Optimize remove_no_need_vars (#59912)

AndSonder · web-flow · commit c0445a47212f · 2023-12-14T11:03:34.000+08:00
* optimize remove no need var

* apply suggestions from code review
diff --git a/python/paddle/distributed/auto_parallel/static/reshard.py b/python/paddle/distributed/auto_parallel/static/reshard.py
@@ -903,7 +903,8 @@ def remove_no_need_vars(
             for var in remove_vars:
                 if var in feed_var_names:
                     continue
-                block._remove_var(var)
+                block._remove_var(var, sync=False)
+            block._sync_with_cpp()
 
     @staticmethod
     def remove_no_need_in_main(
@@ -971,7 +972,8 @@ def remove_no_need_in_startup(
             if var_name not in actual_need_vars:
                 remove_vars.add(var_name)
         for var in remove_vars:
-            startup_block._remove_var(var)
+            startup_block._remove_var(var, sync=False)
+        startup_block._sync_with_cpp()
 
         remove_op_idx = []
         vars = startup_block.vars
