Metadata-Version: 2.4
Name: iroh-tinker
Version: 0.3.0
Summary: Distributed training framework using iroh for 100k+ node scale across PyTorch, JAX, Taichi and more
Author: Research Team
License: MIT OR Apache-2.0
Project-URL: Homepage, https://github.com/scairesearch/iroh_tinker
Project-URL: Repository, https://github.com/scairesearch/iroh_tinker
Project-URL: Issues, https://github.com/scairesearch/iroh_tinker/issues
Keywords: distributed-training,iroh,p2p,allreduce,pytorch,jax,taichi,mlx
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.24
Provides-Extra: torch
Requires-Dist: torch>=2.1; extra == "torch"
Provides-Extra: jax
Requires-Dist: jax>=0.4.20; extra == "jax"
Requires-Dist: jaxlib>=0.4.20; extra == "jax"
Provides-Extra: taichi
Requires-Dist: taichi>=1.6; extra == "taichi"
Provides-Extra: mlx
Requires-Dist: mlx>=0.5; extra == "mlx"
Provides-Extra: iroh
Requires-Dist: iroh>=0.1.0; extra == "iroh"
Provides-Extra: colossalai
Requires-Dist: colossalai>=0.4; extra == "colossalai"
Provides-Extra: cloud
Requires-Dist: pyyaml>=6.0; extra == "cloud"
Provides-Extra: aws
Requires-Dist: boto3>=1.34; extra == "aws"
Provides-Extra: gcs
Requires-Dist: google-cloud-storage>=2.10; extra == "gcs"
Provides-Extra: serve
Requires-Dist: fastapi>=0.110; extra == "serve"
Requires-Dist: uvicorn[standard]>=0.27; extra == "serve"
Requires-Dist: pydantic>=2.5; extra == "serve"
Requires-Dist: httpx>=0.27; extra == "serve"
Provides-Extra: sky
Requires-Dist: skypilot; extra == "sky"
Provides-Extra: spotty
Requires-Dist: spotty>=1.3; extra == "spotty"
Provides-Extra: dev
Requires-Dist: pytest>=7.4; extra == "dev"
Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
Requires-Dist: ruff>=0.4; extra == "dev"
Requires-Dist: mypy>=1.8; extra == "dev"
Requires-Dist: build>=1.0; extra == "dev"
Requires-Dist: twine>=5.0; extra == "dev"
Provides-Extra: zamlpy
Requires-Dist: zamlpy>=0.2.0; extra == "zamlpy"

# iroh-tinker

