feat: DFlash VLM training support with SGLang backend

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

Author: 2 people

scripts/train_dflash.py

@@ -20,7 +20,7 @@
 from tqdm import tqdm
 from transformers import AutoConfig, AutoTokenizer
 
-from datasets import load_dataset
+from datasets import load_dataset, load_from_disk
 from specforge.args import SGLangBackendArgs, TrackerArgs
 from specforge.core.dflash import OnlineDFlashModel
 from specforge.data import build_eagle3_dataset, prepare_dp_dataloaders
@@ -35,6 +35,8 @@
 from specforge.tracker import create_tracker
 from specforge.utils import get_last_checkpoint, print_on_rank0, print_with_rank
 
+logging.getLogger("sglang.srt.mem_cache.memory_pool").setLevel(logging.WARNING)
+
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Train DFlash Draft Model")
@@ -80,6 +82,13 @@ def parse_args():
         help="Gamma for exponential loss decay weighting (paper Eq.4). "
         "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"
+    )
+    model_group.add_argument("--lm-head-key", type=str, default="lm_head.weight")
+    model_group.add_argument("--is-vlm", action="store_true")
+    model_group.add_argument("--min-pixels", type=int, default=50176)
+    model_group.add_argument("--max-pixels", type=int, default=802816)
 
     dataset_group = parser.add_argument_group("dataset")
     dataset_group.add_argument("--train-data-path", type=str, required=True)
@@ -190,7 +199,16 @@ def build_models(args) -> Tuple[DFlashTargetModel, DFlashDraftModel]:
     return target_model, draft_model
 
 
-def build_dataloader(args, tokenizer) -> Tuple[DataLoader, Optional[DataLoader]]:
+def _load_raw_dataset(data_path: str):
+    """load jsonl"""
+    if os.path.isdir(data_path):
+        return load_from_disk(data_path)
+    return load_dataset("json", data_files=data_path)["train"]
+
+
+def build_dataloader(
+    args, tokenizer, processor=None
+) -> Tuple[DataLoader, Optional[DataLoader]]:
     """Build train and eval dataloaders."""
     import hashlib
 
@@ -202,7 +220,7 @@ def build_dataloader(args, tokenizer) -> Tuple[DataLoader, Optional[DataLoader]]
     )
     cache_key = hashlib.md5(cache_params_string.encode()).hexdigest()
 
-    train_dataset = load_dataset("json", data_files=args.train_data_path)["train"]
+    train_dataset = _load_raw_dataset(args.train_data_path)
     train_eagle3_dataset = build_eagle3_dataset(
         dataset=train_dataset,
         tokenizer=tokenizer,
@@ -212,8 +230,9 @@ def build_dataloader(args, tokenizer) -> Tuple[DataLoader, Optional[DataLoader]]
         cache_dir=os.path.join(args.cache_dir, "processed_dataset"),
         cache_key=cache_key,
         num_proc=args.build_dataset_num_proc,
+        is_vlm=args.is_vlm,
+        processor=processor,
     )
-
     min_loss_tokens = 2 * args.block_size
     original_size = len(train_eagle3_dataset)
     train_eagle3_dataset = train_eagle3_dataset.filter(
@@ -229,24 +248,28 @@ def build_dataloader(args, tokenizer) -> Tuple[DataLoader, Optional[DataLoader]]
         num_workers=args.dataloader_num_workers,
         shuffle=True,
         process_group=get_dp_group(),
+        is_vlm=args.is_vlm,
     )
 
     eval_dataloader = None
     if args.eval_data_path:
-        eval_dataset = load_dataset("json", data_files=args.eval_data_path)["train"]
+        eval_dataset = _load_raw_dataset(args.eval_data_path)
         eval_eagle3_dataset = build_eagle3_dataset(
             dataset=eval_dataset,
             tokenizer=tokenizer,
             chat_template=args.chat_template,
             max_length=args.max_length,
             is_preformatted=args.is_preformatted,
+            is_vlm=args.is_vlm,
+            processor=processor,
         )
         eval_dataloader = prepare_dp_dataloaders(
             eval_eagle3_dataset,
             args.batch_size,
             num_workers=args.dataloader_num_workers,
             shuffle=False,
             process_group=get_dp_group(),
+            is_vlm=args.is_vlm,
         )
 
     return train_dataloader, eval_dataloader
@@ -353,11 +376,18 @@ def main():
                 f"Provided ckpt dir {args.ckpt_dir} is not a valid directory."
             )
 
+    start_epoch = 0
+    global_step = 0
+    ckpt_info = None
+
     if args.resume and os.path.isdir(args.output_dir):
         draft_model_last_checkpoint, ckpt_info = get_last_checkpoint(
             args.output_dir, prefix=r"epoch_\d+_step"
         )
