broccoli-ml 6.0.0__py3-none-any.whl → 9.2.2__py3-none-any.whl

broccoli/linear.py

@@ -1,4 +1,8 @@
 import math
+import random
+import warnings
+from typing import Union, List, Iterable
+
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -136,3 +140,178 @@ class WeightNormedLinear(nn.Module):
             f"WeightNormedLinear(in_features={self.in_features},"
             f"out_features={self.out_features}, bias={self.use_bias})"
         )
+
+
+class RecyclingLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        row_recycling_rate: float = 0.0,
+        column_recycling_rate: float = 0.0,
+        adaptive=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.bias = bias
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        self.row_recycling_rate = row_recycling_rate
+        self.column_recycling_rate = column_recycling_rate
+        self.adaptive = adaptive
+        self.optimisers = []
+        self.initial_learning_rates = []
+        self._warned_about_registration = False
+
+    def register_optimiser(self, optimiser: torch.optim.Optimizer):
+        self.optimisers.append(optimiser)
+        self.initial_learning_rates.append(self._get_learning_rate(optimiser))
+        if self.initial_learning_rates[-1] == 0.0:
+            warnings.warn(
+                "Learning rate of registered optimiser was 0.0 - make sure "
+                "you haven't initialised a scheduler before registering the "
+                "optimiser",
+                stacklevel=2,
+            )
+
+    def _get_learning_rate(self, optimiser: torch.optim.Optimizer):
+        for group in optimiser.param_groups:
+            for param in group["params"]:
+                if param is self.linear.weight:
+                    return group["lr"]
+
+    def _get_multiplier(self):
+        if not self.adaptive or not self.optimisers:
+            return 1.0
+        else:
+            init = self.initial_learning_rates
+            current = [self._get_learning_rate(o) for o in self.optimisers]
+            pairs = zip(current, init, strict=True)
+            multipliers = [a / b for a, b in pairs if b != 0.0]
+            return min(multipliers) if multipliers else 0.0
+
+    def forward(self, x):
+        multiplier = self._get_multiplier()
+        col_recycling_rate = self.column_recycling_rate * multiplier
+        row_recycling_rate = self.row_recycling_rate * multiplier
+
+        if self.training and self.optimisers:
+
+            if row_recycling_rate > 0:
+                probs = torch.rand(self.linear.out_features, device=x.device)
+                mask = probs < row_recycling_rate
+                if mask.any():
+                    # nonzero returns [N, 1], squeeze to get [N]
+                    indices = torch.nonzero(mask).squeeze(-1)
+                    self.reset_rows(indices, self.optimisers)
+
+            if col_recycling_rate > 0:
+                probs = torch.rand(self.linear.in_features, device=x.device)
+                mask = probs < col_recycling_rate
+                if mask.any():
+                    indices = torch.nonzero(mask).squeeze(-1)
+                    self.reset_columns(indices, self.optimisers)
+
+        elif self.training and not self._warned_about_registration:
+            warnings.warn(
+                "RecyclingLinear: No optimiser registered. Recycling disabled.",
+                stacklevel=2,
+            )
+            self._warned_about_registration = True
+
+        return self.linear(x)
+
+    def reset_rows(
+        self,
+        indices: Iterable[int],
+        optimisers: Union[
+            List[torch.optim.Optimizer], torch.optim.Optimizer, None
+        ] = None,
+    ):
+        """
+        Update some of the weight rows to be equal to the mean of all weight rows.
+        """
+        if optimisers is None:
+            optimisers = []
+        if not isinstance(optimisers, list):
+            optimisers = [optimisers]
+
+        device = self.linear.weight.device
+        idx_tensor = torch.as_tensor(list(indices), dtype=torch.long, device=device)
+
+        if idx_tensor.numel() == 0:
+            return
+
+        with torch.no_grad():
+            # Calculate mean of all rows including the rows to be reset
+            mean_vector = self.linear.weight.data.mean(
+                dim=0, keepdim=True
+            )  # [1, in_features]
+            update_data = mean_vector.expand(idx_tensor.size(0), -1)
+            self.linear.weight.data[idx_tensor] = update_data
+
+            if self.linear.bias is not None:
+                self.linear.bias.data[idx_tensor] = 0.0
+
+            self._reset_optim_state(self.linear.weight, idx_tensor, optimisers, dim=0)
+            if self.linear.bias is not None:
+                self._reset_optim_state(self.linear.bias, idx_tensor, optimisers, dim=0)
+
+    def reset_columns(
+        self,
+        indices: Iterable[int],
+        optimisers: Union[
+            List[torch.optim.Optimizer], torch.optim.Optimizer, None
+        ] = None,
+    ):
+        """
+        Update some of the weight columns to be random as though reinitialised.
+        """
+        if optimisers is None:
+            optimisers = []
+        if not isinstance(optimisers, list):
+            optimisers = [optimisers]
+
+        device = self.linear.weight.device
+        idx_tensor = torch.as_tensor(list(indices), dtype=torch.long, device=device)
+
+        if idx_tensor.numel() == 0:
+            return
+
+        with torch.no_grad():
+            # 1. Generate Random Columns
+            # Shape: [out_features, N_indices]
+            weights = self.linear.weight.data
+            stdv = 1.0 / math.sqrt(weights.size(1))
+
+            # Generate [Rows, N] block
+            random_weights = torch.rand(
+                weights.size(0), idx_tensor.size(0), device=device
+            )
+            random_weights = (random_weights - 0.5) * 2.0 * stdv
+
+            # 2. Update Weights (One-shot)
+            # We assign into the columns specified by idx_tensor
+            self.linear.weight.data[:, idx_tensor] = random_weights
+
+            # 3. Update Optimizers
+            # Bias is untouched by column resets (bias is shape [Out], cols are [In])
+            self._reset_optim_state(self.linear.weight, idx_tensor, optimisers, dim=1)
+
+    def _reset_optim_state(self, param, idx_tensor, optimisers, dim):
+        """
+        Zeroes out the optimizer state for the given indices in a single operation.
+        """
+        for optimiser in optimisers:
+            if param not in optimiser.state:
+                continue
+            state = optimiser.state[param]
+
+            for _, buffer in state.items():
+                if torch.is_tensor(buffer) and buffer.shape == param.shape:
+                    # Vectorized zeroing
+                    if dim == 0:
+                        buffer[idx_tensor] = 0.0
+                    else:
+                        buffer[:, idx_tensor] = 0.0

