Add CallbackGroup & Metadata factory function

nemo/collections/nlp/parts/megatron_trainer_builder.py

@@ -32,6 +32,7 @@
     NLPFSDPStrategy,
     PipelineMixedPrecisionPlugin,
 )
+from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
 from nemo.utils import logging
 from nemo.utils.callbacks.dist_ckpt_io import (
     AsyncFinalizableCheckpointIO,
@@ -199,6 +200,9 @@ def create_trainer(self, callbacks=None) -> Trainer:
         precision = self.cfg.trainer.precision
         strategy = self._training_strategy()
         plugins = self._plugins()
+        if callbacks is None:
+            callbacks = []
+        callbacks.extend(CallbackGroup.get_instance().callbacks)
         callbacks = self._callbacks(callbacks)
         trainer = Trainer(plugins=plugins, strategy=strategy, **self.cfg.trainer, callbacks=callbacks)
         # Restore the precision value after Trainer is built.
@@ -227,6 +231,7 @@ def _callbacks(self, callbacks: Optional[list]) -> list:
     def create_trainer(self, callbacks=None) -> Trainer:
         strategy = self._training_strategy()
         plugins = self._plugins()
+        callbacks.extend(CallbackGroup.get_instance().callbacks)
         callbacks = self._callbacks(callbacks)
         return Trainer(plugins=plugins, strategy=strategy, **self.cfg.trainer, callbacks=callbacks)
 

nemo/core/classes/common.py

@@ -15,6 +15,7 @@
 
 """Interfaces common to all Neural Modules and Models."""
 from __future__ import annotations
+
 import copy
 import hashlib
 import inspect
@@ -42,6 +43,7 @@
 from nemo.core.config.templates.model_card import NEMO_DEFAULT_MODEL_CARD_TEMPLATE
 from nemo.core.connectors.save_restore_connector import SaveRestoreConnector
 from nemo.core.neural_types import NeuralType, NeuralTypeComparisonResult
+from nemo.lightning.pytorch.callbacks.callback_group import CallbackGroup
 from nemo.utils import logging
 from nemo.utils.cloud import maybe_download_from_cloud
 from nemo.utils.data_utils import resolve_cache_dir
@@ -739,6 +741,7 @@ def from_pretrained(
         Returns:
             A model instance of a particular model class or its underlying config (if return_config is set).
         """
+        CallbackGroup.get_instance().on_load_checkpoint_start()
         if save_restore_connector is None:
             save_restore_connector = SaveRestoreConnector()
 
@@ -772,6 +775,7 @@ def from_pretrained(
             trainer=trainer,
             save_restore_connector=save_restore_connector,
         )
+        CallbackGroup.get_instance().on_load_checkpoint_end()
         return instance
 
     @classmethod

nemo/core/classes/modelPT.py

@@ -47,6 +47,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, wrap_methods_with_callbacks
 from nemo.utils import logging, model_utils
 from nemo.utils.app_state import AppState
 from nemo.utils.debug_hook import register_debug_hooks
@@ -224,6 +225,7 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
 
     def __init_subclass__(cls) -> None:
         cls._save_restore_connector = SaveRestoreConnector()
+        wrap_methods_with_callbacks(cls)
 
     def on_fit_start(self) -> None:
         if self.cfg.get("dump_debug_info", False):

nemo/lightning/__init__.py

@@ -27,6 +27,7 @@
 from nemo.lightning.fabric.plugins import FabricMegatronMixedPrecision
 from nemo.lightning.fabric.strategies import FabricMegatronStrategy
 from nemo.lightning.nemo_logger import NeMoLogger
+from nemo.lightning.one_logger_callback import OneLoggerNeMoCallback
 from nemo.lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint
 from nemo.lightning.pytorch.optim import (
     LRSchedulerModule,
@@ -72,6 +73,7 @@ def _is_slurm_interactive_mode():
     "lr_scheduler",
     "NeMoLogger",
     "ModelCheckpoint",
+    "OneLoggerNeMoCallback",
     "OptimizerModule",
     "Trainer",
     "configure_no_restart_validation_training_loop",

nemo/lightning/nemo_logger.py

@@ -165,6 +165,31 @@ def setup(self, trainer: Union[pl.Trainer, fl.Fabric], resume_if_exists: bool =
             self._setup_trainer_loggers(trainer, _dir, version)
             self._setup_trainer_model_checkpoint(trainer, log_dir=log_dir, ckpt=self.ckpt)
 
+            # Configure OneLogger callback
+            try:
+                from omegaconf import OmegaConf
+
+                from nemo.utils.exp_manager import configure_onelogger
+
+                # Create a minimal config for OneLogger
+                cfg = OmegaConf.create(
+                    {
+                        "exp_manager": {
+                            "wandb_logger_kwargs": {
+                                "project": "nemo_experiments",
+                                "name": self.name,
+                                "id": version or None,
+                            }
+                        }
+                    }
+                )
+
+                # Configure OneLogger
+                configure_onelogger(cfg, trainer)
+                logging.info("OneLogger configured successfully")
+            except Exception as e:
+                logging.warning(f"Failed to configure OneLogger: {e}")
+
         self._setup_files_to_move(log_dir, app_state)
         self._setup_file_logging(log_dir)
 

nemo/lightning/one_logger_callback.py

@@ -0,0 +1,150 @@
+"""
+OneLogger callback for NeMo training.
+
+This module provides a callback that integrates OneLogger telemetry with NeMo training.
+"""
+
+import functools
+import logging
+from typing import Any, Dict, List, Optional, Type
+
+import nv_one_logger.training_telemetry.api.callbacks as CB
+import pytorch_lightning as pl
+import torch
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.core import LightningModule
+from pytorch_lightning.plugins.io import AsyncCheckpointIO
+from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+
+
+class OneLoggerNeMoCallback(Callback):
+    """
+    NeMo callback that integrates with OneLogger v2 for tracking metrics.
+
+    This callback implements NeMo's callback group API and internally
+    uses OneLogger's training telemetry functionality to track metrics.
+    """
+
+    def __init__(
+        self,
+        callback_config: Optional[Dict[str, Any]] = None,
+        log_interval: int = 1,
+        async_io_checkpoint_classes: List[Type[Any]] | None = None,
+    ):
+        """
+        Initialize the OneLogger NeMo callback.
+
+        Args:
+            callback_config (dict): Configuration dictionary with metadata
+                from MetaInfoManager(cfg).get_metadata()
+            log_interval (int): How often to log metrics
+            async_io_checkpoint_classes (List[Type]): Additional classes to identify as async checkpoints
+        """
+        super().__init__()
+        self.log_interval = log_interval
+        self.async_io_checkpoint_classes = async_io_checkpoint_classes or []
+        self.state = {
+            "is_async_checkpoint": None,
+        }
+
+        # Extract configuration values
+        if callback_config is not None:
+            self.app_name = callback_config.get("app_name", "")
+            self.perf_tag = callback_config.get("perf_tag", "")
+            self.session_tag = callback_config.get("session_tag", "")
+            self.global_batch_size = callback_config.get("global_batch_size", 0)
+        else:
+            self.app_name = ""
+            self.perf_tag = ""
+            self.session_tag = ""
+            self.global_batch_size = 0
+
+    def __getattr__(self, name: str) -> Any:
+        """Automatically forward any undefined method calls to the OneLogger v2 callbacks mainly for non-trainer methods.
+
+        This eliminates the need for manually writing pass-through methods for each OneLogger API.
+        Only methods that need custom logic (like those interacting with the trainer) need to be
+        explicitly defined in this class.
+
+        Args:
+            name: The name of the method being called
+        Returns:
+            The method from the OneLogger v2 callbacks
+        Raises:
+            AttributeError: If the method is not found in the OneLogger callbacks
+        """
+        # Check if the method exists in the OneLogger callbacks module
+        if hasattr(CB, name):
+            # Get the original method
+            original_method = getattr(CB, name)
+
+            # Create a wrapper that adds rank_zero_only decorator
+            @functools.wraps(original_method)
+            def wrapper(*args, **kwargs):
+                return rank_zero_only(original_method)(*args, **kwargs)
+
+            return wrapper
+
+        # If not found, raise AttributeError as normal
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    @rank_zero_only
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Called when training begins."""
+        # Extract necessary information from the trainer
+        current_step = trainer.global_step
+        max_steps = trainer.max_steps if hasattr(trainer, 'max_steps') else 0
+
+        CB.on_train_start(train_iterations_start=current_step, train_iterations_target_or_fn=max_steps)
+
+    @rank_zero_only
+    def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        CB.on_train_end()
+
+    @rank_zero_only
+    def on_train_batch_start(self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int) -> None:
+        CB.on_training_single_iteration_start()
+
+    @rank_zero_only
+    def on_train_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        CB.on_training_single_iteration_end()
+
+    @rank_zero_only
+    def on_validation_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        CB.on_validation_start()
+
+    @rank_zero_only
+    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        CB.on_validation_end()
+
+    @rank_zero_only
+    def on_validation_batch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        CB.on_validation_single_iteration_start()
+
+    @rank_zero_only
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        CB.on_validation_single_iteration_end()