1use std::collections::{BTreeMap, BTreeSet};
4
5use eredu_checkpoint::{
6 recipe::RecipeError,
7 store::{CheckpointSource, StoreError, TensorSelection},
8};
9
10use crate::{
11 LocalModelLayout, LocalTensorLayout, ResidencyDeclarationError, TensorPlacement, WeightBinding,
12 WeightBindingSelectionError,
13};
14
15#[derive(Debug, Clone, Copy, Eq, PartialEq)]
17pub struct PlacementRank {
18 world_size: usize,
19 global_rank: usize,
20}
21
22impl PlacementRank {
23 pub fn new(world_size: usize, global_rank: usize) -> Result<Self, PlacementPlanError> {
25 if world_size == 0 || global_rank >= world_size {
26 return Err(PlacementPlanError::InvalidRank {
27 world_size,
28 global_rank,
29 });
30 }
31 Ok(Self {
32 world_size,
33 global_rank,
34 })
35 }
36
37 pub const fn world_size(self) -> usize {
39 self.world_size
40 }
41
42 pub const fn global_rank(self) -> usize {
44 self.global_rank
45 }
46}
47
48#[derive(Debug, Clone, Eq, PartialEq)]
50pub struct TensorSlice {
51 axis: usize,
52 start: usize,
53 end: usize,
54 index: usize,
55 parts: usize,
56}
57
58impl TensorSlice {
59 pub fn for_shape(
61 shape: &[usize],
62 axis: usize,
63 index: usize,
64 parts: usize,
65 ) -> Result<Self, PlacementPlanError> {
66 let dimension = *shape.get(axis).ok_or(PlacementPlanError::AxisOutOfBounds {
67 axis,
68 rank: shape.len(),
69 })?;
70 if parts == 0 || index >= parts || dimension == 0 || !dimension.is_multiple_of(parts) {
71 return Err(PlacementPlanError::InvalidShard {
72 axis,
73 index,
74 parts,
75 dimension,
76 });
77 }
78 let width = dimension / parts;
79 let start = index
80 .checked_mul(width)
81 .ok_or(PlacementPlanError::ArithmeticOverflow)?;
82 Ok(Self {
83 axis,
84 start,
85 end: start + width,
86 index,
87 parts,
88 })
89 }
90
91 pub const fn axis(&self) -> usize {
93 self.axis
94 }
95 pub const fn start(&self) -> usize {
97 self.start
98 }
99 pub const fn end(&self) -> usize {
101 self.end
102 }
103 pub const fn index(&self) -> usize {
105 self.index
106 }
107 pub const fn parts(&self) -> usize {
109 self.parts
110 }
111 pub fn local_shape(&self, source_shape: &[usize]) -> Vec<usize> {
113 let mut shape = source_shape.to_vec();
114 shape[self.axis] = self.end - self.start;
115 shape
116 }
117}
118
119#[derive(Debug, Clone)]
120struct TensorPlan {
121 placement: TensorPlacement,
122 expected_source_shape: Option<Vec<usize>>,
123}
124
125#[derive(Debug, Clone)]
127pub struct PlacementPlan {
128 rank: PlacementRank,
129 tensors: BTreeMap<String, TensorPlan>,
130 default: Option<TensorPlacement>,
131}
132
133impl PlacementPlan {
134 pub const fn new(rank: PlacementRank) -> Self {
136 Self {
137 rank,
138 tensors: BTreeMap::new(),
139 default: None,
140 }
141 }
142
143 pub fn replicated(rank: PlacementRank) -> Self {
145 Self::new(rank).with_default(TensorPlacement::Replicated)
146 }
147
148 pub const fn rank(&self) -> PlacementRank {
150 self.rank
151 }
152
153 pub fn with_default(mut self, placement: TensorPlacement) -> Self {
155 self.default = Some(placement);
156 self
157 }
158
159 pub fn insert(&mut self, source: impl Into<String>, placement: TensorPlacement) {
161 self.tensors.insert(
162 source.into(),
163 TensorPlan {
164 placement,
165 expected_source_shape: None,
166 },
167 );
168 }
169
170 pub fn insert_expected(
172 &mut self,
173 source: impl Into<String>,
174 expected_source_shape: impl Into<Vec<usize>>,
175 placement: TensorPlacement,
176 ) -> Result<(), PlacementPlanError> {
177 let shape = expected_source_shape.into();
178 validate_tensor_placement(&placement, &shape, self.rank)?;
179 self.tensors.insert(
180 source.into(),
181 TensorPlan {
182 placement,
183 expected_source_shape: Some(shape),
184 },
185 );
186 Ok(())
187 }
188
189 pub fn insert_quantized_companions(
191 &mut self,
192 prefix: &str,
193 placement: TensorPlacement,
194 has_biases: bool,
195 ) {
196 self.insert(format!("{prefix}.weight"), placement.clone());
197 self.insert(format!("{prefix}.scales"), placement.clone());
198 if has_biases {
199 self.insert(format!("{prefix}.biases"), placement);
200 }
201 }
202
203 pub fn placement(&self, source: &str) -> Option<&TensorPlacement> {
205 self.tensors.get(source).map(|plan| &plan.placement)
206 }
207
208 pub fn validate(&self) -> Result<(), PlacementPlanError> {
210 for tensor in self.tensors.values() {
211 validate_tensor_plan(tensor, self.rank)?;
212 }
213 if let Some(default) = &self.default {
214 validate_tensor_plan(
215 &TensorPlan {
216 placement: default.clone(),
217 expected_source_shape: None,
218 },
219 self.rank,
220 )?;
221 }
222 Ok(())
223 }
224
225 pub fn potentially_local(&self, source: &str) -> Result<bool, PlacementPlanError> {
227 let plan = self.source_plan(source)?;
228 Ok(!matches!(plan.placement, TensorPlacement::Omit)
229 && !matches!(plan.placement, TensorPlacement::Rank { rank } if rank != self.rank.global_rank))
230 }
231
232 pub fn resolve(
234 &self,
235 source: &str,
236 shape: &[usize],
237 ) -> Result<ResolvedTensorPlacement, PlacementPlanError> {
238 let plan = self.source_plan(source)?;
239 if plan
240 .expected_source_shape
241 .as_ref()
242 .is_some_and(|expected| expected != shape)
243 {
244 return Err(PlacementPlanError::SourceShapeMismatch {
245 checkpoint_source: source.to_owned(),
246 expected: plan.expected_source_shape.clone().unwrap(),
247 actual: shape.to_vec(),
248 });
249 }
250 validate_tensor_placement(&plan.placement, shape, self.rank)?;
251 Ok(match &plan.placement {
252 TensorPlacement::Replicated | TensorPlacement::Local => {
253 ResolvedTensorPlacement::Materialize
254 }
255 TensorPlacement::Omit => ResolvedTensorPlacement::Omit,
256 TensorPlacement::Rank { rank } if *rank == self.rank.global_rank => {
257 ResolvedTensorPlacement::Materialize
258 }
259 TensorPlacement::Rank { .. } => ResolvedTensorPlacement::Omit,
260 TensorPlacement::Shard { axis, index, parts } => {
261 let slice = TensorSlice::for_shape(shape, *axis, *index, *parts)?;
262 ResolvedTensorPlacement::Selection(TensorSelection::Range {
263 axis: slice.axis,
264 start: slice.start,
265 end: slice.end,
266 })
267 }
268 TensorPlacement::Range { axis, start, end } => {
269 ResolvedTensorPlacement::Selection(TensorSelection::Range {
270 axis: *axis,
271 start: *start,
272 end: *end,
273 })
274 }
275 TensorPlacement::Indices { axis, indices } => {
276 ResolvedTensorPlacement::Selection(TensorSelection::Indices {
277 axis: *axis,
278 indices: indices.clone(),
279 })
280 }
281 })
282 }
283
284 pub fn validate_loaded_sources(
286 &self,
287 loaded: &BTreeSet<String>,
288 mut unexpected: Vec<String>,
289 ) -> Result<(), PlacementPlanError> {
290 let mut missing = self
291 .tensors
292 .iter()
293 .filter_map(|(source, plan)| {
294 let required = !matches!(plan.placement, TensorPlacement::Omit)
295 && !matches!(plan.placement, TensorPlacement::Rank { rank } if rank != self.rank.global_rank);
296 (required && !loaded.contains(source)).then_some(source.clone())
297 })
298 .collect::<Vec<_>>();
299 missing.sort();
300 unexpected.sort();
301 unexpected.dedup();
302 if missing.is_empty() && unexpected.is_empty() {
303 Ok(())
304 } else {
305 Err(PlacementPlanError::Coverage {
306 missing,
307 unexpected,
308 })
309 }
310 }
311
312 fn source_plan(&self, source: &str) -> Result<TensorPlan, PlacementPlanError> {
313 self.tensors
314 .get(source)
315 .cloned()
316 .or_else(|| {
317 self.default.as_ref().map(|placement| TensorPlan {
318 placement: placement.clone(),
319 expected_source_shape: None,
320 })
321 })
322 .ok_or_else(|| PlacementPlanError::UnexpectedSource {
323 checkpoint_source: source.to_owned(),
324 })
325 }
326}
327
328#[derive(Debug, Clone, Eq, PartialEq)]
330pub enum ResolvedTensorPlacement {
331 Materialize,
333 Omit,
335 Selection(TensorSelection),
337}
338
339fn validate_tensor_plan(plan: &TensorPlan, rank: PlacementRank) -> Result<(), PlacementPlanError> {
340 match &plan.placement {
341 TensorPlacement::Rank { rank: owner } if *owner >= rank.world_size => {
342 Err(PlacementPlanError::OwnerOutOfBounds {
343 owner: *owner,
344 world_size: rank.world_size,
345 })
346 }
347 TensorPlacement::Shard { axis, index, parts } if *parts == 0 || *index >= *parts => {
348 Err(PlacementPlanError::InvalidShard {
349 axis: *axis,
350 index: *index,
351 parts: *parts,
352 dimension: 0,
353 })
354 }
355 TensorPlacement::Range { axis, start, end } if start >= end => {
356 Err(PlacementPlanError::InvalidRange {
357 axis: *axis,
358 start: *start,
359 end: *end,
360 dimension: None,
361 })
362 }
363 TensorPlacement::Indices { axis, indices }
364 if indices.is_empty()
365 || indices.iter().collect::<BTreeSet<_>>().len() != indices.len() =>
366 {
367 Err(PlacementPlanError::InvalidIndices { axis: *axis })
368 }
369 placement => {
370 if let Some(shape) = &plan.expected_source_shape {
371 validate_tensor_placement(placement, shape, rank)?;
372 }
373 Ok(())
374 }
375 }
376}
377
378fn validate_tensor_placement(
379 placement: &TensorPlacement,
380 shape: &[usize],
381 rank: PlacementRank,
382) -> Result<(), PlacementPlanError> {
383 match placement {
384 TensorPlacement::Rank { rank: owner } if *owner >= rank.world_size => {
385 Err(PlacementPlanError::OwnerOutOfBounds {
386 owner: *owner,
387 world_size: rank.world_size,
388 })
389 }
390 TensorPlacement::Shard { axis, index, parts } => {
391 TensorSlice::for_shape(shape, *axis, *index, *parts).map(|_| ())
392 }
393 TensorPlacement::Range { axis, start, end } => {
394 let dimension = *shape
395 .get(*axis)
396 .ok_or(PlacementPlanError::AxisOutOfBounds {
397 axis: *axis,
398 rank: shape.len(),
399 })?;
400 if start >= end || *end > dimension {
401 Err(PlacementPlanError::InvalidRange {
402 axis: *axis,
403 start: *start,
404 end: *end,
405 dimension: Some(dimension),
406 })
407 } else {
408 Ok(())
409 }
410 }
411 TensorPlacement::Indices { axis, indices } => {
412 let dimension = *shape
413 .get(*axis)
414 .ok_or(PlacementPlanError::AxisOutOfBounds {
415 axis: *axis,
416 rank: shape.len(),
417 })?;
418 if indices.is_empty()
419 || indices.iter().collect::<BTreeSet<_>>().len() != indices.len()
420 || indices.iter().any(|index| *index >= dimension)
421 {
422 Err(PlacementPlanError::InvalidIndices { axis: *axis })
423 } else {
424 Ok(())
425 }
426 }
427 _ => Ok(()),
428 }
429}
430
431#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
433pub enum PlacementPlanError {
434 #[error("global rank {global_rank} is outside world size {world_size}")]
436 InvalidRank {
437 world_size: usize,
439 global_rank: usize,
441 },
442 #[error("owner rank {owner} is outside world size {world_size}")]
444 OwnerOutOfBounds {
445 owner: usize,
447 world_size: usize,
449 },
450 #[error("tensor axis {axis} is outside rank {rank}")]
452 AxisOutOfBounds {
453 axis: usize,
455 rank: usize,
457 },
458 #[error("shard {index}/{parts} on axis {axis} is invalid for dimension {dimension}")]
460 InvalidShard {
461 axis: usize,
463 index: usize,
465 parts: usize,
467 dimension: usize,
469 },
470 #[error("range {start}..{end} on axis {axis} is invalid for dimension {dimension:?}")]
472 InvalidRange {
473 axis: usize,
475 start: usize,
477 end: usize,
479 dimension: Option<usize>,
481 },
482 #[error("index selection on axis {axis} is invalid")]
484 InvalidIndices {
485 axis: usize,
487 },
488 #[error("tensor slice arithmetic overflowed")]
490 ArithmeticOverflow,
491 #[error("checkpoint source {checkpoint_source:?} is unexpected")]
493 UnexpectedSource {
494 checkpoint_source: String,
496 },
497 #[error("checkpoint source {checkpoint_source:?} has shape {actual:?}, expected {expected:?}")]
499 SourceShapeMismatch {
500 checkpoint_source: String,
502 expected: Vec<usize>,
504 actual: Vec<usize>,
506 },
507 #[error("placement coverage mismatch: missing {missing:?}, unexpected {unexpected:?}")]
509 Coverage {
510 missing: Vec<String>,
512 unexpected: Vec<String>,
514 },
515}
516
517pub fn place_weight_bindings(
519 bindings: Vec<WeightBinding>,
520 source: &dyn CheckpointSource,
521 layout: &LocalModelLayout,
522) -> Result<Vec<WeightBinding>, BindingPlacementError> {
523 apply_binding_placements(bindings, source, layout, false)
524}
525
526pub fn place_addressable_member_bindings(
531 bindings: Vec<WeightBinding>,
532 source: &dyn CheckpointSource,
533 layout: &LocalModelLayout,
534) -> Result<Vec<WeightBinding>, BindingPlacementError> {
535 apply_binding_placements(bindings, source, layout, true)
536}
537
538fn apply_binding_placements(
539 bindings: Vec<WeightBinding>,
540 source: &dyn CheckpointSource,
541 layout: &LocalModelLayout,
542 skip_member_axis: bool,
543) -> Result<Vec<WeightBinding>, BindingPlacementError> {
544 let source_keys = source.source_keys().into_iter().collect::<BTreeSet<_>>();
545 let mut output = Vec::with_capacity(bindings.len());
546 for binding in bindings {
547 if binding.is_alias() {
548 output.push(binding);
549 continue;
550 }
551 let logical_target = binding
552 .logical_target()
553 .ok_or_else(|| BindingPlacementError::MissingLogicalTarget {
554 binding: binding.name().to_owned(),
555 })?
556 .to_owned();
557 let tensor = layout.tensor(&logical_target).ok_or_else(|| {
558 BindingPlacementError::UnknownLogicalTarget {
559 binding: binding.name().to_owned(),
560 target: logical_target.clone(),
561 }
562 })?;
563 let companions = binding.quantization_companions().cloned();
564
565 if !skip_member_axis
566 && binding.recipe().is_none()
567 && !source_keys.contains(binding.checkpoint_key())
568 {
569 if !tensor.additional_placements().is_empty() {
570 return Err(BindingPlacementError::CompoundPlacementRequiresRecipe {
571 binding: binding.name().to_owned(),
572 });
573 }
574 let selection = placement_selection(tensor, tensor.placement(), tensor.global_shape())?;
575 if selection == TensorSelection::Full {
576 output.push(binding);
577 continue;
578 }
579 let global = tensor
580 .global_shape()
581 .iter()
582 .try_fold(1usize, |value, item| value.checked_mul(*item));
583 let local = tensor
584 .local_shape()
585 .iter()
586 .try_fold(1usize, |value, item| value.checked_mul(*item));
587 let expected_bytes = global
588 .zip(local)
589 .and_then(|(global, local)| {
590 binding
591 .expected_bytes()
592 .checked_mul(local as u64)
593 .and_then(|bytes| bytes.checked_div(global as u64))
594 })
595 .ok_or_else(|| BindingPlacementError::ByteGeometry {
596 binding: binding.name().to_owned(),
597 })?;
598 let mut placed = WeightBinding::new(
599 binding.name(),
600 binding.checkpoint_key(),
601 selection,
602 expected_bytes,
603 )?
604 .with_logical_target(logical_target)?;
605 if let Some(companions) = companions {
606 placed = placed.with_quantization_companions(
607 companions.scale(),
608 companions.affine_bias().map(str::to_owned),
609 )?;
610 }
611 output.push(placed);
612 continue;
613 }
614
615 let mut recipe = binding.source_recipe();
616 let mut selected = false;
617 for placement in tensor
618 .additional_placements()
619 .iter()
620 .chain(std::iter::once(tensor.placement()))
621 {
622 let member_axis = matches!(
623 placement,
624 TensorPlacement::Shard { axis: 0, .. }
625 | TensorPlacement::Range { axis: 0, .. }
626 | TensorPlacement::Indices { axis: 0, .. }
627 );
628 if skip_member_axis && member_axis {
629 continue;
630 }
631 let metadata = recipe.infer(source)?;
632 let selection = placement_selection(tensor, placement, metadata.shape())?;
633 if selection != TensorSelection::Full {
634 recipe = recipe.select_bounded(source, selection)?;
635 selected = true;
636 }
637 }
638 if !selected {
639 output.push(binding);
640 continue;
641 }
642 let expected_bytes = recipe.infer(source)?.byte_len();
643 let mut placed = WeightBinding::from_recipe(binding.name(), recipe, expected_bytes)?
644 .with_logical_target(logical_target)?;
645 if let Some(companions) = companions {
646 placed = placed.with_quantization_companions(
647 companions.scale(),
648 companions.affine_bias().map(str::to_owned),
649 )?;
650 }
651 output.push(placed);
652 }
653 Ok(output)
654}
655
656pub fn placement_selection(
658 tensor: &LocalTensorLayout,
659 placement: &TensorPlacement,
660 stored_shape: &[usize],
661) -> Result<TensorSelection, BindingPlacementError> {
662 let scale_boundary = |axis: usize, boundary: usize| -> Result<usize, BindingPlacementError> {
663 let semantic =
664 *tensor
665 .global_shape()
666 .get(axis)
667 .ok_or(BindingPlacementError::AxisOutOfBounds {
668 axis,
669 rank: tensor.global_shape().len(),
670 })?;
671 let stored = *stored_shape
672 .get(axis)
673 .ok_or(BindingPlacementError::AxisOutOfBounds {
674 axis,
675 rank: stored_shape.len(),
676 })?;
677 boundary
678 .checked_mul(stored)
679 .and_then(|value| value.checked_div(semantic))
680 .filter(|scaled| *scaled * semantic == boundary * stored)
681 .ok_or(BindingPlacementError::UnalignedBoundary {
682 axis,
683 boundary,
684 semantic,
685 stored,
686 })
687 };
688
689 Ok(match placement {
690 TensorPlacement::Replicated | TensorPlacement::Local => TensorSelection::Full,
691 TensorPlacement::Shard { axis, index, parts } => {
692 let stored =
693 *stored_shape
694 .get(*axis)
695 .ok_or(BindingPlacementError::AxisOutOfBounds {
696 axis: *axis,
697 rank: stored_shape.len(),
698 })?;
699 if *parts == 0 || *index >= *parts || !stored.is_multiple_of(*parts) {
700 return Err(BindingPlacementError::InvalidShard {
701 axis: *axis,
702 index: *index,
703 parts: *parts,
704 stored,
705 });
706 }
707 let width = stored / *parts;
708 TensorSelection::Range {
709 axis: *axis,
710 start: index * width,
711 end: (index + 1) * width,
712 }
713 }
714 TensorPlacement::Range { axis, start, end } => TensorSelection::Range {
715 axis: *axis,
716 start: scale_boundary(*axis, *start)?,
717 end: scale_boundary(*axis, *end)?,
718 },
719 TensorPlacement::Indices { axis, indices } => {
720 let stored =
721 *stored_shape
722 .get(*axis)
723 .ok_or(BindingPlacementError::AxisOutOfBounds {
724 axis: *axis,
725 rank: stored_shape.len(),
726 })?;
727 let semantic = *tensor.global_shape().get(*axis).ok_or(
728 BindingPlacementError::AxisOutOfBounds {
729 axis: *axis,
730 rank: tensor.global_shape().len(),
731 },
732 )?;
733 if stored != semantic {
734 return Err(BindingPlacementError::IndexedPackedStorage {
735 axis: *axis,
736 semantic,
737 stored,
738 });
739 }
740 TensorSelection::Indices {
741 axis: *axis,
742 indices: indices.clone(),
743 }
744 }
745 TensorPlacement::Omit | TensorPlacement::Rank { .. } => {
746 return Err(BindingPlacementError::NonLocalPlacement {
747 placement: placement.clone(),
748 });
749 }
750 })
751}
752
753#[derive(Debug, thiserror::Error)]
755pub enum BindingPlacementError {
756 #[error("binding {binding:?} has no logical placement target")]
758 MissingLogicalTarget {
759 binding: String,
761 },
762 #[error("binding {binding:?} targets unknown layout entry {target:?}")]
764 UnknownLogicalTarget {
765 binding: String,
767 target: String,
769 },
770 #[error("compound placement for {binding:?} requires an admitted checkpoint recipe")]
772 CompoundPlacementRequiresRecipe {
773 binding: String,
775 },
776 #[error("cannot derive rank-local byte geometry for {binding:?}")]
778 ByteGeometry {
779 binding: String,
781 },
782 #[error("placement axis {axis} is outside tensor rank {rank}")]
784 AxisOutOfBounds {
785 axis: usize,
787 rank: usize,
789 },
790 #[error("semantic boundary {boundary} on axis {axis} ({semantic}) is not aligned to stored width {stored}")]
792 UnalignedBoundary {
793 axis: usize,
795 boundary: usize,
797 semantic: usize,
799 stored: usize,
801 },
802 #[error("shard {index}/{parts} on axis {axis} is invalid for stored width {stored}")]
804 InvalidShard {
805 axis: usize,
807 index: usize,
809 parts: usize,
811 stored: usize,
813 },
814 #[error("indexed axis {axis} has semantic width {semantic} but stored width {stored}")]
816 IndexedPackedStorage {
817 axis: usize,
819 semantic: usize,
821 stored: usize,
823 },
824 #[error("execution-unit binding has non-local placement {placement:?}")]
826 NonLocalPlacement {
827 placement: TensorPlacement,
829 },
830 #[error(transparent)]
832 Store(#[from] StoreError),
833 #[error(transparent)]
835 Recipe(#[from] RecipeError),
836 #[error(transparent)]
838 Declaration(#[from] ResidencyDeclarationError),
839 #[error(transparent)]
841 Selection(#[from] WeightBindingSelectionError),
842}