broccoli/transformer.py

@@ -13,6 +13,7 @@ from .rope import RotaryEmbedding, apply_rotary_emb
 try:
     from flash_attn import flash_attn_func
 
+    print("Using flash-attn.")
     FLASH_ATTN = True
 except ImportError:
     pass
@@ -324,6 +325,8 @@ class FeedforwardBlock(nn.Module):
         activation=nn.ReLU,
         activation_kwargs=None,
         dropout=0.0,
+        inner_dropout=None,
+        outer_dropout=None,
         linear_module_up=nn.Linear,
         linear_module_down=nn.Linear,
         pre_norm=True,
@@ -337,6 +340,7 @@ class FeedforwardBlock(nn.Module):
         self.checkpoint = checkpoint
         self.residual_path = residual_path
         self.post_norm = post_norm
+        self.xglu = activation.__name__.endswith("GLU")
 
         if self.residual_path and (output_features < input_features):
             raise ValueError(
@@ -353,29 +357,75 @@ class FeedforwardBlock(nn.Module):
         else:
             self.activation = activation()
 
-        self.dropout = nn.Dropout(dropout)
+        self.inner_dropout = nn.Dropout(
+            inner_dropout if inner_dropout is not None else dropout
+        )
+        self.outer_dropout = nn.Dropout(
+            outer_dropout if outer_dropout is not None else dropout
+        )
 
         self.max_features = (
-            2 * ratio * output_features
-            if activation.__name__.endswith("GLU")
-            else ratio * output_features
+            2 * ratio * output_features if self.xglu else ratio * output_features
         )
 
+        self.linear_in = linear_module_up(input_features, self.max_features)
+        self.linear_out = linear_module_down(ratio * output_features, output_features)
+
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
-                linear_module_up(input_features, self.max_features),
+                self.linear_in,
                 self.activation,
+                self.inner_dropout,
                 nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
-                linear_module_down(ratio * output_features, output_features),
-                self.dropout,
+                self.linear_out,
+                self.outer_dropout,
             ]
         )
 
+        self.recycling_enabled = False
+        if hasattr(self.linear_in, "row_recycling_rate") and hasattr(
+            self.linear_out, "column_recycling_rate"
+        ):
+            self.recycling_enabled = True
+            self.master_recycling_rate = self.linear_in.row_recycling_rate
+            self.linear_in.row_recycling_rate = 0.0
+            self.linear_out.column_recycling_rate = 0.0
+            if (
+                hasattr(self.linear_in, "column_recycling_rate")
+                and self.linear_in.column_recycling_rate > 0
+            ) or (
+                hasattr(self.linear_out, "row_recycling_rate")
+                and self.linear_out.row_recycling_rate > 0
+            ):
+                raise NotImplementedError(
+                    "At the moment this layer can only support recycling linear "
+                    "layers if the in layer resets only rows and the out layer "
+                    "resets only columns."
+                )
+
         self.reset_parameters()
 
     def forward(self, x):
 
