Repository: hiyouga/LLaMA-Factory
File: src/llamafactory/model/model_utils/visual.py (Python)

Code before:
    r"""
    Freezes vision tower and language model for VLM full/freeze tuning.
    """
    model_type = getattr(config, "model_type", None)
    forbidden_modules = set()
    if model_type in ["llava", "llava_next", "llava_next_video", "paligemma", "video_llava"]:
        if finetuning_args.freeze_vision_tower:
            forbidden_modules.add("vision_tower")

        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")

    elif model_type == "minicpmv":
        if finetuning_args.freeze_vision_tower:
            forbidden_modules.add("vpm")
            forbidden_modules.add("apm")
            forbidden_modules.add("resampler")
            forbidden_modules.add("tts")

    return forbidden_modules


def get_image_seqlen(config: "PretrainedConfig") -> int:
    r"""
    Computes the number of special tokens per image.
    """
    model_type = getattr(config, "model_type", None)
    if model_type == "llava":
        image_seqlen = (config.vision_config.image_size // config.vision_config.patch_size) ** 2
        if getattr(config, "vision_feature_select_strategy", "default") == "full":  # add [CLS] token
            image_seqlen += 1
    elif model_type == "paligemma":
        image_seqlen = config.vision_config.num_image_tokens
    else:
        image_seqlen = -1

    return image_seqlen


Review comment (suggestion):
Please rebase this pr and modify the logic here to adhere https://github.com/hiyouga/LLaMA-Factory/pull/6600

about rebase: https://github.com/hiyouga/LLaMA-Factory/blob/main/.github/CONTRIBUTING.md#create-a-pull-request

Also you can use `git commit --amend` to reduce the number of commits

Code after:
    """

    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 getattr(model, "quantization_method", None):
        model_type = getattr(model.config, "model_type", None)
        if model_type in COMPOSITE_MODELS:
            mm_projector = COMPOSITE_MODELS[model_type].get_projector(model)
        else:
            return

        logger.info_rank0(f"Casting multimodal projector outputs in {model_args.compute_dtype}.")
        mm_projector.register_forward_hook(_mm_projector_forward_post_hook)


def configure_visual_model(config: "PretrainedConfig") -> None:
    r"""
    Patches VLMs before loading them.
    """
    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))

    if getattr(config, "is_yi_vl_derived_model", None):
        logger.info_rank0("Detected Yi-VL model, applying projector patch.")
        transformers.models.llava.modeling_llava.LlavaMultiModalProjector = LlavaMultiModalProjectorForYiVL


def get_forbidden_modules(config: "PretrainedConfig", finetuning_args: "FinetuningArguments") -> Set[str]:
    r"""
    Freezes vision tower and language model for VLM full/freeze tuning.
    """
    model_type = getattr(config, "model_type", None)
    forbidden_modules = set()
    if model_type in COMPOSITE_MODELS:
        if finetuning_args.freeze_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:
            language_model_keys = COMPOSITE_MODELS[model_type].language_model_keys
            logger.info_rank0(f"Set language model not trainable: {language_model_keys}.")