feat(ml): wire TFT safetensors checkpoint save/load

Replace empty placeholder safetensors file with actual model weight
serialization using the VarMap, matching the DQN trainable adapter pattern.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-21 19:23:40 +01:00
parent 828e551f56
commit 0e1c8be186

View File

@@ -438,7 +438,7 @@ impl UnifiedTrainable for TrainableTFT {
/// Save model checkpoint in standardized format
///
/// Format: JSON for metadata (model weights require TFT refactoring)
/// Saves model weights via safetensors and metadata via JSON.
///
/// # Arguments
/// * `checkpoint_path` - Path to save checkpoint (without extension)
@@ -446,9 +446,6 @@ impl UnifiedTrainable for TrainableTFT {
/// # Returns
/// Path to saved checkpoint
fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
// TODO: Implement safetensors checkpoint save when TFT exposes VarMap
// For now, save only metadata
// Create and save checkpoint metadata
let metadata = CheckpointMetadata {
model_type: "TFT".to_string(),
@@ -464,34 +461,93 @@ impl UnifiedTrainable for TrainableTFT {
// Save metadata to JSON
crate::training::unified_trainer::checkpoint::save_metadata(&metadata, checkpoint_path)?;
// Create placeholder safetensors file for compatibility
// Save model weights to safetensors format
let safetensors_path = format!("{}.safetensors", checkpoint_path);
std::fs::write(&safetensors_path, b"")
.map_err(|e| MLError::ModelError(format!("Failed to create checkpoint file: {}", e)))?;
// Extract tensors from VarMap
let vars_data = self
.model
.varmap
.data()
.lock()
.map_err(|e| MLError::TrainingError(format!("Failed to lock VarMap for checkpoint save: {}", e)))?;
let mut tensors: HashMap<String, Tensor> = HashMap::new();
for (name, var) in vars_data.iter() {
tensors.insert(name.clone(), var.as_tensor().clone());
}
// Save using safetensors
candle_core::safetensors::save(&tensors, &safetensors_path)
.map_err(|e| MLError::ModelError(format!("Failed to save safetensors: {}", e)))?;
tracing::info!(
"Saved TFT checkpoint to {} (step {}, {} tensors)",
checkpoint_path,
self.step_count,
tensors.len()
);
Ok(safetensors_path)
}
/// Load model checkpoint from standardized format
///
/// Loads model weights from safetensors and metadata from JSON.
///
/// # Arguments
/// * `checkpoint_path` - Path to checkpoint (without extension)
///
/// # Returns
/// Loaded checkpoint metadata
fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
// TODO: Implement safetensors checkpoint load when TFT exposes VarMap
// For now, load only metadata
// Load metadata from JSON
let metadata =
crate::training::unified_trainer::checkpoint::load_metadata(checkpoint_path)?;
// Validate model type
if metadata.model_type != "TFT" {
return Err(MLError::ModelError(format!(
"Invalid model type in checkpoint: expected TFT, got {}",
metadata.model_type
)));
}
// Load model weights from safetensors
let safetensors_path = format!("{}.safetensors", checkpoint_path);
let tensors = candle_core::safetensors::load(&safetensors_path, &self.model.device)
.map_err(|e| MLError::ModelError(format!("Failed to load safetensors: {}", e)))?;
// Load tensors into VarMap
let vars_data = self
.model
.varmap
.data()
.lock()
.map_err(|e| MLError::TrainingError(format!("Failed to lock VarMap for checkpoint load: {}", e)))?;
for (name, tensor) in &tensors {
if let Some(var) = vars_data.get(name) {
var.set(tensor).map_err(|e| {
MLError::ModelError(format!("Failed to set var {}: {}", name, e))
})?;
} else {
tracing::warn!("Checkpoint contains unknown variable: {}", name);
}
}
// Update model state from metadata
self.step_count = metadata.step;
self.model.is_trained = true;
self.learning_rate = metadata.metrics.learning_rate;
tracing::info!(
"Loaded TFT checkpoint from {} (step {}, {} tensors)",
checkpoint_path,
metadata.step,
tensors.len()
);
Ok(metadata)
}