Flexium.AI - Flexible Resource Allocation for GPU training with live migration
Project description
Flexible Resource Allocation - Seamlessly migrate PyTorch training between GPUs with zero interruption.
Your model continues from exactly where it left off, and the source GPU is completely freed with zero VRAM residue.
Quick Start • Features • How It Works • Documentation • Dashboard
What is Flexium?
Flexium is a GPU orchestration system that enables dynamic device migration for PyTorch training jobs. It allows training processes to be moved between GPUs without leaving any memory traces on the source device.
Key Features
- Seamless Migration: Training continues from the exact batch where it stopped
- Zero VRAM Residue: When a process migrates, the source GPU has 0 MB used
- Minimal Code Changes: Just 2 lines to integrate
- Web Dashboard: Real-time monitoring and one-click migration
- Works Offline: Training continues even if server connection is lost
- No Server Installation: Just
pip install flexium- no agents or daemons needed - Framework Compatibility: Works with PyTorch Lightning, Hugging Face, timm, FastAI, and more
Quick Start
1. Install
pip install flexium
2. Sign Up
Create a free account at app.flexium.ai
3. Add Two Lines to Your Training
import flexium
flexium.init() # That's it!
import torch
model = Net().cuda()
optimizer = torch.optim.Adam(model.parameters())
for epoch in range(100):
for batch in dataloader:
data = batch.cuda()
loss = model(data).sum()
loss.backward()
optimizer.step()
4. Run with Your Workspace
export FLEXIUM_SERVER="app.flexium.ai/myworkspace"
python train.py
Your training is now migration-enabled. Monitor and migrate via the dashboard.
The Problem
Traditional approaches to GPU migration leave memory fragments:
# This doesn't fully free memory!
model = model.to("cuda:1") # Old GPU still has memory residue
torch.cuda.empty_cache() # Doesn't guarantee cleanup
The Solution
Flexium uses driver-level migration that guarantees complete memory release:
┌───────────────────────────────────────┐
│ Training on OLD GPU │
│ │
│ Your PyTorch code runs normally │
│ │
└───────────────────────────────────────┘
│
│ MIGRATE
│ (100% memory freed!)
▼
┌───────────────────────────────────────┐
│ Training on NEW GPU │
│ │
│ Resumes from exact position │
│ No progress lost │
│ │
└───────────────────────────────────────┘
How It Works
┌───────────────────────────────────────────────────────────┐
│ YOUR GPU MACHINE │
│ │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ Training Process │ │
│ │ - Your PyTorch training code │ │
│ │ - Initialized with flexium.init() │ │
│ └─────────────────────────────────────────────────────┘ │
│ │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ GPU 0 │ │ GPU 1 │ │ GPU 2 │ │ GPU 3 │ │
│ └─────────┘ └─────────┘ └─────────┘ └─────────┘ │
└───────────────────────────────────────────────────────────┘
│
│ Communicates with
▼
┌───────────────────────────────────────────────────────────┐
│ FLEXIUM CLOUD (flexium.ai) │
│ │
│ Web dashboard for monitoring and control │
└───────────────────────────────────────────────────────────┘
- Sign Up: Create a free account at app.flexium.ai
- Connect: Set
FLEXIUM_SERVERand run your training - Monitor & Migrate: Use the web dashboard for one-click migration
Use Cases
GPU Sharing
import flexium
flexium.init()
# Free up GPUs for teammates without killing your job
train_model() # Can be migrated via dashboard anytime
Pause for Priority Jobs
import flexium
flexium.init()
# Pause your job to free GPU, resume later on any available GPU
train_model() # Pause via dashboard, resume when ready
GPU Error Recovery
import flexium
flexium.init()
for batch in dataloader:
with flexium.auto.recoverable():
# If OOM occurs, migrates to healthy GPU and continues
loss = model(batch.cuda()).sum()
loss.backward()
Requirements
- Python 3.8+
- PyTorch 2.0+ with CUDA support
- Linux x86_64
- NVIDIA Driver:
- 550+ for pause/resume (same GPU)
- 580+ for GPU migration (different GPU)
Documentation
Features
Become a Design Partner
We're looking for design partners to explore advanced capabilities:
- Automatic migration based on resource optimization
- Distributed training support (DDP/FSDP)
- Integration with job schedulers (Slurm/Kubernetes)
- Multi-node GPU orchestration
If you're managing multi-GPU servers and want to shape the future of GPU orchestration, contact us.
License
MIT License - see LICENSE for details.
Contributing
Contributions welcome! Please see our GitHub repository to report issues or submit pull requests.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file flexium-0.2.5.tar.gz.
File metadata
- Download URL: flexium-0.2.5.tar.gz
- Upload date:
- Size: 71.5 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
4b8b3849249166f50bc8bca2657edafa07bfa6345a152b4d86fb7b5ffa3babbc
|
|
| MD5 |
5d45438e531df50a2b98a7dd92f2ba0f
|
|
| BLAKE2b-256 |
d348187c6fd628127e5e63f27002139b9bb54341f41da9182c883c63849c1319
|
Provenance
The following attestation bundles were made for flexium-0.2.5.tar.gz:
Publisher:
publish.yml on flexiumai/flexium
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
flexium-0.2.5.tar.gz -
Subject digest:
4b8b3849249166f50bc8bca2657edafa07bfa6345a152b4d86fb7b5ffa3babbc - Sigstore transparency entry: 1166725271
- Sigstore integration time:
-
Permalink:
flexiumai/flexium@8349331e7e584f9637d7b9892747595eaeeed3e2 -
Branch / Tag:
refs/tags/v0.2.5 - Owner: https://github.com/flexiumai
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@8349331e7e584f9637d7b9892747595eaeeed3e2 -
Trigger Event:
push
-
Statement type:
File details
Details for the file flexium-0.2.5-py3-none-any.whl.
File metadata
- Download URL: flexium-0.2.5-py3-none-any.whl
- Upload date:
- Size: 75.4 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
afef83b471012ca2fb57d3fb63ab56eee46b3b25e944e86ced7a61f60711c7bc
|
|
| MD5 |
f6530cb5bdd8c76cf978acd2798216b6
|
|
| BLAKE2b-256 |
35bad12890d69720d935c2a229796d3f67b612edba33be84a8fd4acecbe69a8a
|
Provenance
The following attestation bundles were made for flexium-0.2.5-py3-none-any.whl:
Publisher:
publish.yml on flexiumai/flexium
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
flexium-0.2.5-py3-none-any.whl -
Subject digest:
afef83b471012ca2fb57d3fb63ab56eee46b3b25e944e86ced7a61f60711c7bc - Sigstore transparency entry: 1166725330
- Sigstore integration time:
-
Permalink:
flexiumai/flexium@8349331e7e584f9637d7b9892747595eaeeed3e2 -
Branch / Tag:
refs/tags/v0.2.5 - Owner: https://github.com/flexiumai
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@8349331e7e584f9637d7b9892747595eaeeed3e2 -
Trigger Event:
push
-
Statement type: