Abdoulaye coumbassa
openhands
feat(vision): trainable stage-1 vision projector with frozen encoder and LM
ba3bbc7 | """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 | |
| 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) | |
| 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 |