[Feature] Multimodal (image+text) DFlash training via server-side capture (Qwen3.5)

docs/advanced_features/vlm_dflash.md

@@ -0,0 +1,99 @@
+# VLM DFlash Support - Status (add_vl_support)
+
+This branch ports sgl-project/SpecForge PR #585 (commit `9323a510`, author zyk42)
+onto the server-only unified runtime **and implements end-to-end multimodal
+(image+text) DFlash training on top of it**. This document records what landed,
+what was deliberately not ported, and the validation status.
+
+> **Rope decision**: multimodal DFlash drafts use the plain 1D rope convention -
+> the same as the text-only path - on purpose. Visual information reaches the
+> draft exclusively through the captured target hidden states, never through the
+> draft's own positional embedding, so the draft has no use for the target's
+> (3, L) mRoPE positions. Staying on one position convention keeps training and
+> serving byte-identical to the text pipeline and reuses the stock
+> `configs/qwen3.5-4b-dflash.json` draft geometry.
+
+## What this branch contains
+
+### Foundation (ported from PR #585)
+
+- **Draft model (VLM-capable)** - `specforge/modeling/draft/dflash.py`:
+  partial rotation in `apply_rotary_pos_emb` (`rotary_dim < head_dim`, for
+  Qwen3.5/Qwen3.6 `partial_rotary_factor=0.25`). The draft always uses the
+  stock `Qwen3RotaryEmbedding`.
+- **Weight-key resolution** - `resolve_target_weight_keys()` in
+  `specforge/modeling/target/target_utils.py` auto-selects
+  `model.language_model.embed_tokens.weight` for VLM targets;
+  `populate_dflash_generated_config` reads language-model depth via
+  `text_config`.
+
+### Multimodal capture (new in this branch)
+
+End-to-end data flow: **JSONL (+ image) -> expanded ids/loss mask -> capture
+request (single-placeholder ids + base64 image) -> patched SGLang server
+expands, runs the ViT, captures aux hidden states -> Mooncake -> collator ->
+training forward on plain 1D positions**.
+
+- `model.input_modality: multimodal` (DFlash only): a `FeatureContract`
+  (`{input_ids, loss_mask, hidden_states}`) and a
+  `ServerStreamingProvider` with a VLM `ServerInputAdapter`
+  (`specforge/algorithms/common/vlm_input.py`). Multimodal capture stores the
+  same three tensors as text capture - no `position_ids` artifact is requested
+  or consumed.
+- `specforge/data/vlm_preprocessing.py`: ShareGPT-style records with an
+  optional `image` field (path or base64); the target's own chat template and
+  HF processor produce the expanded `input_ids`/`loss_mask` (image region
+  expanded in id space, mask zeros). One image per sample max (v1); text-only
+  samples work in the same run.
+- `patches/sglang/v0.5.14/spec-capture.patch`: tracks upstream's rewritten
+  async streaming sink. Multimodal capture requests ride the stock
+  `input_ids` + `image_data` `/generate` path with
+  `SGLANG_MM_AVOID_RETOKENIZE=1` (set by the managed launcher for
+  `input_modality=multimodal`), so the server re-expands placeholders in id
+  space with zero retokenization drift - and the passthrough/seq-len checks
+  fail loudly if client and server expansions ever disagree.
+- Training: `OnlineDFlashModel._forward_draft_blocks` builds positions with
+  the unconditional text-path 1D `arange` convention; multimodal batches flow
+  through the identical forward as text batches. Text runs are
+  byte-identical to before.
+- Recipe: `examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml`
+  (single-node Ascend NPU managed stack; draft config
+  `configs/qwen3.5-4b-dflash.json`).
+
+## Not ported (by design)
+
+- HF-backend VLM capture (`dflash_target_model.py`, `_build_vlm_reqs`,
+  `mm_token_type_ids`) and the `train_dflash.py --is-vlm` plumbing from the
+  pre-#678 script stack - superseded by server capture.
+- `QwenVLOnlineDFlashModel` wiring - PR #585 referenced this class but never
+  defined it; the unified runtime needs no separate VLM wrapper class.
+- mRoPE draft support (`Qwen3InterleavedMultiRotaryEmbedding`, the
+  `rope_scaling.mrope_interleaved` switch, and the server `position_ids`
+  capture artifact): the draft consumes visual information only through the
+  captured target hidden states, so the (3, L) target positions carry no
+  signal for it. Retired in favor of the single plain-rope convention above.
+- Two accidental reverts in the original #585 diff (domino projector code,
+  D-PACE CLI args) were dropped during the cherry-pick.
+- Offline (precomputed hidden states) multimodal capture: the offline path
+  stays text-only for now.
+- Online evaluation for multimodal runs.
+
+## Validation status
+
+- **Verified (CPU, this repo)**: registration parity and provider gates,
+  request/payload construction, expansion math, collator, golden
+  topology/recipe tests - `tests/test_algorithms/test_dflash_multimodal.py`
+  plus the updated `test_config` suites.
+- **Verified statically**: the regenerated patch applies cleanly both ways to
+  pristine sglang v0.5.14 (`git apply --check` / `--reverse --check`).
+- **Not yet verified (needs GPU/NPU)**: a live multimodal capture run
+  (Qwen3-VL / Qwen3.5 target + ViT) and an end-to-end training run. This is
+  the next step; see the recipe above.
+
+## Reference results from PR #585 (HF stack, author-validated)
+
+- Qwen3-VL-30B-A3B-Thinking, 278K target-regenerated samples, 5-layer draft,
+  block_size=8: accept length 3.52, +35.8% inference speedup (4x RTX 5090,
+  TP=4, SGLang 0.5.12).
+- Data must be target-model greedy-regenerated; system prompt must match
+  between training and inference; <10K samples overfit severely.

