PaddlePaddle
diff --git a/‎python/paddle/distributed/auto_tuner/prune.py
Copy file name to clipboard
+11-8Lines changed: 11 additions & 8 deletions b/‎python/paddle/distributed/auto_tuner/prune.py
Copy file name to clipboard
+11-8Lines changed: 11 additions & 8 deletions
diff --git a/‎python/paddle/distributed/auto_tuner/recorder.py
Copy file name to clipboard
+17-16Lines changed: 17 additions & 16 deletions b/‎python/paddle/distributed/auto_tuner/recorder.py
Copy file name to clipboard
+17-16Lines changed: 17 additions & 16 deletions
@@ -54,7 +54,10 @@ def same_cfgs_beside(attr, cur_cfg, history_cfgs=[]):
         for key in cur_cfg:
             if key == attr:
                 continue
-            if key not in cfg or cfg[key] != cur_cfg[key]:
+            if key not in cfg or (
+                cfg[key] != cur_cfg[key]
+                and key not in ["estimated_memory_usage"]
+            ):
                 same = False
                 break
         if same:
@@ -467,7 +470,7 @@ def prune_by_recompute_history(tuner_cfg, cur_cfg, history_cfgs=[]):
     if not use_recompute:
         cfgs = same_cfgs_beside("recompute_granularity", cur_cfg, history_cfgs)
         if cfgs:
-            pruned_reason = f"recompute_granularity {cfg['recompute_granularity']} invalid because use_recompute is {use_recompute}."
+            pruned_reason = f"recompute_granularity invalid because use_recompute is {use_recompute}."
             log_pruned_info(cur_cfg, pruned_reason)
             return True
     return False
@@ -567,21 +570,21 @@ def prune_by_memory_estimation(tuner_cfg, cur_cfg, history_cfgs=[]):
     )
 
     if result.returncode == 0:
-        cur_memory_usage = round(float(result.stdout), 2)
+        cur_memory_usage = int(round(float(result.stdout), 2))
         cur_cfg["estimated_memory_usage"] = cur_memory_usage
-        msg = f"estimated memory usage: {cur_memory_usage} GB"
-        memory_exceeded = cur_memory_usage > max_memory_usage
+        msg = f"Estimated {cur_cfg} memory usage: {cur_memory_usage} MB"
+        memory_exceeded = cur_memory_usage > (max_memory_usage * 1024)
         if memory_exceeded:
-            msg += ", WILL BE PRUNED!"
+            msg += ", it will be pruned!"
         logger.info(msg)
-        return cur_memory_usage > max_memory_usage
+        return memory_exceeded
     else:
         raise ValueError(
             f"memory_estimation_tool failed with error: {result.stderr}"
         )
 
 
-@register_prune
+@register_prune_history
 def prune_by_sharding_overlap(tuner_cfg, cur_cfg, history_cfgs=[]):
     """Prune by sharding overlap for single dp estimation"""
     if "sharding_overlap" in cur_cfg:
 
@@ -100,22 +100,23 @@ def store_history(self, path="./history.csv"):
         # get enhanced report in dp-estimation mode
         if self.search_algo == "dp_estimation":
             metric_name = self.tuner_cfg['metric_cfg']['name']
-            _history = []
-            for cfg in self.history:
-                if (
-                    "sharding_overlap" not in cfg.keys()
-                    or cfg["sharding_overlap"] is None
-                ) and cfg["error_info"] is None:
-                    _history.append(copy.deepcopy(cfg))
-            _history.sort(
-                key=lambda x: x[self.additional_metric_key]
-                if x[self.additional_metric_key] is not None
-                else float('-inf'),
-                reverse=True,
-            )
-            self._store_history_impl(
-                data=_history, path=path.split('.csv')[0] + '_enhanced.csv'
-            )
+            if self.additional_metric_key:
+                _history = []
+                for cfg in self.history:
+                    if (
+                        "sharding_overlap" not in cfg.keys()
+                        or cfg["sharding_overlap"] is None
+                    ) and cfg["error_info"] is None:
+                        _history.append(copy.deepcopy(cfg))
+                _history.sort(
+                    key=lambda x: x[self.additional_metric_key]
+                    if x[self.additional_metric_key] is not None
+                    else float('-inf'),
+                    reverse=True,
+                )
+                self._store_history_impl(
+                    data=_history, path=path.split('.csv')[0] + '_enhanced.csv'
+                )
 
         """Store history to csv file."""
         self.store_path = path