diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..245c571 Binary files /dev/null and b/.DS_Store differ diff --git a/.gitignore b/.gitignore index b7faf40..01be6f0 100644 --- a/.gitignore +++ b/.gitignore @@ -205,3 +205,6 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ + +#Folder for storing the generated images +Metrics/ diff --git a/README.md b/README.md index 445986a..38bc2cf 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,24 @@ export NCCL_P2P_LEVEL=NVL # if needed PYTHONUNBUFFERED=1 nohup deepspeed --master_port 29501 qwenvl_run_sqa.py args/qwen.yaml --deepspeed --deepspeed_config ds_config.json > qwenvl.log 2>&1 & ``` +To train the Qwen2-VL-2B model with IVT-LR on M3CoT: + +``` +cd qwen_vl +export CUDA_VISIBLE_DEVICES=0,1,2,3 +export NCCL_P2P_LEVEL=NVL # if needed +PYTHONUNBUFFERED=1 nohup deepspeed --master_port 29502 qwenvl_run_2b.py args/qwen2vl_2b.yaml --deepspeed --deepspeed_config ds_config.json > qwenvl_2b.log 2>&1 & +``` + +To train the Qwen2-VL-2B model with IVT-LR on ScienceQA: + +``` +cd qwen_vl +export CUDA_VISIBLE_DEVICES=0,1,2,3 +export NCCL_P2P_LEVEL=NVL # if needed +PYTHONUNBUFFERED=1 nohup deepspeed --master_port 29502 qwenvl_run_2b_sqa.py args/qwen2vl_2b.yaml --deepspeed --deepspeed_config ds_config.json > qwenvl_2b_sqa.log 2>&1 & +``` + #### Chameleon For Chameleon on M3CoT: @@ -184,6 +202,18 @@ export CUDA_VISIBLE_DEVICES=0 nohup python infer_sqa.py > infer.log 2>&1 & ``` +Qwen2-VL-2B on M3CoT: +``` +export CUDA_VISIBLE_DEVICES=0 +nohup python infer_2b.py --checkpoint_path your_2b_pth_path > infer_2b.log 2>&1 & +``` + +Qwen2-VL-2B on ScienceQA: +``` +export CUDA_VISIBLE_DEVICES=0 +nohup python infer_2b_sqa.py --checkpoint_path your_2b_pth_path > infer_2b_sqa.log 2>&1 & +``` + Chameleon on M3CoT: ``` export CUDA_VISIBLE_DEVICES=0 diff --git a/chameleon/args/chameleon.yaml b/chameleon/args/chameleon.yaml index 12bb09c..a310f62 100644 --- a/chameleon/args/chameleon.yaml +++ b/chameleon/args/chameleon.yaml @@ -15,4 +15,5 @@ batch_size_training: 2 debug: False gradient_accumulation_steps: 8 num_epochs: 16 -lr: !!float "4e-5" \ No newline at end of file +lr: !!float "4e-5" +patch_reuse_policy: never \ No newline at end of file diff --git a/chameleon/chameleon_ivtlr.py b/chameleon/chameleon_ivtlr.py index 915b783..6d5115c 100644 --- a/chameleon/chameleon_ivtlr.py +++ b/chameleon/chameleon_ivtlr.py @@ -29,6 +29,7 @@ def __init__( eos_token_id, image_token_id, num_selected_patches: int = 64, + patch_reuse_policy: str = "never", ): super(IVTLR, self).__init__() @@ -40,6 +41,10 @@ def __init__( self.end_latent_id = end_latent_id self.image_token_id = image_token_id self.num_selected_patches = num_selected_patches + valid_policies = {"never", "next_step_only", "always"} + if patch_reuse_policy not in valid_policies: + raise ValueError(f"Invalid patch_reuse_policy={patch_reuse_policy}. Expected one of {valid_policies}.") + self.patch_reuse_policy = patch_reuse_policy if isinstance(self.base_causallm, GPT2LMHeadModel): self.embedding = self.base_causallm.transformer.get_input_embeddings() @@ -136,6 +141,7 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, original_mask = torch.ones((B, new_S), dtype=torch.bool, device=device) # image_mask no repeated True image_mask = torch.zeros((B, 3000), dtype=torch.bool, device=device) + recently_selected_mask = torch.zeros((B, 3000), dtype=torch.bool, device=device) for b in range(B): s = img_starts[b] image_mask[b, s:s+1024] = True @@ -181,6 +187,7 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, avg_attn = torch.cat(attentions, dim=1).mean(dim=1) # (B, seq_len) current_seq_len = avg_attn.size(1) select_image_embeds = [] + current_selected_mask = torch.zeros_like(image_mask) for b in range(B): @@ -191,6 +198,8 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, scores = last_attn.clone() allowed_positions = image_mask[b, :current_seq_len] # shape=(S,) + if self.patch_reuse_policy == "next_step_only": + allowed_positions = allowed_positions & (~recently_selected_mask[b, :current_seq_len]) invalid = ~allowed_positions scores[invalid] = float("-inf") @@ -201,7 +210,10 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, logging.debug(f"topk_rel: {topk_rel}") logging.debug(f"abs idx: {abs_idxs}") - image_mask[b, abs_idxs] = False + if self.patch_reuse_policy == "never": + image_mask[b, abs_idxs] = False + elif self.patch_reuse_policy == "next_step_only": + current_selected_mask[b, abs_idxs] = True picked = inputs_embeds[b, abs_idxs, :] # (K, D) select_image_embeds.append(picked) @@ -223,6 +235,7 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, new_position_ids = [] new_original_mask = [] new_image_mask = [] + new_recently_selected_mask = [] batch_max_len = 0 for b in range(B): @@ -258,6 +271,14 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, merged_img = torch.cat([img_pref, img_v, img_suf], dim=0) new_image_mask.append(merged_img) + # recently_selected_mask (for next_step_only) + if self.patch_reuse_policy == "next_step_only": + recent_pref = current_selected_mask[b, :end_b] + recent_suf = current_selected_mask[b, end_b:] + recent_v = torch.zeros(self.num_selected_patches, device=input_ids.device, dtype=torch.bool) + merged_recent = torch.cat([recent_pref, recent_v, recent_suf], dim=0) + new_recently_selected_mask.append(merged_recent) + batch_max_len = max(batch_max_len, merged_b.size(0)) @@ -266,6 +287,7 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, padded_pos = [] padded_orig = [] padded_img = [] + padded_recent = [] for b in range(B): emb_b = new_inputs_embeds[b] @@ -279,12 +301,17 @@ def forward(self, input_ids, attention_mask, labels, position_ids, pixel_values, padded_pos.append(pos_b.unsqueeze(0)) padded_orig.append(orig_b.unsqueeze(0)) padded_img.append(img_b.unsqueeze(0)) + if self.patch_reuse_policy == "next_step_only": + recent_b = new_recently_selected_mask[b] + padded_recent.append(recent_b.unsqueeze(0)) inputs_embeds = torch.cat(padded_embeds, dim=0) attention_mask = torch.cat(padded_att, dim=0) position_ids = torch.cat(padded_pos, dim=0) original_mask = torch.cat(padded_orig, dim=0) image_mask = torch.cat(padded_img, dim=0) # (B, new_S) + if self.patch_reuse_policy == "next_step_only": + recently_selected_mask = torch.cat(padded_recent, dim=0) K = self.num_selected_patches for b in range(B): diff --git a/chameleon/chameleon_run.py b/chameleon/chameleon_run.py index 9048057..dff65b6 100644 --- a/chameleon/chameleon_run.py +++ b/chameleon/chameleon_run.py @@ -71,6 +71,8 @@ def main(): parser.add_argument("--deepspeed", action="store_true", help="Enable DeepSpeed") parser.add_argument("--deepspeed_config", default="ds_config.json", help="DeepSpeed config path") parser.add_argument("--local_rank", type=int, default=-1, help="Local rank passed by DeepSpeed") + parser.add_argument("--patch_reuse_policy", choices=["never", "next_step_only", "always"], default=None, + help="Patch selection reuse policy across latent reasoning steps") args = parser.parse_args() # Initialize DeepSpeed @@ -85,6 +87,7 @@ def main(): config_dict = yaml.safe_load(f) configs = Config(config_dict) + patch_reuse_policy = args.patch_reuse_policy or getattr(configs, "patch_reuse_policy", "never") set_seed(configs.seed) save_dir = os.path.join(configs.save_path, configs.name) @@ -144,7 +147,15 @@ def main(): model.print_trainable_parameters() - model = IVTLR(model, latent_id, start_id, end_id, tokenizer.eos_token_id, image_token_id) + model = IVTLR( + model, + latent_id, + start_id, + end_id, + tokenizer.eos_token_id, + image_token_id, + patch_reuse_policy=patch_reuse_policy, + ) print(f"Running Deepspeed on rank = {rank}, world size = {world_size}") model = model.to(rank) diff --git a/chameleon/chameleon_run_sqa.py b/chameleon/chameleon_run_sqa.py index 5acc6ba..7c581f8 100644 --- a/chameleon/chameleon_run_sqa.py +++ b/chameleon/chameleon_run_sqa.py @@ -72,6 +72,8 @@ def main(): parser.add_argument("--deepspeed", action="store_true", help="Enable DeepSpeed") parser.add_argument("--deepspeed_config", default="ds_config.json", help="DeepSpeed config path") parser.add_argument("--local_rank", type=int, default=-1, help="Local rank passed by DeepSpeed") + parser.add_argument("--patch_reuse_policy", choices=["never", "next_step_only", "always"], default=None, + help="Patch selection reuse policy across latent reasoning steps") args = parser.parse_args() # Initialize DeepSpeed @@ -86,6 +88,7 @@ def main(): config_dict = yaml.safe_load(f) configs = Config(config_dict) + patch_reuse_policy = args.patch_reuse_policy or getattr(configs, "patch_reuse_policy", "never") set_seed(configs.seed) save_dir = os.path.join(configs.save_path, configs.name) @@ -146,7 +149,15 @@ def main(): model.print_trainable_parameters() - model = IVTLR(model, latent_id, start_id, end_id, tokenizer.eos_token_id, image_token_id) + model = IVTLR( + model, + latent_id, + start_id, + end_id, + tokenizer.eos_token_id, + image_token_id, + patch_reuse_policy=patch_reuse_policy, + ) print(f"Running Deepspeed on rank = {rank}, world size = {world_size}") model = model.to(rank) diff --git a/chameleon/infer_chameleon.py b/chameleon/infer_chameleon.py index 787c61c..2a66114 100644 --- a/chameleon/infer_chameleon.py +++ b/chameleon/infer_chameleon.py @@ -18,8 +18,9 @@ ) device = "cuda" if torch.cuda.is_available() else "cpu" +PATCH_REUSE_POLICY = "never" -def load_inference_model(checkpoint_path): +def load_inference_model(checkpoint_path, patch_reuse_policy="never"): processor = ChameleonProcessor.from_pretrained("facebook/chameleon-7b") tokenizer = processor.tokenizer tokenizer.padding_side = "right" @@ -65,6 +66,7 @@ def load_inference_model(checkpoint_path): end_latent_id=end_id, eos_token_id=tokenizer.eos_token_id, image_token_id=image_token_id, + patch_reuse_policy=patch_reuse_policy, ) state_dict = torch.load(checkpoint_path, map_location="cpu") @@ -78,7 +80,7 @@ def load_inference_model(checkpoint_path): model.eval() return model, processor, tokenizer -model, processor, tokenizer = load_inference_model("your_pth_path") +model, processor, tokenizer = load_inference_model("your_pth_path", patch_reuse_policy=PATCH_REUSE_POLICY) os.makedirs("output", exist_ok=True) diff --git a/chameleon/infer_chameleon_scienceqa.py b/chameleon/infer_chameleon_scienceqa.py index 9c59f94..e655842 100644 --- a/chameleon/infer_chameleon_scienceqa.py +++ b/chameleon/infer_chameleon_scienceqa.py @@ -22,8 +22,13 @@ ) device = "cuda" if torch.cuda.is_available() else "cpu" +# #In that file, PATCH_REUSE_POLICY = "never". +# "never" means selected patches are masked out for all subsequent reasoning steps. +# If you want no masking, set it to "always". +# If you want only one-step blocking, set it to "next_step_only". +PATCH_REUSE_POLICY = "always" -def load_inference_model(checkpoint_path): +def load_inference_model(checkpoint_path, patch_reuse_policy="never"): print("Loading Chameleon model...") @@ -75,6 +80,7 @@ def load_inference_model(checkpoint_path): end_latent_id=end_id, eos_token_id=tokenizer.eos_token_id, image_token_id=image_token_id, + patch_reuse_policy=patch_reuse_policy, ) state_dict = torch.load(checkpoint_path, map_location="cpu") @@ -88,7 +94,7 @@ def load_inference_model(checkpoint_path): model.eval() return model, processor, tokenizer -model, processor, tokenizer = load_inference_model("your_pth_path") +model, processor, tokenizer = load_inference_model("your_pth_path", patch_reuse_policy=PATCH_REUSE_POLICY) os.makedirs("sqa_output", exist_ok=True) diff --git a/configs/adaptive_controller_qwen2b.yaml b/configs/adaptive_controller_qwen2b.yaml new file mode 100644 index 0000000..c69730b --- /dev/null +++ b/configs/adaptive_controller_qwen2b.yaml @@ -0,0 +1,50 @@ +project: ivtlr_adaptive_controller +model_id: Qwen/Qwen2-VL-2B-Instruct + +# Set this to your fully trained IVT-LR epoch-16 fp32 checkpoint. +ivtlr_checkpoint_path: /path/IVT-LR/qwen_vl_2b/qwen2b_IVTLR/epoch_16_full_model_fp32.pth +teacher_checkpoint_path: null + +output_dir: adaptive_controller_runs/qwen2vl_2b + +dataset_name: LightChen2333/M3CoT +dataset_split: train +max_train_examples: 3000 +num_proc: 16 +num_workers: 2 + +seed: 0 +bf16: true +use_lora: true +lora_r: 64 +lora_alpha: 16 +lora_dropout: 0.05 + +epochs_per_stage: 4 +max_latent_stage: 5 +scheduled_stage: 5 +pad_latent_to_max: true +batch_size_training: 1 +gradient_accumulation_steps: 4 +patch_reuse_policy: never + +teacher_k: 10 +budget_candidates: [2, 4, 6, 8, 10] +max_controller_steps: 10 + +lambda_patch: 0.001 +reward_temperature: 0.7 +advantage_mode: softmax + +freeze_base_model: true +train_controller_only: true +use_step_embedding: true +controller_hidden_dim: 1024 +controller_lr: 0.0002 +weight_decay: 0.01 +grad_clip_norm: 1.0 + +log_every: 5 +save_every: 250 +num_debug_traces: 20 +max_train_steps: 1500 \ No newline at end of file diff --git a/configs/adaptive_lora_stage2_qwen2b.yaml b/configs/adaptive_lora_stage2_qwen2b.yaml new file mode 100644 index 0000000..04e5134 --- /dev/null +++ b/configs/adaptive_lora_stage2_qwen2b.yaml @@ -0,0 +1,46 @@ +project: ivtlr_adaptive_lora_stage2 +model_id: Qwen/Qwen2-VL-2B-Instruct + +ivtlr_checkpoint_path: /content/checkpoints/epoch_16_full_model_fp32.pth +teacher_checkpoint_path: null +controller_checkpoint_path: /content/drive/MyDrive/adaptive_controller_runs/qwen2vl_2b_grpo_stage1_min2_no_patch_penalty/controller_grpo_final.pt +lora_stage2_checkpoint_path: null + +output_dir: /content/drive/MyDrive/adaptive_controller_runs/qwen2vl_2b_stage2_lora + +dataset_name: LightChen2333/M3CoT +dataset_split: train +max_train_examples: 3000 +num_proc: 16 +num_workers: 1 + +seed: 4 +bf16: true +use_lora: true +lora_r: 64 +lora_alpha: 16 +lora_dropout: 0.05 + +epochs_per_stage: 3 +max_latent_stage: 3 +scheduled_stage: 3 +pad_latent_to_max: true +batch_size_training: 1 +gradient_accumulation_steps: 4 + +patch_reuse_policy: always +teacher_k: 10 +max_controller_steps: 10 +controller_hidden_dim: 1024 +use_step_embedding: true + +freeze_base_model: false +train_controller_only: false + +stage2_lora_lr: 0.00001 +weight_decay: 0.0 +grad_clip_norm: 1.0 + +log_every: 5 +save_every: 250 +max_train_steps: 1000 diff --git a/qwen_vl/args/qwen.yaml b/qwen_vl/args/qwen.yaml index cf4c3a8..f3281c3 100644 --- a/qwen_vl/args/qwen.yaml +++ b/qwen_vl/args/qwen.yaml @@ -14,4 +14,5 @@ batch_size_training: 2 debug: False gradient_accumulation_steps: 8 num_epochs: 16 -lr: !!float "4e-5" \ No newline at end of file +lr: !!float "4e-5" +patch_reuse_policy: never \ No newline at end of file diff --git a/qwen_vl/args/qwen2vl_2b.yaml b/qwen_vl/args/qwen2vl_2b.yaml new file mode 100644 index 0000000..86129da --- /dev/null +++ b/qwen_vl/args/qwen2vl_2b.yaml @@ -0,0 +1,21 @@ +project: ivtlr +save_path: /path/IVT-LR/qwen_vl_2b/ +name: qwen2b_IVTLR + +epochs_per_stage: 4 +max_latent_stage: 5 +pad_latent_to_max: True + +load_model_path: None +seed: 0 +resume: 0 +bf16: True +batch_size_training: 4 +debug: False +gradient_accumulation_steps: 8 +num_epochs: 16 +lr: !!float "4e-5" +patch_reuse_policy: never +enable_nvt_loss: False +nvt_loss_weight: 0.1 +nvt_loss_epsilon: !!float "1e-8" \ No newline at end of file diff --git a/qwen_vl/controller.py b/qwen_vl/controller.py new file mode 100644 index 0000000..ad53622 --- /dev/null +++ b/qwen_vl/controller.py @@ -0,0 +1,334 @@ +from dataclasses import dataclass +from typing import Dict, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +@dataclass +class ControllerSequenceStats: + loss: torch.Tensor + mean_logprob: torch.Tensor + patch_top1_accuracy: torch.Tensor + stop_accuracy: torch.Tensor + token_count: int + + +class ControllerStateUpdater(nn.Module): + def __init__( + self, + controller_dim: int, + patch_dim: int, + max_steps: int = 10, + use_step_embedding: bool = True, + ): + super().__init__() + self.controller_dim = controller_dim + self.max_steps = max_steps + self.use_step_embedding = use_step_embedding + step_dim = controller_dim if use_step_embedding else 0 + self.step_embedding = ( + nn.Embedding(max_steps, controller_dim) if use_step_embedding else None + ) + self.mlp = nn.Sequential( + nn.Linear(controller_dim + patch_dim + step_dim, controller_dim * 4), + nn.GELU(), + nn.Linear(controller_dim * 4, controller_dim), + ) + self.norm = nn.LayerNorm(controller_dim) + + def forward( + self, + controller_state: torch.Tensor, + selected_patch: torch.Tensor, + step_idx: int, + ) -> torch.Tensor: + param_dtype = self.mlp[0].weight.dtype + controller_state = controller_state.to(dtype=param_dtype) + selected_patch = selected_patch.to(dtype=param_dtype) + pieces = [controller_state, selected_patch] + if self.use_step_embedding: + step = min(step_idx, self.max_steps - 1) + step_ids = torch.full( + (controller_state.size(0),), + step, + dtype=torch.long, + device=controller_state.device, + ) + pieces.append(self.step_embedding(step_ids)) + update = self.mlp(torch.cat(pieces, dim=-1)) + return self.norm(controller_state + update) + + +class PatchPointerController(nn.Module): + """Pointer controller over per-example image patch embeddings plus STOP. + + STOP is represented only as the last controller logit. It is never a text + tokenizer id and should never be appended to the IVT-LR embedding stream. + """ + + def __init__( + self, + model_dim: int, + controller_dim: Optional[int] = None, + max_steps: int = 10, + use_step_embedding: bool = True, + ): + super().__init__() + self.model_dim = model_dim + self.controller_dim = controller_dim or model_dim + self.max_steps = max_steps + self.state_proj = nn.Linear(model_dim, self.controller_dim) + self.query_proj = nn.Linear(self.controller_dim, self.controller_dim) + self.key_proj = nn.Linear(model_dim, self.controller_dim) + self.stop_proj = nn.Linear(self.controller_dim, 1) + self.updater = ControllerStateUpdater( + controller_dim=self.controller_dim, + patch_dim=model_dim, + max_steps=max_steps, + use_step_embedding=use_step_embedding, + ) + self.scale = self.controller_dim ** -0.5 + + def initial_state(self, reasoning_state: torch.Tensor) -> torch.Tensor: + return self.state_proj(reasoning_state.to(dtype=self.state_proj.weight.dtype)) + + def forward( + self, + controller_state: torch.Tensor, + patch_embeddings: torch.Tensor, + patch_valid_mask: Optional[torch.Tensor] = None, + selected_mask: Optional[torch.Tensor] = None, + allow_stop: bool = True, + ) -> torch.Tensor: + if patch_embeddings.dim() != 3: + raise ValueError("patch_embeddings must have shape [B, N, D]") + param_dtype = self.query_proj.weight.dtype + controller_state = controller_state.to(dtype=param_dtype) + patch_embeddings = patch_embeddings.to(dtype=param_dtype) + bsz, n_patches, _ = patch_embeddings.shape + q = self.query_proj(controller_state).unsqueeze(-1) + k = self.key_proj(patch_embeddings) + patch_logits = torch.bmm(k, q).squeeze(-1) * self.scale + + if patch_valid_mask is not None: + patch_logits = patch_logits.masked_fill(~patch_valid_mask.bool(), float("-inf")) + if selected_mask is not None: + patch_logits = patch_logits.masked_fill(selected_mask.bool(), float("-inf")) + + stop_logit = self.stop_proj(controller_state) + if not allow_stop: + stop_logit = stop_logit.fill_(float("-inf")) + return torch.cat([patch_logits, stop_logit], dim=-1) + + def update_state( + self, + controller_state: torch.Tensor, + selected_patch: torch.Tensor, + step_idx: int, + ) -> torch.Tensor: + return self.updater(controller_state, selected_patch, step_idx) + + def gather_patch( + self, + patch_embeddings: torch.Tensor, + patch_indices: torch.Tensor, + ) -> torch.Tensor: + gather_idx = patch_indices.view(-1, 1, 1).expand(-1, 1, patch_embeddings.size(-1)) + return patch_embeddings.gather(1, gather_idx).squeeze(1) + + def teacher_forced_sequence_loss( + self, + reasoning_state: torch.Tensor, + patch_embeddings: torch.Tensor, + patch_valid_mask: torch.Tensor, + target_actions: torch.Tensor, + sequence_weights: Optional[torch.Tensor] = None, + ) -> ControllerSequenceStats: + """Sequentially train p1, p2, ..., STOP with teacher forcing. + + target_actions has shape [B, L]. STOP must be encoded as N, where N is + the padded patch dimension for patch_embeddings. + """ + bsz, n_patches, _ = patch_embeddings.shape + stop_index = n_patches + state = self.initial_state(reasoning_state) + selected_mask = torch.zeros( + (bsz, n_patches), dtype=torch.bool, device=patch_embeddings.device + ) + logprob_sum = torch.zeros(bsz, device=patch_embeddings.device) + token_counts = torch.zeros(bsz, device=patch_embeddings.device) + patch_correct = torch.zeros(bsz, device=patch_embeddings.device) + patch_total = torch.zeros(bsz, device=patch_embeddings.device) + stop_correct = torch.zeros(bsz, device=patch_embeddings.device) + stop_total = torch.zeros(bsz, device=patch_embeddings.device) + + for step_idx in range(target_actions.size(1)): + target = target_actions[:, step_idx] + active = target >= 0 + if not active.any(): + break + + logits = self.forward( + state, + patch_embeddings, + patch_valid_mask=patch_valid_mask, + selected_mask=selected_mask, + ) + log_probs = F.log_softmax(logits, dim=-1) + safe_target = target.clamp(min=0) + step_logprob = log_probs.gather(-1, safe_target.unsqueeze(-1)).squeeze(-1) + logprob_sum = logprob_sum + step_logprob.masked_fill(~active, 0.0) + token_counts = token_counts + active.float() + + pred = logits.argmax(dim=-1) + is_stop = target == stop_index + is_patch = active & ~is_stop + patch_correct = patch_correct + ((pred == target) & is_patch).float() + patch_total = patch_total + is_patch.float() + stop_correct = stop_correct + ((pred == stop_index) & is_stop).float() + stop_total = stop_total + is_stop.float() + + update_mask = is_patch + if update_mask.any(): + patch_target = target.clamp(max=n_patches - 1) + selected_patch = self.gather_patch(patch_embeddings, patch_target) + updated_state = self.update_state(state, selected_patch, step_idx) + state = torch.where(update_mask.unsqueeze(-1), updated_state, state) + selected_update = torch.zeros_like(selected_mask) + selected_update = selected_update.scatter(1, patch_target.unsqueeze(1), True) + selected_mask = selected_mask | (selected_update & update_mask.unsqueeze(1)) + + if sequence_weights is None: + sequence_weights = torch.ones_like(logprob_sum) + sequence_weights = sequence_weights.detach() + loss = -(sequence_weights * logprob_sum).mean() + denom = token_counts.clamp(min=1.0) + mean_logprob = (logprob_sum / denom).mean() + patch_acc = (patch_correct.sum() / patch_total.sum().clamp(min=1.0)).detach() + stop_acc = (stop_correct.sum() / stop_total.sum().clamp(min=1.0)).detach() + return ControllerSequenceStats( + loss=loss, + mean_logprob=mean_logprob.detach(), + patch_top1_accuracy=patch_acc, + stop_accuracy=stop_acc, + token_count=int(token_counts.sum().item()), + ) + + @torch.no_grad() + def greedy_select( + self, + reasoning_state: torch.Tensor, + patch_embeddings: torch.Tensor, + patch_valid_mask: torch.Tensor, + max_steps: Optional[int] = None, + ) -> Dict[str, torch.Tensor]: + max_steps = max_steps or self.max_steps + bsz, n_patches, _ = patch_embeddings.shape + stop_index = n_patches + state = self.initial_state(reasoning_state) + selected_mask = torch.zeros( + (bsz, n_patches), dtype=torch.bool, device=patch_embeddings.device + ) + selected = torch.full( + (bsz, max_steps), -1, dtype=torch.long, device=patch_embeddings.device + ) + lengths = torch.zeros(bsz, dtype=torch.long, device=patch_embeddings.device) + stopped = torch.zeros(bsz, dtype=torch.bool, device=patch_embeddings.device) + + for step_idx in range(max_steps): + logits = self.forward( + state, + patch_embeddings, + patch_valid_mask=patch_valid_mask, + selected_mask=selected_mask, + ) + action = logits.argmax(dim=-1) + is_stop = action == stop_index + take_patch = (~stopped) & (~is_stop) + if take_patch.any(): + patch_action = action.clamp(max=n_patches - 1) + selected[:, step_idx] = torch.where(take_patch, patch_action, selected[:, step_idx]) + selected_patch = self.gather_patch(patch_embeddings, patch_action) + updated_state = self.update_state(state, selected_patch, step_idx) + state = torch.where(take_patch.unsqueeze(-1), updated_state, state) + selected_update = torch.zeros_like(selected_mask) + selected_update = selected_update.scatter(1, patch_action.unsqueeze(1), True) + selected_mask = selected_mask | (selected_update & take_patch.unsqueeze(1)) + lengths = lengths + take_patch.long() + stopped = stopped | is_stop + if stopped.all(): + break + + return {"selected_indices": selected, "lengths": lengths, "stopped": stopped} + + def sample_select( + self, + reasoning_state: torch.Tensor, + patch_embeddings: torch.Tensor, + patch_valid_mask: torch.Tensor, + max_steps: Optional[int] = None, + temperature: float = 1.0, + min_patches: int = 0, + ) -> Dict[str, torch.Tensor]: + max_steps = max_steps or self.max_steps + temperature = max(float(temperature), 1e-6) + bsz, n_patches, _ = patch_embeddings.shape + stop_index = n_patches + state = self.initial_state(reasoning_state) + selected_mask = torch.zeros( + (bsz, n_patches), dtype=torch.bool, device=patch_embeddings.device + ) + selected = torch.full( + (bsz, max_steps), -1, dtype=torch.long, device=patch_embeddings.device + ) + lengths = torch.zeros(bsz, dtype=torch.long, device=patch_embeddings.device) + stopped = torch.zeros(bsz, dtype=torch.bool, device=patch_embeddings.device) + logprob_sum = torch.zeros(bsz, dtype=state.dtype, device=patch_embeddings.device) + entropy_sum = torch.zeros(bsz, dtype=state.dtype, device=patch_embeddings.device) + action_count = torch.zeros(bsz, dtype=state.dtype, device=patch_embeddings.device) + + for step_idx in range(max_steps): + logits = self.forward( + state, + patch_embeddings, + patch_valid_mask=patch_valid_mask, + selected_mask=selected_mask, + allow_stop=step_idx >= min_patches, + ) + dist = torch.distributions.Categorical(logits=logits / temperature) + action = dist.sample() + active = ~stopped + is_stop = action == stop_index + take_patch = active & (~is_stop) + + logprob = dist.log_prob(action) + entropy = dist.entropy() + logprob_sum = logprob_sum + logprob.masked_fill(~active, 0.0) + entropy_sum = entropy_sum + entropy.masked_fill(~active, 0.0) + action_count = action_count + active.float() + + if take_patch.any(): + patch_action = action.clamp(max=n_patches - 1) + selected[:, step_idx] = torch.where(take_patch, patch_action, selected[:, step_idx]) + selected_patch = self.gather_patch(patch_embeddings, patch_action) + updated_state = self.update_state(state, selected_patch, step_idx) + state = torch.where(take_patch.unsqueeze(-1), updated_state, state) + selected_update = torch.zeros_like(selected_mask) + selected_update = selected_update.scatter(1, patch_action.unsqueeze(1), True) + selected_mask = selected_mask | (selected_update & take_patch.unsqueeze(1)) + lengths = lengths + take_patch.long() + + stopped = stopped | (active & is_stop) + if stopped.all(): + break + + return { + "selected_indices": selected, + "lengths": lengths, + "stopped": stopped, + "logprob_sum": logprob_sum, + "entropy_sum": entropy_sum, + "action_count": action_count, + } diff --git a/qwen_vl/dataset.py b/qwen_vl/dataset.py index 68d0d33..f01c2f2 100644 --- a/qwen_vl/dataset.py +++ b/qwen_vl/dataset.py @@ -24,7 +24,7 @@ ) -def get_dataset(dataset, tokenizer, processor, max_size=1000000000): +def get_dataset(dataset, tokenizer, processor, max_size=1000000000, num_proc=32): def tokenize_sample(sample, max_length=3400): image = sample["image"] @@ -87,7 +87,7 @@ def tokenize_sample(sample, max_length=3400): if dist.get_rank() == 0: processed_dataset = [ dataset.map( - tokenize_sample, remove_columns=list(dataset.features), num_proc=32 + tokenize_sample, remove_columns=list(dataset.features), num_proc=num_proc ) ] else: @@ -97,7 +97,7 @@ def tokenize_sample(sample, max_length=3400): else: dataset = dataset.map( - tokenize_sample, remove_columns=list(dataset.features), num_proc=32 + tokenize_sample, remove_columns=list(dataset.features), num_proc=num_proc ) return dataset @@ -139,6 +139,8 @@ def __call__(self, features, return_tensors=None): feature["labels"] = [self.label_pad_token_id] * n_tok_pad + feature[ "labels" ] + if "answer_mask" in feature: + feature["answer_mask"] = [0] * n_tok_pad + feature["answer_mask"] feature["attention_mask"] = [0] * n_tok_pad + feature["attention_mask"] return_tensors = "pt" @@ -149,7 +151,7 @@ def __call__(self, features, return_tensors=None): { k: v for k, v in feature.items() - if k != label_name and k != "position_ids" + if k != label_name and k != "position_ids" and k != "answer_mask" } for feature in features ] @@ -196,6 +198,14 @@ def __call__(self, features, return_tensors=None): batch["position_ids"], dtype=torch.int64 ) + if "answer_mask" in features[0]: + answer_masks = [feature["answer_mask"] for feature in features] + max_answer_len = max(len(m) for m in answer_masks) + batch["answer_mask"] = [ + mask + [0] * (max_answer_len - len(mask)) for mask in answer_masks + ] + batch["answer_mask"] = torch.tensor(batch["answer_mask"], dtype=torch.int64) + return batch def get_cot_latent_dataset( @@ -229,14 +239,19 @@ def process_dataset(sample): scheduled_stage_to_train, ) + steps_tokenized = sample["steps_tokenized"][n_skip_steps:] + steps_len = sum(len(step) for step in steps_tokenized) + tokens = ( sample["question_tokenized"] + [latent_id] * n_latent_tokens - + list( - itertools.chain.from_iterable(sample["steps_tokenized"][n_skip_steps:]) - ) + + list(itertools.chain.from_iterable(steps_tokenized)) + sample["answer_tokenized"] ) + + answer_start = len(sample["question_tokenized"]) + n_latent_tokens + steps_len + answer_len = len(sample["answer_tokenized"]) + answer_mask = [0] * answer_start + [1] * answer_len return { "input_ids": tokens, @@ -250,6 +265,7 @@ def process_dataset(sample): + len(sample["question_tokenized"]) : ], "attention_mask": [1] * len(tokens), + "answer_mask": answer_mask, "idx": sample["idx"], "position_ids": list(range(len(tokens))), "pixel_values": torch.tensor(sample["pixel_values"]), diff --git a/qwen_vl/ds_config.json b/qwen_vl/ds_config.json index c4dbbfd..31b6790 100644 --- a/qwen_vl/ds_config.json +++ b/qwen_vl/ds_config.json @@ -1,6 +1,6 @@ { - "train_batch_size": 64, - "train_micro_batch_size_per_gpu": 2, + "train_batch_size": 32, + "train_micro_batch_size_per_gpu": 4, "gradient_accumulation_steps": 8, "bf16": { "enabled": true diff --git a/qwen_vl/infer.py b/qwen_vl/infer.py index 4065a3b..395337d 100644 --- a/qwen_vl/infer.py +++ b/qwen_vl/infer.py @@ -12,6 +12,8 @@ import os import time from datetime import timedelta +import argparse +from tqdm import tqdm logging.basicConfig( filename='qwenvl_32_infer_time.log', level=logging.DEBUG, @@ -21,8 +23,9 @@ import pdb device = "cuda" if torch.cuda.is_available() else "cpu" +DEFAULT_PATCH_REUSE_POLICY = "always" -def load_inference_model(checkpoint_path): +def load_inference_model(checkpoint_path, patch_reuse_policy="never", patch_sampling_strategy="attention_topk"): processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") tokenizer = AutoTokenizer.from_pretrained( "Qwen/Qwen2-VL-7B-Instruct", @@ -75,7 +78,9 @@ def load_inference_model(checkpoint_path): eos_token_id=tokenizer.eos_token_id, image_token_id=image_token_id, visual_start_id=visual_start_id, - visual_end_id=visual_end_id + visual_end_id=visual_end_id, + patch_reuse_policy=patch_reuse_policy, + patch_sampling_strategy=patch_sampling_strategy, ) state_dict = torch.load(checkpoint_path, map_location="cpu") @@ -91,10 +96,6 @@ def load_inference_model(checkpoint_path): model.eval() return model, processor, tokenizer -model, processor, tokenizer = load_inference_model("your_path") - -os.makedirs("output", exist_ok=True) - def format_prompt(example): question = example["question"].strip() rationale = example["rationale"].replace("\n", " ").strip() @@ -123,20 +124,30 @@ def process_func(example): "topic": example["topic"] } -dataset = load_dataset("LightChen2333/M3CoT") -val_dataset = dataset["test"] -val_dataset = val_dataset.filter(lambda e: e["image"] is not None).map(process_func) +def build_eval_dataset(): + dataset = load_dataset("LightChen2333/M3CoT") + val_dataset = dataset["test"] + return val_dataset.filter(lambda e: e["image"] is not None).map(process_func) -def evaluate_and_save(eval_dataset, model, processor): + +def evaluate_and_save(eval_dataset, model, processor, output_path, latent_n=3, max_new_tokens=512): model.eval() correct = 0 total = 0 total_generated_tokens = 0 total_generate_time = 0.0 - - output_path = "output/qwen2vl_32.jsonl" + + output_dir = os.path.dirname(output_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + with open(output_path, "a", encoding="utf-8") as f_out: - for ex in eval_dataset: + for ex in tqdm( + eval_dataset, + total=len(eval_dataset), + desc="Evaluating M3CoT", + dynamic_ncols=True, + ): input_text = ex["question_raw"] messages = [{ "role": "user", @@ -146,7 +157,7 @@ def evaluate_and_save(eval_dataset, model, processor): ] }] text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - text = text + "<|latent|>" + "<|latent|>" + "<|latent|>" + text = text + ("<|latent|>" * latent_n) image_inputs, video_inputs = process_vision_info(messages) inputs = processor( text=[text], @@ -165,7 +176,7 @@ def evaluate_and_save(eval_dataset, model, processor): attention_mask=torch.tensor(inputs["attention_mask"]), pixel_values=torch.tensor(inputs["pixel_values"]), image_grid_thw=torch.tensor(inputs["image_grid_thw"]), - max_new_tokens=512 + max_new_tokens=max_new_tokens ) generate_end_time = time.time() sample_generate_time = generate_end_time - generate_start_time @@ -224,5 +235,40 @@ def evaluate_and_save(eval_dataset, model, processor): logging.info(f"[FINAL] Avg generated tokens per sample: {avg_generated_tokens:.1f}") logging.info(f"[FINAL] Total generate time: {total_generate_time:.2f}s ({timedelta(seconds=int(total_generate_time))})") logging.info(f"[FINAL] Avg generate time per sample: {avg_time_per_sample:.3f}s") - -evaluate_and_save(val_dataset, model, processor) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Qwen2-VL IVTLR inference on M3CoT") + parser.add_argument("--checkpoint_path", type=str, required=True, help="Path to model state_dict checkpoint (.pth)") + parser.add_argument("--latent_n", type=int, default=3, help="Number of <|latent|> tokens appended to the prompt") + parser.add_argument("--patch_reuse_policy", type=str, default=DEFAULT_PATCH_REUSE_POLICY, + choices=["never", "next_step_only", "always"], + help="Patch selection reuse policy during generation") + parser.add_argument("--patch_sampling_strategy", type=str, default="attention_topk", + choices=["attention_topk", "random_image_only", "all_image_patches"], + help="Patch sampling strategy for selecting visual tokens") + parser.add_argument("--output_path", type=str, default="output/qwen2vl_32.jsonl", help="Path to write JSONL predictions") + parser.add_argument("--max_new_tokens", type=int, default=512, help="Maximum generated tokens per sample") + return parser.parse_args() + + +def main(): + args = parse_args() + model, processor, _ = load_inference_model( + args.checkpoint_path, + patch_reuse_policy=args.patch_reuse_policy, + patch_sampling_strategy=args.patch_sampling_strategy, + ) + val_dataset = build_eval_dataset() + evaluate_and_save( + val_dataset, + model, + processor, + output_path=args.output_path, + latent_n=args.latent_n, + max_new_tokens=args.max_new_tokens, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen_vl/infer_2b.py b/qwen_vl/infer_2b.py new file mode 100644 index 0000000..76528db --- /dev/null +++ b/qwen_vl/infer_2b.py @@ -0,0 +1,402 @@ +from transformers import AutoTokenizer, AutoProcessor +from qwen_ivtlr import IVTLR +from transformers import Qwen2VLForConditionalGeneration +import torch +import deepspeed +from peft import LoraConfig,get_peft_model +from qwen_vl_utils import process_vision_info +from datasets import load_dataset +from utils import set_seed +import re +import logging +import json +import os +import time +from datetime import timedelta +import argparse +from tqdm import tqdm +logging.basicConfig( + filename='qwenvl_2b_infer_time.log', + level=logging.DEBUG, + format='[%(asctime)s] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' +) +import pdb + +device = "cuda" if torch.cuda.is_available() else "cpu" +DEFAULT_PATCH_REUSE_POLICY = "always" + +def load_inference_model(checkpoint_path, patch_reuse_policy="never", patch_sampling_strategy="attention_topk"): + processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") + tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen2-VL-2B-Instruct", + use_fast=False, + trust_remote_code=True, + padding_side="right" + ) + + tokenizer.add_special_tokens({ + "additional_special_tokens": [ + "<|start-latent|>", + "<|end-latent|>", + "<|latent|>" + ] + }) + + base_model = Qwen2VLForConditionalGeneration.from_pretrained( + "Qwen/Qwen2-VL-2B-Instruct", + device_map="cuda", + torch_dtype=torch.bfloat16, + trust_remote_code=True, + attn_implementation="eager" + ) + base_model.resize_token_embeddings(len(tokenizer)) + processor.tokenizer = tokenizer + + lora_config = LoraConfig( + task_type="CAUSAL_LM", + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + r=64, + lora_alpha=16, + lora_dropout=0.05, + bias="none", + inference_mode=False + ) + base_model = get_peft_model(base_model, lora_config) + + latent_id = tokenizer.convert_tokens_to_ids("<|latent|>") + start_id = tokenizer.convert_tokens_to_ids("<|start-latent|>") + end_id = tokenizer.convert_tokens_to_ids("<|end-latent|>") + image_token_id = tokenizer.convert_tokens_to_ids(processor.image_token) + visual_start_id = tokenizer.convert_tokens_to_ids("<|vision_start|>") + visual_end_id = tokenizer.convert_tokens_to_ids("<|vision_end|>") + + model = IVTLR( + base_model, + latent_token_id=latent_id, + start_latent_id=start_id, + end_latent_id=end_id, + eos_token_id=tokenizer.eos_token_id, + image_token_id=image_token_id, + visual_start_id=visual_start_id, + visual_end_id=visual_end_id, + patch_reuse_policy=patch_reuse_policy, + patch_sampling_strategy=patch_sampling_strategy, + processor_model_id="Qwen/Qwen2-VL-2B-Instruct", + ) + + state_dict = torch.load(checkpoint_path, map_location="cpu") + print(state_dict.keys()) + if any(k.startswith("module.") for k in state_dict.keys()): + state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} + + model.load_state_dict(state_dict, strict=True) + print(model) + print("Successfully load") + + model = model.to(device) + model.eval() + return model, processor, tokenizer + +def format_prompt(example): + question = example["question"].strip() + rationale = example["rationale"].replace("\n", " ").strip() + answer = example["answer"].strip() + choices = example["choices"] + image = example["image"] + + choices_str = "\n".join([f"{chr(65+i)}.{{{choice.strip()}}}" for i, choice in enumerate(choices)]) + user_prompt = ( + f"[Question]:{{{question}}}\n" + f"[Options]:\n{choices_str}\n" + f"Answer:" + ) + return user_prompt, rationale, answer, image + +def process_func(example): + prompt, rationale, answer, image = format_prompt(example) + + return { + "question_raw": prompt, + "image_raw": image, + "gt_answer": answer, + "id": example["id"], + "choices": example["choices"], + "domain": example["domain"], + "topic": example["topic"] + } + +def build_eval_dataset(data_percent=100.0, sample_seed=42): + dataset = load_dataset("LightChen2333/M3CoT") + val_dataset = dataset["test"] + val_dataset = val_dataset.filter(lambda e: e["image"] is not None).map(process_func) + if data_percent >= 100: + return val_dataset + if data_percent <= 0: + raise ValueError("data_percent must be in (0, 100].") + sample_size = max(1, int(len(val_dataset) * (data_percent / 100.0))) + val_dataset = val_dataset.shuffle(seed=sample_seed).select(range(sample_size)) + return val_dataset + + +def compute_latent_attention_trace(model, inputs, attn_threshold=None, attn_threshold_multiplier=5.0): + seq_len = inputs["input_ids"].shape[1] + position_ids = torch.arange(seq_len, device=inputs["input_ids"].device).unsqueeze(0) + labels = inputs["input_ids"].clone() + + with torch.no_grad(): + outputs = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + labels=labels, + position_ids=position_ids, + pixel_values=inputs["pixel_values"], + image_grid_thw=inputs["image_grid_thw"], + return_latent_attn=True, + latent_attn_threshold=attn_threshold, + latent_attn_threshold_multiplier=attn_threshold_multiplier, + ) + + return outputs.latent_attn_trace + + +def compute_token_norms(model, inputs): + seq_len = inputs["input_ids"].shape[1] + position_ids = torch.arange(seq_len, device=inputs["input_ids"].device).unsqueeze(0) + labels = inputs["input_ids"].clone() + + with torch.no_grad(): + outputs = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + labels=labels, + position_ids=position_ids, + pixel_values=inputs["pixel_values"], + image_grid_thw=inputs["image_grid_thw"], + return_token_norms=True, + ) + + if not outputs.token_norms: + return None + return outputs.token_norms[0] + + +def evaluate_and_save( + eval_dataset, + model, + processor, + output_path, + latent_n=3, + max_new_tokens=512, + attn_trace_path=None, + attn_threshold=None, + attn_threshold_multiplier=5.0, + analyze_token_norms=False, + token_norms_path=None, +): + model.eval() + correct = 0 + total = 0 + total_generated_tokens = 0 + total_generate_time = 0.0 + attn_traces = [] + + output_dir = os.path.dirname(output_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + norms_file = None + if analyze_token_norms and token_norms_path: + token_norms_dir = os.path.dirname(token_norms_path) + if token_norms_dir: + os.makedirs(token_norms_dir, exist_ok=True) + norms_file = open(token_norms_path, "a", encoding="utf-8") + + try: + with open(output_path, "a", encoding="utf-8") as f_out: + for ex in tqdm( + eval_dataset, + total=len(eval_dataset), + desc="Evaluating M3CoT", + dynamic_ncols=True, + ): + input_text = ex["question_raw"] + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": ex["image_raw"], "resized_height": 280, "resized_width": 280}, + {"type": "text", "text": input_text} + ] + }] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text = text + ("<|latent|>" * latent_n) + image_inputs, video_inputs = process_vision_info(messages) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt" + ).to(device) + if attn_trace_path: + trace = compute_latent_attention_trace( + model, + inputs, + attn_threshold=attn_threshold, + attn_threshold_multiplier=attn_threshold_multiplier, + ) + attn_traces.append({ + "id": ex["id"], + "latent_attn": trace[0] if trace else [], + }) + if analyze_token_norms and norms_file is not None: + norms = compute_token_norms(model, inputs) + if norms is not None: + norms_result = { + "id": ex["id"], + "patch_norms": norms.get("patch_norms", []), + "reasoning_norms": norms.get("reasoning_norms", []), + "aggregate_stats": norms.get("aggregate_stats", {}), + } + norms_file.write(json.dumps(norms_result, ensure_ascii=False) + "\n") + norms_file.flush() + + input_ids = inputs["input_ids"] + prompt_length = input_ids.shape[1] + + generate_start_time = time.time() + with torch.no_grad(): + outputs = model.generate( + input_ids=torch.tensor(inputs["input_ids"]), + attention_mask=torch.tensor(inputs["attention_mask"]), + pixel_values=torch.tensor(inputs["pixel_values"]), + image_grid_thw=torch.tensor(inputs["image_grid_thw"]), + max_new_tokens=max_new_tokens + ) + generate_end_time = time.time() + sample_generate_time = generate_end_time - generate_start_time + total_generate_time += sample_generate_time + + generated_tokens = outputs[0, prompt_length:] + new_generated_text = processor.decode(generated_tokens, skip_special_tokens=True) + output_text = processor.decode(outputs[0], skip_special_tokens=True) + logging.debug(f"[OUTPUT] {output_text}") + + num_generated_tokens = len(generated_tokens) + total_generated_tokens += num_generated_tokens + + cleaned_text = re.sub( + r'(?<=answer:)\s*(\n+\s*)?assistant\b', + '', + output_text, + flags=re.IGNORECASE + ) + matches = re.finditer( + r'(?:the\s+answer\s+is|Answer:)\s*[\n\s]*([A-Z])', + cleaned_text, + flags=re.IGNORECASE | re.DOTALL + ) + candidates = {match.group(1).upper() for match in matches} + gt_answer = ex["gt_answer"].strip().upper() + + if gt_answer in candidates: + correct += 1 + logging.debug(f"correct: True") + total += 1 + logging.debug(f"[TOTAL] {total}") + + # pdb.set_trace() + message_question = ex["question_raw"] + message_question = message_question.replace("", "", 1).replace("Answer:", "", 1).strip() + message_question = message_question.split("Answer:")[0].strip() + + result = { + "id": ex["id"], + "choices": ex["choices"], + "answer": ex["gt_answer"], + "domain": ex["domain"], + "topic": ex["topic"], + "messages": [ + message_question, + new_generated_text + ] + } + f_out.write(json.dumps(result, ensure_ascii=False) + "\n") + f_out.flush() + + avg_generated_tokens = total_generated_tokens / total if total > 0 else 0 + avg_time_per_sample = total_generate_time / total if total > 0 else 0 + + logging.info(f"[FINAL] Avg generated tokens per sample: {avg_generated_tokens:.1f}") + logging.info(f"[FINAL] Total generate time: {total_generate_time:.2f}s ({timedelta(seconds=int(total_generate_time))})") + logging.info(f"[FINAL] Avg generate time per sample: {avg_time_per_sample:.3f}s") + finally: + if norms_file is not None: + norms_file.close() + + if attn_trace_path: + attn_dir = os.path.dirname(attn_trace_path) + if attn_dir: + os.makedirs(attn_dir, exist_ok=True) + with open(attn_trace_path, "w", encoding="utf-8") as f: + json.dump(attn_traces, f, ensure_ascii=False, indent=2) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Qwen2-VL IVTLR inference on M3CoT") + parser.add_argument("--checkpoint_path", type=str, required=True, help="Path to model state_dict checkpoint (.pth)") + parser.add_argument("--latent_n", type=int, default=3, help="Number of <|latent|> tokens appended to the prompt") + parser.add_argument("--patch_reuse_policy", type=str, default=DEFAULT_PATCH_REUSE_POLICY, + choices=["never", "next_step_only", "always"], + help="Patch selection reuse policy during generation") + parser.add_argument("--patch_sampling_strategy", type=str, default="attention_topk", + choices=["attention_topk", "random_image_only", "all_image_patches"], + help="Patch sampling strategy for selecting visual tokens") + parser.add_argument("--output_path", type=str, default="output/qwen2vl_2b.jsonl", help="Path to write JSONL predictions") + parser.add_argument("--max_new_tokens", type=int, default=512, help="Maximum generated tokens per sample") + parser.add_argument("--attn_trace_path", type=str, default=None, help="Path to write JSON attention traces") + parser.add_argument("--attn_threshold", type=float, default=None, help="Absolute attention threshold for patch logging") + parser.add_argument("--attn_threshold_multiplier", type=float, default=5.0, help="Threshold multiplier for 1/seq_len baseline") + parser.add_argument("--data_percent", type=float, default=100.0, help="Percentage of dataset to use for inference") + parser.add_argument("--sample_seed", type=int, default=42, help="Random seed for dataset sampling") + parser.add_argument( + "--analyze_patch_reasoning_norms", + action="store_true", + help="Compute patch vs reasoning token norms and store per-example JSON", + ) + parser.add_argument( + "--token_norms_path", + type=str, + default="output/qwen2vl_2b_token_norms.jsonl", + help="Path to write token norm JSONL", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + model, processor, _ = load_inference_model( + args.checkpoint_path, + patch_reuse_policy=args.patch_reuse_policy, + patch_sampling_strategy=args.patch_sampling_strategy, + ) + if not (0 < args.data_percent <= 100): + raise ValueError("--data_percent must be in (0, 100].") + set_seed(args.sample_seed) + val_dataset = build_eval_dataset(data_percent=args.data_percent, sample_seed=args.sample_seed) + evaluate_and_save( + val_dataset, + model, + processor, + output_path=args.output_path, + latent_n=args.latent_n, + max_new_tokens=args.max_new_tokens, + attn_trace_path=args.attn_trace_path, + attn_threshold=args.attn_threshold, + attn_threshold_multiplier=args.attn_threshold_multiplier, + analyze_token_norms=args.analyze_patch_reasoning_norms, + token_norms_path=args.token_norms_path, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen_vl/infer_2b_sqa.py b/qwen_vl/infer_2b_sqa.py new file mode 100644 index 0000000..29a75bb --- /dev/null +++ b/qwen_vl/infer_2b_sqa.py @@ -0,0 +1,361 @@ +from transformers import AutoTokenizer, AutoProcessor +from qwen_ivtlr import IVTLR +from transformers import Qwen2VLForConditionalGeneration +import torch +import deepspeed +from peft import LoraConfig,get_peft_model +from qwen_vl_utils import process_vision_info +from datasets import load_dataset +import re +import logging +import json +import os +import time +from datetime import timedelta +import argparse +from tqdm import tqdm +logging.basicConfig( + filename='qwenvl_2b_infer_time.log', + level=logging.DEBUG, + format='[%(asctime)s] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' +) +import pdb + +device = "cuda" if torch.cuda.is_available() else "cpu" +DEFAULT_PATCH_REUSE_POLICY = "always" + +def load_inference_model(checkpoint_path, patch_reuse_policy="never", patch_sampling_strategy="attention_topk"): + processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") + tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen2-VL-2B-Instruct", + use_fast=False, + trust_remote_code=True, + padding_side="right" + ) + + tokenizer.add_special_tokens({ + "additional_special_tokens": [ + "<|start-latent|>", + "<|end-latent|>", + "<|latent|>" + ] + }) + + base_model = Qwen2VLForConditionalGeneration.from_pretrained( + "Qwen/Qwen2-VL-2B-Instruct", + device_map="cuda", + torch_dtype=torch.bfloat16, + trust_remote_code=True, + attn_implementation="eager" + ) + base_model.resize_token_embeddings(len(tokenizer)) + processor.tokenizer = tokenizer + + lora_config = LoraConfig( + task_type="CAUSAL_LM", + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + r=64, + lora_alpha=16, + lora_dropout=0.05, + bias="none", + inference_mode=False + ) + base_model = get_peft_model(base_model, lora_config) + + latent_id = tokenizer.convert_tokens_to_ids("<|latent|>") + start_id = tokenizer.convert_tokens_to_ids("<|start-latent|>") + end_id = tokenizer.convert_tokens_to_ids("<|end-latent|>") + image_token_id = tokenizer.convert_tokens_to_ids(processor.image_token) + visual_start_id = tokenizer.convert_tokens_to_ids("<|vision_start|>") + visual_end_id = tokenizer.convert_tokens_to_ids("<|vision_end|>") + + model = IVTLR( + base_model, + latent_token_id=latent_id, + start_latent_id=start_id, + end_latent_id=end_id, + eos_token_id=tokenizer.eos_token_id, + image_token_id=image_token_id, + visual_start_id=visual_start_id, + visual_end_id=visual_end_id, + patch_reuse_policy=patch_reuse_policy, + patch_sampling_strategy=patch_sampling_strategy, + processor_model_id="Qwen/Qwen2-VL-2B-Instruct", + ) + + state_dict = torch.load(checkpoint_path, map_location="cpu") + print(state_dict.keys()) + if any(k.startswith("module.") for k in state_dict.keys()): + state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} + + model.load_state_dict(state_dict, strict=True) + print(model) + print("Successfully load") + + model = model.to(device) + model.eval() + return model, processor, tokenizer + +def format_prompt(example): + question = example["question"].strip() + answer = example["answer"] + choices = example.get("choices", []) + image = example["image"] + + if choices: + choices_str = "\n".join([f"({chr(65+i)}).{{{choice.strip()}}}" for i, choice in enumerate(choices)]) + user_prompt = ( + f"[Question]:{{{question}}}\n" + f"[Options]:\n{choices_str}\n" + f"Answer:" + ) + else: + user_prompt = f"[Question]:{{{question}}}\nAnswer:" + + return user_prompt, answer, image + +def process_func(example, idx): + prompt, answer, image = format_prompt(example) + + return { + "idx": idx, + "question_raw": prompt, + "image_raw": image, + "gt_answer": answer, + } + +def has_image(example): + return "image" in example and example["image"] is not None + + +def build_eval_dataset(): + dataset = load_dataset("derek-thomas/ScienceQA") + test_dataset = dataset["test"] + test_dataset = test_dataset.map(lambda example, idx: {"original_idx": idx, **example}, with_indices=True) + test_dataset = test_dataset.filter(has_image) + test_dataset = test_dataset.map(lambda example: process_func(example, example["original_idx"])) + return test_dataset + + +def compute_latent_attention_trace(model, inputs, attn_threshold=None, attn_threshold_multiplier=5.0): + seq_len = inputs["input_ids"].shape[1] + position_ids = torch.arange(seq_len, device=inputs["input_ids"].device).unsqueeze(0) + labels = inputs["input_ids"].clone() + + with torch.no_grad(): + outputs = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + labels=labels, + position_ids=position_ids, + pixel_values=inputs["pixel_values"], + image_grid_thw=inputs["image_grid_thw"], + return_latent_attn=True, + latent_attn_threshold=attn_threshold, + latent_attn_threshold_multiplier=attn_threshold_multiplier, + ) + + return outputs.latent_attn_trace + + +def evaluate_and_save( + eval_dataset, + model, + processor, + output_json_path, + latent_n=3, + max_new_tokens=512, + attn_trace_path=None, + attn_threshold=None, + attn_threshold_multiplier=5.0, +): + model.eval() + correct = 0 + total = 0 + results = {} + total_generated_tokens = 0 + total_generate_time = 0.0 + attn_traces = [] + + output_dir = os.path.dirname(output_json_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + + for ex in tqdm( + eval_dataset, + total=len(eval_dataset), + desc="Evaluating ScienceQA", + dynamic_ncols=True, + ): + idx = str(ex["idx"]) + input_text = ex["question_raw"] + + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": ex["image_raw"], "resized_height": 280, "resized_width": 280}, + {"type": "text", "text": input_text} + ] + }] + + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text = text + ("<|latent|>" * latent_n) + + image_inputs, video_inputs = process_vision_info(messages) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt" + ).to(device) + + if attn_trace_path: + trace = compute_latent_attention_trace( + model, + inputs, + attn_threshold=attn_threshold, + attn_threshold_multiplier=attn_threshold_multiplier, + ) + attn_traces.append({ + "idx": idx, + "latent_attn": trace[0] if trace else [], + }) + + prompt_length = inputs["input_ids"].shape[1] + + generate_start_time = time.time() + + with torch.no_grad(): + outputs = model.generate( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + pixel_values=inputs["pixel_values"], + image_grid_thw=inputs["image_grid_thw"], + max_new_tokens=max_new_tokens + ) + generate_end_time = time.time() + sample_generate_time = generate_end_time - generate_start_time + total_generate_time += sample_generate_time + generated_tokens = outputs[0, prompt_length:] + generated_text = processor.decode(generated_tokens, skip_special_tokens=True) + num_generated_tokens = len(generated_tokens) + total_generated_tokens += num_generated_tokens + + pred_answer = extract_answer(generated_text) + + + results[idx] = pred_answer + + + gt_answer = ex["gt_answer"] + if pred_answer == gt_answer: + correct += 1 + + total += 1 + + + output_data = {"results": results} + with open(output_json_path, "w", encoding="utf-8") as f: + json.dump(output_data, f, ensure_ascii=False, indent=2) + + if attn_trace_path: + attn_dir = os.path.dirname(attn_trace_path) + if attn_dir: + os.makedirs(attn_dir, exist_ok=True) + with open(attn_trace_path, "w", encoding="utf-8") as f: + json.dump(attn_traces, f, ensure_ascii=False, indent=2) + + accuracy = correct / total if total > 0 else 0 + avg_generated_tokens = total_generated_tokens / total if total > 0 else 0 + avg_time_per_sample = total_generate_time / total if total > 0 else 0 + + + logging.info(f"[FINAL] Total: {total}, Correct: {correct}, Accuracy: {accuracy:.2%}") + logging.info(f"[FINAL] Avg generated tokens per sample: {avg_generated_tokens:.1f}") + logging.info(f"[FINAL] Total generate time: {total_generate_time:.2f}s ({timedelta(seconds=int(total_generate_time))})") + logging.info(f"[FINAL] Avg generate time per sample: {avg_time_per_sample:.3f}s") + + + print(f"[FINAL] Total: {total}, Correct: {correct}, Accuracy: {accuracy:.2%}") + print(f"Results saved to: {output_json_path}") + + return accuracy + +def extract_answer(text): + digit_patterns = [ + r'Therefore,?\s*the\s+answer\s+is\s+(\d)', + r'the\s+answer\s+is\s+(\d)', + r'answer\s+is:?\s*(\d)', + ] + + for pattern in digit_patterns: + match = re.search(pattern, text, re.IGNORECASE) + if match: + answer_idx = int(match.group(1)) + logging.debug(f"Extracted answer (digit): {answer_idx}") + return answer_idx + + + letter_patterns = [ + r'Therefore,?\s*the\s+answer\s+is\s+([A-Z])', + r'the\s+answer\s+is\s+([A-Z])', + r'answer\s+is:?\s*([A-Z])', + ] + + for pattern in letter_patterns: + match = re.search(pattern, text, re.IGNORECASE) + if match: + letter = match.group(1).upper() + + answer_idx = ord(letter) - ord('A') + logging.debug(f"Extracted answer (letter): {letter} -> index {answer_idx}") + return answer_idx + + + logging.warning(f"No answer pattern found in text: {text[:200]}") + return -1 + + +def parse_args(): + parser = argparse.ArgumentParser(description="Qwen2-VL IVTLR inference on ScienceQA") + parser.add_argument("--checkpoint_path", type=str, required=True, help="Path to model state_dict checkpoint (.pth)") + parser.add_argument("--latent_n", type=int, default=3, help="Number of <|latent|> tokens appended to the prompt") + parser.add_argument("--patch_reuse_policy", type=str, default=DEFAULT_PATCH_REUSE_POLICY, + choices=["never", "next_step_only", "always"], + help="Patch selection reuse policy during generation") + parser.add_argument("--patch_sampling_strategy", type=str, default="attention_topk", + choices=["attention_topk", "random_image_only", "all_image_patches"], + help="Patch sampling strategy for selecting visual tokens") + parser.add_argument("--output_path", type=str, default="sqa_output/qwen2vl_2b_scienceqa.json", help="Path to write JSON output") + parser.add_argument("--max_new_tokens", type=int, default=512, help="Maximum generated tokens per sample") + parser.add_argument("--attn_trace_path", type=str, default=None, help="Path to write JSON attention traces") + parser.add_argument("--attn_threshold", type=float, default=None, help="Absolute attention threshold for patch logging") + parser.add_argument("--attn_threshold_multiplier", type=float, default=5.0, help="Threshold multiplier for 1/seq_len baseline") + return parser.parse_args() + + +def main(): + args = parse_args() + model, processor, _ = load_inference_model( + args.checkpoint_path, + patch_reuse_policy=args.patch_reuse_policy, + patch_sampling_strategy=args.patch_sampling_strategy, + ) + test_dataset = build_eval_dataset() + evaluate_and_save( + test_dataset, + model, + processor, + output_json_path=args.output_path, + latent_n=args.latent_n, + max_new_tokens=args.max_new_tokens, + attn_trace_path=args.attn_trace_path, + attn_threshold=args.attn_threshold, + attn_threshold_multiplier=args.attn_threshold_multiplier, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen_vl/infer_sqa.py b/qwen_vl/infer_sqa.py index 31c5af5..ba05b0c 100644 --- a/qwen_vl/infer_sqa.py +++ b/qwen_vl/infer_sqa.py @@ -12,6 +12,8 @@ import os import time from datetime import timedelta +import argparse +from tqdm import tqdm logging.basicConfig( filename='qwenvl_32_infer_time.log', level=logging.DEBUG, @@ -21,8 +23,9 @@ import pdb device = "cuda" if torch.cuda.is_available() else "cpu" +DEFAULT_PATCH_REUSE_POLICY = "always" -def load_inference_model(checkpoint_path): +def load_inference_model(checkpoint_path, patch_reuse_policy="never", patch_sampling_strategy="attention_topk"): processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") tokenizer = AutoTokenizer.from_pretrained( "Qwen/Qwen2-VL-7B-Instruct", @@ -75,7 +78,9 @@ def load_inference_model(checkpoint_path): eos_token_id=tokenizer.eos_token_id, image_token_id=image_token_id, visual_start_id=visual_start_id, - visual_end_id=visual_end_id + visual_end_id=visual_end_id, + patch_reuse_policy=patch_reuse_policy, + patch_sampling_strategy=patch_sampling_strategy, ) state_dict = torch.load(checkpoint_path, map_location="cpu") @@ -91,10 +96,6 @@ def load_inference_model(checkpoint_path): model.eval() return model, processor, tokenizer -model, processor, tokenizer = load_inference_model("your_path") - -os.makedirs("output", exist_ok=True) - def format_prompt(example): question = example["question"].strip() answer = example["answer"] @@ -123,28 +124,37 @@ def process_func(example, idx): "gt_answer": answer, } -dataset = load_dataset("derek-thomas/ScienceQA") -test_dataset = dataset["test"] - def has_image(example): return "image" in example and example["image"] is not None -test_dataset = test_dataset.map(lambda example, idx: {"original_idx": idx, **example}, with_indices=True) -test_dataset = test_dataset.filter(has_image) -test_dataset = test_dataset.map(lambda example: process_func(example, example["original_idx"])) +def build_eval_dataset(): + dataset = load_dataset("derek-thomas/ScienceQA") + test_dataset = dataset["test"] + test_dataset = test_dataset.map(lambda example, idx: {"original_idx": idx, **example}, with_indices=True) + test_dataset = test_dataset.filter(has_image) + test_dataset = test_dataset.map(lambda example: process_func(example, example["original_idx"])) + return test_dataset + -def evaluate_and_save(eval_dataset, model, processor): +def evaluate_and_save(eval_dataset, model, processor, output_json_path, latent_n=3, max_new_tokens=512): model.eval() correct = 0 total = 0 results = {} total_generated_tokens = 0 total_generate_time = 0.0 + + output_dir = os.path.dirname(output_json_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) - output_json_path = "sqa_output/qwen_2_scienceqa.json" - - for ex in eval_dataset: + for ex in tqdm( + eval_dataset, + total=len(eval_dataset), + desc="Evaluating ScienceQA", + dynamic_ncols=True, + ): idx = str(ex["idx"]) input_text = ex["question_raw"] @@ -157,7 +167,7 @@ def evaluate_and_save(eval_dataset, model, processor): }] text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - text = text + "<|latent|>" + "<|latent|>" + "<|latent|>" + text = text + ("<|latent|>" * latent_n) image_inputs, video_inputs = process_vision_info(messages) inputs = processor( @@ -178,7 +188,7 @@ def evaluate_and_save(eval_dataset, model, processor): attention_mask=inputs["attention_mask"], pixel_values=inputs["pixel_values"], image_grid_thw=inputs["image_grid_thw"], - max_new_tokens=512 + max_new_tokens=max_new_tokens ) generate_end_time = time.time() sample_generate_time = generate_end_time - generate_start_time @@ -255,4 +265,39 @@ def extract_answer(text): logging.warning(f"No answer pattern found in text: {text[:200]}") return -1 -evaluate_and_save(test_dataset, model, processor) + +def parse_args(): + parser = argparse.ArgumentParser(description="Qwen2-VL IVTLR inference on ScienceQA") + parser.add_argument("--checkpoint_path", type=str, required=True, help="Path to model state_dict checkpoint (.pth)") + parser.add_argument("--latent_n", type=int, default=3, help="Number of <|latent|> tokens appended to the prompt") + parser.add_argument("--patch_reuse_policy", type=str, default=DEFAULT_PATCH_REUSE_POLICY, + choices=["never", "next_step_only", "always"], + help="Patch selection reuse policy during generation") + parser.add_argument("--patch_sampling_strategy", type=str, default="attention_topk", + choices=["attention_topk", "random_image_only", "all_image_patches"], + help="Patch sampling strategy for selecting visual tokens") + parser.add_argument("--output_path", type=str, default="sqa_output/qwen_2_scienceqa.json", help="Path to write JSON output") + parser.add_argument("--max_new_tokens", type=int, default=512, help="Maximum generated tokens per sample") + return parser.parse_args() + + +def main(): + args = parse_args() + model, processor, _ = load_inference_model( + args.checkpoint_path, + patch_reuse_policy=args.patch_reuse_policy, + patch_sampling_strategy=args.patch_sampling_strategy, + ) + test_dataset = build_eval_dataset() + evaluate_and_save( + test_dataset, + model, + processor, + output_json_path=args.output_path, + latent_n=args.latent_n, + max_new_tokens=args.max_new_tokens, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen_vl/qwen_adaptive_ivtlr.py b/qwen_vl/qwen_adaptive_ivtlr.py new file mode 100644 index 0000000..a18f339 --- /dev/null +++ b/qwen_vl/qwen_adaptive_ivtlr.py @@ -0,0 +1,819 @@ +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import CrossEntropyLoss +from transformers.models.gpt2 import GPT2LMHeadModel +from transformers import AutoProcessor + +try: + from controller import PatchPointerController +except ImportError: + from qwen_vl.controller import PatchPointerController + + +@dataclass +class LatentStepTrace: + latent_step_idx: int + reasoning_state: torch.Tensor + image_positions: torch.Tensor + patch_embeddings: torch.Tensor + patch_valid_mask: torch.Tensor + attention_scores: torch.Tensor + ranked_patch_indices: torch.Tensor + selected_patch_embeddings: torch.Tensor + end_position: int + selected_count: int + + +@dataclass +class AdaptiveIVTLROutput: + loss: Optional[torch.Tensor] + ce_loss: Optional[torch.Tensor] + logits: torch.Tensor + inputs_embeds: torch.Tensor + answer_logprob: Optional[torch.Tensor] = None + answer_token_counts: Optional[torch.Tensor] = None + total_inserted_patches: int = 0 + latent_traces: List[LatentStepTrace] = field(default_factory=list) + controller_trace: List[Dict[str, torch.Tensor]] = field(default_factory=list) + + +class QwenAdaptiveIVTLR(nn.Module): + """Adaptive-controller IVT-LR path. + + The existing qwen_ivtlr.IVTLR baseline stays untouched. This wrapper + mirrors its latent insertion mechanics, then adds teacher traces, + forced-budget replay, and controller-driven selection. STOP exists only in + the controller action space and is never inserted into inputs_embeds. + """ + + def __init__( + self, + base_causallm, + latent_token_id: int, + start_latent_id: int, + end_latent_id: int, + eos_token_id: int, + image_token_id: int, + visual_start_id: int, + visual_end_id: int, + controller: Optional[PatchPointerController] = None, + teacher_k: int = 10, + max_controller_steps: int = 10, + patch_reuse_policy: str = "never", + processor_model_id: str = "Qwen/Qwen2-VL-2B-Instruct", + ): + super().__init__() + self.base_causallm = base_causallm + self.latent_token_id = latent_token_id + self.start_latent_id = start_latent_id + self.end_latent_id = end_latent_id + self.eos_token_id = eos_token_id + self.image_token_id = image_token_id + self.visual_start_id = visual_start_id + self.visual_end_id = visual_end_id + self.teacher_k = int(teacher_k) + self.max_controller_steps = int(max_controller_steps) + if patch_reuse_policy not in {"never", "next_step_only", "always"}: + raise ValueError(f"Invalid patch_reuse_policy={patch_reuse_policy}") + self.patch_reuse_policy = patch_reuse_policy + self.processor = AutoProcessor.from_pretrained(processor_model_id) + + if isinstance(self.base_causallm, GPT2LMHeadModel): + self.embedding = self.base_causallm.transformer.get_input_embeddings() + else: + self.embedding = self.base_causallm.get_input_embeddings() + + model_dim = self.embedding.embedding_dim + self.controller = controller or PatchPointerController( + model_dim=model_dim, + max_steps=max_controller_steps, + ) + + def freeze_base_model(self): + for param in self.base_causallm.parameters(): + param.requires_grad = False + return self + + def train_controller_only(self): + self.freeze_base_model() + for param in self.controller.parameters(): + param.requires_grad = True + return self + + def _prepare_inputs_embeds(self, input_ids, pixel_values, image_grid_thw): + inputs_embeds = self.embedding(input_ids) + if pixel_values is None: + image_mask_init = torch.zeros_like(input_ids, dtype=torch.bool) + return inputs_embeds, image_mask_init + + pixel_values = pixel_values.type(self.base_causallm.visual.get_dtype()) + image_embeds = self.base_causallm.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.image_token_id).sum().item() + if n_image_tokens != image_embeds.shape[0]: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, " + f"features {image_embeds.shape[0]}" + ) + image_mask_init = input_ids == self.image_token_id + expand_mask = image_mask_init.unsqueeze(-1).expand(-1, -1, inputs_embeds.size(-1)) + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(expand_mask, image_embeds) + return inputs_embeds, image_mask_init + + @staticmethod + def _pad_1d(items: List[torch.Tensor], value=0, dtype=None): + max_len = max((x.numel() for x in items), default=0) + if max_len == 0: + device = items[0].device + return torch.empty((len(items), 0), device=device, dtype=dtype or items[0].dtype) + out = [] + for item in items: + pad_len = max_len - item.numel() + if pad_len: + pad = torch.full( + (pad_len,), + value, + device=item.device, + dtype=dtype or item.dtype, + ) + item = torch.cat([item.to(dtype=dtype or item.dtype), pad], dim=0) + out.append(item) + return torch.stack(out, dim=0) + + @staticmethod + def _pad_2d(items: List[torch.Tensor]): + max_len = max((x.size(0) for x in items), default=0) + width = items[0].size(-1) + out = [] + for item in items: + pad_len = max_len - item.size(0) + if pad_len: + pad = torch.zeros( + (pad_len, width), + device=item.device, + dtype=item.dtype, + ) + item = torch.cat([item, pad], dim=0) + out.append(item) + return torch.stack(out, dim=0) + + def _rank_image_patches( + self, + inputs_embeds, + image_mask, + attn_scores, + top_k: int, + ): + patch_positions = [] + patch_scores = [] + patch_embeds = [] + ranked_indices = [] + valid_masks = [] + selected_embeds = [] + for b in range(inputs_embeds.size(0)): + positions = image_mask[b, : attn_scores.size(1)].nonzero(as_tuple=True)[0] + scores = attn_scores[b, positions] if positions.numel() else torch.empty(0, device=attn_scores.device) + order = torch.argsort(scores, descending=True) + ranked = order[: min(top_k, order.numel())] + embeds = inputs_embeds[b, positions, :] + patch_positions.append(positions) + patch_scores.append(scores) + patch_embeds.append(embeds) + ranked_indices.append(ranked) + valid_masks.append(torch.ones(positions.numel(), device=inputs_embeds.device, dtype=torch.bool)) + selected_embeds.append(embeds[ranked] if ranked.numel() else embeds[:0]) + + return { + "positions": self._pad_1d(patch_positions, value=-1, dtype=torch.long), + "scores": self._pad_1d(patch_scores, value=float("-inf")), + "embeddings": self._pad_2d(patch_embeds), + "ranked_indices": self._pad_1d(ranked_indices, value=-1, dtype=torch.long), + "valid_mask": self._pad_1d(valid_masks, value=False, dtype=torch.bool), + "selected_embeddings": self._pad_2d(selected_embeds), + } + + def _select_for_step( + self, + mode: str, + step_idx: int, + reasoning_state: torch.Tensor, + ranked: Dict[str, torch.Tensor], + forced_budget: Optional[int], + teacher_trace: Optional[Sequence[LatentStepTrace]], + ): + positions = ranked["positions"] + if mode == "teacher": + rel_indices = ranked["ranked_indices"][:, : self.teacher_k] + elif mode == "forced_budget": + if teacher_trace is None: + raise ValueError("forced_budget mode requires teacher_trace") + rel_indices = teacher_trace[step_idx].ranked_patch_indices[:, : forced_budget] + elif mode == "adaptive": + if reasoning_state.size(0) != 1: + raise ValueError("adaptive controller inference currently expects batch_size=1") + selection = self.controller.greedy_select( + reasoning_state, + ranked["embeddings"], + ranked["valid_mask"], + max_steps=self.max_controller_steps, + ) + rel_indices = selection["selected_indices"][:, : int(selection["lengths"].max().item())] + elif mode == "adaptive_sample": + if reasoning_state.size(0) != 1: + raise ValueError("sampled adaptive controller currently expects batch_size=1") + selection = self.controller.sample_select( + reasoning_state, + ranked["embeddings"], + ranked["valid_mask"], + max_steps=self.max_controller_steps, + temperature=getattr(self, "_controller_sample_temperature", 1.0), + min_patches=getattr(self, "_controller_min_patches", 0), + ) + rel_indices = selection["selected_indices"][:, : int(selection["lengths"].max().item())] + else: + raise ValueError(f"Unknown adaptive IVT-LR mode: {mode}") + + selected_positions = [] + selected_embeds = [] + selected_counts = [] + for b in range(positions.size(0)): + rel = rel_indices[b] + rel = rel[rel >= 0] + valid_rel = rel[rel < ranked["embeddings"].size(1)] + pos = positions[b, valid_rel] if valid_rel.numel() else positions[b, :0] + pos = pos[pos >= 0] + embeds = ( + ranked["embeddings"][b, valid_rel, :] + if valid_rel.numel() + else ranked["embeddings"][b, :0, :] + ) + selected_positions.append(pos) + selected_embeds.append(embeds) + selected_counts.append(embeds.size(0)) + return selected_positions, selected_embeds, selected_counts, rel_indices, selection if mode in {"adaptive", "adaptive_sample"} else None + + def _merge_selected_embeddings( + self, + inputs_embeds, + attention_mask, + position_ids, + original_mask, + image_mask, + trace_mask, + latent_lists, + end, + selected_positions, + selected_embeds, + selected_counts, + hidden_states, + pass_idx, + ): + inputs_embeds_detached = inputs_embeds.detach().clone() + for b in range(inputs_embeds.size(0)): + if len(latent_lists[b]) > pass_idx: + t_idx = latent_lists[b][pass_idx] + rel_pos = max(0, min(t_idx - 1, hidden_states.size(1) - 1)) + inputs_embeds_detached[b, t_idx, :] = hidden_states[b, rel_pos, :] + inputs_embeds = inputs_embeds_detached + + new_embeds = [] + new_att = [] + new_pos = [] + new_orig = [] + new_img = [] + new_trace = [] + max_len = 0 + for b in range(inputs_embeds.size(0)): + k_b = selected_counts[b] + merged = torch.cat( + [inputs_embeds[b, :end, :], selected_embeds[b], inputs_embeds[b, end:, :]], + dim=0, + ) + att = torch.cat( + [ + attention_mask[b, :end], + torch.ones(k_b, device=attention_mask.device, dtype=attention_mask.dtype), + attention_mask[b, end:], + ], + dim=0, + ) + pos = torch.arange(merged.size(0), device=position_ids.device) + orig = torch.cat( + [ + original_mask[b, :end], + torch.zeros(k_b, device=original_mask.device, dtype=torch.bool), + original_mask[b, end:], + ], + dim=0, + ) + img = torch.cat( + [ + image_mask[b, :end], + torch.zeros(k_b, device=image_mask.device, dtype=torch.bool), + image_mask[b, end:], + ], + dim=0, + ) + trace = torch.cat( + [ + trace_mask[b, :end], + torch.ones(k_b, device=trace_mask.device, dtype=torch.bool), + trace_mask[b, end:], + ], + dim=0, + ) + new_embeds.append(merged) + new_att.append(att) + new_pos.append(pos) + new_orig.append(orig) + new_img.append(img) + new_trace.append(trace) + max_len = max(max_len, merged.size(0)) + + def pad_vec(item, value=0): + if item.size(0) == max_len: + return item + return F.pad(item, (0, max_len - item.size(0)), value=value) + + def pad_embed(item): + if item.size(0) == max_len: + return item + pad = torch.zeros( + (max_len - item.size(0), item.size(1)), + device=item.device, + dtype=item.dtype, + ) + return torch.cat([item, pad], dim=0) + + inputs_embeds = torch.stack([pad_embed(x) for x in new_embeds], dim=0) + attention_mask = torch.stack([pad_vec(x) for x in new_att], dim=0) + position_ids = torch.stack([pad_vec(x) for x in new_pos], dim=0) + original_mask = torch.stack([pad_vec(x).bool() for x in new_orig], dim=0) + image_mask = torch.stack([pad_vec(x).bool() for x in new_img], dim=0) + trace_mask = torch.stack([pad_vec(x).bool() for x in new_trace], dim=0) + + if self.patch_reuse_policy == "never": + for b, pos in enumerate(selected_positions): + image_mask[b, pos] = False + + for b, k_b in enumerate(selected_counts): + for i, pos in enumerate(latent_lists[b]): + if pos > end: + latent_lists[b][i] = pos + k_b + + return inputs_embeds, attention_mask, position_ids, original_mask, image_mask, trace_mask + + @staticmethod + def _answer_logprob(logits, labels, answer_mask): + if answer_mask is None: + answer_mask = labels != -100 + seq_len = min(logits.size(1), labels.size(1), answer_mask.size(1)) + logits = logits[:, -seq_len:, :] + labels = labels[:, -seq_len:] + answer_mask = answer_mask[:, -seq_len:].bool() + shift_logits = logits[:, :-1, :] + shift_labels = labels[:, 1:] + shift_mask = answer_mask[:, 1:] & (shift_labels != -100) + log_probs = F.log_softmax(shift_logits, dim=-1) + safe_labels = shift_labels.masked_fill(shift_labels == -100, 0) + gathered = log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1) + gathered = gathered.masked_fill(~shift_mask, 0.0) + counts = shift_mask.sum(dim=-1) + avg_logprob = gathered.sum(dim=-1) / counts.clamp(min=1) + return avg_logprob, counts + + def forward( + self, + input_ids, + attention_mask, + labels, + position_ids, + pixel_values, + image_grid_thw=None, + answer_mask=None, + mode: str = "teacher", + forced_budget: Optional[int] = None, + teacher_trace: Optional[Sequence[LatentStepTrace]] = None, + return_trace: bool = False, + ) -> AdaptiveIVTLROutput: + bsz, seq_len = input_ids.shape + inputs_embeds, image_mask_init = self._prepare_inputs_embeds( + input_ids, pixel_values, image_grid_thw + ) + original_mask = torch.ones((bsz, seq_len), dtype=torch.bool, device=input_ids.device) + image_mask = image_mask_init.clone() + trace_mask = torch.zeros_like(image_mask) + + vs_indices = (input_ids == self.visual_start_id).nonzero(as_tuple=True) + ve_indices = (input_ids == self.visual_end_id).nonzero(as_tuple=True) + vs_pos_per_batch = {b.item(): vs_indices[1][i].item() for i, b in enumerate(vs_indices[0])} + ve_pos_per_batch = {b.item(): ve_indices[1][i].item() for i, b in enumerate(ve_indices[0])} + for b in range(bsz): + if b in vs_pos_per_batch and b in ve_pos_per_batch: + image_mask[b, vs_pos_per_batch[b] + 1 : ve_pos_per_batch[b]] = True + + latent_indices = (input_ids == self.latent_token_id).nonzero() + latent_lists = [ + [idx[1].item() for idx in latent_indices if idx[0] == b] + for b in range(bsz) + ] + max_latents = max((len(x) for x in latent_lists), default=0) + end = min((lst[0] for lst in latent_lists if lst), default=seq_len) + all_logits = [] + traces = [] + controller_trace = [] + total_inserted = 0 + + for pass_idx in range(max_latents): + outputs = self.base_causallm( + inputs_embeds=inputs_embeds[:, :end, :], + attention_mask=attention_mask[:, :end], + position_ids=position_ids[:, :end], + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + output_hidden_states=True, + output_attentions=True, + ) + all_logits.append(outputs.logits) + hidden_states = outputs.hidden_states[-1] + attentions = outputs.attentions + avg_attn = torch.cat(attentions, dim=1).mean(dim=1) + attn_scores = avg_attn[:, end - 1, :] + reasoning_state = [] + for b in range(bsz): + t_idx = latent_lists[b][pass_idx] + rel_pos = max(0, min(t_idx - 1, hidden_states.size(1) - 1)) + reasoning_state.append(hidden_states[b, rel_pos, :]) + reasoning_state = torch.stack(reasoning_state, dim=0) + + ranked = self._rank_image_patches( + inputs_embeds, + image_mask[:, : attn_scores.size(1)], + attn_scores, + top_k=max(self.teacher_k, forced_budget or 0), + ) + selected_positions, selected_embeds, selected_counts, selected_rel_indices, selection_info = self._select_for_step( + mode, + pass_idx, + reasoning_state, + ranked, + forced_budget, + teacher_trace, + ) + total_inserted += sum(selected_counts) + if return_trace or mode == "teacher": + traces.append( + LatentStepTrace( + latent_step_idx=pass_idx, + reasoning_state=reasoning_state.detach(), + image_positions=ranked["positions"].detach(), + patch_embeddings=ranked["embeddings"].detach(), + patch_valid_mask=ranked["valid_mask"].detach(), + attention_scores=ranked["scores"].detach(), + ranked_patch_indices=ranked["ranked_indices"].detach(), + selected_patch_embeddings=ranked["selected_embeddings"].detach(), + end_position=end, + selected_count=max(selected_counts) if selected_counts else 0, + ) + ) + if mode in {"adaptive", "adaptive_sample"}: + step_trace = { + "latent_step_idx": pass_idx, + "selected_counts": torch.tensor(selected_counts), + "selected_patch_indices": selected_rel_indices.detach().cpu(), + } + if selection_info is not None and "logprob_sum" in selection_info: + step_trace.update( + { + "logprob_sum": selection_info["logprob_sum"], + "entropy_sum": selection_info["entropy_sum"], + "action_count": selection_info["action_count"], + } + ) + controller_trace.append(step_trace) + + ( + inputs_embeds, + attention_mask, + position_ids, + original_mask, + image_mask, + trace_mask, + ) = self._merge_selected_embeddings( + inputs_embeds, + attention_mask, + position_ids, + original_mask, + image_mask, + trace_mask, + latent_lists, + end, + selected_positions, + selected_embeds, + selected_counts, + hidden_states, + pass_idx, + ) + if pass_idx + 1 >= max_latents: + end = inputs_embeds.size(1) + else: + end = end + 1 + max(selected_counts) + + outputs = self.base_causallm( + inputs_embeds=inputs_embeds[:, :end, :], + attention_mask=attention_mask[:, :end], + position_ids=position_ids[:, :end], + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + output_hidden_states=False, + output_attentions=False, + ) + all_logits.append(outputs.logits) + logits = torch.cat(all_logits, dim=1) + + final_s = logits.size(1) + new_labels = torch.full((bsz, final_s), -100, device=input_ids.device, dtype=labels.dtype) + new_answer_mask = torch.zeros((bsz, final_s), device=input_ids.device, dtype=torch.bool) + for b in range(bsz): + label_len = labels.size(1) + new_labels[b, -label_len:] = labels[b] + if answer_mask is not None: + new_answer_mask[b, -label_len:] = answer_mask[b].bool() + else: + new_answer_mask[b, -label_len:] = labels[b] != -100 + + shift_logits = logits[:, :-1, :].contiguous() + shift_labels = new_labels[:, 1:].contiguous() + ce_loss = CrossEntropyLoss(ignore_index=-100)( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1), + ) + answer_logprob, answer_counts = self._answer_logprob(logits, new_labels, new_answer_mask) + return AdaptiveIVTLROutput( + loss=ce_loss, + ce_loss=ce_loss, + logits=logits, + inputs_embeds=inputs_embeds, + answer_logprob=answer_logprob, + answer_token_counts=answer_counts, + total_inserted_patches=total_inserted, + latent_traces=traces, + controller_trace=controller_trace, + ) + + def controller_teacher_forcing_loss( + self, + teacher_trace: Sequence[LatentStepTrace], + budgets: Sequence[int], + budget_weights: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + losses = [] + patch_accs = [] + stop_accs = [] + device = budget_weights.device + for step in teacher_trace: + reasoning = step.reasoning_state + patch_embeddings = self._gather_trace_patch_embeddings(step) + patch_valid = step.patch_valid_mask.bool() + n_patches = patch_embeddings.size(1) + for budget_idx, budget in enumerate(budgets): + k = min(int(budget), step.ranked_patch_indices.size(1)) + targets = step.ranked_patch_indices[:, :k] + stop = torch.full( + (targets.size(0), 1), + n_patches, + device=targets.device, + dtype=torch.long, + ) + target_actions = torch.cat([targets, stop], dim=1) + weights = budget_weights[:, budget_idx].to(device) + stats = self.controller.teacher_forced_sequence_loss( + reasoning, + patch_embeddings, + patch_valid, + target_actions, + sequence_weights=weights, + ) + losses.append(stats.loss) + patch_accs.append(stats.patch_top1_accuracy) + stop_accs.append(stats.stop_accuracy) + if not losses: + zero = torch.tensor(0.0, device=device, requires_grad=True) + return {"loss": zero, "patch_top1_accuracy": zero.detach(), "stop_accuracy": zero.detach()} + return { + "loss": torch.stack(losses).mean(), + "patch_top1_accuracy": torch.stack(patch_accs).mean(), + "stop_accuracy": torch.stack(stop_accs).mean(), + } + + @staticmethod + def _gather_trace_patch_embeddings(step: LatentStepTrace) -> torch.Tensor: + return step.patch_embeddings + + @torch.no_grad() + def generate( + self, + input_ids, + attention_mask, + pixel_values, + image_grid_thw, + max_new_tokens: int = 128, + output_controller_trace: bool = False, + ): + if input_ids.size(0) != 1: + raise ValueError("Adaptive controller generation currently supports batch_size=1.") + + self.eval() + position_ids = torch.arange( + input_ids.size(1), + dtype=torch.long, + device=input_ids.device, + ).unsqueeze(0) + adaptive_out = self.forward( + input_ids=input_ids, + attention_mask=attention_mask, + labels=input_ids.clone(), + position_ids=position_ids, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + mode="adaptive", + ) + + tokens = input_ids[0].detach().tolist() + next_token = torch.argmax(adaptive_out.logits[0, -1]).item() + tokens.append(next_token) + + current_inputs_embeds = adaptive_out.inputs_embeds + current_attention_mask = torch.ones( + (1, current_inputs_embeds.size(1)), + device=current_inputs_embeds.device, + dtype=attention_mask.dtype, + ) + next_token_embedding = self.embedding( + torch.tensor([[next_token]], device=current_inputs_embeds.device) + ) + current_inputs_embeds = torch.cat([current_inputs_embeds, next_token_embedding], dim=1) + current_attention_mask = torch.cat( + [ + current_attention_mask, + torch.ones((1, 1), device=current_inputs_embeds.device, dtype=attention_mask.dtype), + ], + dim=1, + ) + + past_key_values = None + for _ in range(max_new_tokens - 1): + if past_key_values is None: + inputs_embeds_for_forward = current_inputs_embeds + attention_mask_for_forward = current_attention_mask + position_ids = torch.arange( + current_inputs_embeds.size(1), + dtype=torch.long, + device=current_inputs_embeds.device, + ).unsqueeze(0) + else: + inputs_embeds_for_forward = next_token_embedding + attention_mask_for_forward = current_attention_mask + position_ids = torch.tensor( + [[current_inputs_embeds.size(1) - 1]], + dtype=torch.long, + device=current_inputs_embeds.device, + ) + + outputs = self.base_causallm( + inputs_embeds=inputs_embeds_for_forward, + attention_mask=attention_mask_for_forward, + position_ids=position_ids, + pixel_values=pixel_values if past_key_values is None else None, + image_grid_thw=image_grid_thw if past_key_values is None else None, + past_key_values=past_key_values, + use_cache=True, + ) + past_key_values = outputs.past_key_values + next_token = torch.argmax(outputs.logits[0, -1]).item() + tokens.append(next_token) + if next_token == self.eos_token_id: + break + + next_token_embedding = self.embedding( + torch.tensor([[next_token]], device=current_inputs_embeds.device) + ) + current_inputs_embeds = torch.cat([current_inputs_embeds, next_token_embedding], dim=1) + current_attention_mask = torch.cat( + [ + current_attention_mask, + torch.ones((1, 1), device=current_inputs_embeds.device, dtype=attention_mask.dtype), + ], + dim=1, + ) + + output_ids = torch.tensor(tokens, dtype=torch.long, device=input_ids.device).unsqueeze(0) + if output_controller_trace: + return output_ids, adaptive_out.controller_trace + return output_ids + + def generate_with_sampled_controller( + self, + input_ids, + attention_mask, + pixel_values, + image_grid_thw, + max_new_tokens: int = 128, + controller_temperature: float = 1.0, + min_patches: int = 0, + ): + if input_ids.size(0) != 1: + raise ValueError("Sampled controller generation currently supports batch_size=1.") + + self._controller_sample_temperature = controller_temperature + self._controller_min_patches = min_patches + position_ids = torch.arange( + input_ids.size(1), + dtype=torch.long, + device=input_ids.device, + ).unsqueeze(0) + adaptive_out = self.forward( + input_ids=input_ids, + attention_mask=attention_mask, + labels=input_ids.clone(), + position_ids=position_ids, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + mode="adaptive_sample", + ) + + tokens = input_ids[0].detach().tolist() + next_token = torch.argmax(adaptive_out.logits[0, -1].detach()).item() + tokens.append(next_token) + + current_inputs_embeds = adaptive_out.inputs_embeds.detach() + current_attention_mask = torch.ones( + (1, current_inputs_embeds.size(1)), + device=current_inputs_embeds.device, + dtype=attention_mask.dtype, + ) + next_token_embedding = self.embedding( + torch.tensor([[next_token]], device=current_inputs_embeds.device) + ).detach() + current_inputs_embeds = torch.cat([current_inputs_embeds, next_token_embedding], dim=1) + current_attention_mask = torch.cat( + [ + current_attention_mask, + torch.ones((1, 1), device=current_inputs_embeds.device, dtype=attention_mask.dtype), + ], + dim=1, + ) + + past_key_values = None + with torch.no_grad(): + for _ in range(max_new_tokens - 1): + if past_key_values is None: + inputs_embeds_for_forward = current_inputs_embeds + attention_mask_for_forward = current_attention_mask + position_ids = torch.arange( + current_inputs_embeds.size(1), + dtype=torch.long, + device=current_inputs_embeds.device, + ).unsqueeze(0) + else: + inputs_embeds_for_forward = next_token_embedding + attention_mask_for_forward = current_attention_mask + position_ids = torch.tensor( + [[current_inputs_embeds.size(1) - 1]], + dtype=torch.long, + device=current_inputs_embeds.device, + ) + + outputs = self.base_causallm( + inputs_embeds=inputs_embeds_for_forward, + attention_mask=attention_mask_for_forward, + position_ids=position_ids, + pixel_values=pixel_values if past_key_values is None else None, + image_grid_thw=image_grid_thw if past_key_values is None else None, + past_key_values=past_key_values, + use_cache=True, + ) + past_key_values = outputs.past_key_values + next_token = torch.argmax(outputs.logits[0, -1]).item() + tokens.append(next_token) + if next_token == self.eos_token_id: + break + + next_token_embedding = self.embedding( + torch.tensor([[next_token]], device=current_inputs_embeds.device) + ).detach() + current_inputs_embeds = torch.cat([current_inputs_embeds, next_token_embedding], dim=1) + current_attention_mask = torch.cat( + [ + current_attention_mask, + torch.ones((1, 1), device=current_inputs_embeds.device, dtype=attention_mask.dtype), + ], + dim=1, + ) + + output_ids = torch.tensor(tokens, dtype=torch.long, device=input_ids.device).unsqueeze(0) + return output_ids, adaptive_out.controller_trace diff --git a/qwen_vl/qwen_ivtlr.py b/qwen_vl/qwen_ivtlr.py index 327cf3f..9b3cbfb 100644 --- a/qwen_vl/qwen_ivtlr.py +++ b/qwen_vl/qwen_ivtlr.py @@ -1,5 +1,6 @@ import torch import torch.nn as nn +import torch.nn.functional as F from torch.nn import CrossEntropyLoss from collections import namedtuple from transformers.models.gpt2 import GPT2LMHeadModel @@ -14,8 +15,21 @@ import pdb from transformers.cache_utils import DynamicCache -Outputs = namedtuple("Outputs", ["loss", "inputs_embeds", "logits"]) -MAX_N_LATENT = 4 +Outputs = namedtuple( + "Outputs", + [ + "loss", + "ce_loss", + "nvt_loss", + "qvr_loss", + "causal_loss", + "inputs_embeds", + "logits", + "latent_attn_trace", + "token_norms", + ], +) +MAX_N_LATENT = 4 class IVTLR(nn.Module): @@ -31,6 +45,19 @@ def __init__( visual_start_id, visual_end_id, num_selected_patches: int = 32, + patch_reuse_policy: str = "never", + patch_sampling_strategy: str = "attention_topk", + processor_model_id: str = "Qwen/Qwen2-VL-7B-Instruct", + enable_nvt_loss: bool = False, + nvt_loss_weight: float = 0.0, + nvt_loss_epsilon: float = 1e-8, + enable_qvr_loss: bool = False, + qvr_loss_weight: float = 0.0, + qvr_loss_epsilon: float = 1e-8, + qvr_num_layers: int = 4, + enable_causal_loss: bool = False, + causal_loss_weight: float = 0.0, + causal_loss_epsilon: float = 1e-8, ): super(IVTLR, self).__init__() @@ -44,6 +71,27 @@ def __init__( self.visual_start_id = visual_start_id self.visual_end_id = visual_end_id self.num_selected_patches = num_selected_patches + valid_policies = {"never", "next_step_only", "always"} + if patch_reuse_policy not in valid_policies: + raise ValueError(f"Invalid patch_reuse_policy={patch_reuse_policy}. Expected one of {valid_policies}.") + self.patch_reuse_policy = patch_reuse_policy + valid_sampling_strategies = {"attention_topk", "random_image_only", "all_image_patches"} + if patch_sampling_strategy not in valid_sampling_strategies: + raise ValueError( + f"Invalid patch_sampling_strategy={patch_sampling_strategy}. " + f"Expected one of {valid_sampling_strategies}." + ) + self.patch_sampling_strategy = patch_sampling_strategy + self.enable_nvt_loss = enable_nvt_loss and nvt_loss_weight > 0 + self.nvt_loss_weight = nvt_loss_weight + self.nvt_loss_epsilon = float(nvt_loss_epsilon) if nvt_loss_epsilon is not None else 1e-8 + self.enable_qvr_loss = enable_qvr_loss and qvr_loss_weight > 0 + self.qvr_loss_weight = qvr_loss_weight + self.qvr_loss_epsilon = float(qvr_loss_epsilon) if qvr_loss_epsilon is not None else 1e-8 + self.qvr_num_layers = max(int(qvr_num_layers), 1) + self.enable_causal_loss = enable_causal_loss and causal_loss_weight > 0 + self.causal_loss_weight = causal_loss_weight + self.causal_loss_epsilon = float(causal_loss_epsilon) if causal_loss_epsilon is not None else 1e-8 # tested with GPT2 and Llama3 if isinstance(self.base_causallm, GPT2LMHeadModel): @@ -52,7 +100,115 @@ def __init__( self.embedding = self.base_causallm.get_input_embeddings() # self.processor = ChameleonProcessor.from_pretrained("facebook/chameleon-7b") - self.processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") + self.processor = AutoProcessor.from_pretrained(processor_model_id) + + def _compute_nvt_loss(self, attentions, query_index, inserted_spans): + if not inserted_spans: + return None + + per_batch_losses = [] + for batch_index, span in enumerate(inserted_spans): + if span is None: + continue + + span_start, span_end = span + if span_end <= span_start: + continue + + layer_masses = [] + for layer_attn in attentions: + if query_index >= layer_attn.size(-2) or span_end > layer_attn.size(-1): + continue + token_to_span = layer_attn[batch_index, :, query_index, span_start:span_end].sum(dim=-1) + layer_masses.append(token_to_span.mean()) + + if not layer_masses: + continue + + mt = torch.stack(layer_masses).mean() + per_batch_losses.append(-torch.log(mt + self.nvt_loss_epsilon)) + + if not per_batch_losses: + return None + + return torch.stack(per_batch_losses).mean() + + def _compute_qvr_loss(self, attentions, query_positions, inserted_spans, question_mask): + if not inserted_spans: + return None + + avg_attn = self._average_last_attentions(attentions, self.qvr_num_layers) + if avg_attn is None: + return None + + attn = avg_attn.mean(dim=1) + seq_len = attn.size(-1) + + per_batch_losses = [] + for batch_index, span in enumerate(inserted_spans): + if span is None: + continue + + query_index = query_positions[batch_index] + if query_index is None or query_index < 0 or query_index >= seq_len: + continue + + span_start, span_end = span + if span_end <= span_start: + continue + + vis_mass = attn[batch_index, query_index, span_start:span_end].sum() + ques_mass = attn[batch_index, query_index, question_mask[batch_index]].sum() + h_val = (2.0 * vis_mass * ques_mass) / (vis_mass + ques_mass + self.qvr_loss_epsilon) + per_batch_losses.append(-torch.log(h_val + self.qvr_loss_epsilon)) + + if not per_batch_losses: + return None + + return torch.stack(per_batch_losses).mean() + + @staticmethod + def _summarize_norms(norms_tensor: torch.Tensor): + if norms_tensor.numel() == 0: + return {"count": 0, "mean": None, "std": None, "min": None, "max": None} + norms_tensor = norms_tensor.float() + return { + "count": int(norms_tensor.numel()), + "mean": float(norms_tensor.mean().item()), + "std": float(norms_tensor.std(unbiased=False).item()), + "min": float(norms_tensor.min().item()), + "max": float(norms_tensor.max().item()), + } + + @staticmethod + def _average_last_attentions(attentions, num_layers): + if not attentions: + return None + n_layers = min(num_layers, len(attentions)) + stacked = torch.stack(attentions[-n_layers:], dim=0) + return stacked.mean(dim=0) + + @staticmethod + def _compute_answer_logprob(logits, labels, answer_mask): + seq_len = min(logits.size(1), labels.size(1), answer_mask.size(1)) + logits = logits[..., -seq_len:, :] + labels = labels[..., -seq_len:] + answer_mask = answer_mask[..., -seq_len:] + + shift_logits = logits[..., :-1, :] + shift_labels = labels[..., 1:] + shift_answer_mask = answer_mask[..., 1:].bool() + log_probs = torch.log_softmax(shift_logits, dim=-1) + safe_labels = shift_labels.clone() + safe_labels[safe_labels == -100] = 0 + gathered = log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1) + valid = shift_answer_mask & (shift_labels != -100) + gathered = gathered.masked_fill(~valid, 0.0) + token_counts = valid.sum(dim=-1) + token_counts_clamped = token_counts.clamp(min=1) + mean_logprob = gathered.sum(dim=-1) / token_counts_clamped + return mean_logprob, token_counts + def forward( self, input_ids: torch.LongTensor, # shape = (B, S) @@ -61,6 +217,11 @@ def forward( position_ids: torch.LongTensor, # shape = (B, S) pixel_values: torch.FloatTensor, # shape = (B, 3, H, W) image_grid_thw: torch.Tensor = None, + answer_mask: torch.LongTensor = None, + return_latent_attn: bool = False, + return_token_norms: bool = False, + latent_attn_threshold: float = None, + latent_attn_threshold_multiplier: float = 5.0, **kwargs ): @@ -99,6 +260,14 @@ def forward( max_len = 3000 image_mask = torch.zeros((B, max_len), dtype=torch.bool, device=input_ids.device) image_mask[:, :S] = image_mask_init + trace_mask = torch.zeros((B, max_len), dtype=torch.bool, device=input_ids.device) + recently_selected_mask = torch.zeros((B, max_len), dtype=torch.bool, device=input_ids.device) + patch_insert_mask = None + reasoning_insert_mask = None + need_insert_origin_masks = return_token_norms or self.enable_qvr_loss or self.enable_causal_loss + if need_insert_origin_masks: + patch_insert_mask = torch.zeros((B, max_len), dtype=torch.bool, device=input_ids.device) + reasoning_insert_mask = torch.zeros((B, max_len), dtype=torch.bool, device=input_ids.device) for b in range(B): @@ -120,6 +289,9 @@ def forward( kv_cache = None all_logits = [] + nvt_losses = [] + qvr_losses = [] + prev_inserted_spans = None if max_n_latents > 0: for pass_idx in range(max_n_latents): @@ -155,28 +327,168 @@ def forward( all_logits.append(logits_this) + if self.enable_nvt_loss and prev_inserted_spans is not None: + nvt_loss = self._compute_nvt_loss(attentions, end - 1, prev_inserted_spans) + if nvt_loss is not None: + nvt_losses.append(nvt_loss) + + if self.enable_qvr_loss and prev_inserted_spans is not None: + if B > 0: + seq_len_this_pass = attention_mask[:, :end].size(1) + positions_this_pass = torch.arange(seq_len_this_pass, device=input_ids.device).unsqueeze(0).expand(B, -1) + first_latent_pos = min(lst[0] for lst in latent_lists if len(lst) > 0) if any(len(lst) > 0 for lst in latent_lists) else None + if first_latent_pos is not None: + question_mask = ( + original_mask[:, :end] + & (~image_mask[:, :end]) + & (positions_this_pass < first_latent_pos) + ) + query_positions = [ + (lst[pass_idx] if pass_idx < len(lst) else (lst[-1] if lst else None)) + for lst in latent_lists + ] + qvr_loss = self._compute_qvr_loss( + attentions, + query_positions, + prev_inserted_spans, + question_mask, + ) + if qvr_loss is not None: + qvr_losses.append(qvr_loss) + # Top-K avg_attn = torch.cat(attentions, dim=1).mean(dim=1) # (B, seq_len) current_seq_len = avg_attn.size(1) select_image_embeds = [] + current_selected_mask = torch.zeros_like(image_mask) + selected_counts = [] + selected_patch_origins = [] if need_insert_origin_masks else None + selected_reasoning_origins = [] if need_insert_origin_masks else None for b in range(B): last_attn = avg_attn[b, end - 1] # shape=(seq_len,) vs, ve = vs_pos_per_batch[b], ve_pos_per_batch[b] scores = last_attn.clone() - allowed_positions = image_mask[b, :current_seq_len] # shape=(S,) - invalid = ~allowed_positions - scores[invalid] = float("-inf") - - rel_scores = scores[vs+1 : ve] # (image_len,) - topk_rel = rel_scores.topk(self.num_selected_patches, sorted=False)[1] # rel idx - abs_idxs = (vs + 1) + topk_rel - logging.debug(f"topk_rel: {topk_rel}") + + image_allowed_positions = image_mask[b, :current_seq_len] + trace_allowed_positions = trace_mask[b, :current_seq_len] + if self.patch_reuse_policy == "next_step_only": + not_recent = ~recently_selected_mask[b, :current_seq_len] + image_allowed_positions = image_allowed_positions & not_recent + trace_allowed_positions = trace_allowed_positions & not_recent + + if self.patch_sampling_strategy == "all_image_patches": + image_abs_idxs = torch.arange(vs + 1, ve, device=input_ids.device) + if image_abs_idxs.numel() == 0: + raise ValueError("No image patch positions available for all_image_patches strategy.") + abs_idxs = image_abs_idxs + trace_abs_idxs = torch.empty(0, dtype=torch.long, device=input_ids.device) + elif self.patch_sampling_strategy == "random_image_only": + image_pool_mask = image_allowed_positions.clone() + image_pool_mask[:vs + 1] = False + image_pool_mask[ve:] = False + image_candidates = torch.nonzero(image_pool_mask, as_tuple=False).squeeze(-1) + + if image_candidates.numel() >= self.num_selected_patches: + rand_order = torch.randperm(image_candidates.numel(), device=input_ids.device) + abs_idxs = image_candidates[rand_order[:self.num_selected_patches]] + elif image_candidates.numel() > 0: + n_to_fill = self.num_selected_patches - image_candidates.numel() + fill_idxs = image_candidates[torch.randint( + low=0, + high=image_candidates.numel(), + size=(n_to_fill,), + device=input_ids.device, + )] + abs_idxs = torch.cat([image_candidates, fill_idxs], dim=0) + else: + image_span = torch.arange(vs + 1, ve, device=input_ids.device) + if image_span.numel() == 0: + raise ValueError("No image patch positions available for random_image_only sampling.") + fill_rand = torch.randint( + low=0, + high=image_span.numel(), + size=(self.num_selected_patches,), + device=input_ids.device, + ) + abs_idxs = image_span[fill_rand] + + image_abs_idxs = abs_idxs + trace_abs_idxs = torch.empty(0, dtype=torch.long, device=input_ids.device) + else: + if pass_idx == 0: + image_quota = self.num_selected_patches + trace_quota = 0 + else: + trace_quota = self.num_selected_patches // 2 + image_quota = self.num_selected_patches - trace_quota + + image_scores = scores.clone() + image_invalid = ~image_allowed_positions + image_scores[image_invalid] = float("-inf") + image_rel_scores = image_scores[vs + 1 : ve] + n_image_candidates = int(image_allowed_positions[vs + 1 : ve].sum().item()) + image_take = min(image_quota, n_image_candidates) + if image_take > 0: + topk_image_rel = image_rel_scores.topk(image_take, sorted=False)[1] + image_abs_idxs = (vs + 1) + topk_image_rel + else: + image_abs_idxs = torch.empty(0, dtype=torch.long, device=input_ids.device) + + trace_scores = scores.clone() + trace_invalid = ~trace_allowed_positions + trace_scores[trace_invalid] = float("-inf") + n_trace_candidates = int(trace_allowed_positions.sum().item()) + trace_take = min(trace_quota, n_trace_candidates) + if trace_take > 0: + trace_abs_idxs = trace_scores.topk(trace_take, sorted=False)[1] + else: + trace_abs_idxs = torch.empty(0, dtype=torch.long, device=input_ids.device) + + abs_idxs = torch.cat([image_abs_idxs, trace_abs_idxs], dim=0) + + if abs_idxs.numel() < self.num_selected_patches: + combined_allowed = image_allowed_positions | trace_allowed_positions + if abs_idxs.numel() > 0: + combined_allowed[abs_idxs] = False + combined_scores = scores.clone() + combined_scores[~combined_allowed] = float("-inf") + n_extra_candidates = int(combined_allowed.sum().item()) + n_to_fill = min(self.num_selected_patches - abs_idxs.numel(), n_extra_candidates) + if n_to_fill > 0: + extra_abs_idxs = combined_scores.topk(n_to_fill, sorted=False)[1] + abs_idxs = torch.cat([abs_idxs, extra_abs_idxs], dim=0) + + if abs_idxs.numel() < self.num_selected_patches: + n_to_fill = self.num_selected_patches - abs_idxs.numel() + if abs_idxs.numel() > 0: + # Keep selection pool restricted to image/trace by padding from selected indices. + repeat_count = (n_to_fill + abs_idxs.numel() - 1) // abs_idxs.numel() + pad_abs_idxs = abs_idxs.repeat(repeat_count)[:n_to_fill] + abs_idxs = torch.cat([abs_idxs, pad_abs_idxs], dim=0) + else: + # Safety fallback: only sample from original image span, never generic context tokens. + image_span_scores = scores.clone() + allowed_image_span = torch.zeros_like(image_span_scores, dtype=torch.bool) + allowed_image_span[vs + 1 : ve] = True + image_span_scores[~allowed_image_span] = float("-inf") + abs_idxs = image_span_scores.topk(self.num_selected_patches, sorted=False)[1] + + logging.debug(f"selected image idx: {image_abs_idxs}") + logging.debug(f"selected trace idx: {trace_abs_idxs}") logging.debug(f"abs idx: {abs_idxs}") - image_mask[b, abs_idxs] = False + if need_insert_origin_masks: + selected_patch_origins.append(image_mask[b, abs_idxs].clone()) + selected_reasoning_origins.append(trace_mask[b, abs_idxs].clone()) + if self.patch_reuse_policy == "never": + image_mask[b, abs_idxs] = False + trace_mask[b, abs_idxs] = False + elif self.patch_reuse_policy == "next_step_only": + current_selected_mask[b, abs_idxs] = True picked = inputs_embeds[b, abs_idxs, :] # (K, D) select_image_embeds.append(picked) + selected_counts.append(abs_idxs.numel()) select_image_embeds = torch.stack(select_image_embeds, dim=0) # (B, K, D) inputs_embeds_detached = inputs_embeds.detach().clone() @@ -193,20 +505,27 @@ def forward( new_position_ids = [] new_original_mask = [] new_image_mask = [] + new_trace_mask = [] + new_recently_selected_mask = [] + new_patch_insert_mask = [] if need_insert_origin_masks else None + new_reasoning_insert_mask = [] if need_insert_origin_masks else None + current_inserted_spans = [] batch_max_len = 0 for b in range(B): + K_b = selected_counts[b] end_b = end prefix_b = inputs_embeds[b, :end_b, :] # (end_b, D) suffix_b = inputs_embeds[b, end_b:, :] # (old_len - end_b, D) v_embed_b = select_image_embeds[b] # (K, D) merged_b = torch.cat([prefix_b, v_embed_b, suffix_b], dim=0) # (old_len+K, D) new_inputs_embeds.append(merged_b) + current_inserted_spans.append((end_b, end_b + K_b)) # attention_mask att_pref = attention_mask[b, :end_b] # (end_b,) att_suf = attention_mask[b, end_b:] # (old_len-end_b,) - att_v = torch.ones(self.num_selected_patches, device=attention_mask.device, dtype=attention_mask.dtype) + att_v = torch.ones(K_b, device=attention_mask.device, dtype=attention_mask.dtype) merged_att = torch.cat([att_pref, att_v, att_suf], dim=0) # (new_len,) new_attention_mask.append(merged_att) @@ -217,17 +536,45 @@ def forward( # original_mask orig_pref = original_mask[b, :end_b] # (end_b,) orig_suf = original_mask[b, end_b:] # (old_len-end_b,) - orig_v = torch.zeros(self.num_selected_patches, device=input_ids.device, dtype=torch.bool) + orig_v = torch.zeros(K_b, device=input_ids.device, dtype=torch.bool) merged_orig = torch.cat([orig_pref, orig_v, orig_suf], dim=0) new_original_mask.append(merged_orig) # image_mask img_pref = image_mask[b, :end_b] img_suf = image_mask[b, end_b:] - img_v = torch.zeros(self.num_selected_patches, device=input_ids.device, dtype=torch.bool) + img_v = torch.zeros(K_b, device=input_ids.device, dtype=torch.bool) merged_img = torch.cat([img_pref, img_v, img_suf], dim=0) new_image_mask.append(merged_img) + # trace_mask + trace_pref = trace_mask[b, :end_b] + trace_suf = trace_mask[b, end_b:] + trace_v = torch.ones(K_b, device=input_ids.device, dtype=torch.bool) + merged_trace = torch.cat([trace_pref, trace_v, trace_suf], dim=0) + new_trace_mask.append(merged_trace) + + # recently_selected_mask (for next_step_only) + if self.patch_reuse_policy == "next_step_only": + recent_pref = current_selected_mask[b, :end_b] + recent_suf = current_selected_mask[b, end_b:] + recent_v = torch.zeros(K_b, device=input_ids.device, dtype=torch.bool) + merged_recent = torch.cat([recent_pref, recent_v, recent_suf], dim=0) + new_recently_selected_mask.append(merged_recent) + + if need_insert_origin_masks: + patch_pref = patch_insert_mask[b, :end_b] + patch_suf = patch_insert_mask[b, end_b:] + patch_v = selected_patch_origins[b].to(torch.bool) + merged_patch = torch.cat([patch_pref, patch_v, patch_suf], dim=0) + new_patch_insert_mask.append(merged_patch) + + reasoning_pref = reasoning_insert_mask[b, :end_b] + reasoning_suf = reasoning_insert_mask[b, end_b:] + reasoning_v = selected_reasoning_origins[b].to(torch.bool) + merged_reasoning = torch.cat([reasoning_pref, reasoning_v, reasoning_suf], dim=0) + new_reasoning_insert_mask.append(merged_reasoning) + batch_max_len = max(batch_max_len, merged_b.size(0)) padded_embeds = [] @@ -235,6 +582,10 @@ def forward( padded_pos = [] padded_orig = [] padded_img = [] + padded_trace = [] + padded_recent = [] + padded_patch = [] + padded_reasoning = [] for b in range(B): emb_b = new_inputs_embeds[b] @@ -242,30 +593,50 @@ def forward( pos_b = new_position_ids[b] orig_b = new_original_mask[b] img_b = new_image_mask[b] + trace_b = new_trace_mask[b] padded_embeds.append(emb_b.unsqueeze(0)) padded_att.append(att_b.unsqueeze(0)) padded_pos.append(pos_b.unsqueeze(0)) padded_orig.append(orig_b.unsqueeze(0)) padded_img.append(img_b.unsqueeze(0)) + padded_trace.append(trace_b.unsqueeze(0)) + if self.patch_reuse_policy == "next_step_only": + recent_b = new_recently_selected_mask[b] + padded_recent.append(recent_b.unsqueeze(0)) + if need_insert_origin_masks: + patch_b = new_patch_insert_mask[b] + reasoning_b = new_reasoning_insert_mask[b] + padded_patch.append(patch_b.unsqueeze(0)) + padded_reasoning.append(reasoning_b.unsqueeze(0)) inputs_embeds = torch.cat(padded_embeds, dim=0) attention_mask = torch.cat(padded_att, dim=0) position_ids = torch.cat(padded_pos, dim=0) original_mask = torch.cat(padded_orig, dim=0) image_mask = torch.cat(padded_img, dim=0) # (B, new_S) - K = self.num_selected_patches + trace_mask = torch.cat(padded_trace, dim=0) + if self.patch_reuse_policy == "next_step_only": + recently_selected_mask = torch.cat(padded_recent, dim=0) + if need_insert_origin_masks: + patch_insert_mask = torch.cat(padded_patch, dim=0) + reasoning_insert_mask = torch.cat(padded_reasoning, dim=0) + prev_inserted_spans = current_inserted_spans for b in range(B): + K_b = selected_counts[b] for i, pos in enumerate(latent_lists[b]): if pos > end: - latent_lists[b][i] = pos + K + latent_lists[b][i] = pos + K_b logging.debug(f"latent pos: {latent_lists[b][i]}") if pass_idx + 1 >= max_n_latents: end = inputs_embeds.size(1) else: - end = end + 1 + K + if B != 1 and self.patch_sampling_strategy == "all_image_patches": + raise ValueError("all_image_patches currently supports batch_size=1 only.") + end = end + 1 + selected_counts[0] + output_attentions = self.enable_nvt_loss or return_latent_attn or self.enable_qvr_loss if kv_cache: outputs = self.base_causallm( inputs_embeds=inputs_embeds[:, :end, :], @@ -274,7 +645,7 @@ def forward( pixel_values=pixel_values, image_grid_thw=image_grid_thw, output_hidden_states=True, - output_attentions=False, + output_attentions=output_attentions, ) else: outputs = self.base_causallm( @@ -284,10 +655,13 @@ def forward( pixel_values=pixel_values, image_grid_thw=image_grid_thw, output_hidden_states=True, - output_attentions=False, + output_attentions=output_attentions, ) all_logits.append(outputs.logits) - + if self.enable_nvt_loss and prev_inserted_spans is not None and outputs.attentions is not None: + nvt_loss = self._compute_nvt_loss(outputs.attentions, end - 1, prev_inserted_spans) + if nvt_loss is not None: + nvt_losses.append(nvt_loss) else: outputs = self.base_causallm( input_ids=input_ids, @@ -300,20 +674,175 @@ def forward( ) all_logits.append(outputs.logits) + if self.enable_qvr_loss and prev_inserted_spans is not None: + final_seq_len = inputs_embeds.size(1) + final_outputs = self.base_causallm( + inputs_embeds=inputs_embeds[:, :final_seq_len, :], + attention_mask=attention_mask[:, :final_seq_len], + position_ids=position_ids[:, :final_seq_len], + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + output_hidden_states=False, + output_attentions=True, + ) + if any(len(lst) > 0 for lst in latent_lists): + positions_final = torch.arange(final_seq_len, device=input_ids.device).unsqueeze(0).expand(B, -1) + first_latent_pos = min(lst[0] for lst in latent_lists if len(lst) > 0) + question_mask_final = ( + original_mask[:, :final_seq_len] + & (~image_mask[:, :final_seq_len]) + & (positions_final < first_latent_pos) + ) + final_query_positions = [lst[-1] if len(lst) > 0 else None for lst in latent_lists] + final_qvr_loss = self._compute_qvr_loss( + final_outputs.attentions, + final_query_positions, + prev_inserted_spans, + question_mask_final, + ) + if final_qvr_loss is not None: + qvr_losses.append(final_qvr_loss) + + latent_attn_trace = None + if return_latent_attn and outputs.attentions is not None and max_n_latents > 0: + final_attn = outputs.attentions[-1].mean(dim=1) + seq_len = final_attn.size(-1) + threshold = latent_attn_threshold + if threshold is None: + threshold = latent_attn_threshold_multiplier / max(seq_len, 1) + latent_attn_trace = [] + for b in range(B): + image_mask_b = image_mask[b, :seq_len] + original_mask_b = original_mask[b, :seq_len] + text_mask_b = original_mask_b & (~image_mask_b) + trace_mask_b = trace_mask[b, :seq_len] + + per_latent = [] + for t_idx in latent_lists[b]: + if t_idx >= seq_len: + continue + row = final_attn[b, t_idx] + image_mass = row[image_mask_b].sum().item() if image_mask_b.any() else 0.0 + text_mass = row[text_mask_b].sum().item() if text_mask_b.any() else 0.0 + trace_mass = row[trace_mask_b].sum().item() if trace_mask_b.any() else 0.0 + image_above = (image_mask_b & (row > threshold)).nonzero(as_tuple=True)[0].tolist() + trace_above = (trace_mask_b & (row > threshold)).nonzero(as_tuple=True)[0].tolist() + per_latent.append({ + "latent_pos": int(t_idx), + "image_mass": float(image_mass), + "text_mass": float(text_mass), + "trace_mass": float(trace_mass), + "threshold": float(threshold), + "image_above_threshold": image_above, + "trace_above_threshold": trace_above, + }) + latent_attn_trace.append(per_latent) + logits = torch.cat(all_logits, dim=-2) # (B, total_len, V) B, final_S, V = logits.size() + token_norms = None + if return_token_norms: + token_norms = [] + seq_len_for_norms = min(end, inputs_embeds.size(1)) + for b in range(B): + patch_mask_b = patch_insert_mask[b, :seq_len_for_norms] + reasoning_mask_b = reasoning_insert_mask[b, :seq_len_for_norms] + patch_embeds = inputs_embeds[b, :seq_len_for_norms, :][patch_mask_b] + reasoning_embeds = inputs_embeds[b, :seq_len_for_norms, :][reasoning_mask_b] + patch_norms = torch.norm(patch_embeds, dim=-1) if patch_embeds.numel() > 0 else torch.tensor([]) + reasoning_norms = ( + torch.norm(reasoning_embeds, dim=-1) if reasoning_embeds.numel() > 0 else torch.tensor([]) + ) + token_norms.append({ + "patch_norms": patch_norms.float().tolist(), + "reasoning_norms": reasoning_norms.float().tolist(), + "aggregate_stats": { + "patch": self._summarize_norms(patch_norms), + "reasoning": self._summarize_norms(reasoning_norms), + }, + }) + + qvr_loss = torch.stack(qvr_losses).mean() if (self.enable_qvr_loss and qvr_losses) else None + new_labels = torch.full((B, final_S), -100, device=input_ids.device, dtype=labels.dtype) for b in range(B): num_labels = labels.size(1) - new_labels[:, -num_labels:] = labels + new_labels[b, -num_labels:] = labels[b] + + new_answer_mask = None + if answer_mask is not None: + new_answer_mask = torch.zeros((B, final_S), device=answer_mask.device, dtype=answer_mask.dtype) + for b in range(B): + num_labels = labels.size(1) + new_answer_mask[b, -num_labels:] = answer_mask[b] shift_logits = logits[..., :-1, :].contiguous() shift_labels = new_labels[..., 1:].contiguous() loss_fct = CrossEntropyLoss(ignore_index=-100) - loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) - - return Outputs(loss=loss, inputs_embeds=inputs_embeds, logits=logits) + ce_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + nvt_loss = torch.stack(nvt_losses).mean() if (self.enable_nvt_loss and nvt_losses) else None + causal_loss = None + if ( + self.enable_causal_loss + and new_answer_mask is not None + and patch_insert_mask is not None + and B > 1 + ): + final_seq_len = inputs_embeds.size(1) + perm = torch.randperm(B, device=inputs_embeds.device) + if torch.any(perm == torch.arange(B, device=inputs_embeds.device)): + perm = torch.arange(B, device=inputs_embeds.device).roll(1) + + inputs_embeds_corrupt = inputs_embeds.clone() + for b in range(B): + tgt_pos = patch_insert_mask[b, :final_seq_len].nonzero(as_tuple=True)[0] + if tgt_pos.numel() == 0: + continue + donor_pos = patch_insert_mask[perm[b], :final_seq_len].nonzero(as_tuple=True)[0] + if donor_pos.numel() == 0: + continue + donor_embeds = inputs_embeds[perm[b], donor_pos, :] + if donor_embeds.size(0) < tgt_pos.numel(): + repeat_count = (tgt_pos.numel() + donor_embeds.size(0) - 1) // donor_embeds.size(0) + donor_embeds = donor_embeds.repeat(repeat_count, 1)[:tgt_pos.numel()] + elif donor_embeds.size(0) > tgt_pos.numel(): + donor_embeds = donor_embeds[:tgt_pos.numel()] + inputs_embeds_corrupt[b, tgt_pos, :] = donor_embeds + + corrupt_outputs = self.base_causallm( + inputs_embeds=inputs_embeds_corrupt[:, :final_seq_len, :], + attention_mask=attention_mask[:, :final_seq_len], + position_ids=position_ids[:, :final_seq_len], + output_hidden_states=False, + output_attentions=False, + ) + s_real, real_counts = self._compute_answer_logprob(logits, new_labels, new_answer_mask) + s_corrupt, corrupt_counts = self._compute_answer_logprob( + corrupt_outputs.logits, new_labels, new_answer_mask + ) + valid = (real_counts > 0) & (corrupt_counts > 0) + if valid.any(): + causal_loss = F.softplus(-(s_real - s_corrupt))[valid].mean() + loss = ce_loss + if nvt_loss is not None: + loss = loss + self.nvt_loss_weight * nvt_loss + if qvr_loss is not None: + loss = loss + self.qvr_loss_weight * qvr_loss + if causal_loss is not None: + loss = loss + self.causal_loss_weight * causal_loss + + return Outputs( + loss=loss, + ce_loss=ce_loss, + nvt_loss=nvt_loss, + qvr_loss=qvr_loss, + causal_loss=causal_loss, + inputs_embeds=inputs_embeds, + logits=logits, + latent_attn_trace=latent_attn_trace, + token_norms=token_norms, + ) def train(self, mode=True): diff --git a/qwen_vl/qwenvl_run.py b/qwen_vl/qwenvl_run.py index ca6853c..0790c47 100644 --- a/qwen_vl/qwenvl_run.py +++ b/qwen_vl/qwenvl_run.py @@ -69,6 +69,36 @@ def main(): parser.add_argument("--deepspeed", action="store_true", help="Enable DeepSpeed") parser.add_argument("--deepspeed_config", default="ds_config.json", help="DeepSpeed config path") parser.add_argument("--local_rank", type=int, default=-1, help="Local rank passed by DeepSpeed") + parser.add_argument("--patch_reuse_policy", choices=["never", "next_step_only", "always"], default=None, + help="Patch selection reuse policy across latent reasoning steps") + parser.add_argument("--resume_epoch", type=int, default=None, + help="Epoch index to resume from (0-based). Overrides config resume.") + parser.add_argument("--resume_model_path", type=str, default=None, + help="Path to a saved model state_dict (.pth) for resuming training.") + parser.add_argument("--num_proc", type=int, default=None, + help="Number of subprocesses for dataset.map. Overrides config num_proc.") + qvr_group = parser.add_mutually_exclusive_group() + qvr_group.add_argument("--qvr_loss", dest="qvr_loss", action="store_true", + help="Enable question-conditioned visual routing loss.") + qvr_group.add_argument("--no_qvr_loss", dest="qvr_loss", action="store_false", + help="Disable question-conditioned visual routing loss.") + parser.set_defaults(qvr_loss=None) + causal_group = parser.add_mutually_exclusive_group() + causal_group.add_argument("--causal_loss", dest="causal_loss", action="store_true", + help="Enable causal contrastive loss.") + causal_group.add_argument("--no_causal_loss", dest="causal_loss", action="store_false", + help="Disable causal contrastive loss.") + parser.set_defaults(causal_loss=None) + parser.add_argument("--lambda_qvr", type=float, default=None, + help="Weight for question-conditioned visual routing loss.") + parser.add_argument("--lambda_causal", type=float, default=None, + help="Weight for causal contrastive loss.") + parser.add_argument("--qvr_num_layers", type=int, default=None, + help="Number of last layers to average for QVR loss.") + parser.add_argument("--qvr_epsilon", type=float, default=None, + help="Epsilon for QVR loss.") + parser.add_argument("--causal_epsilon", type=float, default=None, + help="Epsilon for causal loss.") args = parser.parse_args() # Initialize DeepSpeed @@ -83,6 +113,31 @@ def main(): config_dict = yaml.safe_load(f) configs = Config(config_dict) + patch_reuse_policy = args.patch_reuse_policy or getattr(configs, "patch_reuse_policy", "never") + start_epoch = args.resume_epoch if args.resume_epoch is not None else int(getattr(configs, "resume", 0)) + resume_model_path = args.resume_model_path or getattr(configs, "load_model_path", None) + num_proc = args.num_proc if args.num_proc is not None else int(getattr(configs, "num_proc", 32)) + enable_qvr_loss = args.qvr_loss if args.qvr_loss is not None else bool( + getattr(configs, "enable_qvr_loss", False) + ) + enable_causal_loss = args.causal_loss if args.causal_loss is not None else bool( + getattr(configs, "enable_causal_loss", False) + ) + qvr_loss_weight = args.lambda_qvr if args.lambda_qvr is not None else float( + getattr(configs, "qvr_loss_weight", 0.01) + ) + causal_loss_weight = args.lambda_causal if args.lambda_causal is not None else float( + getattr(configs, "causal_loss_weight", 0.05) + ) + qvr_num_layers = args.qvr_num_layers if args.qvr_num_layers is not None else int( + getattr(configs, "qvr_num_layers", 4) + ) + qvr_loss_epsilon = args.qvr_epsilon if args.qvr_epsilon is not None else float( + getattr(configs, "qvr_loss_epsilon", 1e-8) + ) + causal_loss_epsilon = args.causal_epsilon if args.causal_epsilon is not None else float( + getattr(configs, "causal_loss_epsilon", 1e-8) + ) set_seed(configs.seed) save_dir = os.path.join(configs.save_path, configs.name) @@ -94,16 +149,21 @@ def main(): cur_ckpts = os.listdir(save_dir) - # check if the job is preempted and resumed. - if len(cur_ckpts) > 0 and rank == 0: + # Non-empty save dir is valid when resuming; block only for fresh runs. + if len(cur_ckpts) > 0 and rank == 0 and start_epoch == 0: raise ValueError( f"Save directory {save_dir} is not empty! " ) - if configs.resume != 0: - # by setting `resume`, we can skip a few epoches at the beginning. - print( - f"Loading from {configs.load_model_path} and skip the first {configs.resume} epochs" + if start_epoch > 0 and rank == 0: + print(f"Resume requested from epoch {start_epoch}") + print(f"Resume checkpoint path: {resume_model_path}") + print(f"Dataset map num_proc: {num_proc}") + elif rank == 0: + print(f"Dataset map num_proc: {num_proc}") + if start_epoch > 0 and not resume_model_path: + raise ValueError( + "Resuming requires a checkpoint path. Set --resume_model_path or load_model_path in the config." ) @@ -148,7 +208,39 @@ def main(): model.print_trainable_parameters() - model = IVTLR(model, latent_id, start_id, end_id, tokenizer.eos_token_id, image_token_id, visual_start_id, visual_end_id) + model = IVTLR( + model, + latent_id, + start_id, + end_id, + tokenizer.eos_token_id, + image_token_id, + visual_start_id, + visual_end_id, + patch_reuse_policy=patch_reuse_policy, + enable_qvr_loss=enable_qvr_loss, + qvr_loss_weight=qvr_loss_weight, + qvr_loss_epsilon=qvr_loss_epsilon, + qvr_num_layers=qvr_num_layers, + enable_causal_loss=enable_causal_loss, + causal_loss_weight=causal_loss_weight, + causal_loss_epsilon=causal_loss_epsilon, + ) + + if start_epoch > 0: + if not os.path.exists(resume_model_path): + raise ValueError(f"Checkpoint not found: {resume_model_path}") + if rank == 0: + print(f"Loading model weights from {resume_model_path}") + state_dict = torch.load(resume_model_path, map_location="cpu") + if any(k.startswith("module.") for k in state_dict.keys()): + state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()} + load_result = model.load_state_dict(state_dict, strict=False) + if rank == 0: + print( + f"Checkpoint loaded. Missing keys: {len(load_result.missing_keys)}, " + f"Unexpected keys: {len(load_result.unexpected_keys)}" + ) print(f"Running Deepspeed on rank = {rank}, world size = {world_size}") model = model.to(rank) @@ -237,12 +329,16 @@ def has_image(example): "image" in example and example["image"] is not None ) - train_dataset = dataset["train"].select(range(20)).filter(has_image) - train_dataset = train_dataset.map(process_example, num_proc=32) + train_dataset = dataset["train"].filter(has_image) + train_dataset = train_dataset.map(process_example, num_proc=num_proc) base_dataset_train = get_dataset( - train_dataset, tokenizer, processor, max_size=5000 if configs.debug else 100000000 + train_dataset, + tokenizer, + processor, + max_size=5000 if configs.debug else 100000000, + num_proc=num_proc, ) total_train_steps = 0 @@ -261,7 +357,7 @@ def has_image(example): collator = MyCollator(tokenizer, latent_id=latent_id, label_pad_token_id=-100) - for epoch in range(configs.resume, configs.num_epochs): + for epoch in range(start_epoch, configs.num_epochs): scheduled_stage = epoch // configs.epochs_per_stage @@ -336,6 +432,10 @@ def has_image(example): "train/loss": loss.detach().float() # * configs.gradient_accumulation_steps, } + if outputs.qvr_loss is not None: + log_dict["train/qvr_loss"] = outputs.qvr_loss.detach().float() + if outputs.causal_loss is not None: + log_dict["train/causal_loss"] = outputs.causal_loss.detach().float() wandb_run.log(log_dict) # print("line432") pbar.set_description( diff --git a/qwen_vl/qwenvl_run_2b.py b/qwen_vl/qwenvl_run_2b.py new file mode 100644 index 0000000..862c0fa --- /dev/null +++ b/qwen_vl/qwenvl_run_2b.py @@ -0,0 +1,497 @@ +import torch +import torch.distributed +import torch.optim as optim +from transformers import AutoModelForCausalLM, AutoTokenizer +from datetime import timedelta +import deepspeed +from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint +from torch.optim import AdamW +import shutil +import numpy as np +from torch.utils.data import Subset +from collections import OrderedDict +import re +import wandb + +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +import torch.distributed as dist +from torch.utils.data.distributed import DistributedSampler +from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy +from transformers.models.llama.modeling_llama import LlamaDecoderLayer +from transformers.models.gpt2.modeling_gpt2 import GPT2Block +from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor +from qwen_vl_utils import process_vision_info +from datasets import load_dataset +import logging +logging.basicConfig( + filename='qwenvl_2b.log', + level=logging.DEBUG, + format='[%(asctime)s] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' +) + +from qwen_ivtlr import IVTLR +from dataset import ( + get_dataset, + get_cot_latent_dataset, + MyCollator, +) + +from tqdm import tqdm +from copy import copy +import itertools +import os, sys +import yaml +import json +import gc +import argparse +import functools +from utils import Config, set_seed +import pdb +from peft import LoraConfig, get_peft_model + +# LoRA +lora_config = LoraConfig( + task_type="CAUSAL_LM", + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + r=64, + lora_alpha=16, + lora_dropout=0.05, + bias="none", + inference_mode=False +) + +def main(): + print("Initializing DeepSpeed Training!") + parser = argparse.ArgumentParser(description="ivtlr") + parser.add_argument("config_file") + parser.add_argument("--deepspeed", action="store_true", help="Enable DeepSpeed") + parser.add_argument("--deepspeed_config", default="ds_config.json", help="DeepSpeed config path") + parser.add_argument("--local_rank", type=int, default=-1, help="Local rank passed by DeepSpeed") + parser.add_argument("--patch_reuse_policy", choices=["never", "next_step_only", "always"], default=None, + help="Patch selection reuse policy across latent reasoning steps") + parser.add_argument("--resume_epoch", type=int, default=None, + help="Epoch index to resume from (0-based). Overrides config resume.") + parser.add_argument("--resume_model_path", type=str, default=None, + help="Path to a saved model state_dict (.pth) for resuming training.") + parser.add_argument("--num_proc", type=int, default=None, + help="Number of subprocesses for dataset.map. Overrides config num_proc.") + qvr_group = parser.add_mutually_exclusive_group() + qvr_group.add_argument("--qvr_loss", dest="qvr_loss", action="store_true", + help="Enable question-conditioned visual routing loss.") + qvr_group.add_argument("--no_qvr_loss", dest="qvr_loss", action="store_false", + help="Disable question-conditioned visual routing loss.") + parser.set_defaults(qvr_loss=None) + causal_group = parser.add_mutually_exclusive_group() + causal_group.add_argument("--causal_loss", dest="causal_loss", action="store_true", + help="Enable causal contrastive loss.") + causal_group.add_argument("--no_causal_loss", dest="causal_loss", action="store_false", + help="Disable causal contrastive loss.") + parser.set_defaults(causal_loss=None) + parser.add_argument("--lambda_qvr", type=float, default=None, + help="Weight for question-conditioned visual routing loss.") + parser.add_argument("--lambda_causal", type=float, default=None, + help="Weight for causal contrastive loss.") + parser.add_argument("--qvr_num_layers", type=int, default=None, + help="Number of last layers to average for QVR loss.") + parser.add_argument("--qvr_epsilon", type=float, default=None, + help="Epsilon for QVR loss.") + parser.add_argument("--causal_epsilon", type=float, default=None, + help="Epsilon for causal loss.") + args = parser.parse_args() + + # Initialize DeepSpeed + deepspeed.init_distributed() + local_rank = args.local_rank + rank = int(os.environ['RANK']) + world_size = int(os.environ['WORLD_SIZE']) + torch.cuda.set_device(local_rank) + print("line 57") + # load the configuration file + with open(args.config_file) as f: + config_dict = yaml.safe_load(f) + + configs = Config(config_dict) + patch_reuse_policy = args.patch_reuse_policy or getattr(configs, "patch_reuse_policy", "never") + start_epoch = args.resume_epoch if args.resume_epoch is not None else int(getattr(configs, "resume", 0)) + resume_model_path = args.resume_model_path or getattr(configs, "load_model_path", None) + num_proc = args.num_proc if args.num_proc is not None else int(getattr(configs, "num_proc", 32)) + enable_qvr_loss = args.qvr_loss if args.qvr_loss is not None else bool( + getattr(configs, "enable_qvr_loss", False) + ) + enable_causal_loss = args.causal_loss if args.causal_loss is not None else bool( + getattr(configs, "enable_causal_loss", False) + ) + qvr_loss_weight = args.lambda_qvr if args.lambda_qvr is not None else float( + getattr(configs, "qvr_loss_weight", 0.01) + ) + causal_loss_weight = args.lambda_causal if args.lambda_causal is not None else float( + getattr(configs, "causal_loss_weight", 0.05) + ) + qvr_num_layers = args.qvr_num_layers if args.qvr_num_layers is not None else int( + getattr(configs, "qvr_num_layers", 4) + ) + qvr_loss_epsilon = args.qvr_epsilon if args.qvr_epsilon is not None else float( + getattr(configs, "qvr_loss_epsilon", 1e-8) + ) + causal_loss_epsilon = args.causal_epsilon if args.causal_epsilon is not None else float( + getattr(configs, "causal_loss_epsilon", 1e-8) + ) + set_seed(configs.seed) + save_dir = os.path.join(configs.save_path, configs.name) + + if not os.path.exists(save_dir) and rank == 0: + os.makedirs(save_dir) + + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + + cur_ckpts = os.listdir(save_dir) + + + # Non-empty save dir is valid when resuming; block only for fresh runs. + if len(cur_ckpts) > 0 and rank == 0 and start_epoch == 0: + raise ValueError( + f"Save directory {save_dir} is not empty! " + ) + + if start_epoch > 0 and rank == 0: + print(f"Resume requested from epoch {start_epoch}") + print(f"Resume checkpoint path: {resume_model_path}") + print(f"Dataset map num_proc: {num_proc}") + elif rank == 0: + print(f"Dataset map num_proc: {num_proc}") + if start_epoch > 0 and not resume_model_path: + raise ValueError( + "Resuming requires a checkpoint path. Set --resume_model_path or load_model_path in the config." + ) + + + + print("start loading model") + + model = Qwen2VLForConditionalGeneration.from_pretrained( + "Qwen/Qwen2-VL-2B-Instruct", device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation="eager" + ) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=configs.lr) + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", use_fast=False, trust_remote_code=True) + tokenizer.padding_side = "right" + tokenizer.pad_token = tokenizer.eos_token + tokenizer.add_tokens("<|start-latent|>") + tokenizer.add_tokens("<|end-latent|>") + tokenizer.add_tokens("<|latent|>") + processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", tokenizer=tokenizer) + latent_id = tokenizer.convert_tokens_to_ids("<|latent|>") + print("latent_id: ", latent_id) + start_id = tokenizer.convert_tokens_to_ids("<|start-latent|>") + end_id = tokenizer.convert_tokens_to_ids("<|end-latent|>") + image_token_id = tokenizer.convert_tokens_to_ids(processor.image_token) + visual_start_id = tokenizer.convert_tokens_to_ids("<|vision_start|>") + visual_end_id = tokenizer.convert_tokens_to_ids("<|vision_end|>") + + model = get_peft_model(model, lora_config) + + loaded = False + + model.resize_token_embeddings(len(tokenizer)) + embeddings = model.get_input_embeddings() + target_id = tokenizer.convert_tokens_to_ids("<<") + # initialize the new token embeddings with a known token + # it helps stablize the training + for token_id in [latent_id, start_id, end_id]: + target_embedding = embeddings.weight.data[token_id] + embeddings.weight.data[token_id] = target_embedding + + lm_head = model.lm_head + lm_head.weight.data[token_id] = lm_head.weight.data[target_id] + + model.print_trainable_parameters() + + model = IVTLR( + model, + latent_id, + start_id, + end_id, + tokenizer.eos_token_id, + image_token_id, + visual_start_id, + visual_end_id, + patch_reuse_policy=patch_reuse_policy, + processor_model_id="Qwen/Qwen2-VL-2B-Instruct", + enable_nvt_loss=getattr(configs, "enable_nvt_loss", False), + nvt_loss_weight=getattr(configs, "nvt_loss_weight", 0.0), + nvt_loss_epsilon=getattr(configs, "nvt_loss_epsilon", 1e-8), + enable_qvr_loss=enable_qvr_loss, + qvr_loss_weight=qvr_loss_weight, + qvr_loss_epsilon=qvr_loss_epsilon, + qvr_num_layers=qvr_num_layers, + enable_causal_loss=enable_causal_loss, + causal_loss_weight=causal_loss_weight, + causal_loss_epsilon=causal_loss_epsilon, + ) + + if start_epoch > 0: + if not os.path.exists(resume_model_path): + raise ValueError(f"Checkpoint not found: {resume_model_path}") + if rank == 0: + print(f"Loading model weights from {resume_model_path}") + state_dict = torch.load(resume_model_path, map_location="cpu") + if any(k.startswith("module.") for k in state_dict.keys()): + state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()} + load_result = model.load_state_dict(state_dict, strict=False) + if rank == 0: + print( + f"Checkpoint loaded. Missing keys: {len(load_result.missing_keys)}, " + f"Unexpected keys: {len(load_result.unexpected_keys)}" + ) + + print(f"Running Deepspeed on rank = {rank}, world size = {world_size}") + model = model.to(rank) + + if configs.bf16: + model.to(torch.bfloat16) + + model_engine, optimizer, _, _ = deepspeed.initialize( + model=model, + config=args.deepspeed_config, + # optimizer = optimizer, + model_parameters=filter(lambda p: p.requires_grad, model.parameters()) + ) + + del model + + dataset = load_dataset("LightChen2333/M3CoT") + + def process_example(example): + rationale = example["rationale"].replace("\n", " ").strip() + example["steps"] = rationale.split(". ") + if example["steps"][-1] == "": + example["steps"].pop() + + if len(example["steps"]) > 3: + total_steps = len(example["steps"]) + step_size = total_steps // 3 + remainder = total_steps % 3 + + new_steps = [] + start = 0 + + for i in range(3): + end = start + step_size + (1 if i < remainder else 0) + new_steps.append(". ".join(example["steps"][start:end])) + start = end + + example["steps"] = new_steps + + + question = example["question"] + choices = example["choices"] + + + choices_str = "[Options]:\n"+"\n".join([ + f"({chr(65 + i)}).{{{choice.strip()}}}" + for i, choice in enumerate(choices) + ]) + question = question + question_with_braces = f"{{{question.strip()}}}" + prefix_str = "Answer:" + + example["question"] = f"[Question]:{question_with_braces}\n{choices_str}\n{prefix_str}\n" + + del example["rationale"] + del example["choices"] + + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": example["image"], "resized_height": 280, "resized_width": 280}, + {"type": "text", "text": example["question"]} + ] + }] + + example["question"] = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + image_inputs, video_inputs = process_vision_info(messages) + inputs = processor( + text=[example["question"]], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt" + ) + inputs = {k: v.tolist() for k, v in inputs.items()} + example["input_ids"] = torch.tensor(inputs["input_ids"][0]) + example["image_grid_thw"] = torch.tensor(inputs["image_grid_thw"]).squeeze(0) + example["pixel_values"] = torch.tensor(inputs["pixel_values"]) + + return example + + print("start dataset") + + def has_image(example): + return ( + "image" in example and example["image"] is not None + ) + + train_dataset = dataset["train"].filter(has_image) + train_dataset = train_dataset.map(process_example, num_proc=num_proc) + + + base_dataset_train = get_dataset( + train_dataset, + tokenizer, + processor, + max_size=5000 if configs.debug else 100000000, + num_proc=num_proc, + ) + + total_train_steps = 0 + + if not configs.debug and rank == 0: + wandb_run = wandb.init(project=configs.project, name=configs.name) + wandb_run.config.update(configs, allow_val_change=True) + text_table = wandb.Table(columns=["step", "text"]) + + else: + wandb_run = None + + + best_acc = 0 + + collator = MyCollator(tokenizer, latent_id=latent_id, label_pad_token_id=-100) + + + for epoch in range(start_epoch, configs.num_epochs): + + scheduled_stage = epoch // configs.epochs_per_stage + + np.random.seed(epoch) + + dataset_train = get_cot_latent_dataset( + scheduled_stage, + base_dataset_train, + configs, + start_id, + latent_id, + end_id, + no_special_marker=True, + shuffle=True, + ) + + train_dataloader = torch.utils.data.DataLoader( + dataset_train, + num_workers=1, + shuffle=False, + pin_memory=True, + batch_size=configs.batch_size_training, + collate_fn=collator, + sampler=DistributedSampler(dataset_train, shuffle=True), + ) + + model_engine.train() + total_length = len(train_dataloader) // configs.gradient_accumulation_steps + pbar = tqdm( + colour="blue", + desc=f"Training Epoch: {epoch+1}", + total=total_length, + dynamic_ncols=True, + ) + for step, batch in enumerate(train_dataloader): + print("start") + if step == 0 and wandb_run and rank == 0: + print("logging training data") + cur_bs = len(batch["input_ids"]) + text_str = "" + for data_idx in range(cur_bs): + for token_idx in range(len(batch["input_ids"][data_idx])): + text_str += ( + str(batch["input_ids"][data_idx][token_idx].item()) + + " " + + str(batch["labels"][data_idx][token_idx].item()) + + " " + + tokenizer.decode( + batch["input_ids"][data_idx][token_idx] + ) + + "\n" + ) + text_str += "====" * 10 + "\n" + + text_table.add_data(total_train_steps, text_str) + + total_train_steps += 1 + batch = { + key: batch[key].to(rank) for key in batch.keys() if key != "idx" + } + + outputs = model_engine(**batch) + loss = outputs.loss + print(f"loss: {loss}") + if rank == 0 and (step + 1) % 50 == 0: + ce_loss = outputs.ce_loss.detach().float() + nvt_loss = outputs.nvt_loss.detach().float() if outputs.nvt_loss is not None else torch.tensor(0.0) + qvr_loss = outputs.qvr_loss.detach().float() if outputs.qvr_loss is not None else torch.tensor(0.0) + causal_loss = ( + outputs.causal_loss.detach().float() + if outputs.causal_loss is not None + else torch.tensor(0.0) + ) + total_loss = loss.detach().float() + print( + f"[step {step + 1}] ce_loss={float(ce_loss):.4f} " + f"nvt_loss={float(nvt_loss):.4f} qvr_loss={float(qvr_loss):.4f} " + f"causal_loss={float(causal_loss):.4f} total_loss={float(total_loss):.4f}" + ) + model_engine.backward(loss) + model_engine.step() + + if wandb_run and rank == 0: + log_dict = { + "train/epoch": epoch + 1, + "train/step": epoch * len(train_dataloader) + step, + "train/loss": loss.detach().float() + # * configs.gradient_accumulation_steps, + } + if outputs.qvr_loss is not None: + log_dict["train/qvr_loss"] = outputs.qvr_loss.detach().float() + if outputs.causal_loss is not None: + log_dict["train/causal_loss"] = outputs.causal_loss.detach().float() + wandb_run.log(log_dict) + # print("line432") + pbar.set_description( + f"Training Epoch: {epoch+1}/{configs.num_epochs}, batch {step}/{len(train_dataloader)} " + f"completed (loss: {round(float(loss.detach().float()), 4)}" + ) + print("finish") + pbar.close() + dist.barrier() + + if ( + not configs.debug + and (epoch + 1) % 4 == 0 + ): + + epoch_save_dir = os.path.join(save_dir, f"epoch_{epoch+1}_checkpoint") + + model_engine.save_checkpoint( + save_dir=epoch_save_dir, + tag=f"epoch_{epoch+1}_zero3_bf32", + client_state={"best_acc": best_acc, "current_epoch": epoch+1} + ) + + if rank == 0: + fp32_state_dict = get_fp32_state_dict_from_zero_checkpoint(epoch_save_dir, tag=f"epoch_{epoch+1}_zero3_bf32") + fp32_output = os.path.join(save_dir, f"epoch_{epoch+1}_full_model_fp32.pth") + + torch.save(fp32_state_dict, fp32_output) + + print(f"Epoch {epoch+1} FP32 save to {fp32_output}") + + if os.path.exists(epoch_save_dir): + shutil.rmtree(epoch_save_dir) + + dist.barrier() + gc.collect() + torch.cuda.empty_cache() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/qwen_vl/qwenvl_run_2b_sqa.py b/qwen_vl/qwenvl_run_2b_sqa.py new file mode 100644 index 0000000..e4d41c4 --- /dev/null +++ b/qwen_vl/qwenvl_run_2b_sqa.py @@ -0,0 +1,492 @@ +import torch +import torch.distributed +import torch.optim as optim +from transformers import AutoModelForCausalLM, AutoTokenizer +from datetime import timedelta +import deepspeed +from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint +from torch.optim import AdamW +import shutil +import numpy as np +from torch.utils.data import Subset +from collections import OrderedDict +import re +import wandb + +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +import torch.distributed as dist +from torch.utils.data.distributed import DistributedSampler +from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy +from transformers.models.llama.modeling_llama import LlamaDecoderLayer +from transformers.models.gpt2.modeling_gpt2 import GPT2Block +from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor +from qwen_vl_utils import process_vision_info +from datasets import load_dataset +import logging +logging.basicConfig( + filename='qwenvl_2b_sqa.log', + level=logging.DEBUG, + format='[%(asctime)s] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' +) + +from qwen_ivtlr import IVTLR +from dataset import ( + get_dataset, + get_cot_latent_dataset, + MyCollator, +) + +from tqdm import tqdm +from copy import copy +import itertools +import os, sys +import yaml +import json +import gc +import argparse +import functools +from utils import Config, set_seed +import pdb +from peft import LoraConfig, get_peft_model + +# LoRA 配置 +lora_config = LoraConfig( + task_type="CAUSAL_LM", + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + r=64, + lora_alpha=16, + lora_dropout=0.05, + bias="none", + inference_mode=False +) + +def main(): + print("Initializing DeepSpeed Training!") + parser = argparse.ArgumentParser(description="coconut") + parser.add_argument("config_file") + parser.add_argument("--deepspeed", action="store_true", help="Enable DeepSpeed") + parser.add_argument("--deepspeed_config", default="ds_config.json", help="DeepSpeed config path") + parser.add_argument("--local_rank", type=int, default=-1, help="Local rank passed by DeepSpeed") + parser.add_argument("--patch_reuse_policy", choices=["never", "next_step_only", "always"], default=None, + help="Patch selection reuse policy across latent reasoning steps") + qvr_group = parser.add_mutually_exclusive_group() + qvr_group.add_argument("--qvr_loss", dest="qvr_loss", action="store_true", + help="Enable question-conditioned visual routing loss.") + qvr_group.add_argument("--no_qvr_loss", dest="qvr_loss", action="store_false", + help="Disable question-conditioned visual routing loss.") + parser.set_defaults(qvr_loss=None) + causal_group = parser.add_mutually_exclusive_group() + causal_group.add_argument("--causal_loss", dest="causal_loss", action="store_true", + help="Enable causal contrastive loss.") + causal_group.add_argument("--no_causal_loss", dest="causal_loss", action="store_false", + help="Disable causal contrastive loss.") + parser.set_defaults(causal_loss=None) + parser.add_argument("--lambda_qvr", type=float, default=None, + help="Weight for question-conditioned visual routing loss.") + parser.add_argument("--lambda_causal", type=float, default=None, + help="Weight for causal contrastive loss.") + parser.add_argument("--qvr_num_layers", type=int, default=None, + help="Number of last layers to average for QVR loss.") + parser.add_argument("--qvr_epsilon", type=float, default=None, + help="Epsilon for QVR loss.") + parser.add_argument("--causal_epsilon", type=float, default=None, + help="Epsilon for causal loss.") + args = parser.parse_args() + + # Initialize DeepSpeed + deepspeed.init_distributed() + local_rank = args.local_rank + rank = int(os.environ['RANK']) + world_size = int(os.environ['WORLD_SIZE']) + torch.cuda.set_device(local_rank) + print("line 57") + # load the configuration file + with open(args.config_file) as f: + config_dict = yaml.safe_load(f) + + configs = Config(config_dict) + patch_reuse_policy = args.patch_reuse_policy or getattr(configs, "patch_reuse_policy", "never") + enable_qvr_loss = args.qvr_loss if args.qvr_loss is not None else bool( + getattr(configs, "enable_qvr_loss", False) + ) + enable_causal_loss = args.causal_loss if args.causal_loss is not None else bool( + getattr(configs, "enable_causal_loss", False) + ) + qvr_loss_weight = args.lambda_qvr if args.lambda_qvr is not None else float( + getattr(configs, "qvr_loss_weight", 0.01) + ) + causal_loss_weight = args.lambda_causal if args.lambda_causal is not None else float( + getattr(configs, "causal_loss_weight", 0.05) + ) + qvr_num_layers = args.qvr_num_layers if args.qvr_num_layers is not None else int( + getattr(configs, "qvr_num_layers", 4) + ) + qvr_loss_epsilon = args.qvr_epsilon if args.qvr_epsilon is not None else float( + getattr(configs, "qvr_loss_epsilon", 1e-8) + ) + causal_loss_epsilon = args.causal_epsilon if args.causal_epsilon is not None else float( + getattr(configs, "causal_loss_epsilon", 1e-8) + ) + set_seed(configs.seed) + save_dir = os.path.join(configs.save_path, configs.name) + + if not os.path.exists(save_dir) and rank == 0: + os.makedirs(save_dir) + + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + + cur_ckpts = os.listdir(save_dir) + + + # check if the job is preempted and resumed. + if len(cur_ckpts) > 0 and rank == 0: + raise ValueError( + f"Save directory {save_dir} is not empty! " + ) + + if configs.resume != 0: + # by setting `resume`, we can skip a few epoches at the beginning. + print( + f"Loading from {configs.load_model_path} and skip the first {configs.resume} epochs" + ) + + + + print("start loading model") + # Todo:modify model and Tokenizer + model = Qwen2VLForConditionalGeneration.from_pretrained( + "Qwen/Qwen2-VL-2B-Instruct", device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation="eager" + ) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=configs.lr) + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", use_fast=False, trust_remote_code=True) + tokenizer.padding_side = "right" + tokenizer.pad_token = tokenizer.eos_token + tokenizer.add_tokens("<|start-latent|>") + tokenizer.add_tokens("<|end-latent|>") + tokenizer.add_tokens("<|latent|>") + processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", tokenizer=tokenizer) + latent_id = tokenizer.convert_tokens_to_ids("<|latent|>") + print("latent_id: ", latent_id) + start_id = tokenizer.convert_tokens_to_ids("<|start-latent|>") + end_id = tokenizer.convert_tokens_to_ids("<|end-latent|>") + image_token_id = tokenizer.convert_tokens_to_ids(processor.image_token) + visual_start_id = tokenizer.convert_tokens_to_ids("<|vision_start|>") + visual_end_id = tokenizer.convert_tokens_to_ids("<|vision_end|>") + print("line159") + model = get_peft_model(model, lora_config) + + loaded = False + + model.resize_token_embeddings(len(tokenizer)) + embeddings = model.get_input_embeddings() + target_id = tokenizer.convert_tokens_to_ids("<<") + # initialize the new token embeddings with a known token + # it helps stablize the training + for token_id in [latent_id, start_id, end_id]: + target_embedding = embeddings.weight.data[token_id] + embeddings.weight.data[token_id] = target_embedding + + lm_head = model.lm_head + lm_head.weight.data[token_id] = lm_head.weight.data[target_id] + + model.print_trainable_parameters() + + model = IVTLR( + model, + latent_id, + start_id, + end_id, + tokenizer.eos_token_id, + image_token_id, + visual_start_id, + visual_end_id, + patch_reuse_policy=patch_reuse_policy, + processor_model_id="Qwen/Qwen2-VL-2B-Instruct", + enable_nvt_loss=getattr(configs, "enable_nvt_loss", False), + nvt_loss_weight=getattr(configs, "nvt_loss_weight", 0.0), + nvt_loss_epsilon=getattr(configs, "nvt_loss_epsilon", 1e-8), + enable_qvr_loss=enable_qvr_loss, + qvr_loss_weight=qvr_loss_weight, + qvr_loss_epsilon=qvr_loss_epsilon, + qvr_num_layers=qvr_num_layers, + enable_causal_loss=enable_causal_loss, + causal_loss_weight=causal_loss_weight, + causal_loss_epsilon=causal_loss_epsilon, + ) + + print(f"Running Deepspeed on rank = {rank}, world size = {world_size}") + model = model.to(rank) + + if configs.bf16: + model.to(torch.bfloat16) + + model_engine, optimizer, _, _ = deepspeed.initialize( + model=model, + config=args.deepspeed_config, + # optimizer = optimizer, + model_parameters=filter(lambda p: p.requires_grad, model.parameters()) + ) + + del model + + dataset = load_dataset("derek-thomas/ScienceQA") + + def process_example(example): + example["answer"] = str(example["answer"]) + lecture = example.get("lecture", "") or "" + solution = example.get("solution", "") or "" + + if lecture and solution: + rationale = (lecture.strip() + " " + solution.strip()).strip() + elif lecture: + rationale = lecture.strip() + elif solution: + rationale = solution.strip() + else: + rationale = example["answer"] + print(f"Warning: Both lecture and solution are empty for question: {example['question']}") + + rationale = rationale.replace("\n", " ").strip() + example["steps"] = rationale.split(". ") + if example["steps"][-1] == "": + example["steps"].pop() + + if len(example["steps"]) > 3: + total_steps = len(example["steps"]) + step_size = total_steps // 3 + remainder = total_steps % 3 + + new_steps = [] + start = 0 + + for i in range(3): + end = start + step_size + (1 if i < remainder else 0) + new_steps.append(". ".join(example["steps"][start:end])) + start = end + + example["steps"] = new_steps + + question = example["question"] + + + if "choices" in example and example["choices"]: + choices = example["choices"] + + choices_str = "[Options]:\n" + "\n".join([ + f"({chr(65 + i)}).{{{choice.strip()}}}" + for i, choice in enumerate(choices) + ]) + else: + choices_str = "" + + question_with_braces = f"{{{question.strip()}}}" + prefix_str = "Answer:" + + if choices_str: + example["question"] = f"[Question]:{question_with_braces}\n{choices_str}\n{prefix_str}\n" + else: + example["question"] = f"[Question]:{question_with_braces}\n{prefix_str}\n" + + if "lecture" in example: + del example["lecture"] + if "solution" in example: + del example["solution"] + if "choices" in example: + del example["choices"] + + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": example["image"], "resized_height": 280, "resized_width": 280}, + {"type": "text", "text": example["question"]} + ] + }] + + example["question"] = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + image_inputs, video_inputs = process_vision_info(messages) + inputs = processor( + text=[example["question"]], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt" + ) + inputs = {k: v.tolist() for k, v in inputs.items()} + example["input_ids"] = torch.tensor(inputs["input_ids"][0]) + example["image_grid_thw"] = torch.tensor(inputs["image_grid_thw"]).squeeze(0) + example["pixel_values"] = torch.tensor(inputs["pixel_values"]) + + return example + + print("start dataset") + + def has_image(example): + return ( + "image" in example and example["image"] is not None + ) + + train_dataset = dataset["train"].filter(has_image) + train_dataset = train_dataset.map(process_example, num_proc=32) + + base_dataset_train = get_dataset( + train_dataset, tokenizer, processor, max_size=5000 if configs.debug else 100000000 + ) + + + total_train_steps = 0 + + if not configs.debug and rank == 0: + wandb_run = wandb.init(project=configs.project, name=configs.name) + wandb_run.config.update(configs, allow_val_change=True) + text_table = wandb.Table(columns=["step", "text"]) + + else: + wandb_run = None + + + best_acc = 0 + + collator = MyCollator(tokenizer, latent_id=latent_id, label_pad_token_id=-100) + + + for epoch in range(configs.resume, configs.num_epochs): + + scheduled_stage = epoch // configs.epochs_per_stage + + np.random.seed(epoch) + + dataset_train = get_cot_latent_dataset( + scheduled_stage, + base_dataset_train, + configs, + start_id, + latent_id, + end_id, + no_special_marker=configs.cot or configs.no_cot or configs.no_thoughts, + shuffle=True, + ) + + train_dataloader = torch.utils.data.DataLoader( + dataset_train, + num_workers=1, + shuffle=False, + pin_memory=True, + batch_size=configs.batch_size_training, + collate_fn=collator, + sampler=DistributedSampler(dataset_train, shuffle=True), + ) + + model_engine.train() + total_length = len(train_dataloader) // configs.gradient_accumulation_steps + pbar = tqdm( + colour="blue", + desc=f"Training Epoch: {epoch+1}", + total=total_length, + dynamic_ncols=True, + ) + for step, batch in enumerate(train_dataloader): + print("start") + # pdb.set_trace() + if step == 0 and wandb_run and rank == 0: + print("logging training data") + cur_bs = len(batch["input_ids"]) + text_str = "" + for data_idx in range(cur_bs): + for token_idx in range(len(batch["input_ids"][data_idx])): + text_str += ( + str(batch["input_ids"][data_idx][token_idx].item()) + + " " + + str(batch["labels"][data_idx][token_idx].item()) + + " " + + tokenizer.decode( + batch["input_ids"][data_idx][token_idx] + ) + + "\n" + ) + text_str += "====" * 10 + "\n" + + text_table.add_data(total_train_steps, text_str) + + total_train_steps += 1 + batch = { + key: batch[key].to(rank) for key in batch.keys() if key != "idx" + } + + outputs = model_engine(**batch) + loss = outputs.loss + print(f"loss: {loss}") + if rank == 0 and (step + 1) % 50 == 0: + ce_loss = outputs.ce_loss.detach().float() + nvt_loss = outputs.nvt_loss.detach().float() if outputs.nvt_loss is not None else torch.tensor(0.0) + qvr_loss = outputs.qvr_loss.detach().float() if outputs.qvr_loss is not None else torch.tensor(0.0) + causal_loss = ( + outputs.causal_loss.detach().float() + if outputs.causal_loss is not None + else torch.tensor(0.0) + ) + total_loss = loss.detach().float() + print( + f"[step {step + 1}] ce_loss={float(ce_loss):.4f} " + f"nvt_loss={float(nvt_loss):.4f} qvr_loss={float(qvr_loss):.4f} " + f"causal_loss={float(causal_loss):.4f} total_loss={float(total_loss):.4f}" + ) + model_engine.backward(loss) + model_engine.step() + + if wandb_run and rank == 0: + log_dict = { + "train/epoch": epoch + 1, + "train/step": epoch * len(train_dataloader) + step, + "train/loss": loss.detach().float() + # * configs.gradient_accumulation_steps, + } + if outputs.qvr_loss is not None: + log_dict["train/qvr_loss"] = outputs.qvr_loss.detach().float() + if outputs.causal_loss is not None: + log_dict["train/causal_loss"] = outputs.causal_loss.detach().float() + wandb_run.log(log_dict) + # print("line432") + pbar.set_description( + f"Training Epoch: {epoch+1}/{configs.num_epochs}, batch {step}/{len(train_dataloader)} " + f"completed (loss: {round(float(loss.detach().float()), 4)}" + ) + print("finish") + pbar.close() + dist.barrier() + + + # save start + if ( + not configs.debug + and (epoch + 1) % 4 == 0 + ): + + epoch_save_dir = os.path.join(save_dir, f"epoch_{epoch+1}_checkpoint") + + model_engine.save_checkpoint( + save_dir=epoch_save_dir, + tag=f"epoch_{epoch+1}_zero3_bf32", + client_state={"best_acc": best_acc, "current_epoch": epoch+1} + ) + + if rank == 0: + fp32_state_dict = get_fp32_state_dict_from_zero_checkpoint(epoch_save_dir, tag=f"epoch_{epoch+1}_zero3_bf32") + fp32_output = os.path.join(save_dir, f"epoch_{epoch+1}_full_model_fp32.pth") + + torch.save(fp32_state_dict, fp32_output) + + print(f"Epoch {epoch+1} save to {fp32_output}") + + if os.path.exists(epoch_save_dir): + shutil.rmtree(epoch_save_dir) + + dist.barrier() + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/qwen_vl/qwenvl_run_sqa.py b/qwen_vl/qwenvl_run_sqa.py index 018dba6..2bca466 100644 --- a/qwen_vl/qwenvl_run_sqa.py +++ b/qwen_vl/qwenvl_run_sqa.py @@ -69,6 +69,30 @@ def main(): parser.add_argument("--deepspeed", action="store_true", help="Enable DeepSpeed") parser.add_argument("--deepspeed_config", default="ds_config.json", help="DeepSpeed config path") parser.add_argument("--local_rank", type=int, default=-1, help="Local rank passed by DeepSpeed") + parser.add_argument("--patch_reuse_policy", choices=["never", "next_step_only", "always"], default=None, + help="Patch selection reuse policy across latent reasoning steps") + qvr_group = parser.add_mutually_exclusive_group() + qvr_group.add_argument("--qvr_loss", dest="qvr_loss", action="store_true", + help="Enable question-conditioned visual routing loss.") + qvr_group.add_argument("--no_qvr_loss", dest="qvr_loss", action="store_false", + help="Disable question-conditioned visual routing loss.") + parser.set_defaults(qvr_loss=None) + causal_group = parser.add_mutually_exclusive_group() + causal_group.add_argument("--causal_loss", dest="causal_loss", action="store_true", + help="Enable causal contrastive loss.") + causal_group.add_argument("--no_causal_loss", dest="causal_loss", action="store_false", + help="Disable causal contrastive loss.") + parser.set_defaults(causal_loss=None) + parser.add_argument("--lambda_qvr", type=float, default=None, + help="Weight for question-conditioned visual routing loss.") + parser.add_argument("--lambda_causal", type=float, default=None, + help="Weight for causal contrastive loss.") + parser.add_argument("--qvr_num_layers", type=int, default=None, + help="Number of last layers to average for QVR loss.") + parser.add_argument("--qvr_epsilon", type=float, default=None, + help="Epsilon for QVR loss.") + parser.add_argument("--causal_epsilon", type=float, default=None, + help="Epsilon for causal loss.") args = parser.parse_args() # Initialize DeepSpeed @@ -83,6 +107,28 @@ def main(): config_dict = yaml.safe_load(f) configs = Config(config_dict) + patch_reuse_policy = args.patch_reuse_policy or getattr(configs, "patch_reuse_policy", "never") + enable_qvr_loss = args.qvr_loss if args.qvr_loss is not None else bool( + getattr(configs, "enable_qvr_loss", False) + ) + enable_causal_loss = args.causal_loss if args.causal_loss is not None else bool( + getattr(configs, "enable_causal_loss", False) + ) + qvr_loss_weight = args.lambda_qvr if args.lambda_qvr is not None else float( + getattr(configs, "qvr_loss_weight", 0.01) + ) + causal_loss_weight = args.lambda_causal if args.lambda_causal is not None else float( + getattr(configs, "causal_loss_weight", 0.05) + ) + qvr_num_layers = args.qvr_num_layers if args.qvr_num_layers is not None else int( + getattr(configs, "qvr_num_layers", 4) + ) + qvr_loss_epsilon = args.qvr_epsilon if args.qvr_epsilon is not None else float( + getattr(configs, "qvr_loss_epsilon", 1e-8) + ) + causal_loss_epsilon = args.causal_epsilon if args.causal_epsilon is not None else float( + getattr(configs, "causal_loss_epsilon", 1e-8) + ) set_seed(configs.seed) save_dir = os.path.join(configs.save_path, configs.name) @@ -148,7 +194,24 @@ def main(): model.print_trainable_parameters() - model = IVTLR(model, latent_id, start_id, end_id, tokenizer.eos_token_id, image_token_id, visual_start_id, visual_end_id) + model = IVTLR( + model, + latent_id, + start_id, + end_id, + tokenizer.eos_token_id, + image_token_id, + visual_start_id, + visual_end_id, + patch_reuse_policy=patch_reuse_policy, + enable_qvr_loss=enable_qvr_loss, + qvr_loss_weight=qvr_loss_weight, + qvr_loss_epsilon=qvr_loss_epsilon, + qvr_num_layers=qvr_num_layers, + enable_causal_loss=enable_causal_loss, + causal_loss_weight=causal_loss_weight, + causal_loss_epsilon=causal_loss_epsilon, + ) print(f"Running Deepspeed on rank = {rank}, world size = {world_size}") model = model.to(rank) @@ -365,6 +428,10 @@ def has_image(example): "train/loss": loss.detach().float() # * configs.gradient_accumulation_steps, } + if outputs.qvr_loss is not None: + log_dict["train/qvr_loss"] = outputs.qvr_loss.detach().float() + if outputs.causal_loss is not None: + log_dict["train/causal_loss"] = outputs.causal_loss.detach().float() wandb_run.log(log_dict) # print("line432") pbar.set_description( diff --git a/scripts/infer_adaptive_controller.py b/scripts/infer_adaptive_controller.py new file mode 100644 index 0000000..2b99166 --- /dev/null +++ b/scripts/infer_adaptive_controller.py @@ -0,0 +1,298 @@ +import argparse +import json +import os +import re +import sys +import time +from datetime import timedelta + +import torch +import yaml +from datasets import load_dataset +from peft import set_peft_model_state_dict +from qwen_vl_utils import process_vision_info +from tqdm import tqdm + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +QWEN_DIR = os.path.join(REPO_ROOT, "qwen_vl") +SCRIPTS_DIR = os.path.join(REPO_ROOT, "scripts") +for path in (QWEN_DIR, SCRIPTS_DIR): + if path not in sys.path: + sys.path.insert(0, path) + +from train_adaptive_controller import build_qwen2vl_adaptive_model +from utils import Config, set_seed + + +def load_config(path): + with open(path, "r", encoding="utf-8") as f: + return Config(yaml.safe_load(f)) + + +def format_m3cot_prompt(example): + question = example["question"].strip() + answer = str(example["answer"]).strip() + choices = example["choices"] + choices_str = "\n".join( + f"{chr(65 + i)}.{{{choice.strip()}}}" for i, choice in enumerate(choices) + ) + prompt = f"[Question]:{{{question}}}\n[Options]:\n{choices_str}\nAnswer:" + return { + "id": example["id"], + "question_raw": prompt, + "image_raw": example["image"], + "gt_answer": answer.upper(), + "choices": choices, + "domain": example.get("domain"), + "topic": example.get("topic"), + } + + +def format_scienceqa_prompt(example, idx): + question = example["question"].strip() + choices = example.get("choices", []) + if choices: + choices_str = "\n".join( + f"({chr(65 + i)}).{{{choice.strip()}}}" for i, choice in enumerate(choices) + ) + prompt = f"[Question]:{{{question}}}\n[Options]:\n{choices_str}\nAnswer:" + else: + prompt = f"[Question]:{{{question}}}\nAnswer:" + return { + "id": str(idx), + "question_raw": prompt, + "image_raw": example["image"], + "gt_answer": int(example["answer"]), + "choices": choices, + } + + +def build_eval_dataset(configs): + task = getattr(configs, "eval_task", "m3cot").lower() + data_percent = float(getattr(configs, "data_percent", 100.0)) + sample_seed = int(getattr(configs, "sample_seed", 42)) + if task == "m3cot": + dataset = load_dataset(getattr(configs, "eval_dataset_name", "LightChen2333/M3CoT")) + split = getattr(configs, "eval_split", "test") + eval_dataset = dataset[split].filter(lambda ex: ex["image"] is not None).map(format_m3cot_prompt) + elif task == "scienceqa": + dataset = load_dataset(getattr(configs, "eval_dataset_name", "derek-thomas/ScienceQA")) + split = getattr(configs, "eval_split", "test") + eval_dataset = dataset[split].map( + lambda ex, idx: {"original_idx": idx, **ex}, + with_indices=True, + ) + eval_dataset = eval_dataset.filter(lambda ex: "image" in ex and ex["image"] is not None) + eval_dataset = eval_dataset.map(lambda ex: format_scienceqa_prompt(ex, ex["original_idx"])) + else: + raise ValueError("eval_task must be 'm3cot' or 'scienceqa'") + + if not (0 < data_percent <= 100): + raise ValueError("data_percent must be in (0, 100].") + if data_percent < 100: + sample_size = max(1, int(len(eval_dataset) * data_percent / 100.0)) + eval_dataset = eval_dataset.shuffle(seed=sample_seed).select(range(sample_size)) + return eval_dataset + + +def extract_m3cot_answer(text): + matches = re.finditer( + r"(?:the\s+answer\s+is|Answer:)\s*[\n\s]*([A-Z])", + text, + flags=re.IGNORECASE | re.DOTALL, + ) + candidates = [m.group(1).upper() for m in matches] + return candidates[-1] if candidates else None + + +def extract_scienceqa_answer(text): + digit_patterns = [ + r"Therefore,?\s*the\s+answer\s+is\s+(\d)", + r"the\s+answer\s+is\s+(\d)", + r"answer\s+is:?\s*(\d)", + ] + for pattern in digit_patterns: + match = re.search(pattern, text, re.IGNORECASE) + if match: + return int(match.group(1)) + letter_patterns = [ + r"Therefore,?\s*the\s+answer\s+is\s+([A-Z])", + r"the\s+answer\s+is\s+([A-Z])", + r"answer\s+is:?\s*([A-Z])", + ] + for pattern in letter_patterns: + match = re.search(pattern, text, re.IGNORECASE) + if match: + return ord(match.group(1).upper()) - ord("A") + return -1 + + +def summarize_controller_trace(trace): + counts = [] + selected_indices = [] + for step in trace: + count = int(step["selected_counts"][0].item()) + counts.append(count) + indices = step["selected_patch_indices"][0] + selected_indices.append([int(x) for x in indices.tolist() if int(x) >= 0]) + return { + "selected_counts_by_latent_step": counts, + "selected_patch_indices_by_latent_step": selected_indices, + "total_selected_patches": int(sum(counts)), + "num_latent_steps": len(counts), + "avg_selected_per_latent_step": float(sum(counts) / len(counts)) if counts else 0.0, + } + + +def main(): + parser = argparse.ArgumentParser(description="Adaptive-controller IVT-LR inference") + parser.add_argument("--config", required=True) + parser.add_argument("--controller_checkpoint_path", default=None) + parser.add_argument("--output_path", default=None) + parser.add_argument("--summary_path", default=None) + parser.add_argument("--max_new_tokens", type=int, default=None) + args = parser.parse_args() + + configs = load_config(args.config) + set_seed(int(getattr(configs, "seed", 0))) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model, _, processor = build_qwen2vl_adaptive_model(configs, device) + + lora_stage2_path = getattr(configs, "lora_stage2_checkpoint_path", None) + if lora_stage2_path: + lora_state = torch.load(lora_stage2_path, map_location=device) + result = set_peft_model_state_dict(model.base_causallm, lora_state) + missing = len(getattr(result, "missing_keys", [])) + unexpected = len(getattr(result, "unexpected_keys", [])) + print( + f"Loaded Stage 2 LoRA checkpoint from {lora_stage2_path}. " + f"missing={missing} unexpected={unexpected}" + ) + + controller_path = args.controller_checkpoint_path or getattr( + configs, "controller_checkpoint_path", None + ) + if not controller_path: + raise ValueError("Set controller_checkpoint_path in config or pass --controller_checkpoint_path.") + controller_state = torch.load(controller_path, map_location=device) + model.controller.load_state_dict(controller_state, strict=True) + model.eval() + + eval_dataset = build_eval_dataset(configs) + task = getattr(configs, "eval_task", "m3cot").lower() + output_path = args.output_path or getattr( + configs, + "prediction_output_path", + os.path.join(getattr(configs, "output_dir", "."), f"adaptive_{task}_predictions.jsonl"), + ) + summary_path = args.summary_path or getattr( + configs, + "summary_output_path", + os.path.join(getattr(configs, "output_dir", "."), f"adaptive_{task}_summary.json"), + ) + max_new_tokens = args.max_new_tokens or int(getattr(configs, "max_new_tokens", 512)) + latent_n = int(getattr(configs, "latent_n", 3)) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + os.makedirs(os.path.dirname(summary_path), exist_ok=True) + + correct = 0 + total = 0 + total_generated_tokens = 0 + total_generate_time = 0.0 + total_selected_patches = 0 + total_latent_steps = 0 + selected_count_hist = {} + + with open(output_path, "w", encoding="utf-8") as f_out: + for ex in tqdm(eval_dataset, total=len(eval_dataset), desc=f"Adaptive {task} inference"): + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": ex["image_raw"], "resized_height": 280, "resized_width": 280}, + {"type": "text", "text": ex["question_raw"]}, + ], + }] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text = text + ("<|latent|>" * latent_n) + image_inputs, video_inputs = process_vision_info(messages) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ).to(device) + prompt_len = inputs["input_ids"].size(1) + + start_time = time.time() + with torch.no_grad(): + output_ids, controller_trace = model.generate( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + pixel_values=inputs["pixel_values"], + image_grid_thw=inputs["image_grid_thw"], + max_new_tokens=max_new_tokens, + output_controller_trace=True, + ) + generate_time = time.time() - start_time + total_generate_time += generate_time + generated_tokens = output_ids[0, prompt_len:] + generated_text = processor.decode(generated_tokens, skip_special_tokens=True) + total_generated_tokens += int(generated_tokens.numel()) + + if task == "scienceqa": + pred = extract_scienceqa_answer(generated_text) + is_correct = pred == int(ex["gt_answer"]) + else: + pred = extract_m3cot_answer(generated_text) + is_correct = pred == str(ex["gt_answer"]).upper() + correct += int(is_correct) + total += 1 + + trace_summary = summarize_controller_trace(controller_trace) + total_selected_patches += trace_summary["total_selected_patches"] + total_latent_steps += trace_summary["num_latent_steps"] + for c in trace_summary["selected_counts_by_latent_step"]: + selected_count_hist[c] = selected_count_hist.get(c, 0) + 1 + + result = { + "id": ex["id"], + "answer": ex["gt_answer"], + "prediction": pred, + "correct": bool(is_correct), + "generated_text": generated_text, + "controller": trace_summary, + } + if "choices" in ex: + result["choices"] = ex["choices"] + f_out.write(json.dumps(result, ensure_ascii=False) + "\n") + f_out.flush() + + accuracy = correct / total if total else 0.0 + avg_tokens = total_generated_tokens / total if total else 0.0 + avg_time = total_generate_time / total if total else 0.0 + avg_selected_per_example = total_selected_patches / total if total else 0.0 + avg_selected_per_latent = total_selected_patches / total_latent_steps if total_latent_steps else 0.0 + summary = { + "task": task, + "total": total, + "correct": correct, + "accuracy": accuracy, + "avg_generated_tokens": avg_tokens, + "total_generate_time_seconds": total_generate_time, + "avg_generate_time_seconds": avg_time, + "avg_selected_patches_per_example": avg_selected_per_example, + "avg_selected_patches_per_latent_step": avg_selected_per_latent, + "selected_count_histogram": selected_count_hist, + "prediction_output_path": output_path, + "controller_checkpoint_path": controller_path, + "lora_stage2_checkpoint_path": lora_stage2_path, + } + with open(summary_path, "w", encoding="utf-8") as f: + json.dump(summary, f, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + print(f"Total generate time: {timedelta(seconds=int(total_generate_time))}") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_adaptive_controller.py b/scripts/train_adaptive_controller.py new file mode 100644 index 0000000..1bb8ccb --- /dev/null +++ b/scripts/train_adaptive_controller.py @@ -0,0 +1,338 @@ +import argparse +import json +import os +import sys +from collections import Counter, defaultdict + +import torch +import torch.nn.functional as F +import yaml +from datasets import load_dataset +from peft import LoraConfig, get_peft_model +from torch.optim import AdamW +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import AutoProcessor, AutoTokenizer, Qwen2VLForConditionalGeneration + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +QWEN_DIR = os.path.join(REPO_ROOT, "qwen_vl") +if QWEN_DIR not in sys.path: + sys.path.insert(0, QWEN_DIR) + +from controller import PatchPointerController +from dataset import MyCollator, get_dataset, get_cot_latent_dataset +from qwen_adaptive_ivtlr import QwenAdaptiveIVTLR +from qwen_vl_utils import process_vision_info +from utils import Config, set_seed + + +def load_yaml(path): + with open(path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + + +def build_qwen2vl_adaptive_model(configs, device): + tokenizer = AutoTokenizer.from_pretrained( + configs.model_id, + use_fast=False, + trust_remote_code=True, + ) + tokenizer.padding_side = "right" + tokenizer.pad_token = tokenizer.eos_token + tokenizer.add_tokens("<|start-latent|>") + tokenizer.add_tokens("<|end-latent|>") + tokenizer.add_tokens("<|latent|>") + processor = AutoProcessor.from_pretrained(configs.model_id, tokenizer=tokenizer) + + base_model = Qwen2VLForConditionalGeneration.from_pretrained( + configs.model_id, + device_map=None, + torch_dtype=torch.bfloat16 if getattr(configs, "bf16", True) else torch.float32, + trust_remote_code=True, + attn_implementation="eager", + ) + base_model.resize_token_embeddings(len(tokenizer)) + if getattr(configs, "use_lora", True): + lora_config = LoraConfig( + task_type="CAUSAL_LM", + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + r=getattr(configs, "lora_r", 64), + lora_alpha=getattr(configs, "lora_alpha", 16), + lora_dropout=getattr(configs, "lora_dropout", 0.05), + bias="none", + inference_mode=False, + ) + base_model = get_peft_model(base_model, lora_config) + + latent_id = tokenizer.convert_tokens_to_ids("<|latent|>") + start_id = tokenizer.convert_tokens_to_ids("<|start-latent|>") + end_id = tokenizer.convert_tokens_to_ids("<|end-latent|>") + image_token_id = tokenizer.convert_tokens_to_ids(processor.image_token) + visual_start_id = tokenizer.convert_tokens_to_ids("<|vision_start|>") + visual_end_id = tokenizer.convert_tokens_to_ids("<|vision_end|>") + + controller = PatchPointerController( + model_dim=base_model.get_input_embeddings().embedding_dim, + controller_dim=getattr(configs, "controller_hidden_dim", None), + max_steps=configs.max_controller_steps, + use_step_embedding=getattr(configs, "use_step_embedding", True), + ) + model = QwenAdaptiveIVTLR( + base_model, + latent_token_id=latent_id, + start_latent_id=start_id, + end_latent_id=end_id, + eos_token_id=tokenizer.eos_token_id, + image_token_id=image_token_id, + visual_start_id=visual_start_id, + visual_end_id=visual_end_id, + controller=controller, + teacher_k=configs.teacher_k, + max_controller_steps=configs.max_controller_steps, + patch_reuse_policy=getattr(configs, "patch_reuse_policy", "never"), + processor_model_id=configs.model_id, + ) + + teacher_path = getattr(configs, "ivtlr_checkpoint_path", None) or getattr( + configs, "teacher_checkpoint_path", None + ) + if teacher_path: + state_dict = torch.load(teacher_path, map_location="cpu") + if any(k.startswith("module.") for k in state_dict.keys()): + state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()} + result = model.load_state_dict(state_dict, strict=False) + print( + f"Loaded teacher checkpoint. missing={len(result.missing_keys)} " + f"unexpected={len(result.unexpected_keys)}" + ) + + if getattr(configs, "freeze_base_model", True) or getattr(configs, "train_controller_only", True): + model.train_controller_only() + model.to(device) + return model, tokenizer, processor + + +def build_m3cot_dataset(configs, tokenizer, processor): + dataset = load_dataset(getattr(configs, "dataset_name", "LightChen2333/M3CoT")) + split = getattr(configs, "dataset_split", "train") + train_dataset = dataset[split].filter(lambda ex: "image" in ex and ex["image"] is not None) + + def process_example(example): + rationale = example["rationale"].replace("\n", " ").strip() + example["steps"] = rationale.split(". ") + if example["steps"] and example["steps"][-1] == "": + example["steps"].pop() + if len(example["steps"]) > 3: + total_steps = len(example["steps"]) + step_size = total_steps // 3 + remainder = total_steps % 3 + new_steps = [] + start = 0 + for i in range(3): + end = start + step_size + (1 if i < remainder else 0) + new_steps.append(". ".join(example["steps"][start:end])) + start = end + example["steps"] = new_steps + + choices_str = "[Options]:\n" + "\n".join( + f"({chr(65 + i)}).{{{choice.strip()}}}" + for i, choice in enumerate(example["choices"]) + ) + question = f"[Question]:{{{example['question'].strip()}}}\n{choices_str}\nAnswer:\n" + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": example["image"], "resized_height": 280, "resized_width": 280}, + {"type": "text", "text": question}, + ], + }] + example["question"] = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + image_inputs, video_inputs = process_vision_info(messages) + inputs = processor( + text=[example["question"]], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + inputs = {k: v.tolist() for k, v in inputs.items()} + example["input_ids"] = torch.tensor(inputs["input_ids"][0]) + example["image_grid_thw"] = torch.tensor(inputs["image_grid_thw"]).squeeze(0) + example["pixel_values"] = torch.tensor(inputs["pixel_values"]) + del example["rationale"] + del example["choices"] + return example + + num_proc = int(getattr(configs, "num_proc", 8)) + train_dataset = train_dataset.map(process_example, num_proc=num_proc) + max_size = int(getattr(configs, "max_train_examples", 100000000)) + return get_dataset(train_dataset, tokenizer, processor, max_size=max_size, num_proc=num_proc) + + +def compute_budget_weights(rewards, mode, temperature, eps=1e-8): + if mode == "softmax": + return F.softmax(rewards / max(temperature, eps), dim=-1) + mean = rewards.mean(dim=-1, keepdim=True) + std = rewards.std(dim=-1, keepdim=True, unbiased=False) + advantages = (rewards - mean) / (std + eps) + if mode == "positive_advantage": + positive = torch.relu(advantages) + return positive / positive.sum(dim=-1, keepdim=True).clamp(min=eps) + if mode == "signed_advantage": + return advantages + raise ValueError(f"Unknown advantage_mode={mode}") + + +def save_debug_trace(path, batch_idx, teacher_out, rewards, weights, budgets): + os.makedirs(os.path.dirname(path), exist_ok=True) + payload = { + "batch_idx": batch_idx, + "budgets": list(budgets), + "rewards": rewards.detach().float().cpu().tolist(), + "weights": weights.detach().float().cpu().tolist(), + "teacher_topk": [ + { + "latent_step_idx": step.latent_step_idx, + "ranked_patch_indices": step.ranked_patch_indices.detach().cpu().tolist(), + "attention_scores": step.attention_scores.detach().float().cpu().tolist(), + } + for step in teacher_out.latent_traces + ], + } + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(payload) + "\n") + + +def main(): + parser = argparse.ArgumentParser(description="Teacher-guided adaptive IVT-LR controller training") + parser.add_argument("--config", default=os.path.join(REPO_ROOT, "configs/adaptive_controller_qwen2b.yaml")) + args = parser.parse_args() + + config_dict = load_yaml(args.config) + configs = Config(config_dict) + set_seed(getattr(configs, "seed", 0)) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model, tokenizer, processor = build_qwen2vl_adaptive_model(configs, device) + + base_dataset = build_m3cot_dataset(configs, tokenizer, processor) + collator = MyCollator( + tokenizer, + latent_id=tokenizer.convert_tokens_to_ids("<|latent|>"), + label_pad_token_id=-100, + ) + dataloader = DataLoader( + get_cot_latent_dataset( + scheduled_stage=int(getattr(configs, "scheduled_stage", configs.max_latent_stage)), + base_dataset=base_dataset, + configs=configs, + start_id=tokenizer.convert_tokens_to_ids("<|start-latent|>"), + latent_id=tokenizer.convert_tokens_to_ids("<|latent|>"), + end_id=tokenizer.convert_tokens_to_ids("<|end-latent|>"), + no_special_marker=True, + shuffle=True, + ), + batch_size=int(getattr(configs, "batch_size_training", 1)), + shuffle=False, + num_workers=int(getattr(configs, "num_workers", 1)), + collate_fn=collator, + ) + + optimizer = AdamW( + [p for p in model.controller.parameters() if p.requires_grad], + lr=float(getattr(configs, "controller_lr", 1e-4)), + weight_decay=float(getattr(configs, "weight_decay", 0.0)), + ) + budgets = [int(x) for x in getattr(configs, "budget_candidates", [2, 4, 6, 8, 10])] + lambda_patch = float(getattr(configs, "lambda_patch", 0.002)) + temperature = float(getattr(configs, "reward_temperature", 1.0)) + advantage_mode = getattr(configs, "advantage_mode", "softmax") + grad_clip_norm = float(getattr(configs, "grad_clip_norm", 1.0)) + log_every = int(getattr(configs, "log_every", 10)) + save_every = int(getattr(configs, "save_every", 500)) + output_dir = getattr(configs, "output_dir", os.path.join(REPO_ROOT, "adaptive_controller_runs")) + os.makedirs(output_dir, exist_ok=True) + debug_trace_path = os.path.join(output_dir, "debug_traces.jsonl") + win_counter = Counter() + avg_advantage = defaultdict(float) + + model.train() + for global_step, batch in enumerate(tqdm(dataloader), start=1): + batch = {k: v.to(device) for k, v in batch.items() if k != "idx"} + with torch.no_grad(): + teacher_out = model( + **batch, + mode="teacher", + return_trace=True, + ) + rewards_by_budget = [] + logprobs_by_budget = [] + patch_counts = [] + for budget in budgets: + rollout = model( + **batch, + mode="forced_budget", + teacher_trace=teacher_out.latent_traces, + forced_budget=budget, + ) + patches_per_example = budget * max(len(teacher_out.latent_traces), 1) + reward = rollout.answer_logprob - lambda_patch * patches_per_example + rewards_by_budget.append(reward) + logprobs_by_budget.append(rollout.answer_logprob) + patch_counts.append(patches_per_example) + rewards = torch.stack(rewards_by_budget, dim=-1) + logprobs = torch.stack(logprobs_by_budget, dim=-1) + weights = compute_budget_weights(rewards, advantage_mode, temperature).detach() + + ctrl_stats = model.controller_teacher_forcing_loss( + teacher_out.latent_traces, + budgets=budgets, + budget_weights=weights, + ) + loss = ctrl_stats["loss"] + optimizer.zero_grad(set_to_none=True) + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(model.controller.parameters(), grad_clip_norm) + optimizer.step() + + winners = rewards.argmax(dim=-1).detach().cpu().tolist() + for winner in winners: + win_counter[budgets[winner]] += 1 + advantages = rewards - rewards.mean(dim=-1, keepdim=True) + for i, budget in enumerate(budgets): + avg_advantage[budget] += float(advantages[:, i].mean().item()) + + if global_step % log_every == 0: + reward_means = rewards.mean(dim=0).detach().float().cpu().tolist() + logprob_means = logprobs.mean(dim=0).detach().float().cpu().tolist() + print( + f"step={global_step} loss={float(loss.detach()):.4f} " + f"grad_norm={float(grad_norm):.4f} " + f"patch_acc={float(ctrl_stats['patch_top1_accuracy']):.3f} " + f"stop_acc={float(ctrl_stats['stop_accuracy']):.3f}" + ) + print(f" reward_by_K={dict(zip(budgets, reward_means))}") + print(f" answer_logprob_by_K={dict(zip(budgets, logprob_means))}") + print(f" best_K_counts={dict(win_counter)} avg_patch_count={sum(patch_counts) / len(patch_counts):.1f}") + + if global_step <= int(getattr(configs, "num_debug_traces", 5)): + save_debug_trace(debug_trace_path, global_step, teacher_out, rewards, weights, budgets) + + if global_step % save_every == 0: + ckpt_path = os.path.join(output_dir, f"controller_step_{global_step}.pt") + torch.save(model.controller.state_dict(), ckpt_path) + + max_steps = int(getattr(configs, "max_train_steps", 0)) + if max_steps and global_step >= max_steps: + break + + torch.save(model.controller.state_dict(), os.path.join(output_dir, "controller_final.pt")) + denom = max(global_step, 1) + avg_advantage = {k: v / denom for k, v in avg_advantage.items()} + print(f"final_best_K_counts={dict(win_counter)}") + print(f"final_avg_advantage_by_K={avg_advantage}") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_adaptive_lora_stage2.py b/scripts/train_adaptive_lora_stage2.py new file mode 100644 index 0000000..ef7718a --- /dev/null +++ b/scripts/train_adaptive_lora_stage2.py @@ -0,0 +1,193 @@ +import argparse +import os +import sys +from collections import deque + +import torch +import yaml +from peft import get_peft_model_state_dict, set_peft_model_state_dict +from torch.optim import AdamW +from torch.utils.data import DataLoader +from tqdm import tqdm + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +QWEN_DIR = os.path.join(REPO_ROOT, "qwen_vl") +if QWEN_DIR not in sys.path: + sys.path.insert(0, QWEN_DIR) + +from dataset import MyCollator, get_cot_latent_dataset +from train_adaptive_controller import build_m3cot_dataset, build_qwen2vl_adaptive_model +from utils import Config, set_seed + + +def load_yaml(path): + with open(path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + + +def load_controller_if_present(model, configs, device): + controller_path = getattr(configs, "controller_checkpoint_path", None) + if not controller_path: + raise ValueError("Stage 2 requires controller_checkpoint_path.") + controller_state = torch.load(controller_path, map_location=device) + model.controller.load_state_dict(controller_state, strict=True) + print(f"Loaded controller checkpoint from {controller_path}") + + +def load_stage2_lora_if_present(model, configs, device): + lora_path = getattr(configs, "lora_stage2_checkpoint_path", None) + if not lora_path: + return + lora_state = torch.load(lora_path, map_location=device) + result = set_peft_model_state_dict(model.base_causallm, lora_state) + missing = len(getattr(result, "missing_keys", [])) + unexpected = len(getattr(result, "unexpected_keys", [])) + print(f"Loaded Stage 2 LoRA checkpoint from {lora_path}. missing={missing} unexpected={unexpected}") + + +def freeze_controller_train_lora_only(model): + for param in model.parameters(): + param.requires_grad = False + for param in model.controller.parameters(): + param.requires_grad = False + + trainable = [] + for name, param in model.base_causallm.named_parameters(): + if "lora_" in name or ".lora_" in name: + param.requires_grad = True + trainable.append(param) + + if not trainable: + raise RuntimeError("No LoRA parameters were found. Set use_lora: true for Stage 2.") + return trainable + + +def save_lora(model, path): + os.makedirs(os.path.dirname(path), exist_ok=True) + torch.save(get_peft_model_state_dict(model.base_causallm), path) + + +def count_selected_patches(controller_trace): + total = 0 + steps = 0 + for step in controller_trace: + total += int(step["selected_counts"].sum().item()) + steps += int(step["selected_counts"].numel()) + return total, steps + + +def main(): + parser = argparse.ArgumentParser(description="Stage 2 adaptive IVT-LR LoRA tuning") + parser.add_argument("--config", required=True) + args = parser.parse_args() + + configs = Config(load_yaml(args.config)) + set_seed(int(getattr(configs, "seed", 0))) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model, tokenizer, processor = build_qwen2vl_adaptive_model(configs, device) + load_controller_if_present(model, configs, device) + load_stage2_lora_if_present(model, configs, device) + + trainable_params = freeze_controller_train_lora_only(model) + trainable_count = sum(p.numel() for p in trainable_params) + print(f"Stage 2 trainable LoRA parameters: {trainable_count:,}") + + model.base_causallm.train() + model.controller.eval() + + base_dataset = build_m3cot_dataset(configs, tokenizer, processor) + collator = MyCollator( + tokenizer, + latent_id=tokenizer.convert_tokens_to_ids("<|latent|>"), + label_pad_token_id=-100, + ) + latent_dataset = get_cot_latent_dataset( + scheduled_stage=int(getattr(configs, "scheduled_stage", configs.max_latent_stage)), + base_dataset=base_dataset, + configs=configs, + start_id=tokenizer.convert_tokens_to_ids("<|start-latent|>"), + latent_id=tokenizer.convert_tokens_to_ids("<|latent|>"), + end_id=tokenizer.convert_tokens_to_ids("<|end-latent|>"), + no_special_marker=True, + shuffle=True, + ) + dataloader = DataLoader( + latent_dataset, + batch_size=int(getattr(configs, "batch_size_training", 1)), + shuffle=False, + num_workers=int(getattr(configs, "num_workers", 1)), + collate_fn=collator, + ) + + optimizer = AdamW( + trainable_params, + lr=float(getattr(configs, "stage2_lora_lr", getattr(configs, "lora_lr", 2e-5))), + weight_decay=float(getattr(configs, "weight_decay", 0.0)), + ) + grad_accum = int(getattr(configs, "gradient_accumulation_steps", 1)) + grad_clip_norm = float(getattr(configs, "grad_clip_norm", 1.0)) + log_every = int(getattr(configs, "log_every", 10)) + save_every = int(getattr(configs, "save_every", 250)) + max_steps = int(getattr(configs, "max_train_steps", 0)) + output_dir = getattr(configs, "output_dir", os.path.join(REPO_ROOT, "adaptive_lora_stage2_runs")) + os.makedirs(output_dir, exist_ok=True) + + loss_window = deque(maxlen=max(log_every, 1)) + logprob_window = deque(maxlen=max(log_every, 1)) + patch_window = deque(maxlen=max(log_every, 1)) + optimizer.zero_grad(set_to_none=True) + update_step = 0 + + for global_step, batch in enumerate(tqdm(dataloader, desc="Stage2 adaptive LoRA"), start=1): + batch = {k: v.to(device) for k, v in batch.items() if k != "idx"} + out = model( + **batch, + mode="adaptive", + ) + answer_logprob = out.answer_logprob.mean() + loss = -answer_logprob / grad_accum + loss.backward() + + selected_total, latent_steps = count_selected_patches(out.controller_trace) + loss_window.append(float((-answer_logprob).detach().item())) + logprob_window.append(float(answer_logprob.detach().item())) + patch_window.append(float(selected_total / max(latent_steps, 1))) + + if global_step % grad_accum == 0: + grad_norm = torch.nn.utils.clip_grad_norm_(trainable_params, grad_clip_norm) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + update_step += 1 + else: + grad_norm = torch.tensor(0.0) + + if global_step % log_every == 0: + print( + f"step={global_step} update={update_step} " + f"answer_nll={sum(loss_window) / len(loss_window):.4f} " + f"answer_logprob={sum(logprob_window) / len(logprob_window):.4f} " + f"avg_selected_per_latent={sum(patch_window) / len(patch_window):.2f} " + f"grad_norm={float(grad_norm):.4f}" + ) + + if global_step % save_every == 0: + save_lora(model, os.path.join(output_dir, f"stage2_lora_step_{global_step}.pt")) + + if max_steps and global_step >= max_steps: + break + + if global_step % grad_accum != 0: + grad_norm = torch.nn.utils.clip_grad_norm_(trainable_params, grad_clip_norm) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + update_step += 1 + print(f"final_partial_update={update_step} grad_norm={float(grad_norm):.4f}") + + final_path = os.path.join(output_dir, "stage2_lora_final.pt") + save_lora(model, final_path) + print(f"Saved Stage 2 LoRA checkpoint to {final_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_controller_grpo_stage1.py b/scripts/train_controller_grpo_stage1.py new file mode 100644 index 0000000..310e50b --- /dev/null +++ b/scripts/train_controller_grpo_stage1.py @@ -0,0 +1,262 @@ +import argparse +import json +import os +import re +import sys +from collections import Counter + +import torch +import yaml +from datasets import load_dataset +from torch.optim import AdamW +from tqdm import tqdm + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +QWEN_DIR = os.path.join(REPO_ROOT, "qwen_vl") +SCRIPTS_DIR = os.path.join(REPO_ROOT, "scripts") +for path in (QWEN_DIR, SCRIPTS_DIR): + if path not in sys.path: + sys.path.insert(0, path) + +from qwen_vl_utils import process_vision_info +from train_adaptive_controller import build_qwen2vl_adaptive_model +from utils import Config, set_seed + + +def load_config(path): + with open(path, "r", encoding="utf-8") as f: + return Config(yaml.safe_load(f)) + + +def format_m3cot_prompt(example): + choices_str = "\n".join( + f"{chr(65 + i)}.{{{choice.strip()}}}" for i, choice in enumerate(example["choices"]) + ) + return { + "id": str(example.get("id", "")), + "question_raw": ( + f"[Question]:{{{example['question'].strip()}}}\n" + f"[Options]:\n{choices_str}\nAnswer:" + ), + "image_raw": example["image"], + "gt_answer": str(example["answer"]).strip().upper(), + } + + +def extract_m3cot_answer(text): + matches = re.finditer( + r"(?:the\s+answer\s+is|Answer:)\s*[\n\s]*([A-Z])", + text, + flags=re.IGNORECASE | re.DOTALL, + ) + candidates = [m.group(1).upper() for m in matches] + if candidates: + return candidates[-1] + fallback = re.search(r"\b([A-Z])\b", text) + return fallback.group(1).upper() if fallback else None + + +def build_train_dataset(configs): + dataset = load_dataset(getattr(configs, "dataset_name", "LightChen2333/M3CoT")) + split = getattr(configs, "dataset_split", "train") + train_dataset = dataset[split].filter(lambda ex: ex["image"] is not None) + train_dataset = train_dataset.map(format_m3cot_prompt) + train_dataset = train_dataset.shuffle(seed=int(getattr(configs, "seed", 0))) + max_examples = int(getattr(configs, "max_train_examples", 100000000)) + if max_examples > 0: + train_dataset = train_dataset.select(range(min(max_examples, len(train_dataset)))) + return train_dataset + + +def encode_example(processor, example, latent_n, device): + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": example["image_raw"], "resized_height": 280, "resized_width": 280}, + {"type": "text", "text": example["question_raw"]}, + ], + }] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text = text + ("<|latent|>" * latent_n) + image_inputs, video_inputs = process_vision_info(messages) + return processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ).to(device) + + +def summarize_trace(controller_trace): + logprob_sum = None + entropy_sum = None + selected_count = 0 + action_count = 0.0 + for step in controller_trace: + selected_count += int(step["selected_counts"][0].item()) + if "logprob_sum" in step: + logprob_sum = step["logprob_sum"] if logprob_sum is None else logprob_sum + step["logprob_sum"] + entropy_sum = step["entropy_sum"] if entropy_sum is None else entropy_sum + step["entropy_sum"] + action_count += float(step["action_count"].detach().sum().item()) + if logprob_sum is None: + raise RuntimeError("Sampled controller trace did not include logprob_sum.") + return logprob_sum.squeeze(0), entropy_sum.squeeze(0), selected_count, action_count + + +def main(): + parser = argparse.ArgumentParser(description="Stage 1 controller-only GRPO for adaptive IVT-LR") + parser.add_argument("--config", required=True) + parser.add_argument("--controller_checkpoint_path", default=None) + args = parser.parse_args() + + configs = load_config(args.config) + set_seed(int(getattr(configs, "seed", 0))) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model, _, processor = build_qwen2vl_adaptive_model(configs, device) + + controller_path = args.controller_checkpoint_path or getattr(configs, "controller_checkpoint_path", None) + if controller_path: + controller_state = torch.load(controller_path, map_location=device) + model.controller.load_state_dict(controller_state, strict=True) + print(f"Loaded controller checkpoint from {controller_path}") + + model.train_controller_only() + model.base_causallm.eval() + model.controller.train() + + output_dir = getattr(configs, "output_dir", os.path.join(REPO_ROOT, "adaptive_controller_grpo_runs")) + os.makedirs(output_dir, exist_ok=True) + rollouts_per_prompt = int(getattr(configs, "grpo_rollouts_per_prompt", 5)) + controller_temperature = float(getattr(configs, "grpo_controller_temperature", 1.0)) + min_patches = int(getattr(configs, "grpo_min_patches", 0)) + max_new_tokens = int(getattr(configs, "max_new_tokens", 64)) + latent_n = int(getattr(configs, "latent_n", 3)) + lambda_patch = float(getattr(configs, "lambda_patch", 0.0002)) + correct_reward = float(getattr(configs, "grpo_correct_reward", 1.0)) + incorrect_reward = float(getattr(configs, "grpo_incorrect_reward", 0.0)) + entropy_coef = float(getattr(configs, "grpo_entropy_coef", 0.0)) + grad_clip_norm = float(getattr(configs, "grad_clip_norm", 1.0)) + log_every = int(getattr(configs, "log_every", 10)) + save_every = int(getattr(configs, "save_every", 250)) + max_steps = int(getattr(configs, "max_train_steps", 0)) + + train_dataset = build_train_dataset(configs) + optimizer = AdamW( + [p for p in model.controller.parameters() if p.requires_grad], + lr=float(getattr(configs, "controller_lr", 5e-5)), + weight_decay=float(getattr(configs, "weight_decay", 0.0)), + ) + + correct_counter = Counter() + reward_window = [] + patch_window = [] + loss_window = [] + trace_path = os.path.join(output_dir, "grpo_stage1_traces.jsonl") + + for global_step, example in enumerate(tqdm(train_dataset, desc="Stage1 controller GRPO"), start=1): + inputs = encode_example(processor, example, latent_n, device) + prompt_len = inputs["input_ids"].size(1) + rollout_logprobs = [] + rollout_entropies = [] + rewards = [] + rollout_payloads = [] + + for rollout_idx in range(rollouts_per_prompt): + output_ids, controller_trace = model.generate_with_sampled_controller( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + pixel_values=inputs["pixel_values"], + image_grid_thw=inputs["image_grid_thw"], + max_new_tokens=max_new_tokens, + controller_temperature=controller_temperature, + min_patches=min_patches, + ) + generated_tokens = output_ids[0, prompt_len:] + generated_text = processor.decode(generated_tokens, skip_special_tokens=True) + pred = extract_m3cot_answer(generated_text) + is_correct = pred == example["gt_answer"] + logprob_sum, entropy_sum, selected_count, action_count = summarize_trace(controller_trace) + reward = (correct_reward if is_correct else incorrect_reward) - lambda_patch * selected_count + + rollout_logprobs.append(logprob_sum) + rollout_entropies.append(entropy_sum / max(action_count, 1.0)) + rewards.append(reward) + correct_counter[int(is_correct)] += 1 + reward_window.append(reward) + patch_window.append(selected_count) + rollout_payloads.append( + { + "rollout": rollout_idx, + "prediction": pred, + "correct": bool(is_correct), + "selected_count": selected_count, + "reward": reward, + "generated_text": generated_text, + } + ) + + rewards_t = torch.tensor(rewards, dtype=torch.float32, device=device) + advantages = rewards_t - rewards_t.mean() + std = rewards_t.std(unbiased=False) + if float(std.item()) > 1e-6: + advantages = advantages / (std + 1e-6) + logprobs_t = torch.stack(rollout_logprobs) + entropies_t = torch.stack(rollout_entropies) + loss = -(advantages.detach() * logprobs_t).mean() - entropy_coef * entropies_t.mean() + + optimizer.zero_grad(set_to_none=True) + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(model.controller.parameters(), grad_clip_norm) + optimizer.step() + loss_window.append(float(loss.detach().item())) + + if global_step <= int(getattr(configs, "num_debug_traces", 5)): + with open(trace_path, "a", encoding="utf-8") as f: + f.write( + json.dumps( + { + "step": global_step, + "id": example["id"], + "answer": example["gt_answer"], + "rewards": rewards, + "advantages": advantages.detach().float().cpu().tolist(), + "rollouts": rollout_payloads, + }, + ensure_ascii=False, + ) + + "\n" + ) + + if global_step % log_every == 0: + total_rollouts = correct_counter[0] + correct_counter[1] + acc = correct_counter[1] / max(total_rollouts, 1) + avg_reward = sum(reward_window[-log_every * rollouts_per_prompt :]) / max( + len(reward_window[-log_every * rollouts_per_prompt :]), 1 + ) + avg_patches = sum(patch_window[-log_every * rollouts_per_prompt :]) / max( + len(patch_window[-log_every * rollouts_per_prompt :]), 1 + ) + avg_loss = sum(loss_window[-log_every:]) / max(len(loss_window[-log_every:]), 1) + print( + f"step={global_step} loss={avg_loss:.4f} reward={avg_reward:.4f} " + f"rollout_acc={acc:.3f} avg_patches={avg_patches:.2f} " + f"grad_norm={float(grad_norm):.4f}" + ) + + if global_step % save_every == 0: + torch.save( + model.controller.state_dict(), + os.path.join(output_dir, f"controller_grpo_step_{global_step}.pt"), + ) + + if max_steps and global_step >= max_steps: + break + + final_path = os.path.join(output_dir, "controller_grpo_final.pt") + torch.save(model.controller.state_dict(), final_path) + print(f"Saved final controller to {final_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_adaptive_controller.py b/tests/test_adaptive_controller.py new file mode 100644 index 0000000..1564f68 --- /dev/null +++ b/tests/test_adaptive_controller.py @@ -0,0 +1,59 @@ +import os +import sys + +import torch + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +QWEN_DIR = os.path.join(REPO_ROOT, "qwen_vl") +if QWEN_DIR not in sys.path: + sys.path.insert(0, QWEN_DIR) + +from controller import PatchPointerController + + +def test_stop_is_extra_controller_logit_not_patch_embedding(): + controller = PatchPointerController(model_dim=8, controller_dim=8, max_steps=4) + state = torch.randn(2, 8) + patches = torch.randn(2, 5, 8) + valid = torch.ones(2, 5, dtype=torch.bool) + logits = controller(controller.initial_state(state), patches, valid) + + assert logits.shape == (2, 6) + stop_index = patches.size(1) + target_sequence = torch.tensor([[1, 2, stop_index], [0, 3, stop_index]]) + inserted_patch_count = (target_sequence != stop_index).sum(dim=1) + assert inserted_patch_count.tolist() == [2, 2] + + +def test_selected_patches_are_masked_from_future_steps(): + controller = PatchPointerController(model_dim=8, controller_dim=8, max_steps=4) + state = controller.initial_state(torch.randn(1, 8)) + patches = torch.randn(1, 4, 8) + valid = torch.ones(1, 4, dtype=torch.bool) + selected = torch.zeros(1, 4, dtype=torch.bool) + selected[0, 2] = True + + logits = controller(state, patches, valid, selected_mask=selected) + assert torch.isneginf(logits[0, 2]) + assert not torch.isneginf(logits[0, 4]) + + +def test_teacher_forcing_is_sequential_and_predicts_stop_after_patches(): + torch.manual_seed(0) + controller = PatchPointerController(model_dim=8, controller_dim=8, max_steps=4) + reasoning = torch.randn(1, 8) + patches = torch.randn(1, 4, 8) + valid = torch.ones(1, 4, dtype=torch.bool) + stop_index = patches.size(1) + targets = torch.tensor([[1, 2, stop_index]]) + + stats = controller.teacher_forced_sequence_loss(reasoning, patches, valid, targets) + assert torch.isfinite(stats.loss) + assert stats.token_count == 3 + + +if __name__ == "__main__": + test_stop_is_extra_controller_logit_not_patch_embedding() + test_selected_patches_are_masked_from_future_steps() + test_teacher_forcing_is_sequential_and_predicts_stop_after_patches() + print("adaptive controller sanity tests passed")