[AutoTuner] get enhanced report in dp estimation mode (#60294)

SylarTiaNII · web-flow · commit d424d15936c4 · 2023-12-25T18:55:19.000+08:00
diff --git a/python/paddle/distributed/auto_tuner/recorder.py b/python/paddle/distributed/auto_tuner/recorder.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import csv
 import os
 from typing import Tuple
@@ -21,9 +22,12 @@
 
 class HistoryRecorder:
     # NOTE increase extenable ablitity
-    def __init__(self) -> None:
+    def __init__(self, tuner_cfg) -> None:
+        self.tuner_cfg = tuner_cfg
+        self.search_algo = self.tuner_cfg['search_algo']['name']
         self.history = []
         self.store_path = None
+        self.additional_metric_key = None
 
     def add_cfg(self, **kwargs):
         cur_configs = {}
@@ -76,11 +80,10 @@ def get_best(self, metric, direction, mode=None) -> Tuple[dict, bool]:
             return (self.history[0], True)
         return (self.history[0], False)
 
-    def store_history(self, path="./history.csv"):
+    def _store_history_impl(self, data, path="./history.csv"):
         """Store history to csv file."""
-        self.store_path = path
         # convert to pd dataframe
-        df = pd.DataFrame(self.history)
+        df = pd.DataFrame(data)
         # move 'job_id' to the first column
         cols = df.columns.tolist()
         cols.insert(0, cols.pop(cols.index('job_id')))
@@ -91,7 +94,32 @@ def store_history(self, path="./history.csv"):
         if 'has_error' in df.columns:
             df = df.drop(columns=['has_error'])
         # write to csv
-        df.to_csv(self.store_path, index=False)
+        df.to_csv(path, index=False)
+
+    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'
+            )
+
+        """Store history to csv file."""
+        self.store_path = path
+        self._store_history_impl(data=self.history, path=path)
 
     def load_history(self, path="./history.csv") -> Tuple[list, bool]:
         """Load history from csv file."""
diff --git a/python/paddle/distributed/launch/main.py b/python/paddle/distributed/launch/main.py
@@ -436,7 +436,7 @@ def launch():
 
         is_first_task = True
         # build history recorder
-        recorder = HistoryRecorder()
+        recorder = HistoryRecorder(tuner_cfg)
 
         job_id = 0
         error_task_nums = 0
@@ -898,6 +898,15 @@ def launch():
                     )
                 )
                 amp = tuner_cfg["search_algo"]["conversion"].get("amp", False)
+                num_gpus = int(cur_cfg["num_gpus"])
+                seq_length = int(
+                    tuner_cfg["model_cfg"].get("max_seq_length", 2048)
+                )
+                cur_cfg[f"unified_{tuner_cfg['metric_cfg']['name']}"] = (
+                    round(single_dp_performance / num_gpus * seq_length, 2)
+                    if single_dp_performance
+                    else None
+                )
                 for bw in comm_bw:
                     if amp:
                         comm_time = model_size_b * (4 + 2) / bw
@@ -916,6 +925,17 @@ def launch():
                     cur_cfg[
                         f"bw_{bw}_{tuner_cfg['metric_cfg']['name']}"
                     ] = multi_dp_performace
+                    cur_cfg[
+                        f"unified_bw_{bw}_{tuner_cfg['metric_cfg']['name']}"
+                    ] = (
+                        round(multi_dp_performace / num_gpus * seq_length, 2)
+                        if multi_dp_performace
+                        else None
+                    )
+                    if recorder.additional_metric_key is None:
+                        recorder.additional_metric_key = (
+                            f"unified_bw_{bw}_{tuner_cfg['metric_cfg']['name']}"
+                        )
 
             error_info = None
             cur_cfg["has_error"] = has_error
