Metadata-Version: 2.4
Name: flextensor
Version: 0.2.0
Summary: FlexTensor: Tensor offloading and management library
License-Expression: Apache-2.0
Keywords: tensor,offloading,machine learning,pytorch,memory management
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: beartype~=0.21
Requires-Dist: numpy>=1.20.0
Requires-Dist: psutil>=5.9.0
Requires-Dist: pydantic>=2.7.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: torch>=2.5
Requires-Dist: scipy>=1.9.0
Requires-Dist: proxytypes3>=1.0.1
Requires-Dist: posix-ipc~=1.3.0
Requires-Dist: shared-memory-dict~=0.7.2
Provides-Extra: test
Requires-Dist: pytest>=9.0.3; extra == "test"
Requires-Dist: pytest-cov>=4.0.0; extra == "test"
Requires-Dist: mypy>=1.10; extra == "test"
Provides-Extra: docs
Requires-Dist: mkdocs~=1.6; extra == "docs"
Requires-Dist: mkdocs-material~=9.7; extra == "docs"
Requires-Dist: mike~=2.1; extra == "docs"
Requires-Dist: mkdocstrings[python]>=0.27; extra == "docs"
Provides-Extra: dev
Requires-Dist: build; extra == "dev"
Requires-Dist: pre-commit; extra == "dev"
Requires-Dist: flextensor[test]; extra == "dev"
Provides-Extra: all
Requires-Dist: flextensor[dev,docs,test]; extra == "all"
Dynamic: license-file

<!--
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-->
# FlexTensor

[![Documentation](https://img.shields.io/badge/Documentation-View-green)](https://github.com/ai-dynamo/flextensor)
[![Dashboard](https://img.shields.io/badge/Dashboard-View-blue)](https://github.com/ai-dynamo/flextensor)

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.

## Features

- **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_block` context 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

## Documentation

For detailed guides, API reference, and more, visit our [Documentation](https://github.com/ai-dynamo/flextensor).

## Quick Installation

To install FlexTensor from PyPI:

```bash
pip install flextensor
```

For more installation options (source, dev, optional dependencies), see the [Installation Guide](https://github.com/ai-dynamo/flextensor/blob/main/docs/installation.md).

## Quick Example

```python
import flextensor
from flextensor import OffloadConfig

# Your existing model
model = YourModel()

# Configure offloading
config = OffloadConfig(
    gpu_device=0,              # GPU to use
    discovery_iters=1,            # Iterations for tensor discovery
    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 - first discovery_iters + profiling_iters iterations are discovery/profiling
for batch in dataloader:
    output = model(batch)  # FlexTensor handles everything
```

See the [Quick Start](https://github.com/ai-dynamo/flextensor/blob/main/docs/quick-start.md) for more examples.

## License

FlexTensor is licensed under the [Apache License 2.0](LICENSE). See [NOTICE](NOTICE) for additional notices and disclaimers regarding external materials.
