Update training dependency compatibility

@@ -15,6 +15,13 @@ To use the Experiment Manager, call :class:`~nemo.utils.exp_manager.exp_manager`  The Experiment Manager is configurable using YAML with Hydra. +NeMo's Lightning integration preserves the previous callback defaults: ``TQDMProgressBar``+and a plain ``ModelSummary``. Explicit callbacks remain unchanged; selecting a+``RichProgressBar`` also selects a Rich model summary unless a summary callback is+already provided. ``trainer.enable_progress_bar: false``,+``trainer.enable_model_summary: false``, and barebones mode retain their usual behavior.+This applies to direct Lightning Trainer construction as well as NeMo's Trainer configuration helper.+ .. code-block:: bash      exp_manager:@@ -33,6 +40,9 @@ Optionally, launch TensorBoard to view the training results in ``exp_dir``, whic  If ``create_checkpoint_callback`` is set to ``True``, then NeMo automatically creates checkpoints during training using PyTorch Lightning's `ModelCheckpoint <https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html>`_.+With ``save_last: true``, NeMo refreshes the last checkpoint at each eligible validation or+training-epoch boundary even when the monitored metric does not improve. Async saves update+scheduling state immediately; unfinished markers remain until writing completes. We can configure the ``ModelCheckpoint`` via YAML or CLI:  .. code-block:: yaml@@ -51,6 +61,14 @@ We can configure the ``ModelCheckpoint`` via YAML or CLI: Resume Training --------------- +New NeMo checkpoints store resolved configuration values without dataclass schema metadata;+the live model configuration and hparams retain their schemas. Previously saved built-in decoding and+adapter configurations remain supported. For an older checkpoint containing an application-defined+schema, import that trusted schema in application code and call+``nemo.utils.serialization.register_checkpoint_schema_types(MyConfig)`` before loading it.+Checkpoint data does not select which schema classes to register.++ To auto-resume training, configure the ``exp_manager``. This feature is important for long training runs that might be interrupted or shut down before the procedure has completed. To auto-resume training, set the following parameters via YAML or CLI: @@ -177,26 +195,18 @@ Neptune  .. _exp_manager_neptune-label: -.. code-block:: yaml--    exp_manager:-        ...-        create_checkpoint_callback: True-        create_neptune_logger: false-        neptune_logger_kwargs:-            project: ${project}-            name: ${name}-            prefix: train-            log_model_checkpoints: false # set to True if checkpoints need to be pushed to Neptune-            tags: null # can specify as an array of strings in yaml array format-            description: null-            <Add any other arguments supported by Neptune logger here>+NeptuneLogger is no longer supported by Lightning 2.6. Keep ``create_neptune_logger: false``.+Setting it to ``true`` raises a configuration error before experiment files are created.+Use another supported Experiment Manager logger, or pass a supported logger directly to the Trainer.  Exponential Moving Average --------------------------  .. _exp_manager_ema-label: +EMA sidecar and optimizer state are restored when resuming training. Validation, testing,+and prediction with a checkpoint use its model weights without requiring an EMA sidecar.+ NeMo supports using exponential moving average (EMA) for model parameters. This can be useful for improving model generalization and stability. To use EMA, set the following parameters via YAML or :class:`~nemo.utils.exp_manager.ExpManagerConfig`. 

@@ -31,5 +31,15 @@     TransformerMultiHeadAttentionAdapter,     TransformerMultiHeadAttentionAdapterConfig, )+from nemo.utils.serialization import register_checkpoint_schema_types  # fmt: on+++register_checkpoint_schema_types(+    MultiHeadAttentionAdapterConfig,+    PositionalEncodingAdapterConfig,+    RelPositionalEncodingAdapterConfig,+    RelPositionMultiHeadAttentionAdapterConfig,+    TransformerMultiHeadAttentionAdapterConfig,+)

@@ -31,6 +31,7 @@ from nemo.collections.common.tokenizers.aggregate_tokenizer import DummyTokenizer from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec from nemo.utils import logging, logging_mode+from nemo.utils.serialization import register_checkpoint_schema_types   def move_dimension_to_the_front(tensor, dim_index):@@ -1389,3 +1390,6 @@ class CTCDecodingConfig: @dataclass class CTCBPEDecodingConfig(CTCDecodingConfig):     pass+++register_checkpoint_schema_types(CTCDecodingConfig, CTCBPEDecodingConfig)

@@ -34,6 +34,7 @@ from nemo.collections.common.tokenizers.aggregate_tokenizer import AggregateTokenizer from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec from nemo.utils import logging+from nemo.utils.serialization import register_checkpoint_schema_types   class AbstractMultiTaskDecoding(ConfidenceMixin):@@ -673,3 +674,6 @@ class MultiTaskDecodingConfig:      # if set to true, return attention scores; ignore them to save memory otherwise     return_xattn_scores: bool = False+++register_checkpoint_schema_types(MultiTaskDecodingConfig)

@@ -34,6 +34,7 @@ from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec from nemo.utils import logging from nemo.utils.enum import PrettyStrEnum+from nemo.utils.serialization import register_checkpoint_schema_types  try:     import kenlm@@ -1927,3 +1928,6 @@ class RNNTBPEDecodingConfig(RNNTDecodingConfig):     """      pass+++register_checkpoint_schema_types(RNNTDecodingConfig, RNNTBPEDecodingConfig)

@@ -22,6 +22,7 @@ import lightning.pytorch as pl import torch from lightning.pytorch import Callback+from lightning.pytorch.trainer.states import TrainerFn from lightning.pytorch.utilities.exceptions import MisconfigurationException from lightning.pytorch.utilities.rank_zero import rank_zero_info @@ -121,6 +122,9 @@ def save_original_optimizer_state(self, trainer: "pl.Trainer"):     def on_load_checkpoint(         self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]     ) -> None:+        # Evaluation restores callback state too, but never resumes optimizer/EMA state.+        if getattr(getattr(trainer, 'state', None), 'fn', TrainerFn.FITTING) != TrainerFn.FITTING:+            return         checkpoint_callback = trainer.checkpoint_callback          # Replace connector._ckpt_path with below to avoid calling into lightning's protected API@@ -137,7 +141,10 @@ def on_load_checkpoint(                 return             ema_path = ckpt_path.replace(ext, f'-EMA{ext}')             if os.path.exists(ema_path):-                ema_state_dict = torch.load(ema_path, map_location=torch.device('cpu'), weights_only=False)+                from nemo.utils.serialization import checkpoint_config_loading++                with checkpoint_config_loading():+                    ema_state_dict = torch.load(ema_path, map_location=torch.device('cpu'), weights_only=True)                  checkpoint['optimizer_states'] = ema_state_dict['optimizer_states']                 del ema_state_dict

@@ -22,6 +22,7 @@ from nemo.collections.common.parts.utils import activation_registry from nemo.core.classes.common import safe_instantiate from nemo.core.classes.mixins import access_mixins, adapter_mixin_strategies+from nemo.utils.serialization import register_checkpoint_schema_types   class AdapterModuleUtil(access_mixins.AccessMixin):@@ -170,3 +171,6 @@ class LinearAdapterConfig:         default_factory=lambda: adapter_mixin_strategies.ResidualAddAdapterStrategyConfig()     )     _target_: str = "{0}.{1}".format(LinearAdapter.__module__, LinearAdapter.__name__)+++register_checkpoint_schema_types(LinearAdapterConfig)

@@ -278,7 +278,7 @@ class AutomodelParallelStrategy(ModelParallelStrategy):     """      @override-    def load_checkpoint(self, checkpoint_path):+    def load_checkpoint(self, checkpoint_path, weights_only: Optional[bool] = None):         """Load DCP optimizer state while preserving lazy-state semantics.          Model tensors remain strict. Optimizer loading alone permits a complete@@ -290,7 +290,7 @@ def load_checkpoint(self, checkpoint_path):          path = Path(self.broadcast(checkpoint_path))         if not _is_sharded_checkpoint(path):-            return super().load_checkpoint(path)+            return super().load_checkpoint(path, weights_only=weights_only)          from torch.distributed.checkpoint import FileSystemReader, load         from torch.distributed.checkpoint.state_dict import (@@ -331,7 +331,7 @@ def load_checkpoint(self, checkpoint_path):         # of resume. Keep this metadata-only payload alive for consumers whose         # restored state can outlive that connector reference; model and         # optimizer tensors were loaded separately through DCP above.-        checkpoint = torch.load(path / _METADATA_FILENAME)+        checkpoint = torch.load(path / _METADATA_FILENAME, weights_only=weights_only)         self._checkpoint_keepalive = checkpoint         return checkpoint 

@@ -20,6 +20,7 @@ import torch  from nemo.core.classes.mixins import AccessMixin+from nemo.utils.serialization import register_checkpoint_schema_types  if TYPE_CHECKING:     from nemo.core.classes.mixins.adapter_mixins import AdapterModuleMixin@@ -263,3 +264,6 @@ class ResidualAddAdapterStrategyConfig:     _target_: str = "{0}.{1}".format(         ResidualAddAdapterStrategy.__module__, ResidualAddAdapterStrategy.__name__     )  # mandatory field+++register_checkpoint_schema_types(ResidualAddAdapterStrategyConfig)

@@ -1718,6 +1718,16 @@ def hparams(self):          return super().hparams +    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:+        """Store a plain configuration snapshot without changing live hparams metadata."""+        super().on_save_checkpoint(checkpoint)+        key = self.CHECKPOINT_HYPER_PARAMS_KEY+        hyper_parameters = checkpoint.get(key)+        if isinstance(hyper_parameters, (dict, DictConfig)) and 'cfg' in hyper_parameters:+            hyper_parameters = copy.deepcopy(hyper_parameters)+            hyper_parameters['cfg'] = OmegaConf.create(OmegaConf.to_container(self.cfg, resolve=True))+            checkpoint[key] = hyper_parameters+     @property     def validation_step_outputs(self):         """

@@ -75,6 +75,8 @@ class TrainerConfig:     enable_model_summary: bool = True     inference_mode: bool = True     barebones: bool = False+    model_registry: Optional[str] = None+    enable_autolog_hparams: bool = True   # Register the trainer config.

@@ -19,7 +19,6 @@  from lightning.pytorch.callbacks import Callback as PTLCallback from nemo.lightning.base_callback import BaseCallback-from nemo.lightning.one_logger_callback import OneLoggerNeMoCallback   class CallbackGroup:@@ -44,7 +43,8 @@ def get_instance(cls) -> 'CallbackGroup':         return cls._instance      def __init__(self) -> None:-        self._callbacks: List[BaseCallback] = [OneLoggerNeMoCallback()]+        # Leave telemetry detached until its callback integration is updated.+        self._callbacks: List[BaseCallback] = []         # Ensure application-end is emitted at most once per process         self._app_end_emitted: bool = False 

@@ -28,11 +28,20 @@ from nemo.utils.nemo_logging import Logger as _Logger from nemo.utils.nemo_logging import LogMode as logging_mode +try:+    from nemo.utils.serialization import register_checkpoint_config_types+except ModuleNotFoundError as error:+    if error.name != 'omegaconf':+        raise+else:+    register_checkpoint_config_types()+ logging = _Logger() try:-    from nemo.utils.lightning_logger_patch import add_memory_handlers_to_pl_logger+    from nemo.utils.lightning_logger_patch import add_memory_handlers_to_pl_logger, preserve_trainer_callback_defaults      add_memory_handlers_to_pl_logger()+    preserve_trainer_callback_defaults() except ModuleNotFoundError:     pass 

