Need help for Applying Visual Prompt Tuning with Qwen2.5-VL vision
Author: davidan208
Pull request: no
Labels: name: Vision

Hi everyone,

I am trying to add soft_prompt for Qwen2.5-VL vision
This is my code for doing this:

```
from transformers import Qwen2_5_VLModel, Qwen2_5_VLForConditionalGeneration, AutoProcessor
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
    Qwen2_5_VisionTransformerPretrainedModel,
    Qwen2_5_VLTextModel,
    Qwen2_5_VLPreTrainedModel
)
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
    Qwen2_5_VLVisionConfig,
    Qwen2_5_VLConfig,
    Qwen2_5_VLTextConfig
)

import torch
import torch.nn as nn
import torch.nn.functional as F
from qwen_vl_utils import process_vision_info
import os


class custom_Qwen2_5_VLVisionConfig(Qwen2_5_VLVisionConfig):
    def __init__(self, soft_prompt = None, **kwargs):
        super().__init__(**kwargs)
        self.soft_prompt = soft_prompt

class custom_Qwen2_5_VLConfig(Qwen2_5_VLConfig):
    sub_configs = {"vision_config": custom_Qwen2_5_VLVisionConfig, "text_config": Qwen2_5_VLTextConfig}
    
class custom_Qwen2_5_VLPreTrainedModel(Qwen2_5_VLPreTrainedModel):
    config: custom_Qwen2_5_VLConfig

class custom_Qwen2_5_VisionTransformerPretrainedModel( custom_Qwen2_5_VLPreTrainedModel, Qwen2_5_VisionTransformerPretrainedModel):
    config: custom_Qwen2_5_VLVisionConfig
    def __init__(self, config, *inputs, **kwargs) -> None:
        super().__init__(config, *inputs, **kwargs)
        visual_soft_prompt = config.soft_prompt

        base_dir = config.recall_path

        if isinstance(visual_soft_prompt, list):
            visual_soft_prompt =  [os.path.normpath(os.path.join(base_dir, path)) for path in visual_soft_prompt]
        
        if len(visual_soft_prompt) != 4:
            raise "This is custom made for Qwen2.5-VL only"
        
        sp7  = torch.load(visual_soft_prompt[0], map_location = "cpu", weights_only= True)
        sp15 = torch.load(visual_soft_prompt[1], map_location = "cpu", weights_only= True)
        sp23 = torch.load(visual_soft_prompt[2], map_location = "cpu", weights_only= True)
        sp31 = torch.load(visual_soft_prompt[3], map_location = "cpu", weights_only= True)

        print("=" * 50)
        print(any(torch.isnan(tensor).any() for tensor in [sp7, sp15, sp23, sp31])) # False
        print("=" * 50)

        temp_shape = sp7.shape
        if any(tensor_.shape != temp_shape for tensor_ in [sp7, sp15, sp23, sp31]):
            raise ValueError("Soft prompts got different shape")
        
        self.register_parameter('soft_prompt_visual_layer_7', nn.Parameter(sp7))
        self.register_parameter('soft_prompt_visual_layer_15', nn.Parameter(sp15))
        self.register_parameter('soft_prompt_visual_layer_23', nn.Parameter(sp23))
        self.register_parameter('soft_prompt_visual_layer_31', nn.Parameter(sp31))

        for name, param in self.named_parameters():
            if 'soft_prompt' in name:
                print(f"  {name}: NaN={torch.isnan(param).any()}")

        self.soft_prompts = {
            "layer_7": self.soft_prompt_visual_layer_7,
            "layer_15": self.soft_prompt_visual_layer_15,
            "layer_23": self.soft_prompt_visual_layer_23,
            "layer_31": self.soft_prompt_visual_layer_31,
        }

        self.num_prompts = temp_shape[0]
    
    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
        hidden_states = self.patch_embed(hidden_states)

        rotary_pos_emb = self.rot_pos_emb(grid_thw)
        window_index, cu_window_seqlens = self.get_window_index(grid_thw)
        
        cu_window_seqlens = torch.tensor(
            cu_window_seqlens, device=hidden_states.device, dtype=torch.int32
        )
        cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)

        seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
        hidden_states = hidden_states[window_index, :, :]
        hidden_states = hidden_states.reshape(seq_len, -1)
        
        rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
        rotary_pos_emb = rotary_pos_emb[window_index, :, :]
        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
        
        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        original_cos, original_sin = emb.cos(), emb.sin()
        
        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(dim=0, dtype=torch.int32)
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

        head_dim = original_cos.shape[-1]
        prompt_cos = torch.ones(self.num_prompts, head_dim, device=hidden_states.device, dtype=hidden_states.dtype)
        prompt_sin = torch.zeros(self.num_prompts, head_dim, device=hidden_states.device, dtype=hidden_states.dtype)

        for layer_num, blk in enumerate(self.blocks):
            if layer_num in self.fullatt_block_indexes:
                match layer_num:
                    case 7:
                        prompts_to_add = self.soft_prompts["layer_7"].squeeze(0).to(hidden_states.device, hidden_states.dtype)
                    case 15:
                        prompts_to_add = self.soft_prompts["layer_15"].squeeze(0).to(hidden_states.device, hidden_states.dtype)
                    case 23:
                        prompts_to_add = self.soft_prompts["layer_23"].squeeze(0).to(hidden_states.device, hidden_states.dtype)
                    case 31:
                        prompts_to_add = self.soft_prompts["layer_31"].squeeze(0).to(hidden_states.device, hidden_states.dtype)

                current_hidden_states = torch.cat([prompts_to_add, hidden_states], dim=0)
                final_cos = torch.cat([prompt_cos, original_cos], dim=0)
                final_sin = torch.cat([prompt_sin, original_sin], dim=0)
                current_pos_embeds = (final_cos, final_sin)

                cu_seqlens_now = cu_seqlens.clone()
                if cu_seqlens_now.numel() > 1:
                    cu_seqlens_now[1:] += self.num_prompts

                output_hidden_states = blk(
                    current_hidden_states,
                    cu_seqlens=cu_seqlens_now,
                    position_embeddings=current_pos_embeds,
                    len_soft_prompts = self.num_prompts,
                    **kwargs,
                )

                hidden_states = output_hidden_states[self.num_prompts:, :]
            else:
                hidden_states = blk(
                    hidden_states,
                    cu_seqlens=cu_window_seqlens,
                    position_embeddings=(original_cos, original_sin),
                    len_soft_prompts = 0,
                    **kwargs,
                )
            
        hidden_states = self.merger(hidden_states)
        reverse_indices = torch.argsort(window_index)
        hidden_states = hidden_states[reverse_indices, :]
        return hidden_states

class custom_Qwen2_5_VLModel(Qwen2_5_VLModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.config.vision_config.recall_path = getattr(self.config, "_name_or_path", None)

        self.visual = custom_Qwen2_5_VisionTransformerPretrainedModel._from_config(config = config.vision_config)
        self.language_model = Qwen2_5_VLTextModel._from_config(config.text_config)

class custom_Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)

        self.model = custom_Qwen2_5_VLModel(config)

model = custom_Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "../2.5/",
    torch_dtype = torch.float32,
    low_cpu_mem_usage = True,
    trust_remote_code = True,
    device_map = "auto"
)


for name, param in model.named_parameters():
    if 'soft_prompt' in name:
        print(torch.isnan(param).any())
    

# tensor(False, device='mps:0')
# tensor(True, device='mps:0')
# tensor(True, device='mps:0')
# tensor(True, device='mps:0')

path = model.config.vision_config.recall_path + model.config.vision_config.soft_prompt[0][2:]
torch.load(path, map_location="cpu", weights_only=True)

# Parameter containing:
# tensor([[ 0.0072, -0.0287, -0.0344,  ...,  0.0534,  0.0571, -0.0062],
#         [-0.0627,  0.0050, -0.0253,  ..., -0.0560,  0.0535,  0.0007],
#         [ 0.0635, -0.0559,  0.0682,  ..., -0.0279, -0.0299, -0.0121],
#         ...,
#         [ 0.0361,  0.0380,  0.0426,  ..., -0.0013, -0.0570, -0.0163],
#         [-0.0547, -0.0663, -0.0223,  ..., -0.0156,  0.0142,  0.0348],
#         [ 0.0061, -0.0403, -0.0181,  ...,  0.0198, -0.0238,  0.0105]],
#        requires_grad=True)

for name, param in model.named_parameters():
    if 'soft_prompt_visual_layer_7' in name:
        print(param)

# Parameter containing:
# tensor([[0., 0., 0.,  ..., 0., 0., 0.],
#         [0., 0., 0.,  ..., 0., 0., 0.],
#         [0., 0., 0.,  ..., 0., 0., 0.],
#         ...,
#         [0., 0., 0.,  ..., 0., 0., 0.],
#         [0., 0., 0.,  ..., 0., 0., 0.],
#         [0., 0., 0.,  ..., 0., 0., 0.]], device='mps:0', requires_grad=True)
```

