scaffold: model (OLMo2-ish + ZBP partition), trainer (DDP/config), data shards, bench

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe

1 files changed, 84 insertions, 0 deletions

diff --git a/src/zbp_scaling/zbp/fa.py b/src/zbp_scaling/zbp/fa.py
new file mode 100644
index 0000000..cd8830b
--- /dev/null
+++ b/src/zbp_scaling/zbp/fa.py
@@ -0,0 +1,84 @@
+"""Standard (layer-wise) feedback alignment: Linear / Conv2d layers whose backward uses a fixed random
+weight B in place of W^T (Lillicrap et al. 2016).  Nonlinearity derivatives are the true ones, taken at
+the forward activations, as in the original algorithm.  Use `apply_fa(module)` to convert all Linear/Conv2d
+layers of a block in place."""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class _FALinearFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W, b, B):
+        ctx.save_for_backward(x, B)
+        ctx.has_bias = b is not None
+        y = x.matmul(W.t())
+        return y + b if b is not None else y
+
+    @staticmethod
+    def backward(ctx, gy):
+        x, B = ctx.saved_tensors
+        gx = gy.matmul(B)                                   # random feedback instead of W
+        gW = gy.reshape(-1, gy.shape[-1]).t().matmul(x.reshape(-1, x.shape[-1]))
+        gb = gy.reshape(-1, gy.shape[-1]).sum(0) if ctx.has_bias else None
+        return gx, gW, gb, None
+
+
+class _FAConvFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W, b, B, stride, padding):
+        ctx.save_for_backward(x, W, B)
+        ctx.stride, ctx.padding = stride, padding
+        ctx.has_bias = b is not None
+        return F.conv2d(x, W, b, stride, padding)
+
+    @staticmethod
+    def backward(ctx, gy):
+        x, W, B = ctx.saved_tensors
+        gx = torch.nn.grad.conv2d_input(x.shape, B, gy, ctx.stride, ctx.padding)     # random feedback kernels
+        gW = torch.nn.grad.conv2d_weight(x, W.shape, gy, ctx.stride, ctx.padding)
+        gb = gy.sum((0, 2, 3)) if ctx.has_bias else None
+        return gx, gW, gb, None, None, None
+
+
+class FALinear(nn.Linear):
+    def __init__(self, *a, **k):
+        super().__init__(*a, **k)
+        self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(self.in_features))
+
+    def forward(self, x):
+        return _FALinearFn.apply(x, self.weight, self.bias, self.B)
+
+
+class FAConv2d(nn.Conv2d):
+    def __init__(self, *a, **k):
+        super().__init__(*a, **k)
+        fan_in = self.in_channels * self.kernel_size[0] * self.kernel_size[1]
+        self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(fan_in))
+
+    def forward(self, x):
+        return _FAConvFn.apply(x, self.weight, self.bias, self.B, self.stride, self.padding)
+
+
+def apply_fa(module):
+    """Replace every nn.Linear / nn.Conv2d inside `module` (recursively) by its FA variant, keeping weights."""
+    for name, child in list(module.named_children()):
+        if type(child) is nn.Linear:
+            new = FALinear(child.in_features, child.out_features, bias=child.bias is not None)
+            new.weight = child.weight
+            if child.bias is not None:
+                new.bias = child.bias
+            new.B = new.B.to(child.weight.device)
+            setattr(module, name, new)
+        elif type(child) is nn.Conv2d:
+            new = FAConv2d(child.in_channels, child.out_channels, child.kernel_size, child.stride, child.padding,
+                           bias=child.bias is not None)
+            new.weight = child.weight
+            if child.bias is not None:
+                new.bias = child.bias
+            new.B = new.B.to(child.weight.device)
+            setattr(module, name, new)
+        else:
+            apply_fa(child)
+    return module