broccoli-ml 3.3.1__tar.gz → 5.1.0__tar.gz

{broccoli_ml-3.3.1 → broccoli_ml-5.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 3.3.1
+Version: 5.1.0
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-3.3.1 → broccoli_ml-5.1.0}/broccoli/transformer.py

@@ -1,14 +1,23 @@
 import math
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
 
 from einops import rearrange
 
 from .rope import RotaryEmbedding, apply_rotary_emb
 
+try:
+    from flash_attn import flash_attn_func
+
+    FLASH_ATTN = True
+except ImportError:
+    pass
+    FLASH_ATTN = False
+
 
 def drop_path(
     x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
@@ -93,6 +102,15 @@ class MHAttention(nn.Module):
 
         self.head_dim = self.embed_dim // self.n_heads
 
+        if self.scaling == "sqrtd":
+            self.scaling_factor = 1 / math.sqrt(self.head_dim)
+        elif self.scaling == "d":
+            # 8/d_model for backwards compatibility,
+            #     per https://github.com/microsoft/mup
+            self.scaling_factor = 8 / self.head_dim
+        else:
+            raise ValueError('`scaling` argument to MHAttention must be "d" or "sqrtd"')
+
         self.q_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
         self.k_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
         self.v_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
@@ -113,6 +131,8 @@ class MHAttention(nn.Module):
         self.source_size = source_size
         self.bos_tokens = bos_tokens
 
+        self.reset_parameters()
+
     @property
     def _kv_distance(self) -> float:
         """
@@ -132,7 +152,71 @@ class MHAttention(nn.Module):
 
         return 1 - similarity
 
-    def forward(self, q, k, v):
+    def add_axial_rope(
+        self, q: torch.Tensor, k: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply Axial RoPE to all tokens except BOS tokens
+        """
+
+        if len(self.source_size) == 1:
+            spatial_dimension_names = "D1"
+            spatial_dimension_values = {"D1": self.source_size[0]}
+        elif len(self.source_size) == 2:
+            spatial_dimension_names = "D1 D2"
+            spatial_dimension_values = {
+                "D1": self.source_size[0],
+                "D2": self.source_size[1],
+            }
+        elif len(self.source_size) == 3:
+            spatial_dimension_names = "D1 D2 D3"
+            spatial_dimension_values = {
+                "D1": self.source_size[0],
+                "D2": self.source_size[1],
+                "D3": self.source_size[2],
+            }
+        else:
+            raise NotImplementedError(
+                "`source_size` must be a tuple of 1, 2 or 3 integers"
+            )
+
+        q_bos, q_img = q[:, : self.bos_tokens, :], q[:, self.bos_tokens :, :]
+        k_bos, k_img = k[:, : self.bos_tokens, :], k[:, self.bos_tokens :, :]
+
+        q_img = rearrange(
+            q_img,
+            f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
+            **spatial_dimension_values,
+        )
+        k_img = rearrange(
+            k_img,
+            f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
+            **spatial_dimension_values,
+        )
+
+        freqs = self.rotary_embedding.get_axial_freqs(*self.source_size)
+
+        q_img = apply_rotary_emb(freqs, q_img)
+        k_img = apply_rotary_emb(freqs, k_img)
+
+        q_img = rearrange(
+            q_img,
+            f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
+        )
+        k_img = rearrange(
+            k_img,
+            f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
+        )
+
+        # Re-combine the BOS tokens and the RoPE-enhanced image tokens
+        q = torch.cat([q_bos, q_img], dim=1)
+        k = torch.cat([k_bos, k_img], dim=1)
+
+        return q, k
+
+    def project_qkv(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         query_batch_size, query_tokens, query_features = q.size()
         key_batch_size, key_tokens, key_features = k.size()
 
@@ -145,66 +229,64 @@ class MHAttention(nn.Module):
 
         if self.causal:
             assert query_tokens == key_tokens
-            assert query_tokens == self.sequence_length
+            assert query_tokens == self.seq_len
 
-        # Project q, k and v
-        q = self.q_proj(q)
-        k = self.k_proj(k)
-        v = self.v_proj(v)
+        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
 
-        # Rearrange dimensions and add RoPE if needed
         if self.rotary_embedding is not None:
+            q, k = self.add_axial_rope(q, k)
 
-            if len(self.source_size) == 1:
-                spatial_dimension_names = "D1"
-                spatial_dimension_values = {"D1": self.source_size[0]}
-            elif len(self.source_size) == 2:
-                spatial_dimension_names = "D1 D2"
-                spatial_dimension_values = {
-                    "D1": self.source_size[0],
-                    "D2": self.source_size[1],
-                }
-            elif len(self.source_size) == 3:
-                spatial_dimension_names = "D1 D2 D3"
-                spatial_dimension_values = {
-                    "D1": self.source_size[0],
-                    "D2": self.source_size[1],
-                    "D3": self.source_size[2],
-                }
-            else:
-                raise NotImplementedError(
-                    "`source_size` must be a tuple of 1, 2 or 3 integers"
-                )
+        return q, k, v
 
-            q_bos, q_img = q[:, : self.bos_tokens, :], q[:, self.bos_tokens :, :]
-            k_bos, k_img = k[:, : self.bos_tokens, :], k[:, self.bos_tokens :, :]
+    def forward(self, q, k, v):
 
-            q_img = rearrange(
-                q_img,
-                f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
-                **spatial_dimension_values,
-            )
-            k_img = rearrange(
-                k_img,
-                f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
-                **spatial_dimension_values,
+        q, k, v = self.project_qkv(q, k, v)
+
+        if FLASH_ATTN:
+            # Divide Q/K/V into heads
+            q = rearrange(q, "b t (h d) -> b t h d", h=self.n_heads)
+            k = rearrange(k, "b t (h d) -> b t h d", h=self.n_heads)
+            v = rearrange(v, "b t (h d) -> b t h d", h=self.n_heads)
+
+            output_with_heads = flash_attn_func(
+                q,
+                k,
+                v,
+                dropout_p=self.dropout.p if self.training else 0.0,
+                softmax_scale=self.scaling_factor,
+                causal=self.causal,
             )
-            freqs = self.rotary_embedding.get_axial_freqs(*self.source_size)
-            q_img = apply_rotary_emb(freqs, q_img)
-            k_img = apply_rotary_emb(freqs, k_img)
 
-            q_img = rearrange(
-                q_img,
-                f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
-            )
-            k_img = rearrange(
-                k_img,
-                f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
-            )
+            output_without_heads = rearrange(output_with_heads, "b t h d -> b t (h d)")
+
+            return self.out_proj(output_without_heads)
+        else:
+            # Divide Q/K/V into heads
+            q = rearrange(q, "b t (h d) -> b h t d", h=self.n_heads)
+            k = rearrange(k, "b t (h d) -> b h t d", h=self.n_heads)
+            v = rearrange(v, "b t (h d) -> b h t d", h=self.n_heads)
+
+            qk_scores = q @ k.transpose(-1, -2)
+
+            qk_scores *= self.scaling_factor
+
+            # Apply mask if causal (must come before softmax)
+            if self.causal:
+                qk_scores.masked_fill_(self.mask, float("-inf"))
+
+            qk_scores = F.softmax(qk_scores, dim=-1)
+
+            qk_scores = self.dropout(qk_scores)
 
-            # Re-combine the BOS tokens and the RoPE-enhanced image tokens
-            q = torch.cat([q_bos, q_img], dim=1)
-            k = torch.cat([k_bos, k_img], dim=1)
+            output_with_heads = qk_scores @ v
+
+            output_without_heads = rearrange(output_with_heads, "b h t d -> b t (h d)")
+
+            return self.out_proj(output_without_heads)
+
+    def attention_scores(self, q, k, v):
+
+        q, k, v = self.project_qkv(q, k, v)
 
         # Divide Q/K/V into heads
         q = rearrange(q, "b t (h d) -> b h t d", h=self.n_heads)
@@ -213,13 +295,7 @@ class MHAttention(nn.Module):
 
         qk_scores = q @ k.transpose(-1, -2)
 
-        if self.scaling == "sqrtd":
-            qk_scores /= math.sqrt(self.head_dim)
-        elif self.scaling == "d":
-            # for backwards compatibility, per https://github.com/microsoft/mup
-            qk_scores *= 8 / self.head_dim
-        else:
-            raise ValueError('`scaling` argument to MHAttention must be "d" or "sqrtd"')
+        qk_scores *= self.scaling_factor
 
         # Apply mask if causal (must come before softmax)
         if self.causal:
@@ -227,11 +303,14 @@ class MHAttention(nn.Module):
 
         qk_scores = F.softmax(qk_scores, dim=-1)
 
-        output_with_heads = qk_scores @ v
+        return qk_scores  # (batch, head, seq_len, seq_len)
 
-        output_without_heads = rearrange(output_with_heads, "b h t d -> b t (h d)")
-
-        return self.out_proj(output_without_heads)
+    def reset_parameters(self):
+        # Default nn.Linear init is kaiming_uniform, which is fine
+        self.q_proj.reset_parameters()
+        self.k_proj.reset_parameters()
+        self.v_proj.reset_parameters()
+        self.out_proj.reset_parameters()
 
 
 class FeedforwardBlock(nn.Module):
@@ -253,9 +332,11 @@ class FeedforwardBlock(nn.Module):
         normformer=False,
         post_norm=True,
         residual_path=True,
+        checkpoint=True,
     ):
         super().__init__()
 
+        self.checkpoint = checkpoint
         self.residual_path = residual_path
         self.post_norm = post_norm
 
@@ -293,13 +374,30 @@ class FeedforwardBlock(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
+
+        if self.checkpoint:
+            processed = checkpoint(self.process, x, use_reentrant=False)
+        else:
+            processed = self.process(x)
+
         if self.residual_path and self.post_norm:
-            return self.layernorm(x + self.process(x))
+            return self.layernorm(x + processed)
         elif self.residual_path:
-            return x + self.process(x)
+            return x + processed
         else:
-            return self.process(x)
+            return processed
+
+    def reset_parameters(self):
+        if self.post_norm:
+            self.layernorm.reset_parameters()
+
+        # Iterate over the sequential block to reset parameters
+        for module in self.process:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
 
 
 class TransformerBlock(nn.Module):
@@ -333,6 +431,7 @@ class TransformerBlock(nn.Module):
         pre_norm=True,
         post_norm=False,
         normformer=False,
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -401,8 +500,11 @@ class TransformerBlock(nn.Module):
             normformer=normformer,
             post_norm=False,  # Handled outside the block
             residual_path=False,  # Handled outside the block
+            checkpoint=checkpoint_ff,
         )
 
+        self.reset_parameters()
+
     @property
     def _kv_distance(self) -> float:
         return self.attn._kv_distance
@@ -410,29 +512,46 @@ class TransformerBlock(nn.Module):
     def forward(self, x):
 
         if self.pre_norm:
-            normx = self.layer_norm_1(x)
-            x = x + self.drop_path(self.attn(normx, normx, normx))
-            normx = self.layer_norm_2(x)
-            x = x + self.drop_path(self.ff(normx))
-        elif self.post_norm:
+            x = self.layer_norm_1(x)
+            x = x + self.drop_path(self.attn(x, x, x))
+            x = self.layer_norm_2(x)
+            x = x + self.drop_path(self.ff(x))
+            if self.post_norm:  # i.e. in addition! Pre and post.
+                x = self.layer_norm_3(x)
+        elif self.post_norm:  # i.e. only, not prenorm, just post
             x = x + self.drop_path(self.attn(x, x, x))
             x = self.layer_norm_1(x)
             x = x + self.drop_path(self.ff(x))
             x = self.layer_norm_2(x)
-        else:
+        else:  # Not pre or post norm. Stand well back.
             x = x + self.drop_path(self.attn(x, x, x))
             x = x + self.drop_path(self.ff(x))
 
-        if self.pre_norm and self.post_norm:
-            x = self.layer_norm_3(x)
-
         return x
 
+    def attention_scores(self, x):
+        """
+        Give back the attention scores used in this layer.
+        """
+        if self.pre_norm:
+            x = self.layer_norm_1(x)
+            return self.attn(x, x, x)
+        else:
+            return self.attn(x, x, x)
+
+    def reset_parameters(self):
+        self.layer_norm_1.reset_parameters()
+        self.layer_norm_2.reset_parameters()
+        self.layer_norm_3.reset_parameters()
+
+        self.attn.reset_parameters()
+        self.ff.reset_parameters()
+
 
 class TransformerEncoder(nn.Module):
     """
     This assumes we already get a sequence of embeddings (e.g. word or image
-        patch embeddings). It uses learned positional embeddings.
+        patch embeddings).
     """
 
     def __init__(
@@ -460,6 +579,7 @@ class TransformerEncoder(nn.Module):
         post_norm=False,
         normformer=False,
         msa_scaling="d",
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -536,16 +656,19 @@ class TransformerEncoder(nn.Module):
                     pre_norm=pre_norm,
                     post_norm=post_norm,
                     normformer=normformer,
+                    checkpoint_ff=checkpoint_ff,
                 )
                 for i in range(n_layers)
             ]
         )
 
