sparse-layers 0.2.0__py3-none-any.whl

sparse_layers/__init__.py

@@ -0,0 +1,69 @@
+"""sparse-layers: structured sparse layers for building memory-efficient neural networks."""
+
+from sparse_layers.layers import (
+    ButterflyLinear,
+    ButterflyMLP,
+    ButterflyMultiHeadAttention,
+    CustomLinear,
+    CustomMLP,
+    MultiHeadAttention,
+    PaddedButterflyLinear,
+    SimpleMLP,
+)
+from sparse_layers.sse import (
+    LinearAttention,
+    LinearAttentionConfig,
+    NaiveMultiPartitionState,
+    NaiveSSEAttention,
+    NaiveSSEMultiHeadAttention,
+    SSEAttention,
+    SSEAttentionAdaptive,
+    SSEAttentionAdaptiveConfig,
+    SSEAttentionConfig,
+    SSEMaskingOps,
+    SSEMaskingOpsConfig,
+    SSEMultiHeadAttention,
+    SSEMultiHeadAttentionConfig,
+    SSEMultiPartitionState,
+    SSEMultiPartitionStateConfig,
+    SSEPartitionSelector,
+    SSEPartitionSelectorConfig,
+    SSESparseSoftmax,
+    SSESparseSoftmaxConfig,
+    SSEVarlenOps,
+    SSEVarlenOpsConfig,
+)
+
+__all__ = [
+    # Layers
+    "ButterflyLinear",
+    "ButterflyMLP",
+    "ButterflyMultiHeadAttention",
+    "CustomLinear",
+    "CustomMLP",
+    "MultiHeadAttention",
+    "PaddedButterflyLinear",
+    "SimpleMLP",
+    # SSE
+    "LinearAttention",
+    "LinearAttentionConfig",
+    "NaiveMultiPartitionState",
+    "NaiveSSEAttention",
+    "NaiveSSEMultiHeadAttention",
+    "SSEAttention",
+    "SSEAttentionAdaptive",
+    "SSEAttentionAdaptiveConfig",
+    "SSEAttentionConfig",
+    "SSEMaskingOps",
+    "SSEMaskingOpsConfig",
+    "SSEMultiHeadAttention",
+    "SSEMultiHeadAttentionConfig",
+    "SSEMultiPartitionState",
+    "SSEMultiPartitionStateConfig",
+    "SSEPartitionSelector",
+    "SSEPartitionSelectorConfig",
+    "SSESparseSoftmax",
+    "SSESparseSoftmaxConfig",
+    "SSEVarlenOps",
+    "SSEVarlenOpsConfig",
+]

sparse_layers/layers/__init__.py

@@ -0,0 +1,23 @@
+"""Core neural network layers with butterfly factorization support."""
+
+from sparse_layers.layers.butterfly_linear import ButterflyLinear
+from sparse_layers.layers.butterfly_mlp import ButterflyMLP
+from sparse_layers.layers.butterfly_multi_head_attention import (
+    ButterflyMultiHeadAttention,
+)
+from sparse_layers.layers.custom_linear import CustomLinear
+from sparse_layers.layers.custom_mlp import CustomMLP
+from sparse_layers.layers.multi_head_attention import MultiHeadAttention
+from sparse_layers.layers.padded_butterfly_linear import PaddedButterflyLinear
+from sparse_layers.layers.simple_mlp import SimpleMLP
+
+__all__ = [
+    "ButterflyLinear",
+    "ButterflyMLP",
+    "ButterflyMultiHeadAttention",
+    "CustomLinear",
+    "CustomMLP",
+    "MultiHeadAttention",
+    "PaddedButterflyLinear",
+    "SimpleMLP",
+]

