Fix the load checkpointing issue -- onelogger callback gets called multiple time in some case. (#14945)

* Update modelPT.py

Signed-off-by: Jiashang Hu <[email protected]>

* Update one_logger_callback.py

Signed-off-by: Jiashang Hu <[email protected]>

* fix the test for error handling

Signed-off-by: Jiashang Hu <[email protected]>

* update the dependency version

Signed-off-by: Jiashang Hu <[email protected]>

---------

Signed-off-by: Jiashang Hu <[email protected]>
Signed-off-by: NeMo Bot <[email protected]>

Author: liquor233

nemo/core/classes/modelPT.py

@@ -478,9 +478,6 @@ def restore_from(
         Returns:
             An instance of type cls or its underlying config (if return_config is set).
         """
-        # Notify OneLogger of checkpoint loading start for telemetry tracking
-        CallbackGroup.get_instance().on_load_checkpoint_start()
-
         if save_restore_connector is None:
             save_restore_connector = SaveRestoreConnector()
 
@@ -514,9 +511,6 @@ def restore_from(
         if isinstance(instance, ModelPT):
             instance._save_restore_connector = save_restore_connector
 
-        # Notify OneLogger of checkpoint loading completion for telemetry tracking
-        CallbackGroup.get_instance().on_load_checkpoint_end()
-
         return instance
 
     @classmethod

nemo/lightning/one_logger_callback.py

@@ -59,9 +59,6 @@ def get_one_logger_init_config() -> Dict[str, Any]:
         # Important fields with defaults - provide if available from config
         "enable_for_current_rank": _should_enable_for_current_rank(),
         "world_size_or_fn": world_size,
-        # Error handling strategy - use DISABLE_QUIETLY_AND_REPORT_METRIC_ERROR to prevent
-        # telemetry errors from crashing the training application
-        "error_handling_strategy": "propagate_exceptions",
     }
 
     return init_config

requirements/requirements_lightning.txt

@@ -8,6 +8,6 @@ torchmetrics>=0.11.0
 transformers~=4.53.0
 wandb
 webdataset>=0.2.86
-nv_one_logger_core>=2.1.0
-nv_one_logger_training_telemetry>=2.1.0
-nv_one_logger_pytorch_lightning_integration>=2.1.0
+nv_one_logger_core>=2.3.0
+nv_one_logger_training_telemetry>=2.3.0
+nv_one_logger_pytorch_lightning_integration>=2.3.0

tests/lightning/test_one_logger_callback.py

@@ -80,7 +80,6 @@ def test_init_configures_provider(
             "session_tag_or_fn": "test-session",
             "enable_for_current_rank": True,
             "world_size_or_fn": 1,
-            "error_handling_strategy": "propagate_exceptions",
         }
         mock_get_config.return_value = mock_init_config
 
@@ -119,7 +118,6 @@ def test_get_one_logger_init_config(self):
             assert config["session_tag_or_fn"] == "test_job"
             assert "enable_for_current_rank" in config
             assert config["world_size_or_fn"] == 4
-            assert config["error_handling_strategy"] == "propagate_exceptions"
 
     @pytest.mark.unit
     def test_get_one_logger_init_config_no_slurm(self):