Metadata-Version: 2.5
Name: gradientmail
Version: 0.0.2
Summary: Distributed training through asynchronous message exchange
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
Author-email: Gabriel Francisco <gabfssilva@gmail.com>
License-Expression: MIT
License-File: LICENSE
Requires-Python: >=3.12
Requires-Dist: safetensors<1,>=0.8
Provides-Extra: jax
Requires-Dist: jax; extra == 'jax'
Provides-Extra: keras
Requires-Dist: keras; extra == 'keras'
Provides-Extra: mlx
Requires-Dist: mlx; extra == 'mlx'
Provides-Extra: s3
Requires-Dist: boto3>=1.36; extra == 's3'
Provides-Extra: tensorflow
Requires-Dist: tensorflow; extra == 'tensorflow'
Provides-Extra: torch
Requires-Dist: torch; extra == 'torch'
Provides-Extra: zstd
Requires-Dist: zstandard>=0.23; extra == 'zstd'
Description-Content-Type: text/markdown

# GradientMail

Distributed training with [Decoupled DiLoCo](https://arxiv.org/abs/2604.21428) over S3-compatible storage. Supports PyTorch and MLX.

![Decoupled DiLoCo training](https://ppgsi-lab.github.io/gradientmail/docs/assets/hero.svg)

## Install

Python 3.12+.

```sh
pip install "gradientmail[s3,torch]"
```

Use `mlx` instead of `torch` for MLX on Apple Silicon.

## Usage

Use your model, optimizer and dataloader:

```python
import gradientmail as gm

session = gm.Session(
    model,
    algorithm=gm.DiLoCo(sync_every=100, fragments=8),
    mailbox=gm.S3Mailbox(bucket="my-training", prefix="experiments/run-1"),
)

def train_step(batch) -> int:
    loss = train(model, optimizer, batch)
    return len(batch)

session.run(train_step, dataloader)
```

Use the same bucket and prefix across machines. A node's `role` picks its part: `"both"`, the default, trains and closes rounds; `"learner"` only trains; `"syncer"` only closes rounds, for a machine that does not train, such as a laptop; `"either"` trains until the run elects it its syncer. With `ranks=N`, each node holds one of `N` fixed slices of the data. See the [PyTorch](https://github.com/ppgsi-lab/gradientmail/blob/main/examples/torch_learner.py) and [MLX](https://github.com/ppgsi-lab/gradientmail/blob/main/examples/mlx_learner.py) examples for complete setups.

[Documentation](https://ppgsi-lab.github.io/gradientmail/docs/guide/) · [Examples](examples/) · [Posts](https://ppgsi-lab.github.io/gradientmail/docs/posts/)

## 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},
}
```

## License

[MIT](LICENSE).
