Эх сурвалжийг харах

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 сар өмнө
parent
commit
8d49ca7159

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

@@ -3,8 +3,10 @@ FROM ${BASE_IMAGE}
 
 
 COPY plugin /opt/dsv4_vision_src
 COPY plugin /opt/dsv4_vision_src
 COPY patch_dspark_skip_bias_vl.py /tmp/patch_dspark_skip_bias_vl.py
 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 \
 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 VLLM_PLUGINS=dsv4_vision
 ENV TORCH_CUDA_ARCH_LIST=12.1a
 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."""
 """Skip Vision-Exp gate tensors the B12X DSpark draft loader cannot bind."""
 from pathlib import Path
 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()
 text = p.read_text()
 old = '''                if name.endswith(".ffn.gate.bias"):
 old = '''                if name.endswith(".ffn.gate.bias"):
                     name = name.replace(
                     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
 Loads official `vision.*` / `aligner.*` / image sentinel tensors and merges
 ViT embeddings on prefill. Language decode stays on the B12X DeepSeek V4 path.
 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
 from __future__ import annotations
 
 
+import os
 from types import SimpleNamespace
 from types import SimpleNamespace
 from typing import Any
 from typing import Any
 
 
 import torch
 import torch
 import torch.nn as nn
 import torch.nn as nn
 
 
-import torch
-import torch.nn as nn
-
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.models.deepseek_v4 import DeepseekV4ForCausalLM
 from vllm.models.deepseek_v4 import DeepseekV4ForCausalLM
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -26,6 +30,19 @@ from .processor import (
 )
 )
 from .vision import Aligner, ViT
 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:
 def _vision_args(hf: Any) -> SimpleNamespace:
     return 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):
 class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
     # ModelConfig.is_multimodal_model keys off this class attribute
     # ModelConfig.is_multimodal_model keys off this class attribute
     # (interfaces.supports_multimodal) — without it the OpenAI layer
     # (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
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> str | None:
     def get_placeholder_str(cls, modality: str, i: int) -> str | None:
         # chat_utils asks the resolved model class for the per-image prompt
         # chat_utils asks the resolved model class for the per-image prompt
         # placeholder; the tokenizer encodes it to <|deepseek_image|> (id 129264).
         # 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
             from .encoding_dsv4 import IMAGE_PLACEHOLDER
 
 
             return 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:
     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)
         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(
     def embed_input_ids(
         self,
         self,
         input_ids: torch.Tensor,
         input_ids: torch.Tensor,
@@ -105,17 +111,25 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
         *,
         *,
         is_multimodal: torch.Tensor | None = None,
         is_multimodal: torch.Tensor | None = None,
     ) -> torch.Tensor:
     ) -> 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:
         if is_multimodal is not None:
             # Sentinel block ids may exceed vocab; zero them before lookup.
             # Sentinel block ids may exceed vocab; zero them before lookup.
             input_ids = input_ids.masked_fill(
             input_ids = input_ids.masked_fill(
                 is_multimodal.to(input_ids.device).unsqueeze(-1), 0
                 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:
     def embed_multimodal(self, **kwargs) -> tuple:
         patches = kwargs.get("image_patches")
         patches = kwargs.get("image_patches")
@@ -124,7 +138,7 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
         grids = kwargs["image_grids"].data
         grids = kwargs["image_grids"].data
         types_all = kwargs["image_types"].data
         types_all = kwargs["image_types"].data
         perm_all = kwargs["image_perm"].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()
         perm_counts = kwargs["image_perm_counts"].data.tolist()
         n_items = int(grids.shape[0])
         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
         # The runner embeds via embed_input_ids (mm merge) and hands us
         # merged inputs_embeds; lm_head and get_mtp_target_hidden_states
         # merged inputs_embeds; lm_head and get_mtp_target_hidden_states
         # stay on the parent class.
         # 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(
         return super().forward(
             input_ids,
             input_ids,
             positions,
             positions,

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

@@ -346,7 +346,7 @@ class DSV4VisionMultiModalProcessor(BaseMultiModalProcessor[DSV4VisionProcessing
         return [
         return [
             PromptReplacement(
             PromptReplacement(
                 modality="image",
                 modality="image",
-                target=token_id,
+                target=[token_id],
                 replacement=replacement,
                 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:
     def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
         x = self.patch_embed(patches)
         x = self.patch_embed(patches)
+        device = x.device
         cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta)
         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:
         for block in self.blocks:
             x = block(x, cos, sin)
             x = block(x, cos, sin)
         return self.norm(x)
         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:
 # DEVIATIONS from the sfxnz run.sh defaults — all forced by our launcher:
 #   * no --host/--port/--nnodes/--node-rank/--master-* (launcher fabric),
 #   * no --host/--port/--nnodes/--node-rank/--master-* (launcher fabric),
 #   * no NCCL/* / IFACE / HCA env (launcher resolves; RoCE on sw 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
 # 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).
 # fused Gram mHC kernel only accepts 1e-6 — upstream README).
 #
 #
@@ -88,6 +90,10 @@ args = [
 [env]
 [env]
 # Vision + engine lane knobs (sfxnz run.sh, verbatim minus fabric)
 # Vision + engine lane knobs (sfxnz run.sh, verbatim minus fabric)
 CUTE_DSL_ARCH = "sm_121a"
 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_AOT_COMPILE = "1"
 VLLM_USE_BREAKABLE_CUDAGRAPH = "0"
 VLLM_USE_BREAKABLE_CUDAGRAPH = "0"
 VLLM_USE_MEGA_AOT_ARTIFACT = "-1"
 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
 # NOTE: extra_mounts stay AFTER args (TOML table headers would swallow the
 # args array otherwise).
 # args array otherwise).
 [[extra_mounts]]
 [[extra_mounts]]
-host = "/mnt/nas/ai/huggingface"
-container = "/mnt/nas/ai/huggingface"
+host = "~/dsv4vision-hf"
+container = "/opt/local-hf"
 env = "HF_HOME"
 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]]
 [[extra_mounts]]
 host = "~/.cache/dsv4vision-jit/triton"
 host = "~/.cache/dsv4vision-jit/triton"
 container = "/root/.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.
   dsv4_vision.image_processor; fork interface = MultimodalRegistry.
   register_processor factory in multimodal/registry.py), or wait for the
   register_processor factory in multimodal/registry.py), or wait for the
   upstream repo to ship it. aiden-fork A/B build launched.
   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.