diff --git a/docs/advanced_features/vlm_dflash.md b/docs/advanced_features/vlm_dflash.md new file mode 100644 index 000000000..23aa530e9 --- /dev/null +++ b/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. diff --git a/examples/configs/README.md b/examples/configs/README.md index 3cbd0cc6a..08969e0a3 100644 --- a/examples/configs/README.md +++ b/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. diff --git a/examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml new file mode 100644 index 000000000..099ac219a --- /dev/null +++ b/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 diff --git a/scripts/regenerate_train_data.py b/scripts/regenerate_train_data.py index f2f737fd2..252c989b2 100644 --- a/scripts/regenerate_train_data.py +++ b/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 placeholder (OpenAI content parts). + content = message.get("content") + if ( + record_images + and not image_attached + and isinstance(content, str) + and "" 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("\n", "").replace("", "") + 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) diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index 2c0bdfe62..2fde4a624 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/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." diff --git a/specforge/algorithms/common/providers.py b/specforge/algorithms/common/providers.py index 12cb0b08d..b3ae2facf 100644 --- a/specforge/algorithms/common/providers.py +++ b/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 diff --git a/specforge/algorithms/common/vlm_input.py b/specforge/algorithms/common/vlm_input.py new file mode 100644 index 000000000..36250d12a --- /dev/null +++ b/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"] diff --git a/specforge/algorithms/dflash/providers.py b/specforge/algorithms/dflash/providers.py index c5bc5837b..98832d5c2 100644 --- a/specforge/algorithms/dflash/providers.py +++ b/specforge/algorithms/dflash/providers.py @@ -26,6 +26,7 @@ TargetDerivedDraftDefaults, make_registration, ) +from specforge.algorithms.common.vlm_input import build_vlm_input_adapter from specforge.algorithms.contracts import ( AlgorithmCapabilities, AlgorithmSpec, @@ -156,6 +157,11 @@ def algorithm_spec() -> AlgorithmSpec: modality="text", required_tensors=ready, ), + FeatureContract( + mode=FeatureMode.STREAMING, + modality="multimodal", + required_tensors=ready, + ), ), capabilities=AlgorithmCapabilities( attention_backends={"eager", "sdpa", "flex_attention"}, @@ -226,6 +232,21 @@ def algorithm_providers() -> AlgorithmProviders: ), build_collator=collator, ), + ServerStreamingProvider( + modality="multimodal", + capture_method="dflash", + target_representation=None, + layout=ServerCaptureLayout( + aux_feature="hidden_states", + last_hidden_feature=None, + passthrough=( + ("input_ids", "input_ids", ()), + ("loss_mask", "loss_mask", ()), + ), + ), + build_collator=collator, + build_input_adapter=build_vlm_input_adapter, + ), ), ) diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index 285433fcd..e09cfe05a 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -458,8 +458,12 @@ def populate_dflash_generated_config( payload: Dict[str, Any], target_config: Any, _cfg: Config ) -> None: from specforge.modeling.draft.dflash import build_target_layer_ids + from specforge.modeling.target.target_utils import target_text_config - target_layers = getattr(target_config, "num_hidden_layers", None) + # VLM targets (Qwen3-VL / Qwen3.5 / Qwen3.6) keep the language-model depth + # under ``text_config``; the top-level value may describe the vision stack. + text_config = target_text_config(target_config) + target_layers = getattr(text_config, "num_hidden_layers", None) if not isinstance(target_layers, int) or target_layers < 1: raise ValueError( "DFlash auto-generation requires target num_hidden_layers, got " diff --git a/specforge/config/schema.py b/specforge/config/schema.py index 27870d0ca..eb955e7c5 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -80,6 +80,11 @@ class ModelConfig(StrictConfigModel): tokenizer_pad_token_id: Optional[int] = Field(default=None, ge=0) #: SGLang target-engine tuning. Ignored by hf/custom backends. sglang_attention_backend: str = "flashinfer" + #: Multimodal (vision encoder) attention backend for capture servers. On + #: Ascend NPU with a non-text input_modality it defaults to ascend_attn + #: (fused vision attention); sdpa materializes N^2 scores and OOMs on + #: large images. + sglang_mm_attention_backend: Optional[str] = None sglang_mem_fraction_static: float = Field(default=0.4, gt=0.0, le=1.0) #: Keep the historical managed-local behavior by default. Hybrid targets #: such as Inkling require the radix tree and can opt back in explicitly. @@ -967,13 +972,24 @@ def apply_overrides(config: Config, overrides: List[str]) -> Config: node = raw keys = path.split(".") for key in keys[:-1]: - if not isinstance(node.get(key), dict): + child = node.get(key) + if child is None: + # Optional sections default to None (e.g. disaggregated + # .managed_local on a server_urls recipe); create them so + # their declared fields can be overridden. Re-validation with + # extra="forbid" still rejects typo'd paths, so this stays safe. + child = {} + node[key] = child + if not isinstance(child, dict): raise ValueError(f"override path {path!r} does not exist") - node = node[key] - if keys[-1] not in node: - raise ValueError(f"override path {path!r} does not exist") - current = node[keys[-1]] - if isinstance(current, (dict, list)) and value.lstrip().startswith(("[", "{")): + node = child + current = node.get(keys[-1]) + # Structured (YAML) values are parsed when the target is a dict/list + # or a yet-unset leaf (None); schema re-validation does the final + # type check either way. + if (current is None or isinstance(current, (dict, list))) and ( + value.lstrip().startswith(("[", "{")) + ): import yaml try: diff --git a/specforge/data/vlm_preprocessing.py b/specforge/data/vlm_preprocessing.py new file mode 100644 index 000000000..fe9a2fdc9 --- /dev/null +++ b/specforge/data/vlm_preprocessing.py @@ -0,0 +1,361 @@ +# 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 +"""VLM (image+text) prompt preparation for server-side multimodal capture. + +The multimodal online path needs, per sample: + +- ``input_ids`` / ``loss_mask``: the FULLY image-expanded token sequence. + These are what the trainer consumes and what the passthrough capture stores, + so they must match the capture server's own expansion one-for-one. +- ``request_input_ids``: the same sequence with each image region collapsed + back to a single ``<|image_pad|>`` placeholder. This is what the capture + request sends; the patched server re-expands it in id space + (``SGLANG_MM_AVOID_RETOKENIZE=1``), which is guaranteed to reproduce the + client expansion because both sides run the target model's own HF processor + on the same image. +- ``image_data``: a ``data:``-URI image string for the capture request (the image + itself never enters the feature store; only text tokens and captured hidden + states do). + +Limitations (v1): at most one image per sample; the image is attached to the +first user turn (matching the Qwen-VL chat layout); text-only samples are +supported in the same run (``image_data=None``). Image references are read +from ``image`` / ``image_path`` (string) or ``images`` (one-element list); +unreadable images and multi-image samples are fatal ``ImageDataError``s, +never a silent text-only downgrade. +""" + +from __future__ import annotations + +import base64 +import io +import os +from typing import Any, Dict, List, Optional, Tuple + +from .prompt_builder import _iter_records +from .template import TEMPLATE_REGISTRY + +_IMAGE_PLACEHOLDER = "<|image_pad|>" +_VISION_PREFIX = "<|vision_start|><|image_pad|><|vision_end|>" + + +def _image_token_count(image_grid_thw, merge_size: int) -> int: + """Merged token count of one image region from its processor grid.""" + + grid = list(image_grid_thw[0]) + count = 1 + for dim in grid: + count *= int(dim) + return count // (merge_size * merge_size) + + +class ImageDataError(ValueError): + """Fatal image-contract violation (unreadable image, multi-image sample). + + Raised instead of silently degrading an image-bearing sample to text-only + training; the prompt-preparation loop re-raises it to abort loudly. + """ + + +def _extract_image_field(record: Dict[str, Any], *, source: str) -> Optional[str]: + """Resolve the single image reference of one record, or None. + + Accepts ``image`` / ``image_path`` (string) and ``images`` (a one-element + list, the convention used by common VLM corpora). A record carrying more + than one image is a fatal error in v1 (single-image contract), and so is + an ``images`` list whose element is not a string. + """ + + image_field = record.get("image") or record.get("image_path") + images = record.get("images") + if image_field and images: + raise ImageDataError( + f"{source}: record has both 'image' and 'images' fields; use one" + ) + if images is None: + return image_field + if not isinstance(images, list): + raise ImageDataError( + f"{source}: 'images' must be a list, got {type(images).__name__}" + ) + if len(images) == 0: + return image_field + if len(images) > 1: + raise ImageDataError( + f"{source}: multi-image samples are not supported (got " + f"{len(images)} images); the v1 contract is one image per sample" + ) + first = images[0] + if not isinstance(first, str): + raise ImageDataError( + f"{source}: 'images[0]' must be a path or base64 string, got " + f"{type(first).__name__}" + ) + return first + + +def _as_data_uri(raw: bytes, *, media_type: str = "image/jpeg") -> str: + """Wrap raw image bytes as a data URI. + + Some SGLang builds (e.g. the Ascend 0.5.14 line) only route ``data:`` + URIs through base64 decoding and treat every other string as a file + path; the data-URI form is accepted by all supported server versions. + """ + + return f"data:{media_type};base64," + base64.b64encode(raw).decode("ascii") + + +def _load_image(image_field: Any, *, source: str): + """Return (pil_image, data_uri) from a path / base64 / data-URI field.""" + + from PIL import Image + + if not isinstance(image_field, str) or not image_field: + raise ImageDataError( + f"{source}: image field must be a file path or base64 string, got " + f"{type(image_field).__name__}" + ) + if os.path.isfile(image_field): + with open(image_field, "rb") as image_file: + raw = image_file.read() + return Image.open(io.BytesIO(raw)).convert("RGB"), _as_data_uri(raw) + encoded = image_field + if encoded.startswith("data:"): + media_type = encoded[5 : encoded.index(";")] if ";" in encoded else "image/jpeg" + encoded = encoded.split(",", 1)[-1] + else: + media_type = "image/jpeg" + try: + raw = base64.b64decode(encoded, validate=True) + except Exception as exc: + raise ImageDataError( + f"{source}: image field is neither an existing file nor valid base64" + ) from exc + return ( + Image.open(io.BytesIO(raw)).convert("RGB"), + _as_data_uri(raw, media_type=media_type), + ) + + +def _render_conversation_text( + tokenizer, + conversations: List[Dict[str, Any]], + *, + with_image: bool, + source: str, +) -> str: + """Render one ShareGPT conversation with the target tokenizer's template. + + The image (if any) becomes an OpenAI-style list content on the first user + turn so the model's own chat template emits the vision markers exactly as + at serving time. + """ + + messages: List[Dict[str, Any]] = [] + image_attached = False + for message in conversations: + role = message.get("role", message.get("from", "")) + content = message.get("content") + if content is None: + content = message.get("value", "") + if role in ("human", "user"): + role = "user" + elif role in ("gpt", "assistant"): + role = "assistant" + if isinstance(content, list): + # OpenAI-style list content: keep text parts; image parts are + # re-attached by this function (single-image v1). + content = " ".join( + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ) + if not isinstance(content, str): + raise ValueError(f"{source}: unsupported message content type") + if with_image and role == "user" and not image_attached: + content = [ + {"type": "image"}, + {"type": "text", "text": content}, + ] + image_attached = True + messages.append({"role": role, "content": content}) + if with_image and not image_attached: + raise ValueError(f"{source}: image sample has no user turn to attach to") + return tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + ) + + +def _expand_image_region( + ids: List[int], + mask: List[int], + *, + pad_token_id: int, + count: int, + source: str, +) -> Tuple[List[int], List[int]]: + """Splice the single image placeholder into ``count`` copies in id space.""" + + positions = [index for index, token in enumerate(ids) if token == pad_token_id] + if len(positions) != 1: + raise ValueError( + f"{source}: expected exactly 1 image placeholder token, found " + f"{len(positions)}" + ) + index = positions[0] + expanded_ids = ids[:index] + [pad_token_id] * count + ids[index + 1 :] + expanded_mask = mask[:index] + [0] * count + mask[index + 1 :] + return expanded_ids, expanded_mask + + +def build_vlm_prompt_payloads( + path: str, + tokenizer, + processor, + *, + chat_template: Optional[str], + max_length: int, + min_loss_tokens: int = 1, + max_prompts: Optional[int] = None, +) -> List[Dict[str, Any]]: + """Build prompt payloads for multimodal server capture from a JSONL file. + + Records follow the ShareGPT shape (``conversations`` column); an optional + top-level ``image`` field carries one image (file path or base64). Returns + control-plane dicts ``{"payload": {...}}`` mirroring + :func:`specforge.data.prompt_builder.prepare_prompt_tasks`. + """ + + template = TEMPLATE_REGISTRY.get(chat_template) if chat_template else None + pad_token_id = tokenizer.convert_tokens_to_ids(_IMAGE_PLACEHOLDER) + if not isinstance(pad_token_id, int): + raise ValueError( + f"tokenizer has no {_IMAGE_PLACEHOLDER!r} special token; is this a " + "Qwen-VL-family tokenizer?" + ) + merge_size = int(getattr(processor.image_processor, "merge_size", 2)) + + payloads: List[Dict[str, Any]] = [] + skipped = 0 + for line_number, record in _iter_records(path): + source = f"{path}:{line_number}" + try: + prepared = _prepare_one_record( + record, + source=source, + tokenizer=tokenizer, + processor=processor, + template=template, + pad_token_id=pad_token_id, + merge_size=merge_size, + max_length=max_length, + min_loss_tokens=min_loss_tokens, + ) + except ImageDataError: + # Image-contract violations are fatal: never degrade an + # image-bearing sample to text-only training silently. + raise + except ValueError as exc: + print(f"WARNING: skipping {source}: {exc}") + skipped += 1 + continue + if prepared is None: + skipped += 1 + continue + payloads.append(prepared) + if max_prompts not in (None, 0) and len(payloads) >= max_prompts: + break + print( + f"VLM prompt preparation done: {len(payloads)} prepared, " + f"{skipped} skipped ({path})" + ) + return payloads + + +def _prepare_one_record( + record: Dict[str, Any], + *, + source: str, + tokenizer, + processor, + template, + pad_token_id: int, + merge_size: int, + max_length: int, + min_loss_tokens: int, +) -> Optional[Dict[str, Any]]: + from .preprocessing import preprocess_conversations + + conversations = record.get("conversations") + if not conversations: + raise ValueError("record has no 'conversations' field") + image_field = _extract_image_field(record, source=source) + + text = _render_conversation_text( + tokenizer, + conversations, + with_image=bool(image_field), + source=source, + ) + # Reuse the text stack: preformatted rendering -> ids + loss mask. + parsed = preprocess_conversations( + tokenizer, + [text], + template, + max_length=max_length, + is_preformatted=True, + ) + if not parsed["input_ids"]: + return None + collapsed = parsed["input_ids"][0][0].tolist() + mask = parsed["loss_mask"][0][0].tolist() + + image_data = None + if image_field: + pil_image, image_data = _load_image(image_field, source=source) + processor_output = processor.image_processor( + images=[pil_image], return_tensors="pt" + ) + count = _image_token_count( + processor_output["image_grid_thw"].tolist(), merge_size + ) + ids, mask = _expand_image_region( + collapsed, + mask, + pad_token_id=pad_token_id, + count=count, + source=source, + ) + if len(ids) > max_length: + # Truncating into an image region would desynchronize the client's + # expanded ids from the server's own expansion; drop the sample. + raise ValueError( + f"expanded sequence length {len(ids)} exceeds max_length " + f"{max_length} (image region cannot be truncated)" + ) + else: + ids = collapsed + + ids = ids[:max_length] + mask = mask[:max_length] + if sum(mask) < min_loss_tokens: + return None + return { + "payload": { + "input_ids": ids, + "loss_mask": mask, + "request_input_ids": collapsed, + "image_data": image_data, + } + } + + +__all__ = ["ImageDataError", "build_vlm_prompt_payloads"] diff --git a/specforge/inference/media.py b/specforge/inference/media.py index 1d2b289dc..f5baf5f35 100644 --- a/specforge/inference/media.py +++ b/specforge/inference/media.py @@ -1,8 +1,11 @@ # coding=utf-8 """Legacy multimodal input types retained for import compatibility. -The current runtime does not support VLM training or media rollout inputs. -These types are not wired into the canonical training path. +Multimodal training now runs through the server-capture path +(``specforge/algorithms/common/vlm_input.py`` with +``model.input_modality="multimodal"``), which transports images inside the +capture request itself. These pixel-tensor types remain unused by the +canonical training path. """ from __future__ import annotations diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index 5633c3af9..74460483f 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -515,6 +515,16 @@ def _managed_local_services( ): # flashinfer does not exist on Ascend; default to the NPU backend. attention_backend = "ascend" + mm_attention_backend = cfg.model.sglang_mm_attention_backend + if ( + mm_attention_backend is None + and cfg.model.input_modality != "text" + and device_visibility_env == "ASCEND_RT_VISIBLE_DEVICES" + ): + # The sdpa vision backend materializes [heads, N, N] attention + # scores and OOMs on large images; ascend_attn is the fused + # (flash-style) vision attention on NPU. + mm_attention_backend = "ascend_attn" argv.extend( _sglang_argv( cfg.model, @@ -526,6 +536,7 @@ def _managed_local_services( else cfg.model.sglang_mem_fraction_static ), "sglang_attention_backend": attention_backend, + "sglang_mm_attention_backend": mm_attention_backend, }, ) ) @@ -536,6 +547,10 @@ def _managed_local_services( "MOONCAKE_GLOBAL_SEGMENT_SIZE": str(mooncake.global_segment_size_bytes), "MOONCAKE_LOCAL_BUFFER_SIZE": str(mooncake.local_buffer_size_bytes), } + if cfg.model.input_modality == "multimodal": + # Multimodal capture requests carry single-placeholder input_ids; + # the server re-expands them in id space (no retokenization drift). + service_env["SGLANG_MM_AVOID_RETOKENIZE"] = "1" capture_services.append( ServiceSpec( command=CommandSpec( diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 4acd30e11..39f86568e 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -74,8 +74,20 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_len = q.size(-2) - q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) - k_embed = (k * cos) + (rotate_half(k) * sin) + rotary_dim = cos.shape[-1] + head_dim = q.shape[-1] + if rotary_dim < head_dim: + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_rot = (q_rot * cos[..., -q_len:, :]) + ( + rotate_half(q_rot) * sin[..., -q_len:, :] + ) + k_rot = (k_rot * cos) + (rotate_half(k_rot) * sin) + q_embed = torch.cat([q_rot, q_pass], dim=-1) + k_embed = torch.cat([k_rot, k_pass], dim=-1) + else: + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/specforge/modeling/target/target_utils.py b/specforge/modeling/target/target_utils.py index c8ca1afcb..3afdf42f4 100644 --- a/specforge/modeling/target/target_utils.py +++ b/specforge/modeling/target/target_utils.py @@ -67,6 +67,45 @@ def target_text_config(config): return getattr(config, "text_config", config) +#: Target ``model_type`` values whose checkpoints nest the language-model +#: weights under ``model.language_model.*`` (Qwen VLM families). +QWEN3_VL_MODEL_TYPES = {"qwen3_vl", "qwen3_vl_moe", "qwen3_5_moe", "qwen3_5"} + + +def is_vlm_target_config(config) -> bool: + """Whether the target config is a vision-language model with a text stack.""" + model_type = getattr(config, "model_type", None) + return model_type in QWEN3_VL_MODEL_TYPES or hasattr(config, "text_config") + + +#: Default (non-VLM) embedding weight key. +DEFAULT_EMBED_KEY = "model.embed_tokens.weight" +#: Embedding weight key inside VLM checkpoints' language-model stack. +VLM_EMBED_KEY = "model.language_model.embed_tokens.weight" + + +def resolve_target_weight_keys( + config, + embed_key: Optional[str] = None, + lm_head_key: Optional[str] = None, +) -> tuple[str, str]: + """Resolve default embedding / lm_head weight keys for a target config. + + VLM checkpoints (Qwen3-VL, Qwen3.5, Qwen3.6) store the language-model + embedding at ``model.language_model.embed_tokens.weight`` instead of the + LLM default ``model.embed_tokens.weight``; the lm_head stays top-level. + An ``embed_key`` that is unset or left at the LLM default is rewritten + automatically for VLM targets; any other explicit key is honored as-is. + """ + if embed_key is None or ( + embed_key == DEFAULT_EMBED_KEY and is_vlm_target_config(config) + ): + embed_key = VLM_EMBED_KEY if is_vlm_target_config(config) else DEFAULT_EMBED_KEY + if lm_head_key is None: + lm_head_key = "lm_head.weight" + return embed_key, lm_head_key + + def target_vocab_size(config) -> int: text_config = target_text_config(config) return int( @@ -113,10 +152,9 @@ def from_pretrained( ) instance = cls(config) - if embed_key is None: - embed_key = "model.embed_tokens.weight" - if lm_head_key is None: - lm_head_key = "lm_head.weight" + embed_key, lm_head_key = resolve_target_weight_keys( + config, embed_key, lm_head_key + ) # 2. Resolve Model Path local_model_path = model_path diff --git a/tests/test_algorithms/test_builtin_providers.py b/tests/test_algorithms/test_builtin_providers.py index 11a99b5fa..cc10fa18f 100644 --- a/tests/test_algorithms/test_builtin_providers.py +++ b/tests/test_algorithms/test_builtin_providers.py @@ -155,12 +155,20 @@ def test_target_derived_defaults_and_overrides_are_provider_owned(self): self.assertIsNone(policy.target_defaults) self.assertIsNone(policy.apply_overrides) - def test_vlm_is_not_registered_as_a_builtin(self): + def test_builtin_modalities_are_pinned_per_algorithm(self): + expected = { + "dflash": {"text", "multimodal"}, + "domino": {"text"}, + "dspark": {"text"}, + "eagle3": {"text"}, + "mtp": {"text"}, + "peagle": {"text"}, + } for registration in self.registry: modalities = { contract.modality for contract in registration.spec.feature_contracts } - self.assertEqual({"text"}, modalities, registration.name) + self.assertEqual(expected[registration.name], modalities, registration.name) def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): training = SimpleNamespace( diff --git a/tests/test_algorithms/test_dflash_multimodal.py b/tests/test_algorithms/test_dflash_multimodal.py new file mode 100644 index 000000000..21c30cd85 --- /dev/null +++ b/tests/test_algorithms/test_dflash_multimodal.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +from specforge.algorithms.builtin import builtin_algorithm_registry +from specforge.algorithms.common.vlm_input import VlmServerInputAdapter +from specforge.algorithms.contracts import FeatureMode +from specforge.data.vlm_preprocessing import _expand_image_region, _image_token_count + +try: + import torch + + TORCH_AVAILABLE = True +except ImportError: # pragma: no cover - torch-free dev boxes + torch = None + TORCH_AVAILABLE = False + + +class MultimodalRegistrationTest(unittest.TestCase): + def setUp(self): + self.registry = builtin_algorithm_registry() + + def test_dflash_registers_the_multimodal_streaming_contract(self): + registration = self.registry.resolve("dflash") + contract = registration.spec.feature_contract( + FeatureMode.STREAMING, "multimodal" + ) + self.assertEqual( + set(contract.required_tensors), + {"input_ids", "loss_mask", "hidden_states"}, + ) + provider = registration.providers.server_streaming_for("multimodal") + self.assertEqual(provider.capture_method, "dflash") + self.assertEqual(provider.layout.aux_feature, "hidden_states") + + def test_other_builtins_have_no_multimodal_contract(self): + for name in ("domino", "dspark", "eagle3", "peagle"): + with self.subTest(algorithm=name): + registration = self.registry.resolve(name) + with self.assertRaises(KeyError): + registration.spec.feature_contract( + FeatureMode.STREAMING, "multimodal" + ) + + def test_input_adapter_factory_builds_a_valid_adapter(self): + registration = self.registry.resolve("dflash") + provider = registration.providers.server_streaming_for("multimodal") + adapter = provider.create_input_adapter( + SimpleNamespace(model=SimpleNamespace(), data=SimpleNamespace()) + ) + self.assertIsInstance(adapter, VlmServerInputAdapter) + + +class VlmExpansionMathTest(unittest.TestCase): + def test_image_token_count_uses_merge_size(self): + self.assertEqual(_image_token_count([[2, 4, 6]], merge_size=2), 12) + self.assertEqual(_image_token_count([[1, 2, 2]], merge_size=1), 4) + + def test_expand_image_region_splices_ids_and_zero_mask(self): + ids, mask = _expand_image_region( + [10, 99, 20], + [0, 0, 1], + pad_token_id=99, + count=4, + source="test", + ) + self.assertEqual(ids, [10, 99, 99, 99, 99, 20]) + self.assertEqual(mask, [0, 0, 0, 0, 0, 1]) + + def test_expand_image_region_requires_exactly_one_placeholder(self): + with self.assertRaises(ValueError): + _expand_image_region([10, 20], [1, 1], pad_token_id=99, count=4, source="t") + with self.assertRaises(ValueError): + _expand_image_region( + [99, 10, 99], [0, 0, 0], pad_token_id=99, count=4, source="t" + ) + + +class ExtractImageFieldTest(unittest.TestCase): + def _extract(self, record): + from specforge.data.vlm_preprocessing import _extract_image_field + + return _extract_image_field(record, source="t") + + def test_image_and_image_path_fields(self): + self.assertEqual(self._extract({"image": "a.jpg"}), "a.jpg") + self.assertEqual(self._extract({"image_path": "b.jpg"}), "b.jpg") + self.assertIsNone(self._extract({})) + + def test_images_list_takes_the_single_element(self): + self.assertEqual(self._extract({"images": ["a.jpg"]}), "a.jpg") + self.assertIsNone(self._extract({"images": []})) + + def test_multi_image_sample_is_fatal(self): + from specforge.data.vlm_preprocessing import ImageDataError + + with self.assertRaises(ImageDataError): + self._extract({"images": ["a.jpg", "b.jpg"]}) + + def test_non_list_images_and_non_string_element_are_fatal(self): + from specforge.data.vlm_preprocessing import ImageDataError + + with self.assertRaises(ImageDataError): + self._extract({"images": "a.jpg"}) + with self.assertRaises(ImageDataError): + self._extract({"images": [123]}) + + def test_conflicting_image_fields_are_fatal(self): + from specforge.data.vlm_preprocessing import ImageDataError + + with self.assertRaises(ImageDataError): + self._extract({"image": "a.jpg", "images": ["b.jpg"]}) + + def test_unreadable_image_is_fatal_not_skipped(self): + from specforge.data.vlm_preprocessing import ImageDataError, _load_image + + with self.assertRaises(ImageDataError): + _load_image("/nonexistent/path/to/image.jpg", source="t") + with self.assertRaises(ImageDataError): + _load_image("not-valid-base64!!!", source="t") + + def test_load_image_returns_data_uri(self): + import base64 as b64mod + + from specforge.data.vlm_preprocessing import _load_image + + # 1x1 white PNG + png_b64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4" + "z8DwHwAFAAH/q842iQAAAABJRU5ErkJggg==" + ) + _, uri = _load_image(png_b64, source="t") + self.assertTrue(uri.startswith("data:image/jpeg;base64,")) + raw = b64mod.b64decode(uri.split(",", 1)[1]) + self.assertEqual(raw, b64mod.b64decode(png_b64)) + + def test_load_image_preserves_input_media_type(self): + from specforge.data.vlm_preprocessing import _load_image + + png_b64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4" + "z8DwHwAFAAH/q842iQAAAABJRU5ErkJggg==" + ) + _, uri = _load_image(f"data:image/png;base64,{png_b64}", source="t") + self.assertTrue(uri.startswith("data:image/png;base64,")) + + +class VlmRequestInputsTest(unittest.TestCase): + def test_build_request_inputs_uses_collapsed_ids_and_image_data(self): + adapter = VlmServerInputAdapter(config=SimpleNamespace()) + tasks = [ + SimpleNamespace( + payload={ + "input_ids": [1, 2, 2, 2, 3], + "request_input_ids": [1, 2, 3], + "image_data": "aGVsbG8=", + } + ), + SimpleNamespace( + payload={ + "input_ids": [7, 8], + "request_input_ids": [7, 8], + "image_data": None, + } + ), + ] + request = adapter.build_request_inputs(tasks) + self.assertEqual(request["input_ids"], [[1, 2, 3], [7, 8]]) + self.assertEqual(request["image_data"], ["aGVsbG8=", None]) + + +@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") +class DraftPositionsTest(unittest.TestCase): + """Drafts always train on the plain 1D position convention, identical to + the text-only path; multimodal capture stores no position ids.""" + + def _build_model(self): + import torch as t + from torch import nn + + from specforge.algorithms.common.dflash_family_model import OnlineDFlashModel + + class _StubDraftModel(nn.Module): + def __init__(self): + super().__init__() + self.recorded = {} + self.sliding_window = None + + def forward( + self, + position_ids=None, + noise_embedding=None, + target_hidden=None, + attention_mask=None, + ): + self.recorded["position_ids"] = position_ids + return t.zeros(1) + + return OnlineDFlashModel( + draft_model=_StubDraftModel(), + target_lm_head=nn.Linear(8, 32, bias=False), + target_embed_tokens=nn.Embedding(32, 8), + mask_token_id=31, + block_size=2, + attention_backend="sdpa", + num_anchors=4, + ) + + def test_positions_follow_the_1d_convention(self): + import torch as t + + model = self._build_model() + b, s = 2, 8 + input_ids = t.randint(0, 31, (b, s)) + hidden_states = t.randn(b, s, 16) + loss_mask = t.ones(b, s) + t.manual_seed(0) + model._forward_draft_blocks(input_ids, hidden_states, loss_mask) + got = model.draft_model.recorded["position_ids"] + self.assertEqual(got.ndim, 2) + self.assertTrue(t.equal(got[:, :s], t.arange(s).unsqueeze(0).expand(b, -1))) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_config/test_schema.py b/tests/test_config/test_schema.py index 03edbbd6a..690835fba 100644 --- a/tests/test_config/test_schema.py +++ b/tests/test_config/test_schema.py @@ -648,6 +648,48 @@ def test_override_bad_path_or_form_raises(self): with self.assertRaises(ValueError): apply_overrides(cfg, ["not-an-assignment"]) + def test_override_creates_none_optional_sections(self): + payload = { + "model": { + "target_model_path": "some/target", + "draft_model_config": "draft.json", + }, + "data": {"train_data_path": "train.jsonl"}, + "deployment": { + "mode": "disaggregated", + "trainer": {"nnodes": 1, "nproc_per_node": 2}, + "disaggregated": { + "control_dir": "outputs/x/control", + "backend": "mooncake", + "server_urls": ["http://127.0.0.1:30000"], + }, + }, + } + cfg = Config.model_validate(payload) + out = apply_overrides( + cfg, + [ + # managed_local and server_urls are mutually exclusive. + "deployment.disaggregated.server_urls=[]", + 'deployment.disaggregated.managed_local.trainer_cuda_visible_devices=["8","9"]', + 'deployment.disaggregated.managed_local.capture_servers=[{port: 40000, cuda_visible_devices: ["0"], tp_size: 1}]', + ], + ) + self.assertEqual(out.deployment.disaggregated.server_urls, []) + managed = out.deployment.disaggregated.managed_local + self.assertIsNotNone(managed) + self.assertEqual(managed.trainer_cuda_visible_devices, ["8", "9"]) + self.assertEqual(len(managed.capture_servers), 1) + self.assertEqual(managed.capture_servers[0].port, 40000) + # Defaults fill in untouched managed sub-configs. + self.assertGreater(managed.mooncake.rpc_port, 0) + # Typo'd paths are still rejected by re-validation. + with self.assertRaises(Exception): + apply_overrides( + cfg, + ["deployment.disaggregated.managed_local.no_such_field=1"], + ) + def test_load_config_applies_overrides(self): path = _write(MINIMAL, ".json") self.addCleanup(os.unlink, path) diff --git a/tests/test_config/test_server_only_online.py b/tests/test_config/test_server_only_online.py index c12c5106b..a5bc1443f 100644 --- a/tests/test_config/test_server_only_online.py +++ b/tests/test_config/test_server_only_online.py @@ -66,6 +66,24 @@ def test_vlm_is_explicitly_unsupported(self): ): resolve_run(config) + def test_multimodal_modality_resolves_for_dflash(self): + payload = _online_payload(modality="multimodal") + payload["training"]["strategy"] = "dflash" + del payload["model"]["vocab_mapping_path"] + resolved = resolve_run(Config.model_validate(payload)) + self.assertEqual(resolved.algorithm.name, "dflash") + streaming = resolved.algorithm.providers.server_streaming_for("multimodal") + self.assertEqual(streaming.layout.aux_feature, "hidden_states") + self.assertIsNotNone(streaming.create_input_adapter(resolved.config)) + + def test_multimodal_modality_is_rejected_for_text_only_algorithms(self): + payload = _online_payload(modality="multimodal") + with self.assertRaisesRegex( + ValueError, + "no streaming feature contract and provider for modality " "'multimodal'", + ): + resolve_run(Config.model_validate(payload)) + def test_offline_configs_reject_retired_backends_instead_of_ignoring_them(self): # Offline consumers never instantiate a target inference backend, but a # config naming a retired backend must fail at load rather than be