From 128a6dc5b9cbce0efe82c1e91badd9d5641251aa Mon Sep 17 00:00:00 2001 From: Slawomir Kierat Date: Tue, 14 Jul 2026 02:54:16 -0700 Subject: [PATCH] feat(speculative): add Qwen3-VL support for DFlash training --- examples/speculative_decoding/eagle_utils.py | 3 + .../torch/export/plugins/hf_spec_export.py | 29 ++- .../torch/speculative/plugins/hf_dflash.py | 202 ++++++++++++++++-- .../torch/export/test_hf_spec_rope_export.py | 13 ++ .../speculative/plugins/test_hf_dflash.py | 81 +++++++ .../plugins/test_hf_speculative_offline.py | 34 +++ 6 files changed, 339 insertions(+), 23 deletions(-) diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index aaf52cca9b3..4b2e78f1c2e 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -129,6 +129,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: diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index a9f4f995540..7ab203f88d7 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/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, diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index a65d4d8a0ac..f8a44cadf86 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/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 @@ -95,6 +96,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: @@ -168,6 +193,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. @@ -531,6 +648,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( @@ -575,12 +701,49 @@ def forward( base_outputs.logits = self._base_model_lm_head(out_hiddens) 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] @@ -589,16 +752,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: @@ -611,8 +773,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 diff --git a/tests/unit/torch/export/test_hf_spec_rope_export.py b/tests/unit/torch/export/test_hf_spec_rope_export.py index 720082bc617..fbeb218793e 100644 --- a/tests/unit/torch/export/test_hf_spec_rope_export.py +++ b/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 diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index cdaba1766e3..4a01501a4ae 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -35,6 +35,7 @@ import modelopt.torch.opt as mto import modelopt.torch.speculative as mtsp +import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG from modelopt.torch.speculative.plugins.hf_dflash import ( DFlashAttention, @@ -119,6 +120,86 @@ def test_convert_sets_mask_token_id(self): assert model.mask_token_id == 0 +def test_qwen3_vl_transformers_53_position_ids_expand_video_grid(monkeypatch): + """Only mRoPE receives a per-frame video grid on Transformers 5.3.""" + original_grid = torch.tensor([[3, 4, 5], [2, 6, 7]]) + expected_position_ids = torch.ones(3, 1, 12, dtype=torch.long) + get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1))) + compute_position_ids = MagicMock() + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace( + get_rope_index=get_rope_index, + compute_3d_position_ids=compute_position_ids, + ), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + position_ids = HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": original_grid, + "mm_token_type_ids": torch.zeros(1, 12, dtype=torch.long), + }, + ) + + assert position_ids is expected_position_ids + assert not compute_position_ids.called + assert torch.equal(original_grid, torch.tensor([[3, 4, 5], [2, 6, 7]])) + assert torch.equal( + get_rope_index.call_args.kwargs["video_grid_thw"], + torch.tensor([[1, 4, 5], [1, 4, 5], [1, 4, 5], [1, 6, 7], [1, 6, 7]]), + ) + + +def test_qwen3_vl_transformers_53_position_ids_require_mm_token_types(monkeypatch): + """Never silently fall back to one-dimensional positions for a visual batch.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="mm_token_type_ids"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={"image_grid_thw": torch.tensor([[1, 4, 4]])}, + ) + + +def test_qwen3_vl_transformers_53_position_ids_reject_bad_mm_token_shape(monkeypatch): + """Keep processor-produced modality ids aligned with the padded text sequence.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="same shape as input_ids"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "image_grid_thw": torch.tensor([[1, 4, 4]]), + "mm_token_type_ids": torch.zeros(1, 11, dtype=torch.long), + }, + ) + + class TestDPaceWeights: """Test the D-PACE position-weighting objective (arXiv:2605.18810).""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py index d6f087d3771..db93583b041 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py +++ b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py @@ -57,6 +57,40 @@ make_speculative_data_module = _eagle_utils.make_speculative_data_module +# --------------------------------------------------------------------------- +# online VLM data-module wiring +# --------------------------------------------------------------------------- + + +def test_vlm_data_module_passes_dflash_label_mode(monkeypatch): + """VLM batches must use unshifted labels for DFlash and preserve OSL settings.""" + data_args = argparse.Namespace( + mode="online", + data_path="unused.jsonl", + vlm_processor="dummy-vlm-processor", + vlm_img_dir="/images", + chat_template=None, + ) + collator = MagicMock() + monkeypatch.setattr(_eagle_utils, "ShardedDataset", MagicMock()) + monkeypatch.setattr(_eagle_utils, "VisionLanguageDataCollator", collator) + + module = make_speculative_data_module( + MagicMock(), data_args, train_len=16, answer_only_loss=True, shift_labels=False + ) + + collator.assert_called_once_with( + processor="dummy-vlm-processor", + train_len=16, + local_image_path="/images", + return_labels=True, + answer_only_loss=True, + shift_labels=False, + chat_template=None, + ) + assert module["data_collator"] is collator.return_value + + # --------------------------------------------------------------------------- # sample_size truncation tests # ---------------------------------------------------------------------------