from __future__ import annotations import math from dataclasses import dataclass from typing import Any, Callable, Optional, Tuple, Union import numpy as np import torch import torch.nn.functional as F from PIL import Image from torch import nn from transformers.activations import ACT2FN from transformers.generation import GenerationMixin from transformers.modeling_outputs import ( MoeCausalLMOutputWithPast, MoeModelOutputWithPast, ModelOutput, ) from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.masking_utils import ( create_causal_mask, create_sliding_window_causal_mask, ) from transformers.modeling_layers import GradientCheckpointingLayer from transformers.processing_utils import Unpack from transformers.utils import TransformersKwargs from transformers.cache_utils import Cache, DynamicCache from transformers.integrations import use_kernel_forward_from_hub try: from .configuration_trinity_vlm import AfmoeConfig, TrinityVLMConfig except Exception: from configuration_trinity_vlm import AfmoeConfig, TrinityVLMConfig def _compute_default_rope_parameters( config=None, device: torch.device | None = None, seq_len: int | None = None, layer_type: str | None = None, ) -> tuple[torch.Tensor, float]: del seq_len, layer_type if config is None: raise ValueError("config is required to compute default RoPE parameters.") base = getattr(config, "rope_theta", 10000.0) partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads dim = int(head_dim * partial_rotary_factor) inv_freq = 1.0 / ( base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) ) return inv_freq, 1.0 if "default" not in ROPE_INIT_FUNCTIONS: ROPE_INIT_FUNCTIONS["default"] = _compute_default_rope_parameters class AfmoeRotaryEmbedding(nn.Module): def __init__(self, config: AfmoeConfig, device=None): super().__init__() # BC: "rope_type" was originally "type" if hasattr(config, "rope_scaling") and config.rope_scaling is not None: self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) else: self.rope_type = "default" self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq def _dynamic_frequency_update(self, position_ids, device): """ dynamic RoPE layers should recompute `inv_freq` in the following situations: 1 - growing beyond the cached sequence length (allow scaling) 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) """ seq_len = torch.max(position_ids) + 1 if seq_len > self.max_seq_len_cached: # growth inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len) self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation self.max_seq_len_cached = seq_len if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset # This .to() is needed if the model has been moved to a device after being initialized (because # the buffer is automatically moved, but not the original copy) self.original_inv_freq = self.original_inv_freq.to(device) self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len def compute_default_rope_parameters( self, config=None, device: torch.device | None = None, seq_len: int | None = None, layer_type: str | None = None, ) -> tuple[torch.Tensor, float]: return _compute_default_rope_parameters( config=config or self.config, device=device, seq_len=seq_len, layer_type=layer_type, ) @torch.no_grad() def forward(self, x, position_ids): if "dynamic" in self.rope_type: self._dynamic_frequency_update(position_ids, device=x.device) # Core RoPE block inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() # Force float32 (see https://github.com/huggingface/transformers/pull/29285) device_type = x.device.type device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" with torch.autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention cos = cos * self.attention_scaling sin = sin * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. position_ids (`torch.Tensor`, *optional*): Deprecated and unused. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_key_value_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) @use_kernel_forward_from_hub("RMSNorm") class AfmoeRMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float): """ AfmoeRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], scaling: float, dropout: float = 0.0, **kwargs, ): key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( query.dtype ) attn_weights = nn.functional.dropout( attn_weights, p=dropout, training=module.training ) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights class AfmoeMLP(nn.Module): def __init__(self, config, intermediate_size=None): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = intermediate_size or config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class AfmoeTokenChoiceRouter(nn.Module): """Token-choice top-K router for MoE routing.""" def __init__(self, config): super().__init__() self.config = config self.top_k = config.num_experts_per_tok self.num_experts = config.num_experts self.score_func = config.score_func self.route_norm = config.route_norm self.route_scale = config.route_scale self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) def forward(self, hidden_states, expert_bias: torch.Tensor | None): _, _, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) scores = self.gate(hidden_states) # Apply scoring function in float32 for stability if self.score_func == "sigmoid": scores = torch.sigmoid(scores.to(torch.float32)) else: scores = F.softmax(scores.to(torch.float32), dim=-1) if expert_bias is not None: _, selected_experts = torch.topk(scores + expert_bias, k=self.top_k, dim=1) top_scores = scores.gather(dim=1, index=selected_experts) else: top_scores, selected_experts = torch.topk(scores, k=self.top_k, dim=1) # Normalize weights if using sigmoid if self.score_func == "sigmoid" and self.route_norm: denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20 top_scores = top_scores / denominator top_scores = top_scores * self.route_scale return top_scores, selected_experts def _can_use_grouped_mm(hidden_states: torch.Tensor) -> bool: return ( hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16 and hasattr(F, "grouped_mm") ) def _router_forward( router: nn.Module, hidden_states: torch.Tensor, expert_bias: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: _, _, hidden_dim = hidden_states.shape hidden_states_flat = hidden_states.view(-1, hidden_dim) router_logits = router.gate(hidden_states_flat) if router.score_func == "sigmoid": router_probs = torch.sigmoid(router_logits.to(torch.float32)) else: router_probs = F.softmax(router_logits.to(torch.float32), dim=-1) if expert_bias is not None: _, selected_experts = torch.topk(router_probs + expert_bias, k=router.top_k, dim=1) top_scores = router_probs.gather(dim=1, index=selected_experts) else: top_scores, selected_experts = torch.topk(router_probs, k=router.top_k, dim=1) if router.score_func == "sigmoid" and router.route_norm: denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20 top_scores = top_scores / denominator top_scores = top_scores * router.route_scale return hidden_states_flat, router_logits, router_probs, top_scores, selected_experts def _router_aux_loss( router_probs: torch.Tensor, selected_experts: torch.Tensor, *, num_experts: int, ) -> torch.Tensor: selected_flat = selected_experts.reshape(-1) top_k = max(1, selected_experts.shape[-1]) token_count = max(1, selected_experts.shape[0]) tokens_per_expert = torch.bincount( selected_flat, minlength=num_experts, ).to(torch.float32) / float(token_count * top_k) router_prob_per_expert = router_probs.mean(dim=0) return num_experts * torch.sum(tokens_per_expert * router_prob_per_expert) def _get_grouped_projection_weights( moe_layer: nn.Module, expert_ids: torch.Tensor, *, projection_name: str, ) -> torch.Tensor: if expert_ids.numel() == 0: raise ValueError("Cannot select grouped weights for an empty expert set.") packed_weights = getattr(moe_layer, f"packed_{projection_name}", None) if isinstance(packed_weights, nn.Parameter): if expert_ids.numel() == packed_weights.shape[0]: full_expert_ids = torch.arange( packed_weights.shape[0], device=expert_ids.device, dtype=expert_ids.dtype, ) if torch.equal(expert_ids, full_expert_ids): return packed_weights return packed_weights.index_select(0, expert_ids).contiguous() experts = getattr(moe_layer, "experts", None) if experts is None: raise ValueError(f"Layer has neither packed experts nor per-expert modules for {projection_name}.") projection_weights = [] requires_grad = False for expert_id in expert_ids.tolist(): weight = getattr(experts[expert_id], projection_name).weight requires_grad = requires_grad or weight.requires_grad projection_weights.append(weight.transpose(0, 1)) if not requires_grad: projection_weights = [weight.detach() for weight in projection_weights] return torch.stack(projection_weights, dim=0).contiguous() def _dense_packed_moe_forward( moe_layer: nn.Module, routed_input: torch.Tensor, token_to_expert: torch.Tensor, ) -> torch.Tensor: routed_output = torch.zeros( routed_input.shape[0], moe_layer.config.hidden_size, device=routed_input.device, dtype=routed_input.dtype, ) packed_gate_proj = getattr(moe_layer, "packed_gate_proj", None) packed_up_proj = getattr(moe_layer, "packed_up_proj", None) packed_down_proj = getattr(moe_layer, "packed_down_proj", None) if not all( isinstance(weight, nn.Parameter) for weight in (packed_gate_proj, packed_up_proj, packed_down_proj) ): for expert_id in range(moe_layer.config.num_experts): mask = token_to_expert == expert_id if not mask.any(): continue expert_input = routed_input[mask] expert_out = moe_layer.experts[expert_id](expert_input) routed_output[mask] = expert_out return routed_output act_fn = ACT2FN[moe_layer.config.hidden_act] for expert_id in range(moe_layer.config.num_experts): mask = token_to_expert == expert_id if not mask.any(): continue expert_input = routed_input[mask] gate_proj = F.linear(expert_input, packed_gate_proj[expert_id].transpose(0, 1)) up_proj = F.linear(expert_input, packed_up_proj[expert_id].transpose(0, 1)) activated = act_fn(gate_proj) * up_proj expert_out = F.linear(activated, packed_down_proj[expert_id].transpose(0, 1)) routed_output[mask] = expert_out return routed_output def _accumulate_routed_output( shared_output: torch.Tensor, routed_output: torch.Tensor, top_scores_sorted: torch.Tensor, token_indices_sorted: torch.Tensor, ) -> torch.Tensor: output = shared_output.to(torch.float32) if routed_output.numel() == 0: return output hidden_dim = routed_output.shape[-1] bytes_per_row = max(1, hidden_dim * 4) target_chunk_bytes = 16 * 1024 * 1024 rows_per_chunk = max(1, target_chunk_bytes // bytes_per_row) for start in range(0, routed_output.shape[0], rows_per_chunk): end = min(start + rows_per_chunk, routed_output.shape[0]) weighted_chunk = routed_output[start:end].to(torch.float32) weighted_chunk.mul_(top_scores_sorted[start:end].unsqueeze(-1)) output.index_add_(0, token_indices_sorted[start:end], weighted_chunk) return output class AfmoeMoE(nn.Module): def __init__(self, config): super().__init__() self.config = config self.use_packed_experts = bool(getattr(config, "packed_experts", False)) self.router = AfmoeTokenChoiceRouter(config) self.shared_experts = None if config.num_shared_experts > 0: self.shared_experts = AfmoeMLP( config, config.moe_intermediate_size * config.num_shared_experts ) self.expert_bias = nn.Parameter(torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False) if self.use_packed_experts: self.experts = None self.packed_gate_proj = nn.Parameter( torch.empty(config.num_experts, config.hidden_size, config.moe_intermediate_size) ) self.packed_up_proj = nn.Parameter( torch.empty(config.num_experts, config.hidden_size, config.moe_intermediate_size) ) self.packed_down_proj = nn.Parameter( torch.empty(config.num_experts, config.moe_intermediate_size, config.hidden_size) ) self.reset_parameters() else: self.experts = nn.ModuleList( [ AfmoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts) ] ) def reset_parameters(self): std = float(getattr(self.config, "initializer_range", 0.02)) nn.init.normal_(self.packed_gate_proj, mean=0.0, std=std) nn.init.normal_(self.packed_up_proj, mean=0.0, std=std) nn.init.normal_(self.packed_down_proj, mean=0.0, std=std) with torch.no_grad(): self.expert_bias.zero_() def forward(self, hidden_states): batch_size, seq_len, hidden_dim = hidden_states.shape hidden_states_flat, _router_logits, router_probs, top_scores, selected_experts = _router_forward( self.router, hidden_states, self.expert_bias, ) if self.shared_experts is not None: shared_output = self.shared_experts(hidden_states_flat) else: shared_output = torch.zeros_like(hidden_states_flat) token_indices_sorted = torch.argsort(selected_experts.view(-1), stable=True) top_scores_sorted = top_scores.view(-1)[token_indices_sorted] token_to_expert = selected_experts.view(-1)[token_indices_sorted] token_indices_sorted = token_indices_sorted // self.config.num_experts_per_tok token_indices_expanded = token_indices_sorted.unsqueeze(-1).expand(-1, hidden_dim) routed_input = torch.gather(hidden_states_flat, dim=0, index=token_indices_expanded).contiguous() routed_output: torch.Tensor | None = None use_grouped_mm = bool(getattr(self.config, "enable_grouped_moe", True)) and _can_use_grouped_mm( routed_input ) if use_grouped_mm: expert_counts = torch.bincount( token_to_expert, minlength=self.config.num_experts, ) grouped_offsets = torch.cumsum( expert_counts, dim=0, dtype=torch.int32, ) packed_gate_proj = getattr(self, "packed_gate_proj", None) packed_up_proj = getattr(self, "packed_up_proj", None) packed_down_proj = getattr(self, "packed_down_proj", None) if all( isinstance(weight, nn.Parameter) for weight in (packed_gate_proj, packed_up_proj, packed_down_proj) ): gate_weights = packed_gate_proj up_weights = packed_up_proj down_weights = packed_down_proj else: active_expert_ids = torch.nonzero(expert_counts > 0, as_tuple=False).flatten() if active_expert_ids.numel() == 0: routed_output = torch.zeros_like(routed_input) gate_weights = up_weights = down_weights = None else: grouped_offsets = torch.cumsum( expert_counts.index_select(0, active_expert_ids), dim=0, dtype=torch.int32, ) gate_weights = _get_grouped_projection_weights( self, active_expert_ids, projection_name="gate_proj", ) up_weights = _get_grouped_projection_weights( self, active_expert_ids, projection_name="up_proj", ) down_weights = _get_grouped_projection_weights( self, active_expert_ids, projection_name="down_proj", ) if routed_output is None: gate_proj = F.grouped_mm(routed_input, gate_weights, offs=grouped_offsets) up_proj = F.grouped_mm(routed_input, up_weights, offs=grouped_offsets) activated = ACT2FN[self.config.hidden_act](gate_proj) * up_proj routed_output = F.grouped_mm(activated, down_weights, offs=grouped_offsets) else: routed_output = _dense_packed_moe_forward( self, routed_input, token_to_expert, ) if routed_output is None: raise RuntimeError("MoE forward did not produce routed output.") if use_grouped_mm: del expert_counts, grouped_offsets if "active_expert_ids" in locals(): del active_expert_ids if "gate_weights" in locals(): del gate_weights, up_weights, down_weights if "gate_proj" in locals(): del gate_proj, up_proj, activated output = _accumulate_routed_output( shared_output=shared_output, routed_output=routed_output, top_scores_sorted=top_scores_sorted, token_indices_sorted=token_indices_sorted, ) aux_loss_coef = float(getattr(self.config, "router_aux_loss_coef", 0.0) or 0.0) self._last_router_aux_loss = None if aux_loss_coef > 0.0: self._last_router_aux_loss = _router_aux_loss( router_probs, selected_experts, num_experts=self.config.num_experts, ) self._last_router_logits = None if getattr(self.config, "output_router_logits", False): self._last_router_logits = router_probs.view(batch_size, seq_len, self.config.num_experts) return output.to(hidden_states.dtype).view(batch_size, seq_len, hidden_dim) class AfmoeAttention(nn.Module): """Multi-headed attention with local/global pattern and gating.""" def __init__(self, config: AfmoeConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention" self.sliding_window = config.sliding_window if self.is_local_attention else None self.q_proj = nn.Linear( config.hidden_size, self.num_heads * self.head_dim, bias=False ) self.k_proj = nn.Linear( config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False ) self.v_proj = nn.Linear( config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False ) self.o_proj = nn.Linear( self.num_heads * self.head_dim, config.hidden_size, bias=False ) self.q_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.gate_proj = nn.Linear( config.hidden_size, self.num_heads * self.head_dim, bias=False ) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape) key_states = self.k_proj(hidden_states).view(hidden_shape) value_states = self.v_proj(hidden_states).view(hidden_shape) gate_states = self.gate_proj(hidden_states) query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) if self.is_local_attention: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: cache_kwargs = {"cache_position": cache_position} key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, cache_kwargs, ) attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": attention_interface = ALL_ATTENTION_FUNCTIONS[ self.config._attn_implementation ] output, _ = attention_interface( self, query_states, key_states, value_states, attention_mask=attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, sliding_window=self.sliding_window, **kwargs, ) output = output.view(*input_shape, -1).contiguous() output = output * F.sigmoid(gate_states) return self.o_proj(output) class AfmoeDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: AfmoeConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.layer_idx = layer_idx self.self_attn = AfmoeAttention(config=config, layer_idx=layer_idx) self.attention_type = config.layer_types[layer_idx] # Dual normalization for attention self.input_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) # Dual normalization for FFN self.pre_mlp_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_mlp_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) # MoE or dense FFN self.moe_enabled = layer_idx >= config.num_dense_layers if self.moe_enabled: self.mlp = AfmoeMoE(config) else: self.mlp = AfmoeMLP(config) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.FloatTensor: residual = hidden_states # Self Attention with dual normalization hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states # FFN with dual normalization residual = hidden_states hidden_states = self.pre_mlp_layernorm(hidden_states) if self.moe_enabled: hidden_states = self.mlp(hidden_states) else: hidden_states = self.mlp(hidden_states) hidden_states = self.post_mlp_layernorm(hidden_states) hidden_states = residual + hidden_states return hidden_states class AfmoePreTrainedModel(PreTrainedModel): config_class = AfmoeConfig base_model_prefix = "model" _no_split_modules = ["AfmoeDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _keep_in_fp32_modules = [ "input_layernorm", "post_attention_layernorm", "pre_mlp_layernorm", "post_mlp_layernorm", "q_norm", "k_norm", "norm", ] _supports_sdpa = True _supports_attention_backend = True supports_gradient_checkpointing = True class AfmoeModel(AfmoePreTrainedModel): _no_split_modules = ["AfmoeDecoderLayer"] def __init__(self, config: AfmoeConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) self.layers = nn.ModuleList( [ AfmoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self.norm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = AfmoeRotaryEmbedding(config=config) self.gradient_checkpointing = False self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> MoeModelOutputWithPast: if input_ids is None and inputs_embeds is None: raise ValueError("You must specify at least one of input_ids or inputs_embeds.") if use_cache and past_key_values is None: past_key_values = DynamicCache() if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, ) if position_ids is None: position_ids = cache_position.unsqueeze(0) if not isinstance(causal_mask_mapping := attention_mask, dict): mask_kwargs = { "config": self.config, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "cache_position": cache_position, "past_key_values": past_key_values, } causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), } hidden_states = inputs_embeds if self.config.mup_enabled: hidden_states = hidden_states * (self.config.hidden_size**0.5) position_embeddings = self.rotary_emb(hidden_states, position_ids) router_aux_losses = [] collect_router_logits = bool(getattr(self.config, "output_router_logits", False)) router_logits = [] if collect_router_logits else None for decoder_layer in self.layers: hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask_mapping[decoder_layer.attention_type], position_ids=position_ids, past_key_value=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) if not getattr(decoder_layer, "moe_enabled", False): continue layer_aux_loss = getattr(decoder_layer.mlp, "_last_router_aux_loss", None) if layer_aux_loss is not None: router_aux_losses.append(layer_aux_loss) if router_logits is not None: layer_router_logits = getattr(decoder_layer.mlp, "_last_router_logits", None) if layer_router_logits is not None: router_logits.append(layer_router_logits) decoder_layer.mlp._last_router_aux_loss = None decoder_layer.mlp._last_router_logits = None hidden_states = self.norm(hidden_states) self._last_router_aux_loss = ( torch.stack(router_aux_losses).mean() if router_aux_losses else None ) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, router_logits=tuple(router_logits) if router_logits else None, ) class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): super().__init__(config) self.model = AfmoeModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_type_ids: Optional[torch.Tensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> Union[Tuple, MoeCausalLMOutputWithPast]: del token_type_ids outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, cache_position=cache_position, **kwargs, ) hidden_states = outputs.last_hidden_state aux_loss = getattr(self.model, "_last_router_aux_loss", None) aux_loss_coef = float(getattr(self.config, "router_aux_loss_coef", 0.0) or 0.0) slice_indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep ) logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None: loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) if loss is not None and aux_loss is not None and aux_loss_coef > 0.0: loss = loss + (aux_loss * aux_loss_coef) return MoeCausalLMOutputWithPast( loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, router_logits=outputs.router_logits, ) @dataclass(frozen=True) class MoondreamVisionConfig: enc_dim: int = 1152 enc_patch_size: int = 14 enc_n_layers: int = 27 enc_ff_dim: int = 4304 enc_n_heads: int = 16 proj_out_dim: int = 2048 crop_size: int = 378 in_channels: int = 3 max_crops: int = 12 overlap_margin: int = 4 proj_inner_dim: int = 8192 @property def image_seq_len(self) -> int: return (self.crop_size // self.enc_patch_size) ** 2 def select_tiling(height: int, width: int, crop_size: int, max_crops: int) -> tuple[int, int]: if height <= crop_size or width <= crop_size: return (1, 1) min_h = math.ceil(height / crop_size) min_w = math.ceil(width / crop_size) if min_h * min_w > max_crops: ratio = math.sqrt(max_crops / (min_h * min_w)) return (max(1, math.floor(min_h * ratio)), max(1, math.floor(min_w * ratio))) h_tiles = math.floor(math.sqrt(max_crops * height / width)) w_tiles = math.floor(math.sqrt(max_crops * width / height)) h_tiles = max(h_tiles, min_h) w_tiles = max(w_tiles, min_w) if h_tiles * w_tiles > max_crops: if w_tiles > h_tiles: w_tiles = math.floor(max_crops / h_tiles) else: h_tiles = math.floor(max_crops / w_tiles) return (max(1, h_tiles), max(1, w_tiles)) def overlap_crop_image( image: np.ndarray, overlap_margin: int, max_crops: int, base_size: tuple[int, int] = (378, 378), patch_size: int = 14, ) -> tuple[np.ndarray, tuple[int, int]]: original_h, original_w = image.shape[:2] margin_pixels = patch_size * overlap_margin total_margin_pixels = margin_pixels * 2 crop_patches = base_size[0] // patch_size crop_window_patches = crop_patches - (2 * overlap_margin) crop_window_size = crop_window_patches * patch_size tiling = select_tiling( original_h - total_margin_pixels, original_w - total_margin_pixels, crop_window_size, max_crops, ) n_crops = tiling[0] * tiling[1] + 1 crops = np.zeros((n_crops, base_size[0], base_size[1], image.shape[2]), dtype=np.uint8) target_size = ( tiling[0] * crop_window_size + total_margin_pixels, tiling[1] * crop_window_size + total_margin_pixels, ) pil_img = Image.fromarray(image) resized = pil_img.resize( (int(target_size[1]), int(target_size[0])), resample=Image.Resampling.LANCZOS, ) image = np.asarray(resized) global_pil = pil_img.resize( (int(base_size[1]), int(base_size[0])), resample=Image.Resampling.LANCZOS, ) crops[0] = np.asarray(global_pil) for i in range(tiling[0]): for j in range(tiling[1]): y0 = i * crop_window_size x0 = j * crop_window_size y_end = min(y0 + base_size[0], image.shape[0]) x_end = min(x0 + base_size[1], image.shape[1]) crop_region = image[y0:y_end, x0:x_end] crops[1 + i * tiling[1] + j, : crop_region.shape[0], : crop_region.shape[1]] = crop_region return crops, tiling @torch.compiler.disable def reconstruct_from_crops( crops: torch.Tensor, tiling: tuple[int, int], overlap_margin: int, patch_size: int = 14, ) -> torch.Tensor: tiling_h, tiling_w = tiling crop_height, crop_width = crops[0].shape[:2] margin_pixels = overlap_margin * patch_size output_h = (crop_height - 2 * margin_pixels) * tiling_h + 2 * margin_pixels output_w = (crop_width - 2 * margin_pixels) * tiling_w + 2 * margin_pixels reconstructed = torch.zeros( (output_h, output_w, crops[0].shape[2]), device=crops[0].device, dtype=crops[0].dtype, ) for i, crop in enumerate(crops): tile_y = i // tiling_w tile_x = i % tiling_w x_start = 0 if tile_x == 0 else margin_pixels x_end = crop_width if tile_x == tiling_w - 1 else crop_width - margin_pixels y_start = 0 if tile_y == 0 else margin_pixels y_end = crop_height if tile_y == tiling_h - 1 else crop_height - margin_pixels out_x = tile_x * (crop_width - 2 * margin_pixels) out_y = tile_y * (crop_height - 2 * margin_pixels) reconstructed[ out_y + y_start : out_y + y_end, out_x + x_start : out_x + x_end, ] = crop[y_start:y_end, x_start:x_end] return reconstructed @torch.compiler.disable def prepare_crops( image: Image.Image, config: MoondreamVisionConfig, device: torch.device, dtype: torch.dtype, ) -> tuple[torch.Tensor, tuple[int, int]]: np_image = np.array(image.convert("RGB")) crops, tiling = overlap_crop_image( np_image, max_crops=config.max_crops, overlap_margin=config.overlap_margin, base_size=(config.crop_size, config.crop_size), patch_size=config.enc_patch_size, ) crops = np.transpose(crops, (0, 3, 1, 2)) crops_tensor = torch.from_numpy(crops).to(device=device, dtype=dtype) crops_tensor = crops_tensor.div_(255.0).sub_(0.5).div_(0.5) return crops_tensor, tiling def create_patches(x: torch.Tensor, patch_size: int) -> torch.Tensor: batch, channels, height, width = x.shape x = x.reshape(batch, channels, height // patch_size, patch_size, width // patch_size, patch_size) x = x.permute(0, 2, 4, 1, 3, 5) x = x.reshape(batch, (height // patch_size) * (width // patch_size), channels * patch_size * patch_size) return x class MoondreamAttention(nn.Module): def __init__(self, dim: int, n_heads: int, dtype: torch.dtype) -> None: super().__init__() self.n_heads = n_heads self.qkv = nn.Linear(dim, 3 * dim, dtype=dtype) self.proj = nn.Linear(dim, dim, dtype=dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: batch, seq_len, dim = x.shape head_dim = dim // self.n_heads q, k, v = self.qkv(x).chunk(3, dim=-1) q = q.view(batch, seq_len, self.n_heads, head_dim).transpose(1, 2) k = k.view(batch, seq_len, self.n_heads, head_dim).transpose(1, 2) v = v.view(batch, seq_len, self.n_heads, head_dim).transpose(1, 2) out = F.scaled_dot_product_attention(q, k, v) out = out.transpose(1, 2).reshape(batch, seq_len, dim) return self.proj(out) class MoondreamMLP(nn.Module): def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, dtype: torch.dtype) -> None: super().__init__() self.fc1 = nn.Linear(in_dim, hidden_dim, dtype=dtype) self.fc2 = nn.Linear(hidden_dim, out_dim, dtype=dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: x = F.gelu(self.fc1(x), approximate="tanh") return self.fc2(x) class MoondreamVisionBlock(nn.Module): def __init__(self, config: MoondreamVisionConfig, dtype: torch.dtype) -> None: super().__init__() self.ln1 = nn.LayerNorm(config.enc_dim, dtype=dtype) self.attn = MoondreamAttention(config.enc_dim, config.enc_n_heads, dtype) self.ln2 = nn.LayerNorm(config.enc_dim, dtype=dtype) self.mlp = MoondreamMLP(config.enc_dim, config.enc_ff_dim, config.enc_dim, dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class MoondreamVisionTower(nn.Module): def __init__( self, config: MoondreamVisionConfig | None = None, dtype: torch.dtype = torch.bfloat16, ) -> None: super().__init__() self.config = config or MoondreamVisionConfig() self.patch_emb = nn.Linear( self.config.enc_patch_size * self.config.enc_patch_size * self.config.in_channels, self.config.enc_dim, dtype=dtype, ) self.blocks = nn.ModuleList( [MoondreamVisionBlock(self.config, dtype=dtype) for _ in range(self.config.enc_n_layers)] ) self.post_ln = nn.LayerNorm(self.config.enc_dim, dtype=dtype) self.proj_mlp = MoondreamMLP( self.config.enc_dim * 2, self.config.proj_inner_dim, self.config.proj_out_dim, dtype, ) self.pos_emb = nn.Parameter( torch.zeros(1, self.config.image_seq_len, self.config.enc_dim, dtype=dtype) ) @property def image_seq_len(self) -> int: return self.config.image_seq_len def encode_crops(self, inputs_bchw: torch.Tensor) -> torch.Tensor: x = create_patches(inputs_bchw, self.config.enc_patch_size) x = self.patch_emb(x) x = x + self.pos_emb for block in self.blocks: x = block(x) return self.post_ln(x) def project_features(self, global_features: torch.Tensor, reconstructed: torch.Tensor) -> torch.Tensor: reconstructed = reconstructed.permute(2, 0, 1) reconstructed = F.adaptive_avg_pool2d( reconstructed, output_size=(self.config.enc_n_layers, self.config.enc_n_layers), ) reconstructed = reconstructed.permute(1, 2, 0).reshape(self.image_seq_len, self.config.enc_dim) return self.proj_mlp(torch.cat([global_features, reconstructed], dim=-1)) def encode_image(self, image: Image.Image) -> torch.Tensor: if not isinstance(image, Image.Image): raise TypeError(f"Expected PIL image, got {type(image)!r}") device = self.pos_emb.device dtype = self.pos_emb.dtype crops, tiling = prepare_crops(image, self.config, device=device, dtype=dtype) outputs = self.encode_crops(crops) global_features = outputs[0] local_features = outputs[1:].view( -1, self.config.enc_n_layers, self.config.enc_n_layers, self.config.enc_dim, ) reconstructed = reconstruct_from_crops( local_features, tiling, patch_size=1, overlap_margin=self.config.overlap_margin, ) return self.project_features(global_features, reconstructed) def encode_images(self, images: list[Image.Image]) -> torch.Tensor: encoded = [self.encode_image(image) for image in images] return torch.stack(encoded, dim=0) def build_image_token_span( *, image_start_token_id: int | None, image_token_id: int | None, image_end_token_id: int | None, image_seq_len: int, bos_token_id: int | None = None, ) -> list[int]: if image_start_token_id is None: raise ValueError("image_start_token_id is not configured.") if image_token_id is None: raise ValueError("image_token_id is not configured.") if image_end_token_id is None: raise ValueError("image_end_token_id is not configured.") token_ids: list[int] = [] if bos_token_id is not None: token_ids.append(bos_token_id) token_ids.append(image_start_token_id) token_ids.extend([image_token_id] * image_seq_len) token_ids.append(image_end_token_id) return token_ids @dataclass class TrinityVLMCausalLMOutputWithPast(ModelOutput): loss: torch.FloatTensor | None = None aux_loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None past_key_values: Cache | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None router_logits: tuple[torch.FloatTensor] | None = None image_hidden_states: torch.FloatTensor | None = None class VisionBridge(nn.Module): def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, dtype: torch.dtype) -> None: super().__init__() self.norm = nn.LayerNorm(in_dim, dtype=dtype) self.fc1 = nn.Linear(in_dim, hidden_dim, dtype=dtype) self.fc2 = nn.Linear(hidden_dim, out_dim, dtype=dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.norm(x) x = torch.nn.functional.gelu(self.fc1(x), approximate="tanh") return self.fc2(x) class TrinityVLMForConditionalGeneration(PreTrainedModel, GenerationMixin): config_class = TrinityVLMConfig base_model_prefix = "trinity_vlm" main_input_name = "input_ids" def __init__(self, config: TrinityVLMConfig) -> None: super().__init__(config) torch_dtype = self._resolve_torch_dtype(config) vision_config = dict(config.vision_config) projector_hidden_dim = int(vision_config.pop("projector_hidden_dim", config.projector_hidden_dim)) self.torch_dtype = torch_dtype self.language_model = AfmoeForCausalLM(self._load_trinity_text_config(config)) self.vision_tower = MoondreamVisionTower( config=MoondreamVisionConfig(**vision_config), dtype=torch_dtype, ) self.multi_modal_projector = VisionBridge( in_dim=config.vision_feature_dim, hidden_dim=projector_hidden_dim, out_dim=self.language_model.config.hidden_size, dtype=torch_dtype, ) self.config.hidden_size = self.language_model.config.hidden_size self.config.vocab_size = self.language_model.config.vocab_size self.config.bos_token_id = self.language_model.config.bos_token_id self.config.eos_token_id = self.language_model.config.eos_token_id self.config.pad_token_id = self.language_model.config.pad_token_id self.post_init() @staticmethod def _resolve_torch_dtype(config: TrinityVLMConfig) -> torch.dtype: dtype_value = getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) if isinstance(dtype_value, str) and hasattr(torch, dtype_value): return getattr(torch, dtype_value) if isinstance(dtype_value, torch.dtype): return dtype_value return torch.bfloat16 @staticmethod def _load_trinity_text_config(config: TrinityVLMConfig) -> AfmoeConfig: if not getattr(config, "text_config", None): raise ValueError("TrinityVLMConfig.text_config must be present in config.json.") text_config = AfmoeConfig(**config.text_config) text_config.vocab_size = config.vocab_size text_config.bos_token_id = config.bos_token_id text_config.eos_token_id = config.eos_token_id text_config.pad_token_id = config.pad_token_id text_config.packed_experts = True text_config.enable_grouped_moe = getattr(config, "enable_grouped_moe", True) text_config.output_router_logits = getattr(config, "output_router_logits", False) text_config._attn_implementation = "sdpa" return text_config @property def device(self) -> torch.device: return self.get_input_embeddings().weight.device @property def dtype(self) -> torch.dtype: return self.get_input_embeddings().weight.dtype def get_input_embeddings(self): return self.language_model.get_input_embeddings() def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) def get_output_embeddings(self): return self.language_model.get_output_embeddings() def set_output_embeddings(self, new_embeddings): self.language_model.set_output_embeddings(new_embeddings) def get_decoder(self): return self.language_model.get_decoder() def set_decoder(self, decoder): self.language_model.set_decoder(decoder) def build_image_token_span(self, *, include_bos: bool = True) -> list[int]: return build_image_token_span( image_start_token_id=self.config.image_start_token_id, image_token_id=self.config.image_token_id, image_end_token_id=self.config.image_end_token_id, image_seq_len=self.config.image_seq_len, bos_token_id=self.config.bos_token_id if include_bos else None, ) def _project_image_feature_tensor(self, image_features: torch.Tensor) -> torch.Tensor: if image_features.shape[-1] == self.language_model.config.hidden_size: return image_features.to(device=self.device, dtype=self.dtype) if image_features.shape[-1] == self.config.vision_feature_dim: return self.multi_modal_projector(image_features.to(device=self.device, dtype=self.dtype)) raise ValueError("Tensor image features must already be in text hidden size or vision feature size.") @torch.compiler.disable def get_image_features( self, images: list[Any] | list[list[Any]] | torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor] | None: if images is None: return None if isinstance(images, torch.Tensor): if images.ndim == 3: projected = self._project_image_feature_tensor(images) counts = torch.ones(projected.size(0), device=projected.device, dtype=torch.long) return projected, counts if images.ndim == 4: batch, num_images = images.shape[:2] projected = self._project_image_feature_tensor(images.flatten(0, 1)) counts = torch.full((batch,), num_images, device=projected.device, dtype=torch.long) return projected, counts raise ValueError("Tensor images must have shape [n_images, seq, dim] or [batch, images, seq, dim].") if not isinstance(images, (list, tuple)): raise TypeError(f"Unsupported image batch type: {type(images)!r}") if not images: empty_features = torch.empty( 0, self.config.image_seq_len, self.language_model.config.hidden_size, device=self.device, dtype=self.dtype, ) empty_counts = torch.empty(0, device=self.device, dtype=torch.long) return empty_features, empty_counts first_item = images[0] if isinstance(first_item, Image.Image): image_batches = [[image] for image in images] elif isinstance(first_item, (list, tuple)): image_batches = [list(sample_images) for sample_images in images] else: raise TypeError(f"Unsupported image batch type: {type(first_item)!r}") flat_images: list[Image.Image] = [] image_counts = [] for sample_images in image_batches: for image in sample_images: if not isinstance(image, Image.Image): raise TypeError(f"Expected PIL images, got {type(image)!r}") flat_images.extend(sample_images) image_counts.append(len(sample_images)) if not flat_images: empty_features = torch.empty( 0, self.config.image_seq_len, self.language_model.config.hidden_size, device=self.device, dtype=self.dtype, ) return empty_features, torch.tensor(image_counts, device=self.device, dtype=torch.long) image_features = self.vision_tower.encode_images(flat_images).to(device=self.device, dtype=self.dtype) projected = self.multi_modal_projector(image_features) return projected, torch.tensor(image_counts, device=self.device, dtype=torch.long) def _get_placeholder_mask( self, input_ids: torch.LongTensor | None, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor, ) -> torch.BoolTensor: if input_ids is None: image_token = torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) special_image_mask = inputs_embeds == self.get_input_embeddings()(image_token) special_image_mask = special_image_mask.all(-1) else: special_image_mask = input_ids == self.config.image_token_id n_image_tokens = special_image_mask.sum() n_image_features = image_features.shape[0] * image_features.shape[1] torch._assert( n_image_tokens == n_image_features, f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}", ) return special_image_mask def _merge_image_features( self, input_ids: torch.LongTensor | None, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor, image_counts: torch.LongTensor, ) -> torch.FloatTensor: if input_ids is not None: num_images = image_counts.to(dtype=torch.long) expected_image_token_counts = num_images * image_features.shape[1] actual_image_token_counts = (input_ids == self.config.image_token_id).sum(dim=1) torch._assert( torch.all(actual_image_token_counts == expected_image_token_counts), "Image placeholder count mismatch.", ) if self.config.image_start_token_id is not None: start_counts = (input_ids == self.config.image_start_token_id).sum(dim=1) torch._assert(torch.all(start_counts == num_images), "image_start token count mismatch.") if self.config.image_end_token_id is not None: end_counts = (input_ids == self.config.image_end_token_id).sum(dim=1) torch._assert(torch.all(end_counts == num_images), "image_end token count mismatch.") special_image_mask = self._get_placeholder_mask(input_ids, inputs_embeds, image_features) special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) flat_image_features = image_features.reshape(-1, image_features.shape[-1]).to( inputs_embeds.device, inputs_embeds.dtype, ) return inputs_embeds.masked_scatter(special_image_mask, flat_image_features) def forward( self, input_ids: torch.LongTensor | None = None, images: list[Any] | list[list[Any]] | torch.Tensor | None = None, attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, use_cache: bool | None = None, **kwargs, ) -> tuple | TrinityVLMCausalLMOutputWithPast: if (input_ids is None) == (inputs_embeds is None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) image_hidden_states = None if images is not None: image_outputs = self.get_image_features(images) if image_outputs is not None: image_features, image_counts = image_outputs if image_features.numel() > 0: image_hidden_states = image_features inputs_embeds = self._merge_image_features( input_ids=input_ids, inputs_embeds=inputs_embeds, image_features=image_features, image_counts=image_counts, ) outputs = self.language_model.model( input_ids=None, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs, ) hidden_states = outputs.last_hidden_state aux_loss = getattr(self.language_model.model, "_last_router_aux_loss", None) aux_loss_coef = float(getattr(self.config, "router_aux_loss_coef", 0.0) or 0.0) slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.language_model.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if loss is not None and aux_loss is not None and aux_loss_coef > 0.0: loss = loss + (aux_loss * aux_loss_coef) return TrinityVLMCausalLMOutputWithPast( loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, router_logits=outputs.router_logits, image_hidden_states=image_hidden_states, ) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, inputs_embeds=None, images=None, attention_mask=None, logits_to_keep=None, is_first_iteration=False, **kwargs, ): model_inputs = super().prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, attention_mask=attention_mask, logits_to_keep=logits_to_keep, is_first_iteration=is_first_iteration, **kwargs, ) if is_first_iteration or not kwargs.get("use_cache", True): model_inputs["images"] = images return model_inputs __all__ = [ "AfmoeConfig", "AfmoeForCausalLM", "AfmoeModel", "AfmoePreTrainedModel", "MoondreamVisionConfig", "MoondreamVisionTower", "TrinityVLMForConditionalGeneration", ]