feat(vision): trainable stage-1 vision projector with frozen encoder and LM

Implements the first genuinely trainable CORTEX vision pipeline, isolated from the
DeepSeek-V4-Pro checkpoint:

IMAGE -> real SigLIP encoder (frozen) -> VisionProjector (trained)
-> PrefixFusion -> small real causal LM (frozen) -> logits

Only VisionProjector is trained. On the CPU sandbox the run converges from loss
7.5445 to 5.1408 over 60 steps on real Flickr30k image/caption pairs, with the
encoder and LM parameter fingerprints unchanged at exactly 0.

- cortex_ai/vision/train.py: config, freezing helpers, fused labels/mask, training
loop, evidence report and separate projector checkpoint.
- examples/train_vision_projector.py + configs/vision_projector.json: reproducible
CPU entry point; every config field is also a CLI flag.
- VisionProjector gains a two-layer 'mlp' type; 'linear' stays the default.
- Fix the same BatchFeature/UserDict bug in VisionAdapter.encode_images that was
fixed in native.py, and extend it with encode_pixels for preprocessed batches.
- tests/test_projector_training.py: 9 tests covering gradient flow, frozen-weight
preservation, projector change, before/after execution, image sensitivity and
mask/label alignment. One real-model test is opt-in via CORTEX_REAL_VISION_TEST=1.
- tools/shard_integrity.py: hashes the 66 shard LFS pointers so weight integrity is
checkable without downloading ~893 GB.
- docs/vision_projector_training.md: explicit PROVEN vs NOT PROVEN scope.

Not proven: that this projector works with the real 1.65T DeepSeek-V4-Pro core
(7168-dim), native multimodal understanding, or final CORTEX performance.

Co-authored-by: openhands <openhands@all-hands.dev>

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ipynb_checkpoints/
+
+# Stage-1 projector checkpoints are reproducible artefacts, not source.
+checkpoints/
+*.pt

AGENTS.md

@@ -88,8 +88,46 @@ tower): 196 image tokens + 8 text tokens -> 204-length fused sequence, logits sh
 ## Test commands
 
 ```bash
-PYTHONPATH=/workspace/project python3 -m pytest tests/ -q           # 6 tests
-cd encoding && PYTHONPATH=$PWD python3 -m pytest test_encoding_dsv4.py -q   # 4 tests
+PYTHONPATH=/workspace/project python3 -m pytest tests/ -q                   # 14 passed, 1 skipped
+PYTHONPATH=/workspace/project/encoding python3 -m pytest encoding/test_encoding_dsv4.py -q  # 4 tests
+CORTEX_REAL_VISION_TEST=1 python3 -m pytest tests/ -q -k real               # real SigLIP + pythia
+python tools/shard_integrity.py --from-worktree --out /tmp/after.json       # 66 shard hashes
 ```
 
-Requires `torch` (CPU wheel), `torchvision`, `pillow`, `transformers`.
+Requires `torch` (CPU wheel), `torchvision`, `pillow`, `transformers`, and `datasets` for training.
+
+## Stage-1 projector training (implemented and verified)
+
+`cortex_ai/vision/train.py` + `examples/train_vision_projector.py` train **only** the projector:
+
+```
+IMAGE -> frozen SigLIP -> VisionProjector (trained) -> PrefixFusion -> frozen causal LM -> logits
+```
+
+Verified numbers on this CPU sandbox (4 cores, ~0.8 s/step at batch size 2):
+
+| Quantity | Value |
+|---|---|
+| vision_hidden_size | 768 (`google/siglip-base-patch16-224` vision tower, 196 tokens/image) |
+| cortex_hidden_size | 128 (`EleutherAI/pythia-14m`) |
+| projector | `mlp`, 768 -> 128, **114,944 params** |
+| fused sequence | 220 = 196 image + 24 text |
+| 60-step run | loss 7.5445 -> 5.1408; projector delta 669.6; encoder/LM delta exactly 0 |
+
+Invariants covered by `tests/test_projector_training.py` (9 tests): projector gets gradients and
+frozen modules get none; projector params change while encoder/LM are byte-identical; the pipeline
+runs before and after training; different images give different fused representations; the fused
+attention mask and labels align with the fused sequence.
+
+Gotchas found while building this:
+
+- `pythia-14m` loads as **float16** by default, which breaks CPU LayerNorm against an fp32
+  projector. Always pin `dtype=torch.float32` explicitly.
+- Passing the **text** attention mask into a fused forward pass fails with
+  `size of tensor a (204) must match tensor b (8)`. Image-token positions need their own mask
+  entries; see `fused_attention_mask`.
+- Causal-LM loss is reported at *every* position, so image-token positions must be set to `-100`
+  or the projector is trained to predict text from its own prefix.
+
+NOT proven: that this projector works with the real DeepSeek-V4-Pro core (7168-dim), native
+multimodal understanding of the 1.65 T checkpoint, or final CORTEX AI performance.

configs/vision_projector.json

@@ -0,0 +1,19 @@
+{
+  "vision_model": "google/siglip-base-patch16-224",
+  "language_model": "EleutherAI/pythia-14m",
+  "projector_type": "mlp",
+  "dtype": "float32",
+  "device": "cpu",
+  "learning_rate": 0.001,
+  "weight_decay": 0.0,
+  "epochs": 1,
+  "batch_size": 2,
+  "max_steps": 40,
+  "max_text_tokens": 24,
+  "dataset": "clip-benchmark/wds_flickr30k",
+  "dataset_split": "test",
+  "seed": 0,
+  "output_checkpoint": "checkpoints/vision_projector.pt",
+  "log_every": 10,
+  "limit": 128
+}

cortex_ai/vision/__init__.py

