gradientmail 0.0.1__py3-none-any.whl

gradientmail/syncer.py

@@ -0,0 +1,528 @@
+"""The syncer: collects a round's snapshots, closes the fragment and commits the round.
+
+Every node runs one, hosted by ``Session`` next to its learner; the node that closes a round
+is elected through the mailbox. Round ``r`` targets fragment ``r mod P`` at global step
+``(r div P)·H + ⌊((r mod P) + 1)·H / P⌋``, so every fragment is synchronized once per ``H``
+steps. The round record is the commit point: a node rebuilds its state from the newest record
+of every fragment, and two nodes that close the same round are resolved at the record write.
+Timings are durations on this process's monotonic clock plus the step durations learners
+report, never wall clocks compared across machines.
+
+Election
+--------
+Round ``r`` is closed by the node that created ``claims/{r}/001``, a create-once object naming
+the claimant. The node that commits ``r`` claims ``r + 1`` in the same breath, before anyone
+else can have seen the record, so the closer stays put while it lives and a fresh run's first
+claimant keeps the role. A node that finds a claim by someone else follows: it polls for the
+round record and continues from it. It creates the next claim and closes the round itself as
+soon as the claimant's lease is gone or stale (``lease`` module: released when the claimant
+closed, unrenewed since it crashed), or, for a claimant that is live but silent, once no record
+has appeared within a timeout measured on its own clock (three times the round duration the
+last record's estimates predict, never less than thirty poll intervals). The leases a follower
+observes carry into its ``close``, so a lease first met while following is not trusted afresh
+at takeover. A node that restarts finds its own claim and resumes. Claims only avoid redundant
+work: whoever commits
+first wins, and the other reads the committed record back (``Superseded``) and continues from
+it. A follower keeps no tensors in memory; once elected it reloads the fragments from the
+records (``recover``), so leadership changes cost one reload and steady state costs followers
+one listing per round.
+
+Quorum
+------
+A round closes once ``quorum_size`` snapshots for it are accepted, or once every live learner
+has reported and at least one snapshot is accepted. A snapshot is accepted iff its ``steps``
+are positive, its schema digest is the run's, its ``base_version`` is the fragment's current
+version, its payload is present and hashes to the record and its tensors match the fragment's
+delta schema (the parameters in the algorithm's delta dtype, with their scales when it is
+quantized); otherwise it stays in
+the round record as rejected (``zero_steps``, ``schema``,
+``stale``, ``missing``, ``corrupt``) with weight 0, the paper's treatment of learners outside
+the round. A live learner is never ``stale``: it advances ``rounds_seen`` only after aligning
+with every fragment it fetched, and ``in_flight ≤ P``; a ``stale`` snapshot comes from a learner
+that skipped a swept payload and published before aligning with the version its next poll listed.
+Liveness is the lease every learner renews while it runs (``lease`` module): a learner whose
+lease has not changed for ``lease_timeout`` seconds on this node's clock, or whose lease is
+gone, is not waited for. Learners therefore join and leave at will, and a quorum above the
+number of live learners closes with the live ones instead of waiting forever; a crashed learner
+holds up the rounds it would have taken part in for one timeout, a learner that closes releases
+its lease at once. Snapshots and leases are polled every ``poll_interval``, one LIST each, and
+payloads are fetched as their snapshots appear, overlapping the wait.
+
+Grace
+-----
+After quorum the closer keeps collecting for
+``min(0.5 · max(ahead·ξ_step - (ξ_quorum + ξ_sync), 0), ξ_round)`` seconds, ``ξ_round`` being the
+round's gap of steps times the step duration learners report and ``ahead`` the steps a learner
+trains before it would stall, ``in_flight`` rounds of them or the algorithm's ``overlap_steps``
+when that is less: half the slack those leave, capped at one round of training. The three
+estimates are EMAs (decay 0.1) of that step duration, taken from the snapshots, of the time
+from round start to quorum and of the time from closing to a durable commit. They travel
+in every round record, so the node that closes the next round continues them; the sync
+estimate lags one round when the closer changes, since a commit's duration is known only to
+the node that made it. The window is zero while any estimate is missing (the first round of a
+run) and ends early once every live learner has reported. With ``in_flight`` 1 the slack is
+negative, since reaching quorum alone takes about a round, and the window stays zero: rounds
+close at quorum and a learner late for a round is left out of it, though its next snapshot
+carries the steps anyway. ``in_flight`` 2 leaves about half a round. The realized
+``quorum_seconds`` and ``grace_seconds`` are in every round record. A snapshot that arrives
+after the close is not counted and is swept with the round.
+
+Recovery
+--------
+``recover`` lists ``rounds/`` and takes the newest record per fragment: its theta payload is
+the anchor and its outer payload holds the fp32 momentum and, when the published tensor is not
+the master (F16 and BF16 schemas, commits that carry an update), the fp32 master; a fragment
+never committed starts from the manifest's Θ⁰ with zero momentum. The
+next round is the one after the last record and the timing estimates are the last record's.
+Nothing else is persisted, so a crash
+before the record write recomputes the round from the snapshots still in the store and a crash
+after it continues with the next round; the commit sequence is the same either way, except
+that a snapshot which landed in between may now be included. A round that fails for any other
+reason is reported (``RoundFailed``) and retried the same way, from the store, one poll
+later. ``run`` raises ``RunConfigMismatch`` when the manifest was created with another
+algorithm configuration. A ``MailboxUnavailable`` is retried indefinitely with backoff. No
+``asyncio.timeout`` wraps a store call, because a cancelled ``to_thread`` would leak its
+thread; botocore's own timeouts bound them.
+
+Garbage collection
+------------------
+After committing round ``r`` the closer deletes ``snapshots/{r}/``, ``blobs/snapshots/{r}/``
+and ``claims/{r}/``, sweeps round ``r - 1`` again (for snapshots that landed after its close,
+for claims created after its sweep, and for a predecessor that crashed between commit and
+sweep) and deletes ``blobs/rounds/{r - 2P}/``, which also removes the payloads a superseded
+closer left behind. Two cycles of theta and outer payloads are kept so that the newest version
+of every fragment, at most ``P`` rounds old, always survives and a learner has a cycle to fetch
+a payload whose record it listed. Round records, the manifest and the init payloads are never
+deleted; ``rounds/`` is the tape of the run.
+"""
+
+import asyncio
+import math
+import time
+from collections.abc import Mapping
+from dataclasses import dataclass
+from statistics import fmean
+
+from gradientmail.adapters import Backend, default_backend
+from gradientmail.algorithms import (
+    Aggregator,
+    Algorithm,
+    Contribution,
+    Estimates,
+    FragmentState,
+    RoundSchedule,
+    layout,
+)
+from gradientmail.algorithms.diloco import grace_window
+from gradientmail.clock import Clock, MonotonicClock
+from gradientmail.compression import Compression
+from gradientmail.events import (
+    Callback,
+    FollowingRound,
+    Recovered,
+    RecoveryFailed,
+    RoundClaimed,
+    RoundClosed,
+    RoundCommitted,
+    RoundFailed,
+    RoundSuperseded,
+    TookOverRound,
+    log_events,
+)
+from gradientmail.lease import Liveness
+from gradientmail.mailbox import Mailbox
+from gradientmail.protocol import (
+    Exchange,
+    Manifest,
+    PayloadCorrupt,
+    PayloadRef,
+    ProtocolError,
+    Rejection,
+    RoundRecord,
+    RoundRecordDraft,
+    RunConfigMismatch,
+    Snapshot,
+    Superseded,
+)
+from gradientmail.serialization import PayloadDType, Schema, SchemaMismatch, TensorBytes, validate
+
+__all__ = ["Following", "Progress", "Syncer", "SyncerState"]
+
+
+@dataclass(frozen=True)
+class SyncerState:
+    """A closer's knowledge between rounds.
+
+    ``round`` is the next round to run, ``fragments`` the last committed state of every
+    fragment and ``estimates`` the timing EMAs behind the grace window.
+    """
+
+    round: int
+    fragments: tuple[FragmentState, ...]
+    estimates: Estimates
+
+
+@dataclass(frozen=True)
+class Following:
+    """A follower's knowledge: the next round, and how long to wait for its record before
+    claiming it, unless the claimant's lease ends first."""
+
+    round: int
+    timeout: float
+
+
+type Progress = SyncerState | Following
+
+
+class Syncer:
+    """Drives a run's rounds against a mailbox; ``run`` is the whole loop, ``run_round`` one step.
+
+    Everything the manifest fixes (schema, plan, schedule, one aggregator per fragment) is read
+    once on first use; the manifest is a create-once object. The outer step runs in ``backend``,
+    by default the installed one (torch, then MLX); constructing a syncer with neither installed
+    raises ``MissingExtra``. Round records are byte-identical across restarts on one backend;
+    another backend may differ in the last bit. ``compression`` is the codec this node's
+    commits are stored with. ``on_event`` receives every event the syncer reports (``events``
+    module), on the thread running ``run``.
+    """
+
+    def __init__(
+        self,
+        *,
+        algorithm: Algorithm,
+        mailbox: Mailbox,
+        learner_id: str,
+        backend: Backend | None = None,
+        clock: Clock | None = None,
+        poll_interval: float = 0.5,
+        lease_timeout: float = 120.0,
+        compression: Compression = "off",
+        on_event: Callback = log_events,
+    ) -> None:
+        self._algorithm = algorithm
+        self._learner_id = learner_id
+        self._backend = backend if backend is not None else default_backend()
+        self._clock = clock if clock is not None else MonotonicClock()
+        self._exchange = Exchange(mailbox, clock=self._clock, compression=compression)
+        self._poll_interval = poll_interval
+        self._on_event = on_event
+        self._liveness = Liveness(lease_timeout)
+        self._run: _Run | None = None
+        self._main: asyncio.Task[None] | None = None
+        self._stopped = False
+
+    async def run(self, *, until_round: int | None = None) -> None:
+        """Wait for the manifest, check the configuration, recover and run round after round.
+
+        Returns once round ``until_round`` is committed; with ``None`` it runs until cancelled
+        or ``stop()``. Raises ``RunConfigMismatch`` when the run was created with another
+        algorithm configuration. A round that fails is reported and retried as after a restart:
+        the claim this node holds, if any, is found again and the round recomputed.
+        """
+        self._main = asyncio.current_task()
+        if self._stopped:
+            return
+        while (manifest := await self._exchange.manifest()) is None:
+            await self._clock.sleep(self._poll_interval)
+        expected = self._algorithm.to_json()
+        if dict(manifest.algorithm) != dict(expected):
+            raise RunConfigMismatch(expected, manifest.algorithm)
+        progress: Progress | None = None
+        while True:
+            try:
+                if progress is None:
+                    progress = await self.recover()
+                if until_round is not None and progress.round > until_round:
+                    return
+                progress = await self.run_round(progress)
+            except Exception as exc:
+                self._on_event(
+                    RecoveryFailed(self._learner_id, exc, self._poll_interval)
+                    if progress is None
+                    else RoundFailed(self._learner_id, progress.round, exc, self._poll_interval)
+                )
+                await self._clock.sleep(self._poll_interval)
+                if progress is not None:
+                    progress = Following(progress.round, self._timeout_of(progress))
+
+    def stop(self) -> None:
+        """Cancel ``run`` from any thread; a ``run`` that has not started yet returns at once."""
+        self._stopped = True
+        if self._main is not None and not self._main.done():
+            loop = self._main.get_loop()
+            if not loop.is_closed():
+                loop.call_soon_threadsafe(self._main.cancel)
+
+    async def recover(self) -> SyncerState:
+        """Rebuild a closer's state from the store.
+
+        Every fragment comes from its newest round record (theta and outer payloads), or from the
+        manifest's Θ⁰ with zero momentum when it was never committed. The next round is the one
+        after the last record; the timing estimates are that record's.
+        """
+
+        async def required(ref: PayloadRef) -> Mapping[str, TensorBytes]:
+            tensors = await self._exchange.fetch(ref)
+            if tensors is None:
+                raise ProtocolError(f"{ref.key} is referenced by a live record but is gone")
+            return tensors
+
+        async def fragment_state(fragment: int) -> FragmentState:
+            record = newest.get(fragment)
+            if record is None:
+                momentum = {
+                    f"momentum.{name}": TensorBytes(
+                        "F32", shapes[name], bytes(4 * math.prod(shapes[name]))
+                    )
+                    for name in run.manifest.plan[fragment]
+                }
+                return FragmentState(0, await required(run.manifest.init[fragment]), momentum)
+            return FragmentState(
+                record.version, await required(record.theta), await required(record.outer)
+            )
+
+        run = await self._resolved()
+        records = await self._exchange.rounds_after(-1)
+        # Records come in round order, so the last one per fragment wins.
+        newest = {record.fragment: record for record in records}
+        shapes = {spec.name: spec.shape for spec in run.manifest.schema.parameters}
+        count = len(run.manifest.plan)
+        async with asyncio.TaskGroup() as group:
+            loads = [group.create_task(fragment_state(p)) for p in range(count)]
+        last = records[-1] if records else None
+        state = SyncerState(
+            round=0 if last is None else last.round + 1,
+            fragments=tuple(load.result() for load in loads),
+            estimates=Estimates() if last is None else last.estimates,
+        )
+        self._on_event(Recovered(self._learner_id, state.round))
+        return state
+
+    async def run_round(self, progress: Progress) -> Progress:
+        """Advance one round: close it when elected, otherwise follow it to its record."""
+        committed = await self._exchange.rounds_after(progress.round - 1)
+        if committed:
+            # Closed already, its claims swept with it; a node arriving now must not reclaim it.
+            return Following(committed[0].round + 1, self._timeout(committed[0]))
+        if not await self.elected(progress.round):
+            return await self.follow(progress)
+        state = progress if isinstance(progress, SyncerState) else await self.recover()
+        if state.round != progress.round:
+            # Committed meanwhile by an earlier claimant that was slow rather than dead.
+            return state
+        return await self.close(state)
+
+    async def elected(self, round: int) -> bool:
+        """Whether this node closes ``round``: it holds a claim, or it makes the first one."""
+        claimants = await self._exchange.claims(round)
+        if claimants:
+            return self._learner_id in claimants
+        claimed = await self._exchange.claim(round, 1, self._learner_id)
+        if claimed:
+            self._on_event(RoundClaimed(self._learner_id, round))
+        return claimed
+
+    async def follow(self, progress: Progress) -> Progress:
+        """Wait for the round's record; take the next claim and close the round once the
+        claimant's lease is gone or stale, or after the timeout."""
+        round = progress.round
+        timeout = self._timeout_of(progress)
+        claimants = await self._exchange.claims(round)
+        self._on_event(FollowingRound(self._learner_id, round, tuple(claimants), timeout))
+        seen_at = self._clock.now()
+        while True:
+            records = await self._exchange.rounds_after(round - 1)
+            if records:
+                return Following(records[0].round + 1, self._timeout(records[0]))
+            leases = await self._exchange.leases()
+            now = self._clock.now()
+            self._liveness = self._liveness.observed(leases, now)
+            gone = bool(claimants) and claimants[-1] not in self._liveness.live(now)
+            if gone or now - seen_at >= timeout:
+                if await self._exchange.claim(round, len(claimants) + 1, self._learner_id):
+                    self._on_event(
+                        TookOverRound(
+                            self._learner_id,
+                            round,
+                            tuple(claimants),
+                            "lease_gone" if gone else "no_record",
+                            timeout,
+                        )
+                    )
+                    state = await self.recover()
+                    return state if state.round != round else await self.close(state)
+                claimants = await self._exchange.claims(round)
+                seen_at = self._clock.now()
+            await self._clock.sleep(self._poll_interval)
+
+    async def close(self, state: SyncerState) -> Progress:
+        """Close round ``state.round`` through its commit and return the state for the next one.
+
+        Collects snapshots until quorum, or until every live learner has reported, keeps
+        collecting through the grace window unless every live learner has reported, closes the
+        fragment, commits, claims the next round and sweeps. When another node committed the
+        round first, its record is adopted and this node follows from there.
+        """
+
+        async def judge(snapshot: Snapshot) -> Contribution | Rejection:
+            if snapshot.steps == 0:
+                return "zero_steps"
+            if snapshot.schema_digest != run.digest:
+                return "schema"
+            if snapshot.base_version != current.version:
+                return "stale"
+            try:
+                tensors = await self._exchange.fetch(snapshot.payload)
+            except PayloadCorrupt:
+                return "corrupt"
+            if tensors is None:
+                return "missing"
+            try:
+                validate(deltas, tensors, [spec.name for spec in deltas.parameters])
+            except SchemaMismatch:
+                return "schema"
+            return Contribution(
+                snapshot.learner,
+                snapshot.local_step,
+                snapshot.steps,
+                snapshot.tokens,
+                snapshot.step_seconds,
+                snapshot.final,
+                tensors,
+            )
+
+        async def poll() -> None:
+            leases = await self._exchange.leases()
+            self._liveness = self._liveness.observed(leases, self._clock.now())
+            snapshots = await self._exchange.snapshots(round)
+            new = tuple(s for s in snapshots if s.learner not in verdicts)
+            async with asyncio.TaskGroup() as group:
+                judged = [(s.learner, group.create_task(judge(s))) for s in new]
+            verdicts.update((learner, task.result()) for learner, task in judged)
+
+        def accepted() -> tuple[Contribution, ...]:
+            contributions = (v for v in verdicts.values() if isinstance(v, Contribution))
+            return tuple(sorted(contributions, key=lambda c: c.learner))
+
+        def reported() -> bool:
+            """Every live learner has a snapshot in the round, accepted or rejected."""
+            return self._liveness.live(self._clock.now()) <= verdicts.keys()
+
+        def closable() -> bool:
+            return len(accepted()) >= self._algorithm.quorum() or (
+                len(accepted()) >= 1 and reported()
+            )
+
+        run = await self._resolved()
+        round = state.round
+        fragment = run.schedule.fragment_of(round)
+        current = state.fragments[fragment]
+        deltas = run.deltas[fragment]
+        verdicts: dict[str, Contribution | Rejection] = {}
+
+        started = self._clock.now()
+        await poll()
+        while not closable():
+            await self._clock.sleep(self._poll_interval)
+            await poll()
+        quorum_at = self._clock.now()
+
+        deadline = quorum_at + grace_window(
+            state.estimates,
+            run.schedule.in_flight,
+            run.schedule.gap(round),
+            run.schedule.overlap_steps,
+        )
+        while (now := self._clock.now()) < deadline and not reported():
+            await self._clock.sleep(min(self._poll_interval, deadline - now))
+            await poll()
+        grace_ended = self._clock.now()
+
+        contributions = accepted()
+        result = run.aggregators[fragment].close_round(current, contributions)
+        updated = result.state
+        estimates = state.estimates.updated(
+            step=fmean(c.step_seconds for c in contributions), quorum=quorum_at - started
+        )
+        draft = RoundRecordDraft(
+            round=round,
+            step=run.schedule.step_of(round),
+            fragment=fragment,
+            version=updated.version,
+            participants=result.participants,
+            rejected=tuple(
+                sorted((learner, why) for learner, why in verdicts.items() if isinstance(why, str))
+            ),
+            cosine=result.cosine,
+            quorum_seconds=quorum_at - started,
+            grace_seconds=grace_ended - quorum_at,
+            estimates=estimates,
+        )
+        self._on_event(RoundClosed(self._learner_id, draft))
+        committing = time.perf_counter()
+        outcome = await self._exchange.commit(draft, updated.theta, updated.outer, result.update)
+        if isinstance(outcome, Superseded):
+            self._on_event(RoundSuperseded(self._learner_id, outcome.record))
+            return Following(outcome.record.round + 1, self._timeout(outcome.record))
+        sync_seconds = self._clock.now() - grace_ended
+        self._on_event(
+            RoundCommitted(self._learner_id, outcome.record, time.perf_counter() - committing)
+        )
+
+        await self._exchange.claim(round + 1, 1, self._learner_id)
+        count = len(state.fragments)
+        for swept in (round, round - 1):
+            if swept >= 0:
+                await self._exchange.sweep_snapshots(swept)
+        await self._exchange.sweep_blobs(round - 2 * count)
+        return SyncerState(
+            round=round + 1,
+            fragments=(*state.fragments[:fragment], updated, *state.fragments[fragment + 1 :]),
+            estimates=estimates.updated(sync=sync_seconds),
+        )
+
+    async def _resolved(self) -> "_Run":
+        if self._run is None:
+            manifest = await self._exchange.manifest()
+            if manifest is None:
+                raise ProtocolError("manifest.json does not exist: no learner has created the run")
+            schema, plan = manifest.schema, manifest.plan
+            specs = {spec.name: spec for spec in schema.parameters}
+            self._run = _Run(
+                manifest=manifest,
+                digest=schema.digest(),
+                deltas=tuple(
+                    layout(
+                        Schema(tuple(specs[name] for name in names)), self._algorithm.delta_dtype
+                    )
+                    for names in plan
+                ),
+                schedule=self._algorithm.schedule(len(plan)),
+                aggregators=tuple(
+                    self._algorithm.aggregator(schema, names, self._backend) for names in plan
+                ),
+            )
+        return self._run
+
+    def _timeout(self, record: RoundRecord) -> float:
+        """How long to wait for the round after ``record`` before taking it over."""
+        return record.estimates.timeout(self._poll_interval)
+
+    def _timeout_of(self, progress: Progress) -> float:
+        return (
+            progress.timeout
+            if isinstance(progress, Following)
+            else progress.estimates.timeout(self._poll_interval)
+        )
+
+
+@dataclass(frozen=True)
+class _Run:
+    """What the manifest fixes for the life of the run; ``deltas`` is what a snapshot of each
+    fragment must hold."""
+
+    manifest: Manifest
+    digest: str
+    deltas: tuple[Schema[PayloadDType], ...]
+    schedule: RoundSchedule
+    aggregators: tuple[Aggregator, ...]

gradientmail-0.0.1.dist-info/METADATA

@@ -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},
+}
+```