Переглянути джерело

dsv4vision: env-gated openai mm branch (VLLM_DSV4_MM), worker shape-follow patches, local-NVMe weights (dgx3+dgx4), cache-clear runbook; text lane verified exact on new boots

kilork 1 місяць тому
батько
коміт
8d49ca7159

+ 3 - 1
docker-images/dsv4-flash-vision/Dockerfile

@@ -3,8 +3,10 @@ FROM ${BASE_IMAGE}
 
 COPY plugin /opt/dsv4_vision_src
 COPY patch_dspark_skip_bias_vl.py /tmp/patch_dspark_skip_bias_vl.py
+COPY patch_mm_hc_streams.py /tmp/patch_mm_hc_streams.py
 RUN python3 -m pip install --no-cache-dir --no-deps /opt/dsv4_vision_src \
-    && python3 /tmp/patch_dspark_skip_bias_vl.py
+    && python3 /tmp/patch_dspark_skip_bias_vl.py \
+    && python3 /tmp/patch_mm_hc_streams.py
 
 ENV VLLM_PLUGINS=dsv4_vision
 ENV TORCH_CUDA_ARCH_LIST=12.1a

+ 9 - 1
docker-images/dsv4-flash-vision/patch_dspark_skip_bias_vl.py

@@ -2,7 +2,15 @@
 """Skip Vision-Exp gate tensors the B12X DSpark draft loader cannot bind."""
 from pathlib import Path
 
