deep_delta_learning/

checkpoint.rs

use std::fmt::{Display, Formatter};
use std::path::{Path, PathBuf};

use burn::module::Module;
use burn::prelude::*;
use burn::record::DefaultRecorder;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};

use crate::config::DdlConfig;
use crate::training::{
    TrainingComparisonReport, TrainingOutcome, TrainingReport, TrainingSweepOutcome,
};
use crate::variant::{ModelInstance, ModelVariant};

const MANIFEST_FILE: &str = "manifest.json";
const WEIGHTS_STEM: &str = "model";
const TRAINING_REPORT_FILE: &str = "training_report.json";
const TRAINING_COMPARISON_REPORT_FILE: &str = "training_comparison.json";
const OPTIMIZER_STATE_FILE: &str = "optimizer_state.bin";
const BEST_VALIDATION_DIR: &str = "best_validation";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointManifest {
    pub variant: ModelVariant,
    pub config: DdlConfig,
    pub num_params: usize,
    pub weights_stem: String,
}

impl CheckpointManifest {
    pub fn new(variant: ModelVariant, config: DdlConfig, num_params: usize) -> Self {
        Self {
            variant,
            config,
            num_params,
            weights_stem: WEIGHTS_STEM.to_string(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct LoadedTrainingArtifact<B: Backend> {
    pub manifest: CheckpointManifest,
    pub report: TrainingReport,
    pub model: ModelInstance<B>,
    pub best_validation_model: Option<ModelInstance<B>>,
    pub optimizer_state: Option<Vec<u8>>,
}

#[derive(Debug)]
pub enum CheckpointError {
    Io(std::io::Error),
    Serde(serde_json::Error),
    Recorder(burn::record::RecorderError),
    ArtifactMismatch(String),
    MissingBestValidationArtifact(PathBuf),
    ModelVariantMismatch {
        variant: ModelVariant,
        actual_model: &'static str,
    },
}

impl Display for CheckpointError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(error) => write!(f, "checkpoint I/O failed: {error}"),
            Self::Serde(error) => write!(f, "checkpoint manifest serialization failed: {error}"),
            Self::Recorder(error) => write!(f, "checkpoint weight serialization failed: {error}"),
            Self::ArtifactMismatch(message) => {
                write!(f, "training artifact is inconsistent: {message}")
            }
            Self::MissingBestValidationArtifact(path) => write!(
                f,
                "training artifact is missing the best-validation checkpoint at {}",
                path.display()
            ),
            Self::ModelVariantMismatch {
                variant,
                actual_model,
            } => write!(
                f,
                "model kind {actual_model} does not match checkpoint variant {}",
                variant.slug()
            ),
        }
    }
}

impl std::error::Error for CheckpointError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Io(error) => Some(error),
            Self::Serde(error) => Some(error),
            Self::Recorder(error) => Some(error),
            Self::ArtifactMismatch(_) => None,
            Self::MissingBestValidationArtifact(_) => None,
            Self::ModelVariantMismatch { .. } => None,
        }
    }
}

impl From<std::io::Error> for CheckpointError {
    fn from(error: std::io::Error) -> Self {
        Self::Io(error)
    }
}

impl From<serde_json::Error> for CheckpointError {
    fn from(error: serde_json::Error) -> Self {
        Self::Serde(error)
    }
}

impl From<burn::record::RecorderError> for CheckpointError {
    fn from(error: burn::record::RecorderError) -> Self {
        Self::Recorder(error)
    }
}