I found that the parameter or register seems to create new all 0 tensors, or sometimes NaN which lead to error in.
I really need to apply this for a final year project, please help.
Thank you all 

Comments:
Looking at your code, seem like the issue is that when you register parameters with register_parameter() and then move the model to a device, PyTorch reinitializes the parameters

Hi @LawJarp-A, I believe so. Do you have any solutions or keywords that I could look into to solve this problem. Thank you.

@davidan208 
I think a simple fix is to first load everything on CPU (if you have the RAM to do that), then move the model
```
model = custom_Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "../2.5/",
    torch_dtype=torch.float32,
    low_cpu_mem_usage=False,  
    device_map=None,         
    trust_remote_code=True,
)

# move to target device
model.to("mps (or) cuda")  
```

OR
 If you want to keep the mem usage to tru, 

```
model = custom_Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "../2.5/",
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="auto",
    trust_remote_code=True,
)

# re-inject your soft prompt weights AFTER load
with torch.no_grad():
    vis = model.model.visua

    def reload_param(param, path):
        t = torch.load(path, map_location="cpu", weights_only=True)
        t = t.to(device=param.device, dtype=param.dtype)
        param.data.copy_(t)

    reload_param(vis.soft_prompt_visual_layer_7,  vis.config.soft_prompt[0])
    reload_param(vis.soft_prompt_visual_layer_15, vis.config.soft_prompt[1])
    reload_param(vis.soft_prompt_visual_layer_23, vis.config.soft_prompt[2])
    reload_param(vis.soft_prompt_visual_layer_31, vis.config.soft_prompt[3])

# sanity check
for n, p in model.named_parameters():
    if "soft_prompt" in n:
        print(n, p.device, torch.isnan(p).any().item(), p.abs().mean().item())
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.