broccoli-ml 3.3.0__tar.gz → 4.0.0__tar.gz

{broccoli_ml-3.3.0 → broccoli_ml-4.0.0}/PKG-INFO

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

{broccoli_ml-3.3.0 → broccoli_ml-4.0.0}/broccoli/transformer.py

@@ -4,11 +4,20 @@ from typing import Optional
 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
@@ -206,32 +215,53 @@ class MHAttention(nn.Module):
             q = torch.cat([q_bos, q_img], dim=1)
             k = torch.cat([k_bos, k_img], dim=1)
 
-        # 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)
-
         if self.scaling == "sqrtd":
-            qk_scores /= math.sqrt(self.head_dim)
+            scaling_factor = 1 / math.sqrt(self.head_dim)
         elif self.scaling == "d":
             # for backwards compatibility, per https://github.com/microsoft/mup
-            qk_scores *= 8 / self.head_dim
+            scaling_factor = 8 / self.head_dim
         else:
             raise ValueError('`scaling` argument to MHAttention must be "d" or "sqrtd"')
 
-        # Apply mask if causal (must come before softmax)
-        if self.causal:
-            qk_scores.masked_fill_(self.mask, float("-inf"))
+        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 if self.training else 0.0,
+                softmax_scale=scaling_factor,
+                causal=self.causal,
+            )
+
+            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 = F.softmax(qk_scores, dim=-1)
+            qk_scores *= scaling_factor
 
-        output_with_heads = qk_scores @ v
+            # Apply mask if causal (must come before softmax)
+            if self.causal:
+                qk_scores.masked_fill_(self.mask, float("-inf"))
 
-        output_without_heads = rearrange(output_with_heads, "b h t d -> b t (h d)")
+            qk_scores = F.softmax(qk_scores, dim=-1)
 
-        return self.out_proj(output_without_heads)
+            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)
 
 
 class FeedforwardBlock(nn.Module):
@@ -259,6 +289,13 @@ class FeedforwardBlock(nn.Module):
         self.residual_path = residual_path
         self.post_norm = post_norm
 
+        if self.residual_path and (output_features < input_features):
+            raise ValueError(
+                "If the number of output features will be less than "
+                "the number of input features, then `residual_path` "
+                "should be set to False."
+            )
+
         if self.post_norm:
             self.layernorm = nn.LayerNorm(output_features)
 
@@ -403,21 +440,20 @@ 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(checkpoint(self.ff, x, use_reentrant=False))
+            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 = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
             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)
+            x = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
 
         return x
 

broccoli_ml-4.0.0/broccoli/utils.py

@@ -0,0 +1,15 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PadTensor(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.args = args
+        self.kwargs = kwargs
+
+    def forward(self, x):
+        if sum(self.args[0]) == 0:
+            return x
+        else:
+            return F.pad(x, *self.args, **self.kwargs)

{broccoli_ml-3.3.0 → broccoli_ml-4.0.0}/broccoli/vit.py

@@ -4,25 +4,12 @@ from typing import Optional
 from .transformer import TransformerEncoder, FeedforwardBlock
 from .cnn import SpaceToDepth, calculate_output_spatial_size, spatial_tuple
 from .activation import ReLU, SquaredReLU, GELU, SwiGLU
+from .utils import PadTensor
 
 from einops import einsum
 from einops.layers.torch import Rearrange
 
 import torch.nn as nn
-import torch.nn.functional as F
-
-
-class PadTensor(nn.Module):
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-        self.args = args
-        self.kwargs = kwargs
-
-    def forward(self, x):
-        if sum(self.args[0]) == 0:
-            return x
-        else:
-            return F.pad(x, *self.args, **self.kwargs)
 
 
 class GetCLSToken(nn.Module):

{broccoli_ml-3.3.0 → broccoli_ml-4.0.0}/pyproject.toml

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