rlx_distributed/cluster/
placement.rs

1// RLX — versatile ML compiler + runtime.
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! **Hardware-aware placement** — split a layer stack into contiguous per-node
5//! stages that (a) FIT each node's RAM/VRAM budget and (b) balance the pipeline.
6//! This replaces hand-tuned `--layers 0:15` splits: give it the model's resident
7//! cost per layer and the probed [`NodeCaps`], get an assignment back.
8
9use super::caps::NodeCaps;
10use super::config::{NodeConfig, PlacementPolicy};
11use anyhow::{Result, bail};
12use std::ops::Range;
13
14/// Resident-memory + compute cost model for a layer-stack model. Bytes are the
15/// ACTUAL per-node resident footprint (packed weights + kept scales), so the
16/// planner's budgets are meaningful.
17#[derive(Debug, Clone)]
18pub struct ModelCost {
19    pub n_layers: usize,
20    /// Average resident bytes per transformer layer.
21    pub per_layer_bytes: u64,
22    /// Extra bytes on the FIRST stage (token embedding).
23    pub embed_bytes: u64,
24    /// Extra bytes on the LAST stage (final norm + LM head).
25    pub head_bytes: u64,
26    /// Relative FLOPs per layer (for the throughput policy; 1.0 = uniform).
27    pub per_layer_flops: f64,
28}
29
30impl ModelCost {
31    pub fn total_bytes(&self) -> u64 {
32        self.embed_bytes + self.head_bytes + self.per_layer_bytes * self.n_layers as u64
33    }
34}
35
36/// One node's stage.
37#[derive(Debug, Clone)]
38pub struct Assignment {
39    pub addr: String,
40    pub ssh: Option<String>,
41    pub layers: Range<usize>,
42    pub first: bool,
43    pub last: bool,
44    /// Estimated resident bytes for this stage (weights + embed/head).
45    pub est_bytes: u64,
46    /// The node's stage RAM budget the planner used.
47    pub budget_bytes: u64,
48    /// Primary device label for the monitor.
49    pub device: String,
50}
51
52/// Fraction of a resident GPU's reported memory usable for one stage's largest
53/// single allocation — the rest is driver reserve + arena/weight-buffer slack.
54const GPU_ALLOC_SAFETY: f64 = 0.85;
55
56/// A node's usable stage-RAM budget. `max_ram_gb`, when set, is the EXPLICIT
57/// budget (the operator's control — overrides a pessimistic `ram_avail`, e.g.
58/// macOS reporting reclaimable cache as used); otherwise the probed available
59/// RAM is used. A *resident* discrete-GPU primary further caps at its VRAM; a
60/// *paged* GPU (CUDA managed memory, or Apple/iGPU unified memory) migrates from
61/// host RAM and so stays bounded by RAM. The reserve headroom is subtracted last.
62fn budget_bytes(caps: &NodeCaps, cfg: &NodeConfig, reserve: u64) -> u64 {
63    let mut b = match cfg.max_ram_gb {
64        Some(gb) => ((gb * 1e9) as u64).min(caps.ram_total),
65        None => caps.ram_avail,
66    };
67    let dev = cfg.primary_device();
68    // CUDA runs on managed (paged) memory here — the stage migrates over PCIe from
69    // host RAM, so it is bounded by RAM, not VRAM. Other discrete GPUs are resident.
70    let paged = matches!(dev, rlx_runtime::Device::Cuda);
71    if dev != rlx_runtime::Device::Cpu && !paged {
72        let vram = caps.accel_mem();
73        if vram > 0 && vram < caps.ram_total {
74            // A single VkDeviceMemory (or discrete cudaMalloc) can't consume the
75            // *whole* reported ceiling: the driver reserves some, the mapped
76            // activation arena + weight buffer are separate allocations, and
77            // alignment adds slack. Cap the stage at 85% of the GPU's memory so
78            // its largest single allocation actually succeeds instead of OOM-ing
79            // right at the ceiling (an amdgpu APU's GTT is especially tight).
80            b = b.min((vram as f64 * GPU_ALLOC_SAFETY) as u64);
81        }
82    }
83    b.saturating_sub(reserve)
84}
85
86/// Max whole layers a node can hold given its byte budget and the per-stage
87/// overhead it must also carry (embed on the first node, head on the last).
88fn layer_cap(budget: u64, per_layer: u64, overhead: u64) -> usize {
89    if per_layer == 0 {
90        return usize::MAX;
91    }
92    (budget.saturating_sub(overhead) / per_layer) as usize
93}
94
95/// Plan contiguous stages. `nodes` is in pipeline order (first node runs the
96/// embedding, last runs the head). Returns one [`Assignment`] per node.
97pub fn plan_placement(
98    model: &ModelCost,
99    nodes: &[(NodeCaps, NodeConfig)],
100    policy: PlacementPolicy,
101    reserve_bytes: u64,
102) -> Result<Vec<Assignment>> {
103    let k = nodes.len();
104    if k == 0 {
105        bail!("no nodes");
106    }
107    // Manual: honour each node's explicit range.
108    if policy == PlacementPolicy::Manual {
109        let mut out = Vec::new();
110        for (i, (caps, cfg)) in nodes.iter().enumerate() {
111            let r = cfg.manual_range().ok_or_else(|| {
112                anyhow::anyhow!("node {} has no `layers` for manual policy", cfg.addr)
113            })?;
114            out.push(mk_assignment(
115                model,
116                caps,
117                cfg,
118                r,
119                i == 0,
120                i == k - 1,
121                reserve_bytes,
122            ));
123        }
124        return Ok(out);
125    }
126
127    let budgets: Vec<u64> = nodes
128        .iter()
129        .map(|(c, cfg)| budget_bytes(c, cfg, reserve_bytes))
130        .collect();
131    // Per-node capacity in layers (first pays embed, last pays head).
132    let caps_layers: Vec<usize> = (0..k)
133        .map(|i| {
134            let overhead = if i == 0 { model.embed_bytes } else { 0 }
135                + if i == k - 1 { model.head_bytes } else { 0 };
136            layer_cap(budgets[i], model.per_layer_bytes, overhead)
137        })
138        .collect();
139    let total_cap: usize = caps_layers
140        .iter()
141        .copied()
142        .map(|c| c.min(model.n_layers))
143        .sum();
144    if total_cap < model.n_layers {
145        bail!(
146            "model ({} layers, {:.1} GB) does not fit the cluster (capacity {} layers). \
147             Lower precision, add a node, raise max_ram, or reduce reserve.",
148            model.n_layers,
149            model.total_bytes() as f64 / 1e9,
150            total_cap
151        );
152    }
153
154    // Weight each node by budget (ram_balanced) or throughput (gflops).
155    let weight: Vec<f64> = match policy {
156        PlacementPolicy::Throughput => nodes
157            .iter()
158            .map(|(c, _)| c.gflops.max(1.0) * model.per_layer_flops.max(0.001))
159            .collect(),
160        _ => budgets.iter().map(|&b| b as f64).collect(),
161    };
162    let wsum: f64 = weight.iter().sum();
163
164    // Proportional target, clamped to per-node layer cap, remainder water-filled.
165    let mut counts: Vec<usize> = (0..k)
166        .map(|i| ((model.n_layers as f64 * weight[i] / wsum).round() as usize).min(caps_layers[i]))
167        .collect();
168    let mut assigned: usize = counts.iter().sum();
169    // Add under-target: give leftover layers to nodes with remaining capacity,
170    // most spare first. Remove over-target similarly.
171    while assigned < model.n_layers {
172        let i = (0..k)
173            .filter(|&i| counts[i] < caps_layers[i])
174            .max_by_key(|&i| caps_layers[i] - counts[i]);
175        match i {
176            Some(i) => {
177                counts[i] += 1;
178                assigned += 1;
179            }
180            None => bail!("placement: could not distribute remaining layers"),
181        }
182    }
183    while assigned > model.n_layers {
184        let i = (0..k).filter(|&i| counts[i] > 0).max_by_key(|&i| counts[i]);
185        match i {
186            Some(i) => {
187                counts[i] -= 1;
188                assigned -= 1;
189            }
190            None => break,
191        }
192    }
193
194    let mut out = Vec::with_capacity(k);
195    let mut start = 0usize;
196    for (i, ((caps, cfg), &cnt)) in nodes.iter().zip(&counts).enumerate() {
197        let r = start..start + cnt;
198        start += cnt;
199        out.push(mk_assignment(
200            model,
201            caps,
202            cfg,
203            r,
204            i == 0,
205            i == k - 1,
206            reserve_bytes,
207        ));
208    }
209    Ok(out)
210}
211
212fn mk_assignment(
213    model: &ModelCost,
214    caps: &NodeCaps,
215    cfg: &NodeConfig,
216    layers: Range<usize>,
217    first: bool,
218    last: bool,
219    reserve: u64,
220) -> Assignment {
221    let n = (layers.end - layers.start) as u64;
222    let est = n * model.per_layer_bytes
223        + if first { model.embed_bytes } else { 0 }
224        + if last { model.head_bytes } else { 0 };
225    Assignment {
226        addr: cfg.addr.clone(),
227        ssh: cfg.ssh.clone(),
228        layers,
229        first,
230        last,
231        est_bytes: est,
232        budget_bytes: budget_bytes(caps, cfg, reserve),
233        device: rlx_runtime::device_label(cfg.primary_device()).to_string(),
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    fn caps(addr: &str, ram_gb: f64, gflops: f64) -> NodeCaps {
242        NodeCaps {
243            addr: addr.into(),
244            os: "linux".into(),
245            cores: 8,
246            ram_total: (ram_gb * 1e9) as u64,
247            ram_avail: (ram_gb * 1e9) as u64,
248            disk_free: 500_000_000_000,
249            devices: vec![],
250            gflops,
251            io_mbps: 0.0,
252        }
253    }
254    fn node(addr: &str) -> NodeConfig {
255        NodeConfig {
256            addr: addr.into(),
257            ssh: None,
258            ckpt_dir: "/x".into(),
259            device: "cpu".into(),
260            precision: "bf16".into(),
261            kv_cache: Default::default(),
262            rng_seed: None,
263            max_ram_gb: None,
264            layers: None,
265        }
266    }
267
268    #[test]
269    fn ram_balanced_fits_and_covers_all_layers() {
270        // 43 layers @ ~3GB each ≈ 129GB across 3 uneven nodes.
271        let model = ModelCost {
272            n_layers: 43,
273            per_layer_bytes: 3_000_000_000,
274            embed_bytes: 200_000_000,
275            head_bytes: 200_000_000,
276            per_layer_flops: 1.0,
277        };
278        let nodes = vec![
279            (caps("a", 60.0, 100.0), node("a")),
280            (caps("b", 55.0, 90.0), node("b")),
281            (caps("c", 44.0, 60.0), node("c")),
282        ];
283        let plan =
284            plan_placement(&model, &nodes, PlacementPolicy::RamBalanced, 5_000_000_000).unwrap();
285        // Contiguous + complete cover of 0..43.
286        assert_eq!(plan[0].layers.start, 0);
287        assert_eq!(plan.last().unwrap().layers.end, 43);
288        for w in plan.windows(2) {
289            assert_eq!(w[0].layers.end, w[1].layers.start);
290        }
291        // Every stage within budget.
292        for a in &plan {
293            assert!(
294                a.est_bytes <= a.budget_bytes,
295                "{} over budget: {} > {}",
296                a.addr,
297                a.est_bytes,
298                a.budget_bytes
299            );
300        }
301    }
302
303    #[test]
304    fn rejects_when_too_big() {
305        let model = ModelCost {
306            n_layers: 100,
307            per_layer_bytes: 5_000_000_000,
308            embed_bytes: 0,
309            head_bytes: 0,
310            per_layer_flops: 1.0,
311        };
312        let nodes = vec![(caps("a", 40.0, 100.0), node("a"))];
313        assert!(
314            plan_placement(&model, &nodes, PlacementPolicy::RamBalanced, 5_000_000_000).is_err()
315        );
316    }
317}
rlx_distributed/cluster/placement.rs

rlx_distributed/cluster/
placement.rs
