311 lines
7.6 KiB
Markdown
311 lines
7.6 KiB
Markdown
# Autograd Engine
|
||
|
||
How automatic differentiation works in RustyTorch++.
|
||
|
||
## Overview
|
||
|
||
The autograd engine implements reverse-mode automatic differentiation (backpropagation) using a tape-based approach.
|
||
|
||
## Key Concepts
|
||
|
||
### Computation Graph
|
||
|
||
Every tensor operation builds a computation graph:
|
||
|
||
```
|
||
┌───────────────┐
|
||
│ Input X │ (leaf, requires_grad=true)
|
||
└───────┬───────┘
|
||
│
|
||
┌───────▼───────┐
|
||
│ MatMul(W) │ W is another leaf
|
||
└───────┬───────┘
|
||
│
|
||
┌───────▼───────┐
|
||
│ ReLU │
|
||
└───────┬───────┘
|
||
│
|
||
┌───────▼───────┐
|
||
│ Loss │ (output)
|
||
└───────────────┘
|
||
```
|
||
|
||
### Gradient Tape
|
||
|
||
Operations are recorded on a tape:
|
||
|
||
```rust
|
||
use rtx_autograd::{AutogradTape, TapeNode};
|
||
|
||
// Tape records operations during forward pass
|
||
let tape = AutogradTape::new();
|
||
|
||
// Each operation adds a node
|
||
tape.record_operation(
|
||
&[input_id, weight_id], // inputs
|
||
output_id, // output
|
||
Box::new(MatMulBackward { /* saved tensors */ }),
|
||
);
|
||
```
|
||
|
||
## Using Variables
|
||
|
||
### Basic Gradient Computation
|
||
|
||
```rust
|
||
use rtx_autograd::{Variable, backward};
|
||
use rtx_tensor::{Tensor, Device};
|
||
|
||
let device = Device::cpu();
|
||
|
||
// Create tracked variables
|
||
let x = Variable::from_tensor(Tensor::from_slice(&[2.0, 3.0], [2], &device)?);
|
||
let w = Variable::from_tensor(Tensor::from_slice(&[1.0, 2.0], [2], &device)?);
|
||
|
||
// Forward pass builds graph
|
||
let y = x.multiply(&w)?; // [2.0, 6.0]
|
||
let z = y.sum()?; // 8.0
|
||
|
||
// Backward pass computes gradients
|
||
z.backward()?;
|
||
|
||
// Access gradients
|
||
println!("dx: {:?}", x.grad()); // [1.0, 2.0] (dy/dx = w)
|
||
println!("dw: {:?}", w.grad()); // [2.0, 3.0] (dy/dw = x)
|
||
```
|
||
|
||
### Gradient Context Managers
|
||
|
||
```rust
|
||
use rtx_autograd::{no_grad, enable_grad, is_grad_enabled};
|
||
|
||
// Check current state
|
||
println!("Grad enabled: {}", is_grad_enabled());
|
||
|
||
// Disable gradients (inference mode)
|
||
let result = no_grad(|| {
|
||
let a = Tensor::randn([100, 100], &device)?;
|
||
let b = Tensor::randn([100, 100], &device)?;
|
||
a.matmul(&b) // No graph built
|
||
});
|
||
|
||
// Re-enable within nested scope
|
||
no_grad(|| {
|
||
enable_grad(|| {
|
||
// Gradients tracked here
|
||
});
|
||
// Gradients not tracked here
|
||
});
|
||
```
|
||
|
||
### Detaching from Graph
|
||
|
||
```rust
|
||
let x = Variable::from_tensor(tensor);
|
||
let y = x.multiply(&w)?;
|
||
|
||
// Detach creates a new variable without graph connection
|
||
let y_detached = y.detach();
|
||
assert!(!y_detached.requires_grad());
|
||
|
||
// Useful for stop-gradient operations
|
||
let target = predictions.detach(); // Don't backprop through target
|
||
```
|
||
|
||
## Backward Functions
|
||
|
||
Each operation has a corresponding backward function:
|
||
|
||
```rust
|
||
use rtx_autograd::{BackwardFunction, BackwardContext};
|
||
|
||
/// Backward function for matrix multiplication: C = A @ B
|
||
pub struct MatMulBackward {
|
||
a: Tensor, // Saved input A
|
||
b: Tensor, // Saved input B
|
||
}
|
||
|
||
impl BackwardFunction for MatMulBackward {
|
||
fn backward(&self, grad_output: &Tensor) -> Vec<Option<Tensor>> {
|
||
// dL/dA = dL/dC @ B^T
|
||
let grad_a = grad_output.matmul(&self.b.transpose(0, 1)?)?;
|
||
|
||
// dL/dB = A^T @ dL/dC
|
||
let grad_b = self.a.transpose(0, 1)?.matmul(grad_output)?;
|
||
|
||
vec![Some(grad_a), Some(grad_b)]
|
||
}
|
||
}
|
||
```
|
||
|
||
### Built-in Backward Functions
|
||
|
||
| Operation | Backward | Notes |
|
||
|-----------|----------|-------|
|
||
| `Add` | [1, 1] | Gradients pass through unchanged |
|
||
| `Mul` | [y, x] | Swap inputs |
|
||
| `MatMul` | [grad @ B^T, A^T @ grad] | Standard matrix calculus |
|
||
| `Sum` | expand | Broadcast gradient |
|
||
| `ReLU` | grad * (x > 0) | Zero where input was negative |
|
||
| `Sigmoid` | grad * σ * (1 - σ) | Uses saved output |
|
||
| `Softmax` | Jacobian-vector product | Complex formula |
|
||
|
||
## Gradient Checkpointing
|
||
|
||
Trade memory for compute by recomputing activations:
|
||
|
||
```rust
|
||
use rtx_autograd::{CheckpointManager, CheckpointStrategy};
|
||
|
||
// Create checkpoint manager
|
||
let mut checkpoint = CheckpointManager::new(CheckpointStrategy::SqrtN);
|
||
|
||
// Checkpoint expensive operations
|
||
let hidden = checkpoint.checkpoint(|| {
|
||
model.layer1.forward(&input)?
|
||
.relu()?
|
||
.layer2.forward()
|
||
})?;
|
||
|
||
// During backward, activations are recomputed
|
||
hidden.backward()?; // Recomputes forward pass for gradients
|
||
```
|
||
|
||
### Checkpointing Strategies
|
||
|
||
| Strategy | Memory | Compute | Use Case |
|
||
|----------|--------|---------|----------|
|
||
| `None` | O(n) | O(1) | Small models |
|
||
| `SqrtN` | O(√n) | O(√n) | Large transformers |
|
||
| `Every(k)` | O(n/k) | O(k) | Custom balance |
|
||
| `Selective` | Varies | Varies | Manual selection |
|
||
|
||
## Higher-Order Derivatives
|
||
|
||
### Hessian Computation
|
||
|
||
```rust
|
||
use rtx_autograd::{HessianComputer};
|
||
|
||
let f = |x: &Variable| -> Variable {
|
||
x.pow(2)?.sum()? // f(x) = sum(x^2)
|
||
};
|
||
|
||
let x = Variable::from_tensor(Tensor::from_slice(&[1.0, 2.0, 3.0], [3], &device)?);
|
||
|
||
// Compute Hessian (second derivative matrix)
|
||
let hessian = HessianComputer::compute(&f, &x)?;
|
||
// For sum(x^2), Hessian is 2*I (diagonal of 2s)
|
||
```
|
||
|
||
### Jacobian Computation
|
||
|
||
```rust
|
||
use rtx_autograd::JacobianComputer;
|
||
|
||
let f = |x: &Variable| -> Variable {
|
||
// f: R^3 -> R^2
|
||
let a = x.narrow(0, 0, 2)?;
|
||
let b = x.narrow(0, 1, 2)?;
|
||
a.add(&b)?
|
||
};
|
||
|
||
let x = Variable::from_tensor(Tensor::randn([3], &device)?);
|
||
let jacobian = JacobianComputer::compute(&f, &x)?; // [2, 3] matrix
|
||
```
|
||
|
||
## Numerical Gradient Checking
|
||
|
||
Verify gradient implementations with finite differences:
|
||
|
||
```rust
|
||
use rtx_autograd::GradientChecker;
|
||
|
||
let checker = GradientChecker::new(1e-5); // epsilon for finite diff
|
||
|
||
let f = |x: &Tensor| {
|
||
x.pow(3)?.sum() // f(x) = sum(x^3)
|
||
};
|
||
|
||
let x = Tensor::randn([10], &device)?;
|
||
let analytical_grad = compute_analytical_grad(&f, &x)?;
|
||
|
||
// Check matches numerical gradient
|
||
let is_correct = checker.check(&f, &x, &analytical_grad)?;
|
||
assert!(is_correct, "Gradient implementation is incorrect");
|
||
```
|
||
|
||
## Graph Caching
|
||
|
||
Cache computation graphs for repeated forward passes:
|
||
|
||
```rust
|
||
use rtx_autograd::{init_graph_cache, cache_graph, get_cached_graph};
|
||
|
||
// Initialize graph cache
|
||
init_graph_cache(1024); // Cache up to 1024 graphs
|
||
|
||
// First forward pass - builds and caches graph
|
||
let output = model.forward(&input)?;
|
||
|
||
// Subsequent passes with same shapes - reuse graph
|
||
let cached = get_cached_graph(&input_shape)?;
|
||
if let Some(graph) = cached {
|
||
// Execute cached graph
|
||
output = graph.execute(&input)?;
|
||
}
|
||
```
|
||
|
||
## Best Practices
|
||
|
||
### 1. Use no_grad for Inference
|
||
|
||
```rust
|
||
// Wrong - builds unnecessary graph
|
||
let output = model.forward(&input)?;
|
||
|
||
// Right - skips graph construction
|
||
let output = no_grad(|| model.forward(&input))?;
|
||
```
|
||
|
||
### 2. Zero Gradients Before Each Step
|
||
|
||
```rust
|
||
for epoch in 0..num_epochs {
|
||
for (x, y) in &dataloader {
|
||
// Clear accumulated gradients
|
||
optimizer.zero_grad();
|
||
|
||
let output = model.forward(&x)?;
|
||
let loss = loss_fn(&output, &y)?;
|
||
|
||
loss.backward()?;
|
||
optimizer.step()?;
|
||
}
|
||
}
|
||
```
|
||
|
||
### 3. Detach When Needed
|
||
|
||
```rust
|
||
// Stop gradient flow to targets
|
||
let target = teacher_model.forward(&x)?.detach();
|
||
let student_out = student_model.forward(&x)?;
|
||
let loss = mse_loss(&student_out, &target)?;
|
||
```
|
||
|
||
### 4. Use Checkpointing for Large Models
|
||
|
||
```rust
|
||
// Memory-efficient transformer
|
||
let attention = checkpoint.checkpoint(|| {
|
||
self.attention.forward(&hidden)?
|
||
})?;
|
||
```
|
||
|
||
## Next Steps
|
||
|
||
- [GPU Runtime](./gpu-runtime.md) - Device and memory management
|
||
- [Performance](../performance/gpu-optimization.md) - Optimization
|
||
- [Troubleshooting](../troubleshooting/memory.md) - Memory issues
|