kieve-moe

Expert-parallel Mixture-of-Experts training over iroh-tinker, built for the GCP fleet (CPU/GPU mix) and integrating with 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

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

Train (local, in-process simulated world)

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:

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.