| """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 <image> 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,
|
| )
|
|
|