@@ -11,6 +11,24 @@ __all__ = [
 if "NativeVisionFusion" in globals():
     __all__ += ["NativeFusionResult", "NativeVisionFusion", "build_native_fusion"]
 
+try:
+    from .train import (
+        ProjectorTrainingConfig,
+        TrainingReport,
+        build_fusion_pipeline,
+        run_training,
+        save_projector,
+        train_step,
+    )
+
+    __all__ += [
+        "ProjectorTrainingConfig", "TrainingReport", "build_fusion_pipeline",
+        "run_training", "save_projector", "train_step",
+    ]
+except ImportError:
+    # Training helpers require the optional dataset/transformers stack.
+    pass
+
 try:
     from .core import (
         ImageProcessor,

cortex_ai/vision/core.py

@@ -6,6 +6,7 @@ separate approved checkpoint). The adapter can therefore be tested with a
 small local encoder while remaining honest about deployment readiness.
 """
 
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any, Protocol, Sequence
 
@@ -47,13 +48,31 @@ class ImageProcessor(Protocol):
 
 
 class VisionProjector(nn.Module):
-    """Map encoder features into CORTEX hidden space."""
+    """Map encoder features into CORTEX hidden space.
+
+    ``projector_type='linear'`` is a single linear map. ``projector_type='mlp'`` is the
+    LLaVA-style two-layer gated projection, whose intermediate width equals the target
+    hidden size. Only this module is trained in the first experimental stage.
+    """
 
     def __init__(self, config: VisionConfig) -> None:
         super().__init__()
-        if config.projector_type != "linear":
-            raise ValueError("Only projector_type='linear' is implemented")
-        self.projection = nn.Linear(config.vision_hidden_size, config.cortex_hidden_size)
+        if config.projector_type == "linear":
+            self.projection: nn.Module = nn.Linear(
+                config.vision_hidden_size, config.cortex_hidden_size
+            )
+        elif config.projector_type == "mlp":
+            self.projection = nn.Sequential(
+                nn.Linear(config.vision_hidden_size, config.cortex_hidden_size),
+                nn.GELU(),
+                nn.Linear(config.cortex_hidden_size, config.cortex_hidden_size),
+            )
+        else:
+            raise ValueError(
+                f"Unknown projector_type={config.projector_type!r}; "
+                "expected 'linear' or 'mlp'"
+            )
+        self.config = config
 
     def forward(self, features: Tensor) -> Tensor:
         if features.ndim != 3:
@@ -73,7 +92,12 @@ class VisionAdapter(nn.Module):
 
     def encode_images(self, images: Any) -> Tensor:
         model_inputs = self.processor(images, return_tensors="pt") if self.processor else images
-        features = self.encoder(**model_inputs) if isinstance(model_inputs, dict) else self.encoder(model_inputs)
+        # Same BatchFeature caveat as native.py: HF processors return a UserDict, which is
+        # not a dict subclass, so it must be tested as a Mapping before unpacking by keyword.
+        if isinstance(model_inputs, Mapping):
+            features = self.encoder(**model_inputs)
+        else:
+            features = self.encoder(model_inputs)
         if hasattr(features, "last_hidden_state"):
             features = features.last_hidden_state
         if not isinstance(features, Tensor):

cortex_ai/vision/native.py

@@ -55,6 +55,18 @@ class NativeVisionFusion(nn.Module):
             encoded = self.vision_encoder(**model_input)
         else:
             encoded = self.vision_encoder(model_input)
+        return self._project(encoded)
+
+    def encode_pixels(self, pixel_values: Tensor, **encoder_kwargs: Any) -> Tensor:
+        """Project already-preprocessed ``pixel_values`` into the core's hidden space.
+
+        Training loops preprocess whole batches once, so they use this instead of
+        ``encode_image`` to avoid re-running the image processor on every step. The encoder is
+        called by keyword here, which suits Transformers vision towers.
+        """
+        return self._project(self.vision_encoder(pixel_values=pixel_values, **encoder_kwargs))
+
+    def _project(self, encoded: Any) -> Tensor:
         if hasattr(encoded, "last_hidden_state"):
             encoded = encoded.last_hidden_state
         if not isinstance(encoded, Tensor) or encoded.ndim != 3:
@@ -90,7 +102,12 @@ def build_native_fusion(
     vision_hidden_size: int,
     cortex_hidden_size: int,
     image_processor: Callable[..., Any] | None = None,
+    projector_type: str = "linear",
 ) -> NativeVisionFusion:
     """Construct the native path around externally supplied encoder and core modules."""
-    config = VisionConfig(vision_hidden_size=vision_hidden_size, cortex_hidden_size=cortex_hidden_size)
+    config = VisionConfig(
+        vision_hidden_size=vision_hidden_size,
+        cortex_hidden_size=cortex_hidden_size,
+        projector_type=projector_type,
+    )
     return NativeVisionFusion(vision_encoder, config, core, image_processor)

cortex_ai/vision/train.py

@@ -0,0 +1,310 @@
+"""Trainable vision projector for CORTEX AI (experimental, stage 1).
+
+Pipeline exercised here::
+
+    IMAGE -> real vision encoder -> VisionProjector -> PrefixFusion
+          -> small real frozen causal LM -> logits -> causal LM loss
+
+The vision encoder and the causal LM stay frozen; **only** ``VisionProjector`` is trained.
+This module is deliberately independent of the 1.65 T DeepSeek-V4-Pro checkpoint: it neither
+reads nor modifies the model-*-of-*.safetensors shards, and the projector checkpoint it writes
+is a separate file.
+
+Scope note: a projector trained against a small stand-in LM demonstrates that the projector is
+genuinely trainable and that the frozen-weight invariants hold. It does **not** demonstrate that
+the trained projector works with the real DeepSeek-V4-Pro core, nor native multimodal
+understanding of that checkpoint.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Iterable, Sequence
+
+import torch
+from torch import Tensor, nn
+
+from .core import PrefixFusion, VisionConfig, VisionProjector
+from .native import NativeVisionFusion, build_native_fusion
+
+IGNORE_INDEX = -100
+
+
+@dataclass
+class ProjectorTrainingConfig:
+    """Explicit dimensions, dtype and optimisation settings for projector training."""
+
+    vision_model: str = "google/siglip-base-patch16-224"
+    language_model: str = "EleutherAI/pythia-14m"
+    projector_type: str = "mlp"
+    dtype: str = "float32"
+    device: str = "cpu"
+    learning_rate: float = 1e-3
+    weight_decay: float = 0.0
+    epochs: int = 1
+    batch_size: int = 2
+    max_steps: int = 40
+    max_text_tokens: int = 24
+    dataset: str = "clip-benchmark/wds_flickr30k"
+    dataset_split: str = "test"
+    seed: int = 0
+    output_checkpoint: str = "checkpoints/vision_projector.pt"
+    log_every: int = 10
+    limit: int | None = 128
+
+    def resolved_dtype(self) -> torch.dtype:
+        return {"float32": torch.float32, "float16": torch.float16,
+                "bfloat16": torch.bfloat16}[self.dtype]
+
+    def resolved_device(self) -> torch.device:
+        return torch.device(self.device)
+
+
+@dataclass
+class TrainingReport:
+    """Recorded evidence from a training run."""
+
+    vision_hidden_size: int
+    lm_hidden_size: int
+    image_token_count: int
+    text_token_count: int
+    fused_sequence_length: int
+    projector_parameters: int
+    dtype: str
+    device: str
+    steps: int
+    loss_first: float
+    loss_last: float
+    loss_history: list[float] = field(default_factory=list)
+    projector_changed: bool = False
+    projector_delta: float = 0.0
+    encoder_unchanged: bool = False
+    lm_unchanged: bool = False
+    checkpoint_path: str = ""
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), indent=2)
+
+
+def freeze_module(module: nn.Module) -> int:
+    """Set requires_grad=False everywhere and return the frozen parameter count."""
+    n = 0
+    for p in module.parameters():
+        p.requires_grad_(False)
+        n += p.numel()
+    module.eval()
+    return n
+
+
+def count_parameters(module: nn.Module, trainable_only: bool = False) -> int:
+    return sum(p.numel() for p in module.parameters() if (p.requires_grad or not trainable_only))
+
+
+def parameter_fingerprint(module: nn.Module) -> dict[str, Tensor]:
+    """Detached clone of every parameter, for before/after comparison."""
+    return {k: v.detach().clone() for k, v in module.state_dict().items()}
+
+
+def fingerprint_delta(before: dict[str, Tensor], after: dict[str, Tensor]) -> float:
+    """Total absolute change across all tensors in two fingerprints."""
+    total = 0.0
+    for k, v in before.items():
+        if k not in after:
+            continue
+        total += (after[k].to(torch.float32) - v.to(torch.float32)).abs().sum().item()
+    return total
+
+
+def build_fusion_pipeline(
+    vision_encoder: nn.Module,
+    core: nn.Module,
+    config: ProjectorTrainingConfig,
+    image_processor: Callable[..., Any] | None = None,
+) -> tuple[NativeVisionFusion, int, int]:
+    """Assemble IMAGE -> encoder -> projector -> PrefixFusion -> frozen core."""
+    vision_hidden = int(vision_encoder.config.hidden_size)
+    lm_hidden = int(core.config.hidden_size)
+    fusion = build_native_fusion(
+        vision_encoder,
+        core,
+        vision_hidden_size=vision_hidden,
+        cortex_hidden_size=lm_hidden,
+        image_processor=image_processor,
+        projector_type=config.projector_type,
+    )
+    # The encoder and the core are frozen; the projector is the only trained part.
+    freeze_module(fusion.vision_encoder)
+    freeze_module(fusion.core)
+    fusion.projector.train()
+    return fusion, vision_hidden, lm_hidden
+
+
+def fused_labels(image_token_count: int, input_ids: Tensor) -> Tensor:
+    """Labels aligned to the fused sequence.
+
+    Causal-LM loss is reported at every position, so the image-token positions would otherwise
+    be trained to predict ordinary text tokens from a vision prefix. They are masked out, so the
+    gradient on the projector comes from predicting the caption tokens conditioned on the image.
+    """
+    ignore = torch.full(
+        (input_ids.shape[0], image_token_count), IGNORE_INDEX, dtype=torch.long, device=input_ids.device
+    )
+    return torch.cat((ignore, input_ids), dim=1)
+
+
+def fused_attention_mask(image_token_count: int, attention_mask: Tensor | None, batch_size: int) -> Tensor:
+    """Attention mask aligned to the fused sequence.
+
+    Image tokens are always attended (1), so the text mask must be extended rather than passed
+    through unchanged. Text padding keeps its own 0 entries.
+    """
+    if attention_mask is None:
+        return torch.ones((batch_size, image_token_count), dtype=torch.long)
+    prefix = torch.ones(
+        (attention_mask.shape[0], image_token_count), dtype=attention_mask.dtype, device=attention_mask.device
+    )
+    return torch.cat((prefix, attention_mask), dim=1)
+
+
+def train_step(
+    fusion: NativeVisionFusion,
+    optimizer: torch.optim.Optimizer,
+    pixel_values: Tensor,
+    input_ids: Tensor,
+    attention_mask: Tensor | None = None,
+) -> float:
+    """One optimisation step. Returns the scalar loss."""
+    image_embeds = fusion.encode_pixels(pixel_values)
+    text_embeds = fusion.core.get_input_embeddings()(input_ids)
+    fused = fusion.fusion.fuse(text_embeds, image_embeds)
+    mask = fused_attention_mask(image_embeds.shape[1], attention_mask, fused.shape[0])
+    labels = fused_labels(image_embeds.shape[1], input_ids)
+
+    out = fusion.core(
+        inputs_embeds=fused, attention_mask=mask, labels=labels, use_cache=False
+    )
+    loss = out.loss
+    optimizer.zero_grad(set_to_none=True)
+    loss.backward()
+    optimizer.step()
+    return float(loss.detach())
+
+
+def build_optimizer(projector: nn.Module, config: ProjectorTrainingConfig) -> torch.optim.Optimizer:
+    params = [p for p in projector.parameters() if p.requires_grad]
+    if not params:
+        raise RuntimeError("projector has no trainable parameters")
+    return torch.optim.AdamW(params, lr=config.learning_rate, weight_decay=config.weight_decay)
+
+
+def iterate_caption_batches(
+    images: Sequence[Any],
+    captions: Sequence[str],
+    tokenizer: Any,
+    image_processor: Any,
+    config: ProjectorTrainingConfig,
+) -> Iterable[tuple[Tensor, Tensor, Tensor]]:
+    """Yield pre-processed (pixel_values, input_ids, attention_mask) batches."""
+    bs = config.batch_size
+    for start in range(0, len(images), bs):
+        batch_images = list(images[start:start + bs])
+        batch_captions = list(captions[start:start + bs])
+        if not batch_images:
+            continue
+        pixel_values = image_processor(batch_images, return_tensors="pt").pixel_values
+        enc = tokenizer(
+            batch_captions,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=config.max_text_tokens,
+        )
+        yield pixel_values, enc["input_ids"], enc["attention_mask"]
+
+
+def run_training(
+    fusion: NativeVisionFusion,
+    batches: Iterable[tuple[Tensor, Tensor, Tensor]],
+    config: ProjectorTrainingConfig,
+    projector_parameters: int,
+    vision_hidden_size: int,
+    lm_hidden_size: int,
+    on_step: Callable[[int, float], None] | None = None,
+) -> TrainingReport:
+    """Train the projector over the supplied batches and report the evidence."""
+    dtype = config.resolved_dtype()
+    device = config.resolved_device()
+
+    before_proj = parameter_fingerprint(fusion.projector)
+    before_enc = parameter_fingerprint(fusion.vision_encoder)
+    before_lm = parameter_fingerprint(fusion.core)
+
+    optimizer = build_optimizer(fusion.projector, config)
+    losses: list[float] = []
+    steps = 0
+    first_shape: tuple[int, int, int] | None = None
+
+    for _epoch in range(config.epochs):
+        for pixel_values, input_ids, attention_mask in batches:
+            if steps >= config.max_steps:
+                break
+            pixel_values = pixel_values.to(device=device, dtype=dtype)
+            input_ids = input_ids.to(device)
+            attention_mask = attention_mask.to(device)
+
+            if first_shape is None:
+                with torch.no_grad():
+                    n_img = fusion.encode_pixels(pixel_values).shape[1]
+                first_shape = (n_img, input_ids.shape[1], n_img + input_ids.shape[1])
+
+            loss = train_step(fusion, optimizer, pixel_values, input_ids, attention_mask)
+            losses.append(loss)
+            steps += 1
+            if on_step is not None:
+                on_step(steps, loss)
+        if steps >= config.max_steps:
+            break
+
+    after_proj = parameter_fingerprint(fusion.projector)
+    after_enc = parameter_fingerprint(fusion.vision_encoder)
+    after_lm = parameter_fingerprint(fusion.core)
+
+    n_img, n_txt, n_fused = first_shape or (0, 0, 0)
+    report = TrainingReport(
+        vision_hidden_size=vision_hidden_size,
+        lm_hidden_size=lm_hidden_size,
+        image_token_count=n_img,
+        text_token_count=n_txt,
+        fused_sequence_length=n_fused,
+        projector_parameters=projector_parameters,
+        dtype=config.dtype,
+        device=config.device,
+        steps=steps,
+        loss_first=losses[0] if losses else float("nan"),
+        loss_last=losses[-1] if losses else float("nan"),
+        loss_history=losses,
+        projector_delta=fingerprint_delta(before_proj, after_proj),
+        encoder_unchanged=fingerprint_delta(before_enc, after_enc) == 0.0,
+        lm_unchanged=fingerprint_delta(before_lm, after_lm) == 0.0,
+    )
+    report.projector_changed = report.projector_delta > 0.0
+    return report
+
+
+def save_projector(projector: nn.Module, config: ProjectorTrainingConfig, report: TrainingReport) -> Path:
+    """Write the projector checkpoint plus its config, separate from any base model."""
+    path = Path(config.output_checkpoint)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    report.checkpoint_path = str(path)
+    torch.save(
+        {
+            "state_dict": projector.state_dict(),
+            "config": asdict(config),
+            "report": json.loads(report.to_json()),
+            "scope": "experimental stage 1: projector trained against a small stand-in LM",
+        },
+        path,
+    )
+    return path

