broccoli-ml 9.0.0__py3-none-any.whl → 9.2.0__py3-none-any.whl

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/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.dist-info → broccoli_ml-9.2.0.dist-info}/METADATA

@@ -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.dist-info → broccoli_ml-9.2.0.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=XaGHZguvK-7hvtIt07zo8uQZBQvS7oMD2K9nPvyYJLE,9769
+broccoli/linear.py,sha256=7uN7zVPJ6Ptec31O8a-GvWT5nZk56Wf1RLJRvUAT0yo,11406
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=Rozh0hExHjwGvvKbMeZfLoB95dDKyDn3X6o1Ms26aAI,24241
+broccoli/transformer.py,sha256=ckwFpNTAeYB_V8F-_DwT0Z5--QUEepM4r6xlSxhCY68,26079
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
 broccoli/vit.py,sha256=sC6K3FK3a8ojOgvNWSWhuZHBtnFrrTQbsDdlagcKJH4,22224
-broccoli_ml-9.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-9.0.0.dist-info/METADATA,sha256=ecQ2BRxtzmNSO2CMAp2rcNRq9L37urE_pKdsPf-jJKs,1368
-broccoli_ml-9.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-9.0.0.dist-info/RECORD,,
+broccoli_ml-9.2.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-9.2.0.dist-info/METADATA,sha256=MnrReG_teIm7YceDS9nfvXnelEXHGP4JFhoz9dhkJWo,1368
+broccoli_ml-9.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-9.2.0.dist-info/RECORD,,