fix: change logging level to info and print flat arguments (#582)

Signed-off-by: Dushyant Behl <[email protected]>

Author: dushyantbehl

tests/utils/test_logging.py

@@ -36,9 +36,9 @@ def test_set_log_level_for_logger_default():
 
     with mock.patch.dict(os.environ, {}, clear=True):
         train_args = copy.deepcopy(TRAIN_ARGS)
-        training_args, logger = set_log_level(train_args)
-        assert logger.getEffectiveLevel() == logging.WARNING
-        assert training_args.log_level == "passive"
+        logger, log_level = set_log_level(level=train_args.log_level)
+        assert logger.getEffectiveLevel() == logging.INFO
+        assert log_level == "info"
 
 
 def test_set_log_level_for_logger_with_env_var():
@@ -48,10 +48,10 @@ def test_set_log_level_for_logger_with_env_var():
     """
 
     with mock.patch.dict(os.environ, {"LOG_LEVEL": "info"}, clear=True):
-        train_args_env = copy.deepcopy(TRAIN_ARGS)
-        training_args, logger = set_log_level(train_args_env)
+        train_args = copy.deepcopy(TRAIN_ARGS)
+        logger, log_level = set_log_level(level=train_args.log_level)
         assert logger.getEffectiveLevel() == logging.INFO
-        assert training_args.log_level == "info"
+        assert log_level == "info"
 
 
 def test_set_log_level_for_logger_with_set_verbosity_and_cli():
@@ -64,9 +64,9 @@ def test_set_log_level_for_logger_with_set_verbosity_and_cli():
     with mock.patch.dict(os.environ, {"TRANSFORMERS_VERBOSITY": "info"}, clear=True):
         train_args = copy.deepcopy(TRAIN_ARGS)
         train_args.log_level = "error"
-        training_args, logger = set_log_level(train_args)
+        logger, log_level = set_log_level(level=train_args.log_level)
         assert logger.getEffectiveLevel() == logging.ERROR
-        assert training_args.log_level == "error"
+        assert log_level == "error"
 
 
 def test_set_log_level_for_logger_with_env_var_and_cli():
@@ -79,6 +79,6 @@ def test_set_log_level_for_logger_with_env_var_and_cli():
     with mock.patch.dict(os.environ, {"LOG_LEVEL": "info"}, clear=True):
         train_args = copy.deepcopy(TRAIN_ARGS)
         train_args.log_level = "error"
-        training_args, logger = set_log_level(train_args)
+        logger, log_level = set_log_level(level=train_args.log_level)
         assert logger.getEffectiveLevel() == logging.ERROR
-        assert training_args.log_level == "error"
+        assert log_level == "error"

tuning/sft_trainer.py

@@ -63,7 +63,7 @@
     USER_ERROR_EXIT_CODE,
     write_termination_log,
 )
-from tuning.utils.logging import set_log_level
+from tuning.utils.logging import pretty_print_args, set_log_level
 
 
 def train(
@@ -120,7 +120,9 @@ def train(
         Tuple: Instance of SFTTrainer , some metadata in a dict
             Metadata contains information on number of added tokens while tuning.
     """
-    train_args, logger = set_log_level(train_args, "sft_trainer_train")
+    logger, train_args.log_level = set_log_level(
+        logger_name="sft_trainer_train", level=train_args.log_level
+    )
     USE_ALORA = False
     try:
         # Third Party
@@ -529,7 +531,6 @@ def get_parser():
         choices=["pt", "lora", "alora", None, "none"],
         default="none",
     )
