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:
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user