@@ -13,51 +13,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -import shutil from abc import ABC, abstractmethod from contextlib import contextmanager from time import time-from typing import Any, Dict, Optional, Union+from typing import Any, Dict, Optional  import lightning.pytorch as pl-import torch from lightning.fabric.plugins import CheckpointIO-from lightning.fabric.utilities.cloud_io import get_filesystem from lightning.fabric.utilities.types import _PATH from lightning.pytorch import Callback from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO  from nemo.utils import logging  try:-    from megatron.core import dist_checkpointing-    from megatron.core.dist_checkpointing.dict_utils import extract_matching_values-    from megatron.core.dist_checkpointing.mapping import ShardedBase-    from megatron.core.dist_checkpointing.serialization import (-        get_default_load_sharded_strategy,-        get_default_save_sharded_strategy,-    )-    from megatron.core.dist_checkpointing.strategies import tensorstore     from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest-    from megatron.core.dist_checkpointing.strategies.base import SaveShardedStrategy-    from megatron.core.dist_checkpointing.strategies.fully_parallel import (-        FullyParallelLoadStrategyWrapper,-        FullyParallelSaveStrategyWrapper,-    )-    from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy-    from megatron.core.dist_checkpointing.validation import StrictHandling-    from megatron.core.parallel_state import get_data_parallel_group      HAVE_MEGATRON_CORE = True  except (ImportError, ModuleNotFoundError) as e:      HAVE_MEGATRON_CORE = False-    IMPORT_ERROR = (-        "megatron-core was not found. "-        "Please see the NeMo README for installation instructions: https://github.com/NVIDIA-NeMo/Speech#megatron-gpt."-        f" Exact error: {e}"-    )+    IMPORT_ERROR = f"Async checkpoint saving requires the optional megatron-core async utilities. Exact error: {e}"   @contextmanager@@ -198,288 +175,3 @@ def _get_checkpoint_io(self, trainer) -> AsyncFinalizableCheckpointIO:                 f'Async finalizer requires an async compatible CheckpointIO, got: {checkpoint_io.__class__}'             )         return checkpoint_io---class DistributedCheckpointIO(AsyncCompatibleCheckpointIO):-    """CheckpointIO for a distributed checkpoint format.--    Args:-        save_ckpt_format (str): Distributed checkpoint format to use for checkpoint saving.-        load_directly_on_device (bool, optional): if True, loads the weights directly-            on GPU. Has effect only for `zarr` based checkpoints (PyT Distributed-            always loads on device). Defaults to True.-        load_strictness (StrictHandling, optional): defines loading strictness.-            If not None, overwrites the `strict` flag passed to `load_checkpoint`.-            Defaults to None.-        async_save (bool): whether to save asynchronously. Should be set to True if-            this class will be wrapped with AsyncFinalizableCheckpointIO.-        torch_dist_multiproc (int, optional): number of extra processes per rank-            used during ckpt save with PyTorch distributed format. Defaults, to None-            which means using an MCore default (2).-        parallel_save (bool): parallelizes the save across ranks. Defaults to True-        parallel_load (bool): parallelizes the load across ranks (followed by params all gather).-            Defaults to False due to some extra memory usage requirement.-    """--    def __init__(-        self,-        save_ckpt_format: str,-        load_directly_on_device: bool = True,-        load_strictness: Optional['StrictHandling'] = None,-        async_save: bool = False,-        torch_dist_multiproc: Optional[int] = None,-        assume_constant_structure: bool = False,-        parallel_save: bool = False,-        parallel_save_within_dp: bool = False,-        parallel_load: bool = False,-    ):-        super().__init__()-        if not HAVE_MEGATRON_CORE:-            raise ImportError(IMPORT_ERROR)--        self.save_ckpt_format = save_ckpt_format-        self.load_directly_on_device = load_directly_on_device-        self.load_strictness = load_strictness-        self.async_save = async_save-        self.torch_dist_multiproc = torch_dist_multiproc-        self.assume_constant_structure = assume_constant_structure-        self.parallel_save = parallel_save-        self.parallel_save_within_dp = parallel_save_within_dp-        self.parallel_load = parallel_load--        self._save_sharded_strategy = None-        self.validated_consistency = False--    @classmethod-    def from_config(cls, model_cfg: dict, async_save: bool = False):-        """Instantiates a DistributedCheckpointIO from a config dict.--        Args:-            model_cfg (dict): model config dict. Most of the configuration-                is extracted from this config.-            async_save (bool, optional): async_save flag is not part of the model config,-                it should be provided separately. Defaults to False.-        """-        return cls(-            save_ckpt_format=model_cfg.get('dist_ckpt_format', 'torch_dist'),-            load_directly_on_device=model_cfg.get('dist_ckpt_load_on_device', True),-            load_strictness=model_cfg.get('dist_ckpt_load_strictness', None),-            async_save=async_save,-            torch_dist_multiproc=model_cfg.get('dist_ckpt_torch_dist_multiproc', None),-            parallel_save=model_cfg.get('dist_ckpt_parallel_save', False),-            parallel_save_within_dp=model_cfg.get('dist_ckpt_parallel_save_within_dp', False),-            parallel_load=model_cfg.get('dist_ckpt_parallel_load', False),-        )--    @_debug_time('DistributedCheckpointIO.save_checkpoint')-    def save_checkpoint(-        self, checkpoint: Dict[str, Any], path: _PATH, storage_options: Optional[Any] = None-    ) -> Optional['AsyncRequest']:-        """Saves a distributed checkpoint. Creates the checkpoint root directory if doesn't exist.--        Args:-            checkpoint (Dict[str, Any]): sharded state dict to save-            path (_PATH): checkpoint directory-            storage_options (Any, optional): Optional parameters when saving the checkpoint-        """-        fs = get_filesystem(path)-        fs.makedirs(path, exist_ok=True)--        validate_sharding_integrity = not (self.validated_consistency and self.assume_constant_structure)-        self.validated_consistency = True--        rank = torch.distributed.get_rank()-        iteration = _get_iteration_from_checkpoint(checkpoint)-        start_time = time()-        async_save_request = dist_checkpointing.save(-            sharded_state_dict=checkpoint,-            checkpoint_dir=path,-            sharded_strategy=self.save_sharded_strategy,-            validate_access_integrity=validate_sharding_integrity,-            async_sharded_save=self.async_save,-        )-        end_time = time()-        log_parts = (-            "Global Checkpoint Save",-            f"Rank: {rank}",-            f"Iteration: {iteration}" if iteration is not None else None,-            f"Start time: {start_time:.3f}s",-            f"Save duration: {end_time - start_time:.3f}s",-        )-        log_message = " : ".join(part for part in log_parts if part is not None)-        logging.info(log_message)--        def iter_finalize_fn():-            logging.info(f'Successfully saved checkpoint from iteration {int(iteration):7d} to {path}')--        if self.async_save:-            assert async_save_request is not None-            async_save_request.add_finalize_fn(iter_finalize_fn)--        return async_save_request--    @_debug_time('DistributedCheckpointIO.load_checkpoint')-    def load_checkpoint(-        self,-        path: _PATH,-        map_location: Optional[Any] = None,-        sharded_state_dict: Dict[str, Any] = None,-        strict: Union[None, bool, 'StrictHandling'] = None,-        validate_access_integrity: Optional[bool] = True,-    ) -> Dict[str, Any]:-        """Loads a distributed checkpoint.--        Args:-            path (_PATH): checkpoint directory-            map_location (Any, optional): required to be None in this implementation-            sharded_state_dict (Dict[str, Any], optional): state dict which-                defines the loading procedure for the distributed checkpoint.-                Defaults to None to comply with the CheckpointIO interface,-                but it's a required argument.-            strict (bool, StrictHandling, optional): adjust load strictness. bool value-                is translated to StrictHandling instance. Gets overwritten by-                `self.load_strictness`. Defaults to None. If `self.load_strictness`-                is also None, strict becomes StrictHandling.ASSUME_OK_UNEXPECTED.--        Returns:-            Dist[str, Any]: loaded checkpoint.-        """-        if sharded_state_dict is None:-            raise ValueError('DistributedCheckpointIO requires passing sharded_state_dict argument to load_checkpoint')-        if map_location is not None:-            raise ValueError('DistributedCheckpointIO doesnt handle map_location argument')--        if self.save_ckpt_format == 'zarr' and self.load_directly_on_device:-            sharded_strategy = tensorstore.TensorStoreLoadShardedStrategy(load_directly_on_device=True)-        else:-            sharded_strategy = None--        if self.parallel_load:-            if sharded_strategy is None:-                sharded_strategy = get_default_load_sharded_strategy(path)-            sharded_strategy = FullyParallelLoadStrategyWrapper(-                sharded_strategy, get_data_parallel_group(with_context_parallel=True)-            )--        if sharded_strategy is not None:-            logging.info(f'Using {sharded_strategy} dist-ckpt load strategy.')--        if isinstance(strict, bool):-            # For backward-compatibility reasons and a bug in MCore (strict check not applied to factories)-            # we must apply a simple strict check here.-            if not strict:-                sharded_state_dict = self.adjust_non_strict_load(path, sharded_state_dict)-            strict = StrictHandling.ASSUME_OK_UNEXPECTED if strict else StrictHandling.LOG_ALL-        if self.load_strictness is not None:-            # Overwrites function argument-            strict = self.load_strictness-        if strict is None:-            # Default behavior-            strict = StrictHandling.ASSUME_OK_UNEXPECTED--        logging.debug(f'Dist ckpt load strictness: {strict}')--        start_time = time()-        ret = dist_checkpointing.load(-            sharded_state_dict=sharded_state_dict,-            checkpoint_dir=path,-            sharded_strategy=sharded_strategy,-            validate_access_integrity=validate_access_integrity,-            strict=strict,-        )-        end_time = time()-        duration = end_time - start_time-        logging.info(-            "Global Checkpoint Load : "-            f"Rank : {torch.distributed.get_rank()} : "-            f"Start time : {start_time:.3f}s : "-            f"Time spent in load_checkpoint: {duration:.3f}s"-        )-        return ret--    def adjust_non_strict_load(self, path: _PATH, sharded_state_dict: Dict[str, Any]):-        """Remove unexpected keys from being loaded into the state dict."""-        ckpt_sharded_metadata = dist_checkpointing.load_tensors_metadata(path)-        loaded_keys = []-        unexpected_keys = []--        def should_remove_missing_sharded_base(x: Any):-            if isinstance(x, ShardedBase):-                if x.key in ckpt_sharded_metadata:-                    loaded_keys.append(x.key)-                    return False-                else:-                    unexpected_keys.append(x.key)-                    return True-            return False--        _, sharded_state_dict = extract_matching_values(sharded_state_dict, should_remove_missing_sharded_base)-        logging.info(f'The following keys are not in the checkpoint and will not be loaded: {unexpected_keys}')--        # TODO: compute missing_keys by:-        #  1. all_gather_object of loaded_keys-        #  2. missing_keys = ckpt_sharded_metadata.keys() - loaded_keys-        return sharded_state_dict--    @_debug_time('DistributedCheckpointIO.remove_checkpoint')-    def remove_checkpoint(self, path: _PATH) -> None:-        """Remove a distributed checkpoint.--        Due to potentially large number of files, the implementation remove the whole directory at once.-        """-        shutil.rmtree(path, ignore_errors=True)--    @property-    def save_sharded_strategy(self) -> 'SaveShardedStrategy':-        """Conditionally initialize and get the sharded strategy to use for saving."""-        if self._save_sharded_strategy is None:-            self._save_sharded_strategy = self._determine_dist_ckpt_save_strategy()-        return self._save_sharded_strategy--    def _determine_dist_ckpt_save_strategy(self):-        """Determine the saving strategy based on constructor args.--        Relies on the default MCore strategy unless extra PyT Distributed format arguments-        are passed in config or in case of a fully parallel save in which case-        a parallelization wrapper is applied.-        """-        if self.save_ckpt_format == 'zarr':-            logging.warning(-                '`zarr` distributed checkpoint backend is deprecated.'-                ' Distributed optimizer checkpoint saving might be extremely slow.'-                ' Please switch to PyTorch Distributed format (model.dist_ckpt_format=torch_dist).'-            )--        if self.async_save and self.save_ckpt_format != 'torch_dist':-            raise ValueError('Async dist-ckpt save supported only for torch_dist format')--        torch_dist_kwargs = {} if self.torch_dist_multiproc is None else dict(thread_count=self.torch_dist_multiproc)-        if self.save_ckpt_format == 'torch_dist' and torch_dist_kwargs:-            save_strategy = TorchDistSaveShardedStrategy(self.save_ckpt_format, 1, **torch_dist_kwargs)-        else:-            save_strategy = get_default_save_sharded_strategy(self.save_ckpt_format, 1)--        # MCore v0.8 introduces `use_cached_ckpt_structure` attribute-        if hasattr(save_strategy, 'use_cached_ckpt_structure'):-            save_strategy.use_cached_ckpt_structure = self.assume_constant_structure--        if self.parallel_save:-            parallelization_group = (-                get_data_parallel_group(with_context_parallel=True) if self.parallel_save_within_dp else None-            )-            save_strategy = FullyParallelSaveStrategyWrapper(-                save_strategy, parallelization_group, self.assume_constant_structure-            )--        logging.info(f'Using {save_strategy} dist-ckpt save strategy.')-        return save_strategy---def _get_iteration_from_checkpoint(checkpoint: Dict[str, Any]) -> Optional[int]:-    return (-        checkpoint.get("loops", {})-        .get("fit_loop", {})-        .get("epoch_loop.batch_progress", {})-        .get("total", {})-        .get("completed", None)-    )