+        self.reset_parameters()
+
     @property
     def _kv_distances(self) -> float:
         return ",".join([str(block._kv_distance) for block in self.blocks])
 
-    def forward(self, x):
+    def preprocess(self, x):
         if self._bos_tokens:
             x = torch.cat([self._bos_embedding.expand(x.size(0), -1, -1), x], dim=1)
         else:
@@ -560,6 +683,10 @@ class TransformerEncoder(nn.Module):
                 )  # to shape (1, seq_len) to broadcast over batch
             )
 
+    def forward(self, x):
+
+        x = self.preprocess(x)
+
         for block in self.blocks:
             x = block(x)
 
@@ -567,3 +694,27 @@ class TransformerEncoder(nn.Module):
             return x[:, self._bos_tokens :, :]
         else:
             return x
+
+    def attention_scores(self, x):
+
+        x = self.preprocess(x)
+
+        layer_scores = []
+
+        for block in self.blocks:
+            # Get attention scores with shape (batch, 1, head, seq_len, seq_len)
+            layer_attention_scores = block.attention_scores(x).unsqueeze(1)
+            layer_scores.append(layer_attention_scores)
+            x = block(x)
+
+        return torch.cat(layer_scores, dim=1)  # (batch, layer, head, seq_len, seq_len)
+
+    def reset_parameters(self):
+        if self._bos_embedding is not None:
+            nn.init.normal_(self._bos_embedding, mean=0.0, std=1.0)
+
+        if self.absolute_position_embedding is not None:
+            self.absolute_position_embedding.reset_parameters()
+
+        for block in self.blocks:
+            block.reset_parameters()

