diff --git a/torchtitan/__init__.py b/torchtitan/__init__.py index 176bce9b60..52c3ff3e22 100644 --- a/torchtitan/__init__.py +++ b/torchtitan/__init__.py @@ -4,5 +4,12 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from importlib.metadata import version + # Import to register quantization modules. import torchtitan.components.quantization # noqa: F401 + +try: + __version__ = version("torchtitan") +except Exception as e: + __version__ = "0.0.0+unknown" diff --git a/torchtitan/train.py b/torchtitan/train.py index 6f039b3c04..1b4096abb6 100644 --- a/torchtitan/train.py +++ b/torchtitan/train.py @@ -718,6 +718,14 @@ def main(trainer_class: type[Trainer]) -> None: trainer_class: The trainer class to instantiate (e.g., Trainer, FluxTrainer, TorchCommsTrainer) """ init_logger() + + import torchtitan + + logger.info( + "torchtitan version: %s (0.0.0 means __version__ is not defined correctly).", + torchtitan.__version__, + ) + config_manager = ConfigManager() config = config_manager.parse_args() trainer: Trainer | None = None