Skip to main content

One command to find why your PyTorch model is slow — and fix it.

Project description

slowai

One command to find why your PyTorch model is slow — and fix it.

slowai diagnoses which performance regime your workload is stuck in (compute-bound, memory-bound, or overhead-bound), prescribes the right fix, auto-applies it, and proves the speedup with before/after measurements. No guesswork. No manual profiler interpretation.

$ slowai fix model.py

==============================================================
  BASELINE: model.py: COMPUTE_BOUND (confidence: 0.85)
  wall time: 7.523s
==============================================================

  Tried 4 remedies:

  1. [10.00x] bf16_autocast  ** BEST **
     Run under bfloat16 automatic mixed precision
     7.523s >>> 0.752s
     regime: compute (confidence: 0.85)

  2. [6.32x] tf32_tensor_cores
     Enable TF32 tensor cores (~2x matmul throughput on Ampere+)
     7.523s >>> 1.191s

  3. [6.22x] high_matmul_precision
     Set float32 matmul precision to 'high'
     7.523s >>> 1.210s

  4. [1.31x] cudnn_benchmark
     Enable cuDNN auto-tuner for conv kernels
     7.523s >>> 5.719s

--------------------------------------------------------------
  WINNER: bf16_autocast
  7.523s >>> 0.752s  (10.00x, +900% faster)
  How: Run under bfloat16 automatic mixed precision
--------------------------------------------------------------

Why this exists

Every deep learning workload is stuck in one of three performance regimes (Horace He, 2022):

Regime What's happening Wrong fix = no speedup
Compute-bound GPU is saturated doing math (matmuls, convolutions) Fusing ops won't help — the math itself is the bottleneck
Memory-bound GPU is waiting for data (pointwise ops, activations) Smaller model won't help — you need less data movement
Overhead-bound GPU is idle waiting for Python/dispatcher (tiny ops) Lower precision won't help — you need fewer, bigger ops

The fix for each regime is different. Applying a compute-bound fix to a memory-bound workload does nothing. Engineers waste hours in profiler UIs figuring this out manually.

slowai does it in one command.

How it works

slowai diagnose model.py    # Classify the regime + prescribe fixes
slowai fix model.py          # ^ plus auto-apply fixes and measure speedup

Under the hood:

  1. Profile — Runs your workload under torch.profiler with CUDA timing, warmup pass, and op-level statistics
  2. Classify — A heuristic classifier analyzes op shares (matmul, normalization, pointwise, tiny-op fraction) to determine the dominant regime
  3. Prescribe — Returns a ranked list of fixes for that regime, cheapest first
  4. Remediate — Auto-applies each applicable fix (TF32 tensor cores, bf16/fp16 autocast, cuDNN benchmark, matmul precision), re-profiles, and ranks by measured speedup

No code changes required. Remedies are environment-level transforms — they modify PyTorch's runtime settings, not your model code.

Benchmarks

Tested on NVIDIA Jetson Orin Nano Super (Ampere GPU, 1024 CUDA cores, JetPack 6.2, PyTorch 2.8.0):

Workload Regime Baseline Best remedy After Speedup
Dense GEMM (4096x4096) Compute 7.523s bf16_autocast 0.752s 10.00x
Pointwise chain (8192x8192) Memory 2.400s tf32_tensor_cores 0.575s 4.17x
Tiny ops (5000 micro-ops) Overhead 3.281s tf32_tensor_cores 1.141s 2.88x
MobileNetV2 inference Compute 1.778s cudnn_benchmark 1.681s 1.06x
ResNet-50 inference Compute 2.105s bf16_autocast 1.969s 1.07x

Production models (MobileNet, ResNet) show modest gains because PyTorch already optimizes common architectures well. The real value is that slowai finds the right fix automatically — cuDNN benchmark wins for convolution-heavy MobileNet, bf16 autocast wins for matmul-heavy ResNet. Different models, different winners, zero guesswork.

Installation

git clone https://github.com/ricojallen37-sketch/slowai.git
cd slowai
pip install -e .

Requires Python 3.10+ and PyTorch >= 2.1 with CUDA support.

Writing a workload

slowai profiles any Python script that exposes a main() function:

# my_model.py
import torch
from torchvision.models import resnet50

model = resnet50().cuda().eval()
data = torch.randn(8, 3, 224, 224, device="cuda")

def main():
    with torch.no_grad():
        for _ in range(30):
            model(data)
slowai fix my_model.py

Architecture

slowai/
  schema.py      # Regime enum, Diagnosis dataclass — the product thesis in types
  profiler.py    # torch.profiler wrapper → ProfileResult (op stats + wall time)
  diagnose.py    # Heuristic classifier → Diagnosis (regime + confidence + prescriptions)
  remediate.py   # Auto-fix engine → FixReport (before/after speedup per remedy)
  cli.py         # CLI entry points: diagnose, fix

The classifier is a pure function of profiler output — no torch dependency, fully unit-testable. The remediate engine applies fixes as environment transforms (global flags, autocast context managers) so it never modifies user code.

What's different

Other tools in this space are profiler UIs that show you data and leave interpretation to you. slowai is the only tool that goes from raw workload to regime classification to ranked prescriptions to auto-applied fixes to measured speedup in a single CLI command.

Tool Profiles Classifies regime Prescribes fixes Auto-applies Measures speedup
PyTorch Profiler Yes No No No No
NVIDIA Nsight Yes No No No No
torch.utils.bottleneck Yes No No No No
DeepSpeed Flops Profiler Yes No No No No
slowai Yes Yes Yes Yes Yes

Roadmap

  • V1 (shipped) — Profile + classify regime for synthetic workloads
  • V2 (shipped) — Noise filtering (sync ops, init ops), normalization-aware classification, real model support
  • V3 (shipped) — Auto-remediate: apply fixes and measure before/after speedup
  • V4 (next) — torch.compile integration, channels_last auto-transform, batch size search, export optimized config

Built by

Rico Allen — @ricojallen37-sketch

Built and tested on NVIDIA Jetson Orin Nano Super Developer Kit.

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

slowai-0.3.0.tar.gz (21.4 kB view details)

Uploaded Source

Built Distribution

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

slowai-0.3.0-py3-none-any.whl (18.7 kB view details)

Uploaded Python 3

File details

Details for the file slowai-0.3.0.tar.gz.

File metadata

  • Download URL: slowai-0.3.0.tar.gz
  • Upload date:
  • Size: 21.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.3

File hashes

Hashes for slowai-0.3.0.tar.gz
Algorithm Hash digest
SHA256 9381687d4224a5b0db09de75b3f9b94ace05f6c0439bb58ec3a13316f9da9452
MD5 523ff4cb14972c9d6a4f2c0dfd786172
BLAKE2b-256 1f19971d1c11307d08ba2d8a1011de228f7f871745e5dd9c4cbfdd9f85aa2eb9

See more details on using hashes here.

File details

Details for the file slowai-0.3.0-py3-none-any.whl.

File metadata

  • Download URL: slowai-0.3.0-py3-none-any.whl
  • Upload date:
  • Size: 18.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.3

File hashes

Hashes for slowai-0.3.0-py3-none-any.whl
Algorithm Hash digest
SHA256 8028d5c43c1abedbb2650896ec8e3483b02b369f30422a11c135abee5981cf09
MD5 6e925a79cbab201ec8501ee953908381
BLAKE2b-256 d902b25a173ac8334a705214f89bd374e7514d89b2a598bcd7bf0f66aa1226a3

See more details on using hashes here.

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