rlx_distributed/
partition.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Pipeline-parallel layer partitioning.
6//!
7//! Layers are split into `world_size` contiguous blocks and assigned to
8//! ranks in **reverse**: rank 0 owns the *last* block (so it produces the
9//! logits and samples), rank `world-1` owns the *first* block (so it does
10//! the token embedding). This matches mlx-lm's `PipelineMixin` and means a
11//! token's activations flow rank `world-1 → … → 1 → 0`.
12//!
13//! The split is exact (every layer assigned exactly once) and as even as
14//! possible: the first `num_layers % world` forward-blocks get one extra
15//! layer.
16
17use std::ops::Range;
18
19/// Where a rank sits in the pipeline. Drives whether it embeds tokens,
20/// just transforms hidden states, or produces logits.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum BlockRole {
23    /// `world == 1`: this rank runs the whole model (embed → layers → head).
24    Single,
25    /// Highest rank: owns the first layers, embeds the input tokens, and
26    /// sends hidden states toward rank 0.
27    First,
28    /// Interior rank: receives hidden states, runs its layers, forwards.
29    Middle,
30    /// Rank 0: owns the last layers, runs the final norm + LM head, and
31    /// produces logits to sample from.
32    Last,
33}
34
35/// Role of `rank` in a `world`-sized pipeline.
36pub fn block_role(rank: u32, world: u32) -> BlockRole {
37    if world <= 1 {
38        BlockRole::Single
39    } else if rank == world - 1 {
40        BlockRole::First
41    } else if rank == 0 {
42        BlockRole::Last
43    } else {
44        BlockRole::Middle
45    }
46}
47
48/// The contiguous, half-open range of layer indices `[start, end)` that
49/// `rank` owns in a `world`-sized pipeline.
50///
51/// rank 0 → last block, rank `world-1` → first block. Panics if
52/// `rank >= world`.
53pub fn pipeline_layer_range(num_layers: usize, rank: u32, world: u32) -> Range<usize> {
54    let world = world.max(1) as usize;
55    let rank = rank as usize;
56    assert!(rank < world, "rank {rank} out of range for world {world}");
57
58    let per = num_layers / world;
59    let extra = num_layers % world;
60    // This rank owns forward-block `j` (reverse assignment).
61    let j = world - 1 - rank;
62    // Forward block i has `per + (i < extra)` layers; the first `extra`
63    // blocks carry the remainder.
64    let size_of = |i: usize| per + if i < extra { 1 } else { 0 };
65    let start: usize = (0..j).map(size_of).sum();
66    let len = size_of(j);
67    start..(start + len)
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn roles_map_correctly() {
76        assert_eq!(block_role(0, 1), BlockRole::Single);
77        assert_eq!(block_role(0, 4), BlockRole::Last);
78        assert_eq!(block_role(3, 4), BlockRole::First);
79        assert_eq!(block_role(1, 4), BlockRole::Middle);
80        assert_eq!(block_role(2, 4), BlockRole::Middle);
81    }
82
83    #[test]
84    fn even_split_is_reversed() {
85        // 16 layers / 4 ranks: rank 0 = last 4, rank 3 = first 4.
86        assert_eq!(pipeline_layer_range(16, 0, 4), 12..16);
87        assert_eq!(pipeline_layer_range(16, 1, 4), 8..12);
88        assert_eq!(pipeline_layer_range(16, 2, 4), 4..8);
89        assert_eq!(pipeline_layer_range(16, 3, 4), 0..4);
90    }
91
92    #[test]
93    fn uneven_split_tiles_exactly() {
94        // 10 layers / 4 ranks: forward block sizes 3,3,2,2.
95        // reverse → rank0=[8,10), rank1=[6,8), rank2=[3,6), rank3=[0,3).
96        assert_eq!(pipeline_layer_range(10, 0, 4), 8..10);
97        assert_eq!(pipeline_layer_range(10, 1, 4), 6..8);
98        assert_eq!(pipeline_layer_range(10, 2, 4), 3..6);
99        assert_eq!(pipeline_layer_range(10, 3, 4), 0..3);
100
101        // Reassembled in rank order (world-1 down to 0) covers [0,10) once.
102        let mut covered = [0u8; 10];
103        for r in 0..4u32 {
104            for i in pipeline_layer_range(10, r, 4) {
105                covered[i] += 1;
106            }
107        }
108        assert!(covered.iter().all(|&c| c == 1), "every layer covered once");
109    }
110
111    #[test]
112    fn single_rank_owns_all() {
113        assert_eq!(pipeline_layer_range(28, 0, 1), 0..28);
114    }
115
116    #[test]
117    fn more_ranks_than_layers_is_safe() {
118        // 2 layers, 4 ranks: two ranks get a layer, two get empty ranges.
119        let total: usize = (0..4u32).map(|r| pipeline_layer_range(2, r, 4).len()).sum();
120        assert_eq!(total, 2);
121    }
122}
rlx_distributed/partition.rs

rlx_distributed/
partition.rs