+        # Recycle weights if using recycling linear layers
+        if self.training and self.recycling_enabled:
+            multiplier = self.linear_in._get_multiplier()
+            rate = self.master_recycling_rate * multiplier
+            if rate > 0:
+                probs = torch.rand(self.linear_out.in_features, device=x.device)
+                mask = probs < rate
+                if mask.any():
+                    indices = torch.nonzero(mask).squeeze(-1)
+                    self.linear_out.reset_columns(indices, self.linear_out.optimisers)
+                    if self.xglu:
+                        indices_in = torch.cat(
+                            [indices, indices + self.linear_out.in_features]
+                        )
+                        self.linear_in.reset_rows(indices_in, self.linear_in.optimisers)
+                    else:
+                        self.linear_in.reset_rows(indices, self.linear_in.optimisers)
+
         if self.checkpoint:
             processed = checkpoint(self.process, x, use_reentrant=False)
         else:
@@ -421,7 +471,9 @@ class TransformerBlock(nn.Module):
         ff_linear_module_up=None,
         ff_linear_module_down=None,
         msa_scaling="d",
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         identity_probability=0.0,
         causal=False,
@@ -483,7 +535,9 @@ class TransformerBlock(nn.Module):
             d_model,
             activation=activation,
             activation_kwargs=activation_kwargs,
-            dropout=mlp_dropout,
+            dropout=ff_dropout,
+            inner_dropout=ff_inner_dropout,
+            outer_dropout=ff_outer_dropout,
             linear_module_up=(
                 ff_linear_module_up
                 if ff_linear_module_up is not None
@@ -566,7 +620,9 @@ class TransformerEncoder(nn.Module):
         activation_kwargs: Optional[dict] = None,
         ff_linear_module_up=None,
         ff_linear_module_down=None,
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         stochastic_depth=0.0,
         causal=False,
@@ -590,7 +646,13 @@ class TransformerEncoder(nn.Module):
         if relative_position_embedding and (source_size is None):
             raise ValueError(
                 "`source_size` for TransformerEncoder cannot be None if"
-                " `position_embedding_type` is relative"
+                " `relative_position_embedding` is True"
+            )
+
+        if absolute_position_embedding and (seq_len is None):
+            raise ValueError(
+                "`seq_len` for TransformerEncoder cannot be None if"
+                " `absolute_position_embedding` is True"
             )
 
         super().__init__()
@@ -605,9 +667,12 @@ class TransformerEncoder(nn.Module):
                 torch.empty(self._utility_tokens, d_model)
             )
             nn.init.normal_(self._utility_token_embedding, mean=0.0, std=1.0)
-            self.full_sequence_length = self.seq_len + self._utility_tokens
         else:
             self._utility_token_embedding = None
+
+        if self._utility_tokens and (self.seq_len is not None):
+            self.full_sequence_length = self.seq_len + self._utility_tokens
+        else:
             self.full_sequence_length = self.seq_len
 
         self.d_model = d_model
@@ -619,7 +684,7 @@ class TransformerEncoder(nn.Module):
         else:
             self.absolute_position_embedding = None
 
-        self.mlp_dropout = mlp_dropout
+        self.mlp_dropout = ff_dropout
         self.msa_dropout = msa_dropout
         self.stochastic_depth = stochastic_depth
 
@@ -648,7 +713,9 @@ class TransformerEncoder(nn.Module):
                     ff_linear_module_up=ff_linear_module_up,
                     ff_linear_module_down=ff_linear_module_down,
                     msa_scaling=msa_scaling,
-                    mlp_dropout=mlp_dropout,
+                    ff_dropout=ff_dropout,
+                    ff_inner_dropout=ff_inner_dropout,
+                    ff_outer_dropout=ff_outer_dropout,
                     msa_dropout=msa_dropout,
                     identity_probability=self.stochastic_depth_probabilities[i],
                     causal=causal,

broccoli/vit.py

@@ -161,7 +161,9 @@ class ViTEncoder(nn.Module):
         transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
-        transformer_initial_ff_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -178,7 +180,9 @@ class ViTEncoder(nn.Module):
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
@@ -333,7 +337,9 @@ class ViTEncoder(nn.Module):
                 ff_linear_module_up=transformer_ff_linear_module_up,
                 ff_linear_module_down=transformer_ff_linear_module_down,
                 msa_scaling=transformer_msa_scaling,
