add Qwen3-VL support for DFlash training - #1975
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe changes update DFlash Qwen3-VL multimodal training, Transformers 5 compatibility, RoPE export configuration, and vision-language data collation. Tests cover version-specific position IDs, multimodal forwarding, label handling, nested RoPE parameters, and VLM model loading fallback behavior. ChangesDFlash and RoPE updates
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant HFDFlashModel
participant Qwen3VLModel
participant get_rope_index
HFDFlashModel->>Qwen3VLModel: recompute position_ids from multimodal inputs
Qwen3VLModel->>get_rope_index: pass expanded video_grid_thw and mm_token_type_ids
get_rope_index-->>HFDFlashModel: return position_ids
HFDFlashModel->>Qwen3VLModel: forward multimodal inputs
Qwen3VLModel-->>HFDFlashModel: return hidden_states
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — 1 CRITICAL, 0 IMPORTANT, 0 SUGGESTION.
Most impactful finding
[CRITICAL Algorithm] shift_labels crashes the new VLM online path (examples/speculative_decoding/eagle_utils.py:145)
The PR passes shift_labels=shift_labels into VisionLanguageDataCollator(...), but that class (modelopt/torch/utils/plugins/transformers_dataset.py:321) declares (processor, train_len, chat_template, add_generation_prompt, answer_only_loss, local_image_path, return_labels) — no shift_labels parameter and no **kwargs. Any real Qwen3-VL online run therefore fails immediately with a TypeError about an unexpected keyword argument shift_labels.
DFlash specifically requires shift_labels=False (unshifted labels), so this is exactly the value the feature must pass — the crash blocks the PR headline capability end-to-end. The text-only branch (LanguageDataCollator, line 135) is unaffected because that class does declare shift_labels.
The new test_vlm_data_module_passes_dflash_label_mode does not catch this: it swaps VisionLanguageDataCollator for a MagicMock, which accepts arbitrary kwargs and hides the real signature mismatch.
Fix: add shift_labels to VisionLanguageDataCollator.init and forward it to super().init(...) (the parent already stores/uses it). Consider exercising the real collator signature in a test.
Assessment
The RoPE-theta extraction, mRoPE frame-expansion, top-level multimodal forward, loss-mask intersection with attention_mask, and the all-parameter DDP dummy-loss change all look correct and well-reasoned. Risk is concentrated in the single collator signature mismatch, which is a hard crash but a small, contained fix. The PR is also explicitly marked draft pending an e2e smoke test and regression coverage. Blocking on the one CRITICAL.
| local_image_path=data_args.vlm_img_dir, | ||
| return_labels=True, | ||
| answer_only_loss=answer_only_loss, | ||
| shift_labels=shift_labels, |
There was a problem hiding this comment.
[CRITICAL Algorithm] VisionLanguageDataCollator does not accept shift_labels, so this call raises TypeError: __init__() got an unexpected keyword argument 'shift_labels' at runtime — crashing the exact VLM online path this PR adds.
VisionLanguageDataCollator.__init__ (modelopt/torch/utils/plugins/transformers_dataset.py:321) is declared as (self, processor, train_len, chat_template, add_generation_prompt, answer_only_loss, local_image_path, return_labels) — it has neither a shift_labels parameter nor **kwargs, and it never forwards shift_labels to super().__init__. The text-only branch above (LanguageDataCollator, line 135) is fine because that class does declare shift_labels, but the VLM subclass does not.
This isn't caught by test_vlm_data_module_passes_dflash_label_mode because that test replaces VisionLanguageDataCollator with a MagicMock, which accepts any kwargs — so the real signature mismatch is masked.
Impact: DFlash for Qwen3-VL requires shift_labels=False, and passing it is the only way to get the unshifted-label behavior DFlash needs — the crash blocks the feature entirely for any real run.
Fix: add shift_labels to VisionLanguageDataCollator.__init__ and forward it to super().__init__(...) (the parent already stores and uses it). For example, in transformers_dataset.py:
def __init__(
self,
processor: str,
train_len: int = 8192,
chat_template: str | None = None,
add_generation_prompt: bool = False,
answer_only_loss: bool = False,
shift_labels: bool = True,
local_image_path: str = "",
return_labels: bool = False,
):
...
super().__init__(
tokenizer=self.processor.tokenizer,
train_len=train_len,
chat_template=chat_template,
add_generation_prompt=add_generation_prompt,
answer_only_loss=answer_only_loss,
shift_labels=shift_labels,
return_labels=return_labels,
)Consider having the test construct the real collator (or assert against its actual signature) so this class of mismatch is caught.
| if ( | ||
| position_ids is not None | ||
| or getattr(self.config, "model_type", None) != "qwen3_vl" | ||
| or not transformers.__version__.startswith("5.3.") |
There was a problem hiding this comment.
question: is this patch only needed for transformers 5.3, or for all transformers 5.x or >5.3?
There was a problem hiding this comment.
Only for Transformers 5.3.x.
Starting with 5.4.0, upstream added that same repeat_interleave(...); grid[:, 0] = 1 expansion internally in get_rope_index(), and it remains in 5.12.
h-guo18
left a comment
h-guo18
left a comment
There was a problem hiding this comment.
[AI review] Supplementing the earlier claude[bot] review (which flagged the shift_labels TypeError): a few remaining issues, details inline.
Also noting a test-coverage gap for the pending regression work: nothing currently exercises the multimodal top-level-forward branch, the loss-mask ∩ attention_mask change, or the all-parameter DDP dummy-loss path.
| local_image_path=data_args.vlm_img_dir, | ||
| return_labels=True, | ||
| answer_only_loss=answer_only_loss, | ||
| shift_labels=shift_labels, |
There was a problem hiding this comment.
[AI review] The PR description says VisionLanguageDataCollator was extended (propagating answer_only_loss/chat-template/label-alignment settings, VLM_MIN_PIXELS/VLM_MAX_PIXELS limits, ChatML-boundary assistant masks, fixed training_seq_len enforcement), but modelopt/torch/utils/plugins/transformers_dataset.py is not part of this diff — on the current head the class still has its old signature. It looks like that file may not have been committed/pushed.
Note this is broader than the shift_labels TypeError already flagged: even after adding that one parameter, the described mask/pixel/seq-len behaviors would still be missing from the PR.
| if ( | ||
| position_ids is not None | ||
| or getattr(self.config, "model_type", None) != "qwen3_vl" | ||
| or not transformers.__version__.startswith("5.3.") |
There was a problem hiding this comment.
[AI review] Two robustness concerns with this gate:
-
Version pin. Tracing transformers 5.3.0, the per-video vs per-frame grid mismatch looks like a 5.3 bug rather than a stable contract:
get_rope_indexconsumes one grid row per rendered visual group (next(grid_iters[...])), while the processor renders T timestamp-separated frame groups per video and the video processor emits one[T, H, W]row per video. Even the vanilla generation path (_prepare_position_ids_for_generation) passes the raw grid, so upstream video inference appears equally broken on 5.3 — meaning upstream will likely fix this insideget_rope_index(as in 4.x, where the expansion was internal). If that fix lands in a 5.3.x patch release, this external expansion becomes a double expansion and silently corrupts positions; on 5.4+ the helper silently disables itself and multimodal batches fall back to the broken internal path. Suggest linking the upstream issue/commit here, and failing loudly (frame-group count ≠ grid-row count) instead of silently returningNonewhen the gate doesn't match. -
model_type != "qwen3_vl"exact match. Family variants likeqwen3_vl_moeskip the precomputation but still enter the multimodal top-level forward below — landing exactly in the broken path this helper works around. Consider covering the family or at least warning.
| - Label alignment: position k predicts token at anchor+k | ||
| - Optional loss decay weighting | ||
| """ | ||
| position_ids = self._qwen3_vl_position_ids( |
There was a problem hiding this comment.
[AI review] This runs before the if not self.training branch, so it also fires on eval prefill (multimodal inputs, no cache). Passing explicit position_ids makes Qwen3VLModel.forward skip compute_3d_position_ids, so self.rope_deltas is never established; a subsequent cached decode step (position_ids=None, past_key_values not None → this helper returns None) then falls into the rope_deltas-based branch with a stale/None delta and computes wrong positions.
Since the docstring says this construction is only needed for DFlash training, consider gating on self.training.
| attention_mask=attention_mask, | ||
| output_hidden_states=True, | ||
| if use_top_level_forward: | ||
| base_forward_kwargs = dict(kwargs) |
There was a problem hiding this comment.
[AI review] base_forward_kwargs forwards everything except assistant_masks/loss_mask. Anything else that reaches forward's **kwargs — e.g. Trainer-injected num_items_in_batch (Trainer detects that this forward accepts **kwargs and adds loss kwargs), or stray dataset columns — gets passed to the HF multimodal forward and can raise TypeError or be silently swallowed depending on version. The text-only branch below is immune because it forwards a fixed argument set.
Consider an allowlist (the multimodal keys + known model kwargs) or filtering against the base forward's signature.
| import logging | ||
|
|
||
| import torch | ||
| import transformers |
There was a problem hiding this comment.
[AI review] nit: import transformers lands between import torch and import torch.nn.functional as F; ruff/isort (I001) will flag this in pre-commit/CI.
|
/claude review |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/unit/torch/speculative/plugins/test_hf_dflash.py (1)
123-340: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSolid, well-verified coverage of the Qwen3-VL mRoPE workaround.
Cross-checked each assertion (video-frame-group expansion math, patch-release
RuntimeError, 5.4.0 native delegation, malformedmm_token_type_idsvalidation,_multimodal_forward_kwargsfiltering, eval-time no-op) against the_qwen3_vl_position_ids/_multimodal_forward_kwargsimplementation and all are consistent.One minor nit: the
fake_model = SimpleNamespace(config=SimpleNamespace(model_type="qwen3_vl"), model=SimpleNamespace(get_rope_index=MagicMock(...)))construction is repeated near-verbatim across 7 test functions (lines 129-135, 166-169, 194-197, 218-221, 243-246, 300-303, 320-323). Consider extracting a smallfake_qwen3_vl_model(get_rope_index=None, model_type="qwen3_vl")helper/fixture to reduce duplication.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/torch/speculative/plugins/test_hf_dflash.py` around lines 123 - 340, Extract the repeated Qwen3-VL SimpleNamespace construction from the affected tests into a small fake_qwen3_vl_model helper or fixture, accepting optional get_rope_index and model_type arguments. Replace each duplicated fake_model setup with this helper while preserving each test’s configured model type and mocked rope method behavior.tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py (1)
95-111: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winTest only checks attribute assignment, not the actual unshifted-label behavior it claims to validate.
The docstring says "must support DFlash's unshifted labels," but the assertion only checks
collator.shift_labels is False— it never calls the collator on sample messages and verifies the resultinglabelsactually equalinput_ids(unshifted), as opposed to the shifted default. As per path instructions, "Tests must exercise the behavior they claim to validate... include an end-to-end test using the real implementation rather than replacing it with a fake."Consider extending this test (with
return_labels=Trueand a minimal message list plus a trivial/mockapply_chat_template) to assertlabelsequalsinput_idsfor the unshifted path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py` around lines 95 - 111, The test test_vlm_data_collator_accepts_unshifted_labels only verifies configuration, not label generation. Invoke the real VisionLanguageDataCollator with return_labels=True using minimal messages and a trivial apply_chat_template setup, then assert the produced labels equal input_ids when shift_labels=False; retain the existing processor setup without replacing the collator implementation.Source: Path instructions
modelopt/torch/utils/plugins/transformers_dataset.py (1)
433-474: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNative
{% generation %}fallback path can silently zero the entire loss for multimodal batches on transformers 5.3.0.When
derive_masks_from_markersisFalse(i.e. the chat template uses HF's native{% generation %}tags rather than ChatML markers),return_assistant_tokens_mask=Trueis passed straight toself.processor.apply_chat_template. There is an open transformers bug (issue#44521, "transformers==5.3.0") whereapply_chat_templatereturns all-zeroassistant_masksfor multimodal inputs because extra placeholder tokens such as<|image_pad|>are inserted later by the processor/tokenizer path, so offset_mapping corresponds to the expanded multimodal text while generation_indices still refers to the pre-expanded text, making the assistant span lookup fail so assistant_masks ends up all zeros.With this collator's current handling, an all-zero
assistant_maskssilently falls intolabels[:] = IGNORE_TOKEN_ID— masking the entire batch's loss with no warning, indistinguishable from the legitimate "fully truncated" case. For a training run with real multimodal data and a native-tag template on an affected transformers version, this would silently produce zero-loss batches instead of surfacing the misconfiguration.Consider at least logging a distinct warning (or raising) when
assistant_masksis all-zero but the batch clearly contains assistant turns, so this doesn't fail silently in production training.🐛 Example: warn distinctly when native tag mask looks suspiciously empty
assistant_mask = tokenized_messages["assistant_masks"] if not isinstance(assistant_mask, torch.Tensor) or not assistant_mask.any(): + print_rank_0( + "WARNING: assistant_masks is empty for a multimodal batch; masking all " + "loss. If this is unexpected, check for HF apply_chat_template offset " + "misalignment with multimodal placeholder tokens (transformers>=5.3.0)." + ) labels[:] = IGNORE_TOKEN_IDAdditionally, this label-construction block (lines 451-472) duplicates the logic in
LanguageDataCollator._process_chat_sample(lines 236-265) almost verbatim (shift vs. unshifted labels, assistant-mask application, all-masked fallback). Consider extracting a shared helper (e.g._labels_from_assistant_mask(input_ids, assistant_mask, shift_labels)) on the base class to avoid the two copies drifting apart.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/utils/plugins/transformers_dataset.py` around lines 433 - 474, Update _process_multimodal_sample to detect an all-zero assistant_masks result when the batch contains assistant turns, and emit a distinct warning or raise instead of silently treating it as legitimate full truncation. Extract the duplicated label construction into a shared base-class helper such as _labels_from_assistant_mask, including shift/unshifted handling and all-masked behavior, then reuse it from both _process_multimodal_sample and LanguageDataCollator._process_chat_sample.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/torch/speculative/plugins/test_fakebase.py`:
- Around line 155-156: Update the transformer auto-class setup in the test
around load_vlm_or_llm() to assign transformers.AutoModelForVision2Seq to None
instead of deleting the attribute, while preserving the _FakeVLM assignment for
AutoModelForImageTextToText. This must prevent the lazy lookup from recreating
the legacy class and ensure the fake model is used deterministically.
---
Nitpick comments:
In `@modelopt/torch/utils/plugins/transformers_dataset.py`:
- Around line 433-474: Update _process_multimodal_sample to detect an all-zero
assistant_masks result when the batch contains assistant turns, and emit a
distinct warning or raise instead of silently treating it as legitimate full
truncation. Extract the duplicated label construction into a shared base-class
helper such as _labels_from_assistant_mask, including shift/unshifted handling
and all-masked behavior, then reuse it from both _process_multimodal_sample and
LanguageDataCollator._process_chat_sample.
In `@tests/unit/torch/speculative/plugins/test_hf_dflash.py`:
- Around line 123-340: Extract the repeated Qwen3-VL SimpleNamespace
construction from the affected tests into a small fake_qwen3_vl_model helper or
fixture, accepting optional get_rope_index and model_type arguments. Replace
each duplicated fake_model setup with this helper while preserving each test’s
configured model type and mocked rope method behavior.
In `@tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py`:
- Around line 95-111: The test test_vlm_data_collator_accepts_unshifted_labels
only verifies configuration, not label generation. Invoke the real
VisionLanguageDataCollator with return_labels=True using minimal messages and a
trivial apply_chat_template setup, then assert the produced labels equal
input_ids when shift_labels=False; retain the existing processor setup without
replacing the collator implementation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cc216def-bb97-4a19-8858-0bcc96f2b68b
📒 Files selected for processing (6)
modelopt/torch/speculative/plugins/hf_dflash.pymodelopt/torch/speculative/utils.pymodelopt/torch/utils/plugins/transformers_dataset.pytests/unit/torch/speculative/plugins/test_fakebase.pytests/unit/torch/speculative/plugins/test_hf_dflash.pytests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/speculative/plugins/hf_dflash.py
4db6e60 to
76ed5d1
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1975 +/- ##
===========================================
+ Coverage 66.83% 77.77% +10.94%
===========================================
Files 519 519
Lines 58916 59720 +804
===========================================
+ Hits 39376 46449 +7073
+ Misses 19540 13271 -6269
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
### What does this PR do? Type of change: Bug fix (CI) The `deploy-preview` job in the `Docs` workflow fails on **every pull request opened from a fork**, which blocks merging for all external contributors. **Root cause.** `deploy-preview` runs `rossjrw/pr-preview-action@v1`, which pushes the built HTML to the `gh-pages` branch. The workflow declares `permissions: contents: write`, but for a `pull_request` event originating from a forked repository GitHub caps the `GITHUB_TOKEN` at **read-only** — the `permissions:` block cannot elevate above that cap. The push is therefore rejected: ``` remote: Permission to NVIDIA/Model-Optimizer.git denied to github-actions[bot]. fatal: unable to access 'https://github.com/NVIDIA/Model-Optimizer.git/': The requested URL returned error: 403 ``` The job's `if:` condition gated on `github.event_name`, `github.event.action` and the `changes` path filter, but never on whether the PR came from a fork — so it always ran and always failed. Because the `changes` filter matches `docs/**`, `modelopt/**` and `.github/workflows/pages.yml`, essentially any substantive fork PR trips this. **Fix.** Restrict `deploy-preview` to PRs whose head branch lives in this repository: ```yaml github.event.pull_request.head.repo.full_name == github.repository ``` A skipped job is not a failed job, so fork PRs are no longer blocked by it. ### Usage N/A — CI-only change. ### Testing Behaviour by scenario: | Scenario | Before | After | | --- | --- | --- | | PR from a branch in this repo | preview deployed | preview deployed (**unchanged**) | | PR from a fork | ❌ fails with 403 | ⏭️ skipped | | Fork deleted (`head.repo` is `null`) | ❌ fails | ⏭️ skipped | - Confirmed against workflow history: recent `Docs` runs on in-repo branches (`main`, `chenjiel/nvfp4-act-headroom`, `mxin/qad-skill`, `haoguo/dspark-ptq-script`) all succeed, while fork-branch runs fail with the 403 above. - `build-docs` was already passing on the affected PRs — only the deploy step failed, so documentation builds are unaffected either way. - YAML parses; `pre-commit run --files .github/workflows/pages.yml` passes. (`yamlfmt` excludes `^.github/workflows/`, so this file is not auto-formatted.) - This PR edits `.github/workflows/pages.yml`, which is itself in the `changes` filter, so it exercises `deploy-preview` on the in-repo path — the preview deploy on this PR passing is a self-check that the unchanged path still works. The `closed` cleanup path is gated by the same condition. That is intentional: a fork PR never deployed a preview directory, so there is nothing to remove. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A — workflow-condition change; not unit-testable - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — CI infrastructure, not user-facing - Did you get Claude approval on this PR?: ❌ — not yet run ### Additional Information Currently blocking #1975 (`add Qwen3-VL support for DFlash training`), which is approved with every other check green and sits at `mergeStateStatus: BLOCKED` solely because of this job. Note that a PR only picks up this fix once its branch contains it, since workflows run from the PR branch's own definitions. A follow-up option, if doc previews for external contributors are wanted: build in the `pull_request` workflow and deploy from a separate `workflow_run`-triggered workflow, which executes in the base-repo context and does get a write token. Deliberately not using `pull_request_target` here — that would run unreviewed PR code with write permissions. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
9f59816 to
19c23eb
Compare
3 participants
What does this PR do?
Type of change: new feature
Adds online DFlash training support for Qwen3-VL–style vision-language models.
Changes include:
AutoModelForImageTextToTextAPI, while retaining compatibility with the legacy VLM auto-model API.VisionLanguageDataCollatorto:answer_only_loss, chat-template, and DFlash label-alignment settings;VLM_MIN_PIXELS/VLM_MAX_PIXELSprocessor limits;training_seq_lenrequired by DFlash block training.Usage