docs/vision_projector_training.md

@@ -0,0 +1,105 @@
+# CORTEX AI — vision projector training (stage 1)
+
+Experimental, reproducible proof that the CORTEX vision projector can be trained on its own,
+against a real vision encoder and a real causal language model, while the base model weights
+stay completely untouched.
+
+## Pipeline
+
+```
+IMAGE
+  -> real vision encoder      (SigLIP, frozen)
+  -> VisionProjector          (TRAINED — the only trainable part)
+  -> PrefixFusion             (image tokens prepended to text embeddings)
+  -> small real causal LM     (frozen)
+  -> logits -> causal LM loss
+```
+
+## Quick start
+
+```bash
+pip install torch torchvision transformers datasets pillow
+
+python examples/train_vision_projector.py --config configs/vision_projector.json
+python examples/train_vision_projector.py --max-steps 20 --batch-size 2   # quick run
+```
+
+The whole script runs on CPU. Measured cost on 4 cores: about 0.8 s per step at batch size 2.
+
+## Exact dimensions
+
+| Quantity | Value |
+|---|---|
+| Vision encoder | `google/siglip-base-patch16-224` (`SiglipVisionModel`) |
+| Vision encoder output (`vision_hidden_size`) | **768** |
+| Image tokens per 224x224 image | **196** |
+| Causal LM | `EleutherAI/pythia-14m` |
+| Causal LM embedding size (`cortex_hidden_size`) | **128** |
+| Projector in -> out | **768 -> 128** |
+| Projector type | `mlp` (Linear 768->128, GELU, Linear 128->128) |
+| Projector parameters | **114,944** |
+| Text tokens per caption | up to 24 |
+| Fused sequence (196 image + 24 text) | **220** |
+| dtype / device | `float32` / `cpu` |
+
+The projector output width must equal the language model's embedding size; that is the
+`cortex_hidden_size` argument in `VisionConfig`.
+
+## What is PROVEN
+
+- Real image -> real encoder -> projector -> PrefixFusion -> frozen LM executes end to end.
+- The projector receives gradients; its parameters change measurably after training.
+- The vision encoder and the causal LM parameters do **not** change (fingerprint delta exactly 0),
+  and receive **no** gradients.
+- Two different images produce different fused representations; identical images produce identical ones.
+- Training loss falls on real Flickr30k image/caption pairs.
+- Reproducible via `python examples/train_vision_projector.py` and the test suite.
+
+## What is NOT PROVEN
+
+- That this trained projector works with the real **DeepSeek-V4-Pro** core. The stand-in LM has
+  hidden size 128; DeepSeek-V4-Pro's is 7168, so the projector would need retraining at that width.
+- Native multimodal understanding of the 1.65 T checkpoint.
+- Final CORTEX AI performance.
+
+Training a projector against a small stand-in LM demonstrates **trainability and the frozen-weight
+invariants**, not end-to-end multimodal capability.
+
+## Loss and masking
+
+Causal LM loss is computed on the fused sequence. Image-token positions are masked to `-100`
+(`IGNORE_INDEX`), so the projector's gradient comes from predicting the caption tokens
+conditioned on the image, rather than from predicting text at the image positions.
+
+## Checkpoint
+
+`checkpoints/vision_projector.pt` — written by the script, git-ignored, and entirely separate
+from the DeepSeek-V4-Pro shards. It contains the projector `state_dict`, the training config,
+and the run report.
+
+## Tests
+
+```bash
+PYTHONPATH=$PWD python3 -m pytest tests/ -q                              # 14 passed, 1 skipped
+CORTEX_REAL_VISION_TEST=1 python3 -m pytest tests/ -q -k real            # real models
+```
+
+The `test_real_models_*` test is opt-in because it downloads real weights.
+
+## Shard integrity
+
+The 66 DeepSeek-V4-Pro shards are never downloaded, modified or read as payloads. Verify with:
+
+```bash
+python tools/shard_integrity.py --from-worktree --out /tmp/after.json
+python tools/shard_integrity.py --compare /tmp/before.json /tmp/after.json
+```
+
+All 66 shard LFS pointer hashes are unchanged by training, and the recorded LFS `oid` values
+match the upstream `deepseek-ai/DeepSeek-V4-Pro-0813` checkpoint.
+
+## Configuration
+
+See `configs/vision_projector.json`. Every field of `ProjectorTrainingConfig` is also a CLI flag
+(`--learning-rate`, `--epochs`, `--batch-size`, `--output-checkpoint`, ...), with the JSON file
+acting as defaults.