@@ -252,6 +252,28 @@ def on_save_checkpoint(self, trainer, pl_module, checkpoint):             get_filesystem(backup_path).rm(backup_path)         return output +    def on_validation_end(self, trainer, pl_module):+        # Lightning 2.6.0 stopped refreshing last when top-k does not improve; preserve NeMo's prior behavior.+        refresh_last = (+            self.save_last is True+            and not self._should_skip_saving_checkpoint(trainer)+            and not self._should_save_on_train_epoch_end(trainer)+        )+        super().on_validation_end(trainer, pl_module)+        if refresh_last and self._last_global_step_saved != trainer.global_step:+            self._save_last_checkpoint(trainer, self._monitor_candidates(trainer))++    def on_train_epoch_end(self, trainer, pl_module):+        # Lightning 2.6.0 stopped refreshing last when top-k does not improve; preserve NeMo's prior behavior.+        refresh_last = (+            self.save_last is True+            and not self._should_skip_saving_checkpoint(trainer)+            and self._should_save_on_train_epoch_end(trainer)+        )+        super().on_train_epoch_end(trainer, pl_module)+        if refresh_last and self._last_global_step_saved != trainer.global_step:+            self._save_last_checkpoint(trainer, self._monitor_candidates(trainer))+     def on_train_end(self, trainer, pl_module):         """         Save the checkpoint on train end.@@ -573,6 +595,10 @@ def _save_checkpoint(self, trainer: 'lightning.pytorch.Trainer', filepath: str)                 storage_options = None             logging.info(f'Checkpoint save for step {trainer.global_step} started at {time.time()}.')             trainer.save_checkpoint(filepath, self.save_weights_only, storage_options=storage_options)+            # Lightning uses these to schedule subsequent saves in this hook.+            # Async completion may happen several steps later and must not rewind them.+            self._last_global_step_saved = trainer.global_step+            self._last_checkpoint_saved = filepath             if self.async_save:                 logging.info(f'Scheduled async checkpoint save for {filepath}')             else:@@ -588,9 +614,6 @@ def _get_finalize_save_checkpoint_callback(          def _cb():             logging.debug(f'Finalize callback called for step {global_step}, filepath {filepath}')-            self._last_global_step_saved = global_step-            self._last_checkpoint_saved = filepath-             # notify loggers             if trainer.is_global_zero:                 for logger in trainer.loggers:

@@ -175,6 +175,7 @@ def _check_uploading_results_so_far(self):     def save_checkpoint(         self, checkpoint: Dict[str, Any], path: Union[str, Path], storage_options: Optional[Any] = None     ) -> None:+        """Serialize a checkpoint and upload it to the configured S3 path."""         # if we have a shared memory directory, we can serialize as a file to shared memory instead of as bytes.         if os.path.exists(SHARED_MEM_DIR):             localfile = self._serialize_checkpoint_to_shm(checkpoint, path)@@ -206,8 +207,12 @@ def save_checkpoint(                 _upload_bytes_to_s3(bytes, path, self.chunk_size_MB, self.max_write_concurrency)      def load_checkpoint(-        self, path: Union[str, Path], map_location: Optional[Callable] = lambda storage, loc: storage+        self,+        path: Union[str, Path],+        map_location: Optional[Callable] = lambda storage, loc: storage,+        weights_only: Optional[bool] = None,     ) -> Dict[str, Any]:+        """Download a checkpoint and forward device mapping and unpickling policy to PyTorch."""         if os.path.exists(SHARED_MEM_DIR):             with NamedTemporaryFile(dir=SHARED_MEM_DIR, delete=True) as tempfile:                 logging.info(@@ -220,33 +225,31 @@ def load_checkpoint(                     chunk_size_MB=self.chunk_size_MB,                     max_concurrency=self.max_read_concurrency,                 )-                checkpoint = torch.load(tempfile.name)+                checkpoint = torch.load(tempfile.name, map_location=map_location, weights_only=weights_only)         else:             file_stream: BytesIO = S3Utils.download_s3_file_to_stream(                 s3_path=path, chunk_size_MB=self.chunk_size_MB, max_concurrency=self.max_read_concurrency             )-            checkpoint = torch.load(file_stream)+            checkpoint = torch.load(file_stream, map_location=map_location, weights_only=weights_only)         return checkpoint      def remove_checkpoint(self, path: Union[str, Path]) -> None:+        """Remove an S3 checkpoint or delegate local removal to the base plugin."""         if S3Utils.is_s3_url(path):             S3Utils.remove_object(path)         else:             super().remove_checkpoint(path)      def teardown(self) -> None:-        # this ensure we wait for final checkpoint to finish uploading at train end.+        """Wait for uploads to finish and remove remaining temporary checkpoint files."""         rank = torch.distributed.get_rank()         if self.async_checkpointing:             logging.info(f'Entering teardown, waiting for all jobs to finish, rank {rank}')             start_time = time.perf_counter()             self._executor.shutdown(wait=True)             logging.info(f'executor shut down after {(time.perf_counter() - start_time):.2f} seconds, rank {rank}') -        '''-        This will be non-empty at the end of training with asynchronous uploading because the futures are not-        processed with _check_uploading_results_so_far. Therefore, check that the path exists before deleting.-        '''+        # Async upload futures may leave temporary files until teardown.         if self._temp_files:             for tfile in self._temp_files:                 if os.path.exists(tfile):