examples/configs/README.md

@@ -69,6 +69,11 @@ target hidden states are concatenated and projected. It remains registered as
 the `eagle3` strategy; EAGLE3.1 is a draft-model configuration variant, not a
 second runtime or launch path.
 
+Multimodal (image+text) training is supported for `training.strategy=dflash`
+via `model.input_modality: multimodal` (see
+`online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml`); the
+remaining catalog is text-only.
+
 Recipes under `online/disaggregated/managed-local` are opt-in, single-node
 full-stack examples. Their typed `managed_local` blocks own Mooncake, one or
 more patched SGLang capture servers, and the trainer GPU allocation; the same
@@ -179,7 +184,7 @@ should make their training strategy and topology explicit.
 | `model.draft_num_hidden_layers` | `null` | Positive fresh-architecture override where the strategy permits it. EAGLE3 remains one layer; P-EAGLE and DFlash may override their generated defaults. |
 | `model.draft_block_size` | `null` | Positive DFlash block-size override; generated DFlash configs default to 16. |
 | `model.target_backend` | `sglang` | `sglang` is the only accepted value; retired `hf`/`custom` names fail at config load. Offline feature consumers do not instantiate a target inference backend. |
-| `model.input_modality` | `text` | The provider modality. The unified runtime supports text only; VLM modalities such as `qwen2_5_vl` are rejected. |
+| `model.input_modality` | `text` | The provider modality. Built-ins support `text`; DFlash additionally supports `multimodal` (image+text server capture). Other identifiers are rejected at application resolution. |
 | `model.shard_target_output` | `false` | Retained for config migration; leave it false on the online disaggregated path. |
 | `model.trust_remote_code` | `false` | Enable only for model repositories that require custom loading code. |
 | `model.use_liger_kernel` | `false` | Enable Liger Qwen3 RMSNorm/SwiGLU kernels for DFlash training. Requires the `specforge[liger]` extra. |
@@ -193,6 +198,7 @@ should make their training strategy and topology explicit.
 | `model.mask_token_id` | `null` | DFlash-family/P-EAGLE mask token override. Otherwise it resolves from the draft config and then the tokenizer. |
 | `model.tokenizer_pad_token_id` | `null` | Explicit non-negative tokenizer pad ID. Use it for released tokenizers that omit padding metadata. |
 | `model.sglang_attention_backend` | `flashinfer` | SGLang attention implementation for an in-process or managed capture server. |
+| `model.sglang_mm_attention_backend` | `null` | Vision-encoder attention backend for capture servers. On Ascend NPU with a non-text `model.input_modality` it defaults to `ascend_attn` (fused); `sdpa` materializes N² vision scores and can OOM on large images. |
 | `model.sglang_mem_fraction_static` | `0.4` | SGLang static-memory fraction in `(0, 1]`; inherited by managed capture servers unless they override it. |
 | `model.sglang_disable_radix_cache` | `true` | Preserve the historical managed-capture behavior. Set `false` for hybrid targets such as Inkling that require the radix tree. Unique per-attempt cache namespaces still force complete capture prefills. |
 | `model.sglang_context_length` | `null` | Positive explicit context limit. Managed capture requires at least `data.max_length + 7`; omitting it derives that value. |
@@ -453,9 +459,13 @@ unless tuning throughput or memory pressure.
   `sp_ulysses_size * sp_ring_size > 1`. Non-USP runs keep both SP sizes at 1.
 - P-EAGLE reuses the EAGLE3 server feature schema, uses `flex_attention`, and
   requires batch size 1.
-- VLM training, including Qwen2.5-VL, is not supported. Online capture accepts
-  text inputs only.
 - `training.compact_teacher` is offline text EAGLE3 only.
