iroh-tinker 0.1.1__tar.gz

iroh_tinker-0.1.1/PKG-INFO

@@ -0,0 +1,392 @@
+Metadata-Version: 2.4
+Name: iroh-tinker
+Version: 0.1.1
+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 and top-k sparsification
+  with error feedback to cut bandwidth by up to 100×.
+- **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]"
+```
+
+## 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 |
+
+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