deep_delta_learning/

training.rs

use std::fmt::{Display, Formatter};

use burn::config::Config;
use burn::module::{AutodiffModule, Module};
use burn::nn::loss::CrossEntropyLossConfig;
use burn::optim::{AdamWConfig, GradientsParams, Optimizer};
use burn::prelude::*;
use burn::record::{BinBytesRecorder, FullPrecisionSettings, Recorder};
use burn::tensor::activation::log_softmax;
use burn::tensor::backend::AutodiffBackend;
use serde::{Deserialize, Serialize};

use crate::baseline::BaselineTransformer;
use crate::checkpoint::LoadedTrainingArtifact;
use crate::config::DdlConfig;
use crate::data::{TokenDataset, TokenDatasetSummary};
use crate::lm::{
    CausalLmMetrics, aggregate_causal_lm_summaries, causal_language_model_summary_with_lengths,
};
use crate::spectral::{BetaHistogram, DeltaRegime, SpectralDiagnostics};
use crate::transformer::DdlTransformer;
use crate::variant::{ModelInstance, ModelVariant};

#[derive(Config, Debug, PartialEq)]
pub struct TrainingConfig {
    #[config(default = 32)]
    pub max_steps: usize,
    #[config(default = 4)]
    pub eval_interval: usize,
    #[config(default = 1e-3)]
    pub learning_rate: f64,
    #[config(default = 0.0)]
    pub min_learning_rate: f64,
    #[config(default = 0)]
    pub warmup_steps: usize,
    #[config(default = 0.1)]
    pub weight_decay: f64,
    #[config(default = 0.9)]
    pub beta1: f64,
    #[config(default = 0.95)]
    pub beta2: f64,
    #[config(default = 1e-5)]
    pub epsilon: f64,
    #[config(default = 1.0)]
    pub grad_clip: f64,
}

impl TrainingConfig {
    pub fn validate(&self) -> Result<(), TrainingError> {
        if self.max_steps == 0 {
            return Err(TrainingError::InvalidConfig(
                "max_steps must be greater than zero",
            ));
        }
        if self.eval_interval == 0 {
            return Err(TrainingError::InvalidConfig(
                "eval_interval must be greater than zero",
            ));
        }
        if !(self.learning_rate.is_finite() && self.learning_rate > 0.0) {
            return Err(TrainingError::InvalidConfig(
                "learning_rate must be finite and positive",
            ));
        }
        if !(self.min_learning_rate.is_finite() && self.min_learning_rate >= 0.0) {
            return Err(TrainingError::InvalidConfig(
                "min_learning_rate must be finite and non-negative",
            ));
        }
        if self.min_learning_rate > self.learning_rate {
            return Err(TrainingError::InvalidConfig(
                "min_learning_rate cannot exceed learning_rate",
            ));
        }
        if self.warmup_steps > self.max_steps {
            return Err(TrainingError::InvalidConfig(
                "warmup_steps cannot exceed max_steps",
            ));
        }
        if !(self.weight_decay.is_finite() && self.weight_decay >= 0.0) {
            return Err(TrainingError::InvalidConfig(
                "weight_decay must be finite and non-negative",
            ));
        }
        if !(self.beta1.is_finite() && (0.0..1.0).contains(&self.beta1)) {
            return Err(TrainingError::InvalidConfig("beta1 must be in [0, 1)"));
        }
        if !(self.beta2.is_finite() && (0.0..1.0).contains(&self.beta2)) {
            return Err(TrainingError::InvalidConfig("beta2 must be in [0, 1)"));
        }
        if !(self.epsilon.is_finite() && self.epsilon > 0.0) {
            return Err(TrainingError::InvalidConfig(
                "epsilon must be finite and positive",
            ));
        }
        if !(self.grad_clip.is_finite() && self.grad_clip > 0.0) {
            return Err(TrainingError::InvalidConfig(
                "grad_clip must be finite and positive",
            ));
        }

        Ok(())
    }

