Provide support and guidance for training vision-LLM with a customized vision encoder, projector, and LLM.
Author: lucasjinreal
Pull request: no
Labels: name: Feature request

### Feature request

Provide support and guidance for training vision-LLM with a customized vision encoder, projector, and LLM.

### Motivation

Provide support and guidance for training vision-LLM with a customized vision encoder, projector, and LLM.

### Your contribution

Provide support and guidance for training vision-LLM with a customized vision encoder, projector, and LLM.

Comments:
hello @lucasjinreal 

To train a vision-LLM with a customized vision encoder, projector, and language model (LLM), you need to:

-Define and integrate a custom vision encoder.
-Implement a projector to map vision features to the text embedding space.
-Use a pre-existing language model (LLM) or customize it for the task.
-Ensure compatibility between the components (e.g., embeddings dimensions).
-Provide guidance for loading datasets, fine-tuning, and training.

Here’s how you can proceed:


Step 1: Custom Vision Encoder
Implement or load a custom vision encoder using transformers or torchvision:
If you’re using a pretrained vision model from transformers, you can modify the existing AutoModel for vision models (e.g., ViT, Swin).
Alternatively, you can use models from torchvision.models (e.g., ResNet, EfficientNet).
Ensure the vision encoder outputs feature embeddings suitable for the projector.

Example:

```
from transformers import AutoModel
from torchvision.models import resnet50

# Using a pretrained vision model from transformers
vision_encoder = AutoModel.from_pretrained("google/vit-base-patch16-224")

# Or using a model from torchvision
vision_encoder = resnet50(pretrained=True)
```
Step 2: Custom Projector

Create a projector module to map the vision encoder’s output to the LLM’s input embedding space.
Use a simple feedforward network with optional activation (e.g., ReLU) to perform this mapping.

Example:

```
import torch.nn as nn

class VisionToTextProjector(nn.Module):
    def __init__(self, vision_dim, text_dim):
        super().__init__()
        self.projector = nn.Sequential(
            nn.Linear(vision_dim, text_dim),
            nn.ReLU(),
            nn.Linear(text_dim, text_dim)
        )

    def forward(self, vision_features):
        return self.projector(vision_features)
```
        
Step 3: Language Model (LLM)
Load a language model (e.g., LLaMA, GPT, BERT) using the transformers library.
Ensure the LLM can accept embeddings from the projector.

Example:

```
from transformers import AutoModelForCausalLM

# Load a pretrained language model
llm = AutoModelForCausalLM.from_pretrained("llava-hf/llava-1.5-7b-hf")
```

Step 4: Combine the Modules
Create a wrapper class to integrate the vision encoder, projector, and LLM.
Ensure the forward pass aligns vision and text inputs.

Example:

```
class VisionLLM(nn.Module):
    def __init__(self, vision_encoder, projector, llm):
        super().__init__()
        self.vision_encoder = vision_encoder
        self.projector = projector
        self.llm = llm

    def forward(self, vision_inputs, text_inputs):
        # Extract vision features
        vision_features = self.vision_encoder(vision_inputs).last_hidden_state
        # Project vision features to text embedding space
        projected_features = self.projector(vision_features)
        # Pass text and vision embeddings to the language model
        outputs = self.llm(inputs_embeds=projected_features, labels=text_inputs)
        return outputs.loss, outputs.logits
```
        
Step 5: Dataset Preparation
Use a multimodal dataset containing both images and text.
Tokenize the text using the LLM’s tokenizer and preprocess the images using the vision encoder’s processor.

Example:

```
from transformers import AutoTokenizer, AutoImageProcessor

# Tokenizer for the LLM
tokenizer = AutoTokenizer.from_pretrained("llava-hf/llava-1.5-7b-hf")
# Image processor for the vision encoder
image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")

# Example preprocessing
image_inputs = image_processor(images, return_tensors="pt")
text_inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
```

Step 6: Training
Define a loss function and optimizer.
Use the Trainer API from transformers or write a custom PyTorch training loop.

Example with Trainer:

```
from transformers import Trainer, TrainingArguments

# Define training arguments
training_args = TrainingArguments(
    output_dir="./vision-llm",
    per_device_train_batch_size=8,
    learning_rate=5e-5,
    num_train_epochs=3,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    logging_dir="./logs",
    logging_steps=10
)

# Define the model
vision_llm = VisionLLM(vision_encoder, projector, llm)

# Initialize the Trainer
trainer = Trainer(
    model=vision_llm,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer
)

# Train the model
trainer.train()
```


@elxandre Hi, thank u so much for the detailed answer.

However, I already made a training script similar to your suggestions, but I found the training didn't goes well. Still can't found a clue. So I am just looking for something inside transformers would do this job.

This is current my not working code:

