Add fused weights

attention.py

@@ -1,9 +1,12 @@
 import torch
 from torch import Tensor as T
 from torch.nn.attention.flex_attention import (
+    BlockMask,
     _mask_mod_signature,
+    and_masks,
     create_block_mask,
     flex_attention,
+    or_masks,
 )
 
 # ---------------------------------------------------------------------------
@@ -106,9 +109,38 @@ _compiled_create_block_mask = torch.compile(
 
 
 @torch.inference_mode()
-def create_attention_mask(*args, **kwargs):
+def create_attention_mask(*args, **kwargs) -> BlockMask:
     """
     NOTE: We compile this for performance/memory reasons in large masks. To reduce
     recompiles due to grad_mode flips, we always run mask creation under inference_mode.
     """
     return _compiled_create_block_mask(*args, **kwargs)
+
+
+def create_batch_attention_mask(
+    input_batch: T,
+    *,
+    pad_token_id: int,
+    eos_token_id: int,
+    soi_token_id: int,
+    eoi_token_id: int,
+    max_len: int | None = None,
+) -> BlockMask:
+    """Build the combined FlexAttention mask for the batch engine.
+
+    Composes causal + document + non-left-pad + image-prefix masks.
+    """
+    B, S = input_batch.size()
+    block_causal_mask_mod = and_masks(
+        get_causal_mask_mod(),
+        get_document_mask_mod(input_batch, eos_token_id),
+        get_non_left_pad_mask_mod(input_batch, pad_token_id),
+    )
+    image_prefix_mask_mod = get_image_prefix_mask_mod(
+        batch=input_batch,
+        soi_id=soi_token_id,
+        eoi_id=eoi_token_id,
+    )
+    mask_mod = or_masks(image_prefix_mask_mod, block_causal_mask_mod)
+    max_len = max_len or S
+    return create_attention_mask(mask_mod, B, None, max_len, max_len)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d842c991349d997852c99ebf0dc6d368fc70b73c658d117a3088132a6cbb68ca
-size 2529523048
+oid sha256:c680d5b1a834a2df61baae9dd694b95856872609d0a769702faf5ad658297641
+size 2529514176

modeling_falcon_perception.py

@@ -1,5 +1,4 @@
 import math
-import time
 from pathlib import Path
 
 import einops as E
@@ -15,8 +14,6 @@ from torch import nn
 from torch.nn.attention.flex_attention import (
     AuxRequest,
     BlockMask,
-    and_masks,
-    or_masks,
 )
 from transformers import AutoTokenizer, PreTrainedModel
 
@@ -25,10 +22,7 @@ from .attention import (
     compiled_flex_attn_decode,
     compiled_flex_attn_prefill,
     create_attention_mask,
-    get_causal_mask_mod,
-    get_document_mask_mod,
-    get_image_prefix_mask_mod,
-    get_non_left_pad_mask_mod,
+    create_batch_attention_mask,
     offset_mask_mod,
 )
 from .configuration_falcon_perception import FalconPerceptionConfig
@@ -99,19 +93,12 @@ class Attention(nn.Module):
         self.q_dim = config.n_heads * self.head_dim
         self.kv_dim = self.n_kv_heads * self.head_dim
 
-        self.wq = nn.Linear(config.dim, self.q_dim, bias=False)
-        self.wk = nn.Linear(config.dim, self.kv_dim, bias=False)
-        self.wv = nn.Linear(config.dim, self.kv_dim, bias=False)
+        self.wqkv = nn.Linear(config.dim, self.q_dim + 2 * self.kv_dim, bias=False)
         self.wo = nn.Linear(config.n_heads * self.head_dim, config.dim, bias=False)
         self.sinks = nn.Parameter(torch.empty((config.n_heads,)))
 
-    def _fuse_weights(self):
-        wqkv_weight = torch.cat([self.wq.weight.data, self.wk.weight.data, self.wv.weight.data], dim=0)
-        self.register_buffer("_wqkv_weight", wqkv_weight)
-        del self.wq, self.wk, self.wv
-
     def _pre_attention_qkv(self, x) -> tuple[T, T, T]:
-        qkv = F.linear(F.rms_norm(x, (x.size(-1),)), self._wqkv_weight)
+        qkv = self.wqkv(F.rms_norm(x, (x.size(-1),)))
         xq, xk, xv = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1)
         xq = E.rearrange(xq, "b s (h d) -> b s h d", d=self.head_dim)
         xk = E.rearrange(xk, "b s (h d) -> b s h d", d=self.head_dim)
