refactor mllm param logic

Former-commit-id: b895c190945cf5d991cb4e4dea2ae73cc9c8d246
This commit is contained in:
hiyouga
2025-01-10 15:41:54 +00:00
parent 1675712a4c
commit dc65ecdf09
10 changed files with 198 additions and 62 deletions

View File

@@ -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,
)

View File

@@ -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."

View File

@@ -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

View File

@@ -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():

View File

@@ -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"],
)