{broccoli_ml-3.3.1 → broccoli_ml-5.1.0}/broccoli/vit.py

@@ -9,7 +9,9 @@ from .utils import PadTensor
 from einops import einsum
 from einops.layers.torch import Rearrange
 
+import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 
 class GetCLSToken(nn.Module):
@@ -31,10 +33,18 @@ class SequencePool(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         weights = self.attention(x)
         return einsum(weights, x, "batch seq, batch seq d_model -> batch d_model")
 
+    def reset_parameters(self):
+        # Iterate over modules in the sequential block
+        for module in self.attention:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class ClassificationHead(nn.Module):
     """
@@ -71,9 +81,16 @@ class ClassificationHead(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         return self.classification_process(x)
 
+    def reset_parameters(self):
+        for module in self.classification_process:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class SequencePoolClassificationHead(ClassificationHead):
     """
@@ -106,6 +123,8 @@ class SequencePoolClassificationHead(ClassificationHead):
             ]
         )
 
+        self.reset_parameters()
+
 
 class ViTEncoder(nn.Module):
     """
@@ -160,6 +179,7 @@ class ViTEncoder(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         linear_module=nn.Linear,
     ):
         super().__init__()
@@ -321,6 +341,7 @@ class ViTEncoder(nn.Module):
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
+                checkpoint_ff=transformer_checkpoint_ff,
             )
         else:
             self.transformer = nn.Identity()
@@ -354,6 +375,7 @@ class ViTEncoder(nn.Module):
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
                 residual_path=transformer_initial_ff_residual_path,
+                checkpoint=transformer_checkpoint_ff,
             )
         else:
             self.initial_ff = nn.Identity()
@@ -373,9 +395,20 @@ class ViTEncoder(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         return self.encoder(x)
 
+    def attention_scores(self, x):
+        x = self.encoder[:-1](x)
+        return self.encoder[-1].attention_scores(x)
+
+    def reset_parameters(self):
+        for module in self.encoder:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class ViT(nn.Module):
     """
