eredu_evaluation/
checkpoint.rs

1//! Adapter from checkpoint-probe artifacts to general parity observations.
2
3use std::{
4    fs,
5    path::{Path, PathBuf},
6};
7
8use eredu_core::{
9    ObservationSelector, ObservationSet, ObservationValue, TensorObservation, TensorObservationData,
10};
11use safetensors::{tensor::SafeTensors, Dtype};
12use serde::{Deserialize, Serialize};
13
14use crate::{
15    compare_observations, LogitTolerance, ParityComparison, ParityPolicy, ParityReport, ParityRule,
16};
17
18const LOGIT_TENSORS: [&str; 2] = ["prefill.logits", "decode.logits"];
19
20/// Configuration for text-checkpoint logit parity.
21#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
22pub struct CheckpointParityOptions {
23    /// Accepted vocabulary-logit differences.
24    pub logits: LogitTolerance,
25}
26
27impl Default for CheckpointParityOptions {
28    fn default() -> Self {
29        Self {
30            logits: LogitTolerance {
31                relative_l2_max: 0.02,
32                cosine_similarity_min: 0.999,
33                top_k: 5,
34                top_k_overlap_min: 4,
35                require_unambiguous_argmax_match: true,
36                argmax_margin_min: 0.0,
37            },
38        }
39    }
40}
41
42/// General parity report plus source artifact identities.
43#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
44pub struct CheckpointParityReport {
45    /// Report schema version.
46    pub format_version: u32,
47    /// Stable evaluation kind.
48    pub kind: String,
49    /// Whether every identity and tensor comparison passed.
50    pub passed: bool,
51    /// Actual backend artifact report.
52    pub actual: PathBuf,
53    /// Reference artifact report.
54    pub reference: PathBuf,
55    /// Shared observation parity result.
56    pub parity: ParityReport,
57}
58
59/// Compares two checkpoint-probe report/tensor pairs using general parity.
60pub fn compare_checkpoint_artifacts(
61    actual: impl AsRef<Path>,
62    reference: impl AsRef<Path>,
63    options: CheckpointParityOptions,
64) -> Result<CheckpointParityReport, CheckpointParityError> {
65    let actual = actual.as_ref();
66    let reference = reference.as_ref();
67    let actual_observations = read_checkpoint_observations(actual)?;
68    let reference_observations = read_checkpoint_observations(reference)?;
69    let policy = ParityPolicy {
70        default: ParityComparison::Exact,
71        rules: vec![ParityRule {
72            selector: ObservationSelector::Prefix("logits".into()),
73            comparison: ParityComparison::Logits {
74                tolerance: options.logits,
75            },
76        }],
77        require_same_paths: true,
78    };
79    let parity = compare_observations(&actual_observations, &reference_observations, &policy)?;
80    Ok(CheckpointParityReport {
81        format_version: 1,
82        kind: "checkpoint_parity".into(),
83        passed: parity.passed,
84        actual: actual.into(),
85        reference: reference.into(),
86        parity,
87    })
88}
89
90fn read_checkpoint_observations(path: &Path) -> Result<ObservationSet, CheckpointParityError> {
91    let report: serde_json::Value = serde_json::from_slice(&fs::read(path)?)?;
92    let input = json_integer_array(&report, &["input", "token_ids"], path)?;
93    let fed = json_integer_array(&report, &["output", "fed_token_ids"], path)?;
94    let tensor_path = report
95        .pointer("/output/tensor_file")
96        .and_then(serde_json::Value::as_str)
97        .ok_or_else(|| CheckpointParityError::InvalidArtifact {
98            path: path.into(),
99            message: "missing string output.tensor_file".into(),
100        })?;
101    let tensor_path = resolve_tensor_path(path, Path::new(tensor_path));
102    let bytes = fs::read(&tensor_path)?;
103    let tensors = SafeTensors::deserialize(&bytes)?;
104
105    let mut observations = ObservationSet::new();
106    observations.insert(
107        "input.token_ids",
108        ObservationValue::Tensor(TensorObservation::new(
109            vec![input.len()],
110            TensorObservationData::I64(input),
111        )?),
112    )?;
113    observations.insert(
114        "output.fed_token_ids",
115        ObservationValue::Tensor(TensorObservation::new(
116            vec![fed.len()],
117            TensorObservationData::I64(fed),
118        )?),
119    )?;
120    for name in LOGIT_TENSORS {
121        let tensor = tensors.tensor(name)?;
122        if tensor.dtype() != Dtype::F32 {
123            return Err(CheckpointParityError::InvalidArtifact {
124                path: tensor_path.clone(),
125                message: format!("tensor {name:?} must be F32, got {:?}", tensor.dtype()),
126            });
127        }
128        let (chunks, remainder) = tensor.data().as_chunks::<4>();
129        if !remainder.is_empty() {
130            return Err(CheckpointParityError::InvalidArtifact {
131                path: tensor_path.clone(),
132                message: format!("tensor {name:?} has a partial F32 value"),
133            });
134        }
135        let values = chunks
136            .iter()
137            .map(|bytes| f32::from_le_bytes(*bytes))
138            .collect();
139        observations.insert(
140            format!("logits.{name}"),
141            ObservationValue::Tensor(TensorObservation::new(
142                tensor.shape().to_vec(),
143                TensorObservationData::F32(values),
144            )?),
145        )?;
146    }
147    Ok(observations)
148}
149
150fn json_integer_array(
151    report: &serde_json::Value,
152    path: &[&str],
153    source: &Path,
154) -> Result<Vec<i64>, CheckpointParityError> {
155    let pointer = format!("/{}", path.join("/"));
156    report
157        .pointer(&pointer)
158        .and_then(serde_json::Value::as_array)
159        .ok_or_else(|| CheckpointParityError::InvalidArtifact {
160            path: source.into(),
161            message: format!("missing integer array {}", path.join(".")),
162        })?
163        .iter()
164        .map(|value| {
165            value
166                .as_i64()
167                .ok_or_else(|| CheckpointParityError::InvalidArtifact {
168                    path: source.into(),
169                    message: format!("{} contains a non-integer value", path.join(".")),
170                })
171        })
172        .collect()
173}
174
175fn resolve_tensor_path(report: &Path, tensor: &Path) -> PathBuf {
176    if tensor.is_absolute() || tensor.exists() {
177        return tensor.into();
178    }
179    let relative = report
180        .parent()
181        .unwrap_or_else(|| Path::new("."))
182        .join(tensor);
183    if relative.exists() {
184        return relative;
185    }
186    report
187        .parent()
188        .unwrap_or_else(|| Path::new("."))
189        .join(tensor.file_name().unwrap_or_default())
190}
191
192/// Invalid checkpoint evidence or parity policy.
193#[derive(Debug, thiserror::Error)]
194pub enum CheckpointParityError {
195    /// Artifact I/O failed.
196    #[error(transparent)]
197    Io(#[from] std::io::Error),
198    /// Artifact JSON is invalid.
199    #[error(transparent)]
200    Json(#[from] serde_json::Error),
201    /// SafeTensors data is invalid or incomplete.
202    #[error(transparent)]
203    SafeTensors(#[from] safetensors::SafeTensorError),
204    /// Portable observation data is invalid.
205    #[error(transparent)]
206    Observation(#[from] eredu_core::ObservationError),
207    /// General comparison failed before a report could be produced.
208    #[error(transparent)]
209    Parity(#[from] crate::ParityError),
210    /// A required checkpoint artifact field is absent or malformed.
211    #[error("invalid checkpoint artifact {}: {message}", path.display())]
212    InvalidArtifact {
213        /// Artifact containing the invalid data.
214        path: PathBuf,
215        /// Specific failure.
216        message: String,
217    },
218}
219
220#[cfg(test)]
221mod tests {
222    use std::collections::BTreeMap;
223
224    use safetensors::tensor::{serialize_to_file, TensorView};
225    use tempfile::tempdir;
226
227    use super::*;
228
229    fn artifact(
230        root: &Path,
231        name: &str,
232        input: &[i64],
233        fed: &[i64],
234        prefill: &[f32],
235        decode: &[f32],
236    ) -> PathBuf {
237        let tensors = root.join(format!("{name}.safetensors"));
238        let prefill_bytes = prefill
239            .iter()
240            .flat_map(|value| value.to_le_bytes())
241            .collect::<Vec<_>>();
242        let decode_bytes = decode
243            .iter()
244            .flat_map(|value| value.to_le_bytes())
245            .collect::<Vec<_>>();
246        let mut views = BTreeMap::new();
247        views.insert(
248            "prefill.logits",
249            TensorView::new(Dtype::F32, vec![1, prefill.len()], &prefill_bytes).unwrap(),
250        );
251        views.insert(
252            "decode.logits",
253            TensorView::new(Dtype::F32, vec![1, decode.len()], &decode_bytes).unwrap(),
254        );
255        serialize_to_file(views, None, &tensors).unwrap();
256        let report = root.join(format!("{name}.json"));
257        fs::write(
258            &report,
259            serde_json::to_vec(&serde_json::json!({
260                "input": {"token_ids": input},
261                "output": {"fed_token_ids": fed, "tensor_file": tensors},
262            }))
263            .unwrap(),
264        )
265        .unwrap();
266        report
267    }
268
269    #[test]
270    fn checkpoint_adapter_uses_general_identity_and_logit_parity() {
271        let root = tempdir().unwrap();
272        let actual = artifact(
273            root.path(),
274            "actual",
275            &[1, 2],
276            &[3],
277            &[0.1, 0.8, 0.2],
278            &[0.7, 0.2, 0.1],
279        );
280        let reference = artifact(
281            root.path(),
282            "reference",
283            &[1, 2],
284            &[3],
285            &[0.1, 0.8, 0.2],
286            &[0.7, 0.2, 0.1],
287        );
288        assert!(
289            compare_checkpoint_artifacts(actual, reference, CheckpointParityOptions::default())
290                .unwrap()
291                .passed
292        );
293    }
294
295    #[test]
296    fn checkpoint_adapter_reports_token_and_numeric_failures_together() {
297        let root = tempdir().unwrap();
298        let actual = artifact(
299            root.path(),
300            "actual",
301            &[1, 9],
302            &[3],
303            &[0.9, 0.1, 0.0],
304            &[0.7, 0.2, 0.1],
305        );
306        let reference = artifact(
307            root.path(),
308            "reference",
309            &[1, 2],
310            &[3],
311            &[0.1, 0.8, 0.2],
312            &[0.7, 0.2, 0.1],
313        );
314        let report =
315            compare_checkpoint_artifacts(actual, reference, CheckpointParityOptions::default())
316                .unwrap();
317        assert!(!report.passed);
318        assert!(report
319            .parity
320            .failures
321            .iter()
322            .any(|failure| failure.contains("input.token_ids")));
323        assert!(report
324            .parity
325            .failures
326            .iter()
327            .any(|failure| failure.contains("prefill.logits")));
328    }
329}
eredu_evaluation/checkpoint.rs

eredu_evaluation/
checkpoint.rs
