modeling_llmjpvl.py

# --------------------------------------------------------
# LLM-jp-VL
# Copyright (c) 2026 LLM-jp
# Licensed under The Apache License 2.0 [see LICENSE for details]
#
# Originally based on InternVL
# Copyright (c) 2024 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

from typing import List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import AutoModelForCausalLM, GenerationConfig, SiglipVisionModel
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,
    MoeCausalLMOutputWithPast,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.models.gpt_oss.modeling_gpt_oss import load_balancing_loss_func
from transformers.utils import logging

from .configuration_llmjpvl import LLMjpVLConfig
from .constants import IMG_CONTEXT_TOKEN

try:
    import flash_attn  # noqa: F401

    has_flash_attn = True
except ImportError:
    has_flash_attn = False

logger = logging.get_logger(__name__)


class LLMjpVLModel(PreTrainedModel):
    config_class = LLMjpVLConfig
    main_input_name = "pixel_values"
    base_model_prefix = "language_model"
    _supports_flash_attn_2 = True
    supports_gradient_checkpointing = True
    accepts_loss_kwargs = False
    # support transformers 4.51.+
    _tp_plan = ""

    @classmethod
    def can_generate(cls):
        return True

    def __init__(
        self,
        config: LLMjpVLConfig,
        vision_backbone=None,
        language_model=None,
        use_flash_attn=True,
    ):
        super().__init__(config)

        image_size = config.force_image_size or config.vision_config.image_size
        patch_size = config.vision_config.patch_size
        self.image_size = image_size
        self.patch_size = patch_size
        self.select_layer = config.select_layer
        self.template = config.template
        self.num_image_token = int(
            (image_size // patch_size) ** 2 * (config.downsample_ratio**2)
        )
        self.downsample_ratio = config.downsample_ratio
        use_flash_attn = use_flash_attn if has_flash_attn else False
        config.vision_config.use_flash_attn = True if use_flash_attn else False
        config.vision_config._attn_implementation = (
            "flash_attention_2" if use_flash_attn else "eager"
        )
        config.llm_config._attn_implementation = (
            "flash_attention_2" if use_flash_attn else "eager"
        )

        logger.info(f"num_image_token: {self.num_image_token}")
        if vision_backbone is not None:
            self.vision_backbone = vision_backbone
        else:
            self.vision_backbone = SiglipVisionModel(config.vision_config)

        if language_model is not None:
            self.language_model = language_model
        else:
            self.language_model = AutoModelForCausalLM.from_config(config.llm_config)
            logger.info(f"language_model type: {type(self.language_model)}")

        vit_hidden_size = config.vision_config.hidden_size
        llm_hidden_size = config.llm_config.hidden_size

        self.mlp1 = nn.Sequential(
            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
            nn.Linear(
                vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size
            ),
            nn.GELU(),
            nn.Linear(llm_hidden_size, llm_hidden_size),
        ).to(torch.bfloat16)

        self.img_context_token_id = getattr(config, "img_context_token_id", None)
        self.tokenizer = None

    def _resolve_img_context_token_id(self):
        """Lazily resolve img_context_token_id from the tokenizer."""
        if self.img_context_token_id is None and self.tokenizer is not None:
            self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(
                IMG_CONTEXT_TOKEN
            )
        return self.img_context_token_id

    def forward(
        self,
        pixel_values: torch.FloatTensor,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        image_flags: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        loss_weight: Optional[List] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        image_flags = image_flags.squeeze(-1)
        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()

        ignore = False
        has_images = pixel_values.shape[0] > 0 and (image_flags == 1).any()
        if has_images:
            vit_embeds = self.extract_feature(pixel_values)
            vit_embeds = vit_embeds[image_flags == 1]

            B, N, C = input_embeds.shape
            input_embeds = input_embeds.reshape(B * N, C)

            input_ids = input_ids.reshape(B * N)
            selected = input_ids == self._resolve_img_context_token_id()
            try:
                input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(
                    -1, C
                )
            except Exception as e:
                vit_embeds = vit_embeds.reshape(-1, C)
                print(
                    f"warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, "
                    f"vit_embeds.shape={vit_embeds.shape}"
                )
                n_token = selected.sum()
                input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]
                ignore = True

            input_embeds = input_embeds.reshape(B, N, C)
        else:
            # No image in this (micro)batch (e.g. pure-text FineVision samples,
            # ~9% of that set). Still run the vision encoder + projector on a
            # dummy patch and add its output with a 0.0 multiplier: this keeps
            # the vision tower in the autograd graph on *every* rank, so its
            # FSDP gradient reduce-scatter collectives fire consistently. If a
            # rank whose whole microbatch is pure-text skipped them, it would
            # desync the data-parallel group -> NCCL watchdog timeout at the
            # next collective (the grad-norm all-reduce in clip_grad_norm_).
            dummy_pixel_values = input_embeds.new_zeros(
                1, 3, self.image_size, self.image_size
            )
            vit_embeds = self.extract_feature(dummy_pixel_values)
            input_embeds = input_embeds + vit_embeds.sum() * 0.0

        outputs = self.language_model(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )
        logits = outputs.logits

        loss = None
        aux_loss = None
        if labels is not None and loss_weight is not None:
            loss_weight = torch.tensor(
                loss_weight, dtype=torch.float32, device=labels.device
            )
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            shift_weights = loss_weight[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss(reduction="none")

            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            shift_weights = shift_weights.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            shift_weights = shift_weights.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

            shift_weights_sum = shift_weights.sum()

            loss = loss * shift_weights
            # clamp_min avoids 0/0 -> NaN when a whole micro-batch is fully
            # masked (no supervised answer tokens, e.g. a long prompt whose
            # answer got truncated at model_max_length). Such a batch then
            # yields loss 0, and its square-avg weight (denom) is also 0, so it
            # contributes nothing to the gradient — instead of poisoning every
            # parameter with NaN (nan*0 == nan, so the `ignore` guard below
            # could not rescue it).
            loss = loss.sum() / shift_weights_sum.clamp_min(1e-8)

        elif labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()

            if (shift_labels == -100).all():
                ignore = True
                shift_labels = shift_labels * 0

            # Flatten the tokens
            loss_fct = CrossEntropyLoss(reduction="none")
            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

            loss_weight = (labels != -100).sum(dim=-1).float()
            loss_weight = 1 / loss_weight.sqrt()
            loss_weight = torch.where(labels != -100, loss_weight.unsqueeze(1), 0.0)

            shift_weights = loss_weight[..., 1:].contiguous()
            shift_weights = shift_weights.view(-1)
            shift_weights = shift_weights.to(shift_logits.device)
            shift_weights_sum = shift_weights.sum()
            loss = loss * shift_weights
            # clamp_min avoids 0/0 -> NaN when a whole micro-batch is fully
            # masked (no supervised answer tokens, e.g. a long prompt whose
            # answer got truncated at model_max_length). Such a batch then
            # yields loss 0, and its square-avg weight (denom) is also 0, so it
            # contributes nothing to the gradient — instead of poisoning every
            # parameter with NaN (nan*0 == nan, so the `ignore` guard below
            # could not rescue it).
            loss = loss.sum() / shift_weights_sum.clamp_min(1e-8)

        if getattr(outputs, "router_logits", None) is not None:
            aux_loss = load_balancing_loss_func(
                outputs.router_logits,
                self.language_model.num_experts,
                self.language_model.num_experts_per_tok,
                attention_mask,
            )

            if loss is not None:
                loss = loss + self.language_model.router_aux_loss_coef * aux_loss.to(
                    loss.device
                )

        if ignore and loss is not None:
            print("[Debug] ignore curr loss")
            loss = loss * 0.0

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        if aux_loss is not None:
            return MoeCausalLMOutputWithPast(
                loss=loss,
                aux_loss=aux_loss,
                logits=logits,
                past_key_values=outputs.past_key_values,
                hidden_states=outputs.hidden_states,
                attentions=outputs.attentions,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def pixel_shuffle(self, x, scale_factor=0.5):
        n, w, h, c = x.size()
        # N, W, H, C --> N, W, H * scale, C // scale
        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
        x = x.permute(0, 2, 1, 3).contiguous()
        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
        x = x.view(
            n,
            int(h * scale_factor),
            int(w * scale_factor),
            int(c / (scale_factor * scale_factor)),
        )
        x = x.permute(0, 2, 1, 3).contiguous()
        return x

    def extract_feature(self, pixel_values):
        if self.select_layer == -1:
            vit_embeds = self.vision_backbone(
                pixel_values=pixel_values
            ).last_hidden_state
        else:
            vit_embeds = self.vision_backbone(
                pixel_values=pixel_values, output_hidden_states=True, return_dict=True
            ).hidden_states[self.select_layer]
        h = w = int(vit_embeds.shape[1] ** 0.5)
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
        vit_embeds = self.mlp1(vit_embeds)
        return vit_embeds

    @torch.no_grad()
    def generate(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        input_ids: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        visual_features: Optional[torch.FloatTensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        output_hidden_states: Optional[bool] = None,
        **generate_kwargs,
    ) -> torch.LongTensor:
        generate_kwargs.pop("token_type_ids", None)
        if generation_config is None:
            generation_config = self.generation_config
        img_context_token_id = self._resolve_img_context_token_id()
        assert img_context_token_id is not None
        if pixel_values is not None:
            if visual_features is not None:
                vit_embeds = visual_features
            else:
                vit_embeds = self.extract_feature(pixel_values)
            input_embeds = self.language_model.get_input_embeddings()(input_ids)
            B, N, C = input_embeds.shape
            input_embeds = input_embeds.reshape(B * N, C)

            input_ids = input_ids.reshape(B * N)
            selected = input_ids == img_context_token_id
            if selected.sum() == 0:
                print(
                    "warning: pixel_values provided but no image context tokens "
                    "found in input_ids, falling back to text-only"
                )
                input_embeds = input_embeds.reshape(B, N, C)
            else:
                input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
                input_embeds = input_embeds.reshape(B, N, C)
        else:
            input_embeds = self.language_model.get_input_embeddings()(input_ids)

        outputs = self.language_model.generate(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            generation_config=generation_config,
            output_hidden_states=output_hidden_states,
            use_cache=True,
            **generate_kwargs,
        )

        return outputs

    @property
    def lm_head(self):
        return self.language_model.get_output_embeddings()

    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        return self.language_model.set_input_embeddings(value)

    def set_output_embeddings(self, value):
        return self.language_model.set_output_embeddings(value)