Check the following code thoroughly. Suggest the best way to handle GPU parallelism for torch.cat of inputs_embeds (from lm_model) showing batch size of 4 and image_embeds (from clip_model) showing batch size of 16, in a 4 gpu setup.
#preprocess_image2.py
import json
import torch
from torch.utils.data import Dataset
from PIL import Image
import os
from torchvision import transforms

class LlavaDataset(Dataset):
    def __init__(self, json_file, image_dir, tokenizer, max_length=512):
        with open(json_file, 'r') as f:
            self.data = json.load(f)
        self.image_dir = image_dir
        self.tokenizer = tokenizer
        self.max_length = max_length

        # Define image transformations
        self.image_transforms = transforms.Compose([
            transforms.ToTensor()
        ])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]

        # Load and preprocess image
        if 'image' in item:
            image_file = os.path.join(self.image_dir, item['image'])
            if not os.path.exists(image_file):
                print(f"Warning: Image file does not exist: {image_file}")
                return None
            image = Image.open(image_file).convert('RGB')
            # image = self.image_transforms(image)
        else:
            print(f"Warning: 'image' key not found in item {idx}")
            return None

        # Construct prompt from conversation
        conversation = item.get('conversations', [])
        prompt = ""
        for turn in conversation:
            if turn['from'] == 'human':
                prompt += f"Human: {turn['value']}\n"
            elif turn['from'] == 'gpt':
                prompt += f"Assistant: {turn['value']}\n"

        # Replace '<image>' token with a special token that your model understands
        prompt = prompt.replace('<image>', '[IMAGE]')

        # Tokenize the prompt
        encoded_prompt = self.tokenizer(
            prompt,
            truncation=True,
            max_length=self.max_length,
            padding='max_length',
            return_tensors='pt',
        )

        return {
            'image': image,
            'input_ids': encoded_prompt['input_ids'].squeeze(0),
            'attention_mask': encoded_prompt['attention_mask'].squeeze(0)
        }

def load_llava_dataset(json_file, image_dir, tokenizer, max_length=2048):
    return LlavaDataset(json_file, image_dir, tokenizer, max_length)

#train_img2.py
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3,5,6,7"
import torch
from torch import nn
from transformers import CLIPModel, CLIPProcessor, Trainer, TrainingArguments, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model
from peft.optimizers import create_loraplus_optimizer
import bitsandbytes as bnb
from torch.amp import autocast
from preprocess_image2 import load_llava_dataset
import tempfile

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
class ImageModel(nn.Module):
    def __init__(self, clip_model_name, lm_model_name, tokenizer):
        super().__init__()
        # Load CLIP model
        self.clip_model = CLIPModel.from_pretrained(
            clip_model_name,
            torch_dtype=torch.bfloat16,
        ) 
        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_name, torch_dtype=torch.bfloat16,)
        #bnb quantized model config
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4"
        )
        
        # Load language model
        self.lm_model = AutoModelForCausalLM.from_pretrained(
            lm_model_name, 
            quantization_config=bnb_config,
            device_map=None,
            torch_dtype=torch.bfloat16,
        )
        self.tokenizer = tokenizer

        # Freeze CLIP model
        for param in self.clip_model.parameters():
            param.requires_grad = False

        # Optionally, freeze the language model to train only the adapter
        for param in self.lm_model.parameters():
            param.requires_grad = False
        # max_length = self.lm_model.config.max_position_embeddings
        # new_max_length = max_length + 1  # Increase by 1 for the image token
        # self.lm_model.resize_position_embeddings(new_max_length)

        # Define a projection layer to map image features to language model's hidden size
        self.image_proj = nn.Linear(
            self.clip_model.config.projection_dim,
            self.lm_model.config.hidden_size,
        )
        # self.image_proj = self.image_proj.to(dtype=torch.bfloat16)

        # For the loss function
        self.loss_fn = nn.CrossEntropyLoss()
        self.lm_model = prepare_model_for_kbit_training(self.lm_model)
        self.lm_model.enable_input_require_grads()
        # Define LoRA configuration tailored for QLoRA
        peft_config = LoraConfig(
            r=8,
            lora_alpha=32,
            lora_dropout=0.1,
            bias="none",
            # Target the appropriate modules in the language model for LoRA
            target_modules=['q_proj', 'k_proj', 'v_proj', 'out_proj', 'fc1', 'fc2'],
            # Specify modules to save when saving the adapter
            modules_to_save=["lm_head", "embed_tokens"],
        )
        # Apply PEFT (QLoRA) to the language model
        self.lm_model = get_peft_model(self.lm_model, peft_config)
        # Optional: Print trainable parameters to verify
        # print(self.lm_model.print_trainable_parameters())

    def forward(self, image, input_ids, attention_mask, labels=None):
        # inputs_embeds = self.lm_model.get_input_embeddings()(input_ids)
        # target_device = inputs_embeds.device
        with autocast(device_type="cuda",dtype=torch.bfloat16):
            target_device = input_ids.device
            # Process images
            with torch.no_grad():
                processed = self.clip_processor(images=image, return_tensors="pt", do_rescale=False, padding=True)
                pixel_values = processed['pixel_values'].to(target_device)
                image_features = self.clip_model.get_image_features(
                    pixel_values=pixel_values
                )

            # Project image features to match LM's hidden size
            projected_image_features = self.image_proj(image_features)

            # Prepare inputs for language model
            # Assuming you want to prepend the image features to the text embeddings
            batch_size = input_ids.size(0)
            print("Batch Size of input_ids:", batch_size)
            seq_length = input_ids.size(1)
            print("Batch Size of input_ids:", seq_length)
            print("Input ids shape: ",input_ids.shape)
            # Expand image features to match sequence length if needed
            # Here, we'll concatenate the image features as a new token at the beginning
            # Get embeddings from the language model
            # self.lm_model = nn.DataParallel(self.lm_model)
            inputs_embeds = self.lm_model.get_input_embeddings()(input_ids)
            # Reshape text embeddings
            # inputs_embeds = inputs_embeds.view(8, 1024, 2560)  # Reshape to [batch_size, feature_dim, hidden_dim]
            print("Input embeds shape: ", inputs_embeds.shape)
            # Concatenate image features to the embeddings
            # Reshape projected_image_features to (batch_size, 1, hidden_size)
            image_embeds = projected_image_features.unsqueeze(1)
            print("Image embeds shape:",image_embeds.shape)
            inputs_embeds = torch.cat([image_embeds, inputs_embeds], dim=1)

            # Adjust attention_mask accordingly
            extended_attention_mask = torch.cat(
                [torch.ones(batch_size, 1, device=target_device), attention_mask.to(target_device)],
                dim=1
            )

            # Pass through language model
            outputs = self.lm_model(
                inputs_embeds=inputs_embeds,
                attention_mask=extended_attention_mask,
                labels=labels.to(target_device) if labels is not None else None
            )

        return outputs