sparse_layers/layers/butterfly_linear.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+def _is_power_of_two(value: int) -> bool:
+    return value > 0 and (value & (value - 1) == 0)
+
+
+class ButterflyLinear(nn.Module):
+    """Butterfly factorization based linear layer.
+
+    This layer follows the structure described in section 2.3.1 of the
+    `reducing_memory_requirements_ipu_butterfly.md` reference. It consists of
+    log2(N) stages of block-diagonal 2x2 butterfly factors, each stored as a
+    learnable parameter. The layer currently supports square power-of-two
+    dimensions and acts as a drop-in replacement for :class:`torch.nn.Linear`.
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super().__init__()
+
+        if in_features <= 0 or out_features <= 0:
+            raise ValueError("in_features and out_features must be positive integers")
+
+        if in_features != out_features:
+            raise ValueError("ButterflyLinear requires in_features == out_features")
+
+        if not _is_power_of_two(in_features):
+            raise ValueError(
+                "ButterflyLinear requires dimensions that are a power of two"
+            )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self._depth = int(math.log2(in_features))
+
+        num_blocks = in_features // 2
+        default_dtype = torch.get_default_dtype()
+        identity_block = torch.eye(2, dtype=default_dtype).unsqueeze(0)
+
+        self.factors = nn.ParameterList()
+        for _ in range(self._depth):
+            factor = identity_block.repeat(num_blocks, 1, 1)
+            factor += 0.01 * torch.randn_like(factor)
+            self.factors.append(nn.Parameter(factor))
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features, dtype=default_dtype))
+        else:
+            self.register_parameter("bias", None)
+
+    @classmethod
+    def from_linear(
+        cls,
+        layer: nn.Linear,
+        *,
+        optimization_steps: int = 4000,
+        learning_rate: float = 0.1,
+        tolerance: float = 1e-7,
+        seed: int | None = None,
+    ) -> "ButterflyLinear":
+        """Construct a butterfly layer approximating a dense :class:`nn.Linear`.
+
+        The method uses gradient-based fitting over the canonical basis to match
+        the transformation represented by ``layer``. The bias term is copied
+        directly before optimisation to accelerate convergence.
+        """
+
+        if layer.in_features != layer.out_features:
+            raise ValueError("ButterflyLinear.from_linear requires a square nn.Linear")
+
+        if not _is_power_of_two(layer.in_features):
+            raise ValueError(
+                "ButterflyLinear.from_linear requires dimensions that are a power of two"
+            )
+
+        device = layer.weight.device
+        dtype = layer.weight.dtype
+
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        result = cls(layer.in_features, layer.out_features, bias=layer.bias is not None)
+        result = result.to(device=device, dtype=dtype)
+
+        if result.bias is not None and layer.bias is not None:
+            with torch.no_grad():
+                result.bias.copy_(layer.bias)
+
+        params = list(result.factors.parameters())
+        if not params:
+            return result
+
+        optimizer_adam = torch.optim.Adam(params, lr=learning_rate)
+
+        eye_input = torch.eye(layer.in_features, device=device, dtype=dtype)
+        with torch.no_grad():
+            target = layer(eye_input)
+
+        best_loss = float("inf")
+
+        for step in range(optimization_steps):
+            optimizer_adam.zero_grad()
+            output = result(eye_input)
+            loss = F.mse_loss(output, target)
+            loss.backward()
+            optimizer_adam.step()
+
+            current_loss = loss.item()
+            if current_loss < best_loss:
+                best_loss = current_loss
+
+            if best_loss <= tolerance:
+                break
+
+        if best_loss > tolerance:
+            optimizer_lbfgs = torch.optim.LBFGS(
+                params,
+                lr=1.0,
+                max_iter=100,
+                tolerance_grad=1e-12,
+                tolerance_change=1e-12,
+                line_search_fn="strong_wolfe",
+            )
+
+            def closure() -> Tensor:
+                optimizer_lbfgs.zero_grad()
+                output_lbfgs = result(eye_input)
+                lbfgs_loss = F.mse_loss(output_lbfgs, target)
+                lbfgs_loss.backward()
+                return lbfgs_loss
+
+            for _ in range(20):
+                loss = optimizer_lbfgs.step(closure)
+                best_loss = min(best_loss, loss.item())
+                if best_loss <= tolerance:
+                    break
+
+        return result
+
+    def forward(self, input: Tensor) -> Tensor:
+        if input.shape[-1] != self.in_features:
+            raise ValueError(
+                f"Expected last dimension {self.in_features}, got {input.shape[-1]}"
+            )
+
+        original_shape = input.shape[:-1]
+        x = input.reshape(-1, self.in_features)
+
+        for stage_index, factor in enumerate(self.factors):
+            x = self._apply_stage(x, factor, stage_index)
+
+        if self.bias is not None:
+            x = x + self.bias
+
+        return x.reshape(*original_shape, self.out_features)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"bias={self.bias is not None}"
+        )
+
+    def _apply_stage(self, x: Tensor, factor: Tensor, stage: int) -> Tensor:
+        batch = x.shape[0]
+        n = self.in_features
+        block = 1 << (stage + 1)
+        half = block >> 1
+
+        staged = x.reshape(batch, -1, block)
+        staged = staged.reshape(batch, -1, half, 2)
+        staged = staged.permute(0, 1, 3, 2).contiguous()
+        pairs = staged.reshape(batch, -1, 2)
+
+        transformed = torch.einsum("bnc,ncd->bnd", pairs, factor)
+
+        transformed = transformed.reshape(batch, -1, 2, half)
+        transformed = transformed.permute(0, 1, 3, 2).contiguous()
+        transformed = transformed.reshape(batch, -1, block)
+        return transformed.reshape(batch, n)
+
+    def to_linear(self) -> nn.Linear:
+        """Return a dense :class:`nn.Linear` with identical behaviour."""
+
+        factor_tensor = self.factors[0] if len(self.factors) > 0 else None
+        if factor_tensor is not None:
+            device = factor_tensor.device
+            dtype = factor_tensor.dtype
+        elif self.bias is not None:
+            device = self.bias.device
+            dtype = self.bias.dtype
+        else:
+            device = torch.device("cpu")
+            dtype = torch.get_default_dtype()
+
+        linear = nn.Linear(
+            self.in_features,
+            self.out_features,
+            bias=self.bias is not None,
+            device=device,
+            dtype=dtype,
+        )
+
+        with torch.no_grad():
+            if self.bias is not None:
+                bias_backup = self.bias.data.clone()
+                self.bias.zero_()
+            else:
+                bias_backup = None
+
+            identity = torch.eye(self.in_features, device=device, dtype=dtype)
+            weight_matrix = self(identity)
+
+            if bias_backup is not None:
+                self.bias.copy_(bias_backup)
+                linear.bias.copy_(bias_backup)
+
+            linear.weight.copy_(weight_matrix.t())
+
+        return linear

sparse_layers/layers/butterfly_mlp.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+
+from sparse_layers.layers.butterfly_linear import ButterflyLinear, _is_power_of_two
+from sparse_layers.layers.simple_mlp import SimpleMLP
+
+
+class ButterflyMLP(nn.Module):
+    """A multi-layer perceptron composed of :class:`ButterflyLinear` layers."""
+
+    def __init__(self, input_dim: int, hidden_dims: Sequence[int], output_dim: int) -> None:
+        super().__init__()
+
+        self._validate_dimensions(input_dim, hidden_dims, output_dim)
+
+        self.input_dim = input_dim
+        self.hidden_dims = list(hidden_dims)
+        self.output_dim = output_dim
+
+        layer_dims = [input_dim, *self.hidden_dims, output_dim]
+
+        modules: list[nn.Module] = []
+        for idx in range(len(layer_dims) - 1):
+            in_features = layer_dims[idx]
+            out_features = layer_dims[idx + 1]
+
+            linear = ButterflyLinear(in_features, out_features)
+            modules.append(linear)
+
+            if idx < len(layer_dims) - 2:
+                modules.append(nn.ReLU())
+
+        self.network = nn.Sequential(*modules)
+
+    @staticmethod
+    def _validate_dimensions(
+        input_dim: int, hidden_dims: Sequence[int], output_dim: int
+    ) -> None:
+        if input_dim <= 0:
+            raise ValueError("input_dim must be a positive integer")
+
+        if not hidden_dims:
+            raise ValueError("hidden_dims must contain at least one positive integer")
+
+        if any(dim <= 0 for dim in hidden_dims):
+            raise ValueError("hidden_dims must contain only positive integers")
+
+        if output_dim <= 0:
+            raise ValueError("output_dim must be a positive integer")
+
+        all_dims = [input_dim, *hidden_dims, output_dim]
+        if len(set(all_dims)) != 1:
+            raise ValueError("ButterflyMLP requires all layer dimensions to be identical")
+
+        if not all(_is_power_of_two(dim) for dim in all_dims):
+            raise ValueError("ButterflyMLP requires power-of-two layer dimensions")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() != 2:
+            raise ValueError("ButterflyMLP expects a 2D input tensor of shape (batch, features)")
+
+        if x.shape[1] != self.input_dim:
+            raise ValueError(
+                f"Expected input with {self.input_dim} features, received {x.shape[1]}"
+            )
+
+        return self.network(x)
+
+    def to_simple_mlp(self) -> SimpleMLP:
+        """Return a :class:`SimpleMLP` with identical behaviour."""
+
+        simple = SimpleMLP(self.input_dim, self.hidden_dims, self.output_dim)
+
+        sparse_layers = [module for module in self.network if isinstance(module, ButterflyLinear)]
+        dense_layers = [module for module in simple.modules() if isinstance(module, nn.Linear)]
+
+        if len(sparse_layers) != len(dense_layers):
+            raise RuntimeError("Unexpected layer mismatch during conversion to SimpleMLP")
+
+        for butterfly_layer, dense_layer in zip(sparse_layers, dense_layers):
+            converted = butterfly_layer.to_linear()
+            with torch.no_grad():
+                dense_layer.weight.copy_(converted.weight)
+                if dense_layer.bias is not None and converted.bias is not None:
+                    dense_layer.bias.copy_(converted.bias)
+
+        return simple
+
+    @classmethod
+    def from_simple_mlp(
+        cls,
+        model: SimpleMLP,
+        *,
+        seed: int | None = None,
+        optimization_steps: int = 4000,
+        learning_rate: float = 0.1,
+        tolerance: float = 1e-7,
+    ) -> "ButterflyMLP":
+        """Construct a :class:`ButterflyMLP` from a compatible :class:`SimpleMLP`."""
+
+        linear_layers = [module for module in model.modules() if isinstance(module, nn.Linear)]
+
+        if not linear_layers:
+            raise ValueError("SimpleMLP must contain at least one linear layer")
+
+        layer_dims = []
+        for layer in linear_layers:
+            if layer.in_features != layer.out_features:
+                raise ValueError(
+                    "ButterflyMLP.from_simple_mlp requires square nn.Linear layers"
+                )
+            if not _is_power_of_two(layer.in_features):
+                raise ValueError(
+                    "ButterflyMLP.from_simple_mlp requires power-of-two dimensions"
+                )
+            layer_dims.append(layer.in_features)
+
+        input_dim = linear_layers[0].in_features
+        hidden_dims = [layer.out_features for layer in linear_layers[:-1]]
+        output_dim = linear_layers[-1].out_features
+
+        result = cls(input_dim, hidden_dims, output_dim)
+
+        butterfly_indices = [
+            idx for idx, module in enumerate(result.network) if isinstance(module, ButterflyLinear)
+        ]
+
+        if len(butterfly_indices) != len(linear_layers):
+            raise RuntimeError("Unexpected layer mismatch during reconstruction from SimpleMLP")
+
+        for index, (linear_layer, target_idx) in enumerate(zip(linear_layers, butterfly_indices)):
+            converted = ButterflyLinear.from_linear(
+                linear_layer,
+                seed=seed,
+                optimization_steps=optimization_steps,
+                learning_rate=learning_rate,
+                tolerance=tolerance,
+            )
+            result.network[target_idx] = converted
+
+        return result

sparse_layers/layers/butterfly_multi_head_attention.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor, nn
+
+from sparse_layers.layers.butterfly_linear import ButterflyLinear
+
+
+class ButterflyMultiHeadAttention(nn.Module):
+    """Multi-head self-attention using ButterflyLinear projections."""
+
+    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.0) -> None:
+        super().__init__()
+
+        if d_model <= 0:
+            raise ValueError("d_model must be a positive integer")
+        if num_heads <= 0:
+            raise ValueError("num_heads must be a positive integer")
+        if d_model % num_heads != 0:
+            raise ValueError("d_model must be divisible by num_heads")
+        if not (0.0 <= dropout < 1.0):
+            raise ValueError("dropout must satisfy 0.0 <= dropout < 1.0")
+
+        if not _is_power_of_two(d_model):
+            raise ValueError("ButterflyMultiHeadAttention requires d_model to be a power of two")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+
+        self.query = ButterflyLinear(d_model, d_model)
+        self.key = ButterflyLinear(d_model, d_model)
+        self.value = ButterflyLinear(d_model, d_model)
+        self.out = ButterflyLinear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+        self._scaling = 1.0 / math.sqrt(self.head_dim)
+
+    def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
+        if x.dim() != 3:
+            raise ValueError("expected input of shape (batch, seq_len, d_model)")
+        batch_size, seq_len, feature_dim = x.shape
+        if feature_dim != self.d_model:
+            raise ValueError(
+                f"expected last dimension {self.d_model}, received {feature_dim}"
+            )
+
+        query = self.query(x)
+        key = self.key(x)
+        value = self.value(x)
+
+        query_heads = self._split_heads(query)
+        key_heads = self._split_heads(key)
+        value_heads = self._split_heads(value)
+
+        scores = torch.matmul(query_heads, key_heads.transpose(-2, -1)) * self._scaling
+
+        if mask is not None:
+            if mask.shape != (batch_size, seq_len):
+                raise ValueError("mask shape must match (batch, seq_len)")
+            if mask.dtype != torch.bool:
+                raise ValueError("mask must have dtype torch.bool")
+            expanded_mask = mask.unsqueeze(1).unsqueeze(2)
+            scores = scores.masked_fill(expanded_mask, float("-inf"))
+
+        attention = torch.softmax(scores, dim=-1)
+        attention = self.dropout(attention)
+
+        context = torch.matmul(attention, value_heads)
+        merged_context = self._merge_heads(context)
+
+        output = self.out(merged_context)
+        return output
+
+    def _split_heads(self, tensor: Tensor) -> Tensor:
+        batch_size, seq_len, _ = tensor.shape
+        reshaped = tensor.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        return reshaped.permute(0, 2, 1, 3)
+
+    def _merge_heads(self, tensor: Tensor) -> Tensor:
+        batch_size, _, seq_len, _ = tensor.shape
+        transposed = tensor.permute(0, 2, 1, 3)
+        return transposed.reshape(batch_size, seq_len, self.d_model)
+
+
+def _is_power_of_two(value: int) -> bool:
+    return value > 0 and (value & (value - 1) == 0)
+
+
+__all__ = ["ButterflyMultiHeadAttention"]
+

sparse_layers/layers/custom_linear.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor, nn
+from torch.nn.parameter import Parameter
+
+
+class CustomLinear(nn.Module):
+    """Drop-in replacement for nn.Linear with identical behavior."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super().__init__()
+
+        if in_features <= 0:
+            raise ValueError("in_features must be a positive integer")
+        if out_features <= 0:
+            raise ValueError("out_features must be a positive integer")
+
+        self.in_features = in_features
+        self.out_features = out_features
+
+        self.weight: Parameter = Parameter(torch.empty(out_features, in_features))
+        if bias:
+            self.bias: Parameter | None = Parameter(torch.empty(out_features))
+        else:
+            self.bias = None
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in = self.weight.size(1)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input: Tensor) -> Tensor:
+        return nn.functional.linear(input, self.weight, self.bias)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"bias={self.bias is not None}"
+        )

