rlx_distributed/lib.rs
1// RLX models — distributed inference.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Multi-node distributed inference for RLX models.
17//!
18//! Layered on top of `rlx-driver`'s transports (`TcpTransport`,
19//! `ThunderboltTransport`, and `rlx-mlx`'s `MlxTransport`) and
20//! `ProcessGroup`:
21//!
22//! - [`config`] — `hosts.json` parsing + process-group construction
23//! ([`DistConfig::connect`]).
24//! - [`partition`] — pipeline-parallel layer assignment
25//! ([`pipeline_layer_range`], [`block_role`]).
26//! - [`pipeline`] — the model-agnostic relay ([`PipelineCoordinator`])
27//! driving per-rank [`BlockRunner`]s.
28//!
29//! A model family provides a [`BlockRunner`] for its layer block (e.g.
30//! `rlx_qwen3`'s `Qwen3PipelineStage`); everything else here is reusable.
31//!
32//! ```no_run
33//! use rlx_distributed::{DistConfig, ParallelMode, PipelineCoordinator};
34//!
35//! # fn demo(rank: u32, mut runner: impl rlx_distributed::BlockRunner) -> anyhow::Result<()> {
36//! let cfg = DistConfig::load("hosts.json", Some(rank), ParallelMode::Pipeline)?;
37//! let group = cfg.connect()?; // blocks until the mesh forms
38//! let coord = PipelineCoordinator::new(group);
39//!
40//! let mut tokens = vec![/* prompt ids */];
41//! for _ in 0..32 {
42//! let tok = coord.forward_step(&mut runner, &tokens, |logits| argmax(logits))?;
43//! tokens.push(tok);
44//! }
45//! coord.barrier()?; // before dropping the group
46//! # Ok(()) }
47//! # fn argmax(_l: &[f32]) -> u32 { 0 }
48//! ```
49
50pub mod config;
51pub mod launch;
52pub mod partition;
53pub mod pipeline;
54
55pub use config::{DistConfig, Hostfile, ParallelMode, TransportBackend};
56pub use launch::{LocalCluster, WorkerArgs, free_loopback_ports, worker_args};
57pub use partition::{BlockRole, block_role, pipeline_layer_range};
58pub use pipeline::{BlockInput, BlockOutput, BlockRunner, PipelineCoordinator};
59
60// Re-export the transport primitives so model crates depend only on
61// `rlx-distributed`, not `rlx-driver` directly.
62pub use rlx_driver::{
63 NetTransport, ProcessGroup, ReduceKind, TcpTransport, ThunderboltTransport, Transport,
64};