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.6.tar.gz (71.9 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.6-py3-none-any.whl (75.9 kB view details)

Uploaded Python 3

File details

Details for the file flexium-0.2.6.tar.gz.

File metadata

  • Download URL: flexium-0.2.6.tar.gz
  • Upload date:
  • Size: 71.9 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.6.tar.gz
Algorithm Hash digest
SHA256 0e8d5509d4e088010fa34e0b81d497a538396b33b5166f3946e9ec24df78374e
MD5 d4ebdb22cdedb172de9ab85ec0105e29
BLAKE2b-256 0ffb5d0cc944617c570b188c679a6f08e745cc1d17ca2394073d4f4a6456098f

See more details on using hashes here.

Provenance

The following attestation bundles were made for flexium-0.2.6.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.6-py3-none-any.whl.

File metadata

  • Download URL: flexium-0.2.6-py3-none-any.whl
  • Upload date:
  • Size: 75.9 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.6-py3-none-any.whl
Algorithm Hash digest
SHA256 26b719ac8e66daed2b00c54a35df7760da32d06575a7361e5d7495de9ed13b67
MD5 5d891e417db6bbdf238c8bae129ee7b7
BLAKE2b-256 7608bfe1298e14945b3ac6782effd50f0cc5a9fd8e053ba79d52361dab1c84c5

See more details on using hashes here.

Provenance

The following attestation bundles were made for flexium-0.2.6-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