pub fn save_checkpoint<B: Backend>(
    model: &ModelInstance<B>,
    variant: ModelVariant,
    config: &DdlConfig,
    directory: impl AsRef<Path>,
) -> Result<CheckpointManifest, CheckpointError> {
    ensure_variant_matches_model(model, variant)?;

    let directory = directory.as_ref();
    let resolved_config = variant.resolve_config(config);
    let manifest = CheckpointManifest::new(variant, resolved_config, model.num_params());
    let recorder = DefaultRecorder::new();

    std::fs::create_dir_all(directory)?;
    match model {
        ModelInstance::Baseline(model) => {
            model
                .as_ref()
                .clone()
                .save_file(weights_path(directory, &manifest.weights_stem), &recorder)?;
        }
        ModelInstance::Ddl(model) => {
            model
                .as_ref()
                .clone()
                .save_file(weights_path(directory, &manifest.weights_stem), &recorder)?;
        }
    }

    let manifest_json = serde_json::to_string_pretty(&manifest)?;
    std::fs::write(manifest_path(directory), manifest_json)?;

    Ok(manifest)
}

pub fn load_checkpoint<B: Backend>(
    directory: impl AsRef<Path>,
    device: &B::Device,
) -> Result<(CheckpointManifest, ModelInstance<B>), CheckpointError> {
    let directory = directory.as_ref();
    let manifest = load_manifest(directory)?;
    let recorder = DefaultRecorder::new();
    let model = match manifest.variant.init_model(&manifest.config, device) {
        ModelInstance::Baseline(model) => {
            ModelInstance::Baseline(Box::new(model.as_ref().clone().load_file(
                weights_path(directory, &manifest.weights_stem),
                &recorder,
                device,
            )?))
        }
        ModelInstance::Ddl(model) => {
            ModelInstance::Ddl(Box::new(model.as_ref().clone().load_file(
                weights_path(directory, &manifest.weights_stem),
                &recorder,
                device,
            )?))
        }
    };

    Ok((manifest, model))
}

pub fn load_manifest(directory: impl AsRef<Path>) -> Result<CheckpointManifest, CheckpointError> {
    let manifest = std::fs::read_to_string(manifest_path(directory.as_ref()))?;
    Ok(serde_json::from_str(&manifest)?)
}

pub fn load_training_artifact<B: Backend>(
    directory: impl AsRef<Path>,
    device: &B::Device,
) -> Result<LoadedTrainingArtifact<B>, CheckpointError> {
    let directory = directory.as_ref();
    let report = load_training_report(directory)?;
    let (manifest, model) = load_checkpoint::<B>(directory, device)?;

    if report.variant != manifest.variant {
        return Err(CheckpointError::ArtifactMismatch(format!(
            "report variant {} does not match manifest variant {}",
            report.variant.slug(),
            manifest.variant.slug()
        )));
    }
    if report.config != manifest.config {
        return Err(CheckpointError::ArtifactMismatch(
            "report config does not match manifest config".to_string(),
        ));
    }
    if report.num_params != manifest.num_params {
        return Err(CheckpointError::ArtifactMismatch(format!(
            "report num_params {} does not match manifest {}",
            report.num_params, manifest.num_params
        )));
    }

    let best_validation_directory = best_validation_path(directory);
    let best_validation_model = if report.best_validation.is_none() {
        None
    } else if best_validation_directory.join(MANIFEST_FILE).is_file() {
        Some(load_checkpoint::<B>(&best_validation_directory, device)?.1)
    } else if report.best_validation_step == Some(report.steps_completed) {
        Some(model.clone())
    } else {
        return Err(CheckpointError::MissingBestValidationArtifact(
            best_validation_directory,
        ));
    };
    let optimizer_state = optimizer_state_path(directory)
        .is_file()
        .then(|| std::fs::read(optimizer_state_path(directory)))
        .transpose()?;

    Ok(LoadedTrainingArtifact {
        manifest,
        report,
        model,
        best_validation_model,
        optimizer_state,
    })
}

pub fn save_training_report(
    report: &TrainingReport,
    directory: impl AsRef<Path>,
) -> Result<(), CheckpointError> {
    let directory = directory.as_ref();
    std::fs::create_dir_all(directory)?;
    write_json(training_report_path(directory), report)
}

pub fn load_training_report(
    directory: impl AsRef<Path>,
) -> Result<TrainingReport, CheckpointError> {
    read_json(training_report_path(directory.as_ref()))
}