sparse_layers/layers/custom_mlp.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+
+from sparse_layers.layers.custom_linear import CustomLinear
+
+
+class CustomMLP(nn.Module):
+    """MLP that uses CustomLinear layers as drop-in replacements for nn.Linear."""
+
+    def __init__(self, input_dim: int, hidden_dims: Sequence[int], output_dim: int) -> None:
+        super().__init__()
+
+        self._validate_dimensions(input_dim, hidden_dims, output_dim)
+
+        hidden_dims = list(hidden_dims)
+        layer_dims = [input_dim, *hidden_dims, output_dim]
+
+        modules: list[nn.Module] = []
+        for idx in range(len(layer_dims) - 1):
+            in_features = layer_dims[idx]
+            out_features = layer_dims[idx + 1]
+
+            linear = CustomLinear(in_features, out_features)
+            modules.append(linear)
+
+            if idx < len(layer_dims) - 2:
+                modules.append(nn.ReLU())
+
+        self.network = nn.Sequential(*modules)
+
+    @staticmethod
+    def _validate_dimensions(input_dim: int, hidden_dims: Sequence[int], output_dim: int) -> None:
+        if input_dim <= 0:
+            raise ValueError("input_dim must be a positive integer")
+
+        if not hidden_dims:
+            raise ValueError("hidden_dims must contain at least one positive integer")
+
+        if any(dim <= 0 for dim in hidden_dims):
+            raise ValueError("hidden_dims must contain only positive integers")
+
+        if output_dim <= 0:
+            raise ValueError("output_dim must be a positive integer")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() != 2:
+            raise ValueError("CustomMLP expects a 2D input tensor of shape (batch, features)")
+
+        return self.network(x)
+