fix #2649

Former-commit-id: 1c850de660c671d92f0bc63f230d338b60b7c0bd

src/llmtuner/model/loader.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 from trl import AutoModelForCausalLMWithValueHead
@@ -19,38 +19,48 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
-def load_model_and_tokenizer(
-    model_args: "ModelArguments",
-    finetuning_args: "FinetuningArguments",
-    is_trainable: Optional[bool] = False,
-    add_valuehead: Optional[bool] = False,
-) -> Tuple["PreTrainedModel", "PreTrainedTokenizer"]:
-    r"""
-    Loads pretrained model and tokenizer.
-
-    Support both training and inference.
-    """
-
-    try_download_model_from_ms(model_args)
-
-    config_kwargs = {
+def _get_init_kwargs(model_args: "ModelArguments") -> Dict[str, Any]:
+    return {
         "trust_remote_code": True,
         "cache_dir": model_args.cache_dir,
         "revision": model_args.model_revision,
         "token": model_args.hf_hub_token,
     }
 
+
+def load_tokenizer(model_args: "ModelArguments") -> "PreTrainedTokenizer":
+    r"""
+    Loads pretrained tokenizer. Must before load_model.
+
+    Note: including inplace operation of model_args.
+    """
+    try_download_model_from_ms(model_args)
+    init_kwargs = _get_init_kwargs(model_args)
+
     tokenizer = AutoTokenizer.from_pretrained(
         model_args.model_name_or_path,
         use_fast=model_args.use_fast_tokenizer,
         split_special_tokens=model_args.split_special_tokens,
         padding_side="right",
-        **config_kwargs,
+        **init_kwargs,
     )
     patch_tokenizer(tokenizer)
+    return tokenizer
 
-    config = AutoConfig.from_pretrained(model_args.model_name_or_path, **config_kwargs)
-    patch_config(config, tokenizer, model_args, config_kwargs, is_trainable)
+
+def load_model(
+    tokenizer: "PreTrainedTokenizer",
+    model_args: "ModelArguments",
+    finetuning_args: "FinetuningArguments",
+    is_trainable: Optional[bool] = False,
+    add_valuehead: Optional[bool] = False,
+) -> "PreTrainedModel":
+    r"""
+    Loads pretrained model. Must after load_tokenizer.
+    """
+    init_kwargs = _get_init_kwargs(model_args)
+    config = AutoConfig.from_pretrained(model_args.model_name_or_path, **init_kwargs)
+    patch_config(config, tokenizer, model_args, init_kwargs, is_trainable)
 
     model = None
     if is_trainable and model_args.use_unsloth:
@@ -76,7 +86,7 @@ def load_model_and_tokenizer(
             logger.warning("Unsloth does not support loading adapters.")
 
     if model is None:
-        model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, config=config, **config_kwargs)
+        model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, config=config, **init_kwargs)
 
     patch_model(model, tokenizer, model_args, is_trainable)
     register_autoclass(config, model, tokenizer)
@@ -105,14 +115,13 @@ def load_model_and_tokenizer(
         model.train()
 
     trainable_params, all_param = count_parameters(model)
-    logger.info(
-        "trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
+    if is_trainable:
+        param_stats = "trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
             trainable_params, all_param, 100 * trainable_params / all_param
         )
-    )
-
-    if not is_trainable:
-        logger.info("This IS expected that the trainable params is 0 if you are using model for inference only.")
+    else:
+        param_stats = "all params: {:d}".format(all_param)
+    logger.info(param_stats)
 
     if model_args.print_param_status:
         for name, param in model.named_parameters():
@@ -122,4 +131,18 @@ def load_model_and_tokenizer(
                 )
             )
 
+    return model
+
+
+def load_model_and_tokenizer(
+    model_args: "ModelArguments",
+    finetuning_args: "FinetuningArguments",
+    is_trainable: Optional[bool] = False,
+    add_valuehead: Optional[bool] = False,
+) -> Tuple["PreTrainedModel", "PreTrainedTokenizer"]:
+    r"""
+    Loads pretrained model and tokenizer.
+    """
+    tokenizer = load_tokenizer(model_args)
+    model = load_model(tokenizer, model_args, finetuning_args, is_trainable, add_valuehead)
     return model, tokenizer