dev

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Author: h-guo18

examples/speculative_decoding/main.py

@@ -170,7 +170,7 @@ def train():
         # To avoid OOM for large models, we load and convert model on CPU first.
         # Model will be moved to GPU during HF trainer.init().
         offline_kwargs = {"num_hidden_layers": 0} if use_offline_training else {}
-        model = transformers.Qwen3VLMoeForConditionalGeneration.from_pretrained(
+        model = transformers.Qwen3VLForConditionalGeneration.from_pretrained(
             model_args.model_name_or_path,
             torch_dtype="auto",
             device_map="cpu",

modelopt/torch/speculative/plugins/transformers.py

@@ -425,16 +425,23 @@ def _base_model_lm_head(self):
     @property
     def _base_llm_config(self):
         """Return the llm config for the base model, from LLM or VLM."""
-        return self.config.llm_config if hasattr(self.config, "llm_config") else self.config
+        # return self.config.llm_config if hasattr(self.config, "llm_config") else self.config
+        return self.config.text_config
 
     def _find_base_model_parts(self):
         """Find model parts from different models and set base_{part}_path attributes."""
         base_model_parts_mapping = {
-            "base_model_path": ["model", "backbone", "language_model.backbone"],
+            "base_model_path": [
+                "model.language_model",
+                "model",
+                "backbone",
+                "language_model.backbone",
+            ],
             "base_model_embeddings_path": [
                 "model.embed_tokens",
                 "backbone.embeddings",
                 "language_model.backbone.embeddings",
+                "model.language_model.embed_tokens",
             ],
             "base_model_lm_head_path": ["lm_head", "language_model.lm_head"],
         }
@@ -747,7 +754,8 @@ def _llm_or_vlm_embedding(self, input_ids, kwargs):
             del vit_embeds
             return tok_embeds.reshape(bs, seq_len, hid_size)
         else:
-            raise ValueError(f"VLM model type {self.config.model_type} not supported")
+            breakpoint()
+            # raise ValueError(f"VLM model type {self.config.model_type} not supported")
 
     def _base_model_forward(
         self,
@@ -769,6 +777,7 @@ def _base_model_forward(
                 **kwargs,
             )
             past_key_values = getattr(outputs, "past_key_values", None)
+            input_embeds = outputs.hidden_states[0]
             base_model_hidden_states = outputs.hidden_states[-1]
             base_model_logits = outputs.logits
 
@@ -780,7 +789,13 @@ def _base_model_forward(
                 labels = labels.view(-1)
                 base_model_loss = loss_fct(loss_logits, labels)
 
-        return base_model_hidden_states, base_model_logits, base_model_loss, past_key_values
+        return (
+            input_embeds,
+            base_model_hidden_states,
+            base_model_logits,
+            base_model_loss,
+            past_key_values,
+        )
 
     def _map_logits_to_draft_vocab(self, full_logits):
         reverse_mapping = (
@@ -872,16 +887,20 @@ def forward(
                 base_model_logits = self.lm_head(base_model_hidden_states)
             base_model_loss, past_key_values = None, None
         else:
-            base_model_hidden_states, base_model_logits, base_model_loss, past_key_values = (
-                self._base_model_forward(
-                    input_ids,
-                    attention_mask,
-                    position_ids,
-                    past_key_values,
-                    self.eagle_freeze_base_model,
-                    labels,
-                    **kwargs,
-                )
+            (
+                base_input_embeds,
+                base_model_hidden_states,
+                base_model_logits,
+                base_model_loss,
+                past_key_values,
+            ) = self._base_model_forward(
+                input_ids,
+                attention_mask,
+                position_ids,
+                past_key_values,
+                self.eagle_freeze_base_model,
+                labels,
+                **kwargs,
             )
 
         if not isinstance(past_key_values, Cache):
@@ -912,7 +931,8 @@ def forward(
             eagle_cache,
         )
         with torch.no_grad():
-            inputs_embeds = self._llm_or_vlm_embedding(eagle_input_ids, kwargs)
+            # inputs_embeds = self._llm_or_vlm_embedding(eagle_input_ids, kwargs)
+            inputs_embeds = base_input_embeds.roll(-1, 1)
 
         past_key_values.eagle_cache = eagle_cache
 

modelopt/torch/utils/plugins/transformers_dataset.py

@@ -17,6 +17,7 @@
 
 import copy
 import itertools
+import os
 
 import torch
 import transformers
@@ -44,8 +45,8 @@ def _sharegpt_to_openai_messages(conversations: list[dict]):
     }
     messages = []
     for msg in conversations:
-        role = role_mapping[msg["from"]]
-        content = msg["value"]
+        role = role_mapping[msg["role"]]
+        content = msg["content"]
         messages.append({"role": role, "content": content})
     return messages
 
@@ -225,7 +226,7 @@ def __init__(
         chat_template: str | None = None,
         add_generation_prompt: bool = False,
         answer_only_loss: bool = False,
-        local_image_path: str | None = None,
+        local_image_path: str = "",
         return_labels: bool = False,
     ):
         """Initialize the VisionLanguageDataset."""
@@ -242,8 +243,6 @@ def __init__(
         )
 
     def _process_multimodal_sample(self, examples):
-        print(examples)
-        breakpoint()
         tokenized_messages = self.processor.apply_chat_template(
             examples,
             tokenize=True,
@@ -279,9 +278,17 @@ def __call__(self, examples):
             for msg in copy_messages:
                 if isinstance(msg["content"], str):
                     msg["content"] = [{"type": "text", "text": msg["content"]}]
+
                 for ctn in msg["content"]:
                     if ctn["type"] == "image" and "image" in ctn:
-                        ctn["image"] = self.local_image_path + "/" + ctn["image"]
+                        ctn["image"] = os.path.abspath(
+                            os.path.join(self.local_image_path, ctn["image"])
+                        )
+                    # If any value in ctn is None, delete that key
+                    # HF dataloader add Nones to align keys. Leads to error in processor.
+                    keys_to_delete = [k for k, v in ctn.items() if v is None]
+                    for k in keys_to_delete:
+                        del ctn[k]
 
             batch.append(copy_messages)