Pre-train Multi-Modal Language model LLaVA

We will look at the steps needed to pre-train a multi modal language model like LLaVA.

If you would like to know how to run the model refer to my earlier article the link to which is here.

Also if you like to build the wheel file for FlashAttention refer to my other article in here.

Pre-requisties:

1. A100 Nvidia GPU
2. Hard disk of atleast 150 GB

Software requirements are(not to mentiona Nvidia eco-system to be in place):

Nvidia driver of atleast 580 (sudo apt-get -y install nvidia-driver-580)

1. nccl library from Nvidia (pip install nvidia-nccl-cu13)
2. huggingface_hub = 0.36.2
3. bitsandbytes (pip3 install bitsandbytes — upgrade)
4. LLaVA version 1.5
5. torch = 2.7.1
6. Flash-Attention version 2.8.3 built for Pytorch 2.7.1, CUDA 12.9
7. Addition pip versions that worked for me( nvtx=0.2.15, deepspeed 0.19.0)

The above is what worked for me.

As part of the preparation one needs to download the Model files like llama,VICUNA , CLIP etc… The below script does the job for us…

#!/usr/bin/env bash
#auhtor:rangapv@yahoo.com
#20-05-2026

chkpt_dwn(){

lavatmp="lava-temp1/LLaVA"

chk0=`pip3 install --upgrade huggingface_hub`
chk1=`hf auth login --no-add-to-git-credential`

env1=`export HF_HUB_DISABLE_SYMLINKS=1`
env2=`export HF_HUB_DISABLE_SYMLINKS_WARNING=True`
chk2=`hf download meta-llama/Meta-Llama-3-8B-Instruct --local-dir ./$lavatmp/checkpoints/Meta-Llama-3-8B-Instruct`

chk3=`hf download openai/clip-vit-large-patch14-336 --local-dir ./$lavatmp/checkpoints/clip-vit-large-patch14-336`

chk4=`hf download liuhaotian/llava-v1.5-mlp2x-336px-pretrain-vicuna-7b-v1.5 --local-dir ./$lavatmp/checkpoints/llava-pretrain-projector`

chk5=`hf download lmsys/vicuna-7b-v1.5 --local-dir ./checkpoints/vicuna-7b-v1.5 --local-dir ./$lavatmp/checkpoints/ `

chk5=`cd ./$lavatmp/playground/data; wget https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain/resolve/main/blip_laion_cc_sbu_558k.json`

}

chkpt_dwn

Now get the main package with all the pre-training and fine tuning scripts..

$ wget https://github.com/haotian-liu/LLaVA/archive/refs/tags/v1.1.3.tar.gz

$ tar -xvf  ./v1.1.3.tar.gz

$ cd LLaVA-1.1.3

$ pip3 install --upgrade pip

$ pip3 install -e .

$ pip3 install -e ".[train]"

For the verisons that we have chosen to make it work in the eco-ssytem the following files needs to be updated in the particular line numbers…

ubuntu@ip-172-31-47-180:~/temp1/LLaVA-1.1.3/llava/train$ ls -l
total 64
-rw-rw-r-- 1 ubuntu ubuntu  4404 Oct 26  2023 llama_flash_attn_monkey_patch.py
-rw-rw-r-- 1 ubuntu ubuntu 11445 Oct 26  2023 llava_trainer.py
-rw-rw-r-- 1 ubuntu ubuntu 37091 Oct 26  2023 train.py
-rw-rw-r-- 1 ubuntu ubuntu   498 Oct 26  2023 train_mem.py

In the file : llama_flash_attn_monkey_patch.py

`change at line 85:

qkv, indices, cu_q_lens, max_s = unpad_input(qkv, key_padding_mask)

qkv, indices, cu_q_lens, maxs, * = unpad_input(qkv, key_padding_mask)`

In the file: train.py:

line 260:

#tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()

`tokenized.input_ids.ne(tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0).sum().item()

line 723:

padding_value=self.tokenizer.pad_token_id)

padding_value=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0)

Line: 732`

#attention_mask=input_ids.ne(self.tokenizer.pad_token_id),

attention_mask=input_ids.ne(self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0),

Finally we would need Flash-Attention version 2.8.3.

$ wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.7cxx11abiTRUE-cp310-cp310-linux_x86_64.whl

$ pip3 install ./flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl

Store the images in the “/images” folder and the annotations as “blip.json” under the “playground/data”. A sample is shown below…

  {
    "id": "006577897",
    "image": "00657/006577897.jpg",
    "conversations": [
      {
        "from": "human",
        "value": "Share a concise interpretation of the image provided.\n<image>"
      },
      {
        "from": "gpt",
        "value": "the xbox 360 game shift 3 unleashed"
      }
    ]
  }
]

Now everything in place we run the pre-training…

deepspeed llava/train/train_mem.py \
    --deepspeed ./scripts/zero2.json \
    --model_name_or_path ./checkpoints/Meta-Llama-3-8B-Instruct \
    --version plain \
    --data_path ./playground/data/blip.json \
    --image_folder ./images \
    --vision_tower openai/clip-vit-large-patch14-336 \
    --mm_projector_type mlp2x_gelu \
    --tune_mm_mlp_adapter True \
    --mm_vision_select_layer -2 \
    --mm_use_im_start_end False \
    --mm_use_im_patch_token False \
    --bf16 True \
    --output_dir ./checkpoints/llava-v1.5-13b-pretrain \
    --num_train_epochs 1 \
    --per_device_train_batch_size 32 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 24000 \
    --save_total_limit 1 \
    --learning_rate 1e-3 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --tf32 True \
    --model_max_length 2048 \
    --gradient_checkpointing True \
    --dataloader_num_workers 4 \
    --lazy_preprocess True \
    --report_to wandb