-
     parser.add_argument(
         "--exp_metadata",
         type=str,
@@ -603,7 +604,6 @@ def parse_arguments(parser, json_config=None):
                 raise ValueError(
                     "invocation_string is not passed required for aLoRA usage"
                 )
-
     else:
         (
             model_args,
@@ -687,26 +687,28 @@ def main():
         ) = parse_arguments(parser, job_config)
 
         # Function to set log level for python native logger and transformers training logger
-        training_args, logger = set_log_level(training_args, __name__)
+        logger, training_args.log_level = set_log_level(
+            logger_name=__name__, level=training_args.log_level
+        )
 
-        logger.info(
-            "Flat arguments parsed: \
-            model_args %s, data_args %s, training_args %s, trainer_controller_args %s, \
-            tune_config %s, quantized_lora_config %s, fusedops_kernels_config %s, \
-            attention_and_distributed_packing_config, %s,\
-            fast_moe_config %s, tracker_config %s, exp_metadata %s",
-            model_args,
-            data_args,
-            training_args,
-            trainer_controller_args,
-            tune_config,
-            quantized_lora_config,
-            fusedops_kernels_config,
-            attention_and_distributed_packing_config,
-            fast_moe_config,
-            tracker_configs,
-            exp_metadata,
+        logger.info("fms-hf-tuning execution start")
+        args_dump = pretty_print_args(
+            {
+                "Model Arguments": model_args,
+                "Data Arguments": data_args,
+                "Training Arguments": training_args,
+                "Tune Config": tune_config,
+                "QLoRA Config": quantized_lora_config,
+                "Tracker Config": tracker_configs,
+                "AADP (fms-acceleration) Config": attention_and_distributed_packing_config,
+                "Fused Ops Kernels Config": fusedops_kernels_config,
+                "Fast MoE Config": fast_moe_config,
+                "Trainer Controller Config": trainer_controller_args,
+                "Extra Metadata": exp_metadata,
+            }
         )
+        logger.info(args_dump)
+
     except Exception as e:  # pylint: disable=broad-except
         logger.error(traceback.format_exc())
         write_termination_log(

tuning/utils/logging.py

@@ -13,24 +13,33 @@
 # limitations under the License.
 
 # Standard
+from typing import Dict
 import logging
 import os
 
+# Third Party
+from accelerate.state import PartialState
+import datasets
+import transformers
 
-def set_log_level(train_args, logger_name=None):
+DEFAULT_LOG_LEVEL_MAIN = "INFO"
+DEFAULT_LOG_LEVEL_WORKERS = "WARNING"
+
+
+def set_log_level(logger_name="fms-hf-tuning", level=None):
     """Set log level of python native logger and TF logger via argument from CLI or env variable.
 
     Args:
-        train_args
-            Training arguments for training model.
         logger_name
             Logger name with which the logger is instantiated.
+        level
+            Requested level of the logger
 
     Returns:
-        train_args
-            Updated training arguments for training model.
         train_logger
             Logger with updated effective log level
+        level
+            Level of the logger initialized
     """
 
     # Clear any existing handlers if necessary
@@ -39,26 +48,58 @@ def set_log_level(train_args, logger_name=None):
 
     # Configure Python native logger and transformers log level
     # If CLI arg is passed, assign same log level to python native logger
-    log_level = "WARNING"
-    if train_args.log_level != "passive":
-        log_level = train_args.log_level
-
-    # If CLI arg not is passed and env var LOG_LEVEL is set,
-    # assign same log level to both logger
+    lowest_log_level = DEFAULT_LOG_LEVEL_MAIN
+    if level != "passive":
+        lowest_log_level = level
     elif os.environ.get("LOG_LEVEL"):
-        log_level = os.environ.get("LOG_LEVEL")
-        train_args.log_level = (
-            log_level.lower()
+        # If CLI arg not is passed and env var LOG_LEVEL is set,
+        # assign same log level to both logger
+        lowest_log_level = (
+            os.environ.get("LOG_LEVEL").lower()
             if not os.environ.get("TRANSFORMERS_VERBOSITY")
             else os.environ.get("TRANSFORMERS_VERBOSITY")
         )
 
+    state = PartialState()
+    rank = state.process_index
+
+    log_on_all = os.environ.get("LOG_ON_ALL_PROCESSES")
+    if log_on_all:
+        log_level = lowest_log_level or DEFAULT_LOG_LEVEL_MAIN
+    else:
+        if state.is_local_main_process:
+            log_level = lowest_log_level or DEFAULT_LOG_LEVEL_MAIN
+            datasets.utils.logging.set_verbosity_warning()
+            transformers.utils.logging.set_verbosity_info()
+        else:
+            log_level = DEFAULT_LOG_LEVEL_WORKERS
+            datasets.utils.logging.set_verbosity_error()
+            transformers.utils.logging.set_verbosity_error()
+
+    log_format = f"Rank-{rank} [%(levelname)s]:%(filename)s:%(funcName)s: %(message)s"
+
     logging.basicConfig(
-        format="%(levelname)s:%(filename)s:%(message)s", level=log_level.upper()
+        format=log_format,
+        level=log_level.upper(),
     )
 
     if logger_name:
         train_logger = logging.getLogger(logger_name)
     else:
         train_logger = logging.getLogger()
-    return train_args, train_logger
+
+    return train_logger, log_level.lower()
+
+
+def pretty_print_args(args: Dict):
+    dump = "\n========================= Flat Arguments =========================\n"
+    for name, arg in args.items():
+        if arg:
+            dump += f"---------------------------- {name} -----------------------\n"
+            if hasattr(arg, "__dict__"):
+                arg = vars(arg)
+            max_len = max(len(k) for k in arg.keys())
+            for k, v in sorted(arg.items()):
+                dump += f"  {k:<{max_len}} : {v}\n"
+    dump += "========================= Arguments Done =========================\n"
+    return dump