Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,242 @@
+# Memory Management
+
+Efficient memory usage in RustyTorch++.
+
+## GPU Memory Hierarchy
+
+```
+┌─────────────────────────────────────────┐
+│       Global Memory (HBM) - 24-80 GB    │ ← Tensor storage
+├─────────────────────────────────────────┤
+│       L2 Cache - 4-50 MB                │ ← Auto-managed
+├─────────────────────────────────────────┤
+│       Shared Memory - 48-164 KB/SM      │ ← Kernel local
+├─────────────────────────────────────────┤
+│       Registers - 64K/SM                │ ← Thread local
+└─────────────────────────────────────────┘
+```
+
+## Memory Monitoring
+
+```rust
+use rtx_tensor::Device;
+
+let device = Device::cuda(0)?;
+
+// Query memory
+let total = device.total_memory()?;
+let used = device.used_memory()?;
+let free = device.free_memory()?;
+
+println!("Total: {:.2} GB", total as f64 / 1e9);
+println!("Used: {:.2} GB", used as f64 / 1e9);
+println!("Free: {:.2} GB", free as f64 / 1e9);
+```
+
+## Tensor Pool
+
+### How It Works
+
+```rust
+use rtx_tensor::{TensorPool, pool_stats};
+
+// Request a buffer from the pool
+let buffer = TensorPool::get([1024, 1024], &device)?;
+
+// When buffer is dropped, memory returns to pool (not freed)
+drop(buffer);
+
+// Next allocation with same size reuses memory
+let buffer2 = TensorPool::get([1024, 1024], &device)?;  // No allocation!
+```
+
+### Pool Statistics
+
+```rust
+let stats = pool_stats();
+println!("Pool size: {} MB", stats.total_size / 1_000_000);
+println!("In use: {} MB", stats.used_size / 1_000_000);
+println!("Cached: {} MB", stats.cached_size / 1_000_000);
+println!("Hit rate: {:.1}%", stats.hit_rate * 100.0);
+println!("Fragmentation: {:.1}%", stats.fragmentation * 100.0);
+```
+
+### Pool Management
+
+```rust
+use rtx_tensor::clear_pool;
+
+// Clear unused cached memory
+clear_pool();
+
+// Configure pool behavior
+TensorPool::configure(PoolConfig {
+    max_cached_bytes: 4 * 1024 * 1024 * 1024,  // 4 GB max cache
+    max_split_size: 512 * 1024 * 1024,          // Max block to split
+});
+```
+
+## Gradient Checkpointing
+
+Trade compute for memory during training.
+
+### Automatic Checkpointing
+
+```rust
+use rtx_autograd::{CheckpointManager, CheckpointStrategy};
+
+// √n checkpoints - optimal memory/compute tradeoff
+let checkpoint = CheckpointManager::new(CheckpointStrategy::SqrtN);
+
+// Wrap expensive forward passes
+let hidden = checkpoint.checkpoint(|| {
+    let h1 = layer1.forward(&x)?;
+    let h2 = layer2.forward(&h1)?;
+    layer3.forward(&h2)
+})?;
+
+// During backward, activations are recomputed
+hidden.backward()?;
+```
+
+### Memory Savings
+
+| Strategy | Memory | Compute Overhead |
+|----------|--------|------------------|
+| None | O(n) | 1x |
+| √n checkpoints | O(√n) | ~1.5x |
+| Every-k | O(n/k) | ~kx |
+
+## Memory-Efficient Attention
+
+### Flash Attention
+
+```rust
+use rtx_flash_attention::FlashAttention;
+
+// Standard attention: O(N²) memory
+// Flash attention: O(N) memory
+
+let flash = FlashAttention::new(config)?;
+let output = flash.forward(&q, &k, &v)?;
+```
+
+### Chunked Attention
+
+```rust
+// Process attention in chunks for very long sequences
+let chunk_size = 1024;
+let mut outputs = vec![];
+
+for chunk_start in (0..seq_len).step_by(chunk_size) {
+    let chunk_end = (chunk_start + chunk_size).min(seq_len);
+    let q_chunk = q.narrow(1, chunk_start, chunk_end - chunk_start)?;
+    let out_chunk = attention.forward(&q_chunk, &k, &v)?;
+    outputs.push(out_chunk);
+}
+
+let output = Tensor::cat(&outputs, 1)?;
+```
+
+## Pinned Memory
+
+Faster CPU-GPU transfers.
+
+```rust
+use rtx_memory::PinnedBuffer;
+
+// Allocate pinned (page-locked) host memory
+let pinned = PinnedBuffer::new(size)?;
+
+// Fill with data
+pinned.copy_from_slice(&data);
+
+// Fast async transfer to GPU
+let gpu_tensor = Tensor::from_pinned_async(&pinned, shape, &device, &stream)?;
+```
+
+## Memory Debugging
+
+### Find Memory Leaks
+
+```rust
+use rtx_memory::MemoryTracker;
+
+// Enable tracking
+MemoryTracker::enable();
+
+// Run your code...
+train_loop(&model, &data)?;
+
+// Report allocations
+let report = MemoryTracker::report();
+for allocation in report.live_allocations() {
+    println!("{}: {} bytes at {}",
+        allocation.id,
+        allocation.size,
+        allocation.location
+    );
+}
+```
+
+### Memory Snapshots
+
+```rust
+use rtx_memory::Snapshot;
+
+let before = Snapshot::capture(&device)?;
+
+// Operations...
+let output = model.forward(&input)?;
+
+let after = Snapshot::capture(&device)?;
+
+// Compare
+let diff = after.diff(&before);
+println!("Memory change: {} MB", diff.total_bytes / 1_000_000);
+```
+
+## Best Practices
+
+### 1. Delete Intermediate Tensors
+
+```rust
+// Bad - keeps all intermediates
+let a = model.layer1(&x)?;
+let b = model.layer2(&a)?;
+let c = model.layer3(&b)?;
+
+// Good - release when not needed
+let mut h = model.layer1(&x)?;
+h = model.layer2(&h)?;
+h = model.layer3(&h)?;
+```
+
+### 2. Use no_grad for Inference
+
+```rust
+// Training: stores activations for backward
+let output = model.forward(&input)?;
+
+// Inference: no activation storage
+let output = no_grad(|| model.forward(&input))?;
+```
+
+### 3. Process in Batches
+
+```rust
+// OOM risk with large batch
+let all_outputs = model.forward(&huge_batch)?;
+
+// Safe - process in chunks
+let mut outputs = vec![];
+for chunk in data.chunks(safe_batch_size) {
+    outputs.push(model.forward(&chunk)?);
+}
+```
+
+## Next Steps
+
+- [GPU Optimization](./gpu-optimization.md) - Compute optimization
+- [Profiling](./profiling.md) - Find memory bottlenecks
+- [Troubleshooting Memory](../troubleshooting/memory.md) - Fix OOM