diff --git a/README.md b/README.md index 4f31418..b01d4d7 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,23 @@ print(logits.shape, float(loss)) `hidden_size` must equal `num_attention_heads * head_dim`. Inputs longer than `block_size` fail explicitly. `crop_block_size(n)` only decreases that limit. +## Allocation-aware update + +The default geometric update uses `torch.addcmul(x, k, delta)` rather than +materializing `delta * k` as a Python-level intermediate before addition. An +explicit reference implementation remains in `model.geometric`, with output and +backward-gradient parity covered by tests. + +Measure the two paths on the target hardware: + +```bash +python benchmarks/geometric_update.py --device cpu --iterations 200 +``` + +On CUDA the script also reports peak allocated memory. The measurable acceptance +target is equal outputs/gradients and no increase in peak allocation; this +repository does not claim a hardware-independent latency improvement. + ## Research status and benchmark policy The spectral statement above is an exact fixed-`k`, fixed-`beta` property. diff --git a/benchmarks/geometric_update.py b/benchmarks/geometric_update.py new file mode 100644 index 0000000..8357804 --- /dev/null +++ b/benchmarks/geometric_update.py @@ -0,0 +1,65 @@ +"""Profile literal and allocation-aware geometric residual updates.""" + +import argparse +import time + +import torch + +from model.geometric import geometric_update, reference_geometric_update + + +def synchronize(device): + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def profile(function, values, iterations): + device = values[0].device + with torch.inference_mode(): + for _ in range(20): + function(*values) + synchronize(device) + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + + started = time.perf_counter() + with torch.inference_mode(): + for _ in range(iterations): + function(*values) + synchronize(device) + latency = (time.perf_counter() - started) * 1000 / iterations + peak = ( + torch.cuda.max_memory_allocated(device) / 1024**2 + if device.type == "cuda" + else None + ) + return latency, peak + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--sequence-length", type=int, default=1024) + parser.add_argument("--hidden-size", type=int, default=768) + parser.add_argument("--iterations", type=int, default=200) + args = parser.parse_args() + + device = torch.device(args.device) + shape = (args.batch_size, args.sequence_length, args.hidden_size) + values = ( + torch.randn(shape, device=device), + torch.randn(shape, device=device), + torch.randn(shape[:-1] + (1,), device=device), + ) + for name, function in ( + ("literal", reference_geometric_update), + ("addcmul", geometric_update), + ): + latency, peak = profile(function, values, args.iterations) + peak_text = f", peak_allocated={peak:.1f} MiB" if peak is not None else "" + print(f"{name}: {latency:.3f} ms/iteration{peak_text}") + + +if __name__ == "__main__": + main() diff --git a/model/geometric.py b/model/geometric.py new file mode 100644 index 0000000..6d68193 --- /dev/null +++ b/model/geometric.py @@ -0,0 +1,19 @@ +"""Auditable implementations of the OUROBOROS elementwise update.""" + +import torch + + +def reference_geometric_update( + x: torch.Tensor, k: torch.Tensor, delta: torch.Tensor +) -> torch.Tensor: + """Literal reference: x + delta * k.""" + + return x + delta * k + + +def geometric_update( + x: torch.Tensor, k: torch.Tensor, delta: torch.Tensor +) -> torch.Tensor: + """Use addcmul so the update is expressed as one elementwise operator.""" + + return torch.addcmul(x, k, delta) diff --git a/model/ouroboros.py b/model/ouroboros.py index c6ce542..e293481 100644 --- a/model/ouroboros.py +++ b/model/ouroboros.py @@ -12,6 +12,7 @@ from .init_utils import init_gpt_weights from .pydantic_config import validate_pretrained_config_kwargs from .rotary import Rotary, apply_rotary_emb +from .geometric import geometric_update def _logit(p: float) -> float: @@ -163,8 +164,7 @@ def forward(self, x: torch.Tensor, *, k_in: torch.Tensor, context: torch.Tensor) v = torch.sigmoid(v) * self.v_sigmoid_scale delta = (beta * (v - proj)).to(dtype=x.dtype) # (B, T, 1) - update = delta * k - return x + update + return geometric_update(x, k, delta) class OuroborosBlock(nn.Module): diff --git a/tests/test_geometric.py b/tests/test_geometric.py new file mode 100644 index 0000000..eec9f8a --- /dev/null +++ b/tests/test_geometric.py @@ -0,0 +1,25 @@ +import torch + +from model.geometric import geometric_update, reference_geometric_update + + +def run_and_grad(function, values): + x, k, delta = [value.detach().clone().requires_grad_(True) for value in values] + output = function(x, k, delta) + gradients = torch.autograd.grad(output.square().mean(), (x, k, delta)) + return output, gradients + + +def test_addcmul_matches_reference_output_and_gradients(): + torch.manual_seed(11) + values = ( + torch.randn(2, 7, 16), + torch.randn(2, 7, 16), + torch.randn(2, 7, 1), + ) + reference_output, reference_gradients = run_and_grad(reference_geometric_update, values) + optimized_output, optimized_gradients = run_and_grad(geometric_update, values) + + torch.testing.assert_close(optimized_output, reference_output) + for optimized, reference in zip(optimized_gradients, reference_gradients): + torch.testing.assert_close(optimized, reference)