Configure MCore logger (NVIDIA#7781)

Signed-off-by: Mikołaj Błaż <[email protected]>
Signed-off-by: Piotr Żelasko <[email protected]>

nemo/utils/exp_manager.py

@@ -46,6 +46,7 @@
 from nemo.utils.get_rank import is_global_rank_zero
 from nemo.utils.lightning_logger_patch import add_filehandlers_to_pl_logger
 from nemo.utils.loggers import ClearMLLogger, ClearMLParams, DLLogger, DLLoggerParams, MLFlowParams
+from nemo.utils.mcore_logger import add_handlers_to_mcore_logger
 from nemo.utils.model_utils import uninject_model_parallel_rank
 
 
@@ -511,6 +512,8 @@ def exp_manager(trainer: 'pytorch_lightning.Trainer', cfg: Optional[Union[DictCo
         # doing the initialization such as moving files
         time.sleep(cfg.seconds_to_sleep)
 
+    add_handlers_to_mcore_logger()
+
     return log_dir
 
 

nemo/utils/mcore_logger.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging as _logging
+
+from nemo.utils import logging as nemo_logger
+
+
+def add_handlers_to_mcore_logger():
+    """Add NeMo handlers to MCore loggers.
+
+    MCore doesn't have and handlers for loggers (see
+    https://docs.python.org/3/howto/logging-cookbook.html#adding-handlers-other-than-nullhandler-to-a-logger-in-a-library
+    for a rationale). We have to add handlers explicitly.
+    """
+    mcore_logger = _logging.getLogger('megatron.core')
+    for handler in nemo_logger._handlers.values():
+        mcore_logger.addHandler(handler)
+    mcore_logger.propagate = False
+    mcore_logger.setLevel(nemo_logger._logger.level)