address reviewer comments.

Co-authored-by: hukongyi <hukongyi@cmbchina.com>

Author: 2 people

scripts/train_dflash.py

@@ -83,7 +83,9 @@ def parse_args():
         "Suggested: 7 for block_size=16, 5 for 10, 4 for 8. None disables.",
     )
     model_group.add_argument(
-        "--embed-key", type=str, default="model.language_model.embed_tokens.weight"
+        "--embed-key", type=str, default=None,
+        help="Key for embedding weights in the target model checkpoint. "
+             "Defaults to auto-detection based on model architecture.",
     )
     model_group.add_argument("--lm-head-key", type=str, default="lm_head.weight")
     model_group.add_argument("--is-vlm", action="store_true")

specforge/data/preprocessing.py

@@ -201,8 +201,6 @@ def preprocess_vlm_conversations(
             - pixel_values: List of pixel values for images in the examples.
             - image_grid_thw: List of image grid tensors.
     """
-    system_prompt = chat_template.system_prompt
-
     # prepare result
     results = {
         "input_ids": [],
@@ -212,16 +210,15 @@ def preprocess_vlm_conversations(
         "image_grid_thw": [],
     }
 
-    # Note: currently, we assume that each example has only one image
-    for i, image in enumerate(examples["image"]):
+    for i, images in enumerate(examples["images"]):
         source = examples["conversations"][i]
         messages = []
         # messages = [{"role": "system", "content": system_prompt}]
         if not source:
             # if the source is None, skip it
             continue
 
-        if not image:
+        if not images:
             text_messages = []
             convroles = ["user", "assistant"]
             for j, sentence in enumerate(source):
@@ -266,26 +263,17 @@ def preprocess_vlm_conversations(
             source = source[1:]
 
         convroles = ["user", "assistant"]
-        has_added_image = False
+        has_added_images = False
         for j, sentence in enumerate(source):
             role = sentence["role"]
             assert role == convroles[j % 2], f"unexpected role {role}"
             if role == "user":
-                # if the message is from user and has image, process the image
-                if not has_added_image:
-                    messages.append(
-                        {
-                            "role": role,
-                            "content": [
-                                {
-                                    "type": "image",
-                                    "image": image,
-                                },
-                                {"type": "text", "text": sentence["content"]},
-                            ],
-                        }
-                    )
-                    has_added_image = True
+                # Insert all images into the first user message
+                if not has_added_images:
+                    content = [{"type": "image", "image": img} for img in images]
+                    content.append({"type": "text", "text": sentence["content"]})
+                    messages.append({"role": role, "content": content})
+                    has_added_images = True
                 else:
                     messages.append({"role": role, "content": sentence["content"]})
             else:
@@ -318,7 +306,7 @@ def preprocess_vlm_conversations(
         input_ids = encoding.input_ids[0]
         offsets = encoding.offset_mapping[0]
         pixel_values = encoding.pixel_values
-        image_grid_thw = encoding.image_grid_thw[0]
+        image_grid_thw = encoding.image_grid_thw  # shape: (num_images, 3)
 
         # get conversation with image info for loss mask generation
         decoded_conversation = processor.tokenizer.decode(
@@ -334,7 +322,7 @@ def preprocess_vlm_conversations(
         results["loss_mask"].append(loss_mask[None, :])
         results["attention_mask"].append(torch.ones_like(loss_mask)[None, :])
         results["pixel_values"].append(pixel_values)
-        results["image_grid_thw"].append(image_grid_thw[None, :])
+        results["image_grid_thw"].append(image_grid_thw)
     return results
 
 

specforge/data/utils.py

@@ -205,6 +205,10 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
                 - attention_mask: torch.Tensor of shape (B, N)
                 - loss_mask: torch.Tensor of shape (B, N)
         """
+        assert len(features) == 1, (
+            f"VlmDataCollatorWithPadding requires batch_size=1, got {len(features)}. "
+            "Set per_device_train_batch_size=1 in your training config."
+        )
         max_length = max(item["input_ids"].shape[1] for item in features)
         batch_input_ids = torch.cat(
             [self.paddingtensor2D(item["input_ids"], max_length) for item in features]

specforge/modeling/target/dflash_target_model.py

@@ -110,12 +110,8 @@ def _init_vlm_attributes(self):
         self.spatial_merge_size = getattr(vision_config, "spatial_merge_size", 2)
         self.vlm_model_type = getattr(vision_config, "model_type", "")
 
-        text_config = getattr(hf_config, "text_config", hf_config)
         self.tokens_per_second = None
 
-        rope_params = getattr(text_config, "rope_parameters", {}) or {}
-        self.mrope_interleaved = rope_params.get("mrope_interleaved", False)
-
     @classmethod
     def from_pretrained(
         cls,
@@ -434,9 +430,9 @@ def generate_dflash_data(
             output_hidden_states=True,
             use_cache=False,
         )
-        if pixel_values:
+        if pixel_values is not None:
             model_kwargs["pixel_values"] = pixel_values
-        if image_grid_thw:
+        if image_grid_thw is not None:
             model_kwargs["image_grid_thw"] = image_grid_thw
         outputs = self.model(**model_kwargs)
 

specforge/modeling/target/sglang_backend/patch.py

@@ -98,7 +98,7 @@ def initialize_model_parallel(
         4 tensor model-parallel groups:
             [g0, g1], [g2, g3], [g4, g5], [g6, g7]
         2 pipeline model-parallel groups:
-            [g0, g2, g4, g6], [b1, g3, g5, g7]
+            [g0, g2, g4, g6], [g1, g3, g5, g7]
     Note that for efficiency, the caller should make sure adjacent ranks
     are on the same DGX box. For example if we are using 2 DGX-1 boxes
     with a total of 16 GPUs, rank 0 to 7 belong to the first box and

specforge/modeling/target/target_utils.py

@@ -52,7 +52,10 @@ def from_pretrained(
         instance = cls(config)
 
         if embed_key is None:
-            embed_key = "model.embed_tokens.weight"
+            if hasattr(config, "text_config") and config.text_config is not None:
+                embed_key = "model.language_model.embed_tokens.weight"
+            else:
+                embed_key = "model.embed_tokens.weight"
         if lm_head_key is None:
             lm_head_key = "lm_head.weight"
 

tests/test_modeling/test_target/test_sglang_backend/test_sglang_backend.py

@@ -32,7 +32,6 @@ def test_dense(rank, world_size, port, tp_size):
         device="cuda",
         attention_backend="fa3",
         mem_fraction_static=0.4,
-        # enable_torch_compile=True,
         enable_nccl_nvls=True,
         enable_symm_mem=False,
         enable_torch_compile=True,