eae-runtime 0.1.0__py3-none-any.whl

eae_runtime/__init__.py

@@ -0,0 +1,80 @@
+"""
+EAE Runtime: a standalone, programmable, block-local reverse-mode execution
+engine that sits *above* PyTorch and orchestrates many local autograd
+computations, instead of replacing PyTorch's autodiff engine.
+
+Quick start:
+
+    import torch.nn as nn
+    from eae_runtime import EAERuntime, RuntimeConfig
+    from eae_runtime.passes import ClipPass, GaussianNoisePass
+
+    model = nn.Sequential(
+        nn.Linear(64, 64), nn.ReLU(),
+        nn.Linear(64, 64), nn.ReLU(),
+        nn.Linear(64, 10),
+    )
+    optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)
+    config = RuntimeConfig(
+        scheduler="sequential",
+        memory="pool",
+        backend="auto",
+        passes=[ClipPass(max_norm=1.0)],
+    )
+    runtime = EAERuntime(model, optimizer, config)
+    loss = runtime.train_step(x, lambda out: nn.functional.mse_loss(out, target))
+"""
+
+from .adjoint import AdjointState
+from .backend import BackendManager
+from .blocks import BlockDecomposer, EAEBlock
+from .boundary_store import BoundaryStore
+from .config import RuntimeConfig
+from .events import Event, EventBus, EventType
+from .forward_executor import ForwardExecutor
+from .memory import MemoryManager, MemoryPolicy, NullMemoryPolicy, PoolMemoryPolicy
+from .pipeline import AdjointPipeline, PassManager
+from .profiler import Profiler
+from .reconstruction import ReconstructionEngine
+from .runtime import EAERuntime
+from .schedulers import (
+    AsyncScheduler,
+    BaseScheduler,
+    DistributedScheduler,
+    PipelineScheduler,
+    SequentialScheduler,
+    build_scheduler,
+)
+
+try:
+    from ._version import version as __version__
+except ImportError:  # pragma: no cover - only hit in a non-built source tree
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "AdjointState",
+    "BackendManager",
+    "BlockDecomposer",
+    "EAEBlock",
+    "BoundaryStore",
+    "RuntimeConfig",
+    "Event",
+    "EventBus",
+    "EventType",
+    "ForwardExecutor",
+    "MemoryManager",
+    "MemoryPolicy",
+    "NullMemoryPolicy",
+    "PoolMemoryPolicy",
+    "AdjointPipeline",
+    "PassManager",
+    "Profiler",
+    "ReconstructionEngine",
+    "EAERuntime",
+    "AsyncScheduler",
+    "BaseScheduler",
+    "DistributedScheduler",
+    "PipelineScheduler",
+    "SequentialScheduler",
+    "build_scheduler",
+]

eae_runtime/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = 'g336fa5cb8'

eae_runtime/adjoint.py

