OneLogger Integration

nemo/collections/llm/api.py

@@ -58,6 +58,7 @@
 from nemo.lightning.base import NEMO_MODELS_CACHE
 from nemo.lightning.ckpt_utils import ckpt_to_context_subdir
 from nemo.lightning.pytorch.callbacks import PEFT, JitTransform, ModelTransform
+from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
 from nemo.utils import logging
 from nemo.utils.get_rank import is_global_rank_zero
 
@@ -135,6 +136,9 @@ def train(
 
     trainer.fit(model, data)
 
+    # Track app end for NeMo v2 recipe-based applications
+    CallbackGroup.get_instance().on_app_end()
+
     return app_state.exp_dir
 
 
@@ -1255,11 +1259,19 @@ def _setup(
         resume_if_exists=getattr(resume, "resume_if_exists", False),
         task_config=getattr(train, "__io__", None),
     )
+
+    # Configure telemetry via CallbackGroup
+    CallbackGroup.get_instance().update_config(nemo_version='v2', trainer=trainer, data=data)
+
     if resume is not None:
+        CallbackGroup.get_instance().on_load_checkpoint_start()
         resume.setup(trainer, model)
+        CallbackGroup.get_instance().on_load_checkpoint_end()
 
     if optim:
+        CallbackGroup.get_instance().on_optimizer_init_start()
         optim.connect(model)
+        CallbackGroup.get_instance().on_optimizer_init_end()
     if tokenizer:  # TODO: Improve this
         _use_tokenizer(model, data, tokenizer)
 

nemo/collections/llm/fn/mixin.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import lightning.pytorch as pl
 from torch import nn
 from typing_extensions import Self
 
@@ -50,6 +51,15 @@ class FNMixin:
         True
     """
 
+    def __init_subclass__(cls, **kwargs):
+        # Add OneLogger timing hooks for LightningModule subclasses to enable telemetry tracking
+        if issubclass(cls, pl.LightningModule):
+            from nemo.lightning.pytorch.callbacks.callback_group import hook_class_init_with_callbacks
+
+            hook_class_init_with_callbacks(cls, "on_model_init_start", "on_model_init_end")
+
+        super().__init_subclass__(**kwargs)
+
     def forall(self, func: fn.ModulePredicate, recurse: bool = False) -> bool:
         """
         Evaluates a predicate for all modules in the container, optionally recursively.

nemo/collections/llm/gpt/data/mock.py

@@ -68,6 +68,10 @@ def __init__(
         vocab_file: Optional[str] = None,
         merges_file: Optional[str] = None,
     ):
+        from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
+
+        CallbackGroup.get_instance().on_dataloader_init_start()
+
         super().__init__()
         self.seq_length = seq_length
         self.micro_batch_size = micro_batch_size
@@ -96,6 +100,8 @@ def __init__(
             rampup_batch_size=rampup_batch_size,
         )
 
+        CallbackGroup.get_instance().on_dataloader_init_end()
+
     def setup(self, stage: str = "") -> None:
         """
         Setup the data module.

nemo/collections/llm/modelopt/speculative/model_transform.py

@@ -24,7 +24,7 @@
 
 try:
     ALGORITHMS = {
-        "eagle3": mtsp.EAGLE3_DEFAULT_CFG,
+        "eagle3": mtsp.EAGLE3_DEFAULT_CFG if hasattr(mtsp, "EAGLE3_DEFAULT_CFG") else None,
         # more TBD
     }
 except UnavailableError:

nemo/collections/llm/t5/data/mock.py

@@ -49,6 +49,10 @@ def __init__(
         persistent_workers: bool = False,
         create_attention_mask: bool = False,
     ):
+        from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
+
+        CallbackGroup.get_instance().on_dataloader_init_start()
+
         super().__init__()
         self.seq_length = seq_length
         self.seq_length_dec = seq_length_dec
@@ -72,6 +76,8 @@ def __init__(
             rampup_batch_size=rampup_batch_size,
         )
 
+        CallbackGroup.get_instance().on_dataloader_init_end()
+
     def setup(self, stage: str = "") -> None:
         """Setup the datasets"""
         self._train_ds = _MockT5Dataset(

nemo/collections/speechlm2/parts/optim_setup.py

@@ -88,7 +88,7 @@ def freeze_and_subset(
 
         >>> model = MyModel()
         ... # freeze all LLM parameters in "model.llm"
-        ... params = freeze_and_subset(model.named_parameters(), ['^llm\..+$'])
+        ... params = freeze_and_subset(model.named_parameters(), [r'^llm\\.\\..+$'])
         ... optimizer = torch.optim.AdamW(params, lr=1e-3)
 
     """

nemo/collections/vlm/neva/data/preloaded.py

@@ -516,6 +516,10 @@ def __init__(
         num_image_embeddings_per_tile: int = 576,
         seed: int = 1234,
     ) -> None:
+        from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
+
+        CallbackGroup.get_instance().on_dataloader_init_start()
+
         super().__init__()
         if not isinstance(paths, (list, tuple)):
             paths = [paths]
@@ -576,6 +580,8 @@ def custom_on_megatron_step_start(self, step):
             dataloader_type="cyclic",
         )
 