@@ -195,27 +182,13 @@ def squared_relu_gate(packed: T, hidden_dim: int) -> T:
 class FeedForward(nn.Module):
     def __init__(self, dim: int, hidden_dim: int):
         super().__init__()
-        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w13 = nn.Linear(dim, 2 * hidden_dim, bias=False)
         self.w2 = nn.Linear(hidden_dim, dim, bias=False)
-        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
         self.hidden_dim = hidden_dim
 
-    def _fuse_weights(self):
-        if hasattr(self, "_w13_weight"):
-            return
-        w1_weight_fused = self.w1.weight.data * math.sqrt(2.0)
-        w13_weight = torch.empty(
-            (2 * self.hidden_dim, self.w1.weight.shape[1]),
-            device=w1_weight_fused.device, dtype=w1_weight_fused.dtype,
-        )
-        w13_weight[0::2] = w1_weight_fused
-        w13_weight[1::2] = self.w3.weight.data
-        self.register_buffer("_w13_weight", w13_weight)
-        del self.w1, self.w3
-
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = F.rms_norm(x, (x.size(-1),))
-        w13_out = F.linear(x, self._w13_weight)
+        w13_out = self.w13(x)
         return self.w2(squared_relu_gate(w13_out, self.hidden_dim))
 
 
@@ -357,31 +330,17 @@ class FalconPerceptionForSegmentation(PreTrainedModel):
 
     # -- Weight management ---------------------------------------------------
 
-    def _fuse_weights(self):
+    def _ensure_device_buffers(self):
+        """Recompute non-persistent buffers that HF meta-device loading may discard."""
         if self._weights_fused:
             return
-
         device = self.tok_embeddings.weight.device
         c = self.config
-
-        # Recompute freqs_cis on the actual device — non-persistent buffers
-        # get replaced with empty tensors by transformers' meta-device loading.
         rope_dim = c.head_dim // 2
         freqs_cis = precompute_freqs_cis(rope_dim, c.max_seq_len, c.rope_theta).to(device)
         self.register_buffer("freqs_cis", freqs_cis, persistent=False)
-
-        # Ensure freqs_cis_golden is on the right device (loaded from safetensors)
         if self.freqs_cis_golden.device != device:
             self.freqs_cis_golden = self.freqs_cis_golden.to(device)
-
-        for layer in self.layers.values():
-            layer.attention._fuse_weights()
-            layer.feed_forward._fuse_weights()
-        self.coord_decoder.w1.weight.mul_(math.sqrt(2))
-        self.size_decoder.w1.weight.mul_(math.sqrt(2))
-        if self.config.do_segmentation:
-            for layer in self.proj_segm.layers:
-                layer.weight.mul_(math.sqrt(2))
         self._weights_fused = True
 
     def compile_model(self):
@@ -418,19 +377,14 @@ class FalconPerceptionForSegmentation(PreTrainedModel):
     # -- Attention mask ------------------------------------------------------
 
     def get_attention_mask(self, input_batch: T, max_len: int | None = None):
-        B, S = input_batch.size()
-        c = self.config
-        block_causal_mask_mod = and_masks(
-            get_causal_mask_mod(),
-            get_document_mask_mod(input_batch, c.eos_id),
-            get_non_left_pad_mask_mod(input_batch, self._pad_token_id),
-        )
-        image_prefix_mask_mod = get_image_prefix_mask_mod(
-            batch=input_batch, soi_id=c.image_cls_token_id, eoi_id=c.img_end_id,
+        return create_batch_attention_mask(
+            input_batch,
+            pad_token_id=self._pad_token_id,
+            eos_token_id=self.config.eos_id,
+            soi_token_id=self.config.image_cls_token_id,
+            eoi_token_id=self.config.img_end_id,
+            max_len=max_len,
         )
-        mask_mod = or_masks(image_prefix_mask_mod, block_causal_mask_mod)
-        max_len = max_len or S
-        return create_attention_mask(mask_mod, B, None, max_len, max_len)
 
     def get_upsampler_attn_mask(self, H, W, h, w, device):
         return create_attention_mask(
@@ -699,7 +653,7 @@ class FalconPerceptionForSegmentation(PreTrainedModel):
                     "mask_rle": {"counts": str, "size": [H, W]},
                 }
         """
-        self._fuse_weights()
+        self._ensure_device_buffers()
         if compile:
             self.compile_model()