Skip to main content

Flexium.AI - Flexible Resource Allocation for GPU training with live migration

Project description

Flexium.AI Logo

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.

PyPI version Python versions License

Quick StartFeaturesHow It WorksDocumentationDashboard


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          │
└───────────────────────────────────────────────────────────┘
  1. Sign Up: Create a free account at app.flexium.ai
  2. Connect: Set FLEXIUM_SERVER and run your training
  3. 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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

flexium-0.2.5.tar.gz (71.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

flexium-0.2.5-py3-none-any.whl (75.4 kB view details)

Uploaded Python 3

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

Hashes for flexium-0.2.5.tar.gz
Algorithm Hash digest
SHA256 4b8b3849249166f50bc8bca2657edafa07bfa6345a152b4d86fb7b5ffa3babbc
MD5 5d45438e531df50a2b98a7dd92f2ba0f
BLAKE2b-256 d348187c6fd628127e5e63f27002139b9bb54341f41da9182c883c63849c1319

See more details on using hashes here.

Provenance

The following attestation bundles were made for flexium-0.2.5.tar.gz:

Publisher: publish.yml on flexiumai/flexium

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

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

Hashes for flexium-0.2.5-py3-none-any.whl
Algorithm Hash digest
SHA256 afef83b471012ca2fb57d3fb63ab56eee46b3b25e944e86ced7a61f60711c7bc
MD5 f6530cb5bdd8c76cf978acd2798216b6
BLAKE2b-256 35bad12890d69720d935c2a229796d3f67b612edba33be84a8fd4acecbe69a8a

See more details on using hashes here.

Provenance

The following attestation bundles were made for flexium-0.2.5-py3-none-any.whl:

Publisher: publish.yml on flexiumai/flexium

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page