+        CallbackGroup.get_instance().on_dataloader_init_end()
+
     def setup(self, stage: str = "") -> None:
         assert len(self.paths) == 1, "not yet support blend dataset in Neva 2.0!"
         self._train_ds = NevaDataset(

nemo/core/classes/modelPT.py

@@ -48,6 +48,7 @@
 from nemo.core.classes.common import Model
 from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
 from nemo.core.optim import McoreDistributedOptimizer, prepare_lr_scheduler
+from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
 from nemo.utils import logging, model_utils
 from nemo.utils.app_state import AppState
 from nemo.utils.debug_hook import register_debug_hooks
@@ -86,6 +87,10 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
                 f"trainer constructor argument must be either None or lightning.pytorch.Trainer. "
                 f"But got {type(trainer)} instead."
             )
+
+        # Track model init start
+        CallbackGroup.get_instance().on_model_init_start()
+
         super().__init__()
 
         """
@@ -152,6 +157,8 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
         if torch.cuda.is_available() and torch.cuda.current_device() is not None:
             app_state.device_id = torch.cuda.current_device()
 
+        CallbackGroup.get_instance().on_model_init_end()
+        CallbackGroup.get_instance().on_dataloader_init_start()
         if self._cfg is not None and not self._is_model_being_restored():
             # Setup data loaders now (default) or defer setup to `self.setup()`
             # if `defer_setup` is set in the config of the corresponding dataloader.
@@ -198,6 +205,8 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
                     f"Test config : \n{OmegaConf.to_yaml(self._cfg.test_ds)}"
                 )
 
+        CallbackGroup.get_instance().on_dataloader_init_end()
+
         # Create list of lists for val and test outputs to support multiple dataloaders
         # Initialize an empty list as sometimes self._validation_dl can be None at this stage
         self._validation_step_outputs = None
@@ -469,6 +478,8 @@ def restore_from(
         Returns:
             An instance of type cls or its underlying config (if return_config is set).
         """
+        # Notify OneLogger of checkpoint loading start for telemetry tracking
+        CallbackGroup.get_instance().on_load_checkpoint_start()
 
         if save_restore_connector is None:
             save_restore_connector = SaveRestoreConnector()
@@ -502,6 +513,10 @@ def restore_from(
         )
         if isinstance(instance, ModelPT):
             instance._save_restore_connector = save_restore_connector
+
+        # Notify OneLogger of checkpoint loading completion for telemetry tracking
+        CallbackGroup.get_instance().on_load_checkpoint_end()
+
         return instance
 
     @classmethod