Distributed training that scales from a laptop to **1,000,000 nodes** using
[iroh](https://www.iroh.computer/) for peer-to-peer QUIC gradient
synchronization — no cluster scheduler, no rendezvous server, no static IPs.

You write a normal training loop against a Tinker-style API; iroh-tinker
handles the distributed gradient allreduce, checkpoint sharding, fault
tolerance, and selects the right communication algorithm for your cluster size.

```python
import iroh_tinker as it

client = it.TrainingClient(base_model="Qwen/Qwen3-8B", rank=0, world_size=8,
                           backend="pytorch", lora_rank=32)

for batch in dataloader:
    client.forward_backward(batch, loss_fn="cross_entropy")
    client.optim_step(it.AdamParams(learning_rate=1e-4))
    if client.step % 100 == 0:
        client.save_state()
```

## Quickstart

```bash
# install (from a checkout)
pip install -e .

# or with optional backends
pip install -e ".[torch,jax]" 
```

```python
import iroh_tinker as it

client = it.TrainingClient(base_model="tiny", backend="numpy", rank=0, world_size=1)
```

Use the CLI:

```bash
iroh-tinker backends            # list backends
iroh-tinker device              # detect CUDA/ROCm/RISC-V/CPU
iroh-tinker topology 100000     # show fanout/levels for 100k nodes
iroh-tinker train --steps 20    # short local training loop
iroh-tinker bench --world-size 8  # allreduce micro-benchmarks
iroh-tinker --json device       # machine-readable output
```

## Why

Existing distributed trainers (DDP, FSDP, ColossalAI) assume a tightly-coupled
data center: shared filesystem, NCCL, a rendezvous host. iroh-tinker instead
routes gradient exchange over iroh's QUIC connections, which hole-punch through
NATs and fall back to relays, so ranks can live in different clouds or regions.
The communication layer is *hierarchical*, so no node ever talks to more than a
bounded number of peers regardless of total cluster size.

## Features

- **Multi-framework.** First-class backends for PyTorch (LoRA on HuggingFace
  models), JAX, Taichi (CPU/CUDA/Metal/Vulkan/OpenGL), MLX, and a pure-NumPy
  fallback for CI and CPU-only nodes.
- **Real P2P transport.** `IrohTransport` runs gradient exchange over iroh QUIC:
  ticket-based rendezvous, one connection per rank, and hub-relay so any-to-any
  collectives work on a bounded-fan-out overlay. An in-process transport mirrors
  it for tests and simulations.
- **Hierarchical allreduce.** Ring (< 100 ranks), Tree (100–10k), and Butterfly
  (≥ 10k) with a hierarchical wrapper, selected automatically.
- **Tinker-compatible API.** `forward_backward`, `optim_step`, `sample`,
  `save_state` / `load_state`, plus `ServiceClient` / `SamplingClient`.
- **Gradient compression.** 8-bit block quantization, top-k sparsification and
  scheme-coupled error feedback, plus varint-indexed top-k schemes that cut
  bandwidth by up to ~1863x.
- **Sharded checkpointing.** Content-addressed shards (SHA-256) with a manifest,
  so a recovering rank fetches only its own shard.
- **Rust core (optional).** A high-performance core for networking, allreduce,
  and scheduling lives under `src/` and is exposed behind optional Cargo
  features.
- **Multi-cloud launch.** A SkyPilot-compatible `task.yaml` and thin native
  drivers for GCP/AWS/Azure/SSH, with spot failover and reading the same
  rendezvous as training — no VPC peering required.
- **OpenAI-compatible serving.** A vLLM-style engine (continuous batching,
  `/v1/chat/completions`, SSE streaming) with tensor parallelism over iroh.

## Install

```bash
# Python package (NumPy only, works everywhere)
pip install -e python

# with a specific framework backend
pip install -e "python[torch]"
pip install -e "python[jax]"
pip install -e "python[taichi]"
pip install -e "python[mlx]"
```

From PyPI, the latest release is `0.3.0`:

```bash
pip install "iroh-tinker==0.3.0"
```

`0.3.0` adds the varint compression schemes, the real TCP transport and the
50M-parameter benchmark harness. To work against an unreleased commit:

```bash
pip install "iroh-tinker @ git+https://github.com/scairesearch/iroh_tinker.git"
```

## Quickstarts

### Single process, any backend

```bash
python examples/train_all_backends.py     # runs every backend that's installed
```

### Simulated multi-rank cluster (in-process, no network)

```bash
python examples/train_local_cluster.py --world-size 8 --backend numpy
```

After synchronization every rank reports the **same parameter checksum** — proof
that the allreduce is exact.

### Real distributed run

Run one process per rank (each on its own machine or container). Ranks discover
each other over iroh and synchronize gradients peer-to-peer:

```bash
# on each node i of 8 (shared rendezvous dir, e.g. NFS or a KV store):
RANK=$i WORLD=8 USE_IROH=1 python examples/train_pytorch.py
```

To try the real iroh path in one process (requires `pip install iroh`):

```bash
python examples/train_iroh_cluster.py --world-size 4 --steps 10
```

### Topology scaling

```bash
python examples/scale_topology.py
```

```
nodes=        8  levels=1  fanout=8      max_children=0
nodes=      100  levels=4  fanout=100    max_children=99
nodes=   10,000  levels=4  fanout=1,000  max_children=1,000
nodes=  100,000  levels=4  fanout=1,000  max_children=1,000
nodes=1,000,000  levels=4  fanout=10,000 max_children=10,000
```

## Multi-cloud launch (SkyPilot-compatible)

Launch the same job across providers with a thin, native driver per cloud —
no VPC peering, no central scheduler. Tasks are described with a SkyPilot
`task.yaml` superset, so existing task files work unchanged:

```yaml
name: train
resources:
  ordered:                 # try in order, fail over automatically
    - {infra: gcp, accelerators: A100:8, use_spot: true}
    - {infra: aws, accelerators: p5.48xlarge, use_spot: true}
workdir: .
setup: pip install -e ".[torch]"
run: torchrun --nproc-per-node 8 train.py
```

```bash
iroh-tinker launch task.yaml          # provisions, installs, runs, renders
iroh-tinker launch task.yaml --dry-run
iroh-tinker jobs                      # list your clusters
iroh-tinker logs train --rank 0
iroh-tinker stop train

iroh-tinker sky launch task.yaml -c my-cluster   # SkyPilot CLI verbs
iroh-tinker sky status | down my-cluster | logs my-cluster
iroh-tinker check                     # which engines / CLIs are available
```

Native engines wrap `gcloud` / `aws` / `az` (and SSH hosts); SkyPilot and Spotty
are optional fallback engines. Ranks rendezvous through a shared ticket store
(local dir, S3, or GCS) and then talk directly over iroh — the same code path
whether the nodes share a data center or not. Set
`IROH_TINKER_RENDEZVOUS_SECRET` on every rank to HMAC-sign tickets (wrong or
tampered tickets raise `PermissionError`). Down ranks are detected via pid
checks on `status` and surfaced in `last_error`. Install the optional
`sky`/`vllm` PATH shims with `iroh-tinker --install-shims`.

## Serving (OpenAI-compatible, vLLM-style)

An inference server with continuous batching and a vLLM-shaped offline API:

```bash
iroh-tinker serve --model tiny --tensor-parallel-size 1 --port 8000
```

```bash
curl localhost:8000/v1/chat/completions -H 'content-type: application/json' \
  -d '{"model":"tiny","messages":[{"role":"user","content":"hi"}],"max_tokens":16}'
```

Works with `openai`, LangChain, and llama-index clients (`base_url=http://host:8000/v1`).
Streaming (`"stream": true`, SSE), `/v1/models`, and `/v1/completions` are
supported. Offline generation mirrors vLLM:

```python
from iroh_tinker.serving import LLM, SamplingParams

llm = LLM("tiny", tensor_parallel_size=1)
out = llm.generate(["hello"], SamplingParams(max_tokens=16, temperature=0.7))
print(out[0].outputs[0].text)
```

Tensor parallelism shards the LM head across ranks and all-gathers logits over
iroh; `serving/tp.py` is validated for bit-equality against single-rank output.

## Architecture

See [`ARCHITECTURE.md`](ARCHITECTURE.md) for the full design. In brief:

```
                    ┌─────────────────────────────────────────┐
                    │            TrainingClient               │
                    │  forward_backward → optim_step → save   │
                    └───────────────┬─────────────────────────┘
                                    │
             ┌──────────────────────┼──────────────────────┐
             ▼                      ▼                      ▼
      ┌────────────┐        ┌──────────────┐       ┌──────────────┐
      │  Backends  │        │  Allreduce   │       │ Checkpointing│
      │ torch/jax/ │        │ ring/tree/   │       │ sha-256      │
      │ taichi/mlx │        │ butterfly    │       │ shards       │
      └────────────┘        └──────┬───────┘       └──────────────┘
                                   ▼
                          ┌──────────────────┐
                          │    Transport     │
                          │ iroh QUIC │ local │
                          └──────────────────┘
```

## Backend support

| Backend   | Models                          | Autodiff       | Device               |
|-----------|---------------------------------|----------------|----------------------|
| PyTorch   | HuggingFace causal LMs + LoRA   | `autograd`     | CUDA / MPS / CPU     |
| JAX       | custom transformer (Flax-ready) | `jax.grad`     | CPU / GPU / TPU      |
| Taichi    | tied-embedding MLP              | hand-written   | CPU/CUDA/Metal/Vulkan|
| MLX       | tied-embedding MLP              | `value_and_grad`| Apple silicon       |
| NumPy     | tied-embedding MLP              | hand-written   | CPU                  |

All backends implement the same `TrainingBackend` interface
(`python/iroh_tinker/backends/base.py`), so the distributed loop is identical
regardless of framework.

## Testing

```bash
# Python
PYTHONPATH=python pytest tests -q

# Rust core
cargo test
```

The transport/allreduce suite runs the real `IrohTransport` code path against an
in-process fake of the `iroh` API (`tests/_fake_iroh.py`), so the relay protocol,
barrier, and end-to-end distributed training are covered without a network.

## Performance

At 100k+ nodes the allreduce is bandwidth-bound. `CompressedTransport` compresses
each gradient buffer on the wire using self-describing frames
(`python/iroh_tinker/comm/compression.py`):

| Scheme    | Wire reduction (1M params, 1% keep) | Typical use              |
|-----------|------------------------------------|--------------------------|
| `8bit`    | ~3.8x                              | accuracy-sensitive       |
| `topk`    | ~50x                               | sparse gradients         |
| `topk8bit`| ~80x                               | aggressive bandwidth cut |
| `topk_varint` | ~1000x (0.1% keep)        | extreme bandwidth cut (fp16 values, varint indices) |
| `topk4_varint` | ~1670x (0.1% keep)       | extreme bandwidth cut (4-bit values, varint indices) |
| `topk2_varint` | ~1863x (0.1% keep)       | extreme bandwidth cut (2-bit values, varint indices) |

The three `*_varint` schemes keep a sparse top-k selection but encode the
selected indices as varint deltas and the values at low precision. Their
error-feedback path is scheme-coupled (`ErrorFeedbackCodec`), so the residual
compensates for exactly the loss the wire format introduces.

On a simulated 1 Gbps / 10us-latency link (`benchmarks/bench_bandwidth.py`), ring
allreduce of a 1M-element gradient shows ~2.6x (`8bit`), ~6.5x (`topk`) and
~6.3x (`topk8bit`) wall-clock speedup vs. uncompressed fp32. Real RDMA/InfiniBand
links push this higher since compression CPU cost becomes negligible relative to
the bandwidth saved on multi-GB models.

```bash
# compression ratio vs. precision
PYTHONPATH=python python benchmarks/bench_compression.py

# allreduce algorithm scaling (star/ring/tree/butterfly) with/without compression
PYTHONPATH=python python benchmarks/bench_allreduce.py

# bandwidth-limited end-to-end speedup
PYTHONPATH=python python benchmarks/bench_bandwidth.py
```

## Hardware & topology

`python/iroh_tinker/device.py` detects CUDA (NVIDIA), ROCm/HIP (AMD), RISC-V and
CPU devices. `python/iroh_tinker/hw.py` derives an interconnect hint (NVLink/PCIe
vs. Ethernet) and adjusts allreduce fanout/group-size parameters accordingly.
`TrainingConfig.device` auto-resolves to the best available device (or accepts an
explicit `cuda`/`rocm`/`cpu` string), and `TrainingClient.device` exposes it.

`python/iroh_tinker/comm/rdma.py` provides a pluggable verbs transport:
`PyverbsBinding` talks to real IB/RoCE hardware via `pyverbs` (open device, PD,
MR, CQ, QP), while `LoopbackVerbsBinding` covers in-process tests with the same
framing. `build_loopback_cluster(n)` creates `n` linked transports for CI.

`python/iroh_tinker/ofo.py` implements OFA-style elastic subnet sampling:
`SubnetSampler` draws width/depth/kernel/resolution multipliers per hardware
tier (CUDA/ROCm/CPU/RISC-V) and narrows them against latency/memory/throughput
targets, deterministically with a seed.

## Observability

`iroh_tinker.metrics` is a dependency-free metrics registry (counters, gauges,
histograms) with Prometheus text exposition. `TrainingClient` records allreduce
latency, gradient bytes (raw vs. on-the-wire), compression ratio, and step
counts automatically. The serving stack exposes `GET /metrics` with
request/token/time-to-first-token series, and `iroh-tinker metrics` dumps the
in-process registry (`--snapshot` for JSON).

## API stability & deprecation

The v1 public surface is: `TrainingClient`, `TrainingConfig`, `Transport`,
`AllreduceAlgorithm`, and `TrainingBackend`. These follow semver: breaking
changes only in major releases. Experimental surfaces (`VerbsRdmaTransport`
hardware path, `ofo` hooks, `compat.sky_cli`) may change in minor releases and
will emit `FutureWarning` at least one minor release before removal. Internal
modules (`_*`) are not covered by the stability guarantee.

## License

MIT OR Apache-2.0
