Merge pull request #198 from Modalities/rank_0_logging

feat: moved print_rank_0 function towards a better fitting place

src/modalities/dataloader/open_gptx_dataset/mmap_dataset.py

@@ -24,14 +24,7 @@
 from numpy._typing import NDArray
 from torch.utils.data import Dataset
 
-
-def print_rank_0(message: str):
-    """If distributed is initialized, print only on rank 0."""
-    if torch.distributed.is_initialized():
-        if torch.distributed.get_rank() == 0:
-            print(message, flush=True)
-    else:
-        print(message, flush=True)
+from modalities.util import print_rank_0
 
 
 def get_best_fitting_dtype(vocab_size: Optional[int] = None) -> np.dtype:

src/modalities/util.py

@@ -17,6 +17,15 @@
 from modalities.running_env.fsdp.reducer import Reducer
 
 
+def print_rank_0(message: str):
+    """If distributed is initialized, print only on rank 0."""
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            print(message, flush=True)
+    else:
+        print(message, flush=True)
+
+
 def parse_enum_by_name(name: str, enum_type: Type[Enum]) -> Enum:
     try:
         return enum_type[name]