perf: use multithread to accelarate stat computing and loading

Author: OutisLi

deepmd/pt/utils/stat.py

@@ -3,6 +3,9 @@
 from collections import (
     defaultdict,
 )
+from concurrent.futures import (
+    ThreadPoolExecutor,
+)
 from typing import (
     Any,
     Callable,
@@ -39,7 +42,7 @@
 def make_stat_input(
     datasets: list[Any], dataloaders: list[Any], nbatches: int
 ) -> dict[str, Any]:
-    """Pack data for statistics.
+    """Pack data for statistics in parallel.
 
     Args:
     - dataset: A list of dataset to analyze.
@@ -49,49 +52,83 @@ def make_stat_input(
     -------
     - a list of dicts, each of which contains data from a system
     """
-    lst = []
     log.info(f"Packing data for statistics from {len(datasets)} systems")
-    for i in range(len(datasets)):
-        sys_stat = {}
-        with torch.device("cpu"):
-            iterator = iter(dataloaders[i])
-            numb_batches = min(nbatches, len(dataloaders[i]))
-            for _ in range(numb_batches):
-                try:
-                    stat_data = next(iterator)
-                except StopIteration:
-                    iterator = iter(dataloaders[i])
-                    stat_data = next(iterator)
-                if (
-                    "find_fparam" in stat_data
-                    and "fparam" in stat_data
-                    and stat_data["find_fparam"] == 0.0
-                ):
-                    # for model using default fparam
-                    stat_data.pop("fparam")
-                    stat_data.pop("find_fparam")
-                for dd in stat_data:
-                    if stat_data[dd] is None:
-                        sys_stat[dd] = None
-                    elif isinstance(stat_data[dd], torch.Tensor):
-                        if dd not in sys_stat:
-                            sys_stat[dd] = []
-                        sys_stat[dd].append(stat_data[dd])
-                    elif isinstance(stat_data[dd], np.float32):
-                        sys_stat[dd] = stat_data[dd]
-                    else:
-                        pass
-
-        for key in sys_stat:
-            if isinstance(sys_stat[key], np.float32):
-                pass
-            elif sys_stat[key] is None or sys_stat[key][0] is None:
+    dataloader_lens = [len(dl) for dl in dataloaders]
+    args_list = [
+        (dataloaders[i], nbatches, dataloader_lens[i]) for i in range(len(datasets))
+    ]
+
+    lst = []
+    # I/O intensive, set a larger number of workers
+    with ThreadPoolExecutor(max_workers=256) as executor:
+        lst = list(executor.map(_process_one_dataset, args_list))
+    log.info("Finished packing data.")
+    return lst
+
+
+def _process_one_dataset(args: tuple[Any, int, int]) -> dict[str, Any]:
+    """
+    Helper function to process a single dataset's dataloader for statistics.
+    Designed to be called in parallel by a ThreadPoolExecutor.
+
+    Parameters
+    ----------
+    args : tuple(Any, int, int)
+        A tuple containing (dataloader, nbatches, dataloader_len)
+
+    Returns
+    -------
+    dict[str, Any]
+        The processed sys_stat dictionary for one dataset.
+    """
+    dataloader, nbatches, dataloader_len = args
+    sys_stat = {}
+
+    with torch.device("cpu"):
+        iterator = iter(dataloader)
+        numb_batches = min(nbatches, dataloader_len)
+
+        for _ in range(numb_batches):
+            try:
+                stat_data = next(iterator)
+            except StopIteration:
+                iterator = iter(dataloader)
+                stat_data = next(iterator)
+
+            if (
+                "find_fparam" in stat_data
+                and "fparam" in stat_data
+                and stat_data["find_fparam"] == 0.0
+            ):
+                # for model using default fparam
+                stat_data.pop("fparam")
+                stat_data.pop("find_fparam")
+
+            for dd in stat_data:
+                if stat_data[dd] is None:
+                    sys_stat[dd] = None
+                elif isinstance(stat_data[dd], torch.Tensor):
+                    if dd not in sys_stat:
+                        sys_stat[dd] = []
+                    sys_stat[dd].append(stat_data[dd])
+                elif isinstance(stat_data[dd], np.float32):
+                    sys_stat[dd] = stat_data[dd]
+                else:
+                    pass
+
+    for key in sys_stat:
+        if isinstance(sys_stat[key], np.float32):
+            pass
+        elif isinstance(sys_stat[key], list):
+            if sys_stat[key][0] is None:
                 sys_stat[key] = None
-            elif isinstance(stat_data[dd], torch.Tensor):
+            else:
                 sys_stat[key] = torch.cat(sys_stat[key], dim=0)
-        dict_to_device(sys_stat)
-        lst.append(sys_stat)
-    return lst
+        elif sys_stat[key] is None:
+            pass
+
+    dict_to_device(sys_stat)
+    return sys_stat
 
 
 def _restore_from_file(

deepmd/utils/env_mat_stat.py

@@ -10,6 +10,9 @@
 from collections.abc import (
     Iterator,
 )
+from concurrent.futures import (
+    ThreadPoolExecutor,
+)
 from typing import (
     Optional,
 )
@@ -142,7 +145,7 @@ def save_stats(self, path: DPPath) -> None:
             (path / kk).save_numpy(np.array([vv.number, vv.sum, vv.squared_sum]))
 
     def load_stats(self, path: DPPath) -> None:
-        """Load the statistics of the environment matrix.
+        """Load the statistics of the environment matrix in parallel.
 
         Parameters
         ----------
@@ -151,13 +154,18 @@ def load_stats(self, path: DPPath) -> None:
         """
         if len(self.stats) > 0:
             raise ValueError("The statistics has already been computed.")
-        for kk in path.glob("*"):
-            arr = kk.load_numpy()
-            self.stats[kk.name] = StatItem(
-                number=arr[0],
-                sum=arr[1],
-                squared_sum=arr[2],
-            )
+
+        files_to_load = list(path.glob("*"))
+
+        if not files_to_load:
+            raise ValueError(f"No statistics files found in {path}.")
+
+        with ThreadPoolExecutor(max_workers=128) as executor:
+            results = executor.map(self._load_stat_file, files_to_load)
+
+        for name, stat_item in results:
+            if stat_item is not None:
+                self.stats[name] = stat_item
 
     def load_or_compute_stats(
         self, data: list[dict[str, np.ndarray]], path: Optional[DPPath] = None
@@ -216,3 +224,19 @@ def get_std(
             kk: vv.compute_std(default=default, protection=protection)
             for kk, vv in self.stats.items()
         }
+
+    @staticmethod
+    def _load_stat_file(file_path: DPPath) -> tuple[str, StatItem]:
+        """Helper function for parallel loading of stat files."""
+        try:
+            arr = file_path.load_numpy()
+            if arr.shape == (3,):
+                return file_path.name, StatItem(
+                    number=arr[0], sum=arr[1], squared_sum=arr[2]
+                )
+            else:
+                log.warning(f"Skipping malformed stat file: {file_path.name}")
+                return file_path.name, None
+        except Exception as e:
+            log.warning(f"Failed to load stat file {file_path.name}: {e}")
+            return file_path.name, None