From 25ebec23b2d25e8ccb731a2c236d725525bff3bf Mon Sep 17 00:00:00 2001 From: njucckevin <827023266@qq.com> Date: Wed, 4 Jun 2025 04:45:47 +0000 Subject: [PATCH] feat: Add Qwen2.5-VL support --- README.md | 7 +- eval/screenSpot.py | 37 ++- eval/screenSpot_pro.py | 37 ++- eval/screenSpot_v2.py | 37 ++- pyproject.toml | 5 +- scripts/train.sh | 1 + src/gui_actor/modeling.py | 4 +- src/gui_actor/modeling_qwen25vl.py | 374 +++++++++++++++++++++++++++++ train.py | 28 ++- 9 files changed, 494 insertions(+), 36 deletions(-) create mode 100644 src/gui_actor/modeling_qwen25vl.py diff --git a/README.md b/README.md index f2089c4..96c1072 100644 --- a/README.md +++ b/README.md @@ -127,7 +127,12 @@ bash scripts/train.sh ``` ## :checkered_flag: Evaluation on GUI Grounding Benchmarks -For evaluation on ScreenSpot and ScreenSpot-v2, you can directly run the scripts under the `scripts/` folder like `python eval/screenSpot.py` or `python eval/screenSpot_v2.py`. +For evaluation on ScreenSpot and ScreenSpot-v2, you can directly run the scripts under the `scripts/` folder like: +```bash +# model_type: qwen2vl or qwen25vl +python eval/screenSpot.py --model_type qwen2vl --model_name_or_path microsoft/GUI-Actor-2B-Qwen2-VL +python eval/screenSpot_v2.py --model_type qwen25vl --model_name_or_path microsoft/GUI-Actor-3B-Qwen2.5-VL +``` For evaluation on ScreenSpot-Pro, you first need to download the data from [here](https://huggingface.co/datasets/likaixin/ScreenSpot-Pro), then run the following command: ```bash diff --git a/eval/screenSpot.py b/eval/screenSpot.py index 0cc6a64..81b48a2 100644 --- a/eval/screenSpot.py +++ b/eval/screenSpot.py @@ -9,6 +9,7 @@ from gui_actor.constants import chat_template from gui_actor.modeling import Qwen2VLForConditionalGenerationWithPointer +from gui_actor.modeling_qwen25vl import Qwen2_5_VLForConditionalGenerationWithPointer from gui_actor.inference import inference, ForceFollowTokensLogitsProcessor from gui_actor.utils import do_boxes_overlap from gui_actor.constants import DEFAULT_POINTER_PAD_TOKEN, DEFAULT_POINTER_END_TOKEN, grounding_system_message @@ -27,19 +28,32 @@ def normalize_bbox(bbox_x1y1x2y2, img_width, img_height): y2 = y2 / img_height return x1, y1, x2, y2 -def evaluate(model_name_or_path, use_placeholder, topk): +def evaluate(model_name_or_path, use_placeholder, topk, model_type="qwen2vl"): # initialize model data_processor = Qwen2VLProcessor.from_pretrained(model_name_or_path) tokenizer = data_processor.tokenizer for k, v in tokenizer.added_tokens_encoder.items(): print(v, k) - model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( - model_name_or_path, - torch_dtype=torch.bfloat16, - device_map="cuda:0", - attn_implementation="flash_attention_2" - ).eval() + if model_type == "qwen2vl": + print("Loading Qwen2-VL") + model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( + model_name_or_path, + torch_dtype=torch.bfloat16, + device_map="cuda:0", + attn_implementation="flash_attention_2" + ).eval() + elif model_type == "qwen25vl": + print("Loading Qwen2.5-VL") + model = Qwen2_5_VLForConditionalGenerationWithPointer.from_pretrained( + model_name_or_path, + torch_dtype=torch.bfloat16, + device_map="cuda:0", + attn_implementation="flash_attention_2" + ).eval() + else: + print("undefine model type") + input() print(f"Loaded model from {model_name_or_path}") logits_processor_pointer = ForceFollowTokensLogitsProcessor( @@ -61,6 +75,12 @@ def evaluate(model_name_or_path, use_placeholder, topk): "forum": "web" } + # special system message for our different models + if model_type == "qwen2vl": + grounding_system_message = "You are a GUI agent. You are given a task and a screenshot of the screen. You need to perform a series of pyautogui actions to complete the task." + elif model_type == "qwen25vl": + grounding_system_message = "You are a GUI agent. Given a screenshot of the current GUI and a human instruction, your task is to locate the screen element that corresponds to the instruction. You should output a PyAutoGUI action that performs a click on the correct position. To indicate the click location, we will use some special tokens, which is used to refer to a visual patch later. For example, you can output: pyautogui.click()." + results = [] for i, example in tqdm(enumerate(dataset), total=len(dataset)): ele = { @@ -248,6 +268,7 @@ def format_cell(cell): """ if __name__ == "__main__": parser = argparse.ArgumentParser() + parser.add_argument("--model_type", type=str, default="qwen2vl", choices=["qwen2vl", "qwen25vl"]) parser.add_argument("--model_name_or_path", type=str, default="microsoft/GUI-Actor-2B-Qwen2-VL") parser.add_argument("--save_path", type=str, default="./") parser.add_argument('--topk', type=int, default=3, help='Topk') @@ -271,7 +292,7 @@ def format_cell(cell): results = json.load(f) else: print(f"Evaluating {args.model_name_or_path}...") - results = evaluate(args.model_name_or_path, args.use_placeholder, args.topk) + results = evaluate(args.model_name_or_path, args.use_placeholder, args.topk, args.model_type) with open(pred_path, "w") as f: json.dump(results, f) print(f"Saved {len(results)} predictions to {pred_path}") diff --git a/eval/screenSpot_pro.py b/eval/screenSpot_pro.py index 0b80469..2c924e4 100644 --- a/eval/screenSpot_pro.py +++ b/eval/screenSpot_pro.py @@ -9,6 +9,7 @@ from PIL import Image from gui_actor.constants import chat_template from gui_actor.modeling import Qwen2VLForConditionalGenerationWithPointer +from gui_actor.modeling_qwen25vl import Qwen2_5_VLForConditionalGenerationWithPointer from gui_actor.inference import inference, ForceFollowTokensLogitsProcessor from gui_actor.utils import do_boxes_overlap from gui_actor.constants import DEFAULT_POINTER_PAD_TOKEN, DEFAULT_POINTER_END_TOKEN, grounding_system_message @@ -27,19 +28,32 @@ def normalize_bbox(bbox_x1y1x2y2, img_width, img_height): y2 = y2 / img_height return x1, y1, x2, y2 -def evaluate(model_name_or_path, data_fn, image_dir, use_placeholder, topk, resize_to_pixels=None): +def evaluate(model_name_or_path, data_fn, image_dir, use_placeholder, topk, resize_to_pixels=None, model_type="qwen2vl"): # initialize model data_processor = Qwen2VLProcessor.from_pretrained(model_name_or_path) tokenizer = data_processor.tokenizer for k, v in tokenizer.added_tokens_encoder.items(): print(v, k) - model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( - model_name_or_path, - torch_dtype=torch.bfloat16, - device_map="cuda:0", - attn_implementation="flash_attention_2" - ).eval() + if model_type == "qwen2vl": + print("Loading Qwen2-VL") + model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( + model_name_or_path, + torch_dtype=torch.bfloat16, + device_map="cuda:0", + attn_implementation="flash_attention_2" + ).eval() + elif model_type == "qwen25vl": + print("Loading Qwen2.5-VL") + model = Qwen2_5_VLForConditionalGenerationWithPointer.from_pretrained( + model_name_or_path, + torch_dtype=torch.bfloat16, + device_map="cuda:0", + attn_implementation="flash_attention_2" + ).eval() + else: + print("undefine model type") + input() print(f"Loaded model from {model_name_or_path}") logits_processor_pointer = ForceFollowTokensLogitsProcessor( @@ -54,6 +68,12 @@ def evaluate(model_name_or_path, data_fn, image_dir, use_placeholder, topk, resi data = json.load(f) print(f"Loaded {len(data)} examples from {data_fn}") + # special system message for our different models + if model_type == "qwen2vl": + grounding_system_message = "You are a GUI agent. You are given a task and a screenshot of the screen. You need to perform a series of pyautogui actions to complete the task." + elif model_type == "qwen25vl": + grounding_system_message = "You are a GUI agent. Given a screenshot of the current GUI and a human instruction, your task is to locate the screen element that corresponds to the instruction. You should output a PyAutoGUI action that performs a click on the correct position. To indicate the click location, we will use some special tokens, which is used to refer to a visual patch later. For example, you can output: pyautogui.click()." + results = [] for i, example in tqdm(enumerate(data), total=len(data)): ele = { @@ -253,6 +273,7 @@ def format_cell(cell): """ if __name__ == "__main__": parser = argparse.ArgumentParser() + parser.add_argument("--model_type", type=str, default="qwen2vl", choices=["qwen2vl", "qwen25vl"]) parser.add_argument("--model_name_or_path", type=str, default="microsoft/GUI-Actor-2B-Qwen2-VL") parser.add_argument("--save_path", type=str, default="./") parser.add_argument("--data_path", type=str, default="/mnt/data/ScreenSpot-Pro") @@ -281,7 +302,7 @@ def format_cell(cell): results = json.load(f) else: print(f"Evaluating {args.model_name_or_path}...") - results = evaluate(args.model_name_or_path, data_fn, image_dir, args.use_placeholder, args.topk, resize_to_pixels) + results = evaluate(args.model_name_or_path, data_fn, image_dir, args.use_placeholder, args.topk, resize_to_pixels, args.model_type) with open(pred_path, "w") as f: json.dump(results, f) print(f"Saved {len(results)} predictions to {pred_path}") diff --git a/eval/screenSpot_v2.py b/eval/screenSpot_v2.py index 9b4126d..e79ddf7 100644 --- a/eval/screenSpot_v2.py +++ b/eval/screenSpot_v2.py @@ -9,6 +9,7 @@ from gui_actor.constants import chat_template from gui_actor.modeling import Qwen2VLForConditionalGenerationWithPointer +from gui_actor.modeling_qwen25vl import Qwen2_5_VLForConditionalGenerationWithPointer from gui_actor.inference import inference, ForceFollowTokensLogitsProcessor from gui_actor.utils import do_boxes_overlap from gui_actor.constants import DEFAULT_POINTER_PAD_TOKEN, DEFAULT_POINTER_END_TOKEN, grounding_system_message @@ -27,19 +28,32 @@ def normalize_bbox(bbox_x1y1x2y2, img_width, img_height): y2 = y2 / img_height return x1, y1, x2, y2 -def evaluate(model_name_or_path, use_placeholder, topk): +def evaluate(model_name_or_path, use_placeholder, topk, model_type="qwen2vl"): # initialize model data_processor = Qwen2VLProcessor.from_pretrained(model_name_or_path) tokenizer = data_processor.tokenizer for k, v in tokenizer.added_tokens_encoder.items(): print(v, k) - model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( - model_name_or_path, - torch_dtype=torch.bfloat16, - device_map="cuda:0", - attn_implementation="flash_attention_2" - ).eval() + if model_type == "qwen2vl": + print("Loading Qwen2-VL") + model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( + model_name_or_path, + torch_dtype=torch.bfloat16, + device_map="cuda:0", + attn_implementation="flash_attention_2" + ).eval() + elif model_type == "qwen25vl": + print("Loading Qwen2.5-VL") + model = Qwen2_5_VLForConditionalGenerationWithPointer.from_pretrained( + model_name_or_path, + torch_dtype=torch.bfloat16, + device_map="cuda:0", + attn_implementation="flash_attention_2" + ).eval() + else: + print("undefine model type") + input() print(f"Loaded model from {model_name_or_path}") logits_processor_pointer = ForceFollowTokensLogitsProcessor( @@ -61,6 +75,12 @@ def evaluate(model_name_or_path, use_placeholder, topk): "forum": "web" } + # special system message for our different models + if model_type == "qwen2vl": + grounding_system_message = "You are a GUI agent. You are given a task and a screenshot of the screen. You need to perform a series of pyautogui actions to complete the task." + elif model_type == "qwen25vl": + grounding_system_message = "You are a GUI agent. Given a screenshot of the current GUI and a human instruction, your task is to locate the screen element that corresponds to the instruction. You should output a PyAutoGUI action that performs a click on the correct position. To indicate the click location, we will use some special tokens, which is used to refer to a visual patch later. For example, you can output: pyautogui.click()." + results = [] for i, example in tqdm(enumerate(dataset), total=len(dataset)): ele = { @@ -248,6 +268,7 @@ def format_cell(cell): """ if __name__ == "__main__": parser = argparse.ArgumentParser() + parser.add_argument("--model_type", type=str, default="qwen2vl", choices=["qwen2vl", "qwen25vl"]) parser.add_argument("--model_name_or_path", type=str, default="microsoft/GUI-Actor-2B-Qwen2-VL") parser.add_argument("--save_path", type=str, default="./") parser.add_argument('--topk', type=int, default=3, help='Topk') @@ -271,7 +292,7 @@ def format_cell(cell): results = json.load(f) else: print(f"Evaluating {args.model_name_or_path}...") - results = evaluate(args.model_name_or_path, args.use_placeholder, args.topk) + results = evaluate(args.model_name_or_path, args.use_placeholder, args.topk, args.model_type) with open(pred_path, "w") as f: json.dump(results, f) print(f"Saved {len(results)} predictions to {pred_path}") diff --git a/pyproject.toml b/pyproject.toml index 6ed073e..a6acfec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,9 +14,10 @@ dependencies = [ "accelerate==1.1.1", "qwen-vl-utils==0.0.8", "deepspeed==0.16.0", - "transformers==4.50.0", + "transformers==4.51.3", "flash-attn", - "wandb==0.18.3" + "wandb==0.18.3", + "datasets>=2.18.0" ] requires-python = ">=3.10,<3.13" readme = "README.md" diff --git a/scripts/train.sh b/scripts/train.sh index 07d2173..8f5a609 100644 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -8,6 +8,7 @@ torchrun --nproc_per_node=4 train.py \ --data_path data/data_config.yaml \ --image_folder "" \ --model_name_or_path ${llm_model} \ + --model_type "qwen2vl" \ --group_by_modality_length True \ --bf16 True \ --output_dir ${output_dir} \ diff --git a/src/gui_actor/modeling.py b/src/gui_actor/modeling.py index 1665310..9b43a65 100644 --- a/src/gui_actor/modeling.py +++ b/src/gui_actor/modeling.py @@ -166,7 +166,7 @@ def forward(self, if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) # shape: (batch_size, seq_len, d_model) if pixel_values is not None: - pixel_values = pixel_values.type(self.visual.get_dtype()) + pixel_values = pixel_values.type(self.visual.dtype) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) n_image_tokens = (input_ids == self.config.image_token_id).sum().item() n_image_features = image_embeds.shape[0] @@ -184,7 +184,7 @@ def forward(self, inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) n_video_tokens = (input_ids == self.config.video_token_id).sum().item() n_video_features = video_embeds.shape[0] diff --git a/src/gui_actor/modeling_qwen25vl.py b/src/gui_actor/modeling_qwen25vl.py new file mode 100644 index 0000000..4bbe935 --- /dev/null +++ b/src/gui_actor/modeling_qwen25vl.py @@ -0,0 +1,374 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLCausalLMOutputWithPast, Qwen2_5_VLForConditionalGeneration +from gui_actor.constants import IGNORE_INDEX +from typing import List, Tuple, Union, Optional +from gui_actor.trainer import rank0_print + +class QwenVLwithVisionHeadOutputWithPast(Qwen2_5_VLCausalLMOutputWithPast): + """ + Output class for Qwen2_5_VL with pointer head, extending the base output class. + + Args: + lm_loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Language modeling loss. + pointer_loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Vision pointer network loss. + pointer_scores (`List[torch.FloatTensor]`, *optional*): + Attention scores from the pointer network, one tensor per batch item. + loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Combined loss (weighted sum of lm_loss and pointer_loss). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores from the language modeling head. + past_key_values, hidden_states, attentions, rope_deltas: + Same as parent class. + """ + def __init__(self, lm_loss=None, pointer_loss=None, pointer_scores=None, *args, **kwargs): + super().__init__(*args, **kwargs) + self.lm_loss = lm_loss + self.pointer_loss = pointer_loss + self.pointer_scores = pointer_scores + + +class VisionHead_MultiPatch(nn.Module): + def __init__(self, d_model, projection_dim, num_attention_heads=8, dropout_rate=0.1): + super().__init__() + self.d_model = d_model + + # Note: We omit additional normalization here because Qwen2VL + # already normalizes hidden states using RMSNorm. + self.projection_enc = nn.Sequential( + nn.Linear(d_model, projection_dim), + nn.GELU(), + nn.Linear(projection_dim, d_model) + ) + self.projection_dec = nn.Sequential( + nn.Linear(d_model, projection_dim), + nn.GELU(), + nn.Linear(projection_dim, d_model) + ) + + # Add self-attention layer for visual features + self.self_attention = nn.MultiheadAttention( + embed_dim=d_model, + num_heads=num_attention_heads, + dropout=dropout_rate, + batch_first=True + ) + + # Layer normalization and residual connection + self.layer_norm = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout_rate) + + def forward(self, + hidden_state_enc, # shape: [n_enc, d_model] where n_enc can vary with image size + hidden_state_dec, # shape: [n_dec, d_model] there can be multiple query in one sample + labels: Optional[torch.Tensor] = None, # shape: [n_dec, n_enc], binary mask of patches in bbox + do_single_patch: bool = False, + ): + + enc_input = hidden_state_enc.unsqueeze(0) + attn_output, _ = self.self_attention( + query=enc_input, + key=enc_input, + value=enc_input, + # attn_mask=attention_mask, + need_weights=False + ) + # Residual connection and layer normalization + hidden_state_enc_ctx = self.layer_norm(enc_input + self.dropout(attn_output)) + # Remove batch dimension + hidden_state_enc_ctx = hidden_state_enc_ctx.squeeze(0) # [n_enc, d_model] + + # Apply the projection networks. + proj_enc = self.projection_enc(hidden_state_enc_ctx) # [n_enc, d_model] + proj_dec = self.projection_dec(hidden_state_dec) # [n_dec, d_model] + + # Compute scaled dot-product attention scores. + # Scaling by sqrt(d_model) is critical regardless of variable n_enc. + scaling = self.d_model ** 0.5 + patch_logits = torch.matmul(proj_dec, proj_enc.transpose(0, 1)) / scaling # [n_dec, n_enc] + + # Softmax normalization is applied along the encoder dimension. + attn_weights = F.softmax(patch_logits, dim=-1) + + loss = None + if (labels is not None) and (not do_single_patch): + epsilon = 1e-8 + labels_float = labels.float() + # Normalize each row to get target probability distribution + target_dist = labels_float / (labels_float.sum(dim=-1, keepdim=True) + epsilon) + + # Apply log_softmax to logits + pred_log_probs = F.log_softmax(patch_logits, dim=-1) + # Use KL divergence as loss + loss = F.kl_div(pred_log_probs, target_dist, reduction='batchmean') + + if do_single_patch and (labels is not None): + loss = F.cross_entropy(attn_scores, labels) + + return attn_weights, loss + + +class Qwen2_5_VLForConditionalGenerationWithPointer(Qwen2_5_VLForConditionalGeneration): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.multi_patch_pointer_head = VisionHead_MultiPatch(self.config.hidden_size, self.config.hidden_size) + self.pointer_loss_weight = kwargs.get("pointer_loss_weight", 1.0) + self.lm_loss_weight = kwargs.get("lm_loss_weight", 1.0) + self.post_init() + + def reset_loss_weights(self, pointer_loss_weight, lm_loss_weight): + self.pointer_loss_weight = pointer_loss_weight + self.lm_loss_weight = lm_loss_weight + + def forward(self, + input_ids: torch.LongTensor = None, # (batch_size, seq_len) + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = 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, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + # Grounding + visual_token_indices_of_coordinates: Optional[torch.Tensor] = None, # shape: (batch_size, n_target); each element is the ground-truth index of the visual token that should be attended to for the corresponding target token + multi_patch_labels: Optional[torch.Tensor] = None, # shape: list [(n_target, n_visual), ...]; binary mask of patches in bbox + if_multi_patch: bool = True, + coordinates: Optional[List[Tuple[float, float]]] = None, + verbose: bool = False) -> Union[Tuple, QwenVLwithVisionHeadOutputWithPast]: + + 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 + + if verbose: + rank0_print(f"input_ids: {input_ids.shape}, {input_ids[0][:5]}...") + rank0_print(f"labels: {labels.shape}, {labels[0][:5]}...") + rank0_print(f"pixel_values: {pixel_values.shape}") + rank0_print(f"image_grid_thw: {image_grid_thw.shape}, {image_grid_thw}") + rank0_print(f"coordinates: {coordinates}") + rank0_print(f"visual_token_indices_of_coordinates: {visual_token_indices_of_coordinates}") + rank0_print(f"return_dict: {return_dict}") + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) # shape: (batch_size, seq_len, d_model) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + image_mask = ( + (input_ids == self.config.image_token_id) + .unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + video_mask = ( + (input_ids == self.config.video_token_id) + .unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, attention_mask + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0 + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + delta = delta.to(position_ids.device) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] # shape: (batch_size, seq_len, d_model) + logits = self.lm_head(hidden_states) + + lm_loss = None + if labels is not None and self.lm_loss_weight > 0: + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + lm_loss = loss_fct(shift_logits, shift_labels) + + + # If vision supervision is requested, process the action head. + pointer_loss = None + pointer_scores = [] + if visual_token_indices_of_coordinates is not None: + batch_size = input_ids.shape[0] + pointer_losses = [] + + # Process each sample individually because the number of visual and target tokens may vary. + for i in range(batch_size): + dummy_target = False + + # Get the token ids and corresponding hidden states for sample i. + token_ids = input_ids[i] # shape: (seq_length,) + hs = hidden_states[i] # shape: (seq_length, d_model) + + # Identify visual tokens indices. + visual_mask = (token_ids == self.config.image_token_id) + visual_indices = torch.nonzero(visual_mask, as_tuple=False).squeeze(-1) # shape: (n_visual,) + + # Identify target tokens (the ones that should attend to visual features). + target_mask = (token_ids == self.config.pointer_pad_token_id) + target_indices = torch.nonzero(target_mask, as_tuple=False).squeeze(-1) + + # If either visual or target tokens are missing, skip this sample. + if visual_indices.numel() == 0: + raise ValueError(f"No visual or target tokens found for sample {i}.") + if target_indices.numel() == 0: + target_indices = torch.tensor([hs.shape[0] - 1]) # take the last token as the dummy target token + gt = torch.tensor([0]).to(hs.device) # take the first visual token as the dummy ground truth + if if_multi_patch: # task the first 4 visual tokens as the ground truth + n_t = target_indices.size(0) # 目标 token 个数 + n_v = visual_indices.size(0) + sample_labels = torch.zeros( + (n_t, n_v), device=hs.device, dtype=torch.float + ) + sample_labels[:, :min(4, n_v)] = 1 + dummy_target = True + else: + # For supervision, we assume that visual_token_indices_of_coordinates[i] is a tensor of shape (n_target,) + # where each element is an integer in the range [0, n_visual-1] indicating the ground-truth visual token. + gt = visual_token_indices_of_coordinates[i].to(hs.device) # shape: (n_target,) + if if_multi_patch: + sample_labels = multi_patch_labels[i] + if sample_labels is None: + n_t = target_indices.size(0) # 目标 token 个数 + n_v = visual_indices.size(0) + sample_labels = torch.zeros( + (n_t, n_v), device=hs.device, dtype=torch.float + ) + sample_labels[:, :min(4, n_v)] = 1 + dummy_target = True + + # Gather the corresponding hidden state representations. + # visual_hidden = hs[visual_indices] # shape: (n_visual, d_model) + visual_embeds = inputs_embeds[i][visual_indices] + target_hidden = hs[target_indices] # shape: (n_target, d_model) + + # Calculate loss for multi-patch mode + if if_multi_patch: + # Ensure the number of targets matches between sample and labels + if sample_labels.shape[0] != target_indices.shape[0]: + raise ValueError(f"Sample {i} has mismatched target counts: {sample_labels.shape[0]} labels but found {target_indices.shape[0]} target tokens") + + # Process using VisionHead_MultiPatch + attn_scores, loss_v = self.multi_patch_pointer_head( + visual_embeds, + target_hidden, + labels=sample_labels + ) + + else: + # Deprecated branch - single patch mode is no longer used + # Run the action head to compute the attention (from target tokens to visual tokens) and its loss. + attn_scores, loss_v = self.pointer_head(visual_embeds, target_hidden, labels=gt) + + pointer_scores.append(attn_scores.detach().cpu()) + + pointer_losses.append(loss_v * 0.0 if dummy_target else loss_v) + + pointer_loss = torch.stack(pointer_losses).mean() + + # Combine the LM loss and vision loss using the provided loss weights. + + if lm_loss is None: + total_loss = pointer_loss + elif pointer_loss is None: + total_loss = lm_loss + else: + total_loss = self.lm_loss_weight * lm_loss + self.pointer_loss_weight * pointer_loss + + if return_dict: + return QwenVLwithVisionHeadOutputWithPast( + lm_loss=lm_loss, + pointer_loss=pointer_loss, + pointer_scores=pointer_scores, + loss=total_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + ) + else: + # When labels are provided, parent's forward returns a tuple with loss as the first element. + if labels is not None: + # Replace the LM loss with the combined loss. + output = (lm_loss, pointer_loss, logits, pointer_scores,) + outputs[1:] + print(f"returning: total_loss, logits, pointer_scores, ...") + return (total_loss,) + output if total_loss is not None else output + else: + return outputs \ No newline at end of file diff --git a/train.py b/train.py index acac540..c943a6a 100644 --- a/train.py +++ b/train.py @@ -26,6 +26,7 @@ ) from gui_actor.modeling import Qwen2VLForConditionalGenerationWithPointer +from gui_actor.modeling_qwen25vl import Qwen2_5_VLForConditionalGenerationWithPointer apply_liger_kernel_to_qwen2_vl() @@ -39,6 +40,7 @@ class ModelArguments: model_name_or_path: Optional[str] = field(default="facebook/opt-125m") flash_attn_2_enabled: bool = field(default=True) + model_type: str = field(default="qwen2vl", metadata={"help": "model type: qwen2vl or qwen25vl"}) @dataclass class DataArguments: @@ -227,13 +229,25 @@ def train(): # rank0_print(f"evaluation_args = {vars(evaluation_args)}\n\n") # set up model - model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( - model_args.model_name_or_path, - cache_dir=training_args.cache_dir, - attn_implementation="flash_attention_2" if model_args.flash_attn_2_enabled else None, - torch_dtype=(torch.bfloat16 if training_args.bf16 else None), - low_cpu_mem_usage=False, - ) + if model_args.model_type == "qwen2vl": + model = Qwen2VLForConditionalGenerationWithPointer.from_pretrained( + model_args.model_name_or_path, + cache_dir=training_args.cache_dir, + attn_implementation="flash_attention_2" if model_args.flash_attn_2_enabled else None, + torch_dtype=(torch.bfloat16 if training_args.bf16 else None), + low_cpu_mem_usage=False, + ) + elif model_args.model_type == "qwen25vl": + model = Qwen2_5_VLForConditionalGenerationWithPointer.from_pretrained( + model_args.model_name_or_path, + cache_dir=training_args.cache_dir, + attn_implementation="flash_attention_2" if model_args.flash_attn_2_enabled else None, + torch_dtype=(torch.bfloat16 if training_args.bf16 else None), + low_cpu_mem_usage=False, + ) + else: + raise ValueError(f"Unsupported model type: {model_args.model_type}") + model.config.use_cache = False model.reset_loss_weights(pointer_loss_weight=training_args.pointer_loss_weight, lm_loss_weight=training_args.lm_loss_weight)