refactor mllm param logic
Former-commit-id: b895c190945cf5d991cb4e4dea2ae73cc9c8d246
This commit is contained in:
@@ -396,8 +396,7 @@ _register_template(
|
||||
format_user=StringFormatter(slots=["### Instruction:\n{{content}}\n\n### Response:\n"]),
|
||||
format_assistant=StringFormatter(slots=["{{content}}", {"eos_token"}, "\n\n"]),
|
||||
default_system=(
|
||||
"Below is an instruction that describes a task. "
|
||||
"Write a response that appropriately completes the request.\n\n"
|
||||
"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
|
||||
),
|
||||
replace_jinja_template=True,
|
||||
)
|
||||
|
||||
@@ -364,6 +364,10 @@ class FinetuningArguments(
|
||||
default=True,
|
||||
metadata={"help": "Whether ot not to freeze vision tower in MLLM training."},
|
||||
)
|
||||
freeze_multi_modal_projector: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "Whether or not to freeze the multi modal projector in MLLM training."},
|
||||
)
|
||||
train_mm_proj_only: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Whether or not to train the multimodal projector for MLLM only."},
|
||||
@@ -398,6 +402,7 @@ class FinetuningArguments(
|
||||
self.additional_target: Optional[List[str]] = split_arg(self.additional_target)
|
||||
self.galore_target: List[str] = split_arg(self.galore_target)
|
||||
self.freeze_vision_tower = self.freeze_vision_tower or self.train_mm_proj_only
|
||||
self.freeze_multi_modal_projector = self.freeze_multi_modal_projector and not self.train_mm_proj_only
|
||||
self.use_ref_model = self.stage == "dpo" and self.pref_loss not in ["orpo", "simpo"]
|
||||
|
||||
assert self.finetuning_type in ["lora", "freeze", "full"], "Invalid fine-tuning method."
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict
|
||||
|
||||
import torch
|
||||
@@ -202,12 +203,8 @@ def load_model(
|
||||
|
||||
logger.info_rank0(param_stats)
|
||||
|
||||
if model_args.print_param_status:
|
||||
if model_args.print_param_status and int(os.getenv("LOCAL_RANK", "0")) == 0:
|
||||
for name, param in model.named_parameters():
|
||||
print(
|
||||
"name: {}, dtype: {}, device: {}, trainable: {}".format(
|
||||
name, param.dtype, param.device, param.requires_grad
|
||||
)
|
||||
)
|
||||
print(f"name: {name}, dtype: {param.dtype}, device: {param.device}, trainable: {param.requires_grad}")
|
||||
|
||||
return model
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
from typing import TYPE_CHECKING, List
|
||||
|
||||
from ...extras import logging
|
||||
from .visual import COMPOSITE_MODELS
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -34,18 +35,12 @@ def find_all_linear_modules(model: "PreTrainedModel", freeze_vision_tower: bool)
|
||||
forbidden_modules.add("output_layer")
|
||||
elif model_type == "internlm2":
|
||||
forbidden_modules.add("output")
|
||||
elif model_type in ["llava", "llava_next", "llava_next_video", "mllama", "paligemma", "video_llava"]:
|
||||
forbidden_modules.add("multi_modal_projector")
|
||||
elif model_type == "qwen2_vl":
|
||||
forbidden_modules.add("merger")
|
||||
|
||||
if freeze_vision_tower:
|
||||
if model_type == "mllama":
|
||||
forbidden_modules.add("vision_model")
|
||||
elif model_type == "qwen2_vl":
|
||||
forbidden_modules.add("visual")
|
||||
else:
|
||||
forbidden_modules.add("vision_tower")
|
||||
if model_type in COMPOSITE_MODELS:
|
||||
forbidden_modules.add(COMPOSITE_MODELS[model_type].projector_key)
|
||||
|
||||
if freeze_vision_tower and model_type in COMPOSITE_MODELS:
|
||||
forbidden_modules.update(COMPOSITE_MODELS[model_type].vision_model_keys)
|
||||
|
||||
module_names = set()
|
||||
for name, module in model.named_modules():
|
||||
|
||||
@@ -15,7 +15,8 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import TYPE_CHECKING, List, Sequence, Set, Tuple, Union
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Set, Tuple, Union
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
@@ -35,6 +36,40 @@ logger = logging.get_logger(__name__)
|
||||
transformers_logger = transformers.utils.logging.get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompositeModel:
|
||||
model_type: str
|
||||
projector_key: str
|
||||
vision_model_keys: List[str]
|
||||
language_model_keys: List[str]
|
||||
|
||||
def get_projector(self, module: "torch.nn.Module") -> "torch.nn.Module":
|
||||
for key in self.projector_key.split("."):
|
||||
module = getattr(module, key)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
COMPOSITE_MODELS: Dict[str, "CompositeModel"] = {}
|
||||
|
||||
|
||||
def _register_composite_model(
|
||||
model_type: str,
|
||||
projector_key: Optional[str] = None,
|
||||
vision_model_keys: Optional[List[str]] = None,
|
||||
language_model_keys: Optional[List[str]] = None,
|
||||
):
|
||||
projector_key = projector_key or "multi_modal_projector"
|
||||
vision_model_keys = vision_model_keys or ["vision_tower"]
|
||||
language_model_keys = language_model_keys or ["language_model"]
|
||||
COMPOSITE_MODELS[model_type] = CompositeModel(
|
||||
model_type=model_type,
|
||||
projector_key=projector_key,
|
||||
vision_model_keys=vision_model_keys,
|
||||
language_model_keys=language_model_keys,
|
||||
)
|
||||
|
||||
|
||||
class LlavaMultiModalProjectorForYiVL(torch.nn.Module):
|
||||
def __init__(self, config: "LlavaConfig") -> None:
|
||||
super().__init__()
|
||||
@@ -92,10 +127,8 @@ def autocast_projector_dtype(model: "PreTrainedModel", model_args: "ModelArgumen
|
||||
|
||||
if getattr(model, "quantization_method", None):
|
||||
model_type = getattr(model.config, "model_type", None)
|
||||
if model_type in ["llava", "llava_next", "llava_next_video", "mllama", "paligemma", "video_llava"]:
|
||||
mm_projector: "torch.nn.Module" = getattr(model, "multi_modal_projector")
|
||||
elif model_type == "qwen2_vl":
|
||||
mm_projector: "torch.nn.Module" = getattr(getattr(model, "visual"), "merger")
|
||||
if model_type in COMPOSITE_MODELS:
|
||||
mm_projector = COMPOSITE_MODELS[model_type].get_projector(model)
|
||||
else:
|
||||
return
|
||||
|
||||
@@ -107,8 +140,7 @@ def configure_visual_model(config: "PretrainedConfig") -> None:
|
||||
r"""
|
||||
Patches VLMs before loading them.
|
||||
"""
|
||||
model_type = getattr(config, "model_type", None)
|
||||
if model_type in ["llava", "llava_next", "llava_next_video", "mllama", "paligemma", "video_llava"]:
|
||||
if getattr(config, "text_config", None) and not getattr(config, "hidden_size", None):
|
||||
# required for ds zero3 and valuehead models
|
||||
setattr(config, "hidden_size", getattr(config.text_config, "hidden_size", None))
|
||||
|
||||
@@ -123,25 +155,21 @@ def get_forbidden_modules(config: "PretrainedConfig", finetuning_args: "Finetuni
|
||||
"""
|
||||
model_type = getattr(config, "model_type", None)
|
||||
forbidden_modules = set()
|
||||
if model_type in ["llava", "llava_next", "llava_next_video", "paligemma", "video_llava"]:
|
||||
if model_type in COMPOSITE_MODELS:
|
||||
if finetuning_args.freeze_vision_tower:
|
||||
forbidden_modules.add("vision_tower")
|
||||
vision_model_keys = COMPOSITE_MODELS[model_type].vision_model_keys
|
||||
logger.info_rank0(f"Set vision model not trainable: {vision_model_keys}.")
|
||||
forbidden_modules.update(vision_model_keys)
|
||||
|
||||
if finetuning_args.freeze_multi_modal_projector:
|
||||
projector_key = COMPOSITE_MODELS[model_type].projector_key
|
||||
logger.info_rank0(f"Set multi model projector not trainable: {projector_key}.")
|
||||
forbidden_modules.add(projector_key)
|
||||
|
||||
if finetuning_args.train_mm_proj_only:
|
||||
forbidden_modules.add("language_model")
|
||||
|
||||
elif model_type == "mllama":
|
||||
if finetuning_args.freeze_vision_tower:
|
||||
forbidden_modules.add("vision_model")
|
||||
|
||||
if finetuning_args.train_mm_proj_only:
|
||||
forbidden_modules.add("language_model")
|
||||
|
||||
elif model_type == "qwen2_vl":
|
||||
if finetuning_args.train_mm_proj_only:
|
||||
forbidden_modules.update({"visual.patch_embed", "visual.blocks", "model", "lm_head"})
|
||||
elif finetuning_args.freeze_vision_tower:
|
||||
forbidden_modules.add("visual")
|
||||
language_model_keys = COMPOSITE_MODELS[model_type].language_model_keys
|
||||
logger.info_rank0(f"Set language model not trainable: {language_model_keys}.")
|
||||
forbidden_modules.update(language_model_keys)
|
||||
|
||||
return forbidden_modules
|
||||
|
||||
@@ -190,18 +218,57 @@ def patch_target_modules(
|
||||
model_type = getattr(config, "model_type", None)
|
||||
vit_model_type = getattr(getattr(config, "vision_config", None), "model_type", None)
|
||||
if finetuning_args.freeze_vision_tower:
|
||||
if model_type in ["llava", "llava_next", "llava_next_video", "paligemma", "video_llava"]:
|
||||
return "^(?!.*vision_tower).*(?:{}).*".format("|".join(target_modules))
|
||||
elif model_type == "mllama":
|
||||
return "^(?!.*vision_model).*(?:{}).*".format("|".join(target_modules))
|
||||
elif model_type == "qwen2_vl":
|
||||
return "^(?!.*visual).*(?:{}).*".format("|".join(target_modules))
|
||||
if model_type in COMPOSITE_MODELS:
|
||||
vision_model_keys = COMPOSITE_MODELS[model_type].vision_model_keys
|
||||
logger.info_rank0(f"Set vision model not trainable: {vision_model_keys}.")
|
||||
vision_model_keys = "|".join(vision_model_keys)
|
||||
target_modules = "|".join(target_modules)
|
||||
return f"^(?!.*{vision_model_keys}).*(?:{target_modules}).*"
|
||||
else:
|
||||
return target_modules
|
||||
else:
|
||||
if model_type == "qwen2_vl":
|
||||
if model_type == "qwen2_vl": # avoid attaching lora to Conv3D layer
|
||||
return "^(?!.*patch_embed).*(?:{}).*".format("|".join(target_modules))
|
||||
elif vit_model_type == "pixtral":
|
||||
return "^(?!.*patch_conv).*(?:{}).*".format("|".join(target_modules))
|
||||
else:
|
||||
return target_modules
|
||||
|
||||
|
||||
_register_composite_model(
|
||||
model_type="llava",
|
||||
)
|
||||
|
||||
|
||||
_register_composite_model(
|
||||
model_type="llava_next",
|
||||
)
|
||||
|
||||
|
||||
_register_composite_model(
|
||||
model_type="llava_next_video",
|
||||
)
|
||||
|
||||
|
||||
_register_composite_model(
|
||||
model_type="paligemma",
|
||||
)
|
||||
|
||||
|
||||
_register_composite_model(
|
||||
model_type="video_llava",
|
||||
)
|
||||
|
||||
|
||||
_register_composite_model(
|
||||
model_type="mllama",
|
||||
vision_model_keys=["vision_model"],
|
||||
)
|
||||
|
||||
|
||||
_register_composite_model(
|
||||
model_type="qwen2_vl",
|
||||
projector_key="visual.merger",
|
||||
vision_model_keys=["visual.patch_embed", "visual.blocks"],
|
||||
language_model_keys=["model", "lm_head"],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user