Struct GradientCheckpointManager 

pub struct GradientCheckpointManager<T: GpuDataType> { /* private fields */ }

Gradient checkpoint manager (GPU feature required)

Implementations

impl<T: GpuDataType> GradientCheckpointManager<T>

pub fn new( gpu_context: Arc<GpuContext>, memory_budget: u64, policy: RecomputationPolicy, ) -> Self

Create a new gradient checkpoint manager

pub fn checkpoint_activation( &self, layer_id: usize, activation: &GpuBuffer<T>, recomputation_cost: u32, ) -> Result<()>

Save an activation checkpoint

pub fn get_checkpoint(&self, layer_id: usize) -> Result<Option<GpuBuffer<T>>>

Retrieve a checkpointed activation (removes it from the checkpoint manager)

pub fn has_checkpoint(&self, layer_id: usize) -> bool

Check if a layer has a checkpoint

pub fn remove_checkpoint(&self, layer_id: usize) -> Result<()>

Remove a checkpoint and free its memory

pub fn clear(&self) -> Result<()>

Clear all checkpoints

pub fn memory_usage(&self) -> u64

Get current memory usage in bytes

pub fn memory_budget(&self) -> u64

Get memory budget in bytes

pub fn num_checkpoints(&self) -> usize

Get number of active checkpoints

pub fn get_statistics(&self) -> CheckpointStatistics

Get checkpoint statistics