-p = Path("/usr/local/lib/python3.12/dist-packages/vllm/models/deepseek_v4/nvidia/dspark.py")
+import vllm
+
+base = Path(vllm.__file__).parent / "models" / "deepseek_v4" / "nvidia"
+for name in ("dspark.py", "dspark_v2.py"):
+    if (base / name).exists():
+        p = base / name
+        break
+else:
+    raise SystemExit("no dspark.py/dspark_v2.py under " + str(base))
 text = p.read_text()
 old = '''                if name.endswith(".ffn.gate.bias"):
                     name = name.replace(

+ 83 - 0
docker-images/dsv4-flash-vision/patch_mm_hc_streams.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Make the b12x runner's multimodal `inputs_embeds` buffers follow the
+model's native embed shape instead of a fixed 2-D [T, H].
+
+At boot the warmup probed `x` shapes of [16, 16, 4096] (hc-stream run) and
+[14, 4096] (plain run) from the SAME layer — the dsv4 b12x embedding emits
+varying rank per request. Re-fit the buffer whenever the non-token tail
+changes so both the mm merge and CUDA-graph capture see a stable tensor.
+"""
+from pathlib import Path
+
+import vllm
+
+p = (
+    Path(vllm.__file__).parent
+    / "v1/worker/gpu/mm/encoder_runner.py"
+)
+text = p.read_text()
+old = """        x = self.model.embed_input_ids(
+            input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed
+        )
+        # Copy to the pre-allocated buffer for CUDA graphs.
+        self.inputs_embeds[: x.shape[0]] = x
+        return self.inputs_embeds
+"""
+new = """        x = self.model.embed_input_ids(
+            input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed
+        )
+        # dsv4 b12x embeds vary in rank per request ([T, K, H] hc-stream runs,
+        # [T, H] plain): re-fit the buffer when the non-token tail changes.
+        tail = tuple(x.shape[1:])
+        if (
+            self.inputs_embeds.dim() != x.dim()
+            or tuple(self.inputs_embeds.shape[1:]) != tail
+        ):
+            self.inputs_embeds = torch.zeros(
+                self.max_num_tokens, *tail, dtype=x.dtype, device=x.device
+            )
+        # Copy to the pre-allocated buffer for CUDA graphs.
+        self.inputs_embeds[: x.shape[0]] = x
+        return self.inputs_embeds
+"""
+if old not in text:
+    raise SystemExit(f"encoder_runner patch target missing in {p}")
+p.write_text(text.replace(old, new, 1))
+print("encoder_runner mm shape-follow patch applied")
+
+# Same shape-follow for the DSpark speculator's re-embed buffer: with the
+# multimodal gate on, the MTP speculator re-embeds draft tokens through the
+# target's embed and assigns [B, T, H] rows into a [max, hidden] buffer.
+sp = (
+    Path(vllm.__file__).parent
+    / "v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py"
+)
+sp_text = sp.read_text()
+sp_old = """        if self.inputs_embeds is not None:
+            mm_embeds, is_mm_embed = mm_inputs or (None, None)
+            self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(
+                self.input_buffers.input_ids[:num_tokens],
+                multimodal_embeddings=mm_embeds,
+                is_multimodal=is_mm_embed,
+            )
+"""
+sp_new = """        if self.inputs_embeds is not None:
+            mm_embeds, is_mm_embed = mm_inputs or (None, None)
+            emb = self.model.embed_input_ids(
+                self.input_buffers.input_ids[:num_tokens],
+                multimodal_embeddings=mm_embeds,
+                is_multimodal=is_mm_embed,
+            )
+            if emb.dim() != self.inputs_embeds.dim() or tuple(
+                emb.shape[1:]
+            ) != tuple(self.inputs_embeds.shape[1:]):
+                self.inputs_embeds = torch.zeros_like(
+                    emb, dtype=emb.dtype, device=emb.device
+                )
+            self.inputs_embeds[:num_tokens] = emb
+"""
+if sp_old in sp_text:
+    sp.write_text(sp_text.replace(sp_old, sp_new, 1))
+    print("speculator re-embed shape-follow patch applied")
+else:
+    print("speculator re-embed pattern not found — skipping (harmless)")

+ 51 - 32
docker-images/dsv4-flash-vision/plugin/dsv4_vision/model.py

@@ -2,19 +2,23 @@
 
 Loads official `vision.*` / `aligner.*` / image sentinel tensors and merges
 ViT embeddings on prefill. Language decode stays on the B12X DeepSeek V4 path.
+
+The OpenAI multimodal branch (processor trio + embed-merge hooks) is gated
+behind VLLM_DSV4_MM=1: the fork's V2 "always-embed" mm arm routes through
+this override with a [B, T]-ranks ids slice and currently degenerates
+logits, so with VLLM_DSV4_MM unset (the default) the lane boots text-only —
+the verified state — and VLLM_DSV4_MM=1 opts into the experimental path.
 """
 
 from __future__ import annotations
 
+import os
 from types import SimpleNamespace
 from typing import Any
 
 import torch
 import torch.nn as nn
 
-import torch
-import torch.nn as nn
-
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.models.deepseek_v4 import DeepseekV4ForCausalLM
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -26,6 +30,19 @@ from .processor import (
 )
 from .vision import Aligner, ViT
 
+_MM_ON = os.environ.get("VLLM_DSV4_MM", "0") == "1"
+
+
+def _maybe_mm(cls):
+    """Register the OpenAI multimodal processor trio only when mm is on."""
+    if not _MM_ON:
+        return cls
+    return MULTIMODAL_REGISTRY.register_processor(
+        DSV4VisionMultiModalProcessor,
+        info=DSV4VisionProcessingInfo,
+        dummy_inputs=DSV4VisionDummyInputsBuilder,
+    )(cls)
+
 
 def _vision_args(hf: Any) -> SimpleNamespace:
     return SimpleNamespace(
@@ -44,22 +61,25 @@ def _vision_args(hf: Any) -> SimpleNamespace:
     )
 
 
-@MULTIMODAL_REGISTRY.register_processor(
-    DSV4VisionMultiModalProcessor,
-    info=DSV4VisionProcessingInfo,
-    dummy_inputs=DSV4VisionDummyInputsBuilder,
-)
+@_maybe_mm
 class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
     # ModelConfig.is_multimodal_model keys off this class attribute
     # (interfaces.supports_multimodal) — without it the OpenAI layer
-    # answers "not a multimodal model" for image_url requests.
-    supports_multimodal = True
+    # answers "not a multimodal model" for image_url requests. Gated by
+    # _maybe_mm so the text-only default lane stays verified.
+    supports_multimodal = _MM_ON
+    # The fork's hash MoE routes by token id (tid2eid gather), so the merged
+    # mm path must keep input_ids alongside the merged embeddings.
+    requires_raw_input_tokens = _MM_ON
+    supports_multimodal_raw_input_only = False
+    supports_encoder_tp_data = False
+    supports_mm_device_do_normalize = False
 
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> str | None:
         # chat_utils asks the resolved model class for the per-image prompt
         # placeholder; the tokenizer encodes it to <|deepseek_image|> (id 129264).
-        if modality.startswith("image"):
+        if _MM_ON and modality.startswith("image"):
             from .encoding_dsv4 import IMAGE_PLACEHOLDER
 
             return IMAGE_PLACEHOLDER
@@ -84,20 +104,6 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
     def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
         return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
 
-    def _merge_mm_embeddings(
-        self,
-        inputs_embeds: torch.Tensor,
-        multimodal_embeddings,
-        is_multimodal: torch.Tensor,
-    ) -> torch.Tensor:
-        # Mirrors interfaces._merge_multimodal_embeddings: fill masked rows in
-        # order with the concatenated per-item embedding runs.
-        flat = torch.cat([emb.reshape(-1) for emb in multimodal_embeddings])
-        inputs_embeds.masked_scatter_(
-            is_multimodal.to(inputs_embeds.device).unsqueeze(-1), flat
-        )
-        return inputs_embeds
-
     def embed_input_ids(
         self,
         input_ids: torch.Tensor,
@@ -105,17 +111,25 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
         *,
         is_multimodal: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        # The fork's embed output follows the input rank: [B, T, H] from the
+        # speculator's [B, T] ids, [T, H] from plain [T] ids. mm rows must
+        # land on the same rows either way (they never exceed inner dims).
         if is_multimodal is not None:
             # Sentinel block ids may exceed vocab; zero them before lookup.
             input_ids = input_ids.masked_fill(
                 is_multimodal.to(input_ids.device).unsqueeze(-1), 0
             )
-        inputs_embeds = super().embed_input_ids(input_ids)
-        if not multimodal_embeddings:
-            return inputs_embeds
-        return self._merge_mm_embeddings(
-            inputs_embeds, multimodal_embeddings, is_multimodal
-        )
+        out = super().embed_input_ids(input_ids)
+        if multimodal_embeddings:
+            mm_stack = torch.cat(multimodal_embeddings)  # [N, H]
+            if out.dim() == 3:
+                out[is_multimodal] = mm_stack.to(out.dtype)
+            else:
+                out.masked_scatter_(
+                    is_multimodal.to(out.device).unsqueeze(-1),
+                    mm_stack.reshape(-1),
+                )
+        return out
 
     def embed_multimodal(self, **kwargs) -> tuple:
         patches = kwargs.get("image_patches")
@@ -124,7 +138,7 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
         grids = kwargs["image_grids"].data
         types_all = kwargs["image_types"].data
         perm_all = kwargs["image_perm"].data
-        t_counts = kwargs["image_type_counts"].data.tolist() if hasattr(kwargs["image_type_counts"], "tolist") else kwargs["image_type_counts"].data.tolist()
+        t_counts = kwargs["image_type_counts"].data.tolist()
         perm_counts = kwargs["image_perm_counts"].data.tolist()
         n_items = int(grids.shape[0])
 
@@ -179,6 +193,11 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
         # The runner embeds via embed_input_ids (mm merge) and hands us
         # merged inputs_embeds; lm_head and get_mtp_target_hidden_states
         # stay on the parent class.
+        if not _MM_ON and inputs_embeds is not None:
+            # text-only belt-and-braces: without the mm gate the parent never
+            # sees runner-generated embeds, but if one sneaks in, sync the
+            # token path instead so hidden shapes stay native.
+            inputs_embeds = None
         return super().forward(
             input_ids,
             positions,

+ 1 - 1
docker-images/dsv4-flash-vision/plugin/dsv4_vision/processor.py

@@ -346,7 +346,7 @@ class DSV4VisionMultiModalProcessor(BaseMultiModalProcessor[DSV4VisionProcessing
         return [
             PromptReplacement(
                 modality="image",
-                target=token_id,
+                target=[token_id],
                 replacement=replacement,
             )
         ]

+ 3 - 0
docker-images/dsv4-flash-vision/plugin/dsv4_vision/vision.py

@@ -96,7 +96,10 @@ class ViT(nn.Module):
 
     def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
         x = self.patch_embed(patches)
+        device = x.device
         cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta)
+        # lru_cache builds these on CPU; haul them to the model device per image.
+        cos, sin = cos.to(device, non_blocking=True), sin.to(device, non_blocking=True)
         for block in self.blocks:
             x = block(x, cos, sin)
         return self.norm(x)

+ 19 - 5
recipes/dsv4-flash-vision/recipe.toml

@@ -20,9 +20,11 @@
 # DEVIATIONS from the sfxnz run.sh defaults — all forced by our launcher:
 #   * no --host/--port/--nnodes/--node-rank/--master-* (launcher fabric),
 #   * no NCCL/* / IFACE / HCA env (launcher resolves; RoCE on sw fabric),
-#   * weights stream from the shared NAS hub (their kit uses per-node local
-#     NVMe; if the loader balloons memory like it did on EXL3, flip to
-#     --load-format instanttensor).
+#   * weights from per-node local NVMe copies (`~/dsv4vision-hf`, 167 GiB
+#     each on dgx3/dgx4 — rsync'd from the NAS hub 2026-09-01); shard load
+#     drops from ~7 min (NAS) to well under a minute. Fallback: re-add the
+#     `[[extra_mounts]] host = "/mnt/nas/ai/huggingface" … env = "HF_HOME"`
+#     block below to stream from the NAS hub again.
 # VLLM_USE_B12X_MHC stays OFF (this checkpoint has rms_norm_eps 1e-20; the
 # fused Gram mHC kernel only accepts 1e-6 — upstream README).
 #
@@ -88,6 +90,10 @@ args = [
 [env]
 # Vision + engine lane knobs (sfxnz run.sh, verbatim minus fabric)
 CUTE_DSL_ARCH = "sm_121a"
+# AOT compile stays ON for the text-only lane (verified). The experimental
+# VLLM_DSV4_MM=1 mode traces its embed path with a 2-D dummy and then rejects
+# the hc-stream 3-D form ("wrong number of dimensions3 for op: input") —
+# flip this to "0" when testing that mode.
 VLLM_USE_AOT_COMPILE = "1"
 VLLM_USE_BREAKABLE_CUDAGRAPH = "0"
 VLLM_USE_MEGA_AOT_ARTIFACT = "-1"
@@ -109,10 +115,18 @@ VLLM_NO_USAGE_STATS = "1"
 # NOTE: extra_mounts stay AFTER args (TOML table headers would swallow the
 # args array otherwise).
 [[extra_mounts]]
-host = "/mnt/nas/ai/huggingface"
-container = "/mnt/nas/ai/huggingface"
+host = "~/dsv4vision-hf"
+container = "/opt/local-hf"
 env = "HF_HOME"
 
+# Weights from the shared NAS hub; JIT caches persist per node.
+# NOTE: extra_mounts stay AFTER args (TOML table headers would swallow the
+# args array otherwise).
+# [[extra_mounts]]
+# host = "/mnt/nas/ai/huggingface"
+# container = "/mnt/nas/ai/huggingface"
+# env = "HF_HOME"
+
 [[extra_mounts]]
 host = "~/.cache/dsv4vision-jit/triton"
 container = "/root/.triton"

+ 27 - 0
specs/tech-architecture/glm53-recipe-state.md

@@ -260,3 +260,30 @@ docker-images/dsv4-flash-vision/, bases digest-pinned; sm121 build works,
   dsv4_vision.image_processor; fork interface = MultimodalRegistry.
   register_processor factory in multimodal/registry.py), or wait for the
   upstream repo to ship it. aiden-fork A/B build launched.
+
+## dsv4vision MM-into-OpenAI campaign — 2026-09-01 checkpoint
+
+Built the full OpenAI multimodal stack in our plugin (processor trio,
+grid/sentinel tokens, embed_multimodal + embed_input_ids merge; verified in
+-image: embedding rows byte-exact vs aligner(ViT(patches))[perm]) and got the
+engine through profiling+warmup with mm registered. Two hard walls remain
+(before the first image_url token):
+
+1. The fork's V2 "always-embed" mm arm passes `input_batch.input_ids[:n]` —
+   a wrong-rank slice that arrives as [B, T]-shaped 2-D ids (observed
+   (17,17) for a 17-token prompt, (7,7) on DSpark-6 verify steps) — and
+   downstream logits degenerate (all-`<|begin▁of▁sentence|>`).
+2. The AOT/dynamo trace of inputs_embeds freezes the dummy's 2-D shape and
+   then asserts "wrong number of dimensions3 for op: input" on the native
+   3-D stream (tilelang mhc kernels are rank-strict both ways).
+
+State: the whole mm branch is env-gated (`VLLM_DSV4_MM=1`), default OFF =
+the verified text-only lane — restored EXACT on the new per-node NVMe
+weights (~/dsv4vision-hf on dgx3+dgx4, 167 GiB each, recipe HF_HOME flipped
+from the NAS hub; shard load 7 min → <1 min, boot ~6 min end-to-end).
+Bridges done tonight for the next attempt: encoder-runner + speculator
+buffer shape-follow patches (patch_mm_hc_streams.py), placeholder
+`<|deepseek_image|>` = in-vocab id 129264, `get_placeholder_str` chat
+integration, AOT=0 knob for mm experiments. nid: fork A/B (aiden base)
+still pending its pip-layout patch; dspark_gate patch now auto-finds the
+fork's venv root.