fix llava qlora

Former-commit-id: 01c5a669f6fe598aac1758a700a7607da37db1bc

src/llmtuner/model/patcher.py

@@ -15,6 +15,7 @@ from .utils.longlora import configure_longlora
 from .utils.moe import add_z3_leaf_module, configure_moe
 from .utils.quantization import configure_quantization
 from .utils.rope import configure_rope
+from .utils.visual import autocast_projector_dtype
 
 
 if TYPE_CHECKING:
@@ -92,6 +93,9 @@ def patch_model(
     if model_args.resize_vocab:
         resize_embedding_layer(model, tokenizer)
 
+    if model_args.visual_inputs:
+        autocast_projector_dtype(model, model_args)
+
     if is_trainable:
         prepare_model_for_training(model, model_args)
         add_z3_leaf_module(model)

src/llmtuner/model/utils/visual.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from ...extras.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+    from ...hparams import ModelArguments
+
+
+logger = get_logger(__name__)
+
+
+def autocast_projector_dtype(
+    model: "PreTrainedModel", model_args: "ModelArguments", mm_projector_name: str = "multi_modal_projector"
+) -> None:
+    def _mm_projector_forward_post_hook(
+        module: "torch.nn.Module", args: Tuple["torch.Tensor"], output: "torch.Tensor"
+    ) -> "torch.Tensor":
+        return output.to(model_args.compute_dtype)
+
+    if hasattr(model, mm_projector_name):
+        logger.info("Casting multimodal projector outputs in {}.".format(model_args.compute_dtype))
+        mm_projector: "torch.nn.Module" = getattr(model, mm_projector_name)
+        mm_projector.register_forward_hook(_mm_projector_forward_post_hook)