feat:mult-ranks adaptation(dist)

Author: lijialin03

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 repos:
 # Common hooks
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.1.0
+    rev: v6.0.0
     hooks:
     -   id: check-added-large-files
     -   id: check-merge-conflict

README.md

@@ -56,7 +56,7 @@ from padiff.utils import logger
 from padiff import compare_dumps
 
 if __name__ == "__main__":
-    logger.reset_dir( "./padiff_log")
+    logger.setup( "./padiff_log")
 
     cfg = {
         "atol": 1e-6,

padiff/abstracts/hooks/guard.py

@@ -248,7 +248,7 @@ def PaDiffGuard(
         # set max calls
         calls_context.set_limit(max_calls)
 
-        proxy_model = create_model(model, name=name, reset_dir=reset_flag)
+        proxy_model = create_model(model, name=name, reset=reset_flag)
         model._padiff_proxy = proxy_model
         logger.debug(f"PaDiffGuard: creating proxy model.")
 

padiff/abstracts/proxy/__init__.py

@@ -14,8 +14,9 @@
 
 # this folder is just used to support assign_weight interface
 
+import os
 import paddle
-from ...utils import logger
+from ...utils import logger, reset_dir
 from .model import ProxyModel
 
 
@@ -42,14 +43,15 @@ def remove_inplace(model):
             submodel.inplace = False
 
 
-def create_model(model, name=None, dump_freq=1, reset_dir=True):
+def create_model(model, name=None, dump_freq=1, reset=True):
     retval = ProxyModel.create_from(model, name, dump_freq)
     init_route(retval)
-    if retval.framework == "paddle" and paddle.distributed.get_rank() % 8 == 0 and reset_dir:
+    if retval.framework == "paddle" and retval.rank // 8 == 0 and reset:
         # Only reset the root path once for each machine, here we assume each machine has 8 GPUs
-        logger.reset_dir(retval.dump_path)
+        reset_dir(retval.dump_path)
     if retval.framework == "torch":
-        if reset_dir:
-            logger.reset_dir(retval.dump_path)
+        if reset:
+            reset_dir(retval.dump_path)
         remove_inplace(retval)
+    logger.setup(retval.dump_path)
     return retval

padiff/abstracts/proxy/model.py

@@ -20,7 +20,7 @@
 from ..marker import Marker
 from ..report import Report
 from ...tools import dump_grads, dump_params, dump_report, dump_weights, get_dump_root_path
-from ...utils import deco_iter, logger
+from ...utils import deco_iter, logger, get_rank, reset_dir
 from .params import ProxyParam
 
 
@@ -34,7 +34,8 @@ def __init__(self, model, name, framework, dump_freq=1):
         self.report = Report(self.marker)
         self.step = 0
 
-        self.dump_path = get_dump_root_path() + "/" + self.name
+        self.rank = get_rank(framework)
+        self.dump_path = f"{get_dump_root_path()}/{self.name}/rank_{self.rank}"
 
         self.dump_freq = dump_freq
         if self.dump_freq > 1:
@@ -190,31 +191,27 @@ def clear_report(self):
     def try_dump(self, dump_path=None):
         if self.step % self.dump_freq == 0:
             if dump_path is None:
-                dump_path = f"{self.dump_path}/step_{self.step}/rank_{paddle.distributed.get_rank()}"
-            logger.reset_dir(dump_path)
+                dump_path = f"{self.dump_path}/step_{self.step}/rank_{self.rank}"
+            reset_dir(dump_path)
             self.dump_params(dump_path)
             self.dump_report(dump_path)
         self.clear_report()
         self.step += 1
 
     def dump_report(self, dump_path=None):
-        if dump_path is None:
-            dump_path = f"{self.dump_path}/rank_{paddle.distributed.get_rank()}"
+        dump_path = self.dump_path if dump_path is None else dump_path
         dump_report(self, dump_path)
 
     def dump_params(self, dump_path=None):
-        if dump_path is None:
-            dump_path = f"{self.dump_path}/rank_{paddle.distributed.get_rank()}"
+        dump_path = self.dump_path if dump_path is None else dump_path
         dump_params(self, dump_path)
 
     def dump_weights(self, dump_path=None):
-        if dump_path is None:
-            dump_path = f"{self.dump_path}/rank_{paddle.distributed.get_rank()}"
+        dump_path = self.dump_path if dump_path is None else dump_path
         dump_weights(self, dump_path)
 
     def dump_grads(self, dump_path=None):
-        if dump_path is None:
-            dump_path = f"{self.dump_path}/rank_{paddle.distributed.get_rank()}"
+        dump_path = self.dump_path if dump_path is None else dump_path
         dump_grads(self, dump_path)
 
     """

padiff/cli.py

@@ -228,7 +228,7 @@ def main():
         cli_cfg["base_framework"] = args.base_framework
 
     log_dir = cli_cfg.pop("log_dir", "./padiff_log")
-    logger.reset_dir(log_dir)
+    logger.setup(log_dir)
 
     pt_cmd = cli_cfg.get("pt_cmd")
     pd_cmd = cli_cfg.get("pd_cmd")

padiff/tools/dump.py

@@ -19,7 +19,7 @@
 import numpy
 import paddle
 
-from ..utils import Counter, frames_to_string, logger, save_model_struct, get_numpy_from_tensor
+from ..utils import Counter, frames_to_string, logger, save_model_struct, get_numpy_from_tensor, reset_dir
 
 dump_root_path = os.path.join(sys.path[0], "padiff_dump")
 
@@ -34,7 +34,7 @@ def get_dump_root_path():
 
 
 def numpy_dumper(path, prefix):
-    logger.reset_dir(path)
+    reset_dir(path)
     counter = Counter()
 
     def dumper(value):
@@ -51,8 +51,30 @@ def dumper(value):
 """
 
 
+def report_deduplicate(report):
+    """
+    In some cases, it is necessary to deduplicate report (e.g., when pipeline parallelism > 1
+    and gradient accumulation steps > 1).
+    """
+    if not report.stack.root:
+        return report
+
+    base_root_str = report.stack.root[0].net_str
+    idx_repeat = len(report.stack.root)
+    for i in range(1, len(report.stack.root)):
+        if report.stack.root[i].net_str == base_root_str:
+            idx_repeat = i
+            break
+    report.stack.root = report.stack.root[:idx_repeat]
+    logger.warning_once(
+        "The report contains duplications, which might occur when pipeline parallelism > 1 and gradient accumulation > 1. "
+        "The report is automatically deduplicated, please note if this deduplication was incorrect."
+    )
+    return report
+
+
 def dump_report(model, dump_path):
-    report = model.report
+    report = report_deduplicate(model.report)
     tensor_path = dump_path + "/tensors"
     tensor_dumper = numpy_dumper(tensor_path, "tensor")
 

padiff/utils/log.py

@@ -30,7 +30,6 @@
 class Logger:
     def __init__(self):
         self._logger = None
-        self._is_initialized = False
         self.log_path = "padiff_log"
 
         for key, conf in log_config.items():
@@ -41,26 +40,25 @@ def __init__(self):
             log_colors={key: conf["color"] for key, conf in log_config.items()},
         )
 
-    def setup(self, log_parent_dir):
-        if self._is_initialized:
-            return
-
-        self._logger = logging.getLogger("padiff")
+    def setup(self, log_root_dir):
+        os.makedirs(log_root_dir, exist_ok=True)
 
         silent_flag = os.getenv("PADIFF_SILENT")
         log_level_flag = os.getenv("PADIFF_LOG_LEVEL")
-
         if log_level_flag and log_level_flag.upper() in ("DEBUG", "INFO", "WARNING", "ERROR"):
             log_level = getattr(logging, log_level_flag.upper())
         else:
             log_level = logging.INFO
-        self._logger.setLevel(log_level)
-        self._logger.propagate = False
 
-        if self._logger.handlers:
-            self._logger.handlers.clear()
+        if self._logger is None:
+            self._logger = logging.getLogger("padiff")
+            self._logger.setLevel(log_level)
+            self._logger.propagate = False
+
+            if self._logger.handlers:
+                self._logger.handlers.clear()
 
-        log_file_path = os.path.join(log_parent_dir, "padiff.log")
+        log_file_path = os.path.join(log_root_dir, "padiff.log")
         file_handler = logging.FileHandler(log_file_path, encoding="utf-8")
         file_formatter = logging.Formatter("[AutoDiff] [%(levelname)s] %(message)s")
         file_handler.setFormatter(file_formatter)
@@ -74,7 +72,7 @@ def setup(self, log_parent_dir):
             self._logger.info(f"Logging initialized. Log file: {log_file_path}")
 
         self._is_initialized = True
-        self.log_path = log_parent_dir
+        self.log_path = log_root_dir
 
     def info(self, *args):
         if self._logger is not None:
@@ -108,14 +106,10 @@ def debug(self, *args):
     def debug_once(self, *args):
         self.debug(*args)
 
-    def reset_dir(self, path):
-        if os.path.exists(path):
-            shutil.rmtree(path)
-        os.makedirs(path)
-        self.setup(path)
-
-    def log_file(self, filename, mode, info):
-        filepath = os.path.join(self.log_path, filename)
+    def log_file(self, filename, mode, info, root_dir=None):
+        if root_dir is None:
+            root_dir = self.log_path
+        filepath = os.path.join(root_dir, filename)
         with open(filepath, mode) as f:
             f.write(info)
         return filepath
@@ -129,6 +123,12 @@ def log_file(self, filename, mode, info):
 """
 
 
+def reset_dir(path):
+    if os.path.exists(path):
+        shutil.rmtree(path)
+    os.makedirs(path)
+
+
 def print_report_info(nodes, reports, exc, stage, msg=None):
 
     logger.error(

padiff/utils/utils.py

@@ -57,6 +57,15 @@ def get_numpy_from_tensor(tensor):
     return np_array
 
 
+def get_rank(framework: str):
+    rank = 0
+    if framework == "paddle" and paddle.distributed.is_initialized():
+        rank = paddle.distributed.get_rank()
+    elif framework == "torch" and torch.distributed.is_initialized():
+        rank = torch.distributed.get_rank()
+    return rank
+
+
 """
     clone tensor
 """