| """LFM2-VL and LFM2 on the hybrid stack. |
| |
| HF's classes stay the frame (the SigLIP2 tower, the projector, the image |
| features, loading and saving); only the language model is replaced, with HF's |
| module names, so checkpoints load and save unchanged. A state with any number |
| of questions runs as one `hybrid.Tree`, with the same outputs as HF's own |
| implementation. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import Lfm2Config, Lfm2VlConfig |
| from transformers.models.lfm2.modeling_lfm2 import Lfm2ForCausalLM as HfLfm2 |
| from transformers.models.lfm2_vl.modeling_lfm2_vl import Lfm2VlForConditionalGeneration as HfLfm2Vl |
|
|
| from .hybrid import LanguageModel, RMSNorm, Tree, attention, causal_conv, rope |
|
|
|
|
| class Attention(nn.Module): |
| """GQA with per-head q/k norms and RoPE over the whole head.""" |
|
|
| def __init__(self, cfg): |
| super().__init__() |
| self.heads, self.kv_heads = cfg.num_attention_heads, cfg.num_key_value_heads |
| self.head_dim = getattr(cfg, "head_dim", None) or cfg.hidden_size // self.heads |
| self.q_proj = nn.Linear(cfg.hidden_size, self.heads * self.head_dim, bias=False) |
| self.k_proj = nn.Linear(cfg.hidden_size, self.kv_heads * self.head_dim, bias=False) |
| self.v_proj = nn.Linear(cfg.hidden_size, self.kv_heads * self.head_dim, bias=False) |
| self.out_proj = nn.Linear(self.heads * self.head_dim, cfg.hidden_size, bias=False) |
| self.q_layernorm = RMSNorm(self.head_dim, cfg.norm_eps) |
| self.k_layernorm = RMSNorm(self.head_dim, cfg.norm_eps) |
| self.theta = cfg.rope_parameters["rope_theta"] |
| self.scale = self.head_dim**-0.5 |
|
|
| def forward(self, x: torch.Tensor, pos: torch.Tensor, tree: Tree | None) -> torch.Tensor: |
| batch, length, _ = x.shape |
| q = self.q_layernorm(self.q_proj(x).view(batch, length, self.heads, self.head_dim)) |
| k = self.k_layernorm(self.k_proj(x).view(batch, length, self.kv_heads, self.head_dim)) |
| v = self.v_proj(x).view(batch, length, self.kv_heads, self.head_dim) |
| q, k = rope(q, pos, self.theta, self.head_dim // 2), rope(k, pos, self.theta, self.head_dim // 2) |
| y = attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), tree, None, self.scale) |
| return self.out_proj(y.transpose(1, 2).flatten(2)) |
|
|
|
|
| class ShortConv(nn.Module): |
| """`out(C * conv(B * x))`; HF's `nn.Conv1d` holds the taps, `causal_conv` runs them.""" |
|
|
| def __init__(self, cfg): |
| super().__init__() |
| dim, taps = cfg.hidden_size, cfg.conv_L_cache |
| self.conv = nn.Conv1d(dim, dim, taps, groups=dim, bias=cfg.conv_bias, padding=taps - 1) |
| self.in_proj = nn.Linear(dim, 3 * dim, bias=cfg.conv_bias) |
| self.out_proj = nn.Linear(dim, dim, bias=cfg.conv_bias) |
|
|
| def forward(self, x: torch.Tensor, pos: torch.Tensor, tree: Tree | None) -> torch.Tensor: |
| b, c, u = self.in_proj(x).transpose(1, 2).chunk(3, dim=1) |
| z = causal_conv(b * u, self.conv.weight, tree) |
| if self.conv.bias is not None: |
| z = z + self.conv.bias[:, None] |
| return self.out_proj((c * z).transpose(1, 2)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| hidden = cfg.intermediate_size |
| if cfg.block_auto_adjust_ff_dim: |
| hidden = int(2 * hidden / 3) |
| if cfg.block_ffn_dim_multiplier is not None: |
| hidden = int(cfg.block_ffn_dim_multiplier * hidden) |
| hidden = cfg.block_multiple_of * ((hidden + cfg.block_multiple_of - 1) // cfg.block_multiple_of) |
| self.w1 = nn.Linear(cfg.hidden_size, hidden, bias=False) |
| self.w3 = nn.Linear(cfg.hidden_size, hidden, bias=False) |
| self.w2 = nn.Linear(hidden, cfg.hidden_size, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.w2(F.silu(self.w1(x)) * self.w3(x)) |
|
|
|
|
| def language_model(cfg) -> LanguageModel: |
| return LanguageModel(cfg, lambda kind: ShortConv(cfg) if kind == "conv" else Attention(cfg), |
| lambda: MLP(cfg), cfg.norm_eps) |
|
|
|
|
| class Lfm2VlForConditionalGeneration(HfLfm2Vl): |
| """HF's LFM2-VL with the hybrid language model, under HF's name so a saved |
| checkpoint reads as HF's.""" |
|
|
| def __init__(self, config: Lfm2VlConfig): |
| super().__init__(config) |
| self.model.language_model = language_model(config.text_config) |
|
|
| def answer(self, trunk: torch.Tensor, questions: torch.Tensor, lengths: torch.Tensor, **vision) -> torch.Tensor: |
| """Logits at the last token of every question, the state (and its image) |
| read once: `trunk` `(1, P)` ids, `questions` their ids back to back.""" |
| embeds = self.get_input_embeddings()(trunk) |
| if vision: |
| features = torch.cat(self.model.get_image_features(**vision).pooler_output).to(embeds.dtype) |
| slots = self.model.get_placeholder_mask(trunk, inputs_embeds=embeds, image_features=features) |
| embeds = embeds.masked_scatter(slots, features) |
| return self.lm_head(self.model.language_model.answer(embeds, questions, lengths)) |
|
|
|
|
| class Lfm2ForCausalLM(HfLfm2): |
| """HF's LFM2 with the hybrid language model, under HF's name.""" |
|
|
| def __init__(self, config: Lfm2Config): |
| super().__init__(config) |
| self.model = language_model(config) |
|
|
| def answer(self, trunk: torch.Tensor, questions: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: |
| return self.lm_head(self.model.answer(self.model.embed_tokens(trunk), questions, lengths)) |
|
|
|
|
