[model] update kt code (#9406)

src/llamafactory/model/adapter.py

@@ -167,7 +167,7 @@ def _setup_lora_tuning(
             is_mergeable = False
 
         if model_args.use_kt:
-            assert len(model_args.adapter_name_or_path) == 1, "Up to now, KTransformers model only accepts a single adapter, for more features, you can contact with us."
+            assert len(model_args.adapter_name_or_path) == 1, "KTransformers model only accepts a single adapter"
             is_mergeable = False
 
         if model_args.use_unsloth:
@@ -190,7 +190,9 @@ def _setup_lora_tuning(
 
         if model_args.use_kt:
             if model_args.infer_backend != EngineName.KT:
-                raise ValueError("We should use ktransformers as backend to infer the adapter fine-tuned by ktransformers.")
+                raise ValueError(
+                    "We should use ktransformers as backend to infer the adapter fine-tuned by ktransformers."
+                )
 
         for adapter in adapter_to_merge:
             model: LoraModel = PeftModel.from_pretrained(model, adapter, **init_kwargs)
@@ -218,9 +220,9 @@ def _setup_lora_tuning(
         if model_args.use_kt:
             new_list = []
             for m in target_modules:
-                if m in ('down_proj', 'up_proj', 'gate_proj'):
+                if m in ("down_proj", "up_proj", "gate_proj"):
                     new_list.extend([f"mlp.{m}", f"shared_experts.{m}"])
-                elif m not in ('generate_linear', 'orig_module', 'prefill_linear'):
+                elif m not in ("generate_linear", "orig_module", "prefill_linear"):
                     new_list.append(m)
 
             target_modules[:] = new_list

src/llamafactory/model/loader.py

@@ -146,6 +146,7 @@ def load_model(
     lazy_load = False
     if model_args.use_kt:
         from ktransformers.sft.monkey_patch_torch_module import install_patch
+
         install_patch()
         model = load_kt_pretrained_model(config, model_args)
     elif model_args.use_unsloth:

src/llamafactory/model/model_utils/attention.py

@@ -59,6 +59,7 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
         requested_attn_implementation = "sdpa"
     elif model_args.flash_attn == AttentionFunction.FA2:
         from transformers import is_torch_npu_available
+
         if not (is_flash_attn_2_available() or is_torch_npu_available()):
             logger.warning_rank0("FlashAttention-2 is not installed.")
             return

src/llamafactory/model/model_utils/ktransformers.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import importlib.util as _u
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import torch
 
@@ -43,6 +43,7 @@ if KT_AVAILABLE:
 
 logger = logging.get_logger(__name__)
 
+
 def _get_kt_kwargs(
     config: "PretrainedConfig",
     model_name_or_path: str,
@@ -64,9 +65,7 @@ def _get_kt_kwargs(
     }
 
 
-def load_kt_pretrained_model(
-    config: "PretrainedConfig", model_args: "ModelArguments"
-) -> Optional["PreTrainedModel"]:
+def load_kt_pretrained_model(config: "PretrainedConfig", model_args: "ModelArguments") -> "PreTrainedModel":
     r"""Optionally load pretrained model with KTransformers. Used in training."""
     custom_models = {
         "DeepseekV2ForCausalLM": DeepseekV2ForCausalLM,
@@ -79,7 +78,7 @@ def load_kt_pretrained_model(
     Config().chunk_size = model_args.chunk_size
     config = AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code)
 
-    if model_args.mode == 'long_context':
+    if model_args.mode == "long_context":
         assert config.architectures[0] == "LlamaForCausalLM", "only LlamaForCausalLM support long_context mode"
         torch.set_default_dtype(torch.float16)
     else:
@@ -88,9 +87,7 @@ def load_kt_pretrained_model(
     with torch.device("meta"):
         if config.architectures[0] in custom_models:
             print("using custom modeling_xxx.py.")
-            if (
-                "Qwen2Moe" in config.architectures[0]
-            ):  # Qwen2Moe must use flash_attention_2 to avoid overflow.
+            if "Qwen2Moe" in config.architectures[0]:  # Qwen2Moe must use flash_attention_2 to avoid overflow.
                 config._attn_implementation = "flash_attention_2"
             if "Llama" in config.architectures[0]:
                 config._attn_implementation = "eager"
@@ -115,21 +112,17 @@ def load_kt_pretrained_model(
     return model
 
 
-def get_kt_peft_model(
-    model: "PreTrainedModel", peft_kwargs: dict[str, Any]
-) -> "PreTrainedModel":
+def get_kt_peft_model(model: "PreTrainedModel", peft_kwargs: dict[str, Any]) -> "PreTrainedModel":
     r"""Get the peft model for the pretrained model with KTransformers. Used in training."""
     from ktransformers.sft.peft_utils.mapping import get_peft_model
 
     return get_peft_model(model, peft_kwargs)
 
 
-def load_kt_peft_model(
-    model_args: "ModelArguments", model: "PreTrainedModel",
-) -> "PreTrainedModel":
+def load_kt_peft_model(model_args: "ModelArguments", model: "PreTrainedModel") -> "PreTrainedModel":
     r"""Load peft model with KTransformers. Used in both training and inference."""
     load_adapter_name_or_path = model_args.adapter_name_or_path[0]
-    if load_adapter_name_or_path.endswith('.gguf'):
+    if load_adapter_name_or_path.endswith(".gguf"):
         inject_lora_layer(model, load_adapter_name_or_path)
         adapter_gguf_loader = GGUFLoader(load_adapter_name_or_path)
         load_weights(model, adapter_gguf_loader, adapter_gguf=True)