@@ -36,7 +36,7 @@ from lightning.pytorch.callbacks import Callback, ModelCheckpoint from lightning.pytorch.callbacks.early_stopping import EarlyStopping from lightning.pytorch.callbacks.timer import Interval, Timer-from lightning.pytorch.loggers import MLFlowLogger, NeptuneLogger, TensorBoardLogger, WandbLogger+from lightning.pytorch.loggers import MLFlowLogger, TensorBoardLogger, WandbLogger from lightning.pytorch.loops import _TrainingEpochLoop from lightning.pytorch.strategies.ddp import DDPStrategy from lightning.pytorch.trainer.connectors.checkpoint_connector import _CheckpointConnector@@ -841,6 +841,11 @@ def error_checks(trainer: 'lightning.pytorch.Trainer', cfg: Optional[Union[DictC         - Prints error messages when 1) run on multi-node and not Slurm, and             2) run on multi-gpu without DDP     """+    if cfg.create_neptune_logger:+        raise ValueError(+            'create_neptune_logger is unsupported by the required Lightning version. '+            'Use another exp_manager logger or pass a supported logger to the Trainer.'+        )     if HydraConfig.initialized() and get_original_cwd() != os.getcwd():         raise ValueError(             "Hydra changed the working directory. This interferes with ExpManger's functionality."@@ -1317,19 +1322,7 @@ def configure_loggers(         logging.info("ClearMLLogger has been set up")      if create_neptune_logger:-        if neptune_kwargs is None:-            neptune_kwargs = {}-        if "name" not in neptune_kwargs and "project" not in neptune_kwargs:-            raise ValueError("name and project are required for neptune_logger")-        if "api_key" not in neptune_kwargs and not os.getenv("NEPTUNE_API_TOKEN", None):-            raise ValueError(-                "either api_key should be set in neptune_kwargs or NEPTUNE_API_TOKEN should "-                "be set in environment variable for neptune_logger"-            )-        neptune_logger = NeptuneLogger(**neptune_kwargs)--        logger_list.append(neptune_logger)-        logging.info("NeptuneLogger has been set up")+        raise ValueError('create_neptune_logger is unsupported by the required Lightning version.')      trainer._logger_connector.configure_logger(logger_list) @@ -1340,13 +1333,13 @@ class NeMoCheckpointConnector(_CheckpointConnector):     distributed training settings to pre-load checkpoint.     """ -    def resume_start(self, checkpoint_path=None) -> None:+    def resume_start(self, checkpoint_path=None, weights_only: Optional[bool] = None) -> None:         """resume_start"""         checkpoint_path = self.trainer.ckpt_path         if checkpoint_path is not None:             logging.info(f'Resuming from checkpoint {checkpoint_path}, rank {torch.distributed.get_rank()}')         start_time = time.perf_counter()-        super().resume_start(checkpoint_path)+        super().resume_start(checkpoint_path, weights_only=weights_only)         if checkpoint_path is not None:             logging.info(                 'Time elapsed loading checkpoint/optimizer states: '

@@ -14,6 +14,7 @@ # limitations under the License.  import logging as _logging+from functools import wraps from logging.handlers import MemoryHandler  import lightning.pytorch as pl@@ -57,3 +58,56 @@ def add_filehandlers_to_pl_logger(all_log_file, err_log_file):         HANDLERS["memory_err"].setTarget(HANDLERS["file_err"])         HANDLERS["memory_err"].close()         del HANDLERS["memory_err"]+++def preserve_trainer_callback_defaults():+    """Keep pre-2.6 callback selection for Trainers in a NeMo process.++    NeMo entry points also construct Lightning Trainers directly, so configure+    defaults at the existing Lightning integration boundary. Explicit callbacks,+    disable flags and barebones mode retain Lightning's validation and behavior.+    """+    from importlib import import_module++    for namespace in ('lightning.pytorch', 'pytorch_lightning'):+        try:+            connector = import_module(f'{namespace}.trainer.connectors.callback_connector')+            callbacks = import_module(f'{namespace}.callbacks')+        except ModuleNotFoundError as error:+            if error.name != namespace.split('.')[0]:+                raise+            continue+        _preserve_callback_selection(connector._CallbackConnector, callbacks)+++def _preserve_callback_selection(connector, callbacks):+    """Install an idempotent adapter around one Lightning callback connector."""+    if getattr(connector, '_nemo_callback_defaults_preserved', False):+        return++    configure_progress = connector._configure_progress_bar+    configure_summary = connector._configure_model_summary_callback++    @wraps(configure_progress)+    def progress_bar(self, enable_progress_bar=True):+        """Supply the previous default while retaining upstream flag validation."""+        if enable_progress_bar and not any(isinstance(cb, callbacks.ProgressBar) for cb in self.trainer.callbacks):+            self.trainer.callbacks.append(callbacks.TQDMProgressBar())+        return configure_progress(self, enable_progress_bar)++    @wraps(configure_summary)+    def model_summary(self, enable_model_summary):+        """Match the old summary selection, including explicit Rich progress bars."""+        if enable_model_summary and not any(isinstance(cb, callbacks.ModelSummary) for cb in self.trainer.callbacks):+            summary_type = (+                callbacks.RichModelSummary+                if isinstance(self.trainer.progress_bar_callback, callbacks.RichProgressBar)+                else callbacks.ModelSummary+            )+            self.trainer.callbacks.append(summary_type())+            return None+        return configure_summary(self, enable_model_summary)++    connector._configure_progress_bar = progress_bar+    connector._configure_model_summary_callback = model_summary+    connector._nemo_callback_defaults_preserved = True

@@ -0,0 +1,125 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0+"""Checkpoint configuration types supported by the weights-only loader."""++from collections import defaultdict+from contextlib import contextmanager+from dataclasses import is_dataclass+from enum import Enum+from pathlib import Path, PosixPath, WindowsPath+from pickle import UnpicklingError+from typing import Any, Dict, List, Set, Tuple, Union, get_args, get_origin, get_type_hints++import torch+from omegaconf import DictConfig, ListConfig+from omegaconf.base import ContainerMetadata, Metadata, UnionNode+from omegaconf.nodes import (+    AnyNode,+    BooleanNode,+    BytesNode,+    EnumNode,+    FloatNode,+    IntegerNode,+    InterpolationResultNode,+    PathNode,+    StringNode,+    ValueNode,+)+++def checkpoint_config_types() -> list:+    """Return the fixed data types and restricted annotation loader used by checkpoints."""+    return [+        DictConfig,+        ListConfig,+        ContainerMetadata,+        Metadata,+        UnionNode,+        AnyNode,+        BooleanNode,+        BytesNode,+        EnumNode,+        FloatNode,+        IntegerNode,+        InterpolationResultNode,+        PathNode,+        StringNode,+        ValueNode,+        Path,+        PosixPath,+        WindowsPath,+        Any,+        Dict,+        List,+        Set,+        Tuple,+        Union,+        # Typing aliases pickle through _operator.getitem. The replacement below+        # accepts only fixed typing origins and type arguments, never arbitrary+        # object subscription (including OmegaConf interpolation resolution).+        (_restore_checkpoint_annotation, "_operator.getitem"),+        defaultdict,+        dict,+        list,+        tuple,+        int,+        float,+        str,+        bool,+        bytes,+    ]+++@contextmanager+def checkpoint_config_loading():+    """Add scoped configuration types while retaining existing process registrations."""+    existing = set(torch.serialization.get_safe_globals())+    additions = [entry for entry in checkpoint_config_types() if entry not in existing]+    with torch.serialization.safe_globals(additions):+        yield+++def register_checkpoint_config_types() -> None:+    """Register the fixed configuration container types used by Lightning checkpoints."""+    torch.serialization.add_safe_globals(checkpoint_config_types())+++def register_checkpoint_schema_types(*schemas: type) -> None:+    """Register source-defined dataclass schemas, nested schemas, and annotated enums.++    Call from the schema's module after its definitions. Only trusted source types+    are traversed; checkpoint contents never select imports or additional globals.+    """+    pending = list(schemas)+    registered = set()+    while pending:+        schema = pending.pop()+        arguments = get_args(schema)+        if arguments:+            pending.extend(arguments)+        elif isinstance(schema, type) and schema not in registered:+            if is_dataclass(schema):+                registered.add(schema)+                pending.extend(get_type_hints(schema).values())+            elif issubclass(schema, Enum):+                registered.add(schema)+    torch.serialization.add_safe_globals(list(registered))+++def _restore_checkpoint_annotation(origin, arguments):+    """Rebuild typing annotations without permitting general object item access."""+    if not any(origin is candidate for candidate in (Dict, List, Set, Tuple, Union)):+        raise UnpicklingError('Checkpoint item access is restricted to typing annotations.')+    items = arguments if isinstance(arguments, tuple) else (arguments,)+    if not all(_is_checkpoint_annotation(item) for item in items):+        raise UnpicklingError('Checkpoint annotation arguments must be types.')+    return origin[arguments]+++def _is_checkpoint_annotation(value):+    if value is None or value is Ellipsis or isinstance(value, type):+        return True+    origin = get_origin(value)+    return any(origin is candidate for candidate in (dict, list, set, tuple, Union)) and all(+        _is_checkpoint_annotation(argument) for argument in get_args(value)+    )

@@ -79,9 +79,10 @@ classifiers = [  [project.optional-dependencies] core = [-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",@@ -92,9 +93,10 @@ core = [ ]  lightning = [-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",@@ -147,9 +149,10 @@ tts = [     "whisper_normalizer",     "datasets>=3.2.0",     "sentencepiece<1.0.0",-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",@@ -171,9 +174,10 @@ audio = [     "datasets>=3.2.0",     "pandas",     "sentencepiece<1.0.0",-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",@@ -197,9 +201,10 @@ vllm = [ ]  all = [-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",@@ -281,9 +286,10 @@ common = [     "einops",     "pandas",     "sentencepiece<1.0.0",-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",@@ -307,9 +313,10 @@ asr = [     "datasets>=3.2.0",     "pandas",     "sentencepiece<1.0.0",-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",@@ -336,9 +343,10 @@ speechlm2 = [     "datasets>=3.2.0",     "pandas",     "sentencepiece<1.0.0",-    "hydra-core>1.3,<=1.3.2",-    "lightning>2.2.1,<=2.4.0",-    "omegaconf<=2.3",+    "hydra-core>=1.3.8,<1.4",+    "lightning>=2.6.6,<2.7",+    "pytorch-lightning>=2.6.6,<2.7",+    "omegaconf>=2.3,<2.4",     "torchmetrics>=0.11.0",     "transformers",     "wandb",

@@ -52,7 +52,8 @@ def test_default_init(self):         assert strategy.activation_checkpointing_llm is False         assert strategy.activation_checkpointing_perception is False -    def test_sharded_checkpoint_metadata_is_retained_after_return(self, monkeypatch, tmp_path):+    @pytest.mark.parametrize('weights_only', [None, False, True])+    def test_sharded_checkpoint_metadata_is_retained_after_return(self, monkeypatch, tmp_path, weights_only):         class CheckpointMetadata:             pass @@ -77,9 +78,15 @@ def read_metadata(self):         monkeypatch.setattr(dcp, "FileSystemReader", Reader)         monkeypatch.setattr(dcp, "load", lambda state, checkpoint_id: None)         monkeypatch.setattr(dcp_state_dict, "get_model_state_dict", lambda model: {})-        monkeypatch.setattr("torch.load", lambda path: CheckpointMetadata()) -        checkpoint = strategy.load_checkpoint(tmp_path)+        def load_metadata(path, *, weights_only=None):+            assert weights_only is policy+            return CheckpointMetadata()++        policy = weights_only+        monkeypatch.setattr("torch.load", load_metadata)++        checkpoint = strategy.load_checkpoint(tmp_path, weights_only=weights_only)         checkpoint_ref = weakref.ref(checkpoint)         del checkpoint         gc.collect()

@@ -271,11 +271,12 @@ def test_expand_ids_with_audio_matches_unpad_behavior_for_all_padding_row():     assert expanded.tolist() == [[0, 0], [11, 12]]  -def test_device_falls_back_before_draft_configuration():-    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})+@pytest.mark.parametrize("device", ["cpu", "meta"])+def test_device_falls_back_before_draft_configuration(device):+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1, device=device), {"dflash": {"mask_token_id": 18}}).to(device)      assert module.draft_model is None-    assert module.device == torch.device("cpu")+    assert module.device == next(module.target.parameters()).device   class _BatchTarget(nn.Module):

@@ -0,0 +1,336 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0+"""Compatibility coverage for dependency imports and checkpoint configuration."""++import io+import operator+import pickle+import subprocess+import sys+from dataclasses import dataclass+from importlib import import_module+from unittest.mock import patch++import pytest+import torch+from omegaconf import OmegaConf+++@pytest.mark.unit+@pytest.mark.parametrize('module', ['nemo.core', 'nemo.utils.exp_manager'])+def test_dependency_imports(module):+    result = subprocess.run([sys.executable, '-c', f'import {module}'], capture_output=True, text=True, timeout=60)+    assert result.returncode == 0, result.stderr+++@pytest.mark.unit+@pytest.mark.parametrize('namespace', ['lightning.pytorch', 'pytorch_lightning'])+@pytest.mark.parametrize('schema', ['untyped', 'ctc', 'ctc_bpe', 'rnnt', 'rnnt_bpe'])+def test_checkpoint_configuration_loads_with_weights_only(tmp_path, namespace, schema):+    lightning = import_module(namespace)++    class Minimal(lightning.LightningModule):+        def __init__(self, cfg):+            super().__init__()+            self.cfg = cfg+            self.weight = torch.nn.Parameter(torch.zeros(2))++    config = OmegaConf.create({'name': 'model', 'layers': [2, 4], 'enabled': True, 'dropout': 0.1, 'optional': None})+    if schema != 'untyped':+        from nemo.collections.asr.parts.submodules.ctc_decoding import CTCBPEDecodingConfig, CTCDecodingConfig+        from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTBPEDecodingConfig, RNNTDecodingConfig++        schemas = dict(+            ctc=CTCDecodingConfig,+            ctc_bpe=CTCBPEDecodingConfig,+            rnnt=RNNTDecodingConfig,+            rnnt_bpe=RNNTBPEDecodingConfig,+        )+        config.decoding = OmegaConf.structured(schemas[schema])+    path = tmp_path / 'model.ckpt'+    torch.save(+        {+            'pytorch-lightning_version': lightning.__version__,+            'hyper_parameters': {'cfg': config},+            'state_dict': {'weight': torch.ones(2)},+        },+        path,+    )+    import nemo.utils  # noqa: F401++    restored = Minimal.load_from_checkpoint(path, weights_only=True)+    assert OmegaConf.to_container(restored.cfg) == OmegaConf.to_container(config)+    assert torch.equal(restored.weight, torch.ones(2))+++@pytest.mark.unit+def test_weights_only_still_rejects_unregistered_objects():+    stream = io.BytesIO()+    torch.save(_UnregisteredObject(), stream)+    stream.seek(0)+    import nemo.utils  # noqa: F401++    with pytest.raises(pickle.UnpicklingError):+        torch.load(stream, weights_only=True)+++@pytest.mark.unit+@pytest.mark.parametrize('namespace', ['lightning.pytorch', 'pytorch_lightning'])+def test_checkpoint_instantiators_keep_the_lightning_policy(tmp_path, namespace):+    lightning = import_module(namespace)+    cli = import_module(f'{namespace}.cli')++    class Minimal(lightning.LightningModule):+        pass++    checkpoint = {+        'state_dict': {},+        'hyper_parameters': {'_instantiator': 'builtins.print'},+        'pytorch-lightning_version': lightning.__version__,+    }+    path = tmp_path / 'model.ckpt'+    torch.save(checkpoint, path)+    with patch('builtins.print') as instantiate:+        with pytest.raises(ValueError, match='allowlist'):+            Minimal.load_from_checkpoint(path, weights_only=True)+        instantiate.assert_not_called()+    checkpoint['hyper_parameters']['_instantiator'] = f'{namespace}.cli.instantiate_module'+    torch.save(checkpoint, path)+    with patch.object(cli, 'instantiate_module', side_effect=lambda cls, kwargs: cls()) as instantiate:+        assert isinstance(Minimal.load_from_checkpoint(path, weights_only=True), Minimal)+        instantiate.assert_called_once_with(Minimal, {})+++@pytest.mark.unit+def test_default_lifecycle_callbacks_do_not_import_telemetry():+    script = r"""+import importlib.abc+import sys++class BlockTelemetry(importlib.abc.MetaPathFinder):+    def find_spec(self, fullname, path=None, target=None):+        if fullname.startswith('nv_one_logger'):+            raise AssertionError('telemetry imported')++sys.meta_path.insert(0, BlockTelemetry())+import nemo.core+import nemo.utils.exp_manager+from nemo.lightning.base_callback import BaseCallback+from nemo.lightning.callback_group import CallbackGroup++class CustomCallback(BaseCallback):+    calls = 0+    def on_model_init_start(self):+        self.calls += 1++group = CallbackGroup()+assert group.callbacks == []+callback = CustomCallback()+group.register(callback)+group.on_model_init_start()+assert callback.calls == 1+assert not any(name.startswith('nv_one_logger') for name in sys.modules)+"""+    result = subprocess.run([sys.executable, '-c', script], capture_output=True, text=True, timeout=60)+    assert result.returncode == 0, result.stdout + result.stderr+++@pytest.mark.unit+def test_ctc_model_default_configuration_checkpoint(tmp_path):+    from nemo.collections.asr.models import EncDecCTCModel+    from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig++    cfg = OmegaConf.create(+        {+            'preprocessor': {+                '_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',+                'features': 4,+            },+            'encoder': {+                '_target_': 'nemo.collections.asr.modules.ConvASREncoder',+                'feat_in': 4,+                'activation': 'relu',+                'jasper': [+                    {+                        'filters': 4,+                        'repeat': 1,+                        'kernel': [1],+                        'stride': [1],+                        'dilation': [1],+                        'dropout': 0.0,+                        'residual': False,+                    }+                ],+            },+            'decoder': {+                '_target_': 'nemo.collections.asr.modules.ConvASRDecoder',+                'feat_in': 4,+                'num_classes': 2,+                'vocabulary': ['a', 'b'],+            },+        }+    )+    assert 'decoding' not in cfg+    model = EncDecCTCModel(cfg)+    assert OmegaConf.get_type(model.cfg.decoding) is CTCDecodingConfig+    path = tmp_path / 'ctc.ckpt'+    torch.save(+        {+            'state_dict': model.state_dict(),+            'hyper_parameters': {'cfg': model.cfg},+            'pytorch-lightning_version': import_module('lightning.pytorch').__version__,+        },+        path,+    )+    restored = EncDecCTCModel.load_from_checkpoint(path, map_location='cpu', weights_only=True)+    assert OmegaConf.to_container(restored.cfg, resolve=True, enum_to_str=True) == OmegaConf.to_container(+        model.cfg, resolve=True, enum_to_str=True+    )+    for name, value in model.state_dict().items():+        torch.testing.assert_close(restored.state_dict()[name], value)+++@pytest.mark.unit+def test_checkpoint_annotations_do_not_resolve_config_values(tmp_path):+    calls = []+    OmegaConf.register_new_resolver('checkpoint_probe', lambda: calls.append('resolved'))+    try:+        cfg = OmegaConf.create({'value': '${checkpoint_probe:}'})+        path = tmp_path / 'annotation.ckpt'+        torch.save(_ItemAccess(cfg), path)+        with pytest.raises(pickle.UnpicklingError, match='typing annotations'):+            torch.load(path, weights_only=True)+        assert calls == []+    finally:+        OmegaConf.clear_resolver('checkpoint_probe')+++@pytest.mark.unit+def test_unregistered_structured_metadata_is_rejected(tmp_path):+    path = tmp_path / 'unknown.ckpt'+    torch.save(OmegaConf.structured(_UnregisteredSchema), path)+    with pytest.raises(pickle.UnpicklingError):+        torch.load(path, weights_only=True)+++@pytest.mark.unit+def test_configuration_registration_keeps_asr_imports_lazy():+    script = "import sys; import nemo.utils; assert not any(n.startswith('nemo.collections.asr') for n in sys.modules)"+    result = subprocess.run([sys.executable, '-c', script], capture_output=True, text=True, timeout=60)+    assert result.returncode == 0, result.stderr+++@dataclass+class _UnregisteredSchema:+    name: str = 'model'+++class _ItemAccess:+    def __init__(self, cfg):+        self.cfg = cfg++    def __reduce__(self):+        return operator.getitem, (self.cfg, 'value')+++class _UnregisteredObject:+    pass+++@pytest.mark.unit+@pytest.mark.parametrize('schema', ['aed', 'adapter', 'path'])+def test_additional_legacy_checkpoint_configurations(tmp_path, schema):+    from pathlib import Path++    from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecodingConfig+    from nemo.collections.common.parts.adapter_modules import LinearAdapterConfig++    configs = {+        'aed': OmegaConf.structured(MultiTaskDecodingConfig),+        'adapter': OmegaConf.structured(LinearAdapterConfig(in_features=50, dim=8)),+        'path': OmegaConf.create({'path': Path('models/checkpoint')}),+    }+    path = tmp_path / 'legacy.ckpt'+    torch.save(configs[schema], path)+    restored = torch.load(path, weights_only=True)+    assert OmegaConf.to_container(restored) == OmegaConf.to_container(configs[schema])+++@pytest.mark.unit+def test_aed_default_config_checkpoint(tmp_path, test_data_dir):+    from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecodingConfig+    from tests.collections.asr.test_asr_multitask_model_bpe import asr_model++    initial = asr_model.__wrapped__(test_data_dir)+    cfg = OmegaConf.create(OmegaConf.to_container(initial.cfg))+    del cfg.decoding+    model = type(initial)(cfg)+    assert OmegaConf.get_type(model.cfg.decoding) is MultiTaskDecodingConfig+    _assert_model_checkpoint_roundtrip(model, tmp_path / 'aed.ckpt')+++@pytest.mark.unit+def test_adapter_model_checkpoint_preserves_live_schema(tmp_path):+    from nemo.collections.common.parts.adapter_modules import LinearAdapterConfig+    from tests.core.mixins.adapters.test_adapter_model_mixin import DefaultAdapterModel, get_model_config++    model = DefaultAdapterModel(get_model_config())+    model.add_adapter('linear', LinearAdapterConfig(in_features=50, dim=8))+    assert model.is_adapter_available()+    assert OmegaConf.get_type(model.hparams.cfg.adapters.linear) is LinearAdapterConfig+    _assert_model_checkpoint_roundtrip(model, tmp_path / 'adapter.ckpt')+++@pytest.mark.unit+def test_application_schema_is_preserved_in_model_hparams(tmp_path):+    from omegaconf import open_dict++    from tests.core.test_save_restore import MockModel++    model = MockModel(OmegaConf.create({}))+    with open_dict(model.cfg):+        model.cfg.custom = OmegaConf.structured(_UnregisteredSchema)+        model.cfg.alias = '${custom.name}'+    assert OmegaConf.get_type(model.cfg.custom) is _UnregisteredSchema+    assert OmegaConf.get_type(model.hparams.cfg.custom) is _UnregisteredSchema+    assert OmegaConf.to_container(model.cfg, resolve=False)['alias'] == '${custom.name}'+    _assert_model_checkpoint_roundtrip(model, tmp_path / 'application.ckpt')+++def _assert_model_checkpoint_roundtrip(model, path):+    lightning = import_module('lightning.pytorch')+    trainer = lightning.Trainer(+        default_root_dir=path.parent,+        accelerator='cpu',+        logger=False,+        enable_checkpointing=False,+        enable_progress_bar=False,+        enable_model_summary=False,+    )+    model.set_trainer(trainer)+    trainer.strategy.connect(model)+    trainer.save_checkpoint(path, weights_only=True)+    saved = torch.load(path, map_location='cpu', weights_only=True)+    assert OmegaConf.get_type(saved['hyper_parameters']['cfg']) is dict+    restored = type(model).load_from_checkpoint(path, map_location='cpu', weights_only=True)+    assert OmegaConf.to_container(restored.cfg, resolve=True, enum_to_str=True) == OmegaConf.to_container(+        model.cfg, resolve=True, enum_to_str=True+    )+    for name, value in model.state_dict().items():+        torch.testing.assert_close(restored.state_dict()[name], value)+++@pytest.mark.unit+def test_checkpoint_snapshot_preserves_supported_enum_and_path_values(tmp_path):+    from pathlib import Path++    from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig  # noqa: F401+    from nemo.collections.asr.parts.utils.batched_beam_decoding_utils import PruningMode+    from tests.core.test_save_restore import MockModel++    model = MockModel(OmegaConf.create({'mode': PruningMode.EARLY, 'path': Path('models/checkpoint')}))+    path = tmp_path / 'values.ckpt'+    _assert_model_checkpoint_roundtrip(model, path)+    restored = MockModel.load_from_checkpoint(path, map_location='cpu', weights_only=True)+    assert restored.cfg.mode is PruningMode.EARLY+    assert restored.cfg.path == model.cfg.path == Path('models/checkpoint')

@@ -210,34 +210,19 @@ def test_trainer_loggers(self, tmp_path):         assert isinstance(test_trainer.logger, pl.loggers.WandbLogger)      @pytest.mark.unit-    def test_trainer_neptune_logger(self, tmp_path):-        pytest.importorskip("neptune", reason="could not import `neptune`, use `pip install neptune` to run this test")--        test_trainer = pl.Trainer(accelerator='cpu', logger=False)-        # Check that a create_neptune_logger=True errors out unless neptune_logger_kwargs is passed.-        with pytest.raises(ValueError):-            _ = exp_manager(-                test_trainer,+    @pytest.mark.parametrize('kwargs', [None, {"name": "test", "project": "test", "api_key": "unused"}])+    def test_trainer_neptune_logger(self, tmp_path, kwargs):+        trainer = pl.Trainer(accelerator='cpu', logger=False, enable_checkpointing=False)+        with pytest.raises(ValueError, match='create_neptune_logger.*Lightning'):+            exp_manager(+                trainer,                 {-                    "create_tensorboard_logger": False,-                    "create_checkpoint_callback": False,-                    "exp_dir": str(tmp_path),+                    "explicit_log_dir": str(tmp_path / 'experiment'),                     "create_neptune_logger": True,+                    "neptune_logger_kwargs": kwargs,                 },             )-        # Check that a NeptuneLogger is attached to logger if create_neptune_logger=True and neptune_logger_kwargs has name-        # and project-        _ = exp_manager(-            test_trainer,-            {-                "create_tensorboard_logger": False,-                "create_checkpoint_callback": False,-                "exp_dir": str(tmp_path),-                "create_neptune_logger": True,-                "neptune_logger_kwargs": {"name": "", "project": "", "api_key": ""},-            },-        )-        assert isinstance(test_trainer.logger, pl.loggers.NeptuneLogger)+        assert not (tmp_path / 'experiment').exists()      @pytest.mark.unit     def test_checkpoint_configurations(self):

@@ -0,0 +1,294 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0+"""Resume and callback compatibility with the supported Lightning release."""++import copy+import io+from pathlib import Path+from types import SimpleNamespace+from unittest.mock import Mock, patch++import lightning.pytorch as pl+import pytest+import torch+from lightning.fabric.plugins import TorchCheckpointIO+from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO+from lightning.pytorch.strategies import ModelParallelStrategy+from lightning.pytorch.trainer.connectors.checkpoint_connector import _CheckpointConnector+from lightning.pytorch.trainer.states import TrainerFn+from torch.utils.data import DataLoader++from nemo.collections.common.callbacks import EMA+from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy+from nemo.utils.callbacks.dist_ckpt_io import AsyncFinalizableCheckpointIO+from nemo.utils.callbacks.nemo_model_checkpoint import NeMoModelCheckpoint+from nemo.utils.exp_manager import NeMoCheckpointConnector, exp_manager+++@pytest.mark.unit+@pytest.mark.parametrize('async_save', [False, True])+@pytest.mark.parametrize('at_epoch_end', [False, True])+def test_last_checkpoint_advances_when_metric_plateaus(tmp_path, async_save, at_epoch_end):+    checkpoint = NeMoModelCheckpoint(+        dirpath=tmp_path,+        filename='{step}',+        monitor='val_loss',+        save_top_k=1,+        save_last=True,+        save_on_train_epoch_end=at_epoch_end,+        async_save=async_save,+        save_nemo_on_train_end=False,+    )+    deferred = _DeferredCheckpointIO()+    trainer = pl.Trainer(+        default_root_dir=tmp_path,+        accelerator='cpu',+        devices=1,+        max_epochs=2,+        callbacks=[checkpoint],+        plugins=[deferred] if async_save else None,+        logger=False,+        enable_progress_bar=False,+        enable_model_summary=False,+        num_sanity_val_steps=0,+    )+    trainer.fit(_PlateauModel())+    if async_save:+        # Scheduling updates Lightning bookkeeping; completion must not rewind it.+        assert checkpoint._last_global_step_saved == trainer.global_step+        deferred.flush()+        assert checkpoint._last_global_step_saved == trainer.global_step+    saved = torch.load(checkpoint.last_model_path, weights_only=True)+    assert saved['global_step'] == trainer.global_step == 8+    assert torch.load(checkpoint.best_model_path, weights_only=True)['global_step'] == 4+++@pytest.mark.unit+@pytest.mark.parametrize('weights_only', [None, False, True])+def test_automodel_checkpoint_forwards_load_policy(tmp_path, weights_only):+    strategy = AutomodelParallelStrategy()+    strategy.broadcast = lambda path: path+    with patch.object(ModelParallelStrategy, 'load_checkpoint', return_value={'global_step': 7}) as load:+        assert strategy.load_checkpoint(tmp_path / 'model.ckpt', weights_only=weights_only) == {'global_step': 7}+    load.assert_called_once_with(tmp_path / 'model.ckpt', weights_only=weights_only)+++@pytest.mark.unit+def test_resume_connector_forwards_load_policy():+    trainer = SimpleNamespace(ckpt_path=None)+    connector = NeMoCheckpointConnector(trainer)+    with patch.object(_CheckpointConnector, 'resume_start') as resume:+        connector.resume_start(weights_only=True)+    resume.assert_called_once_with(None, weights_only=True)+++@pytest.mark.unit+@pytest.mark.parametrize('shared_memory', [False, True])+@pytest.mark.parametrize('weights_only', [False, True])+def test_s3_checkpoint_load_policy_and_map_location(tmp_path, monkeypatch, shared_memory, weights_only):+    pytest.importorskip('boto3')+    from nemo.utils.callbacks import s3_checkpoint_io as module++    path = tmp_path / 'model.ckpt'+    torch.save({'tensor': torch.ones(2)}, path)+    monkeypatch.setattr(module, 'SHARED_MEM_DIR', str(tmp_path if shared_memory else tmp_path / 'missing'))+    monkeypatch.setattr(torch.distributed, 'get_rank', lambda: 0)+    monkeypatch.setattr(module.S3Utils, 'download_s3_file_to_stream', lambda **kw: io.BytesIO(path.read_bytes()))+    monkeypatch.setattr(+        module.S3Utils, 'download_s3_file_to_path', lambda **kw: Path(kw['file_path']).write_bytes(path.read_bytes())+    )+    checkpoint_io = module.S3CheckpointIO('s3://bucket')+    map_location = Mock(side_effect=lambda storage, loc: storage)+    with patch.object(module.torch, 'load', wraps=torch.load) as load:+        loaded = checkpoint_io.load_checkpoint(+            's3://bucket/model.ckpt', map_location=map_location, weights_only=weights_only+        )+    assert torch.equal(loaded['tensor'], torch.ones(2))+    assert load.call_args.kwargs['weights_only'] is weights_only+    assert load.call_args.kwargs['map_location'] is map_location+    map_location.assert_called()+++@pytest.mark.unit+@pytest.mark.parametrize('stage', [TrainerFn.VALIDATING, TrainerFn.TESTING, TrainerFn.PREDICTING])+def test_ema_evaluation_without_sidecar(tmp_path, stage):+    trainer = pl.Trainer(+        default_root_dir=tmp_path,+        accelerator='cpu',+        logger=False,+        enable_progress_bar=False,+        enable_model_summary=False,+        callbacks=[NeMoModelCheckpoint(dirpath=tmp_path)],+    )+    trainer.state.fn = stage+    trainer.ckpt_path = str(tmp_path / 'model.ckpt')+    checkpoint = {'optimizer_states': ['unchanged']}+    EMA(decay=0.9).on_load_checkpoint(trainer, pl.LightningModule(), checkpoint)+    assert checkpoint['optimizer_states'] == ['unchanged']+++@pytest.mark.unit+def test_neptune_rejected_before_creating_experiment(tmp_path):+    trainer = pl.Trainer(default_root_dir=tmp_path, accelerator='cpu', logger=False, enable_checkpointing=False)+    directory = tmp_path / 'experiment'+    with pytest.raises(ValueError, match='create_neptune_logger.*Lightning'):+        exp_manager(trainer, {'explicit_log_dir': str(directory), 'create_neptune_logger': True})+    assert not directory.exists()+++class _PlateauModel(pl.LightningModule):+    def __init__(self):+        super().__init__()+        self.weight = torch.nn.Parameter(torch.ones(()))++    def training_step(self, batch, batch_idx):+        return self.weight.square()++    def validation_step(self, batch, batch_idx):+        self.log('val_loss', 1.0)++    def configure_optimizers(self):+        return torch.optim.SGD(self.parameters(), lr=0.1)++    def train_dataloader(self):+        return DataLoader(torch.ones(4), batch_size=1)++    def val_dataloader(self):+        return DataLoader(torch.ones(1), batch_size=1)+++class _DeferredCheckpointIO(AsyncFinalizableCheckpointIO):+    def __init__(self):+        _WrappingCheckpointIO.__init__(self, TorchCheckpointIO())+        self.pending = []++    def save_checkpoint(self, checkpoint, path, storage_options=None):+        self.pending.append((copy.deepcopy(checkpoint), path, storage_options['finalize_fn']))++    def flush(self):+        for checkpoint, path, finalize in self.pending:+            self.checkpoint_io.save_checkpoint(checkpoint, path)+            finalize()+        self.pending.clear()++    def teardown(self):+        pass+++@pytest.mark.unit+def test_automodel_resume_connector_restores_metadata(tmp_path):+    strategy = AutomodelParallelStrategy()+    strategy.broadcast = lambda path: path+    connector = _CheckpointConnector(SimpleNamespace(strategy=strategy))+    state = {'global_step': 7, 'pytorch-lightning_version': pl.__version__}+    with patch.object(ModelParallelStrategy, 'load_checkpoint', return_value=state) as load:+        connector.resume_start(tmp_path / 'model.ckpt', weights_only=True)+    assert connector._loaded_checkpoint['global_step'] == 7+    load.assert_called_once_with(tmp_path / 'model.ckpt', weights_only=True)+++@pytest.mark.unit+@pytest.mark.parametrize('sidecar', ['valid', 'missing', 'unsupported'])+def test_ema_fit_sidecar_policy(tmp_path, sidecar):+    import pickle+    from lightning.pytorch.utilities.exceptions import MisconfigurationException++    trainer = pl.Trainer(+        default_root_dir=tmp_path,+        accelerator='cpu',+        logger=False,+        enable_progress_bar=False,+        enable_model_summary=False,+        callbacks=[NeMoModelCheckpoint(dirpath=tmp_path)],+    )+    trainer.state.fn = TrainerFn.FITTING+    trainer.ckpt_path = str(tmp_path / 'model.ckpt')+    if sidecar != 'missing':+        value = torch.ones(2) if sidecar == 'valid' else _UnsupportedValue()+        torch.save({'optimizer_states': [value]}, tmp_path / 'model-EMA.ckpt')+    checkpoint = {'optimizer_states': ['unchanged']}+    if sidecar == 'valid':+        EMA(decay=0.9).on_load_checkpoint(trainer, pl.LightningModule(), checkpoint)+        torch.testing.assert_close(checkpoint['optimizer_states'][0], torch.ones(2))+    else:+        error = MisconfigurationException if sidecar == 'missing' else pickle.UnpicklingError+        with pytest.raises(error):+            EMA(decay=0.9).on_load_checkpoint(trainer, pl.LightningModule(), checkpoint)+        assert checkpoint['optimizer_states'] == ['unchanged']+++class _UnsupportedValue:+    pass+++@pytest.mark.unit+@pytest.mark.parametrize('accumulate', [1, 2])+def test_default_resume_matches_previous_lightning_behavior(tmp_path, accumulate):+    def make_trainer(directory, steps):+        trainer = pl.Trainer(+            default_root_dir=directory,+            accelerator='cpu',+            devices=1,+            logger=False,+            enable_checkpointing=False,+            enable_progress_bar=False,+            enable_model_summary=False,+            max_steps=steps,+            max_epochs=3,+            val_check_interval=2,+            num_sanity_val_steps=0,+            accumulate_grad_batches=accumulate,+        )+        exp_manager(+            trainer,+            {+                'explicit_log_dir': str(directory),+                'create_tensorboard_logger': False,+                'create_preemption_callback': False,+                'checkpoint_callback_params': {+                    'monitor': 'val_loss',+                    'save_top_k': 1,+                    'save_last': True,+                    'save_nemo_on_train_end': False,+                },+            },+        )+        return trainer++    first = make_trainer(tmp_path / 'first', 2)+    first.fit(_ResumeModel())+    resumed = make_trainer(tmp_path / 'resumed', 4)+    model = _ResumeModel()+    resumed.fit(model, ckpt_path=first.checkpoint_callback.last_model_path)+    # Reference traces independently compared with Lightning 2.4. The default+    # non-stateful dataloader restarts its iterator in both releases.+    assert model.seen == ([1, 2] if accumulate == 1 else [1, 2, 3, 4])+    assert model.validation_steps == ([4] if accumulate == 1 else [3, 4])+    assert resumed.global_step == resumed.fit_loop.epoch_loop._batches_that_stepped == 4+    assert resumed.optimizers[0].param_groups[0]['lr'] == pytest.approx(0.06561)+    assert model.weight.item() == pytest.approx(0.24586007 if accumulate == 1 else 0.23775484)+++class _ResumeModel(_PlateauModel):+    def __init__(self):+        super().__init__()+        self.weight.data.fill_(0.25)+        self.seen = []+        self.validation_steps = []++    def training_step(self, batch, batch_idx):+        self.seen.append(int(batch.item()))+        return (self.weight * batch / 10).square().mean()++    def on_validation_epoch_end(self):+        self.validation_steps.append(self.global_step)++    def configure_optimizers(self):+        optimizer = torch.optim.SGD(self.parameters(), lr=0.1)+        return {+            'optimizer': optimizer,+            'lr_scheduler': {'scheduler': torch.optim.lr_scheduler.StepLR(optimizer, 1, 0.9), 'interval': 'step'},+        }++    def train_dataloader(self):+        return DataLoader(torch.arange(1, 9, dtype=torch.float32), batch_size=1)

@@ -16,6 +16,7 @@ """Unit tests for OneLoggerNeMoCallback."""  import os+from importlib import import_module from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -24,6 +25,14 @@ from lightning.pytorch.callbacks import ModelCheckpoint from omegaconf import OmegaConf +# The legacy integration is inactive; retain these tests for compatible environments.+try:+    import_module('nemo.lightning.one_logger_callback')+except TypeError as error:+    if 'OneLoggerPTLTrainer.save_checkpoint: weights_only' not in str(error):+        raise+    pytest.skip('Legacy telemetry integration requires an update for this Lightning version', allow_module_level=True)+ from nemo.lightning.base_callback import BaseCallback from nemo.lightning.one_logger_callback import (     OneLoggerNeMoCallback,

@@ -0,0 +1,86 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0+"""NeMo preserves the callback defaults used before the Lightning upgrade."""++from importlib import import_module++import pytest++from nemo.utils.lightning_logger_patch import preserve_trainer_callback_defaults+++@pytest.mark.unit+@pytest.mark.parametrize('namespace', ['lightning.pytorch', 'pytorch_lightning'])+@pytest.mark.parametrize('flags', [{}, {'enable_progress_bar': None, 'enable_model_summary': None}])+def test_default_callbacks_are_tqdm_and_plain_summary(tmp_path, namespace, flags):+    lightning = import_module(namespace)+    callbacks = import_module(f'{namespace}.callbacks')+    trainer = lightning.Trainer(default_root_dir=tmp_path, accelerator='cpu', logger=False, **flags)+    assert type(trainer.progress_bar_callback) is callbacks.TQDMProgressBar+    assert (+        type(next(cb for cb in trainer.callbacks if isinstance(cb, callbacks.ModelSummary))) is callbacks.ModelSummary+    )+++@pytest.mark.unit+@pytest.mark.parametrize('namespace', ['lightning.pytorch', 'pytorch_lightning'])+@pytest.mark.parametrize('option', ['enable_progress_bar', 'enable_model_summary', 'barebones'])+def test_disabled_callbacks_stay_disabled(tmp_path, namespace, option):+    lightning = import_module(namespace)+    callbacks = import_module(f'{namespace}.callbacks')+    flags = {option: option == 'barebones'}+    trainer = lightning.Trainer(default_root_dir=tmp_path, accelerator='cpu', logger=False, **flags)+    if option in ('enable_progress_bar', 'barebones'):+        assert trainer.progress_bar_callback is None+    if option in ('enable_model_summary', 'barebones'):+        assert not any(isinstance(cb, callbacks.ModelSummary) for cb in trainer.callbacks)+++@pytest.mark.unit+@pytest.mark.parametrize('namespace', ['lightning.pytorch', 'pytorch_lightning'])+@pytest.mark.parametrize('rich', [False, True])+def test_explicit_progress_bar_keeps_old_summary_selection(tmp_path, namespace, rich):+    pytest.importorskip('rich')+    lightning = import_module(namespace)+    callbacks = import_module(f'{namespace}.callbacks')+    progress = callbacks.RichProgressBar() if rich else callbacks.TQDMProgressBar(refresh_rate=7)+    trainer = lightning.Trainer(default_root_dir=tmp_path, accelerator='cpu', logger=False, callbacks=[progress])+    assert trainer.progress_bar_callback is progress+    summary = next(cb for cb in trainer.callbacks if isinstance(cb, callbacks.ModelSummary))+    assert type(summary) is (callbacks.RichModelSummary if rich else callbacks.ModelSummary)+++@pytest.mark.unit+@pytest.mark.parametrize('namespace', ['lightning.pytorch', 'pytorch_lightning'])+@pytest.mark.parametrize('rich', [False, True])+def test_explicit_summary_is_retained(tmp_path, namespace, rich):+    pytest.importorskip('rich')+    lightning = import_module(namespace)+    callbacks = import_module(f'{namespace}.callbacks')+    summary = callbacks.RichModelSummary(max_depth=3) if rich else callbacks.ModelSummary(max_depth=3)+    trainer = lightning.Trainer(default_root_dir=tmp_path, accelerator='cpu', logger=False, callbacks=[summary])+    assert [cb for cb in trainer.callbacks if isinstance(cb, callbacks.ModelSummary)] == [summary]+++@pytest.mark.unit+def test_callback_patch_is_idempotent():+    from lightning.pytorch.trainer.connectors.callback_connector import _CallbackConnector++    before = _CallbackConnector._configure_progress_bar+    preserve_trainer_callback_defaults()+    preserve_trainer_callback_defaults()+    assert _CallbackConnector._configure_progress_bar is before+++@pytest.mark.unit+@pytest.mark.parametrize('namespace', ['lightning.pytorch', 'pytorch_lightning'])+def test_explicit_progress_bar_validation_is_not_bypassed(tmp_path, namespace):+    lightning = import_module(namespace)+    callbacks = import_module(f'{namespace}.callbacks')+    exceptions = import_module(f'{namespace}.utilities.exceptions')+    for flags, bars in [+        ({'enable_progress_bar': False}, [callbacks.TQDMProgressBar()]),+        ({}, [callbacks.TQDMProgressBar(), callbacks.TQDMProgressBar()]),+    ]:+        with pytest.raises(exceptions.MisconfigurationException):+            lightning.Trainer(default_root_dir=tmp_path, accelerator='cpu', logger=False, callbacks=bars, **flags)

@@ -4241,16 +4241,16 @@ cu13 = [  [[package]] name = "hydra-core"-version = "1.3.2"+version = "1.3.8" source = { registry = "https://pypi.org/simple" } dependencies = [     { name = "antlr4-python3-runtime" },     { name = "omegaconf" },     { name = "packaging" }, ]-sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" }+sdist = { url = "https://files.pythonhosted.org/packages/7c/1d/e1ad315ca93f1b090b0782942535fea8344e8b2d749c92117b2bf45f4407/hydra_core-1.3.8.tar.gz", hash = "sha256:18d5dfb841734a936507e75e7b8390b406ed22494f7fb8b4b7b7c122102de3d5", size = 3278858, upload-time = "2026-10-25T13:42:22.491Z" } wheels = [-    { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" },+    { url = "https://files.pythonhosted.org/packages/13/69/8e5cf7ad847538fc180fbd06e8442b8bede53ee4852458a16e24231395e9/hydra_core-1.3.8-py3-none-any.whl", hash = "sha256:d5f184c291691cba835c7f04bc4e36270ad88b1d7e394a3712f5af084c30b394", size = 168949, upload-time = "2026-10-25T13:42:20.959Z" }, ]  [[package]]@@ -5120,7 +5120,7 @@ wheels = [  [[package]] name = "lightning"-version = "2.4.0"+version = "2.6.6" source = { registry = "https://pypi.org/simple" } dependencies = [     { name = "fsspec", extra = ["http"] },@@ -5137,9 +5137,9 @@ dependencies = [     { name = "tqdm" },     { name = "typing-extensions" }, ]-sdist = { url = "https://files.pythonhosted.org/packages/56/d0/78ea244ac044cd4df15aa8294a50ff3561fb177e7e5ba788aaa542046cae/lightning-2.4.0.tar.gz", hash = "sha256:9156604cc56e4b2b603f34fa7f0fe5107375c8e6d85e74544b319a15faa9ed0e", size = 620632, upload-time = "2024-08-07T09:46:44.399Z" }+sdist = { url = "https://files.pythonhosted.org/packages/39/b8/fbf1469d917c3ed3b22e00ab3a852fb9a5f85b9dc7f7060eedc3449098e7/lightning-2.6.6.tar.gz", hash = "sha256:1c6329e241f8a85f29fedead89c1ee138b36caa08de60c260c6a7eaee81c8295", size = 659902, upload-time = "2026-09-10T09:41:03.009Z" } wheels = [-    { url = "https://files.pythonhosted.org/packages/a3/2c/85eaf42c983b0cd81bcda5876da2c8e2a9fd347908666ea9855724369171/lightning-2.4.0-py3-none-any.whl", hash = "sha256:560163af9711cf59055c448232c473150a299089efce0d2be3cc3288082d8768", size = 810971, upload-time = "2024-08-07T09:46:39.874Z" },+    { url = "https://files.pythonhosted.org/packages/12/43/7da7fe42f254ce57777d1a86eece3f86d1b3329fc568d48e2bc16c8b7ebc/lightning-2.6.6-py3-none-any.whl", hash = "sha256:2bcbd6ee840071cd076c7dc9953d7ac1040df92f1c5fa76583aae644f037ab85", size = 849265, upload-time = "2026-09-10T09:41:00.882Z" }, ]  [[package]]@@ -6840,6 +6840,7 @@ all = [     { name = "pypinyin" },     { name = "pypinyin-dict" },     { name = "pystoi" },+    { name = "pytorch-lightning" },     { name = "sacrebleu" },     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-all') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-compiled') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-cu12') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-cu13') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2-only') or (python_full_version < '3.11' and extra != 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-all') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-compiled') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-cu12') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-cu13') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2-only') or (python_full_version >= '3.11' and extra != 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },@@ -6866,6 +6867,7 @@ asr = [     { name = "omegaconf" },     { name = "packaging" },     { name = "pandas" },+    { name = "pytorch-lightning" },     { name = "sacrebleu" },     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },@@ -6905,6 +6907,7 @@ audio = [     { name = "pandas" },     { name = "pesq", marker = "platform_machine != 'x86_64' or sys_platform != 'darwin' or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },     { name = "pystoi" },+    { name = "pytorch-lightning" },     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },     { name = "sentencepiece" },@@ -6924,6 +6927,7 @@ common = [     { name = "nv-one-logger-training-telemetry" },     { name = "omegaconf" },     { name = "pandas" },+    { name = "pytorch-lightning" },     { name = "sentencepiece" },     { name = "torchmetrics" },     { name = "transformers" },@@ -6964,6 +6968,7 @@ core = [     { name = "nv-one-logger-pytorch-lightning-integration" },     { name = "nv-one-logger-training-telemetry" },     { name = "omegaconf" },+    { name = "pytorch-lightning" },     { name = "torchmetrics" },     { name = "transformers" },     { name = "wandb" },@@ -6988,6 +6993,7 @@ lightning = [     { name = "nv-one-logger-pytorch-lightning-integration" },     { name = "nv-one-logger-training-telemetry" },     { name = "omegaconf" },+    { name = "pytorch-lightning" },     { name = "torchmetrics" },     { name = "transformers" },     { name = "wandb" },@@ -7019,6 +7025,7 @@ speechlm2 = [     { name = "pyopenjtalk" },     { name = "pypinyin" },     { name = "pypinyin-dict" },+    { name = "pytorch-lightning" },     { name = "sacrebleu" },     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-all') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-compiled') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-cu12') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-cu13') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2') or (python_full_version < '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2-only') or (python_full_version < '3.11' and extra != 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-all') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-compiled') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-cu12') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-cu13') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2') or (python_full_version >= '3.11' and extra == 'extra-12-nemo-toolkit-speechlm2-only') or (python_full_version >= '3.11' and extra != 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },@@ -7058,6 +7065,7 @@ tts = [     { name = "pyopenjtalk" },     { name = "pypinyin" },     { name = "pypinyin-dict" },+    { name = "pytorch-lightning" },     { name = "sacrebleu" },     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-12-nemo-toolkit-all' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-compiled-a100' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-cu13') or (extra == 'extra-12-nemo-toolkit-cu12' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-cu13' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2' and extra == 'extra-12-nemo-toolkit-vllm') or (extra == 'extra-12-nemo-toolkit-speechlm2-only' and extra == 'extra-12-nemo-toolkit-vllm')" },@@ -7153,14 +7161,14 @@ requires-dist = [     { name = "flashoptim", marker = "extra == 'speechlm2'" },     { name = "flashoptim", marker = "extra == 'speechlm2-only'" },     { name = "huggingface-hub", specifier = ">=0.24" },-    { name = "hydra-core", marker = "extra == 'all'", specifier = ">1.3,<=1.3.2" },-    { name = "hydra-core", marker = "extra == 'asr'", specifier = ">1.3,<=1.3.2" },-    { name = "hydra-core", marker = "extra == 'audio'", specifier = ">1.3,<=1.3.2" },-    { name = "hydra-core", marker = "extra == 'common'", specifier = ">1.3,<=1.3.2" },-    { name = "hydra-core", marker = "extra == 'core'", specifier = ">1.3,<=1.3.2" },-    { name = "hydra-core", marker = "extra == 'lightning'", specifier = ">1.3,<=1.3.2" },-    { name = "hydra-core", marker = "extra == 'speechlm2'", specifier = ">1.3,<=1.3.2" },-    { name = "hydra-core", marker = "extra == 'tts'", specifier = ">1.3,<=1.3.2" },+    { name = "hydra-core", marker = "extra == 'all'", specifier = ">=1.3.8,<1.4" },+    { name = "hydra-core", marker = "extra == 'asr'", specifier = ">=1.3.8,<1.4" },+    { name = "hydra-core", marker = "extra == 'audio'", specifier = ">=1.3.8,<1.4" },+    { name = "hydra-core", marker = "extra == 'common'", specifier = ">=1.3.8,<1.4" },+    { name = "hydra-core", marker = "extra == 'core'", specifier = ">=1.3.8,<1.4" },+    { name = "hydra-core", marker = "extra == 'lightning'", specifier = ">=1.3.8,<1.4" },+    { name = "hydra-core", marker = "extra == 'speechlm2'", specifier = ">=1.3.8,<1.4" },+    { name = "hydra-core", marker = "extra == 'tts'", specifier = ">=1.3.8,<1.4" },     { name = "janome", marker = "extra == 'all'" },     { name = "janome", marker = "extra == 'speechlm2'" },     { name = "janome", marker = "extra == 'tts'" },@@ -7188,14 +7196,14 @@ requires-dist = [     { name = "librosa", marker = "extra == 'speechlm2'", specifier = ">=0.10.1" },     { name = "librosa", marker = "extra == 'tts'" },     { name = "librosa", marker = "extra == 'tts'", specifier = ">=0.10.1" },-    { name = "lightning", marker = "extra == 'all'", specifier = ">2.2.1,<=2.4.0" },-    { name = "lightning", marker = "extra == 'asr'", specifier = ">2.2.1,<=2.4.0" },-    { name = "lightning", marker = "extra == 'audio'", specifier = ">2.2.1,<=2.4.0" },-    { name = "lightning", marker = "extra == 'common'", specifier = ">2.2.1,<=2.4.0" },-    { name = "lightning", marker = "extra == 'core'", specifier = ">2.2.1,<=2.4.0" },-    { name = "lightning", marker = "extra == 'lightning'", specifier = ">2.2.1,<=2.4.0" },-    { name = "lightning", marker = "extra == 'speechlm2'", specifier = ">2.2.1,<=2.4.0" },-    { name = "lightning", marker = "extra == 'tts'", specifier = ">2.2.1,<=2.4.0" },+    { name = "lightning", marker = "extra == 'all'", specifier = ">=2.6.6,<2.7" },+    { name = "lightning", marker = "extra == 'asr'", specifier = ">=2.6.6,<2.7" },+    { name = "lightning", marker = "extra == 'audio'", specifier = ">=2.6.6,<2.7" },+    { name = "lightning", marker = "extra == 'common'", specifier = ">=2.6.6,<2.7" },+    { name = "lightning", marker = "extra == 'core'", specifier = ">=2.6.6,<2.7" },+    { name = "lightning", marker = "extra == 'lightning'", specifier = ">=2.6.6,<2.7" },+    { name = "lightning", marker = "extra == 'speechlm2'", specifier = ">=2.6.6,<2.7" },+    { name = "lightning", marker = "extra == 'tts'", specifier = ">=2.6.6,<2.7" },     { name = "mamba-ssm", marker = "extra == 'compiled'", specifier = "==2.3.2.post1" },     { name = "mamba-ssm", marker = "extra == 'compiled-a100'", specifier = "==2.3.2.post1" },     { name = "matplotlib", marker = "extra == 'all'" },@@ -7241,14 +7249,14 @@ requires-dist = [     { name = "nv-one-logger-training-telemetry", marker = "extra == 'lightning'", specifier = ">=2.3.1" },     { name = "nv-one-logger-training-telemetry", marker = "extra == 'speechlm2'", specifier = ">=2.3.1" },     { name = "nv-one-logger-training-telemetry", marker = "extra == 'tts'", specifier = ">=2.3.1" },-    { name = "omegaconf", marker = "extra == 'all'", specifier = "<=2.3" },-    { name = "omegaconf", marker = "extra == 'asr'", specifier = "<=2.3" },-    { name = "omegaconf", marker = "extra == 'audio'", specifier = "<=2.3" },-    { name = "omegaconf", marker = "extra == 'common'", specifier = "<=2.3" },-    { name = "omegaconf", marker = "extra == 'core'", specifier = "<=2.3" },-    { name = "omegaconf", marker = "extra == 'lightning'", specifier = "<=2.3" },-    { name = "omegaconf", marker = "extra == 'speechlm2'", specifier = "<=2.3" },-    { name = "omegaconf", marker = "extra == 'tts'", specifier = "<=2.3" },+    { name = "omegaconf", marker = "extra == 'all'", specifier = ">=2.3,<2.4" },+    { name = "omegaconf", marker = "extra == 'asr'", specifier = ">=2.3,<2.4" },+    { name = "omegaconf", marker = "extra == 'audio'", specifier = ">=2.3,<2.4" },+    { name = "omegaconf", marker = "extra == 'common'", specifier = ">=2.3,<2.4" },+    { name = "omegaconf", marker = "extra == 'core'", specifier = ">=2.3,<2.4" },+    { name = "omegaconf", marker = "extra == 'lightning'", specifier = ">=2.3,<2.4" },+    { name = "omegaconf", marker = "extra == 'speechlm2'", specifier = ">=2.3,<2.4" },+    { name = "omegaconf", marker = "extra == 'tts'", specifier = ">=2.3,<2.4" },     { name = "onnx", specifier = ">=1.7.0" },     { name = "onnx-ir", marker = "extra == 'compiled'", specifier = "==0.2.1" },     { name = "onnx-ir", marker = "extra == 'compiled-a100'", specifier = "==0.2.1" },@@ -7283,6 +7291,14 @@ requires-dist = [     { name = "pypinyin-dict", marker = "extra == 'tts'" },     { name = "pystoi", marker = "extra == 'all'" },     { name = "pystoi", marker = "extra == 'audio'" },+    { name = "pytorch-lightning", marker = "extra == 'all'", specifier = ">=2.6.6,<2.7" },+    { name = "pytorch-lightning", marker = "extra == 'asr'", specifier = ">=2.6.6,<2.7" },+    { name = "pytorch-lightning", marker = "extra == 'audio'", specifier = ">=2.6.6,<2.7" },+    { name = "pytorch-lightning", marker = "extra == 'common'", specifier = ">=2.6.6,<2.7" },+    { name = "pytorch-lightning", marker = "extra == 'core'", specifier = ">=2.6.6,<2.7" },+    { name = "pytorch-lightning", marker = "extra == 'lightning'", specifier = ">=2.6.6,<2.7" },+    { name = "pytorch-lightning", marker = "extra == 'speechlm2'", specifier = ">=2.6.6,<2.7" },+    { name = "pytorch-lightning", marker = "extra == 'tts'", specifier = ">=2.6.6,<2.7" },     { name = "sacrebleu", marker = "extra == 'all'" },     { name = "sacrebleu", marker = "extra == 'asr'" },     { name = "sacrebleu", marker = "extra == 'asr-only'" },@@ -11525,7 +11541,7 @@ wheels = [  [[package]] name = "pytorch-lightning"-version = "2.6.1"+version = "2.6.6" source = { registry = "https://pypi.org/simple" } dependencies = [     { name = "fsspec", extra = ["http"] },@@ -11541,9 +11557,9 @@ dependencies = [     { name = "tqdm" },     { name = "typing-extensions" }, ]-sdist = { url = "https://files.pythonhosted.org/packages/8b/ac/ebd5f6f58691cbd4f73836e43e1727f3814311b960c41f88e259606ca2b2/pytorch_lightning-2.6.1.tar.gz", hash = "sha256:ba08f8901cf226fcca473046ad9346f414e99117762dc869c76e650d5b3d7bdc", size = 665563, upload-time = "2026-01-30T14:59:11.636Z" }+sdist = { url = "https://files.pythonhosted.org/packages/97/e4/5025cb42311b4f8cc6d1b12cd81399ca433159f8e3831ab229c5a6c5177e/pytorch_lightning-2.6.6.tar.gz", hash = "sha256:52c52b20522dd3d3feb3bb3aae2dbe3f291e8e836d4111a74cfd72e9df018834", size = 662507, upload-time = "2026-09-10T09:41:16.684Z" } wheels = [-    { url = "https://files.pythonhosted.org/packages/0e/93/c8c361bf0a2fe50f828f32def460e8b8a14b93955d3fd302b1a9b63b19e4/pytorch_lightning-2.6.1-py3-none-any.whl", hash = "sha256:1f8118567ec829e3055f16cf1aa320883a86a47c836951bfd9dcfa34ec7ffd59", size = 857273, upload-time = "2026-01-30T14:59:10.141Z" },+    { url = "https://files.pythonhosted.org/packages/4f/cf/ba8e4cc1b982ff6dc6f9fdef0e4074605438b82640f6270888a10e7e42c3/pytorch_lightning-2.6.6-py3-none-any.whl", hash = "sha256:71f95c7b22f25c4f91cc659e1734bcdc6c69c8b66543f2ed2d78dbe29ce2f167", size = 853045, upload-time = "2026-09-10T09:41:14.923Z" }, ]  [[package]]