Add debug log support (modelscope#1226)

Author: tastelikefeet

swift/llm/utils/model.py

@@ -910,9 +910,14 @@ def get_model_tokenizer_from_repo(model_dir: str,
                 trust_remote_code=True,
             )
         else:
+            logger.info(f'Model loading with args: model_dir: {model_dir},'
+                        f'torch_dtype: {torch_dtype},'
+                        f'model_kwargs: {model_kwargs}')
             with context:
                 model = automodel_class.from_pretrained(
                     model_dir, config=model_config, torch_dtype=torch_dtype, trust_remote_code=True, **model_kwargs)
+        if hasattr(model, 'hf_device_map'):
+            logger.debug(f'Model hf_device_map: {model.hf_device_map}')
     return model, tokenizer
 
 

swift/llm/utils/utils.py

@@ -45,11 +45,10 @@
 
 logger.handlers[0].setFormatter(logger_format)
 ms_logger.handlers[0].setFormatter(logger_format)
+log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
 if is_local_master():
-    logger.setLevel(logging.INFO)
-    ms_logger.setLevel(logging.INFO)
+    ms_logger.setLevel(log_level)
 else:
-    logger.setLevel(logging.ERROR)
     ms_logger.setLevel(logging.ERROR)
 
 os.environ['TOKENIZERS_PARALLELISM'] = 'true'

swift/utils/logger.py

@@ -9,12 +9,7 @@
 formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
 
-def is_master():
-    rank = int(os.getenv('RANK', -1))
-    return rank in {-1, 0}
-
-
-def get_logger(log_file: Optional[str] = None, log_level: int = logging.INFO, file_mode: str = 'w'):
+def get_logger(log_file: Optional[str] = None, log_level: Optional[int] = None, file_mode: str = 'w'):
     """ Get logging logger
 
     Args:
@@ -24,7 +19,9 @@ def get_logger(log_file: Optional[str] = None, log_level: int = logging.INFO, fi
         file_mode: Specifies the mode to open the file, if filename is
             specified (if filemode is unspecified, it defaults to 'w').
     """
-
+    if log_level is None:
+        log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
+        log_level = getattr(logging, log_level, logging.INFO)
     logger_name = __name__.split('.')[0]
     logger = logging.getLogger(logger_name)
     logger.propagate = False
@@ -47,7 +44,7 @@ def get_logger(log_file: Optional[str] = None, log_level: int = logging.INFO, fi
     handlers = [stream_handler]
 
     if importlib.util.find_spec('torch') is not None:
-        is_worker0 = is_master()
+        is_worker0 = int(os.getenv('LOCAL_RANK', -1)) in {-1, 0}
     else:
         is_worker0 = True
 
@@ -76,7 +73,7 @@ def add_file_handler_if_needed(logger, log_file, file_mode, log_level):
             return
 
     if importlib.util.find_spec('torch') is not None:
-        is_worker0 = is_master()
+        is_worker0 = int(os.getenv('LOCAL_RANK', -1)) in {-1, 0}
     else:
         is_worker0 = True
 

swift/utils/torch_utils.py

@@ -13,7 +13,7 @@
 from torch.nn import Module
 from transformers.utils import is_torch_npu_available, strtobool
 
-from .logger import get_logger, is_master
+from .logger import get_logger
 
 logger = get_logger()
 
@@ -82,6 +82,11 @@ def is_local_master():
     return local_rank in {-1, 0}
 
 
+def is_master():
+    rank = int(os.getenv('RANK', -1))
+    return rank in {-1, 0}
+
+
 def use_torchacc() -> bool:
     return strtobool(os.getenv('USE_TORCHACC', '0'))