@@ -426,6 +459,7 @@ class ViT(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         head=SequencePoolClassificationHead,
         batch_norm_logits=True,
         logit_projection_layer=nn.Linear,
@@ -492,6 +526,7 @@ class ViT(nn.Module):
             transformer_mlp_dropout=transformer_mlp_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
+            transformer_checkpoint_ff=transformer_checkpoint_ff,
             linear_module=linear_module,
         )
 
@@ -502,9 +537,26 @@ class ViT(nn.Module):
             batch_norm_logits=batch_norm_logits,
         )
 
+        self.reset_parameters()
+
     @property
     def sequence_length(self):
         return self.encoder.sequence_length
 
     def forward(self, x):
         return self.pool(self.encoder(x))
+
+    def attention_scores(self, x):
+        return self.encoder.attention_scores(x)
+
+    def head_to_bos_token_attention(self, x):
+        all_attention = self.attention_scores(x)
+        batch_averages = torch.mean(all_attention, dim=0, keepdim=False)
+        sequence_averages = torch.mean(batch_averages, dim=-1, keepdim=False)
+        n_bos_tokens = self.encoder.encoder._bos_tokens
+        just_bos = sequence_averages[:, :, :n_bos_tokens]
+        return F.softmax(just_bos, dim=-1)  # (layer, head, bos_token)
+
+    def reset_parameters(self):
+        self.encoder.reset_parameters()
+        self.pool.reset_parameters()

{broccoli_ml-3.3.1 → broccoli_ml-5.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "3.3.1"
+version = "5.1.0"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}