"""Multimodal wrapper architecture that prefixes visual embeddings into an LLM.""" from collections import OrderedDict from typing import Any, Iterable import torch import torch.nn as nn from taoTrain.config import TrainingModeEnum from taoTrain.core import BaseModel from .cnn_encoder import CNNEncoder from .registry import register_architecture, get_model @register_architecture("multimodal_wrapper") class MultimodalWrapper(BaseModel): """Wrap a text-only LLM with a CNN vision tower and visual prefix projector.""" def __init__(self, config: Any): """Initialize the multimodal wrapper from a full VLM config.""" super().__init__(config) if getattr(config, "mode", None) not in {TrainingModeEnum.VLM, TrainingModeEnum.VLM_SFT}: raise ValueError("MultimodalWrapper expects a VLMConfig or VLMSFTConfig") self.train_config = config self.model_config = config.model self.vision_prefix_tokens = config.vision_prefix_tokens self.image_token = config.image_token llm_config = self.model_config.model_copy(deep=True) llm_architecture = llm_config.llm_architecture_type if llm_architecture is None: raise ValueError("model.llm_architecture_type must be set for multimodal_wrapper") llm_config.architecture_type = llm_architecture self.llm = get_model(llm_config) self.hidden_dim = llm_config.hidden_dim if self.model_config.vision_encoder_type != "cnn": raise ValueError(f"Unsupported vision encoder type: {self.model_config.vision_encoder_type}") self.vision_encoder = CNNEncoder( image_size=config.image_size, output_dim=self.model_config.vision_output_dim, channels=self.model_config.cnn_channels, kernel_size=self.model_config.cnn_kernel_size, ) self.vision_projector = nn.Linear( self.model_config.vision_output_dim, self.hidden_dim * self.vision_prefix_tokens, ) self._configure_trainable_parameters() def get_num_layers(self) -> int: """Return the number of underlying LLM blocks.""" if hasattr(self.llm, "get_num_layers"): return self.llm.get_num_layers() return getattr(self.llm.config, "num_layers", 0) def _embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: """Lookup token embeddings from the wrapped LLM.""" if hasattr(self.llm, "token_embedding"): return self.llm.token_embedding(input_ids) if hasattr(self.llm, "embed_tokens"): return self.llm.embed_tokens(input_ids) raise AttributeError("Wrapped LLM does not expose a supported token embedding layer") def _encode_visual_prefix(self, pixel_values: torch.Tensor) -> torch.Tensor: """Encode images and expand them into a sequence of visual prefix embeddings.""" features = self.vision_encoder(pixel_values) projected = self.vision_projector(features) return projected.view(pixel_values.size(0), self.vision_prefix_tokens, self.hidden_dim) def _build_attention_mask(self, attention_mask: torch.Tensor) -> torch.Tensor: """Expand the single image placeholder into a visual prefix attention mask.""" prefix_mask = torch.ones( attention_mask.size(0), self.vision_prefix_tokens, dtype=attention_mask.dtype, device=attention_mask.device, ) return torch.cat([prefix_mask, attention_mask[:, 1:]], dim=1) def _build_labels(self, labels: torch.Tensor | None, device: torch.device) -> torch.Tensor | None: """Mask visual prefix positions out of the language-modeling loss.""" if labels is None: return None prefix_labels = torch.full( (labels.size(0), self.vision_prefix_tokens), -100, dtype=labels.dtype, device=device, ) return torch.cat([prefix_labels, labels[:, 1:]], dim=1) def _validate_image_placeholder(self, attention_mask: torch.Tensor) -> None: """Validate the dataset contract that the image placeholder is the first active token.""" first_token_active = attention_mask[:, 0] if not torch.all(first_token_active == 1): raise ValueError("Multimodal samples must place the placeholder at the first non-padding position") def _configure_trainable_parameters(self) -> None: """Apply freezing rules for the wrapped LLM.""" if not self.train_config.freeze_llm: for param in self.llm.parameters(): param.requires_grad = True return for param in self.llm.parameters(): param.requires_grad = False num_layers = self.get_num_layers() unfreeze_last_n = self.train_config.unfreeze_last_n_layers if unfreeze_last_n > num_layers: raise ValueError( f"unfreeze_last_n_layers ({unfreeze_last_n}) cannot exceed LLM layers ({num_layers})" ) if unfreeze_last_n == 0: return llm_blocks = getattr(self.llm, "blocks", None) if llm_blocks is None: raise AttributeError("Wrapped LLM does not expose `blocks`, cannot selectively unfreeze trailing layers") for block in llm_blocks[-unfreeze_last_n:]: for param in block.parameters(): param.requires_grad = True for attr_name in ("final_norm", "output_head", "lm_head"): module = getattr(self.llm, attr_name, None) if module is not None: for param in module.parameters(): param.requires_grad = True @staticmethod def _split_decay_groups(named_parameters: Iterable[tuple[str, nn.Parameter]]) -> tuple[list[nn.Parameter], list[nn.Parameter]]: """Split trainable params into decay and no-decay groups.""" decay_params: list[nn.Parameter] = [] no_decay_params: list[nn.Parameter] = [] for name, param in named_parameters: if not param.requires_grad: continue if "bias" in name or "norm" in name: no_decay_params.append(param) else: decay_params.append(param) return decay_params, no_decay_params def get_optimizer_param_groups(self, config) -> list[dict[str, Any]] | None: """Provide separate parameter groups for the vision tower and unfrozen LLM subset.""" vision_named_params = list(self.vision_encoder.named_parameters()) + list(self.vision_projector.named_parameters()) llm_named_params = list(self.llm.named_parameters()) vision_decay, vision_no_decay = self._split_decay_groups(vision_named_params) llm_decay, llm_no_decay = self._split_decay_groups(llm_named_params) param_groups: list[dict[str, Any]] = [] if vision_decay: param_groups.append({ "params": vision_decay, "lr": config.vision_learning_rate, "weight_decay": config.optimizer.weight_decay, }) if vision_no_decay: param_groups.append({ "params": vision_no_decay, "lr": config.vision_learning_rate, "weight_decay": 0.0, }) if llm_decay: param_groups.append({ "params": llm_decay, "lr": config.llm_learning_rate, "weight_decay": config.optimizer.weight_decay, }) if llm_no_decay: param_groups.append({ "params": llm_no_decay, "lr": config.llm_learning_rate, "weight_decay": 0.0, }) return param_groups or None def load_state_dict(self, state_dict, strict: bool = True): """Load either full multimodal checkpoints or text-only LLM checkpoints.""" current_state = self.state_dict() remapped_state = OrderedDict() for key, value in state_dict.items(): if key in current_state: remapped_state[key] = value elif f"llm.{key}" in current_state: remapped_state[f"llm.{key}"] = value else: remapped_state[key] = value return super().load_state_dict(remapped_state, strict=strict) def forward( self, input_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Inject visual prefix embeddings and delegate the language loss to the wrapped LLM.""" if inputs_embeds is not None: return self.llm( input_ids=None, attention_mask=attention_mask, labels=labels, inputs_embeds=inputs_embeds, ) if input_ids is None: raise ValueError("input_ids must be provided for multimodal forward passes") if pixel_values is None: raise ValueError("pixel_values must be provided for multimodal forward passes") if attention_mask is None: attention_mask = torch.ones_like(input_ids) self._validate_image_placeholder(attention_mask) visual_prefix = self._encode_visual_prefix(pixel_values) token_embeddings = self._embed_input_ids(input_ids) combined_embeddings = torch.cat([visual_prefix, token_embeddings[:, 1:, :]], dim=1) combined_attention_mask = self._build_attention_mask(attention_mask) combined_labels = self._build_labels(labels, combined_embeddings.device) if combined_embeddings.size(1) > self.llm.config.max_seq_length: raise ValueError( f"Expanded multimodal sequence length ({combined_embeddings.size(1)}) exceeds max_seq_length " f"({self.llm.config.max_seq_length})" ) return self.llm( input_ids=None, attention_mask=combined_attention_mask, labels=combined_labels, inputs_embeds=combined_embeddings, )