Codex/adaptive controller training

.gitignore

@@ -205,3 +205,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+#Folder for storing the generated images
+Metrics/

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 <span id="chameleon"></span>
 
 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

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"
+lr: !!float "4e-5"
+patch_reuse_policy: never

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):

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)

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)

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)
 

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)
 

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

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