kieve-moe 0.1.0__tar.gz

kieve_moe-0.1.0/PKG-INFO

@@ -0,0 +1,77 @@
+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.

kieve_moe-0.1.0/README.md

@@ -0,0 +1,62 @@
+# 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.

kieve_moe-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "kieve-moe"
+version = "0.1.0"
+description = "Expert-parallel MoE training over iroh-tinker: fractal tiling + procedural data + discrete diffusion"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+dependencies = [
+    "numpy>=1.24",
+    "torch>=2.0",
+    "iroh-tinker>=0.5.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0",
+    "ruff>=0.5",
+    "mypy>=1.10",
+]
+
+[project.scripts]
+kieve-moe = "kieve_moe.cli:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+extend-select = ["E", "F", "I", "UP", "B", "SIM", "T20"]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["T20"]
+
+[tool.mypy]
+python_version = "3.10"
+ignore_missing_imports = true

kieve_moe-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

kieve_moe-0.1.0/src/kieve_moe/__init__.py

@@ -0,0 +1,16 @@
+"""kieve-moe: expert-parallel MoE training over iroh-tinker.
+
+This is the scaling layer for BertBM-style Mixture-of-Experts models: the
+dense GeGLU MoE block (BertBM ``moe.py``) is split across ranks so each
+process owns a shard of the experts, a router sends tokens to the rank that
+owns their selected experts, and activations flow back over the
+``iroh_tinker.comm.Transport`` interface (in-process ``LocalTransport`` for
+tests, ``SocketTransport``/native transports for real hosts). Gradients are
+synchronized with the ring allreduce on the same transports, so the whole
+training loop is decentralized — no NCCL/gloo process group required.
+"""
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__"]

kieve_moe-0.1.0/src/kieve_moe/cli.py