The output will be in the output directory as “mm_projector.bin” . This is the projector output that we have for the newly trained images.

$ cd results/
ubuntu@ip-172-31-47-180:~/LLaVA-1.1.3/results$ ls -l
total 4
drwxrwxr-x 3 ubuntu ubuntu 4096 May 24 08:39 llava-llama3-8b-pretrain
ubuntu@ip-172-31-47-180:~/LLaVA-1.1.3/results$ cd llava-llama3-8b-pretrain/
ubuntu@ip-172-31-47-180:~/LLaVA-1.1.3/results/llava-llama3-8b-pretrain$ ls
config.json  merge.py  merge2.py  merged_model  mm_projector.bin  test.py  trainer_state.json
ubuntu@ip-172-31-47-180:~/LLaVA-1.1.3/results/llava-llama3-8b-pretrain$ ls -l
total 41004
-rw-rw-r-- 1 ubuntu ubuntu     1177 May 24 08:08 config.json
-rwxrwxrwx 1 ubuntu ubuntu      579 May 24 08:20 merge.py
-rwxrwxrwx 1 ubuntu ubuntu     1406 May 24 08:30 merge2.py
drwxrwxr-x 2 ubuntu ubuntu     4096 May 24 08:32 merged_model
-rw-rw-r-- 1 ubuntu ubuntu 41961935 May 24 08:08 mm_projector.bin
-rwxrwxrwx 1 ubuntu ubuntu      423 May 24 08:24 test.py
-rw-rw-r-- 1 ubuntu ubuntu      663 May 24 08:08 trainer_state.json

We can then merge the weights to the base model then use it for inferencing. The script to merge the bin to the Base model weights file is as shown below…

#!/usr/bin/env python3

import torch
from llava.model import LlavaLlamaForCausalLM
from transformers import AutoTokenizer

model_path = "../../checkpoints/Meta-Llama-3-8B-Instruct"
mm_projector_path = "./mm_projector.bin"
output_path = "./merged_model"

# Load model
model = LlavaLlamaForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
)

# Load mm_projector weights
mm_weights = torch.load(mm_projector_path, map_location="cpu")

# Check model has these keys
model_keys = [k for k in model.state_dict().keys() if "mm_projector" in k]
print("Model projector keys:", model_keys)

if len(model_keys) == 0:
    # Model not initialized with vision — manually add the weights
    model.model.mm_projector = torch.nn.Sequential(
        torch.nn.Linear(mm_weights["model.mm_projector.0.weight"].shape[1],
                        mm_weights["model.mm_projector.0.weight"].shape[0]),
        torch.nn.GELU(),
        torch.nn.Linear(mm_weights["model.mm_projector.2.weight"].shape[1],
                        mm_weights["model.mm_projector.2.weight"].shape[0]),
    )
    model.load_state_dict(mm_weights, strict=False)
else:
    # Keys exist, just load
    model.load_state_dict(mm_weights, strict=False)

# Save
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)
print("Saved merged model to:", output_path)

After running the script , we get the final model with our trained images…

$ cd merged_model/
ubuntu@ip-172-31-47-180:~/LLaVA-1.1.3/results/llava-llama3-8b-pretrain/merged_model$ ls -l
total 15775152
-rw-rw-r-- 1 ubuntu ubuntu        739 May 24 08:31 config.json
-rw-rw-r-- 1 ubuntu ubuntu        194 May 24 08:31 generation_config.json
-rw-rw-r-- 1 ubuntu ubuntu 9976557522 May 24 08:32 pytorch_model-00001-of-00002.bin
-rw-rw-r-- 1 ubuntu ubuntu 6168005941 May 24 08:32 pytorch_model-00002-of-00002.bin
-rw-rw-r-- 1 ubuntu ubuntu      27068 May 24 08:32 pytorch_model.bin.index.json
-rw-rw-r-- 1 ubuntu ubuntu         73 May 24 08:32 special_tokens_map.json
-rw-rw-r-- 1 ubuntu ubuntu    9085671 May 24 08:32 tokenizer.json
-rw-rw-r-- 1 ubuntu ubuntu      50977 May 24 08:32 tokenizer_config.json

If you have any issues, you can open an issue in the github repo page or alternatively contact me at rangapv@yahoo.com and you can also find me on X.com(twitter) @rangapv . My linkedin ***https://linkedin.com/in/rangapv***

References:

If you like to know the variuos components and functionality of the LLaVA based Multi model:

[embed]MLX & CUDA examples with Vision encoder for MultiModal Model like LLaVA to perform as Visual... LLaVA - Large Language and Vision Assistant is an end-to-end trained large multimodal model that connects a vision…rangapv.medium.com

[embed][LLaVA/scripts/v1_5/pretrain.sh at main · haotian-liu/LLaVA NeurIPS'23 Oral] Visual Instruction Tuning (LLaVA) built towards GPT-4V level capabilities and beyond. …github.com](https://github.com/haotian-liu/LLaVA/blob/main/scripts/v1_5/pretrain.sh)

Some of my other related article on Medium:

[embed]GPU based Decision Optimization to solve Linear Programming using cuOpt: Calculate CVaR and VaR GPU based Decision Optimization to solve Linear Programming using cuOpt: Calculate CVaR and VaR In this article we will…rangapv.medium.com

TensorRT_LLM Recipes(Quantization & KV Cache):

[embed]TensorRT_LLM Recipes-1 In this article we are going to look at some of the ways to call the TenosrRT_LLM API with parameters like…rangapv.medium.com

[embed]Nvidia cuda and TensorRT install on a Tesla GPU for Deeplearning Model inference In this article we are going to install TensorRT on a Tesla T4 GPU with all the required packages like the cuda…rangapv.medium.com