A library for accelerating PyTorch models using ONNX Runtime:

• torch-ort to train PyTorch models faster with ONNX Runtime
• moe to scale large models and improve their quality
• torch-ort-infer to perform inference on PyTorch models with ONNX Runtime and Intel® OpenVINO™

You need a machine with at least one NVIDIA or AMD GPU to run ONNX Runtime for PyTorch.

You can install and run torch-ort in your local environment, or with Docker.

1. Install CUDA

2. Install CuDNN

3. Install torch-ort

   • pip install torch-ort

4. Run post-installation script for ORTModule

   • python -m torch_ort.configure

Get install instructions for other combinations in the Get Started Easily section at https://www.onnxruntime.ai/ under the Optimize Training tab.

1. Clone this repo

   • git clone [email protected]:pytorch/ort.git

2. Install extra dependencies

   • pip install wget pandas sklearn transformers

3. Run a test training script

   • python ./ort/tests/bert_for_sequence_classification.py

Mixture of Experts layer implementation is available in the ort_moe folder.

Clone this repo

git clone https://github.com/pytorch/ort.git

Build MoE

cd ort_moe
python setup.py install

• Ubuntu 18.04, 20.04
• Python* 3.7, 3.8 or 3.9

• pip install torch-ort-infer[openvino]
• Run post installation configuration script python -m torch_ort.configure

1. Clone this repo

   • git clone [email protected]:pytorch/ort.git

2. Install extra dependencies

   • pip install wget pandas sklearn transformers

3. Run a test script

   • python ./torch_ort_inference/tests/bert_for_sequence_classification.py

The torch-ort library accelerates training of large transformer PyTorch models to reduce the training time and GPU cost with a few lines of code change. It is built on top of highly successful and proven technologies of ONNX Runtime and ONNX format and includes the ONNX Runtime Optimizer and Data Sampler.

from torch_ort import ORTModule
model = ORTModule(model)
# PyTorch training script follows

import torch
from torch_ort.optim import FusedAdam
class NeuralNet(torch.nn.Module):
    ...
# Only supports GPU Currently.
device = "cuda"
model = NeuralNet(...).to(device)
ort_fused_adam_optimizer = FusedAdam(
    model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01, eps=1e-8
)

# To use FP16_Optimizer, Add these lines : 
from torch_ort.optim import FP16_Optimizer
ort_fused_adam_optimizer = FP16_Optimizer(ort_fused_adam_optimizer)


loss = model(...).sum()
loss.backward()
ort_fused_adam_optimizer.step()
ort_fused_adam_optimizer.zero_grad()

For detailed documentation see FusedAdam

For a full working example see FusedAdam Test Example

FP16_Optimizer is a simple wrapper to replace inefficient FP16_Optimizer function calls implemented by libraries for example Apex, DeepSpeed, Megatron-LM.

For detailed documentation see FP16 Optimizer

import torch
from torch.utils.data import DataLoader 
from torch_ort.utils.data import LoadBalancingDistributedSampler
class MyDataset(torch.utils.data.Dataset):
   ...
   
def collate_fn(data): 
    ...
    return samples, label_list 
samples = [...] 
labels = [...] 
dataset = MyDataset(samples, labels) 
data_sampler = sampler.LoadBalancingDistributedSampler( 
    dataset, complexity_fn=complexity_fn, world_size=2, rank=0, shuffle=False 
) 
train_dataloader = DataLoader(dataset, batch_size=2, sampler=data_sampler, collate_fn=collate_fn) 
for batched_data, batched_label in train_dataloader: 
    optimizer.zero_grad() 
    loss = loss_fn(model(batched_data) , batched_labels) 
    loss.backward() 
    optimizer.step() 
    

For detailed documentation see LoadBalancingDistributedSampler

For a full working example see LoadBalancingDistributedSampler Test Example

To see torch-ort in action, see https://github.com/microsoft/onnxruntime-training-examples, which shows you how to train the most popular HuggingFace models.

To run MoE, add the layer to your model as described in the tutorial: ort_moe/docs/tutorials/moe_tutorial.py

For more details, see ort_moe/docs/moe.md

Note: ONNX Runtime is not required to run the MoE layer. It is integrated in standalone PyTorch.

ONNX Runtime for PyTorch supports PyTorch model inference using ONNX Runtime and Intel® OpenVINO™.

It is available via the torch-ort-infer python package. This package enables OpenVINO™ Execution Provider for ONNX Runtime by default for accelerating inference on various Intel® CPUs, Intel® integrated GPUs, and Intel® Movidius™ Vision Processing Units - referred to as VPU.

Execution Providers
OpenVINO

Users can configure different options for a given Execution Provider to run inference. As an example, OpenVINO™ Execution Provider options can be configured as shown below:

from torch_ort import ORTInferenceModule, OpenVINOProviderOptions
provider_options = OpenVINOProviderOptions(backend = "GPU", precision = "FP16")
model = ORTInferenceModule(model, provider_options = provider_options)

# PyTorch inference script follows

Supported backend-precision combinations:

If no provider options are specified by user, OpenVINO™ Execution Provider is enabled with following options by default:

backend = "CPU"
precision = "FP32"

For more details on APIs, see usage.md.

Below is an example of how you can leverage OpenVINO™ integration with Torch-ORT in a simple NLP usecase.

A pretrained BERT model fine-tuned on the CoLA dataset from HuggingFace model hub is used to predict grammar correctness on a given input text.

from transformers 
import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
from torch_ort import ORTInferenceModule
tokenizer = AutoTokenizer.from_pretrained(
            "textattack/bert-base-uncased-CoLA")
model = AutoModelForSequenceClassification.from_pretrained(
        "textattack/bert-base-uncased-CoLA")
# Wrap model in ORTInferenceModule to prepare the model for inference using OpenVINO Execution Provider on CPU
model = ORTInferenceModule(model)
text = "Replace me any text by you'd like ."
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)
# Post processing
logits = output.logits
logits = logits.detach().cpu().numpy()
# predictions
pred = np.argmax(logits, axis=1).flatten()
print("Grammar correctness label (0=unacceptable, 1=acceptable)")
print(pred)

To see OpenVINO™ integration with Torch-ORT in action, see demos, which shows you how to run inference on some of the most popular Deep Learning models.

Please refer to our contributing guide for more information on how to contribute!

This project has an MIT license, as found in the LICENSE file.