@@ -0,0 +1,100 @@
+"""``kieve-moe`` command line interface.
+
+Subcommands:
+
+``train``       run the distributed MoE training loop
+``encode``      not yet implemented (reserved)
+
+Every command supports ``--json``.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from .config import MoEConfig
+from .distributed import DistributedTrainer, TrainConfig
+
+__all__ = ["build_parser", "main"]
+
+
+def _load_config(path: str) -> MoEConfig:
+    raw = json.loads(Path(path).read_text())
+    known = {f for f in MoEConfig.__dataclass_fields__ if f != "extra"}  # type: ignore[attr-defined]
+    return MoEConfig(**{k: v for k, v in raw.items() if k in known})
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="kieve-moe",
+        description="Expert-parallel MoE training over iroh-tinker.",
+    )
+    parser.add_argument("--json", action="store_true", help="emit JSON on stdout")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    train = sub.add_parser("train", help="run the distributed MoE training loop")
+    train.add_argument("--config", required=True, help="config JSON file")
+    train.add_argument("--world", type=int, default=1, help="world size")
+    train.add_argument("--rank", type=int, default=0, help="this rank")
+    train.add_argument("--steps", type=int, default=20, help="number of steps")
+    train.add_argument("--lr", type=float, default=3e-4, help="learning rate")
+    train.add_argument(
+        "--transport", default="local", choices=["local", "socket"], help="transport"
+    )
+    train.add_argument("--rdv", default=".cache/iroh_tinker", help="rendezvous dir (socket)")
+    train.add_argument("--host", default="127.0.0.1", help="bind host (socket)")
+    train.add_argument("--json", action="store_true", help="emit JSON on stdout")
+    return parser
+
+
+def run_train(args: argparse.Namespace) -> int:
+    cfg = _load_config(args.config)
+    cfg.world_size = args.world
+    train_cfg = TrainConfig(steps=args.steps, lr=args.lr)
+
+    if args.world == 1 or args.transport == "local":
+        from iroh_tinker.comm import LocalNetwork
+
+        net = LocalNetwork(args.world)
+        transport: Any = net.transport(args.rank)
+    else:
+        from iroh_tinker.comm import SocketTransport
+
+        transport = SocketTransport(
+            rank=args.rank, world_size=args.world, rendezvous_dir=args.rdv, host=args.host
+        )
+
+    try:
+        trainer = DistributedTrainer(args.rank, args.world, transport, cfg, train_cfg)
+        out = trainer.run()
+    finally:
+        transport.close()
+    if args.json:
+        sys.stdout.write(json.dumps(out, indent=2) + "\n")
+    else:
+        pass
+    return 0
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(list(argv) if argv is not None else None)
+    try:
+        if args.command == "train":
+            return run_train(args)
+    except (FileNotFoundError, ValueError, RuntimeError) as exc:
+        if getattr(args, "json", False):
+            sys.stdout.write(json.dumps({"error": str(exc)}) + "\n")
+        else:
+            pass
+        return 2
+    parser.error(f"unknown command {args.command!r}")
+    return 2
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

kieve_moe-0.1.0/src/kieve_moe/config.py

@@ -0,0 +1,90 @@
+"""Model configuration for kieve-moe.
+
+The dataclass mirrors the BertBM/ModernBERT MoE fields (``moe_experts``,
+``moe_top_k``, ``moe_aux_loss``) so a config written for BertBM can be reused
+verbatim; a local ``MoEConfig`` keeps the package self-contained when the
+``bertbm`` package is not installed.
+"""
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+__all__ = ["MoEConfig", "model_from_config", "as_dict"]
+
+
+@dataclass
+class MoEConfig:
+    # architecture
+    vocab_size: int = 256
+    hidden_size: int = 64
+    intermediate_size: int = 128
+    num_layers: int = 2
+    num_experts: int = 8
+    top_k: int = 2
+    aux_loss_coef: float = 0.01
+    max_seq_len: int = 64
+    # training
+    lr: float = 3e-4
+    weight_decay: float = 0.0
+    grad_clip: float = 1.0
+    batch_size: int = 8
+    # architecture enhancements
+    fractal_levels: int = 0      # >0: fractal-tiled hierarchical routing (branch**levels experts)
+    fractal_branch: int = 2
+    diffusion: bool = False       # add discrete (masked) diffusion denoising head
+    procedural: bool = True       # generate training tokens procedurally (no dataset needed)
+    # expert placement: expert i is owned by rank (i % world_size)
+    world_size: int = 1
+    # hidden knobs
+    layernorm_eps: float = 1e-5
+    seed: int = 42
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def validate(self) -> None:
+        if self.num_experts < 1:
+            raise ValueError("num_experts must be >= 1")
+        if not 0 < self.top_k <= self.num_experts:
+            raise ValueError(f"top_k must be in 1..{self.num_experts}")
+        if self.hidden_size < 1 or self.intermediate_size < 1:
+            raise ValueError("hidden/intermediate sizes must be positive")
+        if self.num_layers < 1:
+            raise ValueError("num_layers must be >= 1")
+        if self.world_size < 1:
+            raise ValueError("world_size must be >= 1")
+        if self.aux_loss_coef < 0:
+            raise ValueError("aux_loss_coef must be non-negative")
+        if self.fractal_levels > 0:
+            if self.fractal_levels > 0 and self.fractal_branch < 2:
+                raise ValueError("fractal_branch must be >= 2")
+            total = self.fractal_branch ** self.fractal_levels
+            if self.num_experts != total:
+                raise ValueError(
+                    f"fractal tiling requires num_experts == branch**levels "
+                    f"({self.fractal_branch}**{self.fractal_levels}={total})"
+                )
+
+    @classmethod
+    def from_bertbm(cls, cfg: Any, **overrides: Any) -> MoEConfig:
+        """Adopt a BertBM ModernBERTConfig (or any object with the fields)."""
+        return cls(
+            vocab_size=int(getattr(cfg, "vocab_size", 256)),
+            hidden_size=int(getattr(cfg, "hidden_size", 64)),
+            intermediate_size=int(getattr(cfg, "intermediate_size", 128)),
+            num_layers=int(getattr(cfg, "num_hidden_layers", 2)),
+            num_experts=int(getattr(cfg, "moe_experts", 8)),
+            top_k=int(getattr(cfg, "moe_top_k", 2)),
+            aux_loss_coef=float(getattr(cfg, "moe_aux_loss", 0.01)),
+            **overrides,
+        )
+
+
+def model_from_config(cfg: MoEConfig, transport: Any | None = None, rank: int = 0):
+    """Build the MoE encoder (licensed import to avoid the torch import here)."""
+    from .model import MoEEncoder
+
+    return MoEEncoder(cfg, transport=transport, rank=rank)
+
+
+def as_dict(cfg: MoEConfig) -> dict[str, Any]:
+    return asdict(cfg)

kieve_moe-0.1.0/src/kieve_moe/diffusion.py

@@ -0,0 +1,73 @@
+"""Discrete diffusion auxiliary head for generative token modeling.
+
+Implements an absorbing-state (masked) discrete diffusion over the token
+vocabulary — a lightweight MDLM. During training the head is offered a batch
+of sequences, a random timestep is sampled, a fraction of tokens are masked and
+the head predicts the original tokens from the masked sequence (a denoising/
+unmasking objective added to the MoE cross-entropy). Sampling iteratively
+unmasks from a seed, giving a procedural generative extension on top of the
+MoE encoder.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ["DiscreteDiffusionHead"]
+
+
+class DiscreteDiffusionHead(nn.Module):
+    def __init__(self, vocab_size: int, hidden_size: int, mask_token: int | None = None):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.mask_token = vocab_size if mask_token is None else mask_token
+        self.classifier = nn.Linear(hidden_size, vocab_size + 1, bias=False)
+
+    @property
+    def mask_id(self) -> int:
+        return self.mask_token
+
+    def loss(self, hidden: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
+        """Masked-diffusion denoising loss over a batch."""
+        b, t, d = hidden.shape
+        rng = torch.ones((b, t), device=hidden.device).bernoulli_(0.15)
+        masked = ids.clone()
+        masked[rng.bool()] = self.mask_token
+        logits = self.classifier(hidden)  # [B,T,V+1]
+        target = ids[rng.bool()] if rng.any() else ids
+        if rng.any():
+            loss = F.cross_entropy(logits[rng.bool()], target)
+        else:
+            loss = torch.zeros((), device=hidden.device)
+        return loss
+
+    def sample(
+        self,
+        hidden_fn,
+        seed: torch.Tensor,
+        steps: int,
+        temperature: float = 1.0,
+    ) -> torch.Tensor:
+        """Iteratively unmask seed tokens using the encoder hidden states."""
+        ids = seed.clone()
+        schedule = torch.linspace(0.15, 0.0, steps)
+        with torch.no_grad():
+            for frac in schedule:
+                hidden = hidden_fn(ids)
+                logits = self.classifier(hidden)
+                probs = F.softmax(logits / temperature, dim=-1)
+                pred = probs.argmax(-1)
+                mask = (ids == self.mask_token)
+                if not bool(mask.any()):
+                    break
+                # progressively unmask a fraction
+                n_unmask = int(mask.sum() * float(frac))
+                candidates = mask.nonzero(as_tuple=False)
+                if n_unmask > 0 and candidates.numel() > 0:
+                    pick = candidates[:n_unmask]
+                    ids[pick[:, 0], pick[:, 1]] = pred[pick[:, 0], pick[:, 1]]
+                else:
+                    ids[mask] = pred[mask]
+        ids[ids == self.mask_token] = pred[ids == self.mask_token]
+        return ids

kieve_moe-0.1.0/src/kieve_moe/distributed.py

@@ -0,0 +1,139 @@
+"""Distributed training loop for kieve-moe.
+
+Data-parallel gradient synchronization over iroh-tinker transports plus the
+expert-parallel MoE layers. The loop is deliberately small and transport-
+agnostic: any ``Transport`` implementation (in-process ``LocalTransport``,
+TCP ``SocketTransport``, native mesh) can be used, and gradients are averaged
+with the ring allreduce from ``iroh_tinker.comm``.
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+from .config import MoEConfig
+from .model import MoEEncoder
+
+__all__ = ["TrainConfig", "DistributedTrainer"]
+
+
+@dataclass
+class TrainConfig:
+    steps: int = 20
+    lr: float = 3e-4
+    log_every: int = 5
+    seed: int = 42
+    checkpoint_dir: str = ".cache/kieve-moe"
+    resume: bool = True
+
+
+class DistributedTrainer:
+    """Owner of a rank's model, transport and optimizer loop."""
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        transport: Any,
+        cfg: MoEConfig,
+        train_cfg: TrainConfig | None = None,
+        moe_network: Any | None = None,
+    ):
+        self.rank = rank
+        self.world_size = world_size
+        self.transport = transport
+        self.cfg = cfg
+        cfg.world_size = world_size
+        self.train_cfg = train_cfg or TrainConfig()
+        torch.manual_seed(cfg.seed + rank)
+        # Dedicated MoE dispatch transport so the expert RPC never collides
+        # with the gradient collective frames on the shared transport.
+        self._moe_transport: Any | None = None
+        if world_size > 1:
+            from iroh_tinker.comm import LocalNetwork
+
+            if moe_network is None:
+                moe_network = LocalNetwork(world_size)
+            self._moe_net = moe_network
+            self._moe_transport = self._moe_net.transport(rank)
+        self.model = MoEEncoder(
+            cfg,
+            transport=transport,
+            rank=rank,
+            moe_transport=self._moe_transport,
+        )
+        self._param_map = list(self.model.parameters())
+
+    def _step_data(self, step: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """Local data shard: a deterministic stream unique per rank."""
+        if self.cfg.procedural:
+            from .procedural import make_stream
+
+            stream = make_stream(self.cfg, rank=self.rank, world_size=self.world_size)
+            return stream.batch()
+        g = torch.Generator().manual_seed(
+            self.train_cfg.seed + step * 1000 + self.rank
+        )
+        ids = torch.randint(
+            0, self.cfg.vocab_size, (self.cfg.batch_size, self.cfg.max_seq_len), generator=g
+        )
+        targets = torch.randint(
+            0, self.cfg.vocab_size, (self.cfg.batch_size, self.cfg.max_seq_len), generator=g
+        )
+        return ids, targets
+
+    def _allreduce_grads(self) -> None:
+        """Average every parameter gradient across ranks (ring allreduce)."""
+        from iroh_tinker.comm import RingAllreduce
+
+        if self.world_size <= 1:
+            return
+        flat: dict[str, np.ndarray] = {}
+        shapes: dict[str, tuple[int, ...]] = {}
+        for i, p in enumerate(self._param_map):
+            if p.grad is None:
+                continue
+            key = f"p{i}"
+            shapes[key] = tuple(p.grad.shape)
+            flat[key] = p.grad.detach().cpu().numpy().reshape(-1).astype(np.float32)
+        reduced = RingAllreduce(self.transport).allreduce(flat)
+        for key, arr in reduced.items():
+            self._param_map[int(key[1:])].grad = (
+                torch.from_numpy(arr).reshape(shapes[key]).to(self._param_map[int(key[1:])].device)
+            )
+
+    def run(self) -> dict[str, Any]:
+        t0 = time.perf_counter()
+        start_step = 0
+        if self.train_cfg.resume and self.rank == 0:
+            import json
+            from pathlib import Path
+
+            state_path = Path(self.train_cfg.checkpoint_dir) / "state.json"
+            if state_path.exists():
+                try:
+                    start_step = int(json.loads(state_path.read_text()).get("step", 0))
+                except Exception:
+                    start_step = 0
+
+        for step in range(start_step, self.train_cfg.steps):
+            ids, targets = self._step_data(step)
+            self.model.step(ids, targets, self.train_cfg.lr)
+            self._allreduce_grads()
+            self.transport.barrier()
+            if step % self.train_cfg.log_every == 0 and self.rank == 0:
+                pass
+
+        self.transport.barrier()
+        elapsed = time.perf_counter() - t0
+        return {
+            "rank": self.rank,
+            "world_size": self.world_size,
+            "steps": self.train_cfg.steps,
+            "elapsed_s": round(elapsed, 3),
+            "status": "ok",
+        }