show model size

allenai · Nov 19, 2024 · 5660720 · 5660720
1 parent 0face65
commit 5660720
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/src/olmo_core/internal/experiment.py b/src/olmo_core/internal/experiment.py
@@ -4,6 +4,7 @@
 from typing import Callable, Dict, List, Optional, cast
 
 from beaker import Beaker
+from rich import print
 from torch.distributed.device_mesh import DeviceMesh
 
 from olmo_core.config import Config, StrEnum
@@ -111,10 +112,17 @@ def prepare_environment(self):
             raise NotADirectoryError(self)
 
     def run(self, config: ExperimentConfig):
+        print(config)
+        print(
+            "\n"
+            f"[b blue]Total parameters:[/]                {config.model.num_params:,d}\n"
+            f"[b blue]Non-embedding parameters:[/]        {config.model.num_non_embedding_params:,d}"
+        )
+
         if self == SubCmd.launch:
             launch(config)
         elif self == SubCmd.dry_run:
-            log.info(config)
+            pass
         elif self == SubCmd.train:
             try:
                 train(config)