Merge branch 'main' into fix_requirements

Author: jiapingW

configs/qwen2.5-vl-32b-eagle3.json

@@ -268,7 +268,7 @@ def build_target_model(
         if (
             args.is_vlm
             and draft_model_config.target_model_type == "qwen2_5_vl"
-            and args.target_model_backend == "custom"
+            and args.tp_size == 1
         ):
             from transformers import Qwen2_5_VLForConditionalGeneration
 
@@ -456,6 +456,7 @@ def build_dataloaders(
         ),
         is_vlm=args.is_vlm,
     )
+
     if args.eval_data_path is not None or args.eval_hidden_states_path is not None:
         if args.eval_data_path is not None:
             eval_dataset = load_dataset("json", data_files=args.eval_data_path)["train"]
@@ -546,7 +547,7 @@ def run_forward(
     target_model: Optional[Eagle3TargetModel] = None,
     is_online: bool = True,
 ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
-    if args.is_vlm and args.target_model_backend == "custom":
+    if args.is_vlm:
         plosses, _, acces = eagle3_model(
             input_ids=data["input_ids"].cuda(),
             attention_mask=data["attention_mask"].cuda(),
@@ -557,20 +558,10 @@ def run_forward(
     else:
         if is_online:
             # we generate the eagle3 using the target model in an online fashion
-            # Handle VLM data: pixel_values and image_grid_thw are lists
-            # pixel_values = [pv.cuda() for pv in data["pixel_values"]] if args.is_vlm else None
-            image_grid_thw = (
-                [thw.cuda().squeeze() for thw in data["image_grid_thw"]]
-                if args.is_vlm
-                else None
-            )
             eagle3_data = target_model.generate_eagle3_data(
                 input_ids=data["input_ids"].cuda(),
                 attention_mask=data["attention_mask"].cuda(),
                 loss_mask=data["loss_mask"].cuda(),
-                is_vlm=args.is_vlm,
-                pixel_values=data["pixel_values"].cuda(),
-                image_grid_thw=image_grid_thw,
             )
 
             input_ids = get_dp_data_shard_from_tp(eagle3_data.input_ids)
@@ -588,14 +579,13 @@ def run_forward(
             input_ids, target, loss_mask = target_model.preprocess(
                 input_ids, target, loss_mask
             )
+
         plosses, _, acces = eagle3_model(
             input_ids=input_ids,
             attention_mask=attention_mask,
             loss_mask=loss_mask,
             target=target,
             hidden_states=hidden_states,
-            image_grid_thw=image_grid_thw,
-            is_vlm=args.is_vlm,
         )
     return plosses, acces
 
@@ -757,8 +747,6 @@ def main():
     if (
         args.is_vlm
         and getattr(draft_model_config, "target_model_type", None) == "qwen2_5_vl"
-        and args.tp_size == 1
-        and args.target_model_backend != "sglang"
     ):
         eagle3_model = QwenVLOnlineEagle3Model(
             target_model=target_model,
@@ -769,7 +757,6 @@ def main():
         )
     else:
         eagle3_model = OnlineEagle3Model(
-            target_model=target_model,
             draft_model=draft_model,
             length=args.ttt_length,
             attention_backend=args.attention_backend,
@@ -923,6 +910,7 @@ def main():
                     tracker,
                     mode="eval",
                 )
+
             # ================================================
             # 7.3 Save Checkpoints
             # ================================================
@@ -935,6 +923,7 @@ def main():
 
         if args.max_num_steps is not None and global_step >= args.max_num_steps:
             break
+
     # Save final checkpoint if training ended without saving
     if global_step % args.save_interval != 0:
         print_on_rank0(

examples/run_qwen2.5_32b_vl_eagle3_online.sh

@@ -59,7 +59,6 @@ def __init__(
         draft_model: Eagle3DraftModel,
         length: int = 7,
         attention_backend="sdpa",
-        target_model: Optional[Eagle3Model] = None,
     ):
         """
         Args:
@@ -71,7 +70,6 @@ def __init__(
         self.draft_model = draft_model
         self.length = length
         self.attention_backend = attention_backend
-        self.target_model = target_model
 
         if self.attention_backend == "usp":
             self.extract_func = EXTRACT_FUNC_DICT["basic"]
@@ -100,8 +98,6 @@ def forward(
         hidden_states: torch.Tensor,
         past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         position_ids: Optional[torch.Tensor] = None,
-        image_grid_thw: Optional[torch.Tensor] = None,
-        is_vlm: bool = False,
         **kwargs,
     ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]:
         """
@@ -136,22 +132,14 @@ def forward(
             past_key_values_length = past_key_values[0][0].shape[2]
             seq_length_with_past = seq_length_with_past + past_key_values_length
         if position_ids is None:
-            if is_vlm:
-                mrope_positions_ids, mrope_position_delta = (
-                    self.target_model.get_rope_index(
-                        input_ids=input_ids, image_grid_thw=image_grid_thw
-                    )
-                )
-                position_ids = mrope_positions_ids
-            else:
-                device = hidden_states.device
-                position_ids = torch.arange(
-                    past_key_values_length,
-                    seq_length + past_key_values_length,
-                    dtype=torch.long,
-                    device=device,
-                )
-                position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+            device = hidden_states.device
+            position_ids = torch.arange(
+                past_key_values_length,
+                seq_length + past_key_values_length,
+                dtype=torch.long,
+                device=device,
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
         else:
             position_ids = position_ids.view(-1, seq_length).long()