Repository navigation
Conversation
📝 WalkthroughWalkthroughThis PR extends speculative decoding support for vision-language models by adding JSONL fallback dataset loading, expanding multimodal collation and masking, wiring the new collator options through the example path, broadening VLM auto-model selection, and updating HFDFlashModel’s multimodal forward and zero-loss handling. ChangesVLM training and data pipeline updates
Estimated code review effort: 3 (Moderate) | ~35 minutes Sequence Diagram(s)sequenceDiagram
participant Trainer
participant VisionLanguageDataCollator
participant Processor
participant HFDFlashModel
Trainer->>VisionLanguageDataCollator: collate multimodal batch
VisionLanguageDataCollator->>Processor: from_pretrained with pixel bounds
VisionLanguageDataCollator->>VisionLanguageDataCollator: build assistant masks and labels
VisionLanguageDataCollator-->>Trainer: batch tensors and labels
Trainer->>HFDFlashModel: forward(batch)
HFDFlashModel->>HFDFlashModel: detect multimodal kwargs
HFDFlashModel->>HFDFlashModel: super().forward with hidden_states
HFDFlashModel-->>Trainer: outputs or zero-loss return
🚥 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}
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: 2
🤖 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 `@modelopt/torch/speculative/utils.py`:
- Around line 613-620: The model class resolution in utils.py currently prefers
AutoModelForVision2Seq before AutoModelForImageTextToText, which keeps using the
deprecated alias when both are available. Update the fallback order in the model
class lookup so AutoModelForImageTextToText is tried first, then fall back to
AutoModelForVision2Seq, and keep the architecture-based fallback in the same
resolution chain.
In `@modelopt/torch/utils/plugins/transformers_dataset.py`:
- Around line 537-545: The padding logic in `pad_tokenized_messages` only
extends `input_ids` and `attention_mask`, but it also needs to pad
`assistant_masks` when `apply_chat_template()` returns that tensor. Update the
sequence-padding loop near the existing `tokenized_messages.get(key)` handling
to include `assistant_masks` with the same `pad_width`, preserving tensor shape
consistency for the downstream `answer_only_loss` path in
`transformers_dataset`.
🪄 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: 77e86707-9ddf-4546-8afc-48443831de63
📒 Files selected for processing (4)
examples/speculative_decoding/eagle_utils.pymodelopt/torch/speculative/plugins/hf_dflash.pymodelopt/torch/speculative/utils.pymodelopt/torch/utils/plugins/transformers_dataset.py
| # ``AutoModelForVision2Seq`` was renamed in Transformers 5. Prefer the | ||
| # legacy alias when present, then use the current auto-model entry point. | ||
| model_cls = getattr(transformers, "AutoModelForVision2Seq", None) | ||
| if model_cls is None: | ||
| model_cls = getattr(transformers, "AutoModelForImageTextToText", None) | ||
| if model_cls is None: | ||
| architecture = (getattr(model_config, "architectures", None) or [None])[0] | ||
| model_cls = getattr(transformers, architecture, None) if architecture else None |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant file and nearby logic
git ls-files modelopt/torch/speculative/utils.py
wc -l modelopt/torch/speculative/utils.py
sed -n '580,660p' modelopt/torch/speculative/utils.py
# Search for related auto-model selection logic
rg -n "AutoModelForVision2Seq|AutoModelForImageTextToText|architectures|FutureWarning" modelopt/torch/speculative -SRepository: NVIDIA/Model-Optimizer
Length of output: 4378
🌐 Web query:
Hugging Face AutoModelForVision2Seq deprecated removed in v5 use AutoModelForImageTextToText FutureWarning
💡 Result:
The class AutoModelForVision2Seq was deprecated in earlier versions of the Hugging Face Transformers library and has been officially removed in version 5.0.0 [1][2][3]. To resolve the FutureWarning or import errors resulting from this change, you should update your code to use AutoModelForImageTextToText instead [4][1][5]. This change is part of the broader cleanup and architectural updates introduced in Transformers v5 [6][3]. If you are encountering this in your codebase, replacing all instances of AutoModelForVision2Seq with AutoModelForImageTextToText will restore compatibility with the v5.x ecosystem [1][7].
Citations:
- 1: Remove references to AutoModelForVision2Seq huggingface/transformers#41513
- 2: huggingface/transformers@878d60a
- 3: https://github.com/huggingface/transformers/releases/tag/v5.0.0
- 4: https://raw.githubusercontent.com/huggingface/transformers/main/MIGRATION_GUIDE_V5.md
- 5: AutoModelForVision2Seq.from_config() got an unexpected keyword argument 'torch_dtype' huggingface/transformers#40477
- 6: https://github.com/huggingface/transformers/blob/main/MIGRATION_GUIDE_V5.md
- 7: Remove references to AutoModelForVision2Seq huggingface/transformers#41513
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find the Transformers version constraints used by this repo
rg -n "transformers|tf_latest|AutoModelForImageTextToText|AutoModelForVision2Seq" pyproject.toml requirements*.txt noxfile.py modelopt tests -SRepository: NVIDIA/Model-Optimizer
Length of output: 50378
🌐 Web query:
Transformers 4.56 AutoModelForVision2Seq AutoModelForImageTextToText both available deprecation warning
💡 Result:
In the Hugging Face Transformers library, AutoModelForVision2Seq is deprecated and has been superseded by AutoModelForImageTextToText [1][2][3]. In version 4.56, using AutoModelForVision2Seq triggers a FutureWarning because it is slated for removal in version 5.0 [4][3][5]. The library has explicitly transitioned to AutoModelForImageTextToText as the standard class for handling vision-language (image-text-to-text) models [6][7]. To resolve this deprecation warning, you should update your code to replace any instances of AutoModelForVision2Seq with AutoModelForImageTextToText [1][8]. For example: Instead of: from transformers import AutoModelForVision2Seq model = AutoModelForVision2Seq.from_pretrained("model-name") Use: from transformers import AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained("model-name")
Citations:
- 1: Remove references to AutoModelForVision2Seq huggingface/transformers#41513
- 2: huggingface/transformers@878d60a
- 3: huggingface/transformers@cb54ce4
- 4: Warning when using the "image-to-text" pipeline huggingface/transformers#41509
- 5: AutoModelForVision2Seq.from_config() got an unexpected keyword argument 'torch_dtype' huggingface/transformers#40477
- 6: https://huggingface.co/docs/transformers/v4.56.0/tasks/image_text_to_text
- 7: Add auto model for image-text-to-text huggingface/transformers#32472
- 8: Remove references to AutoModelForVision2Seq huggingface/transformers#41513
Prefer AutoModelForImageTextToText first
AutoModelForVision2Seq is still present in the supported 4.56+ range and emits a FutureWarning, so this order keeps taking the deprecated path whenever both classes exist. Try the new class first, then fall back to the legacy alias for older installs.
🤖 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/speculative/utils.py` around lines 613 - 620, The model class
resolution in utils.py currently prefers AutoModelForVision2Seq before
AutoModelForImageTextToText, which keeps using the deprecated alias when both
are available. Update the fallback order in the model class lookup so
AutoModelForImageTextToText is tried first, then fall back to
AutoModelForVision2Seq, and keep the architecture-based fallback in the same
resolution chain.
| for key, pad_value in ( | ||
| ("input_ids", self.tokenizer.pad_token_id), | ||
| ("attention_mask", 0), | ||
| ): | ||
| value = tokenized_messages.get(key) | ||
| if not isinstance(value, torch.Tensor): | ||
| continue | ||
| pad = value.new_full((value.shape[0], pad_width), pad_value) | ||
| tokenized_messages[key] = torch.cat((value, pad), dim=1) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Pad assistant_masks with the other sequence tensors.
When marker-derived masks are not used, apply_chat_template() can return assistant_masks; if the processor also ignored padding="max_length", this helper pads input_ids/attention_mask but leaves assistant_masks short. Line 595 can then fail with a shape mismatch under answer_only_loss=True.
Proposed fix
for key, pad_value in (
("input_ids", self.tokenizer.pad_token_id),
("attention_mask", 0),
+ ("assistant_masks", 0),
):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for key, pad_value in ( | |
| ("input_ids", self.tokenizer.pad_token_id), | |
| ("attention_mask", 0), | |
| ): | |
| value = tokenized_messages.get(key) | |
| if not isinstance(value, torch.Tensor): | |
| continue | |
| pad = value.new_full((value.shape[0], pad_width), pad_value) | |
| tokenized_messages[key] = torch.cat((value, pad), dim=1) | |
| for key, pad_value in ( | |
| ("input_ids", self.tokenizer.pad_token_id), | |
| ("attention_mask", 0), | |
| ("assistant_masks", 0), | |
| ): | |
| value = tokenized_messages.get(key) | |
| if not isinstance(value, torch.Tensor): | |
| continue | |
| pad = value.new_full((value.shape[0], pad_width), pad_value) | |
| tokenized_messages[key] = torch.cat((value, pad), dim=1) |
🤖 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 537 - 545,
The padding logic in `pad_tokenized_messages` only extends `input_ids` and
`attention_mask`, but it also needs to pad `assistant_masks` when
`apply_chat_template()` returns that tensor. Update the sequence-padding loop
near the existing `tokenized_messages.get(key)` handling to include
`assistant_masks` with the same `pad_width`, preserving tensor shape consistency
for the downstream `answer_only_loss` path in `transformers_dataset`.
|
/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 — findings summary.
Totals: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 0
Scope reviewed: All 4 changed files in the PR (transformers_dataset.py, hf_dflash.py, speculative/utils.py, eagle_utils.py). The base-tip diff also surfaced unrelated main-branch drift in other files (e.g. examples/hf_ptq/README.md); those are not part of this PR and were not reviewed.
Most impactful finding (IMPORTANT — inline on transformers_dataset.py):
_pad_sequence_tensors pads input_ids/attention_mask to train_len but not assistant_masks. On the generation-tag VLM path (return_assistant_tokens_mask=True, i.e. no recognized ChatML/Llama markers) combined with a processor that ignores padding="max_length" — the exact case this method exists for — the mask returns at the pre-pad length while input_ids is padded, producing a shape mismatch in _process_multimodal_sample's masked label assignment. The primary Qwen3-VL marker path is unaffected because _build_assistant_masks builds the mask via torch.zeros_like(input_ids) after padding.
Other observations (no action needed):
- The rebuilt DDP dummy-loss (summing
param.reshape(-1)[0] * 0.0over all trainable draft params) is a correct improvement over the prior single-fc.weightreference — it marks every draft parameter as graph-participating so DDP reduction completes on ranks with no valid answer-only labels. - The multimodal top-level-forward routing and the
AutoModelForVision2Seq/AutoModelForImageTextToText/architecture fallback chain both look sound. (Note: CodeRabbit suggested preferringAutoModelForImageTextToTextfirst as the non-deprecated alias — worth considering, but not a correctness issue.) shift_labelsplumbing is consistent end-to-end (main.pypassesshift_labels=not is_dflash; DFlash uses unshifted labels as documented).
Risk assessment: Low-to-moderate. The core Qwen3-VL path is well-guarded; the one blocking issue is confined to the alternate generation-tag masking path. Note the PR is still marked draft (no automated regression coverage yet) and CodeRabbit's pre-merge gate flags two new # nosec comments that violate SECURITY.md — resolve those before merge.
| for key, pad_value in ( | ||
| ("input_ids", self.tokenizer.pad_token_id), | ||
| ("attention_mask", 0), | ||
| ): | ||
| value = tokenized_messages.get(key) | ||
| if not isinstance(value, torch.Tensor): | ||
| continue | ||
| pad = value.new_full((value.shape[0], pad_width), pad_value) | ||
| tokenized_messages[key] = torch.cat((value, pad), dim=1) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] _pad_sequence_tensors pads only input_ids and attention_mask, but not assistant_masks.
Trace the non-marker path: when a VLM chat template has {% generation %} tags but no recognized ChatML/Llama markers, _apply_chat_template sets derive_masks_from_markers=False and therefore requests return_assistant_tokens_mask=True from the processor. If that processor also ignores padding="max_length" (the exact scenario this method exists to compensate for), assistant_masks comes back at the pre-pad length L, while input_ids gets padded here to train_len. Then in _process_multimodal_sample:
labels[..., :-1][assistant_mask[..., 1:] == 0] = IGNORE_TOKEN_IDindexes a [B, train_len-1] tensor with a [B, L-1] boolean mask → shape-mismatch crash (or silently wrong masking if it happens to broadcast).
The primary Qwen3-VL path is safe because markers are recognized and _build_assistant_masks builds the mask after padding via torch.zeros_like(input_ids). But any VLM relying on generation-tag masks hits this.
Why it matters: breaks answer_only_loss for the generation-tag VLM path the PR still advertises supporting.
Fix: include assistant_masks in the pad loop when present (pad with 0):
for key, pad_value in (
("input_ids", self.tokenizer.pad_token_id),
("attention_mask", 0),
("assistant_masks", 0),
):(guarding, as the loop already does, on the value being a rank-2 tensor).
There was a problem hiding this comment.
@skierat Please take a look if this make sense. Thanks
| "images", | ||
| "videos", | ||
| } | ||
| use_top_level_forward = any( |
There was a problem hiding this comment.
LGTM for now. In the long term I think we want to abstract away the model-specific logics from the algorithm logic. Could you help add a TODO comment here?
There was a problem hiding this comment.
Hi @skierat, thanks for the PR! The changes LGTM. Could you please:
- Add some minimal unit tests for the changes in
VisionLanguageDataCollatorandload_vlm_or_llm - Run a end2end test with
tools/launcherand add a yaml example undertools/launcher/examples, similar to the existingtools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash.yaml. (Also works if you prefer to put in a separate PR).
This is a new place we put per-model examples. It should be much easier to create, test and share with users than our previous shell scripts or notebooks.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unit/torch/utils/test_transformers_dataset.py (1)
33-82: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGood hermetic unit test; consider covering the default
shift_labels=Truepath too.This test only exercises
shift_labels=Falsewith marker-derived masks. The collator's default isshift_labels=True(perVisionLanguageDataCollator.__init__), and the alternate branch where_assistant_marker_specs()is empty (nativereturn_assistant_tokens_mask=Truepath) is also untested. Both are common/default runtime paths for this new masking logic.As per path instructions, checked-in tests should "document expected behavior, protect against regressions, or flag backward-incompatible behavior changes" — the shifted-label branch is the default and most-used configuration, so leaving it uncovered is a real regression-protection gap.
🤖 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/utils/test_transformers_dataset.py` around lines 33 - 82, The current test only covers the non-default masking flow in VisionLanguageDataCollator._process_multimodal_sample with shift_labels set to false and marker-derived assistant masks. Add coverage for the default shift_labels=true behavior, and also exercise the native return_assistant_tokens_mask path by making _assistant_marker_specs() return empty so the processor path is used; verify the resulting input_ids, assistant_masks, and labels match the shifted-label expectations.Source: Path instructions
🤖 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.
Nitpick comments:
In `@tests/unit/torch/utils/test_transformers_dataset.py`:
- Around line 33-82: The current test only covers the non-default masking flow
in VisionLanguageDataCollator._process_multimodal_sample with shift_labels set
to false and marker-derived assistant masks. Add coverage for the default
shift_labels=true behavior, and also exercise the native
return_assistant_tokens_mask path by making _assistant_marker_specs() return
empty so the processor path is used; verify the resulting input_ids,
assistant_masks, and labels match the shifted-label expectations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ec4dace4-149f-43c0-b754-b1f337937f83
📒 Files selected for processing (2)
tests/unit/torch/speculative/plugins/test_fakebase.pytests/unit/torch/utils/test_transformers_dataset.py
h-guo18
left a comment
h-guo18
left a comment
There was a problem hiding this comment.
lgtm. Approve to unblock; Please also take a look at AI review comments. Thanks!
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