covex-neuron

Run your PyTorch and Hugging Face code on AWS Trainium, unchanged from NVIDIA except for two lines

GitHub · PyPI · Notebooks

covex-neuron is built on TorchNeuron, the native PyTorch backend for Trainium2 and Trainium3. It handles the parts that make Trainium different from a GPU (bf16, fixed input shapes, the Neuron compiler) and runs your script on AWS by the second, next to an NVIDIA GPU for comparison.

Installation

pip install covex-neuron

Quickstart

Add two lines to an existing GPU script:

import covex_neuron as cn

device = cn.device()                      # "neuron" on Trainium, "cuda" on GPU
pipe = pipeline("sentiment-analysis", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english", device=device)
pipe = cn.prepare(pipe)                   # no-op on GPU / CPU

cn.prepare takes a transformers model (with its tokenizer), a transformers pipeline, or a diffusers pipeline.

Run it on a GPU and on Trainium from your laptop, using only your AWS credentials:

gpu = cn.run(["sentiment.py"], on="gpu")        # 1x NVIDIA A10G
trn = cn.run(["sentiment.py"], on="trainium")   # 1x Trainium2 chip
cn.compare(gpu, trn)

You pay per second while the job runs, and nothing keeps running afterwards.

Notebooks

Each one runs the same script on a GPU and on Trainium and prints the comparison.

Examples

examples/nvidia has sentiment analysis, text classification and image generation scripts as written for a GPU. examples/trainium has the same scripts with the two lines added. To run both and compare:

python examples/run_demo.py

License

Apache 2.0