rlx_distributed/cluster/
config.rs

1// RLX — versatile ML compiler + runtime.
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! **Declarative cluster config** (TOML) — the DX surface. A coordinator loads
5//! one file describing the model, the placement policy, and per-node device /
6//! precision / RNG / KV-cache / RAM settings, then probes + plans + runs.
7//!
8//! ```toml
9//! model = "mlx-community/DeepSeek-V4-Flash-2bit-DQ"   # HF id or local dir
10//! seq = 6
11//! rng_seed = 42            # global; a node may override
12//! reserve_ram_gb = 6.0     # headroom left free on every node
13//!
14//! [placement]
15//! policy = "ram_balanced"  # ram_balanced | throughput | manual
16//!
17//! [[node]]
18//! addr = "127.0.0.1:9100"
19//! ssh = "node-a"                     # ~/.ssh/config alias — for probe / launch / weight sync
20//! ckpt_dir = "/path/to/DeepSeek-V4-Flash-2bit-DQ"
21//! device = "cpu"                     # cpu | metal | cuda | ane | vulkan | "metal+cpu"
22//! precision = "bf16"                 # f32 | f16 | bf16 | mixed
23//! kv_cache = "host"                  # none | host | device
24//! max_ram_gb = 44                    # cap (e.g. leave room for other apps)
25//! # layers = "0:12"                  # manual override (policy = "manual")
26//! ```
27
28use anyhow::{Context, Result};
29use rlx_runtime::{Device, parse_device_list};
30use serde::{Deserialize, Serialize};
31use std::ops::Range;
32use std::path::Path;
33
34/// Where a node keeps the attention KV cache across decode steps.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
36#[serde(rename_all = "lowercase")]
37pub enum KvPolicy {
38    /// No persistent cache (prefill-only / recompute).
39    #[default]
40    None,
41    /// Keep KV in host RAM (works for any device; PCIe copy on GPU).
42    Host,
43    /// Keep KV resident on the compute device (fastest; device-mem permitting).
44    Device,
45}
46
47/// How the coordinator assigns layers to nodes.
48#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
49#[serde(rename_all = "snake_case")]
50pub enum PlacementPolicy {
51    /// Balance so each node's stage fills a similar fraction of its RAM budget
52    /// (fits first; the default — a model too big for any one node still runs).
53    #[default]
54    RamBalanced,
55    /// Balance per-stage compute time by node throughput (GFLOP/s) — minimizes
56    /// the pipeline's critical path when everything already fits.
57    Throughput,
58    /// Use each node's explicit `layers` range verbatim.
59    Manual,
60}
61
62/// Per-node settings.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct NodeConfig {
65    /// `host:port` the coordinator connects to (and the worker binds).
66    pub addr: String,
67    /// SSH alias/host for probing, launching, and syncing weights (None = local).
68    #[serde(default)]
69    pub ssh: Option<String>,
70    /// Local checkpoint directory ON THAT NODE.
71    pub ckpt_dir: String,
72    /// Device string (`cpu`, `metal`, `cuda`, `ane`, `vulkan`, or `a+b` combo —
73    /// first is primary, rest are CPU/host spill for weights that don't fit).
74    #[serde(default = "default_device")]
75    pub device: String,
76    /// Numeric precision string. Float widths: `f32` | `f16` | `bf16` | `mixed`.
77    /// Scaled-GEMM shorthands: `fp8` (e4m3) | `fp8e5m2` | `mxfp8` | `nvfp4` |
78    /// `mxfp4`. Or ANY minifloat by name — `f8e4m3`, `f6e3m2`, `f4e2m1`,
79    /// `f4e3m0`, … (the `fNeXmY` family). The model crate maps it to compile
80    /// flags; exotic formats change numerics and want backend/tensor-core support.
81    #[serde(default = "default_precision")]
82    pub precision: String,
83    #[serde(default)]
84    pub kv_cache: KvPolicy,
85    /// Per-node RNG seed (overrides the global). Keeps sampling reproducible and
86    /// *distinct* per stage when desired.
87    #[serde(default)]
88    pub rng_seed: Option<u64>,
89    /// Hard cap on this node's stage RAM (GB). The planner never exceeds it.
90    #[serde(default)]
91    pub max_ram_gb: Option<f64>,
92    /// Manual `"a:b"` layer range (used when policy = manual, or to pin a node).
93    #[serde(default)]
94    pub layers: Option<String>,
95}
96
97fn default_device() -> String {
98    "cpu".into()
99}
100fn default_precision() -> String {
101    "bf16".into()
102}
103
104impl NodeConfig {
105    /// Parsed device preference list (primary first). CPU spill devices follow.
106    pub fn devices(&self) -> Result<Vec<Device>> {
107        parse_device_list(&self.device.replace('+', ","))
108            .map_err(|e| anyhow::anyhow!("node {}: bad device `{}`: {e:?}", self.addr, self.device))
109    }
110    /// Primary compute device (first in the list).
111    pub fn primary_device(&self) -> Device {
112        self.devices()
113            .ok()
114            .and_then(|d| d.into_iter().next())
115            .unwrap_or(Device::Cpu)
116    }
117    /// True if the node lists a CPU/host spill target after its primary.
118    pub fn cpu_offload(&self) -> bool {
119        self.devices()
120            .map(|d| d.len() > 1 && d.contains(&Device::Cpu))
121            .unwrap_or(false)
122    }
123    /// Manual layer range if given.
124    pub fn manual_range(&self) -> Option<Range<usize>> {
125        let s = self.layers.as_ref()?;
126        let (a, b) = s.split_once(':')?;
127        Some(a.trim().parse().ok()?..b.trim().parse().ok()?)
128    }
129}
130
131/// Whole-cluster config.
132#[derive(Debug, Clone, Serialize, Deserialize)]
133pub struct ClusterConfig {
134    /// HF repo id or a local directory holding the checkpoint.
135    pub model: String,
136    /// Prompt length the pipeline graphs are built for.
137    #[serde(default = "default_seq")]
138    pub seq: usize,
139    /// Global RNG seed (a node's `rng_seed` overrides it).
140    #[serde(default)]
141    pub rng_seed: Option<u64>,
142    /// RAM (GB) to leave free on every node beyond its stage.
143    #[serde(default = "default_reserve")]
144    pub reserve_ram_gb: f64,
145    #[serde(default)]
146    pub placement: PlacementSection,
147    /// The nodes, in pipeline order.
148    #[serde(rename = "node")]
149    pub nodes: Vec<NodeConfig>,
150}
151
152#[derive(Debug, Clone, Default, Serialize, Deserialize)]
153pub struct PlacementSection {
154    #[serde(default)]
155    pub policy: PlacementPolicy,
156}
157
158fn default_seq() -> usize {
159    8
160}
161fn default_reserve() -> f64 {
162    6.0
163}
164
165impl ClusterConfig {
166    pub fn from_toml_str(s: &str) -> Result<Self> {
167        toml::from_str(s).context("parse cluster TOML")
168    }
169    pub fn from_path(p: impl AsRef<Path>) -> Result<Self> {
170        let s = std::fs::read_to_string(p.as_ref())
171            .with_context(|| format!("read {}", p.as_ref().display()))?;
172        Self::from_toml_str(&s)
173    }
174    /// Effective RNG seed for node `i` (per-node override, else global, else 0).
175    pub fn seed_for(&self, i: usize) -> u64 {
176        self.nodes
177            .get(i)
178            .and_then(|n| n.rng_seed)
179            .or(self.rng_seed)
180            .unwrap_or(0)
181    }
182}
rlx_distributed/cluster/config.rs

rlx_distributed/cluster/
config.rs
