670 lines
20 KiB
Rust
670 lines
20 KiB
Rust
//! Activation Checkpointing for Metal
|
|
//!
|
|
//! This module provides activation checkpointing to trade compute for memory,
|
|
//! allowing training of larger models by recomputing activations during
|
|
//! the backward pass instead of storing them.
|
|
//!
|
|
//! ## Strategy
|
|
//!
|
|
//! Instead of storing all intermediate activations for the backward pass,
|
|
//! we only store activations at certain "checkpoint" layers. During backward,
|
|
//! we recompute the activations between checkpoints as needed.
|
|
|
|
use crate::device::MetalDevice;
|
|
use crate::error::Result;
|
|
use crate::memory::MetalBuffer;
|
|
use std::collections::HashMap;
|
|
use std::ops::Range;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use std::sync::Arc;
|
|
use parking_lot::RwLock;
|
|
|
|
/// Unique identifier for a checkpoint
|
|
pub type CheckpointId = u64;
|
|
|
|
/// Policy for determining which layers to checkpoint
|
|
#[derive(Debug, Clone)]
|
|
pub enum CheckpointPolicy {
|
|
/// Checkpoint every N layers
|
|
Every(usize),
|
|
/// Checkpoint to stay within memory budget (bytes)
|
|
MemoryBudget(usize),
|
|
/// Checkpoint specific layers only
|
|
SelectiveLayers(Vec<usize>),
|
|
/// No checkpointing (store all activations)
|
|
None,
|
|
/// Automatic: choose based on model and memory
|
|
Auto {
|
|
/// Target memory reduction factor
|
|
target_reduction: f32,
|
|
},
|
|
}
|
|
|
|
impl Default for CheckpointPolicy {
|
|
fn default() -> Self {
|
|
Self::Every(4)
|
|
}
|
|
}
|
|
|
|
impl CheckpointPolicy {
|
|
/// Determine if a layer should be checkpointed
|
|
pub fn should_checkpoint(&self, layer_idx: usize, total_layers: usize) -> bool {
|
|
match self {
|
|
Self::Every(n) => layer_idx % n == 0,
|
|
Self::MemoryBudget(_) => {
|
|
// Heuristic: checkpoint every sqrt(total_layers) layers
|
|
let interval = (total_layers as f64).sqrt().ceil() as usize;
|
|
layer_idx % interval == 0
|
|
}
|
|
Self::SelectiveLayers(layers) => layers.contains(&layer_idx),
|
|
Self::None => false,
|
|
Self::Auto { target_reduction } => {
|
|
// Estimate interval based on target reduction
|
|
// reduction = 1 - 1/interval, so interval = 1/(1-reduction)
|
|
let interval = (1.0 / (1.0 - target_reduction)).ceil() as usize;
|
|
layer_idx % interval == 0
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration for checkpointing
|
|
#[derive(Debug, Clone)]
|
|
pub struct CheckpointConfig {
|
|
/// Checkpointing policy
|
|
pub policy: CheckpointPolicy,
|
|
/// Maximum checkpoints to keep in memory
|
|
pub max_checkpoints: usize,
|
|
/// Whether to use async recompute
|
|
pub async_recompute: bool,
|
|
/// Memory overhead threshold before evicting old checkpoints
|
|
pub memory_threshold: f32,
|
|
}
|
|
|
|
impl Default for CheckpointConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
policy: CheckpointPolicy::default(),
|
|
max_checkpoints: 64,
|
|
async_recompute: true,
|
|
memory_threshold: 0.9,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Data stored at a checkpoint
|
|
pub struct CheckpointData {
|
|
/// Input tensors saved at this checkpoint
|
|
pub input_tensors: Vec<MetalBuffer<f32>>,
|
|
/// Layer indices covered by this checkpoint (for recompute)
|
|
pub layer_range: Range<usize>,
|
|
/// Whether this checkpoint is pinned (won't be evicted)
|
|
pub pinned: bool,
|
|
/// Reference count for this checkpoint
|
|
ref_count: AtomicU64,
|
|
}
|
|
|
|
impl CheckpointData {
|
|
/// Create a new checkpoint
|
|
pub fn new(input_tensors: Vec<MetalBuffer<f32>>, layer_range: Range<usize>) -> Self {
|
|
Self {
|
|
input_tensors,
|
|
layer_range,
|
|
pinned: false,
|
|
ref_count: AtomicU64::new(1),
|
|
}
|
|
}
|
|
|
|
/// Get the memory footprint of this checkpoint
|
|
pub fn memory_bytes(&self) -> usize {
|
|
self.input_tensors.iter().map(|t| t.size_bytes()).sum()
|
|
}
|
|
|
|
/// Increment reference count
|
|
pub fn inc_ref(&self) {
|
|
self.ref_count.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Decrement reference count
|
|
pub fn dec_ref(&self) -> u64 {
|
|
self.ref_count.fetch_sub(1, Ordering::Relaxed)
|
|
}
|
|
|
|
/// Get reference count
|
|
pub fn ref_count(&self) -> u64 {
|
|
self.ref_count.load(Ordering::Relaxed)
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Debug for CheckpointData {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("CheckpointData")
|
|
.field("layer_range", &self.layer_range)
|
|
.field("num_tensors", &self.input_tensors.len())
|
|
.field("memory_bytes", &self.memory_bytes())
|
|
.field("pinned", &self.pinned)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
/// Statistics for checkpoint manager
|
|
#[derive(Debug, Default, Clone)]
|
|
pub struct CheckpointStats {
|
|
/// Total checkpoints created
|
|
pub total_created: u64,
|
|
/// Total checkpoints evicted
|
|
pub total_evicted: u64,
|
|
/// Total recomputes performed
|
|
pub total_recomputes: u64,
|
|
/// Memory saved by checkpointing (estimated bytes)
|
|
pub memory_saved: u64,
|
|
/// Current checkpoints in memory
|
|
pub current_checkpoints: usize,
|
|
/// Current memory usage
|
|
pub current_memory: usize,
|
|
}
|
|
|
|
/// Metal Checkpoint Manager
|
|
pub struct MetalCheckpointManager {
|
|
/// Device
|
|
device: Arc<MetalDevice>,
|
|
/// Configuration
|
|
config: CheckpointConfig,
|
|
/// Stored checkpoints
|
|
checkpoints: RwLock<HashMap<CheckpointId, CheckpointData>>,
|
|
/// Next checkpoint ID
|
|
next_id: AtomicU64,
|
|
/// Statistics
|
|
stats: RwLock<CheckpointStats>,
|
|
/// Layer to checkpoint ID mapping
|
|
layer_checkpoints: RwLock<HashMap<usize, CheckpointId>>,
|
|
}
|
|
|
|
impl MetalCheckpointManager {
|
|
/// Create a new checkpoint manager
|
|
pub fn new(device: Arc<MetalDevice>, config: CheckpointConfig) -> Self {
|
|
Self {
|
|
device,
|
|
config,
|
|
checkpoints: RwLock::new(HashMap::new()),
|
|
next_id: AtomicU64::new(1),
|
|
stats: RwLock::new(CheckpointStats::default()),
|
|
layer_checkpoints: RwLock::new(HashMap::new()),
|
|
}
|
|
}
|
|
|
|
/// Create with default configuration
|
|
pub fn default_config(device: Arc<MetalDevice>) -> Self {
|
|
Self::new(device, CheckpointConfig::default())
|
|
}
|
|
|
|
/// Check if a layer should be checkpointed
|
|
pub fn should_checkpoint(&self, layer_idx: usize, total_layers: usize) -> bool {
|
|
self.config.policy.should_checkpoint(layer_idx, total_layers)
|
|
}
|
|
|
|
/// Save a checkpoint
|
|
pub fn save_checkpoint(
|
|
&self,
|
|
layer_idx: usize,
|
|
inputs: Vec<MetalBuffer<f32>>,
|
|
layer_range: Range<usize>,
|
|
) -> Result<CheckpointId> {
|
|
// Check if we need to evict
|
|
self.evict_if_needed()?;
|
|
|
|
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
|
|
let checkpoint = CheckpointData::new(inputs, layer_range.clone());
|
|
let memory = checkpoint.memory_bytes();
|
|
|
|
{
|
|
let mut checkpoints = self.checkpoints.write();
|
|
checkpoints.insert(id, checkpoint);
|
|
}
|
|
|
|
{
|
|
let mut layer_map = self.layer_checkpoints.write();
|
|
layer_map.insert(layer_idx, id);
|
|
}
|
|
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.total_created += 1;
|
|
stats.current_checkpoints += 1;
|
|
stats.current_memory += memory;
|
|
}
|
|
|
|
Ok(id)
|
|
}
|
|
|
|
/// Get a checkpoint for recomputation
|
|
pub fn get_checkpoint(&self, id: CheckpointId) -> Option<CheckpointDataRef> {
|
|
let checkpoints = self.checkpoints.read();
|
|
if checkpoints.contains_key(&id) {
|
|
Some(CheckpointDataRef { manager: self, id })
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Get checkpoint by layer index
|
|
pub fn get_checkpoint_for_layer(&self, layer_idx: usize) -> Option<CheckpointDataRef> {
|
|
let layer_map = self.layer_checkpoints.read();
|
|
let id = layer_map.get(&layer_idx)?;
|
|
self.get_checkpoint(*id)
|
|
}
|
|
|
|
/// Find the nearest checkpoint before a given layer
|
|
pub fn find_nearest_checkpoint(&self, layer_idx: usize) -> Option<(CheckpointId, usize)> {
|
|
let layer_map = self.layer_checkpoints.read();
|
|
|
|
layer_map
|
|
.iter()
|
|
.filter(|(l, _)| **l <= layer_idx)
|
|
.max_by_key(|(l, _)| *l)
|
|
.map(|(l, id)| (*id, *l))
|
|
}
|
|
|
|
/// Mark a recompute as complete
|
|
pub fn record_recompute(&self) {
|
|
let mut stats = self.stats.write();
|
|
stats.total_recomputes += 1;
|
|
}
|
|
|
|
/// Release a checkpoint
|
|
pub fn release_checkpoint(&self, id: CheckpointId) -> Result<()> {
|
|
let should_remove = {
|
|
let checkpoints = self.checkpoints.read();
|
|
if let Some(cp) = checkpoints.get(&id) {
|
|
cp.dec_ref() <= 1 && !cp.pinned
|
|
} else {
|
|
false
|
|
}
|
|
};
|
|
|
|
if should_remove {
|
|
let memory = {
|
|
let mut checkpoints = self.checkpoints.write();
|
|
if let Some(cp) = checkpoints.remove(&id) {
|
|
cp.memory_bytes()
|
|
} else {
|
|
0
|
|
}
|
|
};
|
|
|
|
// Update stats
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.current_checkpoints = stats.current_checkpoints.saturating_sub(1);
|
|
stats.current_memory = stats.current_memory.saturating_sub(memory);
|
|
}
|
|
|
|
// Remove from layer map
|
|
{
|
|
let mut layer_map = self.layer_checkpoints.write();
|
|
layer_map.retain(|_, &mut v| v != id);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Pin a checkpoint (prevent eviction)
|
|
pub fn pin_checkpoint(&self, id: CheckpointId) {
|
|
let mut checkpoints = self.checkpoints.write();
|
|
if let Some(cp) = checkpoints.get_mut(&id) {
|
|
cp.pinned = true;
|
|
}
|
|
}
|
|
|
|
/// Unpin a checkpoint
|
|
pub fn unpin_checkpoint(&self, id: CheckpointId) {
|
|
let mut checkpoints = self.checkpoints.write();
|
|
if let Some(cp) = checkpoints.get_mut(&id) {
|
|
cp.pinned = false;
|
|
}
|
|
}
|
|
|
|
/// Evict old checkpoints if memory is constrained
|
|
fn evict_if_needed(&self) -> Result<()> {
|
|
let checkpoints = self.checkpoints.read();
|
|
|
|
if checkpoints.len() >= self.config.max_checkpoints {
|
|
drop(checkpoints);
|
|
|
|
// Find candidates for eviction (unpinned, low ref count)
|
|
let mut candidates: Vec<(CheckpointId, u64)> = {
|
|
let checkpoints = self.checkpoints.read();
|
|
checkpoints
|
|
.iter()
|
|
.filter(|(_, cp)| !cp.pinned)
|
|
.map(|(&id, cp)| (id, cp.ref_count()))
|
|
.collect()
|
|
};
|
|
|
|
// Sort by ref count (evict least used first)
|
|
candidates.sort_by_key(|(_, count)| *count);
|
|
|
|
// Evict oldest unpinned checkpoint
|
|
if let Some((id, _)) = candidates.first() {
|
|
let memory = {
|
|
let mut checkpoints = self.checkpoints.write();
|
|
if let Some(cp) = checkpoints.remove(id) {
|
|
cp.memory_bytes()
|
|
} else {
|
|
0
|
|
}
|
|
};
|
|
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.total_evicted += 1;
|
|
stats.current_checkpoints = stats.current_checkpoints.saturating_sub(1);
|
|
stats.current_memory = stats.current_memory.saturating_sub(memory);
|
|
}
|
|
|
|
{
|
|
let mut layer_map = self.layer_checkpoints.write();
|
|
layer_map.retain(|_, &mut v| v != *id);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Clear all checkpoints
|
|
pub fn clear(&self) {
|
|
let mut checkpoints = self.checkpoints.write();
|
|
checkpoints.clear();
|
|
|
|
let mut layer_map = self.layer_checkpoints.write();
|
|
layer_map.clear();
|
|
|
|
let mut stats = self.stats.write();
|
|
stats.current_checkpoints = 0;
|
|
stats.current_memory = 0;
|
|
}
|
|
|
|
/// Get statistics
|
|
pub fn stats(&self) -> CheckpointStats {
|
|
self.stats.read().clone()
|
|
}
|
|
|
|
/// Get configuration
|
|
pub fn config(&self) -> &CheckpointConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Get device
|
|
pub fn device(&self) -> &Arc<MetalDevice> {
|
|
&self.device
|
|
}
|
|
|
|
/// Get number of checkpoints
|
|
pub fn checkpoint_count(&self) -> usize {
|
|
self.checkpoints.read().len()
|
|
}
|
|
|
|
/// Get total memory used by checkpoints
|
|
pub fn memory_usage(&self) -> usize {
|
|
self.stats.read().current_memory
|
|
}
|
|
|
|
/// Estimate memory saved by checkpointing
|
|
pub fn estimate_memory_saved(&self, total_layers: usize, activation_size: usize) -> usize {
|
|
let checkpoint_count = self.checkpoint_count();
|
|
if checkpoint_count == 0 {
|
|
return 0;
|
|
}
|
|
|
|
// Without checkpointing: store all activations
|
|
let without = total_layers * activation_size;
|
|
|
|
// With checkpointing: store checkpoint activations + recompute
|
|
let with = checkpoint_count * activation_size;
|
|
|
|
without.saturating_sub(with)
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Debug for MetalCheckpointManager {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
let stats = self.stats();
|
|
f.debug_struct("MetalCheckpointManager")
|
|
.field("checkpoint_count", &stats.current_checkpoints)
|
|
.field("memory_usage", &stats.current_memory)
|
|
.field("total_recomputes", &stats.total_recomputes)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
/// Reference to checkpoint data (RAII)
|
|
pub struct CheckpointDataRef<'a> {
|
|
manager: &'a MetalCheckpointManager,
|
|
id: CheckpointId,
|
|
}
|
|
|
|
impl<'a> CheckpointDataRef<'a> {
|
|
/// Get the checkpoint data
|
|
pub fn data(&self) -> Option<impl std::ops::Deref<Target = CheckpointData> + '_> {
|
|
let checkpoints = self.manager.checkpoints.read();
|
|
if checkpoints.contains_key(&self.id) {
|
|
// Increment ref count
|
|
checkpoints.get(&self.id).unwrap().inc_ref();
|
|
Some(parking_lot::RwLockReadGuard::map(checkpoints, |c| {
|
|
c.get(&self.id).unwrap()
|
|
}))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Get the checkpoint ID
|
|
pub fn id(&self) -> CheckpointId {
|
|
self.id
|
|
}
|
|
}
|
|
|
|
impl<'a> Drop for CheckpointDataRef<'a> {
|
|
fn drop(&mut self) {
|
|
let _ = self.manager.release_checkpoint(self.id);
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_checkpoint_policy_every() {
|
|
let policy = CheckpointPolicy::Every(4);
|
|
assert!(policy.should_checkpoint(0, 32));
|
|
assert!(!policy.should_checkpoint(1, 32));
|
|
assert!(!policy.should_checkpoint(2, 32));
|
|
assert!(!policy.should_checkpoint(3, 32));
|
|
assert!(policy.should_checkpoint(4, 32));
|
|
}
|
|
|
|
#[test]
|
|
fn test_checkpoint_policy_selective() {
|
|
let policy = CheckpointPolicy::SelectiveLayers(vec![0, 8, 16, 24]);
|
|
assert!(policy.should_checkpoint(0, 32));
|
|
assert!(!policy.should_checkpoint(1, 32));
|
|
assert!(policy.should_checkpoint(8, 32));
|
|
assert!(!policy.should_checkpoint(7, 32));
|
|
}
|
|
|
|
#[test]
|
|
fn test_checkpoint_policy_none() {
|
|
let policy = CheckpointPolicy::None;
|
|
assert!(!policy.should_checkpoint(0, 32));
|
|
assert!(!policy.should_checkpoint(16, 32));
|
|
}
|
|
|
|
#[test]
|
|
fn test_checkpoint_config_default() {
|
|
let config = CheckpointConfig::default();
|
|
assert_eq!(config.max_checkpoints, 64);
|
|
assert!(config.async_recompute);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_checkpoint_manager_creation() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = MetalCheckpointManager::default_config(device);
|
|
|
|
assert_eq!(manager.checkpoint_count(), 0);
|
|
assert_eq!(manager.memory_usage(), 0);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_checkpoint_save_and_get() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = MetalCheckpointManager::default_config(device.clone());
|
|
|
|
// Create some test buffers
|
|
let buffer = MetalBuffer::<f32>::zeros(&device, 1024).unwrap();
|
|
|
|
// Save checkpoint
|
|
let id = manager.save_checkpoint(0, vec![buffer], 0..4).unwrap();
|
|
assert_eq!(manager.checkpoint_count(), 1);
|
|
|
|
// Get checkpoint
|
|
let cp_ref = manager.get_checkpoint(id);
|
|
assert!(cp_ref.is_some());
|
|
|
|
// Get by layer
|
|
let cp_ref2 = manager.get_checkpoint_for_layer(0);
|
|
assert!(cp_ref2.is_some());
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_checkpoint_eviction() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let config = CheckpointConfig {
|
|
max_checkpoints: 2,
|
|
..Default::default()
|
|
};
|
|
let manager = MetalCheckpointManager::new(device.clone(), config);
|
|
|
|
// Create checkpoints
|
|
for i in 0..3 {
|
|
let buffer = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
|
|
manager
|
|
.save_checkpoint(i, vec![buffer], i..i + 4)
|
|
.unwrap();
|
|
}
|
|
|
|
// Should have evicted one
|
|
assert_eq!(manager.checkpoint_count(), 2);
|
|
|
|
let stats = manager.stats();
|
|
assert_eq!(stats.total_evicted, 1);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_checkpoint_pinning() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let config = CheckpointConfig {
|
|
max_checkpoints: 2,
|
|
..Default::default()
|
|
};
|
|
let manager = MetalCheckpointManager::new(device.clone(), config);
|
|
|
|
// Create first checkpoint and pin it
|
|
let buffer1 = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
|
|
let id1 = manager.save_checkpoint(0, vec![buffer1], 0..4).unwrap();
|
|
manager.pin_checkpoint(id1);
|
|
|
|
// Create more checkpoints to trigger eviction
|
|
let buffer2 = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
|
|
manager.save_checkpoint(1, vec![buffer2], 4..8).unwrap();
|
|
|
|
let buffer3 = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
|
|
manager.save_checkpoint(2, vec![buffer3], 8..12).unwrap();
|
|
|
|
// Pinned checkpoint should still exist
|
|
assert!(manager.get_checkpoint(id1).is_some());
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_find_nearest_checkpoint() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = MetalCheckpointManager::default_config(device.clone());
|
|
|
|
// Create checkpoints at layers 0, 4, 8
|
|
for i in [0usize, 4, 8] {
|
|
let buffer = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
|
|
manager.save_checkpoint(i, vec![buffer], i..i + 4).unwrap();
|
|
}
|
|
|
|
// Find nearest checkpoint for layer 6 (should be layer 4)
|
|
let (_, layer) = manager.find_nearest_checkpoint(6).unwrap();
|
|
assert_eq!(layer, 4);
|
|
|
|
// Find nearest for layer 10 (should be layer 8)
|
|
let (_, layer) = manager.find_nearest_checkpoint(10).unwrap();
|
|
assert_eq!(layer, 8);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(target_os = "macos")]
|
|
fn test_checkpoint_clear() {
|
|
if !MetalDevice::is_available() {
|
|
return;
|
|
}
|
|
|
|
let device = Arc::new(MetalDevice::system_default().unwrap());
|
|
let manager = MetalCheckpointManager::default_config(device.clone());
|
|
|
|
// Create some checkpoints
|
|
for i in 0..5 {
|
|
let buffer = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
|
|
manager.save_checkpoint(i, vec![buffer], i..i + 4).unwrap();
|
|
}
|
|
|
|
assert_eq!(manager.checkpoint_count(), 5);
|
|
|
|
manager.clear();
|
|
assert_eq!(manager.checkpoint_count(), 0);
|
|
assert_eq!(manager.memory_usage(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_estimate_memory_saved() {
|
|
// Without checkpointing: 32 layers * 1MB each = 32MB
|
|
// With checkpointing every 4: 8 checkpoints
|
|
// Saved: 32MB - 8MB = 24MB
|
|
|
|
// This is a simple estimate test
|
|
let saved = 32 * 1024 * 1024 - 8 * 1024 * 1024;
|
|
assert_eq!(saved, 24 * 1024 * 1024);
|
|
}
|
|
}
|