per_epoch_stats reported from all ranks

argonne-lcf · Feb 21, 2025 · b535005 · b535005
1 parent d695c60
commit b535005
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/dlio_benchmark/checkpointing/pytorch_checkpointing.py b/dlio_benchmark/checkpointing/pytorch_checkpointing.py
@@ -23,6 +23,7 @@
 from dlio_benchmark.common.constants import MODULE_CHECKPOINT
 from dlio_benchmark.common.enumerations import CheckpointLocationType
 from dlio_benchmark.utils.utility import DLIOMPI
+import logging
 
 def get_torch_datatype(datatype):
     if datatype == "fp32":

diff --git a/dlio_benchmark/utils/statscounter.py b/dlio_benchmark/utils/statscounter.py
@@ -421,10 +421,10 @@ def finalize(self):
     def save_data(self):
         # Dump statistic counters to files for postprocessing
         # Overall stats
+        with open(os.path.join(self.output_folder, f'{self.my_rank}_per_epoch_stats.json'), 'w') as outfile:
+            json.dump(self.per_epoch_stats, outfile, indent=4)
+            outfile.flush()
         if self.my_rank == 0:
-            with open(os.path.join(self.output_folder, 'per_epoch_stats.json'), 'w') as outfile:
-                json.dump(self.per_epoch_stats, outfile, indent=4)
-                outfile.flush()
             with open(os.path.join(self.output_folder, 'summary.json'), 'w') as outfile:
                 json.dump(self.summary, outfile, indent=4)
         self.output['hostname'] = socket.gethostname()