pub fn save_training_comparison_report(
    report: &TrainingComparisonReport,
    directory: impl AsRef<Path>,
) -> Result<(), CheckpointError> {
    let directory = directory.as_ref();
    std::fs::create_dir_all(directory)?;
    write_json(training_comparison_report_path(directory), report)
}

pub fn load_training_comparison_report(
    directory: impl AsRef<Path>,
) -> Result<TrainingComparisonReport, CheckpointError> {
    read_json(training_comparison_report_path(directory.as_ref()))
}

pub fn save_training_artifact<B: Backend>(
    outcome: &TrainingOutcome<B>,
    directory: impl AsRef<Path>,
) -> Result<CheckpointManifest, CheckpointError> {
    let directory = directory.as_ref();
    let manifest = save_checkpoint(
        &outcome.model,
        outcome.report.variant,
        &outcome.report.config,
        directory,
    )?;
    save_training_report(&outcome.report, directory)?;
    std::fs::write(optimizer_state_path(directory), &outcome.optimizer_state)?;

    let best_validation_directory = best_validation_path(directory);
    if best_validation_directory.exists() {
        std::fs::remove_dir_all(&best_validation_directory)?;
    }
    if let Some(best_model) = outcome.best_validation_model.as_ref() {
        save_checkpoint(
            best_model,
            outcome.report.variant,
            &outcome.report.config,
            &best_validation_directory,
        )?;
        save_training_report(&outcome.report, &best_validation_directory)?;
    }

    Ok(manifest)
}

pub fn save_training_sweep<B: Backend>(
    sweep: &TrainingSweepOutcome<B>,
    directory: impl AsRef<Path>,
) -> Result<Vec<CheckpointManifest>, CheckpointError> {
    let directory = directory.as_ref();
    std::fs::create_dir_all(directory)?;
    save_training_comparison_report(&sweep.report, directory)?;

    sweep
        .outcomes
        .iter()
        .map(|outcome| {
            save_training_artifact(outcome, directory.join(outcome.report.variant.slug()))
        })
        .collect()
}

fn ensure_variant_matches_model<B: Backend>(
    model: &ModelInstance<B>,
    variant: ModelVariant,
) -> Result<(), CheckpointError> {
    match (model, variant.uses_ddl()) {
        (ModelInstance::Baseline(_), false) | (ModelInstance::Ddl(_), true) => Ok(()),
        (ModelInstance::Baseline(_), true) => Err(CheckpointError::ModelVariantMismatch {
            variant,
            actual_model: "baseline",
        }),
        (ModelInstance::Ddl(_), false) => Err(CheckpointError::ModelVariantMismatch {
            variant,
            actual_model: "ddl",
        }),
    }
}

fn manifest_path(directory: &Path) -> PathBuf {
    directory.join(MANIFEST_FILE)
}

fn weights_path(directory: &Path, weights_stem: &str) -> PathBuf {
    directory.join(weights_stem)
}

fn training_report_path(directory: &Path) -> PathBuf {
    directory.join(TRAINING_REPORT_FILE)
}

fn training_comparison_report_path(directory: &Path) -> PathBuf {
    directory.join(TRAINING_COMPARISON_REPORT_FILE)
}

fn optimizer_state_path(directory: &Path) -> PathBuf {
    directory.join(OPTIMIZER_STATE_FILE)
}

fn best_validation_path(directory: &Path) -> PathBuf {
    directory.join(BEST_VALIDATION_DIR)
}

fn write_json<T: Serialize>(path: PathBuf, value: &T) -> Result<(), CheckpointError> {
    std::fs::write(path, serde_json::to_string_pretty(value)?)?;
    Ok(())
}

fn read_json<T: DeserializeOwned>(path: PathBuf) -> Result<T, CheckpointError> {
    let contents = std::fs::read_to_string(path)?;
    Ok(serde_json::from_str(&contents)?)
}