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,13 +393,11 @@ 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,

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -74,6 +74,7 @@
 import logging
 
 import torch
+import transformers
 import torch.nn.functional as F
 from transformers import PreTrainedModel
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config
@@ -100,6 +101,30 @@
 __all__ = ["HFDFlashModel"]
 
 
+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 +208,98 @@ 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 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``.
+
+        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.
+        """
+        if (
+            position_ids is not None
+            or getattr(self.config, "model_type", None) != "qwen3_vl"
+            or not transformers.__version__.startswith("5.3.")
+            # 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")
+        mm_token_type_ids = model_kwargs.get("mm_token_type_ids")
+        backbone = getattr(self, "model", None)
+        get_rope_index = getattr(backbone, "get_rope_index", None)
+        compute_position_ids = getattr(backbone, "compute_3d_position_ids", None)
+        if not isinstance(image_grid_thw, torch.Tensor) and not isinstance(
+            video_grid_thw, torch.Tensor
+        ):
+            return position_ids
+        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:
+            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 +709,15 @@ def forward(
         - Label alignment: position k predicts token at anchor+k
         - Optional loss decay weighting
         """
+        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 +764,49 @@ 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.
+            multimodal_keys = {
+                "pixel_values",
+                "pixel_values_videos",
+                "image_grid_thw",
+                "video_grid_thw",
+                "image_sizes",
+                "images",
+                "videos",
+            }
+            use_top_level_forward = any(kwargs.get(key) is not None for key in multimodal_keys)
             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:
+                    base_forward_kwargs = dict(kwargs)
+                    # Training-only data keys are not accepted by Hugging Face model forwards.
+                    base_forward_kwargs.pop("assistant_masks", None)
+                    base_forward_kwargs.pop("loss_mask", None)
+                    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 +815,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 +836,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

tests/unit/torch/export/test_hf_spec_rope_export.py

@@ -139,3 +139,16 @@ def test_dflash_rope_theta_inherits_base():
     """rope_theta is inherited from the target/base config (draft drafts for the base)."""
     config = _make_dflash_exporter(base_rope_theta=5000000.0)._export_config()
     assert config["rope_theta"] == 5000000.0
+
+
+def test_dflash_rope_theta_inherits_base_rope_parameters():
+    """Transformers 5 stores the target RoPE base in rope_parameters."""
+    exporter = _make_dflash_exporter(base_rope_theta=None)
+    exporter.model.config.rope_parameters = {
+        "rope_type": "default",
+        "rope_theta": 5000000.0,
+    }
+
+    config = exporter._export_config()
+
+    assert config["rope_theta"] == 5000000.0