    fn optimizer<M, B>(&self) -> impl Optimizer<M, B>
    where
        B: AutodiffBackend,
        M: AutodiffModule<B>,
    {
        AdamWConfig::new()
            .with_beta_1(self.beta1 as f32)
            .with_beta_2(self.beta2 as f32)
            .with_epsilon(self.epsilon as f32)
            .with_weight_decay(self.weight_decay as f32)
            .with_grad_clipping(Some(burn::grad_clipping::GradientClippingConfig::Norm(
                self.grad_clip as f32,
            )))
            .init::<B, M>()
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TrainingStepMetrics {
    pub step: usize,
    pub learning_rate: f64,
    pub train: CausalLmMetrics,
    pub validation: Option<CausalLmMetrics>,
    pub train_spectral: Option<TrainingSpectralSnapshot>,
    pub validation_spectral: Option<TrainingSpectralSnapshot>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TrainingSpectralSnapshot {
    pub beta_per_layer: Vec<f32>,
    pub k_eigenvalue_per_layer: Vec<f32>,
    pub spatial_determinant_per_layer: Vec<f32>,
    pub lifted_determinant_per_layer: Vec<f32>,
    pub regime_per_layer: Vec<DeltaRegime>,
    pub beta_histogram: BetaHistogram,
    pub k_coherence_per_layer: Vec<f32>,
    pub correction_norm_per_layer: Vec<f32>,
}

impl From<&SpectralDiagnostics> for TrainingSpectralSnapshot {
    fn from(spectral: &SpectralDiagnostics) -> Self {
        Self {
            beta_per_layer: spectral.beta_per_layer.clone(),
            k_eigenvalue_per_layer: spectral.k_eigenvalue_per_layer.clone(),
            spatial_determinant_per_layer: spectral.spatial_determinant_per_layer.clone(),
            lifted_determinant_per_layer: spectral.lifted_determinant_per_layer.clone(),
            regime_per_layer: spectral.regime_per_layer.clone(),
            beta_histogram: spectral.beta_histogram,
            k_coherence_per_layer: spectral.k_coherence_per_layer.clone(),
            correction_norm_per_layer: spectral.correction_norm_per_layer.clone(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TrainingReport {
    pub variant: ModelVariant,
    pub config: DdlConfig,
    pub training: TrainingConfig,
    pub num_params: usize,
    pub train_dataset: TokenDatasetSummary,
    pub validation_dataset: Option<TokenDatasetSummary>,
    pub steps_completed: usize,
    pub initial_train: CausalLmMetrics,
    pub initial_validation: Option<CausalLmMetrics>,
    pub final_train: CausalLmMetrics,
    pub final_validation: Option<CausalLmMetrics>,
    pub best_validation: Option<CausalLmMetrics>,
    pub initial_train_spectral: Option<TrainingSpectralSnapshot>,
    pub initial_validation_spectral: Option<TrainingSpectralSnapshot>,
    pub final_train_spectral: Option<TrainingSpectralSnapshot>,
    pub final_validation_spectral: Option<TrainingSpectralSnapshot>,
    pub best_validation_spectral: Option<TrainingSpectralSnapshot>,
    pub best_validation_step: Option<usize>,
    pub history: Vec<TrainingStepMetrics>,
}

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

/// Ranked summary for a multi-variant RFC-007 smoke run.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TrainingComparisonReport {
    pub train_dataset: TokenDatasetSummary,
    pub validation_dataset: Option<TokenDatasetSummary>,
    pub training: TrainingConfig,
    pub reports: Vec<TrainingReport>,
    pub final_train_loss_ranking: Vec<ModelVariant>,
    pub best_final_train_variant: ModelVariant,
    pub final_validation_loss_ranking: Option<Vec<ModelVariant>>,
    pub best_final_validation_variant: Option<ModelVariant>,
}

impl TrainingComparisonReport {
    pub fn report(&self, variant: ModelVariant) -> Option<&TrainingReport> {
        self.reports.iter().find(|report| report.variant == variant)
    }

    pub fn best_final_train(&self) -> Option<&TrainingReport> {
        self.report(self.best_final_train_variant)
    }

    pub fn best_final_validation(&self) -> Option<&TrainingReport> {
        self.best_final_validation_variant
            .and_then(|variant| self.report(variant))
    }
}

#[derive(Debug)]
pub struct TrainingSweepOutcome<B: Backend> {
    pub report: TrainingComparisonReport,
    pub outcomes: Vec<TrainingOutcome<B>>,
}

impl<B: Backend> TrainingSweepOutcome<B> {
    pub fn outcome(&self, variant: ModelVariant) -> Option<&TrainingOutcome<B>> {
        self.outcomes
            .iter()
            .find(|outcome| outcome.report.variant == variant)
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct CosineWarmupSchedule {
    base_learning_rate: f64,
    min_learning_rate: f64,
    warmup_steps: usize,
    total_steps: usize,
}

impl CosineWarmupSchedule {
    pub fn new(
        base_learning_rate: f64,
        min_learning_rate: f64,
        warmup_steps: usize,
        total_steps: usize,
    ) -> Result<Self, TrainingError> {
        if total_steps == 0 {
            return Err(TrainingError::InvalidConfig(
                "total_steps must be greater than zero",
            ));
        }
        if warmup_steps > total_steps {
            return Err(TrainingError::InvalidConfig(
                "warmup_steps cannot exceed total_steps",
            ));
        }
        if min_learning_rate > base_learning_rate {
            return Err(TrainingError::InvalidConfig(
                "min_learning_rate cannot exceed base_learning_rate",
            ));
        }

        Ok(Self {
            base_learning_rate,
            min_learning_rate,
            warmup_steps,
            total_steps,
        })
    }

    pub fn learning_rate(&self, step: usize) -> f64 {
        let step = step.min(self.total_steps.saturating_sub(1));

        if self.warmup_steps > 0 && step < self.warmup_steps {
            let progress = (step + 1) as f64 / self.warmup_steps as f64;
            return self.base_learning_rate * progress;
        }

        let cosine_steps = self.total_steps.saturating_sub(self.warmup_steps);
        if cosine_steps <= 1 {
            return self.base_learning_rate;
        }

        let cosine_step = step.saturating_sub(self.warmup_steps);
        let progress = cosine_step as f64 / (cosine_steps - 1) as f64;
        let cosine = (std::f64::consts::PI * progress).cos();
        self.min_learning_rate
            + (self.base_learning_rate - self.min_learning_rate) * 0.5 * (1.0 + cosine)
    }
}

#[derive(Debug, Clone)]
struct DatasetTelemetry {
    metrics: CausalLmMetrics,
    spectral: Option<TrainingSpectralSnapshot>,
}

#[derive(Debug)]
struct ResumeState<M> {
    report: TrainingReport,
    best_validation_model: Option<M>,
    optimizer_state: Vec<u8>,
}

#[derive(Debug)]
struct ForwardPassOutput<B: Backend> {
    logits: Tensor<B, 3>,
    spectral: Option<TrainingSpectralSnapshot>,
}

#[derive(Debug, Default)]
struct SpectralAccumulator {
    beta_per_layer: Vec<f64>,
    k_eigenvalue_per_layer: Vec<f64>,
    spatial_determinant_per_layer: Vec<f64>,
    lifted_determinant_per_layer: Vec<f64>,
    k_coherence_per_layer: Vec<f64>,
    correction_norm_per_layer: Vec<f64>,
    total_weight: f64,
}

impl SpectralAccumulator {
    fn observe(&mut self, spectral: &TrainingSpectralSnapshot, weight: f64) {
        if self.beta_per_layer.is_empty() {
            self.beta_per_layer = vec![0.0; spectral.beta_per_layer.len()];
            self.k_eigenvalue_per_layer = vec![0.0; spectral.k_eigenvalue_per_layer.len()];
            self.spatial_determinant_per_layer =
                vec![0.0; spectral.spatial_determinant_per_layer.len()];
            self.lifted_determinant_per_layer =
                vec![0.0; spectral.lifted_determinant_per_layer.len()];
            self.k_coherence_per_layer = vec![0.0; spectral.k_coherence_per_layer.len()];
            self.correction_norm_per_layer = vec![0.0; spectral.correction_norm_per_layer.len()];
        }

        accumulate_weighted(
            &mut self.beta_per_layer,
            &spectral.beta_per_layer,
            weight,
            "beta_per_layer",
        );
        accumulate_weighted(
            &mut self.k_eigenvalue_per_layer,
            &spectral.k_eigenvalue_per_layer,
            weight,
            "k_eigenvalue_per_layer",
        );
        accumulate_weighted(
            &mut self.spatial_determinant_per_layer,
            &spectral.spatial_determinant_per_layer,
            weight,
            "spatial_determinant_per_layer",
        );
        accumulate_weighted(
            &mut self.lifted_determinant_per_layer,
            &spectral.lifted_determinant_per_layer,
            weight,
            "lifted_determinant_per_layer",
        );
        accumulate_weighted(
            &mut self.k_coherence_per_layer,
            &spectral.k_coherence_per_layer,
            weight,
            "k_coherence_per_layer",
        );
        accumulate_weighted(
            &mut self.correction_norm_per_layer,
            &spectral.correction_norm_per_layer,
            weight,
            "correction_norm_per_layer",
        );

        self.total_weight += weight;
    }

    fn finish(self) -> Option<TrainingSpectralSnapshot> {
        if self.total_weight == 0.0 {
            return None;
        }

        let beta_per_layer = average_weighted(self.beta_per_layer, self.total_weight);
        let k_eigenvalue_per_layer =
            average_weighted(self.k_eigenvalue_per_layer, self.total_weight);
        let spatial_determinant_per_layer =
            average_weighted(self.spatial_determinant_per_layer, self.total_weight);
        let lifted_determinant_per_layer =
            average_weighted(self.lifted_determinant_per_layer, self.total_weight);
        let k_coherence_per_layer = average_weighted(self.k_coherence_per_layer, self.total_weight);
        let correction_norm_per_layer =
            average_weighted(self.correction_norm_per_layer, self.total_weight);
        let regime_per_layer = beta_per_layer
            .iter()
            .map(|beta| DeltaRegime::from_beta(*beta))
            .collect::<Vec<_>>();
        let beta_histogram = BetaHistogram::from_betas(&beta_per_layer);

        Some(TrainingSpectralSnapshot {
            beta_per_layer,
            k_eigenvalue_per_layer,
            spatial_determinant_per_layer,
            lifted_determinant_per_layer,
            regime_per_layer,
            beta_histogram,
            k_coherence_per_layer,
            correction_norm_per_layer,
        })
    }
}

#[derive(Debug)]
pub enum TrainingError {
    EmptyTrainingDataset,
    EmptyValidationDataset,
    EmptyVariantSet,
    InvalidConfig(&'static str),
    OptimizerState(String),
    ResumeStateMismatch(String),
}

impl Display for TrainingError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::EmptyTrainingDataset => {
                write!(f, "training dataset does not contain any batches")
            }
            Self::EmptyValidationDataset => {
                write!(f, "validation dataset does not contain any batches")
            }
            Self::EmptyVariantSet => write!(f, "at least one model variant is required"),
            Self::InvalidConfig(message) => write!(f, "invalid training configuration: {message}"),
            Self::OptimizerState(message) => {
                write!(f, "optimizer state serialization failed: {message}")
            }
            Self::ResumeStateMismatch(message) => {
                write!(f, "invalid resume state: {message}")
            }
        }
    }
}

impl std::error::Error for TrainingError {}

pub fn train_variant<B>(
    base_config: &DdlConfig,
    variant: ModelVariant,
    device: &B::Device,
    train_dataset: &TokenDataset,
    validation_dataset: Option<&TokenDataset>,
    training_config: &TrainingConfig,
) -> Result<TrainingOutcome<B::InnerBackend>, TrainingError>
where
    B: AutodiffBackend,
{
    training_config.validate()?;
    if train_dataset.batches().is_empty() {
        return Err(TrainingError::EmptyTrainingDataset);
    }
    if validation_dataset.is_some_and(|dataset| dataset.batches().is_empty()) {
        return Err(TrainingError::EmptyValidationDataset);
    }

    let resolved_config = variant.resolve_config(base_config);
    let datasets = DatasetPair {
        train: train_dataset,
        validation: validation_dataset,
    };
    match variant {
        ModelVariant::Baseline => train_model(
            variant,
            resolved_config.clone(),
            BaselineTransformer::<B>::new(&resolved_config, device),
            |model| ModelInstance::Baseline(Box::new(model)),
            device,
            datasets,
            training_config,
            None,
        ),
        _ => train_model(
            variant,
            resolved_config.clone(),
            resolved_config.init::<B>(device),
            |model| ModelInstance::Ddl(Box::new(model)),
            device,
            datasets,
            training_config,
            None,
        ),
    }
}

pub fn resume_training<B>(
    artifact: LoadedTrainingArtifact<B>,
    device: &B::Device,
    train_dataset: &TokenDataset,
    validation_dataset: Option<&TokenDataset>,
    training_config: &TrainingConfig,
) -> Result<TrainingOutcome<B::InnerBackend>, TrainingError>
where
    B: AutodiffBackend,
{
    training_config.validate()?;
    if train_dataset.batches().is_empty() {
        return Err(TrainingError::EmptyTrainingDataset);
    }
    if validation_dataset.is_some_and(|dataset| dataset.batches().is_empty()) {
        return Err(TrainingError::EmptyValidationDataset);
    }

    let train_summary = train_dataset.summary();
    if artifact.report.train_dataset != train_summary {
        return Err(TrainingError::ResumeStateMismatch(format!(
            "training dataset summary {:?} does not match saved artifact {:?}",
            train_summary, artifact.report.train_dataset
        )));
    }

    let validation_summary = validation_dataset.map(TokenDataset::summary);
    if artifact.report.validation_dataset != validation_summary {
        return Err(TrainingError::ResumeStateMismatch(format!(
            "validation dataset summary {:?} does not match saved artifact {:?}",
            validation_summary, artifact.report.validation_dataset
        )));
    }

    if artifact.report.steps_completed != artifact.report.history.len() {
        return Err(TrainingError::ResumeStateMismatch(format!(
            "saved report recorded {} completed steps but {} history entries",
            artifact.report.steps_completed,
            artifact.report.history.len()
        )));
    }

    let datasets = DatasetPair {
        train: train_dataset,
        validation: validation_dataset,
    };
    let optimizer_state = if artifact.report.steps_completed == 0 {
        Vec::new()
    } else {
        artifact.optimizer_state.ok_or_else(|| {
            TrainingError::ResumeStateMismatch(
                "saved training artifact is missing optimizer state".to_string(),
            )
        })?
    };
    match artifact.model {
        ModelInstance::Baseline(model) => {
            if artifact.report.variant.uses_ddl() {
                return Err(TrainingError::ResumeStateMismatch(format!(
                    "saved report expects variant {} but checkpoint contains a baseline model",
                    artifact.report.variant.slug()
                )));
            }
            let best_validation_model = match artifact.best_validation_model {
                Some(ModelInstance::Baseline(model)) => Some(model.valid()),
                Some(ModelInstance::Ddl(_)) => {
                    return Err(TrainingError::ResumeStateMismatch(
                        "best-validation checkpoint kind does not match the baseline artifact"
                            .to_string(),
                    ));
                }
                None => None,
            };
            train_model(
                artifact.report.variant,
                artifact.report.config.clone(),
                *model,
                |model| ModelInstance::Baseline(Box::new(model)),
                device,
                datasets,
                training_config,
                Some(ResumeState {
                    report: artifact.report,
                    best_validation_model,
                    optimizer_state,
                }),
            )
        }
        ModelInstance::Ddl(model) => {
            if !artifact.report.variant.uses_ddl() {
                return Err(TrainingError::ResumeStateMismatch(format!(
                    "saved report expects variant {} but checkpoint contains a DDL model",
                    artifact.report.variant.slug()
                )));
            }
            let best_validation_model = match artifact.best_validation_model {
                Some(ModelInstance::Ddl(model)) => Some(model.valid()),
                Some(ModelInstance::Baseline(_)) => {
                    return Err(TrainingError::ResumeStateMismatch(
                        "best-validation checkpoint kind does not match the DDL artifact"
                            .to_string(),
                    ));
                }
                None => None,
            };
            train_model(
                artifact.report.variant,
                artifact.report.config.clone(),
                *model,
                |model| ModelInstance::Ddl(Box::new(model)),
                device,
                datasets,
                training_config,
                Some(ResumeState {
                    report: artifact.report,
                    best_validation_model,
                    optimizer_state,
                }),
            )
        }
    }
}

/// Train a set of model variants under the same data/configuration and rank them by loss.
pub fn train_variants<B>(
    base_config: &DdlConfig,
    variants: &[ModelVariant],
    device: &B::Device,
    train_dataset: &TokenDataset,
    validation_dataset: Option<&TokenDataset>,
    training_config: &TrainingConfig,
) -> Result<TrainingSweepOutcome<B::InnerBackend>, TrainingError>
where
    B: AutodiffBackend,
{
    if variants.is_empty() {
        return Err(TrainingError::EmptyVariantSet);
    }

    let mut outcomes = Vec::with_capacity(variants.len());
    for variant in variants {
        outcomes.push(train_variant::<B>(
            base_config,
            *variant,
            device,
            train_dataset,
            validation_dataset,
            training_config,
        )?);
    }

    let reports = outcomes
        .iter()
        .map(|outcome| outcome.report.clone())
        .collect::<Vec<_>>();
    let final_train_loss_ranking = rank_variants_by_loss(
        reports
            .iter()
            .map(|report| (report.variant, report.final_train.loss)),
    );
    let final_validation_loss_ranking = reports
        .iter()
        .map(|report| {
            report
                .final_validation
                .map(|metrics| (report.variant, metrics.loss))
        })
        .collect::<Option<Vec<_>>>()
        .map(rank_variants_by_loss);

    Ok(TrainingSweepOutcome {
        report: TrainingComparisonReport {
            train_dataset: train_dataset.summary(),
            validation_dataset: validation_dataset.map(TokenDataset::summary),
            training: training_config.clone(),
            reports,
            best_final_train_variant: final_train_loss_ranking[0],
            final_train_loss_ranking,
            best_final_validation_variant: final_validation_loss_ranking
                .as_ref()
                .and_then(|ranking| ranking.first().copied()),
            final_validation_loss_ranking,
        },
        outcomes,
    })
}

#[derive(Clone, Copy)]
struct DatasetPair<'a> {
    train: &'a TokenDataset,
    validation: Option<&'a TokenDataset>,
}

#[allow(clippy::too_many_arguments)]
fn train_model<B, M, F>(
    variant: ModelVariant,
    resolved_config: DdlConfig,
    mut model: M,
    into_model_instance: F,
    device: &B::Device,
    datasets: DatasetPair<'_>,
    training_config: &TrainingConfig,
    resume: Option<ResumeState<M::InnerModule>>,
) -> Result<TrainingOutcome<B::InnerBackend>, TrainingError>
where
    B: AutodiffBackend,
    M: CausalLmModel<B> + AutodiffModule<B> + Module<B>,
    M::InnerModule: CausalLmModel<B::InnerBackend> + Clone,
    F: Fn(M::InnerModule) -> ModelInstance<B::InnerBackend> + Copy,
{
    let mut optimizer = training_config.optimizer::<M, B>();
    let (
        initial_train,
        initial_validation,
        initial_train_spectral,
        initial_validation_spectral,
        mut best_validation,
        mut best_validation_spectral,
        mut best_validation_step,
        mut best_validation_model,
        mut history,
        step_offset,
        resume_optimizer_state,
    ) = match resume {
        Some(resume) => (
            resume.report.initial_train,
            resume.report.initial_validation,
            resume.report.initial_train_spectral,
            resume.report.initial_validation_spectral,
            resume.report.best_validation,
            resume.report.best_validation_spectral,
            resume.report.best_validation_step,
            resume.best_validation_model,
            resume.report.history,
            resume.report.steps_completed,
            Some(resume.optimizer_state),
        ),
        None => {
            let initial_model = model.valid();
            let initial_train_eval =
                evaluate_dataset_with_telemetry(&initial_model, datasets.train, device);
            let initial_validation_eval = datasets
                .validation
                .map(|dataset| evaluate_dataset_with_telemetry(&initial_model, dataset, device));
            let initial_train = initial_train_eval.metrics;
            let initial_validation = initial_validation_eval
                .as_ref()
                .map(|evaluation| evaluation.metrics);

            (
                initial_train,
                initial_validation,
                initial_train_eval.spectral,
                initial_validation_eval
                    .clone()
                    .and_then(|evaluation| evaluation.spectral),
                initial_validation,
                initial_validation_eval
                    .as_ref()
                    .and_then(|evaluation| evaluation.spectral.clone()),
                initial_validation.map(|_| 0),
                initial_validation.map(|_| initial_model.clone()),
                Vec::with_capacity(training_config.max_steps),
                0,
                None,
            )
        }
    };
    if let Some(optimizer_state) = resume_optimizer_state {
        optimizer = load_optimizer_state::<_, M, B>(optimizer, &optimizer_state, device)?;
    }
    let schedule = schedule_for_phase(
        training_config,
        history.last().map(|step| step.learning_rate),
    )?;

    for step in 0..training_config.max_steps {
        let record_spectral = (step + 1) % training_config.eval_interval == 0
            || step + 1 == training_config.max_steps;
        let batch = &datasets.train.batches()[step % datasets.train.num_batches()];
        let input_ids = batch.to_tensor(device);
        let forward = if record_spectral {
            model.forward_with_spectral_snapshot(input_ids.clone(), None)
        } else {
            ForwardPassOutput {
                logits: model.forward_logits(input_ids.clone(), None),
                spectral: None,
            }
        };
        let logits = forward.logits;
        let loss = causal_language_model_training_loss(logits, input_ids, batch.sequence_lengths());
        let train = CausalLmMetrics {
            loss: scalar_from_autodiff_tensor(loss.clone()),
            perplexity: scalar_from_autodiff_tensor(loss.clone()).exp(),
        };
        let grads = GradientsParams::from_grads(loss.backward(), &model);
        let lr = schedule.learning_rate(step);
        model = optimizer.step(lr, model, grads);

        let mut validation_model = None;
        let validation_eval = if record_spectral {
            let valid_model = model.valid();
            let evaluation = datasets
                .validation
                .map(|dataset| evaluate_dataset_with_telemetry(&valid_model, dataset, device));
            validation_model = Some(valid_model);
            evaluation
        } else {
            None
        };
        let validation = validation_eval
            .as_ref()
            .map(|evaluation| evaluation.metrics);

        if let Some((metrics, spectral)) = validation_eval
            .as_ref()
            .map(|evaluation| (evaluation.metrics, evaluation.spectral.clone()))
        {
            let should_update = match best_validation {
                Some(best) => metrics.loss < best.loss,
                None => true,
            };
            if should_update {
                best_validation = Some(metrics);
                best_validation_spectral = spectral;
                best_validation_step = Some(step_offset + step + 1);
                best_validation_model = validation_model.clone();
            }
        }

        history.push(TrainingStepMetrics {
            step: step_offset + step + 1,
            learning_rate: lr,
            train,
            validation,
            train_spectral: forward.spectral,
            validation_spectral: validation_eval.and_then(|evaluation| evaluation.spectral),
        });
    }

    let model = model.valid();
    let final_train_eval = evaluate_dataset_with_telemetry(&model, datasets.train, device);
    let final_validation_eval = datasets
        .validation
        .map(|dataset| evaluate_dataset_with_telemetry(&model, dataset, device));
    let final_train = final_train_eval.metrics;
    let final_validation = final_validation_eval
        .as_ref()
        .map(|evaluation| evaluation.metrics);
    if let Some((metrics, spectral)) = final_validation_eval
        .as_ref()
        .map(|evaluation| (evaluation.metrics, evaluation.spectral.clone()))
    {
        let should_update = match best_validation {
            Some(best) => metrics.loss < best.loss,
            None => true,
        };
        if should_update {
            best_validation = Some(metrics);
            best_validation_spectral = spectral;
            best_validation_step = Some(step_offset + training_config.max_steps);
            best_validation_model = Some(model.clone());
        }
    }
    let num_params = model.num_params();

    Ok(TrainingOutcome {
        report: TrainingReport {
            variant,
            config: resolved_config,
            training: training_config.clone(),
            num_params,
            train_dataset: datasets.train.summary(),
            validation_dataset: datasets.validation.map(TokenDataset::summary),
            steps_completed: history.len(),
            initial_train,
            initial_validation,
            final_train,
            final_validation,
            best_validation,
            initial_train_spectral,
            initial_validation_spectral,
            final_train_spectral: final_train_eval.spectral,
            final_validation_spectral: final_validation_eval
                .and_then(|evaluation| evaluation.spectral),
            best_validation_spectral,
            best_validation_step,
            history,
        },
        model: into_model_instance(model),
        best_validation_model: best_validation_model.map(into_model_instance),
        optimizer_state: save_optimizer_state::<_, M, B>(&optimizer)?,
    })
}

trait CausalLmModel<B: Backend> {
    fn forward_logits(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> Tensor<B, 3>;

    fn forward_with_spectral_snapshot(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> ForwardPassOutput<B> {
        ForwardPassOutput {
            logits: self.forward_logits(input_ids, mask),
            spectral: None,
        }
    }
}

impl<B: Backend> CausalLmModel<B> for BaselineTransformer<B> {
    fn forward_logits(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> Tensor<B, 3> {
        self.forward_logits(input_ids, mask)
    }
}

impl<B: Backend> CausalLmModel<B> for DdlTransformer<B> {
    fn forward_logits(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> Tensor<B, 3> {
        self.forward_logits(input_ids, mask)
    }

    #[cfg(feature = "spectral")]
    fn forward_with_spectral_snapshot(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> ForwardPassOutput<B> {
        let (logits, _, spectral) = self.forward_with_spectral_diagnostics(input_ids, mask);
        ForwardPassOutput {
            logits,
            spectral: Some(TrainingSpectralSnapshot::from(&spectral)),
        }
    }
}

fn evaluate_dataset_with_telemetry<B, M>(
    model: &M,
    dataset: &TokenDataset,
    device: &B::Device,
) -> DatasetTelemetry
where
    B: Backend,
    M: CausalLmModel<B>,
{
    let mut summaries = Vec::with_capacity(dataset.num_batches());
    let mut spectral = SpectralAccumulator::default();
    for batch in dataset.batches() {
        let input_ids = batch.to_tensor(device);
        let output = model.forward_with_spectral_snapshot(input_ids.clone(), None);
        let summary = causal_language_model_summary_with_lengths(
            &output.logits,
            &input_ids,
            batch.sequence_lengths(),
        );
        if let Some(snapshot) = output.spectral.as_ref() {
            spectral.observe(snapshot, summary.prediction_count.max(1) as f64);
        }
        summaries.push(summary);
    }

    DatasetTelemetry {
        metrics: aggregate_causal_lm_summaries(&summaries).into(),
        spectral: spectral.finish(),
    }
}

fn accumulate_weighted(target: &mut [f64], values: &[f32], weight: f64, label: &str) {
    assert_eq!(
        target.len(),
        values.len(),
        "{label} length must remain stable across spectral snapshots"
    );

    for (slot, value) in target.iter_mut().zip(values.iter()) {
        *slot += f64::from(*value) * weight;
    }
}

fn average_weighted(values: Vec<f64>, total_weight: f64) -> Vec<f32> {
    values
        .into_iter()
        .map(|value| (value / total_weight) as f32)
        .collect()
}

fn causal_language_model_training_loss<B: Backend>(
    logits: Tensor<B, 3>,
    input_ids: Tensor<B, 2, Int>,
    sequence_lengths: &[usize],
) -> Tensor<B, 1> {
    let [batch_size, seq_len, vocab_size] = logits.dims();
    if seq_len < 2 {
        return Tensor::<B, 1>::zeros([1], &logits.device());
    }
    assert_eq!(
        sequence_lengths.len(),
        batch_size,
        "sequence_lengths must match the batch size"
    );

    let positions_per_row = seq_len - 1;
    let num_predictions = batch_size * positions_per_row;
    let valid_predictions = sequence_lengths
        .iter()
        .map(|length| length.saturating_sub(1).min(positions_per_row))
        .sum::<usize>();
    if valid_predictions == 0 {
        return Tensor::<B, 1>::zeros([1], &logits.device());
    }

    let shifted_logits = logits
        .slice([0..batch_size, 0..positions_per_row, 0..vocab_size])
        .reshape([num_predictions, vocab_size]);
    let shifted_targets = input_ids
        .slice([0..batch_size, 1..seq_len])
        .reshape([num_predictions]);
    let losses = CrossEntropyLossConfig::new()
        .init(&shifted_logits.device())
        .forward(shifted_logits.clone(), shifted_targets.clone());
    if valid_predictions == num_predictions {
        return losses;
    }

    let gathered_losses = log_softmax(shifted_logits, 1)
        .gather(1, shifted_targets.reshape([num_predictions, 1]))
        .reshape([num_predictions])
        .neg();
    let mask = prediction_mask(
        sequence_lengths,
        positions_per_row,
        &gathered_losses.device(),
    );
    gathered_losses.mul(mask).sum() / valid_predictions as f32
}

fn prediction_mask<B: Backend>(
    sequence_lengths: &[usize],
    positions_per_row: usize,
    device: &B::Device,
) -> Tensor<B, 1> {
    let values = sequence_lengths
        .iter()
        .flat_map(|length| {
            let valid_predictions = length.saturating_sub(1).min(positions_per_row);
            (0..positions_per_row).map(move |position| {
                if position < valid_predictions {
                    1.0
                } else {
                    0.0
                }
            })
        })
        .collect::<Vec<_>>();
    Tensor::<B, 1>::from_floats(values.as_slice(), device)
}

fn scalar_from_autodiff_tensor<B: AutodiffBackend>(tensor: Tensor<B, 1>) -> f32 {
    let values = tensor
        .inner()
        .into_data()
        .to_vec::<f32>()
        .expect("training loss tensor should convert to f32");
    assert_eq!(
        values.len(),
        1,
        "training loss tensor should contain exactly one scalar"
    );
    values[0]
}

fn schedule_for_phase(
    training_config: &TrainingConfig,
    previous_learning_rate: Option<f64>,
) -> Result<CosineWarmupSchedule, TrainingError> {
    let (base_learning_rate, warmup_steps) = match previous_learning_rate {
        Some(previous_learning_rate) => (
            previous_learning_rate.clamp(
                training_config.min_learning_rate,
                training_config.learning_rate,
            ),
            0,
        ),
        None => (training_config.learning_rate, training_config.warmup_steps),
    };

    CosineWarmupSchedule::new(
        base_learning_rate,
        training_config.min_learning_rate,
        warmup_steps,
        training_config.max_steps,
    )
}

fn save_optimizer_state<O, M, B>(optimizer: &O) -> Result<Vec<u8>, TrainingError>
where
    O: Optimizer<M, B>,
    M: AutodiffModule<B>,
    B: AutodiffBackend,
{
    BinBytesRecorder::<FullPrecisionSettings>::default()
        .record(optimizer.to_record(), ())
        .map_err(|error| TrainingError::OptimizerState(error.to_string()))
}

fn load_optimizer_state<O, M, B>(
    optimizer: O,
    bytes: &[u8],
    device: &B::Device,
) -> Result<O, TrainingError>
where
    O: Optimizer<M, B>,
    M: AutodiffModule<B>,
    B: AutodiffBackend,
{
    let record = BinBytesRecorder::<FullPrecisionSettings>::default()
        .load::<O::Record>(bytes.to_vec(), device)
        .map_err(|error| TrainingError::OptimizerState(error.to_string()))?;
    Ok(optimizer.load_record(record))
}

fn rank_variants_by_loss<I>(entries: I) -> Vec<ModelVariant>
where
    I: IntoIterator<Item = (ModelVariant, f32)>,
{
    let mut ranked = entries.into_iter().collect::<Vec<_>>();
    ranked.sort_by(|left, right| left.1.total_cmp(&right.1));
    ranked.into_iter().map(|(variant, _)| variant).collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cosine_schedule_warms_up_then_decays_to_floor() {
        let schedule = CosineWarmupSchedule::new(1.0, 0.1, 2, 6).unwrap();
        let lrs = (0..6)
            .map(|step| schedule.learning_rate(step))
            .collect::<Vec<_>>();

        assert_eq!(lrs[0], 0.5);
        assert_eq!(lrs[1], 1.0);
        assert!(lrs[2] <= 1.0);
        assert!(lrs[3] < lrs[2]);
        assert_eq!(lrs[5], 0.1);
    }
}