[FEAT][print_num_params]

zeta/utils/__init__.py

@@ -8,11 +8,14 @@
     save_memory_snapshot,
 )
 from zeta.utils.disable_logging import disable_warnings_and_logs
+from zeta.utils.params import print_num_params, print_main
 
 __all__ = [
     "track_cuda_memory_usage",
     "benchmark",
     "print_cuda_memory_usage",
     "save_memory_snapshot",
     "disable_warnings_and_logs",
+    "print_num_params",
+    "print_main",
 ]

zeta/utils/params.py

@@ -0,0 +1,29 @@
+import torch.distributed as dist  # Add this line
+
+
+def print_num_params(model):
+    """Print the number of parameters in a model.
+
+    Args:
+        model (_type_): _description_
+    """
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    if dist.is_available():
+        if dist.get_rank() == 0:
+            print(f"Number of parameters in model: {n_params}")
+    else:
+        print(f"Number of parameters in model: {n_params}")
+
+
+def print_main(msg):
+    """Print the message only on the main process.
+
+    Args:
+        msg (_type_): _description_
+    """
+    if dist.is_available():
+        if dist.get_rank() == 0:
+            print(msg)
+    else:
+        print(msg)