use std::{
collections::BTreeSet,
convert::Infallible,
sync::atomic::{AtomicUsize, Ordering},
};
use eredu_checkpoint::{
recipe::RecipeDtype,
store::{
CheckpointLease, CheckpointSource, MemoryWeightStore, StoreError, TensorMetadata,
TensorReadRequest, TensorSelection, TensorSourceProvenance, WeightStoreDiagnostics,
},
LinearFormat, SourceTensorEncoding, StoredDtype,
};
use eredu_nn::{
ParameterMetadata, ParameterSpec, ParameterVisitor, ParameterVisitorMut, Parameterized,
};
use eredu_runtime::{
build_exact_replicated_text_bindings, BindingPlanError, LocalModelLayout, LocalTensorLayout,
ModuleBindingPlanError, ParameterBindingTarget, ParameterRole,
ReplicatedTextMaterializationTask, ReplicatedTextParameterOwner, ReplicatedTextParameterRole,
ReplicatedTextPhysicalSource, TensorPlacement, WeightLoweringDescriptor, WeightLoweringKind,
};
struct Module(ParameterBindingTarget);
impl Parameterized<ParameterBindingTarget> for Module {
fn visit_parameters<'a, V: ParameterVisitor<'a, ParameterBindingTarget>>(
&'a self,
visitor: &mut V,
) {
visitor.visit(
ParameterMetadata::from_spec(&ParameterSpec::trainable("weight").unwrap(), true),
&self.0,
);
}
fn visit_parameters_mut<'a, V: ParameterVisitorMut<'a, ParameterBindingTarget>>(
&'a mut self,
visitor: &mut V,
) {
visitor.visit_mut(
ParameterMetadata::from_spec(&ParameterSpec::trainable("weight").unwrap(), true),
&mut self.0,
);
}
fn set_trainable(&mut self, _: bool) {}
}
struct Source {
store: MemoryWeightStore,
overlay: bool,
leases: AtomicUsize,
}
impl CheckpointSource for Source {
fn source_keys(&self) -> Vec<String> {
self.store.source_keys()
}
fn source_metadata(&self, key: &str) -> Result<TensorMetadata, StoreError> {
self.store.source_metadata(key)
}
fn source_provenance(&self, key: &str) -> Result<TensorSourceProvenance, StoreError> {
let mut provenance = self.store.source_provenance(key)?;
provenance.backing_shard = Some("/exact/model.safetensors".into());
Ok(provenance)
}
fn acquire_lease(&self, request: TensorReadRequest) -> Result<CheckpointLease, StoreError> {
self.leases.fetch_add(1, Ordering::SeqCst);
self.store.acquire_lease(request)
}
fn source_diagnostics(&self) -> Result<WeightStoreDiagnostics, StoreError> {
self.store.source_diagnostics()
}
fn is_authoritative_materialized_key(&self, key: &str) -> bool {
self.overlay && key == "weight"
}
}
fn task(
shape: &[usize],
lowering: WeightLoweringKind,
executable: LinearFormat,
) -> ReplicatedTextMaterializationTask {
let encoding = SourceTensorEncoding::Safetensors(StoredDtype::F32);
ReplicatedTextMaterializationTask::from_exact_source(
"weight",
ReplicatedTextPhysicalSource::new(
"checkpoint.weight",
"checkpoint.weight",
"/exact/model.safetensors",
"checkpoint.weight",
encoding.clone(),
(shape.iter().product::<usize>() * 4) as u64,
)
.unwrap(),
Vec::new(),
shape.to_vec(),
shape.to_vec(),
ReplicatedTextParameterRole::LinearWeight,
ReplicatedTextParameterOwner::StaticRole("head".into()),
executable,
lowering,
WeightLoweringDescriptor::new(
encoding,
executable,
shape.to_vec(),
shape.to_vec(),
Some(shape.len() - 1),
)
.unwrap(),
)
.unwrap()
}
fn source(shape: Vec<usize>, overlay: bool) -> Source {
Source {
store: MemoryWeightStore::from_safetensors([(
if overlay {
"weight"
} else {
"checkpoint.weight"
}
.into(),
if overlay {
safetensors::Dtype::U32
} else {
safetensors::Dtype::F32
},
shape.clone(),
vec![0; shape.iter().product::<usize>() * 4],
)])
.unwrap(),
overlay,
leases: AtomicUsize::new(0),
}
}
fn layout(global: Vec<usize>, local: Vec<usize>, placement: TensorPlacement) -> LocalModelLayout {
let mut layout = LocalModelLayout::default();
layout.insert(
"weight".into(),
LocalTensorLayout::new(
"head",
ParameterRole::ColumnProjection,
global,
local,
placement,
None,
None,
false,
),
);
layout
}
fn bind(
source: &Source,
task: &ReplicatedTextMaterializationTask,
shape: Vec<usize>,
layout: Option<&LocalModelLayout>,
) -> Result<Vec<eredu_runtime::WeightBinding>, ModuleBindingPlanError> {
let module = Module(ParameterBindingTarget {
shape,
dtype: if source.overlay {
RecipeDtype::U32
} else {
RecipeDtype::F32
},
permitted_source_dtypes: Vec::new(),
});
build_exact_replicated_text_bindings(
&module,
source,
&[task],
&BTreeSet::new(),
layout,
|value| Some(value.clone()),
|_, recipe, _| Ok::<_, Infallible>(recipe),
)
}
#[test]
fn exact_binding_places_each_rank_before_validating_its_local_slot() {
let source = source(vec![4, 2], false);
let task = task(&[4, 2], WeightLoweringKind::Direct, LinearFormat::Dense);
assert!(matches!(
bind(&source, &task, vec![2, 2], None),
Err(ModuleBindingPlanError::Plan(
BindingPlanError::ShapeMismatch { .. }
))
));
for rank in 0..2 {
let layout = layout(
vec![4, 2],
vec![2, 2],
TensorPlacement::Shard {
axis: 0,
index: rank,
parts: 2,
},
);
let bindings = bind(&source, &task, vec![2, 2], Some(&layout)).unwrap();
assert_eq!(bindings.len(), 1);
assert_eq!(bindings[0].logical_target(), Some("weight"));
assert_eq!(bindings[0].checkpoint_key(), "checkpoint.weight");
assert_eq!(
bindings[0].selection(),
&TensorSelection::Range {
axis: 0,
start: rank * 2,
end: rank * 2 + 2
}
);
assert_eq!(
bindings[0]
.source_recipe()
.infer(&source as &dyn CheckpointSource)
.unwrap()
.shape(),
[2, 2]
);
assert_eq!(bindings[0].expected_bytes(), 16);
assert!(matches!(
bind(&source, &task, vec![4, 2], Some(&layout)),
Err(ModuleBindingPlanError::Plan(
BindingPlanError::ShapeMismatch { .. }
))
));
}
assert!(bind(
&source,
&task,
vec![2, 2],
Some(&LocalModelLayout::default())
)
.is_err());
assert_eq!(source.leases.load(Ordering::SeqCst), 0);
}
#[test]
fn exact_binding_applies_compound_member_and_tensor_placements_once() {
let source = source(vec![4, 8, 2], false);
let task = task(&[4, 8, 2], WeightLoweringKind::Direct, LinearFormat::Dense);
let mut layout = LocalModelLayout::default();
layout.insert(
"weight".into(),
LocalTensorLayout::new(
"experts",
ParameterRole::ExpertIntermediate,
vec![4, 8, 2],
vec![2, 4, 2],
TensorPlacement::Shard {
axis: 1,
index: 1,
parts: 2,
},
None,
None,
false,
)
.with_additional_placement(TensorPlacement::Range {
axis: 0,
start: 2,
end: 4,
}),
);
let bindings = bind(&source, &task, vec![2, 4, 2], Some(&layout)).unwrap();
assert_eq!(
bindings[0]
.source_recipe()
.infer(&source as &dyn CheckpointSource)
.unwrap()
.shape(),
[2, 4, 2]
);
assert_eq!(bindings[0].expected_bytes(), 64);
assert_eq!(source.leases.load(Ordering::SeqCst), 0);
}
#[test]
fn exact_binding_does_not_reshard_selected_local_transform_output() {
let source = source(vec![2, 4], true);
let task = task(&[4, 32], WeightLoweringKind::Transform, LinearFormat::MxFp4);
let layout = layout(
vec![4, 32],
vec![2, 32],
TensorPlacement::Shard {
axis: 0,
index: 1,
parts: 2,
},
);
let bindings = bind(&source, &task, vec![2, 4], Some(&layout)).unwrap();
assert_eq!(bindings[0].checkpoint_key(), "weight");
assert_eq!(bindings[0].selection(), &TensorSelection::Full);
assert_eq!(
bindings[0]
.source_recipe()
.infer(&source as &dyn CheckpointSource)
.unwrap()
.shape(),
[2, 4]
);
assert_eq!(source.leases.load(Ordering::SeqCst), 0);
}