gradientmail 0.0.1__tar.gz

gradientmail-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Gabriel Francisco
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

gradientmail-0.0.1/PKG-INFO

@@ -0,0 +1,142 @@
+Metadata-Version: 2.4
+Name: gradientmail
+Version: 0.0.1
+Summary: Distributed training through asynchronous message exchange
+Author: Gabriel Francisco
+Author-email: Gabriel Francisco <gabfssilva@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Requires-Dist: safetensors>=0.8,<1
+Requires-Dist: jax ; extra == 'jax'
+Requires-Dist: keras ; extra == 'keras'
+Requires-Dist: mlx ; extra == 'mlx'
+Requires-Dist: boto3>=1.36 ; extra == 's3'
+Requires-Dist: tensorflow ; extra == 'tensorflow'
+Requires-Dist: torch ; extra == 'torch'
+Requires-Dist: zstandard>=0.23 ; extra == 'zstd'
+Requires-Python: >=3.12
+Project-URL: Homepage, https://ppgsi-lab.github.io/gradientmail/
+Project-URL: Documentation, https://ppgsi-lab.github.io/gradientmail/docs/guide/
+Project-URL: Repository, https://github.com/ppgsi-lab/gradientmail
+Provides-Extra: jax
+Provides-Extra: keras
+Provides-Extra: mlx
+Provides-Extra: s3
+Provides-Extra: tensorflow
+Provides-Extra: torch
+Provides-Extra: zstd
+Description-Content-Type: text/markdown
+
+# GradientMail
+
+Distributed training with [Decoupled DiLoCo](https://arxiv.org/abs/2604.21428) over a shared mailbox. Nodes exchange parameter fragments through S3, R2, MinIO or Hugging Face Storage Buckets, without P2P connections or a separate coordinator.
+
+The cluster is heterogeneous and elastic: nodes with different GPUs and frameworks (PyTorch or MLX) can join an ongoing run, contribute to training for a period and leave. One of the training nodes merges each round and publishes the result. If it fails, another takes over.
+
+## How it works
+
+Each learner keeps a local copy of the model and runs its own optimizer. Parameters are split into fragments, synchronized on a staggered schedule. For each fragment, a learner publishes its update: the difference between the committed parameters it last aligned with and the ones it trained.
+
+![Three nodes and a mailbox. Node A (MLX on Apple silicon), node B (PyTorch on an H100) and node C (PyTorch on an RTX 4090) each run a learner and a syncer. Each learner trains fragment k, puts its delta, the anchor minus its parameters, in the mailbox, polls for commits and aligns with them between its steps. Node B's syncer is elected for the round: it lists and gets the round's deltas, merges them with RDA, applies the Nesterov outer step and commits fragment k with the outer state. The other syncers follow it and take over if node B stalls. The mailbox, an S3, R2, MinIO or Hugging Face bucket, holds the deltas and the commits; every arrow is a request to it.](https://ppgsi-lab.github.io/gradientmail/docs/assets/training-flow.svg)
+
+An elected learner reads the updates for a round, merges them and applies the outer optimizer. It commits the updated fragment through the mailbox; learners poll for new commits and align their local models with them. All training and merge arithmetic runs on the learners.
+
+## Install
+
+Python 3.12+.
+
+```sh
+uv add "gradientmail[s3,torch]"
+# or
+pip install "gradientmail[s3,torch]"
+```
+
+Use `mlx` instead of `torch` for MLX on Apple Silicon.
+
+## Usage
+
+Add a `Session` to your existing training loop:
+
+```python
+import gradientmail as gm
+
+with gm.Session(
+    model,
+    algorithm=gm.DiLoCo(sync_every=100, fragments=8),
+    mailbox=gm.S3Mailbox(bucket="my-training", prefix="experiments/run-1"),
+) as session:
+    for batch in dataloader:
+        loss = train_step(model, optimizer, batch)
+        session.step(samples=len(batch))
+```
+
+Run this on each node with the same bucket, prefix, algorithm configuration and parameter schema (names, shapes and dtypes). `model` can be a PyTorch or MLX module. The first node initializes the run; later nodes load its committed parameters before training. Use a new prefix for each run.
+
+`sync_every=100, fragments=8` splits the model into up to eight parameter groups, each synchronized every 100 local steps. `step(samples=...)` reports the tokens or samples processed by that optimizer step. Leaving the context waits for the node's rounds in flight to commit, stopping the wait if the run stalls, and sets the model to the newest committed parameters. Steps trained since a fragment's last round are dropped, as in the DiLoCo papers; `flush_on_close=True` publishes them instead, and a round made only of those final contributions commits them without the outer momentum.
+
+Training continues during synchronization. Aligning a fragment with its commit replaces local updates made while it was in flight. `overlap_steps=N` on `gm.DiLoCo` bounds that: a node waits once a round it published N steps ago is still not aligned with, and `0` is standard DiLoCo's barrier. `mixing=0.5` keeps that share of the local parameters when aligning, blending them with the commit instead of replacing them, which Streaming DiLoCo found to recover most of what overlapped steps cost.
+
+The defaults follow Decoupled DiLoCo. `fragments=1, in_flight=1, overlap_steps=0, merge=gm.WeightedMean()` with `quorum_size` equal to the number of nodes is standard DiLoCo. `partition=gm.Strided()` deals the transformer blocks round the fragments, Streaming DiLoCo's pattern, instead of balancing the fragments by bytes.
+
+`quorum_size` sets the target number of contributions per round (default 1). Rounds can close with fewer contributions once every live node has reported. Crashed nodes stop being counted after their lease expires (default 120 seconds).
+
+`global_batch_size` fixes the batch the live nodes share. `session.batch_size` is this node's share, the global batch divided by the live nodes and rounded up, for the training loop to size its batches by; it changes as nodes join and leave. Without it, each node's batch is its own and the global batch grows with the node count.
+
+`delta_dtype="BF16"` sends each update in bf16, halving an fp32 model's uploads at a rounding of 8 significant bits per element; `"F8_E4M3"` quarters them and `"F4"` takes an eighth, each element scaled by a power of two per tensor or per block of 32, as Streaming DiLoCo sends its outer gradients; by default it keeps each parameter's own dtype. `commit_dtype` publishes each commit's change in the same formats, so a node holding the previous version downloads that instead of the full parameters. `compression="zstd"` on `Session` stores this node's payloads byte-shuffled and zstandard-compressed (the `zstd` extra), about 16% smaller for fp32 and 31% for bf16; every node reads either format, so nodes of one run may differ.
+
+`S3Mailbox` uses boto3 credentials. For other S3-compatible services, pass `endpoint_url=`; Hugging Face Storage Buckets also need `request_checksums="when_required"`.
+
+For evaluation, `gm.align(model, mailbox=...)` loads the committed model without joining the run. GradientMail does not save your inner optimizer state; checkpoint it separately.
+
+## Examples
+
+Run two synthetic learners locally, with PyTorch or MLX installed:
+
+```sh
+uv run python examples/synthetic.py
+```
+
+For a PyTorch/MLX run across two machines, use [torch_learner.py](https://github.com/ppgsi-lab/gradientmail/blob/main/examples/torch_learner.py) and [mlx_learner.py](https://github.com/ppgsi-lab/gradientmail/blob/main/examples/mlx_learner.py). The PyTorch example includes the MinIO setup and commands.
+
+## Further reading
+
+- [Guide](https://ppgsi-lab.github.io/gradientmail/docs/guide/)
+
+  A collection of explanations for understanding distributed training.
+
+- [Shared Queue](https://ppgsi-lab.github.io/gradientmail/docs/patterns/00-shared-queue.html)
+
+  Distributing work across nodes with different processing speeds.
+
+- [Barrier](https://ppgsi-lab.github.io/gradientmail/docs/patterns/01-barrier.html)
+
+  Holding nodes at a named point until a number of them have reached it.
+
+- [Posts](https://ppgsi-lab.github.io/gradientmail/docs/posts/)
+
+  A Pythia-14M run on four rented GPUs, MLX and PyTorch learners in one run, merging deltas with RDA, and implementing a mailbox backend.
+
+## Development
+
+From a checkout:
+
+```sh
+uv sync --extra s3 --extra torch
+uv run pytest
+uv run ruff check
+uv run ruff format --check
+uv run pyrefly check
+```
+
+Tests use memory and moto stores by default. Set `GRADIENTMAIL_S3_ENDPOINT` to test against an S3-compatible store. Framework tests skip when their dependencies are absent.
+
+## Citation
+
+```bibtex
+@software{gradientmail,
+  author = {dos Santos Silva, Gabriel Francisco},
+  title = {GradientMail: Distributed training through asynchronous message exchange},
+  year = {2026},
+  url = {https://github.com/ppgsi-lab/gradientmail},
+}
+```

gradientmail-0.0.1/README.md

@@ -0,0 +1,113 @@
+# GradientMail
+
+Distributed training with [Decoupled DiLoCo](https://arxiv.org/abs/2604.21428) over a shared mailbox. Nodes exchange parameter fragments through S3, R2, MinIO or Hugging Face Storage Buckets, without P2P connections or a separate coordinator.
+
+The cluster is heterogeneous and elastic: nodes with different GPUs and frameworks (PyTorch or MLX) can join an ongoing run, contribute to training for a period and leave. One of the training nodes merges each round and publishes the result. If it fails, another takes over.
+
+## How it works
+
+Each learner keeps a local copy of the model and runs its own optimizer. Parameters are split into fragments, synchronized on a staggered schedule. For each fragment, a learner publishes its update: the difference between the committed parameters it last aligned with and the ones it trained.
+
+![Three nodes and a mailbox. Node A (MLX on Apple silicon), node B (PyTorch on an H100) and node C (PyTorch on an RTX 4090) each run a learner and a syncer. Each learner trains fragment k, puts its delta, the anchor minus its parameters, in the mailbox, polls for commits and aligns with them between its steps. Node B's syncer is elected for the round: it lists and gets the round's deltas, merges them with RDA, applies the Nesterov outer step and commits fragment k with the outer state. The other syncers follow it and take over if node B stalls. The mailbox, an S3, R2, MinIO or Hugging Face bucket, holds the deltas and the commits; every arrow is a request to it.](https://ppgsi-lab.github.io/gradientmail/docs/assets/training-flow.svg)
+
+An elected learner reads the updates for a round, merges them and applies the outer optimizer. It commits the updated fragment through the mailbox; learners poll for new commits and align their local models with them. All training and merge arithmetic runs on the learners.
+
+## Install
+
+Python 3.12+.
+
+```sh
+uv add "gradientmail[s3,torch]"
+# or
+pip install "gradientmail[s3,torch]"
+```
+
+Use `mlx` instead of `torch` for MLX on Apple Silicon.
+
+## Usage
+
+Add a `Session` to your existing training loop:
+
+```python
+import gradientmail as gm
+
+with gm.Session(
+    model,
+    algorithm=gm.DiLoCo(sync_every=100, fragments=8),
+    mailbox=gm.S3Mailbox(bucket="my-training", prefix="experiments/run-1"),
+) as session:
+    for batch in dataloader:
+        loss = train_step(model, optimizer, batch)
+        session.step(samples=len(batch))
+```
+
+Run this on each node with the same bucket, prefix, algorithm configuration and parameter schema (names, shapes and dtypes). `model` can be a PyTorch or MLX module. The first node initializes the run; later nodes load its committed parameters before training. Use a new prefix for each run.
+
+`sync_every=100, fragments=8` splits the model into up to eight parameter groups, each synchronized every 100 local steps. `step(samples=...)` reports the tokens or samples processed by that optimizer step. Leaving the context waits for the node's rounds in flight to commit, stopping the wait if the run stalls, and sets the model to the newest committed parameters. Steps trained since a fragment's last round are dropped, as in the DiLoCo papers; `flush_on_close=True` publishes them instead, and a round made only of those final contributions commits them without the outer momentum.
+
+Training continues during synchronization. Aligning a fragment with its commit replaces local updates made while it was in flight. `overlap_steps=N` on `gm.DiLoCo` bounds that: a node waits once a round it published N steps ago is still not aligned with, and `0` is standard DiLoCo's barrier. `mixing=0.5` keeps that share of the local parameters when aligning, blending them with the commit instead of replacing them, which Streaming DiLoCo found to recover most of what overlapped steps cost.
+
+The defaults follow Decoupled DiLoCo. `fragments=1, in_flight=1, overlap_steps=0, merge=gm.WeightedMean()` with `quorum_size` equal to the number of nodes is standard DiLoCo. `partition=gm.Strided()` deals the transformer blocks round the fragments, Streaming DiLoCo's pattern, instead of balancing the fragments by bytes.
+
+`quorum_size` sets the target number of contributions per round (default 1). Rounds can close with fewer contributions once every live node has reported. Crashed nodes stop being counted after their lease expires (default 120 seconds).
+
+`global_batch_size` fixes the batch the live nodes share. `session.batch_size` is this node's share, the global batch divided by the live nodes and rounded up, for the training loop to size its batches by; it changes as nodes join and leave. Without it, each node's batch is its own and the global batch grows with the node count.
+
+`delta_dtype="BF16"` sends each update in bf16, halving an fp32 model's uploads at a rounding of 8 significant bits per element; `"F8_E4M3"` quarters them and `"F4"` takes an eighth, each element scaled by a power of two per tensor or per block of 32, as Streaming DiLoCo sends its outer gradients; by default it keeps each parameter's own dtype. `commit_dtype` publishes each commit's change in the same formats, so a node holding the previous version downloads that instead of the full parameters. `compression="zstd"` on `Session` stores this node's payloads byte-shuffled and zstandard-compressed (the `zstd` extra), about 16% smaller for fp32 and 31% for bf16; every node reads either format, so nodes of one run may differ.
+
+`S3Mailbox` uses boto3 credentials. For other S3-compatible services, pass `endpoint_url=`; Hugging Face Storage Buckets also need `request_checksums="when_required"`.
+
+For evaluation, `gm.align(model, mailbox=...)` loads the committed model without joining the run. GradientMail does not save your inner optimizer state; checkpoint it separately.
+
+## Examples
+
+Run two synthetic learners locally, with PyTorch or MLX installed:
+
+```sh
+uv run python examples/synthetic.py
+```
+
+For a PyTorch/MLX run across two machines, use [torch_learner.py](https://github.com/ppgsi-lab/gradientmail/blob/main/examples/torch_learner.py) and [mlx_learner.py](https://github.com/ppgsi-lab/gradientmail/blob/main/examples/mlx_learner.py). The PyTorch example includes the MinIO setup and commands.
+
+## Further reading
+
+- [Guide](https://ppgsi-lab.github.io/gradientmail/docs/guide/)
+
+  A collection of explanations for understanding distributed training.
+
+- [Shared Queue](https://ppgsi-lab.github.io/gradientmail/docs/patterns/00-shared-queue.html)
+
+  Distributing work across nodes with different processing speeds.
+
+- [Barrier](https://ppgsi-lab.github.io/gradientmail/docs/patterns/01-barrier.html)
+
+  Holding nodes at a named point until a number of them have reached it.
+
+- [Posts](https://ppgsi-lab.github.io/gradientmail/docs/posts/)
+
+  A Pythia-14M run on four rented GPUs, MLX and PyTorch learners in one run, merging deltas with RDA, and implementing a mailbox backend.
+
+## Development
+
+From a checkout:
+
+```sh
+uv sync --extra s3 --extra torch
+uv run pytest
+uv run ruff check
+uv run ruff format --check
+uv run pyrefly check
+```
+
+Tests use memory and moto stores by default. Set `GRADIENTMAIL_S3_ENDPOINT` to test against an S3-compatible store. Framework tests skip when their dependencies are absent.
+
+## Citation
+
+```bibtex
+@software{gradientmail,
+  author = {dos Santos Silva, Gabriel Francisco},
+  title = {GradientMail: Distributed training through asynchronous message exchange},
+  year = {2026},
+  url = {https://github.com/ppgsi-lab/gradientmail},
+}
+```

gradientmail-0.0.1/pyproject.toml

@@ -0,0 +1,72 @@
+[project]
+name = "gradientmail"
+version = "0.0.1"
+description = "Distributed training through asynchronous message exchange"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.12"
+dependencies = ["safetensors>=0.8,<1"]
+
+[[project.authors]]
+name = "Gabriel Francisco"
+email = "gabfssilva@gmail.com"
+
+[project.urls]
+Homepage = "https://ppgsi-lab.github.io/gradientmail/"
+Documentation = "https://ppgsi-lab.github.io/gradientmail/docs/guide/"
+Repository = "https://github.com/ppgsi-lab/gradientmail"
+
+[project.optional-dependencies]
+s3 = ["boto3>=1.36"]
+torch = ["torch"]
+mlx = ["mlx"]
+jax = ["jax"]
+keras = ["keras"]
+tensorflow = ["tensorflow"]
+zstd = ["zstandard>=0.23"]
+
+[dependency-groups]
+dev = [
+    "pytest>=9",
+    "pytest-describe>=3",
+    "hypothesis",
+    "ruff",
+    "pyrefly>=1.3",
+    "moto[s3]>=5.1.5",
+    "boto3>=1.36",
+    "boto3-stubs-lite[s3]>=1.36",
+    "zstandard>=0.23",
+]
+
+[build-system]
+requires = ["uv_build>=0.13.0,<0.14.0"]
+build-backend = "uv_build"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+describe_prefixes = [
+    "describe",
+    "when",
+]
+
+[tool.ruff]
+target-version = "py312"
+line-length = 100
+extend-exclude = ["*.md"]
+
+[tool.ruff.lint]
+select = [
+    "E",
+    "F",
+    "I",
+    "UP",
+    "B",
+    "SIM",
+    "RUF",
+]
+
+[tool.pyrefly]
+project-includes = ["src"]
+python-version = "3.12"
+search-path = ["src"]

gradientmail-0.0.1/pyproject.toml.orig

@@ -0,0 +1,62 @@
+[project]
+name = "gradientmail"
+version = "0.0.1"
+description = "Distributed training through asynchronous message exchange"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [
+    { name = "Gabriel Francisco", email = "gabfssilva@gmail.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "safetensors>=0.8,<1",
+]
+
+[project.urls]
+Homepage = "https://ppgsi-lab.github.io/gradientmail/"
+Documentation = "https://ppgsi-lab.github.io/gradientmail/docs/guide/"
+Repository = "https://github.com/ppgsi-lab/gradientmail"
+
+[project.optional-dependencies]
+s3 = ["boto3>=1.36"]
+torch = ["torch"]
+mlx = ["mlx"]
+jax = ["jax"]
+keras = ["keras"]
+tensorflow = ["tensorflow"]
+zstd = ["zstandard>=0.23"]
+
+[dependency-groups]
+dev = [
+    "pytest>=9",
+    "pytest-describe>=3",
+    "hypothesis",
+    "ruff",
+    "pyrefly>=1.3",
+    "moto[s3]>=5.1.5",
+    "boto3>=1.36",
+    "boto3-stubs-lite[s3]>=1.36",
+    "zstandard>=0.23",
+]
+
+[build-system]
+requires = ["uv_build>=0.13.0,<0.14.0"]
+build-backend = "uv_build"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+describe_prefixes = ["describe", "when"]
+
+[tool.ruff]
+target-version = "py312"
+line-length = 100
+extend-exclude = ["*.md"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B", "SIM", "RUF"]
+
+[tool.pyrefly]
+project-includes = ["src"]
+python-version = "3.12"
+search-path = ["src"]

gradientmail-0.0.1/src/gradientmail/__init__.py

@@ -0,0 +1,59 @@
+"""GradientMail: distributed training through asynchronous message exchange.
+
+``import gradientmail`` needs only the standard library and safetensors. ``Session`` wraps a
+training loop; every session also takes part in closing rounds, so a run has no separate
+process to operate. What a node reports while it runs is in ``gradientmail.events``, handed
+to the session's ``callbacks``. ``S3Mailbox`` is resolved on first access and raises
+``MissingExtra`` naming the ``s3`` extra when boto3 is absent. The framework backends live in
+``gradientmail.adapters.torch`` and ``gradientmail.adapters.mlx``.
+"""
+
+from typing import TYPE_CHECKING
+
+from gradientmail.adapters import Adapter, Backend, MemoryAdapter, MissingExtra
+from gradientmail.algorithms import Algorithm
+from gradientmail.algorithms.diloco import ByBytes, DiLoCo, Rda, Strided, WeightedMean
+from gradientmail.events import Callback, Event, log_events
+from gradientmail.mailbox import Mailbox, MailboxUnavailable
+from gradientmail.mailbox.memory import MemoryMailbox
+from gradientmail.protocol import RunConfigMismatch
+from gradientmail.serialization import Schema, SchemaMismatch, TensorBytes
+from gradientmail.session import Session, align
+
+if TYPE_CHECKING:
+    from gradientmail.mailbox.s3 import S3Mailbox as S3Mailbox
+
+__all__ = [
+    "Adapter",
+    "Algorithm",
+    "Backend",
+    "ByBytes",
+    "Callback",
+    "DiLoCo",
+    "Event",
+    "Mailbox",
+    "MailboxUnavailable",
+    "MemoryAdapter",
+    "MemoryMailbox",
+    "MissingExtra",
+    "Rda",
+    "RunConfigMismatch",
+    "Schema",
+    "SchemaMismatch",
+    "Session",
+    "Strided",
+    "TensorBytes",
+    "WeightedMean",
+    "align",
+    "log_events",
+]
+
+
+def __getattr__(name: str) -> object:
+    if name == "S3Mailbox":
+        try:
+            from gradientmail.mailbox.s3 import S3Mailbox
+        except ImportError as exc:
+            raise MissingExtra("s3") from exc
+        return S3Mailbox
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

gradientmail-0.0.1/src/gradientmail/adapters/__init__.py

@@ -0,0 +1,174 @@
+"""Framework backends: how a model's parameters become bytes, and how bytes become arithmetic.
+
+A ``Backend`` is everything GradientMail knows about one framework. On a learner it wraps a
+model in an ``Adapter``, the three-method byte boundary (``schema``, ``snapshot``, ``apply``);
+on a syncer it loads ``TensorBytes`` as ``Value``s, the handful of fp32 operations the outer
+step is made of, which serialize themselves back. ``TorchBackend`` and ``MlxBackend`` live
+in their own modules and are imported only when asked for, so this package needs nothing
+beyond the standard library. Validation against the schema is the
+caller's job (``Learner`` before every ``apply``, ``Syncer`` before merging), so adapters
+stay thin.
+"""
+
+from collections.abc import Iterable, Mapping
+from importlib.util import find_spec
+from typing import Protocol, Self, runtime_checkable
+
+from gradientmail.serialization import (
+    DType,
+    ParameterSpec,
+    Quantized,
+    Schema,
+    SchemaMismatch,
+    TensorBytes,
+)
+
+__all__ = [
+    "Adapter",
+    "Backend",
+    "MemoryAdapter",
+    "MissingExtra",
+    "Value",
+    "adapter_for",
+    "backend_for",
+    "default_backend",
+]
+
+
+class MissingExtra(ImportError):
+    """An optional dependency is not installed; ``extra`` names the extra that provides it."""
+
+    extra: str
+
+    def __init__(self, extra: str) -> None:
+        self.extra = extra
+        super().__init__(f"this needs the {extra!r} extra: pip install 'gradientmail[{extra}]'")
+
+
+@runtime_checkable
+class Adapter(Protocol):
+    def schema(self) -> Schema[DType]:
+        """Canonical trainable parameters, sorted by name."""
+        ...
+
+    def snapshot(self, names: Iterable[str]) -> Mapping[str, TensorBytes]:
+        """Owned copies of ``names``, consistent at call time."""
+        ...
+
+    def apply(self, tensors: Mapping[str, TensorBytes]) -> None:
+        """Overwrite parameters in place; never rebinds them, so optimizer state survives."""
+        ...
+
+
+@runtime_checkable
+class Value(Protocol):
+    """An fp32 tensor held by a backend, with the arithmetic the outer step is made of.
+
+    Every operation returns a new value of the same backend. ``norm`` is the Euclidean norm
+    over all elements and ``dot`` the sum of the elementwise products with ``other``.
+    ``to_bytes`` rounds to ``dtype`` (nearest even) and serializes; ``quantize`` serializes to
+    a ``Quantized`` format as ``(data, scales)`` in the layout ``serialization.quantized_specs``
+    describes, each scale a power of two that puts the largest magnitude of its tensor or
+    block at the top of the format, so dequantizing is exact in fp32.
+    """
+
+    def __add__(self, other: Self) -> Self: ...
+
+    def __sub__(self, other: Self) -> Self: ...
+
+    def __mul__(self, scalar: float) -> Self: ...
+
+    def norm(self) -> float: ...
+
+    def dot(self, other: Self) -> float: ...
+
+    def to_bytes(self, dtype: DType) -> TensorBytes: ...
+
+    def quantize(self, dtype: Quantized) -> tuple[TensorBytes, TensorBytes]: ...
+
+
+@runtime_checkable
+class Backend(Protocol):
+    """One framework: it adapts that framework's models and does the syncer's arithmetic."""
+
+    def adapter(self, model: object) -> Adapter:
+        """An Adapter over ``model``; ``TypeError`` when it is not this framework's model."""
+        ...
+
+    def load(self, tensor: TensorBytes) -> Value:
+        """``tensor`` as an fp32 value of its shape; ``ValueError`` for a quantized dtype,
+        whose bytes mean nothing without their scales."""
+        ...
+
+    def dequantize(self, data: TensorBytes, scale: TensorBytes, shape: tuple[int, ...]) -> Value:
+        """The fp32 value of ``shape`` whose ``quantize`` gave ``(data, scale)``, up to the
+        format's rounding; the same floats on every backend."""
+        ...
+
+
+class MemoryAdapter(Adapter):
+    """An Adapter over a dict of ``TensorBytes``: the Phase 2 learner and the fake in tests."""
+
+    def __init__(self, tensors: Mapping[str, TensorBytes]) -> None:
+        self._tensors = dict(tensors)
+
+    @property
+    def tensors(self) -> Mapping[str, TensorBytes]:
+        return self._tensors
+
+    def schema(self) -> Schema[DType]:
+        def spec(name: str, t: TensorBytes) -> ParameterSpec[DType]:
+            match t.dtype:
+                case "F32" | "F16" | "BF16" as dtype:
+                    return ParameterSpec(name, dtype, t.shape)
+                case _:
+                    raise SchemaMismatch([name], "parameter dtype is not F32, F16 or BF16")
+
+        return Schema(tuple(spec(name, t) for name, t in sorted(self._tensors.items())))
+
+    def snapshot(self, names: Iterable[str]) -> Mapping[str, TensorBytes]:
+        return {name: self._tensors[name] for name in names}
+
+    def apply(self, tensors: Mapping[str, TensorBytes]) -> None:
+        self._tensors.update(tensors)
+
+
+def backend_for(model: object) -> Backend:
+    """The backend of ``model``'s framework, read from the modules of ``type(model)``'s MRO.
+
+    A user's ``class MLP(nn.Module)`` lives in the user's module, so the whole MRO is searched;
+    no framework is imported to decide. Anything else is a ``TypeError``.
+    """
+    roots = {cls.__module__.partition(".")[0] for cls in type(model).__mro__}
+    if "torch" in roots:
+        from gradientmail.adapters.torch import TorchBackend
+
+        return TorchBackend()
+    if "mlx" in roots:
+        from gradientmail.adapters.mlx import MlxBackend
+
+        return MlxBackend()
+    raise TypeError(
+        f"no backend for {type(model).__module__}.{type(model).__qualname__}; "
+        "pass a torch or MLX module, or an Adapter"
+    )
+
+
+def adapter_for(model: object) -> Adapter:
+    """``model`` itself when it is an Adapter; otherwise ``backend_for(model).adapter(model)``."""
+    if isinstance(model, Adapter):
+        return model
+    return backend_for(model).adapter(model)
+
+
+def default_backend() -> Backend:
+    """The installed backend: torch, then MLX. ``MissingExtra("torch")`` when neither is."""
+    if find_spec("torch") is not None:
+        from gradientmail.adapters.torch import TorchBackend
+
+        return TorchBackend()
+    if find_spec("mlx") is not None:
+        from gradientmail.adapters.mlx import MlxBackend
+
+        return MlxBackend()
+    raise MissingExtra("torch")