+- Multimodal (image+text) training requires `training.strategy=dflash` with
+  `model.input_modality: multimodal`, a plain-rope DFlash draft config (e.g.
+  `configs/qwen3.5-4b-dflash.json`; the draft needs no mRoPE - visual
+  information arrives through the captured target hidden states), and the
+  v0.5.14 spec-capture patch. Vendor modalities such as `qwen2_5_vl` remain
+  unsupported.
 - Online evaluation is not supported. Offline `data.eval_hidden_states_path`
   and `training.eval_interval` must be configured together.
 

examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml

@@ -0,0 +1,52 @@
+# Multimodal (image+text) DFlash on Qwen3.5-4B: disaggregated online training
+# against an external patched SGLang capture server. Requires sglang v0.5.14
+# with patches/sglang/v0.5.14/spec-capture.patch and
+# SGLANG_MM_AVOID_RETOKENIZE=1 on the capture server.
+#
+# Data: ShareGPT-style JSONL with an optional top-level "image" field (file
+# path or base64) per record; text-only records are supported in the same run.
+model:
+  target_model_path: Qwen/Qwen3.5-4B
+  draft_model_config: configs/qwen3.5-4b-dflash.json
+  target_backend: sglang
+  trust_remote_code: true
+  input_modality: multimodal
+data:
+  train_data_path: ./cache/dataset/train_vl_regen.jsonl
+  max_length: 4096
+  chat_template: qwen3.5
+  cache_dir: cache
+  build_dataset_num_proc: 32
+training:
+  strategy: dflash
+  num_epochs: 10
+  max_steps: 10000
+  batch_size: 2
+  accumulation_steps: 4
+  learning_rate: 0.0006
+  warmup_ratio: 0.04
+  max_grad_norm: 1.0
+  num_anchors: 512
+  loss_decay_gamma: 7.0
+  save_interval: 1000
+  dist_timeout: 30
+  seed: 42
+tracking:
+  report_to: none
+run_id: qwen3.5-4b-vl-dflash
+output_dir: outputs/qwen3.5-4b-vl-dflash
+
+deployment:
+  mode: disaggregated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 4
+  disaggregated:
+    control_dir: outputs/qwen3.5-4b-vl-dflash/control
+    consumer_state_dir: outputs/qwen3.5-4b-vl-dflash/consumer-state
+    backend: mooncake
+    server_urls:
+      - http://127.0.0.1:30000
+    mooncake_metadata_server: http://127.0.0.1:35880/metadata
+    mooncake_master_server_addr: 127.0.0.1:35551
+    mooncake_protocol: tcp

scripts/regenerate_train_data.py

@@ -244,6 +244,29 @@ def build_query_kwargs(args, messages, max_tokens=None):
     return query_kwargs
 
 
