Gradient Cache is a production-ready PyTorch extension that reduces GPU memory usage by 90%+ during neural network training through intelligent gradient compression and CPU offloading.
- 90%+ Memory Savings: Compress gradients by 100x with minimal accuracy impact
- Larger Batch Sizes: Train with 2-3x larger batches on the same hardware
- Simple Integration: Just 3 lines of code to add to any training loop
- Universal Compatibility: Works with any PyTorch model and optimizer
- Production Ready: Tested on A100 and T4 GPUs with real models
| Model | Parameters | Memory Saved | Compression |
|---|---|---|---|
| GPT-2 Small | 124M | 479 MB/step | 100x |
| GPT-2 Medium | 350M | ~1.3 GB/step | 100x |
| Custom NN | 50M | 144 MB/step | 100x |
pip install gradient-cacheOr install from source:
git clone https://github.com/JonSnow1807/gradient-cache
cd gradient-cache
pip install -e .Add gradient cache to any PyTorch training loop with just 3 lines:
importgradient_cache# Create your modelmodel=create_your_model().cuda()
# Add gradient cache (1 line)hook_manager=gradient_cache.create_gradient_cache(model, compression_ratio=100)
# Normal training loopoptimizer=torch.optim.Adam(model.parameters())
forbatchindataloader:
loss=model(batch).mean()
loss.backward()
# Compress gradients (1 line)hook_manager.compress_and_free_gradients()
# Restore gradients and update (1 line)hook_manager.apply_gradients()
optimizer.step()
optimizer.zero_grad()Use the decorator for automatic integration:
frommetaflowimportFlowSpec, stepimportgradient_cacheclassMyTrainingFlow(FlowSpec):
@step@gradient_cache.optimize(compression_ratio=100)deftrain(self):
# Your training code - no changes needed!model=create_model()
optimizer=torch.optim.Adam(model.parameters())
# ... rest of trainingimportpytorch_lightningasplimportgradient_cacheclassMyModel(pl.LightningModule):
def__init__(self):
super().__init__()
self.model=create_model()
self.hook_manager=gradient_cache.create_gradient_cache(self.model)
deftraining_step(self, batch, batch_idx):
loss=self.model(batch).mean()
returnlossdefon_after_backward(self):
self.hook_manager.compress_and_free_gradients()
defoptimizer_step(self, *args, **kwargs):
self.hook_manager.apply_gradients()
super().optimizer_step(*args, **kwargs)# Conservative - 10x compression (keep 10%)hook_manager=gradient_cache.create_gradient_cache(model, compression_ratio=10)
# Aggressive - 1000x compression (keep 0.1%) hook_manager=gradient_cache.create_gradient_cache(model, compression_ratio=1000)# Don't compress embeddings or output layershook_manager=gradient_cache.GradientCacheHookManager(
model,
compression_ratio=100,
exclude_layers=['embedding', 'lm_head']
)# Enable verbose modehook_manager=gradient_cache.create_gradient_cache(model, verbose=True)
# Get compression statisticsstats=hook_manager.get_compression_summary()
print(f"Compression ratio: {stats['overall_compression_ratio']:.1f}x")
print(f"Memory saved: {stats['memory_saved_mb']:.1f} MB")- Gradient Computation: Normal backward pass computes gradients
- Compression: Keep only top 1% of gradient values by magnitude
- CPU Offload: Move compressed gradients to system RAM
- GPU Memory Release: Free GPU memory for next batch
- Gradient Restoration: Restore gradients for optimizer step
- Cost Savings: Use smaller, cheaper GPU instances
- Larger Models: Train models that don't fit in GPU memory
- Faster Research: Iterate quickly with larger batch sizes
- Easy Integration: No model architecture changes needed
Run the test suite:
python tests/test_gradient_cache.pyIf you use Gradient Cache in your research, please cite:
@software{gradient_cache,
title = {Gradient Cache: GPU Memory-Efficient Training},
author = {Gradient Cache Contributors},
year = {2024},
url = {https://github.com/gradient-cache/gradient-cache}
}Apache License 2.0 - see LICENSE for details.
We welcome contributions! Please submit issues and pull requests on GitHub.
- Issues: GitHub Issues
- Discussions: GitHub Discussions
Built with ❤️ for the ML community