def main():
    # Device setup
    torch.cuda.empty_cache()
    # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # Tokenizer and special tokens
    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
    tokenizer.add_special_tokens({'additional_special_tokens': ['[IMAGE]']})

    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})

    # num_labels = tokenizer.vocab_size

    # Load the image model
    model = ImageModel(
        clip_model_name="openai/clip-vit-large-patch14-336",
        lm_model_name="microsoft/phi-2",
        tokenizer=tokenizer
    )
    # model = prepare_model_for_kbit_training(model)
    model.lm_model.gradient_checkpointing_enable()
    # model = model.to(device)
    # Inspect model's modules to find correct target module names
    # print("Model's named modules:")
    # for name, module in model.named_modules():
    #     print(name)
    # print(model.lm_model)

    # # LoRA configuration and setup
    # peft_config = LoraConfig(
    #     r=8,
    #     lora_alpha=32,
    #     lora_dropout=0.1,
    #     bias="none",
    #     target_modules=['Wqkv','out_proj', 'fc1', 'fc2'],#"lm_model.transformer.h.*.mlp.*", "lm_model.transformer.h.*.attn.*"], # ["lm_model.*"],
    #     modules_to_save=["lm_model.lm_head","lm_model.model.embed_tokens"],
    #     # target_modules=["adapter.0","adapater.2"],  # Target the adapter layers
    # )
    # model = get_peft_model(model, peft_config)
    # # Inside ImageModel after applying get_peft_model
    # print(self.lm_model.print_trainable_parameters())
    for name, param in model.named_parameters():
        if param.device != torch.device("cuda:0"):
            print(f"Parameter {name} is on {param.device}")

    # Data loading
    train_dataset = load_llava_dataset(
        json_file='data/llava_instruct_150k/llava_instruct_150k.json',
        image_dir='data/coco/train2017',
        tokenizer=tokenizer,
    )

    # Training arguments
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=1,
        per_device_train_batch_size=4,  # Adjust based on your GPU memory
        gradient_accumulation_steps=8,
        warmup_steps=100,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
        # save_steps=500,
        bf16=True,
        optim="paged_adamw_8bit",
        report_to="tensorboard",
        # gradient_checkpointing=True,
        remove_unused_columns=False,
        dataloader_num_workers=4,  # Adjust based on your system
        run_name="llava_multi_gpu_training",  # Optional: Specify a run name
        save_total_limit=3,  # Optional: Limit the number of saved checkpoints
        # load_best_model_at_end=True,  # Optional: Load the best model at the end of training
    )

    optimizer = create_loraplus_optimizer(
        model=model,
        optimizer_cls=bnb.optim.PagedAdamW8bit,
        lr=5e-5,
        loraplus_lr_ratio=16,
    )
    scheduler = None

    # Custom collate function
    def collate_fn(batch):
        batch = [item for item in batch if item is not None]
        if not batch:
            return None
        images = [item['image'] for item in batch]
        input_ids = torch.stack([item['input_ids'] for item in batch])
        attention_mask = torch.stack([item['attention_mask'] for item in batch])

        # Adjust labels to account for the image token
        labels = input_ids.clone()
        pad_token_id = tokenizer.pad_token_id
        labels = torch.cat([torch.full((labels.size(0), 1), pad_token_id, dtype=torch.long), labels], dim=1)

        return {
            'image': images,
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'labels': labels,
        }

    # Trainer initialization
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        optimizers=(optimizer, scheduler),
        data_collator=collate_fn,
    )

    # Start training
    trainer.train()

    # Save the final model
    trainer.save_model("./final_model/llava150k_qlora_phi2")
    model.lm_model.save_pretrained("./final_model/llava150k_phi2_adapter")

if __name__ == "__main__":
    main()