examples/train_vision_projector.py

@@ -0,0 +1,158 @@
+"""Stage-1 vision projector training for CORTEX AI.
+
+Trains *only* ``VisionProjector`` on top of a frozen real vision encoder and a frozen small
+causal LM::
+
+    IMAGE -> real vision encoder -> VisionProjector -> PrefixFusion -> frozen causal LM -> logits
+
+Nothing here touches the DeepSeek-V4-Pro shards: the ``model-*-of-*.safetensors`` files are only
+read to hash their LFS pointers, and the checkpoint written by this script is a separate file.
+
+Examples::
+
+    python examples/train_vision_projector.py --config configs/vision_projector.json
+    python examples/train_vision_projector.py --max-steps 20 --batch-size 2
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from dataclasses import asdict, fields
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from cortex_ai.vision.train import (  # noqa: E402
+    ProjectorTrainingConfig,
+    build_fusion_pipeline,
+    count_parameters,
+    iterate_caption_batches,
+    run_training,
+    save_projector,
+)
+
+
+def load_dataset(config: ProjectorTrainingConfig) -> tuple[list, list[str]]:
+    """Load real images and captions from the configured Hub dataset."""
+    from datasets import load_dataset as hf_load_dataset
+
+    ds = hf_load_dataset(config.dataset, split=config.dataset_split, streaming=True)
+    images, captions = [], []
+    limit = config.limit or 0
+    for example in ds:
+        image = example.get("jpg") or example.get("image")
+        text = example.get("txt") or example.get("caption") or example.get("text")
+        if image is None or not text:
+            continue
+        if isinstance(text, list):
+            text = text[0]
+        images.append(image.convert("RGB"))
+        captions.append(str(text))
+        if limit and len(images) >= limit:
+            break
+    if not images:
+        raise SystemExit(f"no usable image/caption pairs found in {config.dataset}")
+    return images, captions
+
+
+def parse_args(argv: list[str] | None = None) -> ProjectorTrainingConfig:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--config", help="JSON file of ProjectorTrainingConfig fields")
+    defaults = ProjectorTrainingConfig()
+    for f in fields(defaults):
+        value = getattr(defaults, f.name)
+        if isinstance(value, bool):
+            p.add_argument(f"--{f.name.replace('_', '-')}", type=lambda s: s.lower() == "true", default=None)
+        else:
+            p.add_argument(
+                f"--{f.name.replace('_', '-')}",
+                type=type(value) if value is not None else str,
+                default=None,
+            )
+    args = p.parse_args(argv)
+
+    values = {}
+    if args.config:
+        values.update(json.loads(Path(args.config).read_text()))
+    for f in fields(defaults):
+        v = getattr(args, f.name.replace("-", "_"))
+        if v is not None:
+            values[f.name] = v
+    return ProjectorTrainingConfig(**values)
+
+
+def main(argv: list[str] | None = None) -> int:
+    config = parse_args(argv)
+    torch.manual_seed(config.seed)
+
+    print("=" * 78)
+    print("CORTEX AI - stage-1 vision projector training (experimental)")
+    print("=" * 78)
+    print(f"config        : {json.dumps(asdict(config), indent=2)}")
+
+    from transformers import AutoImageProcessor, AutoModelForCausalLM, AutoTokenizer, SiglipVisionModel
+
+    dtype = config.resolved_dtype()
+    device = config.resolved_device()
+
+    print(f"\n[1] loading real vision encoder : {config.vision_model}")
+    image_processor = AutoImageProcessor.from_pretrained(config.vision_model)
+    vision_encoder = SiglipVisionModel.from_pretrained(config.vision_model).to(dtype=dtype, device=device)
+
+    print(f"[2] loading real frozen LM      : {config.language_model}")
+    core = AutoModelForCausalLM.from_pretrained(config.language_model, dtype=dtype).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(config.language_model)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print("[3] assembling fusion pipeline (encoder + LM frozen, projector trainable)")
+    fusion, vision_hidden, lm_hidden = build_fusion_pipeline(
+        vision_encoder, core, config, image_processor=image_processor
+    )
+    projector_params = count_parameters(fusion.projector)
+    print(f"    vision_hidden_size   = {vision_hidden}")
+    print(f"    lm_hidden_size       = {lm_hidden}")
+    print(f"    projector_type       = {config.projector_type}")
+    print(f"    projector_parameters = {projector_params:,}")
+    print(f"    dtype / device       = {config.dtype} / {config.device}")
+
+    print(f"\n[4] loading real dataset        : {config.dataset} (split={config.dataset_split})")
+    images, captions = load_dataset(config)
+    print(f"    {len(images)} image/caption pairs")
+    print(f"    sample caption: {captions[0][:70]!r}")
+
+    batches = iterate_caption_batches(images, captions, tokenizer, image_processor, config)
+
+    def on_step(step: int, loss: float) -> None:
+        if step == 1 or step % config.log_every == 0:
+            print(f"    step {step:4d}  loss={loss:.4f}")
+
+    print(f"\n[5] training projector for up to {config.max_steps} steps (projector only)")
+    report = run_training(
+        fusion, batches, config, projector_params, vision_hidden, lm_hidden, on_step=on_step
+    )
+
+    path = save_projector(fusion.projector, config, report)
+    report.checkpoint_path = str(path)
+
+    print("\n" + "=" * 78)
+    print("RESULT")
+    print("=" * 78)
+    print(report.to_json())
+    print(f"\ncheckpoint written: {path}")
+    print("\nPROVEN by this run:")
+    print("  - real image -> real encoder -> projector -> PrefixFusion -> frozen LM executed")
+    print("  - projector received gradients and its parameters changed")
+    print("  - vision encoder and causal LM parameters are unchanged")
+    print("\nNOT PROVEN:")
+    print("  - that this projector works with the real DeepSeek-V4-Pro core")
+    print("  - native multimodal understanding of the 1.65T checkpoint")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_projector_training.py

@@ -0,0 +1,259 @@
+"""Invariants for stage-1 vision projector training.
+
+Two layers of evidence:
+
+* ``test_*``                       : fast, deterministic, real tensor math on small modules.
+* ``test_real_models_*``           : the real recipe (real SigLIP encoder + real causal LM).
+  Opt in with ``CORTEX_REAL_VISION_TEST=1``; this needs network access to fetch weights.
+
+Nothing here reads, downloads or writes the DeepSeek-V4-Pro shards.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+from torch import nn
+
+from cortex_ai.vision.core import PrefixFusion, VisionConfig, VisionProjector
+from cortex_ai.vision.native import build_native_fusion
+from cortex_ai.vision.train import (
+    ProjectorTrainingConfig,
+    build_fusion_pipeline,
+    build_optimizer,
+    count_parameters,
+    fingerprint_delta,
+    fused_attention_mask,
+    fused_labels,
+    parameter_fingerprint,
+    train_step,
+)
+
+VISION_HIDDEN = 16
+LM_HIDDEN = 32
+
+
+class SmallVisionEncoder(nn.Module):
+    """Real (non-mock) tiny encoder exercising the same contract as a HF vision tower."""
+
+    def __init__(self, hidden_size: int = VISION_HIDDEN) -> None:
+        super().__init__()
+        self.config = type("Cfg", (), {"hidden_size": hidden_size})()
+        self.conv = nn.Conv2d(3, hidden_size, kernel_size=4, stride=4)
+        self.norm = nn.LayerNorm(hidden_size)
+
+    def forward(self, pixel_values: torch.Tensor) -> object:
+        feats = self.conv(pixel_values)                     # [B, H, h, w]
+        feats = feats.flatten(2).transpose(1, 2)            # [B, tokens, H]
+        return type("Out", (), {"last_hidden_state": self.norm(feats)})()
+
+
+class SmallCausalLM(nn.Module):
+    """Real tiny causal LM with the HF-shaped contract used by the fusion path."""
+
+    def __init__(self, hidden_size: int = LM_HIDDEN, vocab_size: int = 64) -> None:
+        super().__init__()
+        self.config = type("Cfg", (), {"hidden_size": hidden_size, "vocab_size": vocab_size})()
+        self.embed = nn.Embedding(vocab_size, hidden_size)
+        self.head = nn.Linear(hidden_size, vocab_size)
+        self.lm_head = self.head
+        self.calls: list[tuple[int, int]] = []
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.embed
+
+    def forward(self, inputs_embeds, attention_mask=None, labels=None, use_cache=False, **kw):
+        self.calls.append(tuple(inputs_embeds.shape))
+        hidden = torch.tanh(inputs_embeds)
+        logits = self.head(hidden)
+        loss = None
+        if labels is not None:
+            loss = nn.functional.cross_entropy(
+                logits[:, :-1].reshape(-1, logits.shape[-1]),
+                labels[:, 1:].reshape(-1),
+                ignore_index=-100,
+            )
+        return type("Out", (), {"loss": loss, "logits": logits})()
+
+
+def build_fake_pipeline(config: ProjectorTrainingConfig | None = None):
+    config = config or ProjectorTrainingConfig(projector_type="mlp", max_steps=3)
+    encoder = SmallVisionEncoder()
+    core = SmallCausalLM()
+    fusion, vh, lh = build_fusion_pipeline(encoder, core, config)
+    return fusion, encoder, core, vh, lh
+
+
+def make_batch(batch_size: int = 2, text_tokens: int = 4):
+    pixel_values = torch.rand(batch_size, 3, 16, 16)
+    input_ids = torch.randint(0, 64, (batch_size, text_tokens))
+    attention_mask = torch.ones_like(input_ids)
+    return pixel_values, input_ids, attention_mask
+
+
+def test_vision_encoder_and_lm_are_frozen_projector_is_not():
+    fusion, encoder, core, _, _ = build_fake_pipeline()
+    assert all(not p.requires_grad for p in encoder.parameters())
+    assert all(not p.requires_grad for p in core.parameters())
+    assert all(p.requires_grad for p in fusion.projector.parameters())
+    assert count_parameters(fusion.projector, trainable_only=True) == count_parameters(fusion.projector) > 0
+
+
+def test_projector_receives_gradients_and_frozen_weights_receive_none():
+    fusion, encoder, core, _, _ = build_fake_pipeline()
+    optimizer = build_optimizer(fusion.projector, ProjectorTrainingConfig())
+    pixel_values, input_ids, attention_mask = make_batch()
+    train_step(fusion, optimizer, pixel_values, input_ids, attention_mask)
+
+    assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in fusion.projector.parameters())
+    assert all(p.grad is None for p in encoder.parameters())
+    assert all(p.grad is None for p in core.parameters())
+
+
+def test_training_changes_projector_only():
+    fusion, encoder, core, _, _ = build_fake_pipeline()
+    optimizer = build_optimizer(fusion.projector, ProjectorTrainingConfig(learning_rate=1e-2))
+
+    before_proj = parameter_fingerprint(fusion.projector)
+    before_enc = parameter_fingerprint(encoder)
+    before_lm = parameter_fingerprint(core)
+
+    pixel_values, input_ids, attention_mask = make_batch()
+    for _ in range(3):
+        train_step(fusion, optimizer, pixel_values, input_ids, attention_mask)
+
+    assert fingerprint_delta(before_proj, parameter_fingerprint(fusion.projector)) > 0
+    assert fingerprint_delta(before_enc, parameter_fingerprint(encoder)) == 0.0
+    assert fingerprint_delta(before_lm, parameter_fingerprint(core)) == 0.0
+
+
+def test_pipeline_runs_before_and_after_training():
+    fusion, *_ = build_fake_pipeline()
+    optimizer = build_optimizer(fusion.projector, ProjectorTrainingConfig())
+    pixel_values, input_ids, attention_mask = make_batch()
+
+    out_before = fusion.core(
+        inputs_embeds=fusion.fusion.fuse(
+            fusion.core.get_input_embeddings()(input_ids), fusion.encode_pixels(pixel_values)
+        ),
+        attention_mask=None,
+        use_cache=False,
+    )
+    assert out_before.logits.shape[0] == 2
+
+    for _ in range(3):
+        train_step(fusion, optimizer, pixel_values, input_ids, attention_mask)
+
+    out_after = fusion.core(
+        inputs_embeds=fusion.fusion.fuse(
+            fusion.core.get_input_embeddings()(input_ids), fusion.encode_pixels(pixel_values)
+        ),
+        attention_mask=None,
+        use_cache=False,
+    )
+    assert out_after.logits.shape == out_before.logits.shape
+    assert not torch.allclose(out_before.logits, out_after.logits)
+
+
+def test_two_different_images_produce_different_fused_representations():
+    fusion, *_ = build_fake_pipeline()
+    fusable = fusion.projector
+
+    image_a = torch.zeros(1, 3, 16, 16)
+    image_b = torch.ones(1, 3, 16, 16)
+    emb_a = fusion.encode_pixels(image_a)
+    emb_b = fusion.encode_pixels(image_b)
+
+    assert emb_a.shape == emb_b.shape
+    assert not torch.allclose(emb_a, emb_b)
+
+    text = torch.zeros(1, 3, LM_HIDDEN)
+    fused_a = fusion.fusion.fuse(text, emb_a)
+    fused_b = fusion.fusion.fuse(text, emb_b)
+    assert fused_a.shape[1] == emb_a.shape[1] + 3
+    assert not torch.allclose(fused_a, fused_b)
+
+
+def test_fused_mask_and_labels_align_with_fused_sequence():
+    input_ids = torch.tensor([[5, 6, 7], [8, 9, 0]])
+    attention_mask = torch.tensor([[1, 1, 1], [1, 1, 0]])
+    image_tokens = 4
+
+    mask = fused_attention_mask(image_tokens, attention_mask, batch_size=2)
+    labels = fused_labels(image_tokens, input_ids)
+
+    assert mask.shape == (2, image_tokens + 3)
+    assert torch.all(mask[:, :image_tokens] == 1)        # image tokens always attended
+    assert torch.all(mask[:, image_tokens:] == attention_mask)
+    assert labels.shape == (2, image_tokens + 3)
+    assert torch.all(labels[:, :image_tokens] == -100)   # image positions not trained
+    assert torch.equal(labels[:, image_tokens:], input_ids)
+
+
+def test_mlp_projector_is_two_layer_and_linear_still_supported():
+    mlp = VisionProjector(VisionConfig(vision_hidden_size=8, cortex_hidden_size=16, projector_type="mlp"))
+    assert len(mlp.projection) == 3
+    out = mlp(torch.zeros(2, 5, 8))
+    assert out.shape == (2, 5, 16)
+
+    linear = VisionProjector(VisionConfig(vision_hidden_size=8, cortex_hidden_size=16, projector_type="linear"))
+    assert isinstance(linear.projection, nn.Linear)
+
+    with pytest.raises(ValueError):
+        VisionProjector(VisionConfig(vision_hidden_size=8, cortex_hidden_size=16, projector_type="nope"))
+
+
+def test_projector_checkpoint_is_a_separate_file(tmp_path):
+    from cortex_ai.vision.train import TrainingReport, save_projector
+
+    fusion, *_ = build_fake_pipeline()
+    config = ProjectorTrainingConfig(output_checkpoint=str(tmp_path / "projector.pt"))
+    report = TrainingReport(
+        vision_hidden_size=VISION_HIDDEN, lm_hidden_size=LM_HIDDEN, image_token_count=4,
+        text_token_count=4, fused_sequence_length=8, projector_parameters=count_parameters(fusion.projector),
+        dtype="float32", device="cpu", steps=0, loss_first=0.0, loss_last=0.0,
+    )
+    path = save_projector(fusion.projector, config, report)
+    assert path.exists()
+    payload = torch.load(path, map_location="cpu", weights_only=True)
+    assert set(payload["state_dict"]) == set(fusion.projector.state_dict())
+    assert payload["config"]["vision_model"] == config.vision_model
+
+
+@pytest.mark.skipif(
+    os.environ.get("CORTEX_REAL_VISION_TEST") != "1",
+    reason="set CORTEX_REAL_VISION_TEST=1 to run against real SigLIP + real causal LM",
+)
+def test_real_models_projector_trains_and_frozen_weights_stay_frozen():
+    from transformers import AutoImageProcessor, AutoModelForCausalLM, SiglipVisionModel
+
+    config = ProjectorTrainingConfig(
+        projector_type="mlp", dtype="float32", max_steps=4, batch_size=2,
+        learning_rate=1e-3, dataset="clip-benchmark/wds_flickr30k",
+    )
+    processor = AutoImageProcessor.from_pretrained(config.vision_model)
+    encoder = SiglipVisionModel.from_pretrained(config.vision_model).to(torch.float32)
+    core = AutoModelForCausalLM.from_pretrained(config.language_model, dtype=torch.float32)
+    fusion, _, _ = build_fusion_pipeline(encoder, core, config)
+
+    before_proj = parameter_fingerprint(fusion.projector)
+    before_enc = parameter_fingerprint(encoder)
+    before_lm = parameter_fingerprint(core)
+
+    optimizer = build_optimizer(fusion.projector, config)
+    from PIL import Image
+    images = [Image.new("RGB", (224, 224), (i * 60 % 255, 40, 200)) for i in range(4)]
+    pixel_values = processor(images, return_tensors="pt").pixel_values.to(torch.float32)
+    input_ids = torch.randint(0, 1000, (4, 6))
+    attention_mask = torch.ones_like(input_ids)
+
+    losses = [train_step(fusion, optimizer, pixel_values[:2], input_ids[:2], attention_mask[:2]) for _ in range(4)]
+
+    assert all(torch.isfinite(torch.tensor(losses)))
+    assert fingerprint_delta(before_proj, parameter_fingerprint(fusion.projector)) > 0
+    assert fingerprint_delta(before_enc, parameter_fingerprint(encoder)) == 0.0
+    assert fingerprint_delta(before_lm, parameter_fingerprint(core)) == 0.0
+    assert all(p.grad is None for p in encoder.parameters())
+    assert all(p.grad is None for p in core.parameters())