```python

class MyForConditionalGeneration(MyPretrainedModel, Qwen2ForCausalLM):
    def __init__(self, config: MonoConfig):
        # super().__init__(config)
        MyPretrainedModel.__init__(self, config)
        super(Qwen2ForCausalLM, self).__init__(config)

        self.vision_tower = ViTLarge(config=config.vision_config)
        self._attn_implementation = config._attn_implementation

        self._build_image_projection_layers(config)

        self.model = Qwen2Model(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.pad_token_id = config.pad_token_id
        print(f"==> pad_token_id: {self.pad_token_id}")
        self.post_init()

def _build_image_projection_layers(self, config):
        image_dim_out = config.vision_config.hidden_size
        dim_projection = config.hidden_size
        self.mm_projector = nn.Linear(image_dim_out, dim_projection)
        print(f"==> build mm_projector: {image_dim_out} -> {dim_projection}")

    def get_vision_tower(self):
        return self.vision_tower

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def _encode_image(self, pixel_values):
        # print(f"pixel_values: {pixel_values}")
        batch_size, C, H, W = pixel_values.shape
        x = self.vision_tower(pixel_values, output_hidden_states=True)
        x = x[-2]
        x = self.mm_projector(x)
        # print(f"image features: {x}")
        return x

    def _get_input_embeds_with_image(self, input_ids, image_features):
        # 1. replace image token with features; 2. replace -100 in input_ids into zeroes
        # 3. handling right attention_mask
        batch_size = input_ids.size(0)
        processed_embeds = []
        processed_masks = []

        max_seq_len = 0
        for idx in range(batch_size):
            seq = input_ids[idx]
            im_pos = (seq == -200).nonzero(as_tuple=True)[0]

            if im_pos.numel() > 0:
                im_pos = im_pos.item()
                before = seq[:im_pos]
                after = seq[im_pos + 1 :]
                # Exclude -100 tokens (maybe, input_ids padding with -100 intentionly)
                before = before[before != -100]
                after = after[after != -100]
                # Get embeddings for before and after
                before_embed = self.get_input_embeddings()(before)
                after_embed = self.get_input_embeddings()(after)
                # Concatenate before, image features, and after
                seq_embed = torch.cat(
                    [before_embed, image_features[idx], after_embed], dim=0
                )
                new_seq_len = seq_embed.size(0)
            else:
                # Exclude -100 tokens
                valid_tokens = seq[seq != -100]
                seq_embed = self.get_input_embeddings()(valid_tokens)
                new_seq_len = seq_embed.size(0)

            # Update the maximum sequence length
            if new_seq_len > max_seq_len:
                max_seq_len = new_seq_len

            processed_embeds.append(seq_embed)
            attn_mask = torch.ones(new_seq_len, dtype=torch.bool, device=seq.device)
            processed_masks.append(attn_mask)

        # rest embedding is 0, rest mask is False, just padding it
        inputs_embeds = torch.nn.utils.rnn.pad_sequence(
            processed_embeds, batch_first=True, padding_value=0.0
        )
        attn_masks = torch.nn.utils.rnn.pad_sequence(
            processed_masks, batch_first=True, padding_value=0
        )
        return inputs_embeds, attn_masks

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: torch.FloatTensor = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        inputs_embeds: Optional[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,
        cache_position=None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        image_features = None
        if inputs_embeds is None:
            if pixel_values is not None:
                # (batch_size, num_image_tokens, hidden_size)
                image_features = self._encode_image(pixel_values)

            if input_ids is not None:
                inputs_embeds, attention_mask = self._get_input_embeds_with_image(
                    input_ids, image_features
                )

        # padding all to normal sequence length only train
        if labels is not None:
            input_length = inputs_embeds.shape[1]
            label_length = labels.shape[1]

            if labels is not None:
                labels = F.pad(labels, (input_length, 0), value=-100)

            if inputs_embeds is not None:
                # append embeds and attn_mask to labels length
                padding = torch.zeros(
                    inputs_embeds.shape[0],
                    label_length,
                    inputs_embeds.shape[2],
                    dtype=inputs_embeds.dtype,
                    device=inputs_embeds.device,
                )
                inputs_embeds = torch.cat([inputs_embeds, padding], dim=1)
                attention_mask = attention_mask.to(inputs_embeds.dtype)
                attention_mask = F.pad(attention_mask, (0, label_length), value=0)

        return super().forward(
            input_ids=None,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            labels=labels,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
```
This code is responsible for handling input_ids. In this context, a -200 value for the image position placeholder indicates the replacement of image feature tokens.

The remaining task is simply to input inputs_embeds into the Language Learning Model (LLM).

Nevertheless, the training loss remains unchanging and does not decrease at all.

As a result, the following inquiries arise:

1. Is there any specific aspect where my implementation might have gone wrong in detail?
2. If my code is not functioning as expected, is there any functionality within the transformers framework that could assist me in accomplishing my task? (It is actually similar to llava, but I do not require such complexity. I simply desire a ViT (Vision Transformer) + LLM combination and aim to train a simple caption).


@lucasjinreal hey! This question is better places on the forum (https://discuss.huggingface.co/latest), as we try to reserve GH for bugs and feature requests.

Since your model is similar to llava, feel free to use th `LlavaForConditionalGeneration` class providing your won vision and text configs to change the LM and vision backbones. For further question please ask in the forum :)

@lucasjinreal , are you able to create custom vision LLM model using the above code ? 

Hi, please reference to my latest project: https://github.com/lucasjinreal/Namo-R1

I have sucessfully trained a 500M VLM surpassed SmolVLM

@lucasjinreal Thank you so much. I have a custom encoder model trained on my internal dataset. I want to create a smaller VLM like yours. Do you have a sample script for fine-tuning the custom vision LLM model using my custom encoder?

What is the minimum dataset size required to train the custom vision LLM model?

Yes you can, just follow my code add your vision encoder into it.

I think at least 4M images could be robust.

@lucasjinreal Thank you so much.