foxhunt/crates/ml-core/src/checkpoint.rs

//! Safetensors save/load for `GpuVarStore`.
//!
//! Uses the `safetensors` crate directly (not candle's wrapper) to serialize
//! GPU parameters to disk and reload them.
//!
//! ## Format
//!
//! Each parameter is stored as a named F32 tensor in the safetensors file.
//! Metadata includes an optional JSON blob for model-level information
//! (model type, epoch, step, etc.).

use std::collections::BTreeMap;
use std::path::Path;
use std::sync::Arc;

use safetensors::Dtype as StDtype;
use safetensors::tensor::TensorView;
// Import SafeTensors from tensor module explicitly to avoid name conflicts
use safetensors::tensor::SafeTensors as SafeTensorsReader;
use safetensors::serialize_to_file;

use crate::MLError;

#[cfg(feature = "cuda")]
use cudarc::driver::CudaStream;

#[cfg(feature = "cuda")]
use crate::cuda_autograd::var_store::GpuVarStore;

/// Save all parameters from a `GpuVarStore` to a safetensors file.
///
/// Parameters are read from GPU to host and written as F32 tensors.
/// The optional `metadata` map is stored in the safetensors header.
///
/// # Errors
///
/// Returns `MLError` if any GPU readback or file I/O fails.
#[cfg(feature = "cuda")]
pub fn save_safetensors<P: AsRef<Path>>(
    store: &GpuVarStore,
    path: P,
    stream: &Arc<CudaStream>,
    metadata: Option<&BTreeMap<String, String>>,
) -> Result<(), MLError> {
    let _ = stream; // stream used implicitly by export_to_host (store holds its own stream)

    // Export all parameters to host
    let params = store.export_to_host()?; // gpu-exit: checkpoint export

    // Build safetensors data: name -> (shape_as_vec_usize, f32_bytes)
    // We need to hold the host data in scope while building TensorViews.
    let mut host_data: Vec<(String, Vec<usize>, Vec<u8>)> = Vec::with_capacity(params.len());
    for (name, (shape, values)) in &params {
        // Convert f32 slice to bytes
        let bytes: Vec<u8> = values.iter()
            .flat_map(|v| v.to_le_bytes())
            .collect();
        host_data.push((name.clone(), shape.clone(), bytes));
    }

    // Build tensor views
    let views: Vec<(String, TensorView<'_>)> = host_data.iter()
        .map(|(name, shape, bytes)| {
            let view = TensorView::new(StDtype::F32, shape.clone(), bytes)
                .map_err(|e| MLError::ModelError(format!(
                    "Failed to create TensorView for '{name}': {e}"
                )));
            view.map(|v| (name.clone(), v))
        })
        .collect::<Result<Vec<_>, _>>()?;

    // Build metadata as HashMap (safetensors API requirement)
    let meta: Option<std::collections::HashMap<String, String>> = metadata.map(|m| {
        m.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
    });

    // serialize_to_file expects IntoIterator<Item = (S: AsRef<str>, V: View)>
    serialize_to_file(views.into_iter(), meta, path.as_ref())
        .map_err(|e| MLError::CheckpointError(format!(
            "Failed to write safetensors to {}: {e}",
            path.as_ref().display()
        )))?;

    Ok(())
}

/// Load parameters from a safetensors file into a `GpuVarStore`.
///
/// Only loads parameters whose names match existing entries in the store.
/// Shape must match exactly.
///
/// # Errors
///
/// Returns `MLError` if the file cannot be read, parsed, or if shapes mismatch.
#[cfg(feature = "cuda")]
pub fn load_safetensors<P: AsRef<Path>>(
    store: &mut GpuVarStore,
    path: P,
    _stream: &Arc<CudaStream>,
) -> Result<BTreeMap<String, String>, MLError> {
    let file_bytes = std::fs::read(path.as_ref()).map_err(|e| {
        MLError::CheckpointError(format!(
            "Failed to read safetensors from {}: {e}",
            path.as_ref().display()
        ))
    })?;

    // Parse header to get metadata, then deserialize for tensor access
    let (_header_size, parsed_metadata) = SafeTensorsReader::read_metadata(&file_bytes).map_err(|e| {
        MLError::CheckpointError(format!(
            "Failed to parse safetensors metadata from {}: {e}",
            path.as_ref().display()
        ))
    })?;

    let metadata: BTreeMap<String, String> = match parsed_metadata.metadata() {
        Some(m) => m.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
        None => BTreeMap::new(),
    };

    let tensors = SafeTensorsReader::deserialize(&file_bytes).map_err(|e| {
        MLError::CheckpointError(format!(
            "Failed to parse safetensors from {}: {e}",
            path.as_ref().display()
        ))
    })?;

    // Build host data map
    let mut host_data: BTreeMap<String, (Vec<usize>, Vec<f32>)> = BTreeMap::new();

    for (name, _) in store.iter() {
        if let Ok(view) = tensors.tensor(name) {
            let shape: Vec<usize> = view.shape().to_vec();
            let data_bytes = view.data();
            // Convert bytes to f32 (little-endian)
            let numel = data_bytes.len() / 4;
            let mut values = Vec::with_capacity(numel);
            for chunk_idx in 0..numel {
                let offset = chunk_idx * 4;
                if offset + 4 <= data_bytes.len() {
                    let bytes = [
                        data_bytes[offset],
                        data_bytes[offset + 1],
                        data_bytes[offset + 2],
                        data_bytes[offset + 3],
                    ];
                    values.push(f32::from_le_bytes(bytes));
                }
            }
            host_data.insert(name.to_owned(), (shape, values));
        }
    }

    // Import into store
    store.import_from_host(&host_data)?;

    Ok(metadata)
}