@@ -0,0 +1,128 @@
+"""
+AdjointState: the first-class runtime object that flows through the EAE runtime.
+
+Design principle (see spec, "Explicit Adjoint States"):
+    The runtime never passes raw tensors internally. Everything flows through
+    AdjointState. Passes, schedulers and the reconstruction engine all read
+    and write AdjointState objects, never bare torch.Tensors.
+"""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import torch
+
+
+@dataclass
+class AdjointState:
+    """A gradient ("adjoint") flowing backward through the runtime, plus metadata.
+
+    Attributes:
+        tensor: the actual gradient tensor (d Loss / d activation).
+        layer_id: index of the block boundary this adjoint sits at (L, L-1, ..., 0).
+        block: human readable name of the block that produced/consumes this adjoint.
+        dtype: logical dtype this adjoint should be materialized at.
+        device: logical device this adjoint should be materialized at.
+        metadata: free-form dict for passes to stash information.
+        history: ordered list of pass names that have touched this adjoint,
+                 for debugging / provenance / profiling.
+    """
+
+    tensor: torch.Tensor
+    layer_id: int
+    block: str = ""
+    dtype: Optional[torch.dtype] = None
+    device: Optional[torch.device] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    history: List[str] = field(default_factory=list)
+
+    def __post_init__(self):
+        if self.dtype is None:
+            self.dtype = self.tensor.dtype
+        if self.device is None:
+            self.device = self.tensor.device
+
+    # ------------------------------------------------------------------ #
+    # Rich adjoint API - users should almost never touch .tensor directly
+    # ------------------------------------------------------------------ #
+    def norm(self, p: float = 2.0) -> torch.Tensor:
+        """L-p norm of the underlying gradient tensor."""
+        return torch.linalg.vector_norm(self.tensor.float(), ord=p)
+
+    def statistics(self) -> Dict[str, float]:
+        """Cheap summary statistics, useful for logging/profiling passes."""
+        t = self.tensor.detach().float()
+        return {
+            "mean": t.mean().item(),
+            "std": t.std(unbiased=False).item() if t.numel() > 1 else 0.0,
+            "min": t.min().item(),
+            "max": t.max().item(),
+            "norm": torch.linalg.vector_norm(t).item(),
+            "numel": t.numel(),
+            "has_nan": bool(torch.isnan(t).any().item()),
+            "has_inf": bool(torch.isinf(t).any().item()),
+        }
+
+    def quantize(self, dtype: torch.dtype = torch.float16) -> "AdjointState":
+        """Return a new AdjointState with the tensor cast down to `dtype`."""
+        new = self.clone()
+        new.tensor = new.tensor.to(dtype)
+        new.dtype = dtype
+        return new
+
+    def dequantize(self, dtype: torch.dtype = torch.float32) -> "AdjointState":
+        new = self.clone()
+        new.tensor = new.tensor.to(dtype)
+        new.dtype = dtype
+        return new
+
+    def compress(self, ratio: float = 0.5) -> "AdjointState":
+        """Simple magnitude-pruning compression pass, kept out of the runtime
+        core on purpose - this is exactly the kind of thing a real Pass would
+        do; it lives here only as a convenience primitive on the adjoint."""
+        new = self.clone()
+        flat = new.tensor.flatten()
+        k = max(1, int(flat.numel() * (1.0 - ratio)))
+        if k < flat.numel():
+            threshold = flat.abs().kthvalue(flat.numel() - k + 1).values
+            mask = flat.abs() >= threshold
+            flat = flat * mask
+            new.tensor = flat.view_as(new.tensor)
+        return new
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> "AdjointState":
+        new = self.clone()
+        new.tensor = new.tensor.to(device=device or new.device, dtype=dtype or new.dtype)
+        new.device = new.tensor.device
+        new.dtype = new.tensor.dtype
+        return new
+
+    def clone(self) -> "AdjointState":
+        return AdjointState(
+            tensor=self.tensor,
+            layer_id=self.layer_id,
+            block=self.block,
+            dtype=self.dtype,
+            device=self.device,
+            metadata=dict(self.metadata),
+            history=list(self.history),
+        )
+
+    def detach(self) -> "AdjointState":
+        new = self.clone()
+        new.tensor = new.tensor.detach()
+        return new
+
+    def record(self, pass_name: str) -> None:
+        """Append a pass name to this adjoint's provenance history."""
+        self.history.append(pass_name)
+
+    def __repr__(self) -> str:  # pragma: no cover - cosmetic
+        return (
+            f"AdjointState(layer_id={self.layer_id}, block={self.block!r}, "
+            f"shape={tuple(self.tensor.shape)}, dtype={self.dtype}, "
+            f"device={self.device}, history={self.history})"
+        )

eae_runtime/backend.py

