broccoli-ml 9.0.0__tar.gz → 9.2.0__tar.gz

{broccoli_ml-9.0.0 → broccoli_ml-9.2.0}/PKG-INFO

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

{broccoli_ml-9.0.0 → broccoli_ml-9.2.0}/broccoli/linear.py

@@ -1,5 +1,6 @@
 import math
 import random
+import warnings
 from typing import Union, List, Iterable
 
 import torch
@@ -149,34 +150,73 @@ class RecyclingLinear(nn.Module):
         bias: bool = True,
         row_recycling_rate: float = 0.0,
         column_recycling_rate: float = 0.0,
+        adaptive=False,
     ):
         super().__init__()
         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 self.row_recycling_rate > 0:
+            if row_recycling_rate > 0:
                 probs = torch.rand(self.linear.out_features, device=x.device)
-                mask = probs < self.row_recycling_rate
+                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 self.column_recycling_rate > 0:
+            if col_recycling_rate > 0:
                 probs = torch.rand(self.linear.in_features, device=x.device)
-                mask = probs < self.column_recycling_rate
+                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(

{broccoli_ml-9.0.0 → broccoli_ml-9.2.0}/broccoli/transformer.py

@@ -340,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(
@@ -364,27 +365,63 @@ class FeedforwardBlock(nn.Module):
         )
 
         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.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") or hasattr(
+                self.linear_out, "row_recycling_rate"
+            ):
+                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:

{broccoli_ml-9.0.0 → broccoli_ml-9.2.0}/pyproject.toml

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