eae-runtime 0.1.0__py3-none-any.whl

eae_runtime/contrib/transformer_blocks.py

@@ -0,0 +1,259 @@
+"""
+eae_runtime.contrib.transformer_blocks
+=======================================
+
+Reference, current-generation (2024-2026 era) Transformer building blocks
+that are known-good EAE Runtime blocks: each `nn.Module` here has a
+`forward(x) -> Tensor` that is a pure function of `x` and its own
+parameters/buffers - exactly the contract `ReconstructionEngine` relies on
+(see `eae_runtime.blocks.EAEBlock`).
+
+These are deliberately **not** shipped as part of the runtime core. The
+runtime is architecture-agnostic by design (see README "Non-goals"); this
+module exists to:
+
+  1. Prove, concretely, that the runtime handles modern Transformer
+     internals: RMSNorm, rotary position embeddings (RoPE), grouped-query
+     attention (GQA) via `F.scaled_dot_product_attention` (which
+     transparently dispatches to a fused flash-attention / memory-
+     efficient kernel on supported hardware), and SwiGLU feed-forward
+     layers - the components of e.g. LLaMA / Mistral / Qwen-style models.
+  2. Demonstrate the idiomatic pattern for handling information a block
+     needs besides `x` (a causal mask, RoPE tables) *without* breaking the
+     runtime's single-tensor forward contract: precompute it once in
+     `__init__` / `register_buffer`, never as an extra forward() argument.
+  3. Give you a working, forkable starting point. Copy this file into your
+     own project and edit it; the runtime core never needs to change to
+     support your architecture.
+
+Usage::
+
+    from eae_runtime.contrib import PreNormTransformerBlock
+
+    blocks = nn.ModuleList([
+        PreNormTransformerBlock(dim=512, num_heads=8, num_kv_heads=2)
+        for _ in range(depth)
+    ])
+    runtime = EAERuntime(blocks, optimizer, config)
+
+Deliberately not covered here (fork and extend instead):
+  * Cross-attention / encoder-decoder blocks - attention over a second,
+    externally supplied sequence isn't expressible as a pure function of a
+    single `x`. Attach the memory tensor as an instance attribute the same
+    way RoPE tables are attached below, or override `forward` to unpack a
+    tensor you've concatenated/stacked upstream.
+  * KV-cache incremental decoding - this reference targets full-sequence
+    training; inference-time caching belongs in the caller's generation
+    loop, not the runtime.
+  * Mixture-of-Experts routing.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..blocks import EAEBlock
+
+__all__ = [
+    "RMSNorm",
+    "RotaryPositionalEmbedding",
+    "apply_rotary",
+    "CausalSelfAttention",
+    "SwiGLU",
+    "PreNormTransformerBlock",
+]
+
+
+class RMSNorm(nn.Module):
+    """Root-Mean-Square LayerNorm (Zhang & Sennrich, 2019).
+
+    The default normalization in LLaMA / Mistral / Gemma-style models:
+    cheaper than LayerNorm (no mean-centering, no bias term) and matches
+    or beats it empirically at scale.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        xf = x.float()
+        rms = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return (xf * rms).to(dtype) * self.weight
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    """Rotary position embeddings (Su et al., 2021 / RoFormer) - the de
+    facto standard positional scheme in modern decoder-only Transformers.
+
+    Precomputes cos/sin tables as non-persistent buffers (so they move
+    with `.to(device)` but are never saved to a checkpoint or trained) up
+    to `max_seq_len`; `forward(seq_len)` slices out the prefix needed for
+    the current sequence.
+    """
+
+    cos_cached: torch.Tensor
+    sin_cached: torch.Tensor
+
+    def __init__(self, head_dim: int, max_seq_len: int = 4096, base: float = 10000.0):
+        super().__init__()
+        if head_dim % 2 != 0:
+            raise ValueError(f"RoPE requires an even head_dim, got {head_dim}")
+        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        t = torch.arange(max_seq_len).float()
+        freqs = torch.outer(t, inv_freq)  # (max_seq_len, head_dim // 2)
+        self.register_buffer("cos_cached", freqs.cos(), persistent=False)
+        self.register_buffer("sin_cached", freqs.sin(), persistent=False)
+
+    def forward(self, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
+
+
+def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply RoPE rotation to `x` of shape (batch, heads, seq, head_dim)
+    given per-position `cos`/`sin` tables of shape (seq, head_dim // 2)."""
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    cos = cos[None, None, :, :].to(x.dtype)
+    sin = sin[None, None, :, :].to(x.dtype)
+    rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+    return rotated.flatten(-2)
+
+
+class CausalSelfAttention(nn.Module):
+    """Causal (optionally grouped-query) self-attention built on
+    `torch.nn.functional.scaled_dot_product_attention`, which dispatches
+    to a fused flash-attention / memory-efficient kernel on supported
+    hardware with no extra code here - the current recommended way to
+    write attention in PyTorch (2.x+), rather than a hand-rolled
+    `softmax(QK^T / sqrt(d)) @ V`.
+
+    Set `num_kv_heads < num_heads` for grouped-query attention (GQA, as in
+    LLaMA-2 70B / Mistral) to shrink the KV cache at inference time;
+    `num_kv_heads=1` is multi-query attention (MQA). Leave it as `None`
+    for ordinary multi-head attention.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: Optional[int] = None,
+        max_seq_len: int = 4096,
+        dropout: float = 0.0,
+        rope: bool = True,
+    ):
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})")
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        if num_heads % self.num_kv_heads != 0:
+            raise ValueError("num_heads must be divisible by num_kv_heads for GQA")
+        self.head_dim = dim // num_heads
+        self.dropout = dropout
+
+        self.q_proj = nn.Linear(dim, num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.out_proj = nn.Linear(num_heads * self.head_dim, dim, bias=False)
+
+        self.rope = RotaryPositionalEmbedding(self.head_dim, max_seq_len) if rope else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, seq, _ = x.shape
+        q = self.q_proj(x).view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        if self.rope is not None:
+            cos, sin = self.rope(seq)
+            q = apply_rotary(q, cos, sin)
+            k = apply_rotary(k, cos, sin)
+
+        if self.num_kv_heads != self.num_heads:
+            repeat = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+
+        out = F.scaled_dot_product_attention(
+            q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0
+        )
+        out = out.transpose(1, 2).contiguous().view(batch, seq, self.num_heads * self.head_dim)
+        return self.out_proj(out)
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU feed-forward (Shazeer, 2020) - the standard modern
+    replacement for a plain Linear -> GELU -> Linear FFN, used in
+    LLaMA / PaLM / Mistral. `hidden_dim` defaults to ~8/3 * dim rounded up
+    to a multiple of `multiple_of` for hardware-friendly matmul shapes.
+    """
+
+    def __init__(self, dim: int, hidden_dim: Optional[int] = None, multiple_of: int = 256):
+        super().__init__()
+        if hidden_dim is None:
+            hidden_dim = int(2 * (4 * dim) / 3)
+            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class PreNormTransformerBlock(EAEBlock):
+    """One decoder-only Transformer block in the current dominant
+    architecture family (LLaMA/Mistral/Qwen-style): pre-norm residual
+    stream, RMSNorm, causal (optionally grouped-query) attention with
+    RoPE, RMSNorm, SwiGLU feed-forward.
+
+    Stack `depth` of these in an `nn.ModuleList` and hand that straight to
+    `EAERuntime` - `BlockDecomposer` treats each block as one
+    reconstructable unit::
+
+        blocks = nn.ModuleList([
+            PreNormTransformerBlock(dim=512, num_heads=8, num_kv_heads=2)
+            for _ in range(depth)
+        ])
+        runtime = EAERuntime(blocks, optimizer, config)
+
+    Every sub-module here (`RMSNorm`, `CausalSelfAttention`, `SwiGLU`) is
+    independently reusable if your architecture only needs one piece, and
+    every EAE pass (`ClipPass`, `Int8QuantizationPass`, a custom
+    compression/synthetic-gradient pass, ...) applies to the adjoint
+    flowing out of this block exactly as it would for a plain `nn.Linear`.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: Optional[int] = None,
+        max_seq_len: int = 4096,
+        ffn_hidden_dim: Optional[int] = None,
+        dropout: float = 0.0,
+        norm_eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.attn_norm = RMSNorm(dim, eps=norm_eps)
+        self.attn = CausalSelfAttention(
+            dim,
+            num_heads,
+            num_kv_heads=num_kv_heads,
+            max_seq_len=max_seq_len,
+            dropout=dropout,
+        )
+        self.ffn_norm = RMSNorm(dim, eps=norm_eps)
+        self.ffn = SwiGLU(dim, hidden_dim=ffn_hidden_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.attn_norm(x))
+        x = x + self.ffn(self.ffn_norm(x))
+        return x

eae_runtime/events.py

@@ -0,0 +1,84 @@
+"""
+Event system: everything important that happens inside the runtime emits an
+Event. Researchers should never need to edit runtime code just to observe
+execution - they subscribe to the EventBus instead.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional
+
+
+class EventType:
+    FORWARD_STARTED = "ForwardStarted"
+    FORWARD_FINISHED = "ForwardFinished"
+    BLOCK_RECONSTRUCTED = "BlockReconstructed"
+    ADJOINT_CREATED = "AdjointCreated"
+    ADJOINT_MODIFIED = "AdjointModified"
+    MEMORY_ALLOCATED = "MemoryAllocated"
+    MEMORY_RELEASED = "MemoryReleased"
+    SCHEDULER_STEP = "SchedulerStep"
+    KERNEL_LAUNCH = "KernelLaunch"
+    COMMUNICATION_START = "CommunicationStart"
+    COMMUNICATION_FINISH = "CommunicationFinish"
+    TRAIN_STEP_STARTED = "TrainStepStarted"
+    TRAIN_STEP_FINISHED = "TrainStepFinished"
+    PASS_APPLIED = "PassApplied"
+
+
+@dataclass
+class Event:
+    type: str
+    payload: Dict[str, Any] = field(default_factory=dict)
+    timestamp: float = field(default_factory=time.time)
+
+    def __repr__(self) -> str:  # pragma: no cover - cosmetic
+        return f"Event({self.type}, {self.payload})"
+
+
+class EventBus:
+    """A minimal, synchronous pub/sub bus.
+
+    Subscribers are plain callables `fn(event: Event) -> None`. Subscribing
+    to `"*"` receives every event, regardless of type.
+    """
+
+    def __init__(self) -> None:
+        self._subscribers: Dict[str, List[Callable[[Event], None]]] = {}
+        self._log: List[Event] = []
+        self._recording = False
+
+    def subscribe(self, event_type: str, callback: Callable[[Event], None]) -> None:
+        self._subscribers.setdefault(event_type, []).append(callback)
+
+    def unsubscribe(self, event_type: str, callback: Callable[[Event], None]) -> None:
+        if event_type in self._subscribers and callback in self._subscribers[event_type]:
+            self._subscribers[event_type].remove(callback)
+
+    def emit(self, event_type: str, **payload: Any) -> Event:
+        event = Event(type=event_type, payload=payload)
+        if self._recording:
+            self._log.append(event)
+        for callback in self._subscribers.get(event_type, []):
+            callback(event)
+        for callback in self._subscribers.get("*", []):
+            callback(event)
+        return event
+
+    # -- recording / introspection, handy for tests -------------------- #
+    def start_recording(self) -> None:
+        self._recording = True
+        self._log = []
+
+    def stop_recording(self) -> List[Event]:
+        self._recording = False
+        return self._log
+
+    @property
+    def log(self) -> List[Event]:
+        return list(self._log)
+
+    def events_of_type(self, event_type: str) -> List[Event]:
+        return [e for e in self._log if e.type == event_type]

eae_runtime/forward_executor.py

@@ -0,0 +1,35 @@
+"""
+Forward Executor: detached forward, boundary extraction, boundary storage.
+No autograd graph is kept alive - this is what makes the runtime
+memory-cheap relative to a naive PyTorch backward.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+
+from .boundary_store import BoundaryStore
+from .events import EventBus, EventType
+
+
+class ForwardExecutor:
+    def __init__(self, event_bus: Optional[EventBus] = None, compute_dtype: Optional[torch.dtype] = None):
+        self.event_bus = event_bus or EventBus()
+        self.compute_dtype = compute_dtype
+
+    @torch.no_grad()
+    def run(self, blocks: List[nn.Module], x0: torch.Tensor, boundary_store: BoundaryStore) -> torch.Tensor:
+        self.event_bus.emit(EventType.FORWARD_STARTED, num_blocks=len(blocks))
+
+        x = x0
+        boundary_store.put(0, x)
+        for i, block in enumerate(blocks):
+            compute_x = x.to(self.compute_dtype) if self.compute_dtype is not None else x
+            x = block(compute_x)
+            boundary_store.put(i + 1, x)
+
+        self.event_bus.emit(EventType.FORWARD_FINISHED, num_boundaries=len(boundary_store))
+        return x

eae_runtime/logging_utils.py

@@ -0,0 +1,41 @@
+"""
+Structured logging: every runtime event should be observable. This wires
+the EventBus to Python's logging module using a JSON-ish structured
+formatter, so events can be piped to any log aggregator.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Optional
+
+from .events import Event, EventBus
+
+_LOGGER_NAME = "eae_runtime"
+
+
+def get_logger(level: str = "WARNING") -> logging.Logger:
+    logger = logging.getLogger(_LOGGER_NAME)
+    if not logger.handlers:
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter("%(message)s"))
+        logger.addHandler(handler)
+    logger.setLevel(getattr(logging, level.upper(), logging.WARNING))
+    return logger
+
+
+def attach_structured_logging(event_bus: EventBus, level: str = "WARNING", logger: Optional[logging.Logger] = None) -> None:
+    """Subscribe a structured JSON logger to every event on `event_bus`.
+
+    `level` sets the logger's threshold (e.g. "WARNING" hides events by
+    default). Individual event records are always emitted at DEBUG
+    severity, so setting `level="DEBUG"` is what makes them visible.
+    """
+    logger = logger or get_logger(level)
+
+    def _log(event: Event) -> None:
+        record = {"event": event.type, "timestamp": event.timestamp, **event.payload}
+        logger.debug(json.dumps(record, default=str))
+
+    event_bus.subscribe("*", _log)

eae_runtime/memory.py

@@ -0,0 +1,175 @@
+"""
+Memory Manager: owns temporary tensors, scratch buffers, reconstruction
+workspace. Provides request()/release()/reuse() and exposes runtime
+statistics. The allocator should minimize allocations, fragmentation and
+synchronization.
+
+Two policies are provided out of the box:
+  * PoolMemoryPolicy  - buckets free tensors by (shape, dtype, device) and
+                         reuses them, avoiding repeated cudaMalloc/alloc churn.
+  * NullMemoryPolicy   - always allocates fresh / frees immediately; useful
+                         as a correctness baseline and for leak tests.
+
+Users can implement their own MemoryPolicy (Memory Policy API) without
+touching reconstruction logic.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Dict, Optional, Tuple
+
+import torch
+
+from .events import EventBus, EventType
+
+PoolKey = Tuple[Tuple[int, ...], torch.dtype, str]
+
+
+class MemoryPolicy:
+    """Base class for pluggable memory policies."""
+
+    def allocate(self, shape, dtype, device) -> torch.Tensor:
+        raise NotImplementedError
+
+    def free(self, tensor: torch.Tensor) -> None:
+        raise NotImplementedError
+
+    def stats(self) -> Dict[str, int]:
+        raise NotImplementedError
+
+    def reset(self) -> None:
+        raise NotImplementedError
+
+
+class NullMemoryPolicy(MemoryPolicy):
+    """No pooling: every request is a fresh allocation, every release is a
+    real deallocation (subject to Python/CUDA garbage collection)."""
+
+    def __init__(self):
+        self._active = 0
+        self._peak = 0
+        self._allocations = 0
+
+    def allocate(self, shape, dtype, device) -> torch.Tensor:
+        t = torch.empty(shape, dtype=dtype, device=device)
+        self._active += 1
+        self._allocations += 1
+        self._peak = max(self._peak, self._active)
+        return t
+
+    def free(self, tensor: torch.Tensor) -> None:
+        self._active = max(0, self._active - 1)
+
+    def stats(self) -> Dict[str, int]:
+        return {"active": self._active, "peak": self._peak, "allocations": self._allocations, "pooled": 0}
+
+    def reset(self) -> None:
+        self._active = 0
+        self._peak = 0
+        self._allocations = 0
+
+
+class PoolMemoryPolicy(MemoryPolicy):
+    """Reuses freed tensors of matching (shape, dtype, device) instead of
+    reallocating. Thread-safe (relevant for the AsyncScheduler)."""
+
+    def __init__(self, max_pool_per_key: int = 8):
+        self._pool: Dict[PoolKey, list] = {}
+        self._lock = threading.Lock()
+        self._active = 0
+        self._peak = 0
+        self._allocations = 0
+        self._reuses = 0
+        self._max_pool_per_key = max_pool_per_key
+
+    @staticmethod
+    def _key(shape, dtype, device) -> PoolKey:
+        return (tuple(shape), dtype, str(device))
+
+    def allocate(self, shape, dtype, device) -> torch.Tensor:
+        key = self._key(shape, dtype, device)
+        with self._lock:
+            bucket = self._pool.get(key)
+            if bucket:
+                t = bucket.pop()
+                self._reuses += 1
+                self._active += 1
+                self._peak = max(self._peak, self._active)
+                return t
+            self._allocations += 1
+            self._active += 1
+            self._peak = max(self._peak, self._active)
+        return torch.empty(shape, dtype=dtype, device=device)
+
+    def free(self, tensor: torch.Tensor) -> None:
+        key = self._key(tensor.shape, tensor.dtype, tensor.device)
+        with self._lock:
+            self._active = max(0, self._active - 1)
+            bucket = self._pool.setdefault(key, [])
+            if len(bucket) < self._max_pool_per_key:
+                bucket.append(tensor.detach())
+
+    def stats(self) -> Dict[str, int]:
+        pooled = sum(len(v) for v in self._pool.values())
+        return {
+            "active": self._active,
+            "peak": self._peak,
+            "allocations": self._allocations,
+            "reuses": self._reuses,
+            "pooled": pooled,
+        }
+
+    def reset(self) -> None:
+        with self._lock:
+            self._pool.clear()
+            self._active = 0
+            self._peak = 0
+            self._allocations = 0
+            self._reuses = 0
+
+
+_POLICIES = {
+    "pool": PoolMemoryPolicy,
+    "none": NullMemoryPolicy,
+}
+
+
+class MemoryManager:
+    """Facade used by the rest of the runtime. Wraps a MemoryPolicy and adds
+    event emission + a scratch-buffer convenience API."""
+
+    def __init__(self, policy="pool", event_bus: Optional[EventBus] = None):
+        if isinstance(policy, str):
+            if policy not in _POLICIES:
+                raise ValueError(f"Unknown memory policy '{policy}', choices: {list(_POLICIES)}")
+            policy = _POLICIES[policy]()
+        elif isinstance(policy, type) and issubclass(policy, MemoryPolicy):
+            policy = policy()
+        self.policy: MemoryPolicy = policy
+        self.event_bus = event_bus or EventBus()
+
+    def request(self, shape, dtype=torch.float32, device="cpu") -> torch.Tensor:
+        t = self.policy.allocate(shape, dtype, device)
+        self.event_bus.emit(EventType.MEMORY_ALLOCATED, shape=tuple(shape), dtype=str(dtype), device=str(device))
+        return t
+
+    def release(self, tensor: torch.Tensor) -> None:
+        if tensor is None:
+            return
+        shape = tuple(tensor.shape)
+        self.policy.free(tensor)
+        self.event_bus.emit(EventType.MEMORY_RELEASED, shape=shape)
+
+    def reuse(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Zero out and hand back a tensor for reuse without a full
+        free+allocate round trip."""
+        with torch.no_grad():
+            tensor.zero_()
+        return tensor
+
+    def stats(self) -> Dict[str, int]:
+        return self.policy.stats()
+
+    def reset(self) -> None:
+        self.policy.reset()

eae_runtime/passes/__init__.py

@@ -0,0 +1,18 @@
+from .base import EAEPass
+from .clip import ClipPass
+from .quantize import FP16Pass, FP8Pass, Int8QuantizationPass
+from .synthetic_gradient import SyntheticGradientPass
+from .regularization import RegularizationPass, GaussianNoisePass
+from .logging_pass import LoggingPass
+
+__all__ = [
+    "EAEPass",
+    "ClipPass",
+    "FP16Pass",
+    "FP8Pass",
+    "Int8QuantizationPass",
+    "SyntheticGradientPass",
+    "RegularizationPass",
+    "GaussianNoisePass",
+    "LoggingPass",
+]

eae_runtime/passes/base.py

@@ -0,0 +1,31 @@
+"""
+Pass API: every adjoint passes through Pass1 -> Pass2 -> ... -> Passn before
+reaching the next block. Passes are plugins; the runtime never knows what
+they do.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from ..adjoint import AdjointState
+
+
+class EAEPass:
+    """Base class for a single stage in the Adjoint Pipeline.
+
+    Subclasses implement `process(adjoint, context) -> AdjointState`.
+    `context` is a plain dict with useful ambient info (current block name,
+    layer_id, training step, event_bus, etc.) that a pass may read but
+    should not mutate destructively.
+    """
+
+    name: str = "EAEPass"
+
+    def process(self, adjoint: AdjointState, context: Optional[dict] = None) -> AdjointState:
+        raise NotImplementedError
+
+    def __call__(self, adjoint: AdjointState, context: Optional[dict] = None) -> AdjointState:
+        result = self.process(adjoint, context or {})
+        result.record(self.name)
+        return result

eae_runtime/passes/clip.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import torch
+
+from ..adjoint import AdjointState
+from .base import EAEPass
+
+
+class ClipPass(EAEPass):
+    """Clips the adjoint tensor's norm to `max_norm`, mirroring
+    torch.nn.utils.clip_grad_norm_ but operating on a single adjoint in the
+    pipeline rather than a full parameter list."""
+
+    name = "ClipPass"
+
+    def __init__(self, max_norm: float = 1.0):
+        self.max_norm = max_norm
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        new = adjoint.clone()
+        norm = torch.linalg.vector_norm(new.tensor.float())
+        if norm > self.max_norm:
+            scale = self.max_norm / (norm + 1e-6)
+            new.tensor = new.tensor * scale.to(new.tensor.dtype)
+        return new