-                mlp_dropout=transformer_mlp_dropout,
+                ff_dropout=transformer_ff_dropout,
+                ff_inner_dropout=transformer_ff_inner_dropout,
+                ff_outer_dropout=transformer_ff_outer_dropout,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
                 causal=False,
@@ -357,9 +363,21 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=(
                     # First truthy assigned value
-                    transformer_initial_ff_mlp_dropout
-                    if transformer_initial_ff_mlp_dropout is not None
-                    else transformer_mlp_dropout
+                    transformer_initial_ff_dropout
+                    if transformer_initial_ff_dropout is not None
+                    else transformer_ff_dropout
+                ),
+                inner_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_inner_dropout
+                    if transformer_initial_ff_inner_dropout is not None
+                    else transformer_ff_inner_dropout
+                ),
+                outer_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_outer_dropout
+                    if transformer_initial_ff_outer_dropout is not None
+                    else transformer_ff_outer_dropout
                 ),
                 linear_module_up=(
                     # First truthy assigned value
@@ -441,7 +459,9 @@ class ViT(nn.Module):
         transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
-        transformer_initial_ff_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -458,7 +478,9 @@ class ViT(nn.Module):
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
@@ -508,7 +530,9 @@ class ViT(nn.Module):
             transformer_initial_ff_residual_path=transformer_initial_ff_residual_path,
             transformer_initial_ff_linear_module_up=transformer_initial_ff_linear_module_up,
             transformer_initial_ff_linear_module_down=transformer_initial_ff_linear_module_down,
-            transformer_initial_ff_mlp_dropout=transformer_initial_ff_mlp_dropout,
+            transformer_initial_ff_dropout=transformer_initial_ff_dropout,
+            transformer_initial_ff_inner_dropout=transformer_initial_ff_inner_dropout,
+            transformer_initial_ff_outer_dropout=transformer_initial_ff_outer_dropout,
             transformer_pre_norm=transformer_pre_norm,
             transformer_normformer=transformer_normformer,
             transformer_post_norm=transformer_post_norm,
@@ -525,7 +549,9 @@ class ViT(nn.Module):
             transformer_ff_linear_module_up=transformer_ff_linear_module_up,
             transformer_ff_linear_module_down=transformer_ff_linear_module_down,
             transformer_msa_scaling=transformer_msa_scaling,
-            transformer_mlp_dropout=transformer_mlp_dropout,
+            transformer_ff_dropout=transformer_ff_dropout,
+            transformer_ff_inner_dropout=transformer_ff_inner_dropout,
+            transformer_ff_outer_dropout=transformer_ff_outer_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             transformer_checkpoint_ff=transformer_checkpoint_ff,

{broccoli_ml-6.0.0.dist-info → broccoli_ml-9.2.2.dist-info}/METADATA

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

{broccoli_ml-6.0.0.dist-info → broccoli_ml-9.2.2.dist-info}/RECORD

@@ -1,13 +1,13 @@
 broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
 broccoli/activation.py,sha256=-Jf30C6iGqWCorC9HEGn2oduWwjeaCAxGLUUYIy1zX8,3438
 broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
-broccoli/linear.py,sha256=Y7s-DzcwsOipRboNHc4HTScw4mJRalNoVFsNcxOB6a4,4872
+broccoli/linear.py,sha256=IwvPAMbHqOYyz0g5WZyevPAhC1Pn0RTLniFM4E6lJoI,11511
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=MxIdzoxoWx_IWcq86vDZJIV4tk-dMNivhopZu8zJk90,23293
+broccoli/transformer.py,sha256=r-ggAeNDW5QpBi9As1U9sIfxITBOx0WHk_K4zWpyTM8,26233
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=9oyh76ulmX5lDPMCDicQhhqm8RYCvJIgAJkDbYRVdi4,20873
-broccoli_ml-6.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-6.0.0.dist-info/METADATA,sha256=Sv8nRPb7oCAeoMe3AAHIYDewAETvb0ZDxN8IKFniVHk,1368
-broccoli_ml-6.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-6.0.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=sC6K3FK3a8ojOgvNWSWhuZHBtnFrrTQbsDdlagcKJH4,22224
+broccoli_ml-9.2.2.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-9.2.2.dist-info/METADATA,sha256=8ySQYntl9czgYyEQN5nyPS31tjwYC8M8Mx_iYhvtbzg,1368
+broccoli_ml-9.2.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-9.2.2.dist-info/RECORD,,