eredu_core/
topology.rs

1//! Pure parallel topology coordinates and membership validation.
2
3use serde::{Deserialize, Serialize};
4use std::ops::Range;
5
6/// Logical parallel axis.
7#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
8#[serde(rename_all = "snake_case")]
9#[non_exhaustive]
10pub enum ParallelAxis {
11    /// Tensor parallelism.
12    Tensor,
13    /// Pipeline parallelism.
14    Pipeline,
15    /// Expert parallelism.
16    Expert,
17    /// Data parallelism.
18    Data,
19}
20
21/// Coordinate of one rank in a four-dimensional topology.
22#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
23#[non_exhaustive]
24pub struct ParallelCoordinates {
25    /// Tensor coordinate.
26    tensor: usize,
27    /// Pipeline coordinate.
28    pipeline: usize,
29    /// Expert coordinate.
30    expert: usize,
31    /// Data coordinate.
32    data: usize,
33}
34
35impl ParallelCoordinates {
36    /// Creates an explicit Cartesian rank coordinate.
37    pub const fn new(tensor: usize, pipeline: usize, expert: usize, data: usize) -> Self {
38        Self {
39            tensor,
40            pipeline,
41            expert,
42            data,
43        }
44    }
45    /// Tensor coordinate.
46    pub const fn tensor(self) -> usize {
47        self.tensor
48    }
49    /// Pipeline coordinate.
50    pub const fn pipeline(self) -> usize {
51        self.pipeline
52    }
53    /// Expert coordinate.
54    pub const fn expert(self) -> usize {
55        self.expert
56    }
57    /// Data coordinate.
58    pub const fn data(self) -> usize {
59        self.data
60    }
61    /// Returns coordinates with a replaced tensor coordinate.
62    pub const fn with_tensor(mut self, tensor: usize) -> Self {
63        self.tensor = tensor;
64        self
65    }
66    /// Returns coordinates with a replaced pipeline coordinate.
67    pub const fn with_pipeline(mut self, pipeline: usize) -> Self {
68        self.pipeline = pipeline;
69        self
70    }
71    /// Returns coordinates with a replaced expert coordinate.
72    pub const fn with_expert(mut self, expert: usize) -> Self {
73        self.expert = expert;
74        self
75    }
76    /// Returns coordinates with a replaced data coordinate.
77    pub const fn with_data(mut self, data: usize) -> Self {
78        self.data = data;
79        self
80    }
81}
82
83/// Validated sizes for every parallel axis.
84#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize)]
85#[non_exhaustive]
86pub struct ParallelTopology {
87    /// Tensor-parallel size.
88    tensor: usize,
89    /// Pipeline-parallel size.
90    pipeline: usize,
91    /// Expert-parallel size.
92    expert: usize,
93    /// Data-parallel size.
94    data: usize,
95}
96
97impl<'de> Deserialize<'de> for ParallelTopology {
98    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
99    where
100        D: serde::Deserializer<'de>,
101    {
102        #[derive(Deserialize)]
103        struct RawTopology {
104            tensor: usize,
105            pipeline: usize,
106            expert: usize,
107            data: usize,
108        }
109
110        let raw = RawTopology::deserialize(deserializer)?;
111        Self::new(raw.tensor, raw.pipeline, raw.expert, raw.data).map_err(serde::de::Error::custom)
112    }
113}
114
115impl ParallelTopology {
116    /// Validates positive sizes and a representable world size.
117    pub fn new(
118        tensor: usize,
119        pipeline: usize,
120        expert: usize,
121        data: usize,
122    ) -> Result<Self, TopologyError> {
123        let sizes = [tensor, pipeline, expert, data];
124        if sizes.contains(&0) {
125            return Err(TopologyError::ZeroAxis);
126        }
127        sizes
128            .into_iter()
129            .try_fold(1usize, usize::checked_mul)
130            .ok_or(TopologyError::WorldSizeOverflow)?;
131        Ok(Self {
132            tensor,
133            pipeline,
134            expert,
135            data,
136        })
137    }
138    /// Tensor-parallel size.
139    pub const fn tensor(self) -> usize {
140        self.tensor
141    }
142    /// Pipeline-parallel size.
143    pub const fn pipeline(self) -> usize {
144        self.pipeline
145    }
146    /// Expert-parallel size.
147    pub const fn expert(self) -> usize {
148        self.expert
149    }
150    /// Data-parallel size.
151    pub const fn data(self) -> usize {
152        self.data
153    }
154    /// Total rank count.
155    pub fn world_size(self) -> usize {
156        self.tensor * self.pipeline * self.expert * self.data
157    }
158    /// Returns whether every parallel dimension is a singleton.
159    pub const fn is_replicated(self) -> bool {
160        self.tensor == 1 && self.pipeline == 1 && self.expert == 1 && self.data == 1
161    }
162    /// Returns whether an axis contains more than one rank.
163    pub const fn is_axis_active(self, axis: ParallelAxis) -> bool {
164        match axis {
165            ParallelAxis::Tensor => self.tensor > 1,
166            ParallelAxis::Pipeline => self.pipeline > 1,
167            ParallelAxis::Expert => self.expert > 1,
168            ParallelAxis::Data => self.data > 1,
169        }
170    }
171    /// Converts a pipeline-major, tensor, expert-minor rank to coordinates.
172    ///
173    /// Data parallelism is the outermost dimension. Within one data replica,
174    /// `rank = ((pipeline * tensor_size) + tensor) * expert_size + expert`.
175    pub fn coordinates(self, rank: usize) -> Result<ParallelCoordinates, TopologyError> {
176        if rank >= self.world_size() {
177            return Err(TopologyError::RankOutOfRange {
178                rank,
179                world_size: self.world_size(),
180            });
181        }
182        let expert = rank % self.expert;
183        let rank = rank / self.expert;
184        let tensor = rank % self.tensor;
185        let rank = rank / self.tensor;
186        let pipeline = rank % self.pipeline;
187        Ok(ParallelCoordinates {
188            tensor,
189            pipeline,
190            expert,
191            data: rank / self.pipeline,
192        })
193    }
194    /// Returns all ranks matching the supplied coordinate on non-selected axes.
195    pub fn axis_members(
196        self,
197        rank: usize,
198        axis: ParallelAxis,
199    ) -> Result<Vec<usize>, TopologyError> {
200        let mut coordinates = self.coordinates(rank)?;
201        let size = match axis {
202            ParallelAxis::Tensor => self.tensor,
203            ParallelAxis::Pipeline => self.pipeline,
204            ParallelAxis::Expert => self.expert,
205            ParallelAxis::Data => self.data,
206        };
207        (0..size)
208            .map(|coordinate| {
209                match axis {
210                    ParallelAxis::Tensor => coordinates.tensor = coordinate,
211                    ParallelAxis::Pipeline => coordinates.pipeline = coordinate,
212                    ParallelAxis::Expert => coordinates.expert = coordinate,
213                    ParallelAxis::Data => coordinates.data = coordinate,
214                }
215                self.rank_for(coordinates)
216            })
217            .collect()
218    }
219
220    /// Resolves Cartesian coordinates to the unique global rank.
221    pub fn rank_for(self, coordinates: ParallelCoordinates) -> Result<usize, TopologyError> {
222        for (coordinate, size, axis) in [
223            (coordinates.tensor, self.tensor, ParallelAxis::Tensor),
224            (coordinates.pipeline, self.pipeline, ParallelAxis::Pipeline),
225            (coordinates.expert, self.expert, ParallelAxis::Expert),
226            (coordinates.data, self.data, ParallelAxis::Data),
227        ] {
228            if coordinate >= size {
229                return Err(TopologyError::CoordinateOutOfRange {
230                    axis,
231                    coordinate,
232                    size,
233                });
234            }
235        }
236        coordinates
237            .data
238            .checked_mul(self.pipeline)
239            .and_then(|rank| rank.checked_add(coordinates.pipeline))
240            .and_then(|rank| rank.checked_mul(self.tensor))
241            .and_then(|rank| rank.checked_add(coordinates.tensor))
242            .and_then(|rank| rank.checked_mul(self.expert))
243            .and_then(|rank| rank.checked_add(coordinates.expert))
244            .ok_or(TopologyError::WorldSizeOverflow)
245    }
246}
247
248/// One validated rank in a backend-neutral Cartesian topology.
249#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize)]
250#[non_exhaustive]
251pub struct ParallelRankTopology {
252    /// Number of ranks in the complete topology.
253    world_size: usize,
254    /// Global rank represented by this value.
255    global_rank: usize,
256    /// Tensor-parallel rank count.
257    tensor_parallel_size: usize,
258    /// Tensor-parallel coordinate.
259    tensor_parallel_rank: usize,
260    /// Pipeline-parallel rank count.
261    pipeline_parallel_size: usize,
262    /// Pipeline-parallel coordinate.
263    pipeline_parallel_rank: usize,
264    /// Expert-parallel rank count.
265    expert_parallel_size: usize,
266    /// Expert-parallel coordinate.
267    expert_parallel_rank: usize,
268    /// Data-parallel rank count.
269    data_parallel_size: usize,
270    /// Data-parallel coordinate.
271    data_parallel_rank: usize,
272}
273
274impl ParallelRankTopology {
275    /// Validates and resolves one rank in `topology`.
276    pub fn new(topology: ParallelTopology, global_rank: usize) -> Result<Self, TopologyError> {
277        let topology = ParallelTopology::new(
278            topology.tensor,
279            topology.pipeline,
280            topology.expert,
281            topology.data,
282        )?;
283        let coordinates = topology.coordinates(global_rank)?;
284        Ok(Self {
285            world_size: topology.world_size(),
286            global_rank,
287            tensor_parallel_size: topology.tensor,
288            tensor_parallel_rank: coordinates.tensor,
289            pipeline_parallel_size: topology.pipeline,
290            pipeline_parallel_rank: coordinates.pipeline,
291            expert_parallel_size: topology.expert,
292            expert_parallel_rank: coordinates.expert,
293            data_parallel_size: topology.data,
294            data_parallel_rank: coordinates.data,
295        })
296    }
297
298    /// Number of ranks in the complete topology.
299    pub const fn world_size(self) -> usize {
300        self.world_size
301    }
302    /// Global rank represented by this value.
303    pub const fn global_rank(self) -> usize {
304        self.global_rank
305    }
306    /// Tensor-parallel rank count.
307    pub const fn tensor_parallel_size(self) -> usize {
308        self.tensor_parallel_size
309    }
310    /// Tensor-parallel coordinate.
311    pub const fn tensor_parallel_rank(self) -> usize {
312        self.tensor_parallel_rank
313    }
314    /// Pipeline-parallel rank count.
315    pub const fn pipeline_parallel_size(self) -> usize {
316        self.pipeline_parallel_size
317    }
318    /// Pipeline-parallel coordinate.
319    pub const fn pipeline_parallel_rank(self) -> usize {
320        self.pipeline_parallel_rank
321    }
322    /// Expert-parallel rank count.
323    pub const fn expert_parallel_size(self) -> usize {
324        self.expert_parallel_size
325    }
326    /// Expert-parallel coordinate.
327    pub const fn expert_parallel_rank(self) -> usize {
328        self.expert_parallel_rank
329    }
330    /// Data-parallel rank count.
331    pub const fn data_parallel_size(self) -> usize {
332        self.data_parallel_size
333    }
334    /// Data-parallel coordinate.
335    pub const fn data_parallel_rank(self) -> usize {
336        self.data_parallel_rank
337    }
338
339    /// Returns the complete topology shape.
340    pub fn topology(self) -> ParallelTopology {
341        ParallelTopology {
342            tensor: self.tensor_parallel_size,
343            pipeline: self.pipeline_parallel_size,
344            expert: self.expert_parallel_size,
345            data: self.data_parallel_size,
346        }
347    }
348
349    /// Returns whether every parallel dimension is a singleton.
350    pub const fn is_replicated(self) -> bool {
351        self.world_size == 1
352    }
353
354    /// Returns whether an axis contains more than one rank.
355    pub const fn is_axis_active(self, axis: ParallelAxis) -> bool {
356        match axis {
357            ParallelAxis::Tensor => self.tensor_parallel_size > 1,
358            ParallelAxis::Pipeline => self.pipeline_parallel_size > 1,
359            ParallelAxis::Expert => self.expert_parallel_size > 1,
360            ParallelAxis::Data => self.data_parallel_size > 1,
361        }
362    }
363
364    /// Returns this rank's Cartesian coordinates.
365    pub const fn coordinates(self) -> ParallelCoordinates {
366        ParallelCoordinates {
367            tensor: self.tensor_parallel_rank,
368            pipeline: self.pipeline_parallel_rank,
369            expert: self.expert_parallel_rank,
370            data: self.data_parallel_rank,
371        }
372    }
373
374    /// Resolves Cartesian coordinates to a global rank.
375    pub fn global_rank_for(self, coordinates: ParallelCoordinates) -> Result<usize, TopologyError> {
376        self.topology().rank_for(coordinates)
377    }
378
379    /// Returns topology-derived membership in one communication axis.
380    pub fn subgroup(self, axis: ParallelAxis) -> Result<SubgroupMembership, TopologyError> {
381        let coordinates = self.coordinates();
382        let global_ranks = self.topology().axis_members(self.global_rank, axis)?;
383        let rank = match axis {
384            ParallelAxis::Tensor => coordinates.tensor,
385            ParallelAxis::Pipeline => coordinates.pipeline,
386            ParallelAxis::Expert => coordinates.expert,
387            ParallelAxis::Data => coordinates.data,
388        };
389        let size = global_ranks.len();
390        let color = subgroup_color(self.topology(), coordinates, axis)?;
391        if global_ranks.get(rank).copied() != Some(self.global_rank) {
392            return Err(TopologyError::SubgroupIdentity { axis });
393        }
394        Ok(SubgroupMembership {
395            axis,
396            color,
397            rank,
398            size,
399            global_ranks,
400        })
401    }
402
403    /// Ordered global ranks participating in tensor collectives with this rank.
404    pub fn tensor_parallel_peers(self) -> Result<Vec<usize>, TopologyError> {
405        Ok(self.subgroup(ParallelAxis::Tensor)?.global_ranks)
406    }
407
408    /// Ordered global ranks participating in expert exchange with this rank.
409    pub fn expert_parallel_peers(self) -> Result<Vec<usize>, TopologyError> {
410        Ok(self.subgroup(ParallelAxis::Expert)?.global_ranks)
411    }
412
413    /// Preceding pipeline rank with matching coordinates on other axes.
414    pub fn pipeline_predecessor(self) -> Result<Option<usize>, TopologyError> {
415        if self.pipeline_parallel_rank == 0 {
416            return Ok(None);
417        }
418        self.global_rank_for(ParallelCoordinates {
419            pipeline: self.pipeline_parallel_rank - 1,
420            ..self.coordinates()
421        })
422        .map(Some)
423    }
424
425    /// Succeeding pipeline rank with matching coordinates on other axes.
426    pub fn pipeline_successor(self) -> Result<Option<usize>, TopologyError> {
427        if self.pipeline_parallel_rank + 1 == self.pipeline_parallel_size {
428            return Ok(None);
429        }
430        self.global_rank_for(ParallelCoordinates {
431            pipeline: self.pipeline_parallel_rank + 1,
432            ..self.coordinates()
433        })
434        .map(Some)
435    }
436
437    /// Whether this rank owns the stage-local embedding.
438    pub const fn owns_embedding(self) -> bool {
439        self.pipeline_parallel_rank == 0
440    }
441
442    /// Whether this rank owns the stage-local output head.
443    pub const fn owns_output_head(self) -> bool {
444        self.pipeline_parallel_rank + 1 == self.pipeline_parallel_size
445    }
446
447    /// This pipeline stage's balanced decoder-layer range.
448    pub fn layer_range(self, layers: usize) -> Result<Range<usize>, TopologyError> {
449        balanced_contiguous_range(
450            layers,
451            self.pipeline_parallel_size,
452            self.pipeline_parallel_rank,
453            false,
454        )
455    }
456
457    /// This expert rank's balanced routed-expert range.
458    pub fn expert_range(self, experts: usize) -> Result<Range<usize>, TopologyError> {
459        balanced_contiguous_range(
460            experts,
461            self.expert_parallel_size,
462            self.expert_parallel_rank,
463            false,
464        )
465    }
466
467    /// Whether this rank owns `layer` under balanced pipeline placement.
468    pub fn owns_layer(self, layers: usize, layer: usize) -> Result<bool, TopologyError> {
469        Ok(self.layer_range(layers)?.contains(&layer))
470    }
471
472    /// Whether this rank owns `expert` under balanced expert placement.
473    pub fn owns_expert(self, experts: usize, expert: usize) -> Result<bool, TopologyError> {
474        Ok(self.expert_range(experts)?.contains(&expert))
475    }
476
477    /// Validates optional layer/expert geometry before payload access.
478    pub fn preflight(
479        self,
480        decoder_layers: Option<usize>,
481        routed_experts: Option<usize>,
482    ) -> Result<TopologyPreflightReport, TopologyError> {
483        let local_layer_range = match decoder_layers {
484            Some(layers) => Some(self.layer_range(layers)?),
485            None if self.pipeline_parallel_size > 1 => {
486                return Err(TopologyError::MissingLayerCount)
487            }
488            None => None,
489        };
490        let local_expert_range = match routed_experts {
491            Some(experts) => Some(self.expert_range(experts)?),
492            None if self.expert_parallel_size > 1 => return Err(TopologyError::MissingExpertCount),
493            None => None,
494        };
495        Ok(TopologyPreflightReport {
496            topology: self,
497            tensor_subgroup: self.subgroup(ParallelAxis::Tensor)?,
498            pipeline_subgroup: self.subgroup(ParallelAxis::Pipeline)?,
499            expert_subgroup: self.subgroup(ParallelAxis::Expert)?,
500            data_subgroup: self.subgroup(ParallelAxis::Data)?,
501            local_layer_range,
502            local_expert_range,
503            owns_embedding: self.owns_embedding(),
504            owns_output_head: self.owns_output_head(),
505        })
506    }
507}
508
509impl<'de> Deserialize<'de> for ParallelRankTopology {
510    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
511    where
512        D: serde::Deserializer<'de>,
513    {
514        #[derive(Deserialize)]
515        struct RawRankTopology {
516            world_size: usize,
517            global_rank: usize,
518            tensor_parallel_size: usize,
519            tensor_parallel_rank: usize,
520            pipeline_parallel_size: usize,
521            pipeline_parallel_rank: usize,
522            expert_parallel_size: usize,
523            expert_parallel_rank: usize,
524            data_parallel_size: usize,
525            data_parallel_rank: usize,
526        }
527        let raw = RawRankTopology::deserialize(deserializer)?;
528        let topology = ParallelTopology::new(
529            raw.tensor_parallel_size,
530            raw.pipeline_parallel_size,
531            raw.expert_parallel_size,
532            raw.data_parallel_size,
533        )
534        .map_err(serde::de::Error::custom)?;
535        let value = Self::new(topology, raw.global_rank).map_err(serde::de::Error::custom)?;
536        if value.world_size != raw.world_size
537            || value.tensor_parallel_rank != raw.tensor_parallel_rank
538            || value.pipeline_parallel_rank != raw.pipeline_parallel_rank
539            || value.expert_parallel_rank != raw.expert_parallel_rank
540            || value.data_parallel_rank != raw.data_parallel_rank
541        {
542            return Err(serde::de::Error::custom(
543                "parallel rank topology contains inconsistent derived fields",
544            ));
545        }
546        Ok(value)
547    }
548}
549
550/// Topology-derived membership of one rank in an axis subgroup.
551#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
552#[non_exhaustive]
553pub struct SubgroupMembership {
554    /// Axis represented by this subgroup.
555    axis: ParallelAxis,
556    /// Deterministic subgroup color.
557    color: usize,
558    /// Rank within the subgroup.
559    rank: usize,
560    /// Number of subgroup ranks.
561    size: usize,
562    /// Ordered global ranks.
563    global_ranks: Vec<usize>,
564}
565
566impl SubgroupMembership {
567    /// Axis represented by this subgroup.
568    pub const fn axis(&self) -> ParallelAxis {
569        self.axis
570    }
571    /// Deterministic subgroup color.
572    pub const fn color(&self) -> usize {
573        self.color
574    }
575    /// Rank within the subgroup.
576    pub const fn rank(&self) -> usize {
577        self.rank
578    }
579    /// Number of subgroup ranks.
580    pub const fn size(&self) -> usize {
581        self.size
582    }
583    /// Ordered global ranks.
584    pub fn global_ranks(&self) -> &[usize] {
585        &self.global_ranks
586    }
587}
588
589/// Weight-independent ownership report for one parallel rank.
590#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
591#[non_exhaustive]
592pub struct TopologyPreflightReport {
593    /// Complete rank topology.
594    topology: ParallelRankTopology,
595    /// Tensor subgroup.
596    tensor_subgroup: SubgroupMembership,
597    /// Pipeline subgroup.
598    pipeline_subgroup: SubgroupMembership,
599    /// Expert subgroup.
600    expert_subgroup: SubgroupMembership,
601    /// Data subgroup.
602    data_subgroup: SubgroupMembership,
603    /// Locally owned layer range.
604    local_layer_range: Option<Range<usize>>,
605    /// Locally owned expert range.
606    local_expert_range: Option<Range<usize>>,
607    /// Whether the embedding is local.
608    owns_embedding: bool,
609    /// Whether the output head is local.
610    owns_output_head: bool,
611}
612
613impl TopologyPreflightReport {
614    /// Complete rank topology.
615    pub const fn topology(&self) -> ParallelRankTopology {
616        self.topology
617    }
618    /// Tensor subgroup.
619    pub const fn tensor_subgroup(&self) -> &SubgroupMembership {
620        &self.tensor_subgroup
621    }
622    /// Pipeline subgroup.
623    pub const fn pipeline_subgroup(&self) -> &SubgroupMembership {
624        &self.pipeline_subgroup
625    }
626    /// Expert subgroup.
627    pub const fn expert_subgroup(&self) -> &SubgroupMembership {
628        &self.expert_subgroup
629    }
630    /// Data subgroup.
631    pub const fn data_subgroup(&self) -> &SubgroupMembership {
632        &self.data_subgroup
633    }
634    /// Locally owned layer range.
635    pub const fn local_layer_range(&self) -> Option<&Range<usize>> {
636        self.local_layer_range.as_ref()
637    }
638    /// Locally owned expert range.
639    pub const fn local_expert_range(&self) -> Option<&Range<usize>> {
640        self.local_expert_range.as_ref()
641    }
642    /// Whether the embedding is local.
643    pub const fn owns_embedding(&self) -> bool {
644        self.owns_embedding
645    }
646    /// Whether the output head is local.
647    pub const fn owns_output_head(&self) -> bool {
648        self.owns_output_head
649    }
650}
651
652fn subgroup_color(
653    topology: ParallelTopology,
654    coordinates: ParallelCoordinates,
655    axis: ParallelAxis,
656) -> Result<usize, TopologyError> {
657    let dimensions = [
658        (ParallelAxis::Data, coordinates.data, topology.data),
659        (
660            ParallelAxis::Pipeline,
661            coordinates.pipeline,
662            topology.pipeline,
663        ),
664        (ParallelAxis::Tensor, coordinates.tensor, topology.tensor),
665        (ParallelAxis::Expert, coordinates.expert, topology.expert),
666    ];
667    dimensions
668        .into_iter()
669        .filter(|(candidate, _, _)| *candidate != axis)
670        .try_fold(0usize, |color, (_, coordinate, size)| {
671            color
672                .checked_mul(size)
673                .and_then(|value| value.checked_add(coordinate))
674                .ok_or(TopologyError::WorldSizeOverflow)
675        })
676}
677
678/// Computes a deterministic balanced contiguous range.
679pub fn balanced_contiguous_range(
680    total: usize,
681    parts: usize,
682    index: usize,
683    allow_empty: bool,
684) -> Result<Range<usize>, TopologyError> {
685    if parts == 0 {
686        return Err(TopologyError::ZeroPartitions);
687    }
688    if index >= parts {
689        return Err(TopologyError::PartitionOutOfRange { index, parts });
690    }
691    if !allow_empty && total < parts {
692        return Err(TopologyError::EmptyPartition { total, parts });
693    }
694    let base = total / parts;
695    let extra = total % parts;
696    let start = index
697        .checked_mul(base)
698        .and_then(|value| value.checked_add(index.min(extra)))
699        .ok_or(TopologyError::PartitionOverflow)?;
700    let end = start
701        .checked_add(base + usize::from(index < extra))
702        .ok_or(TopologyError::PartitionOverflow)?;
703    Ok(start..end)
704}
705
706/// Topology validation error.
707#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
708#[non_exhaustive]
709pub enum TopologyError {
710    /// An axis size was zero.
711    #[error("parallel topology axis sizes must be positive")]
712    ZeroAxis,
713    /// Product of axis sizes overflowed.
714    #[error("parallel topology world size overflows usize")]
715    WorldSizeOverflow,
716    /// Rank is outside the world.
717    #[error("rank {rank} is outside world size {world_size}")]
718    RankOutOfRange {
719        /// Invalid rank.
720        rank: usize,
721        /// World size.
722        world_size: usize,
723    },
724    /// A coordinate is outside its axis.
725    #[error("{axis:?} coordinate {coordinate} is outside axis size {size}")]
726    CoordinateOutOfRange {
727        /// Invalid axis.
728        axis: ParallelAxis,
729        /// Invalid coordinate.
730        coordinate: usize,
731        /// Axis size.
732        size: usize,
733    },
734    /// Subgroup membership did not map back to the represented rank.
735    #[error("{axis:?} subgroup geometry does not map back to the represented rank")]
736    SubgroupIdentity {
737        /// Invalid subgroup axis.
738        axis: ParallelAxis,
739    },
740    /// Pipeline preflight omitted layer geometry.
741    #[error("pipeline topology preflight requires the decoder-layer count")]
742    MissingLayerCount,
743    /// Expert preflight omitted expert geometry.
744    #[error("expert topology preflight requires the routed-expert count")]
745    MissingExpertCount,
746    /// Partition count is zero.
747    #[error("partition count must be nonzero")]
748    ZeroPartitions,
749    /// Partition index is outside the partition count.
750    #[error("partition index {index} is outside {parts} parts")]
751    PartitionOutOfRange {
752        /// Invalid index.
753        index: usize,
754        /// Partition count.
755        parts: usize,
756    },
757    /// Non-empty partitions were requested with too few items.
758    #[error("cannot divide {total} items among {parts} non-empty partitions")]
759    EmptyPartition {
760        /// Item count.
761        total: usize,
762        /// Partition count.
763        parts: usize,
764    },
765    /// Partition offset calculation overflowed.
766    #[error("balanced range calculation overflowed usize")]
767    PartitionOverflow,
768}
769
770#[cfg(test)]
771mod tests {
772    use super::*;
773
774    #[test]
775    fn rank_topology_is_the_authoritative_coordinate_and_membership_plan() {
776        let topology = ParallelTopology::new(2, 3, 2, 2).unwrap();
777        let rank = ParallelRankTopology::new(topology, 22).unwrap();
778        assert_eq!(
779            rank.coordinates(),
780            ParallelCoordinates {
781                tensor: 1,
782                pipeline: 2,
783                expert: 0,
784                data: 1,
785            }
786        );
787        assert_eq!(rank.global_rank_for(rank.coordinates()).unwrap(), 22);
788        assert_eq!(
789            rank.subgroup(ParallelAxis::Tensor).unwrap().global_ranks,
790            [20, 22]
791        );
792        assert_eq!(
793            rank.subgroup(ParallelAxis::Pipeline).unwrap().global_ranks,
794            [14, 18, 22]
795        );
796        assert_eq!(
797            rank.subgroup(ParallelAxis::Expert).unwrap().global_ranks,
798            [22, 23]
799        );
800        assert_eq!(
801            rank.subgroup(ParallelAxis::Data).unwrap().global_ranks,
802            [10, 22]
803        );
804        assert_eq!(rank.pipeline_predecessor().unwrap(), Some(18));
805        assert_eq!(rank.pipeline_successor().unwrap(), None);
806        assert!(!rank.owns_embedding());
807        assert!(rank.owns_output_head());
808
809        let report = rank.preflight(Some(7), Some(5)).unwrap();
810        assert_eq!(report.local_layer_range, Some(5..7));
811        assert_eq!(report.local_expert_range, Some(0..3));
812        assert_eq!(report.data_subgroup.global_ranks, [10, 22]);
813    }
814
815    #[test]
816    fn rank_mapping_is_exhaustive_for_all_axes() {
817        let topology = ParallelTopology::new(3, 2, 2, 2).unwrap();
818        for global_rank in 0..topology.world_size() {
819            let rank = ParallelRankTopology::new(topology, global_rank).unwrap();
820            assert_eq!(
821                rank.global_rank_for(rank.coordinates()).unwrap(),
822                global_rank
823            );
824            for axis in [
825                ParallelAxis::Tensor,
826                ParallelAxis::Pipeline,
827                ParallelAxis::Expert,
828                ParallelAxis::Data,
829            ] {
830                let subgroup = rank.subgroup(axis).unwrap();
831                assert_eq!(subgroup.global_ranks[subgroup.rank], global_rank);
832                assert_eq!(
833                    subgroup.global_ranks,
834                    topology.axis_members(global_rank, axis).unwrap()
835                );
836            }
837        }
838    }
839
840    #[test]
841    fn balanced_ranges_and_preflight_fail_closed() {
842        let ranges = (0..3)
843            .map(|index| balanced_contiguous_range(8, 3, index, false).unwrap())
844            .collect::<Vec<_>>();
845        assert_eq!(ranges, [0..3, 3..6, 6..8]);
846        assert!(balanced_contiguous_range(2, 3, 0, false).is_err());
847        assert_eq!(balanced_contiguous_range(2, 3, 2, true).unwrap(), 2..2);
848
849        let pipeline =
850            ParallelRankTopology::new(ParallelTopology::new(1, 4, 1, 1).unwrap(), 0).unwrap();
851        assert_eq!(
852            pipeline.preflight(None, None),
853            Err(TopologyError::MissingLayerCount)
854        );
855        let expert =
856            ParallelRankTopology::new(ParallelTopology::new(1, 1, 4, 1).unwrap(), 0).unwrap();
857        assert_eq!(
858            expert.preflight(None, None),
859            Err(TopologyError::MissingExpertCount)
860        );
861    }
862
863    #[test]
864    fn deserialization_cannot_bypass_shape_or_derived_rank_validation() {
865        assert!(serde_json::from_str::<ParallelTopology>(
866            r#"{"tensor":0,"pipeline":1,"expert":1,"data":1}"#
867        )
868        .is_err());
869
870        let rank =
871            ParallelRankTopology::new(ParallelTopology::new(2, 2, 2, 2).unwrap(), 9).unwrap();
872        let encoded = serde_json::to_string(&rank).unwrap();
873        assert_eq!(
874            serde_json::from_str::<ParallelRankTopology>(&encoded).unwrap(),
875            rank
876        );
877
878        let mut inconsistent = serde_json::to_value(rank).unwrap();
879        inconsistent["tensor_parallel_rank"] = serde_json::json!(1);
880        assert!(serde_json::from_value::<ParallelRankTopology>(inconsistent).is_err());
881    }
882}
eredu_core/topology.rs

eredu_core/
topology.rs