-        print_on_rank0(f"Last checkpoint detected: {draft_model_last_checkpoint}")
+        if ckpt_info:
+            start_epoch = ckpt_info[0]
+            global_step = ckpt_info[1]
+            print_on_rank0(f"Last checkpoint detected: {draft_model_last_checkpoint}")
 
     resume_state = None
     if draft_model_last_checkpoint:
@@ -380,7 +410,23 @@ def main():
                 f"step {resume_state['global_step']}"
             )
 
-    tokenizer = AutoTokenizer.from_pretrained(args.target_model_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.target_model_path, trust_remote_code=args.trust_remote_code
+    )
+
+    processor = None
+    if args.is_vlm:
+        from transformers import AutoProcessor
+
+        processor = AutoProcessor.from_pretrained(
+            args.target_model_path,
+            min_pixels=args.min_pixels,
+            max_pixels=args.max_pixels,
+            trust_remote_code=args.trust_remote_code,
+        )
+        print_on_rank0(
+            f"Loaded VLM processor (min_pixels={args.min_pixels}, max_pixels={args.max_pixels})"
+        )
 
     if args.mask_token_id is not None:
         mask_token_id = args.mask_token_id
@@ -396,18 +442,17 @@ def main():
     draft_model.config.dflash_config["target_layer_ids"] = draft_model.target_layer_ids
     print_on_rank0(f"dflash_config: {draft_model.config.dflash_config}")
 
-    train_dataloader, eval_dataloader = build_dataloader(args, tokenizer)
-
+    train_dataloader, eval_dataloader = build_dataloader(args, tokenizer, processor)
     steps_per_epoch = math.ceil(len(train_dataloader) / args.accumulation_steps)
     total_steps = args.num_epochs * steps_per_epoch
     print_on_rank0(f"Total training steps: {total_steps}")
 
     print_on_rank0("Loading target embeddings and head...")
     target_components = TargetEmbeddingsAndHead.from_pretrained(
         args.target_model_path,
-        embed_key="model.embed_tokens.weight",  # Adjust if Qwen/Llama differs
-        lm_head_key="lm_head.weight",
-        device="cuda",
+        embed_key=args.embed_key,
+        lm_head_key=args.lm_head_key,
+        device=torch.cuda.current_device(),
         trust_remote_code=args.trust_remote_code,
     )
 
@@ -441,8 +486,6 @@ def main():
         total_steps=total_steps,
     )
 
-    start_epoch = ckpt_info[0]
-    global_step = ckpt_info[1]
     if resume_state is not None:
         optimizer.scheduler.load_state_dict(resume_state["scheduler_state_dict"])
         start_epoch = resume_state["epoch"]
@@ -475,11 +518,30 @@ def main():
                 continue
             global_step += 1
 
+            input_ids_cpu = data["input_ids"]
+            attention_mask_cpu = data["attention_mask"]
+            loss_mask_cpu = data["loss_mask"]
+
             input_ids = data["input_ids"].cuda()
-            attention_mask = data["attention_mask"].cuda()
             loss_mask = data["loss_mask"].cuda()
+            pixel_values = None
+            image_grid_thw_cpu = None
+            if (
+                args.is_vlm
+                and "pixel_values" in data
+                and data["pixel_values"] is not None
+            ):
+                pixel_values = data["pixel_values"].cuda()
+                image_grid_thw_cpu = [
+                    thw.squeeze() if thw is not None else None
+                    for thw in data["image_grid_thw"]
+                ]
             target_output = target_model.generate_dflash_data(
-                input_ids, attention_mask, loss_mask
+                input_ids_cpu,
+                attention_mask_cpu,
+                loss_mask_cpu,
+                pixel_values=pixel_values,
+                image_grid_thw=image_grid_thw_cpu,
             )
             hidden_states = target_output.hidden_states.cuda()  # Ensure on GPU
 

specforge/data/preprocessing.py

