Metadata-Version: 2.4
Name: kieve-moe
Version: 0.1.0
Summary: Expert-parallel MoE training over iroh-tinker: fractal tiling + procedural data + discrete diffusion
License: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.24
Requires-Dist: torch>=2.0
Requires-Dist: iroh-tinker>=0.5.0
Provides-Extra: dev
Requires-Dist: pytest>=8.0; extra == "dev"
Requires-Dist: ruff>=0.5; extra == "dev"
Requires-Dist: mypy>=1.10; extra == "dev"

# kieve-moe

Expert-parallel **Mixture-of-Experts** training over
[iroh-tinker](https://github.com/scairesearch/iroh_tinker), built for the
GCP fleet (CPU/GPU mix) and integrating with
[elasticgrid](https://pypi.org/project/elasticgrid/) for launch/placement.

`kieve-moe` scales a BertBM/Kieve-style MoE encoder: experts are sharded
across ranks, tokens are routed to the rank that owns their selected experts
over an iroh-tinker transport, and gradients are averaged with the ring
allreduce on the same transport — no NCCL/gloo process group required.

## Features

- **Expert-parallel MoE** (`expert_parallel.py`): top-k GeGLU MoE with a
  load-balancing auxiliary loss; remote experts are served by an aggregator
  rank over a dedicated transport (requests and replies ride disjoint
  channels).
- **Distributed trainer** (`distributed.py`): data-parallel gradient sync,
  barriers, per-rank procedural data shards, resume-from-step checkpointing.
- **Fractal tiling routing** (`fractal.py`): self-similar hierarchical
  routing (branch^levels expert tree) mirroring the Taichi fractal-tiling
  playground — a cheap router for very large expert pools.
- **Procedural generation** (`procedural.py`): seeded pattern/LFSR token
  streams, so the fleet trains without shipping a dataset.
- **Discrete diffusion head** (`diffusion.py`): masked-absorbing-state (MDLM)
  denoising objective + iterative unmasking sampler, adding generative
  token modeling on top of the MoE encoder.
- **Parametric sizing** (`scaling.py`): the flagship **4.5B** configuration
  (15 layers, H=768, I=4096, E=32) with a live parameter estimate.

## Install

```bash
pip install kieve-moe
pip install -e ".[dev]"     # from a checkout
```

## Train (local, in-process simulated world)

```bash
kieve-moe train --config configs/moe-4-5b.json --world 2 --steps 20
```

## Train on the GCP fleet

Use iroh-tinker's `SocketTransport` (`--transport socket`) with a shared
rendezvous dir, and elasticgrid for placement/launch:

```bash
kieve-moe train --config configs/moe-4-5b.json --world 4 \
  --transport socket --rdv /tmp/rdv --host 0.0.0.0
```

Or build an elasticgrid `LaunchSpec` for the same config (see
`scripts/launch_elasticgrid.py`).

## Tests

`pytest tests -q` — 15 tests covering dense↔parallel MoE equivalence,
procedural streams, fractal routing, diffusion sampling, scale estimates, and
the distributed trainer loop.
