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

rlx_distributed/
partition.rs