+def _extract_record_images(data: Dict[str, Any]) -> List[str]:
+    """Resolve image references of a record (``image``/``image_path`` string or
+    ``images`` list), in insertion order."""
+    refs: List[str] = []
+    single = data.get("image") or data.get("image_path")
+    if isinstance(single, str):
+        refs.append(single)
+    images = data.get("images")
+    if isinstance(images, list):
+        refs.extend(r for r in images if isinstance(r, str))
+    return refs
+
+
+def _image_url_part(path: str) -> Dict[str, Any]:
+    import base64
+    import mimetypes
+
+    mime = mimetypes.guess_type(path)[0] or "image/jpeg"
+    with open(path, "rb") as f:
+        b64 = base64.b64encode(f.read()).decode("ascii")
+    return {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}}
+
+
 def call_sglang(
     args,
     server_address: str,
@@ -260,6 +283,8 @@ def call_sglang(
 
     messages = data["conversations"]
     regenerated_messages = []
+    record_images = _extract_record_images(data)
+    image_attached = False
 
     # ignore data which starts with an assistant message
     if messages[0]["role"] == "assistant":
@@ -273,6 +298,26 @@ def call_sglang(
         elif message["role"] == "assistant":
             continue
         elif message["role"] == "user":
+            # Multimodal records: attach the record's images to the first user
+            # turn that carries the <image> placeholder (OpenAI content parts).
+            content = message.get("content")
+            if (
+                record_images
+                and not image_attached
+                and isinstance(content, str)
+                and "<image>" in content
+            ):
+                try:
+                    parts = [_image_url_part(p) for p in record_images]
+                except OSError as exc:
+                    data["status"] = "error"
+                    data["error"] = f"unreadable image file: {exc}"
+                    return data
+                text = content.replace("<image>\n", "").replace("<image>", "")
+                parts.append({"type": "text", "text": text})
+                message = dict(message)
+                message["content"] = parts
+                image_attached = True
             regenerated_messages.append(message)
 
             query_kwargs = build_query_kwargs(args, regenerated_messages, max_tokens)

specforge/algorithms/common/dflash_family_model.py

@@ -446,7 +446,8 @@ def forward(
         max_valid_anchors: Optional[int] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, object]]:
         """Parallel block-wise training forward pass; returns
-        (loss, accuracy, metrics) — same shape as Domino's forward."""
+        (loss, accuracy, metrics) - same shape as Domino's forward.
+        """
         if self.attention_backend == "flex_attention" and not FLEX_ATTENTION_AVAILABLE:
             raise ValueError(
                 "flex_attention is not available on this device; use sdpa/eager."

specforge/algorithms/common/providers.py

@@ -546,8 +546,8 @@ class ServerStreamingProvider:
     """Algorithm adapter for externally captured streaming features.
 
     ``build_input_adapter`` is deliberately modality-neutral. Text providers
-    can leave it unset; the current runtime does not support VLM registration
-    or media requests.
+    can leave it unset; modalities such as ``multimodal`` provide one to own
+    prompt preparation and request construction (e.g. image inputs).
     """
 
     modality: str

specforge/algorithms/common/vlm_input.py

@@ -0,0 +1,101 @@
+# coding=utf-8
+# Copyright 2024 The SpecForge team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+"""Multimodal (image+text) ServerInputAdapter for server-side capture.
+
+Owns the three modality seams for ``modality="multimodal"``:
+
+- ``load_input_tools``: the target tokenizer (the processor is loaded lazily
+  in ``prepare_prompts`` so training-model construction receives a plain
+  tokenizer, exactly like the text path).
+- ``prepare_prompts``: ShareGPT-style JSONL (+ optional ``image`` field) ->
+  payload dicts with expanded ``input_ids``/``loss_mask`` (what the trainer
+  and the passthrough capture use) plus ``request_input_ids`` (single
+  placeholder) and ``data:``-URI ``image_data`` (what the capture request sends).
+- ``build_request_inputs``: batch payloads -> the ``/generate`` model-input
+  fields ``{"input_ids", "image_data"}``.
+
+The capture server must run with ``SGLANG_MM_AVOID_RETOKENIZE=1`` so its
+multimodal processor re-expands placeholders in id space (no retokenization
+drift); the managed launcher sets this for ``input_modality="multimodal"``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Mapping, Sequence
+
+
+class VlmServerInputAdapter:
+    """Image+text input adapter for the SGLang server-capture transport."""
+
+    def __init__(self, config: Any) -> None:
+        self._config = config
+
+    def load_input_tools(self, config: Any) -> Any:
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            config.model.target_model_path,
+            cache_dir=config.model.cache_dir,
+            trust_remote_code=config.model.trust_remote_code,
+        )
+        if config.model.tokenizer_pad_token_id is not None:
+            tokenizer.pad_token_id = config.model.tokenizer_pad_token_id
+        elif tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+        return tokenizer
+
+    def prepare_prompts(
+        self,
+        config: Any,
+        input_tools: Any,
+        *,
+        draft_config: Any,
+    ) -> list[dict[str, Any]]:
+        from transformers import AutoProcessor
+
+        from specforge.algorithms.model_providers import dflash_min_loss_tokens
+        from specforge.data.vlm_preprocessing import build_vlm_prompt_payloads
+
+        tokenizer = input_tools
+        # The processor must inherit the target's own preprocessor config
+        # (min/max pixels, merge size) so the client expansion matches the
+        # capture server's expansion one-for-one; no overrides are accepted.
+        processor = AutoProcessor.from_pretrained(
+            config.model.target_model_path,
+            cache_dir=config.model.cache_dir,
+            trust_remote_code=config.model.trust_remote_code,
+        )
+        path = config.data.prompts_path or config.data.train_data_path
+        if not path:
+            raise ValueError("multimodal prompt preparation requires a data path")
+        return build_vlm_prompt_payloads(
+            path,
+            tokenizer,
+            processor,
+            chat_template=config.data.chat_template,
+            max_length=config.data.max_length,
+            min_loss_tokens=dflash_min_loss_tokens(config, draft_config),
+            max_prompts=config.data.max_prompts,
+        )
+
+    def build_request_inputs(
+        self,
+        tasks: Sequence[Any],
+    ) -> Mapping[str, Any]:
+        return {
+            "input_ids": [list(task.payload["request_input_ids"]) for task in tasks],
+            "image_data": [task.payload["image_data"] for task in tasks],
+        }
+
+
+def build_vlm_input_adapter(config: Any) -> VlmServerInputAdapter:
+    return VlmServerInputAdapter(config)
+
+
+__all__ = ["VlmServerInputAdapter", "build_vlm_input_adapter"]