diff --git a/src/olmo_core/internal/experiment.py b/src/olmo_core/internal/experiment.py index d3c2b208..407e9a19 100644 --- 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)