broccoli-ml 9.5.1__py3-none-any.whl

broccoli/__init__.py

@@ -0,0 +1,6 @@
+from . import activation
+from . import cnn
+from . import tensor
+from . import transformer
+from . import vit
+from . import rope

broccoli/activation.py

@@ -0,0 +1,118 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class ReLU(nn.Module):
+    """
+    A ReLU activation function with optional clamp and leakiness.
+    """
+
+    def __init__(
+        self, clamp=True, leaky=True, negative_slope=0.01, clamp_max=6.0
+    ) -> None:
+        super().__init__()
+        self.clamp = clamp
+        self.leaky = leaky
+        self.negative_slope = negative_slope
+        self.clamp_max = clamp_max
+
+    def forward(self, x):
+        if self.leaky:
+            relu = F.leaky_relu(x, negative_slope=self.negative_slope)
+        else:
+            relu = F.relu(x)
+        if self.clamp:
+            relu = torch.clamp(relu, max=self.clamp_max)
+        return relu
+
+
+class GELU(nn.Module):
+    """
+    A GELU activation function with optional clamp.
+    """
+
+    def __init__(self, clamp=True) -> None:
+        super().__init__()
+        self.clamp = clamp
+        self.gelu = nn.GELU()
+
+    def forward(self, x):
+        gelu = self.gelu(x)
+        if self.clamp:
+            gelu = torch.clamp(gelu, max=6)
+        return gelu
+
+
+class Swish(nn.Module):
+    """
+    Implementation of (beta) Swish
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        # Learnable parameter is called "swiglu beta" so that it is easy to find
+        #   and exclude from weight decay
+        self.swish_beta = nn.Parameter(torch.tensor([1.0]))
+
+    def forward(self, x):
+        return x * F.sigmoid(self.swish_beta * x)
+
+
+class SquaredReLU(nn.Module):
+    """
+    Squared ReLU, as shown in "ReLU^2 wins" (https://arxiv.org/abs/2402.03804) to
+      be as effective as SwiGLU for training LLMs, possibly because it can allow a
+      NN to learn multyiplication, as noted by
+      https://azizbelaweid.substack.com/p/what-is-swiglu-how-to-implement-it
+    """
+
+    def __init__(
+        self, clamp=True, leaky=True, negative_slope: float = 0.01, clamp_max=6
+    ) -> None:
+        super().__init__()
+        self.clamp = clamp
+        self.leaky = leaky
+        self.negative_slope = negative_slope
+        self.clamp_max = clamp_max
+
+    def forward(self, x):
+        if self.leaky:
+            relu = F.leaky_relu(x, negative_slope=self.negative_slope)
+        else:
+            relu = F.relu(x)
+        relu_squared = relu**2
+        if self.clamp:
+            relu_squared = torch.clamp(relu_squared, max=self.clamp_max)
+        return relu_squared
+
+
+class XGLU(nn.Module):
+    """
+    Generic Gated Linear Unit
+    """
+
+    def __init__(self, activation_module: nn.Module) -> None:
+        super().__init__()
+        self.activation = activation_module
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, value = x.chunk(2, dim=-1)
+        return self.activation(gate) * value
+
+
+def SquaredReGLU(clamp=True, leaky=True, negative_slope=0.01, clamp_max=6.0) -> XGLU:
+    """
+    Factory function that creates a GLU with a SquaredReLU activation.
+    """
+    activation_module = SquaredReLU(
+        clamp=clamp, leaky=leaky, negative_slope=negative_slope, clamp_max=clamp_max
+    )
+    return XGLU(activation_module)
+
+
+def SwiGLU() -> XGLU:
+    """
+    Factory function that creates a GLU with a Swish activation.
+    """
+    return XGLU(Swish())

broccoli/cnn.py

@@ -0,0 +1,157 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Union
+
+from einops.layers.torch import Rearrange
+
+
+def spatial_tuple(size: Union[int, tuple], spatial_dimensions):
+    """
+    Converts an integer x to `tuple([x] * spatial_dimensions)`.
+    Performs no operation (i.e. the identity operation) on tuples of length `spatial_dimensions`.
+    Otherwise
+    """
+    if isinstance(size, int):
+        return tuple([size] * spatial_dimensions)
+    elif isinstance(size, tuple) and (len(size) == spatial_dimensions):
+        return size
+    else:
+        raise ValueError(
+            f"For {spatial_dimensions} spatial dimensions, `size` must be "
+            f"an integer or a tuple of length {spatial_dimensions}."
+        )
+
+
+def padding_tensor(padding: tuple):
+    """
+    Converts a tuple of ints (x, y, z) into a tuple of 2-tuples,
+        like ((x, x), (y, y), (z, z)).
+
+    Performs no operation (i.e. the identity operation) on a tuple of 2-tuples.
+
+    Otherwise raises an error.
+    """
+    if all(isinstance(x, int) for x in padding):
+        return tuple([tuple([p] * 2) for p in padding])
+    elif (
+        all(isinstance(p, tuple) for p in padding)
+        and all(len(p) == 2 for p in padding)
+        and all(all(isinstance(x, int) for x in p) for p in padding)
+    ):
+        return padding
+    else:
+        raise ValueError(
+            "Padding must be a tuple of ints of a tuple of 2-tuples of ints. "
+            f"It was {padding}."
+        )
+
+
+def kd_unfold(t: torch.Tensor, kernel_size=1, stride=1, padding=0, k=2):
+    """
+    Unfold operation with k spatial dimensions.
+    Does not support dilation.
+    Only supports equal padding at top and bottom.
+    """
+    if len(t.size()[2:]) != k:
+        raise ValueError(
+            f"Input tensor size should be (N, channels, spatial dims...), so "
+            f"for k = {k}, t.size() should be a tuple of length {k + 2}."
+        )
+
+    N, C = t.size(0), t.size(1)
+
+    kernel_size = spatial_tuple(kernel_size, k)
+    stride = spatial_tuple(stride, k)
+    padding = padding_tensor(spatial_tuple(padding, k))
+
+    output = t
+    output = F.pad(output, sum(reversed(padding), ()))  # i.e. the empty tuple
+
+    for i, _ in enumerate(kernel_size):
+        output = output.unfold(i + 2, kernel_size[i], stride[i])
+
+    permutation = [0, 1] + [i + k + 2 for i in range(k)] + [i + 2 for i in range(k)]
+
+    return output.permute(*permutation).reshape(N, math.prod(kernel_size) * C, -1)
+
+
+def calculate_output_spatial_size(
+    input_spatial_size, kernel_size=1, stride=1, padding=0, dilation=0
+):
+    """
+    Calculate the output size for the spatial dimensions of a convolutional operation
+    """
+    stride = spatial_tuple(stride, len(input_spatial_size))
+
+    # Handle padding keywords that are sometimes used
+    if padding == "same":
+        output_size = ()
+        for i, in_length in enumerate(input_spatial_size):
+            output_size += (math.ceil(in_length / stride[i]),)
+        return output_size
+    elif padding == "valid":
+        padding = 0
+
+    kernel_size = spatial_tuple(kernel_size, len(input_spatial_size))
+    padding = spatial_tuple(padding, len(input_spatial_size))
+    dilation = spatial_tuple(dilation, len(input_spatial_size))
+
+    output_size = ()
+
+    for i, in_length in enumerate(input_spatial_size):
+        output_size += (
+            math.floor(
+                (in_length + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1)
+                / stride[i]
+                + 1
+            ),
+        )
+    return output_size
+
+
+class SpaceToDepth(nn.Module):
+    """
+    An operation that extracts patches from an image-like tensor and stacks
+        them channel-wise.
+    """
+
+    def __init__(self, kernel_size, stride=1, padding=0, spatial_dimensions=2):
+        """
+        Input shape should be in order (channels, spatial dims...),
+            e.g. (channels, height, width)
+        """
+
+        super().__init__()
+
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.spatial_dimensions = spatial_dimensions
+
+    def forward(self, x):
+
+        N, C, *input_spatial_size = x.size()
+
+        patches = kd_unfold(
+            x,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            k=self.spatial_dimensions,
+        )
+
+        output_spatial_size = calculate_output_spatial_size(
+            input_spatial_size=input_spatial_size,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=1,  # kd_unfold doesn't support dilation
+        )
+
+        output_channels = C * math.prod(
+            spatial_tuple(self.kernel_size, self.spatial_dimensions)
+        )
+
+        return patches.view(N, output_channels, *output_spatial_size)

broccoli/linear.py

@@ -0,0 +1,352 @@
+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
+
+from .tensor import SigmaReparamTensor, AnchoredReparamTensor, NormReparamTensor
+
+
+class SpectralNormLinear(nn.Module):
+    """
+    Inspired by Apple's Spectral Normed Linear Layers
+        (https://github.com/apple/ml-sigma-reparam)
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        weights = torch.empty(self.out_features, self.in_features)
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = SigmaReparamTensor(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        # Optional: A nice representation for printing the module.
+        return (
+            f"SpectralNormFeedForward(in_features={self.in_features},"
+            f"out_features={self.out_features}, bias={self.use_bias})"
+        )
+
+
+class AnchoredLinear(nn.Module):
+    """
+    ...
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        weights = torch.empty(self.out_features, self.in_features)
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = AnchoredReparamTensor(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        # Optional: A nice representation for printing the module.
+        return (
+            f"AnchoredLinear(in_features={self.in_features},"
+            f"out_features={self.out_features}, bias={self.use_bias})"
+        )
+
+
+class WeightNormedLinear(nn.Module):
+    """
+    ...
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        weights = torch.empty(self.out_features, self.in_features)
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = NormReparamTensor(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        return (
+            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,
+        xglu=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.bias = bias
+        self.xglu = xglu
+        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 reset_rows(self, indices):
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.size(0):
+            value_indices = indices
+            centred_value_weights = self._mean_value_weights()
+            centred_value_weights = centred_value_weights.expand(indices.size(0), -1)
+            if self.xglu:
+                gate_indices = indices
+                value_indices = indices + (self.linear.out_features // 2)
+                centred_gate_weights = self._mean_gate_weights()
+                centred_gate_weights = centred_gate_weights.expand(indices.size(0), -1)
+                self._update_weights(
+                    gate_indices, 0, centred_gate_weights, self.optimisers  # dim
+                )
+            self._update_weights(
+                value_indices, 0, centred_value_weights, self.optimisers
+            )
+        else:
+            return
+
+    def reset_columns(self, indices):
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.size(0):
+            random_weights = self._random_weights(
+                self.linear.weight.size(0), indices.size(0)
+            )
+            # Make random col weights quiet so they don't introduce loud noise...
+            # ...but not so quiet that FP16 zeros them and ruins symmetry breaking!
+            random_weights *= 0.1
+            self._update_weights(indices, 1, random_weights, self.optimisers)  # dim
+        else:
+            return
+
+    def forward(self, x):
+        if self.training and self.optimisers:
+            self.reset_rows(self.get_reset_indices(0))
+            self.reset_columns(self.get_reset_indices(1))
+        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 get_reset_indices(self, dim):
+        base_rate = self.row_recycling_rate if dim == 0 else self.column_recycling_rate
+        p = base_rate * self._get_multiplier()
+        if dim == 0:
+            if self.xglu:
+                sample_space = self.linear.out_features // 2
+            else:
+                sample_space = self.linear.out_features
+        elif dim == 1:
+            sample_space = self.linear.in_features
+        else:
+            raise ValueError("`dim` must be 0 or 1")
+
+        # Sample the indices
+        probs = torch.rand(sample_space, device=self.linear.weight.device)
+        mask = probs < p
+        if mask.any():
+            return torch.nonzero(mask).squeeze(-1)
+        else:
+            return torch.tensor([], dtype=torch.long, device=self.linear.weight.device)
+
+    def _random_weights(self, rows, columns):
+        device = self.linear.weight.device
+        weights = self.linear.weight.data
+        stdv = 1.0 / math.sqrt(weights.size(1))
+        random_weights = torch.rand(rows, columns, device=device)
+        random_weights -= 0.5  # Range [-0.5, +0.5]
+        random_weights *= 2.0 * stdv  # Range [-stdv, +stdv]
+        return random_weights
+
+    def _mean_value_weights(self):
+        """
+        Only used when self.xglu
+        """
+        weights = self.linear.weight.data
+        rows = weights.size(0)
+        if self.xglu:
+            return self.linear.weight[int(rows / 2) :].data.mean(dim=0, keepdim=True)
+        else:
+            return self.linear.weight.data.mean(dim=0, keepdim=True)
+
+    def _mean_gate_weights(self):
+        """
+        Only used when self.xglu
+        """
+        weights = self.linear.weight.data
+        rows = weights.size(0)
+        return self.linear.weight[: int(rows / 2)].data.mean(dim=0, keepdim=True)
+
+    def _update_weights(
+        self,
+        indices: Iterable[int],
+        dim: int,
+        data: torch.Tensor,
+        optimisers: Union[
+            List[torch.optim.Optimizer], torch.optim.Optimizer, None
+        ] = None,
+    ):
+        if optimisers is None:
+            optimisers = []
+        if not isinstance(optimisers, list):
+            optimisers = [optimisers]
+
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.numel() == 0:
+            return
+
+        with torch.no_grad():
+            if dim == 0:
+                self.linear.weight.data[idx_tensor] = data
+            elif dim == 1:
+                self.linear.weight.data[:, idx_tensor] = data
+            else:
+                raise ValueError("`dim` must be 0 or 1")
+
+            self._reset_optim_state(self.linear.weight, idx_tensor, optimisers, dim=dim)
+
+    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