FlexTensor is a tensor offloading and management library for PyTorch that enables running large models on limited GPU memory by intelligently offloading tensors between GPU and CPU memory.
- Simplified API: Easy-to-use high-level API for automatic tensor offloading
- Automatic Model Patching: Offload model layers without modifying model code
- Manual Control: Fine-grained control with
offload_blockcontext managers - Smart Profiling: Automatic discovery and profiling for optimal performance
- Wildcard Support: Use patterns like
"layers.*"to offload multiple modules - Profile Persistence: Save and load offloading profiles for faster startup
- Lazy Model Initialization: Load models from saved profiles with optimized weight loading
- Shared Memory: Optional shared memory subsystem for cross-process tensor coordination
For detailed guides, API reference, and more, visit our Documentation.
To install FlexTensor from PyPI:
pip install flextensorFor more installation options (source and development), see the Installation Guide.
import flextensor
from flextensor import OffloadConfig
# Your existing model
model = YourModel()
# Configure offloading
config = OffloadConfig(
gpu_device=0, # GPU to use
profiling_iters=10, # Iterations for timing measurement
include_patterns=["layers.*"], # Which modules to offload
)
# Patch the model
model = flextensor.offload(model, config=config)
# Use normally — the first few iterations warm the manager
# (`discovery_iters` + `profiling_iters` under the default
# `skip_discovery=False`; query
# `flextensor.get_offload_manager().iters_before_inference` for
# the exact path-aware count).
for batch in dataloader:
output = model(batch) # FlexTensor handles everythingSee the Quick Start for more examples.
FlexTensor is licensed under the Apache License 2.0. See NOTICE for the project notice, ATTRIBUTIONS.md for third-party dependency attributions, and EXTERNAL_MATERIALS.md for external materials.