add Qwen3-VL support for DFlash training#1975
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (6)
📝 WalkthroughWalkthroughChangesDFlash 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?
2 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