feat(dspark): add multimodal (image+text) training support - #785
Open
curnane-lab wants to merge 5 commits into
Open
curnane-lab wants to merge 5 commits into
curnane-lab wants to merge 5 commits into
Conversation
Port of the DFlash draft-model changes from sgl-project/SpecForge PR sgl-project#585: partial rotation in apply_rotary_pos_emb (rotary_dim < head_dim, for Qwen3.5/Qwen3.6 partial_rotary_factor=0.25) and Qwen3InterleavedMultiRotaryEmbedding selected by rope_scaling.mrope_interleaved. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> (cherry picked from commit 9323a51; scoped to the draft model only, draft configs land with the recipe)
Follow-up adaptation of 9323a51 to the server-only architecture: - specforge/modeling/target/target_utils.py: add QWEN3_VL_MODEL_TYPES and resolve_target_weight_keys(); TargetEmbeddingsAndHead.from_pretrained now auto-selects model.language_model.embed_tokens.weight for VLM targets (Qwen3-VL / Qwen3.5 / Qwen3.6) when the embed key is unset or left at the LLM default. Explicit keys are honored as-is. - specforge/algorithms/model_providers.py: populate_dflash_generated_config reads the language-model depth via the text_config fallback so VLM draft config auto-generation uses the correct num_hidden_layers. (cherry picked from commit e2280f4)
Implements end-to-end VLM DFlash training on the server-only runtime:
- contracts/providers: dflash registers a STREAMING 'multimodal'
FeatureContract ({input_ids, loss_mask, hidden_states, position_ids}) and a
ServerStreamingProvider with a VLM ServerInputAdapter;
ServerCaptureLayout/ServerCaptureSchema gain position_ids_feature.
- data: specforge/data/vlm_preprocessing.py renders conversations with the
target's own chat template (image attached to the first user turn) and
expands the image region in id space via the target's HF processor, yielding
expanded input_ids/loss_mask (trainer + passthrough) and collapsed
request_input_ids + base64 image_data (capture request). One image per
sample (v1); text-only samples work in the same run.
- inference: the capture request's features map gains 'position_ids'; the
regenerated sglang v0.5.14 spec-capture patch writes the request's mRoPE
positions (1, L, 3) int64 into Mooncake from
req.multimodal_inputs.mrope_positions (arange fallback for text prompts).
Managed capture servers set SGLANG_MM_AVOID_RETOKENIZE=1 for
input_modality=multimodal so the server re-expands placeholders in id space
with no retokenization drift.
- training: build_vlm_collator pads position_ids; OnlineDFlashModel.forward
accepts optional (B, S, 3) position ids and gathers 3D mRoPE positions for
context + anchor-offset draft slots; DFlashTrainStrategy passes the tensor
through. Text paths are unchanged.
- recipe: examples/configs/qwen3.5-4b-vl-dflash-multimodal-disaggregated.yaml
+ configs/qwen3.5-4b-vl-dflash.json (Qwen3.5-4B VLM draft geometry).
- tests: new tests/test_algorithms/test_dflash_multimodal.py; update the
text-only assertions in test_builtin_providers / test_server_only_online /
test_unified_feature_reachability / test_launch_topology /
test_example_draft_config_wiring for the new modality and recipe.
GPU/NPU end-to-end validation (live ViT capture + training) is pending and
tracked in docs/advanced_features/vlm_dflash.md.
Co-authored-by: danaodai <chenjiayuan1077@163.com>
Co-authored-by: danaodai <chenjiayuan1077@163.com>
curnane-lab
changed the title
curnane-lab
marked this pull request as ready for review
September 3, 2026 01:20
curnane-lab
requested review from
FlamingoPg,
FrankLeeeee,
shuaills,
sleepcoo and
zyksir
as code owners
September 3, 2026 01:20
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
DSpark (the DFlash-family Markov-branch variant) currently only supports text-only verifiers. This PR extends it to multimodal (image+text) verifiers such as
Qwen/Qwen3.5-4B, reusing the multimodal server-capture infrastructure introduced in #730: image rows ride the Chat-Completions-style capture path (the patched SGLang server expands placeholders and runs the full VLM forward), and the captured verifier aux hidden states already fuse the vision content, so the draft itself never sees pixels.This PR is stacked on #730 and should be merged after it. The branch contains #730's four commits plus one DSpark commit; once #730 lands, the compare view collapses to the single DSpark commit.
Modifications
specforge/algorithms/dspark/providers.py: register aSTREAMING + multimodalfeature contract andServerStreamingProvider. The capture layout matches DSpark's text streaming path (aux=hidden_states,last=target_last_hidden_states, passthroughinput_ids/loss_mask);build_input_adapterinjects DSpark's ownminimum_loss_tokensrule.specforge/algorithms/common/vlm_input.py:VlmServerInputAdapternow accepts an algorithm-ownedminimum_loss_tokenscallable (Nonekeeps the DFlash default), so DFlash text behavior is unchanged.OnlineDSparkModelinheritsOnlineDFlashModel's block-parallel forward, so multimodal batches flow through the identical plain-1D-position training forward as text batches (byte-identical text runs).configs/qwen3.5-4b-dspark.json: DSpark draft config mirroring the official DFlash draft geometry (5 layers, aux target layers [1, 8, 15, 22, 29], vocab 248320,rope_scaling: null) with DSpark specifics (block_size 7, markov_rank 256, dspark projector, confidence head).examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dspark-disaggregated.yaml: single-node Ascend NPU disaggregated-online recipe (sglang capture server + Mooncake).docs/advanced_features/vlm_dflash.md) and tests updated (see below).Related Issues
Stacked on #730; merge after it.
Accuracy Test
No model-side forward/kernel changes (training-support PR).
Benchmark & Profiling
Not applicable - training-side support only, no serving-path changes.
Checklist