DistribAI, a new training framework

DistribAI, a.k.a. Distributed AI, is a new framework under the NAXIUM OSS foundation for training models.

---

Imagine you have many devices with GPUs. Or a large community, but no hardware. How do you train a model?

The answer is and isn't simple: Distributed Training.

https://github.com/naxium-oss/DistribAI is built to train models on many different devices, whether on hardware scattered across the internet or around the house; or both. Users and team members of teams using the grid submit jobs to the admin to be accepted or declined, an orchestrator splits them into micro-tasks, contributor nodes execute them, and verified work lands in a signed credit ledger.

What DistribAI actually is

DistribAI is a distributed training control plane, and it's refreshingly concrete about what that means. Three processes make a grid:

Operators host the orchestrator and publish a join address. Contributors run a node - a desktop app, a daemon, or a join kit for Colab, Kaggle, or a spare VPS. There's also a full CLI and a Textual-based TUI (distribai tui) over the same admin API, so the whole grid is operable from a terminal or CI pipeline.

The life of a job

1. Submit. An organization posts a job over REST (distribai submit ./mytrainer for a script folder, or --recipe job.yaml for a human spec). Training jobs are native PyTorch, either built from architecture_config (decoder transformers, MoE decoders, LSTMs, hybrid attention/RNNs, and more, with size profiles from tiny to XL) or shipped as a sandboxed run.py tarball with AST validation at the door.
2. Split and schedule. The orchestrator decomposes the job into micro-tasks with priority tiers (P0-P3) and, roughly every 5 seconds, pairs queued work with idle capacity.
3. Assign. A worker on a long-lived bidirectional gRPC stream receives a TaskAssign: model weights URL, batch URL, hyperparameters, a deadline, and - the interesting part - an execution paradigm such as sync_cohort_ddp, elastic_sync, or async_federated_round, plus the usual RANK/WORLD_SIZE environment for multi-node cohorts.
4. Train and report. The node trains, streams throttled progress, uploads its gradient or artifact, and the orchestrator re-verifies the result.

Miss a heartbeat (~30 s degraded, ~50 s offline) or blow a deadline and the micro-task is requeued elsewhere. Consumer hardware disappears mid-step all the time; the grid just shrugs.

The trick: stop talking so much

Classic data-parallel training synchronizes gradients every step - fine inside a datacenter, fatal over home internet. DistribAI's answer is a first-class implementation of DiLoCo-style training (arXiv:2311.08105): each worker runs $HH$ inner AdamW steps entirely locally, then reports only a pseudo-gradient:

$$g_i={\theta }_{\text{start}}-{\theta }_{\text{after H steps}}g\_i\; =\; \backslash theta\_\{\backslash text\{start\}\}\; -\; \backslash theta\_\{\backslash text\{after\; H\; steps\}\}$$

The orchestrator averages the pseudo-gradients, applies an outer Nesterov step, and broadcasts new canonical weights via DiLoCoRoundComplete. Bandwidth drops by roughly a factor of $HH$ ; often ~500× versus all-reduce - which is the difference between "distributed training over the internet" being a slideware bullet point and actually working.

And it's not just pretraining. DistribAI also ships distributed GRPO rounds for RL-style fine-tuning: workers generate candidate responses per prompt, report only reward scalars upstream, the orchestrator normalizes advantages across the whole fleet's group, and workers apply the clipped-surrogate + KL update locally. One weight blob down, a few scalars up - the wire stays nearly silent.

Verification on a grid you don't own

Contributors are anonymous strangers with GPUs, so DistribAI treats trust as an engineering problem:

• Byzantine-aware aggregation: median, trimmed mean, and Multi-Krum-style outlier filtering sit between raw results and model updates, plus gradient-norm screens on both sides of the wire.
• Registration challenges: proof-of-computation challenges and hardware fingerprinting make Sybil joins expensive.
• A ledger, not a vibes counter: credits are append-only, hash-chained rows with signatures (a Merkle root for verification), with reliability and early-adopter multipliers. Credits pay for verified work and weight credit-based priority votes on which jobs run next.
• The boring-but-vital layer: optional TLS/mTLS, short-lived JWTs, rate limiting, and schema-validated everything.

Is it a cryptocurrency? No. Credits are an accounting mechanism that makes contributing legible and abuse expensive.

Spin one up

Three terminals, per the README:

# 1) Orchestrator
python -m services_python.orchestrator_grpc

# 2) Contributor dashboard
npm install && node client/server.js

Or skip the browsers entirely:

pip install -e .
distribai node start
distribai job create my-model 1000
distribai tui   # live Overview / Nodes / Jobs / Credits

Contributors without a dashboard can join from a Colab notebook (examples/colab/join_grid.ipynb) in minutes; a laptop on Wi-Fi is a legitimate grid citizen. If you'd rather not run from source, the packaging pipeline builds per-audience artifacts: a community DistribAI-Node, an operator DistribAI-Server, and a standalone admin CLI.

Some minor limits

DistribAI is at 0.9.0 pre-release, and has some limits: the orchestrator and SQLite are comfortable under roughly a thousand nodes, with documented sharding strategies beyond that. Training is Pytorch only today; JAX/Flax, TensorFlow/Keras, Apple MLX, and an ONNX Runtime fast path are on the roadmap, not shipped.

But the hard parts that usually kill community-compute projects; scheduling over flaky nodes, low-bandwidth optimization, adversarial robustness, and incentive accounting - are real code here, not a whitepaper. If you have spare cycles, or a backlog of experiments and no A100 budget, the grid could use you.

• Code, docs, and RFCs: https://github.com/naxium-oss/DistribAI
• License: Apache 2.0
• Part of the NAXIUM OSS foundation

Give it a try and star our repo(it takes just a press)!

If you run into any issues, or want to suggest a feature(or something isn't supported), open a Github issue!