osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
1304 lines
43 KiB
Rust
1304 lines
43 KiB
Rust
//! Activation Checkpointing for Memory Optimization
|
||
//!
|
||
//! This module implements gradient checkpointing (also known as activation
|
||
//! checkpointing or rematerialization) to reduce memory usage during training.
|
||
//! Instead of storing all activations for the backward pass, we selectively
|
||
//! store checkpoints and recompute activations when needed.
|
||
//!
|
||
//! # Memory Savings
|
||
//! - Without checkpointing: O(n) memory for n layers
|
||
//! - With sqrt checkpointing: O(sqrt(n)) memory
|
||
//! - With selective checkpointing: O(k) memory for k checkpoints
|
||
//!
|
||
//! # Example
|
||
//! ```rust,ignore
|
||
//! use rtx_distributed::activation_checkpointing::{
|
||
//! CheckpointedFunction, CheckpointPolicy, CheckpointManager
|
||
//! };
|
||
//!
|
||
//! let policy = CheckpointPolicy::sqrt(); // Checkpoint every sqrt(n) layers
|
||
//! let manager = CheckpointManager::new(policy);
|
||
//!
|
||
//! // Wrap forward function
|
||
//! let output = manager.checkpoint(|| {
|
||
//! model.forward(&input)
|
||
//! });
|
||
//! ```
|
||
|
||
use crate::error::{DistributedError, Result};
|
||
use parking_lot::{Mutex, RwLock};
|
||
use serde::{Deserialize, Serialize};
|
||
use std::any::Any;
|
||
use std::collections::HashMap;
|
||
use std::sync::Arc;
|
||
|
||
// =============================================================================
|
||
// Checkpoint Policy
|
||
// =============================================================================
|
||
|
||
/// Checkpointing policy determines which layers to checkpoint
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub enum CheckpointPolicy {
|
||
/// No checkpointing (store all activations)
|
||
None,
|
||
/// Checkpoint every N layers
|
||
Every { interval: usize },
|
||
/// Checkpoint at sqrt(n) intervals (optimal for memory/compute tradeoff)
|
||
Sqrt,
|
||
/// Checkpoint specific layers by name or index
|
||
Selective { layers: Vec<String> },
|
||
/// Checkpoint all layers (maximum memory savings, highest compute cost)
|
||
All,
|
||
/// Adaptive based on available memory.
|
||
///
|
||
/// Uses a tiered heuristic based on `target_memory_mb`:
|
||
/// - `>4096 MB` (generous): checkpoint every `ceil(sqrt(n))` layers
|
||
/// - `1024–4096 MB` (moderate): checkpoint every other layer
|
||
/// - `<1024 MB` (tight): checkpoint every layer
|
||
Adaptive { target_memory_mb: usize },
|
||
/// Custom policy with a function
|
||
Custom { name: String },
|
||
/// Checkpoint only attention-family layers (identified by name matching).
|
||
///
|
||
/// FFN/MLP layers are intentionally skipped — they are cheaper to recompute
|
||
/// from a saved input than to store, yielding ~40% memory savings with
|
||
/// ~60% less recompute overhead compared to `CheckpointPolicy::All`.
|
||
AttentionSelective {
|
||
/// Substrings that identify attention layers in their name.
|
||
/// A layer is checkpointed when its name contains any of these patterns.
|
||
attention_patterns: Vec<String>,
|
||
},
|
||
}
|
||
|
||
impl CheckpointPolicy {
|
||
/// Create policy with no checkpointing
|
||
pub fn none() -> Self {
|
||
CheckpointPolicy::None
|
||
}
|
||
|
||
/// Checkpoint every N layers
|
||
pub fn every(n: usize) -> Self {
|
||
CheckpointPolicy::Every { interval: n }
|
||
}
|
||
|
||
/// Optimal sqrt(n) checkpointing
|
||
pub fn sqrt() -> Self {
|
||
CheckpointPolicy::Sqrt
|
||
}
|
||
|
||
/// Checkpoint specific layers
|
||
pub fn selective(layers: Vec<String>) -> Self {
|
||
CheckpointPolicy::Selective { layers }
|
||
}
|
||
|
||
/// Checkpoint all layers (maximum memory savings)
|
||
pub fn all() -> Self {
|
||
CheckpointPolicy::All
|
||
}
|
||
|
||
/// Adaptive based on memory target
|
||
pub fn adaptive(target_memory_mb: usize) -> Self {
|
||
CheckpointPolicy::Adaptive { target_memory_mb }
|
||
}
|
||
|
||
/// Checkpoint only attention-family layers.
|
||
///
|
||
/// Uses the standard set of attention-related name patterns:
|
||
/// `"attn"`, `"attention"`, `"self_attn"`, `"cross_attn"`, `"mha"`.
|
||
/// FFN/MLP blocks are not checkpointed, cutting recompute cost by ~60%
|
||
/// relative to `CheckpointPolicy::All` while still saving ~40% of
|
||
/// transformer activation memory.
|
||
pub fn attention_selective() -> Self {
|
||
CheckpointPolicy::AttentionSelective {
|
||
attention_patterns: vec![
|
||
"attn".to_string(),
|
||
"attention".to_string(),
|
||
"self_attn".to_string(),
|
||
"cross_attn".to_string(),
|
||
"mha".to_string(),
|
||
],
|
||
}
|
||
}
|
||
|
||
/// Check if a layer should be checkpointed
|
||
pub fn should_checkpoint(
|
||
&self,
|
||
layer_index: usize,
|
||
layer_name: &str,
|
||
total_layers: usize,
|
||
) -> bool {
|
||
match self {
|
||
CheckpointPolicy::None => false,
|
||
CheckpointPolicy::Every { interval } => layer_index % interval == 0,
|
||
CheckpointPolicy::Sqrt => {
|
||
let sqrt_n = (total_layers as f64).sqrt().ceil() as usize;
|
||
layer_index % sqrt_n == 0
|
||
}
|
||
CheckpointPolicy::Selective { layers } => {
|
||
layers.contains(&layer_name.to_string())
|
||
|| layers.contains(&layer_index.to_string())
|
||
}
|
||
CheckpointPolicy::All => true,
|
||
CheckpointPolicy::Adaptive { target_memory_mb } => {
|
||
// Tiered heuristic based on memory target.
|
||
// We cannot query live GPU memory here without a device handle, so
|
||
// we use the target as a proxy for how aggressively to checkpoint:
|
||
// - Generous target (>4096 MB): checkpoint every ceil(sqrt(n)) layers
|
||
// - Moderate target (1024–4096 MB): checkpoint every other layer
|
||
// - Tight target (<1024 MB): checkpoint every layer
|
||
let tier = match *target_memory_mb {
|
||
t if t > 4096 => (total_layers as f64).sqrt().ceil() as usize,
|
||
t if t > 1024 => 2,
|
||
_ => 1,
|
||
};
|
||
// tier is always >= 1 (sqrt of at least 1 layer rounds up to 1)
|
||
layer_index % tier == 0
|
||
}
|
||
CheckpointPolicy::Custom { .. } => false, // Custom logic handled elsewhere
|
||
CheckpointPolicy::AttentionSelective { attention_patterns } => {
|
||
// Checkpoint the layer if its name contains any attention pattern.
|
||
// Layer index and total_layers are intentionally not used: the
|
||
// decision is purely structural (attention vs. non-attention).
|
||
attention_patterns
|
||
.iter()
|
||
.any(|pat| layer_name.contains(pat.as_str()))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Estimate memory savings factor
|
||
pub fn memory_savings_factor(&self, total_layers: usize) -> f64 {
|
||
match self {
|
||
CheckpointPolicy::None => 1.0,
|
||
CheckpointPolicy::Every { interval } => *interval as f64,
|
||
CheckpointPolicy::Sqrt => (total_layers as f64).sqrt(),
|
||
CheckpointPolicy::Selective { layers } => {
|
||
total_layers as f64 / (layers.len().max(1) as f64)
|
||
}
|
||
CheckpointPolicy::All => total_layers as f64,
|
||
CheckpointPolicy::Adaptive { .. } => 2.0, // Rough estimate
|
||
CheckpointPolicy::Custom { .. } => 1.5, // Rough estimate
|
||
CheckpointPolicy::AttentionSelective { .. } => {
|
||
// Attention blocks constitute ~40% of transformer activation memory.
|
||
// Checkpointing only them therefore saves ~40% of total activations,
|
||
// expressed here as the number of layers' worth of memory freed.
|
||
total_layers as f64 * 0.4
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// Activation Storage
|
||
// =============================================================================
|
||
|
||
/// Stored activation data
|
||
#[derive(Debug)]
|
||
pub struct StoredActivation {
|
||
/// Layer name or identifier
|
||
pub layer_id: String,
|
||
/// Activation tensor data (serialized)
|
||
pub data: Vec<u8>,
|
||
/// Original shape
|
||
pub shape: Vec<usize>,
|
||
/// Data type (e.g., "f32", "f16")
|
||
pub dtype: String,
|
||
/// Memory size in bytes
|
||
pub size_bytes: usize,
|
||
/// Whether this is a checkpoint (will be recomputed if false)
|
||
pub is_checkpoint: bool,
|
||
/// Timestamp when stored
|
||
pub stored_at: std::time::Instant,
|
||
}
|
||
|
||
impl StoredActivation {
|
||
/// Create new stored activation
|
||
pub fn new(layer_id: String, data: Vec<u8>, shape: Vec<usize>, dtype: String) -> Self {
|
||
let size_bytes = data.len();
|
||
Self {
|
||
layer_id,
|
||
data,
|
||
shape,
|
||
dtype,
|
||
size_bytes,
|
||
is_checkpoint: false,
|
||
stored_at: std::time::Instant::now(),
|
||
}
|
||
}
|
||
|
||
/// Mark as checkpoint
|
||
pub fn as_checkpoint(mut self) -> Self {
|
||
self.is_checkpoint = true;
|
||
self
|
||
}
|
||
}
|
||
|
||
/// Activation storage manager
|
||
#[derive(Debug)]
|
||
pub struct ActivationStorage {
|
||
/// Stored activations by layer ID
|
||
activations: HashMap<String, StoredActivation>,
|
||
/// Total memory used
|
||
total_memory: usize,
|
||
/// Memory limit
|
||
memory_limit: usize,
|
||
/// Statistics
|
||
stats: ActivationStats,
|
||
}
|
||
|
||
/// Activation storage statistics
|
||
#[derive(Debug, Default, Clone)]
|
||
pub struct ActivationStats {
|
||
/// Number of activations stored
|
||
pub stored_count: usize,
|
||
/// Number of activations recomputed
|
||
pub recomputed_count: usize,
|
||
/// Total bytes stored
|
||
pub bytes_stored: usize,
|
||
/// Bytes saved by checkpointing
|
||
pub bytes_saved: usize,
|
||
/// Peak memory usage
|
||
pub peak_memory: usize,
|
||
}
|
||
|
||
impl ActivationStorage {
|
||
/// Create new activation storage
|
||
pub fn new(memory_limit: usize) -> Self {
|
||
Self {
|
||
activations: HashMap::new(),
|
||
total_memory: 0,
|
||
memory_limit,
|
||
stats: ActivationStats::default(),
|
||
}
|
||
}
|
||
|
||
/// Store an activation
|
||
pub fn store(&mut self, activation: StoredActivation) -> Result<()> {
|
||
let size = activation.size_bytes;
|
||
|
||
// Check memory limit
|
||
if self.total_memory + size > self.memory_limit {
|
||
// Try to evict non-checkpoint activations
|
||
self.evict_until_free(size)?;
|
||
}
|
||
|
||
let layer_id = activation.layer_id.clone();
|
||
self.total_memory += size;
|
||
self.stats.stored_count += 1;
|
||
self.stats.bytes_stored += size;
|
||
self.stats.peak_memory = self.stats.peak_memory.max(self.total_memory);
|
||
|
||
self.activations.insert(layer_id, activation);
|
||
Ok(())
|
||
}
|
||
|
||
/// Get an activation
|
||
pub fn get(&self, layer_id: &str) -> Option<&StoredActivation> {
|
||
self.activations.get(layer_id)
|
||
}
|
||
|
||
/// Remove an activation
|
||
pub fn remove(&mut self, layer_id: &str) -> Option<StoredActivation> {
|
||
if let Some(activation) = self.activations.remove(layer_id) {
|
||
self.total_memory -= activation.size_bytes;
|
||
Some(activation)
|
||
} else {
|
||
None
|
||
}
|
||
}
|
||
|
||
/// Clear all activations
|
||
pub fn clear(&mut self) {
|
||
self.activations.clear();
|
||
self.total_memory = 0;
|
||
}
|
||
|
||
/// Evict non-checkpoint activations until enough space is free
|
||
fn evict_until_free(&mut self, needed: usize) -> Result<()> {
|
||
let mut to_remove = Vec::new();
|
||
|
||
for (id, activation) in &self.activations {
|
||
if !activation.is_checkpoint {
|
||
to_remove.push(id.clone());
|
||
self.stats.bytes_saved += activation.size_bytes;
|
||
}
|
||
}
|
||
|
||
for id in to_remove {
|
||
self.remove(&id);
|
||
if self.memory_limit - self.total_memory >= needed {
|
||
break;
|
||
}
|
||
}
|
||
|
||
if self.memory_limit - self.total_memory < needed {
|
||
return Err(DistributedError::configuration(format!(
|
||
"Cannot free enough memory: need {} bytes, have {} available",
|
||
needed,
|
||
self.memory_limit - self.total_memory
|
||
)));
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Get current memory usage
|
||
pub fn memory_usage(&self) -> usize {
|
||
self.total_memory
|
||
}
|
||
|
||
/// Get statistics
|
||
pub fn stats(&self) -> &ActivationStats {
|
||
&self.stats
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// Checkpoint Context
|
||
// =============================================================================
|
||
|
||
/// Context for checkpointed computation
|
||
#[derive(Debug)]
|
||
pub struct CheckpointContext {
|
||
/// Forward function inputs (saved for recomputation)
|
||
pub inputs: Vec<Box<dyn Any + Send + Sync>>,
|
||
/// Saved tensors for backward
|
||
pub saved_tensors: Vec<Box<dyn Any + Send + Sync>>,
|
||
/// Whether we need to recompute on backward
|
||
pub needs_recompute: bool,
|
||
/// Layer identifier
|
||
pub layer_id: String,
|
||
}
|
||
|
||
impl CheckpointContext {
|
||
/// Create new checkpoint context
|
||
pub fn new(layer_id: String) -> Self {
|
||
Self {
|
||
inputs: Vec::new(),
|
||
saved_tensors: Vec::new(),
|
||
needs_recompute: false,
|
||
layer_id,
|
||
}
|
||
}
|
||
|
||
/// Save input for potential recomputation
|
||
pub fn save_input<T: Any + Send + Sync + 'static>(&mut self, input: T) {
|
||
self.inputs.push(Box::new(input));
|
||
}
|
||
|
||
/// Save tensor for backward pass
|
||
pub fn save_for_backward<T: Any + Send + Sync + 'static>(&mut self, tensor: T) {
|
||
self.saved_tensors.push(Box::new(tensor));
|
||
}
|
||
|
||
/// Mark that this context needs recomputation
|
||
pub fn mark_recompute(&mut self) {
|
||
self.needs_recompute = true;
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// Checkpointed Segment
|
||
// =============================================================================
|
||
|
||
/// A segment of the model that can be checkpointed
|
||
#[derive(Debug)]
|
||
pub struct CheckpointedSegment {
|
||
/// Segment identifier
|
||
pub id: String,
|
||
/// Layers in this segment
|
||
pub layers: Vec<String>,
|
||
/// Input tensors (saved for recomputation)
|
||
inputs: RwLock<Vec<Vec<u8>>>,
|
||
/// Output tensors
|
||
outputs: RwLock<Vec<Vec<u8>>>,
|
||
/// Whether outputs are currently stored
|
||
outputs_stored: RwLock<bool>,
|
||
}
|
||
|
||
impl CheckpointedSegment {
|
||
/// Create new checkpointed segment
|
||
pub fn new(id: String, layers: Vec<String>) -> Self {
|
||
Self {
|
||
id,
|
||
layers,
|
||
inputs: RwLock::new(Vec::new()),
|
||
outputs: RwLock::new(Vec::new()),
|
||
outputs_stored: RwLock::new(false),
|
||
}
|
||
}
|
||
|
||
/// Save inputs for potential recomputation
|
||
pub fn save_inputs(&self, inputs: Vec<Vec<u8>>) {
|
||
*self.inputs.write() = inputs;
|
||
}
|
||
|
||
/// Save outputs (only done for checkpoints)
|
||
pub fn save_outputs(&self, outputs: Vec<Vec<u8>>) {
|
||
*self.outputs.write() = outputs;
|
||
*self.outputs_stored.write() = true;
|
||
}
|
||
|
||
/// Get saved inputs
|
||
pub fn get_inputs(&self) -> Vec<Vec<u8>> {
|
||
self.inputs.read().clone()
|
||
}
|
||
|
||
/// Get saved outputs
|
||
pub fn get_outputs(&self) -> Option<Vec<Vec<u8>>> {
|
||
if *self.outputs_stored.read() {
|
||
Some(self.outputs.read().clone())
|
||
} else {
|
||
None
|
||
}
|
||
}
|
||
|
||
/// Clear saved outputs (after recomputation)
|
||
pub fn clear_outputs(&self) {
|
||
self.outputs.write().clear();
|
||
*self.outputs_stored.write() = false;
|
||
}
|
||
|
||
/// Check if outputs are available
|
||
pub fn has_outputs(&self) -> bool {
|
||
*self.outputs_stored.read()
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// Checkpoint Manager
|
||
// =============================================================================
|
||
|
||
/// Manager for activation checkpointing
|
||
#[derive(Debug)]
|
||
pub struct ActivationCheckpointManager {
|
||
/// Checkpointing policy
|
||
policy: CheckpointPolicy,
|
||
/// Activation storage
|
||
storage: Mutex<ActivationStorage>,
|
||
/// Registered segments
|
||
segments: RwLock<HashMap<String, Arc<CheckpointedSegment>>>,
|
||
/// Current layer index (for policy decisions)
|
||
current_layer: RwLock<usize>,
|
||
/// Total number of layers
|
||
total_layers: RwLock<usize>,
|
||
/// Recomputation count
|
||
recompute_count: RwLock<usize>,
|
||
/// Whether checkpointing is enabled
|
||
enabled: RwLock<bool>,
|
||
}
|
||
|
||
impl ActivationCheckpointManager {
|
||
/// Create new checkpoint manager
|
||
pub fn new(policy: CheckpointPolicy) -> Self {
|
||
Self {
|
||
policy,
|
||
storage: Mutex::new(ActivationStorage::new(8 * 1024 * 1024 * 1024)), // 8GB default
|
||
segments: RwLock::new(HashMap::new()),
|
||
current_layer: RwLock::new(0),
|
||
total_layers: RwLock::new(0),
|
||
recompute_count: RwLock::new(0),
|
||
enabled: RwLock::new(true),
|
||
}
|
||
}
|
||
|
||
/// Create with memory limit
|
||
pub fn with_memory_limit(policy: CheckpointPolicy, memory_limit: usize) -> Self {
|
||
Self {
|
||
policy,
|
||
storage: Mutex::new(ActivationStorage::new(memory_limit)),
|
||
segments: RwLock::new(HashMap::new()),
|
||
current_layer: RwLock::new(0),
|
||
total_layers: RwLock::new(0),
|
||
recompute_count: RwLock::new(0),
|
||
enabled: RwLock::new(true),
|
||
}
|
||
}
|
||
|
||
/// Set total number of layers
|
||
pub fn set_total_layers(&self, n: usize) {
|
||
*self.total_layers.write() = n;
|
||
}
|
||
|
||
/// Enable or disable checkpointing
|
||
pub fn set_enabled(&self, enabled: bool) {
|
||
*self.enabled.write() = enabled;
|
||
}
|
||
|
||
/// Check if checkpointing is enabled
|
||
pub fn is_enabled(&self) -> bool {
|
||
*self.enabled.read()
|
||
}
|
||
|
||
/// Register a segment for checkpointing
|
||
pub fn register_segment(&self, segment: CheckpointedSegment) -> Arc<CheckpointedSegment> {
|
||
let segment = Arc::new(segment);
|
||
self.segments
|
||
.write()
|
||
.insert(segment.id.clone(), segment.clone());
|
||
segment
|
||
}
|
||
|
||
/// Get a registered segment
|
||
pub fn get_segment(&self, id: &str) -> Option<Arc<CheckpointedSegment>> {
|
||
self.segments.read().get(id).cloned()
|
||
}
|
||
|
||
/// Check if current layer should be checkpointed
|
||
pub fn should_checkpoint_current(&self, layer_name: &str) -> bool {
|
||
if !*self.enabled.read() {
|
||
return false;
|
||
}
|
||
|
||
let layer_index = *self.current_layer.read();
|
||
let total = *self.total_layers.read();
|
||
|
||
self.policy
|
||
.should_checkpoint(layer_index, layer_name, total)
|
||
}
|
||
|
||
/// Advance to next layer
|
||
pub fn next_layer(&self) {
|
||
*self.current_layer.write() += 1;
|
||
}
|
||
|
||
/// Reset layer counter
|
||
pub fn reset_layers(&self) {
|
||
*self.current_layer.write() = 0;
|
||
}
|
||
|
||
/// Store activation
|
||
pub fn store_activation(
|
||
&self,
|
||
layer_id: String,
|
||
data: Vec<u8>,
|
||
shape: Vec<usize>,
|
||
dtype: String,
|
||
) -> Result<()> {
|
||
let activation = StoredActivation::new(layer_id, data, shape, dtype);
|
||
self.storage.lock().store(activation)
|
||
}
|
||
|
||
/// Store checkpoint activation
|
||
pub fn store_checkpoint(
|
||
&self,
|
||
layer_id: String,
|
||
data: Vec<u8>,
|
||
shape: Vec<usize>,
|
||
dtype: String,
|
||
) -> Result<()> {
|
||
let activation = StoredActivation::new(layer_id, data, shape, dtype).as_checkpoint();
|
||
self.storage.lock().store(activation)
|
||
}
|
||
|
||
/// Get stored activation
|
||
pub fn get_activation(&self, layer_id: &str) -> Option<(Vec<u8>, Vec<usize>, String)> {
|
||
self.storage
|
||
.lock()
|
||
.get(layer_id)
|
||
.map(|a| (a.data.clone(), a.shape.clone(), a.dtype.clone()))
|
||
}
|
||
|
||
/// Record a recomputation
|
||
pub fn record_recompute(&self) {
|
||
*self.recompute_count.write() += 1;
|
||
self.storage.lock().stats.recomputed_count += 1;
|
||
}
|
||
|
||
/// Clear all stored activations
|
||
pub fn clear(&self) {
|
||
self.storage.lock().clear();
|
||
self.reset_layers();
|
||
}
|
||
|
||
/// Get memory usage
|
||
pub fn memory_usage(&self) -> usize {
|
||
self.storage.lock().memory_usage()
|
||
}
|
||
|
||
/// Get statistics
|
||
pub fn stats(&self) -> CheckpointStats {
|
||
let storage_stats = self.storage.lock().stats().clone();
|
||
CheckpointStats {
|
||
activations_stored: storage_stats.stored_count,
|
||
activations_recomputed: storage_stats.recomputed_count,
|
||
bytes_stored: storage_stats.bytes_stored,
|
||
bytes_saved: storage_stats.bytes_saved,
|
||
peak_memory: storage_stats.peak_memory,
|
||
recompute_count: *self.recompute_count.read(),
|
||
memory_savings_factor: self.policy.memory_savings_factor(*self.total_layers.read()),
|
||
}
|
||
}
|
||
|
||
/// Get estimated memory savings
|
||
pub fn estimated_memory_savings(&self) -> f64 {
|
||
let total_layers = *self.total_layers.read();
|
||
if total_layers == 0 {
|
||
return 1.0;
|
||
}
|
||
self.policy.memory_savings_factor(total_layers)
|
||
}
|
||
}
|
||
|
||
/// Checkpoint statistics
|
||
#[derive(Debug, Clone)]
|
||
pub struct CheckpointStats {
|
||
/// Number of activations stored
|
||
pub activations_stored: usize,
|
||
/// Number of activations recomputed
|
||
pub activations_recomputed: usize,
|
||
/// Bytes stored
|
||
pub bytes_stored: usize,
|
||
/// Bytes saved by checkpointing
|
||
pub bytes_saved: usize,
|
||
/// Peak memory usage
|
||
pub peak_memory: usize,
|
||
/// Number of recomputations
|
||
pub recompute_count: usize,
|
||
/// Memory savings factor
|
||
pub memory_savings_factor: f64,
|
||
}
|
||
|
||
// =============================================================================
|
||
// Checkpoint Function Wrapper
|
||
// =============================================================================
|
||
|
||
/// Trait for functions that can be checkpointed
|
||
pub trait CheckpointedFunction {
|
||
/// Output type
|
||
type Output;
|
||
|
||
/// Run the function with checkpointing
|
||
fn run(&self, ctx: &mut CheckpointContext) -> Self::Output;
|
||
|
||
/// Recompute the function (for backward pass)
|
||
fn recompute(&self, ctx: &CheckpointContext) -> Self::Output;
|
||
}
|
||
|
||
/// Wrapper to checkpoint a closure
|
||
pub struct CheckpointedClosure<F, O>
|
||
where
|
||
F: Fn() -> O,
|
||
{
|
||
/// The function to checkpoint
|
||
func: F,
|
||
/// Layer ID for this checkpoint
|
||
layer_id: String,
|
||
}
|
||
|
||
impl<F, O> CheckpointedClosure<F, O>
|
||
where
|
||
F: Fn() -> O,
|
||
{
|
||
/// Create new checkpointed closure
|
||
pub fn new(layer_id: String, func: F) -> Self {
|
||
Self { func, layer_id }
|
||
}
|
||
|
||
/// Execute the closure
|
||
pub fn execute(&self) -> O {
|
||
(self.func)()
|
||
}
|
||
|
||
/// Get layer ID
|
||
pub fn layer_id(&self) -> &str {
|
||
&self.layer_id
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// Shared Checkpoint Manager
|
||
// =============================================================================
|
||
|
||
/// Thread-safe shared checkpoint manager
|
||
pub type SharedActivationCheckpointManager = Arc<ActivationCheckpointManager>;
|
||
|
||
/// Create a shared checkpoint manager
|
||
pub fn shared_activation_checkpoint_manager(
|
||
policy: CheckpointPolicy,
|
||
) -> SharedActivationCheckpointManager {
|
||
Arc::new(ActivationCheckpointManager::new(policy))
|
||
}
|
||
|
||
/// Create a shared checkpoint manager with memory limit
|
||
pub fn shared_activation_checkpoint_manager_with_limit(
|
||
policy: CheckpointPolicy,
|
||
memory_limit: usize,
|
||
) -> SharedActivationCheckpointManager {
|
||
Arc::new(ActivationCheckpointManager::with_memory_limit(
|
||
policy,
|
||
memory_limit,
|
||
))
|
||
}
|
||
|
||
// =============================================================================
|
||
// Utility Functions
|
||
// =============================================================================
|
||
|
||
/// Calculate optimal checkpoint interval for given memory budget
|
||
pub fn optimal_checkpoint_interval(
|
||
num_layers: usize,
|
||
activation_size_per_layer: usize,
|
||
memory_budget: usize,
|
||
) -> usize {
|
||
// With checkpointing every k layers:
|
||
// Memory = k * activation_size (for storing between checkpoints)
|
||
// Recomputation cost = O(k) per backward step
|
||
|
||
// Optimal is sqrt(n) checkpoints for balanced memory/compute
|
||
let sqrt_n = (num_layers as f64).sqrt().ceil() as usize;
|
||
|
||
// Adjust based on memory budget
|
||
let max_interval = memory_budget / activation_size_per_layer.max(1);
|
||
|
||
sqrt_n.min(max_interval).max(1)
|
||
}
|
||
|
||
/// Estimate memory savings for a given configuration
|
||
pub fn estimate_memory_savings(
|
||
num_layers: usize,
|
||
activation_size_per_layer: usize,
|
||
checkpoint_interval: usize,
|
||
) -> (usize, usize) {
|
||
// Without checkpointing
|
||
let without = num_layers * activation_size_per_layer;
|
||
|
||
// With checkpointing: store only checkpoint interval activations
|
||
let num_checkpoints = num_layers / checkpoint_interval.max(1);
|
||
let with = checkpoint_interval * activation_size_per_layer
|
||
+ num_checkpoints * activation_size_per_layer;
|
||
|
||
(without, with)
|
||
}
|
||
|
||
// =============================================================================
|
||
// Memory-Aware Checkpointer
|
||
// =============================================================================
|
||
|
||
/// Applies adaptive checkpointing decisions based on live memory tracking.
|
||
///
|
||
/// `MemoryAwareCheckpointer` wraps a [`CheckpointPolicy`] and augments it with
|
||
/// a running estimate of activation memory consumed during a forward pass.
|
||
/// When that estimate exceeds `target_bytes`, every layer is checkpointed
|
||
/// regardless of the underlying policy, preventing OOM conditions.
|
||
///
|
||
/// The internal counter uses `AtomicUsize` so the checkpointer can be shared
|
||
/// across threads without locking overhead.
|
||
///
|
||
/// # Example
|
||
/// ```rust,ignore
|
||
/// use rtx_distributed::activation_checkpointing::{
|
||
/// CheckpointPolicy, MemoryAwareCheckpointer
|
||
/// };
|
||
///
|
||
/// let checkpointer = MemoryAwareCheckpointer::new(
|
||
/// CheckpointPolicy::attention_selective(),
|
||
/// 24, // 24 transformer layers
|
||
/// 4096, // 4 GB target
|
||
/// );
|
||
///
|
||
/// // During forward pass:
|
||
/// checkpointer.record_activation(64 * 1024 * 1024); // 64 MB for this layer
|
||
/// let do_ckpt = checkpointer.should_checkpoint(3, "self_attn_0");
|
||
/// ```
|
||
pub struct MemoryAwareCheckpointer {
|
||
/// The base policy used when memory is below the target.
|
||
pub policy: CheckpointPolicy,
|
||
/// Total number of layers in the model (used by the base policy).
|
||
pub total_layers: usize,
|
||
/// Estimated bytes of activation memory currently live.
|
||
/// Incremented by [`record_activation`] and decremented by
|
||
/// [`free_activation`]. Uses `Relaxed` ordering — exact byte counts
|
||
/// are not required for correctness, only approximate pressure detection.
|
||
estimated_memory_bytes: std::sync::atomic::AtomicUsize,
|
||
/// Memory target in bytes. When `estimated_memory_bytes` exceeds this,
|
||
/// the checkpointer forces checkpointing on every layer.
|
||
pub target_bytes: usize,
|
||
}
|
||
|
||
impl std::fmt::Debug for MemoryAwareCheckpointer {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
f.debug_struct("MemoryAwareCheckpointer")
|
||
.field("policy", &self.policy)
|
||
.field("total_layers", &self.total_layers)
|
||
.field(
|
||
"estimated_memory_bytes",
|
||
&self
|
||
.estimated_memory_bytes
|
||
.load(std::sync::atomic::Ordering::Relaxed),
|
||
)
|
||
.field("target_bytes", &self.target_bytes)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
impl MemoryAwareCheckpointer {
|
||
/// Create a new `MemoryAwareCheckpointer`.
|
||
///
|
||
/// # Parameters
|
||
/// - `policy`: Base checkpointing policy applied when memory is within budget.
|
||
/// - `total_layers`: Number of layers in the model.
|
||
/// - `target_mb`: Memory budget in mebibytes (MiB = 1024 * 1024 bytes).
|
||
pub fn new(policy: CheckpointPolicy, total_layers: usize, target_mb: usize) -> Self {
|
||
Self {
|
||
policy,
|
||
total_layers,
|
||
estimated_memory_bytes: std::sync::atomic::AtomicUsize::new(0),
|
||
target_bytes: target_mb * 1024 * 1024,
|
||
}
|
||
}
|
||
|
||
/// Record that a layer consumed `bytes` of activation memory.
|
||
///
|
||
/// Call this after each layer's forward pass to keep the pressure
|
||
/// estimate current.
|
||
pub fn record_activation(&self, bytes: usize) {
|
||
self.estimated_memory_bytes
|
||
.fetch_add(bytes, std::sync::atomic::Ordering::Relaxed);
|
||
}
|
||
|
||
/// Free tracked memory when activations are released (e.g. after a
|
||
/// segment's backward pass completes).
|
||
///
|
||
/// Saturates at zero; cannot underflow.
|
||
pub fn free_activation(&self, bytes: usize) {
|
||
// Clamp the subtraction so we never wrap on underflow.
|
||
let current = self
|
||
.estimated_memory_bytes
|
||
.load(std::sync::atomic::Ordering::Relaxed);
|
||
let to_sub = bytes.min(current);
|
||
self.estimated_memory_bytes
|
||
.fetch_sub(to_sub, std::sync::atomic::Ordering::Relaxed);
|
||
}
|
||
|
||
/// Decide whether to checkpoint this layer.
|
||
///
|
||
/// If `estimated_memory_bytes > target_bytes`, returns `true` unconditionally
|
||
/// (memory-pressure override). Otherwise delegates to `self.policy`.
|
||
pub fn should_checkpoint(&self, layer_index: usize, layer_name: &str) -> bool {
|
||
let current = self
|
||
.estimated_memory_bytes
|
||
.load(std::sync::atomic::Ordering::Relaxed);
|
||
if current > self.target_bytes {
|
||
// Memory pressure: checkpoint everything to stay under target.
|
||
return true;
|
||
}
|
||
self.policy
|
||
.should_checkpoint(layer_index, layer_name, self.total_layers)
|
||
}
|
||
|
||
/// Current memory pressure as a fraction of the target (0.0 = empty, 1.0 = at target).
|
||
///
|
||
/// Values above 1.0 indicate the memory budget has been exceeded.
|
||
pub fn memory_pressure(&self) -> f32 {
|
||
let current = self
|
||
.estimated_memory_bytes
|
||
.load(std::sync::atomic::Ordering::Relaxed);
|
||
current as f32 / self.target_bytes as f32
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// Tests
|
||
// =============================================================================
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_checkpoint_policy_none() {
|
||
let policy = CheckpointPolicy::none();
|
||
assert!(!policy.should_checkpoint(0, "layer0", 10));
|
||
assert!(!policy.should_checkpoint(5, "layer5", 10));
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_policy_every() {
|
||
let policy = CheckpointPolicy::every(3);
|
||
assert!(policy.should_checkpoint(0, "layer0", 10));
|
||
assert!(!policy.should_checkpoint(1, "layer1", 10));
|
||
assert!(!policy.should_checkpoint(2, "layer2", 10));
|
||
assert!(policy.should_checkpoint(3, "layer3", 10));
|
||
assert!(policy.should_checkpoint(6, "layer6", 10));
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_policy_sqrt() {
|
||
let policy = CheckpointPolicy::sqrt();
|
||
// For 16 layers, sqrt is 4, so checkpoint every 4 layers
|
||
assert!(policy.should_checkpoint(0, "layer0", 16));
|
||
assert!(!policy.should_checkpoint(1, "layer1", 16));
|
||
assert!(policy.should_checkpoint(4, "layer4", 16));
|
||
assert!(policy.should_checkpoint(8, "layer8", 16));
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_policy_selective() {
|
||
let policy = CheckpointPolicy::selective(vec!["attention".to_string(), "5".to_string()]);
|
||
assert!(policy.should_checkpoint(0, "attention", 10));
|
||
assert!(policy.should_checkpoint(5, "mlp", 10));
|
||
assert!(!policy.should_checkpoint(3, "norm", 10));
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_policy_all() {
|
||
let policy = CheckpointPolicy::all();
|
||
for i in 0..10 {
|
||
assert!(policy.should_checkpoint(i, &format!("layer{}", i), 10));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_memory_savings_factor() {
|
||
let none = CheckpointPolicy::none();
|
||
assert_eq!(none.memory_savings_factor(10), 1.0);
|
||
|
||
let every_2 = CheckpointPolicy::every(2);
|
||
assert_eq!(every_2.memory_savings_factor(10), 2.0);
|
||
|
||
let sqrt = CheckpointPolicy::sqrt();
|
||
let factor = sqrt.memory_savings_factor(16);
|
||
assert!((factor - 4.0).abs() < 0.01); // sqrt(16) = 4
|
||
}
|
||
|
||
#[test]
|
||
fn test_activation_storage() {
|
||
let mut storage = ActivationStorage::new(1024 * 1024); // 1MB limit
|
||
|
||
let activation = StoredActivation::new(
|
||
"layer1".to_string(),
|
||
vec![0u8; 1024],
|
||
vec![32, 32],
|
||
"f32".to_string(),
|
||
);
|
||
|
||
assert!(storage.store(activation).is_ok());
|
||
assert_eq!(storage.memory_usage(), 1024);
|
||
|
||
let retrieved = storage.get("layer1");
|
||
assert!(retrieved.is_some());
|
||
assert_eq!(retrieved.unwrap().shape, vec![32, 32]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_activation_storage_eviction() {
|
||
let mut storage = ActivationStorage::new(2048); // 2KB limit
|
||
|
||
// Store non-checkpoint activation
|
||
let act1 = StoredActivation::new(
|
||
"layer1".to_string(),
|
||
vec![0u8; 1024],
|
||
vec![32, 32],
|
||
"f32".to_string(),
|
||
);
|
||
storage.store(act1).unwrap();
|
||
|
||
// Store checkpoint activation
|
||
let act2 = StoredActivation::new(
|
||
"layer2".to_string(),
|
||
vec![0u8; 512],
|
||
vec![16, 32],
|
||
"f32".to_string(),
|
||
)
|
||
.as_checkpoint();
|
||
storage.store(act2).unwrap();
|
||
|
||
// Try to store another large activation - should evict non-checkpoint
|
||
let act3 = StoredActivation::new(
|
||
"layer3".to_string(),
|
||
vec![0u8; 1024],
|
||
vec![32, 32],
|
||
"f32".to_string(),
|
||
);
|
||
storage.store(act3).unwrap();
|
||
|
||
// layer1 should be evicted, layer2 (checkpoint) should remain
|
||
assert!(storage.get("layer1").is_none());
|
||
assert!(storage.get("layer2").is_some());
|
||
assert!(storage.get("layer3").is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_manager() {
|
||
let manager = ActivationCheckpointManager::new(CheckpointPolicy::every(2));
|
||
manager.set_total_layers(10);
|
||
|
||
assert!(manager.is_enabled());
|
||
assert!(manager.should_checkpoint_current("layer0")); // Index 0
|
||
|
||
manager.next_layer();
|
||
assert!(!manager.should_checkpoint_current("layer1")); // Index 1
|
||
|
||
manager.next_layer();
|
||
assert!(manager.should_checkpoint_current("layer2")); // Index 2
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_manager_disable() {
|
||
let manager = ActivationCheckpointManager::new(CheckpointPolicy::all());
|
||
manager.set_total_layers(10);
|
||
|
||
assert!(manager.should_checkpoint_current("layer0"));
|
||
|
||
manager.set_enabled(false);
|
||
assert!(!manager.should_checkpoint_current("layer0"));
|
||
|
||
manager.set_enabled(true);
|
||
assert!(manager.should_checkpoint_current("layer0"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpointed_segment() {
|
||
let segment = CheckpointedSegment::new(
|
||
"segment1".to_string(),
|
||
vec!["layer1".to_string(), "layer2".to_string()],
|
||
);
|
||
|
||
assert!(!segment.has_outputs());
|
||
|
||
segment.save_inputs(vec![vec![1, 2, 3]]);
|
||
segment.save_outputs(vec![vec![4, 5, 6]]);
|
||
|
||
assert!(segment.has_outputs());
|
||
assert_eq!(segment.get_inputs(), vec![vec![1, 2, 3]]);
|
||
assert_eq!(segment.get_outputs(), Some(vec![vec![4, 5, 6]]));
|
||
|
||
segment.clear_outputs();
|
||
assert!(!segment.has_outputs());
|
||
}
|
||
|
||
#[test]
|
||
fn test_optimal_checkpoint_interval() {
|
||
// For 16 layers, optimal is sqrt(16) = 4
|
||
let interval = optimal_checkpoint_interval(16, 1024 * 1024, 100 * 1024 * 1024);
|
||
assert!(interval >= 4);
|
||
|
||
// With very limited memory, should reduce interval
|
||
let interval = optimal_checkpoint_interval(16, 1024 * 1024, 2 * 1024 * 1024);
|
||
assert!(interval <= 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_estimate_memory_savings() {
|
||
let (without, with) = estimate_memory_savings(16, 1024 * 1024, 4);
|
||
|
||
// Without: 16 * 1MB = 16MB
|
||
assert_eq!(without, 16 * 1024 * 1024);
|
||
|
||
// With: 4 * 1MB (interval) + 4 * 1MB (checkpoints) = 8MB
|
||
// Actually: interval activations + checkpoint activations
|
||
assert!(with < without);
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_context() {
|
||
let mut ctx = CheckpointContext::new("layer1".to_string());
|
||
|
||
ctx.save_input(vec![1.0f32, 2.0, 3.0]);
|
||
ctx.save_for_backward(vec![4.0f32, 5.0, 6.0]);
|
||
ctx.mark_recompute();
|
||
|
||
assert!(ctx.needs_recompute);
|
||
assert_eq!(ctx.inputs.len(), 1);
|
||
assert_eq!(ctx.saved_tensors.len(), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_checkpoint_stats() {
|
||
let manager = ActivationCheckpointManager::new(CheckpointPolicy::sqrt());
|
||
manager.set_total_layers(16);
|
||
|
||
manager
|
||
.store_activation(
|
||
"layer0".to_string(),
|
||
vec![0u8; 1024],
|
||
vec![32, 32],
|
||
"f32".to_string(),
|
||
)
|
||
.unwrap();
|
||
|
||
manager.record_recompute();
|
||
manager.record_recompute();
|
||
|
||
let stats = manager.stats();
|
||
assert_eq!(stats.activations_stored, 1);
|
||
assert_eq!(stats.recompute_count, 2);
|
||
assert!((stats.memory_savings_factor - 4.0).abs() < 0.01);
|
||
}
|
||
|
||
// -------------------------------------------------------------------------
|
||
// AttentionSelective tests
|
||
// -------------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn test_attention_selective_matches_attn_layers() {
|
||
let policy = CheckpointPolicy::attention_selective();
|
||
// Standard attention name patterns must be recognised.
|
||
assert!(policy.should_checkpoint(0, "self_attn", 12));
|
||
assert!(policy.should_checkpoint(0, "cross_attention_layer", 12));
|
||
assert!(policy.should_checkpoint(0, "mha_block", 12));
|
||
// FFN / MLP layers must not be checkpointed.
|
||
assert!(!policy.should_checkpoint(0, "ffn_layer", 12));
|
||
assert!(!policy.should_checkpoint(0, "mlp_dense", 12));
|
||
assert!(!policy.should_checkpoint(0, "layer_norm", 12));
|
||
}
|
||
|
||
#[test]
|
||
fn test_attention_selective_memory_savings_factor() {
|
||
let policy = CheckpointPolicy::attention_selective();
|
||
let factor = policy.memory_savings_factor(12);
|
||
// 12 * 0.4 = 4.8
|
||
assert!((factor - 4.8).abs() < 1e-6, "expected 4.8, got {factor}");
|
||
}
|
||
|
||
#[test]
|
||
fn test_custom_attention_patterns() {
|
||
let policy = CheckpointPolicy::AttentionSelective {
|
||
attention_patterns: vec!["transformer_block".to_string()],
|
||
};
|
||
assert!(policy.should_checkpoint(0, "transformer_block_0", 12));
|
||
assert!(!policy.should_checkpoint(0, "linear_layer", 12));
|
||
}
|
||
|
||
// -------------------------------------------------------------------------
|
||
// Adaptive (fixed) tests
|
||
// -------------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn test_adaptive_tight_memory_checkpoints_all() {
|
||
let policy = CheckpointPolicy::Adaptive {
|
||
target_memory_mb: 512,
|
||
};
|
||
// tier == 1 → every layer is checkpointed
|
||
for i in 0..8 {
|
||
assert!(
|
||
policy.should_checkpoint(i, "any_layer", 8),
|
||
"layer {i} should be checkpointed under tight target"
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_moderate_memory_checkpoints_every_other() {
|
||
let policy = CheckpointPolicy::Adaptive {
|
||
target_memory_mb: 2048,
|
||
};
|
||
// tier == 2 → even indices only
|
||
for i in 0..8usize {
|
||
let expected = i % 2 == 0;
|
||
assert_eq!(
|
||
policy.should_checkpoint(i, "layer", 8),
|
||
expected,
|
||
"layer {i}: expected={expected}"
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_generous_memory_checkpoints_sqrt() {
|
||
let policy = CheckpointPolicy::Adaptive {
|
||
target_memory_mb: 8192,
|
||
};
|
||
let total = 16usize;
|
||
// ceil(sqrt(16)) = 4 → checkpoint at indices 0, 4, 8, 12
|
||
assert!(policy.should_checkpoint(0, "layer", total));
|
||
assert!(!policy.should_checkpoint(1, "layer", total));
|
||
assert!(!policy.should_checkpoint(2, "layer", total));
|
||
assert!(!policy.should_checkpoint(3, "layer", total));
|
||
assert!(policy.should_checkpoint(4, "layer", total));
|
||
assert!(!policy.should_checkpoint(5, "layer", total));
|
||
assert!(policy.should_checkpoint(8, "layer", total));
|
||
assert!(policy.should_checkpoint(12, "layer", total));
|
||
assert!(!policy.should_checkpoint(13, "layer", total));
|
||
}
|
||
|
||
// -------------------------------------------------------------------------
|
||
// MemoryAwareCheckpointer tests
|
||
// -------------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn test_memory_aware_checkpointer_pressure() {
|
||
let checkpointer = MemoryAwareCheckpointer::new(
|
||
CheckpointPolicy::attention_selective(),
|
||
12,
|
||
1024, // 1 GiB target
|
||
);
|
||
// Freshly created — no memory recorded yet.
|
||
assert!(
|
||
checkpointer.memory_pressure() < 0.01,
|
||
"initial pressure should be near zero"
|
||
);
|
||
|
||
checkpointer.record_activation(512 * 1024 * 1024); // 512 MiB
|
||
let pressure = checkpointer.memory_pressure();
|
||
assert!(
|
||
pressure > 0.49 && pressure < 0.51,
|
||
"pressure after 512 MiB / 1024 MiB should be ~0.5, got {pressure}"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_memory_aware_checkpointer_fallback_under_pressure() {
|
||
let checkpointer = MemoryAwareCheckpointer::new(
|
||
CheckpointPolicy::None, // normally checkpoints nothing
|
||
12,
|
||
256, // 256 MiB target
|
||
);
|
||
checkpointer.record_activation(300 * 1024 * 1024); // exceeds target
|
||
// Memory pressure override must force checkpointing even though policy is None.
|
||
assert!(
|
||
checkpointer.should_checkpoint(0, "ffn_layer"),
|
||
"should force-checkpoint under memory pressure regardless of policy"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_memory_aware_free_reduces_pressure() {
|
||
let checkpointer = MemoryAwareCheckpointer::new(CheckpointPolicy::None, 12, 1024);
|
||
checkpointer.record_activation(200 * 1024 * 1024);
|
||
checkpointer.free_activation(200 * 1024 * 1024);
|
||
assert!(
|
||
checkpointer.memory_pressure() < 0.01,
|
||
"pressure should return to near zero after freeing all recorded memory"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_memory_aware_free_does_not_underflow() {
|
||
// Freeing more than was recorded must saturate at zero, not wrap.
|
||
let checkpointer = MemoryAwareCheckpointer::new(CheckpointPolicy::None, 4, 512);
|
||
checkpointer.record_activation(10 * 1024 * 1024);
|
||
checkpointer.free_activation(50 * 1024 * 1024); // larger than recorded
|
||
assert_eq!(
|
||
checkpointer
|
||
.estimated_memory_bytes
|
||
.load(std::sync::atomic::Ordering::Relaxed),
|
||
0,
|
||
"counter must not underflow to usize::MAX"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_memory_aware_checkpointer_delegates_to_policy() {
|
||
// Below pressure threshold the policy should be respected.
|
||
let checkpointer = MemoryAwareCheckpointer::new(
|
||
CheckpointPolicy::attention_selective(),
|
||
12,
|
||
4096, // large target — won't be exceeded
|
||
);
|
||
checkpointer.record_activation(1024); // trivial amount
|
||
// Attention layer — policy says checkpoint.
|
||
assert!(checkpointer.should_checkpoint(0, "self_attn"));
|
||
// FFN layer — policy says don't checkpoint.
|
||
assert!(!checkpointer.should_checkpoint(0, "ffn_dense"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_memory_aware_checkpointer_debug_impl() {
|
||
// Ensure the manual Debug impl doesn't panic.
|
||
let c = MemoryAwareCheckpointer::new(CheckpointPolicy::None, 8, 512);
|
||
let _ = format!("{c:?}");
|
||
}
|
||
}
|