@@ -0,0 +1,52 @@
+"""
+Backend Manager: responsible for choosing the execution backend. Blocks can
+override backend selection by carrying a `.eae_backend` attribute naming a
+preferred backend.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Optional
+
+import torch
+
+
+class BackendManager:
+    SUPPORTED = ("cpu", "cuda", "triton", "rocm", "auto")
+
+    def __init__(self, backend: str = "auto"):
+        if backend not in self.SUPPORTED:
+            raise ValueError(f"Unknown backend '{backend}', choices: {self.SUPPORTED}")
+        self.requested = backend
+        self.resolved = self._resolve(backend)
+
+    @staticmethod
+    def _resolve(backend: str) -> str:
+        if backend == "auto":
+            return "cuda" if torch.cuda.is_available() else "cpu"
+        if backend == "cuda" and not torch.cuda.is_available():
+            return "cpu"  # graceful fallback rather than a hard failure
+        if backend in ("triton", "rocm"):
+            # Triton/ROCm kernels are launched *through* CUDA devices in this
+            # runtime; if no CUDA device is present we fall back to plain CPU
+            # execution so the same code runs everywhere.
+            return "cuda" if torch.cuda.is_available() else "cpu"
+        return backend
+
+    def device_for(self, block=None) -> torch.device:
+        override = getattr(block, "eae_backend", None)
+        backend = self._resolve(override) if override else self.resolved
+        return torch.device(backend if backend in ("cpu", "cuda") else "cpu")
+
+    @contextlib.contextmanager
+    def autocast(self, enabled: bool = False, dtype: Optional[torch.dtype] = None):
+        device_type = "cuda" if self.resolved == "cuda" else "cpu"
+        if not enabled:
+            yield
+            return
+        with torch.autocast(device_type=device_type, dtype=dtype or torch.float16, enabled=True):
+            yield
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return f"BackendManager(requested={self.requested!r}, resolved={self.resolved!r})"

eae_runtime/blocks.py

@@ -0,0 +1,69 @@
+"""
+Block interface. The runtime reconstructs blocks independently - no block
+should know anything about global training.
+
+Per the spec's non-goals, we do NOT attempt automatic partitioning of
+arbitrary nn.Module graphs. Manual / helper-based block specification is
+sufficient: users hand us a `nn.Sequential` (or any list of `nn.Module`s
+that compose linearly, e.g. repeated Transformer blocks) and we treat each
+element as one reconstructable unit.
+"""
+
+from __future__ import annotations
+
+from typing import List, Sequence
+
+import torch.nn as nn
+
+
+class EAEBlock(nn.Module):
+    """Optional marker base class for a block usable by the EAE runtime.
+
+    Subclassing this is *not* required - any `nn.Module` whose `forward`
+    is a pure function of its input and its own parameters works fine.
+    Subclass it when you want to be explicit about intent, or want the
+    stricter contract this class documents:
+
+      * forward(x) -> Tensor            (single tensor in, single tensor out)
+      * parameters() -> Iterator[Tensor] (inherited from nn.Module)
+      * no reliance on global/mutable state outside the block
+    """
+
+    def forward(self, x):  # pragma: no cover - documentation stub
+        raise NotImplementedError
+
+
+class BlockDecomposer:
+    """Turns a user model into an ordered list of reconstructable blocks."""
+
+    @staticmethod
+    def from_sequential(seq: nn.Sequential) -> List[nn.Module]:
+        return list(seq.children())
+
+    @staticmethod
+    def from_list(modules: Sequence[nn.Module]) -> List[nn.Module]:
+        return list(modules)
+
+    @staticmethod
+    def decompose(model) -> List[nn.Module]:
+        """Best-effort dispatcher: nn.Sequential -> children, nn.ModuleList ->
+        children, list-like -> list, single nn.Module -> a one-block model.
+
+        `nn.ModuleList` is handled explicitly (not just via generic
+        `nn.Module` iteration) because it's the idiomatic way to hold a
+        repeated Transformer block stack, e.g.
+        `nn.ModuleList([TransformerBlock(...) for _ in range(depth)])`.
+        """
+        if isinstance(model, nn.Sequential):
+            return BlockDecomposer.from_sequential(model)
+        if isinstance(model, nn.ModuleList):
+            return list(model)
+        if isinstance(model, (list, tuple)):
+            return BlockDecomposer.from_list(model)
+        if isinstance(model, nn.Module):
+            return [model]
+        raise TypeError(
+            f"Cannot decompose object of type {type(model)} into EAE blocks. "
+            "Pass an nn.Sequential, an nn.ModuleList, a list of nn.Module, "
+            "or a single nn.Module."
+        )

eae_runtime/boundary_store.py

@@ -0,0 +1,65 @@
+"""
+Boundary Store: stores only x0, x1, ..., xL (block-boundary activations).
+Nothing else - no internal activations, no autograd graph.
+
+Supports CPU offload, pinned memory, optional compression and configurable
+storage precision, all orthogonal to the reconstruction logic above it.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+import torch
+
+
+class BoundaryStore:
+    def __init__(
+        self,
+        offload: bool = False,
+        precision: Optional[torch.dtype] = None,
+        pin_memory: bool = False,
+    ):
+        self.offload = offload
+        self.precision = precision
+        self.pin_memory = pin_memory and torch.cuda.is_available()
+        self._store: Dict[int, torch.Tensor] = {}
+        self._orig_dtype: Dict[int, torch.dtype] = {}
+        self._orig_device: Dict[int, torch.device] = {}
+
+    def put(self, idx: int, tensor: torch.Tensor) -> None:
+        t = tensor.detach()
+        self._orig_dtype[idx] = t.dtype
+        self._orig_device[idx] = t.device
+
+        if self.precision is not None:
+            t = t.to(self.precision)
+        if self.offload:
+            t = t.to("cpu")
+            if self.pin_memory:
+                t = t.pin_memory()
+        else:
+            t = t.clone()  # own our copy, independent of caller's buffer
+        self._store[idx] = t
+
+    def get(self, idx: int, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        if idx not in self._store:
+            raise KeyError(f"Boundary {idx} was never stored")
+        t = self._store[idx]
+        target_device = device if device is not None else self._orig_device[idx]
+        target_dtype = dtype if dtype is not None else self._orig_dtype[idx]
+        return t.to(device=target_device, dtype=target_dtype)
+
+    def __contains__(self, idx: int) -> bool:
+        return idx in self._store
+
+    def __len__(self) -> int:
+        return len(self._store)
+
+    def clear(self) -> None:
+        self._store.clear()
+        self._orig_dtype.clear()
+        self._orig_device.clear()
+
+    def memory_bytes(self) -> int:
+        return sum(t.element_size() * t.nelement() for t in self._store.values())

eae_runtime/config.py

@@ -0,0 +1,48 @@
+"""
+Single configuration object for the runtime. No scattered options.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, List, Optional, Union
+
+import torch
+
+
+@dataclass
+class RuntimeConfig:
+    # scheduler: name string ("sequential" | "async" | "pipeline" | "distributed")
+    # or a BaseScheduler instance for full user control.
+    scheduler: Union[str, Any] = "sequential"
+
+    # memory policy: name string ("pool" | "none") or a MemoryPolicy instance.
+    memory: Union[str, Any] = "pool"
+
+    # backend: name string ("cpu" | "cuda" | "triton" | "rocm") or a
+    # BackendManager instance.
+    backend: Union[str, Any] = "auto"
+
+    # ordered list of EAEPass instances forming the adjoint pipeline.
+    passes: List[Any] = field(default_factory=list)
+
+    # mixed precision compute dtype used during reconstruction (None = fp32).
+    compute_dtype: Optional[torch.dtype] = None
+
+    # boundary store options
+    boundary_offload: bool = False          # move boundary activations to CPU
+    boundary_precision: Optional[torch.dtype] = None  # e.g. torch.float16 to save memory
+    pin_memory: bool = False
+
+    # pipeline scheduler option
+    num_microbatches: int = 1
+
+    # misc
+    grad_clip_norm: Optional[float] = None
+    seed: Optional[int] = None
+    enable_profiler: bool = True
+    log_level: str = "WARNING"
+
+    def __post_init__(self):
+        if self.seed is not None:
+            torch.manual_seed(self.seed)

eae_runtime/contrib/__init__.py

@@ -0,0 +1,25 @@
+"""
+Optional, opt-in reference implementations built on top of the EAE Runtime
+core. Nothing under `eae_runtime.contrib` is imported by
+`eae_runtime.__init__`, and the runtime core has zero dependency on it -
+this package exists purely to demonstrate and give you a starting point
+for real Transformer architectures. Fork freely.
+"""
+
+from .transformer_blocks import (
+    CausalSelfAttention,
+    PreNormTransformerBlock,
+    RMSNorm,
+    RotaryPositionalEmbedding,
+    SwiGLU,
+    apply_rotary,
+)
+
+__all__ = [
+    "CausalSelfAttention",
+    "PreNormTransformerBlock",
+    "RMSNorm",
+    "RotaryPositionalEmbedding",
+    "SwiGLU",
+    "apply_rotary",
+]