add Qwen3-VL support for DFlash training

examples/speculative_decoding/eagle_utils.py

@@ -141,6 +141,9 @@ def make_speculative_data_module(
                 train_len=train_len,
                 local_image_path=data_args.vlm_img_dir,
                 return_labels=True,
+                answer_only_loss=answer_only_loss,
+                shift_labels=shift_labels,
+                chat_template=chat_template,
             )
 
     else:

modelopt/torch/export/plugins/hf_spec_export.py

@@ -29,6 +29,23 @@
 
 ALL_SPEC_MODES = ["eagle", "dflash"]
 
+
+def _get_rope_theta(config, default=None):
+    """Get RoPE theta from either legacy or Transformers 5 config fields."""
+    rope_theta = getattr(config, "rope_theta", None)
+    if rope_theta is not None:
+        return rope_theta
+
+    # Transformers 5 stores this under rope_parameters (and exposes the same
+    # data through rope_scaling for backwards compatibility).
+    for attr in ("rope_parameters", "rope_scaling"):
+        rope_config = getattr(config, attr, None)
+        if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None:
+            return rope_config["rope_theta"]
+
+    return default
+
+
 LLAMA_EAGLE_SINGLE_LAYER = {
     "required": {
         "layers.0.self_attn.q_proj",
@@ -376,14 +393,10 @@ def _export_config(self):
             "initializer_range": getattr(base_config, "initializer_range", 0.02),
             "attention_bias": getattr(draft_config, "attention_bias", False),
             "attention_dropout": getattr(draft_config, "attention_dropout", 0.0),
-            # Inherit the target's rope_theta: DFlash injects the target's KV into every
-            # draft layer, so the draft's RoPE base must match the target's. (The draft
-            # arch config carries no rope_theta of its own.)
-            "rope_theta": (
-                getattr(base_config, "rope_theta", None)
-                if getattr(base_config, "rope_theta", None) is not None
-                else getattr(draft_config, "rope_theta", 1000000.0)
-            ),
+            # Inherit the target's RoPE base: DFlash injects target KV into every draft
+            # layer, so their RoPE bases must match. Transformers 5 stores rope_theta
+            # in rope_parameters rather than a top-level config attribute.
+            "rope_theta": _get_rope_theta(base_config, _get_rope_theta(draft_config, 1000000.0)),
             # YaRN long-context scaling is injected below (see the rope_scaling block).
             "rope_scaling": None,
             "tie_word_embeddings": False,

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -72,9 +72,11 @@
 """
 
 import logging
+from typing import Any
 
 import torch
 import torch.nn.functional as F
+import transformers
 from transformers import PreTrainedModel
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config
 from transformers.trainer_pt_utils import LabelSmoother
@@ -100,6 +102,54 @@
 __all__ = ["HFDFlashModel"]
 
 
+_QWEN3_VL_MROPE_WORKAROUND_VERSION = "5.3.0"
+_MULTIMODAL_FORWARD_KWARGS = frozenset(
+    {
+        "pixel_values",
+        "pixel_values_videos",
+        "image_grid_thw",
+        "video_grid_thw",
+        "mm_token_type_ids",
+        "image_sizes",
+        "images",
+        "videos",
+    }
+)
+
+
+def _multimodal_forward_kwargs(model_kwargs: dict) -> dict:
+    """Return collator fields accepted by Hugging Face multimodal forwards."""
+    return {
+        name: value
+        for name, value in model_kwargs.items()
+        if name in _MULTIMODAL_FORWARD_KWARGS and value is not None
+    }
+
+
+def _expand_qwen3_video_grid_thw(video_grid_thw: torch.Tensor) -> torch.Tensor:
+    """Return the per-frame video grid representation used by Qwen3-VL RoPE.
+
+    Qwen3-VL's video processor emits one ``[T, H, W]`` row per source video, but
+    its rendered prompt contains a separate visual-token group for every temporal
+    frame.  Transformers 5.3's ``get_rope_index`` consumes one grid row per
+    rendered group, while the vision encoder still requires the original one-row-
+    per-video representation.  This helper is therefore used *only* for mRoPE
+    position construction; callers must keep the original tensor for the model
+    forward.
+    """
+    if video_grid_thw.ndim != 2 or video_grid_thw.shape[-1] != 3:
+        raise ValueError(
+            "Qwen3-VL video_grid_thw must have shape [num_videos, 3], got "
+            f"{tuple(video_grid_thw.shape)}."
+        )
+    if torch.any(video_grid_thw[:, 0] <= 0):
+        raise ValueError("Qwen3-VL video_grid_thw temporal lengths must be positive.")
+
+    expanded_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+    expanded_grid_thw[:, 0] = 1
+    return expanded_grid_thw
+
+
 def _dpace_position_weights(
     confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None
 ) -> torch.Tensor:
@@ -183,6 +233,125 @@ def _base_llm_config(self):
             or self.config
         )
 
+    def _qwen3_vl_position_ids(
+        self,
+        input_ids,
+        attention_mask,
+        position_ids,
+        past_key_values,
+        inputs_embeds,
+        model_kwargs,
+    ):
+        """Precompute Qwen3-VL mRoPE positions for Transformers 5.3.0 batches.
+
+        The video encoder consumes one grid row per source video, whereas mRoPE
+        consumes one row per rendered temporal-frame group.  Calling the
+        top-level model with the original video grid makes the two contracts
+        conflict.  Construct the mRoPE positions with a frame-expanded copy,
+        then pass the original grid to the vision encoder in ``forward``.
+
+        Transformers 5.4.0 performs this frame expansion in ``get_rope_index``
+        itself; only 5.3.0 needs the external workaround. See
+        https://github.com/huggingface/transformers/blob/v5.4.0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py
+
+        Prefer ``get_rope_index`` over ``compute_3d_position_ids``. The latter
+        writes ``rope_deltas`` into the base model even though DFlash training
+        never supplies a cache.  Keeping this calculation side-effect free is
+        important when the frozen target is reused for consecutive training
+        batches or validation.
+        """
+        model_type = str(getattr(self.config, "model_type", ""))
+        if (
+            position_ids is not None
+            or not model_type.startswith("qwen3_vl")
+            # Cached decoding uses the base model's rope_deltas path.  DFlash
+            # training has no cache and is the only path that needs the
+            # frame-expanded construction below.
+            or past_key_values is not None
+        ):
+            return position_ids
+
+        image_grid_thw = model_kwargs.get("image_grid_thw")
+        video_grid_thw = model_kwargs.get("video_grid_thw")
+        if not isinstance(image_grid_thw, torch.Tensor) and not isinstance(
+            video_grid_thw, torch.Tensor
+        ):
+            return position_ids
+
+        if transformers.__version__ != _QWEN3_VL_MROPE_WORKAROUND_VERSION:
+            if transformers.__version__.startswith("5.3."):
+                raise RuntimeError(
+                    "Qwen3-VL DFlash mRoPE supports Transformers 5.3.0 or >=5.4.0; "
+                    f"got {transformers.__version__}. A 5.3.x patch release may already "
+                    "expand video_grid_thw internally."
+                )
+            return position_ids
+
+        mm_token_type_ids = model_kwargs.get("mm_token_type_ids")
+        backbone = getattr(self, "model", None)
+        # Probed dynamically: which one exists depends on the Transformers version.
+        get_rope_index: Any = getattr(backbone, "get_rope_index", None)
+        compute_position_ids: Any = getattr(backbone, "compute_3d_position_ids", None)
+        if (
+            not isinstance(mm_token_type_ids, torch.Tensor)
+            or input_ids is None
+            or (not callable(get_rope_index) and not callable(compute_position_ids))
+        ):
+            raise ValueError(
+                "Qwen3-VL DFlash training requires input_ids, mm_token_type_ids, and "
+                "a Qwen3-VL model with get_rope_index or compute_3d_position_ids. "
+                "Use the Qwen3-VL AutoProcessor without dropping mm_token_type_ids."
+            )
+
+        if mm_token_type_ids.shape != input_ids.shape:
+            raise ValueError(
+                "Qwen3-VL mm_token_type_ids must have the same shape as input_ids, got "
+                f"{tuple(mm_token_type_ids.shape)} and {tuple(input_ids.shape)}."
+            )
+
+        rope_video_grid_thw = video_grid_thw
+        if isinstance(video_grid_thw, torch.Tensor) and video_grid_thw.numel() > 0:
+            video_token_mask = mm_token_type_ids == 2
+            if isinstance(attention_mask, torch.Tensor):
+                video_token_mask = video_token_mask & attention_mask.bool()
+            video_group_starts = video_token_mask.clone()
+            video_group_starts[:, 1:] &= ~video_token_mask[:, :-1]
+            expected_video_groups = int(video_grid_thw[:, 0].sum())
+            actual_video_groups = int(video_group_starts.sum())
+            if actual_video_groups != expected_video_groups:
+                raise ValueError(
+                    "Qwen3-VL video frame groups do not match video_grid_thw: "
+                    f"expected {expected_video_groups}, found {actual_video_groups}."
+                )
+            rope_video_grid_thw = _expand_qwen3_video_grid_thw(video_grid_thw)
+
+        rope_kwargs = {
+            "input_ids": input_ids,
+            "image_grid_thw": image_grid_thw,
+            "video_grid_thw": rope_video_grid_thw,
+            "attention_mask": attention_mask,
+            "mm_token_type_ids": mm_token_type_ids,
+        }
+        if callable(get_rope_index):
+            position_ids, _ = get_rope_index(**rope_kwargs)
+        else:
+            position_ids = compute_position_ids(
+                **rope_kwargs,
+                inputs_embeds=inputs_embeds,
+                past_key_values=past_key_values,
+            )
+
+        expected_shape = (3, *input_ids.shape)
+        valid_position_ids = isinstance(position_ids, torch.Tensor) and (
+            tuple(position_ids.shape) == expected_shape
+        )
+        if not valid_position_ids:
+            raise RuntimeError(
+                "Qwen3-VL produced invalid mRoPE position ids: expected shape "
+                f"{expected_shape}, got {getattr(position_ids, 'shape', None)}."
+            )
+        return position_ids
+
     def _find_base_model_parts(self):
         """Locate base model submodules (backbone, embeddings, lm_head) by probing known paths.
 
@@ -592,6 +761,16 @@ def forward(
         - Label alignment: position k predicts token at anchor+k
         - Optional loss decay weighting
         """
+        if self.training:
+            position_ids = self._qwen3_vl_position_ids(
+                input_ids,
+                attention_mask,
+                position_ids,
+                past_key_values,
+                inputs_embeds,
+                kwargs,
+            )
+
         if not self.training:
             if self.dflash_offline:
                 raise RuntimeError(
@@ -638,12 +817,37 @@ def forward(
             )
             target_hidden = base_outputs.target_hidden
         else:
-            # TODO: For co-training the base model, remove no_grad and eval() switch.
+            # Multimodal models need the top-level conditional-generation forward so their
+            # image/video features are inserted before the language model runs.  Keep the
+            # long-standing narrow call for text-only models.
+            base_forward_kwargs = _multimodal_forward_kwargs(kwargs)
+            use_top_level_forward = bool(base_forward_kwargs)
             with torch.no_grad():
-                raw_outputs = super().forward(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
+                if use_top_level_forward:
+                    raw_outputs = super().forward(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        position_ids=position_ids,
+                        past_key_values=past_key_values,
+                        inputs_embeds=inputs_embeds,
+                        use_cache=False,
+                        output_attentions=output_attentions,
+                        output_hidden_states=True,
+                        cache_position=cache_position,
+                        return_dict=True,
+                        **base_forward_kwargs,
+                    )
+                else:
+                    raw_outputs = super().forward(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        output_hidden_states=True,
+                    )
+
+            if not getattr(raw_outputs, "hidden_states", None):
+                raise RuntimeError(
+                    "The base model did not return hidden states required for DFlash training. "
+                    "Ensure its top-level multimodal forward supports output_hidden_states=True."
                 )
             offset = 1
             selected = [raw_outputs.hidden_states[lid + offset] for lid in self.target_layer_ids]
@@ -652,16 +856,15 @@ def forward(
                 target_hidden=target_hidden, logits=raw_outputs.logits
             )
 
-        # 2. Build loss mask.
-        # When labels are provided (answer_only_loss), they already encode both
-        # assistant masking and padding (-100 for both). When labels are not
-        # provided, fall back to attention_mask for padding only.
+        # 2. Build loss mask. Labels carry optional answer-only masking, but do
+        # not in general mark padded tokens with -100 (the VLM collator creates
+        # them from padded input_ids). Always intersect with attention_mask so
+        # anchor sampling and loss never include the padded tail.
+        loss_mask = torch.ones(bsz, seq_len, device=device)
         if labels is not None:
-            loss_mask = (labels != LabelSmoother.ignore_index).float()
-        elif attention_mask is not None:
-            loss_mask = attention_mask.float()
-        else:
-            loss_mask = torch.ones(bsz, seq_len, device=device)
+            loss_mask = loss_mask * (labels != LabelSmoother.ignore_index).float()
+        if attention_mask is not None:
+            loss_mask = loss_mask * attention_mask.float()
 
         # In offline training, assistant mask is dumped and passed as kwarg.
         if kwargs.get("loss_mask") is not None:
@@ -674,8 +877,16 @@ def forward(
         n_blocks = anchor_positions.shape[1]
 
         if n_blocks == 0 or not block_keep_mask.any():
-            # Zero loss that still flows through dflash_module for DDP gradient sync
-            dummy = self.dflash_module.fc.weight.sum() * 0.0
+            # Keep all trainable draft parameters in the graph so DDP can reduce a rank
+            # that receives an all-masked answer-only batch.
+            dummy = sum(
+                (
+                    parameter.reshape(-1)[0] * 0.0
+                    for parameter in self.dflash_module.parameters()
+                    if parameter.requires_grad
+                ),
+                torch.zeros((), device=device),
+            )
             return ModelOutput(loss=dummy, logits=base_outputs.logits, train_acc=[[0.0]])
 
         # 4. Build draft inputs

modelopt/torch/speculative/utils.py

@@ -610,7 +610,13 @@ def load_vlm_or_llm(
         return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code)
 
     if _is_vlm:
-        model_cls = transformers.AutoModelForVision2Seq
+        # Transformers 5 renamed AutoModelForVision2Seq to
+        # AutoModelForImageTextToText.  Prefer the pre-5 name so this loader
+        # continues to support the Transformers 4 environments used by older
+        # speculative-decoding jobs.
+        model_cls = getattr(transformers, "AutoModelForVision2Seq", None)
+        if model_cls is None:
+            model_cls = transformers.AutoModelForImageTextToText
     else:
         model_cls = transformers.AutoModelForCausalLM