Skip to main content
Pythia profiles your PyTorch code across multiple GPU types and recommends the optimal hardware for your workload. Instead of guessing which GPU you need or over-provisioning expensive hardware, Pythia tests your code on available GPUs and shows you exactly which one fits.

​
What You Get

Run a single command and Pythia will:
  • Profile your code on all available GPU types
  • Predict memory requirements to identify compatible hardware
  • Estimate runtime for each GPU option
  • Recommend the best GPU based on your workload
╭─────────────────────── GPU Selection Results ───────────────────────╮
│                                                                     │
│  ✓ Analysis complete                                                │
│                                                                     │
│  Recommended: A100 (40 GB VRAM)                                     │
│                                                                     │
╰─────────────────────────────────────────────────────────────────────╯

          Profiling Results
┏━━━━━━┳━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┓
┃ GPU  ┃ Status  ┃    Time ┃    Cost ┃
┡━━━━━━╇━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━┩
│ T4   │ success │   8.2hr │   $2.46 │
│ A100 │ success │   3.1hr │   $5.27 │
│ H100 │ success │   1.8hr │  $10.44 │
└──────┴─────────┴─────────┴─────────┘

              Compatible GPUs
┏━━━━━━━┳━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━┓
┃ GPU   ┃  VRAM ┃ GPUs Needed ┃ Utilization ┃
┡━━━━━━━╇━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━┩
│ T4    │ 16 GB │           1 │         68% │
│ A100 ✓│ 40 GB │           1 │         27% │
│ H100  │ 80 GB │           1 │         14% │
└───────┴───────┴─────────────┴─────────────┘

​
CLI Usage

​
Run GPU Selection

lyceum gpu-selection run train.py
Pass script arguments after --: 
lyceum gpu-selection run train.py -- --epochs 10 --lr 0.001

​
Options

OptionDescription
-f, --file-nameName for the execution
-t, --timeoutTimeout per sub-job in seconds (1-600, default: 60)
-r, --requirementsRequirements file path or pip requirements string
--importPre-import modules (can be used multiple times)
--no-configSkip workspace config from .lyceum/config.json
-d, --debugShow detailed debug information

​
Check Status

lyceum gpu-selection status <execution_id>

​
Requirements

  • Single Python script as entry point (executed as main)
  • Must target CUDA/GPU as the PyTorch device
  • Command-line arguments are supported
  • PyTorch and PyTorch Lightning
  • Imported models (Hugging Face, torchvision, timm)
  • Single model, single GPU
  • Jupyter notebooks (not supported at this point)
  • Advanced or custom gradient manipulation
  • CUDA extensions that bypass PyTorch tensors

​
Prediction Details

Predicts training and inference memory for a single model. Accounts for gradients, model weights, activations, optimizers, and mixed-precision.Does not account for data transfer outside the training/inference loop or operator memory overhead (intermediate tensors).
Estimates inference/training loop time based on iteration time.Does not account for VM startup time or initial data downloading and loading.