@@ -518,6 +533,9 @@ def load_from_checkpoint(
         Loads ModelPT from checkpoint, with some maintenance of restoration.
         For documentation, please refer to LightningModule.load_from_checkpoint() documentation.
         """
+        # Notify OneLogger of checkpoint loading start for telemetry tracking
+        CallbackGroup.get_instance().on_load_checkpoint_start()
+
         checkpoint = None
         try:
             cls._set_model_restore_state(is_being_restored=True)
@@ -533,6 +551,10 @@ def load_from_checkpoint(
 
         finally:
             cls._set_model_restore_state(is_being_restored=False)
+
+        # Notify OneLogger of checkpoint loading completion for telemetry tracking
+        CallbackGroup.get_instance().on_load_checkpoint_end()
+
         return checkpoint
 
     @abstractmethod
@@ -729,7 +751,8 @@ def setup_optimization(
 
         if optimizer_cls is None:
             # Try to get optimizer name for dynamic resolution, defaulting to Adam
-            optimizer_name = optim_config.get('name', 'adam')
+            # Use or instead of default as None will also results in default value not used.
+            optimizer_name = optim_config.get('name') or 'adam'
         else:
             if inspect.isclass(optimizer_cls):
                 optimizer_name = optimizer_cls.__name__.lower()
@@ -890,8 +913,12 @@ def configure_optimizers(self):
         """
         Configure the optimizer and scheduler.
         """
+        # Track optimizer init start
+        CallbackGroup.get_instance().on_optimizer_init_start()
         self.setup_optimization()
 
+        CallbackGroup.get_instance().on_optimizer_init_end()
+
         if self._scheduler is None:
             return self._optimizer
         else:
@@ -955,6 +982,9 @@ def setup(self, stage: Optional[str] = None):
             if no_test_dataloader and test_deferred_setup:
                 self.setup_multiple_test_data(test_data_config=self._cfg.test_ds)
 
+        if stage == 'fit':
+            CallbackGroup.get_instance().update_config(nemo_version='v1', trainer=self._trainer)
+
     def train_dataloader(self):
         """
         Get the training dataloader.
@@ -1344,6 +1374,8 @@ def maybe_init_from_pretrained_checkpoint(self, cfg: OmegaConf, map_location: st
                 f"Found : {[args[idx] for idx, arg_present in enumerate(arg_matches) if arg_present]}"
             )
 
+        CallbackGroup.get_instance().on_load_checkpoint_start()
+
         if 'init_from_nemo_model' in cfg and cfg.init_from_nemo_model is not None:
             with open_dict(cfg):
                 if isinstance(cfg.init_from_nemo_model, str):
@@ -1460,6 +1492,9 @@ def maybe_init_from_pretrained_checkpoint(self, cfg: OmegaConf, map_location: st
                 else:
                     raise TypeError("Invalid type: init_from_ptl_ckpt is not a string or a dict!")
 
+        # Track load checkpoint end
+        CallbackGroup.get_instance().on_load_checkpoint_end()
+
     def teardown(self, stage: str):
         """
         Called at the end of fit and test.

nemo/core/config/hydra_runner.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import argparse
 import functools
 import os
 import sys
@@ -103,7 +102,7 @@ def wrapper(cfg_passthrough: Optional[DictConfig] = None) -> Any:
                         # Make sure the path is not set - as this will disable validation scheme.
                         if path != '':
                             sys.stderr.write(
-                                f"ERROR Cannot set config file path using `--config-name` when "
+                                "ERROR Cannot set config file path using `--config-name` when "
                                 "using schema. Please set path using `--config-path` and file name using "
                                 "`--config-name` separately.\n"
                             )
@@ -133,6 +132,10 @@ def parse_args(self, args=None, namespace=None):
                     config_path=config_path,
                     config_name=config_name,
                 )
+                # Import here to avoid circular import
+                from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
+
+                CallbackGroup.get_instance().on_app_end()
 
         return wrapper
 

nemo/core/optim/lr_scheduler.py

@@ -31,7 +31,7 @@
 from omegaconf import DictConfig, OmegaConf
 from torch.optim.lr_scheduler import _LRScheduler
 
-from nemo.core.config import SchedulerParams, get_scheduler_config, register_scheduler_params
+from nemo.core.config.schedulers import SchedulerParams, get_scheduler_config, register_scheduler_params
 from nemo.utils import logging
 from nemo.utils.model_utils import maybe_update_config_version
 

nemo/core/optim/optimizers.py

@@ -23,11 +23,11 @@
 from torch.optim import adadelta, adagrad, adamax, rmsprop, rprop
 from torch.optim.optimizer import Optimizer
 
-from nemo.core.config import OptimizerParams, get_optimizer_config, register_optimizer_params
+from nemo.core.config.optimizers import OptimizerParams, get_optimizer_config, register_optimizer_params
 from nemo.core.optim.adafactor import Adafactor
 from nemo.core.optim.adan import Adan
 from nemo.core.optim.novograd import Novograd
-from nemo.utils import logging
+
 from nemo.utils.model_utils import maybe_update_config_version
 
 AVAILABLE_OPTIMIZERS = {
@@ -195,14 +195,23 @@ def get_optimizer(name: str, **kwargs: Optional[Dict[str, Any]]) -> Optimizer:
         )
     if name == 'fused_adam':
         if not torch.cuda.is_available():
-            raise ValueError(f'CUDA must be available to use fused_adam.')
+            raise ValueError('CUDA must be available to use fused_adam.')
 
     optimizer = AVAILABLE_OPTIMIZERS[name]
     optimizer = partial(optimizer, **kwargs)
     return optimizer
 
 
 def init_optimizer_states(optimizer: Optimizer):
+    """
+    Initialize optimizer states for Adam-based optimizers.
+
+    This function initializes the exponential moving averages (exp_avg and exp_avg_sq)
+    for Adam, AdamW, and FusedAdam optimizers if they haven't been initialized yet.
+
+    Args:
+        optimizer: The optimizer instance to initialize states for
+    """
     adam_nondist_optims = (optim.Adam, optim.AdamW)
     if HAVE_APEX:
         adam_nondist_optims += (FusedAdam,)

nemo/lightning/io/mixin.py

@@ -27,6 +27,7 @@
 
 import fiddle as fdl
 import fiddle._src.experimental.dataclasses as fdl_dc
+import lightning.pytorch as pl
 from cloudpickle import dump
 from cloudpickle import load as pickle_load
 from fiddle._src import config as config_lib
@@ -189,6 +190,13 @@ def __new__(cls, *args, **kwargs):
     def __init_subclass__(cls):
         _io_register_serialization(cls)
 
+        # Add OneLogger timing hooks for data modules to enable telemetry tracking
+        if issubclass(cls, pl.LightningDataModule):
+            from nemo.lightning.pytorch.callbacks.callback_group import hook_class_init_with_callbacks
+
+            hook_class_init_with_callbacks(cls, "on_dataloader_init_start", "on_dataloader_init_end")
+        super().__init_subclass__()
+
     def io_transform_args(self, init_fn, *args, **kwargs) -> Dict[str, Any]:
         """
         Transforms and captures the arguments passed to the `__init__` method, filtering out