ipfrs_tensorlogic/gradient/
checkpoint.rs

1//! Gradient checkpointing — snapshot and restore training state.
2//!
3//! [`GradientCheckpoint`] wraps a [`ComputationGraphStore`] with optimizer
4//! metadata so that a training run can be interrupted and resumed without
5//! recomputing the full forward pass from scratch.
6
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10use super::computation_graph::ComputationGraphStore;
11use super::GradientError;
12
13/// A serializable snapshot of the gradient computation state.
14///
15/// Contains the full [`ComputationGraphStore`] plus training metadata so that
16/// a training run can be interrupted and resumed without recomputing the
17/// forward pass from scratch.
18#[derive(Debug, Serialize, Deserialize)]
19pub struct GradientCheckpoint {
20    /// The computation graph at the time of checkpointing
21    pub graph: ComputationGraphStore,
22    /// Optimizer step counter
23    pub step: u64,
24    /// CID of the loss tensor (None before the first loss is computed)
25    pub loss_cid: Option<String>,
26    /// Raw optimizer state blobs, keyed by parameter name
27    pub optimizer_state: HashMap<String, Vec<u8>>,
28    /// Unix timestamp (seconds) when this checkpoint was created
29    pub timestamp: u64,
30}
31
32impl GradientCheckpoint {
33    /// Create a new checkpoint at the given step.
34    pub fn new(graph: ComputationGraphStore, step: u64) -> Self {
35        let timestamp = std::time::SystemTime::now()
36            .duration_since(std::time::UNIX_EPOCH)
37            .map(|d| d.as_secs())
38            .unwrap_or(0);
39
40        Self {
41            graph,
42            step,
43            loss_cid: None,
44            optimizer_state: HashMap::new(),
45            timestamp,
46        }
47    }
48
49    /// Attach the CID of the current loss tensor.
50    pub fn with_loss_cid(mut self, cid: impl Into<String>) -> Self {
51        self.loss_cid = Some(cid.into());
52        self
53    }
54
55    /// Store a raw optimizer state blob for the given parameter.
56    pub fn set_optimizer_state(&mut self, param: impl Into<String>, state: Vec<u8>) {
57        self.optimizer_state.insert(param.into(), state);
58    }
59
60    /// Persist this checkpoint to a file at `path`.
61    ///
62    /// The file is written as JSON so it is human-inspectable and
63    /// version-control friendly.
64    pub fn save(&self, path: &std::path::Path) -> Result<(), GradientError> {
65        let bytes = serde_json::to_vec_pretty(self).map_err(|e| {
66            GradientError::InvalidGradient(format!("Checkpoint save serialization: {e}"))
67        })?;
68
69        // Use a temp-file-then-rename strategy for atomicity.
70        let parent = path.parent().unwrap_or_else(|| std::path::Path::new("."));
71        let tmp_path = parent.join(format!(".ckpt_tmp_{}.json", uuid::Uuid::new_v4()));
72
73        std::fs::write(&tmp_path, &bytes)
74            .map_err(|e| GradientError::InvalidGradient(format!("Checkpoint write: {e}")))?;
75
76        std::fs::rename(&tmp_path, path)
77            .map_err(|e| GradientError::InvalidGradient(format!("Checkpoint rename: {e}")))?;
78
79        Ok(())
80    }
81
82    /// Load a checkpoint from the file at `path`.
83    pub fn load(path: &std::path::Path) -> Result<Self, GradientError> {
84        let bytes = std::fs::read(path)
85            .map_err(|e| GradientError::InvalidGradient(format!("Checkpoint read: {e}")))?;
86
87        serde_json::from_slice(&bytes)
88            .map_err(|e| GradientError::InvalidGradient(format!("Checkpoint load parse: {e}")))
89    }
90}
ipfrs_tensorlogic/gradient/checkpoint.rs

ipfrs_tensorlogic/gradient/
checkpoint.rs