tools/shard_integrity.py

@@ -0,0 +1,102 @@
+"""Shard integrity check for the 66 DeepSeek-V4-Pro shards.
+
+Needs no network. Modes:
+
+* ``--from-worktree``  : SHA-256 of the LFS pointer file currently on disk.
+* ``--from-git REV``   : SHA-256 of the pointer text stored in the git object store.
+* ``--compare A B``    : diff two manifests produced by the above.
+
+The raw safetensors payloads are ~893 GB and are deliberately never downloaded. Hashing the
+pointer files proves the shards were not touched. The ``oid sha256:`` line inside each pointer
+is the upstream LFS object id, which the hub-storage audit confirms is identical to the source
+DeepSeek checkpoint.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import pathlib
+import subprocess
+import sys
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+
+
+def sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def pointer_oid(text: str) -> str | None:
+    for line in text.splitlines():
+        if line.startswith("oid sha256:"):
+            return line.split(":", 1)[1].strip()
+    return None
+
+
+def shard_names() -> list[str]:
+    names = sorted(p.name for p in ROOT.glob("model-*-of-*.safetensors"))
+    if len(names) != 66:
+        raise SystemExit(f"expected 66 shards, found {len(names)}")
+    return names
+
+
+def from_worktree() -> dict:
+    out = {}
+    for name in shard_names():
+        data = (ROOT / name).read_bytes()
+        out[name] = {
+            "pointer_sha256": sha256_bytes(data),
+            "lfs_oid": pointer_oid(data.decode("utf-8", "replace")),
+            "size": len(data),
+        }
+    return out
+
+
+def from_git(revision: str) -> dict:
+    out = {}
+    for name in shard_names():
+        blob = subprocess.run(
+            ["git", "show", f"{revision}:{name}"],
+            capture_output=True, check=True, cwd=ROOT,
+        ).stdout
+        out[name] = {
+            "pointer_sha256": sha256_bytes(blob),
+            "lfs_oid": pointer_oid(blob.decode("utf-8", "replace")),
+            "size": len(blob),
+        }
+    return out
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--from-git", metavar="REV")
+    ap.add_argument("--from-worktree", action="store_true")
+    ap.add_argument("--out", default=None)
+    ap.add_argument("--compare", nargs=2, metavar=("A", "B"))
+    args = ap.parse_args()
+
+    if args.compare:
+        a = json.loads(pathlib.Path(args.compare[0]).read_text())
+        b = json.loads(pathlib.Path(args.compare[1]).read_text())
+        if a == b:
+            print(f"IDENTICAL: {len(a)} shards match between {args.compare[0]} and {args.compare[1]}")
+            return 0
+        keys = sorted(set(a) | set(b))
+        diffs = [k for k in keys if a.get(k) != b.get(k)]
+        print(f"MISMATCH on {len(diffs)} shard(s): {diffs[:10]}")
+        return 1
+
+    manifest = from_git(args.from_git) if args.from_git else from_worktree()
+    oids = {v["lfs_oid"] for v in manifest.values()}
+    if args.out:
+        pathlib.Path(args.out).write_text(json.dumps(manifest, indent=1, sort_keys=True))
+        print(f"wrote {args.out}: {len(manifest)} shards, {len(oids)} distinct LFS oids")
+    else:
+        print(f"{len(manifest)} shards, {len(oids)} distinct LFS oids")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())