@@ -215,33 +215,79 @@ def preprocess_vlm_conversations(
     # Note: currently, we assume that each example has only one image
     for i, image in enumerate(examples["image"]):
         source = examples["conversations"][i]
-        messages = [{"role": "system", "content": system_prompt}]
+        messages = []
+        # messages = [{"role": "system", "content": system_prompt}]
         if not source:
             # if the source is None, skip it
             continue
 
+        if not image:
+            text_messages = []
+            convroles = ["user", "assistant"]
+            for j, sentence in enumerate(source):
+                role = sentence["role"]
+                assert role == convroles[j % 2], f"unexpected role {role}"
+                text_messages.append({"role": role, "content": sentence["content"]})
+            conversation = processor.apply_chat_template(
+                text_messages,
+                tokenize=False,
+                add_generation_prompt=False,
+            )
+            encoding = processor(
+                text=[conversation],
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+                return_offsets_mapping=True,
+                add_special_tokens=False,
+            )
+
+            input_ids = encoding.input_ids[0]
+            offsets = encoding.offset_mapping[0]
+
+            # get conversation with image info for loss mask generation
+            decoded_conversation = processor.tokenizer.decode(
+                encoding.input_ids[0], skip_special_tokens=False
+            )
+
+            # Apply loss mask
+            loss_mask = _apply_loss_mask_from_chat_template(
+                decoded_conversation, offsets, chat_template
+            )
+            results["input_ids"].append(input_ids[None, :])
+            results["loss_mask"].append(loss_mask[None, :])
+            results["attention_mask"].append(torch.ones_like(loss_mask)[None, :])
+            results["pixel_values"].append(torch.empty(0, 0).float())
+            results["image_grid_thw"].append([])
+            continue
+
         if source[0]["role"] != "user":
             # if the first message is not from user, skip it
             source = source[1:]
 
         convroles = ["user", "assistant"]
+        has_added_image = 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
-                messages.append(
-                    {
-                        "role": role,
-                        "content": [
-                            {
-                                "type": "image",
-                                "image": image,
-                            },
-                            {"type": "text", "text": sentence["content"]},
-                        ],
-                    }
-                )
+                if not has_added_image:
+                    messages.append(
+                        {
+                            "role": role,
+                            "content": [
+                                {
+                                    "type": "image",
+                                    "image": image,
+                                },
+                                {"type": "text", "text": sentence["content"]},
+                            ],
+                        }
+                    )
+                    has_added_image = True
+                else:
+                    messages.append({"role": role, "content": sentence["content"]})
             else:
                 messages.append({"role": role, "content": sentence["content"]})
 

specforge/data/utils.py

@@ -218,12 +218,30 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
         batch_loss_mask = torch.cat(
             [self.paddingtensor2D(item["loss_mask"], max_length) for item in features]
         )
-        batch_pixel_values = torch.cat(
-            [item["pixel_values"] for item in features], dim=0
-        )
-        batch_image_grid_thw = torch.cat(
-            [item["image_grid_thw"] for item in features], dim=0
-        )
+        # Collect pixel_values and image_grid_thw per sample.
+        # Image samples have non-empty pixel_values; text-only samples have empty tensors.
+        all_pixel_values = []
+        all_image_grid_thw = []
+        for item in features:
+            pv = item.get("pixel_values")
+            thw = item.get("image_grid_thw")
+            if pv is not None and isinstance(pv, torch.Tensor) and pv.numel() > 0:
+                all_pixel_values.append(pv)
+                all_image_grid_thw.append(thw)
+            else:
+                all_image_grid_thw.append(None)
+
+        if all_pixel_values:
+            batch_pixel_values = torch.cat(all_pixel_values, dim=0)
+        else:
+            batch_pixel_values = None
+
+        # If all samples are text-only, set image_grid_thw to None
+        if all(thw is None for thw in all_image_grid_thw):
+            batch_image_grid_thw = None
+        else:
+            batch_image_grid_thw = all_image_grid_thw
+
         batch = {
             "input_ids": batch_input_ids,
             "attention_mask": batch_attention_mask,
@@ -304,17 +322,10 @@ def prepare_dp_dataloaders(
 
 
 def parse_harmony_message_content(content):
-    """
-    解析 content 字符串中的 Harmony 格式。
-    如果匹配到 Harmony 格式，返回包含 channel 和 content 的列表；
-    否则，返回原内容并标记为默认 channel。
-    """
-    # 匹配 <|channel|>xxx<|message|>yyy<|end|>
     pattern = r"<\|channel\|>(.*?)<\|message\|>(.*?)<\|end|>"
     matches = re.findall(pattern, content, re.DOTALL)
 
     if not matches:
-        # 如果没有匹配到 Harmony 标签，视作普通文本
         return [{"channel": "text", "content": content}]
 
     results = []
@@ -324,22 +335,17 @@ def parse_harmony_message_content(content):
 
 
 def process_harmony_conversations(conversation):
-    """
-    处理传入的 list[list[dict]] 结构
-    """
     new_conversation = []
     for msg in conversation:
         role = msg.get("role")
         original_content = msg.get("content", "")
 
-        # 解析 content 中的 Harmony 结构
         segments = parse_harmony_message_content(original_content)
 
-        # 为每个解析出的通道生成一个新的消息字典
         for seg in segments:
             new_msg = {
                 "role": role,
-                "channel": seg["channel"],  # 新增字段标识通道
+                "channel": seg["channel"],
                 "content": seg["content"],
             }
             new_conversation.append(new_msg)